Share DeepSeek V4 Metal index queries across four keys

Adapt the query reuse idea from Tarjei Mandt and contributors in #1073 to the 64-head Flash scorer. Each SIMD group retains its original dot-product reduction. Store the rounded contribution from each head before summing them in their original order.

On M5 Max, V4 Flash 0731 Q2 steady decode improved from 39.34 to 40.27 tokens/s at 16K and 35.31 to 37.80 at 66K. All full-vocabulary logits matched on the same 96-token continuation at both lengths. Keep the one-key schedules for short rows, tensor parallelism and strict math.

ds4_metal.m

@@ -554,6 +554,7 @@ static id<MTLComputePipelineState> g_dsv4_topk_mask_scatter_pipeline;
 static id<MTLComputePipelineState> g_dsv4_indexer_weighted_sum_pipeline;
 static id<MTLComputePipelineState> g_dsv4_indexer_score_one_direct_pipeline;
 static id<MTLComputePipelineState> g_dsv4_indexer_score_one_heads_pipeline;
+static bool g_dsv4_indexer_fast_math;
 static id<MTLComputePipelineState> g_dsv4_compressor_store_one_pipeline;
 static id<MTLComputePipelineState> g_dsv4_sort_i32_rows_asc_pipeline;
 static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_pipeline;
@@ -9022,6 +9023,7 @@ int ds4_gpu_init(void) {
             ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_direct");
         g_dsv4_indexer_score_one_heads_pipeline =
             ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_heads");
+        g_dsv4_indexer_fast_math = !drift_math_safe;
         g_dsv4_compressor_store_one_pipeline =
             ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one");
         g_dsv4_compressor_pack_ratio4_pipeline =
@@ -19157,7 +19159,11 @@ int ds4_gpu_indexer_score_one_tensor(
             /* Single-host decode benefits from reducing all 64 heads at once.
              * The four-head schedule is faster on the two-Mac RDMA path. */
             const bool single = g_tp_split_world == 1;
-            id<MTLComputePipelineState> direct_pipeline = single ?
+            /* The grouped kernel matches the default library's rounding order. */
+            id<MTLComputePipelineState> grouped = single && n_comp >= 2048u && g_dsv4_indexer_fast_math &&
+                !(g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) ?
+                ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_rows4") : nil;
+            id<MTLComputePipelineState> direct_pipeline = grouped ? grouped : single ?
                 ds4_gpu_hot_pipeline(g_dsv4_indexer_score_one_heads_pipeline, "kernel_dsv4_indexer_score_one_heads") :
                 ds4_gpu_hot_pipeline(g_dsv4_indexer_score_one_direct_pipeline, "kernel_dsv4_indexer_score_one_direct");
             if (!direct_pipeline) return 0;
@@ -19187,8 +19193,11 @@ int ds4_gpu_indexer_score_one_tensor(
             [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
             [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3];
             [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
-            [enc setThreadgroupMemoryLength:(single ? 64u : 136u) * sizeof(float) atIndex:0];
-            [enc dispatchThreadgroups:MTLSizeMake(n_comp, 1, 1)
+            if (grouped) DS4_TEST_NOTE("kernel_dsv4_indexer_score_one_rows4");
+            const NSUInteger partials = grouped ? 256u : single ? 64u : 136u;
+            const NSUInteger groups = grouped ? (n_comp - 1u) / 4u + 1u : n_comp;
+            [enc setThreadgroupMemoryLength:partials * sizeof(float) atIndex:0];
+            [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
                  threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
             ds4_gpu_end_compute_encoder(cb, enc);
 

metal/dsv4_misc.metal

@@ -546,6 +546,45 @@ kernel void kernel_dsv4_indexer_score_one_heads(
     }
 }
 
+// Four keys share each query load. Store every head's rounded contribution
+// before adding them in head order, as in the one-key kernel.
+kernel void kernel_dsv4_indexer_score_one_rows4(
+        constant ds4_metal_args_dsv4_indexer_scores_fused &args,
+        device const float4 *q, device const float *weights,
+        device const float4 *keys, device float *scores,
+        threadgroup float *partials [[threadgroup(0)]],
+        uint group [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
+        ushort lane [[thread_index_in_simdgroup]], ushort sg [[simdgroup_index_in_threadgroup]]) {
+    const uint first = group * 4u;
+    const uint last = args.n_comp - 1u;
+    const float4 k0 = keys[(ulong)first * 32u + lane];
+    const float4 k1 = keys[(ulong)min(first + 1u, last) * 32u + lane];
+    const float4 k2 = keys[(ulong)min(first + 2u, last) * 32u + lane];
+    const float4 k3 = keys[(ulong)min(first + 3u, last) * 32u + lane];
+    for (uint h = sg; h < 64u; h += 4u) {
+        const float4 query = q[h * 32u + lane];
+        const float s0 = simd_sum(dot(query, k0));
+        const float s1 = simd_sum(dot(query, k1));
+        const float s2 = simd_sum(dot(query, k2));
+        const float s3 = simd_sum(dot(query, k3));
+        if (lane == 0) {
+            // The fast-math one-key kernel rounds score times weight before scaling.
+#pragma clang fp reassociate(off)
+            const float w = weights[h];
+            partials[h] = max(s0, 0.0f) * w * args.scale;
+            partials[64u + h] = max(s1, 0.0f) * w * args.scale;
+            partials[128u + h] = max(s2, 0.0f) * w * args.scale;
+            partials[192u + h] = max(s3, 0.0f) * w * args.scale;
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (tid < 4u && first + tid <= last) {
+        float acc = 0.0f;
+        for (uint h = 0; h < 64u; h++) acc += partials[64u * tid + h];
+        scores[first + tid] = acc;
+    }
+}
+
 // Decode router post-processing for one token. The selected expert ids are
 // already known; this gathers their probabilities, normalizes by the selected
 // sum, clamps the denominator like the reference path, and applies DS4's 1.5

tests/test_v41_indexer_score.c

@@ -208,12 +208,59 @@ done:
     return failed;
 }
 
+/* The V4 Flash kernel reduces 64 heads, with one shared query load per four
+ * keys. Check every tail length, exact float scores and the output boundary. */
+static int check_flash_score_rows4(void) {
+    enum { HEADS = 64, ROWS = 8193 };
+    const uint32_t widths[] = {2047, 2048, 2049, 2050, 2051, ROWS};
+    ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(HEADS * ROW_BYTES);
+    ds4_gpu_tensor *weights = ds4_gpu_tensor_alloc(HEADS * sizeof(float));
+    ds4_gpu_tensor *keys = ds4_gpu_tensor_alloc(ROWS * ROW_BYTES);
+    ds4_gpu_tensor *one = ds4_gpu_tensor_alloc((ROWS + GUARD_ROWS) * sizeof(float));
+    ds4_gpu_tensor *four = ds4_gpu_tensor_alloc((ROWS + GUARD_ROWS) * sizeof(float));
+    int failed = 1;
+    if (!q || !weights || !keys || !one || !four) goto done;
+    float *query = ds4_gpu_tensor_contents(q), *weight = ds4_gpu_tensor_contents(weights);
+    float *key = ds4_gpu_tensor_contents(keys), *a = ds4_gpu_tensor_contents(one);
+    float *b = ds4_gpu_tensor_contents(four);
+    const float scale = 1.0f / sqrtf((float)(HEADS * HEAD_DIM));
+    for (int fp4 = 0; fp4 <= 1; fp4++) {
+        for (size_t i = 0; i < HEADS * HEAD_DIM; i++)
+            query[i] = fp4 ? rng_fp4() : ldexpf(rng_float(), (int)(i % 23) - 11);
+        for (size_t i = 0; i < HEADS; i++) weight[i] = rng_float();
+        for (size_t i = 0; i < ROWS * HEAD_DIM; i++)
+            key[i] = fp4 ? rng_fp4() : ldexpf(rng_float(), (int)(i % 19) - 9);
+        for (size_t i = 0; i < sizeof(widths) / sizeof(widths[0]); i++) {
+            const uint32_t n = widths[i];
+            for (uint32_t j = n; j < n + GUARD_ROWS; j++) b[j] = -123.0f;
+            ds4_gpu_test_set_flags(DS4_GPU_TEST_INDEX_SCORE_DIRECT);
+            if (!ds4_gpu_indexer_score_one_tensor(one, q, weights, keys, n, HEADS, HEAD_DIM, scale)) goto done;
+            ds4_gpu_test_set_flags(DS4_GPU_TEST_COUNT_KERNELS);
+            const char *name = "kernel_dsv4_indexer_score_one_rows4";
+            const uint64_t before = ds4_gpu_test_kernel_runs(name);
+            if (!ds4_gpu_indexer_score_one_tensor(four, q, weights, keys, n, HEADS, HEAD_DIM, scale) ||
+                memcmp(a, b, n * sizeof(*a))) goto done;
+            if (!getenv("DS4_METAL_MATH_SAFE") &&
+                ds4_gpu_test_kernel_runs(name) != before + (n >= 2048u)) goto done;
+            for (uint32_t j = n; j < n + GUARD_ROWS; j++) if (b[j] != -123.0f) goto done;
+        }
+    }
+    failed = 0;
+done:
+    ds4_gpu_test_set_flags(0);
+    ds4_gpu_tensor_free(q); ds4_gpu_tensor_free(weights); ds4_gpu_tensor_free(keys);
+    ds4_gpu_tensor_free(one); ds4_gpu_tensor_free(four);
+    fprintf(stderr, "V4 Flash four-key score, exact scores and tail guards: %s\n", failed ? "FAIL" : "PASS");
+    return failed;
+}
+
 int main(void) {
     if (!ds4_gpu_init()) {
         fprintf(stderr, "metal init failed\n");
         return 1;
     }
     int fail = check_block_selection();
+    fail |= check_flash_score_rows4();
     const uint32_t tails[] = {1, 2, 3, 4, 15, 16, 17, 31, 33, 1023, 1025};
     for (size_t i = 0; i < sizeof(tails) / sizeof(*tails); i++) {
         fail |= run_case("tail-f32", tails[i], 0, 0, 0);