url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

serve mux: keep ragged decode batches off the fused Metal kernels

Browse Source 
With METAL=1 + COLI_METAL=1, run_serve_mux (SERVE_BATCH=1) truncated
every completion to exactly 1 token: step_decode_batch passes per-row
kvs[]/positions[] with pos_base=0, but the two Metal decode fast paths
(attention_rows and the FULL-LAYER CB in layer_forward_rows) ignored
them and dispatched coli_metal_attn_decode/coli_metal_layer_decode with
the model-bound Lc/Rc and the hardcoded pos_base. The kernels' contract
is one sequence, row s at pos_base+s: ragged rows got roped at position
0 and attended over a T=1 window of the wrong cache, so greedy decode
hit a stop token on the first batched step (DONE ... STAT 1).

Gate both fast paths on !kvs. Ragged mux rows now take the CPU absorb
path, which already reads kvs[s]/positions[s]/ks->kv_start per row;
plain serve, chat/run, prefill and MTP verification (kvs==NULL) keep
the fused GPU kernels unchanged.

Verified on GLM-5.2 int4 (M5 Max): mux with COLI_METAL=1 went from
1-token DONE to full 16/16-token greedy completions, byte-identical to
the plain-serve comparator on both test prompts; CPU-only mux was
already correct (bisection); test-c and metal-test pass.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
monotophic
2026-07-15 05:54:50 -04:00
parent 62419af188
commit 0811730845
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/glm.c
+12 -4
View File
@@ -1789,8 +1789,14 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
double ta0=now_s();
#ifdef COLI_METAL
/* Fused decode attention on GPU: whole layer in one command buffer (keeps the GPU hot).
* S<=4 absorption path with st0==0, DSA selection inactive, and GLM-5.2 int4 dims. */
if(g_metal_enabled && S<=4 && (g_absorb==1||(g_absorb<0&&S<=4)) && m->kv_start[layer]==0
* S<=4 absorption path with st0==0, DSA selection inactive, and GLM-5.2 int4 dims.
* RAGGED GUARD (!kvs): the kernel takes ONE Lc/Rc pair and ONE pos_base — it assumes
* row s is token pos_base+s of the SAME sequence. The batched mux decode
* (step_decode_batch) passes per-row kvs[]/positions[] with pos_base=0, so the kernel
* would rope every row at position 0 and attend over a 1-token window of the wrong
* cache -> greedy decode hits EOS at token 2 (mux answers truncated to 1 token).
* Ragged rows take the CPU absorb path below, which reads kvs[s]/positions[s]. */
if(g_metal_enabled && !kvs && S<=4 && (g_absorb==1||(g_absorb<0&&S<=4)) && m->kv_start[layer]==0
&& D==6144 && H==64 && c->q_lora==2048 && c->kv_lora==512 && c->qk_nope==192
&& c->qk_rope==64 && vh==256 && l->kv_b.fmt==2){
int sel_active = m->has_dsa && layer<c->n_layers && c->idx_type[layer] && (pos_base+S) > c->index_topk;
@@ -2925,8 +2931,10 @@ static void layer_forward_rows(Model *m, Layer *l, int li, float *x, int S, int
#ifdef COLI_METAL
/* FULL-LAYER CB: in_ln + attention + residuo + post_ln + shared expert + router/top-K
* in un solo submit GPU; la CPU legge il routing e fa solo resolve/disk/expert-CB.
* Fallback: qualsiasi condizione mancante -> percorso CPU intero qui sotto. */
if(g_metal_enabled && S<=4 && li<c->n_layers && l->sparse
* Fallback: qualsiasi condizione mancante -> percorso CPU intero qui sotto.
* !kvs: ragged mux rows (per-row KV/position) are not expressible in this kernel's
* single Lc/Rc + pos_base contract — see the matching guard in attention_rows. */
if(g_metal_enabled && !kvs && S<=4 && li<c->n_layers && l->sparse
&& (g_absorb==1||(g_absorb<0&&S<=4)) && m->kv_start[li]==0
&& D==6144 && c->n_heads==64 && c->q_lora==2048 && c->kv_lora==512
&& c->qk_nope==192 && c->qk_rope==64 && c->v_head==256 && l->kv_b.fmt==2