Repository navigation
[ROCm][DSv4.1] Support the NVFP4 compressed KV cache on gfx950 - #57463
Open
JohnQinAMD wants to merge 1 commit into
Open
JohnQinAMD wants to merge 1 commit into
JohnQinAMD wants to merge 1 commit into
Conversation
DeepSeek-V4.1 stores its compressed KV as NVFP4 (288 B/token: 512 e2m1 values packed two per byte, then 32 e4m3 scales of 16 dims each) next to an fp8 sliding-window cache. On ROCm neither end of that record worked: the pack is PTX, and reading it was gated on SM100, so `--kv-cache-dtype nvfp4_ds_mla` was rejected outright and ROCm kept the 584 B V4 record for both caches. Writing: `_fp32x2_to_fp4x2` gains a gfx950 branch. `f` is not an AMDGPU asm constraint, so the PTX pack aborted every worker at startup with `unknown asm constraint 'f'`. `v_cvt_scalef32_pk_fp4_f32` replaces it -- it takes the low nibble first, divides by the E8M0 of a third operand (1.0 leaves the value alone) and writes only the destination's low byte. Reading: an NVFP4 row loader in the gfx950 sparse decode tile, selected per cache from the record width the KV-cache spec publishes, decoding a nibble pair with `v_cvt_scalef32_pk_f32_fp4`. That instruction's scale operand is E8M0 too, so the e4m3 scale is applied separately, and its destination is a register pair, so the result is taken as one 64-bit value whose halves are the even and odd dim of the byte. A layer whose kernel can read the record now opts in rather than the gate testing for SM100 alone. The sliding-window record and the fp8 path are unchanged. Measured on MI355X (gfx950). The pack matches `cvt.rn.satfinite.e2m1x2.f32` byte for byte over every level, both signed zeros, the exact ties and the saturating ends. Decode, at the shape this model runs on 4 ranks (16 local heads, index_topk 512, sliding window 128) captured in a HIP graph with prebuilt ragged indices, stays within 1% of the 584 B record from 1 to 64 decode tokens, for a compressed cache half the size, and GSM8K over 1319 questions is 90.30% on the 584 B record against 89.92% on this one. Capacity is a separate matter: the compressed cache is ~10% of this model's per-block bytes, so the smaller record is worth +2.5% max concurrency, and realising even that needs a fix to the KV-cache grouping heuristic that is filed separately. This change was written with AI assistance (Claude Code). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: JohnQinAMD <yanyuan.qin@amd.com>
This was referenced Sep 18, 2026
3 of 4 tasks
JohnQinAMD
marked this pull request as ready for review
September 21, 2026 04:13
JohnQinAMD
requested review from
AndreasKaratzas,
WoosukKwon,
dllehr-amd,
hongxiayang,
mgoin,
tjtanaa,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
September 21, 2026 04:13
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
--kv-cache-dtype nvfp4_ds_mlacannot be selected on ROCm at all today. The pack in_fp32x2_to_fp4x2is PTX,fis not an AMDGPU asm constraint, and every worker aborts during startup withunknown asm constraint 'f'. Reading was missing too: only FlashMLA's SM100 sparse decode could decode the record, so the dtype was rejected everywhere else.The V4.1 compressed record (#56935) packs 512 e2m1 values two per byte, then 32 e4m3 scales of 16 dims each, and sits beside an fp8 sliding-window cache — 288 B/token against the 584 B V4 record.
Approach
_fp32x2_to_fp4x2gains a gfx950 branch usingv_cvt_scalef32_pk_fp4_f32. That instruction takes the low nibble first where the PTX one takes the high nibble, divides by the E8M0 of a third operand (1.0 leaves the value alone), and writes only the destination's low byte._load_v41_nvfp4_gfx950_chunkdecodes a nibble pair withv_cvt_scalef32_pk_f32_fp4. That instruction's scale operand is E8M0 too, so this record's e4m3 scale is applied separately, and its destination is a register pair, so the result comes back as one 64-bit value whose halves are the even and odd dim of the byte.reads_nvfp4_compressed_cachelets a layer whose kernel can decode the record opt in, replacing a gate on SM100 alone. gfx942 still rejects the dtype, now with an accurate message.The PTX pack has two other callers — the indexer Q quantization and the V4 indexer's MXFP4 K cache — so the same fix unblocks those write paths on ROCm as well, though neither is reachable yet for other reasons.
Test plan
The decode test packs known rows into each record, drives
rocm_sparse_attn_decode, and compares against a torch softmax over the same dequantized rows, parametrized over the record pairs the model allocates so the V4 arm is the regression case. The pack test checks the gfx950 byte againstcvt.rn.satfinite.e2m1x2.f32over every level and its negation, both signed zeros, the exact ties, and values past the ends of the format.Test result
MI355X (gfx950): tests pass, the pack is byte-exact, cosine 0.999998 against the torch reference for the NVFP4 arm and unchanged for the fp8 arm.
Decode latency at the shape this model runs on 4 ranks — 16 local heads (
num_attention_heads64 // TP 4),index_topk512, sliding window 128 — captured in a HIP graph with prebuilt ragged indices:Half the compressed-cache bytes at the same decode cost, within 1% across the range.
Accuracy. GSM8K over all 1319 questions, 5-shot, TP=4: 90.30% on the 584 B record against 89.92% here, inside the 0.83 pt standard error at that sample size.
Capacity.
fp8_ds_mlais untouched, so nothing regresses — but the new dtype measures below it as things stand, 123.31x against 135.58x max concurrency at TP=4 and 65,536 context. The compressed cache is only ~10% of this model's per-block bytes (the sliding-window cache is 78%), and the smaller page also shrinks_get_packed_kv_cache_groups' anchor, which splits the unrelated state buckets from 9 groups into 17. That is a pre-existing, platform-independent issue in the grouping heuristic rather than something this PR introduces.What makes the record pay off is narrowing the sliding-window cache alongside it. With the V4.1 MXFP8 record in place the pair reaches 11,020,825 tokens and 168.16x, +24.0% over the V4 baseline.
Known follow-up
nvfp4_ds_mlakeeps a wider sliding-window record beside the 288 B one.AI assistance (Claude Code) was used in preparing this change.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.