@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70ae0b7435

[ROCm][MLA] Enable sparse MLA Gluon kernel from Aiter (#53492)

Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com> Signed-off-by: Mehmet Cagri <mehmet.kaymak@amd.com>

Mehmet Cagri committed Sep 28, 2026 at 20:08 UTC 70ae0b7435bcb1ceef051df28c50fb7ee7d1dfec
7 files changed +989 -22
tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py new
+326
@@ -0,0 +1,326 @@
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 +
11 +pytestmark = pytest.mark.skipif(
12 + not current_platform.is_rocm(), reason="ROCm-only aiter path"
13 +)
14 +
15 +
16 +def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]:
17 + ends = indptr.tolist()
18 + return [indices[s:e].tolist() for s, e in zip(ends, ends[1:])]
19 +
20 +
21 +def test_triton_sparse_mla_gate(monkeypatch) -> None:
22 + import vllm._aiter_ops as aiter_ops
23 + from vllm._aiter_ops import rocm_aiter_ops
24 + from vllm.platforms.rocm import on_gfx950
25 +
26 + warnings: list[str] = []
27 + monkeypatch.setattr(
28 + aiter_ops.logger,
29 + "warning_once",
30 + lambda msg, *args, **kwargs: warnings.append(msg % args),
31 + )
32 +
33 + def enabled(aiter: bool, flag: bool) -> bool:
34 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", str(int(aiter)))
35 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", str(int(flag)))
36 + rocm_aiter_ops.refresh_env_variables()
37 + return bool(rocm_aiter_ops.is_triton_sparse_mla_enabled())
38 +
39 + try:
40 + assert not enabled(aiter=True, flag=False) and not warnings
41 + assert not enabled(aiter=False, flag=True)
42 + assert "VLLM_ROCM_USE_AITER is off" in warnings.pop()
43 + # On gfx950 the flag takes effect; elsewhere it warns and falls back.
44 + assert enabled(aiter=True, flag=True) == on_gfx950()
45 + if not on_gfx950():
46 + assert "gfx950" in warnings.pop()
47 + finally:
48 + monkeypatch.undo()
49 + rocm_aiter_ops.refresh_env_variables()
50 +
51 +
52 +@pytest.mark.parametrize(
53 + "rope_dim",
54 + [0, 64],
55 + ids=["rope_free", "appended_rope"], # GLM-5.3-Flash; GLM-5.1/5.2, V3.2
56 +)
57 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
58 +@torch.inference_mode()
59 +def test_forward_mqa_prepares_triton_sparse_mla_inputs(
60 + monkeypatch, rope_dim: int, kv_cache_dtype: str
61 +) -> None:
62 + """What forward_mqa hands aiter: global slots built from the request-local
63 + top-k rows, q quantized with the layer's scale, and the cache as stored."""
64 + from vllm._aiter_ops import rocm_aiter_ops
65 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
66 + ROCMAiterMLASparseImpl,
67 + )
68 +
69 + device = torch.device("cuda")
70 + gen = torch.Generator().manual_seed(0)
71 + kv_lora_rank, num_heads, topk, block_size = 512, 16, 128, 16
72 + head_size = kv_lora_rank + rope_dim
73 + fp8 = kv_cache_dtype == "fp8"
74 +
75 + # A decode row, then a short prefill.
76 + query_lens, seq_lens = [1, 5], [40, 20]
77 + num_tokens = sum(query_lens)
78 + block_table = torch.randperm(8, generator=gen).view(2, 4).to(device, torch.int32)
79 + topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32)
80 + req_id_per_token: list[int] = []
81 + rows: list[list[int]] = []
82 + for req, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)):
83 + for pos in range(seq_len - query_len, seq_len):
84 + local = torch.randperm(pos + 1, generator=gen)
85 + topk_indices[len(rows), : pos + 1] = local.int()
86 + slots = block_table[req, local // block_size].cpu() * block_size
87 + rows.append((slots + local % block_size).tolist())
88 + req_id_per_token.append(req)
89 + indptr = torch.tensor([0] + [len(row) for row in rows]).cumsum(0)
90 +
91 + impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl)
92 + impl.num_heads = num_heads
93 + impl.kv_lora_rank = kv_lora_rank
94 + impl.qk_rope_head_dim = rope_dim
95 + impl.scale = head_size**-0.5
96 + impl.sinks = torch.randn(num_heads, generator=gen).to(device)
97 + impl.kv_cache_dtype = kv_cache_dtype
98 + impl.topk_indices_buffer = topk_indices.to(device)
99 + impl.use_aiter_sparse_mla = True
100 + layer = SimpleNamespace(
101 + _q_scale=torch.tensor([0.02], device=device),
102 + _k_scale=torch.tensor([0.01], device=device),
103 + )
104 + metadata = SimpleNamespace(
105 + num_actual_tokens=num_tokens,
106 + topk_tokens=topk,
107 + req_id_per_token=torch.tensor(req_id_per_token, dtype=torch.int32).to(device),
108 + block_table=block_table,
109 + block_size=block_size,
110 + paged_kv_indptr=indptr.to(device, torch.int32),
111 + paged_kv_indices=torch.zeros(num_tokens * topk, dtype=torch.int32).to(device),
112 + attn_out_dtype=torch.bfloat16,
113 + )
114 + cache_dtype = current_platform.fp8_dtype() if fp8 else torch.bfloat16
115 + kv_cache = torch.zeros(8, block_size, head_size, device=device).to(cache_dtype)
116 + # Two rows of cudagraph padding past the real tokens.
117 + q = torch.randn(num_tokens + 2, num_heads, head_size, generator=gen).to(
118 + device, torch.bfloat16
119 + )
120 + calls = []
121 + monkeypatch.setattr(
122 + rocm_aiter_ops, "triton_sparse_mla_fwd", lambda *a, **k: calls.append((a, k))
123 + )
124 +
125 + out, lse = impl.forward_mqa(q, kv_cache, metadata, layer)
126 +
127 + assert len(calls) == 1
128 + (q_in, kv_in, o, sm_scale, kv_indptr, kv_indices), kwargs = calls[0]
129 + assert _rows_from_ragged(kv_indices, kv_indptr) == rows
130 + assert q_in.shape[0] == num_tokens
131 + if fp8:
132 + assert q_in.dtype == cache_dtype
133 + torch.testing.assert_close(
134 + q_in.float() * layer._q_scale, q[:num_tokens].float(), atol=1e-3, rtol=0.07
135 + )
136 + else:
137 + assert q_in.data_ptr() == q.data_ptr()
138 + assert kv_in.shape == (8 * block_size, 1, 1, head_size)
139 + assert kv_in.data_ptr() == kv_cache.data_ptr()
140 + assert o.data_ptr() == out.data_ptr() and o.shape[0] == num_tokens
141 + assert sm_scale == impl.scale
142 + assert kwargs["kv_lora_rank"] == kv_lora_rank
143 + assert kwargs["qk_rope_head_dim"] == rope_dim
144 + assert kwargs["q_scale"] is layer._q_scale
145 + assert kwargs["kv_scale"] is layer._k_scale
146 + assert kwargs["attn_sink"] is impl.sinks
147 + assert lse is None
148 +
149 +
150 +@pytest.mark.parametrize(
151 + ("model", "compress_ratio"),
152 + [("deepseek_v41", 1), ("deepseek_v41", 2), ("deepseek_v4", 4)],
153 +)
154 +@torch.inference_mode()
155 +def test_paged_prefill_indices_match_gathered_prefill(
156 + monkeypatch, model: str, compress_ratio: int
157 +) -> None:
158 + """The aiter prefill reads the paged caches through the SWA builder's rows
159 + and build_prefill_topk_ragged_indices, and must attend exactly the keys the
160 + dequantize-and-gather prefill does."""
161 + import importlib
162 +
163 + from tests.v1.attention.utils import create_vllm_config
164 + from vllm._aiter_ops import rocm_aiter_ops
165 + from vllm.v1.attention.backend import CommonAttentionMetadata
166 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
167 + build_prefill_topk_ragged_indices,
168 + )
169 + from vllm.v1.kv_cache_interface import SlidingWindowMLASpec
170 +
171 + rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm")
172 + device = torch.device("cuda")
173 + gen = torch.Generator().manual_seed(compress_ratio)
174 + window, width, block_size, max_model_len = 16, 32, 16, 128
175 + comp_block = block_size // compress_ratio
176 + num_compressed = max_model_len // compress_ratio
177 +
178 + # A decode row, then prefills with and without a prefix.
179 + prefix_lens, query_lens = [20, 0, 13, 40], [1, 7, 20, 9]
180 + seq_lens = [p + q for p, q in zip(prefix_lens, query_lens)]
181 + num_decodes = num_decode_tokens = 1
182 + num_tokens = sum(query_lens)
183 + qsl = torch.tensor([0] + query_lens).cumsum(0).to(torch.int32)
184 + positions = torch.cat([torch.arange(p, s) for p, s in zip(prefix_lens, seq_lens)])
185 + token_to_req = torch.repeat_interleave(
186 + torch.arange(len(query_lens)), torch.tensor(query_lens)
187 + )
188 + swa_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32)
189 + comp_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32)
190 +
191 + vllm_config = create_vllm_config(
192 + model_name="facebook/opt-125m",
193 + max_model_len=max_model_len,
194 + block_size=block_size,
195 + max_num_seqs=4,
196 + max_num_batched_tokens=64,
197 + hf_config_override={
198 + "compress_ratios": [compress_ratio],
199 + "index_topk": width,
200 + "sliding_window": window,
201 + },
202 + )
203 + swa_spec = SlidingWindowMLASpec(
204 + block_size=block_size,
205 + num_kv_heads=1,
206 + head_size=512,
207 + dtype=torch.bfloat16,
208 + sliding_window=window,
209 + cache_dtype_str="fp8_ds_mla",
210 + model_version="deepseek_v4",
211 + )
212 + slots = swa_bt.cpu()[token_to_req, positions // block_size] * block_size + (
213 + positions % block_size
214 + )
215 + common = CommonAttentionMetadata(
216 + query_start_loc=qsl.to(device),
217 + query_start_loc_cpu=qsl,
218 + seq_lens=torch.tensor(seq_lens, dtype=torch.int32).to(device),
219 + seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32),
220 + num_reqs=len(query_lens),
221 + num_actual_tokens=num_tokens,
222 + max_query_len=max(query_lens),
223 + max_seq_len=max(seq_lens),
224 + block_table_tensor=swa_bt,
225 + slot_mapping=slots.to(device, torch.int64),
226 + positions=positions.to(device),
227 + causal=True,
228 + )
229 +
230 + def build(enabled: bool):
231 + monkeypatch.setattr(
232 + rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: enabled
233 + )
234 + return rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
235 + swa_spec, ["swa"], vllm_config, device
236 + ).build(0, common)
237 +
238 + assert build(enabled=False).prefill_swa_ragged_indices is None
239 + md = build(enabled=True)
240 + assert md.num_decode_tokens == num_decode_tokens
241 +
242 + # Local top-k rows. Under each row's causal cap sit a failed candidate and
243 + # one past the pool; past the cap, stale values neither path may read.
244 + num_prefill_tokens = num_tokens - num_decode_tokens
245 + topk = torch.full((num_prefill_tokens, width), -1, dtype=torch.int32)
246 + for row in range(num_prefill_tokens):
247 + visible = (int(positions[row + num_decode_tokens]) + 1) // compress_ratio
248 + cap = min(visible, width)
249 + topk[row, :cap] = torch.randperm(visible, generator=gen)[:cap].int()
250 + if cap >= 2:
251 + topk[row, :2] = torch.tensor([-1, num_compressed + 3])
252 + topk[row, cap:] = torch.randint(
253 + 0, num_compressed, (width - cap,), generator=gen
254 + ).int()
255 + topk = topk.to(device)
256 +
257 + topk_indices, topk_indptr = build_prefill_topk_ragged_indices(
258 + topk,
259 + md.token_to_req_indices,
260 + md.query_start_loc,
261 + md.seq_lens,
262 + md.is_valid_token,
263 + comp_bt,
264 + block_size=comp_block,
265 + compress_ratio=compress_ratio,
266 + num_compressed=num_compressed,
267 + token_offset=num_decode_tokens,
268 + )
269 + paged_keys = [
270 + sorted([("s", x) for x in swa if x >= 0] + [("c", x) for x in top if x >= 0])
271 + for swa, top in zip(
272 + _rows_from_ragged(
273 + md.prefill_swa_ragged_indices, md.prefill_swa_ragged_indptr
274 + ),
275 + _rows_from_ragged(topk_indices, topk_indptr),
276 + )
277 + ]
278 +
279 + # The gathered path's combine, mapped from workspace offsets to cache slots.
280 + M = num_compressed + window + num_tokens
281 + combined, combined_lens = rocm.combine_topk_swa_indices(
282 + topk,
283 + md.query_start_loc[num_decodes:],
284 + md.prefill_seq_lens,
285 + md.prefill_gather_lens,
286 + window,
287 + compress_ratio,
288 + width,
289 + M,
290 + num_compressed,
291 + )
292 + gather_start = (md.prefill_seq_lens - md.prefill_gather_lens).tolist()
293 + swa_bt_cpu, comp_bt_cpu = swa_bt.cpu(), comp_bt.cpu()
294 +
295 + def slot(block_table, req, pos, size):
296 + return int(block_table[req, pos // size]) * size + pos % size
297 +
298 + gathered_keys = []
299 + for row, (values, n) in enumerate(zip(combined.tolist(), combined_lens.tolist())):
300 + req = int(token_to_req[row + num_decode_tokens])
301 + chunk_req = req - num_decodes
302 + keys = []
303 + for value in values[:n]:
304 + if value < 0:
305 + continue
306 + offset = value - chunk_req * M
307 + if offset < num_compressed:
308 + keys.append(("c", slot(comp_bt_cpu, req, offset, comp_block)))
309 + else:
310 + pos = gather_start[chunk_req] + offset - num_compressed
311 + keys.append(("s", slot(swa_bt_cpu, req, pos, block_size)))
312 + gathered_keys.append(sorted(keys))
313 +
314 + assert paged_keys == gathered_keys
315 +
316 +
317 +@pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"])
318 +def test_aiter_sparse_mla_stays_off_with_kv_connector(monkeypatch, model) -> None:
319 + import importlib
320 +
321 + from vllm._aiter_ops import rocm_aiter_ops
322 +
323 + rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm")
324 + monkeypatch.setattr(rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: True)
325 + assert rocm._aiter_sparse_mla_enabled(has_kv_transfer=False)
326 + assert not rocm._aiter_sparse_mla_enabled(has_kv_transfer=True)
vllm/_aiter_ops.py
+73
@@ -1927,6 +1927,7 @@ class rocm_aiter_ops:
1927 VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
1928 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
1929 VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
1930 + VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: Controls Triton sparse MLA (gfx950).
1931 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
1932 VLLM_ROCM_USE_AITER_MOE_SITUV2: SiTUv2 FlyDSL MoE activation
1933 dtype (a16w4 | a8w4 | a4w4).
@@ -1998,6 +1999,7 @@ class rocm_aiter_ops:
1999 _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
2000 # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
2001 _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
2002 + _TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA
2003 _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
2004 _MOE_SITUV2 = _resolve_situv2_activation()
2005 # TODO: Consolidate under _LINEAR_ENABLED
@@ -2027,6 +2029,7 @@ class rocm_aiter_ops:
2029 cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
2030 cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
2031 cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
2032 + cls._TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA
2033 cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
2034 cls._MOE_SITUV2 = _resolve_situv2_activation()
2035 _sync_aiter_situv2_moe_env()
@@ -2391,6 +2394,25 @@ class rocm_aiter_ops:
2394 def is_triton_rotary_embed_enabled(cls) -> bool:
2395 return cls._AITER_ENABLED and cls._TRITON_ROTARY_EMBED
2396
2397 + @classmethod
2398 + @if_aiter_supported
2399 + def is_triton_sparse_mla_enabled(cls) -> bool:
2400 + if not cls._TRITON_SPARSE_MLA:
2401 + return False
2402 + from vllm.platforms.rocm import on_gfx950
2403 +
2404 + if not cls._AITER_ENABLED:
2405 + reason = "VLLM_ROCM_USE_AITER is off"
2406 + elif not on_gfx950():
2407 + reason = "the kernel is gfx950-only"
2408 + else:
2409 + return True
2410 + logger.warning_once(
2411 + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; ignoring it.",
2412 + reason,
2413 + )
2414 + return False
2415 +
2416 @classmethod
2417 @if_aiter_supported
2418 def is_triton_gemm_enabled(cls) -> bool:
@@ -3137,6 +3159,57 @@ class rocm_aiter_ops:
3159 ) -> tuple[torch.Tensor, torch.Tensor]:
3160 return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up)
3161
3162 + @staticmethod
3163 + def triton_sparse_mla_fwd(
3164 + q: torch.Tensor,
3165 + kv_buffer: torch.Tensor,
3166 + o: torch.Tensor,
3167 + sm_scale: float,
3168 + kv_indptr: torch.Tensor,
3169 + kv_indices: torch.Tensor,
3170 + kv_lora_rank: int = 512,
3171 + qk_rope_head_dim: int = 64,
3172 + q_scale: torch.Tensor | None = None,
3173 + kv_scale: torch.Tensor | None = None,
3174 + attn_sink: torch.Tensor | None = None,
3175 + extra_kv_buffer: torch.Tensor | None = None,
3176 + extra_kv_indptr: torch.Tensor | None = None,
3177 + extra_kv_indices: torch.Tensor | None = None,
3178 + has_invalid: bool = True,
3179 + ) -> None:
3180 + """Sparse MLA read straight from the KV cache, for prefill and decode.
3181 +
3182 + kv_indices are ragged global slot ids. With has_invalid they may hold
3183 + -1, which the kernel masks; a caller whose stream never holds a
3184 + negative slot passes False, which lets the fp8 kernel stage K straight
3185 + into LDS. The cache format (bf16, per-tensor fp8, or DeepSeek V4's
3186 + paged fp8_ds_mla) is inferred from kv_buffer and kv_scale. An fp8 q
3187 + must come with its q_scale and runs both dots in fp8. The extra segment
3188 + is DeepSeek V4's second cache: SWA window in kv_buffer, top-k
3189 + compressed tokens here.
3190 + """
3191 + from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd
3192 +
3193 + fp8_q = q.dtype == FP8_DTYPE
3194 + sparse_mla_fwd(
3195 + q,
3196 + kv_buffer,
3197 + kv_indptr,
3198 + kv_indices,
3199 + sm_scale,
3200 + kv_scale=kv_scale,
3201 + kv_lora_rank=kv_lora_rank,
3202 + qk_rope_head_dim=qk_rope_head_dim,
3203 + has_invalid=has_invalid,
3204 + dot_precision="fp8" if fp8_q else "bf16",
3205 + q_scale=q_scale if fp8_q else None,
3206 + out=o,
3207 + attn_sink=attn_sink,
3208 + extra_kv=extra_kv_buffer,
3209 + extra_indptr=extra_kv_indptr,
3210 + extra_indices=extra_kv_indices,
3211 + )
3212 +
3213 @staticmethod
3214 def mla_decode_fwd(
3215 q: torch.Tensor,
vllm/envs.py
+9
@@ -148,6 +148,7 @@ if TYPE_CHECKING:
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_SPARSE_MLA: bool = False
152 VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
153 VLLM_ROCM_USE_AITER_FP8BMM: bool = True
154 VLLM_ROCM_USE_AITER_FP4BMM: bool = True
@@ -1343,6 +1344,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
1344 "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: (
1345 os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1")
1346 ),
1347 + # Whether sparse MLA prefill and decode run on aiter's Triton kernel, which
1348 + # reads the KV cache as stored (bf16 or fp8, paged or flat). Used by the
1349 + # ROCM_AITER_MLA_SPARSE backend (DeepSeek V3.2, GLM-5.x) and DeepSeek
1350 + # V4 / V4.1. gfx950 only. By default is disabled.
1351 + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA": lambda: (
1352 + os.getenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", "False").lower()
1353 + in ("true", "1")
1354 + ),
1355 # Whether to use aiter rope.
1356 # By default is disabled.
1357 "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: (
vllm/models/deepseek_v4/amd/rocm.py
+172 -12
@@ -35,6 +35,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
35 DeepseekSparseSWAMetadataBuilder,
36 )
37 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
38 + build_prefill_topk_ragged_indices,
39 build_ragged_indices_from_dense,
40 rocm_inv_rope_einsum,
41 rocm_inverse_rope_rows_,
@@ -103,6 +104,22 @@ def _trust_dsv4_extra_cache_nan_free(
104 )
105
106
107 +def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool:
108 + from vllm._aiter_ops import rocm_aiter_ops
109 +
110 + if not rocm_aiter_ops.is_triton_sparse_mla_enabled():
111 + return False
112 + # The aiter kernel does not scrub NaNs, so like the trust check above it
113 + # needs a cache only the canonical writer has touched.
114 + if has_kv_transfer:
115 + logger.warning_once(
116 + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV "
117 + "connector; DeepSeek V4 keeps the Triton sparse attention."
118 + )
119 + return False
120 + return True
121 +
122 +
123 def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
124 lengths = lengths.to(dtype=torch.int32).contiguous()
125 indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
@@ -446,6 +463,9 @@ class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata):
463 class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
464 decode_swa_ragged_indices: torch.Tensor | None = None
465 decode_swa_ragged_indptr: torch.Tensor | None = None
466 + # Built only for the aiter prefill, which reads the SWA cache in place.
467 + prefill_swa_ragged_indices: torch.Tensor | None = None
468 + prefill_swa_ragged_indptr: torch.Tensor | None = None
469
470
471 class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
@@ -565,6 +585,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
585 dtype=torch.int32,
586 device=self.device,
587 )
588 + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(
589 + self.vllm_config.kv_transfer_config is not None
590 + )
591
592 def build(
593 self,
@@ -580,6 +603,23 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
603 replay_start=replay_start,
604 )
605
606 + prefill_ragged_indices = None
607 + prefill_ragged_indptr = None
608 + if (
609 + self._use_aiter_sparse_mla
610 + and base.num_prefill_tokens > 0
611 + and base.prefill_swa_indices is not None
612 + and base.prefill_swa_lens is not None
613 + ):
614 + # Paged-coordinate rows (image visibility included), so every
615 + # layer's prefill shares one build.
616 + prefill_ragged_indices, prefill_ragged_indptr = (
617 + build_ragged_indices_from_dense(
618 + base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1),
619 + base.prefill_swa_lens,
620 + )
621 + )
622 +
623 ragged_indices = None
624 ragged_indptr = None
625 if (
@@ -606,6 +646,8 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
646 **vars(base),
647 decode_swa_ragged_indices=ragged_indices,
648 decode_swa_ragged_indptr=ragged_indptr,
649 + prefill_swa_ragged_indices=prefill_ragged_indices,
650 + prefill_swa_ragged_indptr=prefill_ragged_indptr,
651 )
652
653
@@ -623,6 +665,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
665 """ROCm sparse MLA attention layer for DeepSeek V4."""
666
667 backend_cls = DeepseekV4ROCMAiterMLASparseBackend
668 + _use_aiter_sparse_mla = False
669
670 def __init__(self, *args, **kwargs):
671 vllm_config = args[0] if args else kwargs["vllm_config"]
@@ -648,6 +691,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
691 else:
692 # Disable indexer inner overlap.
693 self.indexer.aux_stream = None
694 + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer)
695
696 def _run_sequential_pipeline(
697 self,
@@ -1190,18 +1234,20 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1234 if attn_metadata is None:
1235 # Warmup dummy run: no real metadata. Reserve the same bf16
1236 # gather workspace _forward_prefill would; the dequantize / topk
1193 - # / sparse_fwd kernels are skipped this step.
1194 - swa_only = self.compress_ratio <= 1
1195 - N = (
1196 - 0
1197 - if swa_only
1198 - else (self.max_model_len + self.compress_ratio - 1)
1199 - // self.compress_ratio
1200 - )
1201 - M = N + self.window_size + self.max_num_batched_tokens
1202 - current_workspace_manager().get_simultaneous(
1203 - ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
1204 - )
1237 + # / sparse_fwd kernels are skipped this step. The aiter prefill
1238 + # reads the caches in place and never asks for it.
1239 + if not self._use_aiter_sparse_mla:
1240 + swa_only = self.compress_ratio <= 1
1241 + N = (
1242 + 0
1243 + if swa_only
1244 + else (self.max_model_len + self.compress_ratio - 1)
1245 + // self.compress_ratio
1246 + )
1247 + M = N + self.window_size + self.max_num_batched_tokens
1248 + current_workspace_manager().get_simultaneous(
1249 + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
1250 + )
1251 output.zero_()
1252 return
1253
@@ -1312,6 +1358,22 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1358 topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices
1359 topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr
1360
1361 + if self._use_aiter_sparse_mla:
1362 + assert swa_metadata.decode_swa_ragged_indices is not None
1363 + assert swa_metadata.decode_swa_ragged_indptr is not None
1364 + self._aiter_sparse_mla(
1365 + q,
1366 + output,
1367 + self.swa_cache_layer.kv_cache,
1368 + swa_metadata.decode_swa_ragged_indices,
1369 + swa_metadata.decode_swa_ragged_indptr,
1370 + kv_cache,
1371 + topk_ragged_indices,
1372 + topk_ragged_indptr,
1373 + )
1374 + # The aiter kernel leaves the inverse RoPE to the caller.
1375 + return 0
1376 +
1377 return rocm_sparse_attn_decode(
1378 q=q,
1379 kv_cache=kv_cache,
@@ -1351,6 +1413,17 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1413 attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
1414 swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
1415 ) -> None:
1416 + if self._use_aiter_sparse_mla:
1417 + self._forward_prefill_aiter(
1418 + q=q,
1419 + compressed_k_cache=compressed_k_cache,
1420 + swa_k_cache=swa_k_cache,
1421 + output=output,
1422 + attn_metadata=attn_metadata,
1423 + swa_metadata=swa_metadata,
1424 + )
1425 + return
1426 +
1427 swa_only = attn_metadata is None
1428
1429 num_prefills = swa_metadata.num_prefills
@@ -1476,3 +1549,90 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1549 attn_sink=self.attn_sink,
1550 output=output[query_start:query_end],
1551 )
1552 +
1553 + def _forward_prefill_aiter(
1554 + self,
1555 + q: torch.Tensor,
1556 + compressed_k_cache: torch.Tensor | None,
1557 + swa_k_cache: torch.Tensor,
1558 + output: torch.Tensor,
1559 + attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
1560 + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
1561 + ) -> None:
1562 + """Prefill straight off the paged caches, with no bf16 gather."""
1563 + num_decode_tokens = swa_metadata.num_decode_tokens
1564 + num_prefill_tokens = swa_metadata.num_prefill_tokens
1565 + assert swa_metadata.prefill_swa_ragged_indices is not None
1566 + assert swa_metadata.prefill_swa_ragged_indptr is not None
1567 +
1568 + topk_ragged_indices = None
1569 + topk_ragged_indptr = None
1570 + if attn_metadata is None:
1571 + compressed_k_cache = None
1572 + else:
1573 + assert compressed_k_cache is not None
1574 + assert swa_metadata.token_to_req_indices is not None
1575 + assert swa_metadata.query_start_loc is not None
1576 + assert swa_metadata.seq_lens is not None
1577 + assert swa_metadata.is_valid_token is not None
1578 + if self.compress_ratio == 4:
1579 + assert self.topk_indices_buffer is not None
1580 + topk_indices = self.topk_indices_buffer[
1581 + num_decode_tokens : num_decode_tokens + num_prefill_tokens
1582 + ]
1583 + else:
1584 + topk_indices = attn_metadata.c128a_prefill_topk_indices
1585 + assert topk_indices is not None
1586 + topk_ragged_indices, topk_ragged_indptr = build_prefill_topk_ragged_indices(
1587 + topk_indices,
1588 + swa_metadata.token_to_req_indices,
1589 + swa_metadata.query_start_loc,
1590 + swa_metadata.seq_lens,
1591 + swa_metadata.is_valid_token,
1592 + attn_metadata.block_table,
1593 + block_size=attn_metadata.block_size // self.compress_ratio,
1594 + compress_ratio=self.compress_ratio,
1595 + num_compressed=(self.max_model_len + self.compress_ratio - 1)
1596 + // self.compress_ratio,
1597 + token_offset=num_decode_tokens,
1598 + num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1],
1599 + )
1600 +
1601 + self._aiter_sparse_mla(
1602 + q[:num_prefill_tokens],
1603 + output[:num_prefill_tokens],
1604 + swa_k_cache,
1605 + swa_metadata.prefill_swa_ragged_indices,
1606 + swa_metadata.prefill_swa_ragged_indptr,
1607 + compressed_k_cache,
1608 + topk_ragged_indices,
1609 + topk_ragged_indptr,
1610 + )
1611 +
1612 + def _aiter_sparse_mla(
1613 + self,
1614 + q: torch.Tensor,
1615 + output: torch.Tensor,
1616 + swa_k_cache: torch.Tensor,
1617 + swa_indices: torch.Tensor,
1618 + swa_indptr: torch.Tensor,
1619 + compressed_k_cache: torch.Tensor | None,
1620 + topk_indices: torch.Tensor | None,
1621 + topk_indptr: torch.Tensor | None,
1622 + ) -> None:
1623 + from vllm._aiter_ops import rocm_aiter_ops
1624 +
1625 + # Both fp8 caches are read in place: the SWA window as the main segment,
1626 + # the top-k compressed tokens as the extra one.
1627 + rocm_aiter_ops.triton_sparse_mla_fwd(
1628 + q,
1629 + swa_k_cache,
1630 + output,
1631 + self.scale,
1632 + swa_indptr,
1633 + swa_indices,
1634 + attn_sink=self.attn_sink[: q.shape[1]],
1635 + extra_kv_buffer=compressed_k_cache,
1636 + extra_kv_indptr=topk_indptr,
1637 + extra_kv_indices=topk_indices,
1638 + )
vllm/models/deepseek_v41/amd/rocm.py
+241 -9
@@ -38,6 +38,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
38 DeepseekSparseSWAMetadata,
39 )
40 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
41 + build_prefill_topk_ragged_indices,
42 build_ragged_indices_from_dense,
43 rocm_inv_rope_einsum,
44 rocm_inverse_rope_mxfp8_rows,
@@ -64,6 +65,22 @@ def _trust_dsv4_extra_cache_nan_free(
65 )
66
67
68 +def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool:
69 + from vllm._aiter_ops import rocm_aiter_ops
70 +
71 + if not rocm_aiter_ops.is_triton_sparse_mla_enabled():
72 + return False
73 + # The aiter kernel does not scrub NaNs, so like the trust check above it
74 + # needs a cache only the canonical writer has touched.
75 + if has_kv_transfer:
76 + logger.warning_once(
77 + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV "
78 + "connector; DeepSeek V4 keeps the Triton sparse attention."
79 + )
80 + return False
81 + return True
82 +
83 +
84 def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
85 lengths = lengths.to(dtype=torch.int32).contiguous()
86 indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
@@ -408,6 +425,9 @@ _TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
425 class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
426 decode_swa_ragged_indices: torch.Tensor | None = None
427 decode_swa_ragged_indptr: torch.Tensor | None = None
428 + # Built only for the aiter prefill, which reads the SWA cache in place.
429 + prefill_swa_ragged_indices: torch.Tensor | None = None
430 + prefill_swa_ragged_indptr: torch.Tensor | None = None
431
432
433 class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
@@ -432,6 +452,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
452 dtype=torch.int32,
453 device=self.device,
454 )
455 + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(
456 + self.vllm_config.kv_transfer_config is not None
457 + )
458
459 def build(
460 self,
@@ -447,6 +470,22 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
470 replay_start=replay_start,
471 )
472
473 + prefill_ragged_indices = None
474 + prefill_ragged_indptr = None
475 + if (
476 + self._use_aiter_sparse_mla
477 + and base.num_prefill_tokens > 0
478 + and base.prefill_swa_indices is not None
479 + and base.prefill_swa_lens is not None
480 + ):
481 + # Paged-coordinate rows, so every layer's prefill shares one build.
482 + prefill_ragged_indices, prefill_ragged_indptr = (
483 + build_ragged_indices_from_dense(
484 + base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1),
485 + base.prefill_swa_lens,
486 + )
487 + )
488 +
489 ragged_indices = None
490 ragged_indptr = None
491 if (
@@ -473,6 +512,8 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
512 **vars(base),
513 decode_swa_ragged_indices=ragged_indices,
514 decode_swa_ragged_indptr=ragged_indptr,
515 + prefill_swa_ragged_indices=prefill_ragged_indices,
516 + prefill_swa_ragged_indptr=prefill_ragged_indptr,
517 )
518
519
@@ -497,6 +538,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
538
539 backend_cls = DeepseekV4ROCMAiterMLASparseBackend
540 swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend
541 + _use_aiter_sparse_mla = False
542
543 def __init__(self, *args, **kwargs):
544 vllm_config = args[0] if args else kwargs["vllm_config"]
@@ -518,6 +560,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
560 # same thing. The source owns one cache per compress ratio, refreshed
561 # when it runs; consumers below it read through.
562 self._topk_ragged_cache: dict[int, _TopkRagged] = {}
563 + self._prefill_topk_ragged_cache: dict[
564 + str, tuple[torch.Tensor, torch.Tensor]
565 + ] = {}
566 self._index_source_prefix: str | None = None
567 if self.compress_ratio > 0:
568 assert self.index_source_layer_id is not None
@@ -530,6 +575,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
575 "this rank; PP splits inside a v4.1 index-sharing group "
576 "are not supported."
577 )
578 + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer)
579 if self.compressor is None and self.indexer is None:
580 # Dense layers have nothing to overlap; keep the base serial path.
581 self.aux_stream_list = None
@@ -915,9 +961,10 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
961 attn_metadata = forward_context.attn_metadata
962
963 if attn_metadata is None:
918 - # Warmup dummy run: no real metadata. Reserve the same bf16
919 - # gather workspace _forward_prefill would; the dequantize / topk
920 - # / sparse_fwd kernels are skipped this step.
964 + # Warmup dummy run: no real metadata. Reserve the workspace
965 + # _forward_prefill would; the dequantize / topk / sparse_fwd kernels
966 + # are skipped this step. The aiter prefill reads the caches in place,
967 + # so it asks only for the bf16 rows of the MXFP8 output.
968 swa_only = self.compress_ratio == 0
969 N = (
970 0
@@ -926,9 +973,14 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
973 // self.compress_ratio
974 )
975 M = N + self.window_size + self.max_num_batched_tokens
929 - current_workspace_manager().get_simultaneous(
930 - *self._prefill_workspace_shapes(M, self.max_num_batched_tokens, q)
976 + shapes = self._prefill_workspace_shapes(
977 + M,
978 + self.max_num_batched_tokens,
979 + q,
980 + gather=not self._use_aiter_sparse_mla,
981 )
982 + if shapes:
983 + current_workspace_manager().get_simultaneous(*shapes)
984 if mxfp8_out is None:
985 output.zero_()
986 else:
@@ -1087,6 +1139,38 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1139 num_decode_tokens=num_decode_tokens,
1140 )
1141
1142 + if self._use_aiter_sparse_mla:
1143 + assert swa_metadata.decode_swa_ragged_indices is not None
1144 + assert swa_metadata.decode_swa_ragged_indptr is not None
1145 + out = output
1146 + if output_mxfp8 is not None:
1147 + # The aiter kernel has no MXFP8 epilogue: write bf16 rows, then
1148 + # rotate and quantize them as the decode reduce would have.
1149 + out = self._mxfp8_bf16_rows(q.shape[0], q)
1150 + assert out is not None
1151 + self._aiter_sparse_mla(
1152 + q,
1153 + out,
1154 + self.swa_cache_layer.kv_cache,
1155 + swa_metadata.decode_swa_ragged_indices,
1156 + swa_metadata.decode_swa_ragged_indptr,
1157 + kv_cache,
1158 + topk_ragged_indices,
1159 + topk_ragged_indptr,
1160 + )
1161 + if output_mxfp8 is not None:
1162 + rocm_inverse_rope_mxfp8_rows(
1163 + out,
1164 + positions,
1165 + self.rotary_emb.cos_sin_cache,
1166 + self.rope_head_dim,
1167 + output_mxfp8[0],
1168 + output_mxfp8[1],
1169 + )
1170 + return q.shape[0]
1171 + # The aiter kernel leaves the inverse RoPE to the caller.
1172 + return 0
1173 +
1174 return rocm_sparse_attn_decode(
1175 q=q,
1176 kv_cache=kv_cache,
@@ -1132,6 +1216,19 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1216 With ``mxfp8_out`` the rows go to a bf16 workspace first and are then
1217 inverse-RoPE'd and quantized into it in one pass.
1218 """
1219 + if self._use_aiter_sparse_mla:
1220 + self._forward_prefill_aiter(
1221 + q=q,
1222 + positions=positions,
1223 + compressed_k_cache=compressed_k_cache,
1224 + swa_k_cache=swa_k_cache,
1225 + output=output,
1226 + attn_metadata=attn_metadata,
1227 + swa_metadata=swa_metadata,
1228 + mxfp8_out=mxfp8_out,
1229 + )
1230 + return
1231 +
1232 swa_only = attn_metadata is None
1233
1234 num_prefills = swa_metadata.num_prefills
@@ -1249,15 +1346,17 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1346 )
1347
1348 def _prefill_workspace_shapes(
1252 - self, M: int, num_prefill_tokens: int, q: torch.Tensor
1349 + self, M: int, num_prefill_tokens: int, q: torch.Tensor, gather: bool = True
1350 ) -> list[tuple[tuple[int, ...], torch.dtype]]:
1351 """The prefill gather buffer, plus bf16 rows for the MXFP8 output.
1352
1353 One ``get_simultaneous`` call: separate calls alias the same memory.
1354 + ``gather=False`` leaves out the gather buffer, which the aiter kernel,
1355 + reading the caches in place, never needs.
1356 """
1258 - shapes: list[tuple[tuple[int, ...], torch.dtype]] = [
1259 - ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)
1260 - ]
1357 + shapes: list[tuple[tuple[int, ...], torch.dtype]] = []
1358 + if gather:
1359 + shapes.append(((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16))
1360 if _ON_GFX950:
1361 shapes.append(
1362 (
@@ -1266,3 +1365,136 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1365 )
1366 )
1367 return shapes
1368 +
1369 + def _prefill_topk_ragged(
1370 + self,
1371 + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
1372 + attn_metadata: DeepseekV4FlashMLAMetadata,
1373 + compressed_k_cache: torch.Tensor,
1374 + ) -> tuple[torch.Tensor, torch.Tensor]:
1375 + """Prefill counterpart of ``_decode_topk_ragged``: layers that share an
1376 + index source and a compressed cache reuse one build per step."""
1377 + assert swa_metadata.token_to_req_indices is not None
1378 + assert swa_metadata.query_start_loc is not None
1379 + assert swa_metadata.seq_lens is not None
1380 + assert swa_metadata.is_valid_token is not None
1381 + assert self.topk_indices_buffer is not None
1382 + assert self._index_source_prefix is not None
1383 + assert self.compressed_cache_prefix is not None
1384 +
1385 + source = self._static_forward_context[self._index_source_prefix]
1386 + if source is self:
1387 + source._prefill_topk_ragged_cache = {}
1388 + cached = source._prefill_topk_ragged_cache.get(self.compressed_cache_prefix)
1389 + if cached is not None:
1390 + return cached
1391 +
1392 + num_decode_tokens = swa_metadata.num_decode_tokens
1393 + num_prefill_tokens = swa_metadata.num_prefill_tokens
1394 + built = build_prefill_topk_ragged_indices(
1395 + self.topk_indices_buffer[
1396 + num_decode_tokens : num_decode_tokens + num_prefill_tokens
1397 + ],
1398 + swa_metadata.token_to_req_indices,
1399 + swa_metadata.query_start_loc,
1400 + swa_metadata.seq_lens,
1401 + swa_metadata.is_valid_token,
1402 + attn_metadata.block_table,
1403 + block_size=attn_metadata.block_size // self.compress_ratio,
1404 + compress_ratio=self.compress_ratio,
1405 + num_compressed=(self.max_model_len + self.compress_ratio - 1)
1406 + // self.compress_ratio,
1407 + token_offset=num_decode_tokens,
1408 + num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1],
1409 + )
1410 + source._prefill_topk_ragged_cache[self.compressed_cache_prefix] = built
1411 + return built
1412 +
1413 + def _forward_prefill_aiter(
1414 + self,
1415 + q: torch.Tensor,
1416 + positions: torch.Tensor,
1417 + compressed_k_cache: torch.Tensor | None,
1418 + swa_k_cache: torch.Tensor,
1419 + output: torch.Tensor | None,
1420 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
1421 + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
1422 + mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
1423 + ) -> None:
1424 + """Prefill straight off the paged caches, with no bf16 gather.
1425 +
1426 + With ``mxfp8_out`` the kernel writes bf16 rows to the workspace, which
1427 + are then inverse-RoPE'd and quantized into it as ``_forward_prefill``
1428 + does.
1429 + """
1430 + num_prefill_tokens = swa_metadata.num_prefill_tokens
1431 + assert swa_metadata.prefill_swa_ragged_indices is not None
1432 + assert swa_metadata.prefill_swa_ragged_indptr is not None
1433 +
1434 + topk_indices = None
1435 + topk_indptr = None
1436 + if attn_metadata is None:
1437 + compressed_k_cache = None
1438 + else:
1439 + assert compressed_k_cache is not None
1440 + topk_indices, topk_indptr = self._prefill_topk_ragged(
1441 + swa_metadata, attn_metadata, compressed_k_cache
1442 + )
1443 +
1444 + if mxfp8_out is not None:
1445 + output = self._mxfp8_bf16_rows(num_prefill_tokens, q)
1446 + assert output is not None
1447 + self._aiter_sparse_mla(
1448 + q[:num_prefill_tokens],
1449 + output[:num_prefill_tokens],
1450 + swa_k_cache,
1451 + swa_metadata.prefill_swa_ragged_indices,
1452 + swa_metadata.prefill_swa_ragged_indptr,
1453 + compressed_k_cache,
1454 + topk_indices,
1455 + topk_indptr,
1456 + )
1457 + if mxfp8_out is not None:
1458 + rocm_inverse_rope_mxfp8_rows(
1459 + output[:num_prefill_tokens],
1460 + positions[:num_prefill_tokens],
1461 + self.rotary_emb.cos_sin_cache,
1462 + self.rope_head_dim,
1463 + mxfp8_out[0][:num_prefill_tokens],
1464 + mxfp8_out[1][:num_prefill_tokens],
1465 + )
1466 +
1467 + def _mxfp8_bf16_rows(self, num_tokens: int, q: torch.Tensor) -> torch.Tensor:
1468 + """bf16 attention rows for the aiter kernel to write when the layer's
1469 + output is MXFP8, taken from the workspace the warmup reserved."""
1470 + shapes = self._prefill_workspace_shapes(0, num_tokens, q, gather=False)
1471 + assert shapes, "an MXFP8 attention output is gfx950-only"
1472 + return current_workspace_manager().get_simultaneous(*shapes)[0]
1473 +
1474 + def _aiter_sparse_mla(
1475 + self,
1476 + q: torch.Tensor,
1477 + output: torch.Tensor,
1478 + swa_k_cache: torch.Tensor,
1479 + swa_indices: torch.Tensor,
1480 + swa_indptr: torch.Tensor,
1481 + compressed_k_cache: torch.Tensor | None,
1482 + topk_indices: torch.Tensor | None,
1483 + topk_indptr: torch.Tensor | None,
1484 + ) -> None:
1485 + from vllm._aiter_ops import rocm_aiter_ops
1486 +
1487 + # Both fp8 caches are read in place: the SWA window as the main segment,
1488 + # the top-k compressed tokens as the extra one.
1489 + rocm_aiter_ops.triton_sparse_mla_fwd(
1490 + q,
1491 + swa_k_cache,
1492 + output,
1493 + self.scale,
1494 + swa_indptr,
1495 + swa_indices,
1496 + attn_sink=self.attn_sink[: q.shape[1]],
1497 + extra_kv_buffer=compressed_k_cache,
1498 + extra_kv_indptr=topk_indptr,
1499 + extra_kv_indices=topk_indices,
1500 + )
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+72 -1
@@ -396,9 +396,11 @@ class ROCMAiterMLASparseMetadataBuilder(
396 attention_context = vllm_config.compilation_config.static_forward_context
397 # Sink decode must use AITER's nonpersistent path. In particular,
398 # gfx942 has no persistent+LSE kernel, and its metadata heuristic
399 - # terminates for HY-V4's TP1 H64 shape.
399 + # terminates for HY-V4's TP1 H64 shape. The Triton sparse MLA kernel
400 + # never reads the persistent metadata either.
401 self._use_persistent_metadata = all(
402 getattr(attention_context[name].impl, "sinks", None) is None
403 + and not getattr(attention_context[name].impl, "use_aiter_sparse_mla", False)
404 for name in layer_names
405 )
406 # Bounds the KV-split heuristic (see `_sparse_decode_max_split`).
@@ -716,6 +718,7 @@ class ROCMAiterMLASparseImpl(
718 is_sparse = True
719 supports_dense_mha_prefill = False
720 supports_dcp = False
721 + use_aiter_sparse_mla = False
722
723 def __init__(
724 self,
@@ -765,6 +768,32 @@ class ROCMAiterMLASparseImpl(
768 (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
769 (q_concat_shape, vllm_config.model_config.dtype),
770 )
771 + self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
772 +
773 + if rocm_aiter_ops.is_triton_sparse_mla_enabled():
774 + reason = self._aiter_sparse_mla_unsupported_reason(vllm_config)
775 + if reason is None:
776 + self.use_aiter_sparse_mla = True
777 + else:
778 + logger.warning_once(
779 + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; "
780 + "using the default sparse MLA path instead.",
781 + reason,
782 + )
783 +
784 + def _aiter_sparse_mla_unsupported_reason(
785 + self, vllm_config: VllmConfig
786 + ) -> str | None:
787 + model_dtype = vllm_config.model_config.dtype
788 + if model_dtype != torch.bfloat16:
789 + return f"the model dtype is {model_dtype}, not bfloat16"
790 + cache_dtype = self.kv_cache_dtype
791 + if cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"):
792 + return f"kv-cache dtype {cache_dtype!r} is neither bfloat16 nor e4m3 fp8"
793 + # Context parallelism merges each rank's partials with a per-token LSE.
794 + if self.dcp_world_size > 1 or self.pcp_world_size > 1:
795 + return "context parallelism is not supported"
796 + return None
797
798 def record_logical_topk_ready(self) -> None:
799 # This impl shares the top-k indices buffer via SharedTopkIndicesBuffer
@@ -975,6 +1004,43 @@ class ROCMAiterMLASparseImpl(
1004
1005 return output, lse
1006
1007 + def _forward_mla_aiter(
1008 + self,
1009 + layer: AttentionLayer,
1010 + q: torch.Tensor, # [sq, heads, d_qk], not head-padded
1011 + kv_c_and_k_pe_cache: torch.Tensor,
1012 + attn_metadata: ROCMAiterMLASparseMetadata,
1013 + ) -> torch.Tensor:
1014 + """_forward_mla on aiter's Triton sparse MLA kernel.
1015 +
1016 + Reads the same index stream, but needs no q head padding and no
1017 + persistent metadata, and its launch depends on shapes only, so it is
1018 + CUDA-graph capturable.
1019 + """
1020 + num_actual_tokens = attn_metadata.num_actual_tokens
1021 + output = torch.empty(
1022 + [q.shape[0], self.num_heads, self.kv_lora_rank],
1023 + dtype=attn_metadata.attn_out_dtype,
1024 + device=q.device,
1025 + )
1026 + rocm_aiter_ops.triton_sparse_mla_fwd(
1027 + q[:num_actual_tokens],
1028 + kv_c_and_k_pe_cache.view(-1, 1, 1, q.shape[-1]),
1029 + output[:num_actual_tokens],
1030 + self.scale,
1031 + attn_metadata.paged_kv_indptr,
1032 + attn_metadata.paged_kv_indices,
1033 + kv_lora_rank=self.kv_lora_rank,
1034 + qk_rope_head_dim=self.qk_rope_head_dim,
1035 + q_scale=layer._q_scale,
1036 + kv_scale=layer._k_scale,
1037 + attn_sink=self.sinks,
1038 + # triton_convert_req_index_to_global_index writes 0, never -1,
1039 + # for an invalid top-k entry, so no slot in the stream is negative.
1040 + has_invalid=False,
1041 + )
1042 + return output
1043 +
1044 def forward_mqa(
1045 self,
1046 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -1026,6 +1092,11 @@ class ROCMAiterMLASparseImpl(
1092 original_q_shape = q.shape
1093 q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale)
1094 q = q.view(original_q_shape)
1095 + if self.use_aiter_sparse_mla:
1096 + output = self._forward_mla_aiter(
1097 + layer, q, kv_c_and_k_pe_cache, attn_metadata
1098 + )
1099 + return output, None
1100 mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
1101 return self._forward_mla(
1102 layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+96
@@ -4306,3 +4306,99 @@ def rocm_sparse_attn_decode(
4306 if direct_out is None:
4307 output.copy_(attn_out.to(output.dtype))
4308 return output.shape[0] if inv_rope_positions is not None else 0
4309 +
4310 +
4311 +@triton.jit(do_not_specialize=["token_offset"])
4312 +def _prefill_topk_global_slots_kernel(
4313 + out_ptr,
4314 + lens_ptr,
4315 + topk_ptr,
4316 + topk_stride,
4317 + token_to_req_ptr,
4318 + query_start_loc_ptr,
4319 + seq_lens_ptr,
4320 + is_valid_token_ptr,
4321 + block_table_ptr,
4322 + block_table_stride,
4323 + token_offset,
4324 + num_compressed,
4325 + TOPK: tl.constexpr,
4326 + COMPRESS_RATIO: tl.constexpr,
4327 + BLOCK_SIZE: tl.constexpr,
4328 + BLOCK_W: tl.constexpr,
4329 +):
4330 + row = tl.program_id(0)
4331 + token = row + token_offset
4332 + req = tl.load(token_to_req_ptr + token)
4333 + query_start = tl.load(query_start_loc_ptr + req)
4334 + query_len = tl.load(query_start_loc_ptr + req + 1) - query_start
4335 + pos = tl.load(seq_lens_ptr + req) - query_len + token - query_start
4336 + # Only (pos + 1) // ratio compressed tokens exist yet; the rest of the row
4337 + # can hold stale indices.
4338 + row_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOPK)
4339 + row_len = tl.where(tl.load(is_valid_token_ptr + token), row_len, 0)
4340 +
4341 + for start in range(0, TOPK, BLOCK_W):
4342 + cols = start + tl.arange(0, BLOCK_W)
4343 + live = cols < row_len
4344 + local = tl.load(topk_ptr + row * topk_stride + cols, mask=live, other=-1)
4345 + valid = live & (local >= 0) & (local < num_compressed)
4346 + local = tl.where(valid, local, 0)
4347 + block = tl.load(
4348 + block_table_ptr + req * block_table_stride + local // BLOCK_SIZE,
4349 + mask=valid,
4350 + other=0,
4351 + )
4352 + slot = tl.where(valid, block * BLOCK_SIZE + local % BLOCK_SIZE, -1)
4353 + tl.store(out_ptr + row * TOPK + cols, slot, mask=cols < TOPK)
4354 + tl.store(lens_ptr + row, row_len)
4355 +
4356 +
4357 +def build_prefill_topk_ragged_indices(
4358 + topk_indices: torch.Tensor,
4359 + token_to_req_indices: torch.Tensor,
4360 + query_start_loc: torch.Tensor,
4361 + seq_lens: torch.Tensor,
4362 + is_valid_token: torch.Tensor,
4363 + block_table: torch.Tensor,
4364 + block_size: int,
4365 + compress_ratio: int,
4366 + num_compressed: int,
4367 + token_offset: int,
4368 + num_rows: int = -1,
4369 +) -> tuple[torch.Tensor, torch.Tensor]:
4370 + """Map prefill top-k rows to a ragged stream of compressed-cache slots.
4371 +
4372 + ``topk_indices`` holds local compressed positions for the prefill tokens,
4373 + which sit at ``token_offset`` in the batch; ``token_to_req_indices``,
4374 + ``query_start_loc``, ``seq_lens`` and ``block_table`` are batch-wide.
4375 + ``block_size`` is the compressed cache's, i.e. already divided by the ratio.
4376 + """
4377 + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1)
4378 + num_tokens, width = topk_indices.shape
4379 + dense = torch.empty(
4380 + (num_tokens, width), dtype=torch.int32, device=topk_indices.device
4381 + )
4382 + lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
4383 + if num_tokens > 0 and width > 0:
4384 + _prefill_topk_global_slots_kernel[(num_tokens,)](
4385 + dense,
4386 + lens,
4387 + topk_indices,
4388 + topk_indices.stride(0),
4389 + token_to_req_indices,
4390 + query_start_loc,
4391 + seq_lens,
4392 + is_valid_token,
4393 + block_table,
4394 + block_table.stride(0),
4395 + token_offset,
4396 + num_compressed,
4397 + TOPK=width,
4398 + COMPRESS_RATIO=compress_ratio,
4399 + BLOCK_SIZE=block_size,
4400 + BLOCK_W=min(triton.next_power_of_2(width), 1024),
4401 + )
4402 + else:
4403 + lens.zero_()
4404 + return build_ragged_indices_from_dense(dense, lens, num_rows=num_rows)