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

// scale_t: e8m0 (1 byte) -> bf16 (2 bytes) for V32
using scale_t = std::conditional_t<MODEL_TYPE == ModelType::V32, bf16, e8m0>;

// 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;

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

#!/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()