@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 25624f65cf

[ROCm][MLA] Add an AITER ASM round-robin decode route for DCP multi-token verify (#56861)

Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: seungrokj <144636725+seungrokj@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>

xiaohuguo2023 committed Sep 28, 2026 at 19:52 UTC 25624f65cf40ff6ce62e3261331b1b1ea7fc5e0a
5 files changed +1202 -30
tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Routing/selection tests for the AITER ASM cprr DCP-verify route.
4 +
5 +Companion to the upstream segmented-verify tests. Everything here is pure CPU:
6 +it pins the *decisions* (which route a KV group takes, which head count the
7 +kernel runs at, which configurations are refused) rather than kernel numerics.
8 +The numerics parity test against the segmented route needs a GPU and both
9 +routes built; it is marked and skipped when unavailable.
10 +"""
11 +
12 +from __future__ import annotations
13 +
14 +import pytest
15 +
16 +rocm_aiter_mla = pytest.importorskip(
17 + "vllm.v1.attention.backends.mla.rocm_aiter_mla",
18 + reason="ROCm AITER MLA backend not available",
19 +)
20 +
21 +_heads = rocm_aiter_mla._asm_dcp_verify_heads
22 +_configured = rocm_aiter_mla._asm_dcp_verify_configured
23 +_select_route = rocm_aiter_mla._select_dcp_decode_route
24 +Route = rocm_aiter_mla._DCPDecodeRoute
25 +NATIVE = rocm_aiter_mla._NATIVE_CPRR_HEADS
26 +MIN_QLEN = rocm_aiter_mla._MIN_CPRR_QLEN
27 +
28 +
29 +def test_dcp_verify_env_defaults_to_segmented(monkeypatch):
30 + monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False)
31 + assert rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "segmented"
32 +
33 +
34 +def test_dcp_verify_env_rejects_unknown_route(monkeypatch):
35 + monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", "unknown")
36 + with pytest.raises(ValueError, match="Invalid value"):
37 + _ = rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY
38 +
39 +
40 +# --------------------------------------------------------------------------
41 +# _asm_dcp_verify_heads -- which head count the cprr kernel runs at
42 +# --------------------------------------------------------------------------
43 +
44 +
45 +@pytest.mark.parametrize("n", NATIVE)
46 +def test_native_head_counts_are_used_as_is(n):
47 + assert _heads(n) == n
48 +
49 +
50 +def test_non_native_head_count_pads_up_to_the_next_native():
51 + """K3 at TP8/DCP8 gathers 96 heads for the target; the kernel has no 96
52 + variant, so it must run at 128, not silently fall out of the asm path."""
53 + assert _heads(96) == 128
54 + assert _heads(1) == 16
55 + assert _heads(17) == 32
56 + assert _heads(65) == 128
57 +
58 +
59 +def test_head_count_above_the_largest_native_is_unservable():
60 + """0 means 'cannot serve'; anything else would route a shape the kernel
61 + has no build for."""
62 + assert _heads(max(NATIVE) + 1) == 0
63 + assert _heads(256) == 0
64 +
65 +
66 +# --------------------------------------------------------------------------
67 +# _asm_dcp_verify_configured -- reachability
68 +# --------------------------------------------------------------------------
69 +
70 +
71 +@pytest.fixture
72 +def on_gfx950(monkeypatch):
73 + """The cprr kernels are built only for gfx950, so the gate checks the arch.
74 + Patched rather than detected, so both branches run on any runner."""
75 + import vllm.platforms.rocm as rocm
76 +
77 + monkeypatch.setattr(rocm, "on_gfx950", lambda: True)
78 +
79 +
80 +def test_route_needs_dcp(on_gfx950):
81 + assert _configured(dcp_world_size=1, cp_interleave=1, multi_token_decode=True) is (
82 + False
83 + )
84 + assert _configured(dcp_world_size=8, cp_interleave=1, multi_token_decode=True) is (
85 + True
86 + )
87 +
88 +
89 +def test_round_robin_interleave_other_than_one_is_excluded(on_gfx950):
90 + """The kernel's global-position causal window assumes interleave 1, the
91 + same restriction the segmented gate carries."""
92 + assert _configured(dcp_world_size=8, cp_interleave=2, multi_token_decode=True) is (
93 + False
94 + )
95 + assert _configured(dcp_world_size=8, cp_interleave=16, multi_token_decode=True) is (
96 + False
97 + )
98 +
99 +
100 +def test_route_needs_multi_token_decode(on_gfx950):
101 + """Single-token decode is already served. Enabling the route for it would
102 + pad the head count on every step and buy nothing."""
103 + assert _configured(dcp_world_size=8, cp_interleave=1, multi_token_decode=False) is (
104 + False
105 + )
106 +
107 +
108 +def test_route_needs_gfx950(monkeypatch):
109 + """hsa/gfx942/mla and hsa/gfx1250/mla carry no cprr rows, so off gfx950 the
110 + kernel lookup fails at the first verify step rather than falling back."""
111 + import vllm.platforms.rocm as rocm
112 +
113 + monkeypatch.setattr(rocm, "on_gfx950", lambda: False)
114 + assert _configured(dcp_world_size=8, cp_interleave=1, multi_token_decode=True) is (
115 + False
116 + )
117 +
118 +
119 +# --------------------------------------------------------------------------
120 +# qlen floor
121 +# --------------------------------------------------------------------------
122 +
123 +
124 +def test_min_cprr_qlen_is_above_two():
125 + """Qlen 2 is below CPRR's floor and must use segmented verification."""
126 + assert MIN_QLEN > 2
127 +
128 +
129 +@pytest.mark.parametrize(
130 + "supports,causal,qlen,asm,expected",
131 + [
132 + (True, True, 1, True, Route.PLAIN),
133 + (True, True, 2, True, Route.SEGMENTED),
134 + (True, True, 3, True, Route.CPRR),
135 + (True, True, 4, True, Route.CPRR),
136 + (True, True, 4, False, Route.SEGMENTED),
137 + (True, False, 2, True, Route.PLAIN),
138 + (True, False, 4, False, Route.PLAIN),
139 + (False, True, 3, True, Route.CPRR),
140 + ],
141 +)
142 +def test_dcp_decode_route(supports, causal, qlen, asm, expected):
143 + assert _select_route(supports, causal, qlen, asm) is expected
144 +
145 +
146 +def test_causal_multi_token_batch_without_a_valid_route_fails():
147 + with pytest.raises(RuntimeError, match="requires either segmented MLA"):
148 + _select_route(False, True, 2, True)
149 +
150 +
151 +# Kernel numerics live in test_rocm_aiter_mla_dcp_cprr_numerics.py, which runs
152 +# the real builder and the real asm kernel on a gfx950 GPU.
tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py new
+540
@@ -0,0 +1,540 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""GPU numerics for the AITER ASM cprr DCP-verify route.
4 +
5 +Runs the REAL metadata builder (``AiterMLAMetadataBuilder.build``, so
6 +``get_mla_metadata_v1`` is exercised with the cprr kwargs and the persistent
7 +schedule buffers are produced) and the REAL asm kernel, on ONE GPU with the DCP
8 +ranks simulated one after another.
9 +
10 +That is legitimate because the aiter MLA backend contains no collective on this
11 +path: the query all-gather happens in the *layer*, strictly before
12 +``forward_mqa``, so the head fold, the round-robin kernel call and the un-fold
13 +are pure per-rank local arithmetic in which ``cp_world_size``/``cp_rank`` are
14 +just numbers feeding the causal-mask maths. Restricting a shard to a residue
15 +class of global positions is exactly what DCP rank r sees.
16 +
17 +Three properties are checked:
18 + * each shard matches an exact torch reference over the positions it holds,
19 + with causality applied on the GLOBAL position (the thing cprr reconstructs
20 + in-kernel and the segmented route instead carries in per-row lengths);
21 + * the LSE merge of all shards reproduces full-context attention;
22 + * a shard told the wrong ``cp_rank`` fails both -- without this positive
23 + control a passing run proves nothing, since a kernel that ignored global
24 + positions still produces plausible-looking output.
25 +"""
26 +
27 +from __future__ import annotations
28 +
29 +from types import SimpleNamespace
30 +
31 +import pytest
32 +import torch
33 +
34 +from vllm.platforms import current_platform
35 +
36 +if not current_platform.is_rocm():
37 + pytest.skip("ROCm AITER MLA tests", allow_module_level=True)
38 +
39 +from vllm._aiter_ops import is_aiter_found # noqa: E402
40 +from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402
41 +
42 +KV_LORA_RANK = 512
43 +QK_ROPE_HEAD_DIM = 64
44 +QK_NOPE_HEAD_DIM = 128
45 +HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM
46 +SCALE = (QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM) ** -0.5
47 +
48 +DCP = 8
49 +NUM_REQS = 8
50 +CTX = 256
51 +QLEN = 5 # 1 + num_speculative_tokens; must clear _MIN_CPRR_QLEN
52 +
53 +# fp8-e4m3 keeps 3 mantissa bits, so the relative error floor is ~6.2e-2. Gates
54 +# are relative for the same reason: an absolute bound would just be a bound on
55 +# the output magnitude.
56 +SHARD_RTOL = 1.0e-1
57 +MERGE_RTOL = 5.0e-2
58 +
59 +MODEL = "deepseek-ai/DeepSeek-R1"
60 +
61 +
62 +def _gpu_available() -> bool:
63 + return is_aiter_found() and torch.accelerator.is_available()
64 +
65 +
66 +def _on_gfx950() -> bool:
67 + if not is_aiter_found():
68 + return False
69 + from vllm.platforms.rocm import on_gfx950
70 +
71 + return on_gfx950()
72 +
73 +
74 +requires_cprr = pytest.mark.skipif(
75 + not (_gpu_available() and _on_gfx950()),
76 + reason="cprr DCP verify needs ROCm + AITER on gfx950",
77 +)
78 +
79 +
80 +@pytest.fixture(autouse=True)
81 +def select_asm_route(monkeypatch):
82 + monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", "asm")
83 +
84 +
85 +# --------------------------------------------------------------------------
86 +# reference maths
87 +# --------------------------------------------------------------------------
88 +
89 +
90 +def _torch_reference(q_f32, kv_f32, ctx, qlen, scale, positions):
91 + """Exact absorbed-MLA decode over an explicit set of GLOBAL KV positions.
92 +
93 + A shard sees only ``positions``, but causality still uses the global
94 + position: query token i (global position ctx + i) may attend to any p in
95 + positions with p <= ctx + i. Restricting ``positions`` to a round-robin
96 + residue class is precisely DCP rank r's view.
97 + """
98 + nr, _, heads, _ = q_f32.shape
99 + dev = q_f32.device
100 + o = torch.zeros(nr, qlen, heads, KV_LORA_RANK, dtype=torch.float32, device=dev)
101 + lse = torch.full((nr, qlen, heads), float("-inf"), dtype=torch.float32, device=dev)
102 + for r in range(nr):
103 + kv_sel = kv_f32[r].index_select(0, positions)
104 + for i in range(qlen):
105 + valid = positions <= (ctx + i)
106 + if not bool(valid.any()):
107 + continue
108 + k = kv_sel[valid]
109 + s = (q_f32[r, i] @ k.T) * scale
110 + m = s.max(dim=-1, keepdim=True).values
111 + p = torch.exp(s - m)
112 + denom = p.sum(dim=-1, keepdim=True)
113 + o[r, i] = (p / denom) @ k[:, :KV_LORA_RANK]
114 + lse[r, i] = m.squeeze(-1) + torch.log(denom.squeeze(-1))
115 + return o, lse
116 +
117 +
118 +def _merge_partials(os_, lses):
119 + """Cross-shard softmax merge (natural-log LSE)."""
120 + stack_lse = torch.stack(lses)
121 + m = stack_lse.max(dim=0).values
122 + w = torch.exp(stack_lse - m)
123 + num = (w.unsqueeze(-1) * torch.stack(os_)).sum(dim=0)
124 + den = w.sum(dim=0).unsqueeze(-1)
125 + return num / den, m + torch.log(w.sum(dim=0))
126 +
127 +
128 +# --------------------------------------------------------------------------
129 +# single-process rank simulation
130 +# --------------------------------------------------------------------------
131 +
132 +
133 +class _FakeGroup:
134 + def __init__(self, world_size: int, rank: int):
135 + self.world_size = world_size
136 + self.rank_in_group = rank
137 + self.rank = rank
138 + self.local_rank = rank
139 +
140 + def all_gather(self, tensor, dim=0):
141 + raise RuntimeError(
142 + "collective called in a single-process test -- the path under test "
143 + "is supposed to be collective-free; this indicates a real change"
144 + )
145 +
146 +
147 +def _install_fake_groups(dcp_size: int, cp_rank: int, tp_size: int = 1) -> None:
148 + """Point vLLM's parallel_state at rank-simulating stubs (no NCCL)."""
149 + import vllm.distributed.parallel_state as ps
150 + import vllm.model_executor.layers.attention.mla_attention as mla_attention
151 + import vllm.v1.attention.backends.mla.rocm_aiter_mla as backend
152 +
153 + dcp = _FakeGroup(dcp_size, cp_rank)
154 + tp = _FakeGroup(tp_size, cp_rank)
155 + ps.get_dcp_group = lambda: dcp
156 + ps.get_tp_group = lambda: tp
157 + # Both of these imported the symbol directly at module import time, so
158 + # patching parallel_state alone leaves them resolving the real group (and
159 + # the builder then reports dcp_world_size 1, silently skipping the route).
160 + backend.get_dcp_group = lambda: dcp
161 + mla_attention.get_dcp_group = lambda: dcp
162 +
163 +
164 +@pytest.fixture(autouse=True)
165 +def _restore_parallel_groups():
166 + """Undo _install_fake_groups after every test.
167 +
168 + It patches module globals in three modules, so without this the fake DCP
169 + group leaks into the rest of the pytest session and silently breaks
170 + unrelated tests that expect a real (or absent) group.
171 + """
172 + import vllm.distributed.parallel_state as ps
173 + import vllm.model_executor.layers.attention.mla_attention as mla_attention
174 + import vllm.v1.attention.backends.mla.rocm_aiter_mla as backend
175 +
176 + saved = [
177 + (mod, name, getattr(mod, name))
178 + for mod, name in (
179 + (ps, "get_dcp_group"),
180 + (ps, "get_tp_group"),
181 + (backend, "get_dcp_group"),
182 + (mla_attention, "get_dcp_group"),
183 + )
184 + ]
185 + yield
186 + for mod, name, value in saved:
187 + setattr(mod, name, value)
188 +
189 +
190 +class _PrefillBackendStub:
191 + """The base builder does ``attention_layer.prefill_backend.clone()``.
192 +
193 + prefill_backend is consulted only from the chunked-context prefill path,
194 + which a decode-only test never enters.
195 + """
196 +
197 + def clone(self):
198 + return self
199 +
200 +
201 +def _make_ctx_layer(vllm_config):
202 + """Stand-in for the registered attention layer in static_forward_context.
203 +
204 + ``MLACommonMetadataBuilder.__init__`` reads the MLA dims, ``dcp_manager``
205 + and ``prefill_backend`` off this object. ``dcp_manager`` is the Direct-DCP
206 + symmetric-memory workspace, used ONLY by the chunked-prefill KV all-gather;
207 + it is constructed without running __init__ so no peer pointers are exchanged
208 + and no collective is issued. That is what keeps this test single-process.
209 + """
210 + from vllm.v1.attention.ops.dcp import MLADCPManager
211 +
212 + mgr = object.__new__(MLADCPManager)
213 + mgr.init_kv_gather = lambda *a, **k: None
214 +
215 + return SimpleNamespace(
216 + non_causal_multi_token_decode=False,
217 + q_lora_rank=None,
218 + kv_lora_rank=KV_LORA_RANK,
219 + qk_nope_head_dim=QK_NOPE_HEAD_DIM,
220 + qk_rope_head_dim=QK_ROPE_HEAD_DIM,
221 + v_head_dim=QK_NOPE_HEAD_DIM,
222 + dcp_manager=mgr,
223 + prefill_backend=_PrefillBackendStub(),
224 + )
225 +
226 +
227 +def _ensure_workspace_manager(device: torch.device) -> None:
228 + """Init the global workspace manager, which GPUModelRunner does in production.
229 +
230 + ``AiterMLAMetadataBuilder.__init__`` reserves the fp8 prefill
231 + persistent-scheduling scratch through ``current_workspace_manager()``, so
232 + merely constructing a builder outside a serve trips its assert. Idempotent:
233 + one builder is constructed per DCP shard. Deliberately does NOT lock the
234 + workspace -- production locks only after warmup, and this keeps allocating
235 + across shards.
236 + """
237 + from vllm.v1.worker.workspace import (
238 + init_workspace_manager,
239 + is_workspace_manager_initialized,
240 + )
241 +
242 + if not is_workspace_manager_initialized():
243 + init_workspace_manager(device)
244 +
245 +
246 +class _Layer:
247 + def __init__(self, device):
248 + self._q_scale = torch.tensor(1.0, device=device)
249 + self._k_scale = torch.tensor(1.0, device=device)
250 + self._v_scale = torch.tensor(1.0, device=device)
251 +
252 +
253 +def _build_impl(num_heads: int, dcp: int, cp_rank: int):
254 + """Real AiterMLAImpl.forward_mqa bound to a minimally-populated instance.
255 +
256 + Constructing the full impl would drag in kv_b_proj and the rest of the
257 + weight-dependent MLA setup -- i.e. model loading, which this test exists to
258 + avoid. forward_mqa's fold path reads only the attributes set below.
259 + """
260 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAImpl
261 +
262 + impl = object.__new__(AiterMLAImpl)
263 + impl.kv_lora_rank = KV_LORA_RANK
264 + impl.qk_rope_head_dim = QK_ROPE_HEAD_DIM
265 + impl.num_heads = num_heads
266 + impl.dcp_world_size = dcp
267 + impl.dcp_rank = cp_rank
268 + impl.pcp_world_size = 1
269 + impl.kv_cache_dtype = "fp8"
270 + impl.scale = SCALE
271 + # `_decode_num_heads` is a read-only property deriving num_heads *
272 + # dcp_world_size, both set above; assert rather than assign.
273 + assert impl._decode_num_heads == num_heads * dcp
274 + return impl
275 +
276 +
277 +def _make_vllm_config(heads_per_rank: int):
278 + """A DCP8 MLA config with no weights and no distributed init.
279 +
280 + DeepSeek-R1's config supplies the MLA dims (kv_lora_rank 512, rope 64,
281 + nope 128) that the aiter asm kernel is built for; only the head count is
282 + overridden, to place the GATHERED count (heads_per_rank * DCP) on the
283 + kernel variant under test. ``get_num_attention_heads`` reads
284 + model_arch_config, a snapshot taken at config-creation time, so patching
285 + hf_config alone would be inert.
286 + """
287 + from tests.v1.attention.utils import create_vllm_config
288 +
289 + vllm_config = create_vllm_config(
290 + model_name=MODEL,
291 + tensor_parallel_size=1,
292 + max_model_len=2048,
293 + max_num_seqs=NUM_REQS * 2,
294 + max_num_batched_tokens=4096,
295 + )
296 + vllm_config.model_config.model_arch_config.total_num_attention_heads = (
297 + heads_per_rank
298 + )
299 + # Without a speculative config the batch reorder threshold stays at 1 and a
300 + # qlen-QLEN row is classified as a PREFILL, so the decode path under test is
301 + # never reached. _init_reorder_batch_threshold reads exactly these two
302 + # fields, so a stand-in is enough and avoids pulling in a draft model.
303 + vllm_config.speculative_config = SimpleNamespace(
304 + num_speculative_tokens=QLEN - 1, parallel_drafting=False
305 + )
306 + vllm_config.cache_config.cache_dtype = "fp8"
307 + vllm_config.parallel_config.decode_context_parallel_size = DCP
308 + vllm_config.parallel_config.cp_kv_cache_interleave_size = 1
309 + return vllm_config
310 +
311 +
312 +def _run_shard(
313 + vllm_config,
314 + spec,
315 + layer_name,
316 + device,
317 + cp_rank,
318 + kern_rank,
319 + q_flat,
320 + kv_flat,
321 + nr,
322 + qlen,
323 + ctx,
324 + dcp,
325 + pos,
326 +):
327 + """Build real metadata for one DCP rank and run the real asm kernel.
328 +
329 + ``kern_rank`` is separate from ``cp_rank`` so the positive control can hand
330 + rank r's pages to the kernel while telling it that it is rank r+1.
331 + """
332 + from vllm.v1.attention.backend import CommonAttentionMetadata
333 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder
334 +
335 + L = ctx + qlen
336 + local_len = int(pos.numel())
337 +
338 + # One GLOBAL kv cache: page (req*L + p) holds request req position p. Shards
339 + # select pages by residue class, so every shard reads the same underlying
340 + # bytes and no re-quantisation noise is introduced between arms.
341 + block_table = (
342 + torch.arange(nr, device=device).view(nr, 1) * L + pos.view(1, -1)
343 + ).to(torch.int32)
344 +
345 + # seq_lens is the GLOBAL context; dcp_local_seq_lens is what this rank
346 + # actually holds. Passing the global length as the local one walks the
347 + # kernel off the end of the page table.
348 + seq_lens_cpu = torch.full((nr,), L, dtype=torch.int32)
349 + local_cpu = torch.full((nr,), local_len, dtype=torch.int32)
350 +
351 + common = CommonAttentionMetadata(
352 + query_start_loc=torch.arange(
353 + 0, nr * qlen + 1, qlen, dtype=torch.int32, device=device
354 + ),
355 + query_start_loc_cpu=torch.arange(0, nr * qlen + 1, qlen, dtype=torch.int32),
356 + seq_lens=seq_lens_cpu.to(device),
357 + num_reqs=nr,
358 + num_actual_tokens=nr * qlen,
359 + max_query_len=qlen,
360 + max_seq_len=L,
361 + block_table_tensor=block_table,
362 + slot_mapping=torch.arange(nr * qlen, dtype=torch.int64, device=device),
363 + causal=True,
364 + dcp_local_seq_lens=local_cpu.to(device),
365 + dcp_local_seq_lens_cpu_upper_bound=local_cpu,
366 + )
367 +
368 + _install_fake_groups(dcp, kern_rank)
369 + builder = AiterMLAMetadataBuilder(spec, [layer_name], vllm_config, device)
370 + md = builder.build(0, common)
371 +
372 + impl = _build_impl(builder.num_heads, dcp, kern_rank)
373 + o, lse = impl.forward_mqa(q_flat, kv_flat, md, _Layer(device))
374 + torch.accelerator.synchronize()
375 + return md, o, lse, local_len
376 +
377 +
378 +def _numerics(heads_per_rank: int, sabotage: bool = False):
379 + """Run every DCP shard and return (per-shard max rel err, merged rel err)."""
380 + import aiter
381 +
382 + from vllm.config import set_current_vllm_config
383 +
384 + device = torch.device("cuda", 0)
385 + torch.accelerator.set_device_index(device.index)
386 + torch.manual_seed(0)
387 + _ensure_workspace_manager(device)
388 +
389 + _install_fake_groups(DCP, 0)
390 + vllm_config = _make_vllm_config(heads_per_rank)
391 +
392 + with set_current_vllm_config(vllm_config):
393 + from vllm.v1.kv_cache_interface import MLAAttentionSpec
394 +
395 + layer_name = "model.layers.0.self_attn.attn"
396 + vllm_config.compilation_config.static_forward_context[layer_name] = (
397 + _make_ctx_layer(vllm_config)
398 + )
399 + # block_size 1: one page per token, so page ids are positions and the
400 + # residue-class page table is expressible directly.
401 + spec = MLAAttentionSpec(
402 + block_size=1,
403 + num_kv_heads=1,
404 + head_size=HEAD_SIZE,
405 + dtype=vllm_config.model_config.dtype,
406 + non_causal_multi_token_decode=False,
407 + )
408 +
409 + nr, ctx, qlen = NUM_REQS, CTX, QLEN
410 + L = ctx + qlen
411 + gathered_heads = heads_per_rank * DCP
412 + kv_dtype = aiter.dtypes.fp8
413 +
414 + kv_flat = torch.randn(
415 + nr * L, 1, HEAD_SIZE, dtype=torch.float32, device=device
416 + ).to(kv_dtype)
417 + q_flat = (
418 + torch.randn(
419 + nr * qlen,
420 + gathered_heads,
421 + HEAD_SIZE,
422 + dtype=torch.float32,
423 + device=device,
424 + )
425 + .mul_(0.5)
426 + .to(kv_dtype)
427 + )
428 + # Dequantise for the reference so the comparison measures the kernel,
429 + # not the fp8 cast that both sides share.
430 + kv_f32 = kv_flat.float().view(nr, L, HEAD_SIZE)
431 + q_f32 = q_flat.float().view(nr, qlen, gathered_heads, HEAD_SIZE)
432 +
433 + ref_o, _ = _torch_reference(
434 + q_f32, kv_f32, ctx, qlen, SCALE, torch.arange(L, device=device)
435 + )
436 +
437 + shard_o, shard_lse = [], []
438 + max_shard_rel = 0.0
439 + for r in range(DCP):
440 + pos = torch.arange(r, L, DCP, device=device) # residue class r
441 + kern_rank = (r + 1) % DCP if sabotage else r
442 + md, o, lse, _ = _run_shard(
443 + vllm_config,
444 + spec,
445 + layer_name,
446 + device,
447 + r,
448 + kern_rank,
449 + q_flat,
450 + kv_flat,
451 + nr,
452 + qlen,
453 + ctx,
454 + DCP,
455 + pos,
456 + )
457 + # The route actually under test must have been taken; a silent
458 + # fall-out to Triton/segmented would otherwise pass quietly.
459 + assert md.decode.dcp_route is rocm_aiter_mla._DCPDecodeRoute.CPRR
460 + assert md.decode.asm_decode_num_heads, "cprr asm route was not taken"
461 +
462 + o = o.float().view(nr, qlen, gathered_heads, KV_LORA_RANK)
463 + shard_lse.append(lse.float().view(nr, qlen, gathered_heads))
464 + shard_o.append(o)
465 +
466 + po, _ = _torch_reference(q_f32, kv_f32, ctx, qlen, SCALE, pos)
467 + rel = (o - po).abs().max().item() / max(po.abs().max().item(), 1e-6)
468 + max_shard_rel = max(max_shard_rel, rel)
469 +
470 + merged, _ = _merge_partials(shard_o, shard_lse)
471 + merge_rel = (merged - ref_o).abs().max().item() / max(
472 + ref_o.abs().max().item(), 1e-6
473 + )
474 + return max_shard_rel, merge_rel
475 +
476 +
477 +# --------------------------------------------------------------------------
478 +# tests
479 +# --------------------------------------------------------------------------
480 +
481 +
482 +@requires_cprr
483 +@pytest.mark.parametrize(
484 + "heads_per_rank",
485 + [
486 + 16, # 128 gathered heads -- a NATIVE cprr kernel, no padding
487 + 12, # 96 gathered heads -- K3's DSpark target, must pad up to 128
488 + ],
489 + ids=["native-128", "padded-96-to-128"],
490 +)
491 +@torch.inference_mode()
492 +def test_cprr_shards_match_reference_and_merge_to_full_context(heads_per_rank):
493 + """Every DCP shard must be right, and their LSE merge must reproduce
494 + full-context attention.
495 +
496 + A shard that ignored global positions (applying a local causal tail
497 + instead) still produces plausible-looking output; it only diverges from the
498 + reference on the tokens whose causal window is truncated. Checking the merge
499 + alone would hide that, so both are asserted.
500 + """
501 + shard_rel, merge_rel = _numerics(heads_per_rank)
502 + assert shard_rel < SHARD_RTOL, f"per-shard rel err {shard_rel:.3e}"
503 + assert merge_rel < MERGE_RTOL, f"merged rel err {merge_rel:.3e}"
504 +
505 +
506 +@requires_cprr
507 +@torch.inference_mode()
508 +def test_wrong_cp_rank_breaks_the_result():
509 + """Positive control: without it, a passing run above proves nothing.
510 +
511 + Feeding shard r's pages to the kernel while telling it that it is rank r+1
512 + shifts every reconstructed global position by one, so both the causal mask
513 + and the round-robin reconstruction are wrong. If the error does NOT blow up,
514 + the kernel is not using cp_rank and the test above has no power.
515 + """
516 + shard_rel, merge_rel = _numerics(16, sabotage=True)
517 + assert shard_rel > SHARD_RTOL or merge_rel > MERGE_RTOL, (
518 + f"sabotaged run still passed (shard {shard_rel:.3e}, merge "
519 + f"{merge_rel:.3e}) -- cp_rank is not reaching the kernel"
520 + )
521 +
522 +
523 +@requires_cprr
524 +def test_cprr_kernel_is_selected_for_the_k3_shape():
525 + """The K3 DSpark target shape must actually reach the cprr route.
526 +
527 + Guards against a silent fall-out to segmented/Triton, which is what the
528 + stock configuration does and what this route exists to avoid.
529 + """
530 + from vllm.v1.attention.backends.mla import rocm_aiter_mla as m
531 +
532 + assert (
533 + m._asm_dcp_verify_configured(
534 + dcp_world_size=DCP, cp_interleave=1, multi_token_decode=True
535 + )
536 + is True
537 + )
538 + assert QLEN >= m._MIN_CPRR_QLEN
539 + assert m._select_dcp_decode_route(True, True, QLEN, True) is m._DCPDecodeRoute.CPRR
540 + assert m._asm_dcp_verify_heads(96) in m._NATIVE_CPRR_HEADS
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+242 -7
@@ -87,21 +87,31 @@ def _builder(
87 kv_cache_dtype: str = "auto",
88 dcp_world_size: int = 1,
89 dcp_rank: int = 0,
90 + asm_dcp_verify: bool = False,
91 + decode_causal: bool = True,
92 + supports_segmented_dcp_verify: bool | None = None,
93 ):
94 + decode_num_heads = num_heads * dcp_world_size
95 + asm_heads = (
96 + rocm_aiter_mla._asm_dcp_verify_heads(decode_num_heads) if asm_dcp_verify else 0
97 + )
98 + if supports_segmented_dcp_verify is None:
99 + supports_segmented_dcp_verify = rocm_aiter_mla._segmented_dcp_verify_supported(
100 + dcp_world_size, 1
101 + )
102 stub = SimpleNamespace(
103 device=torch.device("cpu"),
104 num_heads=num_heads,
105 # DCP gathers every rank's head shard before decode, and the routing
106 # predicates read the gathered count.
96 - _decode_num_heads=num_heads * dcp_world_size,
107 + _decode_num_heads=decode_num_heads,
108 dcp_world_size=dcp_world_size,
109 dcp_rank=dcp_rank,
110 cp_kv_cache_interleave_size=1,
111 # Mirrors the production constructor: configuration decides whether the
101 - # segmented route is available, query length decides per batch.
102 - _supports_segmented_dcp_verify=rocm_aiter_mla._segmented_dcp_verify_supported(
103 - dcp_world_size, 1
104 - ),
112 + # segmented route is available, query length decides per batch. CPRR
113 + # preferred still keeps this True so qlen < _MIN_CPRR_QLEN can fall back.
114 + _supports_segmented_dcp_verify=supports_segmented_dcp_verify,
115 # Derived once in the real constructor, so derive it once here too.
116 _segmented_page_size=rocm_aiter_mla._segmented_mla_page_size(kernel_block_size),
117 _dcp_verify_buffers=None,
@@ -122,7 +132,11 @@ def _builder(
132 _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len),
133 _use_persistent_metadata=False,
134 kernel_block_size=kernel_block_size,
125 - _num_attention_heads=AiterMLAHelper.get_actual_mla_num_heads(num_heads),
135 + _num_attention_heads=(
136 + asm_heads
137 + if asm_dcp_verify
138 + else AiterMLAHelper.get_actual_mla_num_heads(num_heads)
139 + ),
140 _mla_work_meta_data=torch.empty(1, dtype=torch.int32),
141 _mla_work_info_set=torch.empty(1, dtype=torch.int32),
142 _mla_work_indptr=torch.empty(1, dtype=torch.int32),
@@ -132,7 +146,19 @@ def _builder(
146 _mla_q_dtype=torch.bfloat16,
147 _mla_kv_dtype=torch.bfloat16,
148 decode_attn_out_dtype=torch.bfloat16,
135 - _decode_causal=True,
149 + _decode_causal=decode_causal,
150 + # Default off so existing cases keep exercising segmented/persistent.
151 + # Pass asm_dcp_verify=True for the CPRR route and its qlen-floor fallback.
152 + _asm_dcp_verify=asm_dcp_verify,
153 + _asm_dcp_verify_heads=asm_heads,
154 + _g_kv_indptr_buf=(
155 + torch.zeros(max_decode_rows + 1, dtype=torch.int32)
156 + if asm_dcp_verify and dcp_world_size > 1
157 + else None
158 + ),
159 + # Only read when the asm route is on; the real builder derives it from
160 + # the device CU count.
161 + _mla_max_split_per_batch=256,
162 )
163 # Bound to the stub rather than faked: the verify flatten's per-row view is
164 # part of what _build_decode is being tested for.
@@ -287,6 +313,209 @@ def test_dcp_fp8_verify_build_uses_segmented(monkeypatch):
313 assert metadata.dcp_verify is not None
314
315
316 +def _cprr_decode_batch(builder, *, qlen: int, num_reqs: int = 1):
317 + seq_lens = torch.full((num_reqs,), 16, dtype=torch.int32)
318 + query_start_loc = torch.arange(0, (num_reqs + 1) * qlen, qlen, dtype=torch.int32)
319 + return AiterMLAMetadataBuilder._build_decode(
320 + builder,
321 + block_table_tensor=torch.zeros(num_reqs, 8, dtype=torch.int32),
322 + seq_lens_device=seq_lens,
323 + max_seq_len=16,
324 + query_start_loc_cpu=query_start_loc,
325 + query_start_loc_device=query_start_loc,
326 + num_decode_tokens=num_reqs * qlen,
327 + max_query_len=qlen,
328 + dcp_tot_seq_lens_device=seq_lens * 8,
329 + )
330 +
331 +
332 +def test_asm_preferred_qlen2_falls_back_to_segmented(monkeypatch):
333 + """K=4 can clamp to qlen 2; CPRR has no kernel, so use segmented MLA."""
334 + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
335 + get_mla_metadata_v1 = mock.MagicMock()
336 + monkeypatch.setitem(
337 + sys.modules,
338 + "aiter",
339 + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
340 + )
341 +
342 + metadata = _cprr_decode_batch(
343 + _builder(
344 + mtp_decode_qlen=5,
345 + dcp_world_size=8,
346 + num_heads=12,
347 + asm_dcp_verify=True,
348 + kernel_block_size=2,
349 + ),
350 + qlen=2,
351 + )
352 +
353 + assert metadata.dcp_verify is not None
354 + assert metadata.g_kv_indptr is None
355 + assert not metadata.has_persistent_metadata
356 + get_mla_metadata_v1.assert_not_called()
357 +
358 +
359 +def test_asm_preferred_qlen_at_cprr_floor_stays_on_cprr(monkeypatch):
360 + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
361 + get_mla_metadata_v1 = mock.MagicMock()
362 + monkeypatch.setitem(
363 + sys.modules,
364 + "aiter",
365 + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
366 + )
367 + monkeypatch.setattr(
368 + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
369 + )
370 +
371 + metadata = _cprr_decode_batch(
372 + _builder(
373 + mtp_decode_qlen=5,
374 + dcp_world_size=8,
375 + num_heads=12,
376 + asm_dcp_verify=True,
377 + ),
378 + qlen=rocm_aiter_mla._MIN_CPRR_QLEN,
379 + )
380 +
381 + assert metadata.dcp_verify is None
382 + assert metadata.g_kv_indptr is not None
383 + assert metadata.has_persistent_metadata
384 + assert get_mla_metadata_v1.call_args.kwargs.get("is_cp_round_robin") is True
385 +
386 +
387 +def test_cprr_capable_single_token_decode_keeps_persistent_metadata(monkeypatch):
388 + """The CPRR qlen floor must not suppress the ordinary qlen-1 schedule."""
389 + get_mla_metadata_v1 = mock.MagicMock()
390 + monkeypatch.setitem(
391 + sys.modules,
392 + "aiter",
393 + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
394 + )
395 + monkeypatch.setattr(
396 + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
397 + )
398 +
399 + metadata = _cprr_decode_batch(
400 + _builder(
401 + mtp_decode_qlen=5,
402 + dcp_world_size=8,
403 + num_heads=12,
404 + asm_dcp_verify=True,
405 + ),
406 + qlen=1,
407 + )
408 +
409 + assert metadata.g_kv_indptr is None
410 + assert metadata.has_persistent_metadata
411 + assert get_mla_metadata_v1.call_args.kwargs.get("is_cp_round_robin") is None
412 +
413 +
414 +def test_non_causal_dcp_block_uses_plain_mask0_metadata(monkeypatch):
415 + """A non-causal block needs no CPRR global-position window or head pad."""
416 + get_mla_metadata_v1 = mock.MagicMock()
417 + monkeypatch.setitem(
418 + sys.modules,
419 + "aiter",
420 + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
421 + )
422 + monkeypatch.setattr(
423 + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
424 + )
425 + builder = _builder(
426 + mtp_decode_qlen=5,
427 + dcp_world_size=8,
428 + num_heads=12,
429 + asm_dcp_verify=True,
430 + )
431 + builder._decode_causal = False
432 +
433 + metadata = _cprr_decode_batch(builder, qlen=4)
434 +
435 + assert metadata.dcp_verify is None
436 + assert metadata.g_kv_indptr is None
437 + assert metadata.has_persistent_metadata
438 + assert get_mla_metadata_v1.call_args.args[3] == 96
439 + assert get_mla_metadata_v1.call_args.args[5] is False
440 + assert get_mla_metadata_v1.call_args.kwargs.get("is_cp_round_robin") is None
441 +
442 +
443 +def test_asm_qlen2_without_segmented_fails_during_build(monkeypatch):
444 + monkeypatch.setitem(
445 + sys.modules,
446 + "aiter",
447 + SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()),
448 + )
449 + monkeypatch.setattr(
450 + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
451 + )
452 +
453 + with pytest.raises(RuntimeError, match="requires either segmented MLA"):
454 + _cprr_decode_batch(
455 + _builder(
456 + mtp_decode_qlen=5,
457 + dcp_world_size=8,
458 + num_heads=12,
459 + asm_dcp_verify=True,
460 + supports_segmented_dcp_verify=False,
461 + ),
462 + qlen=2,
463 + )
464 +
465 +
466 +def test_non_causal_dcp_block_bypasses_cprr_in_forward(monkeypatch):
467 + """Process-level CPRR capability must not pad or route a non-causal batch."""
468 + captured = {}
469 +
470 + def fake_aiter_decode(q, kv_buffer, out, *args, **kwargs):
471 + captured["q_heads"] = q.shape[1]
472 + captured["kwargs"] = kwargs
473 + return None, torch.zeros(q.shape[0], q.shape[1])
474 +
475 + monkeypatch.setattr(
476 + rocm_aiter_mla, "_get_aiter_mla_decode", lambda: fake_aiter_decode
477 + )
478 +
479 + impl = object.__new__(AiterMLAImpl)
480 + impl.num_heads = 12
481 + impl.dcp_world_size = 8
482 + impl.kv_cache_dtype = "auto"
483 + impl.kv_lora_rank = 512
484 + impl.qk_rope_head_dim = 64
485 + impl.scale = 576**-0.5
486 + decode = SimpleNamespace(
487 + max_qo_len=4,
488 + qo_indptr=torch.tensor([0, 4], dtype=torch.int32),
489 + paged_kv_indptr=torch.tensor([0, 1], dtype=torch.int32),
490 + paged_kv_indices=torch.zeros(1, dtype=torch.int32),
491 + paged_kv_last_page_len=torch.ones(1, dtype=torch.int32),
492 + use_gluon_decode=False,
493 + use_gluon_verify=False,
494 + dcp_route=rocm_aiter_mla._DCPDecodeRoute.PLAIN,
495 + dcp_verify=None,
496 + g_kv_indptr=None,
497 + has_persistent_metadata=False,
498 + attn_out_dtype=torch.bfloat16,
499 + asm_decode_num_heads=0,
500 + mla_num_kv_splits=0,
501 + cp_world_size=8,
502 + cp_rank=0,
503 + min_kv_seq_len=1,
504 + )
505 + attn_metadata = SimpleNamespace(decode=decode, causal=False, work_meta_data=None)
506 + layer = SimpleNamespace(_q_scale=torch.tensor(1.0), _k_scale=torch.tensor(1.0))
507 + q = torch.zeros(4, 96, 576, dtype=torch.bfloat16)
508 +
509 + output, lse = impl.forward_mqa(q, torch.zeros(1, 1, 576), attn_metadata, layer)
510 +
511 + assert captured["q_heads"] == 96
512 + assert captured["kwargs"]["causal"] is False
513 + assert "g_kv_indptr" not in captured["kwargs"]
514 + assert "num_kv_splits" not in captured["kwargs"]
515 + assert output.shape[1] == 96
516 + assert lse is not None and lse.shape == (4, 96)
517 +
518 +
519 def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch):
520 """Single-token DCP decode must come back with an LSE the merge can use.
521
@@ -325,9 +554,14 @@ def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch):
554 paged_kv_last_page_len=torch.ones(num_tokens, dtype=torch.int32),
555 use_gluon_decode=False,
556 use_gluon_verify=False,
557 + dcp_route=rocm_aiter_mla._DCPDecodeRoute.PLAIN,
558 dcp_verify=None,
559 has_persistent_metadata=False,
560 attn_out_dtype=torch.bfloat16,
561 + # 0 == the round-robin (cprr) asm route is off, this rank's decode takes
562 + # the existing path. Mirrors the AiterMLADecodeMetadata default.
563 + asm_decode_num_heads=0,
564 + g_kv_indptr=None,
565 )
566 attn_metadata = SimpleNamespace(decode=decode, causal=True, work_meta_data=None)
567 layer = SimpleNamespace(_q_scale=torch.tensor(1.0), _k_scale=torch.tensor(1.0))
@@ -453,6 +687,7 @@ def test_segmented_dcp_verify_matches_causal_attention(monkeypatch):
687 ),
688 use_gluon_decode=False,
689 use_gluon_verify=False,
690 + dcp_route=rocm_aiter_mla._DCPDecodeRoute.SEGMENTED,
691 dcp_verify=view,
692 attn_out_dtype=torch.bfloat16,
693 )
vllm/envs.py
+10
@@ -145,6 +145,7 @@ if TYPE_CHECKING:
145 VLLM_ROCM_USE_AITER_RMSNORM: bool = True
146 VLLM_ROCM_USE_AITER_MLA: bool = True
147 VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"
148 + VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["asm", "segmented"] = "segmented"
149 VLLM_ROCM_USE_AITER_MHA: bool = True
150 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
151 VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
@@ -1309,6 +1310,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
1310 "VLLM_ROCM_USE_AITER_MLA": lambda: (
1311 os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1")
1312 ),
1313 + # Kernel for causal multi-token (spec-decode) verify steps under decode
1314 + # context parallelism on gfx950: "asm" uses AITER's round-robin ASM
1315 + # decode, "segmented" the Triton segmented MLA path.
1316 + "VLLM_ROCM_AITER_MLA_DCP_VERIFY": env_with_choices(
1317 + "VLLM_ROCM_AITER_MLA_DCP_VERIFY",
1318 + "segmented",
1319 + ["asm", "segmented"],
1320 + case_sensitive=False,
1321 + ),
1322 # Small-head (<16) AITER MLA decode kernel selection. Small head counts
1323 # (e.g. Kimi-K3: 12 heads/rank at TP8, 6 at TP16) can decode either through
1324 # the Gluon small-head kernel or through the padded persistent-scheduling
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+258 -23
@@ -3,11 +3,13 @@
3
4 import functools
5 from dataclasses import dataclass
6 +from enum import Enum
7 from pathlib import Path
8 from typing import TYPE_CHECKING, ClassVar, Final
9
10 import torch
11
12 +from vllm import envs
13 from vllm._aiter_ops import rocm_aiter_ops
14 from vllm.config import VllmConfig
15 from vllm.config.cache import CacheDType
@@ -239,6 +241,78 @@ def _segmented_dcp_verify_supported(dcp_world_size: int, cp_interleave: int) ->
241 )
242
243
244 +# DCP decode routing
245 +#
246 +# With decode context parallelism each rank holds a round-robin shard of the KV
247 +# cache (token t on rank t % dcp_world_size; requires interleave size 1) and
248 +# attends over it with all DCP-gathered query heads. Partial outputs are merged
249 +# across ranks by LSE. Each decode batch takes one route:
250 +#
251 +# * plain: qlen == 1 or a non-causal block. Every query sees every local token.
252 +# * cprr: causal qlen >= _MIN_CPRR_QLEN, gfx950, speculative decoding. AITER's
253 +# round-robin ASM decode rebuilds global positions from g_kv_indptr, cp_rank
254 +# and cp_world_size and applies the causal window in-kernel. Kernels exist only
255 +# at _NATIVE_CPRR_HEADS; other gathered head counts are padded up.
256 +# * segmented: causal qlen > 1 when cprr is not selected or qlen is below
257 +# _MIN_CPRR_QLEN. Each verify token becomes a single-query row whose KV length
258 +# encodes its causal window.
259 +#
260 +# The plain kernel must never serve a causal qlen > 1 block: it applies
261 +# causality on local indices, which is wrong over a round-robin shard.
262 +_NATIVE_CPRR_HEADS: Final = (16, 32, 64, 128)
263 +_MIN_CPRR_QLEN: Final = 3
264 +
265 +
266 +class _DCPDecodeRoute(Enum):
267 + PLAIN = "plain"
268 + SEGMENTED = "segmented"
269 + CPRR = "cprr"
270 +
271 +
272 +def _asm_dcp_verify_heads(decode_num_heads: int) -> int:
273 + """Head count the cprr asm decode runs at, or 0 if it cannot serve this shape."""
274 + if decode_num_heads in _NATIVE_CPRR_HEADS:
275 + return decode_num_heads
276 + for nat in _NATIVE_CPRR_HEADS:
277 + if nat >= decode_num_heads:
278 + return nat
279 + return 0
280 +
281 +
282 +def _select_dcp_decode_route(
283 + supports_segmented: bool,
284 + causal: bool,
285 + max_qo_len: int,
286 + asm_selected: bool,
287 +) -> _DCPDecodeRoute:
288 + """Select one DCP decode route for this batch; see DCP decode routing."""
289 + if not causal or max_qo_len <= 1:
290 + return _DCPDecodeRoute.PLAIN
291 + if asm_selected and max_qo_len >= _MIN_CPRR_QLEN:
292 + return _DCPDecodeRoute.CPRR
293 + if supports_segmented:
294 + return _DCPDecodeRoute.SEGMENTED
295 + raise RuntimeError(
296 + "ROCM_AITER_MLA DCP multi-token verify requires either segmented "
297 + "MLA or the round-robin asm decode."
298 + )
299 +
300 +
301 +def _asm_dcp_verify_configured(
302 + dcp_world_size: int, cp_interleave: int, multi_token_decode: bool
303 +) -> bool:
304 + """Whether asm cprr DCP verify is reachable in this process.
305 +
306 + Requires DCP with interleave size 1, speculative decoding, and gfx950 (the
307 + only arch AITER builds cprr kernels for).
308 + """
309 + if not (dcp_world_size > 1 and cp_interleave == 1 and multi_token_decode):
310 + return False
311 + from vllm.platforms.rocm import on_gfx950
312 +
313 + return bool(on_gfx950())
314 +
315 +
316 def _aiter_mla_small_head_mode() -> str:
317 """Small-head (<16) MLA decode kernel selection.
318
@@ -410,8 +484,19 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
484 max_qo_len: int | None = None
485 # Minimum KV length used by Gluon to choose a safe split count.
486 min_kv_seq_len: int = 1
487 + # Exactly one DCP decode route is selected by the metadata builder.
488 + dcp_route: _DCPDecodeRoute = _DCPDecodeRoute.PLAIN
489 # Set exactly when this batch routes to segmented DCP verification.
490 dcp_verify: AiterMLADCPVerifyMetadata | None = None
491 + # cprr only: global per-request page indptr and this rank's position.
492 + g_kv_indptr: torch.Tensor | None = None
493 + cp_world_size: int = 1
494 + cp_rank: int = 0
495 + # Padded head count cprr runs at; 0 when this batch is not on cprr.
496 + asm_decode_num_heads: int = 0
497 + # Must equal the split cap the reduce scratch was sized with, or the kernel
498 + # writes past it and faults.
499 + mla_num_kv_splits: int = 0
500 # Small-head decode uses Gluon (avoids padding to 16).
501 use_gluon_decode: bool = False
502 # Small-head non-DCP multi-token verify uses Gluon's native MTP entry.
@@ -460,6 +545,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
545 # Set from the common metadata every build; a batch is causal unless the
546 # drafter says otherwise.
547 _decode_causal: bool = True
548 + # A non-causal draft needs no cross-shard causal window: the ordinary
549 + # mask0 decode returns per-row LSE for the DCP merge.
550 + supports_non_causal_multi_token_dcp: ClassVar[bool] = True
551
552 @staticmethod
553 def _uniform_padded_mtp_qo_len(
@@ -513,15 +601,48 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
601 parallel_config.decode_context_parallel_size,
602 parallel_config.cp_kv_cache_interleave_size,
603 )
604 + asm_dcp_verify_config = (
605 + _asm_dcp_verify_configured(
606 + parallel_config.decode_context_parallel_size,
607 + parallel_config.cp_kv_cache_interleave_size,
608 + multi_token_decode=vllm_config.speculative_config is not None,
609 + )
610 + and envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "asm"
611 + )
612 super().__init__(
613 kv_cache_spec,
614 layer_names,
615 vllm_config,
616 device,
617 AiterMLAMetadata,
522 - supports_dcp_with_varlen=supports_segmented_dcp_verify,
618 + supports_dcp_with_varlen=(
619 + supports_segmented_dcp_verify or asm_dcp_verify_config
620 + ),
621 )
622 + self._asm_dcp_verify = False
623 + self._asm_dcp_verify_heads = 0
624 + # The head-count half of the route decision needs self.num_heads, so only
625 + # the config half is known before super().__init__().
626 + if asm_dcp_verify_config:
627 + self._asm_dcp_verify_heads = _asm_dcp_verify_heads(
628 + self.num_heads * self.dcp_world_size
629 + )
630 + if not self._asm_dcp_verify_heads:
631 + raise ValueError(
632 + "VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm, but the round-robin asm "
633 + f"decode has no kernel for {self.num_heads * self.dcp_world_size} "
634 + f"DCP-gathered heads (native counts: {_NATIVE_CPRR_HEADS}). "
635 + "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented."
636 + )
637 + self._asm_dcp_verify = True
638 self._supports_segmented_dcp_verify = supports_segmented_dcp_verify
639 + self._mla_max_split_per_batch = 0
640 + if self._asm_dcp_verify:
641 + # Cap on KV splits per batch; at low batch the KV axis is the only
642 + # parallelism. Buffer sizing and every runtime call must see this value.
643 + self._mla_max_split_per_batch = torch.cuda.get_device_properties(
644 + device
645 + ).multi_processor_count
646
647 self.compilation_config = vllm_config.compilation_config
648 self.decode_attn_out_dtype = vllm_config.model_config.dtype
@@ -576,6 +697,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
697 self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
698 self._decode_num_heads
699 )
700 + # the cprr kernel exists only at the native head
701 + # counts, so run (and size the persistent schedule) at the padded count.
702 + # 96 is 16-aligned, so get_actual_mla_num_heads leaves it alone and the
703 + # pad has to be applied here explicitly.
704 + if self._asm_dcp_verify:
705 + self._num_attention_heads = self._asm_dcp_verify_heads
706 kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
707 if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
708 kv_cache_dtype_str = "fp8"
@@ -619,6 +746,13 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
746 kv_dtype,
747 is_sparse=False,
748 fast_mode=True,
749 + # Must match max_split_per_batch in get_mla_metadata_v1; without it
750 + # aiter sizes the reduce scratch from capacity (~9 GiB).
751 + **(
752 + dict(max_split_per_batch=self._mla_max_split_per_batch)
753 + if self._asm_dcp_verify
754 + else {}
755 + ),
756 )
757 self._mla_work_meta_data = torch.empty(
758 work_meta_data_size, dtype=work_meta_data_type, device=device
@@ -664,6 +798,14 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
798 device,
799 )
800
801 + # Full cudagraphs need a stable address for the cprr global page indptr.
802 + # Zero-initialized; element 0 is the indptr base and is never rewritten.
803 + self._g_kv_indptr_buf: torch.Tensor | None = None
804 + if self._asm_dcp_verify:
805 + self._g_kv_indptr_buf = torch.zeros(
806 + max_num_reqs + 1, dtype=torch.int32, device=device
807 + )
808 +
809 # Persistent buffers for segmented DCP verification. Captured graphs
810 # require stable addresses while row lengths vary between replays.
811 self._dcp_verify_buffers: AiterMLADCPVerifyMetadata | None = None
@@ -683,8 +825,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
825 )
826
827 if self._supports_segmented_dcp_verify and self._mtp_decode_qlen > 1:
686 - # A DCP rank's shard of the longest sequence bounds every verify
687 - # row, and full graphs need that bound to be constant.
828 + # Allocate even when CPRR is the preferred route: a later
829 + # replay can fall below _MIN_CPRR_QLEN and needs these
830 + # addresses. A DCP rank's shard of the longest sequence
831 + # bounds every verify row, and full graphs need that bound
832 + # to be constant.
833 num_dcp_partitions = (
834 self.dcp_world_size * self.cp_kv_cache_interleave_size
835 )
@@ -1063,6 +1208,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1208 pad_uniform_mtp = padded_mtp_qo_len > 0
1209
1210 seq_lens_for_kernel = seq_lens_device
1211 + # the global lengths need the same dummy-row
1212 + # treatment as the local ones below, but only for g_kv_indptr -- the
1213 + # metadata field keeps the true values, which is what the segmented
1214 + # route wants (it derives its own row lens and gives a dummy row
1215 + # length 0, correctly).
1216 + g_tot_seq_lens = dcp_tot_seq_lens_device
1217 num_kernel_reqs = num_reqs
1218 if pad_uniform_mtp:
1219 qo_lens_device = (
@@ -1074,6 +1225,19 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1225 seq_lens_for_kernel,
1226 seq_lens_for_kernel.new_full((), max_qo_len),
1227 )
1228 + # A dummy row just got a LOCAL length of max_qo_len, so its GLOBAL
1229 + # length has to say the same thing or the cprr mask is degenerate:
1230 + # global_len 0 puts the row query positions at {-3,-2,-1} while its
1231 + # local shard sits at {rank, W+rank, 2W+rank}, so every entry is
1232 + # masked and the row softmaxes over nothing. max_qo_len * W is
1233 + # exactly the global length whose round-robin shard is max_qo_len
1234 + # on every rank.
1235 + if g_tot_seq_lens is not None:
1236 + g_tot_seq_lens = torch.where(
1237 + qo_lens_device > 0,
1238 + g_tot_seq_lens,
1239 + g_tot_seq_lens.new_full((), max_qo_len * self.dcp_world_size),
1240 + )
1241
1242 if self._graph_seq_lens is not None:
1243 self._graph_seq_lens[:num_kernel_reqs].copy_(
@@ -1106,16 +1270,21 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1270 causal,
1271 self._kv_cache_bytes,
1272 )
1109 - use_segmented_dcp_verify = (
1110 - self._supports_segmented_dcp_verify and max_qo_len > 1
1111 - )
1273 + dcp_route = _DCPDecodeRoute.PLAIN
1274 + if self.dcp_world_size > 1:
1275 + dcp_route = _select_dcp_decode_route(
1276 + self._supports_segmented_dcp_verify,
1277 + causal,
1278 + int(max_qo_len),
1279 + self._asm_dcp_verify,
1280 + )
1281
1282 # Segmented DCP verify carries its own per-row subpage table, so the
1283 # flat per-token view is dead work for it. Leave the buffer alone and
1284 # hand the metadata None, so a future reader cannot pick up whatever
1285 # the previous batch left behind.
1286 paged_kv_indices = None
1118 - if not use_segmented_dcp_verify:
1287 + if dcp_route is not _DCPDecodeRoute.SEGMENTED:
1288 if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
1289 self.paged_kv_indices.fill_(-1)
1290
@@ -1169,6 +1338,19 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1338 else:
1339 qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
1340
1341 + g_kv_indptr = None
1342 + if dcp_route is _DCPDecodeRoute.CPRR and g_tot_seq_lens is not None:
1343 + assert self._g_kv_indptr_buf is not None
1344 + _ngk = g_tot_seq_lens.shape[0]
1345 + g_kv_indptr = self._g_kv_indptr_buf[: _ngk + 1]
1346 + g_kv_indptr[1:].copy_(g_tot_seq_lens.cumsum(dim=0, dtype=torch.int32))
1347 + # A FULL cudagraph replays at its capture-time request count, so a
1348 + # smaller real batch leaves the tail holding the PREVIOUS step's
1349 + # cumsum. Flatten it the way paged_kv_indptr is flattened below, so
1350 + # padding rows read global_len == 0 and agree with their zero-length
1351 + # local shard instead of a stale global length.
1352 + self._g_kv_indptr_buf[_ngk + 1 :].fill_(g_kv_indptr[-1])
1353 +
1354 has_persistent_metadata = False
1355 # Only the asm decode consumes the schedule, so gate on the routing
1356 # rather than on num_heads >= 16, which denies it to a padded rank
@@ -1178,7 +1360,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1360 use_persistent_metadata = (
1361 not use_gluon_decode
1362 and not use_gluon_verify
1181 - and not use_segmented_dcp_verify
1363 + and dcp_route is not _DCPDecodeRoute.SEGMENTED
1364 # A padded rank has no bf16 persistent kernel past qlen 4 where the
1365 # gfx950 fold is absent; the non-persistent entry covers it. fp8
1366 # keeps the schedule -- its fold rejects non-persistent outright.
@@ -1218,13 +1400,28 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1400 uni_qo_len = (
1401 max_qo_len if pad_uniform_mtp or torch.all(qo_len == max_qo_len) else -1
1402 )
1403 + cprr_kwargs: dict = {}
1404 + is_causal = causal
1405 + if self._asm_dcp_verify:
1406 + cprr_kwargs = dict(
1407 + max_split_per_batch=self._mla_max_split_per_batch,
1408 + intra_batch_mode=False,
1409 + )
1410 + if dcp_route is _DCPDecodeRoute.CPRR:
1411 + is_causal = False
1412 + cprr_kwargs["is_cp_round_robin"] = True
1413 + metadata_num_heads = (
1414 + self._asm_dcp_verify_heads
1415 + if dcp_route is _DCPDecodeRoute.CPRR
1416 + else AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads)
1417 + )
1418 get_mla_metadata_v1(
1419 qo_indptr,
1420 paged_kv_indptr,
1421 paged_kv_last_page_len,
1225 - self._num_attention_heads,
1422 + metadata_num_heads,
1423 1,
1227 - causal,
1424 + is_causal,
1425 self._mla_work_meta_data,
1426 self._mla_work_info_set,
1427 self._mla_work_indptr,
@@ -1238,6 +1435,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1435 fast_mode=True,
1436 dtype_q=self._mla_q_dtype,
1437 dtype_kv=self._mla_kv_dtype,
1438 + **cprr_kwargs,
1439 )
1440 has_persistent_metadata = True
1441
@@ -1261,7 +1459,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1459 min_kv_seq_len = int(per_req_len.min().item())
1460
1461 dcp_verify = None
1264 - if use_segmented_dcp_verify:
1462 + if dcp_route is _DCPDecodeRoute.SEGMENTED:
1463 assert dcp_tot_seq_lens_device is not None
1464 dcp_verify = self._build_dcp_verify_row_view(
1465 int(max_qo_len),
@@ -1279,7 +1477,19 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1477 dcp_tot_seq_lens=dcp_tot_seq_lens_device,
1478 max_qo_len=max_qo_len,
1479 min_kv_seq_len=min_kv_seq_len,
1480 + dcp_route=dcp_route,
1481 dcp_verify=dcp_verify,
1482 + g_kv_indptr=g_kv_indptr,
1483 + cp_world_size=self.dcp_world_size,
1484 + cp_rank=self.dcp_rank,
1485 + asm_decode_num_heads=(
1486 + self._asm_dcp_verify_heads if dcp_route is _DCPDecodeRoute.CPRR else 0
1487 + ),
1488 + mla_num_kv_splits=(
1489 + self._mla_max_split_per_batch
1490 + if dcp_route is _DCPDecodeRoute.CPRR
1491 + else 0
1492 + ),
1493 use_gluon_decode=use_gluon_decode,
1494 use_gluon_verify=use_gluon_verify,
1495 attn_out_dtype=self.decode_attn_out_dtype,
@@ -2062,8 +2272,10 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
2272 )
2273 return o, None
2274
2065 - verify = decode.dcp_verify
2066 - if verify is not None:
2275 + dcp_route = decode.dcp_route
2276 + if dcp_route is _DCPDecodeRoute.SEGMENTED:
2277 + verify = decode.dcp_verify
2278 + assert verify is not None
2279 if type(q) is tuple:
2280 q_nope, q_pe = q
2281 else:
@@ -2085,11 +2297,6 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
2297 decode.attn_out_dtype,
2298 )
2299
2088 - if self.dcp_world_size > 1 and int(decode.max_qo_len) > 1:
2089 - raise RuntimeError(
2090 - "ROCM_AITER_MLA DCP multi-token verify requires segmented MLA."
2091 - )
2092 -
2300 if type(q) is tuple:
2301 q = torch.cat(q, dim=-1)
2302
@@ -2102,8 +2309,19 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
2309 "ROCM_AITER_MLA decode expected the DCP-gathered query head count "
2310 f"{self._decode_num_heads}, got {q.shape[1]}"
2311 )
2105 - mla_padded_q = AiterMLAHelper.get_mla_padded_q(self._decode_num_heads, q)
2106 - mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads)
2312 + # cprr has kernels only at native head counts, so pad q up (e.g. 96 ->
2313 + # 128). Heads are independent in MLA; the extras are sliced off below.
2314 + asm_dcp_heads = decode.asm_decode_num_heads
2315 + if asm_dcp_heads > self._decode_num_heads:
2316 + mla_num_heads = asm_dcp_heads
2317 + mla_padded_q = AiterMLAHelper.get_mla_padded_q(
2318 + self._decode_num_heads, q, target_heads=mla_num_heads
2319 + )
2320 + else:
2321 + mla_padded_q = AiterMLAHelper.get_mla_padded_q(self._decode_num_heads, q)
2322 + mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(
2323 + self._decode_num_heads
2324 + )
2325 o = torch.empty(
2326 B,
2327 mla_num_heads,
@@ -2137,6 +2355,16 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
2355
2356 lse = None
2357 if self.dcp_world_size > 1:
2358 + if dcp_route is _DCPDecodeRoute.CPRR:
2359 + # cprr only: global positions for the in-kernel causal window.
2360 + assert decode.g_kv_indptr is not None
2361 + mla_kwargs["g_kv_indptr"] = decode.g_kv_indptr
2362 + mla_kwargs["cp_world_size"] = decode.cp_world_size
2363 + mla_kwargs["cp_rank"] = decode.cp_rank
2364 + # Must match the schedule get_mla_metadata_v1 was built with.
2365 + assert decode.mla_num_kv_splits > 0
2366 + mla_kwargs["num_kv_splits"] = decode.mla_num_kv_splits
2367 + mla_kwargs["intra_batch_mode"] = False
2368 # The vLLM custom-op wrapper exposes only the in-place output and
2369 # drops aiter's final LSE, which the cross-shard merge needs, so go
2370 # through aiter's native entry point on the DCP path.
@@ -2151,6 +2379,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
2379 decode.max_qo_len,
2380 sm_scale=self.scale,
2381 return_lse=True,
2382 + causal=attn_metadata.causal,
2383 **mla_kwargs,
2384 )
2385 assert lse is not None, (
@@ -2172,7 +2401,13 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
2401 **mla_kwargs,
2402 )
2403
2175 - output = AiterMLAHelper.get_mla_unpadded_o(self._decode_num_heads, o)
2176 - if lse is not None:
2177 - lse = AiterMLAHelper.get_mla_unpadded_lse(self._decode_num_heads, lse)
2404 + if asm_dcp_heads > self._decode_num_heads:
2405 + # Keep a view: a post-kernel copy races the symm-mem a2a combine.
2406 + output = o[:, : self._decode_num_heads]
2407 + if lse is not None:
2408 + lse = lse[:, : self._decode_num_heads, ...]
2409 + else:
2410 + output = AiterMLAHelper.get_mla_unpadded_o(self._decode_num_heads, o)
2411 + if lse is not None:
2412 + lse = AiterMLAHelper.get_mla_unpadded_lse(self._decode_num_heads, lse)
2413 return output, lse