From 63b9931b7399c9bfa7a8f7cdd60e2991ba87fbf9 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Date: Tue, 29 Sep 2026 01:37:57 -0500 Subject: [PATCH] [ROCm][DSv4.1] Paged MXFP4 sparse indexer on aiter's MQA-logits kernel (#58671) Signed-off-by: Mehmet Cagri Kaymak Co-authored-by: Claude Opus 5.5 Co-authored-by: Shanshan Shen <467638484@qq.com> --- .../test_rocm_paged_mxfp4_indexer.py | 599 ++++++++++++++++ .../test_rocm_paged_mxfp4_indexer_plan.py | 162 +++++ vllm/config/attention.py | 5 +- .../layers/rocm_paged_mxfp4_indexer.py | 116 +++ vllm/models/deepseek_v41/amd/rocm.py | 150 +++- vllm/models/deepseek_v41/attention.py | 40 +- vllm/v1/attention/backends/mla/indexer.py | 24 +- .../backends/mla/rocm_paged_mxfp4_indexer.py | 513 +++++++++++++ .../attention/ops/rocm_paged_mxfp4_indexer.py | 676 ++++++++++++++++++ 9 files changed, 2276 insertions(+), 9 deletions(-) create mode 100644 tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py create mode 100644 tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py create mode 100644 vllm/model_executor/layers/rocm_paged_mxfp4_indexer.py create mode 100644 vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py create mode 100644 vllm/v1/attention/ops/rocm_paged_mxfp4_indexer.py diff --git a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py new file mode 100644 index 0000000000..311ee89a8d --- /dev/null +++ b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py @@ -0,0 +1,599 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1's ROCm MXFP4 sparse indexer against a torch reference. + +The indexer K cache is a view into vLLM's block-major KV pool: its pages are +a block stride apart with other layers' pages in between, so the tests put +it there. The reference reads the keys back from a second cache the same +writer fills in natural order, so it shares none of the preshuffled +addressing it checks. +""" + +import itertools +import types + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm-only", allow_module_level=True) + +from vllm.config import CUDAGraphMode +from vllm.models.deepseek_v41.amd.rocm import ( + rocm_mxfp4_indexer_k_store, + rocm_mxfp4_indexer_q_quant, +) +from vllm.platforms.rocm import get_cdna_version +from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerPrefillChunkMetadata, +) +from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( + RocmMxfp4IndexerMetadata, + native_decode, + plan_gather_launches, + plan_prefill_chunks, +) +from vllm.v1.attention.ops import rocm_paged_mxfp4_indexer as ops +from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + +if get_cdna_version() != 4: + pytest.skip("CDNA4-only", allow_module_level=True) + +HEADS, HEAD_DIM = 32, 128 +WIDTH = HEAD_DIM // 2 + HEAD_DIM // 32 +# V4.1's 8-token blocks, with the pool (2048 blocks) and top-k (512) scaled +# down so short test contexts still select. +CAND_BLOCK, POOL_BLOCKS, TOPK = 8, 16, 32 +MAX_LEN = 4096 +DEVICE = "cuda" + +_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) +E2M1 = torch.cat([_E2M1, -_E2M1]) + + +def _dequant(packed, e8m0): + """[..., D/2] e2m1 pairs and [..., D/32] e8m0 to [..., D] fp32.""" + table = E2M1.to(packed.device) + vals = torch.stack([table[(packed & 0xF).long()], table[(packed >> 4).long()]], -1) + scale = torch.pow(2.0, e8m0.float() - 127) + return (vals.flatten(-2).unflatten(-1, (-1, 32)) * scale[..., None]).flatten(-2) + + +@pytest.fixture(autouse=True) +def _workspace(): + init_workspace_manager(torch.device(DEVICE)) + yield + reset_workspace_manager() + + +class _Case: + """Requests with their keys written into a block-major pool (ratio-2 and + ratio-1 pages side by side in every block) and into natural-order caches + the reference reads. The bf16 inputs are kept for the end-to-end test.""" + + def __init__(self, seq_lens, block): + from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import ( + k_norm_rope_mxfp4_cache, + ) + + self.seq_lens = seq_lens + num_blocks = sum(cdiv(n, block) for n in seq_lens) + 3 + perm = torch.randperm(num_blocks) + width = max(cdiv(n, block) for n in seq_lens) + self.block_table = torch.zeros(len(seq_lens), width, dtype=torch.int32) + used = 0 + for req, n in enumerate(seq_lens): + self.block_table[req, : cdiv(n, block)] = perm[used : used + cdiv(n, block)] + used += cdiv(n, block) + self.block_table = self.block_table.to(DEVICE) + + def align(n): + return cdiv(n, 576) * 576 + + page2, page1 = block // 2 * WIDTH, block * WIDTH + stride = align(page2) + align(page1) + 576 + self.pool = torch.zeros(num_blocks * stride, dtype=torch.uint8, device=DEVICE) + self.page_bytes = torch.zeros_like(self.pool, dtype=torch.bool) + for i in range(num_blocks): + base = i * stride + self.page_bytes[base : base + page2] = True + base += align(page2) + self.page_bytes[base : base + page1] = True + self.cache = { + 2: torch.as_strided( + self.pool, (num_blocks, block // 2, WIDTH), (stride, WIDTH, 1), 0 + ), + 1: torch.as_strided( + self.pool, + (num_blocks, block, WIDTH), + (stride, WIDTH, 1), + align(page2), + ), + } + self.offsets = [0] + for n in seq_lens: + self.offsets.append(self.offsets[-1] + n) + # the same keys in natural order, one entry per page + self.natural = { + r: torch.zeros(self.offsets[-1], 1, WIDTH, dtype=torch.uint8, device=DEVICE) + for r in (1, 2) + } + self.cos_sin = cos_sin = torch.randn(MAX_LEN, 64, device=DEVICE) + self.norm = norm = torch.rand(HEAD_DIM, device=DEVICE) + 0.5 + self.k_pre = {} + for req, n in enumerate(seq_lens): + if n == 0: + continue + k_pre = torch.randn(n, HEAD_DIM, device=DEVICE).to(torch.bfloat16) + self.k_pre[req] = k_pre + pos = torch.arange(n, device=DEVICE) + for ratio, cache in self.cache.items(): + comp = pos // ratio + boundary = (pos + 1) % ratio == 0 + entries = cache.shape[1] + page = self.block_table[req, comp // entries].long() + slot = torch.where(boundary, page * entries + comp % entries, -1) + nat = torch.where(boundary, self.offsets[req] + comp, -1) + rocm_mxfp4_indexer_k_store( + k_pre, + pos, + cos_sin, + norm, + 1e-6, + cache, + slot, + ratio, + True, + num_heads=HEADS, + ) + k_norm_rope_mxfp4_cache( + k_pre, pos, cos_sin, norm, 1e-6, self.natural[ratio], nat, ratio + ) + + def keys(self, ratio, req, n): + nat = self.natural[ratio][self.offsets[req] : self.offsets[req] + n, 0] + return _dequant(nat[:, : HEAD_DIM // 2], nat[:, HEAD_DIM // 2 :]) + + +def _queries(num_rows): + q = torch.randint(0, 256, (num_rows, HEADS, HEAD_DIM // 2), dtype=torch.uint8) + scales = torch.randint(122, 130, (num_rows, HEADS, HEAD_DIM // 32)) + q, scales = q.to(DEVICE), scales.to(torch.uint8).to(DEVICE) + weights = torch.randn(num_rows, HEADS, device=DEVICE) + # the fused Q kernel returns the scales as one int32 per head + return q, scales.view(torch.int32).squeeze(-1), weights, _dequant(q, scales) + + +def _scores(q, w, keys): + return (torch.relu(q @ keys.T) * w[:, None]).sum(0) + + +def _assert_topk(row, scores, k): + """``row`` is a top-k of the entries of ``scores`` above -inf, up to the + fp32 reduce order.""" + valid = scores > float("-inf") + want = min(k, int(valid.sum())) + sel = row[row >= 0].long() + assert sel.numel() == want and sel.unique().numel() == want + if want: + assert bool(valid[sel].all()) + kth = scores[valid].topk(want).values[-1] + finite = scores[torch.isfinite(scores)] + tol = 1e-4 * float(finite.abs().max()) if finite.numel() else 0.0 + assert bool((scores[sel] >= kth - tol).all()) + + +def _block_maxima(scores): + end = scores.numel() + if end == 0: + return scores + padded = torch.full((cdiv(end, CAND_BLOCK) * CAND_BLOCK,), float("-inf")) + padded[:end] = scores.cpu() + out = padded.view(-1, CAND_BLOCK).amax(1) + out[(end - 1) // CAND_BLOCK] = float("inf") + return out.to(scores.device) + + +def _forward_context(monkeypatch, metadata): + context = types.SimpleNamespace( + attn_metadata={"indexer": metadata}, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + monkeypatch.setattr(ops, "get_forward_context", lambda: context) + + +def _decode_metadata(case, rows, ratio, query_lens): + lens = torch.tensor([(p + 1) // ratio for _, p in rows], dtype=torch.int32) + lens = lens.to(DEVICE) + block_table = case.block_table[[req for req, _ in rows]].contiguous() + entries = case.cache[ratio].shape[1] + next_n = max(query_lens) + context_lens = (torch.tensor(case.seq_lens, dtype=torch.int32) // ratio).to(DEVICE) + native = native_decode(lens, block_table, query_lens, next_n, context_lens) + schedule = torch.empty( + ops.rocm_mxfp4_decode_schedule_words(HEADS, HEAD_DIM, entries, next_n), + dtype=torch.int32, + device=DEVICE, + ) + return RocmMxfp4IndexerMetadata( + seq_lens=None, + max_seq_len=max(case.seq_lens), + slot_mapping=None, + num_decodes=len(case.seq_lens), + num_decode_tokens=len(rows), + num_prefills=0, + num_prefill_tokens=0, + decode=types.SimpleNamespace(block_table=block_table, seq_lens=lens[:, None]), + decode_row_lens=lens, + decode_block_ends=(lens + CAND_BLOCK - 1) // CAND_BLOCK, + decode_native=native, + decode_schedule=ops.build_rocm_mxfp4_decode_schedule( + lens, HEADS, HEAD_DIM, entries, schedule, MAX_LEN // ratio, native + ), + ) + + +def _prefill_metadata( + case, rows, ratio, chunk_bounds, query_start_loc, gather_rows, gather +): + chunks = [] + for req_lo, req_hi, t0, t1 in chunk_bounds: + ends = torch.tensor( + [(rows[t][1] + 1) // ratio for t in range(t0, t1)], dtype=torch.int32 + ).to(DEVICE) + chunks.append( + DeepseekV32IndexerPrefillChunkMetadata( + block_table=case.block_table[req_lo:req_hi], + cu_seqlen_ks=torch.zeros_like(ends), + cu_seqlen_ke=ends, + cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE), + token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE), + total_seq_lens=1, + token_start=t0, + token_end=t1, + num_reqs=req_hi - req_lo, + ) + ) + seq_lens = torch.tensor(case.seq_lens) + context_lens = (seq_lens // ratio).int().to(DEVICE) + plans = plan_prefill_chunks( + chunks, + query_start_loc, + seq_lens, + context_lens, + ratio, + CAND_BLOCK, + 0.0 if gather else None, + torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE), + ) + return RocmMxfp4IndexerMetadata( + seq_lens=None, + max_seq_len=max(case.seq_lens), + slot_mapping=None, + num_decodes=0, + num_decode_tokens=0, + num_prefills=len(case.seq_lens), + num_prefill_tokens=len(rows), + prefill=types.SimpleNamespace(chunks=chunks), + prefill_plans=plans, + gather_launches=plan_gather_launches(chunks, plans, gather_rows), + ) + + +# The kernel prefers 128-token pages. +BLOCKS = pytest.mark.parametrize("block", [128, 64]) + + +@BLOCKS +def test_k_store_stays_inside_its_pages(block): + """The writer touches only its pages of the block-major pool. What they hold + is checked by the layer tests, which read them back through aiter's kernel.""" + set_random_seed(0) + case = _Case([700, 333], block) + assert int(case.pool[~case.page_bytes].count_nonzero()) == 0 + + +@pytest.mark.parametrize( + "rows, width, k", + # vLLM's kernel below 128 rows, aiter's from 128 rows and on every k=2048 + [(8, 4096, 512), (128, 4096, 512), (8, 16384, 2048)], +) +def test_topk_matches_torch_on_both_kernels(rows, width, k): + set_random_seed(0) + logits = torch.randn(rows, width, device=DEVICE) + lengths = torch.randint(1, width + 1, (rows,), dtype=torch.int32, device=DEVICE) + lengths[0] = k // 2 # shorter than k, so the rest of the row pads with -1 + out = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE) + ops._topk(logits, lengths, out, k) + for r in range(rows): + n = min(k, int(lengths[r])) + want = torch.topk(logits[r, : int(lengths[r])], n).indices.sort().values + got = out[r] + assert torch.equal(got[got >= 0].sort().values, want.int()) + assert int((got == -1).sum()) == k - n + + +def _run_layers(monkeypatch, case, rows, metadata): + """Ratio-2 dense layer, the candidate source, and a consumer both ways; + each checked against the torch reference.""" + num_rows = len(rows) + q, q_scale, weights, q_ref = _queries(num_rows) + hidden = torch.empty(num_rows, 1, device=DEVICE) + + def topk_buffer(): + return torch.full((num_rows + 2, TOPK), 7, dtype=torch.int32, device=DEVICE) + + _forward_context(monkeypatch, metadata(2)) + out = topk_buffer() + ops.rocm_mxfp4_sparse_attn_indexer( + hidden, + "indexer", + case.cache[2], + q, + q_scale, + weights, + TOPK, + HEAD_DIM, + MAX_LEN // 2, + out, + compress_ratio=2, + ) + for i, (req, pos) in enumerate(rows): + end = (pos + 1) // 2 + keys = case.keys(2, req, end) + _assert_topk(out[i], _scores(q_ref[i], weights[i], keys), TOPK) + + source = metadata(1) + _forward_context(monkeypatch, source) + pool = torch.full((num_rows + 2, POOL_BLOCKS), 5, dtype=torch.int32, device=DEVICE) + out = topk_buffer() + ops.rocm_mxfp4_sparse_attn_indexer( + hidden, + "indexer", + case.cache[1], + q, + q_scale, + weights, + TOPK, + HEAD_DIM, + MAX_LEN, + out, + compress_ratio=1, + candidate_blocks=pool, + candidate_block_size=CAND_BLOCK, + candidate_write=True, + ) + allowed = [] + for i, (req, pos) in enumerate(rows): + scores = _scores(q_ref[i], weights[i], case.keys(1, req, pos + 1)) + _assert_topk(out[i], scores, TOPK) + _assert_topk(pool[i], _block_maxima(scores), POOL_BLOCKS) + mask = torch.zeros_like(scores, dtype=torch.bool) + for block in pool[i][pool[i] >= 0].tolist(): + mask[block * CAND_BLOCK : (block + 1) * CAND_BLOCK] = True + allowed.append(torch.where(mask, scores, float("-inf"))) + + pa = ops._aiter() + build = pa.build_candidate_gather + builds = [] + + def counting_build(*args, **kwargs): + builds.append(1) + return build(*args, **kwargs) + + monkeypatch.setattr(pa, "build_candidate_gather", counting_build) + for gather in (True, False): + consumer = metadata(1, gather) + consumer.decode_use_gather = gather + _forward_context(monkeypatch, consumer) + builds.clear() + # Two consumer layers of one step: the second reuses the first's pools. + for _ in range(2): + out = topk_buffer() + ops.rocm_mxfp4_sparse_mqa_indexer( + hidden, + "indexer", + case.cache[1], + q, + q_scale, + weights, + TOPK, + HEAD_DIM, + MAX_LEN, + out, + 1, + pool, + CAND_BLOCK, + POOL_BLOCKS * CAND_BLOCK, + ) + for i in range(num_rows): + _assert_topk(out[i], allowed[i], TOPK) + if gather: + launches = len(consumer.gather_launches) + (consumer.decode is not None) + assert len(builds) == launches + + +@BLOCKS +@pytest.mark.parametrize( + "query_lens", + # Uniform steps launch the dense layers on next_n-row sequences, also with + # cudagraph padding (query length 0) after them; a ragged step keeps a row + # per token. + [[6, 6, 6, 6], [2, 2, 2], [2, 2, 2, 0], [2, 1, 2]], + ids=["native6", "native2", "padded", "ragged"], +) +def test_decode_layers_match_reference(monkeypatch, block, query_lens): + set_random_seed(0) + next_n = max(query_lens) + case = _Case([q and n for q, n in zip(query_lens, [900, 333, 610, 1200])], block) + # a padding request has next_n rows of its own and no context + rows = [ + (req, n - q + j if q else -1) + for req, (n, q) in enumerate(zip(case.seq_lens, query_lens)) + for j in range(q or next_n) + ] + _run_layers( + monkeypatch, + case, + rows, + lambda r, gather=False: _decode_metadata(case, rows, r, query_lens), + ) + + +@pytest.mark.parametrize( + "seq_lens,new_tokens,chunks", + [ + # one request past a cached prefix, then one sliced into two chunks + ([760, 500], [300, 500], [(0, 1, 0, 300), (1, 2, 300, 560), (1, 2, 560, 800)]), + # three requests in one chunk, the first two with equal rows + ([600, 450, 900], [200, 200, 120], [(0, 3, 0, 520)]), + # four requests with four different query lengths in one chunk + ([900, 600, 450, 700], [150, 90, 200, 60], [(0, 4, 0, 500)]), + ], + ids=["sliced", "batched", "ragged"], +) +# The consumers rejoin a request's rows across chunks, then cut them at most +# gather_rows a launch. +@pytest.mark.parametrize("gather_rows", [1 << 20, 128], ids=["joined", "cut"]) +@BLOCKS +def test_prefill_layers_match_reference( + monkeypatch, seq_lens, new_tokens, chunks, gather_rows, block +): + set_random_seed(0) + case = _Case(seq_lens, block) + rows = [ + (req, n - q + i) + for req, (n, q) in enumerate(zip(seq_lens, new_tokens)) + for i in range(q) + ] + query_start_loc = [0, *itertools.accumulate(new_tokens)] + _run_layers( + monkeypatch, + case, + rows, + lambda r, gather=True: _prefill_metadata( + case, rows, r, chunks, query_start_loc, gather_rows, gather + ), + ) + + +def _rope_gptj(x, cos_sin, positions): + """GPT-J RoPE on the last 64 dims, as the indexer rotates Q and K: bf16 in + and out.""" + x = x.bfloat16().float() + cos, sin = cos_sin[positions].chunk(2, dim=-1) + if x.dim() == 3: + cos, sin = cos[:, None], sin[:, None] + rot = x[..., -64:] + even, odd = rot[..., 0::2].clone(), rot[..., 1::2].clone() + rot[..., 0::2] = even * cos - odd * sin + rot[..., 1::2] = odd * cos + even * sin + return x.bfloat16() + + +def _mxfp4_round_trip(x): + """MXFP4 round trip of x, rounded as the writers do: a scale of + 2^ceil(log2(amax / 6)) per 32 values, then the nearest e2m1 value, ties to + even.""" + blocks = x.float().unflatten(-1, (-1, 32)) + amax = blocks.abs().amax(-1, keepdim=True).clamp(min=2**-126) + scale = torch.exp2(torch.ceil(torch.log2(amax / 6))) + y = (blocks / scale).clamp(-6, 6) + # Even codes first: argmin keeps the first of two equal distances. + grid = torch.cat([_E2M1[0::2], _E2M1[1::2]]).to(x.device) + idx = (y.abs()[..., None] - grid).abs().argmin(-1) + return (grid[idx] * y.sign() * scale).flatten(-2) + + +def _torch_keys(case, ratio, req): + """A request's compressed keys in torch: k_norm of each group's last + token, rotated at the group's first position, through MXFP4.""" + k = case.k_pre[req][ratio - 1 :: ratio].float() + k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) * case.norm + pos = torch.arange(k.shape[0], device=DEVICE) * ratio + return _mxfp4_round_trip(_rope_gptj(k, case.cos_sin, pos)) + + +def _topk_indices(q, weights, keys, ends, k): + """Top-k of sum_h relu(q_h . key) * weights_h over each row's [0, end).""" + scores = (torch.einsum("thd,cd->thc", q, keys).relu() * weights[..., None]).sum(1) + cols = torch.arange(keys.shape[0], device=DEVICE) + return scores.masked_fill(cols >= ends[:, None], float("-inf")).topk(k).indices + + +def _recall(selected, reference): + """Per row, the fraction of the reference top-k that ``selected`` picked.""" + return (selected[:, :, None] == reference[:, None, :]).any(1).float().mean(1) + + +@BLOCKS +@pytest.mark.parametrize("ratio", [2, 1]) +@pytest.mark.parametrize("step", ["decode", "prefill"]) +def test_selection_from_bf16_matches_torch(monkeypatch, step, ratio, block): + """End to end from bf16 K and Q: vLLM's MXFP4 K store and Q quant, the + paged pool, metadata from the planning functions the builder calls, and + the indexer op pick what the same MXFP4 math picks in torch.""" + set_random_seed(0) + if step == "decode": + case = _Case([3000, 2600, 1900, 3500, 1200, 2200, 3900, 1600], block) + rows = [(req, n - 1) for req, n in enumerate(case.seq_lens)] + metadata = _decode_metadata(case, rows, ratio, [1] * len(rows)) + else: + new_tokens = [300, 200, 250] + case = _Case([1500, 900, 2000], block) + rows = [ + (req, n - q + i) + for req, (n, q) in enumerate(zip(case.seq_lens, new_tokens)) + for i in range(q) + ] + metadata = _prefill_metadata( + case, + rows, + ratio, + [(0, len(new_tokens), 0, len(rows))], + [0, *itertools.accumulate(new_tokens)], + 1 << 20, + gather=False, + ) + num_rows = len(rows) + q_bf16 = torch.randn(num_rows, HEADS, HEAD_DIM, device=DEVICE).to(torch.bfloat16) + w = torch.randn(num_rows, HEADS, device=DEVICE) + pos = torch.tensor([p for _, p in rows], device=DEVICE) + scales = (HEAD_DIM**-0.5, HEADS**-0.5) + + (q, q_scale), weights = rocm_mxfp4_indexer_q_quant( + pos, q_bf16, case.cos_sin, w, *scales + ) + _forward_context(monkeypatch, metadata) + out = torch.full((num_rows, TOPK), 7, dtype=torch.int32, device=DEVICE) + ops.rocm_mxfp4_sparse_attn_indexer( + torch.empty(num_rows, 1, device=DEVICE), + "indexer", + case.cache[ratio], + q, + q_scale, + weights, + TOPK, + HEAD_DIM, + MAX_LEN // ratio, + out, + compress_ratio=ratio, + ) + ends = (pos + 1) // ratio + assert bool(((out >= 0) & (out < ends[:, None])).all()) + assert all(row.unique().numel() == TOPK for row in out) + + q_ref = _mxfp4_round_trip(_rope_gptj(q_bf16, case.cos_sin, pos)) + w_ref = w * scales[0] * scales[1] + agree = [] + for req in range(len(case.seq_lens)): + mine = torch.tensor([r == req for r, _ in rows], device=DEVICE) + ref = _topk_indices( + q_ref[mine], w_ref[mine], _torch_keys(case, ratio, req), ends[mine], TOPK + ) + agree.append(_recall(out[mine], ref)) + assert torch.cat(agree).mean().item() >= 0.99 diff --git a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py new file mode 100644 index 0000000000..0283ed99ee --- /dev/null +++ b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The ROCm MXFP4 indexer reads the paged cache in place, so each prefill +chunk is split back into its requests' rows, which one launch finds through +the chunk's query_start_loc. A chunk is whole requests or one query slice of a +long request, and both have to come back with the right rows and the right +block-table rows. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm-only", allow_module_level=True) + +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerPrefillChunkMetadata, +) +from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( + native_decode, + plan_gather_launches, + plan_prefill_chunks, + split_prefill_chunks, +) +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import rocm_mxfp4_consumer_rows + + +def _chunk(block_table, token_start, token_end, row_ends): + ends = torch.tensor(row_ends, dtype=torch.int32) + return DeepseekV32IndexerPrefillChunkMetadata( + block_table=block_table, + cu_seqlen_ks=torch.full_like(ends, 7), + cu_seqlen_ke=ends + 7, + cu_seq_lens=torch.zeros(1, dtype=torch.int32), + token_to_seq=torch.zeros(1, dtype=torch.int32), + total_seq_lens=1, + token_start=token_start, + token_end=token_end, + num_reqs=block_table.shape[0], + ) + + +# A decode (1 token), then prefills of 3, 3, 5 and 6 query tokens; the last +# one is sliced into two chunks. +QUERY_START_LOC = [0, 1, 4, 7, 12, 18] +SEQ_LENS = torch.tensor([40, 10, 12, 21, 64]) +BLOCK_TABLE = torch.arange(20, dtype=torch.int32).view(5, 4) +CHUNKS = [ + _chunk(BLOCK_TABLE[1:4], 1, 12, list(range(11))), + _chunk(BLOCK_TABLE[4:5], 12, 15, [29, 30, 30]), + _chunk(BLOCK_TABLE[4:5], 15, 18, [31, 31, 32]), +] + + +def _plans(min_gather_width, block=8): + return plan_prefill_chunks( + CHUNKS, + QUERY_START_LOC, + SEQ_LENS, + SEQ_LENS.int() // 2, + 2, + block, + min_gather_width, + torch.tensor(QUERY_START_LOC, dtype=torch.int32), + ) + + +def test_chunks_split_into_requests(): + plans = _plans(20.0) + + assert [p.requests for p in plans] == [ + [(0, 3, 0), (3, 6, 1), (6, 11, 2)], + [(0, 3, 0)], + [(0, 3, 0)], + ] + # every chunk launches once, with offsets local to the chunk + assert [p.query_start_loc.tolist() for p in plans] == [ + [0, 3, 6, 11], + [0, 3], + [0, 3], + ] + # the logits are as wide as the longest compressed context in the chunk + assert [p.width for p in plans] == [10, 32, 32] + assert [p.context_lens.tolist() for p in plans] == [[5, 6, 10], [32], [32]] + torch.testing.assert_close(plans[2].row_ends, torch.tensor([31, 31, 32]).int()) + torch.testing.assert_close(plans[2].block_ends, torch.tensor([4, 4, 4]).int()) + assert [p.use_gather for p in plans] == [False, True, True] + assert [p.first_request for p in plans] == [1, 4, 4] + # no candidate blocks, as in the model-neutral builder + for plan in _plans(None, block=0): + assert plan.block_ends is None and not plan.use_gather + + +def test_launches_follow_the_logits_budget(monkeypatch): + """Each launch's fp32 logits fit in VLLM_SPARSE_INDEXER_MAX_LOGITS_MB. A dense + chunk's logits are [rows, longest context in the chunk]: K is read in place, + not gathered, so requests share a launch until their rows no longer fit, and + a request that does not fit alone is split on its query rows. A consumer's + logits are [rows, pool].""" + budget = 4 * 1000 * 30 + contexts, queries = torch.tensor([1000, 600, 200]), torch.tensor([10, 10, 10]) + assert split_prefill_chunks(contexts, queries, budget, 2) == [ + (slice(2, 5), slice(0, 30)) + ] + # the third request would take the chunk past the budget + queries = torch.tensor([10, 10, 11]) + assert split_prefill_chunks(contexts, queries, budget) == [ + (slice(0, 2), slice(0, 20)), + (slice(2, 3), slice(0, 11)), + ] + assert split_prefill_chunks(torch.tensor([1000]), torch.tensor([100]), budget) == [ + (slice(0, 1), slice(lo, min(lo + 30, 100))) for lo in range(0, 100, 30) + ] + # the default 512 MiB holds 128 rows of a 1M context, or 8192 of a 16K pool + assert split_prefill_chunks( + torch.tensor([1 << 20]), torch.tensor([200]), 512 << 20 + ) == [(slice(0, 1), slice(0, 128)), (slice(0, 1), slice(128, 200))] + monkeypatch.setenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512") + assert rocm_mxfp4_consumer_rows(16384) == 8192 + + +def test_gather_launches_rejoin_a_requests_rows(): + """The consumers' logits are [rows, pool] whatever the context, so a + request's rows go back together across the chunks that cut them, and a + launch never mixes requests.""" + # the first chunk is too narrow to gather, the last request's slices rejoin + (joined,) = plan_gather_launches(CHUNKS, _plans(20.0), 100) + assert (joined.token_start, joined.token_end) == (12, 18) + # every chunk gathers: the first chunk's three requests launch apart + launches = plan_gather_launches(CHUNKS, _plans(0.0), 100) + assert [(g.token_start, g.token_end) for g in launches] == [ + (1, 4), + (4, 7), + (7, 12), + (12, 18), + ] + assert plan_gather_launches(CHUNKS, _plans(None), 100) == [] + + +def test_decode_launches_native_on_uniform_steps(): + """A decode step goes to the kernel as next_n-row sequences only when its + shape says so: each request has next_n rows, and cudagraph padding (query + length 0) comes only after them. A FULL graph captured on a uniform batch + then replays the same launch on a padded one.""" + rows = torch.zeros(8, dtype=torch.int32) + row_block_table = torch.arange(32, dtype=torch.int32).view(8, 4) + context_lens = torch.tensor([9, 4, 7, 0, 5], dtype=torch.int32) + + native = native_decode(rows, row_block_table, [2, 2, 2, 0], 2, context_lens) + assert native is not None and native.next_n == 2 + assert native.context_lens.tolist() == [9, 4, 7, 0] + # every request's first flattened row, as a view + assert native.block_table.tolist() == row_block_table[::2].tolist() + assert native.block_table.data_ptr() == row_block_table.data_ptr() + + assert native_decode(rows, row_block_table, [2, 0, 2, 2], 2, context_lens) is None + # a ragged step, rows padded past the requests, and no speculation + assert native_decode(rows[:5], row_block_table, [2, 1, 2], 2, context_lens) is None + assert native_decode(rows, row_block_table, [2, 2, 2], 2, context_lens) is None + assert native_decode(rows[:4], row_block_table, [1] * 4, 1, context_lens) is None diff --git a/vllm/config/attention.py b/vllm/config/attention.py index a6602ed549..8a9be6336c 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -93,7 +93,10 @@ class AttentionConfig: an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension (the top-k runs on the kernels' bf16 logits). The sparse path costs O(candidate blocks) per query regardless of context length, so it pays off - for long contexts (roughly 32K tokens and beyond) and is slower below.""" + for long contexts (roughly 32K tokens and beyond) and is slower below. + On ROCm gfx950 it runs aiter's paged MXFP4 MQA-logits kernel over the + candidate pool instead, with fp32 logits, and keeps the masked dense walk + for steps whose contexts are too short for the pool to pay.""" hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental diff --git a/vllm/model_executor/layers/rocm_paged_mxfp4_indexer.py b/vllm/model_executor/layers/rocm_paged_mxfp4_indexer.py new file mode 100644 index 0000000000..b42f63ac94 --- /dev/null +++ b/vllm/model_executor/layers/rocm_paged_mxfp4_indexer.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse attention indexer layers on aiter's paged MXFP4 kernels, for ROCm. + +With ``indexer_kv_dtype="mxfp4"`` on gfx950, the model writes the indexer K +cache in the order aiter's MQA-logits kernel reads (see +`rocm_paged_mxfp4_cache_layout`), and the kernel scores the cache in place. +Requires `RocmMxfp4IndexerMetadataBuilder` metadata. +""" + +import torch +from torch import nn + +from vllm.config import get_current_vllm_config +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import ( + rocm_mxfp4_sparse_attn_indexer, + rocm_mxfp4_sparse_mqa_indexer, +) + + +class RocmSparseAttnIndexer(SparseAttnIndexer): + """`SparseAttnIndexer` whose HIP path walks the paged MXFP4 cache with + aiter's kernel. A two-level indexer's source layer also publishes its + candidate pool here.""" + + def forward_hip( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ): + assert isinstance(q_quant, tuple) and self.skip_k_cache_insert, ( + "the ROCm MXFP4 indexer takes (values, scales) Q and a K cache " + "the model already wrote" + ) + q_values, q_scale = q_quant + return rocm_mxfp4_sparse_attn_indexer( + hidden_states, + self.k_cache.prefix, + self.k_cache.kv_cache, + q_values, + q_scale, + weights, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.topk_indices_buffer, + compress_ratio=self.compress_ratio, + candidate_blocks=self.candidate_blocks, + candidate_block_size=self.candidate_block_size, + candidate_write=self.candidate_write, + ) + + +class RocmSparseMQAIndexer(nn.Module): + """DeepSeek-V4.1's candidate-consuming indexer on aiter's paged MXFP4 + MQA-logits kernel: the consumers score only the candidate pool the source + layer published, read straight from the paged cache. + + Built and called as `SparseMQAIndexer` is, so the model can take either. + """ + + weights_dtype = torch.float32 + """The scorer's weights and logits stay fp32: gfx950 has no packed bf16 + arithmetic, so its head reduce costs the same in either dtype.""" + + def __init__( + self, + k_cache, + topk_tokens: int, + head_dim: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + ): + super().__init__() + self.k_cache = k_cache + self.topk_tokens = topk_tokens + self.head_dim = head_dim + # aiter's kernel bounds a row by the model length in compressed + # positions, not by the prefill buffer SparseMQAIndexer sizes. + model_config = get_current_vllm_config().model_config + self.max_model_len = model_config.max_model_len // k_cache.compress_ratio + self.topk_indices_buffer = topk_indices_buffer + self.candidate_blocks = candidate_blocks + self.candidate_block_size = candidate_block_size + self.num_candidate_cols = candidate_blocks.shape[1] * candidate_block_size + + def forward( + self, + hidden_states: torch.Tensor, + q_quant: tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ) -> torch.Tensor: + assert k is None, "the model writes the indexer K cache" + q_values, q_scale = q_quant + return rocm_mxfp4_sparse_mqa_indexer( + hidden_states, + self.k_cache.prefix, + self.k_cache.kv_cache, + q_values, + q_scale, + weights, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.topk_indices_buffer, + self.k_cache.compress_ratio, + self.candidate_blocks, + self.candidate_block_size, + self.num_candidate_cols, + ) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 897b0d12db..fe020f4061 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -2,8 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +from collections.abc import Callable from dataclasses import dataclass -from typing import cast +from typing import Any, cast import torch @@ -15,8 +16,14 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic +from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import ( + RocmSparseAttnIndexer, + RocmSparseMQAIndexer, +) from vllm.models.deepseek_v41.attention import ( DeepseekV4Attention, + DeepseekV4Indexer, + DeepseekV4IndexerCache, _replace_layer_index, ) from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache @@ -32,6 +39,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( CommonAttentionMetadata, + MultipleOf, ) from vllm.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWABackend, @@ -47,6 +55,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_sparse_attn_decode, rocm_sparse_attn_prefill, ) +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import ( + rocm_paged_mxfp4_cache_layout, +) from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) @@ -517,11 +528,141 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu ) +@functools.cache +def _aiter_indexer_cache_ops() -> tuple[Callable[..., None], Callable[..., tuple]]: + """The aiter indexer key writer and query quantizer.""" + from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import ( + k_norm_rope_mxfp4_cache, + ) + from aiter.ops.triton.rope.q_rope_mxfp4_quant import q_rope_mxfp4_quant + + return k_norm_rope_mxfp4_cache, q_rope_mxfp4_quant + + +def rocm_mxfp4_indexer_k_store( + k_pre: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + k_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + use_fp4_cache: bool, + *, + num_heads: int, +) -> None: + """`indexer_k_norm_rope_store` for the ROCm MXFP4 cache: aiter's cache op + writes the key in the order its MQA-logits kernel reads with ``num_heads`` + query heads.""" + assert use_fp4_cache, "the ROCm indexer cache op writes MXFP4 only" + layout = rocm_paged_mxfp4_cache_layout(num_heads, k_pre.shape[1], k_cache.shape[1]) + k_norm_rope_mxfp4_cache, _ = _aiter_indexer_cache_ops() + k_norm_rope_mxfp4_cache( + k_pre, + positions, + cos_sin_cache, + rms_norm_weight, + rms_norm_eps, + k_cache, + kv_slot_mapping, + compress_ratio, + shuffle=layout, + ) + + +def rocm_mxfp4_indexer_q_quant( + positions: torch.Tensor, + index_q: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + use_fp4: bool = True, + weights_out_dtype: torch.dtype = torch.float32, +) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]: + """`fused_indexer_q_rope_quant` for the ROCm MXFP4 indexer, on aiter's op: + ((packed [T, H, D // 2], e8m0 as one int32 per head [T, H]), fp32 + weights).""" + assert use_fp4 and weights_out_dtype == torch.float32, ( + "the ROCm indexer quantizes Q to MXFP4 and scores with fp32 weights" + ) + _, q_rope_mxfp4_quant = _aiter_indexer_cache_ops() + q_packed, q_scale, weights_out = q_rope_mxfp4_quant( + index_q, + positions, + index_q_cos_sin_cache, + index_weights, + index_weights_softmax_scale * index_weights_head_scale, + ) + return (q_packed, q_scale.view(torch.int32).squeeze(-1)), weights_out + + +class DeepseekV41RocmMxfp4Indexer(DeepseekV4Indexer): + """The indexer on aiter's paged MXFP4 cache: its K store and Q quant write + in the order aiter's MQA-logits kernel reads, and its layers score with it. + ``_produce_k`` and ``forward_q`` mirror ``DeepseekV4Indexer``'s except for + those two calls.""" + + mqa_cls = RocmSparseMQAIndexer + attn_cls = RocmSparseAttnIndexer + + def _produce_k( + self, + latent: torch.Tensor | None, + positions: torch.Tensor, + rotary_emb: torch.nn.Module, + ) -> None: + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict) or latent is None: + return + assert self.owns_k + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) + k_pre, _ = self.wk(latent) + rocm_mxfp4_indexer_k_store( + k_pre, + positions, + rotary_emb.cos_sin_cache, + self.k_norm.weight, + self.k_norm.variance_epsilon, + self.k_cache.kv_cache, + indexer_metadata.slot_mapping, + self.compress_ratio, + self.use_fp4_kv, + num_heads=self.n_head, + ) + + def forward_q( + self, + qr: torch.Tensor | QuantizedActivation, + qr_scale: torch.Tensor | None, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: torch.nn.Module, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + q = self._wq_b_proj(qr, qr_scale).view(-1, self.n_head, self.head_dim) + (q, q_scale), weights = rocm_mxfp4_indexer_q_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + use_fp4=self.use_fp4_kv, + weights_out_dtype=self.indexer_weights_dtype, + ) + return q, q_scale, weights + + class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): @staticmethod def get_name() -> str: return "ROCM_FLASHMLA_SPARSE_DSV4" + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: + return [64, 128] + @staticmethod def get_builder_cls() -> type[DeepseekV4SparseMLAMetadataBuilder]: return DeepseekV4SparseMLAMetadataBuilder @@ -540,6 +681,13 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend _use_aiter_sparse_mla = False + def _indexer_cls( + self, k_cache: DeepseekV4IndexerCache | None + ) -> type[DeepseekV4Indexer]: + if k_cache is not None and k_cache.rocm_mxfp4: + return DeepseekV41RocmMxfp4Indexer + return super()._indexer_cls(k_cache) + def __init__(self, *args, **kwargs): vllm_config = args[0] if args else kwargs["vllm_config"] super().__init__(*args, **kwargs) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index c27a8d9baf..73d9c0c723 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -35,6 +35,9 @@ from vllm.models.deepseek_v41.common.ops import ( if TYPE_CHECKING: from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr + from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import ( + RocmSparseMQAIndexer, + ) from vllm.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWAMetadata, ) @@ -261,6 +264,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): """Return whether this instance stores fp8 KV in fp8_ds_mla layout.""" return self.use_fp8_ds_mla_layout + def _indexer_cls( + self, k_cache: "DeepseekV4IndexerCache | None" + ) -> type["DeepseekV4Indexer"]: + """The indexer for this layer's index K cache. A platform subclass + returns one that writes and scores its own cache layout.""" + return DeepseekV4Indexer + def __init__( self, vllm_config: VllmConfig, @@ -449,7 +459,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ) is_candidate_source = layer_id == self.candidate_source_layer uses_candidates = 0 <= self.candidate_source_layer < layer_id - self.indexer = DeepseekV4Indexer( + self.indexer = self._indexer_cls(index_k_cache)( vllm_config, config=config, hidden_size=self.hidden_size, @@ -1152,6 +1162,10 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): self.compress_ratio = compress_ratio vllm_config = get_current_vllm_config() self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits + # aiter's paged MXFP4 kernels write this cache and read it in place. + self.rocm_mxfp4 = current_platform.is_rocm() and dsa_indexer_uses_fp4( + vllm_config + ) compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -1190,13 +1204,21 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): # it 512B-aligned; the main KV record read from the same blocks # needs its own TMA stride, so keep both. block_stride_alignment=( - math.lcm(512, page_alignment) if self.sparse_logits else None + math.lcm(512, page_alignment) + if self.sparse_logits and not self.rocm_mxfp4 + else None ), ) def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: + if self.rocm_mxfp4: + from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( + DeepseekV41RocmMxfp4IndexerBackend, + ) + + return DeepseekV41RocmMxfp4IndexerBackend if self.sparse_logits: return DeepseekV41SparseIndexerBackend return DeepseekV41IndexerBackend @@ -1218,6 +1240,12 @@ class DeepseekV4Indexer(nn.Module): later indexers mask their scores to those blocks before their own top-k. """ + # The scoring layers; a platform subclass swaps in its own. + mqa_cls: ClassVar[type["SparseMQAIndexer | RocmSparseMQAIndexer"]] = ( + SparseMQAIndexer + ) + attn_cls: ClassVar[type[SparseAttnIndexer]] = SparseAttnIndexer + def __init__( self, vllm_config: VllmConfig, @@ -1305,7 +1333,7 @@ class DeepseekV4Indexer(nn.Module): and candidate_block_buffer is not None and not candidate_write ) - self.indexer_op: SparseAttnIndexer | SparseMQAIndexer + self.indexer_op: SparseAttnIndexer | SparseMQAIndexer | RocmSparseMQAIndexer if use_sparse_logits: if not self.use_fp4_kv: raise ValueError( @@ -1314,7 +1342,7 @@ class DeepseekV4Indexer(nn.Module): ) assert candidate_block_buffer is not None assert topk_indices_buffer is not None - self.indexer_op = SparseMQAIndexer( + self.indexer_op = self.mqa_cls( self.k_cache, self.topk_tokens, self.head_dim, @@ -1324,7 +1352,7 @@ class DeepseekV4Indexer(nn.Module): candidate_block_size, ) else: - self.indexer_op = SparseAttnIndexer( + self.indexer_op = self.attn_cls( self.k_cache, self.quant_block_size, self.scale_fmt, @@ -1343,7 +1371,7 @@ class DeepseekV4Indexer(nn.Module): # The fused Q kernel writes the per-head weights in the dtype the # scoring kernels take, so no cast runs per step. self.indexer_weights_dtype = ( - SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32 + self.mqa_cls.weights_dtype if use_sparse_logits else torch.float32 ) def _produce_k( diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index e7eb67b61f..f2f2c51925 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -51,7 +51,7 @@ from vllm.v1.kv_cache_interface import ( logger = init_logger(__name__) # The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout) -# and mxfp4 is the opt-in Blackwell path. +# and mxfp4 is the opt-in Blackwell and gfx950 path. DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4") @@ -64,6 +64,28 @@ def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool: f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})." ) use_fp4 = kv_dtype == "mxfp4" + if use_fp4 and current_platform.is_rocm(): + from vllm._aiter_ops import rocm_aiter_ops + from vllm.platforms.rocm import get_cdna_version + + # Only DeepSeek-V4.1 is wired to the ROCm MXFP4 cache; other DSA models + # would silently keep their FP8 one. + model_config = vllm_config.model_config + if model_config is None or model_config.hf_config.model_type != "deepseek_v41": + raise ValueError( + "indexer_kv_dtype='mxfp4' on ROCm is only supported for " + "DeepSeek-V4.1-Flash." + ) + if get_cdna_version() != 4: + raise ValueError( + "indexer_kv_dtype='mxfp4' on ROCm requires CDNA4 (MI350X/MI355X)." + ) + if not rocm_aiter_ops.is_enabled(): + raise ValueError( + "indexer_kv_dtype='mxfp4' on ROCm runs on aiter's kernels; enable " + "aiter with VLLM_ROCM_USE_AITER=1." + ) + return True if use_fp4 and not current_platform.is_device_capability_family(100): raise ValueError( "indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs " diff --git a/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py b/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py new file mode 100644 index 0000000000..cfc30204da --- /dev/null +++ b/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Indexer metadata for the ROCm paged MXFP4 path. + +Extends the dense indexer metadata with what aiter's paged MXFP4 MQA-logits +kernel needs to read the cache in place: each prefill chunk as its requests, +packed through query_start_loc into one launch, and the decode rows as +next_n-row sequences on uniform steps. `RocmMxfp4IndexerMetadataBuilder` plans +this for any model's indexer, and a two-level indexer's candidate blocks once +a subclass sets their size. DeepSeek-V4.1's builder does, and with +``AttentionConfig.indexer_sparse_logits`` adds the per-step state the +candidate consumers share. +""" + +import bisect +from dataclasses import dataclass, field, fields + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.v1.attention.backend import ( + AttentionCGSupport, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerMetadataBuilder, + DeepseekV32IndexerPrefillChunkMetadata, + DeepseekV41IndexerBackend, +) +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import ( + build_rocm_mxfp4_decode_schedule, + rocm_mxfp4_consumer_rows, + rocm_mxfp4_decode_schedule_words, + rocm_paged_mxfp4_cache_layout, +) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec + +logger = init_logger(__name__) + +# The consumers walk the pool instead of the context once the context is this +# many pools long. Set from kernel timings, not end-to-end runs: the gather +# loads more per key than the dense walk, so it only wins on a large cut. +_GATHER_MIN_CUT_DECODE = 3.0 +_GATHER_MIN_CUT_PREFILL = 1.5 + + +@dataclass +class RocmMxfp4PrefillPlan: + """One prefill chunk as the requests the kernel launches on, all in one + launch: the kernel finds each row's request through query_start_loc, and + a request's rows share its KV tiles.""" + + requests: list[tuple[int, int, int]] + """(first row, end row, request) per request with rows in the chunk; rows + count from ``chunk.token_start``, requests index ``chunk.block_table``.""" + width: int + """Logits columns: an upper bound on the chunk's compressed contexts.""" + row_ends: torch.Tensor + """[rows] int32 exclusive compressed key bound of each row.""" + context_lens: torch.Tensor + """[chunk.num_reqs] int32 compressed context of each request.""" + query_start_loc: torch.Tensor + """[chunk.num_reqs + 1] int32 row offsets of the chunk's requests, local to + the chunk.""" + first_request: int + """The step's index of the chunk's first request.""" + block_ends: torch.Tensor | None = None + """[rows] int32 candidate blocks each row sees, for the source's pool.""" + use_gather: bool = False + + +@dataclass +class RocmMxfp4GatherLaunch: + """Query rows [token_start, token_end) of one request, which the candidate + consumers score against the pool in one launch.""" + + token_start: int + token_end: int + block_table: torch.Tensor + """[1, max_blocks] int32, the request's block table row.""" + context_len: torch.Tensor + """[1] int32 compressed context of the request.""" + row_ends: torch.Tensor + """[rows] int32 exclusive compressed key bound of each row.""" + pool: tuple[dict, torch.Tensor] | None = None + """The resolved candidate lists, built by the first consumer layer and + reused by the rest.""" + + +@dataclass +class RocmMxfp4NativeDecode: + """A decode step whose requests all have next_n query rows, as the + kernel's sequences: a workgroup then walks a KV tile once for up to + next_n rows instead of once per row.""" + + next_n: int + context_lens: torch.Tensor + """[requests] int32 compressed context of each request.""" + block_table: torch.Tensor + """[requests, max_blocks] int32, each request's first flattened row.""" + + +def native_decode( + row_lens: torch.Tensor, + row_block_table: torch.Tensor, + query_lens: list[int], + next_n: int, + context_lens: torch.Tensor, +) -> RocmMxfp4NativeDecode | None: + """The step as next_n-row sequences, or None when a request has fewer rows. + + ``row_lens`` and ``row_block_table`` are the flattened rows' bounds and + block table, ``query_lens`` the decode requests' query lengths. Trailing + empty requests are cudagraph padding: they keep their next_n rows, with no + context. Only the step's shape decides, so a FULL graph captured on a + uniform batch replays the same launch on a padded one. + """ + num_reqs = len(query_lens) + num_full = next((i for i, n in enumerate(query_lens) if n != next_n), num_reqs) + if ( + next_n <= 1 + or row_lens.shape[0] != num_reqs * next_n + or any(query_lens[num_full:]) + ): + return None + return RocmMxfp4NativeDecode( + next_n, context_lens[:num_reqs], row_block_table[::next_n] + ) + + +@dataclass +class RocmMxfp4IndexerMetadata(DeepseekV32IndexerMetadata): + prefill_plans: list[RocmMxfp4PrefillPlan] = field(default_factory=list) + """Parallel to ``prefill.chunks``.""" + decode_row_lens: torch.Tensor | None = None + """[rows] int32 compressed context of each decode query row.""" + decode_block_ends: torch.Tensor | None = None + """[rows] int32 candidate blocks each decode row sees, for the source's + pool.""" + decode_native: RocmMxfp4NativeDecode | None = None + """The same rows as next_n-row sequences, for the dense launches.""" + decode_schedule: torch.Tensor | None = None + """The dense decode launches' work schedule, built once for the step.""" + decode_use_gather: bool = False + """Whether the candidate consumers gather for the decode rows.""" + decode_gather: tuple[dict, torch.Tensor] | None = None + """The pool the first consumer resolved for the decode rows.""" + gather_launches: list[RocmMxfp4GatherLaunch] = field(default_factory=list) + """The prefill rows of the chunks that gather, as the consumers launch + them.""" + + +def plan_prefill_chunks( + chunks: list[DeepseekV32IndexerPrefillChunkMetadata], + query_start_loc: list[int], + seq_lens_cpu: torch.Tensor, + context_lens: torch.Tensor, + compress_ratio: int, + candidate_block_size: int, + min_gather_width: float | None, + query_start_loc_device: torch.Tensor, +) -> list[RocmMxfp4PrefillPlan]: + """Split each chunk into its requests. + + ``query_start_loc`` and ``seq_lens_cpu`` are the step's host copies (the + latter an upper bound), ``context_lens`` the device compressed lengths. + A chunk covers whole requests, or a query slice of one request; either + way it launches once, packed through the chunk's part of the device + ``query_start_loc``. + """ + block = candidate_block_size + plans = [] + for chunk in chunks: + t0, t1 = chunk.token_start, chunk.token_end + first = bisect.bisect_right(query_start_loc, t0) - 1 + requests = [] + for req in range(chunk.num_reqs): + lo = max(query_start_loc[first + req], t0) - t0 + hi = min(query_start_loc[first + req + 1], t1) - t0 + if hi > lo: + requests.append((lo, hi, req)) + reqs = slice(first, first + chunk.num_reqs) + packed = torch.clamp(query_start_loc_device[first : reqs.stop + 1], t0, t1) + packed -= t0 + width = int(seq_lens_cpu[reqs].max()) // compress_ratio + row_ends = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks + plans.append( + RocmMxfp4PrefillPlan( + requests=requests, + width=width, + row_ends=row_ends, + context_lens=context_lens[reqs], + query_start_loc=packed, + first_request=first, + block_ends=(row_ends + block - 1) // block if block else None, + use_gather=min_gather_width is not None and width >= min_gather_width, + ) + ) + return plans + + +def split_prefill_chunks( + context_lens: torch.Tensor, + query_lens: torch.Tensor, + max_logits_bytes: int, + request_offset: int = 0, +) -> list[tuple[slice, slice]]: + """(request slice, query slice) chunks, as the base chunker returns them. + + Nothing gathers K here, so the requests' contexts do not add up: a chunk's + logits are [rows, widest compressed context] fp32, at most + ``max_logits_bytes``. A request too wide to launch whole is cut on its + query rows. + """ + + def max_rows(width: int) -> int: + return max(1, max_logits_bytes // (4 * max(width, 1))) + + chunks: list[tuple[slice, slice]] = [] + end = 0 + while end < len(context_lens): + start, rows, width = end, 0, 0 + while end < len(context_lens): + q = int(query_lens[end]) + w = max(width, int(context_lens[end])) + if rows and rows + q > max_rows(w): + break + rows, width = rows + q, w + end += 1 + step = max_rows(width) + reqs = slice(start + request_offset, end + request_offset) + chunks.extend( + (reqs, slice(lo, min(lo + step, rows))) for lo in range(0, rows, step) + ) + return chunks + + +def plan_gather_launches( + chunks: list[DeepseekV32IndexerPrefillChunkMetadata], + plans: list[RocmMxfp4PrefillPlan], + max_rows: int, +) -> list[RocmMxfp4GatherLaunch]: + """The candidate consumers' launches over the chunks that gather. + + A consumer's logits are [rows, pool] however long the context, so the + dense split, sized for the context, would launch it far more often than + its memory needs. Each request's rows go back together across the chunks + that cut them, then out again at most ``max_rows`` a launch. + """ + # Per request: (first token, [(plan, first row, end row)]), and the chunk + # and request index its block table and context come from. + runs: list[tuple[int, list, DeepseekV32IndexerPrefillChunkMetadata, int]] = [] + last: tuple[int, int] | None = None + for chunk, plan in zip(chunks, plans): + if not plan.use_gather: + last = None + continue + for lo, hi, req in plan.requests: + request, t_lo = plan.first_request + req, chunk.token_start + lo + if last != (request, t_lo): + runs.append((t_lo, [], chunk, req)) + runs[-1][1].append((plan, lo, hi)) + last = (request, chunk.token_start + hi) + launches = [] + for token_start, pieces, chunk, req in runs: + ends = [plan.row_ends[lo:hi] for plan, lo, hi in pieces] + row_ends = ends[0] if len(ends) == 1 else torch.cat(ends) + block_table = chunk.block_table[req : req + 1] + context_len = pieces[0][0].context_lens[req : req + 1] + for lo in range(0, row_ends.shape[0], max_rows): + hi = min(lo + max_rows, row_ends.shape[0]) + launches.append( + RocmMxfp4GatherLaunch( + token_start + lo, + token_start + hi, + block_table, + context_len, + row_ends[lo:hi], + ) + ) + return launches + + +class DeepseekV41RocmMxfp4IndexerBackend(DeepseekV41IndexerBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_V41_ROCM_MXFP4_INDEXER" + + @staticmethod + def get_supported_kernel_block_sizes( + kv_cache_spec: KVCacheSpec | None = None, + ) -> list[int | MultipleOf]: + # --block-size 128 preferred. + return [64, 128] + + @staticmethod + def get_builder_cls() -> type["DeepseekV41RocmMxfp4IndexerMetadataBuilder"]: + return DeepseekV41RocmMxfp4IndexerMetadataBuilder + + +class RocmMxfp4IndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): + """Dense indexer metadata plus the in-place paged launches' layout. + + Requirements are checked here, at engine start. + """ + + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> AttentionCGSupport: + # Decode rows are flattened into persistent buffers, and whether a + # step launches them as next_n-row sequences depends on its shape + # alone (see native_decode). + return AttentionCGSupport.ALWAYS + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + **kwargs, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) + name = "the ROCm MXFP4 indexer" + if self.dcp_world_size > 1 or self.use_pcp: + raise NotImplementedError( + f"{name} does not support decode/prefill context parallel." + ) + hf_config = vllm_config.model_config.hf_text_config + num_heads, head_dim = hf_config.index_n_heads, hf_config.index_head_dim + self.num_heads, self.head_dim = num_heads, head_dim + self.page_entries = kv_cache_spec.block_size // self.compress_ratio + self.cache_layout = rocm_paged_mxfp4_cache_layout( + num_heads, head_dim, self.page_entries + ) + # Ragged steps take a row per query token, as do the candidate + # consumers' gathers. Uniform steps also launch the dense layers on + # next_n-row sequences, see native_decode. + self.use_flattening = True + # A two-level (candidate block) indexer's builder sets these. + self.candidate_block_size = 0 + self.min_gather_width: float | None = None + + self.decode_block_ends_buffer = torch.zeros( + self.arange_buffer.shape[0], dtype=torch.int32, device=device + ) + self.context_lens_buffer = torch.zeros( + self.arange_buffer.shape[0], dtype=torch.int32, device=device + ) + spec_config = vllm_config.speculative_config + # Adaptive verification replays a FULL graph on drafts reallocated + # across requests, so a step's shape does not fix its query lengths. + adaptive = spec_config is not None and spec_config.enable_adaptive_verification + self.native_next_n = 1 if adaptive else self.num_speculative_tokens + 1 + # The indexer's logits width, which the decode schedule sizes slices by. + max_model_len = vllm_config.model_config.max_model_len + self.logits_width = max_model_len // self.compress_ratio + # Every dense layer of this group reads the same rows through the same + # page geometry, so the step builds their schedule once, not per layer. + self.decode_schedule_buffer = torch.empty( + rocm_mxfp4_decode_schedule_words( + num_heads, head_dim, self.page_entries, self.native_next_n + ), + dtype=torch.int32, + device=device, + ) + + def _split_indexer_prefill_chunks( # type: ignore[override] + self, + compressed_seq_lens_cpu: torch.Tensor, + prefill_query_lens_cpu: torch.Tensor, + workspace_size: int, + max_logits_bytes: int, + request_offset: int = 0, + ) -> list[tuple[slice, slice]]: + return split_prefill_chunks( + compressed_seq_lens_cpu, + prefill_query_lens_cpu, + max_logits_bytes, + request_offset, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> RocmMxfp4IndexerMetadata: + base = super().build(common_prefix_len, common_attn_metadata, fast_build) + metadata = RocmMxfp4IndexerMetadata( + **{f.name: getattr(base, f.name) for f in fields(base)} + ) + cm = common_attn_metadata + context_lens = self.context_lens_buffer[: cm.num_reqs] + torch.floor_divide( + cm.seq_lens[: cm.num_reqs], self.compress_ratio, out=context_lens + ) + if base.decode is not None: + lengths = base.decode.seq_lens.view(-1) + metadata.decode_row_lens = lengths + block = self.candidate_block_size + if block: + ends = self.decode_block_ends_buffer[: lengths.shape[0]] + torch.add(lengths, block - 1, out=ends) + ends.floor_divide_(block) + metadata.decode_block_ends = ends + query_start_loc = cm.query_start_loc_cpu[: base.num_decodes + 1] + metadata.decode_native = native_decode( + lengths, + base.decode.block_table, + torch.diff(query_start_loc).tolist(), + self.native_next_n, + context_lens, + ) + metadata.decode_schedule = build_rocm_mxfp4_decode_schedule( + lengths, + self.num_heads, + self.head_dim, + self.page_entries, + self.decode_schedule_buffer, + self.logits_width, + metadata.decode_native, + ) + if base.prefill is not None: + assert cm.seq_lens_cpu_upper_bound is not None + metadata.prefill_plans = plan_prefill_chunks( + base.prefill.chunks, + cm.query_start_loc_cpu.tolist(), + cm.seq_lens_cpu_upper_bound, + context_lens, + self.compress_ratio, + self.candidate_block_size, + self.min_gather_width, + cm.query_start_loc, + ) + return metadata + + +class DeepseekV41RocmMxfp4IndexerMetadataBuilder(RocmMxfp4IndexerMetadataBuilder): + """DeepSeek-V4.1's two-level indexer: sets the candidate block size the + base plans with and, with ``indexer_sparse_logits``, adds the consumers' + gathers.""" + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + **kwargs, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) + hf_config = vllm_config.model_config.hf_text_config + self.candidate_block_size = getattr(hf_config, "candidate_block_size", 0) + self.num_candidate_cols = 0 + if vllm_config.attention_config.indexer_sparse_logits: + prefix = "attention_config.indexer_sparse_logits" + block = self.candidate_block_size + topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0) + layout = self.cache_layout + if block <= 0 or topk_blocks <= 0: + raise ValueError( + f"{prefix} requires a two-level (candidate block) indexer " + "model such as DeepSeek-V4.1." + ) + if topk_blocks & (topk_blocks - 1): + raise ValueError( + f"{prefix} on ROCm resolves the candidate pool with a sort " + f"that needs a power-of-two block count, not {topk_blocks}." + ) + if self.page_entries % block or block > layout.n_per_tile: + raise ValueError( + f"{prefix} on ROCm needs {block}-entry candidate blocks to " + f"tile a {self.page_entries}-entry indexer page and fit in " + f"one {layout.n_per_tile}-entry shuffle group." + ) + self.num_candidate_cols = topk_blocks * block + self.min_gather_width = self.num_candidate_cols * _GATHER_MIN_CUT_PREFILL + self.gather_rows = rocm_mxfp4_consumer_rows(self.num_candidate_cols) + logger.info_once( + "DeepSeek V4.1 indexer: candidate consumers score the %d " + "candidate positions with aiter's paged MXFP4 gather once the " + "context is %.1fx (decode) / %.1fx (prefill) as long.", + self.num_candidate_cols, + _GATHER_MIN_CUT_DECODE, + _GATHER_MIN_CUT_PREFILL, + ) + + def _gather_pays(self, context: int, cut: float) -> bool: + return self.num_candidate_cols > 0 and context >= cut * self.num_candidate_cols + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> RocmMxfp4IndexerMetadata: + metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) + if metadata.decode is not None: + metadata.decode_use_gather = self._gather_pays( + metadata.max_seq_len // self.compress_ratio, _GATHER_MIN_CUT_DECODE + ) + if metadata.prefill is not None and self.num_candidate_cols: + metadata.gather_launches = plan_gather_launches( + metadata.prefill.chunks, metadata.prefill_plans, self.gather_rows + ) + return metadata diff --git a/vllm/v1/attention/ops/rocm_paged_mxfp4_indexer.py b/vllm/v1/attention/ops/rocm_paged_mxfp4_indexer.py new file mode 100644 index 0000000000..008b4729cf --- /dev/null +++ b/vllm/v1/attention/ops/rocm_paged_mxfp4_indexer.py @@ -0,0 +1,676 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse attention indexer on aiter's paged MXFP4 MQA-logits kernel. + +gfx950 only. The kernel reads the preshuffled paged indexer K cache in place, +so no layer gathers K into a contiguous buffer, prefill included. Any model +whose indexer supplies the inputs `rocm_mxfp4_sparse_attn_indexer` documents +can run the dense path. DeepSeek-V4.1's two-level indexer adds to it: the +candidate source takes its block maxima from the same walk that writes its +logits, and the candidate consumers can walk the candidate pool instead of +the whole context: the pool is resolved once per step and shared by all of +them. +""" + +import functools +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, NamedTuple + +import torch + +import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import CUDAGraphMode +from vllm.forward_context import get_forward_context +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask, +) +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_candidate_mask_strided, + _max_decode_logits_rows, +) +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerPrefillChunkMetadata, + ) + from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( + RocmMxfp4GatherLaunch, + RocmMxfp4IndexerMetadata, + RocmMxfp4NativeDecode, + RocmMxfp4PrefillPlan, + ) + +MXFP4_BLOCK_SIZE = 32 + + +@functools.cache +def _aiter(): + """The aiter module with the paged MXFP4 MQA-logits launcher, its schedule + and cache_format.""" + from aiter.ops.triton.attention import pa_mqa_logits_mxfp4 + + return pa_mqa_logits_mxfp4 + + +@functools.cache +def _aiter_topk() -> Callable[..., None]: + """The aiter per-row top-k.""" + from aiter.ops.topk import top_k_per_row_decode + + return top_k_per_row_decode + + +class RocmPagedMxfp4CacheLayout(NamedTuple): + """The byte order of an indexer K page, as the kernel reads it, and the + shuffle pattern aiter's K cache op writes it with. + + A page is cut into runs of ``n_per_tile`` tokens. Inside a run, each token's + packed values are split into ``d_per_tile``-byte chunks along the head + dimension, and the run stores chunk 0 of every token, then chunk 1, and so + on: ``[chunk, token, byte]``. Its e8m0 scales are stored as + ``[scale % scale_lanes, token, scale // scale_lanes]``. + """ + + n_per_tile: int + d_per_tile: int + scale_lanes: int + + +# The scale order the kernel reads: aiter's mode 1, a lane's scales innermost. +# cache_format does not report the order or the lane split yet. +_SCALE_ORDER = 1 +_WAVE_SIZE = 64 + + +@functools.cache +def rocm_paged_mxfp4_cache_layout( + num_heads: int, head_dim: int, page_entries: int +) -> RocmPagedMxfp4CacheLayout: + """The K page layout from aiter's ``cache_format``, the only place vLLM + reads it, so a change there cannot reach the writer unchecked. + + Raises: + ValueError: aiter rejects the page size, or describes an order this + layout cannot express. + + """ + fmt = _aiter().cache_format(num_heads, head_dim, page_entries) + try: + n_per_tile = int(fmt["n_per_tile"]) + d_per_tile = int(fmt["d_per_tile"]) + scale_order = fmt.get("scale_mode", _SCALE_ORDER) + scale_lanes = int(fmt.get("scale_lanes", _WAVE_SIZE // n_per_tile)) + except (KeyError, TypeError, ValueError, ZeroDivisionError) as e: + raise ValueError(f"aiter's cache_format() returned {fmt!r}") from e + if ( + scale_order != _SCALE_ORDER + or min(n_per_tile, d_per_tile, scale_lanes) <= 0 + or page_entries % n_per_tile + or (head_dim // 2) % d_per_tile + or (head_dim // MXFP4_BLOCK_SIZE) % scale_lanes + ): + raise ValueError( + f"aiter's cache_format() returned {fmt!r} for {num_heads} heads of " + f"{head_dim} in {page_entries}-entry pages, a K page order " + "RocmPagedMxfp4CacheLayout cannot express" + ) + return RocmPagedMxfp4CacheLayout(n_per_tile, d_per_tile, scale_lanes) + + +def rocm_mxfp4_decode_schedule_words( + num_heads: int, head_dim: int, page_entries: int, next_n: int = 1 +) -> int: + """int32 words of the largest schedule a decode step can build, flattened or + not. The slice cap can take it past target_wgs, up to SCHED_SLOT_CAP.""" + words = 4 * _aiter().SCHED_SLOT_CAP + for rows in {1, next_n}: + config = _aiter().select_config(num_heads, head_dim, rows, page_entries) + words = max(words, 4 * config["target_wgs"]) + return words + + +def build_rocm_mxfp4_decode_schedule( + row_lens: torch.Tensor, + num_heads: int, + head_dim: int, + page_entries: int, + out: torch.Tensor, + logits_width: int, + native: "RocmMxfp4NativeDecode | None" = None, +) -> torch.Tensor | None: + """Work descriptors that even out a decode step, or None where the static + grid already fills the machine. Depends only on the rows' lengths, the + cache geometry and the logits width, so one serves every layer of a group. + The width sizes the slices, capped the way the static grid caps them.""" + if native is None: + return _aiter().build_schedule( + row_lens, + 1, + num_heads, + head_dim, + page_entries, + out=out, + max_model_len=logits_width, + ) + return _aiter().build_schedule( + native.context_lens, + native.next_n, + num_heads, + head_dim, + page_entries, + out=out, + row_ends=row_lens, + max_model_len=logits_width, + ) + + +def _kv_view(kv_cache: torch.Tensor, head_dim: int) -> torch.Tensor: + """The indexer cache as [pages, entries, 1, bytes]. The page stride is the + block-major pool's, not the page's own size.""" + num_pages, entries, width = kv_cache.shape + assert kv_cache.dtype == torch.uint8 + assert width == head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE, width + return torch.as_strided( + kv_cache, (num_pages, entries, 1, width), (kv_cache.stride(0), width, width, 1) + ) + + +def rocm_mxfp4_consumer_rows(num_candidate_cols: int) -> int: + """Query rows per candidate-consumer launch. Its logits are [rows, pool] + fp32 whatever the context, so the logits budget alone sizes it.""" + budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + return max(1, budget // (4 * num_candidate_cols)) + + +def reserve_rocm_mxfp4_indexer_workspace( + hidden_states: torch.Tensor, + logits_width: int, + candidate_block_size: int = 0, + gather_block_size: int = 0, + num_candidate_cols: int = 0, +) -> None: + """Profiling run: claim the decode logits workspace and the peak prefill + logits, block scores included when the layer writes them, candidate lists + when it gathers.""" + rows = _max_decode_logits_rows(hidden_states.shape[0]) + specs = [((rows, logits_width), torch.float32)] + budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + if candidate_block_size: + nblocks = triton.cdiv(logits_width, candidate_block_size) + specs.append(((rows, nblocks), torch.float32)) + budget += budget // candidate_block_size + if gather_block_size: + # aiter allocates the candidate lists: each gathered row, 8 B per + # candidate block (int32 slot and position). The first consumer builds + # them and they live until the step ends, across the later layers, so + # one reservation is held for the rest of the forward. + held = get_forward_context().additional_kwargs + if "rocm_mxfp4_candidate_lists" not in held: + held["rocm_mxfp4_candidate_lists"] = torch.empty( + hidden_states.shape[0] * (num_candidate_cols // gather_block_size) * 8, + dtype=torch.uint8, + device=hidden_states.device, + ) + current_workspace_manager().get_simultaneous(*specs) + torch.empty(budget, dtype=torch.uint8, device=hidden_states.device) + + +def _topk( + logits: torch.Tensor, lengths: torch.Tensor, out: torch.Tensor, k: int +) -> None: + """Top-k over each row's [0, lengths[row]); -1 pads rows shorter than k. + + aiter's kernel is faster on the candidate top-k and from 128 rows on; below + that, a wide decode step is faster on vLLM's. Shape alone decides, so a FULL + graph replays the same choice at every context length. + """ + rows, width = logits.shape + if rows == 0: + return + if width >= k and (k >= 2048 or rows >= 128): + _aiter_topk()( + logits, + 1, + lengths.reshape(-1), + out, + rows, + logits.stride(0), + logits.stride(1), + k=k, + ) + else: + torch.ops._C.top_k_per_row_decode( + logits, 1, lengths, out, rows, logits.stride(0), logits.stride(1), k + ) + + +@triton.jit +def _remap_compact_topk_kernel( + idx_ptr, + idx_stride, + pos_ptr, + pos_stride, + K: tl.constexpr, + BLOCK: tl.constexpr, + PADDED_K: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, PADDED_K) + live = cols < K + slot = tl.load(idx_ptr + row * idx_stride + cols, mask=live, other=-1) + hit = slot >= 0 + slot = tl.where(hit, slot, 0) + start = tl.load( + pos_ptr + row * pos_stride + slot // BLOCK, mask=live & hit, other=0 + ) + tl.store( + idx_ptr + row * idx_stride + cols, + tl.where(hit, start + slot % BLOCK, -1).to(tl.int32), + mask=live, + ) + + +def _remap_compact_topk( + indices: torch.Tensor, positions: torch.Tensor, block: int +) -> None: + """Candidate slots to context positions, in place: slot j sits in pool + block j // block, which starts at positions[j // block].""" + rows, k = indices.shape + if rows == 0: + return + _remap_compact_topk_kernel[(rows,)]( + indices, + indices.stride(0), + positions, + positions.stride(0), + K=k, + BLOCK=block, + PADDED_K=triton.next_power_of_2(k), + num_warps=4, + ) + + +@dataclass +class _Layer: + """One indexer layer's inputs and outputs for this step.""" + + metadata: "RocmMxfp4IndexerMetadata" + kv: torch.Tensor + q: torch.Tensor + q_scale: torch.Tensor + weights: torch.Tensor + topk_buffer: torch.Tensor + topk_tokens: int + compress_ratio: int + candidates: torch.Tensor | None + block: int + full_graph: bool + + @property + def num_heads(self) -> int: + return self.q.shape[1] + + @property + def head_dim(self) -> int: + return self.q.shape[2] * 2 + + def rows(self, lo: int, hi: int, seqs: int = 1) -> tuple[torch.Tensor, ...]: + """Q, its scales and the weights of token rows [lo, hi), as ``seqs`` + sequences of next_n rows each.""" + return ( + self.q[lo:hi].view(seqs, -1, *self.q.shape[1:]), + self.q_scale[lo:hi].view(seqs, -1, *self.q_scale.shape[1:]), + self.weights[lo:hi], + ) + + def decode_rows(self, n: int) -> tuple[torch.Tensor, ...]: + """The first n token rows, one sequence each.""" + return ( + self.q[:n].unsqueeze(1), + self.q_scale[:n].unsqueeze(1), + self.weights[:n], + ) + + +def _layer( + k_cache_prefix: str, + kv_cache: torch.Tensor, + q_values: torch.Tensor, + q_scale: torch.Tensor, + weights: torch.Tensor, + topk_indices_buffer: torch.Tensor, + topk_tokens: int, + head_dim: int, + compress_ratio: int, + candidate_blocks: torch.Tensor | None, + candidate_block_size: int, +) -> _Layer: + from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( + RocmMxfp4IndexerMetadata, + ) + + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + assert isinstance(attn_metadata, dict) + metadata = attn_metadata[k_cache_prefix] + assert isinstance(metadata, RocmMxfp4IndexerMetadata), ( + "the ROCm MXFP4 indexer needs RocmMxfp4IndexerMetadataBuilder metadata" + ) + kv = _kv_view(kv_cache, head_dim) + num_heads = q_values.shape[1] + return _Layer( + metadata=metadata, + kv=kv, + q=q_values, + # The fused Q kernel returns one int32 per head; the kernel reads the + # four ue8m0 bytes behind it. + q_scale=q_scale.view(torch.uint8).view(q_scale.shape[0], num_heads, -1), + weights=weights, + topk_buffer=topk_indices_buffer, + topk_tokens=topk_tokens, + compress_ratio=compress_ratio, + candidates=candidate_blocks, + block=candidate_block_size, + full_graph=forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL, + ) + + +def _block_scores(layer: _Layer, scores: torch.Tensor | None) -> dict: + if scores is None: + return {} + return dict( + calc_block_scores=True, block_scores=scores, candidate_block_size=layer.block + ) + + +def _dense_prefill( + layer: _Layer, + chunk: "DeepseekV32IndexerPrefillChunkMetadata", + plan: "RocmMxfp4PrefillPlan", + candidate_write: bool, +) -> None: + pa = _aiter() + t0, t1 = chunk.token_start, chunk.token_end + logits = layer.q.new_empty((t1 - t0, plan.width), dtype=torch.float32) + scores = None + if candidate_write: + nblocks = triton.cdiv(plan.width, layer.block) + scores = layer.q.new_empty((t1 - t0, nblocks), dtype=torch.float32) + # One launch for the chunk: the kernel finds each row's request from + # query_start_loc on its own grid. + pa.paged_mxfp4_mqa_logits( + layer.q[t0:t1], + layer.q_scale[t0:t1], + layer.kv, + layer.weights[t0:t1], + plan.context_lens, + chunk.block_table, + plan.width, + out_logits=logits, + clean_logits=False, + row_ends=plan.row_ends, + query_start_loc=plan.query_start_loc, + **_block_scores(layer, scores), + ) + candidates = None if layer.candidates is None else layer.candidates[t0:t1] + if scores is not None: + assert candidates is not None and plan.block_ends is not None + _topk(scores, plan.block_ends, candidates, candidates.shape[1]) + elif candidates is not None: + apply_candidate_mask(logits, None, plan.row_ends, candidates, layer.block) + _topk( + logits, + plan.row_ends, + layer.topk_buffer[t0:t1, : layer.topk_tokens], + layer.topk_tokens, + ) + + +def _dense_decode(layer: _Layer, logits_width: int, candidate_write: bool) -> None: + metadata = layer.metadata + lengths = metadata.decode_row_lens + assert metadata.decode is not None and lengths is not None + rows = lengths.shape[0] + specs = [((rows, logits_width), torch.float32)] + if candidate_write: + specs.append(((rows, triton.cdiv(logits_width, layer.block)), torch.float32)) + logits, *scores = current_workspace_manager().get_simultaneous(*specs) + native = metadata.decode_native + if native is not None: + # A request's next_n rows go in as one sequence, so a workgroup walks + # each KV tile once for all of them. + q, q_scale, weights = layer.rows(0, rows, native.context_lens.shape[0]) + context_lens, block_table, row_ends = ( + native.context_lens, + native.block_table, + lengths, + ) + else: + q, q_scale, weights = layer.decode_rows(rows) + context_lens, block_table, row_ends = lengths, metadata.decode.block_table, None + assert block_table.shape[0] == rows, "one block-table row per query row" + _aiter().paged_mxfp4_mqa_logits( + q, + q_scale, + layer.kv, + weights, + context_lens, + block_table, + logits_width, + out_logits=logits, + clean_logits=False, + row_ends=row_ends, + schedule=metadata.decode_schedule, + **_block_scores(layer, scores[0] if scores else None), + ) + candidates = None if layer.candidates is None else layer.candidates[:rows] + if scores: + assert candidates is not None and metadata.decode_block_ends is not None + _topk(scores[0], metadata.decode_block_ends, candidates, candidates.shape[1]) + elif candidates is not None: + _apply_candidate_mask_strided(logits, None, lengths, candidates, layer.block) + _topk( + logits, + lengths, + layer.topk_buffer[:rows, : layer.topk_tokens], + layer.topk_tokens, + ) + + +def _gather_prefill( + layer: _Layer, launch: "RocmMxfp4GatherLaunch", num_cols: int +) -> None: + pa = _aiter() + assert layer.candidates is not None + t0, t1 = launch.token_start, launch.token_end + if launch.pool is None: + launch.pool = pa.build_candidate_gather( + layer.candidates[t0:t1], + launch.row_ends, + launch.block_table.expand(t1 - t0, -1), + layer.kv, + layer.num_heads, + layer.head_dim, + layer.block, + ) + gather, slot_ends = launch.pool + compact = layer.q.new_empty((t1 - t0, num_cols), dtype=torch.float32) + q, q_scale, weights = layer.rows(t0, t1) + pa.paged_mxfp4_mqa_logits( + q, + q_scale, + layer.kv, + weights, + launch.context_len, + launch.block_table, + num_cols, + out_logits=compact, + clean_logits=False, + row_ends=slot_ends, + use_gather=True, + candidates=gather, + ) + out = layer.topk_buffer[t0:t1, : layer.topk_tokens] + _topk(compact, slot_ends, out, layer.topk_tokens) + _remap_compact_topk(out, gather["positions"], layer.block) + + +def _gather_decode(layer: _Layer, num_cols: int) -> None: + pa = _aiter() + assert layer.candidates is not None + metadata = layer.metadata + lengths = metadata.decode_row_lens + assert metadata.decode is not None and lengths is not None + block_table = metadata.decode.block_table + rows = lengths.shape[0] + if metadata.decode_gather is None: + metadata.decode_gather = pa.build_candidate_gather( + layer.candidates[:rows], + lengths, + block_table, + layer.kv, + layer.num_heads, + layer.head_dim, + layer.block, + ) + gather, slot_ends = metadata.decode_gather + (compact,) = current_workspace_manager().get_simultaneous( + ((rows, num_cols), torch.float32) + ) + q, q_scale, weights = layer.decode_rows(rows) + pa.paged_mxfp4_mqa_logits( + q, + q_scale, + layer.kv, + weights, + lengths, + block_table, + num_cols, + out_logits=compact, + clean_logits=False, + row_ends=slot_ends, + use_gather=True, + candidates=gather, + ) + out = layer.topk_buffer[:rows, : layer.topk_tokens] + _topk(compact, slot_ends, out, layer.topk_tokens) + _remap_compact_topk(out, gather["positions"], layer.block) + + +@eager_break_during_capture +def rocm_mxfp4_sparse_attn_indexer( + hidden_states: torch.Tensor, + k_cache_prefix: str, + kv_cache: torch.Tensor, + q_values: torch.Tensor, + q_scale: torch.Tensor, + weights: torch.Tensor, + topk_tokens: int, + head_dim: int, + max_model_len: int, + topk_indices_buffer: torch.Tensor, + compress_ratio: int = 1, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, +) -> torch.Tensor: + """Dense indexer: every layer that scores the whole context, the + candidate source included. With ``candidate_blocks`` and not + ``candidate_write`` the scores are masked to the pool first, as the + shared path does. + + The model supplies ``q_values`` [T, H, D // 2] packed e2m1, ``q_scale`` + holding each head's D // 32 ue8m0 bytes (one int32 per head at D = 128), + fp32 ``weights`` with the softmax and head scales folded in, and + ``kv_cache`` pages in `rocm_paged_mxfp4_cache_layout` order for H heads. + """ + if not isinstance(get_forward_context().attn_metadata, dict): + reserve_rocm_mxfp4_indexer_workspace( + hidden_states, + max_model_len, + candidate_block_size if candidate_write else 0, + ) + return topk_indices_buffer + layer = _layer( + k_cache_prefix, + kv_cache, + q_values, + q_scale, + weights, + topk_indices_buffer, + topk_tokens, + head_dim, + compress_ratio, + candidate_blocks, + candidate_block_size, + ) + metadata = layer.metadata + topk_indices_buffer[: hidden_states.shape[0]] = -1 + if metadata.prefill is not None: + for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans): + _dense_prefill(layer, chunk, plan, candidate_write) + if metadata.decode is not None: + _dense_decode(layer, max_model_len, candidate_write) + return topk_indices_buffer + + +@eager_break_during_capture +def rocm_mxfp4_sparse_mqa_indexer( + hidden_states: torch.Tensor, + k_cache_prefix: str, + kv_cache: torch.Tensor, + q_values: torch.Tensor, + q_scale: torch.Tensor, + weights: torch.Tensor, + topk_tokens: int, + head_dim: int, + max_model_len: int, + topk_indices_buffer: torch.Tensor, + compress_ratio: int, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + num_candidate_cols: int, +) -> torch.Tensor: + """Candidate consumer: score only the source's pool where the builder's + length gate says it pays, else the dense walk masked to the pool.""" + if not isinstance(get_forward_context().attn_metadata, dict): + reserve_rocm_mxfp4_indexer_workspace( + hidden_states, + max(max_model_len, num_candidate_cols), + gather_block_size=candidate_block_size, + num_candidate_cols=num_candidate_cols, + ) + return topk_indices_buffer + layer = _layer( + k_cache_prefix, + kv_cache, + q_values, + q_scale, + weights, + topk_indices_buffer, + topk_tokens, + head_dim, + compress_ratio, + candidate_blocks, + candidate_block_size, + ) + metadata = layer.metadata + topk_indices_buffer[: hidden_states.shape[0]] = -1 + if metadata.prefill is not None: + for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans): + if not plan.use_gather: + _dense_prefill(layer, chunk, plan, candidate_write=False) + for launch in metadata.gather_launches: + _gather_prefill(layer, launch, num_candidate_cols) + if metadata.decode is not None: + # A FULL graph cannot follow the per-step gate; the gather is the side + # that stays flat in the context length. + if metadata.decode_use_gather or layer.full_graph: + _gather_decode(layer, num_candidate_cols) + else: + _dense_decode(layer, max_model_len, candidate_write=False) + return topk_indices_buffer