fix(qwen): retain prefill attention arithmetic across Metal partitions

A Qwen prefill crossing the 512-block sparse-attention budget leaves three
internal dense rows at absolute positions 2048 through 2050. With chunk 128
or chunk 2048, these rows previously selected the short-tail decode kernel
because their subrange contained fewer than nine rows. A single-pass
prefill selected matrix attention for the same rows. The different
rounding introduced drift despite identical input tokens and queries.

Carry the parent prefill policy through both dense and sparse attention
subranges. Add a Metal prefill entry point backed by the existing shared
implementation so internal short partitions retain matrix attention when
the parent prefill has more than eight rows. Preserve the existing shape
guards and scalar diagnostic override. Genuine short prefill tails,
decode, speculative verification and non-Apple dispatch retain their
previous policies; shader arithmetic and cache storage are unchanged.

Add regression coverage comparing a whole 2051-row attention reference
with the three-row dense partition, and sparse partitions of one through
eight rows. Require bit-identical outputs, bounded tensor views and intact
output guards. Include the previous scalar dispatch as a negative control
and verify that the diagnostic override still selects its arithmetic.

Validation on Apple M1 Max, 32 GiB, Qwen3.8-Flash-Next-Q2, SSD streaming:
- Production and affected test executables build successfully.
- Complete Qwen Metal kernel suite passes with MTL_DEBUG_LAYER=1.
- Memory/workspace, generation and MTP prefill suites pass. Generation
  includes a genuine three-token tail and the legacy prefill reference;
  MTP checks use the suite's existing numerical tolerances.
- For a 5760-token prompt and 128 teacher-forced decode steps, requested
  chunks 128, 2048 and 8192 each match the original single-pass reference
  bit for bit at all 129 full-vocabulary frontiers: 96,099,840 logit-pair
  comparisons, with MTP disabled and a 1024-expert cache.
- The original chunked prefill had maximum absolute logit error 0.597853661
  against that single pass; the candidate has zero on this fixture.
- Per-layer traces locate the first divergence at layer 3 attention output
  for token 2048, after identical inputs and queries. Traced final captures
  reproduce the uninstrumented captures exactly.
- git diff --check passes.

This targets internal attention partitions rather than promising equality
for arbitrary chunk sizes: genuine small tails and very short sparse
indexer scoring retain their existing policies. The single-pass comparison
is a numerical consistency check, not a model-quality or throughput claim.

Refs: https://github.com/antirez/ds4/pull/1056#issuecomment-5731630396

ds4.c

@@ -58451,6 +58451,26 @@ static bool qwen4_graph_linear(ds4_qwen4_gpu_graph *g, const ds4_model *m, const
     return ok;
 }
 
+/* Preserve the parent attention policy across dense/sparse subranges. In a
+ * large prefill the three dense rows at positions 2048..2050 still belong
+ * to matrix attention; their short internal range must not select decode
+ * arithmetic. Real short prefill tails and speculative rows are unchanged. */
+static bool qwen4_graph_attention_dispatch(ds4_qwen4_gpu_graph *g, uint32_t il,
+        ds4_gpu_tensor *out, ds4_gpu_tensor *q, ds4_gpu_tensor *gate,
+        uint32_t rows, uint32_t pos0, bool use_sel, uint32_t parent_rows) {
+    const float scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM);
+#ifdef __APPLE__
+    if (g->projection_phase == QWEN4_PROJECTION_PREFILL && parent_rows > 8u) {
+        return ds4_gpu_qwen4_attn_prefill_tensor(out, q, gate, g->layer_k_cache[il],
+                g->layer_v_cache[il], g->sel_tokens, g->n_sel, rows, DS4_N_HEAD,
+                DS4_N_HEAD_KV, DS4_N_HEAD_DIM, pos0, use_sel, g->sel_stride, scale) != 0;
+    }
+#endif
+    return ds4_gpu_qwen4_attn_decode_tensor(out, q, gate, g->layer_k_cache[il],
+            g->layer_v_cache[il], g->sel_tokens, g->n_sel, parent_rows <= 2u ? g->attn_part : NULL,
+            rows, DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, pos0, use_sel, g->sel_stride, scale) != 0;
+}
+
 /* T tokens at positions pos0.. of one attention layer.  Tokens with no more
  * complete blocks than the budget attend densely; the rest score the pooled
  * block keys and attend their top-k blocks plus the incomplete tail.  The
@@ -58464,13 +58484,9 @@ static bool qwen4_graph_attention_core(ds4_qwen4_gpu_graph *g, uint32_t il,
                                        uint32_t n_blocks_after, uint32_t cpos0, uint32_t cT) {
     const uint32_t ratio = 4u;
     const uint32_t q_dim = DS4_N_HEAD * DS4_N_HEAD_DIM;
-    const float scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM);
     const uint32_t n_dense = qwen4_attention_dense_rows(cpos0, cT, g->k_blocks);
     if (n_dense > 0 &&
-        !ds4_gpu_qwen4_attn_decode_tensor(o_rows, q_rows, gate_rows, g->layer_k_cache[il], g->layer_v_cache[il],
-                                          g->sel_tokens, g->n_sel, cT <= 2u ? g->attn_part : NULL, n_dense,
-                                          DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, cpos0, false, g->sel_stride,
-                                          scale)) {
+        !qwen4_graph_attention_dispatch(g, il, o_rows, q_rows, gate_rows, n_dense, cpos0, false, cT)) {
         return false;
     }
     if (n_dense >= cT) return true;
@@ -58493,10 +58509,7 @@ static bool qwen4_graph_attention_core(ds4_qwen4_gpu_graph *g, uint32_t il,
                          n_blocks_after, n_sparse, g->k_blocks) &&
         ds4_gpu_qwen4_idx_expand_tensor(g->sel_tokens, g->n_sel, g->sel_blocks, n_sparse, g->k_blocks, ratio,
                                         sp0, g->sel_stride) &&
-        ds4_gpu_qwen4_attn_decode_tensor(o, q, gate, g->layer_k_cache[il], g->layer_v_cache[il],
-                                         g->sel_tokens, g->n_sel, cT <= 2u ? g->attn_part : NULL, n_sparse,
-                                         DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, sp0, true, g->sel_stride,
-                                         scale);
+        qwen4_graph_attention_dispatch(g, il, o, q, gate, n_sparse, sp0, true, cT);
     ds4_gpu_tensor_free(iqn);
     ds4_gpu_tensor_free(o);
     ds4_gpu_tensor_free(gate);

ds4_gpu.h

@@ -3468,6 +3468,16 @@ int ds4_gpu_qwen4_attn_decode_tensor(
         const ds4_gpu_tensor *sel_tokens, const ds4_gpu_tensor *n_sel, ds4_gpu_tensor *part,
         uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim,
         uint32_t pos0, bool use_sel, uint32_t sel_stride, float scale);
+#ifdef __APPLE__
+/* Metal: a subrange of a matrix-attention prefill retains its parent kernel
+ * arithmetic even when the dense/sparse partition contains only 1..8 rows. */
+int ds4_gpu_qwen4_attn_prefill_tensor(
+        ds4_gpu_tensor *out, const ds4_gpu_tensor *q, const ds4_gpu_tensor *gate,
+        const ds4_gpu_tensor *k_cache, const ds4_gpu_tensor *v_cache,
+        const ds4_gpu_tensor *sel_tokens, const ds4_gpu_tensor *n_sel,
+        uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim,
+        uint32_t pos0, bool use_sel, uint32_t sel_stride, float scale);
+#endif
 /* Routed experts; shared_type == UINT32_MAX disables the shared-expert slot,
  * otherwise mid/part carry n_slots+1 entries and the reduce weights the last
  * one by sigmoid(shared_gate). */

ds4_metal.m

@@ -49631,12 +49631,12 @@ uint64_t ds4_gpu_qwen4_attn_part_floats(uint32_t n_tokens, uint32_t n_head, uint
     return (uint64_t)n_tokens * n_head * QWEN4_ATTN_MAX_SPLITS * (2u + head_dim);
 }
 
-int ds4_gpu_qwen4_attn_decode_tensor(
+static int qwen4_attn_tensor_impl(
         ds4_gpu_tensor *out, const ds4_gpu_tensor *q, const ds4_gpu_tensor *gate,
         const ds4_gpu_tensor *k_cache, const ds4_gpu_tensor *v_cache,
         const ds4_gpu_tensor *sel_tokens, const ds4_gpu_tensor *n_sel, ds4_gpu_tensor *part,
         uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim,
-        uint32_t pos0, bool use_sel, uint32_t sel_stride, float scale) {
+        uint32_t pos0, bool use_sel, uint32_t sel_stride, float scale, bool prefill_mm) {
     const uint32_t n_keys = use_sel ? sel_stride : pos0 + n_tokens;
     uint32_t n_splits = 1;
     if (part) {
@@ -49677,10 +49677,10 @@ int ds4_gpu_qwen4_attn_decode_tensor(
     } else {
         b[7] = b[6];
     }
-    /* Keep short prefill tails on the decode arithmetic. Rounding queries
-     * to half for just a few rows can flip nearly tied expert selections
-     * relative to one-token continuation; large prefills retain matrix tiles. */
-    if (n_splits == 1 && n_tokens > 8u && head_dim == 256u && n_head / n_head_kv <= 16u &&
+    /* Genuine short tails keep their decode arithmetic. A short dense/sparse
+     * subrange of a larger prefill must instead retain the parent's matrix
+     * arithmetic: its row count is a partition detail, not a phase change. */
+    if (n_splits == 1 && (prefill_mm || n_tokens > 8u) && head_dim == 256u && n_head / n_head_kv <= 16u &&
         getenv("DS4_QWEN4_NO_ATTN_MM") == NULL) {
         return qwen4_dispatch(QWEN4_K_ATTN_MM, &args, sizeof(args), b, 7,
                               MTLSizeMake(n_head_kv, n_tokens, 1), MTLSizeMake(128, 1, 1), 0);
@@ -49709,6 +49709,26 @@ int ds4_gpu_qwen4_attn_decode_tensor(
                           MTLSizeMake(n_head, n_tokens, 1), MTLSizeMake(32, 1, 1), 0);
 }
 
+int ds4_gpu_qwen4_attn_decode_tensor(
+        ds4_gpu_tensor *out, const ds4_gpu_tensor *q, const ds4_gpu_tensor *gate,
+        const ds4_gpu_tensor *k_cache, const ds4_gpu_tensor *v_cache,
+        const ds4_gpu_tensor *sel_tokens, const ds4_gpu_tensor *n_sel, ds4_gpu_tensor *part,
+        uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim,
+        uint32_t pos0, bool use_sel, uint32_t sel_stride, float scale) {
+    return qwen4_attn_tensor_impl(out, q, gate, k_cache, v_cache, sel_tokens, n_sel, part,
+                                 n_tokens, n_head, n_head_kv, head_dim, pos0, use_sel, sel_stride, scale, false);
+}
+
+int ds4_gpu_qwen4_attn_prefill_tensor(
+        ds4_gpu_tensor *out, const ds4_gpu_tensor *q, const ds4_gpu_tensor *gate,
+        const ds4_gpu_tensor *k_cache, const ds4_gpu_tensor *v_cache,
+        const ds4_gpu_tensor *sel_tokens, const ds4_gpu_tensor *n_sel,
+        uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim,
+        uint32_t pos0, bool use_sel, uint32_t sel_stride, float scale) {
+    return qwen4_attn_tensor_impl(out, q, gate, k_cache, v_cache, sel_tokens, n_sel, NULL,
+                                 n_tokens, n_head, n_head_kv, head_dim, pos0, use_sel, sel_stride, scale, true);
+}
+
 /* ---- decode batch: attention over rows ----
  *
  * The caches stay per session, so the rows kernels find each row's through

tests/test_qwen4_kernels.c

@@ -983,6 +983,118 @@ static void test_attn_mm(uint32_t T, uint32_t pos0, bool sparse) {
     test_attn_mm_keys(T, pos0, sparse, 6, 24, false);
 }
 
+#ifdef __APPLE__
+/* A large prefill split at the dense/sparse boundary leaves three dense
+ * rows at positions 2048..2050. Those internal rows must keep the parent's
+ * matrix attention arithmetic, not inherit the short-tail decode policy. */
+static void test_attn_prefill_partition_case(bool sparse) {
+    const uint32_t T = sparse ? 17u : 2051u, pos0 = sparse ? 2051u : 0u;
+    const uint32_t H = 8u, Hkv = 1u, D = 256u, budget = 512u;
+    const uint32_t stride = budget * 4u + 4u, row = H * D, guard = 32u;
+    const uint64_t qn = (uint64_t)T * row, kvn = (uint64_t)(pos0 + T) * Hkv * D;
+    const uint64_t scratch_n = 8u * row + 2u * guard;
+    float *q = rand_vec(qn, 1.0f), *gate = rand_vec(qn, 2.0f);
+    _Float16 *kc = malloc(kvn * sizeof(*kc)), *vc = malloc(kvn * sizeof(*vc));
+    int32_t *sel = sparse ? malloc((uint64_t)T * stride * sizeof(*sel)) : NULL;
+    uint32_t *cnt = sparse ? malloc((uint64_t)T * sizeof(*cnt)) : NULL;
+    require_ok(q && gate && kc && vc && (!sparse || (sel && cnt)), "prefill partition allocation");
+    for (uint64_t i = 0; i < kvn; i++) {
+        kc[i] = (_Float16)(frand() - 0.5f);
+        vc[i] = (_Float16)(frand() - 0.5f);
+    }
+    if (sparse) for (uint32_t t = 0; t < T; t++) {
+        const uint32_t pos = pos0 + t, visible = (pos + 1u) / 4u;
+        /* 509 is coprime with each visible block count in this fixture:
+         * distinct selected blocks, with an order different from dense. */
+        for (uint32_t j = 0; j < budget; j++) for (uint32_t k = 0; k < 4u; k++)
+            sel[(uint64_t)t * stride + 4u * j + k] = (int32_t)(4u * ((509u * j + t) % visible) + k);
+        cnt[t] = 4u * budget;
+        for (uint32_t p = 4u * visible; p <= pos; p++) sel[(uint64_t)t * stride + cnt[t]++] = (int32_t)p;
+        for (uint32_t j = cnt[t]; j < stride; j++) sel[(uint64_t)t * stride + j] = -1;
+    }
+    ds4_gpu_tensor *gq = upload(q, qn), *ggate = upload(gate, qn), *gref = upload(NULL, qn);
+    ds4_gpu_tensor *gk = ds4_gpu_tensor_alloc(kvn * sizeof(*kc));
+    ds4_gpu_tensor *gv = ds4_gpu_tensor_alloc(kvn * sizeof(*vc));
+    ds4_gpu_tensor *gsel = sparse ? ds4_gpu_tensor_alloc((uint64_t)T * stride * sizeof(*sel)) : NULL;
+    ds4_gpu_tensor *gcnt = sparse ? ds4_gpu_tensor_alloc((uint64_t)T * sizeof(*cnt)) : NULL;
+    ds4_gpu_tensor *scratch = upload(NULL, scratch_n);
+    require_ok(gk && gv && ds4_gpu_tensor_write(gk, 0, kc, kvn * sizeof(*kc)) &&
+               ds4_gpu_tensor_write(gv, 0, vc, kvn * sizeof(*vc)), "prefill partition KV upload");
+    require_ok(!sparse || (gsel && gcnt &&
+               ds4_gpu_tensor_write(gsel, 0, sel, (uint64_t)T * stride * sizeof(*sel)) &&
+               ds4_gpu_tensor_write(gcnt, 0, cnt, (uint64_t)T * sizeof(*cnt))), "prefill partition selection upload");
+    /* The existing API chooses MM for a whole prefill. Keep it as the
+     * independent dispatch reference instead of calling the new API twice. */
+    require_ok(ds4_gpu_qwen4_attn_decode_tensor(gref, gq, ggate, gk, gv, gsel, gcnt, NULL,
+               T, H, Hkv, D, pos0, sparse, stride, 0.0625f), "whole prefill attention reference");
+    float *ref = download(gref, qn);
+    for (uint64_t i = 0; i < qn; i++) require_ok(isfinite(ref[i]), "finite whole prefill attention");
+    const uint32_t first = sparse ? 1u : 3u, last = sparse ? 8u : 3u;
+    for (uint32_t n = first; n <= last; n++) {
+        const uint32_t t0 = T - n;
+        const uint64_t bytes = (uint64_t)n * row * sizeof(float);
+        ds4_gpu_tensor *vq = ds4_gpu_tensor_view(gq, (uint64_t)t0 * row * sizeof(float), bytes);
+        ds4_gpu_tensor *vgate = ds4_gpu_tensor_view(ggate, (uint64_t)t0 * row * sizeof(float), bytes);
+        ds4_gpu_tensor *vout = ds4_gpu_tensor_view(scratch, guard * sizeof(float), bytes);
+        ds4_gpu_tensor *vsel = sparse ? ds4_gpu_tensor_view(gsel, (uint64_t)t0 * stride * sizeof(*sel),
+                                                            (uint64_t)n * stride * sizeof(*sel)) : NULL;
+        ds4_gpu_tensor *vcnt = sparse ? ds4_gpu_tensor_view(gcnt, (uint64_t)t0 * sizeof(*cnt),
+                                                            (uint64_t)n * sizeof(*cnt)) : NULL;
+        require_ok(vq && vgate && vout && (!sparse || (vsel && vcnt)), "bounded prefill partition views");
+        require_ok(ds4_gpu_tensor_fill_f32(scratch, 17.25f, scratch_n), "prefill partition output guards");
+        require_ok(ds4_gpu_qwen4_attn_prefill_tensor(vout, vq, vgate, gk, gv, vsel, vcnt,
+                   n, H, Hkv, D, pos0 + t0, sparse, stride, 0.0625f), "partitioned prefill attention");
+        float *got = download(scratch, scratch_n);
+        require_ok(memcmp(got + guard, ref + (uint64_t)t0 * row, bytes) == 0,
+                   "prefill partition keeps whole-pass attention bits");
+        for (uint64_t i = 0; i < scratch_n; i++) if (i < guard || i >= guard + (uint64_t)n * row)
+            require_ok(got[i] == 17.25f, "prefill partition output guard intact");
+        free(got);
+        if (!sparse) {
+            /* Negative control: the old three-row dispatch must reproduce
+             * a difference, so this fixture cannot silently miss the bug. */
+            require_ok(ds4_gpu_qwen4_attn_decode_tensor(vout, vq, vgate, gk, gv, NULL, NULL, NULL,
+                       n, H, Hkv, D, pos0 + t0, false, stride, 0.0625f), "short-tail attention control");
+            float *legacy = download(vout, (uint64_t)n * row);
+            double max_error = 0.0;
+            for (uint64_t i = 0; i < (uint64_t)n * row; i++) {
+                require_ok(isfinite(legacy[i]), "finite short-tail attention control");
+                max_error = fmax(max_error, fabs((double)legacy[i] - ref[(uint64_t)t0 * row + i]));
+            }
+            require_ok(max_error > 0.0, "partition fixture detects old short-tail dispatch");
+            setenv("DS4_QWEN4_NO_ATTN_MM", "1", 1);
+            require_ok(ds4_gpu_qwen4_attn_prefill_tensor(vout, vq, vgate, gk, gv, NULL, NULL,
+                       n, H, Hkv, D, pos0 + t0, false, stride, 0.0625f), "prefill attention diagnostic override");
+            unsetenv("DS4_QWEN4_NO_ATTN_MM");
+            float *diagnostic = download(vout, (uint64_t)n * row);
+            require_ok(memcmp(diagnostic, legacy, bytes) == 0, "prefill attention honors scalar override");
+            printf("  prefill dense partition 2048+3 exact; old dispatch max|d|=%.3e\n", max_error);
+            free(diagnostic);
+            free(legacy);
+        }
+        ds4_gpu_tensor_free(vcnt); ds4_gpu_tensor_free(vsel); ds4_gpu_tensor_free(vout);
+        ds4_gpu_tensor_free(vgate); ds4_gpu_tensor_free(vq);
+    }
+    if (sparse) puts("  prefill sparse partitions 1..8 rows exact; output guards intact");
+    free(ref); free(q); free(gate); free(kc); free(vc); free(sel); free(cnt);
+    ds4_gpu_tensor_free(scratch); ds4_gpu_tensor_free(gref); ds4_gpu_tensor_free(gcnt); ds4_gpu_tensor_free(gsel);
+    ds4_gpu_tensor_free(gv); ds4_gpu_tensor_free(gk); ds4_gpu_tensor_free(ggate); ds4_gpu_tensor_free(gq);
+}
+
+static void test_attn_prefill_partition(void) {
+    const uint32_t saved_rng = g_rng;
+    g_rng = 0x4ad34a12u;
+    const char *old = getenv("DS4_QWEN4_NO_ATTN_MM");
+    char *saved = old ? strdup(old) : NULL;
+    require_ok(!old || saved, "save attention diagnostic override");
+    unsetenv("DS4_QWEN4_NO_ATTN_MM");
+    test_attn_prefill_partition_case(false);
+    test_attn_prefill_partition_case(true);
+    if (saved) { setenv("DS4_QWEN4_NO_ATTN_MM", saved, 1); free(saved); }
+    g_rng = saved_rng;
+}
+#endif
+
 #ifndef __APPLE__
 static void test_attn_groups(void) {
     test_attn_mm_keys(32, 4093, false, 6, 24, false);
@@ -4065,7 +4177,9 @@ int main(void) {
     test_attn_mm(4, 128, false);
     test_attn_mm(8, 128, false);
     test_attn_mm(9, 128, false);
-#ifndef __APPLE__
+#ifdef __APPLE__
+    test_attn_prefill_partition();
+#else
     test_attn_groups();
 #endif
     test_gdn(&arena, 2, 6, 32, 7);