Sparse MLA decode (V3.2 / FP8 KV): throughput cost of the B200 accuracy fix (5aa668c) scales with topk

# Sparse MLA decode (V3.2 / FP8 KV): throughput cost of the B200 accuracy fix (5aa668c) scales with `topk`

Hi FlashMLA team — first, thanks for the excellent kernels.

While validating the sparse MLA decode path (DeepSeek-V3.2, FP8 KV cache) on **B300**, we measured the throughput impact of the recent accuracy fix
[`5aa668c` — "Fix flashmla_kv nsa backend accuracy issue on B200" (#5)](https://github.com/deepseek-ai/FlashMLA/commit/5aa668c9efd514ec7f4363a708c9cadcd71d1e36)
and found a **measurable, reproducible latency cost that grows with `topk`** (up to ~7%). We wanted to share the numbers and a self-contained benchmark for your reference.

## What the fix changes

For the `ModelType::V32` path, `5aa668c` raises the per-token scale precision and halves the index/scale prefetch depth to pay for it (`csrc/sm100/decode/head64/config.h`):

```cpp
// scale_t: e8m0 (1 byte) -> bf16 (2 bytes) for V32
using scale_t = std::conditional_t;

// NUM_INDEX_BUFS: 4 -> 2 for V32 (buffers for tma_coords / is_token_valid / scales)
static constexpr int NUM_INDEX_BUFS = MODEL_TYPE == ModelType::V32 ? 2 : 4;
```

The accuracy improvement is clear and worthwhile. The question is purely about the **throughput** side effect of the shallower prefetch (`NUM_INDEX_BUFS 4 -> 2`).

## Measurements (B300 SXM6, production decode shapes)

Median latency over 100 timed iters (20 warmup), V3.2 geometry `h_q=128, d_qk=576, d_v=512, block_size=64`. "before" = `5aa668c^` (parent), "after" = `5aa668c`. Both builds verified numerically correct against a full-precision bf16 reference (cosine ≥ 0.994 in every config).

| batch | seqlen | topk | before (µs) | after (µs) | Δ latency |
|------:|-------:|-----:|------------:|-----------:|----------:|
| 64 | 8192 | 128 | 27.5 | 27.3 | −0.7 % |
| 128 | 8192 | 128 | 29.5 | 29.4 | −0.3 % |
| 64 | 16384 | 512 | 52.1 | 54.0 | +3.6 % |
| 128 | 16384 | 512 | 50.0 | 50.6 | +1.2 % |
| 64 | 32768 | 2048 | 91.2 | 95.0 | +4.2 % |
| 128 | 32768 | 2048 | 136.0 | 146.1 | +7.4 % |
| 256 | 16384 | 2048 | 267.1 | 281.1 | +5.2 % |

The regression is **driven by `topk`**: ~0 % at `topk=128`, growing to **+5–7 % at `topk=2048`**. We isolated this with one-dimension-at-a-time sweeps — varying `seqlen` alone has no effect, and `batch` is a mild secondary factor; `topk` is the dominant driver. This is consistent with the `NUM_INDEX_BUFS 4 -> 2` change, since the index/scale prefetch depth matters more as more KV tokens are gathered per query.

Each data point was repeated several times; the gap is stable to well under 1 % run-to-run, so this is a real effect rather than measurement noise.

## Reproduce

The benchmark is a single self-contained script (attached below / `bench_sparse_mla_decode.py`). It times `flash_mla_with_kvcache` on the production decode shapes and validates each output against a bf16 reference:

```bash
python bench_sparse_mla_decode.py # full production sweep
python bench_sparse_mla_decode.py --batch 128 --seqlen 32768 --topk 2048
```

To get the before/after numbers, build at `5aa668c^` and `5aa668c` respectively and run the same sweep.

Environment:

- GPU: NVIDIA B300 SXM6 (compute capability 10.3), driver 610.43.02
- CUDA 13.0 (nvcc V13.0.88)
- PyTorch 2.11.0+cu130
- flashinfer 0.6.12 (bf16 correctness reference)

bench_sparse_mla_decode.py (click to expand)

```python
#!/usr/bin/env python
"""
Sparse MLA decode benchmark (DeepSeek-V3.2 / FP8 KV cache).

Measures the latency / throughput of FlashMLA's `flash_mla_with_kvcache` for
token-level sparse decoding. Output correctness is validated against a
full-precision bf16 reference computed with TensorRT-LLM's
`trtllm_batch_decode_with_kv_cache_mla` (flashinfer trtllm-gen backend).

FLOP accounting follows tests/lib.py::count_flop_and_mem_vol_for_decode:
FLOP = 2 * h_q * (b * s_q * topk) * (d_qk + d_v)

Usage:
python bench_sparse_mla_decode.py # full production sweep
python bench_sparse_mla_decode.py --batch 128 --topk 2048 --seqlen 32768
"""

import argparse
import os
import sys

import torch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests"))
import quant # noqa: E402 (tests/quant.py: FP8 V3.2 KV-cache layout helpers)

from flash_mla import flash_mla_with_kvcache, get_mla_metadata # noqa: E402

# V3.2 model geometry
H_Q = 128
H_KV = 1
D_QK = 576 # 512 NoPE (kv_lora_rank) + 64 RoPE (qk_rope_head_dim)
D_V = 512
QK_NOPE_HEAD_DIM = 128
KV_LORA_RANK = 512
QK_ROPE_HEAD_DIM = 64
BLOCK_SIZE = 64 # page block size
B_TOPK = 64 # kernel requires topk % B_TOPK == 0
SM_SCALE = D_QK ** -0.55 # matches tests/lib.py::generate_testcase_for_decode

# (batch, seqlen, topk) from tests/test_flash_mla_sparse_decoding.py production cases
PRODUCTION_CONFIGS = [
(64, 8192, 128),
(128, 8192, 128),
(64, 16384, 512),
(128, 16384, 512),
(64, 32768, 2048),
(128, 32768, 2048),
(256, 16384, 2048),
]

def build_inputs(b, seqlen, topk, dtype=torch.bfloat16):
"""Construct shared inputs for both kernels.

KV cache is sized for `b` INDEPENDENT sequences (num_blocks = ceil(b*seqlen
/ block_size)), matching flashinfer/benchmarks/bench_trtllm_gen_mla.py and
tests/lib.py: ~b*seqlen tokens, so reads hit HBM rather than staying L2-
resident (a single shared sequence fits in L2 and inflates throughput).

Returns a dict with:
q : [b, 1, h_q, d_qk] (shared)
k_fp8 : FlashMLA FP8 656-byte blocked KV cache (FlashMLA)
indices : [b, 1, topk] int32, physical slot indices (shared)
kv_trtllm : [num_blocks, 1, block_size, d_qk] bf16 (bf16 reference)
seq_lens : [b] int32 (reference)
The same physical `indices` select the same tokens in every KV layout.
"""
num_blocks = (b * seqlen + BLOCK_SIZE - 1) // BLOCK_SIZE
blocks_per_seq = (seqlen + BLOCK_SIZE - 1) // BLOCK_SIZE

q = torch.randn(b, 1, H_Q, D_QK, dtype=dtype, device="cuda").clamp_(-1, 1)

# bf16 KV, then quantize -> FP8 (FlashMLA) and dequantize -> bf16 (trtllm),
# so both kernels see the *same* (quantization-rounded) K values.
blocked_k = (torch.randn(num_blocks, BLOCK_SIZE, H_KV, D_QK, dtype=dtype, device="cuda") / 10).clamp_(-1, 1)
k_fp8 = quant.quantize_k_cache(blocked_k, quant.FP8KVCacheLayout.V32_FP8Sparse)
k_bf16 = quant.dequantize_k_cache(k_fp8, quant.FP8KVCacheLayout.V32_FP8Sparse) # [nb, bs, 1, d]

# trtllm-gen 4D layout: [num_blocks, 1, block_size, kv_lora_rank+qk_rope]
# (size(-2) must be the page_size, 32 or 64). Used for the bf16 reference.
kv_trtllm = k_bf16.permute(0, 2, 1, 3).contiguous() # [nb, 1, bs, d]

# Unique random block IDs per sequence (ref bench_trtllm_gen_mla.py:53-72).
perm_blocks = torch.randperm(num_blocks, device="cuda").to(torch.int32)
block_table = perm_blocks[: b * blocks_per_seq].view(b, blocks_per_seq).contiguous()

# abs_indices: [b, 1, topk] logical token positions in [0, seqlen), -1 padded.
abs_indices = torch.full((b, 1, topk), -1, dtype=torch.int32, device="cuda")
for i in range(b):
n = min(topk, seqlen)
abs_indices[i, 0, :n] = torch.randperm(seqlen, device="cuda")[:n].to(torch.int32)
# Map logical -> physical slot in the b-independent cache via the block_table.
indices = quant.abs_indices2indices_in_kvcache(abs_indices, block_table, BLOCK_SIZE)

seq_lens = torch.full((b,), seqlen, dtype=torch.int32, device="cuda")

return {
"q": q, "k_fp8": k_fp8, "indices": indices,
"kv_trtllm": kv_trtllm, "seq_lens": seq_lens, "seqlen": seqlen,
}

def bench_ms(fn, warmup=20, iters=100):
"""Median latency (ms) over `iters` runs after `warmup` warmup runs.

Uses CUDA events so we time only kernel execution, not Python/host overhead.
"""
for _ in range(warmup):
fn()
torch.cuda.synchronize()
starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
for i in range(iters):
starts[i].record()
fn()
ends[i].record()
torch.cuda.synchronize()
times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends))
return times[len(times) // 2]

def flop_count(b, topk):
# FLOP = 2 * h_q * (b * s_q * topk) * (d_qk + d_v).
s_q = 1
return 2 * H_Q * (b * s_q * topk) * (D_QK + D_V)

# ---------------------------------------------------------------------------
# Implementations
# ---------------------------------------------------------------------------

def run_flashmla(inp):
q, k_fp8, indices = inp["q"], inp["k_fp8"], inp["indices"]
sched_meta, _ = get_mla_metadata()

def fn():
return flash_mla_with_kvcache(
q=q, k_cache=k_fp8, block_table=None, cache_seqlens=None,
head_dim_v=D_V, tile_scheduler_metadata=sched_meta, num_splits=None,
softmax_scale=SM_SCALE, causal=False, is_fp8_kvcache=True, indices=indices,
)

out, _ = fn()
return out, bench_ms(fn)

_TRTLLM_WS = None

def run_reference(inp):
"""Full-precision bf16 sparse MLA decode (flashinfer trtllm-gen).

Used only as a correctness oracle for FlashMLA's fp8 output; both consume
the same (quantization-rounded) K values and the same physical `indices`.
"""
global _TRTLLM_WS
from flashinfer.mla import trtllm_batch_decode_with_kv_cache_mla

q = inp["q"] # [b, 1, h_q, d_qk]
indices = inp["indices"] # [b, 1, topk] physical slot indices
seq_lens = inp["seq_lens"] # [b]
b = q.shape[0]
topk = indices.shape[-1]

if _TRTLLM_WS is None:
_TRTLLM_WS = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device="cuda")

out = trtllm_batch_decode_with_kv_cache_mla(
query=q,
kv_cache=inp["kv_trtllm"],
workspace_buffer=_TRTLLM_WS,
qk_nope_head_dim=QK_NOPE_HEAD_DIM,
kv_lora_rank=KV_LORA_RANK,
qk_rope_head_dim=QK_ROPE_HEAD_DIM,
block_tables=indices, # sparse indices in sparse_mla_top_k mode
seq_lens=seq_lens,
max_seq_len=inp["seqlen"],
sparse_mla_top_k=topk,
bmm1_scale=SM_SCALE / (D_QK ** 0.5),
bmm2_scale=1.0,
)
return out.view(b, 1, H_Q, D_V) if out.dim() != 4 else out

def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--batch", type=int, default=None)
ap.add_argument("--seqlen", type=int, default=None)
ap.add_argument("--topk", type=int, default=None)
args = ap.parse_args()

if args.batch or args.seqlen or args.topk:
configs = [(args.batch or 128, args.seqlen or 32768, args.topk or 2048)]
else:
configs = PRODUCTION_CONFIGS
for _, _, tk in configs:
assert tk % B_TOPK == 0, f"topk ({tk}) must be a multiple of {B_TOPK}"

import subprocess
gpu = subprocess.run(["nvidia-smi", "--query-gpu=name,compute_cap",
"--format=csv,noheader"], capture_output=True, text=True).stdout.strip().split("\n")[0]
print(f"GPU: {gpu}")
print(f"Model: V3.2 (h_q={H_Q}, d_qk={D_QK}, d_v={D_V}), block_size={BLOCK_SIZE}\n")

rows = []
for b, seqlen, topk in configs:
inp = build_inputs(b, seqlen, topk)
flop = flop_count(b, topk)
ref = run_reference(inp).float() # bf16 correctness oracle
out, ms = run_flashmla(inp)
out = out.float()
max_err = (out - ref).abs().max().item()
cos = torch.nn.functional.cosine_similarity(
out.flatten(), ref.flatten(), dim=0).item()
rows.append({"b": b, "seqlen": seqlen, "topk": topk,
"us": ms * 1e3, "tflops": flop / 1e12 / (ms / 1e3),
"max_err": max_err, "cos": cos})

print(f"{'batch':>6}{'seqlen':>8}{'topk':>6}{'us':>10}{'TFLOPS':>10}"
f"{'max_err':>12}{'cos':>10}")
print("-" * 66)
for r in rows:
print(f"{r['b']:>6}{r['seqlen']:>8}{r['topk']:>6}{r['us']:>10.1f}"
f"{r['tflops']:>10.1f}{r['max_err']:>12.4f}{r['cos']:>10.5f}")

if __name__ == "__main__":
if "CUDA_VISIBLE_DEVICES" not in os.environ:
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
torch.set_default_device("cuda:0")
torch.cuda.set_device(0)
main()
```