indexer_topk: fix quantized q8_1 scratch sizing on CUDA (#2158)

* indexer_topk: fix quantized q8_1 scratch sizing on CUDA

The quantized-K path under-sized its q8_1 scratch buffer: it allocated q->ne[1]*max_rows blocks using the unpadded head dim, but quantize_mmq_q8_1_cuda writes q_padded/QK8_1 blocks per row and must process all q->ne[1]*nrows rows. Size the buffer by (q_padded/QK8_1) blocks x (q->ne[1]*max_rows) rows and pass the full q->ne[1]*nrows row count so the scratch cannot be overrun and every query row is quantized.

CUDA graph identity: INDEXER_TOPK dispatches a source-type-specialized kernel (dense F16 vs quantized cache, F32 vs F16 mask). Source addresses alone do not identify the captured kernel, so snapshot each source type in ggml_graph_node_properties and force re-capture when an INDEXER_TOPK source type changes, preventing a reused graph from replaying the wrong kernel variant when sources are reallocated at the same address.

CPU backend: report INDEXER_TOPK support via iqk_indexer_topk_supported (guarded by GGML_USE_IQK_MULMAT) so the scheduler places the node on a backend that can run it. Harden the scheduler's pass-5 node-assignment check from assert to GGML_ASSERT so a node no backend supports fails as a defined abort under NDEBUG instead of indexing sched->backends[-1].

* indexer_topk: drop CPU capability predicate

* indexer_topk: fall back when unsupported

* cuda: drop speculative INDEXER_TOPK graph type matching

---------

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>

ggml/src/ggml-backend.cpp

@@ -906,6 +906,12 @@ GGML_CALL static bool ggml_backend_cpu_supports_op(ggml_backend_t backend, const
         case GGML_OP_MUL_MAT:
             return true;
             //return op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == ggml_internal_get_type_traits(op->src[0]->type).vec_dot_type;
+        case GGML_OP_INDEXER_TOPK:
+#ifdef GGML_USE_IQK_MULMAT
+            return true;
+#else
+            return false;
+#endif
         default:
             return true;
     }

ggml/src/ggml-cuda/indexer_topk.cu

@@ -160,7 +160,7 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds
     ggml_cuda_pool_alloc<char>  q_converted(ctx.pool());
     auto q_padded = GGML_PAD(q->ne[0], MATRIX_ROW_PADDING);
     if (ggml_is_quantized(k->type)) {
-        auto nbytes_q = q->ne[1] * max_rows * sizeof(block_q8_1)/QK8_1;
+        auto nbytes_q = (size_t)(q_padded/QK8_1) * ((size_t) q->ne[1] * max_rows) * sizeof(block_q8_1);
         nbytes_q += get_mmq_x_max_host(ggml_cuda_info().devices[ctx.device].cc)*sizeof(block_q8_1_mmq);
         q_converted.alloc(nbytes_q);
     } else {
@@ -177,7 +177,7 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds
         auto q_data = (const char *)q->data + istep*max_rows*q->nb[2];
         auto m_data = (const char *)m->data + istep*max_rows*m->nb[1];
         if (ggml_is_quantized(k->type)) {
-            quantize_mmq_q8_1_cuda((const float *)q_data, q_converted.get(), q->ne[0], nrows, 1, q_padded, k->type, ctx.stream());
+            quantize_mmq_q8_1_cuda((const float *)q_data, q_converted.get(), q->ne[0], q->ne[1]*nrows, 1, q_padded, k->type, ctx.stream());
             CUDA_CHECK(cudaGetLastError());
             mmq_args args{(const char *)k->data, q_converted.get(), kq.get(),
                 k->ne[0], k->ne[1], int64_t(k->nb[1]),

src/graphs/build_deepseek2.cpp

@@ -497,8 +497,10 @@ ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer(
             ggml_build_forward_expand(gf, indexer_score);
         }
         auto topk = ggml_indexer_topk(ctx0, indexer_k_b, indexer_q, indexer_weights, indexer_score, GGML_UNARY_OP_RELU, n_top_k);
-        ggml_build_forward_expand(gf, topk);
-        return topk;
+        if (supports_op(topk)) {
+            ggml_build_forward_expand(gf, topk);
+            return topk;
+        }
     }
 
     if (indexer_q->ne[2] <= 8) {

src/graphs/build_openpangu.cpp

@@ -375,16 +375,18 @@ ggml_tensor * llm_build_context::build_openpangu_attention(
                 !dsa_gather_allowed &&
                 openpangu_att_score_should_chunk(n_kv, hparams.param_sink_number, n_head, n_tokens,
                         OPENPANGU_ATT_SCORE_CHUNK, OPENPANGU_ATT_FULL_KQ_MAX_MIB);
+            ggml_tensor * fused_sel_idx = nullptr;
             if (lctx.cparams.fused_idx_topk) {
-                // Fused indexer top-k (CPU-only op): one op computes sum_g w * relu(q.k) + causal
-                // mask -> top-k without materializing the [n_kv, n_ihead, T] score tensor, so no
-                // score chunking is needed. CUDA backends do not implement GGML_OP_INDEXER_TOPK;
-                // the scheduler runs it on CPU and copies the operands across the backend boundary.
                 // The op reads the mask row-strided, so the raw view suffices.
                 ggml_tensor * idx_mask = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens,
                                                       KQ_mask->nb[1], 0);
-                sel_idx = ggml_indexer_topk(ctx0, k_all_idx, q_idx, w_idx, idx_mask,
-                                            GGML_UNARY_OP_RELU, (int) topk);          // [topk, T] i32
+                fused_sel_idx = ggml_indexer_topk(ctx0, k_all_idx, q_idx, w_idx, idx_mask,
+                                                  GGML_UNARY_OP_RELU, (int) topk);    // [topk, T] i32
+            }
+            if (fused_sel_idx && supports_op(fused_sel_idx)) {
+                // One op computes sum_g w * relu(q.k) + causal mask -> top-k without
+                // materializing the [n_kv, n_ihead, T] score tensor.
+                sel_idx = fused_sel_idx;
                 if (il == 0) ggml_set_name(sel_idx, "opg0_idx_sel");
                 dsa_gather_engaged = dsa_gather_allowed;
                 if (dsa_gather_engaged) {

src/llama-build-context.cpp

@@ -139,6 +139,15 @@ void llm_build_context::free() {
     }
 }
 
+bool llm_build_context::supports_op(const ggml_tensor * op) const {
+    for (ggml_backend_t backend : lctx.backends) {
+        if (ggml_backend_supports_op(backend, op)) {
+            return true;
+        }
+    }
+    return false;
+}
+
 ggml_cgraph * llm_build_context::build_k_shift() {
     struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(n_tokens), false);
 

src/llama-build-context.h

@@ -119,6 +119,8 @@ struct llm_build_context {
 
     void free();
 
+    bool supports_op(const ggml_tensor * op) const;
+
     ggml_cgraph * build_k_shift();
 
     ggml_cgraph * build_s_copy();