fix(deepseek-v4): DSA top-k kernel could not represent the real index_topk — remove the literal bounds (#505)

[localai-bot]

Closes #505. Spec: .agents/specs/dsa-topk-bounds.md. Row: MODEL-TEXT-deepseek-v4-deepseek-v4-for-causal-lm.

The defect

DsaTopkKernel sized two thread-local arrays by literal:

bool chosen[512];    // indexed [0, n), n = candidate-window length
int64_t picked[64];  // written [0, topk)

topk is the caller's index_topk — 512 on V4-Flash, 1024 on V4-Pro — so picked[64] was 8x short for Flash and 16x short for Pro, and chosen[512] overflowed on any window wider than 512. Neither bound was asserted nor derived from config. Every pre-existing device case ran at topk=3, nk=5, which is why it was invisible: that's the gate's shape, not the model's.

Measured, not theoretical

On dgx.casa (GB10, sm_121a), the pre-fix kernel driven at Flash's own index_topk of 512 with a 600-wide window faults on the device:

what():  vt cuda: cudaStreamDestroy: an illegal memory access was encountered
test_cuda_deepseek_v4.cpp:214: FATAL ERROR: test case CRASHED: SIGABRT
[doctest] test cases:   6 |   5 passed | 1 failed | 17 skipped
[doctest] Status: FAILURE!

Line 214 is the new REAL index_topk widths case. After the fix, same flags: 23 cases | 23 passed | 0 failed | 0 skipped, 83913/83913 assertions, Status: SUCCESS!.

Both arms ran under flock $HOME/gpu.lock, each asserting its own kernel identity before building — matching the declarations bool chosen[512]; / int64_t picked[64];, not the tokens, since the fixed kernel's comment cites both by name and a token grep reports the fixed tree as the old one.

origin/main then advanced 17 commits, so the suite was rebuilt and re-run from the merged tree: 23/23, 83913 assertions, SUCCESS, with CUTLASS found … enabling sm120a NVFP4 cutlass GEMM and FlashAttention-2 … ENABLED for arch(es) [121a] hard-verified in that run's own configure log, and --list-test-cases confirming all 3 new cases in the binary.

The fix — two passes, no scratch

A threshold formulation over the same total order the host reference sorts by (logit desc, then index asc — total because candidate indices are distinct):

• pass 1 walks the order downwards topk times to land on the topk-th best element, the threshold;
• pass 2 scans the window once ascending and emits everything outranking-or-equal to it.

Exactly topk elements qualify under a total order, and they come out already ascending, so the O(topk^2) emit sort disappears along with the buffers. No per-thread scratch, no bound, no configurable limit. Pass 1 stays O(topk*n), so this is strictly cheaper overall.

Two defensive additions that are explicitly not capacity limits: pass 1 stops if no strictly-worse element exists, and pass 2 carries w < topk, so a NaN row — where every float comparison is false — cannot write past its own row into the next one, which is the failure class this issue was about.

Gate

Three new device cases, all against the independent DsaTopkSelect oracle (std::stable_sort based — a genuinely separate implementation, so this isn't a shared-helper tautology):

case	shape	what it pins
real index_topk widths	(65,80), (512,600), (1024,1200)	just past the old picked[64], then both shipped widths; also asserts no -1 leaks and strictly ascending emit
tie-heavy rows	topk=128, nk=300, quantized logits	the tie-break, which distinct random logits cannot exercise
offset window	topk=512, nk=900, ws=137	the old code indexed its mask by s - s0 and its picks by absolute s

Plus a local pre-flight while the GPU lock was held by other jobs (scratch, ASan+UBSan): the new body vs an independent transcription of the oracle — 0 mismatched entries over 8 named shapes and a 4000-shape randomized sweep, half coarsely quantized to force tie density. And an ASan proof of the defect itself: the old body at topk=512, nk=600 reports stack-buffer-overflow, WRITE of size 1 on the chosen[] init loop.

CPU V4 suites unmoved: dsa 38, pro_variant 197, scaffold 62, all SUCCESS. scripts/agent-preflight.sh --staged fully clean at the pushed SHA.

Rejected alternatives

Asserting the bounds instead of removing them — which #505 itself proposed. A refusal would be honest but would leave the device DSA path unable to run the real index_topk at all, trading a latent overflow for a guaranteed refusal the moment the real-geometry residual lands. The threshold form needs no bound, so there is nothing left to assert.

Sharing a selection helper between kernel and host reference — would remove the duplication and make a CPU test trivial, but would make the gate prove only self-consistency. The duplication is deliberate and recorded as a stop condition.

Scope

Latent, not shipped: dsa_dense = (be.gguf != nullptr) (deepseek_v4.cpp:668) forces the indexer off on the real keep-quant path. This change makes the kernel able to represent the real widths; it does not put the real path on it. The real-geometry DSA sparse path stays a named residual, as does the compressed-key-space candidate window — upstream builds it as ke = row_start + (pos+1) // COMPRESS_RATIO (v1/attention/backends/mla/indexer.py:270-290) while our synthetic path uses we[t] = t+1 over uncompressed keys. Spec §5 records that for whoever builds it out.

Review note

Authored in a coordinating session on direct developer instruction, so it has not had an independent review pass. A fresh reviewer re-running the RED arm on dgx is still owed.

🤖 Generated with Claude Code