Support DeepSeek V4 Flash 0731

metal/dense.metal

@@ -197,33 +197,6 @@ kernel void kernel_mul_mv_q8_0_f32(
     kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }
 
-[[host_name("kernel_mul_mv_q8_0_f32_r4")]]
-kernel void kernel_mul_mv_q8_0_f32_r4(
-        constant ds4_metal_args_mul_mv & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-    kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
-}
-
-// Output projection alias used by the optimized host dispatch.
-[[host_name("kernel_mul_mv_q8_0_f32_nr4")]]
-kernel void kernel_mul_mv_q8_0_f32_nr4(
-        constant ds4_metal_args_mul_mv & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-    kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>(
-        args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
-}
 
 // Decode Q-A/KV pair. Both projections consume the same activation row but
 // have independent weight ranges and output extents. Keep the standalone Q8_0
@@ -497,25 +470,190 @@ kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0(
             clamp_value, shmem, tgpig, tiisg, sgitg);
 }
 
-[[host_name("kernel_dsv4_shared_gate_up_swiglu_q8_0_r4")]]
+// Decode-only fusion of the router logits matvec (F16, embd -> n_expert)
-kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0_r4(
+// with the shared-expert gate/up SwiGLU (Q8_0, embd -> shared).  Both read
+// the same normalized FFN input back to back; one dispatch removes one
+// launch per decode layer.  Router threadgroups replicate
+// kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2); shared threadgroups host two
+// virtual 4-simdgroup cohorts replicating
+// kernel_dsv4_shared_gate_up_swiglu_q8_0 (nsg=4, nr0=2), including its
+// per-row simd/shmem reduction trees.  Bit-exact by construction.
+kernel void kernel_dsv4_router_shared_gate_up_q8_0(
         constant ds4_metal_args_mul_mv & args,
+        constant ds4_metal_args_mul_mv & sargs,
+        device const char * src0_router,
         device const char * src0_gate,
         device const char * src0_up,
         device const char * src1,
+        device       char * dst_router,
         device       char * dst_gate,
         device       char * dst_up,
         device       char * dst_mid,
         constant     float &clamp_value,
         threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig[[threadgroup_position_in_grid]],
+        uint3  tgpig [[threadgroup_position_in_grid]],
-        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort tiisg [[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
-    kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<4, true>(
+    constexpr short NW = N_SIMDWIDTH;
-            args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid,
+    const uint router_tgs = ((uint)args.ne01 + 1u) / 2u;
-            clamp_value, shmem, tgpig, tiisg, sgitg);
+
+    if (tgpig.x < router_tgs) {
+        // Exact replica of kernel_mul_mv_f16_f32_4 with NSG=8, NR0=2.
+        constexpr short NSG = 8;
+        constexpr short NR0 = 2;
+        constexpr short NB  = 32;
+        constexpr short NF  = 16;
+        constexpr short NF4 = NF/4;
+
+        const int nb = args.ne00/NB;
+        const int r0 = tgpig.x*NR0;
+
+        device const float4 * y4 = (device const float4 *) src1;
+
+        device const half4 * ax4[NR0];
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            ax4[row] = (device const half4 *)
+                (src0_router + (uint64_t)(r0 + row)*args.nb01);
+        }
+
+        float sumf[NR0] = { 0.f };
+
+        const short ix = tiisg/(NW/NF);
+        const short il = tiisg%(NW/NF);
+        const int ib0 = sgitg*NF + ix;
+
+        device const float4 * yb4 = y4 + (ib0*NB + il*NF)/4;
+
+        for (int ib = ib0; ib < nb; ib += NSG*NF) {
+            float4 yl4[NF4];
+            FOR_UNROLL (short i = 0; i < NF4; ++i) {
+                yl4[i] = yb4[i];
+            }
+
+            FOR_UNROLL (short row = 0; row < NR0; row++) {
+                device const half4 * 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]);
+                }
+
+                sumf[row] += sumq;
+            }
+
+            yb4 += NSG*NF*NW/4;
+        }
+
+        device float * dst_f32 = (device float *) dst_router;
+        helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01,
+                                        tiisg, sgitg, shmem);
+        return;
+    }
+
+    // Shared-expert part: two virtual nsg=4 cohorts per threadgroup, each an
+    // exact replica of kernel_dsv4_shared_gate_up_swiglu_q8_0 (NR0=2).
+    constexpr short NSG = 4;
+    constexpr short NR0 = 2;
+    constexpr short NQ  = 8;
+
+    const uint   cohort = sgitg >> 2;
+    const ushort vsg    = sgitg & 3u;
+    const uint   vt     = (tgpig.x - router_tgs) * 2u + cohort;
+
+    const int nb = sargs.ne00 / QK8_0;
+    const int r0 = vt * NR0;
+
+    device const float *y = (device const float *) src1;
+
+    device const block_q8_0 *ag[NR0];
+    device const block_q8_0 *au[NR0];
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        const uint64_t offset0 = (uint64_t)(r0 + row) * sargs.nb01;
+        ag[row] = (device const block_q8_0 *)(src0_gate + offset0);
+        au[row] = (device const block_q8_0 *)(src0_up   + offset0);
+    }
+
+    float sumg[NR0] = { 0.f };
+    float sumu[NR0] = { 0.f };
+
+    const short ix = tiisg / (NW / NQ);
+    const short il = tiisg % (NW / NQ);
+    const int ib0 = vsg * NQ + ix;
+    float yl[NQ];
+    device const float *yb = y + ib0 * QK8_0 + il * NQ;
+
+    for (int ib = ib0; ib < nb; ib += NSG * NQ) {
+        FOR_UNROLL (short i = 0; i < NQ; ++i) {
+            yl[i] = yb[i];
+        }
+
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            device const int8_t *qg = ag[row][ib].qs + il * NQ;
+            device const int8_t *qu = au[row][ib].qs + il * NQ;
+
+            float sg = 0.f;
+            float su = 0.f;
+            FOR_UNROLL (short i = 0; i < NQ; ++i) {
+                sg += qg[i] * yl[i];
+                su += qu[i] * yl[i];
+            }
+
+            sumg[row] += sg * ag[row][ib].d;
+            sumu[row] += su * au[row][ib].d;
+        }
+
+        yb += NSG * NQ * QK8_0;
+    }
+
+    threadgroup float *shmem_f32 = (threadgroup float *)shmem + cohort * (2*NR0*NW);
+    threadgroup float *sh_gate[NR0];
+    threadgroup float *sh_up[NR0];
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        sh_gate[row] = shmem_f32 + NW * row;
+        sh_up[row]   = shmem_f32 + NW * (NR0 + row);
+        if (vsg == 0) {
+            sh_gate[row][tiisg] = 0.0f;
+            sh_up[row][tiisg] = 0.0f;
+        }
+        sumg[row] = simd_sum(sumg[row]);
+        sumu[row] = simd_sum(sumu[row]);
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        if (tiisg == 0) {
+            sh_gate[row][vsg] = sumg[row];
+            sh_up[row][vsg] = sumu[row];
+        }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    device float *gate_f32 = (device float *)dst_gate;
+    device float *up_f32   = (device float *)dst_up;
+    device float *mid_f32  = (device float *)dst_mid;
+
+    FOR_UNROLL (short row = 0; row < NR0 && r0 + row < sargs.ne01; ++row) {
+        const float gate = simd_sum(sh_gate[row][tiisg]);
+        const float up = simd_sum(sh_up[row][tiisg]);
+        if (tiisg == 0 && vsg == 0) {
+            const uint out_row = r0 + row;
+            gate_f32[out_row] = gate;
+            up_f32[out_row] = up;
+            float g = gate;
+            float u = up;
+            if (clamp_value > 1.0e-6f) {
+                g = min(g, clamp_value);
+                u = clamp(u, -clamp_value, clamp_value);
+            }
+            const float silu = g / (1.0f + exp(-g));
+            mid_f32[out_row] = silu * u;
+        }
+    }
 }
 
+
 [[host_name("kernel_dsv4_shared_mid_swiglu_q8_0")]]
 kernel void kernel_dsv4_shared_mid_swiglu_q8_0(
         constant ds4_metal_args_mul_mv & args,
@@ -535,24 +673,6 @@ kernel void kernel_dsv4_shared_mid_swiglu_q8_0(
             clamp_value, shmem, tgpig, tiisg, sgitg);
 }
 
-[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0_r4")]]
-kernel void kernel_dsv4_shared_mid_swiglu_q8_0_r4(
-        constant ds4_metal_args_mul_mv & args,
-        device const char * src0_gate,
-        device const char * src0_up,
-        device const char * src1,
-        device       char * dst_gate,
-        device       char * dst_up,
-        device       char * dst_mid,
-        constant     float &clamp_value,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-    kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<4, false>(
-            args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid,
-            clamp_value, shmem, tgpig, tiisg, sgitg);
-}
 
 template<typename T0, typename T1, short NR0, typename args_t>
 void kernel_mul_mv_t_t_impl(
@@ -974,6 +1094,311 @@ kernel void kernel_mul_mv_f16_f32_pair_compressor_store_4(
     state_score[dst] = projected_score[col] + ape_v;
 }
 
+// Decode compressor + indexer-compressor projection in one dispatch.  Both
+// pairs read the same normalized activation with the same F16 matvec shape,
+// so one launch covers all four matrices: threadgroups below the first
+// range boundary run the exact paired matvec + state store of
+// kernel_mul_mv_f16_f32_pair_compressor_store_4 for the attention
+// compressor, the rest for the indexer compressor.  Per-row reduction trees
+// and the per-threadgroup state stores are unchanged, keeping the fused
+// result bit-identical to the two separate dispatches while removing one
+// dispatch per decode layer.
+kernel void kernel_mul_mv_f16_f32_quad_compressor_store_4(
+        constant ds4_metal_args_mul_mv & args,
+        constant ds4_metal_args_compressor_pair_store & store0,
+        constant ds4_metal_args_compressor_pair_store & store1,
+        device const char * src0_a0,
+        device const char * src0_b0,
+        device const char * src0_a1,
+        device const char * src0_b1,
+        device const char * src1,
+        device       char * dst_a0,
+        device       char * dst_b0,
+        device       char * dst_a1,
+        device       char * dst_b1,
+        device const char * ape0,
+        device const char * ape1,
+        device       float * state0_kv,
+        device       float * state0_score,
+        device       float * state1_kv,
+        device       float * state1_score,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig  [[threadgroup_position_in_grid]],
+        ushort tiitg  [[thread_index_in_threadgroup]],
+        ushort tiisg  [[thread_index_in_simdgroup]],
+        ushort sgitg  [[simdgroup_index_in_threadgroup]]) {
+    constexpr short NR0 = 2;
+    const uint tgs0 = ((uint)store0.width + NR0 - 1u) / NR0;
+    const bool second = tgpig.x >= tgs0;
+
+    uint3 local_tgpig = tgpig;
+    if (second) local_tgpig.x = tgpig.x - tgs0;
+
+    ds4_metal_args_mul_mv largs = args;
+    largs.nr0 = NR0;
+    largs.ne01 = second ? (int32_t)store1.width : (int32_t)store0.width;
+
+    if (!second) {
+        kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
+                largs, src0_a0, src0_b0, src1, dst_a0, dst_b0,
+                shmem, local_tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
+                largs, src0_a1, src0_b1, src1, dst_a1, dst_b1,
+                shmem, local_tgpig, tiisg, sgitg);
+    }
+
+    threadgroup_barrier(mem_flags::mem_device);
+
+    // State append: identical to the paired store kernel, scoped to the
+    // range this threadgroup just projected (its own outputs only).
+    constant ds4_metal_args_compressor_pair_store & store = second ? store1 : store0;
+    if (tiitg >= NR0 || store.width == 0u || store.ratio == 0u) {
+        return;
+    }
+    const uint col = local_tgpig.x * (uint)NR0 + tiitg;
+    if (col >= store.width) return;
+
+    const uint pos_mod = store.pos % store.ratio;
+    const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod;
+    const uint dst = dst_row * store.width + col;
+    const uint ape_i = pos_mod * store.width + col;
+
+    device volatile const float * projected_kv = second
+        ? (device volatile const float *)dst_a1
+        : (device volatile const float *)dst_a0;
+    device volatile const float * projected_score = second
+        ? (device volatile const float *)dst_b1
+        : (device volatile const float *)dst_b0;
+    device const char * ape = second ? ape1 : ape0;
+    device float * state_kv = second ? state1_kv : state0_kv;
+    device float * state_score = second ? state1_score : state0_score;
+
+    float ape_v;
+    if (store.ape_type == 1u) {
+        ape_v = (float)(((device const half *)ape)[ape_i]);
+    } else {
+        ape_v = ((device const float *)ape)[ape_i];
+    }
+
+    state_kv[dst] = projected_kv[col];
+    state_score[dst] = projected_score[col] + ape_v;
+}
+
+/* Decode-only fusion: one dispatch covers the q_a/kv Q8 pair projection and
+ * the four F16 compressor projections (attention + indexer) with their
+ * state-store epilogue.  Both stages read the same normalized attention
+ * input and write disjoint outputs.  The q_a/kv range hosts two virtual
+ * NSG=4 cohorts per threadgroup, each an exact replica of
+ * kernel_mul_mv_q8_0_f32_pair (same per-lane K walk and reduction tree, cf.
+ * kernel_dsv4_router_shared_gate_up_q8_0); the compressor ranges run
+ * kernel_mul_mv_f16_f32_pair_4_impl<2> and the paired store epilogue
+ * verbatim, so every output bit matches the two separate dispatches. */
+kernel void kernel_dsv4_qkv_pair_quad_compressor_store_q8_0(
+        constant ds4_metal_args_mul_mv & args0,
+        constant ds4_metal_args_mul_mv & args1,
+        constant ds4_metal_args_mul_mv & cargs,
+        constant ds4_metal_args_compressor_pair_store & store0,
+        constant ds4_metal_args_compressor_pair_store & store1,
+        constant uint & pair_vtgs,
+        device const char * qw0,
+        device const char * qw1,
+        device const char * cw0a,
+        device const char * cw0b,
+        device const char * cw1a,
+        device const char * cw1b,
+        device const char * src1,
+        device       char * dst0,
+        device       char * dst1,
+        device       char * cdst_a0,
+        device       char * cdst_b0,
+        device       char * cdst_a1,
+        device       char * cdst_b1,
+        device const char * ape0,
+        device const char * ape1,
+        device       float * state0_kv,
+        device       float * state0_score,
+        device       float * state1_kv,
+        device       float * state1_score,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig  [[threadgroup_position_in_grid]],
+        ushort tiitg  [[thread_index_in_threadgroup]],
+        ushort tiisg  [[thread_index_in_simdgroup]],
+        ushort sgitg  [[simdgroup_index_in_threadgroup]]) {
+    constexpr short NW = N_SIMDWIDTH;
+    const uint pair_ctgs = (pair_vtgs + 1u) / 2u;
+
+    if (tgpig.x < pair_ctgs) {
+        /* Q8 pair range: cohort c of threadgroup t runs virtual pair
+         * threadgroup 2t+c with the original NSG=4 mapping. */
+        constexpr short NSG = 4;
+        constexpr short NQ  = 8;
+        constexpr short NR0 = 2;
+        const uint   cohort = sgitg >> 2;
+        const ushort vsg    = sgitg & 3u;
+        const uint   vt     = tgpig.x * 2u + cohort;
+        const bool   valid  = vt < pair_vtgs;
+
+        const int r0 = vt * NR0;
+        const bool active_a = valid && r0 < args0.ne01;
+        const bool active_b = valid && r0 < args1.ne01;
+        const int nb = args0.ne00 / QK8_0;
+
+        device const float *y = (device const float *)src1;
+        device const block_q8_0 *ax_a[NR0];
+        device const block_q8_0 *ax_b[NR0];
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            const int out_row = r0 + row;
+            ax_a[row] = active_a && out_row < args0.ne01
+                ? (device const block_q8_0 *)(qw0 + (uint64_t)out_row * args0.nb01)
+                : (device const block_q8_0 *)qw0;
+            ax_b[row] = active_b && out_row < args1.ne01
+                ? (device const block_q8_0 *)(qw1 + (uint64_t)out_row * args1.nb01)
+                : (device const block_q8_0 *)qw1;
+        }
+
+        float suma[NR0] = { 0.f };
+        float sumb[NR0] = { 0.f };
+        const short ix = tiisg / (NW / NQ);
+        const short il = tiisg % (NW / NQ);
+        const int ib0 = vsg * NQ + ix;
+        float yl[NQ];
+        device const float *yb = y + ib0 * QK8_0 + il * NQ;
+
+        if (valid) {
+            for (int ib = ib0; ib < nb; ib += NSG * NQ) {
+                FOR_UNROLL (short i = 0; i < NQ; ++i) {
+                    yl[i] = yb[i];
+                }
+                FOR_UNROLL (short row = 0; row < NR0; ++row) {
+                    const int out_row = r0 + row;
+                    if (active_a && out_row < args0.ne01) {
+                        device const int8_t *qs = ax_a[row][ib].qs + il * NQ;
+                        float sumq = 0.f;
+                        FOR_UNROLL (short i = 0; i < NQ; ++i) {
+                            sumq += qs[i] * yl[i];
+                        }
+                        suma[row] += sumq * ax_a[row][ib].d;
+                    }
+                    if (active_b && out_row < args1.ne01) {
+                        device const int8_t *qs = ax_b[row][ib].qs + il * NQ;
+                        float sumq = 0.f;
+                        FOR_UNROLL (short i = 0; i < NQ; ++i) {
+                            sumq += qs[i] * yl[i];
+                        }
+                        sumb[row] += sumq * ax_b[row][ib].d;
+                    }
+                }
+                yb += NSG * NQ * QK8_0;
+            }
+        }
+
+        threadgroup float *shared =
+            (threadgroup float *)shmem + cohort * (2 * NR0 * NW);
+        threadgroup float *sha[NR0];
+        threadgroup float *shb[NR0];
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            sha[row] = shared + NW * row;
+            shb[row] = shared + NW * (NR0 + row);
+            if (vsg == 0) {
+                sha[row][tiisg] = 0.0f;
+                if (active_b) shb[row][tiisg] = 0.0f;
+            }
+            suma[row] = simd_sum(suma[row]);
+            if (active_b) sumb[row] = simd_sum(sumb[row]);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            if (tiisg == 0) {
+                sha[row][vsg] = suma[row];
+                if (active_b) shb[row][vsg] = sumb[row];
+            }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        device float *out_a = (device float *)dst0;
+        device float *out_b = (device float *)dst1;
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            const float total_a = simd_sum(sha[row][tiisg]);
+            if (tiisg == 0 && vsg == 0) {
+                const int out_row = r0 + row;
+                if (active_a && out_row < args0.ne01) out_a[out_row] = total_a;
+            }
+            if (active_b) {
+                const float total_b = simd_sum(shb[row][tiisg]);
+                if (tiisg == 0 && vsg == 0) {
+                    const int out_row = r0 + row;
+                    if (out_row < args1.ne01) out_b[out_row] = total_b;
+                }
+            }
+        }
+        return;
+    }
+
+    /* Compressor quad range: verbatim body of
+     * kernel_mul_mv_f16_f32_quad_compressor_store_4 on the shifted grid. */
+    constexpr short NR0 = 2;
+    const uint lx = tgpig.x - pair_ctgs;
+    const uint tgs0 = ((uint)store0.width + NR0 - 1u) / NR0;
+    const bool second = lx >= tgs0;
+
+    uint3 local_tgpig = tgpig;
+    local_tgpig.x = second ? lx - tgs0 : lx;
+
+    ds4_metal_args_mul_mv largs = cargs;
+    largs.nr0 = NR0;
+    largs.ne01 = second ? (int32_t)store1.width : (int32_t)store0.width;
+
+    if (!second) {
+        kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
+                largs, cw0a, cw0b, src1, cdst_a0, cdst_b0,
+                shmem, local_tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
+                largs, cw1a, cw1b, src1, cdst_a1, cdst_b1,
+                shmem, local_tgpig, tiisg, sgitg);
+    }
+
+    threadgroup_barrier(mem_flags::mem_device);
+
+    // State append: identical to the paired store kernel, scoped to the
+    // range this threadgroup just projected (its own outputs only).
+    constant ds4_metal_args_compressor_pair_store & store = second ? store1 : store0;
+    if (tiitg >= NR0 || store.width == 0u || store.ratio == 0u) {
+        return;
+    }
+    const uint col = local_tgpig.x * (uint)NR0 + tiitg;
+    if (col >= store.width) return;
+
+    const uint pos_mod = store.pos % store.ratio;
+    const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod;
+    const uint dst = dst_row * store.width + col;
+    const uint ape_i = pos_mod * store.width + col;
+
+    device volatile const float * projected_kv = second
+        ? (device volatile const float *)cdst_a1
+        : (device volatile const float *)cdst_a0;
+    device volatile const float * projected_score = second
+        ? (device volatile const float *)cdst_b1
+        : (device volatile const float *)cdst_b0;
+    device const char * ape = second ? ape1 : ape0;
+    device float * state_kv = second ? state1_kv : state0_kv;
+    device float * state_score = second ? state1_score : state0_score;
+
+    float ape_v;
+    if (store.ape_type == 1u) {
+        ape_v = (float)(((device const half *)ape)[ape_i]);
+    } else {
+        ape_v = ((device const float *)ape)[ape_i];
+    }
+
+    state_kv[dst] = projected_kv[col];
+    state_score[dst] = projected_score[col] + ape_v;
+}
+
 template<typename T0, typename T1, typename args_t>
 void kernel_mul_mv_t_t_short_impl(
         args_t args,
@@ -1476,125 +1901,6 @@ constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]];
 constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]];
 
 #ifdef DS4_METAL_HAS_TENSOR
-template<
-    short NR0, short NR1,
-    typename SA, typename SA_4x4, typename block_q, short nl,
-    void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &),
-    typename T0, typename T0_4x4, typename T1>
-kernel void kernel_mul_mm_mpp(
-        constant ds4_metal_args_mul_mm & args,
-        device const char * srcA,
-        device const char * srcB,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig [[threadgroup_position_in_grid]],
-        ushort tiitg [[thread_index_in_threadgroup]],
-        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
-    (void) sgitg;
-
-    constexpr int NK  = 32;
-    constexpr int NL  = NK/16;
-    constexpr int NUM_THREADS = 128;
-
-    const int K = args.ne00;
-    const int M = args.ne0;
-    const int N = args.ne1;
-    const int im = tgpig.z;
-    const int i12 = im%args.ne12;
-    const int i13 = im/args.ne12;
-    const int r0 = tgpig.y*NR0;
-    const int r1 = tgpig.x*NR1;
-
-    const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
-
-    threadgroup SA *sa = (threadgroup SA *)shmem;
-    threadgroup SA *sb = sa + NR0*NK;
-    auto tA = tensor(sa, dextents<int32_t, 2>(NK, NR0));
-    auto tB = tensor(sb, dextents<int32_t, 2>(NK, NR1));
-
-    device const T1 *ptrB = (device const T1 *)(srcB + args.nb12*i12 + args.nb13*i13);
-    const int strideB = args.nb11/sizeof(T1);
-
-    matmul2d<
-        matmul2d_descriptor(NR1, NR0, NK, false, true, false,
-            matmul2d_descriptor::mode::multiply_accumulate),
-        execution_simdgroups<4>> mm;
-
-    auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA), float>();
-
-    #pragma unroll
-    for (uint16_t i = 0; i < cT.get_capacity(); ++i) {
-        if (cT.is_valid_element(i)) {
-            cT[i] = 0.0f;
-        }
-    }
-
-    for (int loop_k = 0; loop_k < K; loop_k += NK) {
-        for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) {
-            const int row = work/NL;
-            const int k_chunk = work%NL;
-            const int k_pos = loop_k + k_chunk*16;
-            const short k_base = k_chunk*16;
-
-            if (!FC_mul_mm_bc_out || r0 + row < M) {
-                if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
-                    device const T0 *row_ptr = (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0);
-                    FOR_UNROLL (short i = 0; i < 16; i++) {
-                        sa[row*NK + k_base + i] = (k_pos + i < K) ? (SA)row_ptr[k_pos + i] : (SA)0;
-                    }
-                } else {
-                    const int block_idx = k_pos/(16*nl);
-                    const short il = (k_pos/16)%nl;
-                    device const block_q *row_ptr = (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0);
-
-                    SA_4x4 temp_a;
-                    dequantize_func(row_ptr + block_idx, il, temp_a);
-                    FOR_UNROLL (short i = 0; i < 16; i++) {
-                        sa[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0;
-                    }
-                }
-            } else {
-                FOR_UNROLL (short i = 0; i < 16; i++) {
-                    sa[row*NK + k_base + i] = (SA)0;
-                }
-            }
-        }
-        for (int work = tiitg; work < NK*NR1; work += NUM_THREADS) {
-            const int col = work/NK;
-            const int k = work%NK;
-            if ((!FC_mul_mm_bc_out && !FC_mul_mm_bc_inp) ||
-                (r1 + col < N && loop_k + k < K)) {
-                sb[col*NK + k] = (SA)ptrB[(uint64_t)(r1 + col)*strideB + loop_k + k];
-            } else {
-                sb[col*NK + k] = (SA)0;
-            }
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        auto mA = tA.slice(0, 0);
-        auto mB = tB.slice(0, 0);
-        mm.run(mB, mA, cT);
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-    }
-
-    device float *dst_batch = (device float *)dst + im*N*M;
-    if (!FC_mul_mm_bc_out) {
-        device float *dst_tile = dst_batch + r0 + (uint64_t)r1*M;
-        auto tD = tensor(dst_tile, dextents<int32_t, 2>(NR0, NR1), array<int, 2>({1, M}));
-        cT.store(tD);
-    } else {
-        auto tD = tensor(dst_batch, dextents<int32_t, 2>(M, N), array<int, 2>({1, M}));
-        auto mD = tD.slice(r0, r1);
-        cT.store(mD);
-    }
-}
-
-typedef decltype(kernel_mul_mm_mpp<64, 32, half, half4x4, float4x4, 1, dequantize_f32, float, float4x4, float>) mul_mm_mpp_t;
-
-template [[host_name("kernel_mul_mm_f16_f32_mpp")]]  kernel mul_mm_mpp_t kernel_mul_mm_mpp<64, 32, half, half4x4, half4x4, 1, dequantize_f16,  half,  half4x4,  float>;
-
 // Retained Metal4/TensorOps dense prefill kernel.  The legacy MPP prototype
 // staged both operands in threadgroup memory; this version stages only the
 // model weight tile and lets MPP read the dense RHS activation matrix directly
@@ -2144,242 +2450,6 @@ kernel void kernel_mul_mm_f16_f32_scaled(
     }
 }
 
-kernel void kernel_mul_mm_f16_f32_pair(
-        constant ds4_metal_args_mul_mm & args,
-        device const char * src0_a,
-        device const char * src0_b,
-        device const char * src1,
-        device       char * dst_a,
-        device       char * dst_b,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiitg[[thread_index_in_threadgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-    threadgroup half * sa_a = (threadgroup half *)(shmem);
-    threadgroup half * sa_b = (threadgroup half *)(shmem + 4096);
-    threadgroup half * sb   = (threadgroup half *)(shmem + 8192);
-
-    constexpr int NR0 = 64;
-    constexpr int NR1 = 32;
-    constexpr int NK  = 32;
-    constexpr int NL0 = NK/16;
-    constexpr int NL1 = NK/8;
-
-    const int im = tgpig.z;
-    const int r0 = tgpig.y*NR0;
-    const int r1 = tgpig.x*NR1;
-
-    const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
-    const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
-
-    const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1;
-    const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1;
-
-    const short il0 = (tiitg % NL0);
-    short il = il0;
-
-    const int i12 = im%args.ne12;
-    const int i13 = im/args.ne12;
-
-    const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
-    const short    offset1 = il0;
-
-    device const half4x4 * xa = (device const half4x4 *)(src0_a + args.nb01*(r0 + lr0) + offset0) + offset1;
-    device const half4x4 * xb = (device const half4x4 *)(src0_b + args.nb01*(r0 + lr0) + offset0) + offset1;
-
-    const short iy = 8*(tiitg % NL1);
-
-    device const float * y = (device const float *)(src1
-        + args.nb13*i13
-        + args.nb12*i12
-        + args.nb11*(r1 + lr1)
-        + args.nb10*iy);
-
-    simdgroup_half8x8 ma[4];
-    simdgroup_half8x8 mb[2];
-
-    simdgroup_float8x8 mc_a[8];
-    simdgroup_float8x8 mc_b[8];
-
-    for (short i = 0; i < 8; i++) {
-        mc_a[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
-        mc_b[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
-    }
-
-    for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
-        half4x4 temp_a;
-        half4x4 temp_b;
-        dequantize_f16(xa, il, temp_a);
-        dequantize_f16(xb, il, temp_b);
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        FOR_UNROLL (short i = 0; i < 16; i++) {
-            const short sx = 2*il0 + i/8;
-            const short sy = (tiitg/NL0)/8;
-
-            const short lx = (tiitg/NL0)%8;
-            const short ly = i%8;
-
-            const short ib = 8*sx + sy;
-
-            *(sa_a + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
-            *(sa_b + 64*ib + 8*ly + lx) = temp_b[i/4][i%4];
-        }
-
-        if (FC_mul_mm_bc_inp) {
-            for (short i = 0; i < 8; ++i) {
-                const short sx = (tiitg%NL1);
-                const short sy = (tiitg/NL1)/8;
-
-                const short lx = i;
-                const short ly = (tiitg/NL1)%8;
-
-                const short ib = 4*sx + sy;
-
-                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (half) *((device float *) y + i) : 0;
-            }
-        } else {
-            const short sx = (tiitg%NL1);
-            const short sy = (tiitg/NL1)/8;
-
-            const short ly = (tiitg/NL1)%8;
-
-            const short ib = 4*sx + sy;
-
-            *(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)(*((device float2x4 *) y));
-        }
-
-        il = (il + 2 < 1) ? il + 2 : il % 2;
-        xa = (il < 2) ? xa + 2 : xa;
-        xb = (il < 2) ? xb + 2 : xb;
-
-        y += NK;
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        threadgroup const half * lsma_a = (sa_a + 4*64*(sgitg%2));
-        threadgroup const half * lsma_b = (sa_b + 4*64*(sgitg%2));
-        threadgroup const half * lsmb   = (sb   + 2*64*(sgitg/2));
-
-        FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 2; i++) {
-                simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 4; i++) {
-                simdgroup_load(ma[i], lsma_a + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 8; i++) {
-                simdgroup_multiply_accumulate(mc_a[i], mb[i/4], ma[i%4], mc_a[i]);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 4; i++) {
-                simdgroup_load(ma[i], lsma_b + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 8; i++) {
-                simdgroup_multiply_accumulate(mc_b[i], mb[i/4], ma[i%4], mc_b[i]);
-            }
-
-            lsma_a += 8*64;
-            lsma_b += 8*64;
-            lsmb   += 4*64;
-        }
-    }
-
-    if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
-        device float * C_a = (device float *) dst_a +
-            (r0 + 32*(sgitg &  1)) +
-            (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
-        device float * C_b = (device float *) dst_b +
-            (r0 + 32*(sgitg &  1)) +
-            (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
-
-        for (short i = 0; i < 8; i++) {
-            simdgroup_store(mc_a[i], C_a + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
-            simdgroup_store(mc_b[i], C_b + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
-        }
-    } else {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        threadgroup float * temp_str = (threadgroup float *) shmem;
-
-        for (short i = 0; i < 8; i++) {
-            simdgroup_store(mc_a[i],
-                            temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 8*(i%4) + 8*NR0*(i/4),
-                            NR0,
-                            0,
-                            false);
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (sgitg == 0) {
-            for (int j = tiitg; j < nr1; j += NR1) {
-                device float  * D  = (device float *) dst_a + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
-                device float4 * D4 = (device float4 *) D;
-
-                threadgroup float  * C  = temp_str + (j*NR0);
-                threadgroup float4 * C4 = (threadgroup float4 *) C;
-
-                int i = 0;
-                for (; i < nr0/4; i++) {
-                    *(D4 + i) = *(C4 + i);
-                }
-
-                i *= 4;
-                for (; i < nr0; i++) {
-                    *(D + i) = *(C + i);
-                }
-            }
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        for (short i = 0; i < 8; i++) {
-            simdgroup_store(mc_b[i],
-                            temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 8*(i%4) + 8*NR0*(i/4),
-                            NR0,
-                            0,
-                            false);
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (sgitg == 0) {
-            for (int j = tiitg; j < nr1; j += NR1) {
-                device float  * D  = (device float *) dst_b + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
-                device float4 * D4 = (device float4 *) D;
-
-                threadgroup float  * C  = temp_str + (j*NR0);
-                threadgroup float4 * C4 = (threadgroup float4 *) C;
-
-                int i = 0;
-                for (; i < nr0/4; i++) {
-                    *(D4 + i) = *(C4 + i);
-                }
-
-                i *= 4;
-                for (; i < nr0; i++) {
-                    *(D + i) = *(C + i);
-                }
-            }
-        }
-    }
-}
-
 typedef decltype(kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_t;
 
 // Host-visible prefill matmul variants for F16 and Q8_0 weights.

metal/dsv4_hc.metal

@@ -21,20 +21,6 @@ struct ds4_metal_args_dsv4_hc_weighted_sum {
     uint64_t nb1;
 };
 
-struct ds4_metal_args_dsv4_hc_weighted_sum_norm {
-    int64_t  n_embd;
-    int64_t  n_hc;
-    int64_t  n_tokens;
-    uint64_t nb_x0;
-    uint64_t nb_x1;
-    uint64_t nb_x2;
-    uint64_t nb_w0;
-    uint64_t nb_w1;
-    uint64_t nb0;
-    uint64_t nb1;
-    uint64_t nb_norm1;
-    float    norm_eps;
-};
 
 struct ds4_metal_args_dsv4_output_hc_weights4 {
     float post_scale;
@@ -411,6 +397,68 @@ kernel void kernel_dsv4_hc_split_weighted_sum(
 // kernel_dsv4_hc_split_weighted_sum, stores the HC-pre row for diagnostics, and
 // reuses the just-collapsed values from threadgroup memory for the RMSNorm
 // reduction.
+static __attribute__((always_inline)) inline void ds4_hc_comb_weights4_exact(
+        constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args,
+        device volatile const float *mix,
+        device const float *scale,
+        device const float *base,
+        device float *out) {
+    const float epsv = args.eps;
+    const float comb_scale = scale[2];
+
+    float4 r0 =
+        *((device volatile const float4 *)(mix + 8)) * comb_scale +
+        *((device const float4 *)(base + 8));
+    float4 r1 =
+        *((device volatile const float4 *)(mix + 12)) * comb_scale +
+        *((device const float4 *)(base + 12));
+    float4 r2 =
+        *((device volatile const float4 *)(mix + 16)) * comb_scale +
+        *((device const float4 *)(base + 16));
+    float4 r3 =
+        *((device volatile const float4 *)(mix + 20)) * comb_scale +
+        *((device const float4 *)(base + 20));
+
+    const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w));
+    const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w));
+    const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w));
+    const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w));
+
+    r0 = exp(r0 - m0);
+    r1 = exp(r1 - m1);
+    r2 = exp(r2 - m2);
+    r3 = exp(r3 - m3);
+
+    r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv;
+    r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv;
+    r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv;
+    r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv;
+
+    float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
+    r0 *= col_inv;
+    r1 *= col_inv;
+    r2 *= col_inv;
+    r3 *= col_inv;
+
+    for (int iter = 1; iter < args.sinkhorn_iters; ++iter) {
+        r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv);
+        r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv);
+        r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv);
+        r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv);
+
+        col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
+        r0 *= col_inv;
+        r1 *= col_inv;
+        r2 *= col_inv;
+        r3 *= col_inv;
+    }
+
+    *((device float4 *)(out + 8)) = r0;
+    *((device float4 *)(out + 12)) = r1;
+    *((device float4 *)(out + 16)) = r2;
+    *((device float4 *)(out + 20)) = r3;
+}
+
 kernel void kernel_dsv4_hc_split_weighted_sum_norm4(
         constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args,
         device  const char  * mixes,
@@ -519,7 +567,6 @@ kernel void kernel_dsv4_hc_split_weighted_sum_norm4(
     }
 
     threadgroup_barrier(mem_flags::mem_threadgroup);
-
     float sumf = 0.0f;
     for (uint i = tid; i < n4; i += ntg) {
         device const float4 *x0 = (device const float4 *)(x + 0 * args.nb_x1 + (uint64_t)row * args.nb_x2);
@@ -884,6 +931,119 @@ kernel void kernel_dsv4_q8_hc_expand4_q8_0(
     }
 }
 
+kernel void kernel_dsv4_q8_hc_expand4_q8_0_vec_hc(
+        constant ds4_metal_args_mul_mv        & mv,
+        constant ds4_metal_args_dsv4_hc_expand & hc,
+        device  const char * weight,
+        device  const char * input,
+        device        char * block_out,
+        device  const char * residual,
+        device  const char * post,
+        device  const char * comb,
+        device        char * dst,
+        threadgroup   char * shmem [[threadgroup(0)]],
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    if (hc.n_hc != 4 || hc.n_tokens != 1) {
+        return;
+    }
+
+    const short NSG = FC_mul_mv_nsg;
+    constexpr short NW = N_SIMDWIDTH;
+    constexpr short NQ = 8;
+    constexpr short NR0 = N_R0_Q8_0;
+
+    const int nb = mv.ne00 / QK8_0;
+    const int row0 = tgpig.x * NR0;
+
+    const short ix = tiisg / (NW / NQ);
+    const short il = tiisg % (NW / NQ);
+    const int ib0 = sgitg * NQ + ix;
+
+    device const float *y = (device const float *)(input);
+    device const float *yb = y + ib0 * QK8_0 + il * NQ;
+
+    device const block_q8_0 *ax[NR0];
+    FOR_UNROLL(short row = 0; row < NR0; ++row) {
+        const uint64_t off0 = (uint64_t)(row0 + row) * mv.nb01;
+        ax[row] = (device const block_q8_0 *)(weight + off0);
+    }
+
+    float sumf[NR0] = { 0.0f };
+    float yl[NQ];
+
+    for (int ib = ib0; ib < nb; ib += NSG * NQ) {
+        FOR_UNROLL(short i = 0; i < NQ; ++i) {
+            yl[i] = yb[i];
+        }
+
+        FOR_UNROLL(short row = 0; row < NR0; ++row) {
+            device const int8_t *qs = ax[row][ib].qs + il * NQ;
+
+            float sumq = 0.0f;
+            FOR_UNROLL(short i = 0; i < NQ; ++i) {
+                sumq += qs[i] * yl[i];
+            }
+
+            sumf[row] += sumq * ax[row][ib].d;
+        }
+
+        yb += NSG * NQ * QK8_0;
+    }
+
+    threadgroup float *shmem_f32[NR0];
+    FOR_UNROLL(short row = 0; row < NR0; ++row) {
+        shmem_f32[row] = (threadgroup float *)shmem + NW * row;
+        if (sgitg == 0) {
+            shmem_f32[row][tiisg] = 0.0f;
+        }
+        sumf[row] = simd_sum(sumf[row]);
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    FOR_UNROLL(short row = 0; row < NR0; ++row) {
+        if (tiisg == 0) {
+            shmem_f32[row][sgitg] = sumf[row];
+        }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    FOR_UNROLL(short row = 0; row < NR0; ++row) {
+        const int d = row0 + row;
+        if (d >= mv.ne01) {
+            continue;
+        }
+
+        const float block_v = simd_sum(shmem_f32[row][tiisg]);
+        if (tiisg == 0 && sgitg == 0) {
+            *((device float *)(block_out + (uint64_t)d * sizeof(float))) = block_v;
+
+            const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1));
+            const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1));
+            const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1));
+            const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1));
+
+            const float4 post4 = *((device const float4 *)post);
+            const float4 comb0 = *((device const float4 *)(comb + 0 * hc.nb_comb1));
+            const float4 comb1 = *((device const float4 *)(comb + 1 * hc.nb_comb1));
+            const float4 comb2 = *((device const float4 *)(comb + 2 * hc.nb_comb1));
+            const float4 comb3 = *((device const float4 *)(comb + 3 * hc.nb_comb1));
+            float4 acc = block_v * post4;
+            acc += comb0 * r0;
+            acc += comb1 * r1;
+            acc += comb2 * r2;
+            acc += comb3 * r3;
+            FOR_UNROLL (short dst_hc = 0; dst_hc < 4; ++dst_hc) {
+                *((device float *)(dst + (uint64_t)d * hc.nb0 +
+                                   (uint64_t)dst_hc * hc.nb1)) = acc[dst_hc];
+            }
+        }
+    }
+}
+
 // Reduces HC channels to a normal embedding row with the learned pre weights.
 // This is the input adapter before the attention block and before the FFN block.
 kernel void kernel_dsv4_hc_weighted_sum(
@@ -910,76 +1070,6 @@ kernel void kernel_dsv4_hc_weighted_sum(
     *((device float *) (dst + d*args.nb0 + t*args.nb1)) = acc;
 }
 
-// The one-row output head immediately applies a learned RMSNorm after reducing
-// its four HC streams. Preserve the standalone scalar HC accumulation, write
-// the collapsed row for diagnostics, then reload its F32 values from
-// threadgroup memory using the standalone RMSNorm's float4 reduction mapping.
-kernel void kernel_dsv4_hc_weighted_sum_norm4(
-        constant ds4_metal_args_dsv4_hc_weighted_sum_norm & args,
-        device  const char * x,
-        device  const char * weights,
-        device        char * dst,
-        device  const char * norm_weight,
-        device        char * norm_dst,
-        threadgroup  float * shared [[threadgroup(0)]],
-        ushort tid [[thread_position_in_threadgroup]],
-        ushort sgitg [[simdgroup_index_in_threadgroup]],
-        ushort tiisg [[thread_index_in_simdgroup]],
-        ushort ntg [[threads_per_threadgroup]]) {
-    if (args.n_tokens != 1 || args.n_hc != 4 ||
-        args.n_embd <= 0 || (args.n_embd & 3) != 0) {
-        return;
-    }
-
-    const uint n_embd = uint(args.n_embd);
-    const uint n4 = n_embd >> 2;
-    threadgroup float *row_shmem = shared;
-    threadgroup float *sum_shmem = shared + n_embd;
-
-    if (sgitg == 0) {
-        sum_shmem[tiisg] = 0.0f;
-    }
-
-    for (uint d = tid; d < n_embd; d += ntg) {
-        float acc = 0.0f;
-        for (int64_t h = 0; h < args.n_hc; ++h) {
-            const float xv = *((device const float *)(
-                x + (uint64_t)d*args.nb_x0 + (uint64_t)h*args.nb_x1));
-            const float wv = *((device const float *)(
-                weights + (uint64_t)h*args.nb_w0));
-            acc += xv * wv;
-        }
-        row_shmem[d] = acc;
-        *((device float *)(dst + (uint64_t)d*args.nb0)) = acc;
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    threadgroup const float4 *row4 =
-        (threadgroup const float4 *)row_shmem;
-    float sumf = 0.0f;
-    for (uint i = tid; i < n4; i += ntg) {
-        sumf += dot(row4[i], row4[i]);
-    }
-    sumf = simd_sum(sumf);
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (tiisg == 0) {
-        sum_shmem[sgitg] = sumf;
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    sumf = sum_shmem[tiisg];
-    sumf = simd_sum(sumf);
-
-    const float mean = sumf/args.n_embd;
-    const float scale = 1.0f/sqrt(mean + args.norm_eps);
-    device const float4 *w4 = (device const float4 *)norm_weight;
-    device float4 *norm4 = (device float4 *)norm_dst;
-    for (uint i = tid; i < n4; i += ntg) {
-        norm4[i] = (row4[i]*scale)*w4[i];
-    }
-}
-
 // The one-row HC=4 output head historically materializes four device-F32
 // stages across separate launches. Collapse those launches into one tiny
 // two-thread group while preserving the scalar/vector lane mapping and every
@@ -1015,3 +1105,440 @@ kernel void kernel_dsv4_output_hc_weights4(
             args.post_scale * x + args.eps;
     }
 }
+
+
+struct ds4_metal_args_hc_norm_mix {
+    int32_t n;
+    int32_t out_dim;
+    float   eps;
+};
+
+// Fused unweighted RMSNorm + F16 HC-mix projection for DS4 decode HC-pre.
+// The standalone decode path runs kernel_rms_norm_f32_4 over the flattened
+// 4*embd HC row (1024 threads, one threadgroup) and then
+// kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2) over the normalized row.  Both
+// stages are reproduced bit-exactly in one dispatch: every threadgroup
+// redundantly recomputes the norm partials with the original 1024-thread
+// mapping (each real lane covers one virtual thread of each 256-thread
+// slice, preserving every simd_sum tree), and the matvec keeps the original
+// per-row accumulation order with y = x*scale computed on the fly, which
+// rounds identically to the materialized normalized row.  The host wrapper
+// gates this to n == 16384 && out_dim == 24, where the virtual-thread count
+// is exactly 1024 and the mv tail loop is empty.
+kernel void kernel_dsv4_hc_rms_norm_mix_f16(
+        constant ds4_metal_args_hc_norm_mix & args,
+        device const char  * x,
+        device const char  * weight,
+        device       char  * dst,
+        threadgroup  char  * shmem [[threadgroup(0)]],
+        uint3  tgpig [[threadgroup_position_in_grid]],
+        ushort tiisg [[thread_index_in_simdgroup]],
+        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
+    constexpr short NSG = 8;   // ds4_gpu_make_plain_mv_dispatch(16384)
+    constexpr short NW  = N_SIMDWIDTH;
+    constexpr short NR0 = 2;   // plain mv nr0
+    constexpr short NB  = 32;
+    constexpr short NF  = 16;
+    constexpr short NF4 = NF/4;
+    constexpr uint  VTHREADS = 1024u;                 // rms norm threads at n == 16384
+    constexpr short VSLICES  = VTHREADS/(NSG*NW);     // virtual 256-thread slices
+
+    const uint n  = (uint)args.n;
+    const uint n4 = n >> 2;
+
+    device const float4 *x4 = (device const float4 *)x;
+
+    threadgroup float *norm_shmem = (threadgroup float *)shmem;        // NW slots
+    threadgroup float *mv_shmem   = (threadgroup float *)shmem + NW;   // NW*NR0 slots
+
+    // Phase A: exact replica of kernel_rms_norm_f32_4's reduction tree with
+    // the 1024 virtual threads folded onto this threadgroup's 8 simdgroups.
+    for (short v = 0; v < VSLICES; ++v) {
+        const uint vt = (uint)(sgitg + NSG*v)*NW + tiisg;
+        float sumf = 0.0f;
+        for (uint i00 = vt; i00 < n4; i00 += VTHREADS) {
+            sumf += dot(x4[i00], x4[i00]);
+        }
+        sumf = simd_sum(sumf);
+        if (tiisg == 0) {
+            norm_shmem[sgitg + NSG*v] = sumf;
+        }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    float total = norm_shmem[tiisg];
+    total = simd_sum(total);
+    const float mean  = total/(float)args.n;
+    const float scale = 1.0f/sqrt(mean + args.eps);
+
+    // Phase B: exact replica of kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2) with
+    // the normalized operand recomputed as x*scale instead of reloaded.
+    const int nb = args.n/NB;
+    const int r0 = tgpig.x*NR0;
+
+    device const half4 * 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));
+    }
+
+    float sumf_mv[NR0] = { 0.f };
+
+    const short ix = tiisg/(NW/NF);
+    const short il = tiisg%(NW/NF);
+    const int ib0 = sgitg*NF + ix;
+
+    for (int ib = ib0; ib < nb; ib += NSG*NF) {
+        float4 yl4[NF4];
+        FOR_UNROLL (short i = 0; i < NF4; ++i) {
+            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;
+
+            float sumq = 0.f;
+            FOR_UNROLL (short i = 0; i < NF4; ++i) {
+                sumq += dot(float4(xb4[i]), yl4[i]);
+            }
+
+            sumf_mv[row] += sumq;
+        }
+    }
+
+    // n == 16384 makes the scalar tail loop of the original empty.
+    device float * dst_f32 = (device float *) dst;
+    helper_mv_reduce_and_write<NR0>(dst_f32, sumf_mv, r0, args.out_dim,
+                                    tiisg, sgitg, (threadgroup char *)mv_shmem);
+}
+
+// M5 specialization: pack two exact NR0=2 HC-mix producer groups into one
+// 512-thread group. Two independent eight-simdgroup clusters retain the
+// matvec reductions while the exact RMS scale is redundantly formed six,
+// rather than twelve, times.
+kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2(
+        constant ds4_metal_args_hc_norm_mix & args,
+        device const char  * x,
+        device const char  * weight,
+        device       char  * dst,
+        threadgroup  char  * shmem [[threadgroup(0)]],
+        uint3  tgpig [[threadgroup_position_in_grid]],
+        ushort tiisg [[thread_index_in_simdgroup]],
+        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
+    constexpr short NSG_CLUSTER = 8;
+    constexpr short NCLUSTER = 2;
+    constexpr short NSG_TOTAL = NSG_CLUSTER * NCLUSTER;
+    constexpr short NW = N_SIMDWIDTH;
+    constexpr short NR0 = 2;
+    constexpr short NB = 32;
+    constexpr short NF = 16;
+    constexpr short NF4 = NF/4;
+    constexpr uint VTHREADS = 1024u;
+    constexpr short VSLICES = VTHREADS/(NSG_TOTAL*NW);
+
+    const uint n = (uint)args.n;
+    const uint n4 = n >> 2;
+    device const float4 *x4 = (device const float4 *)x;
+    threadgroup float *norm_shmem = (threadgroup float *)shmem;
+    threadgroup float *mv_shmem = norm_shmem + NW;
+
+    // Exact 1024-virtual-thread RMS reduction, now folded two ways over
+    // the 16 physical simdgroups instead of four ways over eight.
+    for (short v = 0; v < VSLICES; ++v) {
+        const uint vt = (uint)(sgitg + NSG_TOTAL*v)*NW + tiisg;
+        float sumf = 0.0f;
+        for (uint i00 = vt; i00 < n4; i00 += VTHREADS) {
+            sumf += dot(x4[i00], x4[i00]);
+        }
+        sumf = simd_sum(sumf);
+        if (tiisg == 0) {
+            norm_shmem[sgitg + NSG_TOTAL*v] = sumf;
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    float total = norm_shmem[tiisg];
+    total = simd_sum(total);
+    const float mean = total/(float)args.n;
+    const float scale = 1.0f/sqrt(mean + args.eps);
+
+    // Two independent eight-simdgroup clusters reproduce two original
+    // NR0=2 matvec threadgroups inside this 512-thread threadgroup.
+    const short cluster = sgitg / NSG_CLUSTER;
+    const short local_sg = sgitg - cluster*NSG_CLUSTER;
+    const int nb = args.n/NB;
+    const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0;
+
+    device const half4 *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));
+    }
+
+    float sumf_mv[NR0] = { 0.f };
+    const short ix = tiisg/(NW/NF);
+    const short il = tiisg%(NW/NF);
+    const int ib0 = local_sg*NF + ix;
+    for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) {
+        float4 yl4[NF4];
+        FOR_UNROLL (short i = 0; i < NF4; ++i) {
+            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;
+            float sumq = 0.f;
+            FOR_UNROLL (short i = 0; i < NF4; ++i) {
+                sumq += dot(float4(xb4[i]), yl4[i]);
+            }
+            sumf_mv[row] += sumq;
+        }
+    }
+
+    threadgroup float *cluster_shmem[NR0];
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        cluster_shmem[row] = mv_shmem +
+            ((uint)cluster*NR0 + row)*NW;
+        if (local_sg == 0) {
+            cluster_shmem[row][tiisg] = 0.0f;
+        }
+        sumf_mv[row] = simd_sum(sumf_mv[row]);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        if (tiisg == 0) {
+            cluster_shmem[row][local_sg] = sumf_mv[row];
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    device float *mixes_f32 = (device float *)dst;
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        const float tot = simd_sum(cluster_shmem[row][tiisg]);
+        if (tiisg == 0 && local_sg == 0 && r0 + row < args.out_dim) {
+            mixes_f32[r0 + row] = tot;
+        }
+    }
+}
+
+kernel void kernel_dsv4_hc_rms_norm_mix_f16_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]]) {
+    constexpr short NSG_CLUSTER = 8;
+    constexpr short NCLUSTER = 2;
+    constexpr short NSG_TOTAL = NSG_CLUSTER * NCLUSTER;
+    constexpr short NW = N_SIMDWIDTH;
+    constexpr short NR0 = 2;
+    constexpr short NB = 32;
+    constexpr short NF = 16;
+    constexpr short NF4 = NF/4;
+    constexpr uint VTHREADS = 1024u;
+    constexpr short VSLICES = VTHREADS/(NSG_TOTAL*NW);
+
+    const uint n = (uint)args.n;
+    const uint n4 = n >> 2;
+    device const float4 *x4 = (device const float4 *)x;
+    threadgroup float *norm_shmem = (threadgroup float *)shmem;
+    threadgroup float *mv_shmem = norm_shmem + NW;
+
+    // Exact 1024-virtual-thread RMS reduction, now folded two ways over
+    // the 16 physical simdgroups instead of four ways over eight.
+    for (short v = 0; v < VSLICES; ++v) {
+        const uint vt = (uint)(sgitg + NSG_TOTAL*v)*NW + tiisg;
+        float sumf = 0.0f;
+        for (uint i00 = vt; i00 < n4; i00 += VTHREADS) {
+            sumf += dot(x4[i00], x4[i00]);
+        }
+        sumf = simd_sum(sumf);
+        if (tiisg == 0) {
+            norm_shmem[sgitg + NSG_TOTAL*v] = sumf;
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    float total = norm_shmem[tiisg];
+    total = simd_sum(total);
+    const float mean = total/(float)args.n;
+    const float scale = 1.0f/sqrt(mean + args.eps);
+
+    // Two independent eight-simdgroup clusters reproduce two original
+    // NR0=2 matvec threadgroups inside this 512-thread threadgroup.
+    const short cluster = sgitg / NSG_CLUSTER;
+    const short local_sg = sgitg - cluster*NSG_CLUSTER;
+    const int nb = args.n/NB;
+    const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0;
+
+    device const half4 *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));
+    }
+
+    float sumf_mv[NR0] = { 0.f };
+    const short ix = tiisg/(NW/NF);
+    const short il = tiisg%(NW/NF);
+    const int ib0 = local_sg*NF + ix;
+    for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) {
+        float4 yl4[NF4];
+        FOR_UNROLL (short i = 0; i < NF4; ++i) {
+            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;
+            float sumq = 0.f;
+            FOR_UNROLL (short i = 0; i < NF4; ++i) {
+                sumq += dot(float4(xb4[i]), yl4[i]);
+            }
+            sumf_mv[row] += sumq;
+        }
+    }
+
+    threadgroup float *cluster_shmem[NR0];
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        cluster_shmem[row] = mv_shmem +
+            ((uint)cluster*NR0 + row)*NW;
+        if (local_sg == 0) {
+            cluster_shmem[row][tiisg] = 0.0f;
+        }
+        sumf_mv[row] = simd_sum(sumf_mv[row]);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        if (tiisg == 0) {
+            cluster_shmem[row][local_sg] = sumf_mv[row];
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    device volatile float *mixes_f32 =
+        (device volatile float *)dst;
+    if (local_sg == 0) {
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            const float tot = simd_sum(cluster_shmem[row][tiisg]);
+            if (tiisg == 0 && r0 + row < args.out_dim) {
+                mixes_f32[r0 + row] = tot;
+            }
+        }
+    }
+
+    // The first producer group owns mix[0:4].  After materializing and
+    // reloading those values, fold the established 1024-thread HC collapse
+    // and RMS reduction over this group's 512 physical threads as two
+    // independent virtual slices.  This retains the original 32-partial tree.
+    threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
+    const uint tid = (uint)sgitg * (uint)NW + (uint)tiisg;
+    threadgroup float *pre_shmem = norm_shmem + 32u + 4u*NW;
+    threadgroup float *sum_shmem = pre_shmem + 4;
+
+    if (tgpig.x == 0) {
+        device float *out = (device float *)split;
+        if (tid == 0) {
+            const float4 pre_z =
+                *((device volatile const float4 *)mixes_f32) * hc_scale[0] +
+                *((device const float4 *)hc_base);
+            const float4 pre =
+                1.0f / (1.0f + exp(-pre_z)) + split_args.eps;
+            *((device float4 *)out) = pre;
+            pre_shmem[0] = pre.x;
+            pre_shmem[1] = pre.y;
+            pre_shmem[2] = pre.z;
+            pre_shmem[3] = pre.w;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const uint n4_collapse = uint(split_args.n_embd) >> 2;
+        const uint i0 = tid;
+        const uint i1 = tid + 512u;
+        device const float4 *x0 = (device const float4 *)(
+            x + 0 * split_args.nb_x1);
+        device const float4 *x1 = (device const float4 *)(
+            x + 1 * split_args.nb_x1);
+        device const float4 *x2 = (device const float4 *)(
+            x + 2 * split_args.nb_x1);
+        device const float4 *x3 = (device const float4 *)(
+            x + 3 * split_args.nb_x1);
+
+        float4 v0 = 0.0f;
+        v0 += x0[i0] * pre_shmem[0];
+        v0 += x1[i0] * pre_shmem[1];
+        v0 += x2[i0] * pre_shmem[2];
+        v0 += x3[i0] * pre_shmem[3];
+        float sum0 = simd_sum(dot(v0, v0));
+
+        float4 v1 = 0.0f;
+        if (i1 < n4_collapse) {
+            v1 += x0[i1] * pre_shmem[0];
+            v1 += x1[i1] * pre_shmem[1];
+            v1 += x2[i1] * pre_shmem[2];
+            v1 += x3[i1] * pre_shmem[3];
+        }
+        float sum1 = simd_sum(dot(v1, v1));
+        if (tiisg == 0) {
+            sum_shmem[sgitg] = sum0;
+            sum_shmem[sgitg + 16] = sum1;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        float sumf = sum_shmem[tiisg];
+        sumf = simd_sum(sumf);
+        const float norm_arg =
+            sumf / float(split_args.n_embd) + split_args.norm_eps;
+        const float norm_scale = rsqrt(norm_arg);
+        device float4 *dst4 = (device float4 *)collapse_dst;
+        device const float4 *w4 = (device const float4 *)norm_weight;
+        device float4 *norm4 = (device float4 *)norm_dst;
+        dst4[i0] = v0;
+        norm4[i0] = (v0 * norm_scale) * w4[i0];
+        if (i1 < n4_collapse) {
+            dst4[i1] = v1;
+            norm4[i1] = (v1 * norm_scale) * w4[i1];
+        }
+    } else if (tgpig.x == 1 && tid == 0) {
+        device float *out = (device float *)split;
+        const float4 post_z =
+            *((device volatile const float4 *)(mixes_f32 + 4)) * hc_scale[1] +
+            *((device const float4 *)(hc_base + 4));
+        *((device float4 *)(out + 4)) = 2.0f / (1.0f + exp(-post_z));
+    }
+
+    // Groups 2..5 own exactly the comb range consumed by the
+    // continuation.  Their four-way completion overlaps TG0's independent
+    // pre-collapse/RMS epilogue.  Every writer crosses the uniform publish
+    // fence; only lane zero then participates in the completion protocol.
+    atomic_thread_fence(mem_flags::mem_device,
+                        memory_order_seq_cst,
+                        thread_scope_device);
+    if (tgpig.x < 2 || tid != 0) {
+        return;
+    }
+
+    const uint old = atomic_fetch_add_explicit(
+        completion, 1u, memory_order_relaxed);
+    if (old + 1u != 4u) {
+        return;
+    }
+    atomic_thread_fence(mem_flags::mem_device,
+                        memory_order_seq_cst,
+                        thread_scope_device);
+    ds4_hc_comb_weights4_exact(
+        split_args, mixes_f32, hc_scale, hc_base,
+        (device float *)split);
+    atomic_thread_fence(mem_flags::mem_device,
+                        memory_order_seq_cst,
+                        thread_scope_device);
+    atomic_store_explicit(completion, 0u, memory_order_relaxed);
+}

metal/dsv4_kv.metal

@@ -319,6 +319,228 @@ kernel void kernel_dsv4_compressor_pack_ratio4(
     }
 }
 
+// Decode already holds the complete ratio-4 recurrent window in state layout:
+// eight rows of two head_dim planes. Pack the previous plane from rows 0..3
+// and the current plane from rows 4..7 directly into the transposed [head_dim,
+// 8] layout consumed by the exact GGML softmax/multiply/sum sequence. KV and
+// score move together; no arithmetic or reduction order changes.
+kernel void kernel_dsv4_compressor_pack_ratio4_decode_ggml(
+        constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
+        device const uint * state_kv,
+        device const uint * state_score,
+        device       uint * packed_kv,
+        device       uint * packed_score,
+        uint row [[threadgroup_position_in_grid]],
+        uint tid [[thread_index_in_threadgroup]]) {
+    if (row >= 8u || args.head_dim == 0u || args.n_threads == 0u) {
+        return;
+    }
+
+    const uint64_t state_row_stride = 2ull * args.head_dim;
+    const uint64_t src_plane = row >= 4u ? args.head_dim : 0u;
+    for (uint col = tid; col < args.head_dim; col += args.n_threads) {
+        const uint64_t src = (uint64_t)row * state_row_stride +
+                             src_plane + col;
+        const uint64_t dst = (uint64_t)col * 8u + row;
+        packed_kv[dst] = state_kv[src];
+        packed_score[dst] = state_score[src];
+    }
+}
+
+// Exact decode specialization for the first two operations in GGML's
+// softmax -> multiply -> sum_rows compressor reduction. The normalized
+// softmax values are deliberately materialized in device memory and reloaded
+// after a device barrier before the in-place product, preserving the dispatch
+// boundary's float store/load semantics. The final sum remains the standalone
+// eight-thread sum_rows kernel: changing a 32-thread group into the original
+// eight-thread reduction inside this kernel would make its threadgroup
+// barriers non-uniform or alter simd_sum's active-lane topology.
+kernel void kernel_dsv4_compressor_exact_softmax_product_ratio4(
+        constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
+        device const float * packed_kv,
+        device const float * packed_score,
+        device       float * softmax,
+        device       float * product,
+        threadgroup  float * softmax_scratch [[threadgroup(0)]],
+        uint row [[threadgroup_position_in_grid]],
+        uint tid [[thread_position_in_threadgroup]]) {
+    if (row >= args.head_dim || args.n_comp != 1u ||
+        args.n_threads != 32u) {
+        return;
+    }
+
+    device const float4 * score4 =
+        (device const float4 *)(packed_score + (uint64_t)row * 8u);
+    device float4 * softmax4 =
+        (device float4 *)(softmax + (uint64_t)row * 8u);
+    const float scale = (float)args.replay;
+    const float zero = (float)(args.n_comp - 1u);
+
+    // Match kernel_soft_max_f32_4(width=8, nth=32) literally. Only lanes zero
+    // and one own float4s, while all 32 lanes participate in both reductions.
+    float4 lmax4 = -INFINITY;
+    for (int i00 = (int)tid; i00 < 2; i00 += 32) {
+        lmax4 = fmax(lmax4, score4[i00] * scale + (float4)zero);
+    }
+
+    const float lmax =
+        MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
+    const float max_val = simd_max(lmax);
+
+    float4 lsum4 = 0.0f;
+    for (int i00 = (int)tid; i00 < 2; i00 += 32) {
+        const float4 exp_score4 =
+            exp((score4[i00] * scale + (float4)zero) - max_val);
+        lsum4 += exp_score4;
+        softmax4[i00] = exp_score4;
+    }
+
+    const float lsum =
+        lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
+    threadgroup_barrier(mem_flags::mem_none);
+    const float sum = simd_sum(lsum);
+    const float inv_sum = 1.0f / sum;
+
+    for (int i00 = (int)tid; i00 < 2; i00 += 32) {
+        softmax4[i00] *= inv_sum;
+    }
+
+    // Force the same normalized-softmax device store/reload boundary that the
+    // separate multiply dispatch observes.
+    threadgroup_barrier(mem_flags::mem_device);
+    device volatile const float * reloaded_softmax =
+        (device volatile const float *)(softmax + (uint64_t)row * 8u);
+    device const float * kv_row = packed_kv + (uint64_t)row * 8u;
+    device float * product_row = product + (uint64_t)row * 8u;
+
+    // Match kernel_bin_fuse_f32_f32_f32(width=8, nth=4): four lanes each
+    // process their low element followed by the element four positions later.
+    if (tid < 4u) {
+        for (uint i0 = tid; i0 < 8u; i0 += 4u) {
+            float value = kv_row[i0];
+            value *= reloaded_softmax[i0];
+            product_row[i0] = value;
+        }
+    }
+
+    // All 32 lanes reach the final device barrier. The following standalone
+    // sum_rows dispatch performs the required global reload and exact TG8
+    // two-stage simd_sum topology.
+    threadgroup_barrier(mem_flags::mem_device);
+    (void)softmax_scratch;
+}
+
+// Exact one-dispatch ratio-4 decode pool. This specializes the three-dispatch
+// pack -> exact softmax/product -> sum_rows chain above without changing any
+// floating-point operation or reduction topology. The normalized softmax and
+// product are still materialized and volatile-reloaded through device memory.
+// The two simd_sum calls in the final reduction execute under an eight-lane
+// active mask, exactly matching kernel_sum_rows_f32_f32's original TG8.
+kernel void kernel_dsv4_compressor_exact_pool_ratio4_decode_ggml(
+        constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
+        device const float * state_kv,
+        device const float * state_score,
+        device       float * softmax,
+        device       float * product,
+        device       float * dst,
+        threadgroup  float * sum_scratch [[threadgroup(0)]],
+        uint col [[threadgroup_position_in_grid]],
+        uint tid [[thread_position_in_threadgroup]]) {
+    if (col >= args.head_dim || args.n_comp != 1u ||
+        args.n_threads != 32u) {
+        return;
+    }
+
+    const uint64_t state_row_stride = 2ull * args.head_dim;
+    const float scale = (float)args.replay;
+    const float zero = (float)(args.n_comp - 1u);
+
+    // Match the packed float4 ownership: lane 0 owns rows 0..3 and lane 1
+    // rows 4..7. The gather itself is an integer-addressed bit-preserving load.
+    float4 score_values = -INFINITY;
+    if (tid < 2u) {
+        const uint row0 = 4u * tid;
+        for (uint j = 0u; j < 4u; ++j) {
+            const uint row = row0 + j;
+            const uint64_t src = (uint64_t)row * state_row_stride +
+                                 (row >= 4u ? args.head_dim : 0u) + col;
+            score_values[j] = state_score[src];
+        }
+    }
+
+    const uint64_t scratch_base = (uint64_t)col * 8u;
+    device float4 * softmax4 =
+        (device float4 *)(softmax + scratch_base);
+
+    // Verbatim kernel_soft_max_f32_4(width=8, nth=32) arithmetic.
+    float4 lmax4 = -INFINITY;
+    for (int i00 = (int)tid; i00 < 2; i00 += 32) {
+        lmax4 = fmax(lmax4, score_values * scale + (float4)zero);
+    }
+    const float lmax =
+        MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
+    const float max_val = simd_max(lmax);
+
+    float4 lsum4 = 0.0f;
+    for (int i00 = (int)tid; i00 < 2; i00 += 32) {
+        const float4 exp_score4 =
+            exp((score_values * scale + (float4)zero) - max_val);
+        lsum4 += exp_score4;
+        softmax4[i00] = exp_score4;
+    }
+    const float lsum =
+        lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
+    threadgroup_barrier(mem_flags::mem_none);
+    const float sum = simd_sum(lsum);
+    const float inv_sum = 1.0f / sum;
+    for (int i00 = (int)tid; i00 < 2; i00 += 32) {
+        softmax4[i00] *= inv_sum;
+    }
+
+    threadgroup_barrier(mem_flags::mem_device);
+    device volatile const float * reloaded_softmax =
+        (device volatile const float *)(softmax + scratch_base);
+    device float * product_row = product + scratch_base;
+
+    // Verbatim width=8, TG4 multiply ownership: low element, then +4.
+    if (tid < 4u) {
+        for (uint i0 = tid; i0 < 8u; i0 += 4u) {
+            const uint64_t src = (uint64_t)i0 * state_row_stride +
+                                 (i0 >= 4u ? args.head_dim : 0u) + col;
+            float value = state_kv[src];
+            value *= reloaded_softmax[i0];
+            product_row[i0] = value;
+        }
+    }
+
+    // Preserve the product dispatch's device store/reload boundary.
+    threadgroup_barrier(mem_flags::mem_device);
+    device volatile const float * reloaded_product =
+        (device volatile const float *)product_row;
+
+    // Reproduce kernel_sum_rows_f32_f32(width=8, TG8) literally. MSL defines
+    // simdgroup collectives over active lanes, so the branch recreates the
+    // original eight-lane partial SIMD group inside this 32-thread group.
+    sum_scratch[tid] = 0.0f;
+    float row_sum = 0.0f;
+    if (tid < 8u) {
+        row_sum += reloaded_product[tid];
+        row_sum = simd_sum(row_sum);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (tid == 0u) {
+        sum_scratch[0] = row_sum;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (tid < 8u) {
+        row_sum = sum_scratch[tid];
+        row_sum = simd_sum(row_sum);
+        if (tid == 0u) {
+            dst[col] = row_sum;
+        }
+    }
+}
+
 // Ratio-4 compression keeps two 4-row halves of recurrent state. After an
 // emitted compressed row, the second half becomes the next window's previous
 // half. The old encoder expressed this as four generic copies; this DS4-specific

metal/dsv4_misc.metal

@@ -70,7 +70,7 @@ struct ds4_metal_args_dsv4_indexed_attention {
     uint32_t window;
     uint32_t ratio;
     uint32_t comp_kv_f16;
-    uint32_t pad0;
+    uint32_t n_splits;
     uint64_t q_token_stride;
     uint64_t q_head_stride;
     uint64_t raw_row_stride;
@@ -4915,6 +4915,270 @@ kernel void kernel_dsv4_router_finalize_weights_one_simd(
     }
 }
 
+// M3 decode specialization that materializes the probability
+// transform in device memory before running the exact SIMD selection and
+// weight normalization above. The volatile reload after the device barrier
+// pins the same float store/load boundary as the standalone transform dispatch.
+kernel void kernel_dsv4_router_transform_finalize_weights_one_simd(
+        constant ds4_metal_args_dsv4_router_select_one & args,
+        device const float *logits,
+        device float *probs,
+        device const float *bias,
+        device const int32_t *hash,
+        device const int32_t *tokens,
+        device int32_t *selected,
+        device float *weights,
+        threadgroup float *scratch [[threadgroup(0)]],
+        uint tid [[thread_position_in_threadgroup]]) {
+    if (tid >= 256 || args.hash_mode) return;
+
+    if (tid < 64) {
+        device const float4 *s = (device const float4 *)logits;
+        device float4 *d = (device float4 *)probs;
+        const float4 x = s[tid];
+        const float4 sp = select(log(1.0f + exp(x)), x, x > 20.0f);
+        d[tid] = sqrt(sp);
+    }
+    threadgroup_barrier(mem_flags::mem_device);
+    device volatile const float *reloaded_probs =
+        (device volatile const float *)probs;
+
+    (void)hash;
+    (void)tokens;
+    threadgroup float *score0_tg = scratch;
+    threadgroup int32_t *idx0_tg =
+        (threadgroup int32_t *)(scratch + 256);
+    threadgroup float *score1_tg = scratch + 512;
+    threadgroup int32_t *idx1_tg =
+        (threadgroup int32_t *)(scratch + 768);
+    const float p = reloaded_probs[tid];
+    float score = args.has_bias ? p + bias[tid] : p;
+    int32_t idx = (int32_t)tid;
+    uint cross_stage = 0;
+
+    for (uint k = 2; k <= 256; k <<= 1) {
+        for (uint j = k >> 1; j > 0; j >>= 1) {
+            float peer_score;
+            int32_t peer_idx;
+            bool take_peer;
+            const bool lower = (tid & j) == 0;
+            const bool descending = (tid & k) == 0;
+
+            if (j < 32) {
+                peer_score = simd_shuffle_xor(score, (ushort)j);
+                peer_idx = simd_shuffle_xor(idx, (ushort)j);
+                take_peer = descending
+                    ? (lower ? score < peer_score : score > peer_score)
+                    : (lower ? score > peer_score : score < peer_score);
+                if (take_peer) {
+                    score = peer_score;
+                    idx = peer_idx;
+                }
+            } else {
+                threadgroup float *score_tg =
+                    (cross_stage & 1u) != 0u ? score1_tg : score0_tg;
+                threadgroup int32_t *idx_tg =
+                    (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg;
+                score_tg[tid] = score;
+                idx_tg[tid] = idx;
+                threadgroup_barrier(mem_flags::mem_threadgroup);
+
+                const uint other = tid ^ j;
+                peer_score = score_tg[other];
+                peer_idx = idx_tg[other];
+                take_peer = descending
+                    ? (lower ? score < peer_score : score > peer_score)
+                    : (lower ? score > peer_score : score < peer_score);
+                if (take_peer) {
+                    score = peer_score;
+                    idx = peer_idx;
+                }
+                cross_stage++;
+            }
+        }
+    }
+
+    if (tid < 6) {
+        selected[tid] = idx;
+    }
+    threadgroup_barrier(mem_flags::mem_device);
+
+    threadgroup volatile float *norm_scratch =
+        (threadgroup volatile float *)scratch;
+    if (tid == 0) {
+        device const int32_t *s = selected;
+        norm_scratch[0] = 0.0f;
+        for (uint i = 0; i < 6; i++) {
+            norm_scratch[0] =
+                norm_scratch[0] + reloaded_probs[s[i]];
+        }
+        norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f);
+        norm_scratch[1] = 1.5f / norm_scratch[0];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (tid < 6) {
+        device const int32_t *s = selected;
+        weights[tid] = reloaded_probs[s[tid]] * norm_scratch[1];
+    }
+}
+
+kernel void kernel_dsv4_router_project_select_fused(
+        constant ds4_metal_args_mul_mv & args,
+        constant ds4_metal_args_dsv4_router_select_one & select_args,
+        device const char * src0_router,
+        device const char * src1,
+        device float * logits,
+        device float * probs,
+        device const float * bias,
+        device int32_t * selected,
+        device float * weights,
+        device atomic_uint * completion,
+        threadgroup char * shmem_raw [[threadgroup(0)]],
+        uint3 tgpig [[threadgroup_position_in_grid]],
+        uint3 tpitg [[thread_position_in_threadgroup]],
+        ushort tiisg [[thread_index_in_simdgroup]],
+        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
+    constexpr short NSG = 8;
+    constexpr short NR0 = 2;
+    constexpr short NB  = 32;
+    constexpr short NF  = 16;
+    constexpr short NF4 = NF/4;
+    constexpr short NW  = N_SIMDWIDTH;
+    const uint tid = tpitg.x;
+    const int nb = args.ne00/NB;
+    const int r0 = tgpig.x*NR0;
+    device const float4 *y4 = (device const float4 *)src1;
+    device const half4 *ax4[NR0];
+    FOR_UNROLL (short row = 0; row < NR0; ++row) {
+        ax4[row] = (device const half4 *)
+            (src0_router + (uint64_t)(r0 + row)*args.nb01);
+    }
+    float sumf[NR0] = {0.f};
+    const short ix = tiisg/(NW/NF);
+    const short il = tiisg%(NW/NF);
+    const int ib0 = sgitg*NF + ix;
+    device const float4 *yb4 = y4 + (ib0*NB + il*NF)/4;
+    for (int ib = ib0; ib < nb; ib += NSG*NF) {
+        float4 yl4[NF4];
+        FOR_UNROLL (short i = 0; i < NF4; ++i) {
+            yl4[i] = yb4[i];
+        }
+        FOR_UNROLL (short row = 0; row < NR0; ++row) {
+            device const half4 *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]);
+            }
+            sumf[row] += sumq;
+        }
+        yb4 += NSG*NF*NW/4;
+    }
+    helper_mv_reduce_and_write<NR0>(logits, sumf, r0, args.ne01,
+                                    tiisg, sgitg, shmem_raw);
+
+    threadgroup float *scratch = (threadgroup float *)shmem_raw;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    atomic_thread_fence(mem_flags::mem_device,
+                        memory_order_seq_cst,
+                        thread_scope_device);
+    if (tid == 0) {
+        const uint old = atomic_fetch_add_explicit(
+            completion, 1u, memory_order_relaxed);
+        scratch[0] = old == 127u ? 1.0f : 0.0f;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (scratch[0] == 0.0f) return;
+    atomic_thread_fence(mem_flags::mem_device,
+                        memory_order_seq_cst,
+                        thread_scope_device);
+
+    if (tid < 64) {
+        device volatile const float4 *s =
+            (device volatile const float4 *)logits;
+        device float4 *d = (device float4 *)probs;
+        const float4 xv = s[tid];
+        const float4 sp = select(log(1.0f + exp(xv)), xv, xv > 20.0f);
+        d[tid] = sqrt(sp);
+    }
+    threadgroup_barrier(mem_flags::mem_device);
+    device volatile const float *reloaded_probs =
+        (device volatile const float *)probs;
+
+    threadgroup float *score0_tg = scratch;
+    threadgroup int32_t *idx0_tg =
+        (threadgroup int32_t *)(scratch + 256);
+    threadgroup float *score1_tg = scratch + 512;
+    threadgroup int32_t *idx1_tg =
+        (threadgroup int32_t *)(scratch + 768);
+    const float p = reloaded_probs[tid];
+    float score = select_args.has_bias ? p + bias[tid] : p;
+    int32_t idx = (int32_t)tid;
+    uint cross_stage = 0;
+    for (uint k = 2; k <= 256; k <<= 1) {
+        for (uint j = k >> 1; j > 0; j >>= 1) {
+            float peer_score;
+            int32_t peer_idx;
+            bool take_peer;
+            const bool lower = (tid & j) == 0;
+            const bool descending = (tid & k) == 0;
+            if (j < 32) {
+                peer_score = simd_shuffle_xor(score, (ushort)j);
+                peer_idx = simd_shuffle_xor(idx, (ushort)j);
+                take_peer = descending
+                    ? (lower ? score < peer_score : score > peer_score)
+                    : (lower ? score > peer_score : score < peer_score);
+                if (take_peer) {
+                    score = peer_score;
+                    idx = peer_idx;
+                }
+            } else {
+                threadgroup float *score_tg =
+                    (cross_stage & 1u) != 0u ? score1_tg : score0_tg;
+                threadgroup int32_t *idx_tg =
+                    (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg;
+                score_tg[tid] = score;
+                idx_tg[tid] = idx;
+                threadgroup_barrier(mem_flags::mem_threadgroup);
+                const uint other = tid ^ j;
+                peer_score = score_tg[other];
+                peer_idx = idx_tg[other];
+                take_peer = descending
+                    ? (lower ? score < peer_score : score > peer_score)
+                    : (lower ? score > peer_score : score < peer_score);
+                if (take_peer) {
+                    score = peer_score;
+                    idx = peer_idx;
+                }
+                cross_stage++;
+            }
+        }
+    }
+    if (tid < 6) selected[tid] = idx;
+    threadgroup_barrier(mem_flags::mem_device);
+    threadgroup volatile float *norm_scratch =
+        (threadgroup volatile float *)scratch;
+    if (tid == 0) {
+        norm_scratch[0] = 0.0f;
+        for (uint i = 0; i < 6; ++i) {
+            norm_scratch[0] = norm_scratch[0] + reloaded_probs[selected[i]];
+        }
+        norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f);
+        norm_scratch[1] = 1.5f / norm_scratch[0];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (tid < 6) {
+        weights[tid] = reloaded_probs[selected[tid]] * norm_scratch[1];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    atomic_thread_fence(mem_flags::mem_device,
+                        memory_order_seq_cst,
+                        thread_scope_device);
+    if (tid == 0) {
+        atomic_store_explicit(completion, 0u, memory_order_relaxed);
+    }
+}
+
+
 // Fills the dense compressed-attention mask with -inf. The selected top-k rows
 // are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch.
 kernel void kernel_dsv4_topk_mask(
@@ -5319,6 +5583,117 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8(
     dst4[lane + 96] = o3 * inv_s;
 }
 
+// Each simdgroup owns two heads and updates both from one staged K/V row.
+// This doubles row reuse without increasing the 256-thread workgroup.
+kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual(
+        constant ds4_metal_args_dsv4_indexed_attention &args,
+        device const char *q,
+        device const char *raw_kv,
+        device const char *comp_kv,
+        device const char *topk,
+        device const char *sinks,
+        device char *dst,
+        threadgroup half4 *kv_shared [[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]]) {
+    const uint token = tgpig.x;
+    const uint head0 = tgpig.y*16u + (uint)sg;
+    const uint head1 = head0 + 8u;
+    if (token >= args.n_tokens || head0 >= args.n_head) return;
+
+    device const float4 *qa = (device const float4 *)(q +
+        (uint64_t)token*args.q_token_stride +
+        (uint64_t)head0*args.q_head_stride);
+    half4 qa0 = (half4)qa[lane + 0];
+    half4 qa1 = (half4)qa[lane + 32];
+    half4 qa2 = (half4)qa[lane + 64];
+    half4 qa3 = (half4)qa[lane + 96];
+    half4 qb0 = half4(0.0h), qb1 = half4(0.0h);
+    half4 qb2 = half4(0.0h), qb3 = half4(0.0h);
+    if (head1 < args.n_head) {
+        device const float4 *qb = (device const float4 *)(q +
+            (uint64_t)token*args.q_token_stride +
+            (uint64_t)head1*args.q_head_stride);
+        qb0 = (half4)qb[lane + 0];
+        qb1 = (half4)qb[lane + 32];
+        qb2 = (half4)qb[lane + 64];
+        qb3 = (half4)qb[lane + 96];
+    }
+
+    float Ma = -FLT_MAX/2.0f, Sa = 0.0f;
+    float Mb = -FLT_MAX/2.0f, Sb = 0.0f;
+    float4 ao0 = 0.0f, ao1 = 0.0f, ao2 = 0.0f, ao3 = 0.0f;
+    float4 bo0 = 0.0f, bo1 = 0.0f, bo2 = 0.0f, bo3 = 0.0f;
+
+    const uint qpos = args.pos0 + token;
+    const uint last_pos = args.pos0 + args.n_tokens - 1u;
+    const uint first_raw_pos = last_pos + 1u - args.n_raw;
+    const uint raw_last_pos = first_raw_pos + args.n_raw - 1u;
+    const uint window_first = (args.window != 0u && qpos + 1u > args.window) ?
+        qpos + 1u - args.window : 0u;
+    const uint first = max(first_raw_pos, window_first);
+    const uint last = min(qpos, raw_last_pos);
+    if (first <= last) {
+        for (uint pos = first; pos <= last; pos++) {
+            const uint logical = pos - first_raw_pos;
+            const uint row = (args.raw_start + logical)%args.raw_cap;
+            device const float4 *src = (device const float4 *)(raw_kv +
+                (uint64_t)row*args.raw_row_stride);
+            if (tid < 128) kv_shared[tid] = (half4)src[tid];
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3,
+                args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3);
+            if (head1 < args.n_head) {
+                dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3,
+                    args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3);
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+    }
+
+    const uint visible = min((qpos + 1u)/args.ratio, args.n_comp);
+    device const int32_t *row_topk = (device const int32_t *)(topk +
+        (uint64_t)token*args.topk_token_stride);
+    for (uint i = 0; i < args.top_k; i++) {
+        const int32_t idx = row_topk[i];
+        if (idx < 0) continue;
+        if ((uint)idx >= visible) break;
+        if (tid < 128) {
+            kv_shared[tid] = dsv4_load_cache_h4(comp_kv,
+                args.comp_row_stride, (uint)idx, tid, args.comp_kv_f16 != 0u);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3,
+            args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3);
+        if (head1 < args.n_head) {
+            dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3,
+                args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    dsv4_attend_sink(((device const float *)sinks)[head0],
+        Ma, Sa, ao0, ao1, ao2, ao3);
+    const float ia = Sa == 0.0f ? 0.0f : 1.0f/Sa;
+    device float4 *da = (device float4 *)(dst +
+        (uint64_t)token*args.dst_token_stride +
+        (uint64_t)head0*args.dst_head_stride);
+    da[lane + 0] = ao0*ia; da[lane + 32] = ao1*ia;
+    da[lane + 64] = ao2*ia; da[lane + 96] = ao3*ia;
+    if (head1 < args.n_head) {
+        dsv4_attend_sink(((device const float *)sinks)[head1],
+            Mb, Sb, bo0, bo1, bo2, bo3);
+        const float ib = Sb == 0.0f ? 0.0f : 1.0f/Sb;
+        device float4 *db = (device float4 *)(dst +
+            (uint64_t)token*args.dst_token_stride +
+            (uint64_t)head1*args.dst_head_stride);
+        db[lane + 0] = bo0*ib; db[lane + 32] = bo1*ib;
+        db[lane + 64] = bo2*ib; db[lane + 96] = bo3*ib;
+    }
+}
+
 // Decode specialization of kernel_dsv4_indexed_mixed_attention_heads8.
 // Generation attends one token at a time, so the ratio-4 indexed path spends a
 // visible amount of time repeatedly staging the same K/V row for the eight
@@ -5450,6 +5825,190 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb16(
     dst4[lane + 96] = o3 * inv_s;
 }
 
+// Long-context decode specialization of the indexed mixed-attention path.
+//
+// The ordinary heads8 kernel reuses each K/V row across eight heads, but only
+// launches one threadgroup per head group. Long-context decode therefore has
+// too little parallel work while each group scans its raw and selected rows.
+// This kernel retains the same eight-head reuse while splitting that row
+// sequence across args.n_splits workgroups. A second kernel merges the online
+// softmax partials and applies the attention sink.
+kernel void kernel_dsv4_indexed_mixed_attention_heads8_split(
+        constant ds4_metal_args_dsv4_indexed_attention & args,
+        device const char *q,
+        device const char *raw_kv,
+        device const char *comp_kv,
+        device const char *topk,
+        device       char *tmp,
+        threadgroup half4 *kv_shared [[threadgroup(0)]],
+        uint3  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 rows_per_block = 16u;
+    constexpr uint vecs_per_row = 128u;
+
+    const uint token = tgpig.x;
+    const uint head = tgpig.y * 8u + (uint)sg;
+    const uint split = tgpig.z;
+    const uint n_splits = args.n_splits;
+    if (token >= args.n_tokens || head >= args.n_head ||
+        n_splits < 2u || n_splits > 31u || split >= n_splits) {
+        return;
+    }
+
+    device const float4 *q4 = (device const float4 *)(q +
+        (uint64_t)token * args.q_token_stride +
+        (uint64_t)head  * args.q_head_stride);
+    const half4 q0 = (half4)q4[lane +  0];
+    const half4 q1 = (half4)q4[lane + 32];
+    const half4 q2 = (half4)q4[lane + 64];
+    const half4 q3 = (half4)q4[lane + 96];
+
+    float M = -FLT_MAX/2.0f;
+    float S = 0.0f;
+    float4 o0 = 0.0f;
+    float4 o1 = 0.0f;
+    float4 o2 = 0.0f;
+    float4 o3 = 0.0f;
+
+    const uint qpos = args.pos0 + token;
+    const uint last_pos = args.pos0 + args.n_tokens - 1u;
+    const uint first_raw_pos = last_pos + 1u - args.n_raw;
+    const uint raw_last_pos = first_raw_pos + args.n_raw - 1u;
+    const uint window_first = (args.window != 0u && qpos + 1u > args.window) ?
+        qpos + 1u - args.window : 0u;
+    const uint raw_first = max(first_raw_pos, window_first);
+    const uint raw_last = min(qpos, raw_last_pos);
+    const uint raw_count = raw_first <= raw_last ?
+        raw_last - raw_first + 1u : 0u;
+    const uint total_rows = raw_count + args.top_k;
+    const uint rows_per_split =
+        (total_rows + n_splits - 1u) / n_splits;
+    const uint split_first = min(split * rows_per_split, total_rows);
+    const uint split_last = min(split_first + rows_per_split, total_rows);
+    const uint visible = min((qpos + 1u) / args.ratio, args.n_comp);
+    device const int32_t *row_topk = (device const int32_t *)(topk +
+        (uint64_t)token * args.topk_token_stride);
+
+    for (uint seq0 = split_first; seq0 < split_last;
+         seq0 += rows_per_block) {
+        const uint n_rows = min(rows_per_block, split_last - seq0);
+        for (uint off = (uint)tid;
+             off < n_rows * vecs_per_row;
+             off += 256u) {
+            const uint r = off / vecs_per_row;
+            const uint c = off - r * vecs_per_row;
+            const uint seq = seq0 + r;
+            half4 value = half4(0.0h);
+            if (seq < raw_count) {
+                const uint pos = raw_first + seq;
+                const uint logical = pos - first_raw_pos;
+                const uint row = (args.raw_start + logical) % args.raw_cap;
+                device const float4 *src = (device const float4 *)(raw_kv +
+                    (uint64_t)row * args.raw_row_stride);
+                value = (half4)src[c];
+            } else {
+                const int32_t idx = row_topk[seq - raw_count];
+                if (idx >= 0 && (uint)idx < visible) {
+                    value = dsv4_load_cache_h4(comp_kv,
+                                               args.comp_row_stride,
+                                               (uint)idx,
+                                               c,
+                                               args.comp_kv_f16 != 0u);
+                }
+            }
+            kv_shared[off] = value;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (uint r = 0; r < n_rows; r++) {
+            const uint seq = seq0 + r;
+            bool valid = true;
+            if (seq >= raw_count) {
+                const int32_t idx = row_topk[seq - raw_count];
+                valid = idx >= 0 && (uint)idx < visible;
+            }
+            if (valid) {
+                dsv4_attend_shared_h4_row_at(kv_shared,
+                                             r,
+                                             q0, q1, q2, q3,
+                                             args.scale,
+                                             lane,
+                                             M, S,
+                                             o0, o1, o2, o3);
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    const uint64_t n_rows = (uint64_t)args.n_tokens * args.n_head;
+    const uint64_t row = (uint64_t)token * args.n_head + head;
+    device float4 *partials = (device float4 *)tmp;
+    partials[(row * vecs_per_row + lane +  0u) * n_splits + split] = o0;
+    partials[(row * vecs_per_row + lane + 32u) * n_splits + split] = o1;
+    partials[(row * vecs_per_row + lane + 64u) * n_splits + split] = o2;
+    partials[(row * vecs_per_row + lane + 96u) * n_splits + split] = o3;
+
+    if (lane == 0u) {
+        device float *stats = (device float *)(partials +
+            n_rows * vecs_per_row * n_splits);
+        const uint64_t stat = (row * n_splits + split) * 2u;
+        stats[stat + 0u] = S;
+        stats[stat + 1u] = M;
+    }
+}
+
+kernel void kernel_dsv4_indexed_mixed_attention_heads8_split_reduce(
+        constant ds4_metal_args_dsv4_indexed_attention & args,
+        device const char *tmp,
+        device const char *sinks,
+        device       char *dst,
+        uint tgpig [[threadgroup_position_in_grid]],
+        ushort lane [[thread_index_in_simdgroup]],
+        ushort sg   [[simdgroup_index_in_threadgroup]]) {
+    constexpr uint vecs_per_row = 128u;
+    const uint n_splits = args.n_splits;
+    const uint64_t n_rows = (uint64_t)args.n_tokens * args.n_head;
+    const uint64_t row = tgpig;
+    if (row >= n_rows || n_splits < 2u || n_splits > 31u) {
+        return;
+    }
+
+    device const float4 *partials = (device const float4 *)tmp;
+    device const float *stats = (device const float *)(partials +
+        n_rows * vecs_per_row * n_splits);
+    float part_sum = 0.0f;
+    float part_max = -FLT_MAX/2.0f;
+    if ((uint)lane < n_splits) {
+        const uint64_t stat = (row * n_splits + (uint)lane) * 2u;
+        part_sum = stats[stat + 0u];
+        part_max = stats[stat + 1u];
+    } else if ((uint)lane == n_splits) {
+        const uint head = (uint)(row % args.n_head);
+        part_sum = 1.0f;
+        part_max = ((device const float *)sinks)[head];
+    }
+
+    const float global_max = simd_max(part_max);
+    const float part_scale = part_sum > 0.0f ?
+        exp(part_max - global_max) : 0.0f;
+    const float total_sum = simd_sum(part_sum * part_scale);
+    const float inv_sum = total_sum > 0.0f ? 1.0f / total_sum : 0.0f;
+
+    device float4 *out = (device float4 *)dst + row * vecs_per_row;
+    for (uint i = (uint)sg; i < vecs_per_row; i += 4u) {
+        float4 value = float4(0.0f);
+        if ((uint)lane < n_splits) {
+            value = partials[(row * vecs_per_row + i) * n_splits +
+                             (uint)lane] * part_scale;
+        }
+        value = simd_sum(value);
+        if (lane == 0u) {
+            out[i] = value * inv_sum;
+        }
+    }
+}
+
 static inline float dsv4_indexer_dot128_shared_q(
         float4 c0,
         float4 c1,

metal/dsv4_rope.metal

@@ -41,6 +41,23 @@ struct ds4_metal_args_dsv4_rope_affine_pair {
     float beta_slow;
 };
 
+struct ds4_metal_args_dsv4_head_norm_rope {
+    int32_t n_head;
+    int32_t head_dim;
+    int32_t head_dim4;
+    int32_t n_dims;
+    int32_t n_ctx_orig;
+    int32_t pos0;
+    int32_t inverse;
+    float eps;
+    float freq_base;
+    float freq_scale;
+    float ext_factor;
+    float attn_factor;
+    float beta_fast;
+    float beta_slow;
+};
+
 static float rope_yarn_ramp(const float low, const float high, const int i0) {
     const float y = (i0 / 2 - low) / max(0.001f, high - low);
     return 1.0f - min(1.0f, max(0.0f, y));
@@ -327,36 +344,109 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4(
     *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
 }
 
-// DS4 positions are always affine within one RoPE dispatch. This variant
+// Fuses the per-head RMSNorm and partial Q RoPE while retaining the standalone
-// reconstructs the same wrapped int32 position in-kernel, avoiding the host
+// norm reduction tree and the mode-0 RoPE lane mapping.
-// position array and its buffer binding while preserving the pair lane mapping
+kernel void kernel_dsv4_head_rms_norm_rope_tail_f32(
-// and all floating-point operations of the specialization above.
+        constant ds4_metal_args_dsv4_head_norm_rope & args,
-kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
+        device char * xraw,
-        constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
+        threadgroup float * shmem_f32 [[threadgroup(0)]],
-        device const char * src0 [[buffer(1)]],
+        uint3 tgpig [[threadgroup_position_in_grid]],
-        device       char * dst  [[buffer(4)]],
+        ushort3 tpitg [[thread_position_in_threadgroup]],
-        uint  tid   [[thread_index_in_threadgroup]],
+        ushort sgitg [[simdgroup_index_in_threadgroup]],
-        ushort3 ntg [[threads_per_threadgroup]],
+        ushort tiisg [[thread_index_in_simdgroup]],
-        uint3 tgpig [[threadgroup_position_in_grid]]) {
+        ushort3 ntg [[threads_per_threadgroup]]) {
-    const int i1 = tgpig[0];
+    if (sgitg == 0) {
-    const int i2 = tgpig[1];
+        shmem_f32[tiisg] = 0.0f;
+    }
+
+    const uint head = tgpig.x;
+    const uint tok = tgpig.y;
+    device float4 * x4 = (device float4 *)xraw +
+        ((uint64_t)tok * (uint64_t)args.n_head + head) *
+        (uint64_t)args.head_dim4;
+
+    float sumf = 0.0f;
+    for (int i00 = tpitg.x; i00 < args.head_dim4; i00 += ntg.x) {
+        sumf += dot(x4[i00], x4[i00]);
+    }
+    sumf = simd_sum(sumf);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    sumf = simd_sum(shmem_f32[tiisg]);
+    const float scale = 1.0f / sqrt(sumf / args.head_dim + args.eps);
     const int n_nope = args.head_dim - args.n_dims;
     if (n_nope < 0) {
         return;
     }
 
+    float corr_dims[2];
+    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base,
+                        args.beta_fast, args.beta_slow, corr_dims);
+    const float theta_base = (float)(args.pos0 + (int)tok);
+    const float inv_ndims = -1.0f / args.n_dims;
+    device float * xs = (device float *)x4;
+
+    for (int i0 = tpitg.x; i0 < args.head_dim; i0 += ntg.x) {
+        if (i0 < n_nope) {
+            xs[i0] = xs[i0] * scale;
+            continue;
+        }
+        const int r = i0 - n_nope;
+        if ((r & 1) != 0) {
+            continue;
+        }
+#ifdef DS4_METAL_ROPE_EXP2_LOG2
+        const float theta =
+            theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
+#else
+        const float theta =
+            theta_base * pow(args.freq_base, inv_ndims * r);
+#endif
+        float cos_theta;
+        float sin_theta;
+        rope_yarn(theta, args.freq_scale, corr_dims, r,
+                  args.ext_factor, args.attn_factor,
+                  &cos_theta, &sin_theta);
+        if (args.inverse) {
+            sin_theta = -sin_theta;
+        }
+
+        const float x0 = xs[i0] * scale;
+        const float x1 = xs[i0 + 1] * scale;
+        xs[i0] = x0 * cos_theta - x1 * sin_theta;
+        xs[i0 + 1] = x0 * sin_theta + x1 * cos_theta;
+    }
+}
+
+// DS4 positions are always affine within one RoPE dispatch. This variant
+// reconstructs the same wrapped int32 position in-kernel, avoiding the host
+// position array and its buffer binding while preserving the pair lane mapping
+// and all floating-point operations of the specialization above.
+
+/* Shared, deliberately noinline so that every caller gets bit-identical
+ * trigonometric codegen. The header note about tiny trig codegen changes
+ * flipping sampled tokens is exactly why this body must be compiled once and
+ * shared rather than inlined separately into each kernel. */
+static __attribute__((noinline)) void ds4_rope_tail_pair_affine_row(
+        constant ds4_metal_args_dsv4_rope_affine_pair & args,
+        device const char * src_base,
+        device char * dst_base,
+        int n_nope,
+        uint raw_pos,
+        uint tid,
+        uint nthreads) {
     float corr_dims[2];
     rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
 
-    const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
         const float theta_base = (float)as_type<int>(raw_pos);
     const float inv_ndims = -1.f/args.n_dims;
-    device const char * src_base =
-        src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
-    device char * dst_base =
-        dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
 
-    for (int r = tid; r < args.n_dims; r += ntg.x) {
+    for (int r = tid; r < args.n_dims; r += nthreads) {
         if ((r & 1) != 0) {
             continue;
         }
@@ -381,5 +471,413 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
 
         *((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta;
         *((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta;
+    }}
+
+kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
+        constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
+        device const char * src0 [[buffer(1)]],
+        device       char * dst  [[buffer(4)]],
+        uint  tid   [[thread_index_in_threadgroup]],
+        ushort3 ntg [[threads_per_threadgroup]],
+        uint3 tgpig [[threadgroup_position_in_grid]]) {
+    const int i1 = tgpig[0];
+    const int i2 = tgpig[1];
+    const int n_nope = args.head_dim - args.n_dims;
+    if (n_nope < 0) {
+        return;
+    }
+    const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
+    device const char * src_base =
+        src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
+    device char * dst_base =
+        dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
+    ds4_rope_tail_pair_affine_row(args, src_base, dst_base, n_nope, raw_pos, tid, ntg.x);
+
+}
+
+// Decode-only fusion of the KV RoPE tail with the FP8/raw finalizer. Both were
+// already single 64-thread threadgroups on the same row, back to back, so the
+// pair cost two dispatches (~12.4 us) to touch 2 KB. The RoPE body below is a
+// verbatim copy of kernel_dsv4_rope_tail_f32_inplace_pair_affine specialised to
+// the decode grid (one head, one token, so i1 = i2 = 0) and the finalizer body
+// is a verbatim copy of kernel_dsv4_kv_fp8_store_f32. The barrier between them
+// is required because RoPE writes element pairs across lanes while the raw copy
+// reads them per lane. Arithmetic, order and rounding are unchanged; the header
+// warning above about trigonometric codegen still applies, so this kernel is
+// gated and verified against full-vocabulary logits before promotion.
+kernel void kernel_dsv4_kv_rope_fp8_store_f32(
+        constant ds4_metal_args_dsv4_kv_fp8_store & args,
+        constant ds4_metal_args_dsv4_rope_affine_pair & rope,
+        device        float * kv,
+        device        float * raw_cache,
+        threadgroup   float * scratch [[threadgroup(0)]],
+        uint tid [[thread_index_in_threadgroup]]) {
+    {
+        const int rope_n_nope = rope.head_dim - rope.n_dims;
+        if (rope_n_nope < 0) {
+            return;
+        }
+        ds4_rope_tail_pair_affine_row(rope,
+                                      (device const char *)kv,
+                                      (device char *)kv,
+                                      rope_n_nope,
+                                      rope.pos0,
+                                      tid,
+                                      64u);
+    }
+    /* The RoPE helper writes device-memory pairs that different lanes read
+     * below. A threadgroup-only fence does not make those cross-lane device
+     * writes visible. */
+    threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
+    {
+
+    const int head_dim = args.head_dim;
+    const int n_rot = args.n_rot;
+    const int n_nope = head_dim - n_rot;
+    if (head_dim <= 0 || n_rot < 0 || n_nope < 0 || tid >= 64) {
+        return;
+    }
+
+    device float * raw = raw_cache + (int64_t)args.raw_row * head_dim;
+
+    for (int off = 0; off < n_nope; off += 64) {
+        float v = 0.0f;
+        if (off + (int)tid < n_nope) {
+            v = kv[off + tid];
+            scratch[tid] = abs(v);
+        } else {
+            scratch[tid] = 0.0f;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        for (uint stride = 32; stride > 0; stride >>= 1) {
+            if (tid < stride) {
+                scratch[tid] = max(scratch[tid], scratch[tid + stride]);
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+
+        const float amax = max(scratch[0], 1.0e-4f);
+        const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
+        if (off + (int)tid < n_nope) {
+            const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
+            kv[off + tid] = q;
+            // Diagnostic only: skip the FP16 round-trip that normally matches the
+            // half-typed FlashAttention KV buffer's precision. With this enabled the
+            // indexer will see higher-precision raw values than FlashAttention does,
+            // which is informative but not a production-ready setting.
+#ifdef DS4_METAL_KV_RAW_F32
+            raw[off + tid] = q;
+#else
+            raw[off + tid] = (float)((half)q);
+#endif
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    for (int i = n_nope + tid; i < head_dim; i += 64) {
+#ifdef DS4_METAL_KV_RAW_F32
+        raw[i] = kv[i];
+#else
+        raw[i] = (float)((half)kv[i]);
+#endif
+    }
     }
 }
+
+/* Decode-only sibling of kernel_flash_attn_ext_vec_reduce that also applies the
+ * inverse RoPE tail to the row it just produced, removing a whole dispatch per
+ * layer. Each threadgroup owns one head's entire 512-float row, so the RoPE is
+ * an intra-threadgroup dependency: reduce, barrier, rotate. Both halves call the
+ * same shared noinline helpers the standalone kernels use, so the arithmetic and
+ * its codegen are identical to running the two dispatches back to back. */
+kernel void kernel_flash_attn_ext_vec_reduce_rope(
+        constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
+        device  const char * htmp,
+        device        char * dst,
+        constant ds4_metal_args_dsv4_rope_affine_pair & rope,
+        uint   tgpig[[threadgroup_position_in_grid]],
+        ushort tiitg[[thread_index_in_threadgroup]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    ds4_flash_attn_vec_reduce_row(args, htmp, dst, tgpig, tiisg, sgitg,
+                                  (short)FC_flash_attn_ext_vec_reduce_NWG,
+                                  (short)FC_flash_attn_ext_vec_reduce_DV);
+
+    threadgroup_barrier(mem_flags::mem_device);
+
+    const int n_nope = rope.head_dim - rope.n_dims;
+    if (n_nope < 0) {
+        return;
+    }
+    device char * row = dst + (uint64_t)tgpig * rope.row_bytes;
+    ds4_rope_tail_pair_affine_row(rope,
+                                  (device const char *)row,
+                                  row,
+                                  n_nope,
+                                  rope.pos0,
+                                  tiitg,
+                                  (uint)(32 * FC_flash_attn_ext_vec_reduce_NWG));
+}
+
+struct ds4_metal_args_dsv4_comp_finalize {
+    ds4_metal_args_dsv4_rope_affine_pair rope;
+    float    rms_eps;
+    uint32_t pad0;
+};
+
+/* Decode-only emit-path fusion.  Every ratio-th token, each layer finalizes
+ * one freshly pooled compressor row per compressor: RMS norm, RoPE tail, and
+ * then the FP8 round-trip + F16 commit copy (attention, 512 floats) or the
+ * Hadamard+FP4 QAT (indexer, 128 floats).  Those were seven single-row
+ * dispatches; this kernel is one dispatch with two threadgroups.
+ *
+ * Each phase reproduces its standalone kernel bit-exactly:
+ *  - norm: kernel_rms_norm_mul_f32_4's tree (float4 lanes, simd_sum, zero-
+ *    padded 32-slot cross-simdgroup reduce); 512 uses 128 virtual threads on
+ *    simdgroups 0-3, 128 uses 32 virtual threads on simdgroup 0.
+ *  - rope: ds4_rope_tail_pair_affine_row verbatim (lanes 0-63, nthreads=64).
+ *  - fp8: kernel_dsv4_fp8_kv_quantize_f32's 64-lane shmem max tree and
+ *    round-trip, src==dst so the verbatim tail copy is a no-op and dropped.
+ *  - commit: per-element f32->f16 conversion (value-wise exact).
+ *  - qat: kernel_dsv4_indexer_hadamard_fp4_f32's butterfly and per-32 amax
+ *    tree on lanes 0-127.
+ * Threads outside a phase's virtual width still execute every barrier, so
+ * threadgroup barriers stay uniform across the 256-thread threadgroup. */
+kernel void kernel_dsv4_comp_row_finalize_f32(
+        constant ds4_metal_args_dsv4_comp_finalize & args [[buffer(0)]],
+        device       float * attn_row     [[buffer(1)]],
+        device const float * attn_norm_w  [[buffer(2)]],
+        device       char  * attn_cache   [[buffer(3)]],
+        device       float * index_row    [[buffer(4)]],
+        device const float * index_norm_w [[buffer(5)]],
+        device       float * attn_state_kv    [[buffer(6)]],
+        device       float * attn_state_score [[buffer(7)]],
+        device       float * index_state_kv    [[buffer(8)]],
+        device       float * index_state_score [[buffer(9)]],
+        threadgroup  float * shmem        [[threadgroup(0)]],
+        uint  tgpig [[threadgroup_position_in_grid]],
+        ushort tiitg [[thread_index_in_threadgroup]],
+        ushort tiisg [[thread_index_in_simdgroup]],
+        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
+    constant ds4_metal_args_dsv4_rope_affine_pair & rope_args = args.rope;
+
+    if (tgpig == 0) {
+        /* -------- attention compressor row (512 floats) -------- */
+        {
+            device float4 * y4 = (device float4 *)attn_row;
+            device const float4 * x4 = (device const float4 *)attn_row;
+            device const float4 * w4 = (device const float4 *)attn_norm_w;
+            if (sgitg == 0) {
+                shmem[tiisg] = 0.0f;
+            }
+            float sumf = 0.0f;
+            if (tiitg < 128) {
+                sumf = dot(x4[tiitg], x4[tiitg]);
+            }
+            sumf = simd_sum(sumf);
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            if (tiitg < 128 && tiisg == 0) {
+                shmem[sgitg] = sumf;
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            float total = 0.0f;
+            if (tiitg < 128) {
+                total = simd_sum(shmem[tiisg]);
+            }
+            const float mean  = total / 512.0f;
+            const float scale = 1.0f/sqrt(mean + args.rms_eps);
+            if (tiitg < 128) {
+                y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg];
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_device);
+        ds4_rope_tail_pair_affine_row(rope_args,
+                                      (device const char *)attn_row,
+                                      (device char *)attn_row,
+                                      512 - rope_args.n_dims,
+                                      rope_args.pos0,
+                                      tiitg,
+                                      64u);
+        threadgroup_barrier(mem_flags::mem_device);
+        for (int off = 0; off < 512 - rope_args.n_dims; off += 64) {
+            float v = 0.0f;
+            if (tiitg < 64) {
+                v = attn_row[off + tiitg];
+                shmem[tiitg] = abs(v);
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            for (uint stride = 32; stride > 0; stride >>= 1) {
+                if (tiitg < stride) {
+                    shmem[tiitg] = max(shmem[tiitg], shmem[tiitg + stride]);
+                }
+                threadgroup_barrier(mem_flags::mem_threadgroup);
+            }
+            const float amax = max(shmem[0], 1.0e-4f);
+            const float scale = exp2(ceil(log2(amax / 448.0f)));
+            if (tiitg < 64) {
+                const float q = dsv4_e4m3fn_dequant(clamp(v / scale, -448.0f, 448.0f)) * scale;
+                attn_row[off + tiitg] = q;
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+        threadgroup_barrier(mem_flags::mem_device);
+        if (tiitg < 128) {
+            device const float4 * x4 = (device const float4 *)attn_row;
+            device half4 * o4 = (device half4 *)attn_cache;
+            const float4 v = x4[tiitg];
+            o4[tiitg] = half4(v);
+        }
+        return;
+    }
+
+    if (tgpig >= 2u) {
+        /* Ratio-4 state shifts for both compressors (elementwise row move,
+         * so the flat gid mapping is bit-exact): 4*1024 attention elements
+         * then 4*256 indexer elements. */
+        const uint gid = (tgpig - 2u) * 256u + tiitg;
+        const uint n0 = 4u * 1024u;
+        if (gid < n0) {
+            attn_state_kv[gid] = attn_state_kv[n0 + gid];
+            attn_state_score[gid] = attn_state_score[n0 + gid];
+            return;
+        }
+        const uint gid1 = gid - n0;
+        const uint n1 = 4u * 256u;
+        if (gid1 >= n1) return;
+        index_state_kv[gid1] = index_state_kv[n1 + gid1];
+        index_state_score[gid1] = index_state_score[n1 + gid1];
+        return;
+    }
+
+    /* -------- indexer compressor row (128 floats) -------- */
+    {
+        device float4 * y4 = (device float4 *)index_row;
+        device const float4 * x4 = (device const float4 *)index_row;
+        device const float4 * w4 = (device const float4 *)index_norm_w;
+        if (sgitg == 0) {
+            shmem[tiisg] = 0.0f;
+        }
+        float sumf = 0.0f;
+        if (tiitg < 32) {
+            sumf = dot(x4[tiitg], x4[tiitg]);
+        }
+        sumf = simd_sum(sumf);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (tiisg == 0) {
+            shmem[sgitg] = sumf;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        float total = 0.0f;
+        if (tiitg < 32) {
+            total = simd_sum(shmem[tiisg]);
+        }
+        const float mean  = total / 128.0f;
+        const float scale = 1.0f/sqrt(mean + args.rms_eps);
+        if (tiitg < 32) {
+            y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg];
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_device);
+    ds4_rope_tail_pair_affine_row(rope_args,
+                                  (device const char *)index_row,
+                                  (device char *)index_row,
+                                  128 - rope_args.n_dims,
+                                  rope_args.pos0,
+                                  tiitg,
+                                  64u);
+    threadgroup_barrier(mem_flags::mem_device);
+    {
+        threadgroup float *vals = shmem;
+        threadgroup float *absbuf = shmem + 128;
+        if (tiitg < 128) {
+            vals[tiitg] = index_row[tiitg];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (uint stride = 1u; stride < 128u; stride <<= 1u) {
+            if (tiitg < 128 && (tiitg & stride) == 0u) {
+                const uint base = (tiitg & ~(2u * stride - 1u)) + (tiitg & (stride - 1u));
+                const float a = vals[base];
+                const float b = vals[base + stride];
+                vals[base] = a + b;
+                vals[base + stride] = a - b;
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+        float v = 0.0f;
+        if (tiitg < 128) {
+            v = vals[tiitg] * 0.08838834764831845f;
+            absbuf[tiitg] = abs(v);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        const uint block = tiitg >> 5u;
+        const uint lane = tiitg & 31u;
+        const uint block_base = block * 32u;
+        for (uint stride = 16u; stride > 0u; stride >>= 1u) {
+            if (tiitg < 128 && lane < stride) {
+                absbuf[block_base + lane] = max(absbuf[block_base + lane],
+                                                absbuf[block_base + lane + stride]);
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+        if (tiitg < 128) {
+            const float amax = max(absbuf[block_base], 7.052966104933725e-38f);
+            const float scale = exp2(ceil(log2(amax / 6.0f)));
+            index_row[tiitg] = dsv4_e2m1fn_dequant(clamp(v / scale, -6.0f, 6.0f)) * scale;
+        }
+    }
+}
+
+// Host-visible packed FlashAttention + exact inverse-RoPE decode kernel.
+kernel void kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512(
+        constant ds4_metal_args_flash_attn_ext_vec & args [[buffer(0)]],
+        device const char * q                         [[buffer(1)]],
+        device const char * k                         [[buffer(2)]],
+        device const char * v                         [[buffer(3)]],
+        device const char * mask                      [[buffer(4)]],
+        device const char * sinks                     [[buffer(5)]],
+        device const char * pad                       [[buffer(6)]],
+        device       char * dst                       [[buffer(7)]],
+        constant ds4_metal_args_dsv4_rope_affine_pair & rope
+                                                        [[buffer(8)]],
+        threadgroup char * shmem [[threadgroup(0)]],
+        uint head       [[threadgroup_position_in_grid]],
+        ushort tiitg    [[thread_index_in_threadgroup]],
+        ushort tiisg    [[thread_index_in_simdgroup]],
+        ushort sgitg    [[simdgroup_index_in_threadgroup]]) {
+    /* Uniform specialization guard; host applies the same eligibility gate. */
+    if (!FC_flash_attn_ext_vec_has_mask ||
+        !FC_flash_attn_ext_vec_has_sinks ||
+         FC_flash_attn_ext_vec_has_bias ||
+         FC_flash_attn_ext_vec_has_scap ||
+         FC_flash_attn_ext_vec_nsg != 1 ||
+         FC_flash_attn_ext_vec_nwg != 32 ||
+         FC_flash_attn_ext_vec_ns10 != 512 ||
+         FC_flash_attn_ext_vec_ns20 != 512 ||
+         args.ne01 != 1 || args.ne02 != 64 || args.ne03 != 1 ||
+         args.ne_12_2 != 1 || args.ne_12_3 != 1 ||
+         args.ne31 != 1 || args.ne32 != 1 || args.ne33 != 1 ||
+         args.ne11 <= 0 || args.ne11 > 1024 || head >= (uint)args.ne02 ||
+         args.nb02 != 2048 || args.nb11 != 1024 || args.nb21 != 1024 ||
+         rope.head_dim != 512 || rope.n_dims != 64 ||
+         rope.row_bytes != 2048 || rope.inverse == 0) {
+        return;
+    }
+
+    ds4_flash_attn_vec_packed8_reduce_f16_512(
+        args, q, k, v, mask, sinks, pad, dst, shmem,
+        head, tiisg, sgitg);
+
+    /* Same producer/consumer boundary as the current reduce+RoPE kernel. */
+    threadgroup_barrier(mem_flags::mem_device);
+
+    const int n_nope = rope.head_dim - rope.n_dims;
+    device char * row = dst + (uint64_t)head * rope.row_bytes;
+    ds4_rope_tail_pair_affine_row(rope,
+                                  (device const char *)row,
+                                  row,
+                                  n_nope,
+                                  rope.pos0,
+                                  tiitg,
+                                  32u * 32u);
+}

metal/flash_attn.metal

@@ -1398,24 +1398,26 @@ constant int32_t FC_flash_attn_ext_vec_reduce_NWG [[function_constant(FC_FLASH_A
 
 // Reduces split-K decode FlashAttention partials. It combines each workgroup's
 // output vector and softmax (sum,max) pair into the final attention result.
-kernel void kernel_flash_attn_ext_vec_reduce(
+/* Shared and deliberately noinline so the split-K reduction is compiled once and
+ * every caller gets identical codegen. The RoPE-fused sibling in dsv4_rope.metal
+ * calls this same body, which is what keeps the fusion bit-exact. */
+static __attribute__((noinline)) void ds4_flash_attn_vec_reduce_row(
         constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
         device  const char * htmp,
         device        char * dst,
-        uint   tgpig[[threadgroup_position_in_grid]],
+        uint   tgpig,
-        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort tiisg,
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+        ushort sgitg,
-#define NWG (FC_flash_attn_ext_vec_reduce_NWG)
+        short  NWG_,
-#define DV  (FC_flash_attn_ext_vec_reduce_DV)
+        short  DV_) {
-
     const uint64_t rid = tgpig;
 
     const short iwg = tiisg;
 
-    device const float  * ss    = (device const float  *) htmp + (uint64_t)args.nrows*DV*NWG;
+    device const float  * ss    = (device const float  *) htmp + (uint64_t)args.nrows*DV_*NWG_;
 
-    float S = ss[rid*(2*NWG) + 2*iwg + 0];
+    float S = ss[rid*(2*NWG_) + 2*iwg + 0];
-    float M = ss[rid*(2*NWG) + 2*iwg + 1];
+    float M = ss[rid*(2*NWG_) + 2*iwg + 1];
 
     const float m  = simd_max(M);
     const float ms = exp(M - m);
@@ -1423,19 +1425,268 @@ kernel void kernel_flash_attn_ext_vec_reduce(
     S = simd_sum(S*ms);
     S = S == 0.0f ? 0.0f : 1.0f/S;
 
-    const short DV4 = DV/4;
+    const short DV4 = DV_/4;
 
-    device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG;
+    device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG_;
     device       float4 * dst4  = (device       float4 *) dst  + rid*DV4;
 
-    for (short i = sgitg; i < DV4; i += NWG) {
+    for (short i = sgitg; i < DV4; i += NWG_) {
-        const float4 v = simd_sum(htmp4[i*NWG + iwg]*ms);
+        const float4 v = simd_sum(htmp4[i*NWG_ + iwg]*ms);
 
         if (iwg == 0) {
             dst4[i] = v*S;
         }
     }
-
+}
-#undef NWG
+
-#undef DV
+kernel void kernel_flash_attn_ext_vec_reduce(
+        constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
+        device  const char * htmp,
+        device        char * dst,
+        uint   tgpig[[threadgroup_position_in_grid]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    ds4_flash_attn_vec_reduce_row(args, htmp, dst, tgpig, tiisg, sgitg,
+                                  (short)FC_flash_attn_ext_vec_reduce_NWG,
+                                  (short)FC_flash_attn_ext_vec_reduce_DV);
+}
+
+// M5 decode specialization: time-slice all 32 split-K workgroups through eight
+// physical simdgroups, then reduce through the same 32-lane topology without a
+// device partial buffer. The host gate fixes the exact F16 512-wide geometry.
+static inline void ds4_flash_attn_vec_packed8_reduce_f16_512(
+        constant ds4_metal_args_flash_attn_ext_vec & args,
+        device const char * q,
+        device const char * k,
+        device const char * v,
+        device const char * mask,
+        device const char * sinks,
+        device const char * pad,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint head,
+        ushort tiisg,
+        ushort sgitg) {
+    constexpr short NW = 32;
+    constexpr short C = 32;
+    constexpr short NSG = 8;
+    constexpr short NWG = 32;
+    constexpr short DK4 = 128;
+    constexpr short DV4 = 128;
+    constexpr short SH = 128;
+
+    /* 24,448 dynamic bytes: shared Q, eight score/mask banks, all 32
+     * split-local weights and stats, sink scales, and a padded 32x33 F32
+     * float4 partial plane. */
+    threadgroup half4 *q_shared = (threadgroup half4 *)shmem;
+    threadgroup half *score_banks =
+        (threadgroup half *)(q_shared + DK4);
+    threadgroup volatile float *weights =
+        (threadgroup volatile float *)(score_banks + NSG * SH);
+    threadgroup volatile float *stats = weights + NWG * C;
+    threadgroup volatile float *sink_scale = stats + 2 * NWG;
+    threadgroup volatile float4 *partial_plane =
+        (threadgroup volatile float4 *)(sink_scale + NWG);
+
+    const short lane = (short)tiisg;
+    threadgroup half *bank = score_banks + (short)sgitg * SH;
+    threadgroup float *ss = (threadgroup float *)bank;
+    threadgroup half *sm = bank + 2 * C;
+
+    device const float4 *q4 =
+        (device const float4 *)(q + (uint64_t)head * args.nb02);
+    if (sgitg == 0) {
+        for (short i = lane; i < DK4; i += NW) {
+            q_shared[i] = (half4)q4[i];
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    /* Eight physical simdgroups time-slice the exact 32 legacy split-K
+     * workgroups. The official <=1024-key gate gives each virtual split at
+     * most one 32-row block, so its value partial can be formed later from
+     * these materialized weights without changing online-softmax order. */
+    for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) {
+        float S = 0.0f;
+        float M = -FLT_MAX / 2;
+        float out_scale = 1.0f;
+        const int ic_original = (int)iwg * C;
+
+        weights[(uint)iwg * C + (uint)lane] = 0.0f;
+        ss[lane] = 0.0f;
+        sm[lane] = (half)0.0h;
+        simdgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (ic_original < args.ne11) {
+            device const char *k_block = k;
+            device const char *v_block = v;
+            device const half *pm = (device const half *)mask;
+            int ic = ic_original;
+
+            if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
+                k_block = pad;
+                const uint64_t k_pad_bytes =
+                    args.nb11 * (uint64_t)C *
+                    (uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
+                const uint64_t v_pad_bytes =
+                    args.nb21 * (uint64_t)C *
+                    (uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
+                if (FC_flash_attn_ext_vec_shared_kvpad) {
+                    v_block = k_block;
+                    pm = (device const half *)(k_block +
+                                                k_pad_bytes + v_pad_bytes);
+                } else {
+                    v_block = k_block + k_pad_bytes;
+                    pm = (device const half *)(v_block + v_pad_bytes);
+                }
+                ic = 0;
+            }
+
+            sm[lane] = pm[ic + lane];
+            if (simd_max(sm[lane]) > -MAXHALF) {
+                device const half4 *pk4 =
+                    (device const half4 *)(k_block +
+                                           (uint64_t)ic * args.nb11);
+                threadgroup const half4 *pq4 = q_shared;
+                pk4 += lane;
+                pq4 += lane;
+
+                float lane_mqk = 0.0f;
+                FOR_UNROLL (short cc = 0; cc < C; ++cc) {
+                    float mqk = 0.0f;
+                    FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) {
+                        mqk += dot((float4)pk4[cc * DK4 + ii * NW],
+                                   (float4)pq4[ii * NW]);
+                    }
+                    mqk = simd_sum(mqk);
+                    if (lane == cc) {
+                        lane_mqk = mqk;
+                    }
+                }
+
+                ss[lane] = fma(lane_mqk, args.scale,
+                               (float)sm[lane]);
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+
+                const float old_m = M;
+                const float score = ss[lane];
+                M = simd_max(max(M, score));
+                const float ms = exp(old_m - M);
+                const float vs = exp(score - M);
+                S = S * ms + simd_sum(vs);
+                ss[lane] = vs;
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+
+                weights[(uint)iwg * C + (uint)lane] = ss[lane];
+            }
+
+            if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) {
+                const float old_m = M;
+                const float sink = lane == 0
+                    ? ((device const float *)sinks)[head]
+                    : -FLT_MAX / 2;
+                M = simd_max(max(M, sink));
+                const float ms = exp(old_m - M);
+                const float vs = exp(sink - M);
+                S = S * ms + simd_sum(vs);
+                out_scale = ms;
+            }
+        } else if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) {
+            const float old_m = M;
+            const float sink = lane == 0
+                ? ((device const float *)sinks)[head]
+                : -FLT_MAX / 2;
+            M = simd_max(max(M, sink));
+            const float ms = exp(old_m - M);
+            const float vs = exp(sink - M);
+            S = S * ms + simd_sum(vs);
+            out_scale = ms;
+        }
+
+        if (lane == 0) {
+            stats[2 * (uint)iwg + 0] = S;
+            stats[2 * (uint)iwg + 1] = M;
+            sink_scale[(uint)iwg] = out_scale;
+        }
+        simdgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    /* Recreate the legacy 32-lane reducer exactly: lane is the virtual split
+     * index, including neutral idle splits in their original tree positions. */
+    const short split = lane;
+    float reduce_S = stats[2 * (uint)split + 0];
+    float reduce_M = stats[2 * (uint)split + 1];
+    const float reduce_max = simd_max(reduce_M);
+    const float reduce_ms = exp(reduce_M - reduce_max);
+    reduce_S = simd_sum(reduce_S * reduce_ms);
+    const float reduce_inv =
+        reduce_S == 0.0f ? 0.0f : 1.0f / reduce_S;
+
+    device float4 *dst4 =
+        (device float4 *)(dst +
+                          (uint64_t)head * 512u * sizeof(float));
+
+    /* Form one 32-float4 output quadrant at a time. During production each
+     * physical simdgroup time-slices four virtual splits while SIMD lanes are
+     * contiguous output columns, exactly matching the legacy V loads and
+     * cc-major accumulation. A padded 33-column plane avoids a 32-way TG-bank
+     * conflict when the reducer transposes lanes back to virtual splits. */
+    for (short quadrant = 0; quadrant < 4; ++quadrant) {
+        for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) {
+            float4 lo = float4(0.0f);
+            const int ic_original = (int)iwg * C;
+            if (ic_original < args.ne11) {
+                device const char *v_block = v;
+                int ic = ic_original;
+                if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
+                    device const char *k_block = pad;
+                    const uint64_t k_pad_bytes =
+                        args.nb11 * (uint64_t)C *
+                        (uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
+                    if (FC_flash_attn_ext_vec_shared_kvpad) {
+                        v_block = k_block;
+                    } else {
+                        v_block = k_block + k_pad_bytes;
+                    }
+                    ic = 0;
+                }
+
+                device const half4 *pv4 =
+                    (device const half4 *)(v_block +
+                                           (uint64_t)ic * args.nb21);
+                threadgroup volatile float *split_weights =
+                    weights + (uint)iwg * C;
+                const short oc = quadrant * NW + lane;
+                FOR_UNROLL (short cc = 0; cc < C; ++cc) {
+                    lo += float4(pv4[cc * DV4 + oc]) *
+                          float4(split_weights[cc]);
+                }
+
+                float4 acc = float4(0.0f);
+                acc += lo;
+                if (iwg == 0) {
+                    acc *= sink_scale[0];
+                }
+                lo = acc;
+            }
+            partial_plane[(uint)iwg * 33u + (uint)lane] = lo;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        /* lane is now the legacy split index and each physical simdgroup
+         * reduces four output columns through the identical simd_sum tree. */
+        for (short out_lane = (short)sgitg; out_lane < NW; out_lane += NSG) {
+            const float4 materialized =
+                (float4)partial_plane[(uint)lane * 33u + (uint)out_lane];
+            const float4 reduced = simd_sum(materialized * reduce_ms);
+            if (lane == 0) {
+                dst4[quadrant * NW + out_lane] = reduced * reduce_inv;
+            }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
 }

metal/norm.metal

@@ -241,3 +241,145 @@ kernel void kernel_dsv4_qkv_rms_norm_f32_4(
         y[i] = (x[i] * scale) * w[i];
     }
 }
+
+// Decode-only triple fusion: the q/kv RMS norm, the KV RoPE tail, and the
+// FP8/raw finalizer were three back-to-back dispatches on the same rows.
+// The q threadgroup is byte-identical to kernel_dsv4_qkv_rms_norm_f32_4.
+// The kv threadgroup continues with the shared affine-row RoPE helper (lane
+// mapping preserved: r == lane on the first 64 lanes) and a verbatim copy of
+// kernel_dsv4_kv_fp8_store_f32 with its work predicated to the first 64
+// lanes (barriers stay uniform across the whole threadgroup).  Arithmetic,
+// order and rounding are unchanged; gated and verified against
+// full-vocabulary logits before promotion.
+kernel void kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32(
+        constant ds4_metal_args_qkv_rms_norm & args,
+        constant ds4_metal_args_dsv4_rope_affine_pair & rope,
+        constant ds4_metal_args_dsv4_kv_fp8_store & store,
+        device const float4 * q_src,
+        device const float4 * q_weight,
+        device       float4 * q_dst,
+        device const float4 * kv_src,
+        device const float4 * kv_weight,
+        device       float4 * kv_dst,
+        device       float  * raw_cache,
+        threadgroup float * shmem_f32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3 ntg[[threads_per_threadgroup]]) {
+    if (sgitg == 0) {
+        shmem_f32[tiisg] = 0.0f;
+    }
+
+    const uint row = tgpig.x;
+    const bool kv_task = tgpig.y != 0;
+    const int n = kv_task ? args.kv_n : args.q_n;
+    const int n4 = kv_task ? args.kv_n4 : args.q_n4;
+    const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4);
+
+    device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4;
+    device const float4 * w = kv_task ? kv_weight : q_weight;
+    device       float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4;
+
+    float sumf = 0.0f;
+    for (int i = tpitg.x; i < n4; i += ntg.x) {
+        const float4 v = x[i];
+        sumf += dot(v, v);
+    }
+    sumf = simd_sum(sumf);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    sumf = shmem_f32[tiisg];
+    sumf = simd_sum(sumf);
+
+#ifdef DS4_METAL_NORM_RSQRT_DISABLE
+    const float scale = 1.0f / sqrt(sumf / float(n) + args.eps);
+#else
+    const float scale = rsqrt(sumf / float(n) + args.eps);
+#endif
+
+    for (int i = tpitg.x; i < n4; i += ntg.x) {
+        y[i] = (x[i] * scale) * w[i];
+    }
+
+    if (!kv_task) {
+        return;
+    }
+
+    // KV RoPE tail in place, then the FP8/raw finalizer (verbatim bodies).
+    threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
+
+    device char *kv_row = (device char *)(kv_dst + row * row_stride4);
+    const int rope_n_nope = rope.head_dim - rope.n_dims;
+    if (rope_n_nope < 0) {
+        return;
+    }
+    ds4_rope_tail_pair_affine_row(rope,
+                                  (device const char *)kv_row,
+                                  kv_row,
+                                  rope_n_nope,
+                                  rope.pos0,
+                                  tpitg.x,
+                                  ntg.x);
+
+    threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
+
+    const int head_dim = store.head_dim;
+    const int n_rot = store.n_rot;
+    const int n_nope = head_dim - n_rot;
+    if (head_dim <= 0 || n_rot < 0 || n_nope < 0) {
+        return;
+    }
+    const uint tid = tpitg.x;
+
+    device float *kv = (device float *)kv_row;
+    device float *raw = raw_cache + (int64_t)store.raw_row * head_dim;
+    threadgroup float *scratch = shmem_f32 + 32;
+
+    for (int off = 0; off < n_nope; off += 64) {
+        float v = 0.0f;
+        if (tid < 64u && off + (int)tid < n_nope) {
+            v = kv[off + tid];
+            scratch[tid] = abs(v);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        for (uint stride = 32; stride > 0; stride >>= 1) {
+            if (tid < stride) {
+                scratch[tid] = max(scratch[tid], scratch[tid + stride]);
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+
+        const float amax = max(scratch[0], 1.0e-4f);
+        const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
+        if (tid < 64u && off + (int)tid < n_nope) {
+            const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
+            kv[off + tid] = q;
+#ifdef DS4_METAL_KV_RAW_F32
+            raw[off + tid] = q;
+#else
+            raw[off + tid] = (float)((half)q);
+#endif
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    if (tid < 64u) {
+        for (int i = n_nope + tid; i < head_dim; i += 64) {
+#ifdef DS4_METAL_KV_RAW_F32
+            raw[i] = kv[i];
+#else
+            raw[i] = (float)((half)kv[i]);
+#endif
+        }
+    }
+}

native/metal/README.md

@@ -1,7 +1,7 @@
 # Vendored DS4 Metal boundary
 
-These files are a one-time snapshot of the DS4 Metal boundary from commit
+These files track the DS4 Metal boundary audited at commit
-`efdadd41e20134af4f3381e1ed90e96fe4faef6f`:
+`8db89fe083ae4d17c9a2428ccd29803d3ae8f577` (2026-08-28):
 
 - `ds4_metal.m`
 - `ds4.h`

native/metal/ds4.h

@@ -150,6 +150,7 @@ typedef struct {
     bool glm_mtp_timing;
     bool dspark;
     bool dspark_strict;
+    bool dspark_exact_sampling;
     bool dspark_confidence_threshold_set;
     bool cuda_tensor_parallel;
     bool ssd_streaming;
@@ -158,6 +159,8 @@ typedef struct {
     bool inspect_only;
     /* Multi-GPU placement uses this to price per-layer KV storage. */
     int placement_ctx_hint;
+    /* Number of independently allocated session graphs/caches to reserve. */
+    int placement_session_count_hint;
     /* Server batch mode serializes execution and can share prefill scratch. */
     bool share_session_prefill_workspace;
     bool first_token_test;
@@ -371,6 +374,30 @@ int ds4_test_sample_logits(const float *logits, uint32_t n_vocab,
                            float temperature, int top_k,
                            float top_p, float min_p, uint64_t *rng,
                            float *prob_scratch);
+int ds4_test_sampling_probabilities(const float *logits, uint32_t n_vocab,
+                                    float temperature, int top_k,
+                                    float top_p, float min_p, float *probs);
+int ds4_test_speculative_sample(const float *target_logits,
+                                const float *draft_logits,
+                                uint32_t n_vocab,
+                                float temperature,
+                                int top_k,
+                                float top_p,
+                                float min_p,
+                                uint64_t *rng,
+                                float *target_probs,
+                                float *draft_probs);
+int ds4_test_speculative_delta_sample(const float *target_logits,
+                                      uint32_t n_vocab,
+                                      int draft_token,
+                                      float temperature,
+                                      int top_k,
+                                      float top_p,
+                                      float min_p,
+                                      uint64_t *rng,
+                                      float *target_probs);
+int ds4_test_argmax_excluding_logits(const float *logits, uint32_t n_vocab,
+                                     int excluded_id);
 uint64_t ds4_test_mixed_native_count(void);
 #endif
 int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k);
@@ -403,6 +430,15 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token,
                                         int max_tokens, int eos_token,
                                         int *accepted, int accepted_cap,
                                         char *err, size_t errlen);
+/* Evaluate one already-sampled target token and speculatively extend it.
+ * Positive-temperature DSpark normally commits greedily verified draft
+ * tokens; dspark_exact_sampling selects exact stochastic p/q acceptance. */
+int ds4_session_eval_speculative(ds4_session *s, int first_token,
+                                 int max_tokens, int eos_token,
+                                 float temperature, int top_k,
+                                 float top_p, float min_p, uint64_t *rng,
+                                 int *accepted, int accepted_cap,
+                                 char *err, size_t errlen);
 /* TP worker side of a mirrored speculative-verify block: run its half of the
  * batch verify for KV side effects, then obey the leader's commit frame
  * (keep, or roll back and replay). Only called from ds4_tp_worker_run. */

native/metal/ds4_gpu.h

@@ -78,6 +78,24 @@ int ds4_gpu_begin_commands(void);
 int ds4_gpu_flush_encoder(void);
 int ds4_gpu_flush_commands(void);
 int ds4_gpu_commands_active(void);
+#ifdef __APPLE__
+int ds4_gpu_parallel_ffn_finish(void);
+void ds4_gpu_parallel_ffn_abort(void);
+int ds4_gpu_parallel_ffn_start(
+        ds4_gpu_tensor       *gate,
+        ds4_gpu_tensor       *up,
+        ds4_gpu_tensor       *mid,
+        ds4_gpu_tensor       *shared_out,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              gate_offset,
+        uint64_t              up_offset,
+        uint64_t              down_offset,
+        uint32_t              model_dim,
+        uint32_t              shared_dim,
+        const ds4_gpu_tensor *x,
+        float                 clamp);
+#endif
 int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value);
 int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label);
 int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label);
@@ -95,6 +113,10 @@ int ds4_gpu_synchronize(void);
 int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size);
 int ds4_gpu_set_model_fd(int fd);
 int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map);
+int ds4_gpu_build_derived_artifacts(const void *model_map, uint64_t model_size,
+                                    const char *model_path);
+int ds4_gpu_model_range_replaced(const void *model_map, uint64_t offset,
+                                 uint64_t bytes);
 int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes);
 int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes);
 int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label);
@@ -141,7 +163,30 @@ void ds4_gpu_set_glm_model(bool enabled);
 void ds4_gpu_set_ssd_streaming(bool enabled);
 void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled);
 #ifdef __APPLE__
+int ds4_gpu_device_is_pre_m5_apple_silicon(void);
+int ds4_gpu_device_is_m5_apple_silicon(void);
+int ds4_gpu_set_decode_pipeline_fast_lookup(int enabled);
+/* Strict test oracle for the fixed decode mul_mv pipeline lookup cache. */
+int ds4_gpu_test_decode_pipeline_fast_lookup(void);
+/* Strict test oracle for the extended decode mul_mv_ext (nsg + nxpsg) cache. */
+int ds4_gpu_test_decode_pipeline_fast_lookup_ext(void);
+/* Strict test oracle for the generated resident-prefill MXFP4 half LUT. */
+int ds4_gpu_test_mxfp4_down_half_lut(uint16_t *legacy_bits,
+                                     uint16_t *lut_bits);
+enum {
+    DS4_GPU_TEST_MXFP4_PAIR_TAIL_CULL = 1u << 0,
+    DS4_GPU_TEST_MXFP4_PAIR_COMPACT_TILE = 1u << 1,
+    DS4_GPU_TEST_MXFP4_MAP_SCATTER = 1u << 2,
+    DS4_GPU_TEST_MXFP4_DOWN_TAIL_CULL = 1u << 3,
+    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,
+};
+void ds4_gpu_test_set_flags(uint32_t flags);
 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; }
+static inline int ds4_gpu_device_is_m5_apple_silicon(void) { return 0; }
 #endif
 void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts);
 void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes);
@@ -216,6 +261,7 @@ int ds4_gpu_stream_expert_cache_seed_from_layer_selected(
         uint32_t                          n_tokens,
         uint32_t                          n_seed_tokens,
         uint32_t                          n_selected);
+int ds4_gpu_stream_expert_cache_finish_pending_batch(void);
 int ds4_gpu_stream_expert_cache_release_layer_cache(void);
 #endif
 int ds4_gpu_stream_expert_cache_seed_experts(
@@ -223,6 +269,14 @@ int ds4_gpu_stream_expert_cache_seed_experts(
         const int32_t                     *expert_ids,
         const uint32_t                    *expert_priorities,
         uint32_t                           n_experts);
+#ifdef __APPLE__
+/* Seed from mapped weights with blits appended to the active command buffer. */
+int ds4_gpu_stream_expert_cache_seed_experts_gpu_copy(
+        const ds4_gpu_stream_expert_table *table,
+        const int32_t                     *expert_ids,
+        const uint32_t                    *expert_priorities,
+        uint32_t                           n_experts);
+#endif
 void ds4_gpu_print_memory_report(const char *label);
 
 /* Tensor-parallel per-layer gates (Metal only).  The encoder calls
@@ -659,6 +713,36 @@ int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
         uint64_t                out_dim,
         const ds4_gpu_tensor *x,
         float                   clamp);
+
+int ds4_gpu_router_shared_gate_up_q8_0_tensor(
+        ds4_gpu_tensor       *router_logits,
+        ds4_gpu_tensor       *gate,
+        ds4_gpu_tensor       *up,
+        ds4_gpu_tensor       *mid,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              router_weight_offset,
+        uint64_t              gate_offset,
+        uint64_t              up_offset,
+        uint64_t              in_dim,
+        uint64_t              router_out_dim,
+        uint64_t              out_dim,
+        const ds4_gpu_tensor *x,
+        float                 clamp,
+        bool                  router_only);
+#ifdef __APPLE__
+int ds4_gpu_router_project_select_fused_tensor(
+        ds4_gpu_tensor       *router_logits,
+        ds4_gpu_tensor       *probs,
+        ds4_gpu_tensor       *selected,
+        ds4_gpu_tensor       *weights,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              router_weight_offset,
+        uint64_t              bias_offset,
+        bool                  has_bias,
+        const ds4_gpu_tensor *x);
+#endif
 int ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor(
         ds4_gpu_tensor       *mid,
         const void             *model_map,
@@ -736,6 +820,20 @@ int ds4_gpu_matmul_f16_tensor(
         const ds4_gpu_tensor *x,
         uint64_t                n_tok);
 
+/* CUDA batch path: fold an input RMS normalization into the FP16 activation
+ * conversion used by the following projection. Returns 0 without touching
+ * out when the optimized path is unavailable. */
+int ds4_gpu_matmul_f16_rms_fold_tensor(
+        ds4_gpu_tensor       *out,
+        const void             *model_map,
+        uint64_t                model_size,
+        uint64_t                weight_offset,
+        uint64_t                in_dim,
+        uint64_t                out_dim,
+        const ds4_gpu_tensor *x,
+        uint64_t                n_tok,
+        float                   norm_eps);
+
 /* Exact multi-row form of the DeepSeek 4096x256 F16 router projection. */
 int ds4_gpu_matmul_f16_router_rows_exact_tensor(
         ds4_gpu_tensor       *out,
@@ -777,6 +875,95 @@ int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
         uint32_t                ratio,
         uint32_t                pos);
 
+int ds4_gpu_matmul_f16_quad_compressor_store_tensor(
+        ds4_gpu_tensor       *out0_kv,
+        ds4_gpu_tensor       *out0_score,
+        ds4_gpu_tensor       *out1_kv,
+        ds4_gpu_tensor       *out1_score,
+        ds4_gpu_tensor       *state0_kv,
+        ds4_gpu_tensor       *state0_score,
+        ds4_gpu_tensor       *state1_kv,
+        ds4_gpu_tensor       *state1_score,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              weight0_kv_offset,
+        uint64_t              weight0_score_offset,
+        uint64_t              weight1_kv_offset,
+        uint64_t              weight1_score_offset,
+        uint64_t              ape0_offset,
+        uint32_t              ape0_type,
+        uint64_t              ape1_offset,
+        uint32_t              ape1_type,
+        uint64_t              in_dim,
+        uint32_t              width0,
+        uint32_t              width1,
+        const ds4_gpu_tensor *x,
+        uint32_t              ratio,
+        uint32_t              pos);
+
+/* Decode-only M5 fusion: emit-path compressor row finalize (norm + rope +
+ * fp8/commit + indexer qat) in one dispatch.  Bit-exact vs the separate
+ * dispatches.  Returns 1 when fused, 0 to fall back. */
+int ds4_gpu_dsv4_comp_row_finalize_tensor(
+        ds4_gpu_tensor       *attn_stage,
+        ds4_gpu_tensor       *attn_cache,
+        uint32_t              attn_comp_row,
+        uint64_t              attn_norm_offset,
+        ds4_gpu_tensor       *index_cache,
+        uint32_t              index_comp_row,
+        uint64_t              index_norm_offset,
+        ds4_gpu_tensor       *attn_state_kv,
+        ds4_gpu_tensor       *attn_state_score,
+        ds4_gpu_tensor       *index_state_kv,
+        ds4_gpu_tensor       *index_state_score,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint32_t              pos,
+        uint32_t              n_rot,
+        uint32_t              n_ctx_orig,
+        float                 freq_base,
+        float                 freq_scale,
+        float                 ext_factor,
+        float                 attn_factor,
+        float                 beta_fast,
+        float                 beta_slow,
+        float                 rms_eps);
+
+/* Decode-only M5 fusion: q_a/kv Q8 pair projection + F16 quad compressor
+ * projection/store in one dispatch.  Bit-exact vs the separate dispatches.
+ * Returns 1 when fused, 0 to fall back, -1 on error. */
+int ds4_gpu_qkv_pair_quad_compressor_store_tensor(
+        ds4_gpu_tensor       *qr,
+        ds4_gpu_tensor       *kv_raw,
+        ds4_gpu_tensor       *out0_kv,
+        ds4_gpu_tensor       *out0_score,
+        ds4_gpu_tensor       *out1_kv,
+        ds4_gpu_tensor       *out1_score,
+        ds4_gpu_tensor       *state0_kv,
+        ds4_gpu_tensor       *state0_score,
+        ds4_gpu_tensor       *state1_kv,
+        ds4_gpu_tensor       *state1_score,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              q_a_offset,
+        uint64_t              kv_offset,
+        uint64_t              weight0_kv_offset,
+        uint64_t              weight0_score_offset,
+        uint64_t              weight1_kv_offset,
+        uint64_t              weight1_score_offset,
+        uint64_t              ape0_offset,
+        uint32_t              ape0_type,
+        uint64_t              ape1_offset,
+        uint32_t              ape1_type,
+        uint32_t              in_dim,
+        uint32_t              q_rank,
+        uint32_t              kv_dim,
+        uint32_t              width0,
+        uint32_t              width1,
+        const ds4_gpu_tensor *x,
+        uint32_t              ratio,
+        uint32_t              pos);
+
 int ds4_gpu_matmul_f32_tensor(
         ds4_gpu_tensor       *out,
         const void             *model_map,
@@ -857,6 +1044,31 @@ int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
         uint32_t                rows,
         float                   eps);
 
+int ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_tensor(
+        ds4_gpu_tensor       *q_out,
+        const ds4_gpu_tensor *q,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              q_weight_offset,
+        uint32_t              q_n,
+        ds4_gpu_tensor       *kv_out,
+        const ds4_gpu_tensor *kv,
+        uint64_t              kv_weight_offset,
+        uint32_t              kv_n,
+        ds4_gpu_tensor       *raw_cache,
+        uint64_t              raw_cap,
+        uint32_t              raw_row,
+        uint32_t              n_rot,
+        uint32_t              pos0,
+        uint32_t              n_ctx_orig,
+        float                 freq_base,
+        float                 freq_scale,
+        float                 ext_factor,
+        float                 attn_factor,
+        float                 beta_fast,
+        float                 beta_slow,
+        float                 eps);
+
 int ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor(
         ds4_gpu_tensor       *q_out,
         const ds4_gpu_tensor *q,
@@ -942,6 +1154,8 @@ int ds4_gpu_dsv4_indexer_qat_tensor(
         uint32_t          n_rows,
         uint32_t          head_dim);
 
+
+
 int ds4_gpu_rope_tail_tensor(
         ds4_gpu_tensor *x,
         uint32_t          n_tok,
@@ -1609,7 +1823,9 @@ int ds4_gpu_compressor_update_tensor(
         float                   beta_fast,
         float                   beta_slow,
         float                   rms_eps,
-        bool                    state_already_stored);
+        bool                    state_already_stored,
+        bool                    decode_one_token,
+        bool                    defer_finalize);
 
 int ds4_gpu_compressor_store_batch_tensor(
         const ds4_gpu_tensor *kv,
@@ -2409,18 +2625,6 @@ int ds4_gpu_hc_weighted_sum_tensor(
         uint32_t                n_embd,
         uint32_t                n_hc);
 
-int ds4_gpu_hc_weighted_sum_norm_tensor(
-        ds4_gpu_tensor       *out,
-        ds4_gpu_tensor       *norm_out,
-        const ds4_gpu_tensor *residual_hc,
-        const ds4_gpu_tensor *weights,
-        const void             *model_map,
-        uint64_t                model_size,
-        uint64_t                norm_weight_offset,
-        uint32_t                n_embd,
-        uint32_t                n_hc,
-        float                   norm_eps);
-
 int ds4_gpu_hc_weighted_sum_split_tensor(
         ds4_gpu_tensor       *out,
         const ds4_gpu_tensor *residual_hc,
@@ -2461,6 +2665,17 @@ int ds4_gpu_hc_split_weighted_sum_norm_tensor(
         float                   eps,
         float                   norm_eps);
 
+int ds4_gpu_hc_rms_norm_mix_f16_available(void);
+int ds4_gpu_hc_rms_norm_mix_f16_tensor(
+        ds4_gpu_tensor       *out,
+        const ds4_gpu_tensor *x,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              weight_offset,
+        uint32_t              n,
+        uint32_t              out_dim,
+        float                 eps);
+
 /* Batched HC RMSNorm followed by its narrow F16 mixer projection. On the
  * tuned Metal path, scale_scratch stores one float per row instead of the
  * full normalized HC tensor; other shapes retain the established fallback. */
@@ -2476,6 +2691,29 @@ int ds4_gpu_hc_rms_scale_project_f16_tensor(
         uint32_t                n_rows,
         float                   eps);
 
+#ifdef __APPLE__
+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);
+
+#endif
 int ds4_gpu_output_hc_weights_tensor(
         ds4_gpu_tensor       *out,
         const ds4_gpu_tensor *pre,
@@ -2612,6 +2850,34 @@ int ds4_gpu_matmul_q8_0_hc_expand_tensor(
         uint32_t                n_embd,
         uint32_t                n_hc);
 
+/* Decode-island CUDA graph capture (CUDA backend; Metal/ROCm/CPU stub it
+ * out and stay eager).  Design ported from the Entrpi/ds4 batched-serving
+ * fork's per-layer decode graph capture.  The key identifies a captured
+ * island: layer, island index, and the activation buffers whose addresses
+ * the captured kernels bake in.  ds4_cuda.cu mirrors this struct
+ * byte-for-byte (it does not include this header); keep both in sync. */
+typedef struct ds4_decode_graph_key {
+    uint32_t il;
+    uint32_t island;    /* 0: layer top to pre-rope; 1: attn-out to layer end */
+    uint32_t variant;
+    uint32_t _pad;
+    void    *cur_hc;
+    void    *after_attn_hc;
+    void    *after_ffn_hc;
+    void    *attn_norm;
+} ds4_decode_graph_key;
+
+int  ds4_gpu_decode_graphs_supported(void);
+/* 1: replayed (island already executed; skip encoding it)
+ * 0: capturing (encode the island, then call _end)
+ * -1: run eagerly */
+int  ds4_gpu_decode_graph_begin(const ds4_decode_graph_key *key);
+/* 0: capture committed and launched; -1: capture failed (entry retired;
+ * the caller must re-encode the island eagerly -- no work was executed). */
+int  ds4_gpu_decode_graph_end(const ds4_decode_graph_key *key);
+void ds4_gpu_decode_graph_abort(const ds4_decode_graph_key *key);
+void ds4_gpu_decode_graphs_invalidate(void);
+
 #ifdef __cplusplus
 }
 #endif

src/app.rs

@@ -398,6 +398,7 @@ pub(crate) enum Message {
     PreferenceGlmMtpTimingChanged(bool),
     PreferenceDsparkConfidenceChanged(String),
     PreferenceDsparkStrictChanged(bool),
+    PreferenceDsparkExactSamplingChanged(bool),
     PreferenceSsdChanged(bool),
     PreferenceSsdColdChanged(bool),
     PreferenceSsdCacheChanged(String),
@@ -2719,7 +2720,7 @@ mod tests {
             Some(Message::FocusPrevious)
         ));
         assert!(ModelChoice::DeepSeekV4Flash.supports_dspark());
-        assert!(!ModelChoice::DeepSeekV4Flash0731.supports_dspark());
+        assert!(ModelChoice::DeepSeekV4Flash0731.supports_dspark());
         assert!(!ModelChoice::DeepSeekV4Pro.supports_dspark());
         assert!(!ModelChoice::Glm52.supports_dspark());
     }

src/app/preferences.rs

@@ -45,6 +45,7 @@ pub(super) struct PreferenceDraft {
     pub(super) glm_mtp_timing: bool,
     pub(super) dspark_confidence_threshold: String,
     pub(super) dspark_strict: bool,
+    pub(super) dspark_exact_sampling: bool,
     pub(super) ssd_streaming: bool,
     pub(super) ssd_streaming_cold: bool,
     pub(super) ssd_cache: String,
@@ -111,6 +112,7 @@ impl PreferenceDraft {
             glm_mtp_timing: speculative.glm_mtp_timing,
             dspark_confidence_threshold: optional_string(speculative.dspark_confidence_threshold),
             dspark_strict: speculative.dspark_strict,
+            dspark_exact_sampling: speculative.dspark_exact_sampling,
             ssd_streaming: runtime.ssd.enabled,
             ssd_streaming_cold: runtime.ssd.cold,
             ssd_cache: optional_string(runtime.ssd.cache),
@@ -226,6 +228,7 @@ impl PreferenceDraft {
                 &self.dspark_confidence_threshold,
             )?,
             dspark_strict: self.dspark_strict,
+            dspark_exact_sampling: self.dspark_exact_sampling,
         })
     }
 
@@ -267,6 +270,7 @@ impl PreferenceDraft {
         self.glm_mtp_timing = speculative.glm_mtp_timing;
         self.dspark_confidence_threshold = optional_string(speculative.dspark_confidence_threshold);
         self.dspark_strict = speculative.dspark_strict;
+        self.dspark_exact_sampling = speculative.dspark_exact_sampling;
         self.ssd_streaming = ssd.enabled;
         self.ssd_streaming_cold = ssd.cold;
         self.ssd_cache = optional_string(ssd.cache);
@@ -693,12 +697,16 @@ impl App {
                 self.preference_error = None;
             }
             Message::PreferenceLegacyMtpChanged(enabled) => {
-                self.preference_draft.legacy_mtp_enabled =
+                self.preference_draft.legacy_mtp_enabled = self
-                    self.preference_draft.acceleration_model.supports_dspark() && enabled;
+                    .preference_draft
+                    .acceleration_model
+                    .supports_legacy_mtp()
+                    && enabled;
                 if self.preference_draft.legacy_mtp_enabled {
                     self.preference_draft.dspark_enabled = false;
                     self.preference_draft.dspark_confidence_threshold.clear();
                     self.preference_draft.dspark_strict = false;
+                    self.preference_draft.dspark_exact_sampling = false;
                 }
                 self.preference_error = None;
             }
@@ -708,6 +716,7 @@ impl App {
                 if !self.preference_draft.dspark_enabled {
                     self.preference_draft.dspark_confidence_threshold.clear();
                     self.preference_draft.dspark_strict = false;
+                    self.preference_draft.dspark_exact_sampling = false;
                 } else {
                     self.preference_draft.legacy_mtp_enabled = false;
                 }
@@ -903,6 +912,15 @@ impl App {
                 }
                 self.preference_error = None;
             }
+            Message::PreferenceDsparkExactSamplingChanged(value) => {
+                self.preference_draft.dspark_exact_sampling =
+                    self.preference_draft.acceleration_model.supports_dspark() && value;
+                if self.preference_draft.dspark_exact_sampling {
+                    self.preference_draft.dspark_enabled = true;
+                    self.preference_draft.legacy_mtp_enabled = false;
+                }
+                self.preference_error = None;
+            }
             Message::PreferenceSsdChanged(value) => {
                 self.preference_draft.ssd_streaming = value;
                 self.preference_error = None;
@@ -1011,14 +1029,14 @@ mod tests {
         assert_eq!(draft.context_tokens, "32768");
         draft.context_tokens = "456".into();
         draft
-            .select_generation(ModelChoice::DeepSeekV4Flash, ReasoningMode::High)
+            .select_generation(ModelChoice::DeepSeekV4Flash0731, ReasoningMode::High)
             .unwrap();
         assert_eq!(draft.context_tokens, "123");
 
         draft.select_acceleration(ModelChoice::Glm52).unwrap();
         assert!(!draft.ssd_streaming);
         draft
-            .select_acceleration(ModelChoice::DeepSeekV4Flash)
+            .select_acceleration(ModelChoice::DeepSeekV4Flash0731)
             .unwrap();
         assert!(draft.ssd_streaming);
     }

src/app/view/preferences.rs

@@ -6,7 +6,7 @@ impl App {
         let legacy_mtp_toggle: Option<fn(bool) -> Message> = self
             .preference_draft
             .acceleration_model
-            .supports_dspark()
+            .supports_legacy_mtp()
             .then_some(Message::PreferenceLegacyMtpChanged);
         let legacy_mtp = hint(
             toggle(self.preference_draft.legacy_mtp_enabled)
@@ -36,6 +36,11 @@ impl App {
             .acceleration_model
             .supports_dspark()
             .then_some(Message::PreferenceDsparkStrictChanged);
+        let dspark_exact_toggle: Option<fn(bool) -> Message> = self
+            .preference_draft
+            .acceleration_model
+            .supports_dspark()
+            .then_some(Message::PreferenceDsparkExactSamplingChanged);
         let effective = self
             .preference_draft
             .effective_for(
@@ -74,7 +79,7 @@ impl App {
             text_input("Automatic", &self.preference_draft.directional_steering_ffn);
         let mut steering_attn = text_input("0", &self.preference_draft.directional_steering_attn);
         let mut dspark_confidence = text_input(
-            "0.9 (DS4 default)",
+            "0.6 (DS4 default)",
             &self.preference_draft.dspark_confidence_threshold,
         );
         if self.preference_draft.model != ModelChoice::Glm52 {
@@ -509,7 +514,7 @@ impl App {
                 dspark,
                 preference_input_row(
                     "DSpark confidence threshold",
-                    "How sure the draft model must be, from 0 to 1, before its token is handed to the verifier. Lower forwards more guesses for more speed and more rejected work; blank uses DS4's 0.9.",
+                    "How sure the draft model must be, from 0 to 1, before its token is handed to the verifier. Lower forwards more guesses for more speed and more rejected work; blank uses DS4's 0.6, or 0.8 for exact sampling.",
                     dspark_confidence,
                 ),
                 hint(
@@ -518,6 +523,12 @@ impl App {
                     .on_toggle_maybe(dspark_strict_toggle),
                     "Lets the draft model only propose, never decide: every token is sampled by the full model. Gives up some of the speedup in exchange for output identical to non-speculative decoding.",
                 ),
+                hint(
+                    toggle(self.preference_draft.dspark_exact_sampling)
+                    .label("Use exact DSpark sampling")
+                    .on_toggle_maybe(dspark_exact_toggle),
+                    "For non-zero temperatures, applies DS4's exact acceptance and corrected rejection sampling. Off uses the faster opportunistic mode: sample a boundary token, then accept DSpark tokens only while they match the target's greedy path.",
+                ),
                 text(if self.preference_draft.acceleration_model.supports_dspark() {
                     "Legacy MTP and DSpark use separate managed support artifacts; entering a DSpark threshold or enabling strict mode selects DSpark."
                 } else if self.preference_draft.acceleration_model == ModelChoice::Glm52 {
@@ -532,7 +543,7 @@ impl App {
                     |engine| {
                         let settings = engine.speculative;
                         format!(
-                            "Engine: MTP draft {} • margin {} • legacy MTP {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}",
+                            "Engine: MTP draft {} • margin {} • legacy MTP {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {} • exact sampling {}",
                             settings.mtp_draft_tokens,
                             settings.mtp_margin,
                             if self.preference_draft.legacy_mtp_enabled { "on" } else { "off" },
@@ -542,6 +553,7 @@ impl App {
                             settings.dspark_confidence_threshold,
                             if settings.dspark_confidence_threshold_set { " explicit" } else { " default" },
                             if settings.dspark_strict { "on" } else { "off" },
+                            if settings.dspark_exact_sampling { "on" } else { "off" },
                         )
                     },
                 ))

src/engine.rs

@@ -8,7 +8,7 @@ mod validation;
 
 #[cfg(target_os = "macos")]
 use crate::metrics::{KvLookup, Metrics, SsdStats};
-use crate::model::ModelChoice;
+use crate::model::{ModelChoice, validate_engine_artifacts};
 #[cfg(target_os = "macos")]
 use crate::settings::TurnSettings;
 use crate::settings::{EngineSettings, ReasoningMode};
@@ -232,6 +232,12 @@ pub(crate) struct ModelSummary {
 impl Model {
     #[allow(dead_code)]
     pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
+        validate_engine_artifacts(
+            settings.model,
+            settings.artifacts.mtp.is_some() && !settings.speculative.dspark,
+            settings.speculative.dspark,
+            &settings.artifacts,
+        )?;
         let mut model = Self::open_main(&settings.artifacts.model, settings.model)?;
         if settings.execution.warm_weights {
             model.main.warm()?;
@@ -1225,8 +1231,17 @@ impl Generator {
                     cancelled,
                 )?
             } else {
-                self.executor.eval(token)?;
+                self.executor.eval_speculative_sampled(
-                vec![token]
+                    token,
+                    generation_limit - generated_tokens,
+                    settings.reasoning_mode,
+                    settings.temperature,
+                    settings.top_p,
+                    settings.min_p,
+                    settings.top_k,
+                    &mut rng,
+                    cancelled,
+                )?
             };
             self.publish_execution_stats();
             for token in cycle {
@@ -1545,12 +1560,30 @@ fn sample(
     top_k: i32,
     rng: &mut Rng,
 ) -> i32 {
-    if temperature <= 0.0 {
+    let probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k);
-        return logits
+    sample_probabilities(&probabilities, rng, None)
+}
+
+#[cfg(any(target_os = "macos", test))]
+fn sampling_probabilities(
+    logits: &[f32],
+    temperature: f32,
+    top_p: f32,
+    min_p: f32,
+    top_k: i32,
+) -> Vec<(usize, f32)> {
+    let greedy = || {
+        vec![(
+            logits
                 .iter()
                 .enumerate()
                 .max_by(|a, b| a.1.total_cmp(b.1))
-            .map_or(0, |(index, _)| index as i32);
+                .map_or(0, |(index, _)| index),
+            1.0,
+        )]
+    };
+    if temperature <= 0.0 {
+        return greedy();
     }
     let maximum = logits
         .iter()
@@ -1558,7 +1591,7 @@ fn sample(
         .filter(|value| value.is_finite())
         .fold(f32::NEG_INFINITY, f32::max);
     if !maximum.is_finite() {
-        return 0;
+        return greedy();
     }
     let top_p = if top_p <= 0.0 || top_p > 1.0 {
         1.0
@@ -1571,45 +1604,96 @@ fn sample(
         .enumerate()
         .filter(|(_, logit)| logit.is_finite())
         .map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp()))
-        .filter(|(_, probability)| *probability >= min_p)
         .collect();
     if probabilities.is_empty() {
-        return logits
+        return greedy();
-            .iter()
-            .enumerate()
-            .max_by(|a, b| a.1.total_cmp(b.1))
-            .map_or(0, |(index, _)| index as i32);
     }
-    if top_p < 1.0 || top_k > 0 {
+    if top_p < 1.0 || top_k > 0 || min_p > 0.0 {
         probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
         if top_k > 0 {
-            probabilities.truncate(probabilities.len().min(top_k as usize));
+            probabilities.truncate(probabilities.len().min((top_k as usize).min(1024)));
         }
     }
-    if top_p < 1.0 {
     let total: f32 = probabilities
         .iter()
         .map(|(_, probability)| probability)
         .sum();
     let mut kept = 0.0;
-        let count = probabilities
+    let mut count = 0;
-            .iter()
+    for (_, probability) in &probabilities {
-            .position(|(_, probability)| {
+        if count > 0 && *probability < min_p {
-                kept += *probability;
+            break;
-                kept / total >= top_p
+        }
-            })
+        kept += *probability;
-            .map_or(probabilities.len(), |index| index + 1);
+        count += 1;
-        probabilities.truncate(count);
+        if kept / total >= top_p {
+            break;
+        }
+    }
+    probabilities.truncate(count);
+    if probabilities.is_empty() || !kept.is_finite() || kept <= 0.0 {
+        return greedy();
+    }
+    for (_, probability) in &mut probabilities {
+        *probability /= kept;
+    }
+    if top_p >= 1.0 && top_k <= 0 {
+        probabilities.sort_unstable_by_key(|(token, _)| *token);
+    }
+    probabilities
+}
+
+#[cfg(any(target_os = "macos", test))]
+fn sample_probabilities(
+    probabilities: &[(usize, f32)],
+    rng: &mut Rng,
+    excluded: Option<usize>,
+) -> i32 {
+    let total: f32 = probabilities
+        .iter()
+        .filter(|(token, _)| Some(*token) != excluded)
+        .map(|(_, probability)| probability)
+        .sum();
+    let mut choice = rng.unit() * total;
+    for (token, probability) in probabilities {
+        if Some(*token) == excluded {
+            continue;
         }
-    let kept_total: f32 = probabilities.iter().map(|(_, p)| p).sum();
-    let mut choice = rng.unit() * kept_total;
-    for (token, probability) in &probabilities {
         choice -= probability;
         if choice <= 0.0 {
             return *token as i32;
         }
     }
-    probabilities.last().map_or(0, |(token, _)| *token as i32)
+    probabilities
+        .iter()
+        .rev()
+        .find(|(token, _)| Some(*token) != excluded)
+        .map_or(0, |(token, _)| *token as i32)
+}
+
+#[cfg(any(target_os = "macos", test))]
+fn exact_delta_sample(
+    logits: &[f32],
+    draft: i32,
+    temperature: f32,
+    top_p: f32,
+    min_p: f32,
+    top_k: i32,
+    rng: &mut Rng,
+) -> (i32, bool) {
+    let mut probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k);
+    let draft_probability = probabilities
+        .iter()
+        .find(|(token, _)| *token == draft as usize)
+        .map_or(0.0, |(_, probability)| *probability);
+    if rng.unit() <= draft_probability {
+        return (draft, true);
+    }
+    probabilities.sort_unstable_by_key(|(token, _)| *token);
+    (
+        sample_probabilities(&probabilities, rng, Some(draft as usize)),
+        false,
+    )
 }
 
 #[cfg(any(target_os = "macos", test))]
@@ -1665,6 +1749,37 @@ mod sampling_tests {
         assert_eq!(chunks, ["hello "]);
     }
 
+    #[test]
+    fn exact_delta_sampling_accepts_or_corrects_the_draft() {
+        let mut accept_rng = Rng::new(2);
+        let (accepted, was_draft) =
+            exact_delta_sample(&[10.0, 0.0], 0, 1.0, 1.0, 0.0, 0, &mut accept_rng);
+        assert_eq!((accepted, was_draft), (0, true));
+
+        let mut reject_rng = Rng::new(1);
+        let (replacement, was_draft) =
+            exact_delta_sample(&[0.0, 10.0], 0, 1.0, 1.0, 0.0, 0, &mut reject_rng);
+        assert_eq!((replacement, was_draft), (1, false));
+    }
+
+    #[test]
+    fn sampling_probabilities_match_ds4_filter_order() {
+        let probabilities = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 0.8, 0.2, 0);
+        assert_eq!(probabilities.len(), 2);
+        assert_eq!(probabilities[0].0, 1);
+        assert_eq!(probabilities[1].0, 2);
+        assert!((probabilities.iter().map(|(_, value)| value).sum::<f32>() - 1.0).abs() < 1e-6);
+
+        let min_p_only = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 1.0, 0.2, 0);
+        assert_eq!(
+            min_p_only
+                .iter()
+                .map(|(token, _)| *token)
+                .collect::<Vec<_>>(),
+            [1, 2]
+        );
+    }
+
     #[test]
     fn split_utf8_token_bytes_are_joined_before_decoding() {
         let mut generated = ChatTurn {

src/engine/metal.rs

@@ -10,7 +10,7 @@ use profile::ExpertProfile;
 
 use super::gguf::{F16, F32, Gguf, IQ2_XXS, Q4_K, Q8_0, Tensor as GgufTensor};
 use super::validation::{DsparkConfig, SupportKind, dspark_config};
-use super::{Model, ModelFamily};
+use super::{Model, ModelFamily, Rng, exact_delta_sample};
 use crate::model::ModelChoice;
 use crate::settings::{
     EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
@@ -407,7 +407,13 @@ impl Dspark {
             capture_mask: 0,
             cache_start: 0,
             cache_len: 0,
-            confidence_threshold: settings.dspark_confidence_threshold,
+            confidence_threshold: if settings.dspark_exact_sampling
+                && !settings.dspark_confidence_threshold_set
+            {
+                settings.dspark_confidence_threshold.max(0.8)
+            } else {
+                settings.dspark_confidence_threshold
+            },
             strict: settings.dspark_strict || quality,
             drafted: 0,
             accepted: 0,
@@ -2600,6 +2606,48 @@ struct SpecFrontier {
     dspark_cache_len: u32,
 }
 
+struct SpecPrefixFrontier {
+    layers: Vec<LayerFrontier>,
+}
+
+struct BatchVerification {
+    tops: Vec<i32>,
+    logits: Vec<Vec<f32>>,
+    prefixes: Vec<SpecPrefixFrontier>,
+}
+
+fn capture_compression_frontier(
+    state: &CompressionState,
+    bytes: u64,
+    purpose: &str,
+) -> Result<CompressionFrontier, String> {
+    let state_kv = Buffer::bytes(bytes)?;
+    let state_score = Buffer::bytes(bytes)?;
+    state_kv.copy_from(0, &state.state_kv, 0, bytes, purpose)?;
+    state_score.copy_from(0, &state.state_score, 0, bytes, purpose)?;
+    Ok(CompressionFrontier {
+        state_kv,
+        state_score,
+        bytes,
+        rows: state.rows,
+    })
+}
+
+fn restore_compression_frontier(
+    state: &mut CompressionState,
+    saved: &CompressionFrontier,
+    purpose: &str,
+) -> Result<(), String> {
+    state
+        .state_kv
+        .copy_from(0, &saved.state_kv, 0, saved.bytes, purpose)?;
+    state
+        .state_score
+        .copy_from(0, &saved.state_score, 0, saved.bytes, purpose)?;
+    state.rows = saved.rows;
+    Ok(())
+}
+
 impl LayerState {
     fn allocate(model: &Model, index: u32, context: u32, raw_cap: u32) -> Result<Self, String> {
         let shape = model.shape;
@@ -2671,7 +2719,7 @@ impl Session {
 // SAFETY: declaration order is required because Rust drops fields in order.
 // `session` must release every Buffer before `_context` calls ds4_gpu_cleanup(),
 // and `_context` must drop before `model` unmaps memory wrapped without copying
-// by native/metal/ds4_metal.m:10329. This intentionally differs from
+// by `ds4_gpu_cleanup` in native/metal/ds4_metal.m. This intentionally differs from
 // DS4's `ds4.c` consumes this exact field order; do not reorder it.
 #[derive(Clone, Copy, Default)]
 pub(super) struct ExecutionStats {
@@ -2934,31 +2982,6 @@ impl DeepSeekExecutor {
     }
 
     fn snapshot_spec_frontier(&self) -> Result<SpecFrontier, String> {
-        fn snapshot(state: &CompressionState, bytes: u64) -> Result<CompressionFrontier, String> {
-            let state_kv = Buffer::bytes(bytes)?;
-            let state_score = Buffer::bytes(bytes)?;
-            state_kv.copy_from(
-                0,
-                &state.state_kv,
-                0,
-                bytes,
-                "saving speculative compressor KV state",
-            )?;
-            state_score.copy_from(
-                0,
-                &state.state_score,
-                0,
-                bytes,
-                "saving speculative compressor score state",
-            )?;
-            Ok(CompressionFrontier {
-                state_kv,
-                state_score,
-                bytes,
-                rows: state.rows,
-            })
-        }
-
         let shape = self.model.shape;
         let commands = Commands::begin()?;
         let layers = self
@@ -2971,9 +2994,10 @@ impl DeepSeekExecutor {
                     .as_ref()
                     .map(|state| {
                         let coefficient = if state.ratio == 4 { 2 } else { 1 };
-                        snapshot(
+                        capture_compression_frontier(
                             state,
                             coefficient * coefficient * state.ratio as u64 * shape.head_dim * 4,
+                            "saving speculative compressor state",
                         )
                     })
                     .transpose()?;
@@ -2981,7 +3005,11 @@ impl DeepSeekExecutor {
                     .indexer
                     .as_ref()
                     .map(|state| {
-                        snapshot(state, 4 * state.ratio as u64 * shape.indexer_head_dim * 4)
+                        capture_compression_frontier(
+                            state,
+                            4 * state.ratio as u64 * shape.indexer_head_dim * 4,
+                            "saving speculative indexer state",
+                        )
                     })
                     .transpose()?;
                 Ok::<_, String>(LayerFrontier {
@@ -3020,40 +3048,26 @@ impl DeepSeekExecutor {
     }
 
     fn restore_spec_frontier(&mut self, frontier: &SpecFrontier) -> Result<(), String> {
-        fn restore(
-            state: &mut CompressionState,
-            saved: &CompressionFrontier,
-        ) -> Result<(), String> {
-            state.state_kv.copy_from(
-                0,
-                &saved.state_kv,
-                0,
-                saved.bytes,
-                "restoring speculative compressor KV state",
-            )?;
-            state.state_score.copy_from(
-                0,
-                &saved.state_score,
-                0,
-                saved.bytes,
-                "restoring speculative compressor score state",
-            )?;
-            state.rows = saved.rows;
-            Ok(())
-        }
-
         if frontier.layers.len() != self.session.layers.len() {
             return Err("speculative frontier layer count changed".into());
         }
         let commands = Commands::begin()?;
         for (layer, saved) in self.session.layers.iter_mut().zip(&frontier.layers) {
             match (&mut layer.compression, &saved.compression) {
-                (Some(state), Some(saved)) => restore(state, saved)?,
+                (Some(state), Some(saved)) => restore_compression_frontier(
+                    state,
+                    saved,
+                    "restoring speculative compressor state",
+                )?,
                 (None, None) => {}
                 _ => return Err("speculative compressor layout changed".into()),
             }
             match (&mut layer.indexer, &saved.indexer) {
-                (Some(state), Some(saved)) => restore(state, saved)?,
+                (Some(state), Some(saved)) => restore_compression_frontier(
+                    state,
+                    saved,
+                    "restoring speculative indexer state",
+                )?,
                 (None, None) => {}
                 _ => return Err("speculative indexer layout changed".into()),
             }
@@ -3078,6 +3092,63 @@ impl DeepSeekExecutor {
         Ok(())
     }
 
+    fn commit_spec_prefix(
+        &mut self,
+        baseline: &SpecFrontier,
+        prefix: &SpecPrefixFrontier,
+        proposals: &[i32],
+        logits: &[f32],
+    ) -> Result<(), String> {
+        let count =
+            u32::try_from(proposals.len()).map_err(|_| "speculative prefix is too large")?;
+        if count == 0 || prefix.layers.len() != self.session.layers.len() {
+            return Err("invalid speculative prefix frontier".into());
+        }
+        let commands = Commands::begin()?;
+        for (layer, saved) in self.session.layers.iter_mut().zip(&prefix.layers) {
+            match (&mut layer.compression, &saved.compression) {
+                (Some(state), Some(saved)) => restore_compression_frontier(
+                    state,
+                    saved,
+                    "committing speculative compressor prefix",
+                )?,
+                (None, None) => {}
+                _ => return Err("speculative compressor prefix layout changed".into()),
+            }
+            match (&mut layer.indexer, &saved.indexer) {
+                (Some(state), Some(saved)) => restore_compression_frontier(
+                    state,
+                    saved,
+                    "committing speculative indexer prefix",
+                )?,
+                (None, None) => {}
+                _ => return Err("speculative indexer prefix layout changed".into()),
+            }
+        }
+        if let Some(dspark) = &mut self.dspark {
+            let row = u64::from(count - 1);
+            for slot in 0..dspark.config.target_layers.len() as u64 {
+                dspark.target_hidden.copy_from(
+                    slot * self.model.shape.embd * 4,
+                    &dspark.target_hidden_batch,
+                    (slot * u64::from(self.session.prefill_cap) + row) * self.model.shape.embd * 4,
+                    self.model.shape.embd * 4,
+                    "committing speculative DSpark target prefix",
+                )?;
+            }
+            dspark.capture_mask = (1_u32 << dspark.config.target_layers.len()) - 1;
+            dspark.cache_start = baseline.dspark_cache_start;
+            dspark.cache_len = baseline.dspark_cache_len;
+            dspark.commit_proposed_prefix(count, self.session.raw_cap);
+        }
+        commands.finish()?;
+        self.session.position = baseline.position + count;
+        self.tokens.truncate(baseline.token_len);
+        self.tokens.extend_from_slice(proposals);
+        self.logits.clone_from_slice(logits);
+        Ok(())
+    }
+
     fn verify_target_suffix(
         &mut self,
         proposals: &[i32],
@@ -3115,10 +3186,10 @@ impl DeepSeekExecutor {
         }
 
         let frontier = self.snapshot_spec_frontier()?;
-        let row_tops = match self.eval_batch_tops(proposals) {
+        let verification = match self.eval_batch_tops(proposals) {
-            Ok(tops) => {
+            Ok(verification) => {
                 self.verifier_passes += 1;
-                tops
+                verification
             }
             Err(error) => {
                 self.restore_spec_frontier(&frontier)?;
@@ -3129,7 +3200,7 @@ impl DeepSeekExecutor {
             }
         };
         let mut commit = 1_usize;
-        while commit < proposals.len() && row_tops[commit - 1] == proposals[commit] {
+        while commit < proposals.len() && verification.tops[commit - 1] == proposals[commit] {
             commit += 1;
         }
         if commit == proposals.len() {
@@ -3139,6 +3210,17 @@ impl DeepSeekExecutor {
             return Ok(proposals.to_vec());
         }
 
+        if let (Some(prefix), Some(logits)) = (
+            verification.prefixes.get(commit - 1),
+            verification.logits.get(commit - 1),
+        ) {
+            self.commit_spec_prefix(&frontier, prefix, &proposals[..commit], logits)?;
+            self.verifier_ns = self
+                .verifier_ns
+                .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
+            return Ok(proposals[..commit].to_vec());
+        }
+
         self.restore_spec_frontier(&frontier)?;
         if let Some(dspark) = &mut self.dspark {
             dspark.commit_proposed_prefix(1, self.session.raw_cap);
@@ -3157,6 +3239,133 @@ impl DeepSeekExecutor {
         Ok(proposals[..commit].to_vec())
     }
 
+    #[allow(clippy::too_many_arguments)]
+    fn verify_target_suffix_stochastic(
+        &mut self,
+        proposals: &[i32],
+        temperature: f32,
+        top_p: f32,
+        min_p: f32,
+        top_k: i32,
+        rng: &mut Rng,
+        cancelled: &std::sync::atomic::AtomicBool,
+    ) -> Result<(Vec<i32>, usize), String> {
+        if proposals.is_empty() || cancelled.load(std::sync::atomic::Ordering::Relaxed) {
+            return Ok((Vec::new(), 0));
+        }
+        let started = Instant::now();
+        if self.quality
+            || proposals.len() == 1
+            || self
+                .ssd
+                .as_ref()
+                .is_some_and(|ssd| u64::from(ssd.cache_experts) < self.model.shape.experts)
+        {
+            let mut emitted = Vec::new();
+            let mut accepted = 0;
+            for &proposal in proposals {
+                let (token, was_draft) = exact_delta_sample(
+                    &self.logits,
+                    proposal,
+                    temperature,
+                    top_p,
+                    min_p,
+                    top_k,
+                    rng,
+                );
+                self.eval_target(token)?;
+                self.verifier_passes += 1;
+                emitted.push(token);
+                if !was_draft {
+                    break;
+                }
+                accepted += 1;
+                if cancelled.load(std::sync::atomic::Ordering::Relaxed) {
+                    break;
+                }
+            }
+            self.verifier_ns = self
+                .verifier_ns
+                .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
+            return Ok((emitted, accepted));
+        }
+        let (first, accepted_first) = exact_delta_sample(
+            &self.logits,
+            proposals[0],
+            temperature,
+            top_p,
+            min_p,
+            top_k,
+            rng,
+        );
+        if !accepted_first {
+            if let Some(dspark) = &mut self.dspark {
+                dspark.commit_proposed_prefix(1, self.session.raw_cap);
+            }
+            self.eval_target(first)?;
+            self.verifier_passes += 1;
+            self.verifier_ns = self
+                .verifier_ns
+                .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
+            return Ok((vec![first], 0));
+        }
+
+        let frontier = self.snapshot_spec_frontier()?;
+        let verification = match self.eval_batch_tops(proposals) {
+            Ok(verification) => {
+                self.verifier_passes += 1;
+                verification
+            }
+            Err(error) => {
+                self.restore_spec_frontier(&frontier)?;
+                return Err(error);
+            }
+        };
+        let mut accepted = 1;
+        let mut replacement = None;
+        for (index, &proposal) in proposals.iter().enumerate().skip(1) {
+            let (token, was_draft) = exact_delta_sample(
+                &verification.logits[index - 1],
+                proposal,
+                temperature,
+                top_p,
+                min_p,
+                top_k,
+                rng,
+            );
+            if !was_draft {
+                replacement = Some(token);
+                break;
+            }
+            accepted += 1;
+        }
+        if replacement.is_none() {
+            self.verifier_ns = self
+                .verifier_ns
+                .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
+            return Ok((proposals.to_vec(), accepted));
+        }
+
+        let prefix = verification
+            .prefixes
+            .get(accepted - 1)
+            .ok_or("missing stochastic verifier prefix")?;
+        let logits = verification
+            .logits
+            .get(accepted - 1)
+            .ok_or("missing stochastic verifier logits")?;
+        self.commit_spec_prefix(&frontier, prefix, &proposals[..accepted], logits)?;
+        let replacement = replacement.expect("replacement disappeared");
+        self.eval_target(replacement)?;
+        self.verifier_passes += 1;
+        let mut emitted = proposals[..accepted].to_vec();
+        emitted.push(replacement);
+        self.verifier_ns = self
+            .verifier_ns
+            .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
+        Ok((emitted, accepted))
+    }
+
     pub(super) fn eval_speculative_greedy(
         &mut self,
         first_token: i32,
@@ -3269,6 +3478,81 @@ impl DeepSeekExecutor {
         Ok(accepted)
     }
 
+    #[allow(clippy::too_many_arguments)]
+    fn eval_speculative_sampled(
+        &mut self,
+        first_token: i32,
+        max_tokens: u32,
+        reasoning: ReasoningMode,
+        temperature: f32,
+        top_p: f32,
+        min_p: f32,
+        top_k: i32,
+        rng: &mut Rng,
+        cancelled: &std::sync::atomic::AtomicBool,
+    ) -> Result<Vec<i32>, String> {
+        if self.dspark.is_none() {
+            self.eval_target(first_token)?;
+            return Ok(vec![first_token]);
+        }
+        if !self.speculative.dspark_exact_sampling {
+            return self.eval_speculative_greedy(first_token, max_tokens, reasoning, cancelled);
+        }
+        self.speculative_cycles += 1;
+        self.eval_target(first_token)?;
+        let mut emitted = vec![first_token];
+        if self.dspark.as_ref().is_some_and(|dspark| dspark.strict)
+            || max_tokens <= 1
+            || cancelled.load(std::sync::atomic::Ordering::Relaxed)
+        {
+            return Ok(emitted);
+        }
+        if self.ssd.is_some() {
+            install_speculative_model_maps(&self.model, "DSpark support mapping")?;
+        }
+        let mut dspark = self.dspark.take().expect("DSpark disappeared");
+        let proposals = dspark.propose(
+            &self.model,
+            &self.weights,
+            first_token,
+            self.session.position.saturating_sub(1),
+            self.session.raw_cap,
+        );
+        self.dspark = Some(dspark);
+        let mut proposals = proposals?;
+        proposals.truncate(
+            max_tokens
+                .saturating_sub(1)
+                .min(self.session.context.saturating_sub(self.session.position))
+                as usize,
+        );
+        if let Some(stop) = proposals
+            .iter()
+            .position(|token| self.model.is_stop_token_for_reasoning(*token, reasoning))
+        {
+            proposals.truncate(stop + 1);
+        }
+        if proposals.len() < 2 {
+            self.dspark
+                .as_mut()
+                .expect("DSpark disappeared")
+                .commit_proposed_prefix(1, self.session.raw_cap);
+            return Ok(emitted);
+        }
+        let (verified, accepted) = self.verify_target_suffix_stochastic(
+            &proposals,
+            temperature,
+            top_p,
+            min_p,
+            top_k,
+            rng,
+            cancelled,
+        )?;
+        emitted.extend_from_slice(&verified);
+        self.dspark.as_mut().expect("DSpark disappeared").accepted += accepted as u64;
+        Ok(emitted)
+    }
+
     fn legacy_mtp_draft(&mut self, token: i32, target_hc: bool) -> Result<(i32, f32), String> {
         let support = self
             .model
@@ -3511,11 +3795,15 @@ impl DeepSeekExecutor {
         self.eval_batch_inner(tokens, false).map(|_| ())
     }
 
-    fn eval_batch_tops(&mut self, tokens: &[i32]) -> Result<Vec<i32>, String> {
+    fn eval_batch_tops(&mut self, tokens: &[i32]) -> Result<BatchVerification, String> {
         self.eval_batch_inner(tokens, true)
     }
 
-    fn eval_batch_inner(&mut self, tokens: &[i32], collect_tops: bool) -> Result<Vec<i32>, String> {
+    fn eval_batch_inner(
+        &mut self,
+        tokens: &[i32],
+        collect_tops: bool,
+    ) -> Result<BatchVerification, String> {
         let rows = u32::try_from(tokens.len()).map_err(|_| "prefill batch is too large")?;
         if rows == 0 || rows > self.session.prefill_cap {
             return Err("prefill batch exceeds the configured prefill workspace".into());
@@ -3535,17 +3823,16 @@ impl DeepSeekExecutor {
         let size = self.model.main.len();
         let shape = self.model.shape;
         let pos = self.session.position;
-        let batch_selected_addr = self.ssd.is_some()
+        let mut prefixes = (0..if collect_tops { rows } else { 0 })
-            && self.weights.layers.first().is_some_and(|layer| unsafe {
+            .map(|_| SpecPrefixFrontier {
-                ds4_gpu_stream_prefill_batch_selected_addr_enabled(
+                layers: (0..shape.layers)
-                    rows,
+                    .map(|_| LayerFrontier {
-                    shape.experts as u32,
+                        compression: None,
-                    shape.experts_used as u32,
+                        indexer: None,
-                    layer.expert_gate.kind,
+                    })
-                    layer.expert_down.kind,
+                    .collect(),
-                ) != 0
+            })
-            });
+            .collect::<Vec<_>>();
-
         if self.ssd.is_some() {
             install_deepseek_model_spans(
                 &self.model,
@@ -3580,6 +3867,16 @@ impl DeepSeekExecutor {
             .enumerate()
         {
             let started = Instant::now();
+            let layer_selected_addr = self.ssd.is_some()
+                && unsafe {
+                    ds4_gpu_stream_prefill_batch_selected_addr_enabled(
+                        rows,
+                        shape.experts as u32,
+                        shape.experts_used as u32,
+                        weights.expert_gate.kind,
+                        weights.expert_down.kind,
+                    ) != 0
+                };
             if let Some(ssd) = &self.ssd {
                 install_deepseek_model_spans(
                     &self.model,
@@ -3587,7 +3884,7 @@ impl DeepSeekExecutor {
                         &self.model,
                         weights,
                         index as u32,
-                        batch_selected_addr,
+                        layer_selected_addr,
                         ssd.per_expert_bytes,
                     )?,
                     "DeepSeek prefill layer mapping",
@@ -3606,6 +3903,7 @@ impl DeepSeekExecutor {
                 rows,
                 self.session.raw_cap,
                 self.steering.as_ref(),
+                collect_tops.then_some(prefixes.as_mut_slice()),
             )?;
             if let Some(profile) = &mut self.profile {
                 profile.record(
@@ -3659,6 +3957,7 @@ impl DeepSeekExecutor {
         let output_rows = if collect_tops { rows } else { 1 };
         let first_output = rows - output_rows;
         let mut tops = Vec::with_capacity(output_rows as usize);
+        let mut output_logits = Vec::with_capacity(output_rows as usize);
         for row in first_output..rows {
             let commands = Commands::begin()?;
             self.session.scratch.current_hc.copy_from(
@@ -3672,13 +3971,20 @@ impl DeepSeekExecutor {
             commands.finish()?;
             self.session.scratch.logits.read_f32(&mut self.logits)?;
             tops.push(argmax(&self.logits));
+            if collect_tops {
+                output_logits.push(self.logits.clone());
+            }
         }
         self.session.position += rows;
         self.tokens.extend_from_slice(tokens);
         if let Some(profile) = &self.profile {
             profile.write()?;
         }
-        Ok(tops)
+        Ok(BatchVerification {
+            tops,
+            logits: output_logits,
+            prefixes,
+        })
     }
 
     pub(super) fn logits(&self) -> &[f32] {
@@ -4003,6 +4309,7 @@ impl Executor {
                 dspark_confidence_threshold: 0.9,
                 dspark_confidence_threshold_set: false,
                 dspark_strict: false,
+                dspark_exact_sampling: false,
             },
             crate::settings::EngineSsdSettings {
                 enabled: false,
@@ -4061,6 +4368,38 @@ impl Executor {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub(super) fn eval_speculative_sampled(
+        &mut self,
+        token: i32,
+        max_tokens: u32,
+        reasoning: ReasoningMode,
+        temperature: f32,
+        top_p: f32,
+        min_p: f32,
+        top_k: i32,
+        rng: &mut Rng,
+        cancelled: &std::sync::atomic::AtomicBool,
+    ) -> Result<Vec<i32>, String> {
+        match self {
+            Self::DeepSeek(executor) => executor.eval_speculative_sampled(
+                token,
+                max_tokens,
+                reasoning,
+                temperature,
+                top_p,
+                min_p,
+                top_k,
+                rng,
+                cancelled,
+            ),
+            Self::Glm(executor) => {
+                executor.eval(token)?;
+                Ok(vec![token])
+            }
+        }
+    }
+
     pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
         match self {
             Self::DeepSeek(executor) => executor.eval(token),
@@ -4272,10 +4611,13 @@ fn compress_attention_batch(
     freq_scale: f32,
     ext: f32,
     attn_factor: f32,
+    mut prefixes: Option<&mut [SpecPrefixFrontier]>,
+    layer: usize,
 ) -> Result<u32, String> {
     let ratio = state.ratio;
     let chunk = rows / ratio;
-    if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) {
+    if prefixes.is_none() && (pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)))
+    {
         let before = if pos == 0 { 0 } else { state.rows };
         let target = s
             .compressed_stage
@@ -4439,6 +4781,15 @@ fn compress_attention_batch(
                 )?;
                 state.rows += 1;
             }
+            if let Some(prefixes) = prefixes.as_deref_mut() {
+                let coefficient = if ratio == 4 { 2 } else { 1 };
+                prefixes[row as usize].layers[layer].compression =
+                    Some(capture_compression_frontier(
+                        state,
+                        coefficient * coefficient * u64::from(ratio) * shape.head_dim * 4,
+                        "capturing speculative compressor prefix",
+                    )?);
+            }
         }
     }
     Ok(state.rows)
@@ -4459,9 +4810,12 @@ fn compress_index_batch(
     freq_scale: f32,
     ext: f32,
     attn_factor: f32,
+    mut prefixes: Option<&mut [SpecPrefixFrontier]>,
+    layer: usize,
 ) -> Result<(), String> {
     let ratio = state.ratio;
-    if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) {
+    if prefixes.is_none() && (pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)))
+    {
         let before = if pos == 0 { 0 } else { state.rows };
         let chunk = rows / ratio;
         let target = state.cache.view(
@@ -4613,6 +4967,13 @@ fn compress_index_batch(
                 )?;
                 state.rows += 1;
             }
+            if let Some(prefixes) = prefixes.as_deref_mut() {
+                prefixes[row as usize].layers[layer].indexer = Some(capture_compression_frontier(
+                    state,
+                    4 * u64::from(ratio) * shape.indexer_head_dim * 4,
+                    "capturing speculative indexer prefix",
+                )?);
+            }
         }
     }
     Ok(())
@@ -4631,6 +4992,7 @@ fn encode_batch_layer(
     rows: u32,
     raw_cap: u32,
     steering: Option<&Steering>,
+    mut prefixes: Option<&mut [SpecPrefixFrontier]>,
 ) -> Result<(), String> {
     let hc_dim = shape.hc * shape.embd;
     let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
@@ -4897,6 +5259,8 @@ fn encode_batch_layer(
             freq_scale,
             ext,
             attn_factor,
+            prefixes.as_deref_mut(),
+            layer as usize,
         )?;
     }
 
@@ -4943,6 +5307,8 @@ fn encode_batch_layer(
             freq_scale,
             ext,
             attn_factor,
+            prefixes,
+            layer as usize,
         )?;
         matmul_rows(
             &s.indexer_q,
@@ -6558,21 +6924,6 @@ fn encode_output(
         },
         "output HC weights",
     )?;
-    let fused_sum_norm = unsafe {
-        ds4_gpu_hc_weighted_sum_norm_tensor(
-            s.output_embedding.raw(),
-            s.output_norm.raw(),
-            s.current_hc.raw(),
-            s.output_weights.raw(),
-            map,
-            size,
-            w.output_norm.offset,
-            shape.embd as u32,
-            shape.hc as u32,
-            shape.rms_epsilon,
-        )
-    } != 0;
-    if !fused_sum_norm {
     call(
         unsafe {
             ds4_gpu_hc_weighted_sum_tensor(
@@ -6599,7 +6950,6 @@ fn encode_output(
         },
         "output norm",
     )?;
-    }
     q8(
         &s.logits,
         w.output,
@@ -7186,6 +7536,7 @@ mod tests {
                 dspark_confidence_threshold: 0.9,
                 dspark_confidence_threshold_set: false,
                 dspark_strict: false,
+                dspark_exact_sampling: false,
             },
             EngineSsdSettings {
                 enabled: false,
@@ -7323,6 +7674,7 @@ mod tests {
                 dspark_confidence_threshold: 0.9,
                 dspark_confidence_threshold_set: false,
                 dspark_strict: false,
+                dspark_exact_sampling: false,
             },
             EngineSsdSettings {
                 enabled: false,
@@ -7414,6 +7766,94 @@ mod tests {
         );
     }
 
+    #[test]
+    #[ignore = "requires the installed 0731 target and checkpoint-specific DSpark GGUF fixtures"]
+    fn flash_0731_runs_exact_sampled_dspark() {
+        use super::{DeepSeekExecutor, argmax, configure_sources};
+        use crate::engine::gguf::Gguf;
+        use crate::engine::validation::validate_support;
+        use crate::engine::{Model, Rng};
+        use crate::model::{ModelChoice, validate_engine_artifacts};
+        use crate::settings::{
+            EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
+        };
+        use std::sync::atomic::AtomicBool;
+
+        configure_sources().unwrap();
+        let artifacts = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true);
+        validate_engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true, &artifacts)
+            .unwrap();
+        let mut model =
+            Model::open_main(&artifacts.model, ModelChoice::DeepSeekV4Flash0731).unwrap();
+        let support = Gguf::open(artifacts.mtp.as_ref().unwrap()).unwrap();
+        model.support_kind = Some(validate_support(&support, &model.shape).unwrap());
+        model.support = Some(support);
+        let prompt = model.render_conversation(
+            "",
+            &[crate::engine::ChatTurn {
+                user: true,
+                tool: false,
+                system: false,
+                skip_previous_eos: false,
+                reasoning: None,
+                reasoning_complete: true,
+                content: "hi".into(),
+            }],
+            ReasoningMode::Direct,
+        );
+        let mut executor = DeepSeekExecutor::open(
+            model,
+            64,
+            false,
+            64,
+            100,
+            EngineSpeculativeSettings {
+                mtp_draft_tokens: 1,
+                mtp_margin: 3.0,
+                glm_mtp: false,
+                glm_mtp_timing: false,
+                dspark: true,
+                dspark_confidence_threshold: 0.6,
+                dspark_confidence_threshold_set: false,
+                dspark_strict: false,
+                dspark_exact_sampling: true,
+            },
+            EngineSsdSettings {
+                enabled: false,
+                cold: false,
+                cache_experts: 0,
+                cache_bytes: 0,
+                full_layers: 0,
+                full_layers_set: false,
+                preload_experts: 0,
+            },
+            EngineSteeringSettings {
+                file: None,
+                ffn_scale: 0.0,
+                attention_scale: 0.0,
+            },
+        )
+        .unwrap();
+        executor.prefill(&prompt, |_| true).unwrap();
+        let first = argmax(executor.logits());
+        let cycle = executor
+            .eval_speculative_sampled(
+                first,
+                4,
+                ReasoningMode::Direct,
+                0.8,
+                0.95,
+                0.0,
+                0,
+                &mut Rng::new(7),
+                &AtomicBool::new(false),
+            )
+            .unwrap();
+        assert!(!cycle.is_empty());
+        assert!(executor.logits().iter().all(|logit| logit.is_finite()));
+        assert!(executor.session.position >= prompt.len() as u32 + cycle.len() as u32);
+    }
+
     #[test]
     #[ignore = "requires the installed Flash, legacy MTP, and DSpark GGUF fixtures"]
     fn ssd_streaming_supports_legacy_mtp_and_dspark() {
@@ -7468,6 +7908,7 @@ mod tests {
                     dspark_confidence_threshold: 0.9,
                     dspark_confidence_threshold_set: false,
                     dspark_strict: false,
+                    dspark_exact_sampling: false,
                 },
                 EngineSsdSettings {
                     enabled: true,
@@ -7544,6 +7985,7 @@ mod tests {
                 dspark_confidence_threshold: 0.9,
                 dspark_confidence_threshold_set: false,
                 dspark_strict: false,
+                dspark_exact_sampling: false,
             };
             let mut executor = DeepSeekExecutor::open(
                 model,
@@ -7627,6 +8069,7 @@ mod tests {
                 dspark_confidence_threshold: 0.9,
                 dspark_confidence_threshold_set: false,
                 dspark_strict: false,
+                dspark_exact_sampling: false,
             },
             EngineSsdSettings {
                 enabled: true,
@@ -7691,6 +8134,7 @@ mod tests {
                 dspark_confidence_threshold: 0.9,
                 dspark_confidence_threshold_set: false,
                 dspark_strict: false,
+                dspark_exact_sampling: false,
             },
             EngineSsdSettings {
                 enabled: false,
@@ -7775,6 +8219,7 @@ mod tests {
                     dspark_confidence_threshold: 0.9,
                     dspark_confidence_threshold_set: false,
                     dspark_strict: false,
+                    dspark_exact_sampling: false,
                 },
                 EngineSsdSettings {
                     enabled: false,
@@ -7850,6 +8295,7 @@ mod tests {
                     dspark_confidence_threshold: 0.9,
                     dspark_confidence_threshold_set: false,
                     dspark_strict: false,
+                    dspark_exact_sampling: false,
                 },
                 EngineSsdSettings {
                     enabled: true,

src/engine/metal/glm.rs

@@ -401,6 +401,7 @@ impl GlmExecutor {
                 dspark_confidence_threshold: 0.9,
                 dspark_confidence_threshold_set: false,
                 dspark_strict: false,
+                dspark_exact_sampling: false,
             },
             None,
         )
@@ -3412,6 +3413,7 @@ mod tests {
                     dspark_confidence_threshold: 0.9,
                     dspark_confidence_threshold_set: false,
                     dspark_strict: false,
+                    dspark_exact_sampling: false,
                 },
                 None,
             )

src/engine/metal/gpu.rs

@@ -1283,18 +1283,6 @@ unsafe extern "C" {
         hc: u32,
         eps: f32,
     ) -> i32;
-    pub(super) fn ds4_gpu_hc_weighted_sum_norm_tensor(
-        out: *mut GpuTensor,
-        norm: *mut GpuTensor,
-        residual: *const GpuTensor,
-        weights: *const GpuTensor,
-        map: *const c_void,
-        size: u64,
-        norm_weight: u64,
-        embd: u32,
-        hc: u32,
-        eps: f32,
-    ) -> i32;
     pub(super) fn ds4_gpu_hc_weighted_sum_tensor(
         out: *mut GpuTensor,
         residual: *const GpuTensor,
@@ -1443,9 +1431,9 @@ impl Buffer {
     }
 
     pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
-        // SAFETY: native/metal/ds4_metal.m:7916-7940 bounds-checks the view,
+        // SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks
-        // ARC-retains base_obj.buffer at :7926, and marks the view non-owning at
+        // the view, ARC-retains base_obj.buffer, and marks the view non-owning, so
-        // :7929, so it may outlive and be freed independently of this wrapper.
+        // it may outlive and be freed independently of this wrapper.
         // Recheck those guarantees whenever the vendored Metal file is re-synced.
         NonNull::new(unsafe { ds4_gpu_tensor_view(self.raw(), offset, bytes) })
             .map(Self)

src/engine/validation.rs

@@ -7,7 +7,14 @@ pub(crate) fn validate_model_artifact(
 ) -> Result<(), String> {
     if support {
         let model = Gguf::open(path)?;
-        validate_support(&model, &FLASH).map(|_| ())
+        let shape = match expected {
+            ModelChoice::DeepSeekV4Flash => FLASH,
+            ModelChoice::DeepSeekV4Flash0731 => FLASH_0731,
+            ModelChoice::DeepSeekV4Pro | ModelChoice::Glm52 => {
+                return Err(format!("{expected} does not use an external support GGUF"));
+            }
+        };
+        validate_support(&model, &shape).map(|_| ())
     } else {
         let model = Model::open_main(path, expected)?;
         let summary = model.summary();
@@ -759,7 +766,10 @@ fn validate_glm_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
 }
 
 pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String> {
-    if shape.model != ModelChoice::DeepSeekV4Flash {
+    if !matches!(
+        shape.model,
+        ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Flash0731
+    ) {
         return Err("DSpark support is available only for DeepSeek V4 Flash".into());
     }
     let DsparkConfig {

src/model.rs

@@ -10,22 +10,22 @@ use std::fs;
 use std::path::{Path, PathBuf};
 
 pub(crate) const MODEL_CHOICES: [ModelChoice; 4] = [
-    ModelChoice::DeepSeekV4Flash,
     ModelChoice::DeepSeekV4Flash0731,
+    ModelChoice::DeepSeekV4Flash,
     ModelChoice::DeepSeekV4Pro,
     ModelChoice::Glm52,
 ];
-pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 6] = [
+pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [
+    ManagedArtifactId::DeepSeekV4Flash0731,
+    ManagedArtifactId::DeepSeekV4Flash0731Dspark,
     ManagedArtifactId::DeepSeekV4Flash,
     ManagedArtifactId::DeepSeekV4FlashMtp,
     ManagedArtifactId::DeepSeekV4FlashDspark,
-    ManagedArtifactId::DeepSeekV4Flash0731,
     ManagedArtifactId::DeepSeekV4Pro,
     ManagedArtifactId::Glm52,
 ];
 
 const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf";
-const DEEPSEEK_FLASH_0731_REPOSITORY: &str = "Rednalreden/DeepSeek-V4-Flash-0731-dwarfstar-q2-gguf";
 const GLM_REPOSITORY: &str = "antirez/glm-5.2-gguf";
 
 const FLASH: Artifact = Artifact {
@@ -54,18 +54,26 @@ const FLASH_MTP: Artifact = Artifact {
 };
 const FLASH_0731: Artifact = Artifact {
     label: "DeepSeek V4 Flash 0731 model",
-    file_name: "DeepSeek-V4-Flash-0731-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-imatrix.gguf",
+    file_name: "DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf",
-    repository: DEEPSEEK_FLASH_0731_REPOSITORY,
+    repository: DEEPSEEK_REPOSITORY,
-    size: 86_720_111_520,
+    size: 86_720_111_488,
-    sha256: "0b39f9c337d6b49c77db2190556b8563abf3c5fbb98be3b58cf8d3a1db191e5f",
+    sha256: "ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0",
     support: Some(false),
 };
+const FLASH_0731_DSPARK: Artifact = Artifact {
+    label: "DeepSeek V4 Flash 0731 DSpark support",
+    file_name: "DeepSeek-V4-Flash-DSpark-support-0731.gguf",
+    repository: DEEPSEEK_REPOSITORY,
+    size: 5_989_114_272,
+    sha256: "7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360",
+    support: Some(true),
+};
 const PRO: Artifact = Artifact {
-    label: "DeepSeek V4 Pro model",
+    label: "DeepSeek V4 Pro 0813 model",
-    file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf",
+    file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf",
     repository: DEEPSEEK_REPOSITORY,
     size: 464_627_334_560,
-    sha256: "a0314d9c0e16122cd60071079124a2d17185d317c55a8f95ecb3ed3506278a96",
+    sha256: "c4d997ab9894b6c78b759f7869fe1726b6314b6515f6ff82607df3797c5eb193",
     support: Some(false),
 };
 const GLM: Artifact = Artifact {
@@ -79,9 +87,9 @@ const GLM: Artifact = Artifact {
 
 #[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
 pub(crate) enum ModelChoice {
-    #[default]
     #[serde(rename = "deepseek-v4-flash")]
     DeepSeekV4Flash,
+    #[default]
     #[serde(rename = "deepseek-v4-flash-0731")]
     DeepSeekV4Flash0731,
     #[serde(rename = "deepseek-v4-pro")]
@@ -104,10 +112,14 @@ impl ModelChoice {
         MODEL_CHOICES.into_iter().find(|model| model.id() == id)
     }
 
-    pub(crate) fn supports_dspark(self) -> bool {
+    pub(crate) fn supports_legacy_mtp(self) -> bool {
         self == Self::DeepSeekV4Flash
     }
 
+    pub(crate) fn supports_dspark(self) -> bool {
+        matches!(self, Self::DeepSeekV4Flash | Self::DeepSeekV4Flash0731)
+    }
+
     fn main_artifact(self) -> &'static Artifact {
         match self {
             Self::DeepSeekV4Flash => &FLASH,
@@ -117,6 +129,14 @@ impl ModelChoice {
         }
     }
 
+    fn dspark_artifact(self) -> Option<&'static Artifact> {
+        match self {
+            Self::DeepSeekV4Flash => Some(&FLASH_DSPARK),
+            Self::DeepSeekV4Flash0731 => Some(&FLASH_0731_DSPARK),
+            Self::DeepSeekV4Pro | Self::Glm52 => None,
+        }
+    }
+
     #[cfg(test)]
     fn artifacts(
         self,
@@ -125,8 +145,8 @@ impl ModelChoice {
     ) -> impl Iterator<Item = &'static Artifact> {
         [
             Some(self.main_artifact()),
-            (self.supports_dspark() && legacy_mtp_enabled).then_some(&FLASH_MTP),
+            (self.supports_legacy_mtp() && legacy_mtp_enabled).then_some(&FLASH_MTP),
-            (self.supports_dspark() && dspark_enabled).then_some(&FLASH_DSPARK),
+            dspark_enabled.then(|| self.dspark_artifact()).flatten(),
         ]
         .into_iter()
         .flatten()
@@ -147,22 +167,61 @@ pub(crate) fn engine_artifacts(
 ) -> EngineArtifacts {
     EngineArtifacts {
         model: model.main_artifact().path(model, models_path),
-        mtp: if model.supports_dspark() && legacy_mtp_enabled {
+        mtp: if model.supports_legacy_mtp() && legacy_mtp_enabled {
             Some(FLASH_MTP.path(model, models_path))
-        } else if model.supports_dspark() && dspark_enabled {
+        } else if dspark_enabled {
-            Some(FLASH_DSPARK.path(model, models_path))
+            model
+                .dspark_artifact()
+                .map(|artifact| artifact.path(model, models_path))
         } else {
             None
         },
     }
 }
 
+pub(crate) fn validate_engine_artifacts(
+    model: ModelChoice,
+    legacy_mtp_enabled: bool,
+    dspark_enabled: bool,
+    artifacts: &EngineArtifacts,
+) -> Result<(), String> {
+    if legacy_mtp_enabled && dspark_enabled {
+        return Err("Legacy MTP and DSpark cannot be enabled together".into());
+    }
+    if legacy_mtp_enabled && !model.supports_legacy_mtp() {
+        return Err(format!("Legacy MTP is not compatible with {model}"));
+    }
+    if dspark_enabled && !model.supports_dspark() {
+        return Err(format!("DSpark is not compatible with {model}"));
+    }
+    model
+        .main_artifact()
+        .validate_installed_path(&artifacts.model)?;
+    let expected_support = if legacy_mtp_enabled {
+        model.supports_legacy_mtp().then_some(&FLASH_MTP)
+    } else if dspark_enabled {
+        model.dspark_artifact()
+    } else {
+        None
+    };
+    match (expected_support, artifacts.mtp.as_deref()) {
+        (Some(expected), Some(path)) => expected.validate_installed_path(path),
+        (None, None) => Ok(()),
+        (Some(_), None) => Err(format!("{model} is missing its required support GGUF")),
+        (None, Some(path)) => Err(format!(
+            "{} is not compatible with the selected {model} checkpoint",
+            path.display()
+        )),
+    }
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub(crate) enum ManagedArtifactId {
     DeepSeekV4Flash,
     DeepSeekV4FlashMtp,
     DeepSeekV4FlashDspark,
     DeepSeekV4Flash0731,
+    DeepSeekV4Flash0731Dspark,
     DeepSeekV4Pro,
     Glm52,
 }
@@ -173,7 +232,9 @@ impl ManagedArtifactId {
             Self::DeepSeekV4Flash | Self::DeepSeekV4FlashMtp | Self::DeepSeekV4FlashDspark => {
                 ModelChoice::DeepSeekV4Flash
             }
-            Self::DeepSeekV4Flash0731 => ModelChoice::DeepSeekV4Flash0731,
+            Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => {
+                ModelChoice::DeepSeekV4Flash0731
+            }
             Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro,
             Self::Glm52 => ModelChoice::Glm52,
         }
@@ -185,6 +246,7 @@ impl ManagedArtifactId {
             Self::DeepSeekV4FlashMtp => &FLASH_MTP,
             Self::DeepSeekV4FlashDspark => &FLASH_DSPARK,
             Self::DeepSeekV4Flash0731 => &FLASH_0731,
+            Self::DeepSeekV4Flash0731Dspark => &FLASH_0731_DSPARK,
             Self::DeepSeekV4Pro => &PRO,
             Self::Glm52 => &GLM,
         }
@@ -278,9 +340,9 @@ pub(crate) enum DownloadOutcome {
 impl fmt::Display for ModelChoice {
     fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
         formatter.write_str(match self {
-            Self::DeepSeekV4Flash => "DeepSeek V4 Flash",
+            Self::DeepSeekV4Flash => "DeepSeek V4 Flash (deprecated preview)",
             Self::DeepSeekV4Flash0731 => "DeepSeek V4 Flash 0731",
-            Self::DeepSeekV4Pro => "DeepSeek V4 Pro",
+            Self::DeepSeekV4Pro => "DeepSeek V4 Pro 0813",
             Self::Glm52 => "GLM 5.2",
         })
     }
@@ -296,6 +358,36 @@ struct Artifact {
 }
 
 impl Artifact {
+    fn validate_installed_path(&self, path: &Path) -> Result<(), String> {
+        if path.file_name().and_then(|name| name.to_str()) != Some(self.file_name) {
+            return Err(format!(
+                "{} is not the expected {} artifact",
+                path.display(),
+                self.label
+            ));
+        }
+        let size = path
+            .metadata()
+            .map_err(|error| format!("Could not inspect {}: {error}", path.display()))?
+            .len();
+        if size != self.size {
+            return Err(format!(
+                "{} has {size} bytes, expected {}",
+                path.display(),
+                self.size
+            ));
+        }
+        let verification = fs::read_to_string(path.with_extension("gguf.sha256"))
+            .map_err(|_| format!("{} has not passed checksum verification", path.display()))?;
+        if verification.trim() != self.sha256 {
+            return Err(format!(
+                "{} has the wrong checkpoint identity",
+                path.display()
+            ));
+        }
+        Ok(())
+    }
+
     fn path(&self, model: ModelChoice, models_path: &Path) -> PathBuf {
         models_path.join(model.id()).join(self.file_name)
     }

src/model/transfer.rs

@@ -337,7 +337,7 @@ mod tests {
         assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
         assert_eq!(
             ModelChoice::DeepSeekV4Flash0731.main_artifact().size,
-            86_720_111_520
+            86_720_111_488
         );
         assert_eq!(
             ModelChoice::DeepSeekV4Flash.artifacts(true, true).count(),
@@ -348,7 +348,7 @@ mod tests {
             ModelChoice::DeepSeekV4Flash0731
                 .artifacts(true, true)
                 .count(),
-            1
+            2
         );
 
         let id = SystemTime::now()
@@ -365,8 +365,14 @@ mod tests {
             engine.mtp.as_deref().and_then(Path::file_name),
             Some(std::ffi::OsStr::new(FLASH_DSPARK.file_name))
         );
+        let flash_0731 =
+            engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true, &models_path);
+        assert_eq!(
+            flash_0731.mtp.as_deref().and_then(Path::file_name),
+            Some(std::ffi::OsStr::new(FLASH_0731_DSPARK.file_name))
+        );
         assert!(
-            engine_artifacts(ModelChoice::DeepSeekV4Flash0731, true, true, &models_path)
+            engine_artifacts(ModelChoice::DeepSeekV4Flash0731, true, false, &models_path)
                 .mtp
                 .is_none()
         );
@@ -389,6 +395,36 @@ mod tests {
                 .unwrap(),
             empty.sha256
         );
+        let installed = empty.path(ModelChoice::DeepSeekV4Flash, &models_path);
+        assert!(empty.validate_installed_path(&installed).is_ok());
+        let wrong_name = installed.with_file_name("wrong-checkpoint.gguf");
+        fs::write(&wrong_name, []).unwrap();
+        fs::write(wrong_name.with_extension("gguf.sha256"), empty.sha256).unwrap();
+        assert!(empty.validate_installed_path(&wrong_name).is_err());
+        assert!(
+            validate_engine_artifacts(
+                ModelChoice::DeepSeekV4Flash0731,
+                true,
+                false,
+                &EngineArtifacts {
+                    model: installed.clone(),
+                    mtp: None,
+                },
+            )
+            .is_err()
+        );
+        assert!(
+            validate_engine_artifacts(
+                ModelChoice::DeepSeekV4Pro,
+                false,
+                true,
+                &EngineArtifacts {
+                    model: installed,
+                    mtp: None,
+                },
+            )
+            .is_err()
+        );
         fs::remove_dir_all(models_path).unwrap();
     }
 

src/settings.rs

@@ -30,6 +30,7 @@ pub(crate) struct SpeculativePreferences {
     pub(crate) dspark_enabled: bool,
     pub(crate) dspark_confidence_threshold: Option<f32>,
     pub(crate) dspark_strict: bool,
+    pub(crate) dspark_exact_sampling: bool,
 }
 
 impl Default for SpeculativePreferences {
@@ -43,6 +44,7 @@ impl Default for SpeculativePreferences {
             dspark_enabled: false,
             dspark_confidence_threshold: None,
             dspark_strict: false,
+            dspark_exact_sampling: false,
         }
     }
 }
@@ -62,13 +64,15 @@ impl SpeculativePreferences {
         if self.dspark_enabled && !model.supports_dspark() {
             return Err("DSpark is not available for the selected model.".into());
         }
-        if self.legacy_mtp_enabled && !model.supports_dspark() {
+        if self.legacy_mtp_enabled && !model.supports_legacy_mtp() {
             return Err("Legacy MTP is not available for the selected model.".into());
         }
         if self.legacy_mtp_enabled && self.dspark_enabled {
             return Err("Legacy MTP and DSpark use different support artifacts.".into());
         }
-        if (self.dspark_confidence_threshold.is_some() || self.dspark_strict)
+        if (self.dspark_confidence_threshold.is_some()
+            || self.dspark_strict
+            || self.dspark_exact_sampling)
             && !self.dspark_enabled
         {
             return Err("DSpark tuning requires DSpark to be enabled.".into());
@@ -86,9 +90,10 @@ impl SpeculativePreferences {
             glm_mtp: self.glm_mtp,
             glm_mtp_timing: self.glm_mtp_timing,
             dspark: self.dspark_enabled,
-            dspark_confidence_threshold: self.dspark_confidence_threshold.unwrap_or(0.9),
+            dspark_confidence_threshold: self.dspark_confidence_threshold.unwrap_or(0.6),
             dspark_confidence_threshold_set: self.dspark_confidence_threshold.is_some(),
             dspark_strict: self.dspark_strict,
+            dspark_exact_sampling: self.dspark_exact_sampling,
         }
     }
 }
@@ -103,6 +108,7 @@ pub(crate) struct EngineSpeculativeSettings {
     pub(crate) dspark_confidence_threshold: f32,
     pub(crate) dspark_confidence_threshold_set: bool,
     pub(crate) dspark_strict: bool,
+    pub(crate) dspark_exact_sampling: bool,
 }
 
 /// An expert count, or a whole GiB budget. Written as `4` or `64GB`, the same
@@ -785,7 +791,7 @@ mod tests {
         let defaults = SpeculativePreferences::default();
         let engine = defaults.engine_settings();
         assert_eq!((engine.mtp_draft_tokens, engine.mtp_margin), (1, 3.0));
-        assert_eq!(engine.dspark_confidence_threshold, 0.9);
+        assert_eq!(engine.dspark_confidence_threshold, 0.6);
         assert!(!engine.dspark_confidence_threshold_set);
 
         let tuned = SpeculativePreferences {
@@ -793,6 +799,7 @@ mod tests {
             dspark_enabled: true,
             dspark_confidence_threshold: Some(0.7),
             dspark_strict: true,
+            dspark_exact_sampling: true,
             ..defaults
         };
         assert!(tuned.validate(ModelChoice::DeepSeekV4Flash).is_ok());
@@ -812,6 +819,7 @@ mod tests {
             ..SpeculativePreferences::default()
         };
         assert!(legacy.validate(ModelChoice::DeepSeekV4Flash).is_ok());
+        assert!(legacy.validate(ModelChoice::DeepSeekV4Flash0731).is_err());
         assert!(legacy.validate(ModelChoice::DeepSeekV4Pro).is_err());
         assert!(
             SpeculativePreferences {
@@ -821,6 +829,14 @@ mod tests {
             .validate(ModelChoice::DeepSeekV4Flash)
             .is_err()
         );
+        assert!(
+            SpeculativePreferences {
+                dspark_exact_sampling: true,
+                ..SpeculativePreferences::default()
+            }
+            .validate(ModelChoice::DeepSeekV4Flash0731)
+            .is_err()
+        );
     }
 
     #[test]