@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e476556189

[ROCm][CI] Add attention-sink support to ROCm AITER sparse MLA (#54404)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Andreas Karatzas committed Sep 7, 2026 at 04:26 UTC e476556189506df612218c0888dc675e6ecbfc06
8 files changed +997 -67
tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
+43 -15
@@ -1,11 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""ROCm custom op schema test for AITER MLA decode.
3 +"""ROCm custom op schema tests for AITER MLA decode.
4
5 -A single ``opcheck`` call on ``torch.ops.vllm.rocm_aiter_mla_decode_fwd``
6 -verifies that the custom op is registered and that its schema and fake
7 -implementation are consistent with the real kernel: fake-tensor support for
8 -torch.compile tracing and the ``mutates_args=["o"]`` in-place output aliasing.
5 +``opcheck`` verifies that the decode ops are registered and that their schemas
6 +and fake implementations are consistent with the real kernels: fake-tensor
7 +support for torch.compile tracing and ``mutates_args=["o"]`` in-place output
8 +aliasing.
9 """
10
11 import pytest
@@ -14,16 +14,9 @@ import torch
14 from tests.kernels.utils import opcheck
15 from vllm.platforms import current_platform
16
17 -_SKIP_NON_MI3XX = True
18 -if current_platform.is_rocm():
19 - from vllm.platforms.rocm import on_mi3xx
20 -
21 - _SKIP_NON_MI3XX = not on_mi3xx()
22 -
23 -pytestmark = [
24 - pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
25 - pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
26 -]
17 +pytestmark = pytest.mark.skipif(
18 + not current_platform.is_rocm(), reason="ROCm-specific tests"
19 +)
20
21 Q_HEAD_DIM = 576 # kv_lora_rank + qk_rope_head_dim
22 V_HEAD_DIM = 512 # kv_lora_rank
@@ -31,6 +24,10 @@ V_HEAD_DIM = 512 # kv_lora_rank
24
25 def _require_aiter():
26 from vllm._aiter_ops import is_aiter_found_and_supported
27 + from vllm.platforms.rocm import get_cdna_version
28 +
29 + if get_cdna_version() not in (3, 4):
30 + pytest.skip("AITER MLA requires CDNA 3 or 4")
31
32 if not is_aiter_found_and_supported():
33 pytest.skip("aiter is required on supported ROCm hardware for this test")
@@ -77,3 +74,34 @@ def test_mla_decode_fwd_op_schema() -> None:
74 "reduce_partial_map": None,
75 },
76 )
77 +
78 +
79 +@torch.inference_mode()
80 +def test_mla_decode_fwd_lse_op_schema() -> None:
81 + """Validate graph registration and mutation schema for LSE decode."""
82 + _require_aiter()
83 + # Import ensures the custom op is registered.
84 + from vllm._aiter_ops import rocm_aiter_ops # noqa: F401
85 +
86 + batch_size, nhead = 2, 16
87 + q = torch.randn(batch_size, nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
88 + kv_buffer = torch.randn(32, 1, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
89 + o = torch.zeros(batch_size, nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
90 + qo_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda")
91 + kv_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") * 16
92 + kv_indices = torch.arange(32, dtype=torch.int32, device="cuda")
93 + kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
94 +
95 + opcheck(
96 + torch.ops.vllm.rocm_aiter_mla_decode_fwd_lse,
97 + (q, kv_buffer, o, qo_indptr, 1),
98 + {
99 + "kv_indptr": kv_indptr,
100 + "kv_indices": kv_indices,
101 + "kv_last_page_lens": kv_last_page_lens,
102 + "sm_scale": Q_HEAD_DIM**-0.5,
103 + "logit_cap": 0.0,
104 + "q_scale": None,
105 + "kv_scale": None,
106 + },
107 + )
tests/kernels/attention/test_rocm_aiter_mla_sink.py new
+512
@@ -0,0 +1,512 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for ROCm AITER sparse MLA attention sinks."""
4 +
5 +from types import SimpleNamespace
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.platforms import current_platform
11 +from vllm.utils.torch_utils import set_random_seed
12 +
13 +pytestmark = pytest.mark.skipif(
14 + not current_platform.is_rocm(), reason="ROCm-specific tests"
15 +)
16 +
17 +Q_HEAD_DIM = 576
18 +V_HEAD_DIM = 512
19 +SM_SCALE = Q_HEAD_DIM**-0.5
20 +
21 +
22 +def _gamma_fp32(operations: int) -> float:
23 + u = torch.finfo(torch.float32).eps / 2
24 + return operations * u / (1 - operations * u)
25 +
26 +
27 +def _sink_reference(q, keys, sinks, scale):
28 + # Reference the stored operands, so input quantization is not kernel error.
29 + q, keys, sinks = q.double(), keys.double(), sinks.double()
30 + scores = q @ keys.T * scale
31 + logits = torch.cat((scores, sinks[:, None]), dim=-1)
32 + probabilities = logits.softmax(dim=-1)[:, :-1]
33 + values = keys[:, :V_HEAD_DIM]
34 + expected = probabilities @ values
35 + value_scale = probabilities @ values.abs()
36 + # FP32 dot accumulation and conversion/multiplication of the scale.
37 + score_error = _gamma_fp32(q.shape[-1] + 2) * (q.abs() @ keys.abs().T) * abs(scale)
38 + score_error = score_error.amax(dim=-1) if keys.shape[0] else torch.zeros_like(sinks)
39 + return expected, logits.logsumexp(dim=-1), value_scale, score_error
40 +
41 +
42 +def _assert_sink_output_close(
43 + actual, expected, value_scale, score_error, probability_dtype, native
44 +):
45 + assert actual.shape == expected.shape
46 + # Budget the identifiable rounding stages against the absolute value
47 + # contributions: output-relative ULPs become unbounded when values cancel.
48 + # Rounding P before P@V costs u(P) * sum(P*abs(V)).
49 + # AITER rounds output before and after sink scaling; Triton rounds once.
50 + u_probability = torch.finfo(probability_dtype).eps / 2
51 + u_output = torch.finfo(actual.dtype).eps / 2
52 + output_rounds = 2 if native else 1
53 + amplification = (1 + u_output) ** output_rounds
54 + relative_error = (1 + u_probability) * torch.exp(2 * score_error) - 1
55 + allowance = (
56 + relative_error.unsqueeze(-1) * value_scale * amplification
57 + + (amplification - 1) * expected.abs()
58 + + output_rounds * torch.finfo(actual.dtype).tiny * u_output
59 + )
60 + actual = actual.to(device=expected.device, dtype=torch.float64)
61 + assert torch.isfinite(actual).all()
62 + assert torch.count_nonzero(actual[value_scale == 0]) == 0
63 + error = (actual - expected).abs()
64 + max_ratio = (error / allowance.clamp_min(1e-300)).max().item()
65 + assert max_ratio <= 1, f"exceeded rounding budget by {max_ratio:.3f}x"
66 +
67 +
68 +def _require_aiter() -> None:
69 + from vllm._aiter_ops import is_aiter_found_and_supported
70 + from vllm.platforms.rocm import get_cdna_version
71 +
72 + if get_cdna_version() not in (3, 4):
73 + pytest.skip("AITER MLA requires CDNA 3 or 4")
74 +
75 + if not is_aiter_found_and_supported():
76 + pytest.skip("aiter is required on supported ROCm hardware for this test")
77 +
78 +
79 +@pytest.mark.parametrize(
80 + ("real_heads", "cache_kind"),
81 + [
82 + (1, "bf16"),
83 + (4, "bf16"),
84 + (8, "bf16"),
85 + (8, "fp8"),
86 + (32, "fp8"),
87 + (64, "fp8"),
88 + (12, "bf16"),
89 + (16, "bf16"),
90 + (24, "bf16"),
91 + (32, "bf16"),
92 + (40, "bf16"),
93 + (48, "bf16"),
94 + (64, "bf16"),
95 + (80, "bf16"),
96 + ],
97 +)
98 +@torch.inference_mode()
99 +def test_sparse_mla_sink_matches_ragged_reference(
100 + real_heads: int, cache_kind: str
101 +) -> None:
102 + """Exercise ragged sink decode, head padding, and FP8 scale forwarding."""
103 + _require_aiter()
104 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper
105 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
106 + ROCMAiterMLASparseImpl,
107 + )
108 +
109 + set_random_seed(real_heads * 17 + (cache_kind == "fp8"))
110 + device = torch.device("cuda")
111 + seq_lens = [0, 1, 5, 23, 17]
112 + batch_size = len(seq_lens)
113 +
114 + q_source = torch.randn(batch_size, real_heads, Q_HEAD_DIM, device=device) * 0.2
115 +
116 + pool_size = sum(seq_lens) + 52
117 + kv_source = torch.randn(pool_size, 1, Q_HEAD_DIM, device=device) * 0.2
118 + if cache_kind == "fp8":
119 + fp8_dtype = current_platform.fp8_dtype()
120 + q_scale = torch.tensor(0.5, dtype=torch.float32, device=device)
121 + kv_scale = torch.tensor(0.25, dtype=torch.float32, device=device)
122 + q_real = (q_source / q_scale).to(fp8_dtype)
123 + kv = (kv_source / kv_scale).to(fp8_dtype)
124 + q_ref = q_real.float() * q_scale
125 + kv_ref = kv.float() * kv_scale
126 + else:
127 + q_scale = kv_scale = None
128 + q_real = q_source.to(torch.bfloat16)
129 + kv = kv_source.to(torch.bfloat16)
130 + q_ref = q_real.float()
131 + kv_ref = kv.float()
132 + q = AiterMLAHelper.get_mla_padded_q(real_heads, q_real)
133 +
134 + indices = torch.randperm(pool_size, device=device)[: sum(seq_lens)].to(torch.int32)
135 + kv_indptr = torch.tensor(
136 + [0] + [sum(seq_lens[:i]) for i in range(1, len(seq_lens) + 1)],
137 + dtype=torch.int32,
138 + device=device,
139 + )
140 +
141 + # Non-None garbage proves the sink path cannot accidentally select the
142 + # gfx942 persistent kernel, which has no return-LSE code object.
143 + metadata = SimpleNamespace(
144 + attn_out_dtype=torch.bfloat16,
145 + qo_indptr=torch.arange(batch_size + 1, dtype=torch.int32, device=device),
146 + paged_kv_indptr=kv_indptr,
147 + paged_kv_indices=indices,
148 + paged_kv_last_page_len=torch.ones(batch_size, dtype=torch.int32, device=device),
149 + work_meta_data=torch.tensor([123], dtype=torch.int32),
150 + work_indptr=None,
151 + work_info_set=None,
152 + reduce_indptr=None,
153 + reduce_final_map=None,
154 + reduce_partial_map=None,
155 + num_prefills=0,
156 + num_decodes=batch_size,
157 + num_decode_tokens=batch_size,
158 + max_query_len=1,
159 + )
160 + sinks = torch.linspace(-2.0, 6.0, real_heads, device=device)
161 +
162 + impl = object.__new__(ROCMAiterMLASparseImpl)
163 + impl.num_heads = real_heads
164 + impl.kv_lora_rank = V_HEAD_DIM
165 + impl.kv_cache_dtype = cache_kind
166 + impl.scale = SM_SCALE
167 + impl.sinks = sinks
168 +
169 + output, lse = impl._forward_mla(
170 + SimpleNamespace(_q_scale=q_scale, _k_scale=kv_scale), q, kv, metadata
171 + )
172 + kv_flat = kv_ref[:, 0]
173 + references = []
174 + start = 0
175 + for batch_idx, seq_len in enumerate(seq_lens):
176 + rows = kv_flat[indices[start : start + seq_len].long()]
177 + references.append(_sink_reference(q_ref[batch_idx], rows, sinks, SM_SCALE))
178 + start += seq_len
179 + expected, expected_lse, value_scale, score_error = (
180 + torch.stack(values) for values in zip(*references)
181 + )
182 +
183 + assert output.dtype == torch.bfloat16
184 + if lse is not None:
185 + assert lse.dtype == torch.float32
186 + assert lse.shape == expected_lse.shape
187 + lse_error = (lse.double() - expected_lse).abs()
188 + rounded_lse = expected_lse.float()
189 + positive_inf = torch.full_like(rounded_lse, float("inf"))
190 + ulp = torch.maximum(
191 + torch.nextafter(rounded_lse, positive_inf) - rounded_lse,
192 + rounded_lse - torch.nextafter(rounded_lse, -positive_inf),
193 + ).double()
194 + lse_allowance = score_error + 2 * ulp
195 + assert torch.all(lse_error <= lse_allowance)
196 + torch.testing.assert_close(lse[0].double(), expected_lse[0], atol=0, rtol=0)
197 + else:
198 + from vllm.platforms.rocm import on_gfx942
199 +
200 + assert cache_kind == "bf16" and real_heads in (40, 48, 64) and on_gfx942()
201 + _assert_sink_output_close(
202 + output,
203 + expected,
204 + value_scale,
205 + score_error,
206 + q_real.dtype,
207 + native=lse is not None,
208 + )
209 +
210 +
211 +@pytest.mark.parametrize(
212 + ("q_dtype", "kv_dtype"),
213 + [
214 + (torch.float16, torch.bfloat16),
215 + (torch.bfloat16, torch.float16),
216 + ],
217 +)
218 +def test_sparse_mla_sink_rejects_unsupported_aiter_dtypes(
219 + q_dtype: torch.dtype, kv_dtype: torch.dtype
220 +) -> None:
221 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
222 + ROCMAiterMLASparseImpl,
223 + )
224 +
225 + impl = object.__new__(ROCMAiterMLASparseImpl)
226 + impl.num_heads = 16
227 + impl.kv_lora_rank = V_HEAD_DIM
228 + impl.kv_cache_dtype = "auto"
229 + impl.scale = SM_SCALE
230 + impl.sinks = torch.zeros(16, dtype=torch.float32)
231 + metadata = SimpleNamespace(
232 + attn_out_dtype=torch.bfloat16,
233 + num_prefills=0,
234 + num_decodes=1,
235 + num_decode_tokens=1,
236 + max_query_len=1,
237 + )
238 +
239 + with pytest.raises(ValueError, match="both use BF16 or both use FP8"):
240 + impl._forward_mla(
241 + SimpleNamespace(_q_scale=None, _k_scale=None),
242 + torch.empty(1, 16, Q_HEAD_DIM, dtype=q_dtype),
243 + torch.empty(1, 1, Q_HEAD_DIM, dtype=kv_dtype),
244 + metadata,
245 + )
246 +
247 +
248 +def _make_noncontiguous_sink() -> torch.Tensor:
249 + return torch.empty(8, dtype=torch.float32)[::2]
250 +
251 +
252 +@pytest.mark.parametrize(
253 + ("sinks", "match"),
254 + [
255 + (torch.empty(4, dtype=torch.bfloat16), "must be float32"),
256 + (torch.empty(2, 2, dtype=torch.float32), "must have shape"),
257 + (_make_noncontiguous_sink(), "must be contiguous"),
258 + ],
259 +)
260 +def test_sparse_mla_sink_validation(sinks: torch.Tensor, match: str) -> None:
261 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
262 + ROCMAiterMLASparseImpl,
263 + )
264 +
265 + with pytest.raises(ValueError, match=match):
266 + ROCMAiterMLASparseImpl(
267 + num_heads=4,
268 + head_size=Q_HEAD_DIM,
269 + scale=SM_SCALE,
270 + num_kv_heads=1,
271 + alibi_slopes=None,
272 + sliding_window=None,
273 + kv_cache_dtype="bfloat16",
274 + logits_soft_cap=None,
275 + attn_type="decoder",
276 + kv_sharing_target_layer_name=None,
277 + sinks=sinks,
278 + kv_lora_rank=V_HEAD_DIM,
279 + )
280 +
281 +
282 +def test_sparse_mla_backend_reports_sink_support_for_current_hardware() -> None:
283 + from vllm.platforms.rocm import get_cdna_version
284 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
285 + ROCMAiterMLASparseBackend,
286 + )
287 +
288 + assert ROCMAiterMLASparseBackend.supports_sink() == (get_cdna_version() in (3, 4))
289 +
290 +
291 +def test_sparse_mla_backend_rejects_dcp() -> None:
292 + from vllm.platforms.rocm import RocmPlatform
293 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
294 + from vllm.v1.attention.selector import AttentionSelectorConfig
295 +
296 + selector_config = AttentionSelectorConfig(
297 + head_size=Q_HEAD_DIM,
298 + dtype=torch.bfloat16,
299 + kv_cache_dtype="bfloat16",
300 + block_size=16,
301 + use_mla=True,
302 + has_sink=True,
303 + use_sparse=True,
304 + use_mm_prefix=False,
305 + use_per_head_quant_scales=False,
306 + attn_type="decoder",
307 + use_dcp=True,
308 + )
309 +
310 + with pytest.raises(ValueError, match="DCP not supported"):
311 + RocmPlatform.get_attn_backend_cls(
312 + selected_backend=AttentionBackendEnum.ROCM_AITER_MLA_SPARSE,
313 + attn_selector_config=selector_config,
314 + )
315 +
316 +
317 +@pytest.mark.parametrize("num_heads,block_size", [(8, 16), (20, 32)])
318 +@pytest.mark.parametrize("interleaved_pages", [False, True])
319 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
320 +@torch.inference_mode()
321 +def test_sparse_mla_sink_matches_dense_attention_with_empty_rows_and_paged_cache(
322 + num_heads: int,
323 + block_size: int,
324 + interleaved_pages: bool,
325 + dtype: torch.dtype,
326 +) -> None:
327 + """Preserve head-specific sinks and latent values across ragged cache pages."""
328 + _require_aiter()
329 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper
330 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
331 + ROCMAiterMLASparseImpl,
332 + )
333 + from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view
334 +
335 + head_dim, value_dim = 576, 512
336 + generator = torch.Generator().manual_seed(421)
337 + # The empty row exercises sink-only attention; the last row crosses both
338 + # page boundaries and the kernel's 16-key reduction tile.
339 + selected_rows = [
340 + [],
341 + [(2, 7)],
342 + [(3, 1), (0, block_size - 1), (1, 0), (2, 4)],
343 + [(i % 4, (i * 7) % block_size) for i in range(23)],
344 + ]
345 + q_cpu = torch.randn(
346 + len(selected_rows), num_heads, head_dim, generator=generator
347 + ).to(dtype)
348 + cache_cpu = torch.randn(4, block_size, head_dim, generator=generator).to(dtype)
349 + sinks_cpu = torch.linspace(-4.0, 6.0, num_heads)
350 +
351 + # Layers can share a backing allocation, leaving unused rows between pages.
352 + num_layers = 3 if interleaved_pages else 1
353 + backing = torch.full(
354 + (4, num_layers, block_size, head_dim),
355 + float("nan"),
356 + dtype=dtype,
357 + device="cuda",
358 + )
359 + cache = backing[:, num_layers - 1]
360 + cache.copy_(cache_cpu)
361 + indices = [
362 + page * num_layers * block_size + offset
363 + for rows in selected_rows
364 + for page, offset in rows
365 + ]
366 + lengths = torch.tensor([0, *(len(rows) for rows in selected_rows)])
367 + metadata = SimpleNamespace(
368 + block_size=block_size,
369 + num_prefills=0,
370 + num_decodes=len(selected_rows),
371 + num_decode_tokens=len(selected_rows),
372 + max_query_len=1,
373 + qo_indptr=torch.arange(
374 + len(selected_rows) + 1, dtype=torch.int32, device="cuda"
375 + ),
376 + paged_kv_last_page_len=torch.ones(
377 + len(selected_rows), dtype=torch.int32, device="cuda"
378 + ),
379 + work_meta_data=None,
380 + attn_out_dtype=dtype,
381 + paged_kv_indices=torch.tensor(indices, dtype=torch.int32, device="cuda"),
382 + paged_kv_indptr=lengths.cumsum(0).to(device="cuda", dtype=torch.int32),
383 + )
384 + impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl)
385 + impl.num_heads = num_heads
386 + impl.head_size = head_dim
387 + impl.kv_lora_rank = value_dim
388 + impl.scale = head_dim**-0.5
389 + impl.kv_cache_dtype = "auto"
390 + impl.sinks = sinks_cpu.cuda()
391 + padded_q = AiterMLAHelper.get_mla_padded_q(num_heads, q_cpu.cuda())
392 +
393 + kv_rows, _ = flat_kv_row_view(cache, block_size)
394 + actual, _ = impl._forward_mla(
395 + SimpleNamespace(_q_scale=None, _k_scale=None),
396 + padded_q,
397 + kv_rows.unsqueeze(1),
398 + metadata,
399 + )
400 +
401 + references = []
402 + for query_idx, rows in enumerate(selected_rows):
403 + keys = (
404 + torch.stack([cache_cpu[page, offset] for page, offset in rows])
405 + if rows
406 + else cache_cpu.new_empty((0, head_dim))
407 + )
408 + references.append(
409 + _sink_reference(q_cpu[query_idx], keys, sinks_cpu, impl.scale)
410 + )
411 + expected, _, value_scale, score_error = (
412 + torch.stack(values) for values in zip(*references)
413 + )
414 + assert actual.shape == (len(selected_rows), num_heads, value_dim)
415 + assert actual.dtype == dtype
416 + _assert_sink_output_close(
417 + actual,
418 + expected,
419 + value_scale,
420 + score_error,
421 + dtype,
422 + native=dtype == torch.bfloat16,
423 + )
424 +
425 +
426 +@pytest.mark.parametrize("layout", ["LBNHC", "LBHNC", "BLNHC"])
427 +def test_sparse_mla_backend_resolves_only_contiguous_layer_layouts(monkeypatch, layout):
428 + from vllm.config import CacheConfig
429 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
430 + ROCMAiterMLASparseBackend,
431 + )
432 + from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
433 +
434 + monkeypatch.setenv("VLLM_KV_CACHE_LAYOUT", layout)
435 + config = SimpleNamespace(cache_config=CacheConfig())
436 + supported = [
437 + [x.name for x in ROCMAiterMLASparseBackend.supported_kv_cache_layouts()]
438 + ]
439 + if layout == "BLNHC":
440 + with pytest.raises(ValueError, match="does not satisfy"):
441 + resolve_kv_cache_layout(config, supported)
442 + else:
443 + assert resolve_kv_cache_layout(config, supported).name == layout
444 +
445 +
446 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
447 +@torch.inference_mode()
448 +def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype):
449 + """The public sparse forward joins latent/RoPE queries in the model dtype."""
450 + _require_aiter()
451 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
452 + ROCMAiterMLASparseImpl,
453 + )
454 +
455 + set_random_seed(412)
456 + num_tokens, num_heads, block_size = 2, 8, 16
457 + q = torch.randn(num_tokens, num_heads, Q_HEAD_DIM, device="cuda").to(dtype)
458 + kv = torch.randn(2 * block_size, 1, Q_HEAD_DIM, device="cuda").to(dtype)
459 + selected = torch.tensor([[1, 7, 18], [0, 3, 20]], dtype=torch.int32, device="cuda")
460 + sinks = torch.linspace(-3.0, 5.0, num_heads, device="cuda")
461 + impl = object.__new__(ROCMAiterMLASparseImpl)
462 + impl.num_heads = num_heads
463 + impl.kv_lora_rank = V_HEAD_DIM
464 + impl.kv_cache_dtype = "auto"
465 + impl.scale = SM_SCALE
466 + impl.sinks = sinks
467 + impl.topk_indices_buffer = torch.full(
468 + (num_tokens, 128), -1, dtype=torch.int32, device="cuda"
469 + )
470 + impl.topk_indices_buffer[:, : selected.shape[1]] = selected
471 + impl.q_concat_buffer = torch.empty_like(q)
472 + metadata = SimpleNamespace(
473 + attn_out_dtype=dtype,
474 + num_actual_tokens=num_tokens,
475 + num_prefills=0,
476 + num_decodes=num_tokens,
477 + num_decode_tokens=num_tokens,
478 + max_query_len=1,
479 + block_size=block_size,
480 + topk_tokens=impl.topk_indices_buffer.shape[1],
481 + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device="cuda"),
482 + block_table=torch.tensor([[0, 1]], dtype=torch.int32, device="cuda"),
483 + qo_indptr=torch.arange(num_tokens + 1, dtype=torch.int32, device="cuda"),
484 + paged_kv_indptr=torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda"),
485 + paged_kv_indices=torch.empty(
486 + selected.numel(), dtype=torch.int32, device="cuda"
487 + ),
488 + paged_kv_last_page_len=torch.ones(num_tokens, dtype=torch.int32, device="cuda"),
489 + work_meta_data=None,
490 + )
491 + actual, _ = impl.forward_mqa(
492 + (q[..., :V_HEAD_DIM], q[..., V_HEAD_DIM:]),
493 + kv,
494 + metadata,
495 + SimpleNamespace(_q_scale=None, _k_scale=None),
496 + )
497 + references = [
498 + _sink_reference(q[i], kv[:, 0][selected[i].long()], sinks, SM_SCALE)
499 + for i in range(num_tokens)
500 + ]
501 + expected, _, value_scale, score_error = (
502 + torch.stack(values) for values in zip(*references)
503 + )
504 + assert actual.dtype == dtype
505 + _assert_sink_output_close(
506 + actual,
507 + expected,
508 + value_scale,
509 + score_error,
510 + dtype,
511 + native=dtype == torch.bfloat16,
512 + )
tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
+21
@@ -59,6 +59,7 @@ def _make_builder():
59 builder.paged_kv_indptr = torch.zeros(
60 max_num_batched_tokens + 1, dtype=torch.int32, device="cpu"
61 )
62 + builder._use_persistent_metadata = True
63 builder._num_attention_heads = 16
64 builder._num_compute_units = current_platform.num_compute_units()
65 builder._mla_work_meta_data = torch.empty(1, dtype=torch.int32, device="cpu")
@@ -130,6 +131,7 @@ def _patch_build_deps(monkeypatch, events=None):
131 synchronize=lambda: events.append("sync") if events is not None else None
132 ),
133 )
134 + return fake_aiter
135
136
137 def test_build_populates_decode_only_split_fields(monkeypatch):
@@ -165,6 +167,25 @@ def test_build_populates_mixed_split_fields(monkeypatch):
167 assert md.prefill is None
168
169
170 +def test_sink_build_skips_persistent_metadata(monkeypatch):
171 + builder = _make_builder()
172 + builder._use_persistent_metadata = False
173 + fake_aiter = _patch_build_deps(monkeypatch)
174 +
175 + md = builder.build(
176 + common_prefix_len=0, common_attn_metadata=_make_common_metadata()
177 + )
178 +
179 + fake_aiter.get_mla_metadata_v1.assert_not_called()
180 + assert md.work_meta_data is None
181 + assert md.work_indptr is None
182 + assert md.work_info_set is None
183 + assert md.reduce_indptr is None
184 + assert md.reduce_final_map is None
185 + assert md.reduce_partial_map is None
186 + assert builder._prev_metadata_key is None
187 +
188 +
189 def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch):
190 builder = _make_builder()
191 common_metadata = _make_common_metadata()
tests/v1/attention/test_rocm_glm5next_sparse.py
+52
@@ -1,11 +1,14 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from types import SimpleNamespace
5 +
6 import pytest
7 import torch
8
9 from vllm.platforms import current_platform
10 from vllm.triton_utils import tl, triton
11 +from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
12 from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
13 _use_rocm_sparse_triton,
14 fit_kpool_indices_to_aiter,
@@ -99,6 +102,55 @@ def test_rocm_sparse_triton_route(
102 )
103
104
105 +@pytest.mark.parametrize("num_heads", [8, 12])
106 +def test_rocm_sparse_triton_route_preserves_padded_sinks(monkeypatch, num_heads):
107 + captured = {}
108 +
109 + def fake_rocm_sparse_attn_prefill(**kwargs):
110 + output = kwargs["output"]
111 + captured["attn_sink"] = kwargs["attn_sink"]
112 + output.copy_(
113 + captured["attn_sink"].to(output.dtype).view(1, -1, 1).expand_as(output)
114 + )
115 +
116 + monkeypatch.setattr(
117 + sparse_mod, "rocm_sparse_attn_prefill", fake_rocm_sparse_attn_prefill
118 + )
119 +
120 + impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl)
121 + impl.num_heads = num_heads
122 + impl.kv_lora_rank = 512
123 + impl.kv_cache_dtype = "auto"
124 + impl.scale = 512**-0.5
125 + impl.sinks = torch.arange(num_heads, dtype=torch.float32)
126 +
127 + q = torch.zeros(2, 16, 512, dtype=torch.bfloat16)
128 + kv = torch.zeros(4, 1, 512, dtype=torch.bfloat16)
129 + metadata = SimpleNamespace(
130 + attn_out_dtype=torch.bfloat16,
131 + num_prefills=1,
132 + num_decodes=0,
133 + num_decode_tokens=0,
134 + max_query_len=2,
135 + paged_kv_indices=torch.empty(0, dtype=torch.int32),
136 + paged_kv_indptr=torch.zeros(3, dtype=torch.int32),
137 + )
138 +
139 + output, lse = impl._forward_mla(SimpleNamespace(), q, kv, metadata)
140 +
141 + if num_heads == 8:
142 + expected_sinks = impl.sinks.repeat_interleave(2)
143 + else:
144 + expected_sinks = torch.cat((impl.sinks, impl.sinks[:4]))
145 + torch.testing.assert_close(captured["attn_sink"], expected_sinks)
146 + assert output.shape == (2, num_heads, 512)
147 + torch.testing.assert_close(
148 + output[:, :, 0].float(),
149 + impl.sinks.expand(2, -1),
150 + )
151 + assert lse is None
152 +
153 +
154 def test_rocm_sparse_attention_accepts_glm_nope_dimensions():
155 _validate_sparse_dims(512, 512, 0, "test")
156
vllm/_aiter_ops.py
+142
@@ -624,6 +624,110 @@ def _rocm_aiter_mla_decode_fwd_impl(
624 )
625
626
627 +def _rocm_aiter_mla_decode_fwd_lse_impl(
628 + q: torch.Tensor,
629 + kv_buffer: torch.Tensor,
630 + o: torch.Tensor,
631 + qo_indptr: torch.Tensor,
632 + max_seqlen_qo: int,
633 + kv_indptr: torch.Tensor | None = None,
634 + kv_indices: torch.Tensor | None = None,
635 + kv_last_page_lens: torch.Tensor | None = None,
636 + sm_scale: float = 1.0,
637 + logit_cap: float = 0.0,
638 + q_scale: torch.Tensor | None = None,
639 + kv_scale: torch.Tensor | None = None,
640 +) -> torch.Tensor:
641 + """Run non-persistent AITER MLA decode and return natural-log LSE.
642 +
643 + gfx942 persistent MLA kernels do not provide an LSE code object. Keeping
644 + this as a separate op makes that constraint structural: callers that need
645 + LSE cannot accidentally forward persistent work metadata.
646 + """
647 + from aiter.mla import get_meta_param, mla_decode_fwd
648 +
649 + kwargs: dict[str, float | int | torch.Tensor | None | bool] = {
650 + "sm_scale": sm_scale,
651 + "logit_cap": logit_cap,
652 + "return_lse": True,
653 + }
654 + if _check_aiter_mla_fp8_support():
655 + kwargs["q_scale"] = q_scale
656 + kwargs["kv_scale"] = kv_scale
657 +
658 + if (
659 + q.dtype == torch.bfloat16
660 + and kv_buffer.dtype == torch.bfloat16
661 + and q.shape[1] == 32
662 + ):
663 + from vllm.platforms.rocm import on_gfx950
664 +
665 + if on_gfx950():
666 + # The gfx950 H32 BF16 kernel mishandles ragged tails with split KV.
667 + # Its single-split path returns correct output and natural-log LSE.
668 + kwargs["num_kv_splits"] = 1
669 + kwargs["num_kv_splits_indptr"] = torch.arange(
670 + qo_indptr.shape[0], dtype=torch.int32, device=q.device
671 + )
672 +
673 + if q.dtype == FP8_DTYPE:
674 + assert kv_indices is not None
675 + num_kv_splits, num_kv_splits_indptr = get_meta_param(
676 + None,
677 + qo_indptr.shape[0] - 1,
678 + kv_indices.numel(),
679 + q.shape[1],
680 + max_seqlen_qo,
681 + q.dtype,
682 + )
683 + if num_kv_splits == 1:
684 + # gfx942's one-split FP8 asm writes the final output directly but
685 + # does not write either LSE buffer. Force the normal split reducer,
686 + # which produces both the same output and an accurate natural LSE.
687 + num_kv_splits = 2
688 + num_kv_splits_indptr = torch.arange(
689 + 0,
690 + (qo_indptr.shape[0]) * num_kv_splits,
691 + num_kv_splits,
692 + dtype=torch.int32,
693 + device=q.device,
694 + )
695 + kwargs["num_kv_splits"] = num_kv_splits
696 + kwargs["num_kv_splits_indptr"] = num_kv_splits_indptr
697 +
698 + _, final_lse = mla_decode_fwd(
699 + q,
700 + kv_buffer.view(-1, 1, 1, q.shape[-1]),
701 + o,
702 + qo_indptr,
703 + kv_indptr,
704 + kv_indices,
705 + kv_last_page_lens,
706 + max_seqlen_qo,
707 + **kwargs,
708 + )
709 + if final_lse is None:
710 + raise RuntimeError("AITER MLA decode did not return the requested LSE")
711 + return final_lse
712 +
713 +
714 +def _rocm_aiter_mla_decode_fwd_lse_fake(
715 + q: torch.Tensor,
716 + kv_buffer: torch.Tensor,
717 + o: torch.Tensor,
718 + qo_indptr: torch.Tensor,
719 + max_seqlen_qo: int,
720 + kv_indptr: torch.Tensor | None = None,
721 + kv_indices: torch.Tensor | None = None,
722 + kv_last_page_lens: torch.Tensor | None = None,
723 + sm_scale: float = 1.0,
724 + logit_cap: float = 0.0,
725 + q_scale: torch.Tensor | None = None,
726 + kv_scale: torch.Tensor | None = None,
727 +) -> torch.Tensor:
728 + return torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device)
729 +
730 +
731 def _rocm_aiter_w8a8_gemm_impl(
732 A: torch.Tensor,
733 B: torch.Tensor,
@@ -2170,6 +2274,13 @@ class rocm_aiter_ops:
2274 mutates_args=["o"],
2275 )
2276
2277 + direct_register_custom_op(
2278 + op_name="rocm_aiter_mla_decode_fwd_lse",
2279 + op_func=_rocm_aiter_mla_decode_fwd_lse_impl,
2280 + mutates_args=["o"],
2281 + fake_impl=_rocm_aiter_mla_decode_fwd_lse_fake,
2282 + )
2283 +
2284 direct_register_custom_op(
2285 op_name="rocm_aiter_w8a8_gemm",
2286 op_func=_rocm_aiter_w8a8_gemm_impl,
@@ -2756,6 +2867,37 @@ class rocm_aiter_ops:
2867 reduce_partial_map=reduce_partial_map,
2868 )
2869
2870 + @staticmethod
2871 + def mla_decode_fwd_lse(
2872 + q: torch.Tensor,
2873 + kv_buffer: torch.Tensor,
2874 + o: torch.Tensor,
2875 + sm_scale: float,
2876 + qo_indptr: torch.Tensor,
2877 + max_seqlen_qo: int,
2878 + kv_indptr: torch.Tensor | None = None,
2879 + kv_indices: torch.Tensor | None = None,
2880 + kv_last_page_lens: torch.Tensor | None = None,
2881 + logit_cap: float = 0.0,
2882 + q_scale: torch.Tensor | None = None,
2883 + kv_scale: torch.Tensor | None = None,
2884 + ) -> torch.Tensor:
2885 + """Run MLA decode without persistent metadata and return its LSE."""
2886 + return torch.ops.vllm.rocm_aiter_mla_decode_fwd_lse(
2887 + q,
2888 + kv_buffer.view(-1, 1, 1, q.shape[-1]),
2889 + o,
2890 + qo_indptr,
2891 + max_seqlen_qo,
2892 + kv_indptr,
2893 + kv_indices,
2894 + kv_last_page_lens,
2895 + sm_scale=sm_scale,
2896 + logit_cap=logit_cap,
2897 + q_scale=q_scale,
2898 + kv_scale=kv_scale,
2899 + )
2900 +
2901 @staticmethod
2902 def per_tensor_quant(
2903 x: torch.Tensor,
vllm/v1/attention/backend.py
+9
@@ -226,6 +226,13 @@ class AttentionBackend(ABC):
226 except NotImplementedError:
227 return False
228
229 + @classmethod
230 + def supports_dcp(cls) -> bool:
231 + try:
232 + return cls.get_impl_cls().supports_dcp
233 + except NotImplementedError:
234 + return False
235 +
236 @classmethod
237 def supports_non_causal_dcp(cls) -> bool:
238 builder_cls = cls.get_builder_cls()
@@ -324,6 +331,8 @@ class AttentionBackend(ABC):
331 invalid_reasons.append("KV connector not supported")
332 if use_pcp and not cls.supports_pcp():
333 invalid_reasons.append("PCP not supported")
334 + if use_dcp and not cls.supports_dcp():
335 + invalid_reasons.append("DCP not supported")
336 if (
337 use_adaptive_verification
338 and not cls.supports_device_cpu_query_lens_mismatch()
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+216 -50
@@ -35,7 +35,7 @@ from vllm.v1.attention.backends.utils import split_decodes_and_prefills
35 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
36 rocm_sparse_attn_prefill,
37 )
38 -from vllm.v1.kv_cache_interface import AttentionSpec
38 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
39 from vllm.v1.worker.workspace import current_workspace_manager
40
41 if TYPE_CHECKING:
@@ -350,6 +350,17 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
350 def is_sparse(cls) -> bool:
351 return True
352
353 + @classmethod
354 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
355 + # Global index conversion assumes contiguous pages within each layer.
356 + return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC)
357 +
358 + @classmethod
359 + def supports_sink(cls) -> bool:
360 + from vllm.platforms.rocm import on_mi3xx
361 +
362 + return on_mi3xx()
363 +
364
365 @dataclass
366 class ROCMAiterMLASparseMetadata(AttentionMetadata):
@@ -418,6 +429,14 @@ class ROCMAiterMLASparseMetadataBuilder(
429 self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
430 self.mla_dims = get_mla_dims(self.model_config)
431 self.topk_tokens = vllm_config.model_config.hf_config.index_topk
432 + attention_context = vllm_config.compilation_config.static_forward_context
433 + # Sink decode must use AITER's nonpersistent path. In particular,
434 + # gfx942 has no persistent+LSE kernel, and its metadata heuristic
435 + # terminates for HY-V4's TP1 H64 shape.
436 + self._use_persistent_metadata = all(
437 + getattr(attention_context[name].impl, "sinks", None) is None
438 + for name in layer_names
439 + )
440 # Bounds the KV-split heuristic (see `_sparse_decode_max_split`).
441 self._num_compute_units = current_platform.num_compute_units()
442 self.max_model_len_tensor = torch.tensor(
@@ -591,14 +610,13 @@ class ROCMAiterMLASparseMetadataBuilder(
610 paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens]
611
612 # ----- Compute persistent MLA metadata -----
594 - # The aiter sparse decode kernel uses qseqlen=1 (each query token is
595 - # treated as its own batch entry), so persistent metadata can always
596 - # be precomputed here. The kernel switches to the persistent
597 - # work-stealing path automatically when work_meta_data is non-None.
598 - # The output is a deterministic function of the per-request query and
599 - # context lengths (both clamped to topk_tokens, past which per-token KV
600 - # length saturates) and num_heads; fingerprint those CPU-side and skip
601 - # the launch when nothing changed.
613 + # The AITER sparse decode kernel uses qseqlen=1 (each query token is
614 + # treated as its own batch entry). Build its persistent work metadata
615 + # only when AITER is selected and no layer needs the nonpersistent LSE
616 + # path for attention sinks. The output is a deterministic function of
617 + # the per-request query and context lengths (both clamped to
618 + # topk_tokens, past which per-token KV length saturates) and num_heads;
619 + # fingerprint those CPU-side and skip the launch when nothing changed.
620 head_size = self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim
621 use_triton_sparse = _use_rocm_sparse_triton(
622 kv_cache_dtype=self.kv_cache_dtype,
@@ -615,7 +633,7 @@ class ROCMAiterMLASparseMetadataBuilder(
633 reduce_indptr = None
634 reduce_final_map = None
635 reduce_partial_map = None
618 - if not use_triton_sparse:
636 + if self._use_persistent_metadata and not use_triton_sparse:
637 num_reqs = common_attn_metadata.num_reqs
638 clamped_seq_lens = np.minimum(
639 common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(),
@@ -729,6 +747,7 @@ def reference_mla_sparse_prefill(
747 class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
748 is_sparse = True
749 supports_dense_mha_prefill = False
750 + supports_dcp = False
751
752 def __init__(
753 self,
@@ -753,6 +772,20 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
772 self.head_size = head_size
773 self.scale = float(scale)
774 self.num_kv_heads = num_kv_heads
775 + sinks = mla_args.pop("sinks", None)
776 + if sinks is not None:
777 + if sinks.dtype != torch.float32:
778 + raise ValueError(
779 + f"ROCm AITER MLA sinks must be float32, got {sinks.dtype}"
780 + )
781 + if sinks.ndim != 1 or sinks.numel() != num_heads:
782 + raise ValueError(
783 + "ROCm AITER MLA sinks must have shape "
784 + f"({num_heads},), got {tuple(sinks.shape)}"
785 + )
786 + if not sinks.is_contiguous():
787 + raise ValueError("ROCm AITER MLA sinks must be contiguous")
788 + self.sinks: torch.Tensor | None = sinks
789 self.kv_cache_dtype = kv_cache_dtype
790 self.kv_lora_rank: int = mla_args["kv_lora_rank"]
791 self.softmax_scale = scale
@@ -776,16 +809,31 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
809 q: torch.Tensor, # [sq, heads, d_qk]
810 kv_c_and_k_pe_cache: torch.Tensor, # [blocks, heads, d_qk]
811 attn_metadata: ROCMAiterMLASparseMetadata,
779 - ) -> torch.Tensor:
812 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
813 num_tokens = q.shape[0]
781 - mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
782 - output = torch.empty(
783 - [num_tokens, mla_num_heads, self.kv_lora_rank],
784 - dtype=attn_metadata.attn_out_dtype,
785 - device=q.device,
814 + base_mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
815 + mla_num_heads = base_mla_num_heads
816 + need_lse = self.sinks is not None
817 + from vllm.platforms.rocm import on_gfx942
818 +
819 + # Keep sink attention available for dtypes/head shapes without an
820 + # AITER return-LSE kernel. Sink layers never need persistent metadata.
821 + triton_sink_fallback = (
822 + need_lse
823 + and q.dtype == kv_c_and_k_pe_cache.dtype
824 + and (
825 + q.dtype == torch.float16
826 + or (
827 + q.dtype == torch.bfloat16
828 + and (
829 + mla_num_heads > 128
830 + or (on_gfx942() and 32 < mla_num_heads <= 64)
831 + )
832 + )
833 + )
834 )
835
788 - if _use_rocm_sparse_triton(
836 + if triton_sink_fallback or _use_rocm_sparse_triton(
837 kv_cache_dtype=self.kv_cache_dtype,
838 head_size=q.shape[-1],
839 kv_lora_rank=self.kv_lora_rank,
@@ -794,6 +842,18 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
842 num_decode_tokens=attn_metadata.num_decode_tokens,
843 max_query_len=attn_metadata.max_query_len,
844 ):
845 + output = torch.empty(
846 + [num_tokens, q.shape[1], self.kv_lora_rank],
847 + dtype=attn_metadata.attn_out_dtype,
848 + device=q.device,
849 + )
850 + triton_sinks = None
851 + if self.sinks is not None:
852 + triton_sinks = AiterMLAHelper.get_mla_padded_q(
853 + self.num_heads,
854 + self.sinks.reshape(1, self.num_heads, 1),
855 + q.shape[1],
856 + ).reshape(-1)
857 rocm_sparse_attn_prefill(
858 q=q,
859 kv=kv_c_and_k_pe_cache.view(-1, 1, q.shape[-1]),
@@ -803,44 +863,149 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
863 head_dim=q.shape[-1],
864 nope_head_dim=self.kv_lora_rank,
865 rope_head_dim=q.shape[-1] - self.kv_lora_rank,
806 - attn_sink=None,
866 + attn_sink=triton_sinks,
867 output=output,
868 ragged_indices=attn_metadata.paged_kv_indices,
869 ragged_indptr=attn_metadata.paged_kv_indptr,
870 )
811 - return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)
812 -
813 - # Build kwargs and forward the persistent MLA metadata when it has
814 - # been computed. The aiter mla_decode_fwd switches to its
815 - # work-stealing persistent kernel path when work_meta_data is given.
816 - mla_kwargs: dict = dict(
817 - q_scale=layer._q_scale,
818 - kv_scale=layer._k_scale,
819 - )
820 - if attn_metadata.work_meta_data is not None:
821 - mla_kwargs.update(
822 - work_meta_data=attn_metadata.work_meta_data,
823 - work_indptr=attn_metadata.work_indptr,
824 - work_info_set=attn_metadata.work_info_set,
825 - reduce_indptr=attn_metadata.reduce_indptr,
826 - reduce_final_map=attn_metadata.reduce_final_map,
827 - reduce_partial_map=attn_metadata.reduce_partial_map,
828 - )
871 + output = AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)
872 + return output, None
873 +
874 + # AITER's nonpersistent return-LSE dispatch has discrete head kernels.
875 + supported_head_buckets: tuple[int, ...] | None = None
876 + head_dtype_name = ""
877 + if need_lse:
878 + if (
879 + q.dtype == torch.bfloat16
880 + and kv_c_and_k_pe_cache.dtype == torch.bfloat16
881 + ):
882 + supported_head_buckets = (16, 32, 64, 128)
883 + head_dtype_name = "BF16"
884 + elif (
885 + q.dtype == current_platform.fp8_dtype()
886 + and kv_c_and_k_pe_cache.dtype == current_platform.fp8_dtype()
887 + ):
888 + supported_head_buckets = (16, 128)
889 + head_dtype_name = "FP8"
890 + else:
891 + raise ValueError(
892 + "ROCm AITER MLA attention sinks require query and KV to "
893 + "both use BF16 or both use FP8, got "
894 + f"query={q.dtype}, KV={kv_c_and_k_pe_cache.dtype}"
895 + )
896
830 - rocm_aiter_ops.mla_decode_fwd(
831 - q,
832 - kv_c_and_k_pe_cache,
833 - output,
834 - self.scale,
835 - attn_metadata.qo_indptr,
836 - 1,
837 - attn_metadata.paged_kv_indptr,
838 - attn_metadata.paged_kv_indices,
839 - attn_metadata.paged_kv_last_page_len,
840 - **mla_kwargs,
897 + if supported_head_buckets is not None:
898 + from vllm.platforms.rocm import on_mi3xx
899 +
900 + if on_mi3xx():
901 + supported_heads = next(
902 + (
903 + heads
904 + for heads in supported_head_buckets
905 + if heads >= mla_num_heads
906 + ),
907 + None,
908 + )
909 + if supported_heads is None:
910 + raise ValueError(
911 + "ROCm AITER MLA attention sinks support at most 128 "
912 + f"padded local {head_dtype_name} heads; increase "
913 + "tensor_parallel_size"
914 + )
915 + if supported_heads != mla_num_heads:
916 + q = AiterMLAHelper.get_mla_padded_q(
917 + mla_num_heads, q, supported_heads
918 + )
919 + mla_num_heads = supported_heads
920 + output = torch.empty(
921 + [num_tokens, mla_num_heads, self.kv_lora_rank],
922 + dtype=attn_metadata.attn_out_dtype,
923 + device=q.device,
924 )
925
843 - return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)
926 + if need_lse:
927 + # gfx942 has no persistent MLA code object that writes final LSE.
928 + # The split-KV path consumes the same ragged indices and is exact.
929 + lse = rocm_aiter_ops.mla_decode_fwd_lse(
930 + q,
931 + kv_c_and_k_pe_cache,
932 + output,
933 + self.scale,
934 + attn_metadata.qo_indptr,
935 + 1,
936 + attn_metadata.paged_kv_indptr,
937 + attn_metadata.paged_kv_indices,
938 + attn_metadata.paged_kv_last_page_len,
939 + q_scale=layer._q_scale,
940 + kv_scale=layer._k_scale,
941 + )
942 + else:
943 + # Preserve the persistent work-stealing fast path for models that
944 + # do not use sinks.
945 + mla_kwargs: dict = dict(
946 + q_scale=layer._q_scale,
947 + kv_scale=layer._k_scale,
948 + )
949 + if attn_metadata.work_meta_data is not None:
950 + mla_kwargs.update(
951 + work_meta_data=attn_metadata.work_meta_data,
952 + work_indptr=attn_metadata.work_indptr,
953 + work_info_set=attn_metadata.work_info_set,
954 + reduce_indptr=attn_metadata.reduce_indptr,
955 + reduce_final_map=attn_metadata.reduce_final_map,
956 + reduce_partial_map=attn_metadata.reduce_partial_map,
957 + )
958 +
959 + rocm_aiter_ops.mla_decode_fwd(
960 + q,
961 + kv_c_and_k_pe_cache,
962 + output,
963 + self.scale,
964 + attn_metadata.qo_indptr,
965 + 1,
966 + attn_metadata.paged_kv_indptr,
967 + attn_metadata.paged_kv_indices,
968 + attn_metadata.paged_kv_last_page_len,
969 + **mla_kwargs,
970 + )
971 + lse = None
972 +
973 + if mla_num_heads != base_mla_num_heads:
974 + if mla_num_heads % base_mla_num_heads == 0:
975 + head_stride = mla_num_heads // base_mla_num_heads
976 + output = output[:, ::head_stride]
977 + if lse is not None:
978 + lse = lse[:, ::head_stride]
979 + else:
980 + output = output[:, :base_mla_num_heads]
981 + if lse is not None:
982 + lse = lse[:, :base_mla_num_heads]
983 +
984 + output = AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)
985 + if lse is not None:
986 + lse = AiterMLAHelper.get_mla_unpadded_o(
987 + self.num_heads, lse.unsqueeze(-1)
988 + ).squeeze(-1)
989 +
990 + if self.sinks is not None:
991 + assert lse is not None
992 + # Empty ragged rows have only sink mass and no value contribution.
993 + # AITER can return NaN output/LSE for those rows; do not multiply it
994 + # by a zero normalization factor and propagate the NaN.
995 + has_keys = (
996 + attn_metadata.paged_kv_indptr[1:] > attn_metadata.paged_kv_indptr[:-1]
997 + ).unsqueeze(-1)
998 + lse = torch.where(has_keys, lse, float("-inf"))
999 + sink_lse = torch.logaddexp(lse, self.sinks)
1000 + sink_scale = torch.exp(lse - sink_lse)
1001 + output = torch.where(
1002 + has_keys.unsqueeze(-1),
1003 + output.float() * sink_scale.unsqueeze(-1),
1004 + 0.0,
1005 + ).to(output.dtype)
1006 + lse = sink_lse
1007 +
1008 + return output, lse
1009
1010 def forward_mqa(
1011 self,
@@ -863,6 +1028,8 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
1028 q = self.q_concat_buffer[: ql_nope.shape[0]]
1029 if q_pe.shape[-1] == 0:
1030 q.copy_(ql_nope)
1031 + elif q.dtype == torch.float16:
1032 + torch.cat((ql_nope, q_pe), dim=-1, out=q)
1033 else:
1034 ops.concat_mla_q(ql_nope, q_pe, q)
1035
@@ -892,7 +1059,6 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
1059 q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale)
1060 q = q.view(original_q_shape)
1061 mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
895 - attn_out = self._forward_mla(
1062 + return self._forward_mla(
1063 layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata
1064 )
898 - return attn_out, None
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+2 -2
@@ -2610,7 +2610,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
2610 block_d = triton.next_power_of_2(head_dim)
2611 block_k = 16 if head_dim >= 256 else 32
2612 num_warps = 4
2613 - out = torch.empty_like(q, dtype=torch.bfloat16)
2613 + out = torch.empty_like(q)
2614 _sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
2615 q,
2616 kv,
@@ -3195,7 +3195,7 @@ def rocm_sparse_attn_prefill(
3195 rope_head_dim=rope_head_dim,
3196 topk_length=topk_length,
3197 )
3198 - output.copy_(output_chunk.to(output.dtype))
3198 + output.copy_(output_chunk[..., : output.shape[-1]].to(output.dtype))
3199
3200
3201 def rocm_sparse_attn_decode(