[RFC]: fp8 KV cache for the Ampere sparse-MLA path (TRITON_MLA_SPARSE) via software dequant

[nickus]

🚀 The feature, motivation and pitch

TRITON_MLA_SPARSE (#47629) is bf16-KV-only, and on 24 GB cards that is the binding constraint. I have a working sm_80/sm_86 fp8 KV path and would like to know whether the project wants it before I polish it into a mergeable PR.

The gap, in the PR's own words

TritonMLASparseBackend (vllm/v1/attention/backends/mla/triton_mla_sparse.py):

supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "float16", "bfloat16"]

its decode entry point is literally _forward_bf16_kv, and the base it inherits refuses quantized KV outright (xpu_mla_sparse.py:238):

if is_quantized_kv_cache(self.kv_cache_dtype):
    raise NotImplementedError("FP8 kv is not supported with XPU MLA Sparse yet")

The stated reason fp8 is excluded on Ampere is that Triton cannot handle fp8e4nv there. That is true for native fp8 — and irrelevant for fp8 as storage.

The technique is already in this very PR

vllm/v1/attention/ops/mqa_logits_triton.py:191, added by #47629 itself:

# Kernel decodes FP8 from uint8 via LUT (SM80 Triton can't load fp8e4nv).

The indexer already reads fp8 pages on SM80 by loading them as uint8 and decoding in-register. The KV path simply hasn't had the same treatment. (models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py does the same thing for XPU: "dequantize FP8 KV cache pages to BF16 on the fly, then reuse the BF16 sparse MLA attention kernel".) So the precedent is accepted in-tree, twice, and neither conflicts with #43914 / #47060's "native FP8 requires SM89+" — software dequant is a different mechanism.

What I did

An IS_FP8: tl.constexpr branch inside your own _sparse_mla_compute_tile — the single function through which both _sparse_mla_kernel_final and _sparse_mla_kernel_split load KV — so fp8 inherits your split-KV, your merge kernel, your autotune and your masking for free. The standard 656-byte fp8_ds_mla page is loaded as raw uint8 and decoded with bit-math that is bit-exact vs torch.float8_e4m3fn over all 256 byte values. Nothing about the cache layout changes; the existing C++ concat_and_cache_ds_mla writer already runs on sm_86 unmodified (ENABLE_FP8 is CUDA-version-gated, not arch-gated — verified by execution).

Patch (224 lines, applies to bbe2ab4d6, one flag / one branch / one dispatcher arg): https://github.com/nickus/vllm-fp8kv/blob/main/patches/triton_mla_sparse_fp8.patch

Measured on RTX 3090 (sm_86) — including the part that is not good news

fp8 dequant, all 256 byte values	bit-exact vs torch.float8_e4m3fn
Decode vs fp32 reference, through vLLM's own C++ ds_mla writer and own triton_convert_req_index_to_global_index	cosine 0.999996
KV pool capacity	1.756× (656 B vs 1152 B/token); ~1.63× effective once the indexer's own K-cache is counted
Decode speed vs bf16, same kernel, bs=8 / topk=2048	0.92×
Decode speed vs bf16, same kernel, bs=32 / topk=2048	0.45×

So this is a capacity-for-speed trade, not a free win, and I would rather say so up front than have it found in review. At bs=32 fp8 decode is currently ~2.2× slower. My working hypothesis is that the autotune config lists (_FINAL_AUTOTUNE_CONFIGS: BLOCK_N=16 only; _SPLIT_AUTOTUNE_CONFIGS: BLOCK_N=32 only) are tuned for bf16's 1152 B rows and that fp8's 656 B rows plus a scale load want larger blocks — but I have not measured that on the real kernel yet, so I am flagging it as a hypothesis, not a finding. On a 24 GB card the 1.76× capacity may well be worth the decode cost; on an A100 it plainly is not, which is an argument for keeping this opt-in.

What is NOT done

• The patch above is kernel-only and therefore inert on its own: the backend still declares bf16-only dtypes, get_kv_cache_shape returns a 576-wide row for every dtype (fp8_ds_mla needs 656 — flashmla_sparse.py already handles this), and _canonicalize_sparse_mla_kv_cache_dtype maps fp8_e4m3 -> fp8_ds_mla only for FLASHMLA_SPARSE / SM120. I have that wiring working out-of-tree; it is ~40 more lines.
• No engine-level boot yet. All correctness numbers above are kernel-level. I have not yet run LLM(..., kv_cache_dtype="fp8_e4m3") end-to-end and compared logits against a bf16 run.

The questions

1. Does the project want fp8 KV on the Ampere sparse-MLA path, given it buys ~1.76× KV capacity at a decode cost that is currently significant at larger batch?
2. If yes: kernel-only PR first (the diff above, dispatch dormant until the backend declares the dtype), or one PR including the backend wiring?
3. Any objection to the fp8-as-storage / software-dequant approach in the shared _sparse_mla_compute_tile, given mqa_logits_triton.py already does exactly this on SM80?

Happy to do the engine boot + logits parity and the config-list measurement before opening the PR — I just want to know the feature is wanted first.

---

Unrelated bug found while working on this path, already filed: #48364 / #48366 (fully-masked leading index chunks NaN-poison xpu_mla_sparse).

Alternatives considered

• Native fp8 (fp8e4nv) on Ampere — impossible; that is the constraint the PR itself documents.
• Dequantize whole pages to bf16 before the kernel — what the XPU path does; costs a full bf16 materialization of the gathered KV, which throws away the bandwidth saving that is the entire point on a 24 GB card.
• Keep bf16 KV and shrink the model instead — what 3090/4090 owners do today; it is why GLM-5.2-class DSA models are effectively unusable at long context on consumer Ampere.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the Discord users if needed.

[nickus]

Update: it boots end-to-end, and my autotune hypothesis was wrong

Follow-up on the two open items I listed above. Both are now closed, on hardware (RTX 3090, sm_86, driver 595, vLLM 0.25.0 nightly + this PR at bbe2ab4d6).

1. Engine-level: fp8 KV now boots and matches bf16

Two real engines on a GlmMoeDsa model, identical inputs, only --kv-cache-dtype differing:

	bf16 baseline	fp8_e4m3
backend selected	TRITON_MLA_SPARSE	TRITON_MLA_SPARSE
KV row	576 B	656 B
sparse-MLA decode calls	16 (bf16 path)	16 (fp8 path), 0 bf16
teacher-forced logits vs bf16	—	cosine 0.999999 / 0.999997
teacher-forced top-1 agreement	—	0.80 / 1.00

(Parity is measured teacher-forced. Free-running greedy tokens are not a valid comparison here — on a random-weight toy the logits are near-flat over a 50k vocab, so a ~2 % fp8 perturbation flips a near-tie and the two runs then score different sequences.)

Three things outside the kernel had to change to get an engine up. They may be of interest independent of this RFC, since the first two look like general gaps in the fp8-sparse-MLA path rather than anything specific to Ampere:

1. The backend is selected with the raw --kv-cache-dtype, and canonicalized only afterwards (mla_attention.py: get_attn_backend(...) precedes _canonicalize_sparse_mla_kv_cache_dtype(...)). A backend that declares only fp8_ds_mla is therefore never selected for --kv-cache-dtype fp8_e4m3 — the engine dies with "No valid attention backend found … TRITON_MLA_SPARSE: [kv_cache_dtype not supported]". FLASHMLA_SPARSE sidesteps this by also declaring "fp8" as an alias; I did the same.

2. MLAAttention.get_kv_cache_spec builds MLAAttentionSpec without kv_quant_mode, so it defaults to KVQuantMode.NONE. The worker's reshape path then does layer_cache_dtype = "auto" if spec.kv_quant_mode == NONE else cache_dtype and asks the backend for the unquantized shape — while the pool itself was allocated from MLAAttentionSpec.page_size_bytes, which does special-case fp8_ds_mla at 656 B/token. Result:

   RuntimeError: shape '[128811, 64, 576]' is invalid for input of size 5408001024
   

   (5408001024 = 128811 × 64 × 656.) I fixed it by declaring the quant mode on the spec. If FLASHMLA_SPARSE + fp8_ds_mla goes through the same _reshape_attention_kv_cache path on Hopper, I would expect it to hit this too — worth a look.

3. EngineCore runs out-of-process, so an out-of-tree monkeypatch needs a plugin entry point / sitecustomize. Not upstream's problem; noting it for anyone reproducing.

2. My autotune-config hypothesis was wrong — retracting it

Above I speculated that fp8 might be slow because the autotune config lists (BLOCK_N=16 final / 32 split) are bf16-shaped and fp8's 656-byte rows want bigger blocks. I measured it, and it is false. At bs=32, topk=2048, 64K pool, forcing each config:

BLOCK_N	warps	stages	fp8	bf16
16	4	3	1415 µs	1210 µs
32	4	3	1546 µs	1088 µs
64 / 128	any	any	fails to compile — shared-memory OOM on sm_86	—

BLOCK_N ≥ 64 does not fit sm_86's 100 KB shared-memory budget at all, so "give fp8 bigger blocks" is not available on this hardware. fp8's best config is essentially the one you already ship. The config lists are fine; please disregard that part of my original post.

3. Where that leaves the numbers (correctness-gated at every point: cosine ≥ 0.999996, max_abs ≤ 2.2e-4)

fp8 vs bf16 in your kernel, only the KV dtype differing:

context	bs	fp8/bf16
8 K	1	0.93×
16 K	1	0.86×
64 K	1	0.84×
128 K	1	0.87×
64 K	32	0.88×

So: ~1.76× KV capacity (≈1.63× once the indexer's own K-cache is counted) for ~10–15 % of decode speed. That is the whole trade, and it is not tunable away. On a 24 GB card serving long-context DSA models the alternative to fp8 KV is not "faster decode", it is "does not fit" — but on an A100 the calculus is clearly different, which argues for keeping it opt-in.

The questions from the original post stand: is this wanted, and if so, kernel-only PR first or kernel + backend wiring in one? Happy to send either.

[nickus]

Correction to my table above: the bf16 "KV row" figure is wrong (units)

I wrote KV row | 576 B | 656 B. As written that says fp8 rows are larger than bf16 ones — the opposite of what I measured, and it would make this RFC pointless. 576 is the element count per token per layer (512 kv_lora_rank + 64 qk_rope_head_dim), not bytes.

Per MLAAttentionSpec.real_page_size_bytes:

	bytes / token / layer
bf16	num_kv_heads(1) * head_size(576) * 2 = 1152 B
fp8_ds_mla	656 B (the V3.2 custom layout)

1152 / 656 = 1.76×, which is the capacity figure I quote lower down, and which the 5408001024 = 128811 * 64 * 656 allocation confirms. Correcting it here rather than silently editing the comment.

[ehfd]

We need #47629 merged first.