diff --git a/.gitignore b/.gitignore index e706a1484..9d70f24b5 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ /ds4_test /ds4flash.gguf /TODO.md +glm53flash_prefill*.md /gguf/ /core /core.* diff --git a/Makefile b/Makefile index c58acb939..bd0ebbe6f 100644 --- a/Makefile +++ b/Makefile @@ -365,13 +365,21 @@ tests/test_deepseek4_vision_image: tests/test_deepseek4_vision_image.o ds4_image $(CC) $(CFLAGS) -o $@ $^ -lm ifeq ($(UNAME_S),Darwin) -$(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_metal.o +$(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_metal.o ds4_image.o $(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS) else -$(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_cuda.o $(MMQ_OBJS) +$(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_cuda.o ds4_image.o $(MMQ_OBJS) $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) endif +# Only the Metal build of this test is exercised by `make test`; the CUDA +# variant still builds through `make test-glm53-kda`. +ifeq ($(UNAME_S),Darwin) +GLM53_KDA_DEFAULT_TEST := $(GLM53_KDA_TEST) +else +GLM53_KDA_DEFAULT_TEST := +endif + .PHONY: test-glm53-kda test-glm53-kda: $(GLM53_KDA_TEST) ./$(GLM53_KDA_TEST) @@ -565,7 +573,8 @@ tests/test_prompt_prefix: tests/test_prompt_prefix.o ds4_prompt_prefix.o test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \ tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \ - tests/test_deepseek4_vision_image tests/test_prompt_prefix $(SAMPLING_TEST) ds4 ds4-server ds4-bench ds4-agent + tests/test_deepseek4_vision_image tests/test_prompt_prefix $(SAMPLING_TEST) $(GLM53_KDA_DEFAULT_TEST) \ + ds4 ds4-server ds4-bench ds4-agent ./ds4-eval --self-test-extractors ./ds4_agent_test ./ds4_test @@ -576,6 +585,7 @@ test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \ ./tests/test_prompt_prefix ./tests/test_sampling ./tests/test_deepseek4_vision_image + @if [ -n "$(GLM53_KDA_DEFAULT_TEST)" ]; then ./$(GLM53_KDA_TEST); fi dspark-acceptance: ds4 DS4_DSPARK_MODEL="$(DS4_DSPARK_MODEL)" \ diff --git a/ds4.c b/ds4.c index 91ab214ab..8138462fc 100644 --- a/ds4.c +++ b/ds4.c @@ -37607,6 +37607,46 @@ static uint32_t glm53_graph_resume_prefill_min_tokens(void) { #define DS4_GLM53_INDEX_POOL_SIZE 4u #define DS4_GLM53_PREFILL_CHUNK_TOKENS 2048u +/* These four were compile-time constants with no way to try another value. + * They interact -- the prefill chunk and the layer-flush threshold are both + * 2048 and the flush comparison is a strict >, so raising the chunk to 4096 + * also switches per-layer flushing on across every layer -- so each is + * separately overridable, which is the only way to sweep one at a time. */ +static uint32_t glm_env_u32(const char *name, uint32_t fallback) { + const char *env = getenv(name); + if (!env || !env[0]) return fallback; + char *end = NULL; + errno = 0; + const unsigned long v = strtoul(env, &end, 10); + if (end == env || errno != 0 || v == 0ul || v > UINT32_MAX) return fallback; + return (uint32_t)v; +} + +static uint32_t glm53_prefill_chunk_tokens(void) { + return glm_env_u32("DS4_GLM_PREFILL_CHUNK_TOKENS", + DS4_GLM53_PREFILL_CHUNK_TOKENS); +} + +static uint32_t glm_full_attn_layer_flush_tokens(void) { + return glm_env_u32("DS4_GLM_FULL_ATTN_LAYER_FLUSH_TOKENS", + DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT); +} + +static uint32_t glm_full_attn_resident_cap(void) { + return glm_env_u32("DS4_GLM_FULL_ATTN_CAP", + DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT); +} + +static uint32_t glm_full_attn_streaming_cap(void) { + return glm_env_u32("DS4_GLM_FULL_ATTN_STREAMING_CAP", + DS4_GLM_METAL_STREAMING_FULL_ATTN_CONTEXT); +} + +static uint32_t glm_indexed_prefill_score_scratch_mb(void) { + return glm_env_u32("DS4_GLM_PREFILL_SCORE_SCRATCH_MB", + DS4_GLM_METAL_INDEXED_PREFILL_SCORE_SCRATCH_MB); +} + static uint32_t glm_graph_full_attention_cap(uint32_t ctx_size, bool ssd_streaming); static uint32_t glm_graph_indexed_prefill_chunk_tokens( @@ -37884,8 +37924,9 @@ static uint32_t glm_graph_batch_row_cap( bool expanded_kv) { if (ds4_model_is_glm53()) { uint32_t cap = full_attention_cap; - if (cap > DS4_GLM53_PREFILL_CHUNK_TOKENS) { - cap = DS4_GLM53_PREFILL_CHUNK_TOKENS; + const uint32_t glm53_chunk = glm53_prefill_chunk_tokens(); + if (cap > glm53_chunk) { + cap = glm53_chunk; } if (indexed_prefill_cap != 0 && cap > indexed_prefill_cap) { cap = indexed_prefill_cap; @@ -41071,6 +41112,9 @@ typedef struct ds4_glm_gpu_graph { ds4_gpu_tensor *kda_k; ds4_gpu_tensor *kda_v; ds4_gpu_tensor *kda_lowrank; + /* f_a and g_a run concurrently when the gate chain is paired, so they + * cannot share one low-rank destination the way the serial chain did. */ + ds4_gpu_tensor *kda_lowrank_g; ds4_gpu_tensor *kda_raw_gate; ds4_gpu_tensor *kda_raw_beta; ds4_gpu_tensor *kda_output_gate; @@ -41151,6 +41195,9 @@ typedef struct ds4_glm_gpu_graph { ds4_gpu_tensor *qk_low; ds4_gpu_tensor *attn_partial_lora; ds4_gpu_tensor *attn_partial_ms; + ds4_gpu_tensor *attn_exact_scores; + ds4_gpu_tensor *attn_exact_lora; + ds4_gpu_tensor *attn_exact_denom; ds4_gpu_tensor *batch_indexer_k; ds4_gpu_tensor *batch_indexer_gate; ds4_gpu_tensor *batch_indexer_q; @@ -41834,8 +41881,8 @@ static bool glm_graph_memory_guard_slice_with_transient( static uint32_t glm_graph_full_attention_cap(uint32_t ctx_size, bool ssd_streaming) { uint32_t cap = ssd_streaming ? - DS4_GLM_METAL_STREAMING_FULL_ATTN_CONTEXT : - DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT; + glm_full_attn_streaming_cap() : + glm_full_attn_resident_cap(); if (ctx_size >= DS4_GLM_METAL_LONG_CONTEXT_THRESHOLD && cap > DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT) { cap = DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT; @@ -41853,8 +41900,9 @@ static uint32_t glm_graph_full_prefill_layer_flush_interval( * flush per layer: 76 command-buffer round-trips cost ~35ms while the * whole pass is ~70ms of GPU work. Real prefill chunks keep the * interactive per-layer flush. */ - return (n_tokens > DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT || - command_rows > DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT || + const uint32_t flush_tokens = glm_full_attn_layer_flush_tokens(); + return (n_tokens > flush_tokens || + command_rows > flush_tokens || (logits_requested && n_tokens > 8u)) ? 1u : 0u; } @@ -41921,23 +41969,97 @@ static uint32_t glm_graph_indexed_decode_split_blocks(void) { return (top_k + block_rows - 1u) / block_rows; } +/* Every GLM 5.3 Flash decode optimisation on this branch is behind its own + * rollback switch, and DS4_METAL_DISABLE_GLM53_FLASH_TUNING turns all of them + * off at once, so one variable restores the pre-branch paths for an A/B run. + * The table is the list; each entry is read once and cached. + * + * The prefill kernels pick their variant inside ds4_metal.m, where the shape + * that selects them is known, so their switches live there and read the same + * aggregate: + * DS4_METAL_DISABLE_GLM53_PREFILL_QK_LOW qk-low token tile + * DS4_METAL_DISABLE_GLM53_PREFILL_INDEXED_ATTN indexed attention head width + * DS4_METAL_DISABLE_GLM53_PREFILL_MOE_TAIL_CULL routed-expert tail cull + * DS4_METAL_DISABLE_GLM53_PREFILL_KDA_PREPARE blocked KDA prepare + * DS4_METAL_DISABLE_GLM53_PREFILL_KDA_RECURRENCE two values per SIMDgroup */ +typedef enum { + GLM53_FLASH_HC_PRODUCER_FUSE, + GLM53_FLASH_KDA_GATE_PAIR, + GLM53_FLASH_KDA_GATE_TRIO, + GLM53_FLASH_KDA_OUT_HC_EXPAND, + GLM53_FLASH_ATTN_OUT_HC_EXPAND, + GLM53_FLASH_FFN_HC_EXPAND_ADD, + GLM53_FLASH_SHARED_DOWN_HC_EXPAND, + GLM53_FLASH_DSA_EXACT, + GLM53_FLASH_FEATURE_COUNT +} glm53_flash_feature; + +static bool glm53_flash_feature_enabled(glm53_flash_feature feature) { + static const char *const switches[GLM53_FLASH_FEATURE_COUNT] = { + [GLM53_FLASH_HC_PRODUCER_FUSE] = "DS4_METAL_DISABLE_GLM53_HC_PRODUCER_FUSE", + [GLM53_FLASH_KDA_GATE_PAIR] = "DS4_METAL_DISABLE_GLM53_KDA_GATE_PAIR", + [GLM53_FLASH_KDA_GATE_TRIO] = "DS4_METAL_DISABLE_GLM53_KDA_GATE_TRIO", + [GLM53_FLASH_KDA_OUT_HC_EXPAND] = "DS4_METAL_DISABLE_GLM53_KDA_OUT_HC_EXPAND", + [GLM53_FLASH_ATTN_OUT_HC_EXPAND] = "DS4_METAL_DISABLE_GLM53_ATTN_OUT_HC_EXPAND", + [GLM53_FLASH_FFN_HC_EXPAND_ADD] = "DS4_METAL_DISABLE_GLM53_FFN_HC_EXPAND_ADD", + [GLM53_FLASH_SHARED_DOWN_HC_EXPAND] = "DS4_METAL_DISABLE_GLM53_SHARED_DOWN_HC_EXPAND", + [GLM53_FLASH_DSA_EXACT] = "DS4_METAL_DISABLE_GLM53_DSA_EXACT", + }; + static int8_t state[GLM53_FLASH_FEATURE_COUNT]; /* 0 unread, 1 on, -1 off */ + if (state[feature] == 0) { + state[feature] = + getenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") == NULL && + getenv(switches[feature]) == NULL ? 1 : -1; + } + return state[feature] > 0; +} + +/* Rows per split block for indexed decode attention. The 32/128 step at 1024 + * selected rows was never swept; DS4_GLM_DECODE_SPLIT_BLOCK_ROWS forces one + * value so it can be. A value the split path cannot honour is rejected by the + * availability guard below and falls back, so this cannot select a broken + * configuration. */ static uint32_t glm_graph_indexed_decode_split_block_rows_for(uint32_t n_selected) { + static int forced = -1; + if (forced < 0) { + const char *env = getenv("DS4_GLM_DECODE_SPLIT_BLOCK_ROWS"); + const int v = (env && env[0]) ? atoi(env) : 0; + forced = v > 0 ? v : 0; + } + if (forced > 0) return (uint32_t)forced; return n_selected <= 1024u ? 32u : 128u; } -static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) { +/* The grouped/split kernel scores with lane-split dots and reduces with an + * online softmax across row blocks, so its output is deterministic but not + * bit-identical to the generic kernel's. --quality keeps the generic kernel, + * as it does for every other exact-versus-fast pair; in default mode + * DS4_METAL_DISABLE_GLM53_DSA_SPLIT selects the generic kernel for A/B runs. */ +static bool glm_graph_indexed_decode_split_group8_available( + const ds4_glm_gpu_graph *g, + uint32_t n_selected) { #ifndef __APPLE__ + (void)g; (void)n_selected; return false; #else const uint32_t block_rows = glm_graph_indexed_decode_split_block_rows_for(n_selected); const uint32_t needed_blocks = block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; + static int disabled = -1; + if (disabled < 0) { + disabled = getenv("DS4_METAL_DISABLE_GLM53_DSA_SPLIT") != NULL; + } + if (g->quality || disabled) return false; return n_selected > 512u && block_rows > 0 && needed_blocks > 0 && needed_blocks <= glm_graph_indexed_decode_split_blocks() && - glm_graph_indexed_decode_split_blocks() <= 64u && + /* The reduce kernel walks one thread per block and refuses more + * than 64, so the runtime block count is what has to fit -- not + * split_blocks(), which is the worst-case buffer sizing and is 65 + * for GLM 5.3's 2051-row selection limit. */ + needed_blocks <= 64u && (DS4_N_HEAD % 8u) == 0 && DS4_N_KV_LORA == 512u && DS4_N_ROT == 64u && @@ -41945,6 +42067,41 @@ static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) #endif } +/* GLM 5.3 decode attention runs the phased kernels that reproduce + * kernel_glm_attention_indexed_decode's arithmetic operation for operation + * while sharing each cache row across heads (see the kernel comment in + * metal/dsv4_misc.metal). Their output is bit-identical to the generic + * kernel's, so --quality keeps them; DS4_METAL_DISABLE_GLM53_DSA_EXACT selects + * the generic kernel for A/B runs. The two-host tensor-parallel head split + * keeps the generic kernel until that configuration has been run. + * + * Below 128 selected rows the generic kernel's row traffic is a few megabytes + * per layer and the phased path's three extra dispatches cost more than they + * save: measured -0.4% at 36 rows, +0.3% at 134, +2.2% at 308 and +11% at + * 1,500. Since both kernels are exact, crossing the threshold mid-generation + * changes nothing but speed. */ +static bool glm_graph_indexed_decode_exact_available( + const ds4_glm_gpu_graph *g, + bool tp_split_heads, + uint32_t n_selected) { +#ifndef __APPLE__ + (void)g; + (void)tp_split_heads; + (void)n_selected; + return false; +#else + return glm53_flash_feature_enabled(GLM53_FLASH_DSA_EXACT) && + n_selected >= 128u && + g->glm53 && + !tp_split_heads && + g->attn_exact_scores && g->attn_exact_lora && g->attn_exact_denom && + DS4_N_ROT == 0u && + DS4_N_KV_LORA == 512u && + DS4_N_HEAD <= 64u && + glm_graph_compact_cache_is_f16(); +#endif +} + static bool glm_graph_prefill_stage_sync_boundary(void) { if (ds4_gpu_end_commands() == 0) return false; return ds4_gpu_begin_commands() != 0; @@ -42064,7 +42221,7 @@ static uint32_t glm_graph_indexed_prefill_chunk_tokens( getenv("DS4_GLM53_DISABLE_INDEXED_PREFILL"))) { return 0; } - uint32_t chunk = DS4_GLM53_PREFILL_CHUNK_TOKENS; + uint32_t chunk = glm53_prefill_chunk_tokens(); if (compact_cap != 0 && chunk > compact_cap) chunk = compact_cap; return chunk; } @@ -42078,7 +42235,7 @@ static uint32_t glm_graph_indexed_prefill_score_tokens( uint32_t indexed_prefill_cap, uint32_t compact_cap) { if (indexed_prefill_cap == 0 || compact_cap == 0) return 0; - const uint32_t scratch_mb = DS4_GLM_METAL_INDEXED_PREFILL_SCORE_SCRATCH_MB; + const uint32_t scratch_mb = glm_indexed_prefill_score_scratch_mb(); const uint64_t budget_bytes = (uint64_t)scratch_mb * 1024ull * 1024ull; const uint64_t score_columns = ds4_model_is_glm53() ? glm53_graph_indexer_pool_cap(compact_cap) : compact_cap; @@ -43046,6 +43203,7 @@ static void glm_graph_free(ds4_glm_gpu_graph *g) { ds4_gpu_tensor_free(g->kda_raw_beta); ds4_gpu_tensor_free(g->kda_raw_gate); ds4_gpu_tensor_free(g->kda_lowrank); + ds4_gpu_tensor_free(g->kda_lowrank_g); ds4_gpu_tensor_free(g->kda_v); ds4_gpu_tensor_free(g->kda_k); ds4_gpu_tensor_free(g->kda_q); @@ -43065,6 +43223,9 @@ static void glm_graph_free(ds4_glm_gpu_graph *g) { ds4_gpu_tensor_free(g->k_nope); ds4_gpu_tensor_free(g->kv_norm); ds4_gpu_tensor_free(g->kv_raw); + ds4_gpu_tensor_free(g->attn_exact_denom); + ds4_gpu_tensor_free(g->attn_exact_lora); + ds4_gpu_tensor_free(g->attn_exact_scores); ds4_gpu_tensor_free(g->attn_partial_ms); ds4_gpu_tensor_free(g->attn_partial_lora); ds4_gpu_tensor_free(g->qk_low); @@ -43424,6 +43585,8 @@ static bool glm_graph_alloc_slice( DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_v, kda_projection_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_lowrank, (uint64_t)DS4_N_KDA_HEAD_DIM * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_lowrank_g, + (uint64_t)DS4_N_KDA_HEAD_DIM * sizeof(float)); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_raw_gate, kda_projection_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_raw_beta, (uint64_t)DS4_N_KDA_HEAD * sizeof(float)); @@ -43449,6 +43612,19 @@ static bool glm_graph_alloc_slice( DS4_GLM_GRAPH_ALLOC_TENSOR(g->qk_low, qk_low_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_partial_lora, attn_partial_lora_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_partial_ms, attn_partial_ms_bytes); + if (g->glm53) { + /* Scratch for the phased exact attention kernels: one score per + * (head, selected row), and decode selects at most the dense window + * (ctx_cap) or the pool selector's limit. */ + const uint32_t selected_limit = glm53_graph_indexer_selected_limit(); + const uint32_t exact_rows = + g->ctx_cap > selected_limit ? g->ctx_cap : selected_limit; + DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_exact_scores, + (uint64_t)DS4_N_HEAD * exact_rows * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_exact_lora, qk_low_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_exact_denom, + (uint64_t)DS4_N_HEAD * sizeof(float)); + } DS4_GLM_GRAPH_ALLOC_TENSOR(g->kv_raw, kv_raw_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kv_norm, kv_norm_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->k_nope, k_nope_bytes); @@ -43695,7 +43871,7 @@ static bool glm53_graph_prefill_workspace_ensure( ds4_glm_gpu_graph *g, uint32_t rows) { if (!g || !g->glm53 || rows == 0 || - rows > DS4_GLM53_PREFILL_CHUNK_TOKENS) { + rows > glm53_prefill_chunk_tokens()) { return false; } if (g->glm53_prefill_cap >= rows) return true; @@ -43932,6 +44108,10 @@ static bool glm53_graph_hc_pre_rows( return ok; } +/* stage_t0 is the caller's layer-stage clock, or NULL when the stage profiler + * is off. When set, the four substages below are reported separately and the + * caller's own attn_output line reports only what is left after them, so the + * KDA lump is attributed rather than estimated. */ static bool glm53_graph_kda_attention_rows( ds4_glm_gpu_graph *g, const ds4_model *model, @@ -43939,7 +44119,14 @@ static bool glm53_graph_kda_attention_rows( uint32_t il, uint32_t pos0, uint32_t rows, - ds4_gpu_tensor *attn_out) { + ds4_gpu_tensor *attn_out, + double *stage_t0) { +#define GLM53_KDA_STAGE(name_) do { \ + if (ok && stage_t0) { \ + ok = metal_graph_layer_stage_profile_boundary( \ + "glm53_kda", (name_), il, pos0, rows, stage_t0); \ + } \ + } while (0) if (!g || !model || !l || il >= DS4_MAX_LAYER || rows == 0 || !attn_out || !g->layer_kda_conv_state[il] || !g->layer_kda_recurrent_state[il]) { @@ -43970,6 +44157,7 @@ static bool glm53_graph_kda_attention_rows( if (ok) metal_graph_debug_dump_tensor( "glm53_kda_v_ready", g->batch_kda_v, (uint64_t)rows * projection, il, pos0); + GLM53_KDA_STAGE("kda_qkv"); if (ok) failed_stage = "decay low-rank projection"; if (ok) failed_weight = l->kda_f_a; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_lowrank, model, @@ -44014,6 +44202,7 @@ static bool glm53_graph_kda_attention_rows( if (ok) metal_graph_debug_dump_tensor( "glm53_kda_output_gate_ready", g->batch_kda_output_gate, (uint64_t)rows * projection, il, pos0); + GLM53_KDA_STAGE("kda_gate"); if (ok) failed_stage = "KDA recurrence"; if (ok) failed_weight = NULL; if (ok) ok = ds4_gpu_glm53_kda_prefill( @@ -44041,6 +44230,7 @@ static bool glm53_graph_kda_attention_rows( if (ok) metal_graph_debug_dump_tensor( "glm53_kda_out_ready", g->batch_kda_out, (uint64_t)rows * projection, il, pos0); + GLM53_KDA_STAGE("kda_recur"); if (ok) failed_stage = "output projection"; if (ok) failed_weight = l->kda_output; if (ok) ok = glm53_graph_matmul_rows(attn_out, model, l->kda_output, @@ -44049,6 +44239,8 @@ static bool glm53_graph_kda_attention_rows( if (ok) metal_graph_debug_dump_tensor( "glm53_kda_attn_out_ready", attn_out, (uint64_t)rows * DS4_N_EMBD, il, pos0); + GLM53_KDA_STAGE("kda_out_proj"); +#undef GLM53_KDA_STAGE if (!ok) { if (failed_weight) { fprintf(stderr, @@ -44067,6 +44259,128 @@ static bool glm53_graph_kda_attention_rows( return ok; } +/* Timing-only skip-ablation for the GLM decode layer (comma list in + * DS4_GLM_DECODE_ABLATE): the skipped stage's output buffer keeps stale + * contents, so the run produces garbage text but every remaining dispatch + * (and every TP gate) still executes. Whole-token time deltas against a + * baseline run are the only reliable per-stage cost measurement — the + * stage profiler's per-stage command-buffer splits inflate small stages. */ +#define DS4_GLM_ABLATE_ATTN_OUT (1u << 0) +#define DS4_GLM_ABLATE_ATTN_CORE (1u << 1) +#define DS4_GLM_ABLATE_QPATH (1u << 2) +#define DS4_GLM_ABLATE_INDEXER (1u << 3) +#define DS4_GLM_ABLATE_ROUTED (1u << 4) +#define DS4_GLM_ABLATE_SHARED (1u << 5) +#define DS4_GLM_ABLATE_QKLOW (1u << 6) +/* KDA (linear attention), the whole stage and its four substages. KDA is the + * largest single line in the decode budget and had no ablation arm at all, so + * its cost was estimated rather than measured. */ +#define DS4_GLM_ABLATE_KDA (1u << 7) +#define DS4_GLM_ABLATE_KDA_QKV (1u << 8) +#define DS4_GLM_ABLATE_KDA_GATE (1u << 9) +#define DS4_GLM_ABLATE_KDA_RECUR (1u << 10) +#define DS4_GLM_ABLATE_KDA_OUT (1u << 11) +/* The two largest pieces of what the budget lumps into "norms, hyper- + * connections, residual, LM head": the mHC producer chain that runs twice per + * layer, and the output head. */ +#define DS4_GLM_ABLATE_HC (1u << 12) +#define DS4_GLM_ABLATE_HEAD (1u << 13) + +/* Exact token match against the comma list. A substring test stops working + * as soon as one stage name is a prefix of another: strstr(env, "kda") also + * fires on "kda_qkv", which would ablate the whole stage when only one + * substage was asked for. */ +static bool glm_ablate_names(const char *env, const char *name) { + const size_t n = strlen(name); + for (const char *p = env; *p; ) { + while (*p == ',' || *p == ' ' || *p == '\t') p++; + const char *start = p; + while (*p && *p != ',' && *p != ' ' && *p != '\t') p++; + if ((size_t)(p - start) == n && memcmp(start, name, n) == 0) return true; + } + return false; +} + +/* Non-destructive counterpart to the ablation mask (comma list in + * DS4_GLM_DECODE_REPEAT). A named stage is dispatched one extra time per + * site; because every stage listed here is a pure function of its inputs, + * running it twice writes the same bytes, so the model output is unchanged and + * data-dependent routing downstream cannot shift. The whole-token delta + * against a baseline is then one extra execution of that stage. + * + * Only idempotent stages are offered. The KDA recurrence advances the conv + * and recurrent state, and directional steering updates its input in place, so + * neither can be repeated this way and neither has a bit. + * + * This measures the same thing the ablation arms do from the other side, and + * disagreement between the two is a signal that one of them is lying. */ +#define DS4_GLM_REPEAT_HC_EXPAND (1u << 0) +#define DS4_GLM_REPEAT_HC_PRE (1u << 1) +#define DS4_GLM_REPEAT_HEAD (1u << 2) +#define DS4_GLM_REPEAT_KDA_QKV (1u << 3) +#define DS4_GLM_REPEAT_KDA_GATE (1u << 4) +#define DS4_GLM_REPEAT_KDA_OUT (1u << 5) +/* Router logits and top-k selection. Repeat rather than ablate is the only + * honest instrument here: skipping the router leaves a stale expert selection, + * which changes which experts the routed stage streams and so changes the very + * cost being measured. */ +#define DS4_GLM_REPEAT_ROUTER (1u << 6) +/* qk_low sits inside the attn_core ablation arm, so its cost has only ever + * been measured as part of the 7.7 ms attention stage. */ +#define DS4_GLM_REPEAT_QKLOW (1u << 7) + +static uint32_t glm_decode_repeat_mask(void) { + static int cached = -1; + if (cached < 0) { + uint32_t mask = 0; + const char *env = getenv("DS4_GLM_DECODE_REPEAT"); + if (env) { + if (glm_ablate_names(env, "hc_expand")) mask |= DS4_GLM_REPEAT_HC_EXPAND; + if (glm_ablate_names(env, "hc_pre")) mask |= DS4_GLM_REPEAT_HC_PRE; + if (glm_ablate_names(env, "head")) mask |= DS4_GLM_REPEAT_HEAD; + if (glm_ablate_names(env, "kda_qkv")) mask |= DS4_GLM_REPEAT_KDA_QKV; + if (glm_ablate_names(env, "kda_gate")) mask |= DS4_GLM_REPEAT_KDA_GATE; + if (glm_ablate_names(env, "kda_out")) mask |= DS4_GLM_REPEAT_KDA_OUT; + if (glm_ablate_names(env, "router")) mask |= DS4_GLM_REPEAT_ROUTER; + if (glm_ablate_names(env, "qklow")) mask |= DS4_GLM_REPEAT_QKLOW; + if (mask) { + fprintf(stderr, "ds4: GLM decode stage repeat active (mask 0x%x) — output stays correct, timing only\n", mask); + } + } + cached = (int)mask; + } + return (uint32_t)cached; +} + +static uint32_t glm_decode_ablate_mask(void) { + static int cached = -1; + if (cached < 0) { + uint32_t mask = 0; + const char *env = getenv("DS4_GLM_DECODE_ABLATE"); + if (env) { + if (glm_ablate_names(env, "attn_out")) mask |= DS4_GLM_ABLATE_ATTN_OUT; + if (glm_ablate_names(env, "attn_core")) mask |= DS4_GLM_ABLATE_ATTN_CORE; + if (glm_ablate_names(env, "qpath")) mask |= DS4_GLM_ABLATE_QPATH; + if (glm_ablate_names(env, "indexer")) mask |= DS4_GLM_ABLATE_INDEXER; + if (glm_ablate_names(env, "routed")) mask |= DS4_GLM_ABLATE_ROUTED; + if (glm_ablate_names(env, "shared")) mask |= DS4_GLM_ABLATE_SHARED; + if (glm_ablate_names(env, "qklow")) mask |= DS4_GLM_ABLATE_QKLOW; + if (glm_ablate_names(env, "kda")) mask |= DS4_GLM_ABLATE_KDA; + if (glm_ablate_names(env, "kda_qkv")) mask |= DS4_GLM_ABLATE_KDA_QKV; + if (glm_ablate_names(env, "kda_gate")) mask |= DS4_GLM_ABLATE_KDA_GATE; + if (glm_ablate_names(env, "kda_recur")) mask |= DS4_GLM_ABLATE_KDA_RECUR; + if (glm_ablate_names(env, "kda_out")) mask |= DS4_GLM_ABLATE_KDA_OUT; + if (glm_ablate_names(env, "hc")) mask |= DS4_GLM_ABLATE_HC; + if (glm_ablate_names(env, "head")) mask |= DS4_GLM_ABLATE_HEAD; + if (mask) { + fprintf(stderr, "ds4: GLM decode ablation active (mask 0x%x) — output is garbage, timing only\n", mask); + } + } + cached = (int)mask; + } + return (uint32_t)cached; +} + static bool glm53_graph_hc_pre( ds4_glm_gpu_graph *g, const ds4_model *model, @@ -44083,6 +44397,61 @@ static bool glm53_graph_hc_pre( } const uint32_t hc_dim = DS4_N_HC * DS4_N_EMBD; const uint32_t hc_mix = DS4_N_HC * (DS4_N_HC + 2u); +#if defined(__APPLE__) + /* The compound producer DeepSeek V4 already uses, with the BF16 mix + * weights GLM 5.3 stores instead of F16. It folds the plain RMSNorm, the + * 16384->24 mix matvec, the sinkhorn split/collapse and the weighted + * RMSNorm into one dispatch, so each of the 90 per-token sites costs one + * dispatch instead of four. */ + if (fn->type == DS4_TENSOR_BF16 && + hc_dim == 16384u && hc_mix == 24u && + DS4_N_EMBD == 4096u && DS4_N_HC == 4u && + !metal_graph_use_reference_hc_decode() && + glm53_flash_feature_enabled(GLM53_FLASH_HC_PRODUCER_FUSE) && + /* Same rollback switches as the DeepSeek F16 producer this shares a + * kernel with, so DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS and the two + * producer-specific variables disable both paths rather than leaving + * this one live after the other has been turned off. */ + metal_graph_ported_m5_decode_feature_enabled( + "DS4_METAL_DISABLE_PRE_M5_HC_PRODUCER_PRE_NORM_FUSE", + "DS4_METAL_DISABLE_M5_HC_PRODUCER_PRE_NORM_FUSE")) { + const int fused = ds4_gpu_hc_rms_norm_mix_split_norm_bf16_tensor( + g->hc_mix, + collapsed, + normalized, + g->hc_split, + residual_hc, + model->map, + model->size, + fn->abs_offset, + scale->abs_offset, + base->abs_offset, + norm->abs_offset, + hc_dim, + hc_mix, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_RMS_EPS, + DS4_HC_EPS, + DS4_RMS_EPS); + if (fused < 0) return false; + if (fused > 0) { + if (glm_decode_repeat_mask() & DS4_GLM_REPEAT_HC_PRE) { + if (ds4_gpu_hc_rms_norm_mix_split_norm_bf16_tensor( + g->hc_mix, collapsed, normalized, g->hc_split, + residual_hc, model->map, model->size, fn->abs_offset, + scale->abs_offset, base->abs_offset, norm->abs_offset, + hc_dim, hc_mix, DS4_N_EMBD, DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, DS4_RMS_EPS, DS4_HC_EPS, + DS4_RMS_EPS) < 0) { + return false; + } + } + return true; + } + } +#endif bool ok = ds4_gpu_rms_norm_plain_tensor(g->hc_flat, residual_hc, hc_dim, @@ -44114,17 +44483,22 @@ static bool glm53_graph_kda_attention( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_layer_weights *l, - uint32_t il) { + uint32_t il, + bool *hc_expanded) { + if (hc_expanded) *hc_expanded = false; if (!g || !model || !l || il >= DS4_MAX_LAYER || !g->layer_kda_conv_state[il] || !g->layer_kda_recurrent_state[il]) { return false; } const uint32_t projection = DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; + const uint32_t ablate = glm_decode_ablate_mask(); + const uint32_t repeat = glm_decode_repeat_mask(); bool qk_paired = false; #if defined(__APPLE__) bool qkv_paired = false; - if (getenv("DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE") == NULL && + if (!(ablate & DS4_GLM_ABLATE_KDA_QKV) && + getenv("DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE") == NULL && getenv("DS4_METAL_DISABLE_GLM53_BF16_QKV") == NULL && l->kda_q->type == DS4_TENSOR_BF16 && l->kda_k->type == DS4_TENSOR_BF16 && @@ -44146,7 +44520,8 @@ static bool glm53_graph_kda_attention( const bool qkv_paired = false; #endif #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) - if (l->kda_q->type == DS4_TENSOR_Q4_K && + if (!(ablate & DS4_GLM_ABLATE_KDA_QKV) && + l->kda_q->type == DS4_TENSOR_Q4_K && l->kda_k->type == DS4_TENSOR_Q4_K && getenv("DS4_CUDA_GLM_DISABLE_KDA_QK_PAIR") == NULL) { qk_paired = ds4_gpu_matmul_q4_K_pair_decode_tensor( @@ -44161,33 +44536,158 @@ static bool glm53_graph_kda_attention( g->attn_norm) != 0; } #endif - bool ok = qkv_paired || qk_paired || - glm53_graph_matmul(g->kda_q, model, l->kda_q, - DS4_N_EMBD, projection, g->attn_norm); - if (ok && !qkv_paired && !qk_paired) { - ok = glm53_graph_matmul(g->kda_k, model, l->kda_k, - DS4_N_EMBD, projection, g->attn_norm); - } - if (ok && !qkv_paired) { - ok = glm53_graph_matmul(g->kda_v, model, l->kda_v, - DS4_N_EMBD, projection, g->attn_norm); - } - if (ok) ok = glm53_graph_matmul( - g->kda_lowrank, model, l->kda_f_a, - DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm); - if (ok) ok = glm53_graph_matmul( - g->kda_raw_gate, model, l->kda_f_b, - DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank); - if (ok) ok = glm53_graph_matmul( - g->kda_raw_beta, model, l->kda_beta, - DS4_N_EMBD, DS4_N_KDA_HEAD, g->attn_norm); - if (ok) ok = glm53_graph_matmul( - g->kda_lowrank, model, l->kda_g_a, - DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm); - if (ok) ok = glm53_graph_matmul( - g->kda_output_gate, model, l->kda_g_b, - DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank); - if (ok) ok = ds4_gpu_glm53_kda_decode( + bool ok = true; + /* Each substage is skipped whole; its output tensor then keeps the stale + * contents from the previous token, which is the documented ablation + * contract above -- garbage text, but every other dispatch still runs. */ + if (!(ablate & DS4_GLM_ABLATE_KDA_QKV)) { + ok = qkv_paired || qk_paired || + glm53_graph_matmul(g->kda_q, model, l->kda_q, + DS4_N_EMBD, projection, g->attn_norm); + if (ok && !qkv_paired && !qk_paired) { + ok = glm53_graph_matmul(g->kda_k, model, l->kda_k, + DS4_N_EMBD, projection, g->attn_norm); + } + if (ok && !qkv_paired) { + ok = glm53_graph_matmul(g->kda_v, model, l->kda_v, + DS4_N_EMBD, projection, g->attn_norm); + } + /* Repeat whichever variant actually ran. Re-dispatching the serial + * matvecs when the fused kernel did the work would price a path that + * is not executing. */ + if (ok && (repeat & DS4_GLM_REPEAT_KDA_QKV)) { +#if defined(__APPLE__) + if (qkv_paired) { + ok = ds4_gpu_glm53_matmul_bf16_qkv( + g->kda_q, g->kda_k, g->kda_v, + model->map, model->size, + l->kda_q->abs_offset, l->kda_k->abs_offset, + l->kda_v->abs_offset, + DS4_N_EMBD, projection, g->attn_norm) != 0; + } else +#endif + if (qk_paired) { + ok = glm53_graph_matmul(g->kda_v, model, l->kda_v, + DS4_N_EMBD, projection, g->attn_norm); + } else { + ok = glm53_graph_matmul(g->kda_q, model, l->kda_q, + DS4_N_EMBD, projection, g->attn_norm) && + glm53_graph_matmul(g->kda_k, model, l->kda_k, + DS4_N_EMBD, projection, g->attn_norm) && + glm53_graph_matmul(g->kda_v, model, l->kda_v, + DS4_N_EMBD, projection, g->attn_norm); + } + } + } + bool gate_paired = false; +#if defined(__APPLE__) + /* Five serial matvecs become two paired dispatches plus beta. f_a and g_a + * share the attn_norm row; f_b and g_b read the two low-rank vectors those + * produce, which is why the pair kernel takes separate inputs. beta is a + * different output width and stays on its own. */ + if (ok && !(ablate & DS4_GLM_ABLATE_KDA_GATE) && + g->kda_lowrank_g && + l->kda_f_a->type == DS4_TENSOR_BF16 && + l->kda_g_a->type == DS4_TENSOR_BF16 && + l->kda_f_b->type == DS4_TENSOR_BF16 && + l->kda_g_b->type == DS4_TENSOR_BF16 && + glm53_flash_feature_enabled(GLM53_FLASH_KDA_GATE_PAIR)) { + /* beta reads the same attn_norm row as f_a and g_a, only at a + * shorter output width, so the trio kernel carries all three and the + * chain drops from three dispatches to two. */ + bool beta_fused = l->kda_beta->type == DS4_TENSOR_BF16 && + glm53_flash_feature_enabled(GLM53_FLASH_KDA_GATE_TRIO) && + ds4_gpu_glm53_matmul_bf16_trio( + g->kda_lowrank, g->kda_lowrank_g, g->kda_raw_beta, + model->map, model->size, + l->kda_f_a->abs_offset, l->kda_g_a->abs_offset, + l->kda_beta->abs_offset, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, DS4_N_KDA_HEAD, + g->attn_norm) != 0; + gate_paired = beta_fused || ds4_gpu_glm53_matmul_bf16_pair( + g->kda_lowrank, g->kda_lowrank_g, + model->map, model->size, + l->kda_f_a->abs_offset, l->kda_g_a->abs_offset, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, + g->attn_norm, g->attn_norm) != 0; + if (gate_paired) { + gate_paired = ds4_gpu_glm53_matmul_bf16_pair( + g->kda_raw_gate, g->kda_output_gate, + model->map, model->size, + l->kda_f_b->abs_offset, l->kda_g_b->abs_offset, + DS4_N_KDA_HEAD_DIM, projection, + g->kda_lowrank, g->kda_lowrank_g) != 0; + } + /* A partial failure is safe to fall back from: both halves are pure + * functions of attn_norm, so the serial chain below simply recomputes + * the same values into the same buffers. */ + if (gate_paired && !beta_fused) { + ok = glm53_graph_matmul( + g->kda_raw_beta, model, l->kda_beta, + DS4_N_EMBD, DS4_N_KDA_HEAD, g->attn_norm); + } + if (gate_paired) { + /* The serial fallback's repeat below is unreachable once pairing + * succeeds, so the paired path carries its own. */ + if (ok && (repeat & DS4_GLM_REPEAT_KDA_GATE)) { + if (beta_fused) { + ok = ds4_gpu_glm53_matmul_bf16_trio( + g->kda_lowrank, g->kda_lowrank_g, g->kda_raw_beta, + model->map, model->size, + l->kda_f_a->abs_offset, l->kda_g_a->abs_offset, + l->kda_beta->abs_offset, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, DS4_N_KDA_HEAD, + g->attn_norm) != 0; + } else { + ok = ds4_gpu_glm53_matmul_bf16_pair( + g->kda_lowrank, g->kda_lowrank_g, + model->map, model->size, + l->kda_f_a->abs_offset, l->kda_g_a->abs_offset, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, + g->attn_norm, g->attn_norm) != 0 && + glm53_graph_matmul(g->kda_raw_beta, model, l->kda_beta, + DS4_N_EMBD, DS4_N_KDA_HEAD, g->attn_norm); + } + if (ok) ok = ds4_gpu_glm53_matmul_bf16_pair( + g->kda_raw_gate, g->kda_output_gate, + model->map, model->size, + l->kda_f_b->abs_offset, l->kda_g_b->abs_offset, + DS4_N_KDA_HEAD_DIM, projection, + g->kda_lowrank, g->kda_lowrank_g) != 0; + } + } + } +#endif + if (!gate_paired && !(ablate & DS4_GLM_ABLATE_KDA_GATE)) { + if (ok) ok = glm53_graph_matmul( + g->kda_lowrank, model, l->kda_f_a, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm); + if (ok) ok = glm53_graph_matmul( + g->kda_raw_gate, model, l->kda_f_b, + DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank); + if (ok) ok = glm53_graph_matmul( + g->kda_raw_beta, model, l->kda_beta, + DS4_N_EMBD, DS4_N_KDA_HEAD, g->attn_norm); + if (ok) ok = glm53_graph_matmul( + g->kda_lowrank, model, l->kda_g_a, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm); + if (ok) ok = glm53_graph_matmul( + g->kda_output_gate, model, l->kda_g_b, + DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank); + if (ok && (repeat & DS4_GLM_REPEAT_KDA_GATE)) { + ok = glm53_graph_matmul(g->kda_lowrank, model, l->kda_f_a, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm) && + glm53_graph_matmul(g->kda_raw_gate, model, l->kda_f_b, + DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank) && + glm53_graph_matmul(g->kda_raw_beta, model, l->kda_beta, + DS4_N_EMBD, DS4_N_KDA_HEAD, g->attn_norm) && + glm53_graph_matmul(g->kda_lowrank, model, l->kda_g_a, + DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm) && + glm53_graph_matmul(g->kda_output_gate, model, l->kda_g_b, + DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank); + } + } + if (ok && !(ablate & DS4_GLM_ABLATE_KDA_RECUR)) ok = ds4_gpu_glm53_kda_decode( g->kda_out, g->layer_kda_conv_state[il], g->layer_kda_recurrent_state[il], @@ -44209,12 +44709,56 @@ static bool glm53_graph_kda_attention( 1, DS4_KDA_GATE_LOWER_BOUND, DS4_RMS_EPS) != 0; - if (ok) ok = glm53_graph_matmul(g->attn_out, - model, - l->kda_output, - projection, - DS4_N_EMBD, - g->kda_out); + if (ok && !(ablate & DS4_GLM_ABLATE_KDA_OUT)) { +#if defined(__APPLE__) + /* Fold the HC expansion into this projection's epilogue: the + * simdgroup that finishes output row d already holds it in lane 0, so + * it can write the four HC streams there rather than have a separate + * dispatch read the row straight back. Only valid when nothing sits + * between the two -- directional steering would, so it is required to + * be inactive. */ + if (hc_expanded && g->glm53 && + l->kda_output->type == DS4_TENSOR_BF16 && + g->directional_steering_attn_scale == 0.0f && + g->hc_after_attn && g->hc_cur && g->hc_post && g->hc_comb && + glm53_flash_feature_enabled(GLM53_FLASH_KDA_OUT_HC_EXPAND)) { + if (ds4_gpu_glm53_matmul_bf16_hc_expand4( + g->attn_out, g->hc_after_attn, + model->map, model->size, l->kda_output->abs_offset, + projection, DS4_N_EMBD, + g->kda_out, g->hc_cur, g->hc_post, g->hc_comb, + DS4_N_HC) != 0) { + *hc_expanded = true; + } + } +#endif + if (!(hc_expanded && *hc_expanded)) { + ok = glm53_graph_matmul(g->attn_out, + model, + l->kda_output, + projection, + DS4_N_EMBD, + g->kda_out); + } + if (ok && (repeat & DS4_GLM_REPEAT_KDA_OUT)) { +#if defined(__APPLE__) + if (hc_expanded && *hc_expanded) { + /* Price the projection-plus-expand kernel that is deployed, + * not the bare projection it replaced. */ + ok = ds4_gpu_glm53_matmul_bf16_hc_expand4( + g->attn_out, g->hc_after_attn, + model->map, model->size, l->kda_output->abs_offset, + projection, DS4_N_EMBD, + g->kda_out, g->hc_cur, g->hc_post, g->hc_comb, + DS4_N_HC) != 0; + } else +#endif + { + ok = glm53_graph_matmul(g->attn_out, model, l->kda_output, + projection, DS4_N_EMBD, g->kda_out); + } + } + } return ok; } @@ -44835,42 +45379,6 @@ static double glm_graph_streaming_async_profile_ms(void) { return now_sec() * 1000.0; } -/* Timing-only skip-ablation for the GLM decode layer (comma list in - * DS4_GLM_DECODE_ABLATE): the skipped stage's output buffer keeps stale - * contents, so the run produces garbage text but every remaining dispatch - * (and every TP gate) still executes. Whole-token time deltas against a - * baseline run are the only reliable per-stage cost measurement — the - * stage profiler's per-stage command-buffer splits inflate small stages. */ -#define DS4_GLM_ABLATE_ATTN_OUT (1u << 0) -#define DS4_GLM_ABLATE_ATTN_CORE (1u << 1) -#define DS4_GLM_ABLATE_QPATH (1u << 2) -#define DS4_GLM_ABLATE_INDEXER (1u << 3) -#define DS4_GLM_ABLATE_ROUTED (1u << 4) -#define DS4_GLM_ABLATE_SHARED (1u << 5) -#define DS4_GLM_ABLATE_QKLOW (1u << 6) - -static uint32_t glm_decode_ablate_mask(void) { - static int cached = -1; - if (cached < 0) { - uint32_t mask = 0; - const char *env = getenv("DS4_GLM_DECODE_ABLATE"); - if (env) { - if (strstr(env, "attn_out")) mask |= DS4_GLM_ABLATE_ATTN_OUT; - if (strstr(env, "attn_core")) mask |= DS4_GLM_ABLATE_ATTN_CORE; - if (strstr(env, "qpath")) mask |= DS4_GLM_ABLATE_QPATH; - if (strstr(env, "indexer")) mask |= DS4_GLM_ABLATE_INDEXER; - if (strstr(env, "routed")) mask |= DS4_GLM_ABLATE_ROUTED; - if (strstr(env, "shared")) mask |= DS4_GLM_ABLATE_SHARED; - if (strstr(env, "qklow")) mask |= DS4_GLM_ABLATE_QKLOW; - if (mask) { - fprintf(stderr, "ds4: GLM decode ablation active (mask 0x%x) — output is garbage, timing only\n", mask); - } - } - cached = (int)mask; - } - return (uint32_t)cached; -} - static bool glm_graph_encode_shared_swiglu_one( ds4_gpu_tensor *mid, ds4_gpu_tensor *gate, @@ -44972,6 +45480,11 @@ static bool glm_graph_encode_sparse_ffn_one( ds4_gpu_tensor *ffn_sum, ds4_gpu_tensor *tmp, bool add_residual, + bool defer_final_sum, + /* Out: set when the shared down-projection was fused with the HC + * expand, so the caller must skip the expand entirely rather than + * merely the sum. NULL if the caller cannot honour that. */ + bool *hc_expand_done, bool stage_profile, double *stage_t0) { uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; @@ -45002,6 +45515,21 @@ static bool glm_graph_encode_sparse_ffn_one( DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE) != 0; + if (ok && (glm_decode_repeat_mask() & DS4_GLM_REPEAT_ROUTER)) { + ok = ds4_gpu_matmul_f32_tensor(g->router_logits, + model->map, model->size, + l->ffn_gate_inp->abs_offset, + DS4_N_EMBD, DS4_N_EXPERT, ffn_norm, 1) != 0 && + ds4_gpu_glm_router_select_tensor(g->router_selected, + g->router_weights, + g->router_probs, + model->map, model->size, + l->ffn_exp_probs_b->abs_offset, + g->router_logits, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE) != 0; + } if (ok) ok = glm_graph_profile_stage(stage_profile, "glm_decode_ffn", "router", @@ -45244,16 +45772,53 @@ static bool glm_graph_encode_sparse_ffn_one( g->ssd_streaming, stage_profile, stage_t0); - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_sum, - model, - l->ffn_down_shexp->abs_offset, - DS4_N_FF_EXP, - DS4_N_EMBD, - ffn_mid, - il, - pos, - "shared_down", - g->ssd_streaming) != 0; + bool shared_down_fused = false; +#if defined(__APPLE__) + /* On this ordering the routed stage has already run, so ffn_mid still + * holds the shared mid and ffn_out holds the routed result -- exactly + * the input DeepSeek's fused kernel wants. It does the shared + * down-projection, adds the routed output and expands into the HC + * streams in one dispatch, replacing this matvec and the caller's + * expand together. Metal only, like the other epilogues. */ + if (ok && hc_expand_done && defer_final_sum && g->glm53 && + !g->ssd_streaming && + l->ffn_down_shexp->type == DS4_TENSOR_Q8_0 && + g->hc_next && g->hc_after_attn && g->hc_split && + glm53_flash_feature_enabled(GLM53_FLASH_SHARED_DOWN_HC_EXPAND) && + ds4_gpu_shared_down_hc_expand_q8_0_tensor( + g->hc_next, ffn_sum, + model->map, model->size, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, DS4_N_EMBD, + ffn_mid, ffn_out, + g->hc_after_attn, g->hc_split, + DS4_N_EMBD, DS4_N_HC) != 0) { + shared_down_fused = true; + *hc_expand_done = true; + if (glm_decode_repeat_mask() & DS4_GLM_REPEAT_HC_EXPAND) { + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor( + g->hc_next, ffn_sum, + model->map, model->size, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, DS4_N_EMBD, + ffn_mid, ffn_out, + g->hc_after_attn, g->hc_split, + DS4_N_EMBD, DS4_N_HC) != 0; + } + } +#endif + if (ok && !shared_down_fused) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_sum, + model, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, + DS4_N_EMBD, + ffn_mid, + il, + pos, + "shared_down", + g->ssd_streaming) != 0; + } if (ok) ok = glm_graph_profile_stage(stage_profile, "glm_decode_ffn", "shared_down", @@ -45277,6 +45842,8 @@ static bool glm_graph_encode_sparse_ffn_one( after_attn, tmp, DS4_N_EMBD) != 0; + } else if (ok && defer_final_sum) { + /* caller folds ffn_out + ffn_sum into the HC expand */ } else if (ok) { ok = ds4_gpu_add_tensor(next, ffn_out, @@ -45323,6 +45890,13 @@ static bool glm_graph_encode_ffn_one_normed_from( ds4_gpu_tensor *ffn_sum, ds4_gpu_tensor *tmp, bool add_residual, + /* When set, the routed+shared sum is left undone so the caller can + * fold it into the HC expand's has_add path instead of paying a + * separate add dispatch for it. */ + bool *defer_final_sum, + /* Out: the shared down-projection and the HC expand were fused, so the + * caller must skip the expand as well as the sum. */ + bool *hc_expand_done, bool stage_profile, double *stage_t0) { if (!g || !model || !l || !ffn_norm || !after_attn || !next || @@ -45332,6 +45906,9 @@ static bool glm_graph_encode_ffn_one_normed_from( } if (il < DS4_N_LEADING_DENSE) { + /* Dense layers have no routed/shared split to defer. */ + if (defer_final_sum) *defer_final_sum = false; + if (hc_expand_done) *hc_expand_done = false; const uint64_t hidden = l->ffn_gate->dim[1]; const bool can_fuse_gate_up = glm_graph_weights_are_q8_0(model, @@ -45455,6 +46032,8 @@ static bool glm_graph_encode_ffn_one_normed_from( ffn_sum, tmp, add_residual, + defer_final_sum && *defer_final_sum, + hc_expand_done, stage_profile, stage_t0); } @@ -45467,6 +46046,28 @@ static bool glm53_graph_encode_ffn_tail_one( uint32_t pos, bool stage_profile, double *stage_t0) { + /* The routed+shared sum and the HC expand are adjacent and the expand + * kernel already has a has_add path, so on the decode tail they collapse + * into one dispatch. Directional steering would have to run on the summed + * value in between, so it is required to be inactive. */ + bool hc_expand_done = false; + bool defer_sum = false; +#if defined(__APPLE__) + /* Metal only, like the two attention-side epilogues. ds4_gpu_hc_expand_add_ + * tensor is a stub on ROCm, and while CUDA implements it, changing that + * backend's arithmetic from a change measured only on Metal is not + * something this should do silently. + * + * Also declines while a debug dump of this layer is armed: the deferral + * leaves g->next unwritten, so the "ffn_out" dump below would capture + * whatever the buffer held from a previous token. */ + defer_sum = + g->glm53 && g->ffn_sum && g->ffn_out && g->hc_next && + g->hc_after_attn && g->hc_post && g->hc_comb && + g->directional_steering_ffn_scale == 0.0f && + !metal_graph_debug_wants("ffn_out", il, pos) && + glm53_flash_feature_enabled(GLM53_FLASH_FFN_HC_EXPAND_ADD); +#endif bool ok = glm_graph_encode_ffn_one_normed_from(g, model, l, @@ -45482,6 +46083,8 @@ static bool glm53_graph_encode_ffn_tail_one( g->ffn_sum, g->attn_out, false, + &defer_sum, + &hc_expand_done, stage_profile, stage_t0); if (ok) { @@ -45492,7 +46095,26 @@ static bool glm53_graph_encode_ffn_tail_one( pos); ok = glm_graph_apply_directional_steering_ffn(g, g->next, il, 1); } - if (ok) { + if (ok && hc_expand_done) { + /* shared_down + routed add + HC expand all happened in one dispatch, + * and that dispatch carries its own repeat arm -- re-dispatching the + * standalone expand here would price a path that is not running. */ + } else if (ok && defer_sum) { + ok = ds4_gpu_hc_expand_add_tensor(g->hc_next, + g->ffn_out, + g->ffn_sum, + g->hc_after_attn, + g->hc_post, + g->hc_comb, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok && (glm_decode_repeat_mask() & DS4_GLM_REPEAT_HC_EXPAND)) { + ok = ds4_gpu_hc_expand_add_tensor(g->hc_next, g->ffn_out, + g->ffn_sum, g->hc_after_attn, + g->hc_post, g->hc_comb, + DS4_N_EMBD, DS4_N_HC) != 0; + } + } else if (ok) { ok = ds4_gpu_hc_expand_tensor(g->hc_next, g->next, g->hc_after_attn, @@ -45500,6 +46122,12 @@ static bool glm53_graph_encode_ffn_tail_one( g->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; + if (ok && (glm_decode_repeat_mask() & DS4_GLM_REPEAT_HC_EXPAND)) { + ok = ds4_gpu_hc_expand_tensor(g->hc_next, g->next, + g->hc_after_attn, g->hc_post, + g->hc_comb, DS4_N_EMBD, + DS4_N_HC) != 0; + } } return ok; } @@ -45558,6 +46186,8 @@ static bool glm_graph_encode_ffn_one_from( ffn_sum, tmp, true, + NULL, + NULL, stage_profile, stage_t0); } @@ -48095,7 +48725,7 @@ static bool glm_graph_forward_tokens( uint32_t work_total) { if (!g || !model || !weights || !tokens || n_tokens == 0 || - (g->glm53 && n_tokens > DS4_GLM53_PREFILL_CHUNK_TOKENS) || + (g->glm53 && n_tokens > glm53_prefill_chunk_tokens()) || g->layer_count == 0 || !glm_graph_span_fits_context(g, pos0, n_tokens)) { return false; @@ -48608,7 +49238,8 @@ static bool glm_graph_forward_tokens( il, pos0, n_tokens, - g->batch_attn_out); + g->batch_attn_out, + layer_stage_profile ? &layer_stage_t0 : NULL); if (ok) { const uint64_t projection_rows = (uint64_t)n_tokens * DS4_N_KDA_HEAD * @@ -49803,7 +50434,8 @@ static bool glm_graph_forward_indexed_tokens( il, pos0, n_tokens, - g->batch_attn_out); + g->batch_attn_out, + layer_stage_profile ? &layer_stage_t0 : NULL); goto glm53_indexed_attention_done; } if (ok) { @@ -50749,6 +51381,8 @@ static bool glm_graph_forward_indexed_tokens( g->ffn_sum, g->attn_out, true, + NULL, + NULL, false, NULL); } else if (ok) { @@ -51067,8 +51701,9 @@ static bool glm_graph_prefill_range( while (done < n_tokens) { const uint32_t pos = pos0 + done; uint32_t chunk = n_tokens - done; - if (chunk > DS4_GLM53_PREFILL_CHUNK_TOKENS) { - chunk = DS4_GLM53_PREFILL_CHUNK_TOKENS; + const uint32_t glm53_chunk = glm53_prefill_chunk_tokens(); + if (chunk > glm53_chunk) { + chunk = glm53_chunk; } if (pos < g->ctx_cap) { const uint32_t dense_left = g->ctx_cap - pos; @@ -51665,7 +52300,10 @@ static bool glm_graph_forward_token( &decode_stage_t0); } + const uint32_t decode_ablate = glm_decode_ablate_mask(); + bool attn_hc_expanded = false; DS4_GLM_FT_STAGE("attention mHC pre"); + if (ok && g->glm53 && (decode_ablate & DS4_GLM_ABLATE_HC)) { /* ablate */ } else if (ok && g->glm53) { ok = glm53_graph_hc_pre(g, model, @@ -51688,10 +52326,12 @@ static bool glm_graph_forward_token( DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_norm"); if (ok && glm53_kda) { DS4_GLM_FT_STAGE("KDA attention"); - ok = glm53_graph_kda_attention(g, model, l, il); + if (!(decode_ablate & DS4_GLM_ABLATE_KDA)) { + ok = glm53_graph_kda_attention(g, model, l, il, + &attn_hc_expanded); + } goto glm53_attention_done; } - const uint32_t decode_ablate = glm_decode_ablate_mask(); DS4_GLM_FT_STAGE("DSA q_a projection"); if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q_rank, @@ -52096,6 +52736,16 @@ static bool glm_graph_forward_token( DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_KEY_MLA) != 0; + if (ok && (glm_decode_repeat_mask() & DS4_GLM_REPEAT_QKLOW)) { + ok = ds4_gpu_glm_qk_lowrank_typed_tensor( + tp_split_layer_heads ? tp_qk_low : g->qk_low, + tp_split_layer_heads ? tp_q : g->q, + model->map, model->size, k_weight_offset, + l->attn_k_b->type, + tp_split_layer_heads ? tp_head_count : DS4_N_HEAD, + DS4_N_KV_LORA, (uint32_t)g->q_nope, + DS4_N_KEY_MLA) != 0; + } if (ok) metal_graph_debug_dump_tensor("glm_decode_qk_low", g->qk_low, (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, @@ -52109,7 +52759,31 @@ static bool glm_graph_forward_token( * rest of the layer stays finite (timing-only). */ ok = ds4_gpu_tensor_fill_f32(g->heads, 0.0f, (uint64_t)g->heads_dim) != 0; - } else if (ok && glm_graph_indexed_decode_split_group8_available(last_indexer_selected_count)) { + } else if (ok && l->attn_v_b->type == DS4_TENSOR_Q8_0 && + glm_graph_indexed_decode_exact_available( + g, tp_split_layer_heads, last_indexer_selected_count)) { + ok = ds4_gpu_glm_attention_indexed_decode_exact_typed_tensor( + g->heads, + g->attn_exact_scores, + g->attn_exact_lora, + g->attn_exact_denom, + g->qk_low, + g->layer_kv_lora_cache[il], + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + last_indexer_selected, + last_indexer_selected_count, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA) != 0; + } else if (ok && glm_graph_indexed_decode_split_group8_available( + g, last_indexer_selected_count)) { const uint32_t split_block_rows = glm_graph_indexed_decode_split_block_rows_for(last_indexer_selected_count); const uint32_t split_blocks = @@ -52138,7 +52812,14 @@ static bool glm_graph_forward_token( l->attn_v_b->type, last_indexer_selected, last_indexer_selected_count, - true, + /* GLM 5.2's selections are a dense range or a + * top-k over visible rows, always in range, so + * it keeps the unchecked variant it always ran; + * the checked one costs about 2% of its decode. + * GLM 5.3 pads with UINT32_MAX sentinels and + * must not skip the check, should it ever get + * here. */ + !g->glm53, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), tp_split_layer_heads ? tp_head_count : DS4_N_HEAD, @@ -52286,16 +52967,38 @@ static bool glm_graph_forward_token( g->tp_in[slot], DS4_N_EMBD) != 0; } else { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->attn_out, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->heads, - il, - pos, - "attn_o", - g->ssd_streaming) != 0; +#if defined(__APPLE__) + /* Same epilogue trick as kda_output. This projection is Q8_0, + * and DeepSeek's fused kernel already covers that shape and + * reads post/comb from hc_split at the offsets GLM uses, so no + * new kernel is needed here. */ + if (g->glm53 && !g->ssd_streaming && + l->attn_output->type == DS4_TENSOR_Q8_0 && + g->directional_steering_attn_scale == 0.0f && + g->hc_after_attn && g->hc_cur && g->hc_split && + glm53_flash_feature_enabled(GLM53_FLASH_ATTN_OUT_HC_EXPAND) && + ds4_gpu_matmul_q8_0_hc_expand_tensor( + g->hc_after_attn, g->attn_out, + model->map, model->size, + l->attn_output->abs_offset, + g->heads_dim, DS4_N_EMBD, g->heads, + g->hc_cur, g->hc_split, + DS4_N_EMBD, DS4_N_HC) != 0) { + attn_hc_expanded = true; + } +#endif + if (!attn_hc_expanded) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->attn_out, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->heads, + il, + pos, + "attn_o", + g->ssd_streaming) != 0; + } } } glm53_attention_done: @@ -52311,13 +53014,24 @@ static bool glm_graph_forward_token( g, g->attn_out, il, 1); } if (ok && g->glm53) { - ok = ds4_gpu_hc_expand_tensor(g->hc_after_attn, - g->attn_out, - g->hc_cur, - g->hc_post, - g->hc_comb, - DS4_N_EMBD, - DS4_N_HC) != 0; + /* Skip only the expand when kda_output already folded it in; the + * FFN-side mHC producer below is in this same block and must still + * run. */ + if (!attn_hc_expanded) { + ok = ds4_gpu_hc_expand_tensor(g->hc_after_attn, + g->attn_out, + g->hc_cur, + g->hc_post, + g->hc_comb, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok && (glm_decode_repeat_mask() & DS4_GLM_REPEAT_HC_EXPAND)) { + ok = ds4_gpu_hc_expand_tensor(g->hc_after_attn, g->attn_out, + g->hc_cur, g->hc_post, g->hc_comb, + DS4_N_EMBD, DS4_N_HC) != 0; + } + } + if (ok && (decode_ablate & DS4_GLM_ABLATE_HC)) { /* ablate */ } else if (ok) ok = glm53_graph_hc_pre(g, model, l->hc_ffn_fn, @@ -52420,6 +53134,8 @@ static bool glm_graph_forward_token( g->ffn_sum, g->attn_out, true, + NULL, + NULL, decode_stage_profile, decode_stage_profile ? &decode_stage_t0 : NULL); } @@ -52489,7 +53205,12 @@ static bool glm_graph_forward_token( } if (ok) ok = glm_graph_begin_commands_if_needed(); } - ok = glm_graph_encode_output_head(g, model, weights); + if (!(glm_decode_ablate_mask() & DS4_GLM_ABLATE_HEAD)) { + ok = glm_graph_encode_output_head(g, model, weights); + if (ok && (glm_decode_repeat_mask() & DS4_GLM_REPEAT_HEAD)) { + ok = glm_graph_encode_output_head(g, model, weights); + } + } if (g->ssd_streaming) { if (ok) ok = glm_graph_end_commands_if_active(); else (void)ds4_gpu_synchronize(); @@ -52524,7 +53245,12 @@ static bool glm_graph_forward_token( ok = glm_graph_stream_map_output(g, model, weights); } if (ok) ok = glm_graph_begin_commands_if_needed(); - if (ok) ok = glm_graph_encode_output_head(g, model, weights); + if (ok && !(glm_decode_ablate_mask() & DS4_GLM_ABLATE_HEAD)) { + ok = glm_graph_encode_output_head(g, model, weights); + if (ok && (glm_decode_repeat_mask() & DS4_GLM_REPEAT_HEAD)) { + ok = glm_graph_encode_output_head(g, model, weights); + } + } if (ok) ok = glm_graph_end_commands_if_active(); else (void)ds4_gpu_synchronize(); if (decode_output_profile) { diff --git a/ds4_gpu.h b/ds4_gpu.h index 5b866dde2..3ccf5f43d 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -252,8 +252,22 @@ enum { DS4_GPU_TEST_MXFP4_DOWN_HALF_LUT = 1u << 4, DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4 = 1u << 5, DS4_GPU_TEST_HC_RMS_SCALE_PROJ = 1u << 6, + /* Exercise GLM prefill kernels on synthetic shapes/devices, while keeping + * TP and streaming exclusions. The second flag forces the last KDA block + * to finish before block 0 to test incoming-state ownership. */ + DS4_GPU_TEST_GLM53_PREFILL = 1u << 7, + DS4_GPU_TEST_GLM53_KDA_LAST_BLOCK_FIRST = 1u << 8, }; void ds4_gpu_test_set_flags(uint32_t flags); +enum { + DS4_GPU_GLM53_PREFILL_QK_LOW = 1u << 0, + DS4_GPU_GLM53_PREFILL_INDEXED_ATTN = 1u << 1, + DS4_GPU_GLM53_PREFILL_MOE_TAIL_CULL = 1u << 2, + DS4_GPU_GLM53_PREFILL_KDA_PREPARE = 1u << 3, + DS4_GPU_GLM53_PREFILL_KDA_RECURRENCE = 1u << 4, +}; +/* Returns and clears dispatch coverage recorded only in GLM prefill test mode. */ +uint32_t ds4_gpu_test_glm53_prefill_take_dispatches(void); void ds4_gpu_release_zero_prefix_prefill_mask_cache(void); #else static inline int ds4_gpu_device_is_pre_m5_apple_silicon(void) { return 0; } @@ -1641,6 +1655,47 @@ int ds4_gpu_glm_attention_indexed_decode_typed_tensor( float beta_fast, float beta_slow); +int ds4_gpu_glm_attention_indexed_decode_exact_typed_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *scores, + ds4_gpu_tensor *lora, + ds4_gpu_tensor *denom, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim); + +int ds4_gpu_glm_attention_indexed_decode_exact_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *scores, + ds4_gpu_tensor *lora, + ds4_gpu_tensor *denom, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim); + int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( ds4_gpu_tensor *heads, ds4_gpu_tensor *partial_lora, @@ -2927,6 +2982,27 @@ int ds4_gpu_hc_expand_add_rms_norm_mix_split_norm_f16_tensor( float hc_eps, float norm_eps); +int ds4_gpu_hc_rms_norm_mix_split_norm_bf16_tensor( + ds4_gpu_tensor *mix, + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t mix_weight_offset, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n, + uint32_t mix_dim, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps, + float hc_eps, + float norm_eps); + #endif int ds4_gpu_output_hc_weights_tensor( ds4_gpu_tensor *out, @@ -3097,6 +3173,48 @@ int ds4_gpu_glm53_matmul_bf16_qkv( uint32_t out_dim, const ds4_gpu_tensor *x); +int ds4_gpu_glm53_matmul_bf16_pair( + ds4_gpu_tensor *out_a, + ds4_gpu_tensor *out_b, + const void *model_map, + uint64_t model_size, + uint64_t weight_a_offset, + uint64_t weight_b_offset, + uint32_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x_a, + const ds4_gpu_tensor *x_b); + +uint64_t ds4_gpu_encoder_count(void); + +int ds4_gpu_glm53_matmul_bf16_trio( + ds4_gpu_tensor *out_a, + ds4_gpu_tensor *out_b, + ds4_gpu_tensor *out_c, + const void *model_map, + uint64_t model_size, + uint64_t weight_a_offset, + uint64_t weight_b_offset, + uint64_t weight_c_offset, + uint32_t in_dim, + uint32_t out_dim_ab, + uint32_t out_dim_c, + const ds4_gpu_tensor *x); + +int ds4_gpu_glm53_matmul_bf16_hc_expand4( + ds4_gpu_tensor *out, + ds4_gpu_tensor *hc_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_hc); + #ifndef DS4_GLM53_VISION_TYPES_DEFINED #define DS4_GLM53_VISION_TYPES_DEFINED #define DS4_GLM53_VISION_LAYERS 24u diff --git a/ds4_metal.m b/ds4_metal.m index 47bea0821..abb81bcf7 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -417,9 +417,23 @@ static void ds4_gpu_timeline_attach(id cb) { static id g_dsv4_hc_expand_producer_pre_norm_pipeline; static NSMutableDictionary> *g_dsv4_hc_barrier_cache; static NSMutableDictionary *g_dsv4_hc_barrier_gen; +static id g_dsv4_hc_producer_pre_norm_bf16_pipeline; static id g_hc_weighted_sum_pipeline; static id g_output_hc_weights4_pipeline; static uint32_t g_test_flags; +static uint32_t g_test_glm53_prefill_dispatches; + +uint32_t ds4_gpu_test_glm53_prefill_take_dispatches(void) { + const uint32_t result = g_test_glm53_prefill_dispatches; + g_test_glm53_prefill_dispatches = 0; + return result; +} + +static void ds4_gpu_note_glm53_prefill_dispatch(uint32_t feature) { + if (g_test_flags & DS4_GPU_TEST_GLM53_PREFILL) { + g_test_glm53_prefill_dispatches |= feature; + } +} static id g_hc_expand_pipeline; static id g_unary_sigmoid_pipeline; static id g_unary_silu_pipeline; @@ -551,6 +565,9 @@ static void ds4_gpu_timeline_attach(id cb) { static id g_glm_qk_lowrank_glm52_sg_pipeline; static id g_glm_qk_lowrank_batch_pipeline; static id g_glm_qk_lowrank_batch_glm52_t4_pipeline; +static id g_glm_qk_lowrank_batch_t4_pipeline; +static id g_glm_qk_lowrank_batch_t8_pipeline; +static id g_glm_qk_lowrank_batch_t16_pipeline; static id g_glm_value_project_q8_0_pipeline; static id g_glm_value_project_q8_0_batch_heads_pipeline; static id g_glm_value_project_q8_0_batch_heads_mma_pipeline; @@ -559,6 +576,10 @@ static void ds4_gpu_timeline_attach(id cb) { static id g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline; static id g_glm_attention_indexed_decode_split_group8_reduce_pipeline; static id g_glm_attention_indexed_decode_split_group8_reduce16_pipeline; +static id g_glm_attention_indexed_decode_exact_scores_pipeline; +static id g_glm_attention_indexed_decode_exact_weights_pipeline; +static id g_glm_attention_indexed_decode_exact_lora_pipeline; +static id g_glm_attention_indexed_decode_exact_value_pipeline; static id g_glm_attention_indexed_batch_pipeline; static id g_glm_attention_indexed_batch_group2_pipeline; static id g_glm_attention_indexed_batch_q2_group4_pipeline; @@ -566,6 +587,7 @@ static void ds4_gpu_timeline_attach(id cb) { static id g_glm_attention_indexed_batch_lora_group8_vec_pipeline; static id g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline; static id g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline; +static id g_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads_pipeline; static id g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline; static id g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline; static id g_glm_q4_k_pair_swiglu_f32_pipeline; @@ -682,6 +704,8 @@ static void ds4_gpu_timeline_attach(id cb) { static id g_f16_round_scratch_buffer; static id g_raw_store_round_buffer; static id g_moe_gate_scratch_buffer; +static id g_kda_conv_halo_buffer; +static NSUInteger g_kda_conv_halo_capacity; static id g_moe_down_scratch_buffer; static id g_moe_id_map_buffer; static id g_moe_q4_gate_slots_buffer; @@ -1288,7 +1312,34 @@ static NSUInteger ds4_gpu_tensor_offset(const ds4_gpu_tensor *tensor) { return cb; } +/* Encoder acquisitions, as a proxy for dispatch count. Almost every + * primitive here acquires one encoder per dispatch, so the delta between two + * runs of differing decode length divided by the token difference is + * dispatches per token. It counts acquisitions, not encoder objects: inside a + * batch the same encoder is handed back for every dispatch. Multiplying the + * count by one measured launch cost gives an estimate of launch overhead, not + * a floor -- the per-launch cost was measured on one fusion and need not + * transfer to every kernel and command-buffer arrangement. Read with + * ds4_gpu_encoder_count(). */ +static uint64_t g_encoder_count; + +uint64_t ds4_gpu_encoder_count(void) { return g_encoder_count; } + +static void ds4_gpu_encoder_count_print(void) { + fprintf(stderr, "ds4: metal compute encoder acquisitions (~dispatches): %llu\n", + (unsigned long long)g_encoder_count); +} + +static void ds4_gpu_encoder_count_arm(void) { + static int armed = 0; + if (armed) return; + armed = 1; + if (getenv("DS4_METAL_ENCODER_COUNT")) atexit(ds4_gpu_encoder_count_print); +} + static id ds4_gpu_compute_encoder(id cb) { + g_encoder_count++; + ds4_gpu_encoder_count_arm(); if (g_batch_cb && cb == g_batch_cb) { g_batch_has_work = YES; if (g_timeline_enabled && g_timeline_batch) { @@ -6642,6 +6693,19 @@ static int ds4_gpu_encode_rope_tail_inplace( uint32_t value_type; } ds4_gpu_glm_attention_indexed_decode_args; +typedef struct { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t value_dim; + uint32_t value_row_bytes; + uint32_t value_type; + uint32_t stage_rows; + uint32_t heads_per_group; + float scale; +} ds4_gpu_glm_attention_indexed_decode_exact_args; + typedef struct { uint32_t n_selected; uint32_t cache_cap; @@ -8940,6 +9004,12 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch"); g_glm_qk_lowrank_batch_glm52_t4_pipeline = ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_glm52_t4"); + g_glm_qk_lowrank_batch_t4_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_t4"); + g_glm_qk_lowrank_batch_t8_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_t8"); + g_glm_qk_lowrank_batch_t16_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_t16"); g_glm_value_project_q8_0_pipeline = ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0"); g_glm_value_project_q8_0_batch_heads_pipeline = @@ -8956,6 +9026,14 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce"); g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce16"); + g_glm_attention_indexed_decode_exact_scores_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_exact_scores"); + g_glm_attention_indexed_decode_exact_weights_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_exact_weights"); + g_glm_attention_indexed_decode_exact_lora_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_exact_lora"); + g_glm_attention_indexed_decode_exact_value_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_exact_value"); g_glm_attention_indexed_batch_pipeline = ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch"); g_glm_attention_indexed_batch_group2_pipeline = @@ -8970,6 +9048,8 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid"); g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); + g_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads"); g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal"); g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = @@ -9053,6 +9133,9 @@ int ds4_gpu_init(void) { !g_glm_qk_lowrank_glm52_sg_pipeline || !g_glm_qk_lowrank_batch_pipeline || !g_glm_qk_lowrank_batch_glm52_t4_pipeline || + !g_glm_qk_lowrank_batch_t4_pipeline || + !g_glm_qk_lowrank_batch_t8_pipeline || + !g_glm_qk_lowrank_batch_t16_pipeline || !g_glm_value_project_q8_0_pipeline || !g_glm_value_project_q8_0_batch_heads_pipeline || !g_glm_value_project_q8_0_batch_heads_mma_pipeline || @@ -9061,6 +9144,10 @@ int ds4_gpu_init(void) { !g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline || !g_glm_attention_indexed_decode_split_group8_reduce_pipeline || !g_glm_attention_indexed_decode_split_group8_reduce16_pipeline || + !g_glm_attention_indexed_decode_exact_scores_pipeline || + !g_glm_attention_indexed_decode_exact_weights_pipeline || + !g_glm_attention_indexed_decode_exact_lora_pipeline || + !g_glm_attention_indexed_decode_exact_value_pipeline || !g_glm_attention_indexed_batch_pipeline || !g_glm_attention_indexed_batch_group2_pipeline || !g_glm_attention_indexed_batch_q2_group4_pipeline || @@ -9068,6 +9155,7 @@ int ds4_gpu_init(void) { !g_glm_attention_indexed_batch_lora_group8_vec_pipeline || !g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline || !g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline || + !g_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads_pipeline || !g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline || !g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline || !g_glm_q4_k_pair_swiglu_f32_pipeline || @@ -11548,6 +11636,7 @@ void ds4_gpu_cleanup(void) { g_hc_split_weighted_sum_pipeline = nil; g_hc_split_weighted_sum_norm_pipeline = nil; g_dsv4_hc_producer_pre_norm_pipeline = nil; + g_dsv4_hc_producer_pre_norm_bf16_pipeline = nil; g_hc_weighted_sum_pipeline = nil; g_output_hc_weights4_pipeline = nil; g_hc_expand_pipeline = nil; @@ -11672,6 +11761,9 @@ void ds4_gpu_cleanup(void) { g_glm_qk_lowrank_glm52_sg_pipeline = nil; g_glm_qk_lowrank_batch_pipeline = nil; g_glm_qk_lowrank_batch_glm52_t4_pipeline = nil; + g_glm_qk_lowrank_batch_t4_pipeline = nil; + g_glm_qk_lowrank_batch_t8_pipeline = nil; + g_glm_qk_lowrank_batch_t16_pipeline = nil; g_glm_value_project_q8_0_pipeline = nil; g_glm_value_project_q8_0_batch_heads_pipeline = nil; g_glm_value_project_q8_0_batch_heads_mma_pipeline = nil; @@ -11680,6 +11772,10 @@ void ds4_gpu_cleanup(void) { g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = nil; g_glm_attention_indexed_decode_split_group8_reduce_pipeline = nil; g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = nil; + g_glm_attention_indexed_decode_exact_scores_pipeline = nil; + g_glm_attention_indexed_decode_exact_weights_pipeline = nil; + g_glm_attention_indexed_decode_exact_lora_pipeline = nil; + g_glm_attention_indexed_decode_exact_value_pipeline = nil; g_glm_attention_indexed_batch_pipeline = nil; g_glm_attention_indexed_batch_group2_pipeline = nil; g_glm_attention_indexed_batch_q2_group4_pipeline = nil; @@ -11687,6 +11783,7 @@ void ds4_gpu_cleanup(void) { g_glm_attention_indexed_batch_lora_group8_vec_pipeline = nil; g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = nil; g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = nil; + g_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads_pipeline = nil; g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = nil; g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = nil; g_glm_q4_k_pair_swiglu_f32_pipeline = nil; @@ -11735,6 +11832,8 @@ void ds4_gpu_cleanup(void) { g_stream_expert_validate_status_buffer = nil; g_f16_round_scratch_buffer = nil; g_raw_store_round_buffer = nil; + g_kda_conv_halo_buffer = nil; + g_kda_conv_halo_capacity = 0; g_moe_gate_scratch_buffer = nil; g_moe_down_scratch_buffer = nil; g_moe_id_map_buffer = nil; @@ -36011,6 +36110,56 @@ int ds4_gpu_glm_qk_lowrank_q8_0_tensor( qk_dim); } +static bool ds4_gpu_glm53_prefill_tuning_available(void) { + /* Defaults have been measured and checked for exactness on M3 Ultra only. + * Test mode can exercise the same kernels on smaller fixtures; ownership + * exclusions still apply so it cannot silently turn on TP or streaming. */ + return !g_ssd_streaming_mode && g_tp_split_world == 1 && + ((g_test_flags & DS4_GPU_TEST_GLM53_PREFILL) != 0u || + [g_device.name isEqualToString:@"Apple M3 Ultra"]); +} + +/* + * Heads one simdgroup carries in the GLM 5.3 Flash indexed prefill attention + * kernel; 1 is main's kernel. DS4_METAL_DISABLE_GLM53_PREFILL_INDEXED_ATTN + * restores it for an A/B run, as does the branch-wide + * DS4_METAL_DISABLE_GLM53_FLASH_TUNING. + * DS4_METAL_GLM53_PREFILL_INDEXED_ATTN_HEADS_PER_SG forces one of the + * instantiated widths so the sweep can be repeated. Read per call so a test + * can flip it between dispatches. + */ +static uint32_t ds4_gpu_glm53_prefill_indexed_attn_heads_per_sg(void) { + if (getenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") != NULL || + getenv("DS4_METAL_DISABLE_GLM53_PREFILL_INDEXED_ATTN") != NULL) { + return 1u; + } + const char *env = + getenv("DS4_METAL_GLM53_PREFILL_INDEXED_ATTN_HEADS_PER_SG"); + const int forced = (env && env[0]) ? atoi(env) : 0; + if (forced == 1 || forced == 2) return (uint32_t)forced; + return 2u; +} + +/* + * Tokens per threadgroup for the GLM 5.3 Flash prefill qk-low kernel; 0 keeps + * the per-token reference kernel. DS4_METAL_DISABLE_GLM53_PREFILL_QK_LOW + * restores it for an A/B run, and the branch-wide + * DS4_METAL_DISABLE_GLM53_FLASH_TUNING does the same for every GLM 5.3 Flash + * switch at once. DS4_METAL_GLM53_PREFILL_QK_LOW_TILE forces one of the + * instantiated tiles so the sweep can be repeated; anything else keeps the + * measured default. Read per call so a test can flip it between dispatches. + */ +static uint32_t ds4_gpu_glm53_prefill_qk_low_token_tile(void) { + if (getenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") != NULL || + getenv("DS4_METAL_DISABLE_GLM53_PREFILL_QK_LOW") != NULL) { + return 0u; + } + const char *env = getenv("DS4_METAL_GLM53_PREFILL_QK_LOW_TILE"); + const int forced = (env && env[0]) ? atoi(env) : 0; + if (forced == 4 || forced == 8 || forced == 16) return (uint32_t)forced; + return 8u; +} + int ds4_gpu_glm_qk_lowrank_typed_batch_tensor( ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *q, @@ -36072,12 +36221,45 @@ int ds4_gpu_glm_qk_lowrank_typed_batch_tensor( qk_dim == 256u && row_bytes == 204u && weight_type == DS4_METAL_TENSOR_Q8_0; - id pipeline = - use_glm52_t4 ? - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_glm52_t4_pipeline, - "kernel_glm_qk_lowrank_q8_0_batch_glm52_t4") : - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_pipeline, - "kernel_glm_qk_lowrank_q8_0_batch"); + /* + * GLM 5.3 Flash prefill (qk_nope 256) has no token tile in the shape + * above, so it runs the per-token reference kernel and re-reads each + * head's 136 KB K_b slice once per token. kernel_..._batch_t + * keeps TT consecutive tokens in one threadgroup with the reference + * kernel's expression and block order, so it is bit-identical and + * divides the weight traffic by TT. Gated to the resident + * single-device GLM 5.3 Flash shape this was measured on. + */ + const uint32_t glm53_token_tile = + ds4_gpu_glm53_prefill_qk_low_token_tile(); + const int use_glm53_token_tile = + glm53_token_tile != 0u && + n_tokens >= glm53_token_tile && + n_head == 64u && + kv_lora_dim == 512u && + qk_nope == 256u && + qk_dim == 256u && + row_bytes == 272u && + weight_type == DS4_METAL_TENSOR_Q8_0 && + ds4_gpu_glm53_prefill_tuning_available(); + id pipeline = nil; + if (use_glm53_token_tile) { + ds4_gpu_note_glm53_prefill_dispatch(DS4_GPU_GLM53_PREFILL_QK_LOW); + pipeline = glm53_token_tile == 4u ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_t4_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch_t4") : + glm53_token_tile == 8u ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_t8_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch_t8") : + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_t16_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch_t16"); + } else if (use_glm52_t4) { + pipeline = ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_glm52_t4_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch_glm52_t4"); + } else { + pipeline = ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch"); + } if (!pipeline) return 0; int owned = 0; @@ -36104,7 +36286,13 @@ int ds4_gpu_glm_qk_lowrank_typed_batch_tensor( [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3]; - if (use_glm52_t4) { + if (use_glm53_token_tile) { + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, + ((NSUInteger)n_tokens + glm53_token_tile - 1u) / + glm53_token_tile, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } else if (use_glm52_t4) { [enc setThreadgroupMemoryLength:4u * 192u * sizeof(float) atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, ((NSUInteger)n_tokens + 3u) / 4u, @@ -36474,6 +36662,195 @@ int ds4_gpu_glm_attention_indexed_decode_tensor( beta_slow); } +/* The generic indexed decode attention in four phased dispatches that share + * each cache row across heads; bit-identical to it by construction (see the + * kernel comment in metal/dsv4_misc.metal). f16 compact cache, no RoPE tail, + * Q8_0 value rows, kv_lora_dim a multiple of 64 that stages 16 rows in 32 KiB + * of threadgroup memory, and up to 64 heads. */ +int ds4_gpu_glm_attention_indexed_decode_exact_typed_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *scores, + ds4_gpu_tensor *lora, + ds4_gpu_tensor *denom, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t stage_rows = 16u; + const uint32_t heads_per_group = 8u; + const uint32_t value_threads = 64u; + const uint64_t stage_bytes = + (uint64_t)stage_rows * ((kv_lora_dim / 4u) | 1u) * 4u * sizeof(uint16_t); + if (!heads || !scores || !lora || !denom || !qk_low || !kv_lora_cache || + !model_map || !selected || + n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || n_head * stage_rows > 1024u || + kv_lora_dim == 0 || (kv_lora_dim % 64u) != 0 || stage_bytes > 32768u || + qk_nope == 0 || qk_rope != 0 || value_dim == 0 || + !cache_f16 || value_weight_type != DS4_METAL_TENSOR_Q8_0) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id scoresbuf = ds4_gpu_tensor_buffer(scores); + id lorabuf = ds4_gpu_tensor_buffer(lora); + id denombuf = ds4_gpu_tensor_buffer(denom); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + uint64_t value_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { + fprintf(stderr, "ds4: Metal GLM exact indexed attention received unsupported value type\n"); + return 0; + } + const uint64_t value_weight_bytes = (uint64_t)n_head * value_dim * value_row_bytes; + if (!headsbuf || !scoresbuf || !lorabuf || !denombuf || !lowbuf || + !kvcachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(heads) < (uint64_t)n_head * value_dim * sizeof(float) || + ds4_gpu_tensor_bytes(scores) < (uint64_t)n_head * n_selected * sizeof(float) || + ds4_gpu_tensor_bytes(lora) < (uint64_t)n_head * kv_lora_dim * sizeof(float) || + ds4_gpu_tensor_bytes(denom) < (uint64_t)n_head * sizeof(float) || + ds4_gpu_tensor_bytes(qk_low) < (uint64_t)n_head * kv_lora_dim * sizeof(float) || + ds4_gpu_tensor_bytes(kv_lora_cache) < (uint64_t)cache_cap * kv_lora_dim * sizeof(uint16_t) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_selected * sizeof(uint32_t)) { + fprintf(stderr, "ds4: Metal GLM exact indexed attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: Metal GLM exact indexed attention value range is outside the mapped model\n"); + return 0; + } + uint64_t value_inner = 0; + id valuebuf = + ds4_gpu_wrap_model_range(model_map, model_size, value_weight_offset, + value_weight_bytes, &value_inner); + if (!valuebuf) return 0; + + id scores_pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_exact_scores_pipeline, + "kernel_glm_attention_indexed_decode_exact_scores"); + id weights_pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_exact_weights_pipeline, + "kernel_glm_attention_indexed_decode_exact_weights"); + id lora_pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_exact_lora_pipeline, + "kernel_glm_attention_indexed_decode_exact_lora"); + id value_pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_exact_value_pipeline, + "kernel_glm_attention_indexed_decode_exact_value"); + if (!scores_pipeline || !weights_pipeline || !lora_pipeline || !value_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_decode_exact_args args = { + .n_selected = n_selected, + .cache_cap = cache_cap, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .value_dim = value_dim, + .value_row_bytes = (uint32_t)value_row_bytes, + .value_type = value_weight_type, + .stage_rows = stage_rows, + .heads_per_group = heads_per_group, + /* The generic kernel's scale, computed the same way. */ + .scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)), + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:scores_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:1]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:2]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + [enc setThreadgroupMemoryLength:(NSUInteger)stage_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((n_selected + stage_rows - 1u) / stage_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake((NSUInteger)n_head * stage_rows, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:weights_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; + [enc setBuffer:denombuf offset:ds4_gpu_tensor_offset(denom) atIndex:2]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_head, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:lora_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:1]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:3]; + [enc setBuffer:denombuf offset:ds4_gpu_tensor_offset(denom) atIndex:4]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora) atIndex:5]; + /* Two staging buffers of 32 rows x 128 bytes plus 32 weights per head. */ + [enc setThreadgroupMemoryLength:2u * (32u * 128u + (NSUInteger)heads_per_group * 32u * sizeof(float)) + atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((n_head + heads_per_group - 1u) / heads_per_group, + kv_lora_dim / 64u, 1) + threadsPerThreadgroup:MTLSizeMake((NSUInteger)heads_per_group * 32u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:value_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora) atIndex:1]; + [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:2]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3]; + [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_head, (value_dim + value_threads - 1u) / value_threads, 1) + threadsPerThreadgroup:MTLSizeMake(value_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM exact indexed attention decode")) return 0; + } + return 1; +} + +int ds4_gpu_glm_attention_indexed_decode_exact_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *scores, + ds4_gpu_tensor *lora, + ds4_gpu_tensor *denom, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim) { + return ds4_gpu_glm_attention_indexed_decode_exact_typed_tensor( + heads, scores, lora, denom, qk_low, kv_lora_cache, model_map, model_size, + value_weight_offset, DS4_METAL_TENSOR_Q8_0, selected, n_selected, cache_cap, + cache_f16, n_head, kv_lora_dim, qk_nope, qk_rope, value_dim); +} + int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( ds4_gpu_tensor *heads, ds4_gpu_tensor *partial_lora, @@ -36514,7 +36891,12 @@ int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || n_head == 0 || (n_head % 8u) != 0 || kv_lora_dim != 512u || - qk_nope == 0 || qk_rope != 64u || + /* qk_rope == 0 is GLM 5.3, which has no RoPE tail. The kernel drives + * all its rope work from rope_vecs = qk_rope >> 2 and the scratch size + * below already drops the rope term at 0, so the case is supported -- + * as the freq_base/freq_scale checks below, which are already + * conditioned on qk_rope != 0, imply. */ + qk_nope == 0 || (qk_rope != 64u && qk_rope != 0u) || value_dim == 0 || qk_dim < qk_nope || block_rows == 0u || needed_blocks == 0u || n_blocks < needed_blocks || n_blocks > 64u || @@ -37117,8 +37499,32 @@ static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( cache_f16 && kv_lora_dim == 512u && (qk_rope == 0u || qk_rope == 64u); const bool full_head_groups = (n_head % 8u) == 0u; + uint32_t attn_head_base = 0; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &attn_head_base, &head_count); + /* + * Heads one simdgroup carries, so a threadgroup covers 8 * that many + * and a token needs 64 / (8 * that many) staging passes over its + * selected rows. Gated to the GLM 5.3 DSA shape this was measured on: + * no RoPE tail, 512 latent dimensions, resident, and a head range that + * divides evenly. + */ + const uint32_t heads_per_sg = + ds4_gpu_glm53_prefill_indexed_attn_heads_per_sg(); + const bool use_wide_head_groups = + use_vec_lora && selected_rows_valid && full_head_groups && + heads_per_sg > 1u && + qk_rope == 0u && + n_head == 64u && qk_nope == 256u && + (head_count % (8u * heads_per_sg)) == 0u && + ds4_gpu_glm53_prefill_tuning_available(); id pipeline = nil; - if (use_vec_lora && selected_rows_valid && full_head_groups) { + if (use_wide_head_groups) { + ds4_gpu_note_glm53_prefill_dispatch(DS4_GPU_GLM53_PREFILL_INDEXED_ATTN); + pipeline = ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads_pipeline, + "kernel_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads"); + } else if (use_vec_lora && selected_rows_valid && full_head_groups) { pipeline = ds4_gpu_hot_pipeline( g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline, "kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); @@ -37162,8 +37568,7 @@ static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( .beta_slow = beta_slow, .head_base = 0, }; - uint32_t head_count = n_head; - ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count); + args.head_base = attn_head_base; const NSUInteger scratch_bytes = use_vec_lora ? (16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u + 16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u) : @@ -37186,7 +37591,10 @@ static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:7]; } [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u, + const NSUInteger heads_per_group = + 8u * (NSUInteger)(use_wide_head_groups ? heads_per_sg : 1u); + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + heads_per_group - 1u) / + heads_per_group, (NSUInteger)n_tokens, 1) threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; @@ -38500,9 +38908,37 @@ static int ds4_gpu_glm_routed_moe_batch_grouped_tensor( id map_pipeline = ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); - id gate_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); - id up_pipeline = ds4_gpu_routed_mm_pipeline(up_type); - id down_pipeline = + /* + * Each expert's routed rows are matmul'd in 32-row tiles, and with + * 288 experts sharing 16384 rows the final tile of an expert has + * 16 or fewer rows about half the time. The CULL_TAIL_SIMDGROUPS + * instantiation keeps every thread in staging and at every barrier + * but lets the second row-half skip its MMA and store there; the + * skipped outputs are padding rows nothing reads, so the result is + * unchanged. Gated to the resident single-device Q4_K expert shape + * this was measured on. + */ + const bool use_q4_K_tail_cull = + gate_type == DS4_METAL_TENSOR_Q4_K && + up_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + ds4_gpu_glm53_prefill_tuning_available() && + ((n_total_expert == 288u && n_expert == 8u && + expert_in_dim == 4096u && expert_mid_dim == 2048u && out_dim == 4096u) || + (g_test_flags & DS4_GPU_TEST_GLM53_PREFILL) != 0u) && + getenv("DS4_METAL_DISABLE_GLM53_PREFILL_MOE_TAIL_CULL") == NULL && + getenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") == NULL; + if (use_q4_K_tail_cull) { + ds4_gpu_note_glm53_prefill_dispatch(DS4_GPU_GLM53_PREFILL_MOE_TAIL_CULL); + } + id gate_pipeline = use_q4_K_tail_cull ? + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32_tail_cull", false) : + ds4_gpu_routed_mm_pipeline(gate_type); + id up_pipeline = use_q4_K_tail_cull ? + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32_tail_cull", false) : + ds4_gpu_routed_mm_pipeline(up_type); + id down_pipeline = use_q4_K_tail_cull ? + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f16_tail_cull", false) : ds4_gpu_routed_mm_f16_rhs_pipeline(down_type); if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) { return 0; @@ -44426,7 +44862,11 @@ int ds4_gpu_hc_rms_norm_mix_f16_tensor( } -int ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( +/* The f16 and bf16 producers are the same dispatch with a different mix + * weight type; both are 16 bits per element, so every size, stride and buffer + * binding below is identical and only the pipeline differs. */ +static int ds4_gpu_hc_rms_norm_mix_split_norm_16bit_tensor( + bool mix_is_bf16, ds4_gpu_tensor *mix, ds4_gpu_tensor *out, ds4_gpu_tensor *norm_out, @@ -44501,12 +44941,18 @@ int ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( model_map, model_size, norm_weight_offset, out_bytes, &norm_inner); if (!mix_weight || !scalebuf || !basebuf || !norm_weight) return 0; - if (!g_dsv4_hc_producer_pre_norm_pipeline) { + if (mix_is_bf16) { + if (!g_dsv4_hc_producer_pre_norm_bf16_pipeline) { + g_dsv4_hc_producer_pre_norm_bf16_pipeline = ds4_gpu_get_pipeline( + "kernel_dsv4_hc_rms_norm_mix_bf16_cluster2_pre_norm"); + } + } else if (!g_dsv4_hc_producer_pre_norm_pipeline) { g_dsv4_hc_producer_pre_norm_pipeline = ds4_gpu_get_pipeline( "kernel_dsv4_hc_rms_norm_mix_f16_cluster2_pre_norm"); } id producer = - g_dsv4_hc_producer_pre_norm_pipeline; + mix_is_bf16 ? g_dsv4_hc_producer_pre_norm_bf16_pipeline + : g_dsv4_hc_producer_pre_norm_pipeline; if (!producer || producer.maxTotalThreadsPerThreadgroup < 512u) { return 0; } @@ -44833,6 +45279,37 @@ int ds4_gpu_hc_expand_add_rms_norm_mix_split_norm_f16_tensor( return 1; } +#define DS4_HC_PRODUCER_FORWARD_ARGS \ + mix, out, norm_out, split, residual_hc, model_map, model_size, \ + mix_weight_offset, scale_offset, base_offset, norm_weight_offset, \ + n, mix_dim, n_embd, n_hc, sinkhorn_iters, eps, hc_eps, norm_eps + +int ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( + ds4_gpu_tensor *mix, ds4_gpu_tensor *out, ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, const ds4_gpu_tensor *residual_hc, + const void *model_map, uint64_t model_size, + uint64_t mix_weight_offset, uint64_t scale_offset, + uint64_t base_offset, uint64_t norm_weight_offset, + uint32_t n, uint32_t mix_dim, uint32_t n_embd, uint32_t n_hc, + uint32_t sinkhorn_iters, float eps, float hc_eps, float norm_eps) { + return ds4_gpu_hc_rms_norm_mix_split_norm_16bit_tensor( + false, DS4_HC_PRODUCER_FORWARD_ARGS); +} + +int ds4_gpu_hc_rms_norm_mix_split_norm_bf16_tensor( + ds4_gpu_tensor *mix, ds4_gpu_tensor *out, ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, const ds4_gpu_tensor *residual_hc, + const void *model_map, uint64_t model_size, + uint64_t mix_weight_offset, uint64_t scale_offset, + uint64_t base_offset, uint64_t norm_weight_offset, + uint32_t n, uint32_t mix_dim, uint32_t n_embd, uint32_t n_hc, + uint32_t sinkhorn_iters, float eps, float hc_eps, float norm_eps) { + return ds4_gpu_hc_rms_norm_mix_split_norm_16bit_tensor( + true, DS4_HC_PRODUCER_FORWARD_ARGS); +} + +#undef DS4_HC_PRODUCER_FORWARD_ARGS + int ds4_gpu_output_hc_weights_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *pre, @@ -45773,6 +46250,13 @@ static int glm53_gpu_tensor_has( uint32_t n_rows; } glm53_gpu_bf16_matmul_args; +typedef struct { + uint32_t in_dim; + uint32_t out_dim_ab; + uint32_t out_dim_c; + uint32_t n_rows; +} glm53_gpu_bf16_trio_args; + int ds4_gpu_glm53_embedding_bf16( ds4_gpu_tensor *out, const void *model_map, @@ -45977,6 +46461,222 @@ int ds4_gpu_glm53_matmul_bf16_qkv( } } +int ds4_gpu_glm53_matmul_bf16_pair( + ds4_gpu_tensor *out_a, + ds4_gpu_tensor *out_b, + const void *model_map, + uint64_t model_size, + uint64_t weight_a_offset, + uint64_t weight_b_offset, + uint32_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x_a, + const ds4_gpu_tensor *x_b) { + if (!g_initialized && !ds4_gpu_init()) return 0; + /* Same device scope as the qkv variant this shares a row helper with. */ + if (!ds4_gpu_device_name_contains("M3 Ultra")) return 0; + uint64_t weights = 0; + if (in_dim == 0 || out_dim == 0 || + !glm53_gpu_mul_u64(in_dim, out_dim, &weights) || + !glm53_gpu_tensor_has(x_a, in_dim, sizeof(float)) || + !glm53_gpu_tensor_has(x_b, in_dim, sizeof(float)) || + !glm53_gpu_tensor_has(out_a, out_dim, sizeof(float)) || + !glm53_gpu_tensor_has(out_b, out_dim, sizeof(float))) { + return 0; + } + + @autoreleasepool { + const uint64_t weight_bytes = weights * sizeof(uint16_t); + uint64_t inner_a = 0, inner_b = 0; + id weight_a = glm53_gpu_weight_buffer( + model_map, model_size, weight_a_offset, weight_bytes, + &inner_a, "BF16 pair matrix A"); + id weight_b = glm53_gpu_weight_buffer( + model_map, model_size, weight_b_offset, weight_bytes, + &inner_b, "BF16 pair matrix B"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_mul_mv_bf16_f32_pair"); + if (!weight_a || !weight_b || !pipeline) return 0; + + const uint32_t nsg = glm53_gpu_bf16_mv_nsg(); + glm53_gpu_bf16_matmul_args args = { + .in_dim = in_dim, + .out_dim = out_dim, + .n_rows = 1u, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight_a offset:(NSUInteger)inner_a atIndex:1]; + [enc setBuffer:weight_b offset:(NSUInteger)inner_b atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(x_a) + offset:ds4_gpu_tensor_offset(x_a) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(x_b) + offset:ds4_gpu_tensor_offset(x_b) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_a) + offset:ds4_gpu_tensor_offset(out_a) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_b) + offset:ds4_gpu_tensor_offset(out_b) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake((out_dim + nsg - 1u) / nsg, + 1u, 2u) + threadsPerThreadgroup:MTLSizeMake(32u * nsg, 1u, 1u)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, + "GLM-5.3 BF16 pair matmul"); + } +} + +int ds4_gpu_glm53_matmul_bf16_trio( + ds4_gpu_tensor *out_a, + ds4_gpu_tensor *out_b, + ds4_gpu_tensor *out_c, + const void *model_map, + uint64_t model_size, + uint64_t weight_a_offset, + uint64_t weight_b_offset, + uint64_t weight_c_offset, + uint32_t in_dim, + uint32_t out_dim_ab, + uint32_t out_dim_c, + const ds4_gpu_tensor *x) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!ds4_gpu_device_name_contains("M3 Ultra")) return 0; + uint64_t w_ab = 0, w_c = 0; + if (in_dim == 0 || out_dim_ab == 0 || out_dim_c == 0 || + out_dim_c > out_dim_ab || + !glm53_gpu_mul_u64(in_dim, out_dim_ab, &w_ab) || + !glm53_gpu_mul_u64(in_dim, out_dim_c, &w_c) || + !glm53_gpu_tensor_has(x, in_dim, sizeof(float)) || + !glm53_gpu_tensor_has(out_a, out_dim_ab, sizeof(float)) || + !glm53_gpu_tensor_has(out_b, out_dim_ab, sizeof(float)) || + !glm53_gpu_tensor_has(out_c, out_dim_c, sizeof(float))) { + return 0; + } + + @autoreleasepool { + uint64_t inner_a = 0, inner_b = 0, inner_c = 0; + id weight_a = glm53_gpu_weight_buffer( + model_map, model_size, weight_a_offset, + w_ab * sizeof(uint16_t), &inner_a, "BF16 trio matrix A"); + id weight_b = glm53_gpu_weight_buffer( + model_map, model_size, weight_b_offset, + w_ab * sizeof(uint16_t), &inner_b, "BF16 trio matrix B"); + id weight_c = glm53_gpu_weight_buffer( + model_map, model_size, weight_c_offset, + w_c * sizeof(uint16_t), &inner_c, "BF16 trio matrix C"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_mul_mv_bf16_f32_trio"); + if (!weight_a || !weight_b || !weight_c || !pipeline) return 0; + + const uint32_t nsg = glm53_gpu_bf16_mv_nsg(); + glm53_gpu_bf16_trio_args args = { + .in_dim = in_dim, + .out_dim_ab = out_dim_ab, + .out_dim_c = out_dim_c, + .n_rows = 1u, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight_a offset:(NSUInteger)inner_a atIndex:1]; + [enc setBuffer:weight_b offset:(NSUInteger)inner_b atIndex:2]; + [enc setBuffer:weight_c offset:(NSUInteger)inner_c atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_a) + offset:ds4_gpu_tensor_offset(out_a) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_b) + offset:ds4_gpu_tensor_offset(out_b) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_c) + offset:ds4_gpu_tensor_offset(out_c) atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake((out_dim_ab + nsg - 1u) / nsg, + 1u, 3u) + threadsPerThreadgroup:MTLSizeMake(32u * nsg, 1u, 1u)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, + "GLM-5.3 BF16 trio matmul"); + } +} + +int ds4_gpu_glm53_matmul_bf16_hc_expand4( + ds4_gpu_tensor *out, + ds4_gpu_tensor *hc_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + /* Same device scope as the qkv and pair variants this shares a row helper + * with; every other device keeps the separate matvec and expand. */ + if (!ds4_gpu_device_name_contains("M3 Ultra")) return 0; + if (n_hc != 4u) return 0; + uint64_t weights = 0; + if (in_dim == 0 || out_dim == 0 || + !glm53_gpu_mul_u64(in_dim, out_dim, &weights) || + !glm53_gpu_tensor_has(x, in_dim, sizeof(float)) || + !glm53_gpu_tensor_has(out, out_dim, sizeof(float)) || + !glm53_gpu_tensor_has(hc_out, (uint64_t)n_hc * out_dim, sizeof(float)) || + !glm53_gpu_tensor_has(residual_hc, (uint64_t)n_hc * out_dim, sizeof(float)) || + !glm53_gpu_tensor_has(post, n_hc, sizeof(float)) || + !glm53_gpu_tensor_has(comb, (uint64_t)n_hc * n_hc, sizeof(float))) { + return 0; + } + + @autoreleasepool { + const uint64_t weight_bytes = weights * sizeof(uint16_t); + uint64_t inner = 0; + id weightbuf = glm53_gpu_weight_buffer( + model_map, model_size, weight_offset, weight_bytes, + &inner, "BF16 matrix with HC expand"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_mul_mv_bf16_f32_hc_expand4"); + if (!weightbuf || !pipeline) return 0; + + const uint32_t nsg = glm53_gpu_bf16_mv_nsg(); + glm53_gpu_bf16_matmul_args args = { + .in_dim = in_dim, + .out_dim = out_dim, + .n_rows = 1u, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(residual_hc) + offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(post) + offset:ds4_gpu_tensor_offset(post) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(comb) + offset:ds4_gpu_tensor_offset(comb) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(hc_out) + offset:ds4_gpu_tensor_offset(hc_out) atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake((out_dim + nsg - 1u) / nsg, 1u, 1u) + threadsPerThreadgroup:MTLSizeMake(32u * nsg, 1u, 1u)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer( + cb, owned, "GLM-5.3 BF16 matmul with HC expand"); + } +} + typedef struct { uint32_t width; uint32_t rows; @@ -46952,6 +47652,51 @@ int ds4_gpu_glm53_kda_decode( } } +typedef struct { + uint32_t n_heads; + uint32_t n_rows; + uint32_t block_rows; + uint32_t n_blocks; + uint32_t block_base; + float lower_bound; + float norm_eps; +} ds4_gpu_glm53_kda_blocked_args; + +/* + * Rows one threadgroup of the blocked KDA prepare kernel walks; 0 keeps the + * serial kernel, which runs one threadgroup per head and leaves an 80-core GPU + * about 95% idle. DS4_METAL_DISABLE_GLM53_PREFILL_KDA_PREPARE restores it for + * an A/B run, as does the branch-wide DS4_METAL_DISABLE_GLM53_FLASH_TUNING. + * DS4_METAL_GLM53_PREFILL_KDA_PREPARE_BLOCK forces one block size so the sweep + * can be repeated. Read per call so a test can flip it between dispatches. + */ +static uint32_t ds4_gpu_glm53_prefill_kda_prepare_block(void) { + if (getenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") != NULL || + getenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_PREPARE") != NULL) { + return 0u; + } + const char *env = getenv("DS4_METAL_GLM53_PREFILL_KDA_PREPARE_BLOCK"); + const int forced = (env && env[0]) ? atoi(env) : 0; + if (forced >= 4 && forced <= 4096 && (forced & (forced - 1)) == 0) { + return (uint32_t)forced; + } + return 32u; +} + +static double ds4_gpu_kda_now_sec(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9; +} + +static uint32_t ds4_gpu_glm53_prefill_kda_values_per_sg(void) { + if (getenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") != NULL || + getenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_RECURRENCE") != NULL) return 1u; + const char *env = getenv("DS4_METAL_GLM53_PREFILL_KDA_VALUES_PER_SG"); + const int forced = env ? atoi(env) : 0; + return forced == 1 || forced == 2 || forced == 4 ? (uint32_t)forced : 2u; +} + int ds4_gpu_glm53_kda_prefill( ds4_gpu_tensor *out, ds4_gpu_tensor *conv_state, @@ -47029,8 +47774,13 @@ int ds4_gpu_glm53_kda_prefill( &norm_inner, "KDA output norm"); id prep_pipeline = ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_prepare"); + const uint32_t kda_values = n_heads == 64u && n_tokens >= 32u && + ds4_gpu_glm53_prefill_tuning_available() ? + ds4_gpu_glm53_prefill_kda_values_per_sg() : 1u; id recurrence_pipeline = - ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_recurrence"); + ds4_gpu_get_pipeline(kda_values == 4u ? "kernel_glm53_kda_prefill_recurrence_v4" : + kda_values == 2u ? "kernel_glm53_kda_prefill_recurrence_v2" : + "kernel_glm53_kda_prefill_recurrence"); id output_pipeline = ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_output"); if (!qw || !kw || !vw || !a_log || !dt_bias || !output_norm || @@ -47038,38 +47788,146 @@ int ds4_gpu_glm53_kda_prefill( return 0; } + /* + * Blocked prepare: one threadgroup per (block of rows, head) with the + * convolution history in registers. It needs at least two blocks to + * be worth the halo pass, and a block start of 3 rows or more so the + * halo rows exist. + */ + const uint32_t kda_block_rows = ds4_gpu_glm53_prefill_kda_prepare_block(); + const uint32_t kda_n_blocks = kda_block_rows != 0u ? + 1u + (n_tokens - 1u) / kda_block_rows : 0u; + id halo_pipeline = nil; + id blocked_pipeline = nil; + NSUInteger kda_halo_bytes = 0; + int use_blocked_prepare = + kda_block_rows >= 4u && kda_n_blocks >= 2u && + n_heads == 64u && ds4_gpu_glm53_prefill_tuning_available(); + if (use_blocked_prepare) { + halo_pipeline = + ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_conv_halo"); + blocked_pipeline = + ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_prepare_blocked"); + kda_halo_bytes = (NSUInteger)kda_n_blocks * 3u * + (NSUInteger)projection * 3u * sizeof(float); + use_blocked_prepare = halo_pipeline != nil && blocked_pipeline != nil && + ds4_gpu_ensure_scratch_buffer(&g_kda_conv_halo_buffer, + &g_kda_conv_halo_capacity, + kda_halo_bytes, + "KDA conv halo") != 0; + } + ds4_gpu_glm53_kda_blocked_args blocked_args = { + .n_heads = n_heads, + .n_rows = n_tokens, + .block_rows = kda_block_rows, + .n_blocks = kda_n_blocks, + .block_base = 0, + .lower_bound = gate_lower_bound, + .norm_eps = norm_eps, + }; + glm53_gpu_kda_args args = { .n_heads = n_heads, .n_rows = n_tokens, .lower_bound = gate_lower_bound, .norm_eps = norm_eps, }; + const int kda_profile = getenv("DS4_METAL_PROFILE_KDA_PREFILL") != NULL; + double kda_t0 = 0.0; + if (kda_profile) { + if (g_batch_cb && + (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0)) return 0; + kda_t0 = ds4_gpu_kda_now_sec(); + } int owned = 0; id cb = ds4_gpu_command_buffer(&owned); if (!cb) return 0; id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:prep_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:ds4_gpu_tensor_buffer(q) - offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:ds4_gpu_tensor_buffer(k) - offset:ds4_gpu_tensor_offset(k) atIndex:2]; - [enc setBuffer:ds4_gpu_tensor_buffer(v) - offset:ds4_gpu_tensor_offset(v) atIndex:3]; - [enc setBuffer:ds4_gpu_tensor_buffer(raw_gate) - offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4]; - [enc setBuffer:qw offset:(NSUInteger)qw_inner atIndex:5]; - [enc setBuffer:kw offset:(NSUInteger)kw_inner atIndex:6]; - [enc setBuffer:vw offset:(NSUInteger)vw_inner atIndex:7]; - [enc setBuffer:a_log offset:(NSUInteger)a_inner atIndex:8]; - [enc setBuffer:dt_bias offset:(NSUInteger)dt_inner atIndex:9]; - [enc setBuffer:ds4_gpu_tensor_buffer(conv_state) - offset:ds4_gpu_tensor_offset(conv_state) atIndex:10]; - [enc setThreadgroupMemoryLength:264u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_heads, 1, 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + if (use_blocked_prepare) { + ds4_gpu_note_glm53_prefill_dispatch(DS4_GPU_GLM53_PREFILL_KDA_PREPARE); + [enc setComputePipelineState:halo_pipeline]; + [enc setBytes:&blocked_args length:sizeof(blocked_args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:3]; + [enc setBuffer:g_kda_conv_halo_buffer offset:0 atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(conv_state) + offset:ds4_gpu_tensor_offset(conv_state) atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(kda_n_blocks, 3, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + + [enc setComputePipelineState:blocked_pipeline]; + [enc setBytes:&blocked_args length:sizeof(blocked_args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(raw_gate) + offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4]; + [enc setBuffer:qw offset:(NSUInteger)qw_inner atIndex:5]; + [enc setBuffer:kw offset:(NSUInteger)kw_inner atIndex:6]; + [enc setBuffer:vw offset:(NSUInteger)vw_inner atIndex:7]; + [enc setBuffer:a_log offset:(NSUInteger)a_inner atIndex:8]; + [enc setBuffer:dt_bias offset:(NSUInteger)dt_inner atIndex:9]; + [enc setBuffer:ds4_gpu_tensor_buffer(conv_state) + offset:ds4_gpu_tensor_offset(conv_state) atIndex:10]; + [enc setBuffer:g_kda_conv_halo_buffer offset:0 atIndex:11]; + [enc setThreadgroupMemoryLength:264u * sizeof(float) atIndex:0]; + uint32_t prepare_blocks = kda_n_blocks; + if (g_test_flags & DS4_GPU_TEST_GLM53_KDA_LAST_BLOCK_FIRST) { + /* A legal adversarial scheduling order, made deterministic by + * separate serial dispatches of the unchanged prepare kernel. */ + blocked_args.block_base = kda_n_blocks - 1u; + [enc setBytes:&blocked_args length:sizeof(blocked_args) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, n_heads, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + blocked_args.block_base = 0; + [enc setBytes:&blocked_args length:sizeof(blocked_args) atIndex:0]; + prepare_blocks--; + } + [enc dispatchThreadgroups:MTLSizeMake(prepare_blocks, n_heads, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + } else { + [enc setComputePipelineState:prep_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(raw_gate) + offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4]; + [enc setBuffer:qw offset:(NSUInteger)qw_inner atIndex:5]; + [enc setBuffer:kw offset:(NSUInteger)kw_inner atIndex:6]; + [enc setBuffer:vw offset:(NSUInteger)vw_inner atIndex:7]; + [enc setBuffer:a_log offset:(NSUInteger)a_inner atIndex:8]; + [enc setBuffer:dt_bias offset:(NSUInteger)dt_inner atIndex:9]; + [enc setBuffer:ds4_gpu_tensor_buffer(conv_state) + offset:ds4_gpu_tensor_offset(conv_state) atIndex:10]; + [enc setThreadgroupMemoryLength:264u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_heads, 1, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + } + if (kda_profile) { + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "KDA prepare")) return 0; + if (!owned && ds4_gpu_end_commands() == 0) return 0; + const double t = ds4_gpu_kda_now_sec(); + fprintf(stderr, "ds4: kda stage prepare=%.3f ms\n", (t - kda_t0) * 1000.0); + kda_t0 = t; + if (!owned && ds4_gpu_begin_commands() == 0) return 0; + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + enc = ds4_gpu_compute_encoder(cb); + } [enc setComputePipelineState:recurrence_pipeline]; [enc setBytes:&args length:sizeof(args) atIndex:0]; [enc setBuffer:ds4_gpu_tensor_buffer(q) @@ -47086,9 +47944,24 @@ int ds4_gpu_glm53_kda_prefill( offset:ds4_gpu_tensor_offset(recurrent_state) atIndex:6]; [enc setBuffer:ds4_gpu_tensor_buffer(out) offset:ds4_gpu_tensor_offset(out) atIndex:7]; - [enc dispatchThreadgroups:MTLSizeMake(n_heads, 32, 1) + if (kda_values > 1u) { + ds4_gpu_note_glm53_prefill_dispatch(DS4_GPU_GLM53_PREFILL_KDA_RECURRENCE); + } + [enc dispatchThreadgroups:MTLSizeMake(n_heads, 32u / kda_values, 1) threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + if (kda_profile) { + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "KDA recurrence")) return 0; + if (!owned && ds4_gpu_end_commands() == 0) return 0; + const double t = ds4_gpu_kda_now_sec(); + fprintf(stderr, "ds4: kda stage recurrence=%.3f ms\n", (t - kda_t0) * 1000.0); + kda_t0 = t; + if (!owned && ds4_gpu_begin_commands() == 0) return 0; + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + enc = ds4_gpu_compute_encoder(cb); + } [enc setComputePipelineState:output_pipeline]; [enc setBytes:&args length:sizeof(args) atIndex:0]; [enc setBuffer:ds4_gpu_tensor_buffer(out) diff --git a/gguf-tools/.gitignore b/gguf-tools/.gitignore index d1f1b9e5c..c45693070 100644 --- a/gguf-tools/.gitignore +++ b/gguf-tools/.gitignore @@ -1,4 +1,5 @@ deepseek4-quantize +glm53-requant-bf16 gguf-requantize-dense quality-testing/score_official quality-testing/score_llama diff --git a/gguf-tools/Makefile b/gguf-tools/Makefile index 37d5d2cda..fe6707eea 100644 --- a/gguf-tools/Makefile +++ b/gguf-tools/Makefile @@ -47,11 +47,14 @@ CPPFLAGS ?= -D_GNU_SOURCE .PHONY: all clean quality-score quality-llama-score -all: deepseek4-quantize $(QUANTS_SHARED) +all: deepseek4-quantize glm53-requant-bf16 $(QUANTS_SHARED) deepseek4-quantize: deepseek4-quantize.c quants.c quants.h $(CC) $(CFLAGS) $(CPPFLAGS) -o $@ deepseek4-quantize.c quants.c -lm -pthread +glm53-requant-bf16: glm53_requant_bf16.c quants.c quants.h + $(CC) $(CFLAGS) $(CPPFLAGS) -o $@ glm53_requant_bf16.c quants.c -lm -pthread + $(QUANTS_SHARED): quants.c quants.h $(CC) $(CFLAGS) $(CPPFLAGS) $(SHARED_FLAGS) -fPIC -o $@ quants.c -lm -pthread @@ -63,5 +66,5 @@ quality-llama-score: $(CXX) $(LLAMA_CPP_CXXFLAGS) -o quality-testing/score_llama quality-testing/score_llama.cpp $(LLAMA_CPP_LDLIBS) clean: - rm -f deepseek4-quantize libds4quants.dylib libds4quants.so \ + rm -f deepseek4-quantize glm53-requant-bf16 libds4quants.dylib libds4quants.so \ quality-testing/score_official quality-testing/score_llama diff --git a/gguf-tools/glm53_requant_bf16.c b/gguf-tools/glm53_requant_bf16.c new file mode 100644 index 000000000..dfa9d4ef8 --- /dev/null +++ b/gguf-tools/glm53_requant_bf16.c @@ -0,0 +1,416 @@ +/* + * Requantize GLM 5.3 dense BF16 tensors to a smaller type, in an existing GGUF. + * + * The shipped GLM-5.3-Flash Q4_K artifact stores blk.N.kda_{q,k,v,output}, + * output.weight and token_embd.weight as BF16 while its experts are Q4_K. + * The KDA projections are dense -- every one is read on every decoded token -- + * so at 34 KDA layers they alone account for roughly 8.5 GiB of the per-token + * read traffic, more than all routed experts combined. output.weight is a + * further full matvec per token. + * + * glm53_quantize.py already assigns q8_0 to exactly these groups on its q4 + * artifact (role="linear_attention", "embedding" and "output"), so the result + * is a shape the loader and the generic matmul already accept. This tool + * produces it from an existing GGUF, without needing the source checkpoint. + * + * Everything other than the selected tensors is copied byte for byte, and the + * quantization goes through the same quants.c facade the other tools use, so + * the output differs from the input only in those tensors' type and payload. + */ + +#include "quants.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +enum { + GV_U8 = 0, GV_I8 = 1, GV_U16 = 2, GV_I16 = 3, GV_U32 = 4, GV_I32 = 5, + GV_F32 = 6, GV_BOOL = 7, GV_STR = 8, GV_ARR = 9, GV_U64 = 10, + GV_I64 = 11, GV_F64 = 12 +}; + +typedef struct { + const char *name; + uint64_t name_len; + uint32_t n_dims; + uint64_t dims[4]; + uint64_t ne; + uint32_t type; + uint64_t offset; + uint32_t new_type; + uint64_t new_offset; + uint64_t new_bytes; +} tinfo; + +static const uint8_t *g_base, *g_cur, *g_end; + +/* Set once the scratch output exists, so a die() anywhere below does not leave + * a half-written file sitting next to the real one. */ +static char *g_tmp_path; + +static void die(const char *msg) __attribute__((noreturn)); +static void die(const char *msg) { + fprintf(stderr, "glm53-requant-bf16: %s\n", msg); + if (g_tmp_path) unlink(g_tmp_path); + exit(1); +} + +/* Every size below is derived from attacker-controlled header fields, so each + * multiply and add is checked rather than allowed to wrap into a small, + * plausible-looking value that then passes a range test. */ +static uint64_t mul_or_die(uint64_t a, uint64_t b) { + if (a != 0 && b > UINT64_MAX / a) die("size overflow in the tensor table"); + return a * b; +} + +static uint64_t add_or_die(uint64_t a, uint64_t b) { + if (b > UINT64_MAX - a) die("size overflow in the tensor table"); + return a + b; +} + +static uint64_t pad_or_die(uint64_t x, uint64_t n) { + return add_or_die(x, (n - x % n) % n); +} + +static void need(size_t n) { + if ((size_t)(g_end - g_cur) < n) die("truncated gguf"); +} + +static uint32_t rd_u32(void) { need(4); uint32_t v; memcpy(&v, g_cur, 4); g_cur += 4; return v; } +static uint64_t rd_u64(void) { need(8); uint64_t v; memcpy(&v, g_cur, 8); g_cur += 8; return v; } + +static const char *rd_str(uint64_t *len) { + uint64_t n = rd_u64(); + need(n); + const char *s = (const char *)g_cur; + g_cur += n; + if (len) *len = n; + return s; +} + +static size_t scalar_size(uint32_t t) { + switch (t) { + case GV_U8: case GV_I8: case GV_BOOL: return 1; + case GV_U16: case GV_I16: return 2; + case GV_U32: case GV_I32: case GV_F32: return 4; + case GV_U64: case GV_I64: case GV_F64: return 8; + default: return 0; + } +} + +/* Skips a metadata value, returning its u32 content when it is a plain u32 + * (used only to pick general.alignment out of the stream). */ +static void skip_value(uint32_t t, int *is_u32, uint32_t *u32_out) { + if (is_u32) *is_u32 = 0; + if (t == GV_STR) { rd_str(NULL); return; } + if (t == GV_ARR) { + uint32_t et = rd_u32(); + uint64_t n = rd_u64(); + if (et == GV_STR) { + for (uint64_t i = 0; i < n; i++) rd_str(NULL); + } else { + size_t sz = scalar_size(et); + if (!sz) die("array of unsupported element type"); + const uint64_t span = mul_or_die((uint64_t)sz, n); + if (span > SIZE_MAX) die("metadata array is larger than this address space"); + need((size_t)span); + g_cur += (size_t)span; + } + return; + } + size_t sz = scalar_size(t); + if (!sz) die("unsupported metadata value type"); + if (t == GV_U32 && is_u32) { *is_u32 = 1; *u32_out = rd_u32(); return; } + need(sz); + g_cur += sz; +} + +enum { SEL_KDA = 1u << 0, SEL_HEAD = 1u << 1, SEL_EMBD = 1u << 2 }; + +static int name_is(const char *name, uint64_t len, const char *want) { + size_t wl = strlen(want); + return len == wl && memcmp(name, want, wl) == 0; +} + +static int name_ends(const char *name, uint64_t len, const char *suffix) { + size_t sl = strlen(suffix); + return len >= sl && memcmp(name + len - sl, suffix, sl) == 0; +} + +/* The groups glm53_quantize.py's q4 artifact assigns to Q8_0: the + * linear-attention projections (role="linear_attention") and the embedding and + * output tensors (role="embedding"/"output"). */ +static int selected(const char *name, uint64_t len, unsigned sel) { + if (sel & SEL_KDA) { + if (name_ends(name, len, ".kda_q.weight") || + name_ends(name, len, ".kda_k.weight") || + name_ends(name, len, ".kda_v.weight") || + name_ends(name, len, ".kda_output.weight")) return 1; + } + if ((sel & SEL_HEAD) && name_is(name, len, "output.weight")) return 1; + if ((sel & SEL_EMBD) && name_is(name, len, "token_embd.weight")) return 1; + return 0; +} + +static unsigned parse_selection(const char *spec) { + unsigned sel = 0; + const char *p = spec; + while (*p) { + const char *comma = strchr(p, ','); + size_t n = comma ? (size_t)(comma - p) : strlen(p); + if (n == 3 && !memcmp(p, "kda", 3)) sel |= SEL_KDA; + else if (n == 4 && !memcmp(p, "head", 4)) sel |= SEL_HEAD; + else if (n == 4 && !memcmp(p, "embd", 4)) sel |= SEL_EMBD; + else if (n == 3 && !memcmp(p, "all", 3)) sel |= SEL_KDA | SEL_HEAD | SEL_EMBD; + else die("--tensors takes a comma separated list of kda, head, embd, all"); + if (!comma) break; + p = comma + 1; + } + if (!sel) die("--tensors selected nothing"); + return sel; +} + +int main(int argc, char **argv) { + if (argc < 3) { + fprintf(stderr, + "usage: %s [--type q8_0|q4_K] [--tensors LIST]\n" + " Requantizes BF16 tensors that glm53_quantize.py's q4 artifact\n" + " assigns to q8_0. LIST is a comma separated selection of:\n" + " kda blk.N.kda_{q,k,v,output}.weight (default)\n" + " head output.weight\n" + " embd token_embd.weight\n" + " all all of the above\n", argv[0]); + return 2; + } + const char *in_path = argv[1], *out_path = argv[2]; + const char *want = "q8_0"; + unsigned sel = SEL_KDA; + for (int i = 3; i < argc; i++) { + if (!strcmp(argv[i], "--type") && i + 1 < argc) want = argv[++i]; + else if (!strcmp(argv[i], "--tensors") && i + 1 < argc) sel = parse_selection(argv[++i]); + else die("unrecognised argument; run with no arguments for usage"); + } + ds4q_type target; + if (!strcmp(want, "q8_0")) target = DS4Q_TYPE_Q8_0; + else if (!strcmp(want, "q4_K")) target = DS4Q_TYPE_Q4_K; + else die("target type must be q8_0 or q4_K"); + if (!ds4q_can_quantize(target)) die("quantizer cannot emit that type"); + + int fd = open(in_path, O_RDONLY); + if (fd < 0) die("cannot open input"); + struct stat st; + if (fstat(fd, &st) != 0) die("cannot stat input"); + /* The input stays mmapped for the whole run, so an output that resolves to + * the same file would pull the source out from under every read still to + * come. st_dev/st_ino catches the hard link and the symlink too, which a + * string compare of the two paths would not. */ + struct stat out_st; + if (stat(out_path, &out_st) == 0 && + out_st.st_dev == st.st_dev && out_st.st_ino == st.st_ino) { + die("output resolves to the input; write to a new path instead"); + } + const size_t in_size = (size_t)st.st_size; + if (in_size < 24) die("input is too small to be a gguf file"); + void *map = mmap(NULL, in_size, PROT_READ, MAP_PRIVATE, fd, 0); + if (map == MAP_FAILED) die("cannot mmap input"); + g_base = (const uint8_t *)map; + g_cur = g_base; + g_end = g_base + in_size; + + need(4); + if (memcmp(g_cur, "GGUF", 4) != 0) die("not a gguf file"); + g_cur += 4; + const uint32_t version = rd_u32(); + if (version != 3) fprintf(stderr, "glm53-requant-bf16: warning: gguf version %u\n", version); + const uint64_t n_tensors = rd_u64(); + const uint64_t n_kv = rd_u64(); + /* A tensor-info entry costs at least 8+4+8+4+8 bytes and a kv pair at + * least 8+4+1, so a count past these bounds is a corrupt header. Reject + * it here rather than at the calloc() it would otherwise size. */ + if (n_tensors > in_size / 32) die("implausible tensor count"); + if (n_kv > in_size / 13) die("implausible metadata count"); + + uint32_t alignment = 32; + for (uint64_t i = 0; i < n_kv; i++) { + uint64_t klen; const char *key = rd_str(&klen); + uint32_t vt = rd_u32(); + int is_u32 = 0; uint32_t v = 0; + skip_value(vt, &is_u32, &v); + if (is_u32 && klen == strlen("general.alignment") && + memcmp(key, "general.alignment", klen) == 0) { + alignment = v ? v : 32; + } + } + if (alignment == 0 || (alignment & (alignment - 1)) != 0 || alignment > 65536) { + die("general.alignment is not a power of two in range"); + } + const size_t kv_end = (size_t)(g_cur - g_base); + + tinfo *ts = calloc((size_t)n_tensors, sizeof(*ts)); + if (!ts) die("out of memory"); + for (uint64_t i = 0; i < n_tensors; i++) { + tinfo *t = &ts[i]; + t->name = rd_str(&t->name_len); + t->n_dims = rd_u32(); + if (t->n_dims == 0 || t->n_dims > 4) die("tensor with zero or more than 4 dimensions"); + t->ne = 1; + for (uint32_t d = 0; d < t->n_dims; d++) { + t->dims[d] = rd_u64(); + /* Both guard the ne/dims[0] divisions below and keep the product + * from wrapping into a small, plausible-looking byte count. */ + if (t->dims[d] == 0) die("tensor with a zero-length dimension"); + /* ds4q_row_size takes an int64_t, so a dimension past INT64_MAX + * would be reinterpreted as negative and silently return 0. */ + if (t->dims[d] > (uint64_t)INT64_MAX) die("tensor dimension out of range"); + t->ne = mul_or_die(t->ne, t->dims[d]); + } + t->type = rd_u32(); + t->offset = rd_u64(); + } + const size_t info_end = (size_t)(g_cur - g_base); + const size_t data_start = ds4q_pad(info_end, alignment); + if (data_start > in_size) die("tensor data section starts past the end of the input"); + + /* Plan: pick new types and lay the data section out again. */ + uint64_t cursor = 0, converted = 0, before = 0, after = 0; + for (uint64_t i = 0; i < n_tensors; i++) { + tinfo *t = &ts[i]; + int convert = (t->type == DS4Q_TYPE_BF16) && selected(t->name, t->name_len, sel); + /* row_size() returns 0 for a type this build does not know, for a row + * that is not a whole number of blocks, and for anything out of range. + * Treating that as "copy 0 bytes" would emit a file that still parses + * but has quietly lost the payload, so stop instead. */ + const size_t row_bytes = ds4q_row_size((ds4q_type)t->type, (int64_t)t->dims[0]); + if (row_bytes == 0) { + fprintf(stderr, "glm53-requant-bf16: %.*s is type %" PRIu32 ", which this build cannot size\n", + (int)t->name_len, t->name, t->type); + die("refusing to copy a tensor whose layout is unknown"); + } + const uint64_t nrows = t->ne / t->dims[0]; + const uint64_t old_bytes = mul_or_die((uint64_t)row_bytes, nrows); + if (t->offset > in_size - data_start || + old_bytes > in_size - data_start - t->offset) { + die("tensor data runs past the end of the input"); + } + if (convert && (t->dims[0] % (uint64_t)ds4q_block_size(target)) != 0) { + fprintf(stderr, "glm53-requant-bf16: %.*s row %" PRIu64 " not a multiple of the block size; leaving as is\n", + (int)t->name_len, t->name, t->dims[0]); + convert = 0; + } + t->new_type = convert ? (uint32_t)target : t->type; + t->new_bytes = convert + ? mul_or_die((uint64_t)ds4q_row_size(target, (int64_t)t->dims[0]), nrows) + : old_bytes; + cursor = pad_or_die(cursor, alignment); + t->new_offset = cursor; + cursor = add_or_die(cursor, t->new_bytes); + if (convert) { converted++; before += old_bytes; after += t->new_bytes; } + } + if (!converted) die("no matching BF16 tensors found -- nothing to do"); + fprintf(stderr, + "glm53-requant-bf16: %" PRIu64 " tensors -> %s, %.2f GiB -> %.2f GiB (saves %.2f GiB per full read)\n", + converted, ds4q_type_name(target), + before / 1073741824.0, after / 1073741824.0, (before - after) / 1073741824.0); + + /* Build the file beside its destination and rename it into place at the + * end: out_path then either still holds whatever it held before, or holds + * a complete result, and never a truncated one. + * + * The scratch name comes from mkstemp rather than the pid. A predictable + * name opened with fopen("wb") reintroduces exactly the bug the input/ + * output inode check above closes: if that path is a symlink or hard link + * to the input, the open truncates the mapped source. mkstemp picks an + * unpredictable name and opens O_CREAT|O_EXCL, which neither follows a + * symlink nor reuses an existing file. */ + const size_t tmp_len = strlen(out_path) + 8; + g_tmp_path = malloc(tmp_len); + if (!g_tmp_path) die("out of memory"); + snprintf(g_tmp_path, tmp_len, "%s.XXXXXX", out_path); + const int out_fd = mkstemp(g_tmp_path); + if (out_fd < 0) die("cannot create the scratch output"); + /* Belt and braces: confirm what we hold is a fresh regular file and is not + * the input, before a single byte is written. */ + struct stat tmp_st; + if (fstat(out_fd, &tmp_st) != 0) die("cannot stat the scratch output"); + if (!S_ISREG(tmp_st.st_mode) || + (tmp_st.st_dev == st.st_dev && tmp_st.st_ino == st.st_ino)) { + die("scratch output is not a fresh regular file"); + } + /* mkstemp creates 0600; a model file should follow the umask like any + * other output this tool used to produce. */ + const mode_t mask = umask(0); + (void)umask(mask); + (void)fchmod(out_fd, (mode_t)(0666 & ~mask)); + FILE *out = fdopen(out_fd, "wb"); + if (!out) die("cannot open output"); + /* Header and metadata are copied verbatim; tensor-info entries keep their + * width, so the data section still begins at the same offset. */ + if (fwrite(g_base, 1, kv_end, out) != kv_end) die("write failed"); + for (uint64_t i = 0; i < n_tensors; i++) { + tinfo *t = &ts[i]; + fwrite(&t->name_len, 8, 1, out); + fwrite(t->name, 1, t->name_len, out); + fwrite(&t->n_dims, 4, 1, out); + for (uint32_t d = 0; d < t->n_dims; d++) fwrite(&t->dims[d], 8, 1, out); + fwrite(&t->new_type, 4, 1, out); + if (fwrite(&t->new_offset, 8, 1, out) != 1) die("write failed"); + } + static const uint8_t zeros[4096] = {0}; + size_t here = (size_t)ftello(out); + if (here != info_end) die("tensor info section changed size unexpectedly"); + while (here < data_start) { + size_t n = data_start - here; + if (n > sizeof(zeros)) n = sizeof(zeros); + fwrite(zeros, 1, n, out); + here += n; + } + + ds4q_quantize_init(target); + const int64_t CHUNK = 256; /* rows per pass, keeps the f32 staging small */ + for (uint64_t i = 0; i < n_tensors; i++) { + tinfo *t = &ts[i]; + const size_t want_at = data_start + t->new_offset; + size_t at = (size_t)ftello(out); + while (at < want_at) { + size_t n = want_at - at; + if (n > sizeof(zeros)) n = sizeof(zeros); + fwrite(zeros, 1, n, out); + at += n; + } + const uint8_t *src = g_base + data_start + t->offset; + if (t->new_type == t->type) { + if (fwrite(src, 1, t->new_bytes, out) != t->new_bytes) die("write failed"); + continue; + } + const int64_t ncols = (int64_t)t->dims[0]; + const int64_t nrows = (int64_t)(t->ne / t->dims[0]); + float *f32 = malloc((size_t)ncols * CHUNK * sizeof(float)); + void *qbuf = malloc((size_t)ds4q_row_size(target, ncols) * CHUNK); + if (!f32 || !qbuf) die("out of memory"); + for (int64_t r = 0; r < nrows; r += CHUNK) { + const int64_t rows = (r + CHUNK <= nrows) ? CHUNK : (nrows - r); + const uint16_t *bf = (const uint16_t *)src + (size_t)r * ncols; + for (int64_t k = 0; k < rows * ncols; k++) f32[k] = ds4q_bf16_to_f32(bf[k]); + size_t wrote = ds4q_quantize_chunk(target, f32, qbuf, 0, rows, ncols, NULL); + if (fwrite(qbuf, 1, wrote, out) != wrote) die("write failed"); + } + free(f32); + free(qbuf); + fprintf(stderr, " %.*s -> %s\n", (int)t->name_len, t->name, ds4q_type_name(target)); + } + if (fclose(out) != 0) die("close failed"); + if (rename(g_tmp_path, out_path) != 0) die("cannot move the finished file into place"); + free(g_tmp_path); + g_tmp_path = NULL; + munmap(map, in_size); + close(fd); + fprintf(stderr, "glm53-requant-bf16: wrote %s\n", out_path); + return 0; +} diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal index c161a03a8..fb3a5a92c 100644 --- a/metal/dsv4_hc.metal +++ b/metal/dsv4_hc.metal @@ -1321,6 +1321,26 @@ kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2( } } + +/* Self-contained on purpose. glm53_bf16.metal has an identical helper, but + * depending on it would couple this file to that one across the concatenated + * library: pointing DS4_METAL_GLM53_BF16_SOURCE at an older revision of that + * file would then stop THIS file compiling, which defeats the per-file source + * overrides used for shader A/B runs. */ +static inline float4 ds4_hc_bf16x4_to_f32x4(ushort4 v) { + return float4(as_type((uint)v.x << 16), + as_type((uint)v.y << 16), + as_type((uint)v.z << 16), + as_type((uint)v.w << 16)); +} + +static inline float4 ds4_hc_mix_widen(half4 v) { return float4(v); } +static inline float4 ds4_hc_mix_widen(ushort4 v) { return ds4_hc_bf16x4_to_f32x4(v); } +/* The f16 and bf16 producers differ only in how the mix weights are + * widened; everything else -- the reduction trees, the cluster split, the + * collapse and the pre-norm -- is shared, so the body is a template and the + * kernels below are thin instantiations of it. */ +template static inline void ds4_hc_rms_norm_mix_cluster2_pre_norm_body( constant ds4_metal_args_hc_norm_mix & args, constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & split_args, @@ -1382,10 +1402,10 @@ static inline void ds4_hc_rms_norm_mix_cluster2_pre_norm_body( const int nb = args.n/NB; const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0; - device const half4 *ax4[NR0]; + device const W4 *ax4[NR0]; FOR_UNROLL (short row = 0; row < NR0; ++row) { - ax4[row] = (device const half4 *) - (weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half)); + ax4[row] = (device const W4 *) + (weight + (uint64_t)(r0 + row)*(uint64_t)n*(sizeof(W4)/4)); } float sumf_mv[NR0] = { 0.f }; @@ -1398,10 +1418,10 @@ static inline void ds4_hc_rms_norm_mix_cluster2_pre_norm_body( yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale; } FOR_UNROLL (short row = 0; row < NR0; ++row) { - device const half4 *xb4 = ax4[row] + (ib*NB + il*NF)/4; + device const W4 *xb4 = ax4[row] + (ib*NB + il*NF)/4; float sumq = 0.f; FOR_UNROLL (short i = 0; i < NF4; ++i) { - sumq += dot(float4(xb4[i]), yl4[i]); + sumq += dot(ds4_hc_mix_widen(xb4[i]), yl4[i]); } sumf_mv[row] += sumq; } @@ -1560,7 +1580,7 @@ kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2_pre_norm( uint3 tgpig [[threadgroup_position_in_grid]], ushort tiisg [[thread_index_in_simdgroup]], ushort sgitg [[simdgroup_index_in_threadgroup]]) { - ds4_hc_rms_norm_mix_cluster2_pre_norm_body(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg); + ds4_hc_rms_norm_mix_cluster2_pre_norm_body(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg); } /* Decode-time fusion of the HC post/expand that follows a TP combine with the @@ -1631,5 +1651,28 @@ kernel void kernel_dsv4_hc_expand4_rms_norm_mix_f16_cluster2_pre_norm( thread_scope_device); } threadgroup_barrier(mem_flags::mem_device_and_threadgroup); - ds4_hc_rms_norm_mix_cluster2_pre_norm_body(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg); + ds4_hc_rms_norm_mix_cluster2_pre_norm_body(args, split_args, x, weight, dst, hc_scale, hc_base, split, collapse_dst, norm_weight, norm_dst, completion, shmem, tgpig, tiisg, sgitg); +} + +kernel void kernel_dsv4_hc_rms_norm_mix_bf16_cluster2_pre_norm( + constant ds4_metal_args_hc_norm_mix & args, + constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & split_args, + device const char * x, + device const char * weight, + device char * dst, + device const float * hc_scale, + device const float * hc_base, + device char * split, + device char * collapse_dst, + device const char * norm_weight, + device char * norm_dst, + device atomic_uint * completion, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + ds4_hc_rms_norm_mix_cluster2_pre_norm_body( + args, split_args, x, weight, dst, hc_scale, hc_base, split, + collapse_dst, norm_weight, norm_dst, completion, shmem, + tgpig, tiisg, sgitg); } diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index a9b2397f9..bd8edb757 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -365,6 +365,19 @@ struct ds4_metal_args_glm_attention_indexed_decode_split { uint32_t value_type; }; +struct ds4_metal_args_glm_attention_indexed_decode_exact { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t value_dim; + uint32_t value_row_bytes; + uint32_t value_type; + uint32_t stage_rows; /* rows staged per threadgroup by the score kernel */ + uint32_t heads_per_group; /* heads per threadgroup in the lora kernel */ + float scale; +}; + struct ds4_metal_args_glm_attention_indexed_batch { uint32_t n_tokens; uint32_t n_selected; @@ -2607,6 +2620,123 @@ kernel void kernel_glm_qk_lowrank_q8_0_batch_glm52_t4( } } +/* GLM 5.3 prefill qk-low, one threadgroup per (head, TT tokens). + * + * kernel_glm_qk_lowrank_q8_0_batch gives every (head, token) pair its own + * threadgroup, so each of the 512 Q8_0 rows of a head's K_b slice is re-read + * once per token -- 17.8 GB of weight traffic per layer for 34 GFLOP -- and + * each thread runs a single 256-long dependent FMA chain. + * + * Here a thread owns the same two output rows for TT consecutive tokens: it + * walks its 272-byte row once per TT tokens and keeps TT independent + * accumulator chains. Which threadgroup computes which outputs changes; + * the arithmetic does not. Every output still accumulates + * `acc += d * (float)qs[qi] * x[base + qi]` over ascending blocks and + * ascending qi, written verbatim, so the result is bit-identical to the + * reference kernel above. + */ +template +static inline void glm_qk_lowrank_q8_0_batch_tokens_impl( + constant ds4_metal_args_glm_qk_lowrank_batch & args, + device const char *weight, + device const char *q, + device char *qk_low, + uint tid, + uint nth, + uint3 tgpig) { + constexpr uint kv_lora_dim = 512u; + constexpr uint qk_nope = 256u; + constexpr uint qk_dim = 256u; + constexpr uint row_bytes = 272u; + constexpr uint n_blocks = qk_nope >> 5; + constexpr uint NR = 2u; + + if (args.kv_lora_dim != kv_lora_dim || + args.qk_nope != qk_nope || + args.qk_dim != qk_dim || + args.row_bytes != row_bytes || + args.weight_type != DS4_METAL_GGUF_Q8_0 || + args.n_tokens == 0u || + nth * NR != kv_lora_dim) { + return; + } + + const uint head = tgpig.x + args.head_base; + const uint token0 = tgpig.y * TT; + if (head >= args.n_head || token0 >= args.n_tokens) return; + + const ulong q_token_stride = (ulong)args.n_head * qk_dim; + const ulong low_token_stride = (ulong)args.n_head * kv_lora_dim; + device const float *xbase = (device const float *)q; + + /* A partial tail tile clamps to the last real token instead of branching: + * the clamped columns are read and accumulated, then never stored. */ + ulong xoff[TT]; + FOR_UNROLL (uint t = 0; t < TT; t++) { + const uint token = min(token0 + t, args.n_tokens - 1u); + xoff[t] = token * q_token_stride + head * qk_dim; + } + + const uint j0 = tid; + const uint j1 = tid + nth; + device const char *row0 = + weight + ((uint64_t)head * kv_lora_dim + j0) * row_bytes; + device const char *row1 = + weight + ((uint64_t)head * kv_lora_dim + j1) * row_bytes; + + float acc0[TT]; + float acc1[TT]; + FOR_UNROLL (uint t = 0; t < TT; t++) { + acc0[t] = 0.0f; + acc1[t] = 0.0f; + } + + for (uint block = 0; block < n_blocks; block++) { + device const char *block0 = row0 + (uint64_t)block * 34u; + device const char *block1 = row1 + (uint64_t)block * 34u; + const float d0 = (float)(*((device const half *)block0)); + const float d1 = (float)(*((device const half *)block1)); + device const int8_t *qs0 = (device const int8_t *)(block0 + 2u); + device const int8_t *qs1 = (device const int8_t *)(block1 + 2u); + const uint base = block << 5; + for (uint qi = 0; qi < 32u; qi++) { + const uint col = base + qi; + FOR_UNROLL (uint t = 0; t < TT; t++) { + acc0[t] += d0 * (float)qs0[qi] * xbase[xoff[t] + col]; + acc1[t] += d1 * (float)qs1[qi] * xbase[xoff[t] + col]; + } + } + } + + device float *outbase = (device float *)qk_low; + FOR_UNROLL (uint t = 0; t < TT; t++) { + const uint token = token0 + t; + if (token < args.n_tokens) { + device float *out = + outbase + token * low_token_stride + head * kv_lora_dim; + out[j0] = acc0[t]; + out[j1] = acc1[t]; + } + } +} + +#define DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(TT) \ + kernel void kernel_glm_qk_lowrank_q8_0_batch_t##TT( \ + constant ds4_metal_args_glm_qk_lowrank_batch & args, \ + device const char *weight, \ + device const char *q, \ + device char *qk_low, \ + uint tid [[thread_index_in_threadgroup]], \ + ushort3 ntg_u [[threads_per_threadgroup]], \ + uint3 tgpig [[threadgroup_position_in_grid]]) { \ + glm_qk_lowrank_q8_0_batch_tokens_impl( \ + args, weight, q, qk_low, tid, ntg_u.x, tgpig); \ + } + +DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(4) +DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(8) +DS4_GLM_QK_LOWRANK_BATCH_TOKENS_KERNEL(16) + kernel void kernel_glm_value_project_q8_0( constant ds4_metal_args_glm_qk_lowrank & args, device const char *weight, @@ -2866,7 +2996,12 @@ kernel void kernel_glm_attention_indexed_decode_split_group8_partial_impl( if (args.n_selected == 0u || args.cache_f16 == 0u || args.kv_lora_dim != 512u || - args.qk_rope != 64u || + /* GLM 5.3 has no RoPE tail (n_rot = 0). Everything rope here is + * driven by rope_vecs = qk_rope >> 2, so at 0 the staging loop runs no + * iterations, rope_shared is never touched and the per-lane rope dot + * is skipped -- the kernel is already correct for that case and only + * this guard kept it out. */ + (args.qk_rope != 64u && args.qk_rope != 0u) || args.block_rows == 0u || block >= args.n_blocks) { return; @@ -3175,6 +3310,259 @@ kernel void kernel_glm_attention_indexed_decode_split_group8_reduce16( tid, ntg_u, tgpig); } +/* + * kernel_glm_attention_indexed_decode, in phases, with its arithmetic kept + * operation for operation. + * + * The generic kernel below runs one threadgroup per head and has every head + * walk every selected row twice, so 64 heads re-read each cache row 128 times + * and only 64 threadgroups exist to do it. The four kernels here compute the + * same thing in the same order -- one thread per (head, row) for the + * sequential 512-term score, the generic's 256-thread partition and + * 128/64/../1 tree for the softmax denominator, one thread per column pair + * walking rows 0..n-1 for the weighted sum, and the same quantised value row + * dot -- but stage each cache row once per threadgroup for every head and + * spread the work over hundreds of threadgroups. Every floating-point + * operation, operand and ordering is the generic kernel's, so the output is + * bit-identical to it (tests/test_glm53_kda asserts this at tolerance 0); + * the speed comes from row sharing and parallelism alone. Intermediates live + * in device buffers instead of threadgroup memory: scores[head][s], which the + * weights kernel turns into softmax weights in place, denom[head] and + * lora[head][kv_lora_dim]. f16 compact cache and no RoPE tail only, which is + * GLM 5.3. + * + * Score kernel. Grid: ceil(n_selected / stage_rows) threadgroups of + * n_head * stage_rows threads. Thread t scores row t % stage_rows for head + * t / stage_rows, so a simdgroup holds two heads over the same rows and the + * staged half4 it reads are shared lane pairs. Rows are staged at an odd + * half4 stride so consecutive rows fall in different banks. + */ +kernel void kernel_glm_attention_indexed_decode_exact_scores( + constant ds4_metal_args_glm_attention_indexed_decode_exact & args, + device const char *qk_low, + device const char *kv_lora_cache, + device const uint32_t *selected, + device float *scores, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint stage_rows = args.stage_rows; + const uint kv_vecs = args.kv_lora_dim >> 2; + const uint row_stride = kv_vecs | 1u; + const uint s0 = tgpig.x * stage_rows; + if (s0 >= args.n_selected || stage_rows == 0u) return; + const uint rows = min(stage_rows, args.n_selected - s0); + const uint nthreads = args.n_head * stage_rows; + for (uint off = tid; off < rows * kv_vecs; off += nthreads) { + const uint rr = off / kv_vecs; + const uint vv = off - rr * kv_vecs; + const uint row = selected[s0 + rr]; + kv_shared[rr * row_stride + vv] = row < args.cache_cap + ? ((device const half4 *)kv_lora_cache)[(uint64_t)row * kv_vecs + vv] + : half4(half(0.0f)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint r = tid % stage_rows; + const uint head = tid / stage_rows; + if (r >= rows || head >= args.n_head) return; + const uint s = s0 + r; + const uint row = selected[s]; + device const float *low = + (device const float *)(qk_low + (uint64_t)head * args.kv_lora_dim * sizeof(float)); + threadgroup const half4 *kvrow = kv_shared + r * row_stride; + float score = -INFINITY; + if (row < args.cache_cap) { + float dotv = 0.0f; + uint j = 0; + for (; j + 3u < args.kv_lora_dim; j += 4u) { + threadgroup const half4 *kv4 = kvrow + (j >> 2); + device const float4 *low4 = + (device const float4 *)(low + j); + const float4 kv = (float4)(*kv4); + const float4 qv = *low4; + dotv += qv.x * kv.x + qv.y * kv.y + + qv.z * kv.z + qv.w * kv.w; + } + if (j < args.kv_lora_dim) { + for (; j < args.kv_lora_dim; j++) { + const float kv = (float)((threadgroup const half *)kvrow)[j]; + dotv += low[j] * kv; + } + } + score = dotv * args.scale; + } + scores[(uint64_t)head * args.n_selected + s] = score; +} + +/* Weights kernel. Grid: n_head threadgroups of 256 threads -- the generic + * kernel's threadgroup -- so the per-thread row partition and the reduction + * tree produce its max and denominator. Scores become softmax weights in + * place. */ +kernel void kernel_glm_attention_indexed_decode_exact_weights( + constant ds4_metal_args_glm_attention_indexed_decode_exact & args, + device float *scores, + device float *denom, + threadgroup float *red [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + if (head >= args.n_head || args.n_selected == 0u) return; + const uint nth = ntg_u.x; + device float *sc = scores + (uint64_t)head * args.n_selected; + + float local_max = -INFINITY; + for (uint s = tid; s < args.n_selected; s += nth) { + local_max = max(local_max, sc[s]); + } + red[tid] = local_max; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = max(red[tid], red[tid + step]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_score = red[0]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float local_sum = 0.0f; + for (uint s = tid; s < args.n_selected; s += nth) { + const float w = exp(sc[s] - max_score); + sc[s] = w; + local_sum += w; + } + red[tid] = local_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0u) denom[head] = max(red[0], 1.0e-20f); +} + +/* Lora kernel. Grid: ceil(n_head / heads_per_group) x (kv_lora_dim / 64) + * threadgroups of heads_per_group * 32 threads. A simdgroup is one head over + * 64 consecutive columns, two per lane as in the generic kernel. The rows + * are walked in stages of 32: every thread fetches 16 bytes of the stage's + * 128-byte row slices and one weight, three stages ahead of use, and parks + * them in double-buffered threadgroup memory, so the scattered row reads are + * in flight while the fma chains run. Those chains are the generic kernel's, + * row after row in selection order; a row past cache_cap (or past the end of + * the last stage) contributes fma(0, kv[0], acc), which leaves acc unchanged + * bit for bit, where the generic kernel skips it. */ +#define DS4_GLM_EXACT_LORA_STAGE 32u +#define DS4_GLM_EXACT_LORA_AHEAD 3u +kernel void kernel_glm_attention_indexed_decode_exact_lora( + constant ds4_metal_args_glm_attention_indexed_decode_exact & args, + device const char *kv_lora_cache, + device const uint32_t *selected, + device const float *weights, + device const float *denom, + device float *lora, + threadgroup uchar *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + constexpr uint stage = DS4_GLM_EXACT_LORA_STAGE; + constexpr uint ahead = DS4_GLM_EXACT_LORA_AHEAD; + const uint lane = tid & 31u; + const uint sg = tid >> 5u; + const uint head0 = tgpig.x * args.heads_per_group; + const uint head = head0 + sg; + const uint c0 = tgpig.y * 64u; + const uint j0 = c0 + lane * 2u; + const uint hpg = args.heads_per_group; + const uint n = args.n_selected; + const uint nstages = (n + stage - 1u) / stage; + device const half *cache = (device const half *)kv_lora_cache; + + /* Staging roles: thread t fetches row t / 8, 16-byte chunk t % 8 of the + * kv slice, and the weight of head t / 32, row t % 32. */ + const uint kv_r = tid >> 3u; + const uint kv_chunk = tid & 7u; + const uint w_h = tid >> 5u; + const uint w_r = tid & 31u; + const uint kv_slice_bytes = stage * 64u * sizeof(half); + const uint w_slice_bytes = hpg * stage * sizeof(float); + threadgroup uchar *buf0 = scratch; + threadgroup uchar *buf1 = scratch + kv_slice_bytes + w_slice_bytes; + + uint4 kvq[ahead]; + float wq[ahead]; + #define DS4_GLM_EXACT_LORA_FETCH(k, slot) do { \ + const uint s0_ = (k) * stage; \ + const uint s_ = s0_ + kv_r; \ + const uint row_ = s_ < n ? selected[s_] : 0u; \ + const uint safe_ = row_ < args.cache_cap ? row_ : 0u; \ + kvq[slot] = *(device const uint4 *)(cache + (uint64_t)safe_ * args.kv_lora_dim + c0 + kv_chunk * 8u); \ + const uint sw_ = s0_ + w_r; \ + const uint hh_ = head0 + w_h; \ + wq[slot] = (sw_ < n && hh_ < args.n_head) \ + ? weights[(uint64_t)hh_ * n + sw_] : 0.0f; \ + } while (0) + + for (uint i = 0; i < ahead; i++) { + if (i < nstages) DS4_GLM_EXACT_LORA_FETCH(i, i); + } + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint k = 0; k < nstages; k++) { + threadgroup uchar *buf = (k & 1u) ? buf1 : buf0; + threadgroup uint4 *kv_sh = (threadgroup uint4 *)buf; + threadgroup float *w_sh = (threadgroup float *)(buf + kv_slice_bytes); + kv_sh[kv_r * 8u + kv_chunk] = kvq[0]; + w_sh[w_h * stage + w_r] = wq[0]; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = 1; i < ahead; i++) { + kvq[i - 1u] = kvq[i]; + wq[i - 1u] = wq[i]; + } + if (k + ahead < nstages) DS4_GLM_EXACT_LORA_FETCH(k + ahead, ahead - 1u); + if (head < args.n_head) { + threadgroup const half2 *kv_rows = (threadgroup const half2 *)buf; + threadgroup const float *w_rows = w_sh + sg * stage; + for (uint r = 0; r < stage; r++) { + const float2 kv = (float2)kv_rows[r * 32u + lane]; + const float w = w_rows[r]; + acc0 += w * kv.x; + acc1 += w * kv.y; + } + } + } + #undef DS4_GLM_EXACT_LORA_FETCH + if (head >= args.n_head) return; + const float d = denom[head]; + device float *out = lora + (uint64_t)head * args.kv_lora_dim; + out[j0] = acc0 / d; + out[j0 + 1u] = acc1 / d; +} + +/* Value kernel. Grid: n_head x ceil(value_dim / threads) threadgroups; each + * copies its head's lora vector into threadgroup memory and runs the generic + * kernel's quantised row dot for one output element per thread. */ +kernel void kernel_glm_attention_indexed_decode_exact_value( + constant ds4_metal_args_glm_attention_indexed_decode_exact & args, + device const float *lora, + device const char *value_weight, + device char *heads, + threadgroup float *lora_sum [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + if (head >= args.n_head) return; + const uint nth = ntg_u.x; + device const float *src = lora + (uint64_t)head * args.kv_lora_dim; + for (uint j = tid; j < args.kv_lora_dim; j += nth) lora_sum[j] = src[j]; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint d = tgpig.y * nth + tid; + if (d >= args.value_dim) return; + device float *out = + (device float *)(heads + (uint64_t)head * args.value_dim * sizeof(float)); + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.value_type, row, lora_sum, args.kv_lora_dim); +} + kernel void kernel_glm_attention_indexed_decode( constant ds4_metal_args_glm_attention_indexed_decode & args, device const char *q, @@ -3778,6 +4166,8 @@ kernel void kernel_glm_attention_indexed_batch_group2( } } +/* Original one-head reference from 8969dbb. Keep its expressions and + * specializations independent of the wider-head implementation below. */ template kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl( constant ds4_metal_args_glm_attention_indexed_batch & args, @@ -3979,6 +4369,229 @@ template [[host_name("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_f kernel glm_attention_indexed_batch_lora_group8_vec_t kernel_glm_attention_indexed_batch_lora_group8_vec_impl; +/* Indexed prefill attention over the compact KV cache. + * + * A threadgroup is 8 simdgroups; heads_per_sg is how many heads one simdgroup + * carries, so the threadgroup covers 8 * heads_per_sg heads and every token + * needs 64 / (8 * heads_per_sg) staging passes over its selected rows. At + * heads_per_sg 1 each of the 8 head groups re-stages all 2051 selected rows + * (2 MB per token), which is 34 GB per layer at a 2048-token chunk; carrying + * two heads per simdgroup halves that and each staged row block feeds both. + * + * Nothing a head computes depends on heads_per_sg: the row order, the four + * dot(float4) terms, the simd_sum tree and the online-softmax update are the + * same expressions in the same order, so every head's output is bit-identical + * across the instantiations. */ +template +kernel void kernel_glm_attention_indexed_batch_lora_heads_impl( + constant ds4_metal_args_glm_attention_indexed_batch & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const uint32_t *selected, + device char *lora_out, + threadgroup half4 *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid_u [[thread_index_in_threadgroup]], + ushort lane_u [[thread_index_in_simdgroup]], + ushort sg_u [[simdgroup_index_in_threadgroup]]) { + constexpr uint group_heads = 8u; + constexpr uint stage_rows = 16u; + const uint token = tgpig.y; + const uint tid = (uint)tid_u; + const uint lane = (uint)lane_u; + const uint head_in_group = (uint)sg_u; + const uint head0 = + (tgpig.x * group_heads + head_in_group) * heads_per_sg + args.head_base; + if (token >= args.n_tokens || + args.n_selected == 0u || + args.cache_f16 == 0u || + args.kv_lora_dim != 512u || + (args.qk_rope != 0u && args.qk_rope != 64u)) { + return; + } + + const uint kv_vecs = args.kv_lora_dim >> 2; + const uint rope_vecs = args.qk_rope >> 2; + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = + (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + + threadgroup half4 *kv_shared = scratch; + threadgroup float4 *rope_shared = + (threadgroup float4 *)(kv_shared + stage_rows * kv_vecs); + + bool valid_head[heads_per_sg]; + float4 low[heads_per_sg][4]; + float4 qrope[heads_per_sg]; + FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) { + const uint head = head0 + k; + valid_head[k] = assume_valid_heads || head < args.n_head; + const uint safe_head = valid_head[k] ? head : 0u; + device const float *qh = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)safe_head * qk_dim * sizeof(float)); + device const float4 *low4 = + (device const float4 *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)safe_head * args.kv_lora_dim * sizeof(float)); + low[k][0] = 0.0f; + low[k][1] = 0.0f; + low[k][2] = 0.0f; + low[k][3] = 0.0f; + qrope[k] = 0.0f; + if (valid_head[k]) { + low[k][0] = low4[lane + 0u]; + low[k][1] = low4[lane + 32u]; + low[k][2] = low4[lane + 64u]; + low[k][3] = low4[lane + 96u]; + if (lane < rope_vecs) { + qrope[k] = *((device const float4 *)(qh + args.qk_nope + lane * 4u)); + } + } + } + device const uint32_t *token_selected = + selected + (uint64_t)token * args.n_selected; + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.qk_rope != 0u && args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + float M[heads_per_sg]; + float S[heads_per_sg]; + float4 o[heads_per_sg][4]; + FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) { + M[k] = -FLT_MAX / 2.0f; + S[k] = 0.0f; + o[k][0] = 0.0f; + o[k][1] = 0.0f; + o[k][2] = 0.0f; + o[k][3] = 0.0f; + } + + for (uint base = 0u; base < args.n_selected; base += stage_rows) { + const uint rows = min(stage_rows, args.n_selected - base); + for (uint off = tid; off < rows * kv_vecs; off += 256u) { + const uint rr = off / kv_vecs; + const uint vv = off - rr * kv_vecs; + const uint row = token_selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + if (valid_row) { + device const half4 *src = + (device const half4 *)((device const half *)kv_lora_cache + + (uint64_t)row * args.kv_lora_dim); + kv_shared[off] = src[vv]; + } else { + kv_shared[off] = half4(half(0.0f)); + } + } + for (uint off = tid; off < rows * rope_vecs; off += 256u) { + const uint rr = off / rope_vecs; + const uint vv = off - rr * rope_vecs; + const uint r = vv * 4u; + const uint row = token_selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + if (valid_row) { + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + const float2 y0 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + const float2 y1 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r + 2u, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y); + } else { + rope_shared[off] = float4(0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint rr = 0u; rr < rows; rr++) { + const uint row = token_selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + threadgroup const half4 *kv_row = kv_shared + rr * kv_vecs; + threadgroup const float4 *rope_row = rope_shared + rr * rope_vecs; + /* Verbatim from the one-head kernel, including the repeated + * `(float4)kv_row[...]` in the score and in the output update: + * naming the converted row once is the same value but lets the + * compiler contract the update the other way round, which does + * not round the same. */ + FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) { + float partial = 0.0f; + if (valid_head[k] && valid_row) { + partial += dot(low[k][0], (float4)kv_row[lane + 0u]); + partial += dot(low[k][1], (float4)kv_row[lane + 32u]); + partial += dot(low[k][2], (float4)kv_row[lane + 64u]); + partial += dot(low[k][3], (float4)kv_row[lane + 96u]); + if (lane < rope_vecs) { + partial += dot(qrope[k], rope_row[lane]); + } + } + const float sum = simd_sum(partial); + const float score = + (valid_head[k] && valid_row) ? sum * args.scale : -FLT_MAX / 2.0f; + if (valid_head[k] && valid_row) { + const float new_m = max(M[k], score); + const float old_scale = exp(M[k] - new_m); + const float row_scale = exp(score - new_m); + o[k][0] = o[k][0] * old_scale + (float4)kv_row[lane + 0u] * row_scale; + o[k][1] = o[k][1] * old_scale + (float4)kv_row[lane + 32u] * row_scale; + o[k][2] = o[k][2] * old_scale + (float4)kv_row[lane + 64u] * row_scale; + o[k][3] = o[k][3] * old_scale + (float4)kv_row[lane + 96u] * row_scale; + S[k] = S[k] * old_scale + row_scale; + M[k] = new_m; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + FOR_UNROLL (uint k = 0; k < heads_per_sg; k++) { + if (valid_head[k]) { + const float inv_s = S[k] > 0.0f ? 1.0f / S[k] : 0.0f; + device float4 *out4 = + (device float4 *)(lora_out + + ((uint64_t)token * args.n_head + head0 + k) * + args.kv_lora_dim * sizeof(float)); + out4[lane + 0u] = o[k][0] * inv_s; + out4[lane + 32u] = o[k][1] * inv_s; + out4[lane + 64u] = o[k][2] * inv_s; + out4[lane + 96u] = o[k][3] * inv_s; + } + } +} + +template [[host_name("kernel_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads")]] +kernel glm_attention_indexed_batch_lora_group8_vec_t +kernel_glm_attention_indexed_batch_lora_heads_impl; + template kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl( constant ds4_metal_args_glm_attention_indexed_batch & args, diff --git a/metal/glm53_bf16.metal b/metal/glm53_bf16.metal index a46dd81cd..b7641890d 100644 --- a/metal/glm53_bf16.metal +++ b/metal/glm53_bf16.metal @@ -26,24 +26,20 @@ kernel void kernel_glm53_embedding_bf16( : 0.0f; } -static inline void glm53_mul_mv_bf16_f32_row( - constant glm53_bf16_matmul_args &args, +/* The accumulation, split out unchanged so an epilogue kernel can use the sum + * before it is stored. Callers must range-check out_row and token first. */ +static inline float glm53_mul_mv_bf16_f32_row_sum( + uint in_dim, device const ushort *weights, device const float *x, - device float *out, - uint2 tgpig, - ushort lane, - ushort sg, - ushort nsg) { - const uint out_row = tgpig.x * (uint)nsg + sg; - const uint token = tgpig.y; - if (out_row >= args.out_dim || token >= args.n_rows) return; - - device const ushort *w = weights + (ulong)out_row * args.in_dim; - device const float *xr = x + (ulong)token * args.in_dim; + uint out_row, + uint token, + ushort lane) { + device const ushort *w = weights + (ulong)out_row * in_dim; + device const float *xr = x + (ulong)token * in_dim; float sum = 0.0f; uint k = lane; - for (; k + 224u < args.in_dim; k += 256u) { + for (; k + 224u < in_dim; k += 256u) { const ushort w0 = w[k]; const ushort w1 = w[k + 32u]; const ushort w2 = w[k + 64u]; @@ -69,10 +65,26 @@ static inline void glm53_mul_mv_bf16_f32_row( sum = fma(glm53_bf16_to_f32(w6), x6, sum); sum = fma(glm53_bf16_to_f32(w7), x7, sum); } - for (; k < args.in_dim; k += 32u) { + for (; k < in_dim; k += 32u) { sum = fma(glm53_bf16_to_f32(w[k]), xr[k], sum); } - sum = simd_sum(sum); + return simd_sum(sum); +} + +static inline void glm53_mul_mv_bf16_f32_row( + constant glm53_bf16_matmul_args &args, + device const ushort *weights, + device const float *x, + device float *out, + uint2 tgpig, + ushort lane, + ushort sg, + ushort nsg) { + const uint out_row = tgpig.x * (uint)nsg + sg; + const uint token = tgpig.y; + if (out_row >= args.out_dim || token >= args.n_rows) return; + const float sum = + glm53_mul_mv_bf16_f32_row_sum(args.in_dim, weights, x, out_row, token, lane); if (lane == 0u) out[(ulong)token * args.out_dim + out_row] = sum; } @@ -91,6 +103,96 @@ kernel void kernel_glm53_mul_mv_bf16_f32( tgpig, lane, sg, nsg); } +/* + * BF16 matvec with the HC expansion folded into its epilogue. + * + * The simdgroup that finishes output row d already holds that row's value in + * lane 0, so it can expand it into the four HC streams there instead of + * writing it out and having a second dispatch read it straight back. This is + * the shape kernel_dsv4_q8_hc_expand4_q8_0 already uses for DeepSeek, in BF16. + * + * Decode only: one token, HC = 4. The arithmetic and the operand order match + * kernel_dsv4_hc_expand4 exactly, including that comb is indexed [j][h]. + */ +kernel void kernel_glm53_mul_mv_bf16_f32_hc_expand4( + constant glm53_bf16_matmul_args &args, + device const ushort *weights, + device const float *x, + device float *out, + device const float *residual, + device const float *post, + device const float *comb, + device float *hc_out, + uint2 tgpig [[threadgroup_position_in_grid]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], + ushort nsg [[simdgroups_per_threadgroup]]) { + const uint out_row = tgpig.x * (uint)nsg + sg; + const uint token = tgpig.y; + if (out_row >= args.out_dim || token >= args.n_rows) return; + const float sum = + glm53_mul_mv_bf16_f32_row_sum(args.in_dim, weights, x, out_row, token, lane); + if (lane != 0u) return; + out[(ulong)token * args.out_dim + out_row] = sum; + + const uint n = args.out_dim; + const float r0 = residual[0u * n + out_row]; + const float r1 = residual[1u * n + out_row]; + const float r2 = residual[2u * n + out_row]; + const float r3 = residual[3u * n + out_row]; + for (uint h = 0u; h < 4u; ++h) { + float acc = sum * post[h]; + acc += comb[0u * 4u + h] * r0; + acc += comb[1u * 4u + h] * r1; + acc += comb[2u * 4u + h] * r2; + acc += comb[3u * 4u + h] * r3; + hc_out[h * n + out_row] = acc; + } +} + +struct glm53_bf16_trio_args { + uint in_dim; + uint out_dim_ab; + uint out_dim_c; + uint n_rows; +}; + +/* + * Three matvecs over one shared input in a single dispatch, where the third + * has a shorter output than the first two. GLM 5.3's KDA gate chain is + * exactly that shape: f_a and g_a are [4096 -> 128] and beta is [4096 -> 64], + * all reading attn_norm. The pair kernel could not carry beta because it + * assumes one output width for every slot. + * + * The grid is sized for the wider pair, so the beta slot's upper threadgroups + * exit on the bounds check. + */ +kernel void kernel_glm53_mul_mv_bf16_f32_trio( + constant glm53_bf16_trio_args &args, + device const ushort *weights_a, + device const ushort *weights_b, + device const ushort *weights_c, + device const float *x, + device float *out_a, + device float *out_b, + device float *out_c, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], + ushort nsg [[simdgroups_per_threadgroup]]) { + const uint slot = tgpig.z; + device const ushort *w = slot == 0u ? weights_a + : (slot == 1u ? weights_b : weights_c); + device float *out = slot == 0u ? out_a : (slot == 1u ? out_b : out_c); + const uint out_dim = slot == 2u ? args.out_dim_c : args.out_dim_ab; + const uint out_row = tgpig.x * (uint)nsg + sg; + const uint token = tgpig.y; + if (out_row >= out_dim || token >= args.n_rows) return; + const float sum = + glm53_mul_mv_bf16_f32_row_sum(args.in_dim, w, x, out_row, token, lane); + if (lane == 0u) out[(ulong)token * out_dim + out_row] = sum; +} + kernel void kernel_glm53_mul_mv_bf16_f32_qkv( constant glm53_bf16_matmul_args &args, device const ushort *weights_q, @@ -112,6 +214,32 @@ kernel void kernel_glm53_mul_mv_bf16_f32_qkv( tgpig.xy, lane, sg, nsg); } +/* + * Two independent matvecs of the same shape in one dispatch, selected by + * tgpig.z, exactly as the qkv variant above selects three. The inputs are + * separate pointers rather than one shared row, which lets this serve both + * halves of the GLM 5.3 KDA gate chain: f_a/g_a read the same attn_norm row, + * while f_b/g_b read the two different low-rank vectors those produce. + */ +kernel void kernel_glm53_mul_mv_bf16_f32_pair( + constant glm53_bf16_matmul_args &args, + device const ushort *weights_a, + device const ushort *weights_b, + device const float *x_a, + device const float *x_b, + device float *out_a, + device float *out_b, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], + ushort nsg [[simdgroups_per_threadgroup]]) { + device const ushort *weights = tgpig.z == 0u ? weights_a : weights_b; + device const float *x = tgpig.z == 0u ? x_a : x_b; + device float *out = tgpig.z == 0u ? out_a : out_b; + glm53_mul_mv_bf16_f32_row(args, weights, x, out, + tgpig.xy, lane, sg, nsg); +} + struct glm53_bf16_block16 { ushort v[16]; }; diff --git a/metal/glm53_kda.metal b/metal/glm53_kda.metal index c21d30ecf..b1a2312a9 100644 --- a/metal/glm53_kda.metal +++ b/metal/glm53_kda.metal @@ -48,6 +48,7 @@ kernel void kernel_glm53_kda_decode( threadgroup float *reduce_k = reduce_q + 4u; threadgroup float *reduce_o = reduce_k + 4u; threadgroup float *beta_shared = reduce_o + 4u; + threadgroup float *a_decay_shared = beta_shared + 1u; const uint projection = args.n_heads * D; const uint channel = head * D + tid; @@ -90,16 +91,18 @@ kernel void kernel_glm53_kda_decode( sq[tid] = q_acc / (1.0f + exp(-q_acc)); sk[tid] = k_acc / (1.0f + exp(-k_acc)); sv[tid] = v_acc / (1.0f + exp(-v_acc)); - const float gate = raw_gate[input_base + tid] + dt_bias[channel]; - sd[tid] = exp(args.lower_bound * - (1.0f / (1.0f + exp(-exp(a_log[head]) * gate)))); } if (tid == 0u) { beta_shared[0] = 1.0f / (1.0f + exp(-raw_beta[(ulong)row * args.n_heads + head])); + /* head is uniform over the threadgroup, so exp(a_log[head]) is a + * single value; every one of the D channels used to recompute it. */ + a_decay_shared[0] = exp(a_log[head]); } - threadgroup_barrier(mem_flags::mem_threadgroup | - mem_flags::mem_device); + /* Only threadgroup memory is shared between threads here: the conv-state + * writes above are each thread's own channel and no thread reads another's, + * so the barrier does not need device scope. */ + threadgroup_barrier(mem_flags::mem_threadgroup); float q_sumsq = sq[tid] * sq[tid]; float k_sumsq = sk[tid] * sk[tid]; @@ -119,6 +122,9 @@ kernel void kernel_glm53_kda_decode( if (tid < D) { sq[tid] *= q_scale; sk[tid] *= k_scale; + const float gate = raw_gate[input_base + tid] + dt_bias[channel]; + sd[tid] = exp(args.lower_bound * + (1.0f / (1.0f + exp(-a_decay_shared[0] * gate)))); } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -141,8 +147,9 @@ kernel void kernel_glm53_kda_decode( float hq = simd_sum(dot(h, q4)); if (lane == 0u) so[value] = hq; } - threadgroup_barrier(mem_flags::mem_threadgroup | - mem_flags::mem_device); + /* Likewise: the state writes above are not re-read in this kernel, only + * so[] crosses simdgroups. */ + threadgroup_barrier(mem_flags::mem_threadgroup); float o_sumsq = so[tid] * so[tid]; o_sumsq = simd_sum(o_sumsq); @@ -246,6 +253,184 @@ kernel void kernel_glm53_kda_prefill_prepare( } } +/* + * Boundary rows for the blocked prepare kernel below. + * + * A block starts its causal convolution window on the raw q/k/v of the three + * rows before it, and the block before it overwrites exactly those rows with + * its normalized outputs. Block 0 also needs an immutable copy of the incoming + * conv state: the last block may overwrite that state before block 0 starts. + */ +struct glm53_kda_blocked_args { + uint n_heads; + uint n_rows; + uint block_rows; + uint n_blocks; + uint block_base; + float lower_bound; + float norm_eps; +}; + +kernel void kernel_glm53_kda_prefill_conv_halo( + constant glm53_kda_blocked_args &args, + device const float *q, + device const float *k, + device const float *v, + device float *halo, + device const float *conv_state, + uint2 tgpig [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + constexpr uint D = 128u; + constexpr uint HISTORY = 3u; + constexpr uint NTH = 256u; /* matches the dispatch */ + const uint projection = args.n_heads * D; + const uint block = tgpig.x; + const uint w = tgpig.y; + if (block >= args.n_blocks || w >= HISTORY) return; + const ulong plane = (ulong)args.n_blocks * HISTORY * projection; + const ulong slot = ((ulong)block * HISTORY + w) * projection; + for (uint c = tid; c < projection; c += NTH) { + if (block == 0u) { + const ulong index = (ulong)w * projection + c; + halo[slot + c] = conv_state[index]; + halo[plane + slot + c] = conv_state[HISTORY * projection + index]; + halo[2u * plane + slot + c] = conv_state[2u * HISTORY * projection + index]; + } else { + const uint token = block * args.block_rows + w - HISTORY; + const ulong index = (ulong)token * projection + c; + halo[slot + c] = q[index]; + halo[plane + slot + c] = k[index]; + halo[2u * plane + slot + c] = v[index]; + } + } +} + +/* + * Token-parallel form of kernel_glm53_kda_prefill_prepare. + * + * The serial kernel runs one threadgroup per head -- 64 of them on an 80-core + * GPU -- and walks all the chunk's tokens inside it, so it costs 3.4 ms of a + * 2048-token layer at about 5% occupancy. Nothing in it is actually + * sequential: the causal convolution reads the raw q/k/v of t-3..t, which are + * all known before the kernel starts. + * + * One threadgroup now owns (block of block_rows tokens, head) and keeps the + * three-row convolution history in registers instead of re-reading and + * re-writing the device conv state every token. Its first three rows come + * from the immutable halo above, including the incoming state for block 0. + * The last block leaves the outgoing conv state exactly where the serial kernel + * left it. + * + * Every value keeps the serial kernel's expression and order: the same + * four-term fma chain in the same order, the same silu, the same + * 4-simdgroup RMS reduction, the same decay-gate expression. + */ +kernel void kernel_glm53_kda_prefill_prepare_blocked( + constant glm53_kda_blocked_args &args, + device float *q, + device float *k, + device float *v, + device float *raw_gate, + device const float *q_conv, + device const float *k_conv, + device const float *v_conv, + device const float *a_log, + device const float *dt_bias, + device float *conv_state, + device const float *halo, + threadgroup float *scratch [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint D = 128u; + constexpr uint HISTORY = 3u; + const uint block = tgpig.x + args.block_base; + const uint head = tgpig.y; + if (block >= args.n_blocks || head >= args.n_heads) return; + threadgroup float *sq = scratch; + threadgroup float *sk = sq + D; + threadgroup float *reduce_q = sk + D; + threadgroup float *reduce_k = reduce_q + 4u; + const uint projection = args.n_heads * D; + const uint channel = head * D + tid; + const uint token0 = block * args.block_rows; + const uint token_end = token0 + min(args.block_rows, args.n_rows - token0); + if (token0 >= token_end) return; + + float hq[HISTORY]; + float hk[HISTORY]; + float hv[HISTORY]; + const ulong plane = (ulong)args.n_blocks * HISTORY * projection; + const ulong slot = (ulong)block * HISTORY * projection + channel; + for (uint w = 0; w < HISTORY; w++) { + hq[w] = halo[slot + w * projection]; + hk[w] = halo[plane + slot + w * projection]; + hv[w] = halo[2u * plane + slot + w * projection]; + } + + for (uint token = token0; token < token_end; token++) { + const ulong index = (ulong)token * projection + channel; + float q_acc = 0.0f; + float k_acc = 0.0f; + float v_acc = 0.0f; + for (uint w = 0; w < HISTORY; w++) { + q_acc = fma(hq[w], q_conv[(ulong)channel * 4u + w], q_acc); + k_acc = fma(hk[w], k_conv[(ulong)channel * 4u + w], k_acc); + v_acc = fma(hv[w], v_conv[(ulong)channel * 4u + w], v_acc); + } + const float q_new = q[index]; + const float k_new = k[index]; + const float v_new = v[index]; + q_acc = fma(q_new, q_conv[(ulong)channel * 4u + 3u], q_acc); + k_acc = fma(k_new, k_conv[(ulong)channel * 4u + 3u], k_acc); + v_acc = fma(v_new, v_conv[(ulong)channel * 4u + 3u], v_acc); + hq[0] = hq[1]; hq[1] = hq[2]; hq[2] = q_new; + hk[0] = hk[1]; hk[1] = hk[2]; hk[2] = k_new; + hv[0] = hv[1]; hv[1] = hv[2]; hv[2] = v_new; + + sq[tid] = q_acc / (1.0f + exp(-q_acc)); + sk[tid] = k_acc / (1.0f + exp(-k_acc)); + v[index] = v_acc / (1.0f + exp(-v_acc)); + const float gate = raw_gate[index] + dt_bias[channel]; + raw_gate[index] = exp(args.lower_bound * + (1.0f / (1.0f + exp(-exp(a_log[head]) * gate)))); + threadgroup_barrier(mem_flags::mem_threadgroup | + mem_flags::mem_device); + + float q_sumsq = simd_sum(sq[tid] * sq[tid]); + float k_sumsq = simd_sum(sk[tid] * sk[tid]); + if (lane == 0u) { + reduce_q[sg] = q_sumsq; + reduce_k[sg] = k_sumsq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float q_total = lane < 4u ? reduce_q[lane] : 0.0f; + float k_total = lane < 4u ? reduce_k[lane] : 0.0f; + q_total = simd_sum(q_total); + k_total = simd_sum(k_total); + q[index] = sq[tid] * rsqrt(q_total + 1.0e-6f) * + 0x1.6a09e6p-4f; + k[index] = sk[tid] * rsqrt(k_total + 1.0e-6f); + threadgroup_barrier(mem_flags::mem_threadgroup | + mem_flags::mem_device); + } + + /* The serial kernel leaves the conv state holding the raw q/k/v of the + * last three rows it processed, which is what this block's history is + * once its last token has shifted through. */ + if (token_end == args.n_rows) { + device float *q_state = conv_state; + device float *k_state = q_state + HISTORY * projection; + device float *v_state = k_state + HISTORY * projection; + for (uint w = 0; w < HISTORY; w++) { + q_state[(ulong)w * projection + channel] = hq[w]; + k_state[(ulong)w * projection + channel] = hk[w]; + v_state[(ulong)w * projection + channel] = hv[w]; + } + } +} + kernel void kernel_glm53_kda_prefill_recurrence( constant glm53_kda_args &args, device const float *q, @@ -286,6 +471,59 @@ kernel void kernel_glm53_kda_prefill_recurrence( *state_ptr = h; } +/* Each value row is an independent recurrence. Carrying two or four rows in + * a SIMDgroup reuses q/k/decay and beta across them, while keeping the serial + * token order, dot reductions and FMA expression of the reference above. */ +template +kernel void kernel_glm53_kda_prefill_recurrence_values( + constant glm53_kda_args &args, + device const float *q, + device const float *k, + device const float *v, + device const float *decay, + device const float *raw_beta, + device float *state, + device float *out, + uint2 tgpig [[threadgroup_position_in_grid]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint D = 128u; + const uint head = tgpig.x; + const uint value0 = (tgpig.y * 4u + sg) * VALUES; + if (head >= args.n_heads || value0 + VALUES > D) return; + const uint projection = args.n_heads * D; + const uint k0 = lane * 4u; + float4 h[VALUES]; + FOR_UNROLL (uint i = 0; i < VALUES; i++) { + h[i] = *((device float4 *)(state + ((ulong)head * D + value0 + i) * D + k0)); + } + for (uint token = 0; token < args.n_rows; token++) { + const ulong base = (ulong)token * projection + head * D; + const float4 q4 = *((device const float4 *)(q + base + k0)); + const float4 k4 = *((device const float4 *)(k + base + k0)); + const float4 decay4 = *((device const float4 *)(decay + base + k0)); + const float beta = 1.0f / + (1.0f + exp(-raw_beta[(ulong)token * args.n_heads + head])); + FOR_UNROLL (uint i = 0; i < VALUES; i++) { + h[i] *= decay4; + const float hk = simd_sum(dot(h[i], k4)); + const float delta_v = (v[base + value0 + i] - hk) * beta; + h[i] = fma(k4, float4(delta_v), h[i]); + const float result = simd_sum(dot(h[i], q4)); + if (lane == 0u) out[base + value0 + i] = result; + } + } + FOR_UNROLL (uint i = 0; i < VALUES; i++) { + *((device float4 *)(state + ((ulong)head * D + value0 + i) * D + k0)) = h[i]; + } +} + +typedef decltype(kernel_glm53_kda_prefill_recurrence_values<2u>) glm53_kda_recurrence_values_t; +template [[host_name("kernel_glm53_kda_prefill_recurrence_v2")]] +kernel glm53_kda_recurrence_values_t kernel_glm53_kda_prefill_recurrence_values<2u>; +template [[host_name("kernel_glm53_kda_prefill_recurrence_v4")]] +kernel glm53_kda_recurrence_values_t kernel_glm53_kda_prefill_recurrence_values<4u>; + kernel void kernel_glm53_kda_prefill_output( constant glm53_kda_args &args, device float *out, diff --git a/metal/moe.metal b/metal/moe.metal index 9d0840d7d..f6aa9b13b 100644 --- a/metal/moe.metal +++ b/metal/moe.metal @@ -8819,6 +8819,8 @@ template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull") typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id; typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_f16_rhs; typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_tail_cull; +typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4, true>) mul_mm_id_q4_K_f16_rhs_tail_cull; +typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4, true>) mul_mm_id_q4_K_f32_rhs_tail_cull; typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4>) mul_mm_id_mxfp4_f16_rhs_half_lut; typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_half_lut_tail_cull; typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_addr; @@ -8828,6 +8830,7 @@ typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, hal template [[host_name("kernel_mul_mm_id_q8_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_q2_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_q4_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q4_K_f32_tail_cull")]] kernel mul_mm_id_q4_K_f32_rhs_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4, true>; template [[host_name("kernel_mul_mm_id_q5_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_q6_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>; template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>; @@ -8835,6 +8838,7 @@ template [[host_name("kernel_mul_mm_id_mxfp4_f32")]] kernel mul_mm_id ker template [[host_name("kernel_mul_mm_id_q8_0_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_q2_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_q4_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_q4_K_f16_tail_cull")]] kernel mul_mm_id_q4_K_f16_rhs_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4, true>; template [[host_name("kernel_mul_mm_id_q5_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_q6_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, half, half4x4, half, half2x4>; template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, half, half4x4, half, half2x4>; diff --git a/speed-bench/glm53_decode_findings.md b/speed-bench/glm53_decode_findings.md new file mode 100644 index 000000000..56135086f --- /dev/null +++ b/speed-bench/glm53_decode_findings.md @@ -0,0 +1,959 @@ +# Where GLM 5.3 Flash decode time goes on an M3 Ultra + +Apple M3 Ultra, 80 GPU cores, 512 GB, macOS 26.5.2, Metal backend. +Model: `GLM-5.3-Flash-Q4_K.gguf`, 177.8 GiB, fully resident, no SSD streaming. + +The short version: the routed-expert kernels are already close to the hardware +ceiling, and the largest single consumer of decode bandwidth is not the experts +at all. It is the KDA (linear attention) projections, which this artifact +stores as **BF16** while its experts are Q4_K. Requantizing them to Q8_0 -- +which is what `gguf-tools/glm53_quantize.py` already specifies for +`role == "linear_attention"` -- is worth **+13.4% decode** with no measurable +quality cost. + +## Decode budget + +Measured with `DS4_GLM_DECODE_ABLATE`, which removes a stage and reports the +resulting speed. Baseline 21.19 tok/s, two baseline runs 0.38% apart. + +**The original KDA row was never measured.** `DS4_GLM_DECODE_ABLATE` had no +`kda` bit and the KDA path returned before the mask was read, so the 18.37 ms +and 38.9% recorded here first could not have come from this harness. `kda`, +`kda_qkv`, `kda_gate`, `kda_recur`, `kda_out`, `hc` and `head` now exist, and +the whole budget below was re-measured through them. + +Everything in the table is from **one** run of arms against **one** baseline of +22.375 tok/s = 44.693 ms/token (four interleaved baselines, 0.85% spread). +Do not mix it with the earlier 21.19 tok/s figures; those are superseded. + +| component | ms/token | share | superseded figure | +|---|---:|---:|---| +| KDA attention (34 layers) | 15.99 | 35.8% | was 18.37 / 38.9% | +| DSA attention core (11 layers) | 8.02 | 18.0% | was 7.91 / 16.8% | +| routed MoE (42 layers) | 7.89 | 17.6% | was 8.08 / 17.1% | +| mHC producer chain (90 sites) | 3.99 | 8.9% | was inside the residual row | +| shared expert (42 layers) | 1.98 | 4.4% | was 2.13 / 4.5% | +| output head (norm + logits matvec) | 1.80 | 4.0% | was inside the residual row | +| attn_output projection (11 layers) | 1.33 | 3.0% | was 1.26 / 2.7% | +| q_path (11 layers) | 0.44 | 1.0% | was 0.59 / 1.3% | +| indexer (11 layers) | 0.21 | 0.5% | was 0.35 / 0.7% | +| remaining norms, residual, hc_expand | 3.04 | 6.8% | residual, not ablated | +| **total** | **44.69** | **100%** | | + +Only KDA moved outside noise; every other carried-over row reproduced. + +Bandwidth is deliberately **not** a column here. It is only meaningful where +the byte count is exactly derivable, which is the dense projections; see "What +is left, priced" below for those, computed against 736.9 GB/s. The routed-MoE +and shared-expert byte figures recorded in the first version of this document +(4.43 and 0.55 GiB) depend on which experts a token selects and were never +re-derived here, so they are omitted rather than restated. + +**A per-token bandwidth figure computed over the whole decode step is +misleading.** Dividing total bytes by total decode time gives roughly a fifth +of peak, but the bandwidth-bound kernels only occupy about a fifth of the step. +The kernels themselves are near the ceiling; the rest of the step is other +work. + +**Ablation arms are destructive.** A skipped stage leaves stale contents in +its output buffer, which is fine for timing the dispatches that remain but can +in principle change data-dependent routing downstream (expert selection, index +selection). The arms agree with each other -- `hc,head` measures 5.77 ms +against 5.79 for the two separately, and the KDA substages sum to 15.44 against +15.99 for the whole stage -- so the effect is small here, but these are +skip-ablation estimates, not per-kernel timings. + +## The BF16 KDA projections + +`blk.N.kda_q`, `kda_k`, `kda_v` and `kda_output` are BF16 in this artifact, on +all 34 KDA layers. They are dense -- every one is read on every decoded token: + +| | bytes/token | share of decode traffic | +|---|---:|---:| +| KDA q/k/v/output (BF16) | 8.50 GiB | **60%** | +| routed experts (Q4_K) | 4.43 GiB | 31% | +| shared expert (Q8_0) | 0.55 GiB | 4% | +| attn_output (Q8_0) | 0.73 GiB | 5% | + +The KDA projections alone read nearly twice what all routed experts read. + +This is not what the repo's own quantizer produces. `regular_qtype()` in +`gguf-tools/glm53_quantize.py` maps `role == "linear_attention"` to `Q8_0` for +its default `--artifact q4`, and embedding/output likewise. This artifact has +all of them at BF16, so it was not produced by that path. + +### Why BF16 is not simply a mistake + +Metal has a fused three-way QKV matmul, `ds4_gpu_glm53_matmul_bf16_qkv`, which +requires all three of q/k/v to be BF16 and issues one dispatch instead of +three. It is gated behind `DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE`, so it was +added as an M3 Ultra optimisation. Quantizing KDA forfeits it and falls back +to the generic per-tensor matmul. + +Measured, so the trade is not a guess: + +| | decode tok/s | Δ | +|---|---:|---:| +| baseline (fused BF16 QKV) | 21.12 | -- | +| `DS4_METAL_DISABLE_GLM53_BF16_QKV=1` | 20.97 | -0.7% | +| `DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE=1` | 20.96 | -0.8% | + +The fusion is worth **0.7%**. The BF16 storage it requires costs an order of +magnitude more than that. No fused Q8_0 QKV kernel is needed to capture the +win; the generic fallback is nearly free. + +## Result of requantizing KDA to Q8_0 + +`gguf-tools/glm53-requant-bf16` converts the 136 KDA tensors from BF16 to Q8_0 +into a **new file**, copying every other byte verbatim, through the same +`quants.c` facade the other tools use. It refuses an output that resolves to +the input (same path, hard link or symlink), because the input stays mmapped +for the whole run; there is no in-place mode. + + make -C gguf-tools glm53-requant-bf16 + ./gguf-tools/glm53-requant-bf16 in.gguf out.gguf --type q8_0 --tensors kda + +`--tensors` also takes `head` (`output.weight`), `embd` (`token_embd.weight`) +and `all`; it defaults to `kda`. + +8.50 GiB of KDA weights become 4.52 GiB; the file goes 177.8 -> 173.8 GiB. + +Speed, arms interleaved O-Q-Q-O with the same binary and only the model file +changing, 8 context frontiers: + +| ctx | original | kda Q8_0 | Δ | +|---:|---:|---:|---:| +| 2,048 | 21.05 | 23.91 | +13.54% | +| 8,192 | 20.65 | 23.42 | +13.44% | +| 16,384 | 20.56 | 23.31 | +13.40% | +| **mean (8 ctx)** | | | **+13.37%** | + +Prefill -0.25%. Within-arm drift 0.15-0.19%, so the effect is far outside the +noise, and it is within 0.3 points at every context. + +Quality, teacher-forced over 18,672 tokens of `promessi_sposi.txt`: + +| | avg NLL | perplexity | +|---|---:|---:| +| original (BF16 KDA) | 1.838851 | 6.289309 | +| requantized (Q8_0 KDA) | 1.834773 | **6.263711** | + +No degradation -- marginally better, which at this size is noise. Greedy +generations from both are coherent and track word for word until a late +paraphrase. + +### The projection was too optimistic, and why -- corrected + +The earlier reading of this was wrong, and it mattered, because it set the +priority for the whole KDA path. + +It went: scaling KDA's 497 GB/s by the byte reduction predicts +22%, we +measured +13.4%, therefore only ~62% of KDA was weight streaming and the +remaining **~7 ms/token** is conv1d, gating and the recurrent state update -- +"the next thing to attack, and not a bandwidth problem." + +Direct substage ablation says otherwise. Splitting KDA on the original BF16 +artifact: + +| substage | ms/token | share of decode | +|---|---:|---:| +| qkv projections | 9.68 | 21.7% | +| output projection | 3.16 | 7.1% | +| gate/beta low-rank chain | 1.37 | 3.1% | +| recurrence kernel (conv1d + gating + state) | **1.23** | **2.8%** | +| unattributed (dispatch, interaction) | 0.55 | 1.2% | + +The conv1d, the gating and the recurrent state update together are **1.23 +ms/token**, not ~7. KDA is about 90% weight streaming, not 62%. Requantizing +the same tensors to Q8_0 and re-ablating confirms it directly -- qkv goes 9.68 +-> 5.61 ms and the output projection 3.16 -> 1.85 ms against a pure-bandwidth +prediction of 5.14 and 1.68, so both are ~90% bandwidth-scaled: + +| | original | Q8_0 KDA | +|---|---:|---:| +| decode | 22.375 tok/s | 25.545 tok/s (**+14.2%**) | +| KDA total | 15.99 ms | 10.40 ms | + +The +13.4% headline reproduces (+14.2% here). Only the explanation was wrong. + +Where the original inference went astray: it assumed everything in KDA that +did not scale with weight bytes was recurrence work. Most of it is instead the +projections failing to scale *perfectly* -- they are ~90% bandwidth-bound, not +100% -- plus fixed dispatch cost. Attributing that gap to the recurrence +inflated a 1.23 ms stage into a 7 ms one. + +The practical consequence: **the recurrence kernel is not where the time is.** +Work on `metal/glm53_kda.metal` is capped at 2.8% of decode no matter how good +it gets. The qkv projections, at 21.7%, are the KDA target that matters. + +## Splitting the old "norms, hyper-connections, residual, LM head" row + +That row was a residual -- whatever the other arms did not account for -- and +at ~18% of decode it was the second largest line in the budget with nothing +measured inside it. The `hc` and `head` ablation arms split it: + +| | ms/token | share | +|---|---:|---:| +| mHC producer chain | 3.99 | 8.9% | +| output head | 1.80 | 4.0% | +| everything else in the row | ~3.0 | ~6.7% | + +`hc,head` together measure 5.77 ms against 5.79 for the two separately, so the +split is additive and the arms are not interacting. + +**The mHC producer was the largest unoptimised item in the decode step.** +`glm53_graph_hc_pre` issued four dispatches -- plain RMSNorm, the 16384->24 mix +matvec, the split/mix, and the weighted RMSNorm -- twice per layer over 45 +layers, so 360 small dispatches per token for 3.99 ms of work. DeepSeek V4 +already fused the F16 equivalent; GLM 5.3 now uses the same kernel with BF16 +mix weights, one dispatch per site instead of four: + +| ctx | four dispatches | fused | delta | +|---:|---:|---:|---:| +| 2,048 | 22.282 | 23.545 | **+5.67%** | +| 4,096 | 21.94 | 23.195 | +5.72% | +| 16,384 | 21.795 | 23.00 | +5.53% | + +Bit-exact: all 154,880 logits match to max|delta| = 0. Prefill is unchanged, +since only the decode path is fused. 2.41 ms of the 3.99 ms is gone; the +remaining 1.58 ms is the fused kernel's own arithmetic. + +**This is the largest engine-only decode gain found on this path**, and it is +worth contrasting with the KDA recurrence work the earlier budget pointed at: +that stage is 2.8% of decode in total, while this one change is +5.67%. + +**The output head is nearly all matvec.** 1.80 ms for a [4096 -> 154880] +BF16 matvec is close to what its 1.27 GB costs at this machine's measured +bandwidth, so there is no dispatch overhead worth chasing there. + +### Why GPU-side argmax is not worth doing + +A natural suggestion is to stop reading all 154,880 logits back and scanning +them on the CPU, and instead do a hierarchical argmax/top-k on the GPU and +return only the token. Measured directly on this machine: + + logits readback (memcpy of 605 KiB) 0.0143 ms + CPU argmax scan over 154,880 floats 0.1668 ms + combined 0.1811 ms + +That is **0.40% of a 44.76 ms decode step, below the run-to-run spread**, so +the change could not be shown to work even if it were free. It also would not +remove a synchronisation: the token is needed before the next step can start +either way. Not worth the complexity. + +## A trap in the stage profiler + +`DS4_METAL_DECODE_STAGE_PROFILE` reports a stage named `attn_output` on all 45 +layers, and on KDA layers it is the largest stage in the run. It is **not** +measuring the output projection there. The profile boundary sits after the +`glm53_attention_done:` label, and KDA layers reach that label by `goto`, so on +those layers the `attn_output` sample times the entire KDA attention. + +The real output projection is 2.7% of decode, not 39%. Ablation and the +profiler agree to within 0.3 points once the label is read correctly (2.7% vs +3.0% on the 11 DSA layers, where the label means what it says). + +Two further cautions when using that profiler: it flushes the command buffer at +every boundary, which on this workload adds a uniform ~0.206 ms per boundary +and roughly triples the measured decode time; and because the floor is uniform, +stages that do little real work all read as roughly the floor. Prefer +`DS4_GLM_DECODE_ABLATE` for attribution and use the stage profiler to localise. + +## A non-destructive second instrument + +The ablation arms are destructive: a skipped stage leaves stale contents, so +the run is timing-only and can in principle perturb data-dependent routing. +`DS4_GLM_DECODE_REPEAT` is the other half of the pincer. Every stage it +accepts is a pure function of its inputs, so dispatching it one extra time per +site writes the same bytes; the whole-token delta is then one extra execution +of that stage and **the model output is unchanged**. Verified: all six arms +dump logits identical to the baseline at max|delta| = 0. + +Only idempotent stages get a bit. The KDA recurrence advances conv and +recurrent state, and directional steering updates in place, so neither can be +repeated and neither is offered. + +Where the two instruments agree, the number is trustworthy: + +| stage | ablate | repeat | agreement | +|---|---:|---:|---| +| kda_qkv | 9.69 | 9.76 | 0.7% | +| head | 1.80 | 1.76 | 2% | +| kda_gate | 1.37 | 1.45 | 6% | +| **kda_out** | **3.33** | **2.47** | **35%** | + +The `kda_out` disagreement is reproducible across rounds, and the exact byte +count settles it. Both instruments agree that kda_qkv costs 9.69 ms for +6.845 GB, i.e. 706 GB/s; kda_output is 2.282 GB, which at that rate is 3.23 ms +-- next to the ablation figure, not the repeat one. Repeat underestimates here +because the second dispatch re-reads a 67 MB per-layer weight set that is +partly still resident, while kda_qkv's 201 MB per layer does not survive. + +**So: repeat is the right instrument for dispatch-bound stages and undercounts +cache-friendly bandwidth-bound ones; ablation is the reverse.** Use both, and +let exact bytes arbitrate when they disagree. + +### The remaining bucket, partly split + +`hc_expand` measured **0.55 ms/token, 1.3% of decode** by repeat -- a +dispatch-bound stage, so that figure was the reliable one. Of the 90 sites, 87 are +now folded into whatever produces their input, and between them they returned +0.37 ms of it: + +| site | count | mechanism | gain | +|---|---:|---|---:| +| kda_output (BF16 matvec) | 34 | new epilogue kernel | +0.46% | +| attn_output (Q8_0 matvec) | 11 | DeepSeek's existing fused kernel | +0.11% | +| FFN tail (routed+shared add) | 42 | existing `has_add` path on the expand | +0.14% | + +That is 87 of the 90 sites; the three leading dense FFN layers have no +routed/shared split to defer and keep the separate expand. + +The last two together are +0.48% (t=9.65, n=8), 0.199 ms over 54 sites, or +3.7 us per site -- below the 4.6 us launch cost and the KDA epilogue's 5.6 us, +which fits: those two remove a cheap elementwise add and a Q8_0 matvec rather +than a BF16 matvec plus a 64 KiB round-trip. + +Two of the three needed no new kernel at all. `ds4_gpu_matmul_q8_0_hc_expand_tensor` +already existed for DeepSeek and reads post/comb from `hc_split` at the offsets +GLM uses; `ds4_gpu_hc_expand_add_tensor` already exposed the expand kernel's +`has_add` path. Only the BF16 matvec needed an epilogue written. + +That leaves roughly 2.5 ms in the residual row for the residual adds, +directional steering, the remaining norms and the final HC collapse, none of +which are separated yet. + +With the mHC producer now fused, `hc_pre` measures 1.36 ms by repeat, down from +the 3.99 ms the four-dispatch chain cost. + +## The budget, re-measured after the fusions + +Everything above was measured before the mHC, gate-pairing and HC-expand work. +Re-run on the current tip, baseline 41.598 ms/token: + +| stage | ms | share | +|---|---:|---:| +| KDA attention | 15.39 | 37.0% | +| routed MoE | 7.89 | 19.0% | +| DSA attention core | 7.86 | 18.9% | +| shared expert | 2.07 | 5.0% | +| output head | 1.68 | 4.0% | +| mHC producer | 1.44 | 3.5% | +| attn_output | 1.10 | 2.6% | +| q_path | 0.56 | 1.4% | +| indexer | 0.19 | 0.5% | +| **residual** | **3.42** | **8.2%** | + +The mHC producer is down from 3.99 to 1.44 ms. Note that `kda` now also +covers the HC expansion folded into `kda_output`, so its 15.39 is not directly +comparable with the earlier 15.99. + +### How much launch overhead is left in total + +Chasing the residual stage by stage has diminishing returns, so +`DS4_METAL_ENCODER_COUNT` counts compute-encoder acquisitions instead -- one +per dispatch for essentially every primitive here. (It is acquisitions, not +encoder objects: inside a batch the same encoder is handed back for every +dispatch, so the count is a dispatch proxy.) Differencing two runs of +different decode length removes prefill and setup: + + 6,605 acquisitions over 8 decode tokens + 26,989 over 40 + (26989 - 6605) / 32 = **637 dispatches per decode token** + +If the 4.6 us launch cost measured on the gate pairing transfers to the other +kernels and command-buffer arrangements -- which has not been checked, so treat +this as an estimate rather than a measured floor -- that is about **2.93 +ms/token, 7% of the 41.31 ms step**, spread across every stage rather than +concentrated in the residual. It is roughly what the remaining dispatch-count +work is competing for: no arrangement of the current graph gets under the +launch overhead without removing launches, whatever its exact size. + +For scale, the fusions in this branch have already taken roughly 3.5 ms of +dispatch and intermediate-traffic cost out of the step, so what is left is +smaller than what was found. + +### Splitting the residual + +`DS4_GLM_DECODE_REPEAT` gained a `router` bit. Repeat rather than ablate is +the only honest instrument for it: skipping the router leaves a stale expert +selection, which changes which experts the routed stage streams and therefore +changes the very cost being measured. Verified non-destructive (identical +greedy output). + +| | ms | share of decode | share of the residual | +|---|---:|---:|---:| +| router (logits + top-k, 84 dispatches) | 0.95 | 2.3% | 28% | +| remaining hc_expand (FFN tail, dense attn) | 0.33 | 0.8% | 10% | +| still unattributed | 2.13 | 5.1% | 62% | + +The router reads `ffn_gate_inp`, which is **F32** at [4096, 288] over 42 +layers: 198.3 MB/token, or 0.28 ms at the 707 GB/s the dense projections +achieve. So about a third of the router is weight streaming and the other +~0.66 ms is the top-k select over 288 experts plus launch cost. Requantizing +`ffn_gate_inp` is a model-artifact change and routing precision is the obvious +risk, but it is the only 200 MB/token F32 tensor left in the decode step. + +### The shared expert is not the outlier it looked like + +The original budget recorded the shared expert at 0.55 GiB/token and 279 GB/s, +38% of ceiling -- far below every other kernel, and an obvious target. **That +byte count was under by about 2x.** Summed from the tensor table, the shared +expert reads three Q8_0 [4096, 2048] tensors per layer over 42 layers: + + gate + up + down = 3 x 374.2 MB = 1.123 GB/token + +Against the measured 2.07 ms that is **542 GB/s, 74% of ceiling** -- in the +same band as KDA overall (77%), not an outlier. Its gate/up/SwiGLU is already +fused via `ds4_gpu_shared_mid_swiglu_q8_0_tensor`. Closing the remaining gap +to 707 GB/s would be worth about 0.45 ms, 1.1%, not the large win the 38% +figure implied. + +## The DSA attention core: corrected, then fixed + +An earlier revision priced this stage at 135.4 MB/token and 2.4% of the memory +ceiling. **That was wrong by 23x on bytes**, for three reasons each worth +recording: + +- **`n_rot` is 0 for GLM 5.3**, so a compact cache row is the 512-wide lora + part alone, 1024 B in f16 -- not 1152 B with a 64-wide rope tail. +- **There are 11 DSA layers in the trunk, not 12.** `attn_v_b` appears 12 + times because the MTP layer has one, which is not in the decode path. +- **The cache is not read once per layer.** The generic kernel dispatches one + threadgroup per head -- 64 of them -- each independently walking all selected + rows twice, once to score and once for the weighted sum. + +Corrected traffic: 2051 rows x 1024 B x 2 passes x 64 heads x 11 layers is +2.96 GB, plus 0.20 GB of `attn_k_b`/`attn_v_b`, so 3.15 GB/token. At 7.7 ms +that is 409 GB/s, **56% of ceiling** -- real headroom, but not the collapse the +old figure implied. `qk_low` accounts for 0.55 ms of the stage, leaving 7.23 +ms in the kernel proper. + +### The split kernel was tried first, and is not what ships + +`kernel_glm_attention_indexed_decode_split_group8_partial` is the obvious +candidate: 8 heads per threadgroup so a loaded cache row serves eight of them, +16 rows staged in threadgroup memory so scoring and the weighted sum read +device memory once, and row blocking so the work spreads over many more +threadgroups than the generic kernel's 64. GLM 5.2 decode has been running it +all along; two guards kept GLM 5.3 out (`qk_rope != 64`, every rope path of +which is zero-trip at GLM 5.3's `n_rot = 0`, and a block-count check against +worst-case buffer sizing rather than the runtime count). Relaxing both, at +ctx 2048, interleaved: + +| | tok/s | ms/token | +|---|---:|---:| +| generic kernel | 24.232 | 41.27 | +| split group8 | **28.318** | **35.31** | +| | **+16.86%** | | + +It is not bit-exact against the generic kernel, though: it scores with +lane-split dots and reduces with an online softmax across row blocks, and the +DSA attention outputs differ by 3.06e-05 of range. That is float reordering, +and every quality measurement taken -- greedy generation identical over 128 to +256 tokens on six prompts, long-context NLL within 0.002% -- says it is +harmless. It still fails the standard this branch holds itself to, which is +that a faster path must reproduce the path it replaces, so **on this branch +GLM 5.3 does not use it.** It remains what it was before, GLM 5.2's kernel, +now selectable off under `--quality` (the generic kernel is the exact one) and +via `DS4_METAL_DISABLE_GLM53_DSA_SPLIT`, and covered by `tests/test_glm53_kda` +against a double-precision reference with out-of-range and `UINT32_MAX` rows +in the selection. + +One thing about its call site was wrong and is fixed regardless: it passed +`selected_rows_valid = true`, selecting the kernel variant that skips the `row +< cache_cap` test. GLM 5.2's selections are always in range. GLM 5.3's are +not once more than the 4096-row full-attention window is visible (8192 under +SSD streaming): beyond it the pool selector supplies 2051 rows padded with +`UINT32_MAX` sentinels, and the unchecked variant reads those out of bounds. +On this machine those reads returned values whose effect stayed below the +greedy threshold -- a build with the check skipped was byte-identical over 128 +tokens on prompts of 1,471, 3,841 and 10,352 tokens -- and the 1.04% deviation +an earlier revision attributed to them could not be reproduced. An +out-of-bounds read is a bug whatever it returns, so the call site passes +`false` for every GLM model; on all-valid selections the two variants perform +the same arithmetic in the same order, which the test asserts bit for bit, so +GLM 5.2 is unchanged. + +### The kernel that ships: the generic arithmetic, staged and shared + +The generic kernel's cost is structural, not arithmetic: one threadgroup per +head, each walking every selected row twice. The arithmetic can be kept to +the operation and reorganised around it. +`kernel_glm_attention_indexed_decode_exact_*` computes the same thing in four +phased dispatches: + +- **scores**: one thread per (head, row) running the generic kernel's + sequential 512-term dot, with 16 selected rows staged in threadgroup memory + per threadgroup for all 64 heads at once (1024 threads), so each cache row + is read from device memory once per token instead of 128 times; +- **weights**: one 256-thread threadgroup per head -- the generic kernel's + threadgroup -- running its per-thread row partition and its 128/64/../1 + reduction tree for the max and the denominator, and turning scores into + softmax weights in place; +- **lora**: one thread per (head, column pair) walking rows 0..n-1 in + selection order with the generic kernel's `acc += w * kv` chain, over 8 + heads x 64 columns per threadgroup so a row slice is loaded once for eight + heads. Rows are consumed in stages of 32: every thread fetches 16 bytes and + one weight three stages ahead into double-buffered threadgroup memory, so + the scattered row reads are in flight while the fma chains run. A row past + `cache_cap` contributes `fma(0, kv[0], acc)`, which leaves `acc` unchanged + bit for bit, where the generic kernel skips it; +- **value**: the generic kernel's quantised row dot from threadgroup memory, + one thread per output element, over 256 threadgroups instead of 64. + +Every floating-point operation, operand and ordering is the generic kernel's, +so the output is bit-identical to it, and that is asserted rather than +assumed: + +- `tests/test_glm53_kda` runs the exact kernels and the generic kernel on one + fixture at 8, 513, 1024, 2048, 2051 and 4096 selected rows, with rows at and + past `cache_cap` and `UINT32_MAX` sentinels in the selection, and requires + the outputs to match with `memcmp`; +- greedy generation from this tip is **byte-identical to main** (110afdd, and + b0a147a after the branch was rebased onto the synced main) over 128 + tokens on a 1,471-token prompt at ctx 4096 (dense window, every row valid), + a 3,841-token prompt at ctx 8192 (dense window, 3,841 rows) and a + 10,352-token prompt at ctx 16384 (pool selector, 2051 rows with sentinels). + No switches and no quality mode: the default path reproduces the base + commit. + +What the phases cost at about 1,500 selected rows, each measured by dropping +its dispatch and reading the change in decode time: + +| phase | ms/token | +|---|---:| +| scores | 0.33 | +| weights | 0.18 | +| lora, first version: row ids and weights read from device per row | 2.60 | +| lora, pipelined | 0.64 | +| value | 0.39 | + +The first lora version was slower than the generic kernel it replaced -- two +serialised device loads per row instead of one. Staging and prefetching is +what made the phase cheap; the arithmetic never changed. + +Decode from the CLI, greedy, 128 tokens, single runs (the interleaved +benchmark below is the figure to quote): + +| prompt | ctx | main | this tip | non-exact split kernel | +|---|---:|---:|---:|---:| +| 1,471 tokens | 4096 | 22.21 | **28.25** | 28.67 | +| 3,841 tokens | 8192 | 19.05 | **27.23** | 28.35 | +| 10,352 tokens | 16384 | 20.91 | **27.07** | 27.8 | + +The exact kernels give back nearly all of what the split kernel offered -- +within 1.5% at the short prompt -- while reproducing the base commit's output +bit for bit. Three more dispatches per DSA layer (four instead of one) account +for about 0.15 ms/token of the gap. + +`--quality` keeps the exact kernels, since they are exact; +`DS4_METAL_DISABLE_GLM53_DSA_EXACT` selects the generic kernel for A/B runs. +The two-host tensor-parallel head split keeps the generic kernel until that +configuration has been run. + +Where the phased path starts to pay, decode from the CLI at ctx 4096, 128 +greedy tokens, exact kernels against the generic kernel on the same prompt: + +| selected rows | generic | exact | delta | +|---:|---:|---:|---:| +| ~36 (a one-line chat prompt) | 29.04 | 28.93 | -0.4% | +| ~134 | 28.79 | 28.89 | +0.3% | +| ~207 | 28.51 | 28.83 | +1.1% | +| ~308 | 28.25 | 28.87 | +2.2% | +| ~603 | 27.52 | 28.71 | +4.3% | +| ~992 | 26.57 | 28.57 | +7.5% | +| ~1,500 | 25.51 | 28.25 | +10.7% | + +Below 128 rows the generic kernel's row traffic is a few megabytes per layer +and the three extra dispatches cost more than they save, so the exact path +engages from 128 selected rows. Both kernels are exact, so crossing the +threshold as a generation grows changes nothing but speed. + +## The shared-down fusion, after a second look + +An earlier revision of this document said this could not be done without a +dedicated shared-expert mid buffer, because `glm_graph_routed_moe_one_dispatch` +takes `ffn_mid` as scratch and would clobber the shared mid. + +That is true of only one of the two orderings. `shared_first` is +`streaming_selected_cache`, so the shared expert runs first *only* on the +SSD-streaming path. On the fully-resident path the routed stage has already +finished when the shared expert runs, `ffn_mid` still holds the shared mid and +`ffn_out` holds the routed result -- which is exactly what +`ds4_gpu_shared_down_hc_expand_q8_0_tensor` takes. No extra buffer. + +Fusing the shared down-projection, the routed add and the HC expand into that +one dispatch is worth **+0.77% (t=13.60), 0.320 ms over 42 sites** -- 7.6 us +per site, above both the 4.6 us launch cost and the 3.7 us the plain FFN-tail +fusion returned, because it also removes the `ffn_sum` round-trip. The +streaming path is excluded and keeps the separate dispatches. + +## Two tuning knobs that turn out not to matter + +Both were expected to be worth something on an 80-core GPU and neither is. + +**Decode command-buffer flush cadence.** Indexed decode flushes every 4 +layers, and `DS4_GLM_DECODE_FLUSH_INTERVAL` overrides it. Sweeping 0, 2, 3, 4, +6, 8, 12, 16, 32 at ctx 2048: everything from 3 to 12 lands in 22.25-22.31 +tok/s, inside the run-to-run spread. Only the extremes lose -- 0 (never flush) +at 21.91 and 32 at 22.06. Confirmed at ctx 16384, where 2/4/8 give +21.80/21.81/21.78. **The default of 4 is already right.** + +**DSA split-attention rows per block.** The choice steps straight from 32 to +128 at 1024 selected rows and had never been swept, so +`DS4_GLM_DECODE_SPLIT_BLOCK_ROWS` was added to force one value: + +| rows | ctx 2048 | ctx 16384 | +|---:|---:|---:| +| default (32/128) | 23.50 | 23.00 | +| 32 | 23.51 | 22.97 | +| 64 | 23.51 | 23.02 | +| 96 | -- | 23.01 | +| 128 | 23.51 | 23.02 | +| 256 | 23.49 | 23.02 | + +Flat to within 0.2% at both contexts. The selection count is capped by +`glm53_graph_indexer_selected_limit()`, which does not grow with context, so +this does not become interesting at longer contexts either. The knob is kept +as instrumentation for other GPUs, not because it found anything here. + +## Cumulative engine-only result + +Individual commits report gains against whatever baseline was current when they +landed, which does not compose into a branch number. This is the direct +measurement: the pre-series commit and the branch tip, each built in its own +tree so each reads its own `metal/*.metal`, run against the **same unchanged +GGUF** with the same harness, in the order main / branch / branch / main so +that drift lands on both arms alike. + + ds4-bench, promessi_sposi.txt, 128 greedy tokens per frontier, + frontiers 2048, 4096, 8192, 16384; four runs, main / branch / branch / main + +| frontier | main prefill | branch prefill | prefill | main decode | branch decode | decode | +|---:|---:|---:|---:|---:|---:|---:| +| 2048 | 429.98 / 429.52 | 429.71 / 429.89 | +0.01% | 21.09 / 21.12 | 27.82 / 27.84 | **+31.86%** | +| 4096 | 390.14 / 390.08 | 389.93 / 390.03 | -0.03% | 20.75 / 20.76 | 27.12 / 27.13 | **+30.69%** | +| 8192 | 392.23 / 392.08 | 392.01 / 392.00 | -0.04% | 20.71 / 20.69 | 26.99 / 27.04 | **+30.51%** | +| 16384 | 389.49 / 389.54 | 389.33 / 389.45 | -0.03% | 20.65 / 20.58 | 26.88 / 26.98 | **+30.63%** | + +Both runs of each arm are shown; the deltas compare the means. Prefill is +untouched by this branch's decode work and measures as such. At ctx 2048 the +base arm reproduces the 21.16 tok/s measured at the start of this series, so +machine conditions have not drifted. + +Repeated after the branch was rebased onto the synced main (b0a147a, 24 +upstream commits of Metal tensor-parallel and DSpark work), same protocol: + +| frontier | main prefill | branch prefill | prefill | main decode | branch decode | decode | +|---:|---:|---:|---:|---:|---:|---:| +| 2048 | 429.64 / 430.08 | 429.73 / 429.81 | -0.02% | 20.97 / 21.06 | 27.76 / 27.76 | **+32.10%** | +| 4096 | 390.20 / 390.55 | 390.10 / 390.14 | -0.07% | 20.67 / 20.73 | 27.06 / 27.05 | **+30.70%** | +| 8192 | 392.30 / 392.64 | 392.08 / 391.52 | -0.17% | 20.63 / 20.69 | 26.94 / 26.96 | **+30.45%** | +| 16384 | 389.56 / 389.97 | 389.43 / 389.52 | -0.07% | 20.55 / 20.59 | 26.93 / 26.81 | **+30.63%** | + +The synced main decodes GLM 5.3 Flash at the same rate as 110afdd did and +produces the same greedy output on every prompt used here, so upstream's +changes did not touch this path; the rebase itself conflicted only in the +Makefile's test list and in the mHC producer kernel, which upstream had +refactored into a shared body that the branch now templates. + +Contributions, each measured against the baseline current when it landed: the +mHC producer fusion +5.67%, the KDA gate pairing +0.74%, the three HC-expand +epilogues +0.46% / +0.11% / +0.14%, the shared-down/HC fusion +0.77%, the gate +trio +0.30%, and the exact phased DSA kernels (see above). The widened BF16 +loads (~+5.4%) and the split DSA kernel for GLM 5.3 (+16.86%) were measured on +the way and are not on this branch's default path, for the reason in the next +section. + +Note the base reproduces the 21.19 tok/s of the original budget almost exactly, +which is a useful check that machine conditions have not drifted between the +first measurements in this document and the last. + +Stacking the model-artifact changes on the engine, all at ctx 2048. **This +table predates the exact DSA kernels**: it was taken at d5b7895, when the +engine-only tip measured 23.99 tok/s, and has not been re-measured since, so +its rows are not comparable with the figure above. What it still shows is the +artifact effect on top of one engine state: + +| model file | tok/s at d5b7895 | vs base engine + original artifact | +|---|---:|---:| +| GLM-5.3-Flash-Q4_K | 23.99 | +13.2% | +| GLM-5.3-Flash-Q4_K-kdaQ8 | 27.50 | +29.8% | +| GLM-5.3-Flash-Q4_K-kdaHeadQ8 | 28.23 | +33.2% | + +Only the first row is an engine result. The other two combine it with the +requantized artifacts and should never be quoted as engine tuning. + +## Nothing on the default path is left that is not bit-exact + +Two changes made on the way here were deterministic but not bit-identical to +the paths they replaced. Both are off this branch's default path: + +- **The widened BF16 matvec loads** (about +5.4%) repartitioned which lane + accumulates which k. An exact wide variant would have to redistribute every + lane's strided elements with cross-lane shuffles, two per element, which + costs roughly what the widening saved, so the scalar accumulation -- the + pre-branch kernel -- is the only path. The fused qkv/pair/trio/HC-expand + kernels share its row helper unchanged, so they stay exact. +- **The grouped/split DSA kernel** (+16.86%) is replaced for GLM 5.3 by the + exact phased kernels. GLM 5.2 keeps it as before, with `--quality` + selecting the generic kernel there. + +### Checked end to end against the base commit + +Greedy generation (`--raw-prompt --temp 0`, 128 tokens), the tip and main +(110afdd when first measured; repeated against b0a147a after the rebase, with +the same result and the same main output on every prompt) +each built in its own worktree, byte-compared, both in default mode: + +| prompt | ctx | selection | result | +|---|---:|---|---| +| 1,471 tokens | 4096 | dense, 1,472+ rows | **byte-identical** | +| 3,841 tokens | 8192 | dense, 3,842+ rows | **byte-identical** | +| 10,352 tokens | 16384 | pool top-k, 2051 rows with sentinels | **byte-identical** | + +Encoder counts confirm the exact path ran: it adds three dispatches per DSA +layer per token, 33 x 127 = 4,191 more acquisitions than the generic arm, +which is itself byte-identical to the base. Under `--ssd-streaming`, same +prompt, 32 tokens, the exact and generic arms are byte-identical to each +other and to the resident run. Tensor parallelism was not run; the exact +kernels stay off there. + +## Other models: nothing broke, and one thing had slowed + +Two models that take none of the GLM 5.3 Flash paths were run through the +full test suite and the same main / branch / branch / main `ds4-bench` +protocol, and byte-compared on greedy generation (128 tokens, 1,471- and +3,841-token prompts): + +- **DeepSeek V4 Flash** (`MXFP4Experts-F16HC-...-chat-v2-mxfp4-0731`): every + suite OK on the branch; output byte-identical to main; prefill and decode + within 0.25% of main at every frontier, in both directions. The only + shared code it touches is the templated HC producer, whose f16 + instantiation is the kernel it always ran. + + | frontier | main decode | branch decode | decode | prefill | + |---:|---:|---:|---:|---:| + | 2048 | 42.67 / 42.58 | 42.61 / 42.54 | -0.12% | -0.01% | + | 4096 | 38.87 / 38.73 | 38.69 / 38.72 | -0.24% | -0.04% | + | 8192 | 38.19 / 38.26 | 38.17 / 38.30 | +0.03% | +0.05% | + | 16384 | 37.30 / 37.40 | 37.31 / 37.45 | +0.08% | -0.13% | + +- **GLM 5.3 (`glm-dsa`, `UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K`)**: the full + model, 79 layers with a 64-wide RoPE tail, so it takes the GLM 5.2 path + and the split DSA kernel, not the exact kernels. Output byte-identical to + main. Five suites fail on the branch -- and fail identically on main, with + the same 55 assertions and the same golden-vector statistics: they are + DeepSeek official-vector fixtures and a 30k-token recall test this quant + does not pass on either tree. The benchmark found a real regression: + + | frontier | main decode | branch decode | decode | prefill | + |---:|---:|---:|---:|---:| + | 2048 | 16.32 / 16.26 | 15.96 / 15.95 | **-2.06%** | -0.01% | + | 4096 | 16.24 / 16.22 | 15.91 / 15.90 | **-2.00%** | +0.04% | + | 8192 | 16.02 / 16.03 | 15.68 / 15.70 | **-2.09%** | +0.04% | + | 16384 | 15.64 / 15.62 | 15.32 / 15.33 | **-1.95%** | +0.07% | + + The cause is the split kernel's bounds-checked variant, which the branch had + switched every GLM model to. On Flash it cost 0.24% of decode, with DSA + attention in 11 of 45 layers; here the split kernel runs in 76 of 79 layers + and the same per-call cost is 2% of the step. This model's selections are a + dense range or a top-k over visible rows, always in range, so it goes back + to the unchecked variant it always ran -- main's kernel, bit for bit, as + the all-valid equivalence case in `tests/test_glm53_kda` asserts -- and + only a GLM 5.3 Flash graph, which pads with sentinels, would pass `false` + should it ever reach that call. Re-measured with that change: + + | frontier | main decode | branch decode | decode | prefill | + |---:|---:|---:|---:|---:| + | 2048 | 16.42 / 16.33 | 16.32 / 16.30 | -0.40% | +0.03% | + | 4096 | 16.32 / 16.21 | 16.20 / 16.19 | -0.43% | +0.08% | + | 8192 | 16.01 / 16.01 | 16.01 / 15.93 | -0.25% | +0.01% | + | 16384 | 15.66 / 15.66 | 15.64 / 15.62 | -0.19% | +0.04% | + + What remains is inside main's own run-to-run spread (its two ctx 2048 runs + differ by 0.55%); if any of it is real it is at most 0.4%, against the 2% + before the change. Output stays byte-identical to main. + +## Rollback switches, and what PR #954 does that this branch could use + +antirez/ds4#954 (pre-M5 DeepSeek decode and prefill, bit-exact) puts every +optimisation behind its own `DS4_..._DISABLE_...` switch with an aggregate +that turns the whole set off, and measures each against its rollback. This +branch has the same shape now. Each switch restores the pre-branch path for +one change, and `DS4_METAL_DISABLE_GLM53_FLASH_TUNING` restores all of them: + +| switch | restores | +|---|---| +| `DS4_METAL_DISABLE_GLM53_HC_PRODUCER_FUSE` | four dispatches per mHC producer site instead of the fused BF16 kernel | +| `DS4_METAL_DISABLE_GLM53_KDA_GATE_PAIR` | separate f_a / g_a and f_b / g_b projections | +| `DS4_METAL_DISABLE_GLM53_KDA_GATE_TRIO` | beta as its own projection beside the pair | +| `DS4_METAL_DISABLE_GLM53_KDA_OUT_HC_EXPAND` | a separate HC expand after kda_output | +| `DS4_METAL_DISABLE_GLM53_ATTN_OUT_HC_EXPAND` | a separate HC expand after attn_output | +| `DS4_METAL_DISABLE_GLM53_FFN_HC_EXPAND_ADD` | a separate routed+shared add and HC expand in the FFN tail | +| `DS4_METAL_DISABLE_GLM53_SHARED_DOWN_HC_EXPAND` | the shared down-projection without the routed add and expand | +| `DS4_METAL_DISABLE_GLM53_DSA_EXACT` | the generic DSA attention kernel | +| `DS4_METAL_DISABLE_GLM53_FLASH_TUNING` | every path above at once | + +Not switchable: the KDA decay hoist, a kernel-internal cleanup that is +bit-identical (the KDA prefill/decode consistency test) and worth nothing +measurable, and the prefill constants, which are knobs with their defaults +unchanged. `DS4_METAL_DISABLE_GLM53_DSA_SPLIT` belongs to the GLM 5.2 path. + +With the aggregate set, greedy output is byte-identical to main and decodes the 1,471-token prompt at 22.28 tok/s against main's 22.21, with 146,046 encoder acquisitions over the run against 86,610 on the default path -- the unfused dispatch structure is back, so the switch restores the paths and not just the numbers. + +Two of #954's pieces could in principle apply to GLM 5.3 Flash; neither +does in practice: + +- **Greedy chain decode** keeps the token id on the GPU so the host's + `waitUntilCompleted`, logits readback, argmax and re-encode leave the + per-token critical path; #954 measures the boundary at about 0.5 ms of GPU + idle per DeepSeek token and gains 1.75%. Here `DS4_METAL_GPU_BUSY_PROFILE` + over 16 decode tokens accumulates 35.2 ms of GPU time per 35.3 ms token: + the GLM decode loop already flushes command buffers every four layers, so + the GPU idles about 0.1 ms per token, a 0.3% ceiling. Not worth the + device-resident token ring, GPU argmax and session plumbing it takes. +- **Batch indexer-query pruning** skips the indexer query projection, RoPE, + QAT and weight projection for prefill batches whose attention is entirely + within the dense window; #954 gains 1.3-1.8% of prefill. GLM's indexed + prefill already does this: `use_causal_range_select` is true while the + chunk's rows fit the 4096-row window, and the query projection is inside + `if (!use_causal_range_select)`. + +The rest of #954 is DeepSeek attention and MoE kernels (raw-layer gathered +attention, packed32, RB4-staged prefill rows, sum6/attn-out HC fusions) with +no GLM 5.3 Flash counterpart on the same shapes. + +## A trap when verifying a decode-path change + +`ds4-bench --dump-frontier-logits-dir` writes one file per **frontier**, which +is the logits at the end of prefill. It does not exercise the single-token +decode graph at all. + +This was found the hard way. A change that skipped the FFN-side mHC producer +on every KDA layer -- catastrophic, garbage output after the first token -- +produced frontier logits **bit-identical** to the baseline, because the bug was +entirely in the decode path the dump never touches. A four-token greedy +generation caught it immediately. + +For anything that touches decode, compare **greedy generations** instead: fixed +prompts, `--temp 0`, 128 tokens, byte-compared. Decode is deterministic across +runs (verified), and any bit difference diverges within a few tokens. The +frontier dump is still the right tool for a prefill-path change. + +## A trap when A/B-testing a shader change + +`ds4_gpu_full_source()` reads `metal/*.metal` from disk at run time and there +is no embedded fallback, so building two binaries around a shader edit does +**not** compare two shaders -- both read whatever is on disk when they run. +Use the per-file overrides (`DS4_METAL_GLM53_KDA_SOURCE` and its siblings) with +a single binary instead. A measurement in this file was wrong for exactly +this reason before it was caught. + +## Scope and caveats + +- This is a **model-file** change, not an engine change. It does not speed up + an artifact you already have; it produces a better one. +- Why the shipped artifact is BF16 is not established here. It contradicts the + repo's own quantizer, which suggests the artifact pipeline rather than a + deliberate choice, but if it was deliberate the fix belongs upstream. +- Quality evidence is one perplexity run on one text plus a greedy generation. + That is good evidence for a near-lossless type like Q8_0, not proof. +- `--artifact q4` also specifies Q8_0 for the embedding and output tensors, + which are BF16 here too (~1.2 GiB more per token through the LM head). This + was subsequently done: `--tensors head,embd` covers them, and converting the + head is worth a further +1.80% decode over a KDA-only artifact. `token_embd` + is deliberately not in the default -- it is a single-row lookup per token, so + it saves resident memory rather than decode bandwidth. +- Neither the constants recorded below nor a Q4_K KDA variant were measured. + The tool accepts `q4_K` as a target, which would take KDA to 2.39 GiB, but + Q4_K on attention projections is a materially bigger quality question than + Q8_0 and was not attempted. + +## Prefill, measured + +Prefill had never been swept. Five constants that shape it were compile-time +`#define`s with no override, and two of them interact, so each is now +separately settable -- `DS4_GLM_PREFILL_CHUNK_TOKENS`, +`DS4_GLM_FULL_ATTN_LAYER_FLUSH_TOKENS`, `DS4_GLM_FULL_ATTN_CAP`, +`DS4_GLM_FULL_ATTN_STREAMING_CAP`, `DS4_GLM_PREFILL_SCORE_SCRATCH_MB`. +Defaults are unchanged: logits with the knobs unset and with them set to the +old constants match at max|delta| = 0. + +### Chunk size, with layer flushing held constant + +The document previously warned that the chunk (2048) and the layer-flush +threshold (2048) are the same number against a strict `>`, so raising the chunk +also switches per-layer flushing on -- two changes, not one. Pinning the flush +threshold separates them. Prefill tok/s at ctx 16384: + +| chunk | flush off | flush on | default | +|---:|---:|---:|---:| +| 1024 | 354.40 | 355.00 | 354.34 | +| 2048 | 394.30 | 394.57 | 394.17 | +| 4096 | 393.98 | 394.45 | 394.12 | +| 8192 | 394.05 | 394.43 | 393.97 | + +Two results. **Layer flushing does not matter at all** -- every column agrees +to 0.2%, so the coupling the doc warned about is real in the code and +immaterial in practice. And **the default chunk of 2048 is already optimal**: +1024 costs 10%, while 4096 and 8192 buy nothing. Confirmed at ctx 32768, where +2048/4096/8192 give 388.68/388.31/388.50. + +Raising the chunk is not free elsewhere, either. Context buffers at ctx 4096 +grow 1.62 -> 3.04 -> 5.88 GiB across chunk 1024/2048/4096, so 4096 would cost +nearly 2 GiB for no throughput. The GLM 5.2 path's 4096 is not an argument for +changing this one. + +### The full-attention cap asymmetry is backwards + +`glm_graph_full_attention_cap` gives the SSD-streaming path 8192 and the +fully-resident path 4096, which looked like the memory-constrained machine +getting the larger window. Forcing each value on the resident path, ctx 16384, +interleaved, n=6: + +| cap | prefill | decode | +|---:|---:|---:| +| 4096 | 394.23 (sd 0.03) | 23.17 (sd 0.02) | +| 8192 | 379.23 (sd 0.10) | 23.15 (sd 0.02) | + +**The larger window costs 3.81% of prefill and nothing on decode.** So 4096 is +not the conservative choice, it is the fast one, and the resident default is +right. Whether 8192 pays for itself on the streaming path by reducing +re-streaming is untested here -- `DS4_GLM_FULL_ATTN_STREAMING_CAP` exists to +try it. + +### The 256 MiB score scratch does bind, but not where it hurts yet + +`DS4_GLM_METAL_INDEXED_PREFILL_SCORE_SCRATCH_MB` clamps rows scored per +dispatch. It is not dead: score columns are `compact_cap / 4`, so the budget +starts biting above 131072 allocated context. Observed, chunk 2048 throughout: + +| ctx_alloc | score_rows | scratch | +|---:|---:|---:| +| 65536 | 2048 | 128 MiB | +| 131072 | 2048 | 256 MiB | +| 262144 | **1024** | 256 MiB | +| 524288 | **512** | 256 MiB | + +Raising the budget to 1024 MiB restores 2048 rows at a 524288 allocation. The +model context limit is 1048576, so this is reachable, not theoretical. + +It does not currently cost anything measurable, though: holding the allocation +at 524288 and varying only the budget, a 16384-token prefill runs at 393.84 +tok/s with score_rows=512 against 394.30 with 2048 -- **0.12%, noise**. A +prefill long enough for scoring to dominate was not measured; each run at ctx +65536 with that allocation exceeds ten minutes. So: the clamp is real, the +knob to lift it exists, and nobody has yet shown it matters. + +## Untested constants noticed while reading + +Recorded so the next person does not re-derive them. None were measured. + +All four prefill entries that used to sit here have been measured; see +"Prefill, measured" above. In summary: the chunk default of 2048 is optimal, +per-layer flushing does not matter, the 8192 full-attention cap is 3.81% slower +than 4096 rather than more generous, and the 256 MiB score scratch does clamp +above 131072 allocated context but costs nothing measurable at the prefill +lengths tested. diff --git a/tests/ds4_test.c b/tests/ds4_test.c index cf8bca2c5..19d5c15f8 100644 --- a/tests/ds4_test.c +++ b/tests/ds4_test.c @@ -5385,7 +5385,8 @@ static bool test_logprob_vector_case_disabled(const char *path, static void test_official_logprob_vectors_run(const char *case_filter) { const char *path = getenv("DS4_TEST_VECTOR_FILE"); - if (!path || !path[0]) { + const bool default_fixture = !path || !path[0]; + if (default_fixture) { path = "tests/test-vectors/flash-0731/official.vec"; } FILE *fp = fopen(path, "rb"); @@ -5413,7 +5414,15 @@ static void test_official_logprob_vectors_run(const char *case_filter) { test_vec_case vc; int ran = 0; - while (test_read_vector_case(fp, &vc)) { + /* The default vectors describe DeepSeek V4 Flash, not GLM. An explicitly + * selected fixture remains available for model-specific comparisons. */ + const bool compatible_fixture = !default_fixture || !ds4_engine_is_glm_dsa(engine); + if (!compatible_fixture) { + fprintf(stderr, "ds4-test: DeepSeek API vectors skipped for %s; " + "set DS4_TEST_VECTOR_FILE to a matching fixture\n", + ds4_engine_model_name(engine)); + } + while (compatible_fixture && test_read_vector_case(fp, &vc)) { if (!test_fill_vector_case(fp, &vc)) break; if (case_filter && case_filter[0] && strcmp(vc.id, case_filter)) { continue; @@ -5427,7 +5436,7 @@ static void test_official_logprob_vectors_run(const char *case_filter) { test_logprob_vector_case(engine, &vc); ran++; } - TEST_ASSERT(!case_filter || !case_filter[0] || ran == 1); + TEST_ASSERT(!compatible_fixture || !case_filter || !case_filter[0] || ran == 1); ds4_engine_close(engine); test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill); test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4); @@ -5569,12 +5578,13 @@ static int test_local_golden_overlap(const test_local_golden_case *tc, static float test_local_golden_max_abs(const test_local_golden_case *tc, const float *cand_logits, + int vocab, int n) { float max_abs = 0.0f; if (n > tc->ntop) n = tc->ntop; for (int i = 0; i < n; i++) { const int id = tc->top[i].id; - if (id < 0) continue; + if (id < 0 || id >= vocab) return FLT_MAX; const float abs_delta = fabsf(cand_logits[id] - tc->top[i].logit); if (abs_delta > max_abs) max_abs = abs_delta; } @@ -5634,7 +5644,7 @@ static void test_local_golden_case_run(ds4_engine *engine, const int top20_overlap = test_local_golden_overlap(tc, cand_top, 20); const int top64_overlap = test_local_golden_overlap(tc, cand_top, 64); const float top20_max_abs = - test_local_golden_max_abs(tc, cand_logits, 20); + test_local_golden_max_abs(tc, cand_logits, vocab, 20); fprintf(stderr, "ds4-test: local golden %s top1 ref=%d cand=%d " @@ -5664,7 +5674,8 @@ static void test_local_golden_case_run(ds4_engine *engine, static void test_local_golden_vectors(void) { const char *path = getenv("DS4_TEST_LOCAL_GOLDEN_FILE"); - if (!path || !path[0]) { + const bool default_fixture = !path || !path[0]; + if (default_fixture) { path = "tests/test-vectors/flash-0731/local-golden.vec"; } FILE *fp = fopen(path, "rb"); @@ -5691,7 +5702,13 @@ static void test_local_golden_vectors(void) { } test_local_golden_case tc; - while (test_read_local_golden_case(fp, &tc)) { + const bool compatible_fixture = !default_fixture || !ds4_engine_is_glm_dsa(engine); + if (!compatible_fixture) { + fprintf(stderr, "ds4-test: DeepSeek local golden vectors skipped for %s; " + "set DS4_TEST_LOCAL_GOLDEN_FILE to a matching fixture\n", + ds4_engine_model_name(engine)); + } + while (compatible_fixture && test_read_local_golden_case(fp, &tc)) { if (!test_fill_local_golden_case(fp, &tc)) break; test_local_golden_case_run(engine, &tc); } @@ -6843,6 +6860,7 @@ static void test_print_help(const char *prog) { puts(" DS4_TEST_LONG_PROMPT=FILE Rendered long-context story fact prompt."); puts(" DS4_TEST_VECTOR_FILE=FILE Official fixture. Default: flash-0731/official.vec."); puts(" DS4_TEST_LOCAL_GOLDEN_FILE=FILE Local fixture. Default: flash-0731/local-golden.vec."); + puts(" DeepSeek default fixtures are skipped for GLM; explicit fixtures are always checked."); puts(" DS4_TEST_MPP_EQ_CASE=NAME Run only Tensor equivalence cases whose id contains NAME."); puts(" DS4_TEST_MTP=FILE Legacy MTP support GGUF for --mtp-verify-depth."); puts(" DS4_TEST_DSPARK=FILE DSpark support GGUF for --dspark-verify-depth."); diff --git a/tests/test_glm53_kda.c b/tests/test_glm53_kda.c index 937d6f93d..2b6959e6b 100644 --- a/tests/test_glm53_kda.c +++ b/tests/test_glm53_kda.c @@ -72,6 +72,1116 @@ static float bf16_to_f32(uint16_t value) { return bits.f; } +#ifdef __APPLE__ +/* Normal-range only, and truncating rather than rounding. Both are fine here: + * the compound-producer fixture uses values with at most seven explicit + * mantissa bits, well inside the half normal range, so truncation to ten bits + * is exact and the encoding round-trips. */ +static uint16_t f32_to_f16(float value) { + union { float f; uint32_t u; } b = { .f = value }; + const uint32_t sign = (b.u >> 16) & 0x8000u; + const int32_t exp = (int32_t)((b.u >> 23) & 0xffu) - 127 + 15; + const uint32_t mant = (b.u >> 13) & 0x3ffu; + if (exp <= 0 || exp >= 31) return (uint16_t)sign; + return (uint16_t)(sign | ((uint32_t)exp << 10) | mant); +} + +static float f16_to_f32(uint16_t value) { + const uint32_t sign = (uint32_t)(value & 0x8000u) << 16; + const uint32_t exp = (uint32_t)(value >> 10) & 0x1fu; + const uint32_t mant = (uint32_t)value & 0x3ffu; + union { uint32_t u; float f; } b; + b.u = exp == 0 ? sign : (sign | ((exp - 15u + 127u) << 23) | (mant << 13)); + return b.f; +} +#endif + +/* Exercises ds4_gpu_glm53_matmul_bf16 at one width. The reference is + * accumulated in double and compared with a relative tolerance; a stride or + * indexing error moves a result far more than that, which is what this is + * here to catch. */ +static void check_bf16_matmul(const uint8_t *model, size_t model_bytes, + uint64_t offset, uint32_t in_dim, + uint32_t out_dim, uint32_t rows, + const char *what) { + uint16_t *w = (uint16_t *)(void *)((uint8_t *)(uintptr_t)model + offset); + for (uint32_t o = 0; o < out_dim; o++) { + for (uint32_t i = 0; i < in_dim; i++) { + w[(size_t)o * in_dim + i] = f32_to_bf16( + 0.002f * (float)((int)(o % 11u) - 5) + + 0.001f * (float)((int)(i % 13u) - 6)); + } + } + const size_t x_bytes = (size_t)rows * in_dim * sizeof(float); + const size_t out_bytes = (size_t)rows * out_dim * sizeof(float); + float *x = malloc(x_bytes); + float *expected = malloc(out_bytes); + float *actual = malloc(out_bytes); + require_ok(x && expected && actual, "wide BF16 host allocation"); + for (uint32_t r = 0; r < rows; r++) { + for (uint32_t i = 0; i < in_dim; i++) { + x[(size_t)r * in_dim + i] = + 0.02f * (float)((int)(i % 17u) - 8) + 0.005f * (float)r; + } + for (uint32_t o = 0; o < out_dim; o++) { + double sum = 0.0; + for (uint32_t i = 0; i < in_dim; i++) { + sum += (double)bf16_to_f32(w[(size_t)o * in_dim + i]) * + (double)x[(size_t)r * in_dim + i]; + } + expected[(size_t)r * out_dim + o] = (float)sum; + } + } + ds4_gpu_tensor *gx = ds4_gpu_tensor_alloc(x_bytes); + ds4_gpu_tensor *gout = ds4_gpu_tensor_alloc(out_bytes); + require_ok(gx && gout, "wide BF16 tensor allocation"); + require_ok(ds4_gpu_tensor_write(gx, 0, x, x_bytes), "wide BF16 input write"); + + require_ok(ds4_gpu_glm53_matmul_bf16(gout, model, model_bytes, offset, + in_dim, out_dim, gx, 1), what); + require_ok(ds4_gpu_tensor_read(gout, 0, actual, out_dim * sizeof(float)), + "wide BF16 decode output read"); + for (uint32_t o = 0; o < out_dim; o++) { + require_close(what, actual[o], expected[o], + 2e-5f * (fabsf(expected[o]) + 1.0f)); + } + + require_ok(ds4_gpu_glm53_matmul_bf16(gout, model, model_bytes, offset, + in_dim, out_dim, gx, rows), what); + require_ok(ds4_gpu_tensor_read(gout, 0, actual, out_bytes), + "wide BF16 prefill output read"); + for (uint32_t i = 0; i < rows * out_dim; i++) { + require_close(what, actual[i], expected[i], + 2e-5f * (fabsf(expected[i]) + 1.0f)); + } + ds4_gpu_tensor_free(gx); + ds4_gpu_tensor_free(gout); + free(x); + free(expected); + free(actual); +} + +#ifdef __APPLE__ +/* + * Split-versus-generic indexed decode attention. + * + * GLM 5.3 decode runs kernel_glm_attention_indexed_decode_split_group8 once + * more than 512 rows are selected; the generic kernel is what --quality and + * every other backend run. The two score and reduce in different orders, so + * each is checked against a double-precision reference and they are checked + * against each other with a tolerance rather than bit for bit. + * + * The selection holds what the GLM 5.3 indexer actually emits: rows at and + * past cache_cap and UINT32_MAX tail sentinels, which both kernels must + * exclude. The rows just past cache_cap exist in memory and hold values that + * would dominate every softmax, so a kernel that skips the bounds test fails + * this loudly rather than by luck. The row counts cover one partial block, + * the 17-, 32-, 16- and 33-block reductions decode can request, the + * fixed-count 16-block reduce, and a 65-block request the wrapper must refuse. + */ +static void check_split_dsa_attention(uint8_t *model, size_t model_bytes, + uint64_t value_offset) { + enum { + SA_HEADS = 16, + SA_LORA = 512, + SA_NOPE = 64, + SA_VALUE = 8, + SA_MAX_SELECTED = 4096, + SA_CAP = 4163, /* > SA_MAX_SELECTED and coprime with 7919 */ + SA_POISON_ROWS = 16, /* allocated past cache_cap, never to be read */ + SA_ROWS = SA_CAP + SA_POISON_ROWS, + SA_MAX_BLOCKS = 65, + SA_Q8_ROW_BYTES = (SA_LORA / 32) * 34, + }; + static const struct { + uint32_t n_selected; + uint32_t block_rows; + bool accepted; + } cases[] = { + {8, 32, true}, /* one partial block */ + {513, 32, true}, /* 17 blocks: the first count decode splits */ + {1024, 32, true}, /* 32 blocks */ + {2048, 128, true}, /* 16 blocks: the fixed-count reduce */ + {2051, 128, true}, /* 17 blocks: GLM 5.3's selection limit */ + {2051, 64, true}, /* 33 blocks */ + {4096, 128, true}, /* 32 blocks: the resident dense window */ + {2051, 32, false}, /* 65 blocks: more than the reduce walks */ + }; + + /* Scores need a spread of tens, not a flat softmax, or the running-max + * rescale in the split kernel is never exercised. Each row and head + * carries a multiple of one shared basis vector plus small noise, so + * scores land in about [-17, 17] with many near-maximal rows. */ + float base[SA_LORA]; + for (uint32_t j = 0; j < SA_LORA; j++) { + base[j] = (float)((int)((j * 13u) % 17u) - 8) / 8.0f; + } + uint16_t *kv_bits = malloc((size_t)SA_ROWS * SA_LORA * sizeof(*kv_bits)); + float *kv = malloc((size_t)SA_ROWS * SA_LORA * sizeof(*kv)); + float *low = malloc((size_t)SA_HEADS * SA_LORA * sizeof(*low)); + float *q = calloc((size_t)SA_HEADS * SA_NOPE, sizeof(*q)); + uint32_t *sel = malloc((size_t)SA_MAX_SELECTED * sizeof(*sel)); + double *ref = malloc((size_t)SA_HEADS * SA_VALUE * sizeof(*ref)); + double *lora = malloc((size_t)SA_LORA * sizeof(*lora)); + float *gen = malloc((size_t)SA_HEADS * SA_VALUE * sizeof(*gen)); + float *spl = malloc((size_t)SA_HEADS * SA_VALUE * sizeof(*spl)); + float *spl2 = malloc((size_t)SA_HEADS * SA_VALUE * sizeof(*spl2)); + float *exact = malloc((size_t)SA_HEADS * SA_VALUE * sizeof(*exact)); + require_ok(kv_bits && kv && low && q && sel && ref && lora && + gen && spl && spl2 && exact, "split attention host allocation"); + for (uint32_t row = 0; row < SA_ROWS; row++) { + const float a = row < SA_CAP + ? (float)((int)(row % 23u) - 11) / 22.0f + : 8.0f; /* poison: would dominate any softmax it leaked into */ + for (uint32_t j = 0; j < SA_LORA; j++) { + const float noise = row < SA_CAP + ? (float)((int)((row * 7u + j * 3u + (row ^ j)) % 97u) - 48) / 256.0f + : 0.0f; + const uint16_t bits = f32_to_f16(a * base[j] + noise); + kv_bits[(size_t)row * SA_LORA + j] = bits; + kv[(size_t)row * SA_LORA + j] = f16_to_f32(bits); + } + } + for (uint32_t h = 0; h < SA_HEADS; h++) { + for (uint32_t j = 0; j < SA_LORA; j++) { + low[(size_t)h * SA_LORA + j] = + (0.5f + (float)h / 16.0f) * base[j] + + (float)((int)((h * 11u + j * 5u) % 61u) - 30) / 240.0f; + } + } + /* Q8_0 value rows with unit scales, so a dequantized weight is its int8. */ + require_ok(value_offset + (uint64_t)SA_HEADS * SA_VALUE * SA_Q8_ROW_BYTES <= model_bytes, + "split attention value rows fit the fixture model"); + for (uint32_t h = 0; h < SA_HEADS; h++) { + for (uint32_t d = 0; d < SA_VALUE; d++) { + uint8_t *row = model + value_offset + + (size_t)(h * SA_VALUE + d) * SA_Q8_ROW_BYTES; + for (uint32_t b = 0; b < SA_LORA / 32u; b++) { + const uint16_t one = 0x3c00u; + memcpy(row + b * 34u, &one, sizeof(one)); + int8_t *qs = (int8_t *)(row + b * 34u + 2u); + for (uint32_t i = 0; i < 32u; i++) { + qs[i] = (int8_t)((int)((h * 5u + d * 3u + (b * 32u + i) * 7u) % 15u) - 7); + } + } + } + } + + ds4_gpu_tensor *heads_gpu = ds4_gpu_tensor_alloc((uint64_t)SA_HEADS * SA_VALUE * sizeof(float)); + ds4_gpu_tensor *partial_lora_gpu = ds4_gpu_tensor_alloc( + (uint64_t)SA_MAX_BLOCKS * SA_HEADS * SA_LORA * sizeof(float)); + ds4_gpu_tensor *partial_ms_gpu = ds4_gpu_tensor_alloc( + (uint64_t)SA_MAX_BLOCKS * SA_HEADS * 2u * sizeof(float)); + ds4_gpu_tensor *q_gpu = ds4_gpu_tensor_alloc((uint64_t)SA_HEADS * SA_NOPE * sizeof(float)); + ds4_gpu_tensor *low_gpu = ds4_gpu_tensor_alloc((uint64_t)SA_HEADS * SA_LORA * sizeof(float)); + ds4_gpu_tensor *kv_gpu = ds4_gpu_tensor_alloc((uint64_t)SA_ROWS * SA_LORA * sizeof(uint16_t)); + ds4_gpu_tensor *rope_gpu = ds4_gpu_tensor_alloc(sizeof(float)); + ds4_gpu_tensor *sel_gpu = ds4_gpu_tensor_alloc((uint64_t)SA_MAX_SELECTED * sizeof(uint32_t)); + ds4_gpu_tensor *exact_scores_gpu = ds4_gpu_tensor_alloc( + (uint64_t)SA_HEADS * SA_MAX_SELECTED * sizeof(float)); + ds4_gpu_tensor *exact_lora_gpu = ds4_gpu_tensor_alloc((uint64_t)SA_HEADS * SA_LORA * sizeof(float)); + ds4_gpu_tensor *exact_denom_gpu = ds4_gpu_tensor_alloc((uint64_t)SA_HEADS * sizeof(float)); + require_ok(heads_gpu && partial_lora_gpu && partial_ms_gpu && q_gpu && + low_gpu && kv_gpu && rope_gpu && sel_gpu && exact_scores_gpu && + exact_lora_gpu && exact_denom_gpu, + "split attention GPU allocation"); + require_ok(ds4_gpu_tensor_write(q_gpu, 0, q, (uint64_t)SA_HEADS * SA_NOPE * sizeof(float)) && + ds4_gpu_tensor_write(low_gpu, 0, low, (uint64_t)SA_HEADS * SA_LORA * sizeof(float)) && + ds4_gpu_tensor_write(kv_gpu, 0, kv_bits, (uint64_t)SA_ROWS * SA_LORA * sizeof(uint16_t)), + "split attention input write"); + + for (size_t c = 0; c < sizeof(cases) / sizeof(cases[0]); c++) { + const uint32_t n = cases[c].n_selected; + const uint32_t block_rows = cases[c].block_rows; + const uint32_t n_blocks = (n + block_rows - 1u) / block_rows; + char what[96]; + snprintf(what, sizeof(what), "split attention %u rows x %u/block", + n, block_rows); + + /* A permutation of valid rows, with the indexer's failure shapes + * scattered through it: rows at and just past cache_cap, and the + * UINT32_MAX tail sentinels GLM 5.3's pool expansion emits. */ + for (uint32_t s = 0; s < n; s++) sel[s] = (s * 7919u) % SA_CAP; + sel[0] = SA_CAP; + sel[1] = SA_CAP - 1u; + for (uint32_t s = 50; s < n; s += 97u) sel[s] = SA_CAP + s % 5u; + for (uint32_t s = n >= 3u ? n - 3u : 0u; s < n; s++) sel[s] = UINT32_MAX; + require_ok(ds4_gpu_tensor_write(sel_gpu, 0, sel, (uint64_t)n * sizeof(uint32_t)), + "split attention selection write"); + + /* The reference follows the generic kernel: score valid rows, drop + * the rest, softmax, weighted lora sum, then the value projection. */ + double ref_scale = 0.0; + for (uint32_t h = 0; h < SA_HEADS; h++) { + const float *lh = low + (size_t)h * SA_LORA; + double max_score = -DBL_MAX; + for (uint32_t s = 0; s < n; s++) { + if (sel[s] >= SA_CAP) continue; + const float *row = kv + (size_t)sel[s] * SA_LORA; + double dot = 0.0; + for (uint32_t j = 0; j < SA_LORA; j++) dot += (double)lh[j] * row[j]; + const double score = dot * 0.125; /* 1/sqrt(SA_NOPE) */ + if (score > max_score) max_score = score; + } + double denom = 0.0; + for (uint32_t j = 0; j < SA_LORA; j++) lora[j] = 0.0; + for (uint32_t s = 0; s < n; s++) { + if (sel[s] >= SA_CAP) continue; + const float *row = kv + (size_t)sel[s] * SA_LORA; + double dot = 0.0; + for (uint32_t j = 0; j < SA_LORA; j++) dot += (double)lh[j] * row[j]; + const double w = exp(dot * 0.125 - max_score); + denom += w; + for (uint32_t j = 0; j < SA_LORA; j++) lora[j] += w * row[j]; + } + if (denom < 1e-20) denom = 1e-20; + for (uint32_t d = 0; d < SA_VALUE; d++) { + const uint8_t *row = model + value_offset + + (size_t)(h * SA_VALUE + d) * SA_Q8_ROW_BYTES; + double out = 0.0; + for (uint32_t j = 0; j < SA_LORA; j++) { + const int8_t qv = (int8_t)row[(j / 32u) * 34u + 2u + j % 32u]; + out += (double)qv * (lora[j] / denom); + } + ref[h * SA_VALUE + d] = out; + if (fabs(out) > ref_scale) ref_scale = fabs(out); + } + } + + const int split_rc = ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + heads_gpu, partial_lora_gpu, partial_ms_gpu, q_gpu, low_gpu, + kv_gpu, rope_gpu, model, model_bytes, value_offset, sel_gpu, n, + false, SA_CAP, true, SA_HEADS, SA_LORA, SA_NOPE, 0, SA_VALUE, 0, + block_rows, n_blocks, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f); + if (!cases[c].accepted) { + require_ok(split_rc == 0 && n_blocks > 64u, + "split attention refuses more blocks than the reduce walks"); + continue; + } + require_ok(split_rc, what); + require_ok(ds4_gpu_tensor_read(heads_gpu, 0, spl, (uint64_t)SA_HEADS * SA_VALUE * sizeof(float)), + "split attention output read"); + require_ok(ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + heads_gpu, partial_lora_gpu, partial_ms_gpu, q_gpu, low_gpu, + kv_gpu, rope_gpu, model, model_bytes, value_offset, sel_gpu, n, + false, SA_CAP, true, SA_HEADS, SA_LORA, SA_NOPE, 0, SA_VALUE, 0, + block_rows, n_blocks, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f), what); + require_ok(ds4_gpu_tensor_read(heads_gpu, 0, spl2, (uint64_t)SA_HEADS * SA_VALUE * sizeof(float)), + "split attention repeat read"); + require_ok(ds4_gpu_glm_attention_indexed_decode_tensor( + heads_gpu, q_gpu, low_gpu, kv_gpu, rope_gpu, model, model_bytes, + value_offset, sel_gpu, n, SA_CAP, true, SA_HEADS, SA_LORA, + SA_NOPE, 0, SA_VALUE, 0, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f), + "generic indexed decode attention"); + require_ok(ds4_gpu_tensor_read(heads_gpu, 0, gen, (uint64_t)SA_HEADS * SA_VALUE * sizeof(float)), + "generic attention output read"); + + /* The phased exact kernels claim the generic kernel's arithmetic + * operation for operation, so their output must match it bit for + * bit -- including the excluded rows and sentinels. */ + require_ok(ds4_gpu_glm_attention_indexed_decode_exact_tensor( + heads_gpu, exact_scores_gpu, exact_lora_gpu, exact_denom_gpu, + low_gpu, kv_gpu, model, model_bytes, value_offset, sel_gpu, n, + SA_CAP, true, SA_HEADS, SA_LORA, SA_NOPE, 0, SA_VALUE), + "exact indexed decode attention"); + require_ok(ds4_gpu_tensor_read(heads_gpu, 0, exact, (uint64_t)SA_HEADS * SA_VALUE * sizeof(float)), + "exact attention output read"); + if (memcmp(exact, gen, (size_t)SA_HEADS * SA_VALUE * sizeof(float)) != 0) { + double worst = 0.0; + for (uint32_t i = 0; i < SA_HEADS * SA_VALUE; i++) { + worst = fmax(worst, fabs((double)exact[i] - (double)gen[i])); + } + fprintf(stderr, "%s: exact kernels differ from the generic kernel (max |delta| %.3g)\n", + what, worst); + exit(1); + } + + double gen_err = 0.0, spl_err = 0.0, pair_err = 0.0; + for (uint32_t i = 0; i < SA_HEADS * SA_VALUE; i++) { + if (!isfinite(gen[i]) || !isfinite(spl[i])) { + fprintf(stderr, "%s: non-finite output at %u\n", what, i); + exit(1); + } + gen_err = fmax(gen_err, fabs((double)gen[i] - ref[i])); + spl_err = fmax(spl_err, fabs((double)spl[i] - ref[i])); + pair_err = fmax(pair_err, fabs((double)spl[i] - (double)gen[i])); + } + if (memcmp(spl, spl2, (size_t)SA_HEADS * SA_VALUE * sizeof(float)) != 0) { + fprintf(stderr, "%s: split output changed on repeat\n", what); + exit(1); + } + /* Both kernels accumulate in f32 over up to 2051 rows and 512 lanes; + * a tiling, block or bounds error moves a result by a large fraction + * of ref_scale, orders of magnitude past this. */ + const double tol = 1e-4 * ref_scale; + fprintf(stderr, + "%s: ref_scale %.3g, generic %.3g, split %.3g, split-vs-generic %.3g (tol %.3g), exact == generic\n", + what, ref_scale, gen_err, spl_err, pair_err, tol); + if (gen_err > tol || spl_err > tol || pair_err > tol) { + fprintf(stderr, "%s: attention diverged\n", what); + exit(1); + } + } + + /* GLM 5.2's selections are always in range and it ran the unchecked + * variant before GLM 5.3 was admitted; decode now passes false for every + * GLM model, which is free of numerical consequence only if the two + * variants perform identical arithmetic on valid rows. */ + for (uint32_t s = 0; s < 2048u; s++) sel[s] = (s * 7919u) % SA_CAP; + require_ok(ds4_gpu_tensor_write(sel_gpu, 0, sel, 2048u * sizeof(uint32_t)), + "all-valid selection write"); + for (int assume_valid = 0; assume_valid < 2; assume_valid++) { + require_ok(ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + heads_gpu, partial_lora_gpu, partial_ms_gpu, q_gpu, low_gpu, + kv_gpu, rope_gpu, model, model_bytes, value_offset, sel_gpu, 2048u, + assume_valid != 0, SA_CAP, true, SA_HEADS, SA_LORA, SA_NOPE, 0, + SA_VALUE, 0, 128u, 16u, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f), + "split attention on an all-valid selection"); + require_ok(ds4_gpu_tensor_read(heads_gpu, 0, assume_valid ? spl2 : spl, + (uint64_t)SA_HEADS * SA_VALUE * sizeof(float)), + "all-valid split output read"); + } + if (memcmp(spl, spl2, (size_t)SA_HEADS * SA_VALUE * sizeof(float)) != 0) { + fprintf(stderr, "split attention: bounds-checked and unchecked " + "variants differ on an all-valid selection\n"); + exit(1); + } + + ds4_gpu_tensor_free(exact_denom_gpu); + ds4_gpu_tensor_free(exact_lora_gpu); + ds4_gpu_tensor_free(exact_scores_gpu); + ds4_gpu_tensor_free(sel_gpu); + ds4_gpu_tensor_free(rope_gpu); + ds4_gpu_tensor_free(kv_gpu); + ds4_gpu_tensor_free(low_gpu); + ds4_gpu_tensor_free(q_gpu); + ds4_gpu_tensor_free(partial_ms_gpu); + ds4_gpu_tensor_free(partial_lora_gpu); + ds4_gpu_tensor_free(heads_gpu); + free(exact); + free(spl2); + free(spl); + free(gen); + free(lora); + free(ref); + free(sel); + free(q); + free(low); + free(kv); + free(kv_bits); +} +#endif + +#ifdef __APPLE__ +static void require_prefill_dispatch(uint32_t feature, bool expected, + const char *what) { + const uint32_t dispatched = ds4_gpu_test_glm53_prefill_take_dispatches(); + require_ok(dispatched == (expected ? feature : 0u), what); +} + +/* Exactness oracle for the GLM 5.3 Flash prefill qk-low token tile. + * + * kernel_glm_qk_lowrank_q8_0_batch_t changes only which threadgroup + * computes which outputs and how many tokens one thread carries; every output + * keeps the reference kernel's expression, block order and column order. So + * each tile must reproduce kernel_glm_qk_lowrank_q8_0_batch bit for bit at the + * model's shape, including the partial tail tile that a 1596-token chunk and a + * 33- or 1-token prompt produce. The dispatch runs through the same selection + * the graph uses; DS4_METAL_DISABLE_GLM53_PREFILL_QK_LOW picks the reference + * and DS4_METAL_GLM53_PREFILL_QK_LOW_TILE picks the tile. */ +static void check_glm53_qk_lowrank_token_tile(uint8_t *model, + uint64_t model_bytes, + uint64_t kb_offset) { + enum { + QL_HEADS = 64, + QL_KV_LORA = 512, + QL_QK_NOPE = 256, + QL_QK_DIM = 256, + QL_ROW_BYTES = 272, /* 8 Q8_0 blocks of 34 bytes */ + QL_Q8_0_TYPE = 8, /* GGUF type code for Q8_0 */ + QL_MAX_TOKENS = 2048, + }; + static const uint32_t token_counts[] = { 2048u, 1596u, 33u, 1u }; + static const uint32_t tiles[] = { 4u, 8u, 16u }; + /* Scales that exercise the sign and the subnormal half range, where a + * reassociated product would round differently. */ + static const uint16_t scale_bits[] = { + 0x0001u, 0x8001u, 0x03ffu, 0x83ffu, 0x0000u, 0x8000u, + 0x3c00u, 0xbc00u, 0x1234u, 0x9876u, 0x2c00u, 0xac00u, 0x0400u, 0x8400u, + }; + const uint64_t weight_bytes = + (uint64_t)QL_HEADS * QL_KV_LORA * QL_ROW_BYTES; + require_ok(kb_offset + weight_bytes <= model_bytes, + "qk-low K_b rows fit the fixture model"); + + uint64_t rng = 0x9e3779b97f4a7c15ull; + for (uint64_t row = 0; row < (uint64_t)QL_HEADS * QL_KV_LORA; row++) { + uint8_t *dst = model + kb_offset + row * QL_ROW_BYTES; + for (uint32_t b = 0; b < QL_QK_NOPE / 32u; b++) { + const uint16_t d = + scale_bits[(row * 8u + b) % (sizeof(scale_bits) / sizeof(scale_bits[0]))]; + memcpy(dst + b * 34u, &d, sizeof(d)); + int8_t *qs = (int8_t *)(dst + b * 34u + 2u); + for (uint32_t i = 0; i < 32u; i++) { + rng = rng * 6364136223846793005ull + 1442695040888963407ull; + qs[i] = (int8_t)(uint8_t)(rng >> 33); + } + } + } + + const uint64_t q_elems = (uint64_t)QL_MAX_TOKENS * QL_HEADS * QL_QK_DIM; + const uint64_t out_elems = (uint64_t)QL_MAX_TOKENS * QL_HEADS * QL_KV_LORA; + float *q_host = malloc(q_elems * sizeof(float)); + float *ref_host = malloc(out_elems * sizeof(float)); + float *tile_host = malloc(out_elems * sizeof(float)); + require_ok(q_host && ref_host && tile_host, "qk-low host allocation"); + for (uint64_t i = 0; i < q_elems; i++) { + rng = rng * 6364136223846793005ull + 1442695040888963407ull; + q_host[i] = (float)((int32_t)(uint32_t)(rng >> 32) / 1073741824.0) - 1.0f; + } + + ds4_gpu_tensor *q_gpu = ds4_gpu_tensor_alloc(q_elems * sizeof(float)); + ds4_gpu_tensor *ref_gpu = ds4_gpu_tensor_alloc(out_elems * sizeof(float)); + ds4_gpu_tensor *tile_gpu = ds4_gpu_tensor_alloc(out_elems * sizeof(float)); + require_ok(q_gpu && ref_gpu && tile_gpu, "qk-low GPU allocation"); + require_ok(ds4_gpu_tensor_write(q_gpu, 0, q_host, q_elems * sizeof(float)), + "qk-low q write"); + + for (size_t c = 0; c < sizeof(token_counts) / sizeof(token_counts[0]); c++) { + const uint32_t n_tokens = token_counts[c]; + const uint64_t bytes = + (uint64_t)n_tokens * QL_HEADS * QL_KV_LORA * sizeof(float); + char what[96]; + + require_ok(setenv("DS4_METAL_DISABLE_GLM53_PREFILL_QK_LOW", "1", 1) == 0, + "qk-low reference switch"); + snprintf(what, sizeof(what), "qk-low reference at %u tokens", n_tokens); + require_ok(ds4_gpu_glm_qk_lowrank_typed_batch_tensor( + ref_gpu, q_gpu, model, model_bytes, kb_offset, + QL_Q8_0_TYPE, n_tokens, QL_HEADS, QL_KV_LORA, + QL_QK_NOPE, QL_QK_DIM), what); + require_ok(ds4_gpu_tensor_read(ref_gpu, 0, ref_host, bytes), what); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_QK_LOW, false, what); + require_ok(unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_QK_LOW") == 0, + "qk-low reference switch clear"); + + const size_t tile_count = sizeof(tiles) / sizeof(tiles[0]); + for (size_t t = 0; t <= tile_count; t++) { + const bool rollback = t == tile_count; + const uint32_t tile = tiles[t % tile_count]; + if (rollback) setenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING", "1", 1); + char tile_text[8]; + snprintf(tile_text, sizeof(tile_text), "%u", tile); + require_ok(setenv("DS4_METAL_GLM53_PREFILL_QK_LOW_TILE", tile_text, 1) == 0, + "qk-low tile switch"); + snprintf(what, sizeof(what), "qk-low tile %u at %u tokens rollback=%u", + tile, n_tokens, rollback); + /* A quiet NaN in every output first, so a kernel that skips rows + * fails here rather than matching a stale buffer. Built from bits + * because -ffast-math makes the NAN macro undefined. */ + const uint32_t poison_bits = 0x7fc01234u; + float poison; + memcpy(&poison, &poison_bits, sizeof(poison)); + require_ok(ds4_gpu_tensor_fill_f32(tile_gpu, poison, + (uint64_t)n_tokens * QL_HEADS * QL_KV_LORA), + what); + require_ok(ds4_gpu_glm_qk_lowrank_typed_batch_tensor( + tile_gpu, q_gpu, model, model_bytes, kb_offset, + QL_Q8_0_TYPE, n_tokens, QL_HEADS, QL_KV_LORA, + QL_QK_NOPE, QL_QK_DIM), what); + require_ok(ds4_gpu_tensor_read(tile_gpu, 0, tile_host, bytes), what); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_QK_LOW, + !rollback && n_tokens >= tile, what); + if (memcmp(ref_host, tile_host, (size_t)bytes) != 0) { + for (uint64_t i = 0; i < bytes / sizeof(float); i++) { + if (memcmp(&ref_host[i], &tile_host[i], sizeof(float)) == 0) continue; + fprintf(stderr, + "%s: output %llu is %.9g, reference %.9g\n", + what, (unsigned long long)i, + (double)tile_host[i], (double)ref_host[i]); + break; + } + exit(1); + } + } + require_ok(unsetenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") == 0, + "clear aggregate rollback switch"); + require_ok(unsetenv("DS4_METAL_GLM53_PREFILL_QK_LOW_TILE") == 0, + "qk-low tile switch clear"); + } + + ds4_gpu_tensor_free(tile_gpu); + ds4_gpu_tensor_free(ref_gpu); + ds4_gpu_tensor_free(q_gpu); + free(tile_host); + free(ref_host); + free(q_host); +} + +/* Optional 48-GiB regression for both 32-bit element-offset wrap boundaries. + * Kept out of the ordinary suite so machines with smaller memory can run it. + * The last real token must equal a one-token reference, not retain poison or + * read token zero after a wrapped input offset. */ +static void check_glm53_qk_lowrank_large_offsets(uint8_t *model, + uint64_t model_bytes, + uint64_t kb_offset) { + if (!getenv("DS4_TEST_GLM53_LARGE_QK")) return; + enum { HEADS = 64, NOPE = 256, LORA = 512, TOKENS = 262145 }; + const uint64_t q_row = (uint64_t)HEADS * NOPE * sizeof(float); + const uint64_t out_row = (uint64_t)HEADS * LORA * sizeof(float); + ds4_gpu_tensor *q = ds4_gpu_tensor_alloc((uint64_t)TOKENS * q_row); + ds4_gpu_tensor *out = ds4_gpu_tensor_alloc((uint64_t)TOKENS * out_row); + ds4_gpu_tensor *ref = ds4_gpu_tensor_alloc(out_row); + require_ok(q && out && ref, "large qk-low allocations (48 GiB required)"); + ds4_gpu_tensor *q_last = ds4_gpu_tensor_view(q, (TOKENS - 1ull) * q_row, q_row); + ds4_gpu_tensor *out_last = ds4_gpu_tensor_view(out, (TOKENS - 1ull) * out_row, out_row); + require_ok(q_last && out_last, "large qk-low tail views"); + require_ok(ds4_gpu_tensor_fill_f32(q, 0.0f, (uint64_t)TOKENS * HEADS * NOPE) && + ds4_gpu_tensor_fill_f32(q_last, 1.0f, HEADS * NOPE) && + ds4_gpu_tensor_fill_f32(out_last, 123.0f, HEADS * LORA), "large qk-low inputs and poison"); + require_ok(ds4_gpu_glm_qk_lowrank_typed_batch_tensor(ref, q_last, + model, model_bytes, kb_offset, 8u, 1, HEADS, LORA, NOPE, NOPE), "large qk-low one-token reference"); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_QK_LOW, false, "large qk-low reference coverage"); + require_ok(ds4_gpu_glm_qk_lowrank_typed_batch_tensor(out, q, + model, model_bytes, kb_offset, 8u, TOKENS, HEADS, LORA, NOPE, NOPE), "large qk-low token tile"); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_QK_LOW, true, "large qk-low tile coverage"); + float expected[HEADS * LORA], actual[HEADS * LORA]; + require_ok(ds4_gpu_tensor_read(ref, 0, expected, out_row) && + ds4_gpu_tensor_read(out_last, 0, actual, out_row), "large qk-low readback"); + require_ok(memcmp(expected, actual, out_row) == 0, "large qk-low tail is bit-identical"); + ds4_gpu_tensor_free(out_last); ds4_gpu_tensor_free(q_last); + ds4_gpu_tensor_free(ref); ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(q); + puts("GLM qk-low 64-bit offset regression: PASS"); +} + +/* Exactness oracle for the GLM 5.3 Flash indexed prefill attention head width. + * + * kernel_glm_attention_indexed_batch_lora_group16_vec_valid_fullheads carries + * two heads per simdgroup, so a token stages its selected rows four times + * instead of eight. Each head keeps the one-head kernel's row order, its four + * dot(float4) terms, its simd_sum tree and its online-softmax update, so all + * 512 outputs of every head must match bit for bit. + * DS4_METAL_GLM53_PREFILL_INDEXED_ATTN_HEADS_PER_SG picks the width, and + * DS4_METAL_DISABLE_GLM53_PREFILL_INDEXED_ATTN pins main's one-head kernel. */ +static void check_glm53_indexed_attention_head_width(void) { + enum { + IA_HEADS = 64, + IA_LORA = 512, + IA_NOPE = 256, + IA_CACHE_CAP = 4096, + IA_TOKENS = 8, + IA_MAX_SELECTED = 2051, + }; + /* 2051 is the model's selection limit; the rest land on a partial trailing + * 16-row staging block, which is where a head-width bug would show. */ + static const uint32_t selected_counts[] = { 2051u, 512u, 33u, 16u, 1u }; + + const uint64_t q_elems = (uint64_t)IA_TOKENS * IA_HEADS * IA_NOPE; + const uint64_t low_elems = (uint64_t)IA_TOKENS * IA_HEADS * IA_LORA; + const uint64_t cache_elems = (uint64_t)IA_CACHE_CAP * IA_LORA; + const uint64_t sel_elems = (uint64_t)IA_TOKENS * IA_MAX_SELECTED; + + float *q_host = malloc(q_elems * sizeof(float)); + float *low_host = malloc(low_elems * sizeof(float)); + uint16_t *cache_host = malloc(cache_elems * sizeof(uint16_t)); + uint32_t *sel_host = malloc(sel_elems * sizeof(uint32_t)); + float *ref_host = malloc(low_elems * sizeof(float)); + float *dual_host = malloc(low_elems * sizeof(float)); + require_ok(q_host && low_host && cache_host && sel_host && ref_host && dual_host, + "indexed attention host allocation"); + + uint64_t rng = 0xda3e39cb94b95bdbull; +#define IA_NEXT_UNIT() ( \ + rng = rng * 6364136223846793005ull + 1442695040888963407ull, \ + (float)((int32_t)(uint32_t)(rng >> 32) / 1073741824.0) - 1.0f) + for (uint64_t i = 0; i < q_elems; i++) q_host[i] = IA_NEXT_UNIT(); + for (uint64_t i = 0; i < low_elems; i++) low_host[i] = IA_NEXT_UNIT(); + /* Half values kept in the normal range so the truncating encoder above is + * exact and the fixture round-trips. */ + for (uint64_t i = 0; i < cache_elems; i++) { + const float unit = IA_NEXT_UNIT(); + cache_host[i] = f32_to_f16(unit >= 0.0f ? 0.0625f + unit : -0.0625f + unit); + } + /* Every selected row must be in cache range: this kernel family is the + * "valid rows" instantiation and does not re-check them. */ + for (uint64_t i = 0; i < sel_elems; i++) { + rng = rng * 6364136223846793005ull + 1442695040888963407ull; + sel_host[i] = (uint32_t)((rng >> 33) % (uint64_t)IA_CACHE_CAP); + } + + ds4_gpu_tensor *q_gpu = ds4_gpu_tensor_alloc(q_elems * sizeof(float)); + ds4_gpu_tensor *low_gpu = ds4_gpu_tensor_alloc(low_elems * sizeof(float)); + ds4_gpu_tensor *cache_gpu = ds4_gpu_tensor_alloc(cache_elems * sizeof(uint16_t)); + ds4_gpu_tensor *rope_gpu = ds4_gpu_tensor_alloc(sizeof(float)); + ds4_gpu_tensor *sel_gpu = ds4_gpu_tensor_alloc(sel_elems * sizeof(uint32_t)); + ds4_gpu_tensor *ref_gpu = ds4_gpu_tensor_alloc(low_elems * sizeof(float)); + ds4_gpu_tensor *dual_gpu = ds4_gpu_tensor_alloc(low_elems * sizeof(float)); + require_ok(q_gpu && low_gpu && cache_gpu && rope_gpu && sel_gpu && ref_gpu && dual_gpu, + "indexed attention GPU allocation"); + require_ok(ds4_gpu_tensor_write(q_gpu, 0, q_host, q_elems * sizeof(float)) && + ds4_gpu_tensor_write(low_gpu, 0, low_host, low_elems * sizeof(float)) && + ds4_gpu_tensor_write(cache_gpu, 0, cache_host, cache_elems * sizeof(uint16_t)) && + ds4_gpu_tensor_write(sel_gpu, 0, sel_host, sel_elems * sizeof(uint32_t)), + "indexed attention input write"); + + for (size_t c = 0; c < sizeof(selected_counts) / sizeof(selected_counts[0]); c++) { + const uint32_t n_selected = selected_counts[c]; + const uint64_t bytes = low_elems * sizeof(float); + char what[96]; + + require_ok(setenv("DS4_METAL_GLM53_PREFILL_INDEXED_ATTN_HEADS_PER_SG", "1", 1) == 0, + "indexed attention width switch"); + snprintf(what, sizeof(what), "indexed attention one head at %u rows", n_selected); + require_ok(ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + ref_gpu, q_gpu, low_gpu, cache_gpu, rope_gpu, sel_gpu, + IA_TOKENS, n_selected, IA_CACHE_CAP, true, IA_HEADS, + IA_LORA, IA_NOPE, 0u, 0u, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f), what); + require_ok(ds4_gpu_tensor_read(ref_gpu, 0, ref_host, bytes), what); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_INDEXED_ATTN, false, what); + + require_ok(setenv("DS4_METAL_GLM53_PREFILL_INDEXED_ATTN_HEADS_PER_SG", "2", 1) == 0, + "indexed attention width switch"); + snprintf(what, sizeof(what), "indexed attention two heads at %u rows", n_selected); + const uint32_t poison_bits = 0x7fc01234u; + float poison; + memcpy(&poison, &poison_bits, sizeof(poison)); + require_ok(ds4_gpu_tensor_fill_f32(dual_gpu, poison, low_elems), what); + require_ok(ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + dual_gpu, q_gpu, low_gpu, cache_gpu, rope_gpu, sel_gpu, + IA_TOKENS, n_selected, IA_CACHE_CAP, true, IA_HEADS, + IA_LORA, IA_NOPE, 0u, 0u, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f), what); + require_ok(ds4_gpu_tensor_read(dual_gpu, 0, dual_host, bytes), what); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_INDEXED_ATTN, true, what); + if (memcmp(ref_host, dual_host, (size_t)bytes) != 0) { + for (uint64_t i = 0; i < low_elems; i++) { + if (memcmp(&ref_host[i], &dual_host[i], sizeof(float)) == 0) continue; + fprintf(stderr, + "%s: token %llu head %llu lane element %llu is %.9g, one-head %.9g\n", + what, + (unsigned long long)(i / (IA_HEADS * IA_LORA)), + (unsigned long long)((i / IA_LORA) % IA_HEADS), + (unsigned long long)(i % IA_LORA), + (double)dual_host[i], (double)ref_host[i]); + break; + } + exit(1); + } + } + require_ok(unsetenv("DS4_METAL_GLM53_PREFILL_INDEXED_ATTN_HEADS_PER_SG") == 0, + "indexed attention width switch clear"); +#undef IA_NEXT_UNIT + + ds4_gpu_tensor_free(dual_gpu); + ds4_gpu_tensor_free(ref_gpu); + ds4_gpu_tensor_free(sel_gpu); + ds4_gpu_tensor_free(rope_gpu); + ds4_gpu_tensor_free(cache_gpu); + ds4_gpu_tensor_free(low_gpu); + ds4_gpu_tensor_free(q_gpu); + free(dual_host); + free(ref_host); + free(sel_host); + free(cache_host); + free(low_host); + free(q_host); +} + +/* Invalid selected IDs must be skipped without changing the order of valid + * rows. Compare the guarded API with the same selection compacted through the + * valid API, including the original RoPE and partial-head specializations. */ +static void check_glm53_indexed_attention_invalid_rows(void) { + enum { CAP = 64, MAX_HEADS = 64, LORA = 512, MAX_Q = 320, VALID = 16 }; + float q[MAX_HEADS * MAX_Q], low[MAX_HEADS * LORA]; + uint16_t cache[CAP * LORA], rope[CAP * 64]; + uint32_t compact[VALID], masked[2 * VALID]; + float expected[MAX_HEADS * LORA], actual[MAX_HEADS * LORA]; + for (unsigned i = 0; i < MAX_HEADS * MAX_Q; i++) q[i] = 0.01f * ((int)(i % 17) - 8); + for (unsigned i = 0; i < MAX_HEADS * LORA; i++) low[i] = 0.02f * ((int)(i % 19) - 9); + for (unsigned i = 0; i < CAP * LORA; i++) cache[i] = f32_to_f16(0.125f * ((int)(i % 13) - 6)); + for (unsigned i = 0; i < CAP * 64; i++) rope[i] = f32_to_f16(0.125f * ((int)(i % 7) - 3)); + for (unsigned i = 0; i < VALID; i++) { + compact[i] = (i * 7) % CAP; + masked[2 * i] = compact[i]; + masked[2 * i + 1] = i % 2 ? UINT32_MAX : CAP + i; + } + ds4_gpu_tensor *gq = ds4_gpu_tensor_alloc(sizeof(q)); + ds4_gpu_tensor *glow = ds4_gpu_tensor_alloc(sizeof(low)); + ds4_gpu_tensor *gcache = ds4_gpu_tensor_alloc(sizeof(cache)); + ds4_gpu_tensor *grope = ds4_gpu_tensor_alloc(sizeof(rope)); + ds4_gpu_tensor *gcompact = ds4_gpu_tensor_alloc(sizeof(compact)); + ds4_gpu_tensor *gmasked = ds4_gpu_tensor_alloc(sizeof(masked)); + ds4_gpu_tensor *gout = ds4_gpu_tensor_alloc(sizeof(actual)); + require_ok(gq && glow && gcache && grope && gcompact && gmasked && gout, "invalid attention allocations"); + require_ok(ds4_gpu_tensor_write(gq, 0, q, sizeof(q)) && + ds4_gpu_tensor_write(glow, 0, low, sizeof(low)) && + ds4_gpu_tensor_write(gcache, 0, cache, sizeof(cache)) && + ds4_gpu_tensor_write(grope, 0, rope, sizeof(rope)) && + ds4_gpu_tensor_write(gcompact, 0, compact, sizeof(compact)) && + ds4_gpu_tensor_write(gmasked, 0, masked, sizeof(masked)), "invalid attention uploads"); + for (unsigned h = 0; h < 2; h++) for (unsigned r = 0; r < 2; r++) { + const uint32_t heads = h ? 7 : 64, rot = r ? 64 : 0; + const uint64_t bytes = (uint64_t)heads * LORA * sizeof(float); + require_ok(ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + gout, gq, glow, gcache, grope, gcompact, 1, VALID, CAP, true, + heads, LORA, 256, rot, 4096, 10000.0f, 1.0f, 1.0f, 1.0f, 32.0f, 1.0f), "compact attention reference"); + require_ok(ds4_gpu_tensor_read(gout, 0, expected, bytes), "compact attention read"); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_INDEXED_ATTN, heads == 64 && rot == 0, "compact attention coverage"); + require_ok(ds4_gpu_tensor_fill_f32(gout, 123.0f, heads * LORA), "poison masked output"); + require_ok(ds4_gpu_glm_attention_indexed_batch_lora_tensor( + gout, gq, glow, gcache, grope, gmasked, 1, 2 * VALID, CAP, true, + heads, LORA, 256, rot, 4096, 10000.0f, 1.0f, 1.0f, 1.0f, 32.0f, 1.0f), "masked attention"); + require_ok(ds4_gpu_tensor_read(gout, 0, actual, bytes), "masked attention read"); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_INDEXED_ATTN, false, "masked attention fallback"); + require_ok(memcmp(expected, actual, bytes) == 0, "invalid rows preserve exact valid-row attention"); + } + ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(gmasked); ds4_gpu_tensor_free(gcompact); + ds4_gpu_tensor_free(grope); ds4_gpu_tensor_free(gcache); ds4_gpu_tensor_free(glow); ds4_gpu_tensor_free(gq); +} + +/* Exactness oracle for the Q4_K routed-expert tail cull. + * + * kernel_mul_mm_id_q4_K_{f32,f16}_tail_cull differ from the kernels beside + * them only in that the SIMDgroup pair owning routed rows 16..31 skips its + * MMA and store when the expert's final 32-row tile holds 16 rows or fewer. + * Those outputs are padding rows nothing reads, so both the f16 mid and the + * summed f32 output must be byte-identical. The per-expert row counts below + * cover every final-tile size that matters: exact multiples of 32, 16 or + * fewer, and 17 or more. + * + * Test mode forces the synthetic shape through the cull and records dispatch + * coverage, so unsupported/default-off devices cannot silently compare the + * reference with itself. */ +static void check_glm53_routed_moe_tail_cull(uint8_t *model, + uint64_t model_bytes, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset) { + enum { + MOE_EXPERTS = 36, + MOE_USED = 8, + MOE_DIM = 256, + MOE_TOKENS = 384, + MOE_Q4_K_ROW_BYTES = 144, /* one 256-element Q4_K block */ + MOE_Q4_K_TYPE = 12, /* GGUF type code for Q4_K */ + MOE_EXPERT_BYTES = MOE_DIM * MOE_Q4_K_ROW_BYTES, + }; + uint32_t target_rows[MOE_EXPERTS]; + for (uint32_t i = 0; i < 32u; i++) target_rows[i] = 64u + i; + target_rows[32] = 0u; /* empty expert */ + target_rows[33] = target_rows[34] = 256u; + target_rows[35] = 16u; /* total: 384 tokens * 8 distinct experts */ + static const uint16_t scale_bits[] = { + 0x2c00u, 0xac00u, 0x3400u, 0xb400u, 0x1c00u, 0x9c00u, 0x3800u, 0x2400u, + 0x0001u, 0x8001u, 0x03ffu, 0x83ffu, 0x0000u, 0x8000u, + }; + const uint64_t matrix_bytes = (uint64_t)MOE_EXPERTS * MOE_EXPERT_BYTES; + require_ok(down_offset + matrix_bytes <= model_bytes, + "routed MoE expert weights fit the fixture model"); + + uint64_t rng = 0xc3a5c85c97cb3127ull; +#define MOE_NEXT_BYTE() ( \ + rng = rng * 6364136223846793005ull + 1442695040888963407ull, \ + (uint8_t)(rng >> 33)) + const uint64_t offsets[3] = { gate_offset, up_offset, down_offset }; + for (int m = 0; m < 3; m++) { + for (uint32_t row = 0; row < MOE_EXPERTS * MOE_DIM; row++) { + uint8_t *dst = model + offsets[m] + (uint64_t)row * MOE_Q4_K_ROW_BYTES; + const uint16_t d = scale_bits[(row + (uint32_t)m) % (sizeof(scale_bits) / sizeof(scale_bits[0]))]; + const uint16_t dmin = scale_bits[(row + (uint32_t)m + 3u) % (sizeof(scale_bits) / sizeof(scale_bits[0]))]; + memcpy(dst + 0, &d, sizeof(d)); + memcpy(dst + 2, &dmin, sizeof(dmin)); + for (uint32_t i = 4; i < MOE_Q4_K_ROW_BYTES; i++) dst[i] = MOE_NEXT_BYTE(); + } + } + + const uint64_t x_elems = (uint64_t)MOE_TOKENS * MOE_DIM; + const uint64_t route_elems = (uint64_t)MOE_TOKENS * MOE_USED; + const uint64_t mid_elems = route_elems * MOE_DIM; + const uint64_t out_elems = (uint64_t)MOE_TOKENS * MOE_DIM; + + float *x_host = malloc(x_elems * sizeof(float)); + int32_t *sel_host = malloc(route_elems * sizeof(int32_t)); + float *w_host = malloc(route_elems * sizeof(float)); + float *mid_ref = malloc(mid_elems * sizeof(float)); + float *mid_cull = malloc(mid_elems * sizeof(float)); + float *out_ref = malloc(out_elems * sizeof(float)); + float *out_cull = malloc(out_elems * sizeof(float)); + require_ok(x_host && sel_host && w_host && mid_ref && mid_cull && out_ref && out_cull, + "routed MoE host allocation"); + for (uint64_t i = 0; i < x_elems; i++) { + rng = rng * 6364136223846793005ull + 1442695040888963407ull; + x_host[i] = (float)((int32_t)(uint32_t)(rng >> 32) / 1073741824.0) - 1.0f; + } + for (uint64_t i = 0; i < route_elems; i++) { + rng = rng * 6364136223846793005ull + 1442695040888963407ull; + w_host[i] = 0.05f + (float)(rng >> 40) / 8388608.0f; + } + /* Hand every token the eight experts with the most rows still owed, which + * realizes target_rows exactly and keeps a token's experts distinct. */ + uint32_t remaining[MOE_EXPERTS]; + memcpy(remaining, target_rows, sizeof(remaining)); + for (uint32_t t = 0; t < MOE_TOKENS; t++) { + bool taken[MOE_EXPERTS] = { false }; + for (uint32_t s = 0; s < MOE_USED; s++) { + uint32_t best = MOE_EXPERTS; + for (uint32_t e = 0; e < MOE_EXPERTS; e++) { + if (taken[e]) continue; + if (best == MOE_EXPERTS || remaining[e] > remaining[best]) best = e; + } + require_ok(best < MOE_EXPERTS && remaining[best] > 0, + "routed MoE route construction"); + taken[best] = true; + remaining[best]--; + sel_host[(uint64_t)t * MOE_USED + s] = (int32_t)best; + } + } +#undef MOE_NEXT_BYTE + + ds4_gpu_tensor *x_gpu = ds4_gpu_tensor_alloc(x_elems * sizeof(float)); + ds4_gpu_tensor *sel_gpu = ds4_gpu_tensor_alloc(route_elems * sizeof(int32_t)); + ds4_gpu_tensor *w_gpu = ds4_gpu_tensor_alloc(route_elems * sizeof(float)); + ds4_gpu_tensor *mid_gpu = ds4_gpu_tensor_alloc(mid_elems * sizeof(float)); + ds4_gpu_tensor *out_gpu = ds4_gpu_tensor_alloc(out_elems * sizeof(float)); + require_ok(x_gpu && sel_gpu && w_gpu && mid_gpu && out_gpu, + "routed MoE GPU allocation"); + require_ok(ds4_gpu_tensor_write(x_gpu, 0, x_host, x_elems * sizeof(float)) && + ds4_gpu_tensor_write(sel_gpu, 0, sel_host, route_elems * sizeof(int32_t)) && + ds4_gpu_tensor_write(w_gpu, 0, w_host, route_elems * sizeof(float)), + "routed MoE input write"); + + const uint32_t poison_bits = 0x7fc01234u; + float poison; + memcpy(&poison, &poison_bits, sizeof(poison)); + for (int cull = 0; cull < 2; cull++) { + const char *what = cull ? "routed MoE tail cull" : "routed MoE reference"; + if (cull) { + require_ok(unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_MOE_TAIL_CULL") == 0, what); + } else { + require_ok(setenv("DS4_METAL_DISABLE_GLM53_PREFILL_MOE_TAIL_CULL", "1", 1) == 0, what); + } + require_ok(ds4_gpu_tensor_fill_f32(mid_gpu, poison, mid_elems) && + ds4_gpu_tensor_fill_f32(out_gpu, poison, out_elems), what); + require_ok(ds4_gpu_glm_routed_moe_batch_tensor( + out_gpu, mid_gpu, model, model_bytes, + gate_offset, up_offset, down_offset, + MOE_Q4_K_TYPE, MOE_Q4_K_TYPE, MOE_Q4_K_TYPE, + MOE_EXPERT_BYTES, MOE_Q4_K_ROW_BYTES, + MOE_EXPERT_BYTES, MOE_Q4_K_ROW_BYTES, + MOE_EXPERT_BYTES, MOE_Q4_K_ROW_BYTES, + MOE_DIM, MOE_DIM, MOE_DIM, + sel_gpu, w_gpu, MOE_EXPERTS, MOE_USED, + 10.0f, 0u, x_gpu, MOE_TOKENS, + MOE_USED * MOE_DIM, true), what); + require_prefill_dispatch(DS4_GPU_GLM53_PREFILL_MOE_TAIL_CULL, + cull != 0, what); + require_ok(ds4_gpu_tensor_read(mid_gpu, 0, cull ? mid_cull : mid_ref, + mid_elems * sizeof(float)) && + ds4_gpu_tensor_read(out_gpu, 0, cull ? out_cull : out_ref, + out_elems * sizeof(float)), + what); + } + require_ok(unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_MOE_TAIL_CULL") == 0, + "routed MoE tail cull switch clear"); + + /* The f16 mid occupies the first half of the f32 mid buffer. */ + if (memcmp(mid_ref, mid_cull, (size_t)(mid_elems * sizeof(uint16_t))) != 0 || + memcmp(out_ref, out_cull, (size_t)(out_elems * sizeof(float))) != 0) { + for (uint64_t i = 0; i < out_elems; i++) { + if (memcmp(&out_ref[i], &out_cull[i], sizeof(float)) == 0) continue; + fprintf(stderr, + "routed MoE tail cull: output %llu is %.9g, reference %.9g\n", + (unsigned long long)i, (double)out_cull[i], (double)out_ref[i]); + break; + } + fprintf(stderr, "routed MoE tail cull is not bit-identical\n"); + exit(1); + } + + ds4_gpu_tensor_free(out_gpu); + ds4_gpu_tensor_free(mid_gpu); + ds4_gpu_tensor_free(w_gpu); + ds4_gpu_tensor_free(sel_gpu); + ds4_gpu_tensor_free(x_gpu); + free(out_cull); + free(out_ref); + free(mid_cull); + free(mid_ref); + free(w_host); + free(sel_host); + free(x_host); +} + +/* Exactness oracle for the blocked GLM 5.3 KDA prepare kernel. + * + * kernel_glm53_kda_prefill_prepare_blocked splits the serial kernel's token + * loop across (block, head) threadgroups and carries the three-row causal + * convolution history in registers instead of shifting it through the device + * conv state. Every value keeps the serial kernel's expression and order, so + * the normalized q and k, the silu'd v, the decay gate, the recurrence output + * and the outgoing conv and recurrent states must all match bit for bit -- + * including where the window still reaches into the incoming conv state + * (the first three rows) and on the short trailing block. + * + * This overwrites the KDA convolution fixture weights, so it runs after the + * checks that use them. */ +static void check_glm53_kda_prepare_blocked(uint8_t *model, + uint64_t model_bytes, + uint64_t q_conv_offset, + uint64_t k_conv_offset, + uint64_t v_conv_offset, + uint64_t a_log_offset, + uint64_t dt_bias_offset, + uint64_t norm_offset) { + enum { H = 64, D = 128, P = H * D, MAX_TOKENS = 2048 }; + enum { Q, K, V, GATE, OGATE, BETA, CONV, STATE, OUT, NBUF }; + /* Production head count, partial blocks, and explicit fallback boundaries. */ + static const uint32_t tokens[] = { 2048, 1596, 65, 33, 17, 4, 3, 1 }; + static const struct { + uint32_t block, values; + bool last_first, profile, batch, split, rollback; + } variants[] = { + {0, 1, false, false, false, false, false}, /* serial reference */ + {4, 2, false, false, false, false, false}, + {16, 2, false, false, false, false, false}, + {32, 2, false, false, false, false, false}, + {64, 4, false, false, false, false, false}, + {32, 2, true, false, false, false, false}, /* deterministic race regression */ + {32, 2, false, true, false, false, false}, /* profiler with an owned CB */ + {32, 2, false, true, true, false, false}, /* profiler preserves caller batch */ + {32, 2, true, false, true, true, false}, /* continue 33 + 32 tokens */ + {32, 2, false, false, false, false, true}, /* aggregate rollback */ + }; + const bool inherited_profile = getenv("DS4_METAL_PROFILE_KDA_PREFILL") != NULL; + require_ok(norm_offset + D * sizeof(float) <= model_bytes, + "production KDA weights fit fixture"); + uint64_t rng = 0x2545f4914f6cdd1dull; +#define KP_UNIT() ( \ + rng = rng * 6364136223846793005ull + 1442695040888963407ull, \ + (float)((int32_t)(uint32_t)(rng >> 32) / 1073741824.0) - 1.0f) + const uint64_t conv_offsets[] = { q_conv_offset, k_conv_offset, v_conv_offset }; + for (unsigned m = 0; m < 3; m++) { + float *w = (float *)(model + conv_offsets[m]); + for (unsigned i = 0; i < P * 4u; i++) w[i] = 0.4f * KP_UNIT(); + } + for (unsigned i = 0; i < P; i++) ((float *)(model + dt_bias_offset))[i] = 0.2f * KP_UNIT(); + for (unsigned i = 0; i < H; i++) ((float *)(model + a_log_offset))[i] = 0.3f * KP_UNIT(); + for (unsigned i = 0; i < D; i++) ((float *)(model + norm_offset))[i] = 1.0f + 0.1f * KP_UNIT(); + + uint64_t elements[NBUF]; + float *input[NBUF], *reference[NBUF]; + ds4_gpu_tensor *gpu[NBUF]; + for (unsigned b = 0; b < NBUF; b++) { + elements[b] = b == BETA ? (uint64_t)MAX_TOKENS * H : + b == CONV ? 9u * P : b == STATE ? (uint64_t)P * D : + (uint64_t)MAX_TOKENS * P; + input[b] = malloc(elements[b] * sizeof(float)); + reference[b] = malloc(elements[b] * sizeof(float)); + gpu[b] = ds4_gpu_tensor_alloc(elements[b] * sizeof(float)); + require_ok(input[b] && reference[b] && gpu[b], "KDA oracle allocation"); + for (uint64_t i = 0; i < elements[b]; i++) input[b][i] = KP_UNIT() * (b == STATE ? 0.1f : 1.0f); + } +#undef KP_UNIT + float *actual = malloc((uint64_t)MAX_TOKENS * P * sizeof(float)); + require_ok(actual != NULL, "KDA readback allocation"); + static const unsigned compared[] = { Q, K, V, GATE, CONV, STATE, OUT }; + static const char *const names[] = { "q", "k", "v", "decay", "output gate", "beta", "conv state", "recurrent state", "output" }; + const uint32_t poison_bits = 0x7fc01234u; + float poison; + memcpy(&poison, &poison_bits, sizeof(poison)); + for (unsigned c = 0; c < sizeof(tokens) / sizeof(tokens[0]); c++) { + const uint32_t n = tokens[c]; + uint64_t bytes[NBUF]; + for (unsigned b = 0; b < NBUF; b++) { + bytes[b] = (b == BETA ? (uint64_t)n * H : + (b == CONV || b == STATE) ? elements[b] : (uint64_t)n * P) * sizeof(float); + } + for (unsigned variant = 0; variant < sizeof(variants) / sizeof(variants[0]); variant++) { + if (variants[variant].split && n != 65u) continue; + const uint32_t block = variants[variant].block; + char what[128], text[16]; + snprintf(what, sizeof(what), "KDA n=%u block=%u values=%u last-first=%u profile=%u batch=%u split=%u rollback=%u", + n, block, variants[variant].values, variants[variant].last_first, variants[variant].profile, + variants[variant].batch, variants[variant].split, variants[variant].rollback); + if (variants[variant].rollback) setenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING", "1", 1); + else unsetenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING"); + if (block == 0u) { + setenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_PREPARE", "1", 1); + setenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_RECURRENCE", "1", 1); + } else { + unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_PREPARE"); + unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_RECURRENCE"); + } + snprintf(text, sizeof(text), "%u", block); + setenv("DS4_METAL_GLM53_PREFILL_KDA_PREPARE_BLOCK", text, 1); + snprintf(text, sizeof(text), "%u", variants[variant].values); + setenv("DS4_METAL_GLM53_PREFILL_KDA_VALUES_PER_SG", text, 1); + if (variants[variant].profile) setenv("DS4_METAL_PROFILE_KDA_PREFILL", "1", 1); + else unsetenv("DS4_METAL_PROFILE_KDA_PREFILL"); + ds4_gpu_test_set_flags(DS4_GPU_TEST_GLM53_PREFILL | + (variants[variant].last_first ? DS4_GPU_TEST_GLM53_KDA_LAST_BLOCK_FIRST : 0u)); + for (unsigned b = 0; b < OUT; b++) require_ok(ds4_gpu_tensor_write(gpu[b], 0, input[b], bytes[b]), what); + require_ok(ds4_gpu_tensor_fill_f32(gpu[OUT], poison, (uint64_t)n * P), what); + if (variants[variant].batch) require_ok(ds4_gpu_begin_commands(), what); + const unsigned chunks = variants[variant].split ? 2 : 1; + for (unsigned chunk = 0; chunk < chunks; chunk++) { + const uint32_t offset = chunk == 0 ? 0 : 33; + const uint32_t rows = chunks == 1 ? n : chunk == 0 ? 33 : n - 33; + ds4_gpu_tensor *views[6]; + for (unsigned b = Q; b <= BETA; b++) { + const uint64_t stride = (b == BETA ? H : P) * sizeof(float); + views[b] = ds4_gpu_tensor_view(gpu[b], (uint64_t)offset * stride, (uint64_t)rows * stride); + require_ok(views[b] != NULL, what); + } + ds4_gpu_tensor *out = ds4_gpu_tensor_view(gpu[OUT], (uint64_t)offset * P * sizeof(float), (uint64_t)rows * P * sizeof(float)); + require_ok(out != NULL, what); + require_ok(ds4_gpu_glm53_kda_prefill(out, gpu[CONV], gpu[STATE], + views[Q], views[K], views[V], views[GATE], views[BETA], views[OGATE], + model, model_bytes, q_conv_offset, k_conv_offset, v_conv_offset, + a_log_offset, dt_bias_offset, norm_offset, H, rows, -5.0f, 1e-5f), what); + ds4_gpu_tensor_free(out); + for (unsigned b = Q; b <= BETA; b++) ds4_gpu_tensor_free(views[b]); + } + require_ok(ds4_gpu_end_commands() == (variants[variant].batch ? 1 : 0), "KDA preserves command-batch ownership"); + const uint32_t largest_chunk = chunks == 1 ? n : 33; + const uint32_t expected_dispatches = variants[variant].rollback ? 0u : + (block != 0u && largest_chunk > block ? DS4_GPU_GLM53_PREFILL_KDA_PREPARE : 0u) | + (block != 0u && largest_chunk >= 32u ? DS4_GPU_GLM53_PREFILL_KDA_RECURRENCE : 0u); + require_prefill_dispatch(expected_dispatches, true, what); + for (unsigned j = 0; j < sizeof(compared) / sizeof(compared[0]); j++) { + const unsigned b = compared[j]; + float *dst = variant == 0 ? reference[b] : actual; + require_ok(ds4_gpu_tensor_read(gpu[b], 0, dst, bytes[b]), what); + if (variant == 0 || memcmp(reference[b], actual, bytes[b]) == 0) continue; + for (uint64_t i = 0; i < bytes[b] / sizeof(float); i++) { + if (memcmp(&reference[b][i], &actual[i], sizeof(float)) == 0) continue; + fprintf(stderr, "%s: %s[%llu] %.9g != %.9g\n", what, names[b], + (unsigned long long)i, actual[i], reference[b][i]); + break; + } + exit(1); + } + } + } + unsetenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING"); + unsetenv("DS4_METAL_GLM53_PREFILL_KDA_PREPARE_BLOCK"); + unsetenv("DS4_METAL_GLM53_PREFILL_KDA_VALUES_PER_SG"); + unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_PREPARE"); + unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_KDA_RECURRENCE"); + if (inherited_profile) setenv("DS4_METAL_PROFILE_KDA_PREFILL", "1", 1); + else unsetenv("DS4_METAL_PROFILE_KDA_PREFILL"); + ds4_gpu_test_set_flags(DS4_GPU_TEST_GLM53_PREFILL); + free(actual); + for (unsigned b = 0; b < NBUF; b++) { + ds4_gpu_tensor_free(gpu[b]); + free(reference[b]); + free(input[b]); + } +} + +#endif /* __APPLE__: these oracles exercise Metal-specific dispatch switches. */ + int main(void) { enum { D = 128, @@ -96,7 +1206,41 @@ int main(void) { Q4_OUT = 37, Q4_ROWS = 3, Q8_OFFSET = 60000, - MODEL_BYTES = 65536, + /* Real GLM 5.3 widths: 4096 is kda_{q,k,v}, and 512/1024 the + * low-rank gate projections. */ + WIDE512_OFFSET = 65536, WIDE512_IN = 512, WIDE512_OUT = 4, + WIDE1024_OFFSET = 73728, WIDE1024_IN = 1024, WIDE1024_OUT = 4, + WIDE4096_OFFSET = 90112, WIDE4096_IN = 4096, WIDE4096_OUT = 2, + WIDE_ROWS = 3, + /* Compound HC producer fixture. The f16 and bf16 kernels are two + * instantiations of one template, so they get identical weights in + * both encodings and must agree exactly. */ + HC_N = 16384, HC_MIX = 24, HC_EMBD = 4096, HC_HC = 4, + HC_F16W_OFFSET = 131072, /* HC_N * HC_MIX * 2 = 786432 */ + HC_BF16W_OFFSET = 917504, + HC_SCALE_OFFSET = 1703936, /* 3 floats */ + HC_BASE_OFFSET = 1703968, /* 24 floats */ + HC_NORM_OFFSET = 1704064, /* 4096 floats */ + /* BF16 matvec + HC-expand epilogue fixture */ + FUSED_W_OFFSET = 1720448, /* FUSED_IN * FUSED_OUT * 2 = 131072 */ + FUSED_IN = 1024, FUSED_OUT = 64, FUSED_HC = 4, + /* Q8_0 value rows for the split-vs-generic attention check: + * 16 heads x 8 values x 544 bytes = 69632 */ + SPLIT_V_OFFSET = 1851520, + /* GLM 5.3 attn_k_b for the prefill qk-low oracle: + * 64 heads x 512 rows x 272 bytes = 8912896 */ + QK_LOW_KB_OFFSET = 2097152, + /* Synthetic Q4_K expert matrices for all tail sizes and an empty expert. */ + MOE_GATE_OFFSET = 11010048, + MOE_UP_OFFSET = MOE_GATE_OFFSET + 36 * 256 * 144, + MOE_DOWN_OFFSET = MOE_UP_OFFSET + 36 * 256 * 144, + KP_Q_OFFSET = MOE_DOWN_OFFSET + 36 * 256 * 144, + KP_K_OFFSET = KP_Q_OFFSET + 64 * 128 * 4 * 4, + KP_V_OFFSET = KP_K_OFFSET + 64 * 128 * 4 * 4, + KP_A_OFFSET = KP_V_OFFSET + 64 * 128 * 4 * 4, + KP_DT_OFFSET = KP_A_OFFSET + 64 * 4, + KP_NORM_OFFSET = KP_DT_OFFSET + 64 * 128 * 4, + MODEL_BYTES = KP_NORM_OFFSET + 128 * 4, }; uint8_t *model = mmap(NULL, MODEL_BYTES, PROT_READ | PROT_WRITE, @@ -180,6 +1324,181 @@ int main(void) { for (uint32_t i = 0; i < BF16_ROWS * BF16_OUT; i++) require_close("BF16 prefill matmul", bf16_actual[i], bf16_expected[i], 2e-4f); + /* BF16_IN above is 64; these cover the widths the model actually runs. */ + check_bf16_matmul(model, MODEL_BYTES, WIDE512_OFFSET, WIDE512_IN, + WIDE512_OUT, WIDE_ROWS, "BF16 matmul in_dim=512"); + check_bf16_matmul(model, MODEL_BYTES, WIDE1024_OFFSET, WIDE1024_IN, + WIDE1024_OUT, WIDE_ROWS, "BF16 matmul in_dim=1024"); + check_bf16_matmul(model, MODEL_BYTES, WIDE4096_OFFSET, WIDE4096_IN, + WIDE4096_OUT, WIDE_ROWS, "BF16 matmul in_dim=4096"); + +#ifdef __APPLE__ + /* + * Compound HC producer: the f16 and bf16 kernels share one templated body + * and differ only in how the mix weights are widened. Weights are drawn + * from values with at most seven explicit mantissa bits, so each is exact + * in BOTH half and bfloat16 and the two kernels see bit-identical floats. + * The arithmetic and reduction order are then the same, so the outputs + * must match exactly -- any difference is a bug in one instantiation. + */ + { + static const float exact_both[8] = { + 0.5f, -0.5f, 1.0f, -1.0f, 1.5f, -1.5f, 0.25f, -0.75f + }; + uint16_t *hc_f16 = (uint16_t *)(model + HC_F16W_OFFSET); + uint16_t *hc_bf16 = (uint16_t *)(model + HC_BF16W_OFFSET); + for (uint32_t i = 0; i < (uint32_t)(HC_N * HC_MIX); i++) { + const float w = exact_both[i % 8u] * 0.03125f; + union { float f; uint32_t u; } b = { .f = w }; + hc_f16[i] = f32_to_f16(w); + hc_bf16[i] = (uint16_t)(b.u >> 16); + /* the encodings must round-trip to the same float, or the + * comparison below would be measuring the fixture, not the kernel */ + require_close("HC fixture encoding", f16_to_f32(hc_f16[i]), + bf16_to_f32(hc_bf16[i]), 0.0f); + } + float *hc_scale = (float *)(model + HC_SCALE_OFFSET); + for (int i = 0; i < 3; i++) hc_scale[i] = 0.5f + 0.25f * (float)i; + float *hc_base = (float *)(model + HC_BASE_OFFSET); + for (int i = 0; i < HC_MIX; i++) hc_base[i] = 0.125f * (float)((i % 5) - 2); + float *hc_norm = (float *)(model + HC_NORM_OFFSET); + for (int i = 0; i < HC_EMBD; i++) hc_norm[i] = 1.0f + 0.001f * (float)(i % 7); + + float *hc_x = malloc((size_t)HC_N * sizeof(float)); + require_ok(hc_x != NULL, "HC residual allocation"); + for (int i = 0; i < HC_N; i++) + hc_x[i] = 0.01f * (float)((i % 23) - 11) + 0.002f * (float)(i % 5); + + ds4_gpu_tensor *hc_res = ds4_gpu_tensor_alloc((size_t)HC_N * sizeof(float)); + require_ok(hc_res != NULL, "HC residual tensor"); + require_ok(ds4_gpu_tensor_write(hc_res, 0, hc_x, + (size_t)HC_N * sizeof(float)), + "HC residual write"); + + float out_f16[HC_EMBD], out_bf16[HC_EMBD]; + float nrm_f16[HC_EMBD], nrm_bf16[HC_EMBD]; + float mix_f16[HC_MIX], mix_bf16[HC_MIX]; + for (int pass = 0; pass < 2; pass++) { + ds4_gpu_tensor *mix = ds4_gpu_tensor_alloc(HC_MIX * sizeof(float)); + ds4_gpu_tensor *spl = ds4_gpu_tensor_alloc(HC_MIX * sizeof(float)); + ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(HC_EMBD * sizeof(float)); + ds4_gpu_tensor *nrm = ds4_gpu_tensor_alloc(HC_EMBD * sizeof(float)); + require_ok(mix && spl && out && nrm, "HC output tensors"); + const int rc = pass == 0 + ? ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( + mix, out, nrm, spl, hc_res, model, MODEL_BYTES, + HC_F16W_OFFSET, HC_SCALE_OFFSET, HC_BASE_OFFSET, + HC_NORM_OFFSET, HC_N, HC_MIX, HC_EMBD, HC_HC, + 1u, 1.0e-6f, 1.0e-3f, 1.0e-6f) + : ds4_gpu_hc_rms_norm_mix_split_norm_bf16_tensor( + mix, out, nrm, spl, hc_res, model, MODEL_BYTES, + HC_BF16W_OFFSET, HC_SCALE_OFFSET, HC_BASE_OFFSET, + HC_NORM_OFFSET, HC_N, HC_MIX, HC_EMBD, HC_HC, + 1u, 1.0e-6f, 1.0e-3f, 1.0e-6f); + require_ok(rc > 0, pass == 0 ? "HC producer f16" : "HC producer bf16"); + require_ok(ds4_gpu_tensor_read(mix, 0, + pass == 0 ? mix_f16 : mix_bf16, sizeof(mix_f16)), + "HC mix read"); + require_ok(ds4_gpu_tensor_read(out, 0, + pass == 0 ? out_f16 : out_bf16, sizeof(out_f16)), + "HC collapse read"); + require_ok(ds4_gpu_tensor_read(nrm, 0, + pass == 0 ? nrm_f16 : nrm_bf16, sizeof(nrm_f16)), + "HC pre-norm read"); + ds4_gpu_tensor_free(mix); + ds4_gpu_tensor_free(spl); + ds4_gpu_tensor_free(out); + ds4_gpu_tensor_free(nrm); + } + for (int i = 0; i < HC_MIX; i++) + require_close("HC producer f16 vs bf16 mix", mix_bf16[i], mix_f16[i], 0.0f); + for (int i = 0; i < HC_EMBD; i++) { + require_close("HC producer f16 vs bf16 collapse", out_bf16[i], out_f16[i], 0.0f); + require_close("HC producer f16 vs bf16 pre-norm", nrm_bf16[i], nrm_f16[i], 0.0f); + } + ds4_gpu_tensor_free(hc_res); + free(hc_x); + } + + /* + * BF16 matvec with the HC expansion folded into its epilogue must equal + * the separate matvec followed by ds4_gpu_hc_expand_tensor, exactly. The + * fused kernel reuses the same row accumulation and repeats the expand + * arithmetic in the same operand order, so anything but bit-identical + * output is a bug -- most likely a stride or an index. + */ + { + uint16_t *fw = (uint16_t *)(model + FUSED_W_OFFSET); + for (uint32_t o = 0; o < FUSED_OUT; o++) { + for (uint32_t i = 0; i < FUSED_IN; i++) { + fw[(size_t)o * FUSED_IN + i] = f32_to_bf16( + 0.003f * (float)((int)((o * 7u + i) % 17u) - 8)); + } + } + float fx[FUSED_IN], fres[FUSED_HC * FUSED_OUT]; + float fpost[FUSED_HC], fcomb[FUSED_HC * FUSED_HC]; + for (int i = 0; i < FUSED_IN; i++) + fx[i] = 0.01f * (float)((i % 19) - 9); + for (int i = 0; i < FUSED_HC * FUSED_OUT; i++) + fres[i] = 0.05f * (float)((i % 13) - 6); + for (int i = 0; i < FUSED_HC; i++) fpost[i] = 0.25f + 0.125f * (float)i; + for (int i = 0; i < FUSED_HC * FUSED_HC; i++) + fcomb[i] = 0.1f * (float)((i % 7) - 3); + + ds4_gpu_tensor *tx = ds4_gpu_tensor_alloc(sizeof(fx)); + ds4_gpu_tensor *tres = ds4_gpu_tensor_alloc(sizeof(fres)); + ds4_gpu_tensor *tpost = ds4_gpu_tensor_alloc(sizeof(fpost)); + ds4_gpu_tensor *tcomb = ds4_gpu_tensor_alloc(sizeof(fcomb)); + ds4_gpu_tensor *out_ref = ds4_gpu_tensor_alloc(FUSED_OUT * sizeof(float)); + ds4_gpu_tensor *hc_ref = ds4_gpu_tensor_alloc(sizeof(fres)); + ds4_gpu_tensor *out_fus = ds4_gpu_tensor_alloc(FUSED_OUT * sizeof(float)); + ds4_gpu_tensor *hc_fus = ds4_gpu_tensor_alloc(sizeof(fres)); + require_ok(tx && tres && tpost && tcomb && out_ref && hc_ref && + out_fus && hc_fus, "fused epilogue tensors"); + require_ok(ds4_gpu_tensor_write(tx, 0, fx, sizeof(fx)) && + ds4_gpu_tensor_write(tres, 0, fres, sizeof(fres)) && + ds4_gpu_tensor_write(tpost, 0, fpost, sizeof(fpost)) && + ds4_gpu_tensor_write(tcomb, 0, fcomb, sizeof(fcomb)), + "fused epilogue inputs"); + + require_ok(ds4_gpu_glm53_matmul_bf16( + out_ref, model, MODEL_BYTES, FUSED_W_OFFSET, + FUSED_IN, FUSED_OUT, tx, 1), + "reference BF16 matvec"); + require_ok(ds4_gpu_hc_expand_tensor(hc_ref, out_ref, tres, tpost, tcomb, + FUSED_OUT, FUSED_HC), + "reference HC expand"); + + const int fused = ds4_gpu_glm53_matmul_bf16_hc_expand4( + out_fus, hc_fus, model, MODEL_BYTES, FUSED_W_OFFSET, + FUSED_IN, FUSED_OUT, tx, tres, tpost, tcomb, FUSED_HC); + if (fused == 0) { + fprintf(stderr, + "BF16 matvec + HC expand: not available on this device, skipped\n"); + } else { + float a[FUSED_OUT], b[FUSED_OUT]; + float ha[FUSED_HC * FUSED_OUT], hb[FUSED_HC * FUSED_OUT]; + require_ok(ds4_gpu_tensor_read(out_ref, 0, a, sizeof(a)) && + ds4_gpu_tensor_read(out_fus, 0, b, sizeof(b)) && + ds4_gpu_tensor_read(hc_ref, 0, ha, sizeof(ha)) && + ds4_gpu_tensor_read(hc_fus, 0, hb, sizeof(hb)), + "fused epilogue readback"); + for (int i = 0; i < FUSED_OUT; i++) + require_close("fused epilogue block_out", b[i], a[i], 0.0f); + for (int i = 0; i < FUSED_HC * FUSED_OUT; i++) + require_close("fused epilogue hc stream", hb[i], ha[i], 0.0f); + } + ds4_gpu_tensor_free(tx); + ds4_gpu_tensor_free(tres); + ds4_gpu_tensor_free(tpost); + ds4_gpu_tensor_free(tcomb); + ds4_gpu_tensor_free(out_ref); + ds4_gpu_tensor_free(hc_ref); + ds4_gpu_tensor_free(out_fus); + ds4_gpu_tensor_free(hc_fus); + } +#endif /* __APPLE__: fused HC producers and epilogues are Metal-only. */ + #ifdef DS4_ROCM_BUILD test_block_q4_K *q4_weights = (test_block_q4_K *)(model + Q4_OFFSET); for (uint32_t o = 0; o < Q4_OUT; o++) { @@ -513,6 +1832,8 @@ int main(void) { free(f32_attn_cache); free(f32_attn_q); free(f32_attn_low); + + check_split_dsa_attention(model, MODEL_BYTES, SPLIT_V_OFFSET); #endif #ifdef DS4_ROCM_BUILD @@ -998,6 +2319,26 @@ int main(void) { ds4_gpu_tensor_free(q); ds4_gpu_tensor_free(bf16_out); ds4_gpu_tensor_free(bf16_x); +#ifdef __APPLE__ + /* Never silently compare the reference with itself because the parent + * process has disabled tuning. Each oracle asserts dispatch coverage. */ + require_ok(unsetenv("DS4_METAL_DISABLE_GLM53_FLASH_TUNING") == 0, + "clear inherited aggregate tuning switch for kernel oracles"); + require_ok(unsetenv("DS4_METAL_DISABLE_GLM53_PREFILL_INDEXED_ATTN") == 0, + "clear inherited attention tuning switch for kernel oracles"); + ds4_gpu_test_set_flags(DS4_GPU_TEST_GLM53_PREFILL); + ds4_gpu_test_glm53_prefill_take_dispatches(); + check_glm53_qk_lowrank_token_tile(model, MODEL_BYTES, QK_LOW_KB_OFFSET); + check_glm53_qk_lowrank_large_offsets(model, MODEL_BYTES, QK_LOW_KB_OFFSET); + check_glm53_indexed_attention_head_width(); + check_glm53_indexed_attention_invalid_rows(); + check_glm53_routed_moe_tail_cull(model, MODEL_BYTES, MOE_GATE_OFFSET, + MOE_UP_OFFSET, MOE_DOWN_OFFSET); + check_glm53_kda_prepare_blocked(model, MODEL_BYTES, KP_Q_OFFSET, + KP_K_OFFSET, KP_V_OFFSET, KP_A_OFFSET, + KP_DT_OFFSET, KP_NORM_OFFSET); + ds4_gpu_test_set_flags(0); +#endif ds4_gpu_cleanup(); munmap(model, MODEL_BYTES); puts("GLM-5.3 KDA GPU tests: PASS");