Score long V4.1 decode index rows from bf16 keys

The indexer's keys are FP4 values times a power-of-two scale, and every kernel that writes them rounds them to bf16, so a bf16 copy holds them exactly in half the space of the f32 key cache. Those kernels now also write a bf16 copy, and once there are 8192 or more index keys, decode scores them from the copy. The scoring math and its order are unchanged, so the results are bit-identical. Restoring a saved session rebuilds the copy, and if any restored key isn't exact in bf16, scoring uses the f32 keys. This is Metal only.

DS4_METAL_DISABLE_V41_INDEX_BF16=1 skips the copy entirely.

ds4.c

@@ -40995,6 +40995,10 @@ typedef struct {
     float rows[2][DS4_ENGRAM_COLS * DS4_ENGRAM_DIM];
     ds4_gpu_tensor *window[40];
     ds4_gpu_tensor *compressed[4], *index_cache[4];
+    /* Metal: index_cache's bf16 view for the decode scorer, valid unless a restore brought keys that are not
+     * bf16-exact (then the f32 cache is scored) */
+    ds4_gpu_tensor *index_bf16[4], *index_bf16_flag;
+    bool index_bf16_ok;
     ds4_gpu_tensor *previous_kv[4], *previous_score[4];
     ds4_gpu_tensor *verify_previous_kv[4], *verify_previous_score[4];
     bool verify_prefix;
@@ -41016,6 +41020,10 @@ static bool ds41_read_array(const ds4_model *m, const char *key, uint32_t type,
     return cursor_read(&c, out, bytes);
 }
 
+#ifdef __APPLE__
+static bool ds41_index_bf16_on(void) { return !(ds4_gpu_v41_off() & DS4_GPU_V41_INDEX_BF16); }
+#endif
+
 static void ds41_graph_free(ds41_gpu_graph *g) {
     if (!g) return;
     ds41_dspark_free(g->dspark);
@@ -41047,9 +41055,11 @@ static void ds41_graph_free(ds41_gpu_graph *g) {
     for (uint32_t i = 0; i < 4; i++) {
         ds4_gpu_tensor_free(g->compressed[i]);
         ds4_gpu_tensor_free(g->index_cache[i]);
+        ds4_gpu_tensor_free(g->index_bf16[i]);
         ds4_gpu_tensor_free(g->previous_kv[i]);
         ds4_gpu_tensor_free(g->previous_score[i]);
     }
+    ds4_gpu_tensor_free(g->index_bf16_flag);
 #define DS41_FREE(name, count) ds4_gpu_tensor_free(g->name);
     DS41_SCRATCH(DS41_FREE)
 #undef DS41_FREE
@@ -41068,8 +41078,13 @@ static uint64_t ds41_graph_bytes(uint32_t ctx) {
         .prefill_alias = !getenv("DS4_METAL_DISABLE_V41_PREFILL_ALIAS"),
         .compact_carry = !getenv("DS4_METAL_DISABLE_V41_COMPACT_CARRY")}, *g = &shape;
     uint64_t floats = (uint64_t)40 * 128 * 512;
+#ifdef __APPLE__
+    const uint32_t view = ds41_index_bf16_on() ? 64u : 0u;   /* the index keys' bf16 view */
+#else
+    const uint32_t view = 0;
+#endif
     for (uint32_t i = 0; i < 4; i++)
-        floats += ((uint64_t)ctx / (i < 3 ? 2u : 1u) + 1u) * (512u + 128u) + 2u * 512u;
+        floats += ((uint64_t)ctx / (i < 3 ? 2u : 1u) + 1u) * (512u + 128u + view) + 2u * 512u;
     floats += (uint64_t)2 * 2 * DS4_N_EMBD * DS4_N_HC;
 #define DS41_COUNT(name, count) floats += (count);
     DS41_SCRATCH(DS41_COUNT)
@@ -41109,6 +41124,7 @@ static ds4_context_memory ds41_graph_memory(uint32_t ctx) {
 static void ds41_graph_reset(ds41_gpu_graph *g) {
     g->pos = 0;
     g->valid = true;
+    g->index_bf16_ok = true;   /* every row is rewritten through its producer */
     ds4_engram_history_reset(&g->history);
     ds41_dspark_reset(g->dspark);
     /* Valid lengths, not zeroed storage, determine cache visibility. Each
@@ -41218,7 +41234,14 @@ static DS4_MAYBE_UNUSED bool ds41_graph_alloc(ds41_gpu_graph *g, const ds4_model
         g->previous_kv[i] = ds4_gpu_tensor_alloc(512u * sizeof(float));
         g->previous_score[i] = ds4_gpu_tensor_alloc(512u * sizeof(float));
         if (!g->compressed[i] || !g->index_cache[i] || !g->previous_kv[i] || !g->previous_score[i]) goto fail;
+#ifdef __APPLE__
+        if (ds41_index_bf16_on() && !(g->index_bf16[i] = ds4_gpu_tensor_alloc(cap * 128u * sizeof(uint16_t))))
+            goto fail;
+#endif
     }
+#ifdef __APPLE__
+    if (ds41_index_bf16_on() && !(g->index_bf16_flag = ds4_gpu_tensor_alloc(sizeof(uint32_t)))) goto fail;
+#endif
 #define DS41_ALLOC(name, count) \
     if (!(g->name = (!strcmp(#name, "engram_prefetch") || !strcmp(#name, "image_text_mask") ? \
             ds4_gpu_tensor_alloc_managed : ds4_gpu_tensor_alloc)( \
@@ -41607,6 +41630,39 @@ static bool ds41_short_sweep_off(void) {
     return metal_graph_env_flag("DS4_METAL_DISABLE_V41_SHORT_SWEEP", &cache);
 }
 
+/* Rows [first, first + rows) of an owner's index keys into its bf16 view.
+ * Producers round every key to bf16, so the view holds the keys exactly. */
+static bool ds41_index_bf16_store(ds41_gpu_graph *g, uint32_t owner, uint32_t first, uint32_t rows) {
+#ifdef __APPLE__
+    return !g->index_bf16[owner] ||
+        ds4_gpu_dsv41_index_keys_bf16(g->index_bf16[owner], g->index_cache[owner], first, rows, NULL);
+#else
+    (void)g; (void)owner; (void)first; (void)rows;
+    return true;
+#endif
+}
+/* A restore brings f32 keys from a save: rebuild the views and keep them
+ * only if every live key is bf16-exact; otherwise decode scores the f32
+ * cache until the next reset. */
+static void ds41_index_bf16_restore(ds41_gpu_graph *g, uint32_t pos) {
+    g->index_bf16_ok = false;
+#ifdef __APPLE__
+    const uint32_t zero = 0;
+    if (!g->index_bf16_flag || !ds4_gpu_tensor_write(g->index_bf16_flag, 0, &zero, sizeof(zero))) return;
+    for (uint32_t i = 0; i < 4; i++) {
+        const uint32_t live = pos / (i < 3 ? 2u : 1u);
+        if (live && !ds4_gpu_dsv41_index_keys_bf16(g->index_bf16[i], g->index_cache[i], 0, live, g->index_bf16_flag))
+            return;
+    }
+    uint32_t flag = 1;
+    if (!ds4_gpu_synchronize() || !ds4_gpu_tensor_read(g->index_bf16_flag, 0, &flag, sizeof(flag))) return;
+    g->index_bf16_ok = flag == 0;
+    if (flag) fprintf(stderr, "ds4: V4.1 restored index keys are not bf16-exact; decode scores the f32 keys\n");
+#else
+    (void)pos;
+#endif
+}
+
 /* One block's HC input: the mixer projection, then split, weighted sum with
  * the previous block's pre weights and norm; each half fused where the fused
  * kernel reproduces the separate ones. */
@@ -41778,6 +41834,7 @@ static bool ds41_attention_publish(ds41_gpu_graph *g, const ds4_model *m,
                 !ds41_norm(g->index_k, g->index_k, m, l->indexer_k_norm) ||
                 !ds41_rope_store(g->index_k, g->index_cache[owner], (uint64_t)(n_comp - 1u) * 128u * 4u,
                                  DS4_N_INDEXER_HEAD_DIM, il, rope_pos, DS4_V41_FP4_E8M0) ||
+                !ds41_index_bf16_store(g, owner, n_comp - 1u, 1) ||
                 !ds41_rope_store(g->latent, g->compressed[owner], (uint64_t)(n_comp - 1u) * 512u * 4u,
                                  512, il, rope_pos, DS4_V41_FP4_E4M3)) return false;
         }
@@ -41854,8 +41911,8 @@ static bool ds41_attention_select_published(ds41_gpu_graph *g, const ds4_model *
             ds4_gpu_dsv41_indexer_candidates_tensor(g->selected_comp, g->index_scores, g->block_list,
                 g->index_q, g->index_weights, g->index_cache[owner], g->block_mask, n_comp,
                 1.0f / 64.0f, DS4_N_INDEXER_TOP_K)) return true;
-        if (!ds4_gpu_dsv41_indexer_score_one_tensor(g->index_scores, g->index_q, g->index_weights,
-                g->index_cache[owner], n_comp, DS4_N_INDEXER_HEAD, 128, 1.0f / 64.0f, false)) return false;
+        if (!ds4_gpu_dsv41_indexer_score_row(g->index_scores, g->index_q, g->index_weights, g->index_cache[owner],
+                                             g->index_bf16_ok ? g->index_bf16[owner] : NULL, n_comp)) return false;
 #else
         if (!ds4_gpu_dsv41_indexer_scores_batch(g->index_scores, g->index_q, g->index_weights,
                 g->index_cache[owner], n_comp, 1, pos, ratio)) return false;
@@ -42689,6 +42746,7 @@ static bool ds41_attention_publish_batch(ds41_gpu_graph *g, ds41_prefill_row *b,
         ds4_gpu_dsv41_quantize(b->index_k, DS4_N_INDEXER_HEAD_DIM, rows, DS4_V41_FP4_E8M0) &&
         ds4_gpu_tensor_copy(g->index_cache[owner], (uint64_t)first * DS4_N_INDEXER_HEAD_DIM * sizeof(float),
             b->index_k, 0, (uint64_t)rows * DS4_N_INDEXER_HEAD_DIM * sizeof(float)) &&
+        ds41_index_bf16_store(g, owner, first, rows) &&
         ds4_gpu_dsv41_rope_stride(b->latent, DS4_N_HEAD_DIM, 1, rows,
                                   first * ratio, ratio, true, false) &&
         ds4_gpu_dsv41_quantize(b->latent, DS4_N_HEAD_DIM, rows, DS4_V41_FP4_E4M3) &&
@@ -64597,6 +64655,7 @@ static int ds41_load_payload(ds4_session *s, FILE *fp, const uint32_t *h,
             for (uint32_t i = 0; i < 3 && i < pos; i++)
                 g->history.tail[i] = (int32_t)g->token_map[tokens.v[pos - 1u - i]];
             g->pos = pos;
+            ds41_index_bf16_restore(g, pos);
             ds4_tokens_copy(&s->checkpoint, &tokens);
             s->checkpoint_valid = true;
         }

ds4_deepseek41_gpu.h

@@ -143,6 +143,13 @@ int ds4_gpu_dsv41_indexer_score_one_tensor(ds4_gpu_tensor *scores,
                                            uint32_t n_rows, uint32_t n_head,
                                            uint32_t head_dim, float scale,
                                            bool cache_f16);
+/* Rows [first, first + rows) of the index-key cache into its bf16 view; with flag, a key that is not bf16-exact
+ * sets it nonzero. score_row: one decode index row (32 heads x 128, scale 1/64) over the cache, reading view, when
+ * given, where the wide scorer runs (the same bits). */
+int ds4_gpu_dsv41_index_keys_bf16(ds4_gpu_tensor *view, const ds4_gpu_tensor *cache,
+        uint32_t first, uint32_t rows, ds4_gpu_tensor *flag);
+int ds4_gpu_dsv41_indexer_score_row(ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights,
+        const ds4_gpu_tensor *cache, const ds4_gpu_tensor *view, uint32_t rows);
 /* The top-k of a row filtered by a block mask of more than 2048 blocks, from
  * the scores of its 2048 candidate blocks only: the same ids in the same order
  * (Metal only). Returns 0, leaving the full row to the caller, where unavailable. */

ds4_gpu.h

@@ -55,6 +55,7 @@ enum {
     DS4_GPU_V41_MODEL_TRACKED = 1u << 3,   /* opt-out: DS4_METAL_MODEL_TRACKED */
     DS4_GPU_V41_ROPE_TABLE = 1u << 4,
     DS4_GPU_V41_SPLIT_REDUCE = 1u << 5,
+    DS4_GPU_V41_INDEX_BF16 = 1u << 6,
     DS4_GPU_V41_PLAIN_READERS = 1u << 29,   /* test only: the narrow readers' own kernels */
 };
 uint32_t ds4_gpu_v41_off(void);

ds4_metal.m

@@ -447,6 +447,7 @@ uint32_t ds4_gpu_v41_off(void) {
             "DS4_METAL_MODEL_TRACKED",
             "DS4_METAL_DISABLE_V41_ROPE_TABLE",
             "DS4_METAL_DISABLE_V41_SPLIT_REDUCE",
+            "DS4_METAL_DISABLE_V41_INDEX_BF16",
         };
         for (uint32_t i = 0; i < sizeof(names) / sizeof(*names); i++) {
             const char *v = getenv(names[i]);
@@ -587,6 +588,8 @@ static id<MTLComputePipelineState> g_glm_indexer_rope_tail_pipeline;
 static id<MTLComputePipelineState> g_glm_indexer_score_one_pipeline;
 static id<MTLComputePipelineState> g_glm_indexer_score_one_direct_pipeline;
 static id<MTLComputePipelineState> g_glm_indexer_score_one_wide_pipeline;
+static id<MTLComputePipelineState> g_dsv41_indexer_score_bf16_pipeline;
+static id<MTLComputePipelineState> g_dsv41_index_keys_bf16_pipeline;
 static id<MTLComputePipelineState> g_glm_indexer_scores_batch_pipeline;
 static id<MTLComputePipelineState> g_glm_indexer_scores_tiled_pipeline;
 static id<MTLComputePipelineState> g_glm_indexer_scores_tiled_f32_pipeline;
@@ -9092,6 +9095,10 @@ int ds4_gpu_init(void) {
             ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_direct");
         g_glm_indexer_score_one_wide_pipeline =
             ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_wide");
+        g_dsv41_indexer_score_bf16_pipeline =
+            ds4_gpu_get_pipeline("kernel_dsv41_indexer_score_bf16");
+        g_dsv41_index_keys_bf16_pipeline =
+            ds4_gpu_get_pipeline("kernel_dsv41_index_keys_bf16");
         g_glm_indexer_scores_batch_pipeline =
             ds4_gpu_get_pipeline("kernel_glm_indexer_scores_batch");
         g_glm_indexer_scores_tiled_pipeline =
@@ -9214,6 +9221,8 @@ int ds4_gpu_init(void) {
             !g_glm_indexer_score_one_pipeline ||
             !g_glm_indexer_score_one_direct_pipeline ||
             !g_glm_indexer_score_one_wide_pipeline ||
+            !g_dsv41_indexer_score_bf16_pipeline ||
+            !g_dsv41_index_keys_bf16_pipeline ||
             !g_glm_indexer_scores_batch_pipeline ||
             !g_glm_indexer_scores_tiled_pipeline ||
             !g_glm_indexer_scores_tiled_f32_pipeline ||
@@ -37238,6 +37247,75 @@ int ds4_gpu_dsv41_indexer_score_one_tensor(
                                               n_rows, n_head, head_dim, scale, cache_f16, true);
 }
 
+int ds4_gpu_dsv41_index_keys_bf16(ds4_gpu_tensor *view, const ds4_gpu_tensor *cache,
+        uint32_t first, uint32_t rows, ds4_gpu_tensor *flag) {
+    if (!view || !cache || !rows) return rows == 0;
+    if (!g_initialized && !ds4_gpu_init()) return 0;
+    const uint64_t end = (uint64_t)first + rows;
+    if (ds4_gpu_tensor_bytes(cache) < end * 128u * 4u || ds4_gpu_tensor_bytes(view) < end * 128u * 2u ||
+        (flag && ds4_gpu_tensor_bytes(flag) < 4u) || end * 32u > UINT32_MAX) return 0;
+    @autoreleasepool {
+        id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
+            g_dsv41_index_keys_bf16_pipeline, "kernel_dsv41_index_keys_bf16");
+        if (!pipeline) return 0;
+        const struct { uint32_t n_vec, check; } args = {rows * 32u, flag != NULL};
+        int owned = 0;
+        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
+        if (!cb) return 0;
+        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
+        [enc setComputePipelineState:pipeline];
+        [enc setBytes:&args length:sizeof(args) atIndex:0];
+        [enc setBuffer:ds4_gpu_tensor_buffer(cache) offset:ds4_gpu_tensor_offset(cache) + (uint64_t)first * 128u * 4u atIndex:1];
+        [enc setBuffer:ds4_gpu_tensor_buffer(view) offset:ds4_gpu_tensor_offset(view) + (uint64_t)first * 128u * 2u atIndex:2];
+        const ds4_gpu_tensor *f = flag ? flag : view;   /* unread without check */
+        [enc setBuffer:ds4_gpu_tensor_buffer(f) offset:ds4_gpu_tensor_offset(f) atIndex:3];
+        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)rows * 32u + 255u) / 256u, 1, 1)
+             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+        ds4_gpu_end_compute_encoder(cb, enc);
+        return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 bf16 index keys");
+    }
+}
+
+int ds4_gpu_dsv41_indexer_score_row(ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights,
+        const ds4_gpu_tensor *cache, const ds4_gpu_tensor *view, uint32_t rows) {
+    /* the view where the wide scorer runs (rows >= 8192, as ds4_gpu_glm_indexer_score_one_impl decides) */
+    const bool wide = (g_test_flags & DS4_GPU_TEST_INDEX_SCORE_WIDE) != 0u ||
+        ((g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) == 0u && rows >= 8192u && !ds4_gpu_v41_index_score_wide_off());
+    if (!view || !wide || !ds4_gpu_v41_on(DS4_GPU_V41_INDEX_BF16) || !scores || !q || !weights ||
+        ds4_gpu_tensor_bytes(view) < (uint64_t)rows * 128u * 2u)
+        return ds4_gpu_dsv41_indexer_score_one_tensor(scores, q, weights, cache, rows, 32, 128, 1.0f / 64.0f, false);
+    if (!g_initialized && !ds4_gpu_init()) return 0;
+    if (!rows || ds4_gpu_tensor_bytes(scores) < (uint64_t)rows * 4u || ds4_gpu_tensor_bytes(q) < 4096u * 4u ||
+        ds4_gpu_tensor_bytes(weights) < 32u * 4u) return 0;
+    @autoreleasepool {
+        /* below 65536 rows, S lanes share each key (four below 10240 rows, else two); above, two keys per thread,
+         * 512 threads below 98304 rows and 128 from there (measured) */
+        const char *name = rows < 10240u ? "kernel_dsv41_indexer_score_rows_split4_bf16" :
+            rows < 65536u ? "kernel_dsv41_indexer_score_rows_split2_bf16" : "kernel_dsv41_indexer_score_bf16";
+        const NSUInteger tg = rows < 65536u ? 256 : rows < 98304u ? 512 : 128;
+        const NSUInteger keys_per_tg = rows < 10240u ? tg / 2u : rows < 65536u ? tg : 2u * tg;
+        id<MTLComputePipelineState> pipeline = rows < 65536u ? ds4_gpu_get_pipeline(name) :
+            ds4_gpu_hot_pipeline(g_dsv41_indexer_score_bf16_pipeline, name);
+        if (!pipeline) return 0;
+        DS4_TEST_NOTE(name);
+        ds4_gpu_glm_indexer_score_one_args args = {rows, 32u, 128u, 0u, 1.0f / 64.0f};
+        int owned = 0;
+        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
+        if (!cb) return 0;
+        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
+        [enc setComputePipelineState:pipeline];
+        [enc setBytes:&args length:sizeof(args) atIndex:0];
+        const ds4_gpu_tensor *t[] = {q, weights, view, scores};
+        for (unsigned i = 0; i < 4; i++)
+            [enc setBuffer:ds4_gpu_tensor_buffer(t[i]) offset:ds4_gpu_tensor_offset(t[i]) atIndex:i + 1];
+        [enc setThreadgroupMemoryLength:(4096u + 32u) * sizeof(float) atIndex:0];
+        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)rows + keys_per_tg - 1u) / keys_per_tg, 1, 1)
+             threadsPerThreadgroup:MTLSizeMake(tg, 1, 1)];
+        ds4_gpu_end_compute_encoder(cb, enc);
+        return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 bf16 index score");
+    }
+}
+
 static int ds4_gpu_glm_indexer_scores_batch_grouped_tensor(
         ds4_gpu_tensor       *scores,
         const ds4_gpu_tensor *q,

metal/dsv4_misc.metal

@@ -273,6 +273,11 @@ struct ds4_metal_args_glm_indexer_score_one {
     float    scale;
 };
 
+struct ds4_metal_args_dsv41_index_keys_bf16 {
+    uint32_t n_vec;
+    uint32_t check;
+};
+
 struct ds4_metal_args_glm_indexer_scores_batch {
     uint32_t n_rows;
     uint32_t n_tokens;
@@ -2064,6 +2069,67 @@ kernel void kernel_glm_indexer_score_one_wide(
     scores[row] = score;
 }
 
+/* Index keys as bf16 (FP4 values on a power-of-two scale are bf16-exact); with check, a key that is not raises *flag. */
+kernel void kernel_dsv41_index_keys_bf16(
+        constant ds4_metal_args_dsv41_index_keys_bf16 & args,
+        device const uint4 *src,
+        device ushort4 *dst,
+        device atomic_uint *flag,
+        uint gid [[thread_position_in_grid]]) {
+    if (gid >= args.n_vec) return;
+    const uint4 b = src[gid];
+    dst[gid] = ushort4(b >> 16);
+    if (args.check != 0u && any((b & 0xffffu) != 0u))
+        atomic_fetch_or_explicit(flag, 1u, memory_order_relaxed);
+}
+
+/* The wide scorer over the bf16 keys, two keys per thread: each dot and the head sum (0..31) in its order, its bits. */
+kernel void kernel_dsv41_indexer_score_bf16(
+        constant ds4_metal_args_glm_indexer_score_one & args,
+        device const float *qf,
+        device const float *weights,
+        device const ushort4 *keys,
+        device float *scores,
+        threadgroup float *shared [[threadgroup(0)]],
+        uint3 tgpig [[threadgroup_position_in_grid]],
+        ushort tid [[thread_index_in_threadgroup]],
+        uint3 ntg_u [[threads_per_threadgroup]]) {
+    const uint ntg = ntg_u.x;
+    threadgroup float *qs = shared;           /* [32][128] */
+    threadgroup float *ws = shared + 4096u;   /* [32] */
+    for (uint i = tid; i < 4096u; i += ntg) qs[i] = qf[i];
+    if (tid < 32u) ws[tid] = weights[tid];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (args.n_head != 32u || args.head_dim != 128u) return;
+
+    const uint first = tgpig.x * ntg * 2u + tid;
+    /* the tail's second key rereads the last row and is not written */
+    const uint row[2] = {min(first, args.n_rows - 1u), min(first + ntg, args.n_rows - 1u)};
+    float score[2] = {0.0f, 0.0f};
+    for (uint g = 0; g < 32u; g += 16u) {
+        float acc[2][16];
+        FOR_UNROLL (uint r = 0; r < 2u; r++) FOR_UNROLL (uint h = 0; h < 16u; h++) acc[r][h] = 0.0f;
+        for (uint d0 = 0; d0 < 128u; d0 += 16u) {
+            float4 k[2][4];
+            FOR_UNROLL (uint r = 0; r < 2u; r++) {
+                device const ushort4 *kb = keys + (((uint64_t)row[r] * 128u + d0) >> 2);
+                FOR_UNROLL (uint i = 0; i < 4u; i++) k[r][i] = as_type<float4>(uint4(kb[i]) << 16);
+            }
+            FOR_UNROLL (uint h = 0; h < 16u; h++) {
+                threadgroup const float4 *q4 = (threadgroup const float4 *)(qs + (g + h) * 128u + d0);
+                FOR_UNROLL (uint r = 0; r < 2u; r++)
+                    acc[r][h] += dot(q4[0], k[r][0]) + dot(q4[1], k[r][1]) + dot(q4[2], k[r][2]) + dot(q4[3], k[r][3]);
+            }
+        }
+        FOR_UNROLL (uint r = 0; r < 2u; r++) FOR_UNROLL (uint h = 0; h < 16u; h++) {
+#pragma clang fp contract(off)
+            score[r] += max(acc[r][h] * args.scale, 0.0f) * ws[g + h];
+        }
+    }
+    if (first < args.n_rows) scores[first] = score[0];
+    if (first + ntg < args.n_rows) scores[first + ntg] = score[1];
+}
+
 /* The wide kernel over V4.1 candidate keys: key i is position i % 8 of block
  * list[i / 8], so the scores keep ascending key order. */
 kernel void kernel_dsv41_indexer_score_candidates(
@@ -2114,6 +2180,67 @@ kernel void kernel_dsv41_indexer_score_candidates(
     scores[slot] = score;
 }
 
+/* The same with S lanes per key, the head sum passed lane to lane in order (2T / S keys per T threads). */
+template <uint S>
+kernel void kernel_dsv41_indexer_score_rows_split_bf16(
+        constant ds4_metal_args_glm_indexer_score_one & args,
+        device const float *qf,
+        device const float *weights,
+        device const ushort *keys,
+        device float *scores,
+        threadgroup float *shared [[threadgroup(0)]],
+        uint3 tgpig [[threadgroup_position_in_grid]],
+        ushort tid [[thread_index_in_threadgroup]],
+        uint3 ntg_u [[threads_per_threadgroup]]) {
+    const uint ntg = ntg_u.x;
+    threadgroup float *qs = shared;           /* [32][128] */
+    threadgroup float *ws = shared + 4096u;   /* [32] */
+    for (uint i = tid; i < 4096u; i += ntg) qs[i] = qf[i];
+    if (tid < 32u) ws[tid] = weights[tid];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (args.n_head != 32u || args.head_dim != 128u) return;
+    constexpr uint HS = 32u / S;
+    const uint part = tid % S, lanes = ntg / S;
+    uint slot[2], row[2];
+    FOR_UNROLL (uint r = 0; r < 2u; r++) {
+        slot[r] = tgpig.x * lanes * 2u + r * lanes + tid / S;
+        row[r] = min(slot[r], args.n_rows - 1u);
+    }
+    float acc[2][HS];
+    FOR_UNROLL (uint r = 0; r < 2u; r++) FOR_UNROLL (uint i = 0; i < HS; i++) acc[r][i] = 0.0f;
+    for (uint d0 = 0; d0 < 128u; d0 += 16u) {
+        float4 k[2][4];
+        FOR_UNROLL (uint r = 0; r < 2u; r++) {
+            device const ushort4 *kb = (device const ushort4 *)keys + (((uint64_t)row[r] * 128u + d0) >> 2);
+            FOR_UNROLL (uint i = 0; i < 4u; i++) k[r][i] = as_type<float4>(uint4(kb[i]) << 16);
+        }
+        FOR_UNROLL (uint i = 0; i < HS; i++) {
+            threadgroup const float4 *q4 = (threadgroup const float4 *)(qs + (part * HS + i) * 128u + d0);
+            FOR_UNROLL (uint r = 0; r < 2u; r++)
+                acc[r][i] += dot(q4[0], k[r][0]) + dot(q4[1], k[r][1]) + dot(q4[2], k[r][2]) + dot(q4[3], k[r][3]);
+        }
+    }
+    FOR_UNROLL (uint r = 0; r < 2u; r++) {
+        float score = 0.0f;
+        FOR_UNROLL (uint p = 0; p < S; p++) {
+            if (p) score = simd_shuffle_up(score, (ushort)1);
+            if (part == p) {
+                FOR_UNROLL (uint i = 0; i < HS; i++) {
+#pragma clang fp contract(off)
+                    score += max(acc[r][i] * args.scale, 0.0f) * ws[p * HS + i];
+                }
+            }
+        }
+        if (part == S - 1u && slot[r] < args.n_rows) scores[slot[r]] = score;
+    }
+}
+#define DSV41_ROWS_SPLIT(NAME, S) template [[host_name(NAME)]] \
+    kernel void kernel_dsv41_indexer_score_rows_split_bf16<S>( \
+    constant ds4_metal_args_glm_indexer_score_one &, device const float *, device const float *, device const ushort *, \
+    device float *, threadgroup float *, uint3, ushort, uint3);
+DSV41_ROWS_SPLIT("kernel_dsv41_indexer_score_rows_split2_bf16", 2)
+DSV41_ROWS_SPLIT("kernel_dsv41_indexer_score_rows_split4_bf16", 4)
+
 kernel void kernel_glm_indexer_scores_batch(
         constant ds4_metal_args_glm_indexer_scores_batch & args,
         device const char *q,

tests/test_deepseek41_metal.c

@@ -1780,6 +1780,84 @@ static int check_index_candidates(void) {
     return 1;
 }
 
+/* The bf16 key view and its scorers against the wide scorer on the f32 cache: the same bits for every key. Keys pass
+ * the FP4_E8M0 quantizer with block scales spread over 41 exponents and over 12 (where summation order shows in the
+ * bits); the view converts in two pieces at an odd boundary. The decode row entry reads the view where the wide
+ * scorer runs and scores the f32 keys without it. */
+static int check_index_score_bf16(void) {
+    enum { HEADS = 32, DIM = 128 };
+    const uint32_t counts[] = {8200, 65535, 65536, 98303, 131077};
+    const uint32_t max_keys = 131077;
+    const char *kernels[] = {"kernel_dsv41_indexer_score_rows_split4_bf16", "kernel_dsv41_indexer_score_rows_split2_bf16",
+                             "kernel_dsv41_indexer_score_bf16"};
+    ds4_gpu_tensor *qt = upload(NULL, (size_t)HEADS * DIM * sizeof(float));
+    ds4_gpu_tensor *wt = upload(NULL, (size_t)HEADS * sizeof(float));
+    ds4_gpu_tensor *kt = upload(NULL, (size_t)max_keys * DIM * sizeof(float));
+    ds4_gpu_tensor *vt = upload(NULL, (size_t)max_keys * DIM * sizeof(uint16_t));
+    ds4_gpu_tensor *sw = upload(NULL, ((size_t)max_keys + 1) * sizeof(float));
+    ds4_gpu_tensor *sb = upload(NULL, ((size_t)max_keys + 1) * sizeof(float));
+    ds4_gpu_tensor *ft = upload(NULL, sizeof(uint32_t));
+    CHECK(qt && wt && kt && vt && sw && sb && ft);
+    float *q = ds4_gpu_tensor_contents(qt), *w = ds4_gpu_tensor_contents(wt);
+    float *k = ds4_gpu_tensor_contents(kt);
+    float *wide = ds4_gpu_tensor_contents(sw), *fast = ds4_gpu_tensor_contents(sb);
+    uint32_t *flag = ds4_gpu_tensor_contents(ft);
+    for (int spread = 0; spread < 2; spread++) {
+        for (size_t i = 0; i < (size_t)HEADS * DIM; i++)
+            q[i] = spread == 0 ? ldexpf(random_value(), (int)((i / 32) % 41) - 20) :
+                                 ldexpf(random_value(), (int)((i / 32 * 7) % 13) - 6);
+        for (size_t i = 0; i < HEADS; i++) w[i] = bf16(random_value());
+        for (size_t i = 0; i < (size_t)max_keys * DIM; i++)
+            k[i] = spread == 0 ? ldexpf(random_value(), (int)((i / 32) % 41) - 20) :
+                                 ldexpf(random_value(), (int)((i / 32 * 5 + i / DIM) % 13) - 6);
+        CHECK(ds4_gpu_dsv41_quantize(qt, DIM, HEADS, DS4_V41_FP4_E8M0));
+        CHECK(ds4_gpu_dsv41_quantize(kt, DIM, max_keys, DS4_V41_FP4_E8M0));
+        *flag = 0;
+        CHECK(ds4_gpu_dsv41_index_keys_bf16(vt, kt, 0, 77777, ft));
+        CHECK(ds4_gpu_dsv41_index_keys_bf16(vt, kt, 77777, max_keys - 77777, ft));
+        CHECK(ds4_gpu_synchronize() && *flag == 0);
+        for (size_t ci = 0; ci < sizeof(counts) / sizeof(*counts); ci++) {
+            const uint32_t keys = counts[ci];
+            const char *kernel = kernels[keys < 10240u ? 0 : keys < 65536u ? 1 : 2];
+            wide[keys] = fast[keys] = 12345;
+            ds4_gpu_test_set_flags(DS4_GPU_TEST_INDEX_SCORE_WIDE);
+            CHECK(ds4_gpu_glm_indexer_score_one_tensor(sw, qt, wt, kt, keys, HEADS, DIM, 1.0f / 64.0f, false));
+            ds4_gpu_test_set_flags(DS4_GPU_TEST_COUNT_KERNELS);
+            for (int view = 1; view >= 0; view--) {
+                const uint64_t runs = ds4_gpu_test_kernel_runs(kernel);
+                memset(fast, 0xff, (size_t)keys * sizeof(float));
+                CHECK(ds4_gpu_dsv41_indexer_score_row(sb, qt, wt, kt, view ? vt : NULL, keys) && ds4_gpu_synchronize());
+                CHECK(ds4_gpu_test_kernel_runs(kernel) == runs + (uint64_t)view);
+                CHECK(!memcmp(wide, fast, (size_t)keys * sizeof(float)) && fast[keys] == 12345);
+            }
+            uint32_t positive = 0, rounded = 0;
+            for (uint32_t j = 0; j < keys; j++) positive += wide[j] > 0;
+            /* sensitivity: keys whose f32 score is not the exact sum, so order shows */
+            for (uint32_t j = 0; ci == 0 && j < 2048; j++) {
+                double expected = 0;
+                for (uint32_t h = 0; h < HEADS; h++) {
+                    double dot = 0;
+                    for (uint32_t d = 0; d < DIM; d++)
+                        dot += (double)q[(size_t)h * DIM + d] * k[(size_t)j * DIM + d];
+                    expected += fmax(dot / 64.0, 0) * w[h];
+                }
+                rounded += (double)wide[j] != expected;
+            }
+            CHECK(wide[keys] == 12345 && positive > keys / 16 && (ci != 0 || rounded > 0));
+        }
+    }
+    /* a key that is not bf16-exact raises the flag */
+    uint32_t bits;
+    memcpy(&bits, &k[(size_t)123456 * DIM + 77], 4);
+    bits |= 1u;
+    memcpy(&k[(size_t)123456 * DIM + 77], &bits, 4);
+    CHECK(ds4_gpu_dsv41_index_keys_bf16(vt, kt, 123456, 1, ft) && ds4_gpu_synchronize() && *flag != 0);
+    ds4_gpu_tensor_free(ft); ds4_gpu_tensor_free(sb); ds4_gpu_tensor_free(sw);
+    ds4_gpu_tensor_free(vt); ds4_gpu_tensor_free(kt); ds4_gpu_tensor_free(wt); ds4_gpu_tensor_free(qt);
+    fprintf(stderr, "V4.1 bf16 index score (split4, split2, full) == wide scorer bits PASS\n");
+    return 1;
+}
+
 /* The FP8 copy reader (ds4_gpu_matmul_e4m3_bf16_tensor), with its blend kernels and with the narrow readers, gives
  * the BF16-rounded F16 projection's bits at the Engram (6144 -> 256, two-row blend) and indexer q_b (1280 -> 4096,
  * block-owner blend) shapes: power-of-two 32x32 block scales 2^-15..2^-8 (every weight exact in F16), normal and
@@ -1861,6 +1939,7 @@ static int check_e4m3_copies(void) {
 #else
 static int check_index_score_wide(void) { return 1; }
 static int check_index_candidates(void) { return 1; }
+static int check_index_score_bf16(void) { return 1; }
 static int check_e4m3_copies(void) { return 1; }
 #endif
 
@@ -2115,6 +2194,11 @@ int main(int argc, char **argv) {
         ds4_gpu_cleanup();
         return ok ? 0 : 1;
     }
+    if (argc == 2 && !strcmp(argv[1], "--index-score-bf16")) {
+        const int ok = ds4_gpu_init() && check_index_score_bf16();
+        ds4_gpu_cleanup();
+        return ok ? 0 : 1;
+    }
     if (argc == 2 && !strcmp(argv[1], "--e4m3-copies")) {
         const int ok = ds4_gpu_init() && check_e4m3_copies();
         ds4_gpu_cleanup();
@@ -2151,6 +2235,7 @@ int main(int argc, char **argv) {
     ok = ok && check_e4m3_copies();
     ok = ok && check_rope_tab();
     ok = ok && check_split_collapse();
+    ok = ok && check_index_score_bf16();
     ds4_gpu_cleanup();
     return ok ? 0 : 1;
 }