Repository navigation
[ROCm][DSv4] Fix sparse-indexer logits collapse on gfx950/gfx942 - #50455
Conversation
On gfx950, AITER's deepgemm_fp8_paged_mqa_logits mis-scores the DeepSeek-V4 sparse indexer during decode (paged MQA cache read with the wrong layout), collapsing long-context retrieval (RULER niah_single_1 ~60% at 4K -> ~8% at 32K). Behind VLLM_DSV4_LOGITS_FIX=1 (off by default), reroute decode and MTP paged MQA-logits through a Triton kernel that reads the fp8 KV cache in its true block-flat layout and computes relu(q.k)*weight*scale per key position. Each query's key range is split across programs so there is no cross-program reduction or host sync, keeping it capturable in a full CUDA graph. The write range is clamped to max_model_len and the block-table gather is bounds-guarded. Validated on DeepSeek-V4-Flash (TP4, fp8 KV, gfx950): RULER long-context retrieval recovered across 4K-32K, GSM8K 20-shot +3.4pts flexible-extract, throughput within ~1.6% at long-context decode, verified under MTP (num_speculative_tokens=3). Kernel matches the torch reference to 1.4e-6. Signed-off-by: Frida Andersson <fanderss@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Scale the decode kernel's KV-range split count as min(256, tiles, 1024 / rows) instead of a fixed 64/32, and pipeline the launch (num_warps=4, num_stages=2). All terms are static at launch, so the grid stays fixed under CUDA graph capture. The kernel is memory-bound over the KV range; a fixed split count leaves most CUs idle at low batch / long context. This is 1.7-3.2x faster in isolation and ~3-5% lower TPOT at concurrency 1, collapsing back to 32 splits (identical to before) at high batch, so no regression there. Output is bit-identical (~1e-7 vs torch ref); decode and MTP (next_n>1) both verified. Also drop an unused local and reformat to satisfy ruff. Signed-off-by: Frida Andersson <fanderss@amd.com>
The same paged-MQA-logits collapse reproduces on gfx942, and the Triton reroute fixes it there too: the fp8 KV-cache layout and torch reference are identical across both archs, so the kernel is unchanged and only the arch gate widens (_ON_GFX950 -> _ON_GFX950 or _ON_GFX942). Validated on 8x gfx942 (TP8): RULER NIAH recovers the long-context collapse, MTP (num_speculative_tokens=3) consistent with the non-MTP fix. Signed-off-by: Frida Andersson <fanderss@amd.com>
frida-andersson
changed the title
frida-andersson
changed the title
frida-andersson
changed the title
eab2f97 to
b6a29ad
Compare
Signed-off-by: Frida Andersson <fanderss@amd.com>
b6a29ad to
ac65d1f
Compare
Run the indexer decode QK dot on fp8 operands (fp8 MFMA, f32 accum) instead of upcasting q/kv to f32 with input_precision="ieee". fp8 upcasts exactly, so products are unchanged. Add tl.static_assert(HEAD_SIZE % 4 == 0) for the block_size*D//4 scale offset (D=128, compile-time only). Kernel 4.4-4.7x faster at decode batch 256-512 on gfx942, 8-9x on gfx950. gfx942: bit-identical to the f32 dot. gfx950: ~1e-5 rel (fp8-MFMA accum order), top-k preserved; GSM8K/RULER A/B incl. MTP=3 show no regression. Signed-off-by: Frida Andersson <fanderss@amd.com>
|
I have reviewed the code under ticket.
|
Route the gfx942/gfx950 decode/MTP paged MQA-logits path through the Triton kernel unconditionally instead of behind VLLM_DSV4_LOGITS_FIX, per review feedback to keep the OOTB experience flag-free. Signed-off-by: frida-andersson <fanderss@amd.com>
|
Thanks @tjtanaa for the review! I've addressed the comment now by removing the |
| # gfx942/gfx950: AITER's paged MQA-logits decode kernel scores the sparse | ||
| # indexer incorrectly and collapses long-context retrieval, so route the | ||
| # decode/MTP path through the Triton kernel instead. | ||
| if _ON_GFX950 or _ON_GFX942: |
There was a problem hiding this comment.
So AITER's paged MQA-logits decode kernel has issue for all cases? If we do this, we will be introducingperf regression to all models as it is going to use triton or even torch, rather than gluon kernel from aiter.
There was a problem hiding this comment.
Thanks for the comment! AITER's kernel is fine for the SHUFFLE layout DSv32 uses, the bug is DSv4-specific, where the compressor writes block-flat but deepgemm_fp8_paged_mqa_logits reads shuffled. This only hits sparse-indexer decode, not other models. The gating is too wide though, I'll tighten it so DSv32/GLM keep the current gluon path
f78e56b to
ef2eaae
Compare
f8dfd17 to
eb09a3c
Compare
|
@tjtanaa the Flash NIAH is 100% through 32K, DSv32 GSM8K stays at stock (~0.96). PTAL |
Use _decode_cu_count() so gfx942 (~304 CUs) is not limited to the gfx950 constant 256, and add a Triton vs torch-ref unit test. Signed-off-by: Frida Andersson <fanderss@amd.com>
|
@tjtanaa ping for re-review. AMD CI is red on the NIXL test _failed_recv_reqs, same failure is on today’s main nightly (build 12959) |
5db60ee to
5ff9782
Compare
|
@simondanielsson @tjtanaa - gate change since your last look: Triton/NORMAL only when compress_ratio == 4 (V4.0 C4A). V4.1-Flash is SHUFFLE at ratio 1 and 2; the old >1 check sent ratio-2 layers to the wrong reader. Verified gfx950 TP4: V4-Flash → Triton C4A, V4.1 and DSv3.2 → AITER. PTAL |
|
Hi @frida-andersson, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
V4.1-Flash writes SHUFFLE and has indexer ratios 1 and 2. compress_ratio>1 sent ratio-2 layers down the Triton block-flat reader while ratio-1 stayed on AITER. Only V4.0 C4A (ratio 4) is values-then-scales. Signed-off-by: Frida Andersson <fanderss@amd.com>
5ff9782 to
da0c592
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #89826 for commit |
6 participants
On gfx950 and gfx942, AITER's
deepgemm_fp8_paged_mqa_logitsmis-scores the DSv4 sparse indexer during decode (wrong paged-cache layout), collapsing long-context retrieval.VLLM_DSV4_LOGITS_FIX=1(off by default) reroutes decode + MTP through a Triton kernel insteadThe replacement kernel (
_fp8_paged_mqa_logits_decode_kernel) reads the fp8 KV cache in its true block-flat layout and computesrelu(q·k)·weight·scaleper key position, splitting each query's key range across programs so there is no cross-program reduction or host sync — which keeps it capturable in a full CUDA graphThe decode kernel computes its q·k dot directly in fp8 (fp8 MFMA with f32 accumulation) rather than upcasting q and k to f32 first. Against that f32-upcast path the dot is 4.4–4.7× faster on gfx942 and up to ~9× on gfx950 at serving batch sizes, for the same result — bit-identical on gfx942, and on gfx950 within a ~1e-5 rounding difference that leaves the selected keys unchanged (GSM8K, RULER and MTP all agree).
Env: vLLM
0.26.1rc1.dev124, AITER0.1.16.post5, DeepSeek-V4-Flash, TP=4, fp8 KV,block_size=256,max_model_len=40960, 4x gfx950RULER NIAH
Accuracy % as
stock → fix(50 samples/cell, temp=0, concurrency 16)At 4K the indexer prunes little, so the fix has little to correct there; niah_2
60 → 58is within run-to-run noise. The fix's effect grows with context, where stock collapses and the fix holdsAs a sanity check, I also ran niah_1 through the torch reference - the slow but known-correct version of the decode-logits calculation that the fix's Triton kernel is designed to reproduce. It recovered accuracy to 81% at 16K and 92% at 32K, essentially matching the fix's own niah_1 numbers
GSM8K
flexible-extract accuracy (1319 problems), prefix caching disabled. The gap widens with shot count: more few-shot context means the indexer prunes more, so stock's mis-ranking costs more and the fix recovers more. 5-shot is too short to engage the indexer (within noise)
Throughput
Decode throughput on the committed kernel (with the fp8 dot), prefix caching disabled, 512 output tokens,
ignore_eos, 1 warmup + 2 reps.stock= AITER fused paged-MQA-logits kernel,fix= Triton replacementThe fix runs at parity with AITER's fused kernel — differences are within run-to-run noise, since the logits step is only a few percent of decode
Kernel occupancy tuning
The kernel is memory-bound over the KV range, so the number of KV-range splits per query sets occupancy. A fixed 32 leaves most CUs idle at low batch / long context, so it scales as
min(256, tiles, 1024 / rows)(all static at launch, so still CUDA-graph-safe). Isolated kernel latency vs the fixed-32 version, output unchanged:End-to-end the effect is small — ~3-5% TPOT at concurrency 1, within noise at c8/c32 (where the split count is back to 32). No regression anywhere
gfx942 (MI300-class): same bug, same fix
The identical collapse reproduces on gfx942 and the Triton reroute fixes it there too. The kernel is unchanged — the fp8 KV-cache layout and torch reference are identical across both archs — so only the arch gate widens (
_ON_GFX950→_ON_GFX950 or _ON_GFX942). RULER NIAHstock → fix(TP8, fp8 KV, 8x gfx942, 50 samples/cell, temp=0, concurrency 16):Same signature as gfx950: stock collapses toward ~0 by 32K, the fix recovers and holds. MTP (
num_speculative_tokens=3) verified, consistent with the non-MTP fix