@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 63b9931b73

[ROCm][DSv4.1] Paged MXFP4 sparse indexer on aiter's MQA-logits kernel (#58671)

Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Shanshan Shen <467638484@qq.com>

Mehmet Cagri committed Sep 29, 2026 at 01:37 UTC 63b9931b7399c9bfa7a8f7cdd60e2991ba87fbf9
9 files changed +2276 -9
tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py new
+599
@@ -0,0 +1,599 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek V4.1's ROCm MXFP4 sparse indexer against a torch reference.
4 +
5 +The indexer K cache is a view into vLLM's block-major KV pool: its pages are
6 +a block stride apart with other layers' pages in between, so the tests put
7 +it there. The reference reads the keys back from a second cache the same
8 +writer fills in natural order, so it shares none of the preshuffled
9 +addressing it checks.
10 +"""
11 +
12 +import itertools
13 +import types
14 +
15 +import pytest
16 +import torch
17 +
18 +from vllm.platforms import current_platform
19 +
20 +if not current_platform.is_rocm():
21 + pytest.skip("ROCm-only", allow_module_level=True)
22 +
23 +from vllm.config import CUDAGraphMode
24 +from vllm.models.deepseek_v41.amd.rocm import (
25 + rocm_mxfp4_indexer_k_store,
26 + rocm_mxfp4_indexer_q_quant,
27 +)
28 +from vllm.platforms.rocm import get_cdna_version
29 +from vllm.utils.math_utils import cdiv
30 +from vllm.utils.torch_utils import set_random_seed
31 +from vllm.v1.attention.backends.mla.indexer import (
32 + DeepseekV32IndexerPrefillChunkMetadata,
33 +)
34 +from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
35 + RocmMxfp4IndexerMetadata,
36 + native_decode,
37 + plan_gather_launches,
38 + plan_prefill_chunks,
39 +)
40 +from vllm.v1.attention.ops import rocm_paged_mxfp4_indexer as ops
41 +from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager
42 +
43 +if get_cdna_version() != 4:
44 + pytest.skip("CDNA4-only", allow_module_level=True)
45 +
46 +HEADS, HEAD_DIM = 32, 128
47 +WIDTH = HEAD_DIM // 2 + HEAD_DIM // 32
48 +# V4.1's 8-token blocks, with the pool (2048 blocks) and top-k (512) scaled
49 +# down so short test contexts still select.
50 +CAND_BLOCK, POOL_BLOCKS, TOPK = 8, 16, 32
51 +MAX_LEN = 4096
52 +DEVICE = "cuda"
53 +
54 +_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
55 +E2M1 = torch.cat([_E2M1, -_E2M1])
56 +
57 +
58 +def _dequant(packed, e8m0):
59 + """[..., D/2] e2m1 pairs and [..., D/32] e8m0 to [..., D] fp32."""
60 + table = E2M1.to(packed.device)
61 + vals = torch.stack([table[(packed & 0xF).long()], table[(packed >> 4).long()]], -1)
62 + scale = torch.pow(2.0, e8m0.float() - 127)
63 + return (vals.flatten(-2).unflatten(-1, (-1, 32)) * scale[..., None]).flatten(-2)
64 +
65 +
66 +@pytest.fixture(autouse=True)
67 +def _workspace():
68 + init_workspace_manager(torch.device(DEVICE))
69 + yield
70 + reset_workspace_manager()
71 +
72 +
73 +class _Case:
74 + """Requests with their keys written into a block-major pool (ratio-2 and
75 + ratio-1 pages side by side in every block) and into natural-order caches
76 + the reference reads. The bf16 inputs are kept for the end-to-end test."""
77 +
78 + def __init__(self, seq_lens, block):
79 + from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import (
80 + k_norm_rope_mxfp4_cache,
81 + )
82 +
83 + self.seq_lens = seq_lens
84 + num_blocks = sum(cdiv(n, block) for n in seq_lens) + 3
85 + perm = torch.randperm(num_blocks)
86 + width = max(cdiv(n, block) for n in seq_lens)
87 + self.block_table = torch.zeros(len(seq_lens), width, dtype=torch.int32)
88 + used = 0
89 + for req, n in enumerate(seq_lens):
90 + self.block_table[req, : cdiv(n, block)] = perm[used : used + cdiv(n, block)]
91 + used += cdiv(n, block)
92 + self.block_table = self.block_table.to(DEVICE)
93 +
94 + def align(n):
95 + return cdiv(n, 576) * 576
96 +
97 + page2, page1 = block // 2 * WIDTH, block * WIDTH
98 + stride = align(page2) + align(page1) + 576
99 + self.pool = torch.zeros(num_blocks * stride, dtype=torch.uint8, device=DEVICE)
100 + self.page_bytes = torch.zeros_like(self.pool, dtype=torch.bool)
101 + for i in range(num_blocks):
102 + base = i * stride
103 + self.page_bytes[base : base + page2] = True
104 + base += align(page2)
105 + self.page_bytes[base : base + page1] = True
106 + self.cache = {
107 + 2: torch.as_strided(
108 + self.pool, (num_blocks, block // 2, WIDTH), (stride, WIDTH, 1), 0
109 + ),
110 + 1: torch.as_strided(
111 + self.pool,
112 + (num_blocks, block, WIDTH),
113 + (stride, WIDTH, 1),
114 + align(page2),
115 + ),
116 + }
117 + self.offsets = [0]
118 + for n in seq_lens:
119 + self.offsets.append(self.offsets[-1] + n)
120 + # the same keys in natural order, one entry per page
121 + self.natural = {
122 + r: torch.zeros(self.offsets[-1], 1, WIDTH, dtype=torch.uint8, device=DEVICE)
123 + for r in (1, 2)
124 + }
125 + self.cos_sin = cos_sin = torch.randn(MAX_LEN, 64, device=DEVICE)
126 + self.norm = norm = torch.rand(HEAD_DIM, device=DEVICE) + 0.5
127 + self.k_pre = {}
128 + for req, n in enumerate(seq_lens):
129 + if n == 0:
130 + continue
131 + k_pre = torch.randn(n, HEAD_DIM, device=DEVICE).to(torch.bfloat16)
132 + self.k_pre[req] = k_pre
133 + pos = torch.arange(n, device=DEVICE)
134 + for ratio, cache in self.cache.items():
135 + comp = pos // ratio
136 + boundary = (pos + 1) % ratio == 0
137 + entries = cache.shape[1]
138 + page = self.block_table[req, comp // entries].long()
139 + slot = torch.where(boundary, page * entries + comp % entries, -1)
140 + nat = torch.where(boundary, self.offsets[req] + comp, -1)
141 + rocm_mxfp4_indexer_k_store(
142 + k_pre,
143 + pos,
144 + cos_sin,
145 + norm,
146 + 1e-6,
147 + cache,
148 + slot,
149 + ratio,
150 + True,
151 + num_heads=HEADS,
152 + )
153 + k_norm_rope_mxfp4_cache(
154 + k_pre, pos, cos_sin, norm, 1e-6, self.natural[ratio], nat, ratio
155 + )
156 +
157 + def keys(self, ratio, req, n):
158 + nat = self.natural[ratio][self.offsets[req] : self.offsets[req] + n, 0]
159 + return _dequant(nat[:, : HEAD_DIM // 2], nat[:, HEAD_DIM // 2 :])
160 +
161 +
162 +def _queries(num_rows):
163 + q = torch.randint(0, 256, (num_rows, HEADS, HEAD_DIM // 2), dtype=torch.uint8)
164 + scales = torch.randint(122, 130, (num_rows, HEADS, HEAD_DIM // 32))
165 + q, scales = q.to(DEVICE), scales.to(torch.uint8).to(DEVICE)
166 + weights = torch.randn(num_rows, HEADS, device=DEVICE)
167 + # the fused Q kernel returns the scales as one int32 per head
168 + return q, scales.view(torch.int32).squeeze(-1), weights, _dequant(q, scales)
169 +
170 +
171 +def _scores(q, w, keys):
172 + return (torch.relu(q @ keys.T) * w[:, None]).sum(0)
173 +
174 +
175 +def _assert_topk(row, scores, k):
176 + """``row`` is a top-k of the entries of ``scores`` above -inf, up to the
177 + fp32 reduce order."""
178 + valid = scores > float("-inf")
179 + want = min(k, int(valid.sum()))
180 + sel = row[row >= 0].long()
181 + assert sel.numel() == want and sel.unique().numel() == want
182 + if want:
183 + assert bool(valid[sel].all())
184 + kth = scores[valid].topk(want).values[-1]
185 + finite = scores[torch.isfinite(scores)]
186 + tol = 1e-4 * float(finite.abs().max()) if finite.numel() else 0.0
187 + assert bool((scores[sel] >= kth - tol).all())
188 +
189 +
190 +def _block_maxima(scores):
191 + end = scores.numel()
192 + if end == 0:
193 + return scores
194 + padded = torch.full((cdiv(end, CAND_BLOCK) * CAND_BLOCK,), float("-inf"))
195 + padded[:end] = scores.cpu()
196 + out = padded.view(-1, CAND_BLOCK).amax(1)
197 + out[(end - 1) // CAND_BLOCK] = float("inf")
198 + return out.to(scores.device)
199 +
200 +
201 +def _forward_context(monkeypatch, metadata):
202 + context = types.SimpleNamespace(
203 + attn_metadata={"indexer": metadata},
204 + cudagraph_runtime_mode=CUDAGraphMode.NONE,
205 + )
206 + monkeypatch.setattr(ops, "get_forward_context", lambda: context)
207 +
208 +
209 +def _decode_metadata(case, rows, ratio, query_lens):
210 + lens = torch.tensor([(p + 1) // ratio for _, p in rows], dtype=torch.int32)
211 + lens = lens.to(DEVICE)
212 + block_table = case.block_table[[req for req, _ in rows]].contiguous()
213 + entries = case.cache[ratio].shape[1]
214 + next_n = max(query_lens)
215 + context_lens = (torch.tensor(case.seq_lens, dtype=torch.int32) // ratio).to(DEVICE)
216 + native = native_decode(lens, block_table, query_lens, next_n, context_lens)
217 + schedule = torch.empty(
218 + ops.rocm_mxfp4_decode_schedule_words(HEADS, HEAD_DIM, entries, next_n),
219 + dtype=torch.int32,
220 + device=DEVICE,
221 + )
222 + return RocmMxfp4IndexerMetadata(
223 + seq_lens=None,
224 + max_seq_len=max(case.seq_lens),
225 + slot_mapping=None,
226 + num_decodes=len(case.seq_lens),
227 + num_decode_tokens=len(rows),
228 + num_prefills=0,
229 + num_prefill_tokens=0,
230 + decode=types.SimpleNamespace(block_table=block_table, seq_lens=lens[:, None]),
231 + decode_row_lens=lens,
232 + decode_block_ends=(lens + CAND_BLOCK - 1) // CAND_BLOCK,
233 + decode_native=native,
234 + decode_schedule=ops.build_rocm_mxfp4_decode_schedule(
235 + lens, HEADS, HEAD_DIM, entries, schedule, MAX_LEN // ratio, native
236 + ),
237 + )
238 +
239 +
240 +def _prefill_metadata(
241 + case, rows, ratio, chunk_bounds, query_start_loc, gather_rows, gather
242 +):
243 + chunks = []
244 + for req_lo, req_hi, t0, t1 in chunk_bounds:
245 + ends = torch.tensor(
246 + [(rows[t][1] + 1) // ratio for t in range(t0, t1)], dtype=torch.int32
247 + ).to(DEVICE)
248 + chunks.append(
249 + DeepseekV32IndexerPrefillChunkMetadata(
250 + block_table=case.block_table[req_lo:req_hi],
251 + cu_seqlen_ks=torch.zeros_like(ends),
252 + cu_seqlen_ke=ends,
253 + cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE),
254 + token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE),
255 + total_seq_lens=1,
256 + token_start=t0,
257 + token_end=t1,
258 + num_reqs=req_hi - req_lo,
259 + )
260 + )
261 + seq_lens = torch.tensor(case.seq_lens)
262 + context_lens = (seq_lens // ratio).int().to(DEVICE)
263 + plans = plan_prefill_chunks(
264 + chunks,
265 + query_start_loc,
266 + seq_lens,
267 + context_lens,
268 + ratio,
269 + CAND_BLOCK,
270 + 0.0 if gather else None,
271 + torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE),
272 + )
273 + return RocmMxfp4IndexerMetadata(
274 + seq_lens=None,
275 + max_seq_len=max(case.seq_lens),
276 + slot_mapping=None,
277 + num_decodes=0,
278 + num_decode_tokens=0,
279 + num_prefills=len(case.seq_lens),
280 + num_prefill_tokens=len(rows),
281 + prefill=types.SimpleNamespace(chunks=chunks),
282 + prefill_plans=plans,
283 + gather_launches=plan_gather_launches(chunks, plans, gather_rows),
284 + )
285 +
286 +
287 +# The kernel prefers 128-token pages.
288 +BLOCKS = pytest.mark.parametrize("block", [128, 64])
289 +
290 +
291 +@BLOCKS
292 +def test_k_store_stays_inside_its_pages(block):
293 + """The writer touches only its pages of the block-major pool. What they hold
294 + is checked by the layer tests, which read them back through aiter's kernel."""
295 + set_random_seed(0)
296 + case = _Case([700, 333], block)
297 + assert int(case.pool[~case.page_bytes].count_nonzero()) == 0
298 +
299 +
300 +@pytest.mark.parametrize(
301 + "rows, width, k",
302 + # vLLM's kernel below 128 rows, aiter's from 128 rows and on every k=2048
303 + [(8, 4096, 512), (128, 4096, 512), (8, 16384, 2048)],
304 +)
305 +def test_topk_matches_torch_on_both_kernels(rows, width, k):
306 + set_random_seed(0)
307 + logits = torch.randn(rows, width, device=DEVICE)
308 + lengths = torch.randint(1, width + 1, (rows,), dtype=torch.int32, device=DEVICE)
309 + lengths[0] = k // 2 # shorter than k, so the rest of the row pads with -1
310 + out = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE)
311 + ops._topk(logits, lengths, out, k)
312 + for r in range(rows):
313 + n = min(k, int(lengths[r]))
314 + want = torch.topk(logits[r, : int(lengths[r])], n).indices.sort().values
315 + got = out[r]
316 + assert torch.equal(got[got >= 0].sort().values, want.int())
317 + assert int((got == -1).sum()) == k - n
318 +
319 +
320 +def _run_layers(monkeypatch, case, rows, metadata):
321 + """Ratio-2 dense layer, the candidate source, and a consumer both ways;
322 + each checked against the torch reference."""
323 + num_rows = len(rows)
324 + q, q_scale, weights, q_ref = _queries(num_rows)
325 + hidden = torch.empty(num_rows, 1, device=DEVICE)
326 +
327 + def topk_buffer():
328 + return torch.full((num_rows + 2, TOPK), 7, dtype=torch.int32, device=DEVICE)
329 +
330 + _forward_context(monkeypatch, metadata(2))
331 + out = topk_buffer()
332 + ops.rocm_mxfp4_sparse_attn_indexer(
333 + hidden,
334 + "indexer",
335 + case.cache[2],
336 + q,
337 + q_scale,
338 + weights,
339 + TOPK,
340 + HEAD_DIM,
341 + MAX_LEN // 2,
342 + out,
343 + compress_ratio=2,
344 + )
345 + for i, (req, pos) in enumerate(rows):
346 + end = (pos + 1) // 2
347 + keys = case.keys(2, req, end)
348 + _assert_topk(out[i], _scores(q_ref[i], weights[i], keys), TOPK)
349 +
350 + source = metadata(1)
351 + _forward_context(monkeypatch, source)
352 + pool = torch.full((num_rows + 2, POOL_BLOCKS), 5, dtype=torch.int32, device=DEVICE)
353 + out = topk_buffer()
354 + ops.rocm_mxfp4_sparse_attn_indexer(
355 + hidden,
356 + "indexer",
357 + case.cache[1],
358 + q,
359 + q_scale,
360 + weights,
361 + TOPK,
362 + HEAD_DIM,
363 + MAX_LEN,
364 + out,
365 + compress_ratio=1,
366 + candidate_blocks=pool,
367 + candidate_block_size=CAND_BLOCK,
368 + candidate_write=True,
369 + )
370 + allowed = []
371 + for i, (req, pos) in enumerate(rows):
372 + scores = _scores(q_ref[i], weights[i], case.keys(1, req, pos + 1))
373 + _assert_topk(out[i], scores, TOPK)
374 + _assert_topk(pool[i], _block_maxima(scores), POOL_BLOCKS)
375 + mask = torch.zeros_like(scores, dtype=torch.bool)
376 + for block in pool[i][pool[i] >= 0].tolist():
377 + mask[block * CAND_BLOCK : (block + 1) * CAND_BLOCK] = True
378 + allowed.append(torch.where(mask, scores, float("-inf")))
379 +
380 + pa = ops._aiter()
381 + build = pa.build_candidate_gather
382 + builds = []
383 +
384 + def counting_build(*args, **kwargs):
385 + builds.append(1)
386 + return build(*args, **kwargs)
387 +
388 + monkeypatch.setattr(pa, "build_candidate_gather", counting_build)
389 + for gather in (True, False):
390 + consumer = metadata(1, gather)
391 + consumer.decode_use_gather = gather
392 + _forward_context(monkeypatch, consumer)
393 + builds.clear()
394 + # Two consumer layers of one step: the second reuses the first's pools.
395 + for _ in range(2):
396 + out = topk_buffer()
397 + ops.rocm_mxfp4_sparse_mqa_indexer(
398 + hidden,
399 + "indexer",
400 + case.cache[1],
401 + q,
402 + q_scale,
403 + weights,
404 + TOPK,
405 + HEAD_DIM,
406 + MAX_LEN,
407 + out,
408 + 1,
409 + pool,
410 + CAND_BLOCK,
411 + POOL_BLOCKS * CAND_BLOCK,
412 + )
413 + for i in range(num_rows):
414 + _assert_topk(out[i], allowed[i], TOPK)
415 + if gather:
416 + launches = len(consumer.gather_launches) + (consumer.decode is not None)
417 + assert len(builds) == launches
418 +
419 +
420 +@BLOCKS
421 +@pytest.mark.parametrize(
422 + "query_lens",
423 + # Uniform steps launch the dense layers on next_n-row sequences, also with
424 + # cudagraph padding (query length 0) after them; a ragged step keeps a row
425 + # per token.
426 + [[6, 6, 6, 6], [2, 2, 2], [2, 2, 2, 0], [2, 1, 2]],
427 + ids=["native6", "native2", "padded", "ragged"],
428 +)
429 +def test_decode_layers_match_reference(monkeypatch, block, query_lens):
430 + set_random_seed(0)
431 + next_n = max(query_lens)
432 + case = _Case([q and n for q, n in zip(query_lens, [900, 333, 610, 1200])], block)
433 + # a padding request has next_n rows of its own and no context
434 + rows = [
435 + (req, n - q + j if q else -1)
436 + for req, (n, q) in enumerate(zip(case.seq_lens, query_lens))
437 + for j in range(q or next_n)
438 + ]
439 + _run_layers(
440 + monkeypatch,
441 + case,
442 + rows,
443 + lambda r, gather=False: _decode_metadata(case, rows, r, query_lens),
444 + )
445 +
446 +
447 +@pytest.mark.parametrize(
448 + "seq_lens,new_tokens,chunks",
449 + [
450 + # one request past a cached prefix, then one sliced into two chunks
451 + ([760, 500], [300, 500], [(0, 1, 0, 300), (1, 2, 300, 560), (1, 2, 560, 800)]),
452 + # three requests in one chunk, the first two with equal rows
453 + ([600, 450, 900], [200, 200, 120], [(0, 3, 0, 520)]),
454 + # four requests with four different query lengths in one chunk
455 + ([900, 600, 450, 700], [150, 90, 200, 60], [(0, 4, 0, 500)]),
456 + ],
457 + ids=["sliced", "batched", "ragged"],
458 +)
459 +# The consumers rejoin a request's rows across chunks, then cut them at most
460 +# gather_rows a launch.
461 +@pytest.mark.parametrize("gather_rows", [1 << 20, 128], ids=["joined", "cut"])
462 +@BLOCKS
463 +def test_prefill_layers_match_reference(
464 + monkeypatch, seq_lens, new_tokens, chunks, gather_rows, block
465 +):
466 + set_random_seed(0)
467 + case = _Case(seq_lens, block)
468 + rows = [
469 + (req, n - q + i)
470 + for req, (n, q) in enumerate(zip(seq_lens, new_tokens))
471 + for i in range(q)
472 + ]
473 + query_start_loc = [0, *itertools.accumulate(new_tokens)]
474 + _run_layers(
475 + monkeypatch,
476 + case,
477 + rows,
478 + lambda r, gather=True: _prefill_metadata(
479 + case, rows, r, chunks, query_start_loc, gather_rows, gather
480 + ),
481 + )
482 +
483 +
484 +def _rope_gptj(x, cos_sin, positions):
485 + """GPT-J RoPE on the last 64 dims, as the indexer rotates Q and K: bf16 in
486 + and out."""
487 + x = x.bfloat16().float()
488 + cos, sin = cos_sin[positions].chunk(2, dim=-1)
489 + if x.dim() == 3:
490 + cos, sin = cos[:, None], sin[:, None]
491 + rot = x[..., -64:]
492 + even, odd = rot[..., 0::2].clone(), rot[..., 1::2].clone()
493 + rot[..., 0::2] = even * cos - odd * sin
494 + rot[..., 1::2] = odd * cos + even * sin
495 + return x.bfloat16()
496 +
497 +
498 +def _mxfp4_round_trip(x):
499 + """MXFP4 round trip of x, rounded as the writers do: a scale of
500 + 2^ceil(log2(amax / 6)) per 32 values, then the nearest e2m1 value, ties to
501 + even."""
502 + blocks = x.float().unflatten(-1, (-1, 32))
503 + amax = blocks.abs().amax(-1, keepdim=True).clamp(min=2**-126)
504 + scale = torch.exp2(torch.ceil(torch.log2(amax / 6)))
505 + y = (blocks / scale).clamp(-6, 6)
506 + # Even codes first: argmin keeps the first of two equal distances.
507 + grid = torch.cat([_E2M1[0::2], _E2M1[1::2]]).to(x.device)
508 + idx = (y.abs()[..., None] - grid).abs().argmin(-1)
509 + return (grid[idx] * y.sign() * scale).flatten(-2)
510 +
511 +
512 +def _torch_keys(case, ratio, req):
513 + """A request's compressed keys in torch: k_norm of each group's last
514 + token, rotated at the group's first position, through MXFP4."""
515 + k = case.k_pre[req][ratio - 1 :: ratio].float()
516 + k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) * case.norm
517 + pos = torch.arange(k.shape[0], device=DEVICE) * ratio
518 + return _mxfp4_round_trip(_rope_gptj(k, case.cos_sin, pos))
519 +
520 +
521 +def _topk_indices(q, weights, keys, ends, k):
522 + """Top-k of sum_h relu(q_h . key) * weights_h over each row's [0, end)."""
523 + scores = (torch.einsum("thd,cd->thc", q, keys).relu() * weights[..., None]).sum(1)
524 + cols = torch.arange(keys.shape[0], device=DEVICE)
525 + return scores.masked_fill(cols >= ends[:, None], float("-inf")).topk(k).indices
526 +
527 +
528 +def _recall(selected, reference):
529 + """Per row, the fraction of the reference top-k that ``selected`` picked."""
530 + return (selected[:, :, None] == reference[:, None, :]).any(1).float().mean(1)
531 +
532 +
533 +@BLOCKS
534 +@pytest.mark.parametrize("ratio", [2, 1])
535 +@pytest.mark.parametrize("step", ["decode", "prefill"])
536 +def test_selection_from_bf16_matches_torch(monkeypatch, step, ratio, block):
537 + """End to end from bf16 K and Q: vLLM's MXFP4 K store and Q quant, the
538 + paged pool, metadata from the planning functions the builder calls, and
539 + the indexer op pick what the same MXFP4 math picks in torch."""
540 + set_random_seed(0)
541 + if step == "decode":
542 + case = _Case([3000, 2600, 1900, 3500, 1200, 2200, 3900, 1600], block)
543 + rows = [(req, n - 1) for req, n in enumerate(case.seq_lens)]
544 + metadata = _decode_metadata(case, rows, ratio, [1] * len(rows))
545 + else:
546 + new_tokens = [300, 200, 250]
547 + case = _Case([1500, 900, 2000], block)
548 + rows = [
549 + (req, n - q + i)
550 + for req, (n, q) in enumerate(zip(case.seq_lens, new_tokens))
551 + for i in range(q)
552 + ]
553 + metadata = _prefill_metadata(
554 + case,
555 + rows,
556 + ratio,
557 + [(0, len(new_tokens), 0, len(rows))],
558 + [0, *itertools.accumulate(new_tokens)],
559 + 1 << 20,
560 + gather=False,
561 + )
562 + num_rows = len(rows)
563 + q_bf16 = torch.randn(num_rows, HEADS, HEAD_DIM, device=DEVICE).to(torch.bfloat16)
564 + w = torch.randn(num_rows, HEADS, device=DEVICE)
565 + pos = torch.tensor([p for _, p in rows], device=DEVICE)
566 + scales = (HEAD_DIM**-0.5, HEADS**-0.5)
567 +
568 + (q, q_scale), weights = rocm_mxfp4_indexer_q_quant(
569 + pos, q_bf16, case.cos_sin, w, *scales
570 + )
571 + _forward_context(monkeypatch, metadata)
572 + out = torch.full((num_rows, TOPK), 7, dtype=torch.int32, device=DEVICE)
573 + ops.rocm_mxfp4_sparse_attn_indexer(
574 + torch.empty(num_rows, 1, device=DEVICE),
575 + "indexer",
576 + case.cache[ratio],
577 + q,
578 + q_scale,
579 + weights,
580 + TOPK,
581 + HEAD_DIM,
582 + MAX_LEN // ratio,
583 + out,
584 + compress_ratio=ratio,
585 + )
586 + ends = (pos + 1) // ratio
587 + assert bool(((out >= 0) & (out < ends[:, None])).all())
588 + assert all(row.unique().numel() == TOPK for row in out)
589 +
590 + q_ref = _mxfp4_round_trip(_rope_gptj(q_bf16, case.cos_sin, pos))
591 + w_ref = w * scales[0] * scales[1]
592 + agree = []
593 + for req in range(len(case.seq_lens)):
594 + mine = torch.tensor([r == req for r, _ in rows], device=DEVICE)
595 + ref = _topk_indices(
596 + q_ref[mine], w_ref[mine], _torch_keys(case, ratio, req), ends[mine], TOPK
597 + )
598 + agree.append(_recall(out[mine], ref))
599 + assert torch.cat(agree).mean().item() >= 0.99
tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py new
+162
@@ -0,0 +1,162 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""The ROCm MXFP4 indexer reads the paged cache in place, so each prefill
4 +chunk is split back into its requests' rows, which one launch finds through
5 +the chunk's query_start_loc. A chunk is whole requests or one query slice of a
6 +long request, and both have to come back with the right rows and the right
7 +block-table rows.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.platforms import current_platform
14 +
15 +if not current_platform.is_rocm():
16 + pytest.skip("ROCm-only", allow_module_level=True)
17 +
18 +from vllm.v1.attention.backends.mla.indexer import (
19 + DeepseekV32IndexerPrefillChunkMetadata,
20 +)
21 +from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
22 + native_decode,
23 + plan_gather_launches,
24 + plan_prefill_chunks,
25 + split_prefill_chunks,
26 +)
27 +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import rocm_mxfp4_consumer_rows
28 +
29 +
30 +def _chunk(block_table, token_start, token_end, row_ends):
31 + ends = torch.tensor(row_ends, dtype=torch.int32)
32 + return DeepseekV32IndexerPrefillChunkMetadata(
33 + block_table=block_table,
34 + cu_seqlen_ks=torch.full_like(ends, 7),
35 + cu_seqlen_ke=ends + 7,
36 + cu_seq_lens=torch.zeros(1, dtype=torch.int32),
37 + token_to_seq=torch.zeros(1, dtype=torch.int32),
38 + total_seq_lens=1,
39 + token_start=token_start,
40 + token_end=token_end,
41 + num_reqs=block_table.shape[0],
42 + )
43 +
44 +
45 +# A decode (1 token), then prefills of 3, 3, 5 and 6 query tokens; the last
46 +# one is sliced into two chunks.
47 +QUERY_START_LOC = [0, 1, 4, 7, 12, 18]
48 +SEQ_LENS = torch.tensor([40, 10, 12, 21, 64])
49 +BLOCK_TABLE = torch.arange(20, dtype=torch.int32).view(5, 4)
50 +CHUNKS = [
51 + _chunk(BLOCK_TABLE[1:4], 1, 12, list(range(11))),
52 + _chunk(BLOCK_TABLE[4:5], 12, 15, [29, 30, 30]),
53 + _chunk(BLOCK_TABLE[4:5], 15, 18, [31, 31, 32]),
54 +]
55 +
56 +
57 +def _plans(min_gather_width, block=8):
58 + return plan_prefill_chunks(
59 + CHUNKS,
60 + QUERY_START_LOC,
61 + SEQ_LENS,
62 + SEQ_LENS.int() // 2,
63 + 2,
64 + block,
65 + min_gather_width,
66 + torch.tensor(QUERY_START_LOC, dtype=torch.int32),
67 + )
68 +
69 +
70 +def test_chunks_split_into_requests():
71 + plans = _plans(20.0)
72 +
73 + assert [p.requests for p in plans] == [
74 + [(0, 3, 0), (3, 6, 1), (6, 11, 2)],
75 + [(0, 3, 0)],
76 + [(0, 3, 0)],
77 + ]
78 + # every chunk launches once, with offsets local to the chunk
79 + assert [p.query_start_loc.tolist() for p in plans] == [
80 + [0, 3, 6, 11],
81 + [0, 3],
82 + [0, 3],
83 + ]
84 + # the logits are as wide as the longest compressed context in the chunk
85 + assert [p.width for p in plans] == [10, 32, 32]
86 + assert [p.context_lens.tolist() for p in plans] == [[5, 6, 10], [32], [32]]
87 + torch.testing.assert_close(plans[2].row_ends, torch.tensor([31, 31, 32]).int())
88 + torch.testing.assert_close(plans[2].block_ends, torch.tensor([4, 4, 4]).int())
89 + assert [p.use_gather for p in plans] == [False, True, True]
90 + assert [p.first_request for p in plans] == [1, 4, 4]
91 + # no candidate blocks, as in the model-neutral builder
92 + for plan in _plans(None, block=0):
93 + assert plan.block_ends is None and not plan.use_gather
94 +
95 +
96 +def test_launches_follow_the_logits_budget(monkeypatch):
97 + """Each launch's fp32 logits fit in VLLM_SPARSE_INDEXER_MAX_LOGITS_MB. A dense
98 + chunk's logits are [rows, longest context in the chunk]: K is read in place,
99 + not gathered, so requests share a launch until their rows no longer fit, and
100 + a request that does not fit alone is split on its query rows. A consumer's
101 + logits are [rows, pool]."""
102 + budget = 4 * 1000 * 30
103 + contexts, queries = torch.tensor([1000, 600, 200]), torch.tensor([10, 10, 10])
104 + assert split_prefill_chunks(contexts, queries, budget, 2) == [
105 + (slice(2, 5), slice(0, 30))
106 + ]
107 + # the third request would take the chunk past the budget
108 + queries = torch.tensor([10, 10, 11])
109 + assert split_prefill_chunks(contexts, queries, budget) == [
110 + (slice(0, 2), slice(0, 20)),
111 + (slice(2, 3), slice(0, 11)),
112 + ]
113 + assert split_prefill_chunks(torch.tensor([1000]), torch.tensor([100]), budget) == [
114 + (slice(0, 1), slice(lo, min(lo + 30, 100))) for lo in range(0, 100, 30)
115 + ]
116 + # the default 512 MiB holds 128 rows of a 1M context, or 8192 of a 16K pool
117 + assert split_prefill_chunks(
118 + torch.tensor([1 << 20]), torch.tensor([200]), 512 << 20
119 + ) == [(slice(0, 1), slice(0, 128)), (slice(0, 1), slice(128, 200))]
120 + monkeypatch.setenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512")
121 + assert rocm_mxfp4_consumer_rows(16384) == 8192
122 +
123 +
124 +def test_gather_launches_rejoin_a_requests_rows():
125 + """The consumers' logits are [rows, pool] whatever the context, so a
126 + request's rows go back together across the chunks that cut them, and a
127 + launch never mixes requests."""
128 + # the first chunk is too narrow to gather, the last request's slices rejoin
129 + (joined,) = plan_gather_launches(CHUNKS, _plans(20.0), 100)
130 + assert (joined.token_start, joined.token_end) == (12, 18)
131 + # every chunk gathers: the first chunk's three requests launch apart
132 + launches = plan_gather_launches(CHUNKS, _plans(0.0), 100)
133 + assert [(g.token_start, g.token_end) for g in launches] == [
134 + (1, 4),
135 + (4, 7),
136 + (7, 12),
137 + (12, 18),
138 + ]
139 + assert plan_gather_launches(CHUNKS, _plans(None), 100) == []
140 +
141 +
142 +def test_decode_launches_native_on_uniform_steps():
143 + """A decode step goes to the kernel as next_n-row sequences only when its
144 + shape says so: each request has next_n rows, and cudagraph padding (query
145 + length 0) comes only after them. A FULL graph captured on a uniform batch
146 + then replays the same launch on a padded one."""
147 + rows = torch.zeros(8, dtype=torch.int32)
148 + row_block_table = torch.arange(32, dtype=torch.int32).view(8, 4)
149 + context_lens = torch.tensor([9, 4, 7, 0, 5], dtype=torch.int32)
150 +
151 + native = native_decode(rows, row_block_table, [2, 2, 2, 0], 2, context_lens)
152 + assert native is not None and native.next_n == 2
153 + assert native.context_lens.tolist() == [9, 4, 7, 0]
154 + # every request's first flattened row, as a view
155 + assert native.block_table.tolist() == row_block_table[::2].tolist()
156 + assert native.block_table.data_ptr() == row_block_table.data_ptr()
157 +
158 + assert native_decode(rows, row_block_table, [2, 0, 2, 2], 2, context_lens) is None
159 + # a ragged step, rows padded past the requests, and no speculation
160 + assert native_decode(rows[:5], row_block_table, [2, 1, 2], 2, context_lens) is None
161 + assert native_decode(rows, row_block_table, [2, 2, 2], 2, context_lens) is None
162 + assert native_decode(rows[:4], row_block_table, [1] * 4, 1, context_lens) is None
vllm/config/attention.py
+4 -1
@@ -93,7 +93,10 @@ class AttentionConfig:
93 an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension
94 (the top-k runs on the kernels' bf16 logits). The sparse path costs
95 O(candidate blocks) per query regardless of context length, so it pays off
96 - for long contexts (roughly 32K tokens and beyond) and is slower below."""
96 + for long contexts (roughly 32K tokens and beyond) and is slower below.
97 + On ROCm gfx950 it runs aiter's paged MXFP4 MQA-logits kernel over the
98 + candidate pool instead, with fp32 logits, and keeps the masked dense walk
99 + for steps whose contexts are too short for the pool to pay."""
100
101 hisparse_config: HiSparseConfig | None = None
102 """HiSparse host-resident KV configuration. Setting this enables experimental
vllm/model_executor/layers/rocm_paged_mxfp4_indexer.py new
+116
@@ -0,0 +1,116 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Sparse attention indexer layers on aiter's paged MXFP4 kernels, for ROCm.
4 +
5 +With ``indexer_kv_dtype="mxfp4"`` on gfx950, the model writes the indexer K
6 +cache in the order aiter's MQA-logits kernel reads (see
7 +`rocm_paged_mxfp4_cache_layout`), and the kernel scores the cache in place.
8 +Requires `RocmMxfp4IndexerMetadataBuilder` metadata.
9 +"""
10 +
11 +import torch
12 +from torch import nn
13 +
14 +from vllm.config import get_current_vllm_config
15 +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
16 +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import (
17 + rocm_mxfp4_sparse_attn_indexer,
18 + rocm_mxfp4_sparse_mqa_indexer,
19 +)
20 +
21 +
22 +class RocmSparseAttnIndexer(SparseAttnIndexer):
23 + """`SparseAttnIndexer` whose HIP path walks the paged MXFP4 cache with
24 + aiter's kernel. A two-level indexer's source layer also publishes its
25 + candidate pool here."""
26 +
27 + def forward_hip(
28 + self,
29 + hidden_states: torch.Tensor,
30 + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
31 + k: torch.Tensor | None,
32 + weights: torch.Tensor,
33 + ):
34 + assert isinstance(q_quant, tuple) and self.skip_k_cache_insert, (
35 + "the ROCm MXFP4 indexer takes (values, scales) Q and a K cache "
36 + "the model already wrote"
37 + )
38 + q_values, q_scale = q_quant
39 + return rocm_mxfp4_sparse_attn_indexer(
40 + hidden_states,
41 + self.k_cache.prefix,
42 + self.k_cache.kv_cache,
43 + q_values,
44 + q_scale,
45 + weights,
46 + self.topk_tokens,
47 + self.head_dim,
48 + self.max_model_len,
49 + self.topk_indices_buffer,
50 + compress_ratio=self.compress_ratio,
51 + candidate_blocks=self.candidate_blocks,
52 + candidate_block_size=self.candidate_block_size,
53 + candidate_write=self.candidate_write,
54 + )
55 +
56 +
57 +class RocmSparseMQAIndexer(nn.Module):
58 + """DeepSeek-V4.1's candidate-consuming indexer on aiter's paged MXFP4
59 + MQA-logits kernel: the consumers score only the candidate pool the source
60 + layer published, read straight from the paged cache.
61 +
62 + Built and called as `SparseMQAIndexer` is, so the model can take either.
63 + """
64 +
65 + weights_dtype = torch.float32
66 + """The scorer's weights and logits stay fp32: gfx950 has no packed bf16
67 + arithmetic, so its head reduce costs the same in either dtype."""
68 +
69 + def __init__(
70 + self,
71 + k_cache,
72 + topk_tokens: int,
73 + head_dim: int,
74 + max_total_seq_len: int,
75 + topk_indices_buffer: torch.Tensor,
76 + candidate_blocks: torch.Tensor,
77 + candidate_block_size: int,
78 + ):
79 + super().__init__()
80 + self.k_cache = k_cache
81 + self.topk_tokens = topk_tokens
82 + self.head_dim = head_dim
83 + # aiter's kernel bounds a row by the model length in compressed
84 + # positions, not by the prefill buffer SparseMQAIndexer sizes.
85 + model_config = get_current_vllm_config().model_config
86 + self.max_model_len = model_config.max_model_len // k_cache.compress_ratio
87 + self.topk_indices_buffer = topk_indices_buffer
88 + self.candidate_blocks = candidate_blocks
89 + self.candidate_block_size = candidate_block_size
90 + self.num_candidate_cols = candidate_blocks.shape[1] * candidate_block_size
91 +
92 + def forward(
93 + self,
94 + hidden_states: torch.Tensor,
95 + q_quant: tuple[torch.Tensor, torch.Tensor],
96 + k: torch.Tensor | None,
97 + weights: torch.Tensor,
98 + ) -> torch.Tensor:
99 + assert k is None, "the model writes the indexer K cache"
100 + q_values, q_scale = q_quant
101 + return rocm_mxfp4_sparse_mqa_indexer(
102 + hidden_states,
103 + self.k_cache.prefix,
104 + self.k_cache.kv_cache,
105 + q_values,
106 + q_scale,
107 + weights,
108 + self.topk_tokens,
109 + self.head_dim,
110 + self.max_model_len,
111 + self.topk_indices_buffer,
112 + self.k_cache.compress_ratio,
113 + self.candidate_blocks,
114 + self.candidate_block_size,
115 + self.num_candidate_cols,
116 + )
vllm/models/deepseek_v41/amd/rocm.py
+149 -1
@@ -2,8 +2,9 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import functools
5 +from collections.abc import Callable
6 from dataclasses import dataclass
6 -from typing import cast
7 +from typing import Any, cast
8
9 import torch
10
@@ -15,8 +16,14 @@ from vllm.forward_context import get_forward_context
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
18 from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
19 +from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import (
20 + RocmSparseAttnIndexer,
21 + RocmSparseMQAIndexer,
22 +)
23 from vllm.models.deepseek_v41.attention import (
24 DeepseekV4Attention,
25 + DeepseekV4Indexer,
26 + DeepseekV4IndexerCache,
27 _replace_layer_index,
28 )
29 from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
@@ -32,6 +39,7 @@ from vllm.triton_utils import tl, triton
39 from vllm.utils.multi_stream_utils import execute_in_parallel
40 from vllm.v1.attention.backend import (
41 CommonAttentionMetadata,
42 + MultipleOf,
43 )
44 from vllm.v1.attention.backends.mla.sparse_swa import (
45 DeepseekSparseSWABackend,
@@ -47,6 +55,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
55 rocm_sparse_attn_decode,
56 rocm_sparse_attn_prefill,
57 )
58 +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import (
59 + rocm_paged_mxfp4_cache_layout,
60 +)
61 from vllm.v1.worker.workspace import current_workspace_manager
62
63 logger = init_logger(__name__)
@@ -517,11 +528,141 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
528 )
529
530
531 +@functools.cache
532 +def _aiter_indexer_cache_ops() -> tuple[Callable[..., None], Callable[..., tuple]]:
533 + """The aiter indexer key writer and query quantizer."""
534 + from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import (
535 + k_norm_rope_mxfp4_cache,
536 + )
537 + from aiter.ops.triton.rope.q_rope_mxfp4_quant import q_rope_mxfp4_quant
538 +
539 + return k_norm_rope_mxfp4_cache, q_rope_mxfp4_quant
540 +
541 +
542 +def rocm_mxfp4_indexer_k_store(
543 + k_pre: torch.Tensor,
544 + positions: torch.Tensor,
545 + cos_sin_cache: torch.Tensor,
546 + rms_norm_weight: torch.Tensor,
547 + rms_norm_eps: float,
548 + k_cache: torch.Tensor,
549 + kv_slot_mapping: torch.Tensor,
550 + compress_ratio: int,
551 + use_fp4_cache: bool,
552 + *,
553 + num_heads: int,
554 +) -> None:
555 + """`indexer_k_norm_rope_store` for the ROCm MXFP4 cache: aiter's cache op
556 + writes the key in the order its MQA-logits kernel reads with ``num_heads``
557 + query heads."""
558 + assert use_fp4_cache, "the ROCm indexer cache op writes MXFP4 only"
559 + layout = rocm_paged_mxfp4_cache_layout(num_heads, k_pre.shape[1], k_cache.shape[1])
560 + k_norm_rope_mxfp4_cache, _ = _aiter_indexer_cache_ops()
561 + k_norm_rope_mxfp4_cache(
562 + k_pre,
563 + positions,
564 + cos_sin_cache,
565 + rms_norm_weight,
566 + rms_norm_eps,
567 + k_cache,
568 + kv_slot_mapping,
569 + compress_ratio,
570 + shuffle=layout,
571 + )
572 +
573 +
574 +def rocm_mxfp4_indexer_q_quant(
575 + positions: torch.Tensor,
576 + index_q: torch.Tensor,
577 + index_q_cos_sin_cache: torch.Tensor,
578 + index_weights: torch.Tensor,
579 + index_weights_softmax_scale: float,
580 + index_weights_head_scale: float,
581 + use_fp4: bool = True,
582 + weights_out_dtype: torch.dtype = torch.float32,
583 +) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
584 + """`fused_indexer_q_rope_quant` for the ROCm MXFP4 indexer, on aiter's op:
585 + ((packed [T, H, D // 2], e8m0 as one int32 per head [T, H]), fp32
586 + weights)."""
587 + assert use_fp4 and weights_out_dtype == torch.float32, (
588 + "the ROCm indexer quantizes Q to MXFP4 and scores with fp32 weights"
589 + )
590 + _, q_rope_mxfp4_quant = _aiter_indexer_cache_ops()
591 + q_packed, q_scale, weights_out = q_rope_mxfp4_quant(
592 + index_q,
593 + positions,
594 + index_q_cos_sin_cache,
595 + index_weights,
596 + index_weights_softmax_scale * index_weights_head_scale,
597 + )
598 + return (q_packed, q_scale.view(torch.int32).squeeze(-1)), weights_out
599 +
600 +
601 +class DeepseekV41RocmMxfp4Indexer(DeepseekV4Indexer):
602 + """The indexer on aiter's paged MXFP4 cache: its K store and Q quant write
603 + in the order aiter's MQA-logits kernel reads, and its layers score with it.
604 + ``_produce_k`` and ``forward_q`` mirror ``DeepseekV4Indexer``'s except for
605 + those two calls."""
606 +
607 + mqa_cls = RocmSparseMQAIndexer
608 + attn_cls = RocmSparseAttnIndexer
609 +
610 + def _produce_k(
611 + self,
612 + latent: torch.Tensor | None,
613 + positions: torch.Tensor,
614 + rotary_emb: torch.nn.Module,
615 + ) -> None:
616 + attn_metadata = get_forward_context().attn_metadata
617 + if not isinstance(attn_metadata, dict) or latent is None:
618 + return
619 + assert self.owns_k
620 + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
621 + k_pre, _ = self.wk(latent)
622 + rocm_mxfp4_indexer_k_store(
623 + k_pre,
624 + positions,
625 + rotary_emb.cos_sin_cache,
626 + self.k_norm.weight,
627 + self.k_norm.variance_epsilon,
628 + self.k_cache.kv_cache,
629 + indexer_metadata.slot_mapping,
630 + self.compress_ratio,
631 + self.use_fp4_kv,
632 + num_heads=self.n_head,
633 + )
634 +
635 + def forward_q(
636 + self,
637 + qr: torch.Tensor | QuantizedActivation,
638 + qr_scale: torch.Tensor | None,
639 + indexer_weights: torch.Tensor,
640 + positions: torch.Tensor,
641 + rotary_emb: torch.nn.Module,
642 + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
643 + q = self._wq_b_proj(qr, qr_scale).view(-1, self.n_head, self.head_dim)
644 + (q, q_scale), weights = rocm_mxfp4_indexer_q_quant(
645 + positions,
646 + q,
647 + rotary_emb.cos_sin_cache,
648 + indexer_weights,
649 + self.softmax_scale,
650 + self.n_head**-0.5,
651 + use_fp4=self.use_fp4_kv,
652 + weights_out_dtype=self.indexer_weights_dtype,
653 + )
654 + return q, q_scale, weights
655 +
656 +
657 class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend):
658 @staticmethod
659 def get_name() -> str:
660 return "ROCM_FLASHMLA_SPARSE_DSV4"
661
662 + @staticmethod
663 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
664 + return [64, 128]
665 +
666 @staticmethod
667 def get_builder_cls() -> type[DeepseekV4SparseMLAMetadataBuilder]:
668 return DeepseekV4SparseMLAMetadataBuilder
@@ -540,6 +681,13 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
681 swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend
682 _use_aiter_sparse_mla = False
683
684 + def _indexer_cls(
685 + self, k_cache: DeepseekV4IndexerCache | None
686 + ) -> type[DeepseekV4Indexer]:
687 + if k_cache is not None and k_cache.rocm_mxfp4:
688 + return DeepseekV41RocmMxfp4Indexer
689 + return super()._indexer_cls(k_cache)
690 +
691 def __init__(self, *args, **kwargs):
692 vllm_config = args[0] if args else kwargs["vllm_config"]
693 super().__init__(*args, **kwargs)
vllm/models/deepseek_v41/attention.py
+34 -6
@@ -35,6 +35,9 @@ from vllm.models.deepseek_v41.common.ops import (
35
36 if TYPE_CHECKING:
37 from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
38 + from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import (
39 + RocmSparseMQAIndexer,
40 + )
41 from vllm.v1.attention.backends.mla.sparse_swa import (
42 DeepseekSparseSWAMetadata,
43 )
@@ -261,6 +264,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
264 """Return whether this instance stores fp8 KV in fp8_ds_mla layout."""
265 return self.use_fp8_ds_mla_layout
266
267 + def _indexer_cls(
268 + self, k_cache: "DeepseekV4IndexerCache | None"
269 + ) -> type["DeepseekV4Indexer"]:
270 + """The indexer for this layer's index K cache. A platform subclass
271 + returns one that writes and scores its own cache layout."""
272 + return DeepseekV4Indexer
273 +
274 def __init__(
275 self,
276 vllm_config: VllmConfig,
@@ -449,7 +459,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
459 )
460 is_candidate_source = layer_id == self.candidate_source_layer
461 uses_candidates = 0 <= self.candidate_source_layer < layer_id
452 - self.indexer = DeepseekV4Indexer(
462 + self.indexer = self._indexer_cls(index_k_cache)(
463 vllm_config,
464 config=config,
465 hidden_size=self.hidden_size,
@@ -1152,6 +1162,10 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1162 self.compress_ratio = compress_ratio
1163 vllm_config = get_current_vllm_config()
1164 self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits
1165 + # aiter's paged MXFP4 kernels write this cache and read it in place.
1166 + self.rocm_mxfp4 = current_platform.is_rocm() and dsa_indexer_uses_fp4(
1167 + vllm_config
1168 + )
1169 compilation_config = vllm_config.compilation_config
1170 if prefix in compilation_config.static_forward_context:
1171 raise ValueError(f"Duplicate layer name: {prefix}")
@@ -1190,13 +1204,21 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1204 # it 512B-aligned; the main KV record read from the same blocks
1205 # needs its own TMA stride, so keep both.
1206 block_stride_alignment=(
1193 - math.lcm(512, page_alignment) if self.sparse_logits else None
1207 + math.lcm(512, page_alignment)
1208 + if self.sparse_logits and not self.rocm_mxfp4
1209 + else None
1210 ),
1211 )
1212
1213 def forward(self): ...
1214
1215 def get_attn_backend(self) -> type[AttentionBackend]:
1216 + if self.rocm_mxfp4:
1217 + from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
1218 + DeepseekV41RocmMxfp4IndexerBackend,
1219 + )
1220 +
1221 + return DeepseekV41RocmMxfp4IndexerBackend
1222 if self.sparse_logits:
1223 return DeepseekV41SparseIndexerBackend
1224 return DeepseekV41IndexerBackend
@@ -1218,6 +1240,12 @@ class DeepseekV4Indexer(nn.Module):
1240 later indexers mask their scores to those blocks before their own top-k.
1241 """
1242
1243 + # The scoring layers; a platform subclass swaps in its own.
1244 + mqa_cls: ClassVar[type["SparseMQAIndexer | RocmSparseMQAIndexer"]] = (
1245 + SparseMQAIndexer
1246 + )
1247 + attn_cls: ClassVar[type[SparseAttnIndexer]] = SparseAttnIndexer
1248 +
1249 def __init__(
1250 self,
1251 vllm_config: VllmConfig,
@@ -1305,7 +1333,7 @@ class DeepseekV4Indexer(nn.Module):
1333 and candidate_block_buffer is not None
1334 and not candidate_write
1335 )
1308 - self.indexer_op: SparseAttnIndexer | SparseMQAIndexer
1336 + self.indexer_op: SparseAttnIndexer | SparseMQAIndexer | RocmSparseMQAIndexer
1337 if use_sparse_logits:
1338 if not self.use_fp4_kv:
1339 raise ValueError(
@@ -1314,7 +1342,7 @@ class DeepseekV4Indexer(nn.Module):
1342 )
1343 assert candidate_block_buffer is not None
1344 assert topk_indices_buffer is not None
1317 - self.indexer_op = SparseMQAIndexer(
1345 + self.indexer_op = self.mqa_cls(
1346 self.k_cache,
1347 self.topk_tokens,
1348 self.head_dim,
@@ -1324,7 +1352,7 @@ class DeepseekV4Indexer(nn.Module):
1352 candidate_block_size,
1353 )
1354 else:
1327 - self.indexer_op = SparseAttnIndexer(
1355 + self.indexer_op = self.attn_cls(
1356 self.k_cache,
1357 self.quant_block_size,
1358 self.scale_fmt,
@@ -1343,7 +1371,7 @@ class DeepseekV4Indexer(nn.Module):
1371 # The fused Q kernel writes the per-head weights in the dtype the
1372 # scoring kernels take, so no cast runs per step.
1373 self.indexer_weights_dtype = (
1346 - SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32
1374 + self.mqa_cls.weights_dtype if use_sparse_logits else torch.float32
1375 )
1376
1377 def _produce_k(
vllm/v1/attention/backends/mla/indexer.py
+23 -1
@@ -51,7 +51,7 @@ from vllm.v1.kv_cache_interface import (
51 logger = init_logger(__name__)
52
53 # The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout)
54 -# and mxfp4 is the opt-in Blackwell path.
54 +# and mxfp4 is the opt-in Blackwell and gfx950 path.
55 DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4")
56
57
@@ -64,6 +64,28 @@ def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
64 f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})."
65 )
66 use_fp4 = kv_dtype == "mxfp4"
67 + if use_fp4 and current_platform.is_rocm():
68 + from vllm._aiter_ops import rocm_aiter_ops
69 + from vllm.platforms.rocm import get_cdna_version
70 +
71 + # Only DeepSeek-V4.1 is wired to the ROCm MXFP4 cache; other DSA models
72 + # would silently keep their FP8 one.
73 + model_config = vllm_config.model_config
74 + if model_config is None or model_config.hf_config.model_type != "deepseek_v41":
75 + raise ValueError(
76 + "indexer_kv_dtype='mxfp4' on ROCm is only supported for "
77 + "DeepSeek-V4.1-Flash."
78 + )
79 + if get_cdna_version() != 4:
80 + raise ValueError(
81 + "indexer_kv_dtype='mxfp4' on ROCm requires CDNA4 (MI350X/MI355X)."
82 + )
83 + if not rocm_aiter_ops.is_enabled():
84 + raise ValueError(
85 + "indexer_kv_dtype='mxfp4' on ROCm runs on aiter's kernels; enable "
86 + "aiter with VLLM_ROCM_USE_AITER=1."
87 + )
88 + return True
89 if use_fp4 and not current_platform.is_device_capability_family(100):
90 raise ValueError(
91 "indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs "
vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py new
+513
@@ -0,0 +1,513 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Indexer metadata for the ROCm paged MXFP4 path.
4 +
5 +Extends the dense indexer metadata with what aiter's paged MXFP4 MQA-logits
6 +kernel needs to read the cache in place: each prefill chunk as its requests,
7 +packed through query_start_loc into one launch, and the decode rows as
8 +next_n-row sequences on uniform steps. `RocmMxfp4IndexerMetadataBuilder` plans
9 +this for any model's indexer, and a two-level indexer's candidate blocks once
10 +a subclass sets their size. DeepSeek-V4.1's builder does, and with
11 +``AttentionConfig.indexer_sparse_logits`` adds the per-step state the
12 +candidate consumers share.
13 +"""
14 +
15 +import bisect
16 +from dataclasses import dataclass, field, fields
17 +
18 +import torch
19 +
20 +from vllm.config import VllmConfig
21 +from vllm.logger import init_logger
22 +from vllm.v1.attention.backend import (
23 + AttentionCGSupport,
24 + CommonAttentionMetadata,
25 + MultipleOf,
26 +)
27 +from vllm.v1.attention.backends.mla.indexer import (
28 + DeepseekV32IndexerMetadata,
29 + DeepseekV32IndexerMetadataBuilder,
30 + DeepseekV32IndexerPrefillChunkMetadata,
31 + DeepseekV41IndexerBackend,
32 +)
33 +from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import (
34 + build_rocm_mxfp4_decode_schedule,
35 + rocm_mxfp4_consumer_rows,
36 + rocm_mxfp4_decode_schedule_words,
37 + rocm_paged_mxfp4_cache_layout,
38 +)
39 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
40 +
41 +logger = init_logger(__name__)
42 +
43 +# The consumers walk the pool instead of the context once the context is this
44 +# many pools long. Set from kernel timings, not end-to-end runs: the gather
45 +# loads more per key than the dense walk, so it only wins on a large cut.
46 +_GATHER_MIN_CUT_DECODE = 3.0
47 +_GATHER_MIN_CUT_PREFILL = 1.5
48 +
49 +
50 +@dataclass
51 +class RocmMxfp4PrefillPlan:
52 + """One prefill chunk as the requests the kernel launches on, all in one
53 + launch: the kernel finds each row's request through query_start_loc, and
54 + a request's rows share its KV tiles."""
55 +
56 + requests: list[tuple[int, int, int]]
57 + """(first row, end row, request) per request with rows in the chunk; rows
58 + count from ``chunk.token_start``, requests index ``chunk.block_table``."""
59 + width: int
60 + """Logits columns: an upper bound on the chunk's compressed contexts."""
61 + row_ends: torch.Tensor
62 + """[rows] int32 exclusive compressed key bound of each row."""
63 + context_lens: torch.Tensor
64 + """[chunk.num_reqs] int32 compressed context of each request."""
65 + query_start_loc: torch.Tensor
66 + """[chunk.num_reqs + 1] int32 row offsets of the chunk's requests, local to
67 + the chunk."""
68 + first_request: int
69 + """The step's index of the chunk's first request."""
70 + block_ends: torch.Tensor | None = None
71 + """[rows] int32 candidate blocks each row sees, for the source's pool."""
72 + use_gather: bool = False
73 +
74 +
75 +@dataclass
76 +class RocmMxfp4GatherLaunch:
77 + """Query rows [token_start, token_end) of one request, which the candidate
78 + consumers score against the pool in one launch."""
79 +
80 + token_start: int
81 + token_end: int
82 + block_table: torch.Tensor
83 + """[1, max_blocks] int32, the request's block table row."""
84 + context_len: torch.Tensor
85 + """[1] int32 compressed context of the request."""
86 + row_ends: torch.Tensor
87 + """[rows] int32 exclusive compressed key bound of each row."""
88 + pool: tuple[dict, torch.Tensor] | None = None
89 + """The resolved candidate lists, built by the first consumer layer and
90 + reused by the rest."""
91 +
92 +
93 +@dataclass
94 +class RocmMxfp4NativeDecode:
95 + """A decode step whose requests all have next_n query rows, as the
96 + kernel's sequences: a workgroup then walks a KV tile once for up to
97 + next_n rows instead of once per row."""
98 +
99 + next_n: int
100 + context_lens: torch.Tensor
101 + """[requests] int32 compressed context of each request."""
102 + block_table: torch.Tensor
103 + """[requests, max_blocks] int32, each request's first flattened row."""
104 +
105 +
106 +def native_decode(
107 + row_lens: torch.Tensor,
108 + row_block_table: torch.Tensor,
109 + query_lens: list[int],
110 + next_n: int,
111 + context_lens: torch.Tensor,
112 +) -> RocmMxfp4NativeDecode | None:
113 + """The step as next_n-row sequences, or None when a request has fewer rows.
114 +
115 + ``row_lens`` and ``row_block_table`` are the flattened rows' bounds and
116 + block table, ``query_lens`` the decode requests' query lengths. Trailing
117 + empty requests are cudagraph padding: they keep their next_n rows, with no
118 + context. Only the step's shape decides, so a FULL graph captured on a
119 + uniform batch replays the same launch on a padded one.
120 + """
121 + num_reqs = len(query_lens)
122 + num_full = next((i for i, n in enumerate(query_lens) if n != next_n), num_reqs)
123 + if (
124 + next_n <= 1
125 + or row_lens.shape[0] != num_reqs * next_n
126 + or any(query_lens[num_full:])
127 + ):
128 + return None
129 + return RocmMxfp4NativeDecode(
130 + next_n, context_lens[:num_reqs], row_block_table[::next_n]
131 + )
132 +
133 +
134 +@dataclass
135 +class RocmMxfp4IndexerMetadata(DeepseekV32IndexerMetadata):
136 + prefill_plans: list[RocmMxfp4PrefillPlan] = field(default_factory=list)
137 + """Parallel to ``prefill.chunks``."""
138 + decode_row_lens: torch.Tensor | None = None
139 + """[rows] int32 compressed context of each decode query row."""
140 + decode_block_ends: torch.Tensor | None = None
141 + """[rows] int32 candidate blocks each decode row sees, for the source's
142 + pool."""
143 + decode_native: RocmMxfp4NativeDecode | None = None
144 + """The same rows as next_n-row sequences, for the dense launches."""
145 + decode_schedule: torch.Tensor | None = None
146 + """The dense decode launches' work schedule, built once for the step."""
147 + decode_use_gather: bool = False
148 + """Whether the candidate consumers gather for the decode rows."""
149 + decode_gather: tuple[dict, torch.Tensor] | None = None
150 + """The pool the first consumer resolved for the decode rows."""
151 + gather_launches: list[RocmMxfp4GatherLaunch] = field(default_factory=list)
152 + """The prefill rows of the chunks that gather, as the consumers launch
153 + them."""
154 +
155 +
156 +def plan_prefill_chunks(
157 + chunks: list[DeepseekV32IndexerPrefillChunkMetadata],
158 + query_start_loc: list[int],
159 + seq_lens_cpu: torch.Tensor,
160 + context_lens: torch.Tensor,
161 + compress_ratio: int,
162 + candidate_block_size: int,
163 + min_gather_width: float | None,
164 + query_start_loc_device: torch.Tensor,
165 +) -> list[RocmMxfp4PrefillPlan]:
166 + """Split each chunk into its requests.
167 +
168 + ``query_start_loc`` and ``seq_lens_cpu`` are the step's host copies (the
169 + latter an upper bound), ``context_lens`` the device compressed lengths.
170 + A chunk covers whole requests, or a query slice of one request; either
171 + way it launches once, packed through the chunk's part of the device
172 + ``query_start_loc``.
173 + """
174 + block = candidate_block_size
175 + plans = []
176 + for chunk in chunks:
177 + t0, t1 = chunk.token_start, chunk.token_end
178 + first = bisect.bisect_right(query_start_loc, t0) - 1
179 + requests = []
180 + for req in range(chunk.num_reqs):
181 + lo = max(query_start_loc[first + req], t0) - t0
182 + hi = min(query_start_loc[first + req + 1], t1) - t0
183 + if hi > lo:
184 + requests.append((lo, hi, req))
185 + reqs = slice(first, first + chunk.num_reqs)
186 + packed = torch.clamp(query_start_loc_device[first : reqs.stop + 1], t0, t1)
187 + packed -= t0
188 + width = int(seq_lens_cpu[reqs].max()) // compress_ratio
189 + row_ends = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks
190 + plans.append(
191 + RocmMxfp4PrefillPlan(
192 + requests=requests,
193 + width=width,
194 + row_ends=row_ends,
195 + context_lens=context_lens[reqs],
196 + query_start_loc=packed,
197 + first_request=first,
198 + block_ends=(row_ends + block - 1) // block if block else None,
199 + use_gather=min_gather_width is not None and width >= min_gather_width,
200 + )
201 + )
202 + return plans
203 +
204 +
205 +def split_prefill_chunks(
206 + context_lens: torch.Tensor,
207 + query_lens: torch.Tensor,
208 + max_logits_bytes: int,
209 + request_offset: int = 0,
210 +) -> list[tuple[slice, slice]]:
211 + """(request slice, query slice) chunks, as the base chunker returns them.
212 +
213 + Nothing gathers K here, so the requests' contexts do not add up: a chunk's
214 + logits are [rows, widest compressed context] fp32, at most
215 + ``max_logits_bytes``. A request too wide to launch whole is cut on its
216 + query rows.
217 + """
218 +
219 + def max_rows(width: int) -> int:
220 + return max(1, max_logits_bytes // (4 * max(width, 1)))
221 +
222 + chunks: list[tuple[slice, slice]] = []
223 + end = 0
224 + while end < len(context_lens):
225 + start, rows, width = end, 0, 0
226 + while end < len(context_lens):
227 + q = int(query_lens[end])
228 + w = max(width, int(context_lens[end]))
229 + if rows and rows + q > max_rows(w):
230 + break
231 + rows, width = rows + q, w
232 + end += 1
233 + step = max_rows(width)
234 + reqs = slice(start + request_offset, end + request_offset)
235 + chunks.extend(
236 + (reqs, slice(lo, min(lo + step, rows))) for lo in range(0, rows, step)
237 + )
238 + return chunks
239 +
240 +
241 +def plan_gather_launches(
242 + chunks: list[DeepseekV32IndexerPrefillChunkMetadata],
243 + plans: list[RocmMxfp4PrefillPlan],
244 + max_rows: int,
245 +) -> list[RocmMxfp4GatherLaunch]:
246 + """The candidate consumers' launches over the chunks that gather.
247 +
248 + A consumer's logits are [rows, pool] however long the context, so the
249 + dense split, sized for the context, would launch it far more often than
250 + its memory needs. Each request's rows go back together across the chunks
251 + that cut them, then out again at most ``max_rows`` a launch.
252 + """
253 + # Per request: (first token, [(plan, first row, end row)]), and the chunk
254 + # and request index its block table and context come from.
255 + runs: list[tuple[int, list, DeepseekV32IndexerPrefillChunkMetadata, int]] = []
256 + last: tuple[int, int] | None = None
257 + for chunk, plan in zip(chunks, plans):
258 + if not plan.use_gather:
259 + last = None
260 + continue
261 + for lo, hi, req in plan.requests:
262 + request, t_lo = plan.first_request + req, chunk.token_start + lo
263 + if last != (request, t_lo):
264 + runs.append((t_lo, [], chunk, req))
265 + runs[-1][1].append((plan, lo, hi))
266 + last = (request, chunk.token_start + hi)
267 + launches = []
268 + for token_start, pieces, chunk, req in runs:
269 + ends = [plan.row_ends[lo:hi] for plan, lo, hi in pieces]
270 + row_ends = ends[0] if len(ends) == 1 else torch.cat(ends)
271 + block_table = chunk.block_table[req : req + 1]
272 + context_len = pieces[0][0].context_lens[req : req + 1]
273 + for lo in range(0, row_ends.shape[0], max_rows):
274 + hi = min(lo + max_rows, row_ends.shape[0])
275 + launches.append(
276 + RocmMxfp4GatherLaunch(
277 + token_start + lo,
278 + token_start + hi,
279 + block_table,
280 + context_len,
281 + row_ends[lo:hi],
282 + )
283 + )
284 + return launches
285 +
286 +
287 +class DeepseekV41RocmMxfp4IndexerBackend(DeepseekV41IndexerBackend):
288 + @staticmethod
289 + def get_name() -> str:
290 + return "DEEPSEEK_V41_ROCM_MXFP4_INDEXER"
291 +
292 + @staticmethod
293 + def get_supported_kernel_block_sizes(
294 + kv_cache_spec: KVCacheSpec | None = None,
295 + ) -> list[int | MultipleOf]:
296 + # --block-size 128 preferred.
297 + return [64, 128]
298 +
299 + @staticmethod
300 + def get_builder_cls() -> type["DeepseekV41RocmMxfp4IndexerMetadataBuilder"]:
301 + return DeepseekV41RocmMxfp4IndexerMetadataBuilder
302 +
303 +
304 +class RocmMxfp4IndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder):
305 + """Dense indexer metadata plus the in-place paged launches' layout.
306 +
307 + Requirements are checked here, at engine start.
308 + """
309 +
310 + @classmethod
311 + def get_cudagraph_support(
312 + cls,
313 + vllm_config: VllmConfig,
314 + kv_cache_spec: KVCacheSpec,
315 + ) -> AttentionCGSupport:
316 + # Decode rows are flattened into persistent buffers, and whether a
317 + # step launches them as next_n-row sequences depends on its shape
318 + # alone (see native_decode).
319 + return AttentionCGSupport.ALWAYS
320 +
321 + def __init__(
322 + self,
323 + kv_cache_spec: AttentionSpec,
324 + layer_names: list[str],
325 + vllm_config: VllmConfig,
326 + device: torch.device,
327 + **kwargs,
328 + ) -> None:
329 + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs)
330 + name = "the ROCm MXFP4 indexer"
331 + if self.dcp_world_size > 1 or self.use_pcp:
332 + raise NotImplementedError(
333 + f"{name} does not support decode/prefill context parallel."
334 + )
335 + hf_config = vllm_config.model_config.hf_text_config
336 + num_heads, head_dim = hf_config.index_n_heads, hf_config.index_head_dim
337 + self.num_heads, self.head_dim = num_heads, head_dim
338 + self.page_entries = kv_cache_spec.block_size // self.compress_ratio
339 + self.cache_layout = rocm_paged_mxfp4_cache_layout(
340 + num_heads, head_dim, self.page_entries
341 + )
342 + # Ragged steps take a row per query token, as do the candidate
343 + # consumers' gathers. Uniform steps also launch the dense layers on
344 + # next_n-row sequences, see native_decode.
345 + self.use_flattening = True
346 + # A two-level (candidate block) indexer's builder sets these.
347 + self.candidate_block_size = 0
348 + self.min_gather_width: float | None = None
349 +
350 + self.decode_block_ends_buffer = torch.zeros(
351 + self.arange_buffer.shape[0], dtype=torch.int32, device=device
352 + )
353 + self.context_lens_buffer = torch.zeros(
354 + self.arange_buffer.shape[0], dtype=torch.int32, device=device
355 + )
356 + spec_config = vllm_config.speculative_config
357 + # Adaptive verification replays a FULL graph on drafts reallocated
358 + # across requests, so a step's shape does not fix its query lengths.
359 + adaptive = spec_config is not None and spec_config.enable_adaptive_verification
360 + self.native_next_n = 1 if adaptive else self.num_speculative_tokens + 1
361 + # The indexer's logits width, which the decode schedule sizes slices by.
362 + max_model_len = vllm_config.model_config.max_model_len
363 + self.logits_width = max_model_len // self.compress_ratio
364 + # Every dense layer of this group reads the same rows through the same
365 + # page geometry, so the step builds their schedule once, not per layer.
366 + self.decode_schedule_buffer = torch.empty(
367 + rocm_mxfp4_decode_schedule_words(
368 + num_heads, head_dim, self.page_entries, self.native_next_n
369 + ),
370 + dtype=torch.int32,
371 + device=device,
372 + )
373 +
374 + def _split_indexer_prefill_chunks( # type: ignore[override]
375 + self,
376 + compressed_seq_lens_cpu: torch.Tensor,
377 + prefill_query_lens_cpu: torch.Tensor,
378 + workspace_size: int,
379 + max_logits_bytes: int,
380 + request_offset: int = 0,
381 + ) -> list[tuple[slice, slice]]:
382 + return split_prefill_chunks(
383 + compressed_seq_lens_cpu,
384 + prefill_query_lens_cpu,
385 + max_logits_bytes,
386 + request_offset,
387 + )
388 +
389 + def build(
390 + self,
391 + common_prefix_len: int,
392 + common_attn_metadata: CommonAttentionMetadata,
393 + fast_build: bool = False,
394 + ) -> RocmMxfp4IndexerMetadata:
395 + base = super().build(common_prefix_len, common_attn_metadata, fast_build)
396 + metadata = RocmMxfp4IndexerMetadata(
397 + **{f.name: getattr(base, f.name) for f in fields(base)}
398 + )
399 + cm = common_attn_metadata
400 + context_lens = self.context_lens_buffer[: cm.num_reqs]
401 + torch.floor_divide(
402 + cm.seq_lens[: cm.num_reqs], self.compress_ratio, out=context_lens
403 + )
404 + if base.decode is not None:
405 + lengths = base.decode.seq_lens.view(-1)
406 + metadata.decode_row_lens = lengths
407 + block = self.candidate_block_size
408 + if block:
409 + ends = self.decode_block_ends_buffer[: lengths.shape[0]]
410 + torch.add(lengths, block - 1, out=ends)
411 + ends.floor_divide_(block)
412 + metadata.decode_block_ends = ends
413 + query_start_loc = cm.query_start_loc_cpu[: base.num_decodes + 1]
414 + metadata.decode_native = native_decode(
415 + lengths,
416 + base.decode.block_table,
417 + torch.diff(query_start_loc).tolist(),
418 + self.native_next_n,
419 + context_lens,
420 + )
421 + metadata.decode_schedule = build_rocm_mxfp4_decode_schedule(
422 + lengths,
423 + self.num_heads,
424 + self.head_dim,
425 + self.page_entries,
426 + self.decode_schedule_buffer,
427 + self.logits_width,
428 + metadata.decode_native,
429 + )
430 + if base.prefill is not None:
431 + assert cm.seq_lens_cpu_upper_bound is not None
432 + metadata.prefill_plans = plan_prefill_chunks(
433 + base.prefill.chunks,
434 + cm.query_start_loc_cpu.tolist(),
435 + cm.seq_lens_cpu_upper_bound,
436 + context_lens,
437 + self.compress_ratio,
438 + self.candidate_block_size,
439 + self.min_gather_width,
440 + cm.query_start_loc,
441 + )
442 + return metadata
443 +
444 +
445 +class DeepseekV41RocmMxfp4IndexerMetadataBuilder(RocmMxfp4IndexerMetadataBuilder):
446 + """DeepSeek-V4.1's two-level indexer: sets the candidate block size the
447 + base plans with and, with ``indexer_sparse_logits``, adds the consumers'
448 + gathers."""
449 +
450 + def __init__(
451 + self,
452 + kv_cache_spec: AttentionSpec,
453 + layer_names: list[str],
454 + vllm_config: VllmConfig,
455 + device: torch.device,
456 + **kwargs,
457 + ) -> None:
458 + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs)
459 + hf_config = vllm_config.model_config.hf_text_config
460 + self.candidate_block_size = getattr(hf_config, "candidate_block_size", 0)
461 + self.num_candidate_cols = 0
462 + if vllm_config.attention_config.indexer_sparse_logits:
463 + prefix = "attention_config.indexer_sparse_logits"
464 + block = self.candidate_block_size
465 + topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0)
466 + layout = self.cache_layout
467 + if block <= 0 or topk_blocks <= 0:
468 + raise ValueError(
469 + f"{prefix} requires a two-level (candidate block) indexer "
470 + "model such as DeepSeek-V4.1."
471 + )
472 + if topk_blocks & (topk_blocks - 1):
473 + raise ValueError(
474 + f"{prefix} on ROCm resolves the candidate pool with a sort "
475 + f"that needs a power-of-two block count, not {topk_blocks}."
476 + )
477 + if self.page_entries % block or block > layout.n_per_tile:
478 + raise ValueError(
479 + f"{prefix} on ROCm needs {block}-entry candidate blocks to "
480 + f"tile a {self.page_entries}-entry indexer page and fit in "
481 + f"one {layout.n_per_tile}-entry shuffle group."
482 + )
483 + self.num_candidate_cols = topk_blocks * block
484 + self.min_gather_width = self.num_candidate_cols * _GATHER_MIN_CUT_PREFILL
485 + self.gather_rows = rocm_mxfp4_consumer_rows(self.num_candidate_cols)
486 + logger.info_once(
487 + "DeepSeek V4.1 indexer: candidate consumers score the %d "
488 + "candidate positions with aiter's paged MXFP4 gather once the "
489 + "context is %.1fx (decode) / %.1fx (prefill) as long.",
490 + self.num_candidate_cols,
491 + _GATHER_MIN_CUT_DECODE,
492 + _GATHER_MIN_CUT_PREFILL,
493 + )
494 +
495 + def _gather_pays(self, context: int, cut: float) -> bool:
496 + return self.num_candidate_cols > 0 and context >= cut * self.num_candidate_cols
497 +
498 + def build(
499 + self,
500 + common_prefix_len: int,
501 + common_attn_metadata: CommonAttentionMetadata,
502 + fast_build: bool = False,
503 + ) -> RocmMxfp4IndexerMetadata:
504 + metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
505 + if metadata.decode is not None:
506 + metadata.decode_use_gather = self._gather_pays(
507 + metadata.max_seq_len // self.compress_ratio, _GATHER_MIN_CUT_DECODE
508 + )
509 + if metadata.prefill is not None and self.num_candidate_cols:
510 + metadata.gather_launches = plan_gather_launches(
511 + metadata.prefill.chunks, metadata.prefill_plans, self.gather_rows
512 + )
513 + return metadata
vllm/v1/attention/ops/rocm_paged_mxfp4_indexer.py new
+676
@@ -0,0 +1,676 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Sparse attention indexer on aiter's paged MXFP4 MQA-logits kernel.
4 +
5 +gfx950 only. The kernel reads the preshuffled paged indexer K cache in place,
6 +so no layer gathers K into a contiguous buffer, prefill included. Any model
7 +whose indexer supplies the inputs `rocm_mxfp4_sparse_attn_indexer` documents
8 +can run the dense path. DeepSeek-V4.1's two-level indexer adds to it: the
9 +candidate source takes its block maxima from the same walk that writes its
10 +logits, and the candidate consumers can walk the candidate pool instead of
11 +the whole context: the pool is resolved once per step and shared by all of
12 +them.
13 +"""
14 +
15 +import functools
16 +from collections.abc import Callable
17 +from dataclasses import dataclass
18 +from typing import TYPE_CHECKING, NamedTuple
19 +
20 +import torch
21 +
22 +import vllm.envs as envs
23 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
24 +from vllm.config import CUDAGraphMode
25 +from vllm.forward_context import get_forward_context
26 +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
27 + apply_candidate_mask,
28 +)
29 +from vllm.triton_utils import tl, triton
30 +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
31 + _apply_candidate_mask_strided,
32 + _max_decode_logits_rows,
33 +)
34 +from vllm.v1.worker.workspace import current_workspace_manager
35 +
36 +if TYPE_CHECKING:
37 + from vllm.v1.attention.backends.mla.indexer import (
38 + DeepseekV32IndexerPrefillChunkMetadata,
39 + )
40 + from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
41 + RocmMxfp4GatherLaunch,
42 + RocmMxfp4IndexerMetadata,
43 + RocmMxfp4NativeDecode,
44 + RocmMxfp4PrefillPlan,
45 + )
46 +
47 +MXFP4_BLOCK_SIZE = 32
48 +
49 +
50 +@functools.cache
51 +def _aiter():
52 + """The aiter module with the paged MXFP4 MQA-logits launcher, its schedule
53 + and cache_format."""
54 + from aiter.ops.triton.attention import pa_mqa_logits_mxfp4
55 +
56 + return pa_mqa_logits_mxfp4
57 +
58 +
59 +@functools.cache
60 +def _aiter_topk() -> Callable[..., None]:
61 + """The aiter per-row top-k."""
62 + from aiter.ops.topk import top_k_per_row_decode
63 +
64 + return top_k_per_row_decode
65 +
66 +
67 +class RocmPagedMxfp4CacheLayout(NamedTuple):
68 + """The byte order of an indexer K page, as the kernel reads it, and the
69 + shuffle pattern aiter's K cache op writes it with.
70 +
71 + A page is cut into runs of ``n_per_tile`` tokens. Inside a run, each token's
72 + packed values are split into ``d_per_tile``-byte chunks along the head
73 + dimension, and the run stores chunk 0 of every token, then chunk 1, and so
74 + on: ``[chunk, token, byte]``. Its e8m0 scales are stored as
75 + ``[scale % scale_lanes, token, scale // scale_lanes]``.
76 + """
77 +
78 + n_per_tile: int
79 + d_per_tile: int
80 + scale_lanes: int
81 +
82 +
83 +# The scale order the kernel reads: aiter's mode 1, a lane's scales innermost.
84 +# cache_format does not report the order or the lane split yet.
85 +_SCALE_ORDER = 1
86 +_WAVE_SIZE = 64
87 +
88 +
89 +@functools.cache
90 +def rocm_paged_mxfp4_cache_layout(
91 + num_heads: int, head_dim: int, page_entries: int
92 +) -> RocmPagedMxfp4CacheLayout:
93 + """The K page layout from aiter's ``cache_format``, the only place vLLM
94 + reads it, so a change there cannot reach the writer unchecked.
95 +
96 + Raises:
97 + ValueError: aiter rejects the page size, or describes an order this
98 + layout cannot express.
99 +
100 + """
101 + fmt = _aiter().cache_format(num_heads, head_dim, page_entries)
102 + try:
103 + n_per_tile = int(fmt["n_per_tile"])
104 + d_per_tile = int(fmt["d_per_tile"])
105 + scale_order = fmt.get("scale_mode", _SCALE_ORDER)
106 + scale_lanes = int(fmt.get("scale_lanes", _WAVE_SIZE // n_per_tile))
107 + except (KeyError, TypeError, ValueError, ZeroDivisionError) as e:
108 + raise ValueError(f"aiter's cache_format() returned {fmt!r}") from e
109 + if (
110 + scale_order != _SCALE_ORDER
111 + or min(n_per_tile, d_per_tile, scale_lanes) <= 0
112 + or page_entries % n_per_tile
113 + or (head_dim // 2) % d_per_tile
114 + or (head_dim // MXFP4_BLOCK_SIZE) % scale_lanes
115 + ):
116 + raise ValueError(
117 + f"aiter's cache_format() returned {fmt!r} for {num_heads} heads of "
118 + f"{head_dim} in {page_entries}-entry pages, a K page order "
119 + "RocmPagedMxfp4CacheLayout cannot express"
120 + )
121 + return RocmPagedMxfp4CacheLayout(n_per_tile, d_per_tile, scale_lanes)
122 +
123 +
124 +def rocm_mxfp4_decode_schedule_words(
125 + num_heads: int, head_dim: int, page_entries: int, next_n: int = 1
126 +) -> int:
127 + """int32 words of the largest schedule a decode step can build, flattened or
128 + not. The slice cap can take it past target_wgs, up to SCHED_SLOT_CAP."""
129 + words = 4 * _aiter().SCHED_SLOT_CAP
130 + for rows in {1, next_n}:
131 + config = _aiter().select_config(num_heads, head_dim, rows, page_entries)
132 + words = max(words, 4 * config["target_wgs"])
133 + return words
134 +
135 +
136 +def build_rocm_mxfp4_decode_schedule(
137 + row_lens: torch.Tensor,
138 + num_heads: int,
139 + head_dim: int,
140 + page_entries: int,
141 + out: torch.Tensor,
142 + logits_width: int,
143 + native: "RocmMxfp4NativeDecode | None" = None,
144 +) -> torch.Tensor | None:
145 + """Work descriptors that even out a decode step, or None where the static
146 + grid already fills the machine. Depends only on the rows' lengths, the
147 + cache geometry and the logits width, so one serves every layer of a group.
148 + The width sizes the slices, capped the way the static grid caps them."""
149 + if native is None:
150 + return _aiter().build_schedule(
151 + row_lens,
152 + 1,
153 + num_heads,
154 + head_dim,
155 + page_entries,
156 + out=out,
157 + max_model_len=logits_width,
158 + )
159 + return _aiter().build_schedule(
160 + native.context_lens,
161 + native.next_n,
162 + num_heads,
163 + head_dim,
164 + page_entries,
165 + out=out,
166 + row_ends=row_lens,
167 + max_model_len=logits_width,
168 + )
169 +
170 +
171 +def _kv_view(kv_cache: torch.Tensor, head_dim: int) -> torch.Tensor:
172 + """The indexer cache as [pages, entries, 1, bytes]. The page stride is the
173 + block-major pool's, not the page's own size."""
174 + num_pages, entries, width = kv_cache.shape
175 + assert kv_cache.dtype == torch.uint8
176 + assert width == head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE, width
177 + return torch.as_strided(
178 + kv_cache, (num_pages, entries, 1, width), (kv_cache.stride(0), width, width, 1)
179 + )
180 +
181 +
182 +def rocm_mxfp4_consumer_rows(num_candidate_cols: int) -> int:
183 + """Query rows per candidate-consumer launch. Its logits are [rows, pool]
184 + fp32 whatever the context, so the logits budget alone sizes it."""
185 + budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
186 + return max(1, budget // (4 * num_candidate_cols))
187 +
188 +
189 +def reserve_rocm_mxfp4_indexer_workspace(
190 + hidden_states: torch.Tensor,
191 + logits_width: int,
192 + candidate_block_size: int = 0,
193 + gather_block_size: int = 0,
194 + num_candidate_cols: int = 0,
195 +) -> None:
196 + """Profiling run: claim the decode logits workspace and the peak prefill
197 + logits, block scores included when the layer writes them, candidate lists
198 + when it gathers."""
199 + rows = _max_decode_logits_rows(hidden_states.shape[0])
200 + specs = [((rows, logits_width), torch.float32)]
201 + budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
202 + if candidate_block_size:
203 + nblocks = triton.cdiv(logits_width, candidate_block_size)
204 + specs.append(((rows, nblocks), torch.float32))
205 + budget += budget // candidate_block_size
206 + if gather_block_size:
207 + # aiter allocates the candidate lists: each gathered row, 8 B per
208 + # candidate block (int32 slot and position). The first consumer builds
209 + # them and they live until the step ends, across the later layers, so
210 + # one reservation is held for the rest of the forward.
211 + held = get_forward_context().additional_kwargs
212 + if "rocm_mxfp4_candidate_lists" not in held:
213 + held["rocm_mxfp4_candidate_lists"] = torch.empty(
214 + hidden_states.shape[0] * (num_candidate_cols // gather_block_size) * 8,
215 + dtype=torch.uint8,
216 + device=hidden_states.device,
217 + )
218 + current_workspace_manager().get_simultaneous(*specs)
219 + torch.empty(budget, dtype=torch.uint8, device=hidden_states.device)
220 +
221 +
222 +def _topk(
223 + logits: torch.Tensor, lengths: torch.Tensor, out: torch.Tensor, k: int
224 +) -> None:
225 + """Top-k over each row's [0, lengths[row]); -1 pads rows shorter than k.
226 +
227 + aiter's kernel is faster on the candidate top-k and from 128 rows on; below
228 + that, a wide decode step is faster on vLLM's. Shape alone decides, so a FULL
229 + graph replays the same choice at every context length.
230 + """
231 + rows, width = logits.shape
232 + if rows == 0:
233 + return
234 + if width >= k and (k >= 2048 or rows >= 128):
235 + _aiter_topk()(
236 + logits,
237 + 1,
238 + lengths.reshape(-1),
239 + out,
240 + rows,
241 + logits.stride(0),
242 + logits.stride(1),
243 + k=k,
244 + )
245 + else:
246 + torch.ops._C.top_k_per_row_decode(
247 + logits, 1, lengths, out, rows, logits.stride(0), logits.stride(1), k
248 + )
249 +
250 +
251 +@triton.jit
252 +def _remap_compact_topk_kernel(
253 + idx_ptr,
254 + idx_stride,
255 + pos_ptr,
256 + pos_stride,
257 + K: tl.constexpr,
258 + BLOCK: tl.constexpr,
259 + PADDED_K: tl.constexpr,
260 +):
261 + row = tl.program_id(0).to(tl.int64)
262 + cols = tl.arange(0, PADDED_K)
263 + live = cols < K
264 + slot = tl.load(idx_ptr + row * idx_stride + cols, mask=live, other=-1)
265 + hit = slot >= 0
266 + slot = tl.where(hit, slot, 0)
267 + start = tl.load(
268 + pos_ptr + row * pos_stride + slot // BLOCK, mask=live & hit, other=0
269 + )
270 + tl.store(
271 + idx_ptr + row * idx_stride + cols,
272 + tl.where(hit, start + slot % BLOCK, -1).to(tl.int32),
273 + mask=live,
274 + )
275 +
276 +
277 +def _remap_compact_topk(
278 + indices: torch.Tensor, positions: torch.Tensor, block: int
279 +) -> None:
280 + """Candidate slots to context positions, in place: slot j sits in pool
281 + block j // block, which starts at positions[j // block]."""
282 + rows, k = indices.shape
283 + if rows == 0:
284 + return
285 + _remap_compact_topk_kernel[(rows,)](
286 + indices,
287 + indices.stride(0),
288 + positions,
289 + positions.stride(0),
290 + K=k,
291 + BLOCK=block,
292 + PADDED_K=triton.next_power_of_2(k),
293 + num_warps=4,
294 + )
295 +
296 +
297 +@dataclass
298 +class _Layer:
299 + """One indexer layer's inputs and outputs for this step."""
300 +
301 + metadata: "RocmMxfp4IndexerMetadata"
302 + kv: torch.Tensor
303 + q: torch.Tensor
304 + q_scale: torch.Tensor
305 + weights: torch.Tensor
306 + topk_buffer: torch.Tensor
307 + topk_tokens: int
308 + compress_ratio: int
309 + candidates: torch.Tensor | None
310 + block: int
311 + full_graph: bool
312 +
313 + @property
314 + def num_heads(self) -> int:
315 + return self.q.shape[1]
316 +
317 + @property
318 + def head_dim(self) -> int:
319 + return self.q.shape[2] * 2
320 +
321 + def rows(self, lo: int, hi: int, seqs: int = 1) -> tuple[torch.Tensor, ...]:
322 + """Q, its scales and the weights of token rows [lo, hi), as ``seqs``
323 + sequences of next_n rows each."""
324 + return (
325 + self.q[lo:hi].view(seqs, -1, *self.q.shape[1:]),
326 + self.q_scale[lo:hi].view(seqs, -1, *self.q_scale.shape[1:]),
327 + self.weights[lo:hi],
328 + )
329 +
330 + def decode_rows(self, n: int) -> tuple[torch.Tensor, ...]:
331 + """The first n token rows, one sequence each."""
332 + return (
333 + self.q[:n].unsqueeze(1),
334 + self.q_scale[:n].unsqueeze(1),
335 + self.weights[:n],
336 + )
337 +
338 +
339 +def _layer(
340 + k_cache_prefix: str,
341 + kv_cache: torch.Tensor,
342 + q_values: torch.Tensor,
343 + q_scale: torch.Tensor,
344 + weights: torch.Tensor,
345 + topk_indices_buffer: torch.Tensor,
346 + topk_tokens: int,
347 + head_dim: int,
348 + compress_ratio: int,
349 + candidate_blocks: torch.Tensor | None,
350 + candidate_block_size: int,
351 +) -> _Layer:
352 + from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
353 + RocmMxfp4IndexerMetadata,
354 + )
355 +
356 + forward_context = get_forward_context()
357 + attn_metadata = forward_context.attn_metadata
358 + assert isinstance(attn_metadata, dict)
359 + metadata = attn_metadata[k_cache_prefix]
360 + assert isinstance(metadata, RocmMxfp4IndexerMetadata), (
361 + "the ROCm MXFP4 indexer needs RocmMxfp4IndexerMetadataBuilder metadata"
362 + )
363 + kv = _kv_view(kv_cache, head_dim)
364 + num_heads = q_values.shape[1]
365 + return _Layer(
366 + metadata=metadata,
367 + kv=kv,
368 + q=q_values,
369 + # The fused Q kernel returns one int32 per head; the kernel reads the
370 + # four ue8m0 bytes behind it.
371 + q_scale=q_scale.view(torch.uint8).view(q_scale.shape[0], num_heads, -1),
372 + weights=weights,
373 + topk_buffer=topk_indices_buffer,
374 + topk_tokens=topk_tokens,
375 + compress_ratio=compress_ratio,
376 + candidates=candidate_blocks,
377 + block=candidate_block_size,
378 + full_graph=forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL,
379 + )
380 +
381 +
382 +def _block_scores(layer: _Layer, scores: torch.Tensor | None) -> dict:
383 + if scores is None:
384 + return {}
385 + return dict(
386 + calc_block_scores=True, block_scores=scores, candidate_block_size=layer.block
387 + )
388 +
389 +
390 +def _dense_prefill(
391 + layer: _Layer,
392 + chunk: "DeepseekV32IndexerPrefillChunkMetadata",
393 + plan: "RocmMxfp4PrefillPlan",
394 + candidate_write: bool,
395 +) -> None:
396 + pa = _aiter()
397 + t0, t1 = chunk.token_start, chunk.token_end
398 + logits = layer.q.new_empty((t1 - t0, plan.width), dtype=torch.float32)
399 + scores = None
400 + if candidate_write:
401 + nblocks = triton.cdiv(plan.width, layer.block)
402 + scores = layer.q.new_empty((t1 - t0, nblocks), dtype=torch.float32)
403 + # One launch for the chunk: the kernel finds each row's request from
404 + # query_start_loc on its own grid.
405 + pa.paged_mxfp4_mqa_logits(
406 + layer.q[t0:t1],
407 + layer.q_scale[t0:t1],
408 + layer.kv,
409 + layer.weights[t0:t1],
410 + plan.context_lens,
411 + chunk.block_table,
412 + plan.width,
413 + out_logits=logits,
414 + clean_logits=False,
415 + row_ends=plan.row_ends,
416 + query_start_loc=plan.query_start_loc,
417 + **_block_scores(layer, scores),
418 + )
419 + candidates = None if layer.candidates is None else layer.candidates[t0:t1]
420 + if scores is not None:
421 + assert candidates is not None and plan.block_ends is not None
422 + _topk(scores, plan.block_ends, candidates, candidates.shape[1])
423 + elif candidates is not None:
424 + apply_candidate_mask(logits, None, plan.row_ends, candidates, layer.block)
425 + _topk(
426 + logits,
427 + plan.row_ends,
428 + layer.topk_buffer[t0:t1, : layer.topk_tokens],
429 + layer.topk_tokens,
430 + )
431 +
432 +
433 +def _dense_decode(layer: _Layer, logits_width: int, candidate_write: bool) -> None:
434 + metadata = layer.metadata
435 + lengths = metadata.decode_row_lens
436 + assert metadata.decode is not None and lengths is not None
437 + rows = lengths.shape[0]
438 + specs = [((rows, logits_width), torch.float32)]
439 + if candidate_write:
440 + specs.append(((rows, triton.cdiv(logits_width, layer.block)), torch.float32))
441 + logits, *scores = current_workspace_manager().get_simultaneous(*specs)
442 + native = metadata.decode_native
443 + if native is not None:
444 + # A request's next_n rows go in as one sequence, so a workgroup walks
445 + # each KV tile once for all of them.
446 + q, q_scale, weights = layer.rows(0, rows, native.context_lens.shape[0])
447 + context_lens, block_table, row_ends = (
448 + native.context_lens,
449 + native.block_table,
450 + lengths,
451 + )
452 + else:
453 + q, q_scale, weights = layer.decode_rows(rows)
454 + context_lens, block_table, row_ends = lengths, metadata.decode.block_table, None
455 + assert block_table.shape[0] == rows, "one block-table row per query row"
456 + _aiter().paged_mxfp4_mqa_logits(
457 + q,
458 + q_scale,
459 + layer.kv,
460 + weights,
461 + context_lens,
462 + block_table,
463 + logits_width,
464 + out_logits=logits,
465 + clean_logits=False,
466 + row_ends=row_ends,
467 + schedule=metadata.decode_schedule,
468 + **_block_scores(layer, scores[0] if scores else None),
469 + )
470 + candidates = None if layer.candidates is None else layer.candidates[:rows]
471 + if scores:
472 + assert candidates is not None and metadata.decode_block_ends is not None
473 + _topk(scores[0], metadata.decode_block_ends, candidates, candidates.shape[1])
474 + elif candidates is not None:
475 + _apply_candidate_mask_strided(logits, None, lengths, candidates, layer.block)
476 + _topk(
477 + logits,
478 + lengths,
479 + layer.topk_buffer[:rows, : layer.topk_tokens],
480 + layer.topk_tokens,
481 + )
482 +
483 +
484 +def _gather_prefill(
485 + layer: _Layer, launch: "RocmMxfp4GatherLaunch", num_cols: int
486 +) -> None:
487 + pa = _aiter()
488 + assert layer.candidates is not None
489 + t0, t1 = launch.token_start, launch.token_end
490 + if launch.pool is None:
491 + launch.pool = pa.build_candidate_gather(
492 + layer.candidates[t0:t1],
493 + launch.row_ends,
494 + launch.block_table.expand(t1 - t0, -1),
495 + layer.kv,
496 + layer.num_heads,
497 + layer.head_dim,
498 + layer.block,
499 + )
500 + gather, slot_ends = launch.pool
501 + compact = layer.q.new_empty((t1 - t0, num_cols), dtype=torch.float32)
502 + q, q_scale, weights = layer.rows(t0, t1)
503 + pa.paged_mxfp4_mqa_logits(
504 + q,
505 + q_scale,
506 + layer.kv,
507 + weights,
508 + launch.context_len,
509 + launch.block_table,
510 + num_cols,
511 + out_logits=compact,
512 + clean_logits=False,
513 + row_ends=slot_ends,
514 + use_gather=True,
515 + candidates=gather,
516 + )
517 + out = layer.topk_buffer[t0:t1, : layer.topk_tokens]
518 + _topk(compact, slot_ends, out, layer.topk_tokens)
519 + _remap_compact_topk(out, gather["positions"], layer.block)
520 +
521 +
522 +def _gather_decode(layer: _Layer, num_cols: int) -> None:
523 + pa = _aiter()
524 + assert layer.candidates is not None
525 + metadata = layer.metadata
526 + lengths = metadata.decode_row_lens
527 + assert metadata.decode is not None and lengths is not None
528 + block_table = metadata.decode.block_table
529 + rows = lengths.shape[0]
530 + if metadata.decode_gather is None:
531 + metadata.decode_gather = pa.build_candidate_gather(
532 + layer.candidates[:rows],
533 + lengths,
534 + block_table,
535 + layer.kv,
536 + layer.num_heads,
537 + layer.head_dim,
538 + layer.block,
539 + )
540 + gather, slot_ends = metadata.decode_gather
541 + (compact,) = current_workspace_manager().get_simultaneous(
542 + ((rows, num_cols), torch.float32)
543 + )
544 + q, q_scale, weights = layer.decode_rows(rows)
545 + pa.paged_mxfp4_mqa_logits(
546 + q,
547 + q_scale,
548 + layer.kv,
549 + weights,
550 + lengths,
551 + block_table,
552 + num_cols,
553 + out_logits=compact,
554 + clean_logits=False,
555 + row_ends=slot_ends,
556 + use_gather=True,
557 + candidates=gather,
558 + )
559 + out = layer.topk_buffer[:rows, : layer.topk_tokens]
560 + _topk(compact, slot_ends, out, layer.topk_tokens)
561 + _remap_compact_topk(out, gather["positions"], layer.block)
562 +
563 +
564 +@eager_break_during_capture
565 +def rocm_mxfp4_sparse_attn_indexer(
566 + hidden_states: torch.Tensor,
567 + k_cache_prefix: str,
568 + kv_cache: torch.Tensor,
569 + q_values: torch.Tensor,
570 + q_scale: torch.Tensor,
571 + weights: torch.Tensor,
572 + topk_tokens: int,
573 + head_dim: int,
574 + max_model_len: int,
575 + topk_indices_buffer: torch.Tensor,
576 + compress_ratio: int = 1,
577 + candidate_blocks: torch.Tensor | None = None,
578 + candidate_block_size: int = 0,
579 + candidate_write: bool = False,
580 +) -> torch.Tensor:
581 + """Dense indexer: every layer that scores the whole context, the
582 + candidate source included. With ``candidate_blocks`` and not
583 + ``candidate_write`` the scores are masked to the pool first, as the
584 + shared path does.
585 +
586 + The model supplies ``q_values`` [T, H, D // 2] packed e2m1, ``q_scale``
587 + holding each head's D // 32 ue8m0 bytes (one int32 per head at D = 128),
588 + fp32 ``weights`` with the softmax and head scales folded in, and
589 + ``kv_cache`` pages in `rocm_paged_mxfp4_cache_layout` order for H heads.
590 + """
591 + if not isinstance(get_forward_context().attn_metadata, dict):
592 + reserve_rocm_mxfp4_indexer_workspace(
593 + hidden_states,
594 + max_model_len,
595 + candidate_block_size if candidate_write else 0,
596 + )
597 + return topk_indices_buffer
598 + layer = _layer(
599 + k_cache_prefix,
600 + kv_cache,
601 + q_values,
602 + q_scale,
603 + weights,
604 + topk_indices_buffer,
605 + topk_tokens,
606 + head_dim,
607 + compress_ratio,
608 + candidate_blocks,
609 + candidate_block_size,
610 + )
611 + metadata = layer.metadata
612 + topk_indices_buffer[: hidden_states.shape[0]] = -1
613 + if metadata.prefill is not None:
614 + for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans):
615 + _dense_prefill(layer, chunk, plan, candidate_write)
616 + if metadata.decode is not None:
617 + _dense_decode(layer, max_model_len, candidate_write)
618 + return topk_indices_buffer
619 +
620 +
621 +@eager_break_during_capture
622 +def rocm_mxfp4_sparse_mqa_indexer(
623 + hidden_states: torch.Tensor,
624 + k_cache_prefix: str,
625 + kv_cache: torch.Tensor,
626 + q_values: torch.Tensor,
627 + q_scale: torch.Tensor,
628 + weights: torch.Tensor,
629 + topk_tokens: int,
630 + head_dim: int,
631 + max_model_len: int,
632 + topk_indices_buffer: torch.Tensor,
633 + compress_ratio: int,
634 + candidate_blocks: torch.Tensor,
635 + candidate_block_size: int,
636 + num_candidate_cols: int,
637 +) -> torch.Tensor:
638 + """Candidate consumer: score only the source's pool where the builder's
639 + length gate says it pays, else the dense walk masked to the pool."""
640 + if not isinstance(get_forward_context().attn_metadata, dict):
641 + reserve_rocm_mxfp4_indexer_workspace(
642 + hidden_states,
643 + max(max_model_len, num_candidate_cols),
644 + gather_block_size=candidate_block_size,
645 + num_candidate_cols=num_candidate_cols,
646 + )
647 + return topk_indices_buffer
648 + layer = _layer(
649 + k_cache_prefix,
650 + kv_cache,
651 + q_values,
652 + q_scale,
653 + weights,
654 + topk_indices_buffer,
655 + topk_tokens,
656 + head_dim,
657 + compress_ratio,
658 + candidate_blocks,
659 + candidate_block_size,
660 + )
661 + metadata = layer.metadata
662 + topk_indices_buffer[: hidden_states.shape[0]] = -1
663 + if metadata.prefill is not None:
664 + for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans):
665 + if not plan.use_gather:
666 + _dense_prefill(layer, chunk, plan, candidate_write=False)
667 + for launch in metadata.gather_launches:
668 + _gather_prefill(layer, launch, num_candidate_cols)
669 + if metadata.decode is not None:
670 + # A FULL graph cannot follow the per-step gate; the gather is the side
671 + # that stays flat in the context length.
672 + if metadata.decode_use_gather or layer.full_graph:
673 + _gather_decode(layer, num_candidate_cols)
674 + else:
675 + _dense_decode(layer, max_model_len, candidate_write=False)
676 + return topk_indices_buffer