Score only the candidate blocks of later DeepSeek V4.1 index layers

ds4.c

@@ -40925,7 +40925,7 @@ static uint32_t ds41_carry_cap(uint32_t ctx) {
     X(selected_kv, DS4_N_INDEXER_TOP_K * DS4_N_HEAD_DIM) \
     X(block_scores, DS41_INDEX_BATCH * ((g->ctx + 7u) / 8u)) \
     X(block_selected, DS41_INDEX_BATCH * 2048u) \
-    X(block_mask, (g->ctx + 7u) / 8u) \
+    X(block_mask, (g->ctx + 7u) / 8u) X(block_list, 2048u) \
     X(heads, DS4_N_HEAD * DS4_N_HEAD_DIM) X(low, DS4_N_OUT_GROUP * DS4_N_LORA_O) \
     X(route_logits, DS4_N_EXPERT) X(route_probs, DS4_N_EXPERT) \
     X(selected, DS4_N_EXPERT_USED) X(route_weights, DS4_N_EXPERT_USED) \
@@ -41673,12 +41673,17 @@ static bool ds41_attention_select_published(ds41_gpu_graph *g, const ds4_model *
         if (!ds41_matmul(g->index_q, m, l->indexer_attn_q_b, g->qr, true) ||
             !ds41_rope(g->index_q, DS4_N_INDEXER_HEAD, 128, il, pos, false) ||
             !ds4_gpu_dsv41_quantize(g->index_q, 128, DS4_N_INDEXER_HEAD, DS4_V41_FP4_E8M0) ||
-            !ds41_matmul(g->index_weights, m, l->indexer_proj, g->norm, true) ||
+            !ds41_matmul(g->index_weights, m, l->indexer_proj, g->norm, true)) return false;
 #ifdef __APPLE__
-            !ds4_gpu_dsv41_indexer_score_one_tensor(g->index_scores, g->index_q, g->index_weights,
+        /* past layer 20 only its candidate blocks can be selected */
+        if (il > 20u && (n_comp + 7u) / 8u > 2048u &&
+            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;
 #else
-            !ds4_gpu_dsv41_indexer_scores_batch(g->index_scores, g->index_q, g->index_weights,
+        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;
 #endif
     }

ds4_deepseek41_gpu.h

@@ -136,6 +136,14 @@ 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);
+/* 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. */
+int ds4_gpu_dsv41_indexer_candidates_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *scores,
+                                           ds4_gpu_tensor *list, const ds4_gpu_tensor *q,
+                                           const ds4_gpu_tensor *weights, const ds4_gpu_tensor *keys,
+                                           const ds4_gpu_tensor *block_mask, uint32_t n_comp,
+                                           float scale, uint32_t top_k);
 int ds4_gpu_dsv41_tensor_ops_available(void);
 /* Reuse exact BF16 views of FP4 queries/keys across score tiles. Invalid
  * casts retain the F32 arithmetic for the affected tile. */

ds4_metal.m

@@ -49536,6 +49536,79 @@ int ds4_gpu_dsv41_candidate_mask_batch(ds4_gpu_tensor *mask,
     return ds4_gpu_dsv4_topk_mask_tensor_impl(mask, selected, width, rows, 2048u, mask_stride);
 }
 
+/* The wide scorer and radix select of a filtered row, over its candidates
+ * only: in key order the winners, their order and ties are the same, and the
+ * newest block, always a candidate, ends the list. Returns 0 to keep the
+ * full row where those kernels would not have run. */
+int ds4_gpu_dsv41_indexer_candidates_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *scores,
+                                           ds4_gpu_tensor *list, const ds4_gpu_tensor *q,
+                                           const ds4_gpu_tensor *weights, const ds4_gpu_tensor *keys,
+                                           const ds4_gpu_tensor *block_mask, uint32_t n_comp,
+                                           float scale, uint32_t top_k) {
+    static id<MTLComputePipelineState> list_pipeline, score_pipeline, ids_pipeline;
+    static int env = -1;
+    if (env < 0) env = !getenv("DS4_METAL_DISABLE_V41_INDEX_CANDIDATES");
+    const uint32_t blocks = n_comp <= (1u << 23) ? (n_comp + 7u) / 8u : 0u;
+    const uint32_t count = 16384u - (blocks * 8u - n_comp);
+    if (!env || blocks <= 2048u || !top_k || top_k > 2048u || ds4_gpu_v41_index_score_wide_off() ||
+        getenv("DS4_METAL_DISABLE_V41_TOPK_SELECT") ||
+        (g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) != 0u ||
+        !dsv41_tensor_has_floats(selected, top_k) || !dsv41_tensor_has_floats(scores, count) ||
+        !dsv41_tensor_has_floats(list, 2048u) || !dsv41_tensor_has_floats(block_mask, blocks) ||
+        !dsv41_tensor_has_floats(q, 4096u) || !dsv41_tensor_has_floats(weights, 32u) ||
+        !dsv41_tensor_has_floats(keys, (uint64_t)n_comp * 128u)) return 0;
+    if (!g_initialized && !ds4_gpu_init()) return 0;
+    @autoreleasepool {
+        if (!list_pipeline) {
+            id<MTLComputePipelineState> hist = ds4_gpu_get_pipeline("kernel_topk_select_hist");
+            list_pipeline = ds4_gpu_get_pipeline("kernel_dsv41_candidate_list");
+            score_pipeline = ds4_gpu_get_pipeline("kernel_dsv41_indexer_score_candidates");
+            ids_pipeline = ds4_gpu_get_pipeline("kernel_dsv41_candidate_ids");
+            if (!hist || hist.maxTotalThreadsPerThreadgroup < 1024u || !score_pipeline || !ids_pipeline ||
+                list_pipeline.maxTotalThreadsPerThreadgroup < 1024u) list_pipeline = nil;
+            if (!list_pipeline) env = 0;
+        }
+        if (!list_pipeline) return 0;
+        const ds4_gpu_glm_indexer_score_one_args args = {
+            .n_rows = count, .n_head = 32u, .head_dim = 128u, .cache_f16 = 0u, .scale = scale,
+        };
+        int owned = 0;
+        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
+        id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
+        if (!enc) return 0;
+        [enc setComputePipelineState:list_pipeline];
+        [enc setBytes:&blocks length:sizeof(blocks) atIndex:0];
+        [enc setBuffer:ds4_gpu_tensor_buffer(block_mask) offset:ds4_gpu_tensor_offset(block_mask) atIndex:1];
+        [enc setBuffer:ds4_gpu_tensor_buffer(list) offset:ds4_gpu_tensor_offset(list) atIndex:2];
+        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
+        ds4_gpu_end_compute_encoder(cb, enc);
+        enc = ds4_gpu_compute_encoder(cb);
+        [enc setComputePipelineState:score_pipeline];
+        [enc setBytes:&args length:sizeof(args) atIndex:0];
+        const ds4_gpu_tensor *inputs[] = {q, weights, keys, scores, list};
+        for (NSUInteger i = 0; i < 5; i++)
+            [enc setBuffer:ds4_gpu_tensor_buffer(inputs[i])
+                    offset:ds4_gpu_tensor_offset(inputs[i]) atIndex:i + 1];
+        [enc setThreadgroupMemoryLength:(4096u + 32u) * sizeof(float) atIndex:0];
+        [enc dispatchThreadgroups:MTLSizeMake((count + 255u) / 256u, 1, 1)
+             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+        ds4_gpu_end_compute_encoder(cb, enc);
+        if (!ds4_gpu_finish_command_buffer(cb, owned, "V4.1 candidate index scores") ||
+            !ds4_gpu_indexer_topk_select(ds4_gpu_tensor_buffer(selected), ds4_gpu_tensor_offset(selected),
+                ds4_gpu_tensor_buffer(scores), ds4_gpu_tensor_offset(scores), count, 1, top_k)) return 0;
+        cb = ds4_gpu_command_buffer(&owned);
+        enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
+        if (!enc) return 0;
+        [enc setComputePipelineState:ids_pipeline];
+        [enc setBytes:&top_k length:sizeof(top_k) atIndex:0];
+        [enc setBuffer:ds4_gpu_tensor_buffer(list) offset:ds4_gpu_tensor_offset(list) atIndex:1];
+        [enc setBuffer:ds4_gpu_tensor_buffer(selected) offset:ds4_gpu_tensor_offset(selected) atIndex:2];
+        [enc dispatchThreads:MTLSizeMake(top_k, 1, 1) threadsPerThreadgroup:MTLSizeMake(MIN(top_k, 256u), 1, 1)];
+        ds4_gpu_end_compute_encoder(cb, enc);
+        return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 candidate ids");
+    }
+}
+
 int ds4_gpu_dsv41_hc_mean(uint32_t rows, uint32_t dim, uint32_t hc,
                          const ds4_gpu_tensor *stream, ds4_gpu_tensor *out,
                          uint32_t out_stride, uint32_t out_off) {

metal/dsv41.metal

@@ -375,6 +375,55 @@ kernel void kernel_dsv41_candidate_filter(
         ? scores[offset] : -INFINITY;
 }
 
+/* The 2048 candidate blocks of a 0/-inf mask in ascending order; one
+ * threadgroup of 1024 threads, four mask entries per thread and step. */
+kernel void kernel_dsv41_candidate_list(
+        constant uint &blocks,
+        device const float *mask,
+        device uint *list,
+        uint tid [[thread_index_in_threadgroup]],
+        uint lane [[thread_index_in_simdgroup]],
+        uint simd [[simdgroup_index_in_threadgroup]]) {
+    threadgroup uint part[33];
+    uint base = 0u;
+    for (uint b0 = 0u; b0 < blocks; b0 += 4096u) {
+        const uint b = b0 + tid * 4u;
+        bool keep[4];
+        uint count = 0u;
+        for (uint j = 0u; j < 4u; j++) {
+            keep[j] = b + j < blocks && mask[b + j] == 0.0f;
+            count += keep[j] ? 1u : 0u;
+        }
+        const uint before = simd_prefix_exclusive_sum(count);
+        if (lane == 31u) part[simd] = before + count;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (simd == 0u) {
+            const uint c = part[lane], c_before = simd_prefix_exclusive_sum(c);
+            part[lane] = c_before;
+            if (lane == 31u) part[32] = c_before + c;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        uint at = base + part[simd] + before;
+        for (uint j = 0u; j < 4u; j++) {
+            if (keep[j] && at < 2048u) list[at] = b + j;
+            at += keep[j] ? 1u : 0u;
+        }
+        base += part[32];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+}
+
+/* Positions in the candidate list back to key ids, in place and in order. */
+kernel void kernel_dsv41_candidate_ids(
+        constant uint &count,
+        device const uint *list,
+        device int *selected,
+        uint i [[thread_position_in_grid]]) {
+    if (i >= count) return;
+    const uint p = (uint)selected[i];
+    selected[i] = (int)(list[p >> 3] * 8u + (p & 7u));
+}
+
 struct ds4_metal_args_dsv41_carry {
     uint width, rows, words, format, pack;
 };

metal/dsv4_misc.metal

@@ -2064,6 +2064,56 @@ kernel void kernel_glm_indexer_score_one_wide(
     scores[row] = score;
 }
 
+/* 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(
+        constant ds4_metal_args_glm_indexer_score_one & args,
+        device const char *q,
+        device const float *weights,
+        device const char *indexer_key_cache,
+        device float *scores,
+        device const uint *list,
+        threadgroup float *shared [[threadgroup(0)]],
+        uint3 tgpig [[threadgroup_position_in_grid]],
+        ushort tid [[thread_index_in_threadgroup]],
+        uint3 ntg_u [[threads_per_threadgroup]]) {
+    if (args.n_head != 32u || args.head_dim != 128u) return;
+    const uint ntg = ntg_u.x;
+    threadgroup float *qs = shared;           /* [32][128] */
+    threadgroup float *ws = shared + 4096u;   /* [32] */
+    device const float *qf = (device const float *)q;
+    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);
+
+    const uint slot = tgpig.x * ntg + tid;
+    if (slot >= args.n_rows) return;
+    const uint row = list[slot >> 3] * 8u + (slot & 7u);
+    float acc[32];
+    FOR_UNROLL (uint h = 0; h < 32u; h++) acc[h] = 0.0f;
+    const uint64_t base = (uint64_t)row * 128u;
+    for (uint d0 = 0; d0 < 128u; d0 += 16u) {
+        float4 k[4];
+        if (args.cache_f16 != 0u) {
+            device const half4 *kh = (device const half4 *)indexer_key_cache + ((base + d0) >> 2);
+            FOR_UNROLL (uint i = 0; i < 4u; i++) k[i] = float4(kh[i]);
+        } else {
+            device const float4 *kf = (device const float4 *)indexer_key_cache + ((base + d0) >> 2);
+            FOR_UNROLL (uint i = 0; i < 4u; i++) k[i] = kf[i];
+        }
+        FOR_UNROLL (uint h = 0; h < 32u; h++) {
+            threadgroup const float4 *q4 = (threadgroup const float4 *)(qs + h * 128u + d0);
+            acc[h] += dot(q4[0], k[0]) + dot(q4[1], k[1]) + dot(q4[2], k[2]) + dot(q4[3], k[3]);
+        }
+    }
+    float score = 0.0f;
+    FOR_UNROLL (uint h = 0; h < 32u; h++) {
+#pragma clang fp contract(off)
+        score += max(acc[h] * args.scale, 0.0f) * ws[h];
+    }
+    scores[slot] = score;
+}
+
 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

@@ -1360,8 +1360,102 @@ static int check_index_score_wide(void) {
     fprintf(stderr, "V4.1 wide index score: double-precision oracle PASS\n");
     return 1;
 }
+
+/* Candidate-only index rows against the wide row they replace (wide scores,
+ * candidate filter, radix top-512): the same ids in the same order, with
+ * realistic and wide FP4 block scales, partial newest blocks and ties. */
+static int check_index_candidates(void) {
+    enum { HEADS = 32, DIM = 128, TOP = 512, ROUNDS = 6 };
+    const uint32_t widths[] = {16385, 16391, 20007, 65536, 131072, 131075, 262147};
+    const uint32_t max_keys = 262147, max_blocks = (max_keys + 7) / 8;
+    ds4_gpu_tensor *qt = upload(NULL, (size_t)HEADS * DIM * 4), *wt = upload(NULL, HEADS * 4);
+    ds4_gpu_tensor *kt = upload(NULL, (size_t)max_keys * DIM * 4);
+    ds4_gpu_tensor *full = upload(NULL, (size_t)max_keys * 4), *plain = upload(NULL, (size_t)max_keys * 4);
+    ds4_gpu_tensor *source = upload(NULL, (size_t)max_keys * 4), *cand = upload(NULL, 16384 * 4);
+    ds4_gpu_tensor *list = upload(NULL, 2048 * 4), *blocks_t = upload(NULL, (size_t)max_blocks * 4);
+    ds4_gpu_tensor *picked = upload(NULL, 2048 * 4), *mask = upload(NULL, (size_t)max_blocks * 4);
+    ds4_gpu_tensor *old_sel = upload(NULL, TOP * 4), *new_sel = upload(NULL, TOP * 4);
+    CHECK(qt && wt && kt && full && plain && source && cand && list && blocks_t && picked && mask &&
+          old_sel && new_sel);
+    float *q = ds4_gpu_tensor_contents(qt), *w = ds4_gpu_tensor_contents(wt);
+    float *k = ds4_gpu_tensor_contents(kt), *s = ds4_gpu_tensor_contents(source);
+    const float *p = ds4_gpu_tensor_contents(plain), *c = ds4_gpu_tensor_contents(cand);
+    const float *m = ds4_gpu_tensor_contents(mask);
+    const uint32_t *l = ds4_gpu_tensor_contents(list);
+    const int32_t *a = ds4_gpu_tensor_contents(old_sel), *b = ds4_gpu_tensor_contents(new_sel);
+    for (size_t wi = 0; wi < sizeof(widths) / sizeof(*widths); wi++) {
+        const uint32_t n = widths[wi], blocks = (n + 7) / 8, count = 16384 - (blocks * 8 - n);
+        double seconds[2] = {0, 0};
+        uint32_t ties = 0;
+        for (int mode = 0; mode < 4; mode++) {
+            /* bit 0: block exponents over +-20 instead of +-2; bit 1: ties */
+            const int range = mode & 1 ? 20 : 2;
+            for (size_t i = 0; i < (size_t)HEADS * DIM; i++)
+                q[i] = ldexpf(random_value(), (int)((i / 32) % (2 * range + 1)) - range);
+            for (size_t i = 0; i < (size_t)n * DIM; i++)
+                k[i] = ldexpf(random_value(), (int)((i / 32 + n) % (2 * range + 1)) - range);
+            for (size_t i = 0; i < HEADS; i++) w[i] = random_value() * 0.37139f;
+            for (uint32_t i = 0; i < n; i++) s[i] = (float)((i * 7919u + mode) % 1048573u);
+            CHECK(ds4_gpu_dsv41_quantize(qt, DIM, HEADS, DS4_V41_FP4_E8M0));
+            CHECK(ds4_gpu_dsv41_quantize(kt, DIM, n, DS4_V41_FP4_E8M0));
+            CHECK(ds4_gpu_synchronize());
+            if (mode & 2) {
+                /* copy the best key over 1/16 of all keys: the 512th score ties */
+                CHECK(ds4_gpu_dsv41_indexer_score_one_tensor(plain, qt, wt, kt, n, HEADS, DIM,
+                    1.0f / 64.0f, false) && ds4_gpu_synchronize());
+                uint32_t best = 0;
+                for (uint32_t i = 1; i < n; i++) if (p[i] > p[best]) best = i;
+                for (uint32_t i = 0; i < n; i += 16)
+                    memcpy(k + (size_t)(i ^ 5) % n * DIM, k + (size_t)best * DIM, DIM * 4);
+            }
+            CHECK(ds4_gpu_begin_commands());
+            CHECK(ds4_gpu_dsv41_candidate_blocks(blocks_t, source, n, 1, n - 1, 1));
+            CHECK(ds4_gpu_dsv41_indexer_topk_tensor(picked, blocks_t, blocks, 1, 2048));
+            CHECK(ds4_gpu_dsv4_topk_mask_tensor(mask, picked, blocks, 1, 2048));
+            CHECK(ds4_gpu_dsv41_indexer_score_one_tensor(plain, qt, wt, kt, n, HEADS, DIM, 1.0f / 64.0f, false));
+            CHECK(ds4_gpu_end_commands());
+            for (int round = 0; round < ROUNDS; round++) {
+                for (int path = 0; path < 2; path++) {
+                    const double t0 = monotonic_seconds();
+                    CHECK(ds4_gpu_begin_commands());
+                    if (path) {
+                        CHECK(ds4_gpu_dsv41_indexer_candidates_tensor(new_sel, cand, list, qt, wt, kt, mask,
+                            n, 1.0f / 64.0f, TOP));
+                    } else {
+                        CHECK(ds4_gpu_dsv41_indexer_score_one_tensor(full, qt, wt, kt, n, HEADS, DIM,
+                            1.0f / 64.0f, false));
+                        CHECK(ds4_gpu_dsv41_candidate_filter(full, mask, n, 1, n - 1, 1, 0));
+                        CHECK(ds4_gpu_dsv41_indexer_topk_tensor(old_sel, full, n, 1, TOP));
+                    }
+                    CHECK(ds4_gpu_end_commands() && ds4_gpu_synchronize());
+                    if (round) seconds[path] += monotonic_seconds() - t0;
+                }
+                CHECK(!memcmp(a, b, TOP * sizeof(*a)));
+            }
+            /* the list is the mask's candidates in order; each score is the wide one */
+            uint32_t at = 0;
+            for (uint32_t j = 0; j < blocks; j++) if (m[j] == 0.0f) CHECK(at < 2048 && l[at++] == j);
+            CHECK(at == 2048 && l[2047] == blocks - 1);
+            for (uint32_t i = 0; i < count; i++) {
+                const uint32_t key = l[i / 8] * 8 + i % 8;
+                CHECK(key < n && !memcmp(c + i, p + key, sizeof(*c)));
+            }
+            for (uint32_t i = 0; i < TOP; i++) ties += i && p[a[i]] == p[a[i - 1]];
+        }
+        fprintf(stderr, "V4.1 candidate index rows keys=%u candidates=%u: ids, order and scores "
+                "identical (%u tied neighbours) us per row full=%.0f candidates=%.0f\n",
+                n, count, ties, seconds[0] / (4 * (ROUNDS - 1)) * 1e6, seconds[1] / (4 * (ROUNDS - 1)) * 1e6);
+    }
+    CHECK(!ds4_gpu_dsv41_indexer_candidates_tensor(new_sel, cand, list, qt, wt, kt, mask,
+        16384, 1.0f / 64.0f, TOP));
+    ds4_gpu_tensor *tensors[] = {qt, wt, kt, full, plain, source, cand, list, blocks_t, picked, mask,
+                                 old_sel, new_sel};
+    for (size_t i = 0; i < sizeof(tensors) / sizeof(*tensors); i++) ds4_gpu_tensor_free(tensors[i]);
+    return 1;
+}
 #else
 static int check_index_score_wide(void) { return 1; }
+static int check_index_candidates(void) { return 1; }
 #endif
 
 static int check_compact_carry(void) {
@@ -1585,6 +1679,11 @@ int main(int argc, char **argv) {
         ds4_gpu_cleanup();
         return ok ? 0 : 1;
     }
+    if (argc == 2 && !strcmp(argv[1], "--index-candidates")) {
+        const int ok = ds4_gpu_init() && check_index_candidates();
+        ds4_gpu_cleanup();
+        return ok ? 0 : 1;
+    }
     if (argc == 2 && !strcmp(argv[1], "--index-topk")) {
         const int ok = ds4_gpu_init() && check_causal_topk();
         ds4_gpu_cleanup();
@@ -1609,7 +1708,7 @@ int main(int argc, char **argv) {
     int ok = ds4_gpu_init() && check_router() && check_quantization() && check_engram() && check_rope_stride() &&
              check_rope_quantize() && check_pool() && check_candidates() && check_candidates_batch(16449, 17, 16400) &&
              check_candidates_batch(17017, 17, 17000) && check_sparse_gather() && check_indexer_batch() &&
-             check_index_score_wide() && check_embedding() && check_index_projection() && check_general_topk() &&
+             check_index_score_wide() && check_index_candidates() && check_embedding() && check_index_projection() && check_general_topk() &&
              check_topk_select() && check_causal_topk() && check_indexer_all() && check_compact_carry() &&
              check_attention_output(false) && check_tp_attention();
     ds4_gpu_cleanup();