[ROCm][DSv3.2] Fix FP8 cast in indexer_k_quant_and_cache_triton (top-K accuracy regression for ctx>2048)

[maeehart]

Summary

Fixes the DeepSeek-V3.2 sparse-MLA + MTP top-K accuracy collapse for context_len > 2048 on AMD MI355X by correcting an FP8-cast bug in indexer_k_quant_and_cache_triton.

Root cause

The Triton kernel _indexer_k_quant_and_cache_kernel writes quantized values via:

fp8_val = (val.to(tl.float32) / scale).to(kv_cache_ptr.type.element_ty)

The Python wrapper indexer_k_quant_and_cache_triton previously sliced the uint8-typed kv_cache to extract the K region without re-viewing it as FP8:

kv_cache_value = kv_cache[:, : block_size * head_dim]   # uint8 dtype

so the kernel saw kv_cache_ptr.type.element_ty == tl.uint8. Triton's <float32>.to(tl.uint8) is integer truncation (e.g. 1.7 -> 1, -0.3 -> 255 wrap-around) — not FP8 quantization. The kernel wrote arbitrary integer bytes into the cache, and the downstream _gluon_deepgemm_fp8_paged_mqa_logits_preshuffle kernel then read those bytes as FP8 e4m3fnuz, producing garbled K values.

For context_len <= 2048 the noise happened to preserve top-K rank often enough that the regression was hard to spot, but beyond 2048 the dominant logits no longer aligned with the reference and top-K parity collapsed (~0% match), which is exactly the MTP draft-acceptance cliff observed in production.

The commit message of ab4b563 (which introduced the SHUFFLE indexer) already noted a 3pp GSM8K regression "under investigation" — this is that regression's root cause.

The fix

-    kv_cache_value = kv_cache[:, : block_size * head_dim]
+    kv_cache_value = kv_cache[:, : block_size * head_dim].view(
+        current_platform.fp8_dtype()
+    )

current_platform.fp8_dtype() returns torch.float8_e4m3fnuz on AMD ROCm and torch.float8_e4m3fn on NVIDIA CUDA, matching the kernel's IS_FNUZ constant. The scale region is unchanged — it was already viewed as torch.float32 on the next line.

Validation

Byte-equivalence vs proven-correct path

Same input K (bf16), block=64, head_dim=128:

Block	BUGGY (uint8) vs C+++shuffle_weight	FIXED (fp8 view) vs C+++shuffle_weight
0	8138 / 8192 differ (99.3%)	0 / 8192 differ
1	8146 / 8192 differ (99.4%)	0 / 8192 differ
2	8141 / 8192 differ (99.4%)	0 / 8192 differ

The FIXED path produces byte-identical output to aiter.ops.cache.indexer_k_quant_and_cache followed by aiter.ops.shuffle.shuffle_weight — the exact path used by AITER's official benchmark bench_deepgemm_attention.py.

End-to-end logits parity (MI355X, single GPU)

Setup: heads=64, hd=128, block=64, Preshuffle=True, KVBlockSize=64, ChunkK=256, ue8m0 scales, single batch, single next_n. Reference: bf16-derived torch ground truth with the same per-128-group ue8m0 quantization the kernel uses.

ctx	BUGGY topk@2048	FIXED topk@2048	FIXED max_rel
128	1.000	1.000	1.0e-3
512	1.000	1.000	8.8e-3
1024	1.000	1.000	7.7e-3
2048	1.000	1.000	3.4e-2
2049	0.814	1.000	1.3e-2
4096	0.002	1.000	3.3e-2
8192	0.002	1.000	1.5e-2
16384	0.005	1.000	3.6e-1

The 2049 boundary is the SplitKV/multi-block transition where the BUGGY path's accidental top-K agreement breaks down — exactly matching the user-observed MTP regression threshold.

Production-wrapper validation

Same harness, but mounting the patched rocm_aiter_mla_sparse.py over the original via docker -v and calling indexer_k_quant_and_cache_triton(...) (not the kernel directly):

ctx=  128  PATCHED_WRAPPER: bad/128=     0  topk=1.0000  max_rel=1.030e-03
ctx=  512  PATCHED_WRAPPER: bad/512=     0  topk=1.0000  max_rel=8.782e-03
ctx= 1024  PATCHED_WRAPPER: bad/1024=     0  topk=1.0000  max_rel=7.737e-03
ctx= 2048  PATCHED_WRAPPER: bad/2048=     0  topk=1.0000  max_rel=3.392e-02
ctx= 2049  PATCHED_WRAPPER: bad/2049=     0  topk=1.0000  max_rel=1.272e-02
ctx= 4096  PATCHED_WRAPPER: bad/4096=     0  topk=1.0000  max_rel=3.305e-02
ctx= 8192  PATCHED_WRAPPER: bad/8192=     0  topk=1.0000  max_rel=1.465e-02
ctx=16384  PATCHED_WRAPPER: bad/16384=     1  topk=1.0000  max_rel=3.641e-01

Cross-LLM verification

Verified the bug diagnosis and the fix with a 3-of-3 LLM swarm (gpt-5.1 / Claude-Haiku-4.5 / GPT-oss-20B over the AMD LLM gateway) — unanimous TRUE / HIGH confidence on both.

Why is this triton path used in production at all?

DSv3.2 sparse-MLA on MI355X populates the indexer KV cache via indexer_k_quant_and_cache_triton (selected over the C++ producer in the ROCm path). All non-kernel-side workarounds (e.g. an externally applied aiter.ops.shuffle.shuffle_weight after the producer) are equivalent in correctness but require either an extra pass over the K bytes per token or a global flag that becomes inconsistent on subsequent token writes. Fixing the producer is the cheapest and most local correctness fix; the kernel itself is unchanged.

Test plan

• Run scripts/test_e2e_v3.py — BUGGY vs FIXED vs CPP+SHUFFLE end-to-end logits parity for ctx ∈ {128, 512, 1024, 2048, 2049, 4096, 8192, 16384}
• Run scripts/test_byte_equiv.py — byte equivalence with C++ producer + shuffle_weight
• Run scripts/test_patched_wrapper.py — same harness using the patched production wrapper via Docker volume mount
• Run upstream vllm/tests/kernels/attention/test_deepgemm_attention.py once on a CUDA host (the test does not exercise this code path directly, but should still pass)
• DSv3.2 + MTP end-to-end serving on MI355X (1K input / 100 output, TP4) — confirm GSM8K is back to baseline (~0.9424)

Related

• Introduced in ab4b563 [ROCm][DeepSeek-V3.2][Perf] Enable gluon preshuffle indexer (block_size=64 + SHUFFLE layout)
• Companion memory-safety patch in AITER: [gluon][pa_mqa_logits] memory-safety: mask all OutLogits buffer_store lanes (does NOT fix long-context accuracy) ROCm/aiter#2936 (draft, OOB stores in pa_mqa_logits.py — distinct issue from this accuracy fix)

Signed-off-by: Marko Hartikainen marko.hartikainen@amd.com

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀