indexer_topk: fix quantized q8_1 scratch sizing on CUDA

[joelfarthing]

openPangu-2.0-Flash drives the shared GGML_OP_INDEXER_TOPK op with a quantized indexer K cache (-ictk q8_0), which comes across a scratch-sizing bug in the op's CUDA path. The bug lives in code every INDEXER_TOPK caller shares (DeepSeek2, GLM-DSA), so I'm presenting the fix here separately from the openPangu operator PR #2159, which stacks on it.

The quantized path under-sizes its q8_1 scratch. The old size expression omits the per-row block count q_padded / QK8_1 that quantize_mmq_q8_1_cuda actually writes, and the call passes nrows where the quantizer must process all q->ne[1] * nrows rows, so with the multi-head indexer q the scratch is under-allocated and the matmul consumes rows that were never quantized. The fix sizes the buffer at (q_padded / QK8_1) blocks by (q->ne[1] * max_rows) rows and passes the full row count to the quantizer.

-        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);
         ...
-        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());

The rest of the diff is the not-used-when-unsupported behavior from the review. The graph builders now construct the fused top-k node only when a loaded backend reports support for it, and otherwise keep the unfused chain, via a small supports_op helper on the build context. The CPU backend reports support under GGML_USE_IQK_MULMAT builds with a blanket true, which is exactly base's default: return true, and reports false without IQK. So in every IQK build (the default) node placement and graphs are unchanged from base; the only build whose graph changes is a non-IQK CPU-only build, which on base claimed the op and then could not run it, and now keeps the unfused chain instead.

Scope and blast radius. The CUDA support gate is unchanged: ggml_backend_cuda_supports_op still returns true for INDEXER_TOPK unconditionally. build_deepseek2.cpp produces the same graph as base in any build where a backend supports the op, which is every real configuration. The op's numerical contract is untouched.

Validation, measured at the head this PR opened with (77ab94c4) against main (9d07d8681, still the current base); the fix hunk is byte-identical at the present head and nothing added since touches the quantized path. Config: -ngl 999 -ot exps=CPU -fa off -ctk q8_0 -ictk q8_0 -fidx -dsatk 64, greedy.

A 32K-token needle (planted near 65% depth) does not survive a quantized indexer cache on main: raw completion degenerates to 0 = 0 = 0 =, and through the chat template the model concludes the passage contains no mention of any such code and collapses. With this fix the raw-completion needle is retrieved verbatim.

Perplexity over wikitext (n_ctx 512, batch matched so single-sequence perplexity packs one chunk per batch) pins the failure to the quantized path:

indexer cache	current main	this PR
q8_0	15.49	9.80
f16	9.90	9.90

(19 chunks; the full 109-chunk q8_0 run is 14.91 on main and 8.31 with the fix.) The degradation is specific to a q8_0 indexer cache on main: every other cell sits near 9.8, including main's own f16, so the true perplexity is about 9.8 and the fix restores the quantized path to it. The f16 result is byte-identical between the two (9.8953 both), so the scratch change touches only the quantized path. -ictk defaults to f16.

For GLM-DSA or DeepSeek v4, the A/B to run is -fidx on CUDA: F16 indexer cache, unchanged selections; q8_0 indexer cache, corrected selections.

• I have read the contributing guidelines
• Self-reported review complexity:
  • Low
  • Medium
  • High

[ikawrakow]

This sort of thing can only be written by an LLM. Let me see: suppose we handle correctly size_t overflow. Do we expect that a situation where we are dealing with number if operations greater than the maximum value of size_t on a 64-bit platform to ever complete? On a 10 TFlop CPU this will take several days.

[ikawrakow]

Please discard all changes to IQK. If we cannot compute the indexer topk op on the CPU, then it shall not be used instead of having 160 LOC of checks.

[joelfarthing]

Please discard all changes to IQK. If we cannot compute the indexer topk op on the CPU, then it shall not be used instead of having 160 LOC of checks.

Done. Thanks again. I also added a fallback when unsupported; please let me know if that's not wanted.

[ikawrakow]

You actually ran into this situation? Can you share how?

[joelfarthing]

You actually ran into this situation? Can you share how?

Yes, but this came from a focused INDEXER_TOPK test harness, not from changing cache types inside one openPangu context.

The harness creates one CUDA backend and then runs three same-shape fixtures sequentially: F16 K with an F32 mask, F16 K with an F16 mask, then Q8_0 K with an F32 mask. Each fixture frees its backend buffer before the next one, but the CUDA backend and its graph cache remain alive. The allocator reused the same source addresses.

With CUDA graphs enabled, the matcher accepted the later graph because the output shape/strides, source addresses, and op parameters matched. It did not see that the source types had changed, so INDEXER_TOPK replayed the previously captured type-specialized kernel. The numeric reference/parity check failed. The same sequence passed with CUDA graphs disabled.

Recording the source types forced a recapture and made the graphs-enabled sequence pass. So I did hit the exact same-address/different-type case, but in a deliberate same-process test. I should have made that distinction clear in the description.

[ikawrakow]

OK, theoretically the graph capture/comparison logic can fail in gazillions of different ways. If we wanted to guard against all possible failure modes, but at the same time still be able to reuse captured graphs, it would become way more complicated and possibly so expensive that it wouldn't bring any benefits to use it. But fortunately for us, we are not writing airplane control software here, so we go by actual usage in the inference engine only.