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

[ROCm][MLA][DCP] Support causal multi-token verification (#51705)

Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Signed-off-by: billishyahao <bill.he@amd.com> Co-authored-by: andyluo7 <andy.luo@amd.com> Co-authored-by: billishyahao <bill.he@amd.com>

YukioZzz committed Sep 1, 2026 at 01:42 UTC dbb7fffddbca2815d5d2339c52ca04a6715f2139
7 files changed +1105 -111
tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
+20 -5
@@ -58,6 +58,13 @@ PAGE_SIZE = 1
58 MAX_MODEL_LEN = 1024
59 MAX_NUM_SEQS = 8
60
61 +# The builder resolves the verify routing and records it on the metadata, so
62 +# the gfx950 feature probe has to be forced for the build as well as for the
63 +# forward; ROCm CI also runs gfx942, where the real probe is False.
64 +_GLUON_SUPPORTED_TARGET = (
65 + "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported"
66 +)
67 +
68
69 def _seq_lens() -> list[int]:
70 return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS
@@ -99,6 +106,10 @@ def _run_verify_block():
106 vllm_config.speculative_config = SpeculativeConfig(
107 method="ngram", num_speculative_tokens=QLEN - 1
108 )
109 + vllm_config.model_config.get_num_attention_heads = types.MethodType(
110 + lambda self, parallel_config, arch_config=None: NUM_QUERY_HEADS,
111 + vllm_config.model_config,
112 + )
113
114 spec = MLAAttentionSpec(
115 block_size=PAGE_SIZE,
@@ -138,7 +149,10 @@ def _run_verify_block():
149 captured["page_table"] = kwargs["page_table"].detach().clone()
150 captured["seq_info"] = kwargs["seq_info"].detach().clone()
151
141 - with set_current_vllm_config(vllm_config):
152 + with (
153 + patch(_GLUON_SUPPORTED_TARGET, lambda: True),
154 + set_current_vllm_config(vllm_config),
155 + ):
156 builder = builder_cls(spec, [layer_name], vllm_config, device)
157 common_attn_metadata = create_common_attn_metadata(
158 batch_spec, PAGE_SIZE, device, arange_block_indices=True
@@ -182,10 +196,7 @@ def _run_verify_block():
196 # The Gluon kernel only reads the metadata this test is about, so a spy in
197 # its place keeps the assertions independent of the AITER build.
198 with (
185 - patch(
186 - "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported",
187 - lambda: True,
188 - ),
199 + patch(_GLUON_SUPPORTED_TARGET, lambda: True),
200 patch(
201 "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon",
202 lambda: spy,
@@ -211,6 +222,10 @@ def test_verify_mtp_uses_native_4d_gluon_entry():
222 f"expected a {QLEN}-token verify block, got max_qo_len="
223 f"{decode.max_qo_len}; the MTP path under test was not reached"
224 )
225 + assert decode.use_gluon_verify, (
226 + "the builder owns verify routing, so the metadata -- not the impl -- "
227 + "has to select the Gluon MTP entry for a small-head bf16 verify block"
228 + )
229 assert captured, "forward_mqa did not reach the Gluon kernel"
230
231 seq_lens = _seq_lens()
tests/kernels/attention/test_rocm_aiter_mla_head_padding.py
+2
@@ -324,7 +324,9 @@ def test_h12_aiter_mla_decode_matches_reference():
324
325 impl = object.__new__(AiterMLAImpl)
326 impl.num_heads = NUM_HEADS
327 + impl.dcp_world_size = 1
328 impl.kv_lora_rank = KV_LORA_RANK
329 + impl.qk_rope_head_dim = QK_ROPE_HEAD_DIM
330 impl.scale = SCALE
331 impl.kv_cache_dtype = "auto"
332
tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
+27 -3
@@ -8,6 +8,10 @@ batch_size == 1, so it cannot serve a decode batch. Every fp8 shape therefore
8 has to land on the asm kernels, which ship real fp8 variants for gqa=16. These
9 tests pin that down at the predicates, since the failure it prevents is either a
10 batch assertion or -- worse -- a silently wrong result.
11 +
12 +Gluon is still the route for supported non-DCP decode shapes. DCP multi-token
13 +verification uses segmented MLA instead, so it depends on neither Gluon's
14 +query-head nor its long-context pipeline limits.
15 """
16
17 import pytest
@@ -41,8 +45,8 @@ def gluon_available(monkeypatch):
45 @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES)
46 @pytest.mark.parametrize("num_heads", [1, 2, 4, 5, 6, 8, 12, 16, 32, 128])
47 @pytest.mark.parametrize("max_qo_len", [1, 2, 4, 5, 8, 15])
44 -def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
45 - """No fp8 shape reaches either Gluon entry point.
48 +def test_non_dcp_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
49 + """No non-DCP fp8 shape reaches either Gluon entry point.
50
51 The head count is deliberately swept across divisors of 16 as well as
52 non-divisors: the divisor case (e.g. 8 heads at TP8) is the one that stays
@@ -58,7 +62,7 @@ def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
62 def test_fp8_never_routes_to_gluon_under_any_mode(
63 monkeypatch, kv_cache_dtype, num_heads, mode
64 ):
61 - """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force an fp8 cache onto Gluon.
65 + """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force a non-DCP fp8 cache onto Gluon.
66
67 The dtype guard deliberately precedes the mode knob: honouring an explicit
68 "gluon" request under fp8 would hand Gluon the batch it asserts against, so
@@ -80,6 +84,26 @@ def test_large_head_counts_never_use_gluon(kv_cache_dtype, num_heads, max_qo_len
84 assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
85
86
87 +@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES)
88 +@pytest.mark.parametrize("num_heads", [8, 16, 32, 96])
89 +@pytest.mark.parametrize("max_qo_len", [2, 3, 8])
90 +def test_dcp_multitoken_verify_never_uses_gluon(
91 + gluon_available, kv_cache_dtype, num_heads, max_qo_len
92 +):
93 + assert not AiterMLAHelper.use_gluon_verify(
94 + num_heads, max_qo_len, kv_cache_dtype, dcp_world_size=8
95 + )
96 +
97 +
98 +def test_segmented_dcp_verify_does_not_depend_on_gluon(monkeypatch):
99 + monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False)
100 + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
101 +
102 + assert rocm_aiter_mla._segmented_dcp_verify_supported(8, 1)
103 + # Round-robin interleaving other than 1 is not served by this route.
104 + assert not rocm_aiter_mla._segmented_dcp_verify_supported(8, 4)
105 +
106 +
107 @pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
108 @pytest.mark.parametrize("num_heads", [1, 2, 4, 8])
109 def test_unquantized_divisor_heads_keep_gluon_decode(
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+380 -5
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import math
5 import sys
6 from types import SimpleNamespace
7 from unittest import mock
@@ -16,8 +17,12 @@ if not current_platform.is_rocm():
17 from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402
18 from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402
19 AiterMLAHelper,
20 + AiterMLAImpl,
21 AiterMLAMetadataBuilder,
22 )
23 +from vllm.v1.attention.ops.rocm_aiter_mla_merge import ( # noqa: E402
24 + merge_mla_segments_triton,
25 +)
26
27
28 class _NoOpTritonKernel:
@@ -29,6 +34,19 @@ class _NoOpTritonKernel:
34 pass
35
36
37 +def _lse_combine_natural(
38 + out_a: torch.Tensor,
39 + lse_a: torch.Tensor,
40 + out_b: torch.Tensor,
41 + lse_b: torch.Tensor,
42 + out_dtype: torch.dtype,
43 +) -> tuple[torch.Tensor, torch.Tensor]:
44 + normalizer = torch.logaddexp(lse_a, lse_b)
45 + out = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1)
46 + out += out_b * torch.exp(lse_b - normalizer).unsqueeze(-1)
47 + return out.to(out_dtype), normalizer
48 +
49 +
50 class _ExpandPageIndicesKernel:
51 def __getitem__(self, grid):
52 self.grid = grid
@@ -67,10 +85,27 @@ def _builder(
85 max_decode_rows: int = 32,
86 num_heads: int = 16,
87 kv_cache_dtype: str = "auto",
88 + dcp_world_size: int = 1,
89 + dcp_rank: int = 0,
90 ):
71 - return SimpleNamespace(
91 + stub = SimpleNamespace(
92 device=torch.device("cpu"),
93 num_heads=num_heads,
94 + # DCP gathers every rank's head shard before decode, and the routing
95 + # predicates read the gathered count.
96 + _decode_num_heads=num_heads * dcp_world_size,
97 + dcp_world_size=dcp_world_size,
98 + dcp_rank=dcp_rank,
99 + cp_kv_cache_interleave_size=1,
100 + # Mirrors the production constructor: configuration decides whether the
101 + # segmented route is available, query length decides per batch.
102 + _supports_segmented_dcp_verify=rocm_aiter_mla._segmented_dcp_verify_supported(
103 + dcp_world_size, 1
104 + ),
105 + # Derived once in the real constructor, so derive it once here too.
106 + _segmented_page_size=rocm_aiter_mla._segmented_mla_page_size(kernel_block_size),
107 + _dcp_verify_buffers=None,
108 + _graph_seq_lens=None,
109 _kv_cache_dtype_str=kv_cache_dtype,
110 paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32),
111 paged_kv_indices=torch.empty(1024, dtype=torch.int32),
@@ -96,6 +131,15 @@ def _builder(
131 _mla_kv_dtype=torch.bfloat16,
132 decode_attn_out_dtype=torch.bfloat16,
133 )
134 + # Bound to the stub rather than faked: the verify flatten's per-row view is
135 + # part of what _build_decode is being tested for.
136 + stub._build_dcp_verify_row_view = (
137 + AiterMLAMetadataBuilder._build_dcp_verify_row_view.__get__(stub)
138 + )
139 + stub._fill_dcp_verify_page_table = (
140 + AiterMLAMetadataBuilder._fill_dcp_verify_page_table.__get__(stub)
141 + )
142 + return stub
143
144
145 def test_backend_declares_uniform_batch_support():
@@ -111,6 +155,324 @@ def test_backend_declares_uniform_batch_support():
155 )
156
157
158 +def test_dcp_verify_row_view_is_causal_per_row():
159 + """Row t must cover this rank's share of positions [0, seq_len - qlen + t].
160 +
161 + Regression guard: giving every row of a request the committed prefix only
162 + drops the whole verify block, including each row's own token, since nothing
163 + else adds it back on the target path.
164 + """
165 + qlen = 4
166 + builder = _builder(mtp_decode_qlen=qlen, dcp_world_size=2, kernel_block_size=2)
167 + block_table = torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32)
168 +
169 + view = builder._build_dcp_verify_row_view(
170 + qlen,
171 + block_table,
172 + torch.tensor([10, 12], dtype=torch.int32),
173 + )
174 +
175 + # seq_len 10 -> rows see 7, 8, 9, 10 global positions, of which rank 0 of 2
176 + # holds ceil(n / 2); seq_len 12 -> 9, 10, 11, 12.
177 + assert view.row_lens.tolist() == [4, 4, 5, 5, 5, 5, 6, 6]
178 + assert view.qo_indptr.tolist() == list(range(2 * qlen + 1))
179 + # One query per row, so the page table is the request's ascending shard and
180 + # only the row length selects the causal prefix.
181 + assert view.block_table.tolist() == [[7, 8, 9]] * qlen + [[10, 11, 12]] * qlen
182 + assert view.max_kv_seq_len == 6
183 +
184 +
185 +def test_dcp_verify_rows_sum_to_the_global_window_across_ranks():
186 + """Every rank's rows together must cover each row's window exactly once."""
187 + qlen = 4
188 + dcp_world_size = 2
189 + seq_lens = torch.tensor([10, 12], dtype=torch.int32)
190 + per_rank = []
191 + for dcp_rank in range(dcp_world_size):
192 + builder = _builder(
193 + mtp_decode_qlen=qlen,
194 + dcp_world_size=dcp_world_size,
195 + dcp_rank=dcp_rank,
196 + kernel_block_size=2,
197 + )
198 + view = builder._build_dcp_verify_row_view(
199 + qlen,
200 + torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32),
201 + seq_lens,
202 + )
203 + per_rank.append(view.row_lens.tolist())
204 +
205 + expected = [
206 + int(seq_len) - qlen + t + 1
207 + for seq_len in seq_lens.tolist()
208 + for t in range(qlen)
209 + ]
210 + assert [sum(lens) for lens in zip(*per_rank)] == expected
211 +
212 +
213 +def test_dcp_verify_row_view_uses_static_graph_bound():
214 + """Under full graphs the bound comes from the buffers, not the batch."""
215 + qlen = 3
216 + builder = _builder(
217 + mtp_decode_qlen=qlen,
218 + dcp_world_size=8,
219 + has_full_cudagraphs=True,
220 + kernel_block_size=1536,
221 + )
222 + builder._dcp_verify_buffers = rocm_aiter_mla.AiterMLADCPVerifyMetadata(
223 + row_lens=torch.zeros(qlen, dtype=torch.int32),
224 + block_table=torch.zeros((qlen, 24), dtype=torch.int32),
225 + qo_indptr=torch.zeros(qlen + 1, dtype=torch.int32),
226 + page_size=128,
227 + max_kv_seq_len=3072,
228 + )
229 +
230 + view = builder._build_dcp_verify_row_view(
231 + qlen,
232 + torch.arange(2, dtype=torch.int32).view(1, 2),
233 + torch.tensor([1000], dtype=torch.int32),
234 + )
235 +
236 + buffers = builder._dcp_verify_buffers
237 + assert view.row_lens.tolist() == [125] * qlen
238 + assert view.block_table.data_ptr() == buffers.block_table.data_ptr()
239 + assert view.qo_indptr.data_ptr() == buffers.qo_indptr.data_ptr()
240 + assert view.block_table[0].tolist() == list(range(24))
241 + assert view.max_kv_seq_len == 3072
242 +
243 +
244 +def test_dcp_fp8_verify_build_uses_segmented(monkeypatch):
245 + qlen = 4
246 + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
247 +
248 + metadata = AiterMLAMetadataBuilder._build_decode(
249 + _builder(
250 + mtp_decode_qlen=qlen,
251 + dcp_world_size=2,
252 + kv_cache_dtype="fp8",
253 + kernel_block_size=2,
254 + ),
255 + block_table_tensor=torch.tensor([[0, 1, 2], [10, 11, 12]], dtype=torch.int32),
256 + seq_lens_device=torch.tensor([5, 6], dtype=torch.int32),
257 + max_seq_len=6,
258 + query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
259 + query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
260 + num_decode_tokens=2 * qlen,
261 + dcp_tot_seq_lens_device=torch.tensor([10, 12], dtype=torch.int32),
262 + )
263 +
264 + assert metadata.dcp_verify is not None
265 +
266 +
267 +def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch):
268 + """Single-token DCP decode must come back with an LSE the merge can use.
269 +
270 + This is the path every DCP step takes outside verification. The vLLM custom
271 + op drops aiter's LSE, so the impl calls aiter directly; without an LSE the
272 + MLA layer's cross-rank combine asserts. The head count is deliberately a
273 + non-divisor of 16 so both the output and the LSE have to be unpadded back
274 + to the gathered head count.
275 + """
276 + num_heads, dcp_world_size = 6, 2
277 + decode_heads = num_heads * dcp_world_size # 12, padded to 16 by aiter
278 + num_tokens, head_dim = 3, 576
279 + captured = {}
280 +
281 + def fake_aiter_decode(q, kv_buffer, out, *args, **kwargs):
282 + captured["q_heads"] = q.shape[1]
283 + captured["return_lse"] = kwargs.get("return_lse")
284 + return None, torch.zeros(num_tokens, q.shape[1])
285 +
286 + monkeypatch.setattr(
287 + rocm_aiter_mla, "_get_aiter_mla_decode", lambda: fake_aiter_decode
288 + )
289 +
290 + impl = object.__new__(AiterMLAImpl)
291 + impl.num_heads = num_heads
292 + impl.dcp_world_size = dcp_world_size
293 + impl.kv_cache_dtype = "auto"
294 + impl.kv_lora_rank = 512
295 + impl.qk_rope_head_dim = 64
296 + impl.scale = head_dim**-0.5
297 + decode = SimpleNamespace(
298 + max_qo_len=1,
299 + qo_indptr=torch.arange(num_tokens + 1, dtype=torch.int32),
300 + paged_kv_indptr=torch.zeros(num_tokens + 1, dtype=torch.int32),
301 + paged_kv_indices=torch.zeros(1, dtype=torch.int32),
302 + paged_kv_last_page_len=torch.ones(num_tokens, dtype=torch.int32),
303 + use_gluon_decode=False,
304 + use_gluon_verify=False,
305 + dcp_verify=None,
306 + has_persistent_metadata=False,
307 + attn_out_dtype=torch.bfloat16,
308 + )
309 + attn_metadata = SimpleNamespace(decode=decode, causal=True, work_meta_data=None)
310 + layer = SimpleNamespace(_q_scale=torch.tensor(1.0), _k_scale=torch.tensor(1.0))
311 +
312 + q = torch.zeros(num_tokens, decode_heads, head_dim, dtype=torch.bfloat16)
313 + output, lse = impl.forward_mqa(q, torch.zeros(1, 1, head_dim), attn_metadata, layer)
314 +
315 + assert captured["return_lse"] is True
316 + # aiter is handed the padded head count, but both results come back at the
317 + # gathered count the cross-rank merge expects.
318 + assert captured["q_heads"] == 16
319 + assert lse is not None
320 + assert output.shape[1] == decode_heads
321 + assert lse.shape == (num_tokens, decode_heads)
322 +
323 +
324 +def test_verify_partial_attention_merge():
325 + device = torch.device("cuda")
326 + out_a = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]], device=device)
327 + out_b = torch.tensor([[[5.0, 6.0]], [[7.0, 8.0]]], device=device)
328 + lse_a = torch.tensor([[0.0], [1.0]], device=device)
329 + lse_b = torch.tensor([[1.0], [0.0]], device=device)
330 +
331 + out, lse = _lse_combine_natural(out_a, lse_a, out_b, lse_b, torch.float32)
332 + normalizer = torch.logaddexp(lse_a, lse_b)
333 + expected = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1) + out_b * torch.exp(
334 + lse_b - normalizer
335 + ).unsqueeze(-1)
336 +
337 + torch.testing.assert_close(out, expected)
338 + torch.testing.assert_close(lse, normalizer)
339 +
340 +
341 +def test_segmented_verify_reduce_returns_natural_lse_and_masks_empty_rows():
342 + segment_output = torch.tensor(
343 + [1.0, 3.0, 99.0, 99.0],
344 + device="cuda",
345 + ).reshape(2, 1, 2, 1)
346 + segment_max = torch.tensor(
347 + [[[0.0, 1.0]], [[99.0, 99.0]]],
348 + device="cuda",
349 + )
350 + segment_expsum = torch.ones_like(segment_max)
351 +
352 + output, lse = merge_mla_segments_triton(
353 + segment_output,
354 + segment_max,
355 + segment_expsum,
356 + torch.tensor([2, 0], dtype=torch.int32, device="cuda"),
357 + tile_size=1,
358 + out_dtype=torch.float32,
359 + )
360 +
361 + torch.testing.assert_close(output[0], torch.tensor([[7.0 / 3]], device="cuda"))
362 + assert output[1].item() == 0
363 + torch.testing.assert_close(lse[0], torch.tensor([math.log(3.0)], device="cuda"))
364 + assert lse[1].item() == float("-inf")
365 +
366 +
367 +def test_segmented_dcp_verify_matches_causal_attention(monkeypatch):
368 + """Exercise the target qlen>1 path and merge two rank-local results."""
369 + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
370 +
371 + torch.manual_seed(0)
372 + device = torch.device("cuda")
373 + dcp_world_size = 2
374 + qlen = 3
375 + num_heads = 64
376 + kv_lora_rank = 512
377 + rope_dim = 64
378 + head_dim = kv_lora_rank + rope_dim
379 + global_seq_len = 259
380 + block_size = 128
381 + kv_scale = 0.02
382 + sm_scale = head_dim**-0.5
383 +
384 + q_nope = torch.randn(
385 + qlen, num_heads, kv_lora_rank, dtype=torch.bfloat16, device=device
386 + )
387 + q_pe = torch.randn(qlen, num_heads, rope_dim, dtype=torch.bfloat16, device=device)
388 + kv_source = torch.randn(
389 + global_seq_len, head_dim, dtype=torch.float32, device=device
390 + )
391 + kv_fp8 = (kv_source / kv_scale).to(torch.float8_e4m3fn)
392 + kv_dequant = kv_fp8.float() * kv_scale
393 +
394 + partials = []
395 + for dcp_rank in range(dcp_world_size):
396 + local_kv = kv_fp8[dcp_rank::dcp_world_size]
397 + num_blocks = math.ceil(local_kv.shape[0] / block_size)
398 + kv_cache = torch.zeros(
399 + num_blocks,
400 + block_size,
401 + head_dim,
402 + dtype=torch.float8_e4m3fn,
403 + device=device,
404 + )
405 + kv_cache.view(-1, head_dim)[: local_kv.shape[0]].copy_(local_kv)
406 +
407 + builder = _builder(
408 + mtp_decode_qlen=qlen,
409 + kernel_block_size=block_size,
410 + num_heads=num_heads // dcp_world_size,
411 + kv_cache_dtype="fp8",
412 + dcp_world_size=dcp_world_size,
413 + dcp_rank=dcp_rank,
414 + )
415 + block_table = torch.arange(
416 + num_blocks, dtype=torch.int32, device=device
417 + ).unsqueeze(0)
418 + view = builder._build_dcp_verify_row_view(
419 + qlen,
420 + block_table,
421 + torch.tensor([global_seq_len], dtype=torch.int32, device=device),
422 + )
423 + decode = SimpleNamespace(
424 + max_qo_len=qlen,
425 + paged_kv_indptr=torch.tensor(
426 + [0, local_kv.shape[0]], dtype=torch.int32, device=device
427 + ),
428 + paged_kv_indices=torch.arange(
429 + local_kv.shape[0], dtype=torch.int32, device=device
430 + ),
431 + use_gluon_decode=False,
432 + use_gluon_verify=False,
433 + dcp_verify=view,
434 + attn_out_dtype=torch.bfloat16,
435 + )
436 + attn_metadata = SimpleNamespace(decode=decode, causal=True)
437 + impl = object.__new__(AiterMLAImpl)
438 + impl.num_heads = num_heads // dcp_world_size
439 + impl.dcp_world_size = dcp_world_size
440 + impl.kv_cache_dtype = "fp8"
441 + impl.kv_lora_rank = kv_lora_rank
442 + impl.qk_rope_head_dim = rope_dim
443 + impl.scale = sm_scale
444 + layer = SimpleNamespace(_k_scale=torch.tensor(kv_scale, device=device))
445 +
446 + partials.append(
447 + impl.forward_mqa((q_nope, q_pe), kv_cache, attn_metadata, layer)
448 + )
449 +
450 + output, lse = partials[0]
451 + assert lse is not None
452 + for rank_output, rank_lse in partials[1:]:
453 + assert rank_lse is not None
454 + output, lse = _lse_combine_natural(
455 + output.float(), lse, rank_output.float(), rank_lse, torch.float32
456 + )
457 +
458 + reference = torch.empty_like(output)
459 + q_nope_fp32 = q_nope.float()
460 + q_pe_fp32 = q_pe.float()
461 + for query_pos in range(qlen):
462 + visible = global_seq_len - qlen + query_pos + 1
463 + keys = kv_dequant[:visible]
464 + scores = torch.einsum(
465 + "hd,nd->hn", q_nope_fp32[query_pos], keys[:, :kv_lora_rank]
466 + )
467 + scores += torch.einsum(
468 + "hd,nd->hn", q_pe_fp32[query_pos], keys[:, kv_lora_rank:]
469 + )
470 + probs = torch.softmax(scores * sm_scale, dim=-1)
471 + reference[query_pos] = probs @ keys[:, :kv_lora_rank]
472 +
473 + torch.testing.assert_close(output, reference, rtol=3e-2, atol=3e-2)
474 +
475 +
476 @pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128])
477 @pytest.mark.parametrize(
478 "spec_method, parallel_drafting",
@@ -120,7 +482,7 @@ def test_backend_declares_uniform_batch_support():
482 # and a parallel one, so the threshold is 1 + 2 * num_spec rather than
483 # 1 + num_spec. Sizing the metadata off a method name instead leaves
484 # these at qlen=1 while the router still admits the full range.
123 - ("dspark", True),
485 + ("custom", True),
486 ("eagle", False),
487 ],
488 )
@@ -161,6 +523,8 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
523
524 def init_common_builder(self, *args, **kwargs):
525 self.num_heads = num_heads
526 + self.dcp_world_size = 1
527 + self.cp_kv_cache_interleave_size = 1
528 # Mirror what _init_reorder_batch_threshold would have left behind: the
529 # metadata is sized from the routing threshold, so a stub that skips it
530 # would size for qlen=1 and hide the very mismatch this test covers.
@@ -190,12 +554,23 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
554 num_speculative_tokens=3,
555 parallel_drafting=parallel_drafting,
556 ),
193 - parallel_config=SimpleNamespace(tensor_parallel_size=8),
194 - model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16),
557 + parallel_config=SimpleNamespace(
558 + tensor_parallel_size=8,
559 + decode_context_parallel_size=1,
560 + cp_kv_cache_interleave_size=1,
561 + ),
562 + model_config=SimpleNamespace(
563 + max_model_len=16,
564 + dtype=torch.bfloat16,
565 + get_num_attention_heads=lambda parallel_config: num_heads,
566 + ),
567 scheduler_config=SimpleNamespace(max_num_seqs=2),
568 cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"),
569 compilation_config=SimpleNamespace(
198 - cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False)
570 + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False),
571 + # Empty: the per-layer head-count probe finds no attention layer and
572 + # falls back to the model config above.
573 + static_forward_context={},
574 ),
575 )
576 builder = AiterMLAMetadataBuilder(
vllm/platforms/rocm.py
+10 -16
@@ -900,23 +900,17 @@ class RocmPlatform(Platform):
900 compilation_config = vllm_config.compilation_config
901 parallel_config = vllm_config.parallel_config
902
903 - if compilation_config.cudagraph_mode.has_full_cudagraphs():
904 - # decode context parallel does not support full cudagraphs
905 - if parallel_config.decode_context_parallel_size > 1:
906 - logger.warning_once(
907 - "Decode context parallel (DCP) is enabled, which is "
908 - "incompatible with full CUDA graphs. "
909 - "Overriding cudagraph_mode to PIECEWISE."
910 - )
911 - compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
903 + if (
904 + compilation_config.cudagraph_mode.has_full_cudagraphs()
905 + and parallel_config.prefill_context_parallel_size > 1
906 + ):
907 # prefill context parallel do not support full cudagraphs
913 - elif parallel_config.prefill_context_parallel_size > 1:
914 - logger.warning_once(
915 - "Prefill context parallel (PCP) is enabled, which is "
916 - "incompatible with full CUDA graphs. "
917 - "Overriding cudagraph_mode to PIECEWISE."
918 - )
919 - compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
908 + logger.warning_once(
909 + "Prefill context parallel (PCP) is enabled, which is "
910 + "incompatible with full CUDA graphs. "
911 + "Overriding cudagraph_mode to PIECEWISE."
912 + )
913 + compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
914
915 if parallel_config.worker_cls == "auto":
916 parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+530 -82
@@ -11,6 +11,7 @@ import torch
11 from vllm._aiter_ops import rocm_aiter_ops
12 from vllm.config import VllmConfig
13 from vllm.config.cache import CacheDType
14 +from vllm.distributed import get_dcp_group
15 from vllm.logger import init_logger
16 from vllm.model_executor.layers.attention.mla_attention import (
17 MLACommonBackend,
@@ -21,17 +22,39 @@ from vllm.model_executor.layers.attention.mla_attention import (
22 QueryLenSupport,
23 )
24 from vllm.triton_utils import tl, triton
25 +from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor
26 from vllm.v1.attention.backend import (
27 AttentionCGSupport,
28 AttentionLayer,
29 CommonAttentionMetadata,
30 MultipleOf,
31 )
32 +from vllm.v1.attention.backends.utils import (
33 + get_dcp_local_seq_lens,
34 +)
35 +from vllm.v1.attention.ops.rocm_aiter_mla_merge import (
36 + merge_mla_segments_triton,
37 +)
38 from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
39
40 logger = init_logger(__name__)
41
42
43 +def _segmented_mla_page_size(block_size: int) -> int:
44 + """Largest supported power-of-two subpage dividing a physical KV block.
45 +
46 + The subpage size becomes the segmented kernel's ``TILE_SIZE``: the cache is
47 + reinterpreted as ``(-1, page_size, 1, head_dim)`` before the call. It must
48 + divide the physical block and be a power of two, since the kernel walks a
49 + tile with ``tl.arange(0, TILE_SIZE)``. 128 is the largest tile that tiling
50 + supports -- a tile holds ``TILE_SIZE x kv_lora_rank`` keys and its scores
51 + are ``BLOCK_M x TILE_SIZE`` -- so a larger manager block is split into
52 + several subpages instead of widening the tile.
53 + """
54 + assert block_size > 0
55 + return min(128, largest_power_of_2_divisor(block_size))
56 +
57 +
58 @functools.lru_cache(maxsize=1)
59 def _get_mla_gluon():
60 """Load the small-head Gluon MLA entry point."""
@@ -137,6 +160,43 @@ def _gluon_mla_decode_supported() -> bool:
160 return on_gfx950()
161
162
163 +@functools.lru_cache(maxsize=1)
164 +def _get_segmented_mla_decode():
165 + """Load AITER's segmented MLA decode with unreduced partial output."""
166 + from aiter.ops.triton.attention.mla import mla_decode_fwd
167 +
168 + return mla_decode_fwd
169 +
170 +
171 +@functools.lru_cache(maxsize=1)
172 +def _get_aiter_mla_decode():
173 + from aiter.mla import mla_decode_fwd
174 +
175 + return mla_decode_fwd
176 +
177 +
178 +@functools.lru_cache(maxsize=1)
179 +def _segmented_mla_decode_supported() -> bool:
180 + """Whether AITER exposes the segmented MLA decode used by DCP verify."""
181 + try:
182 + _get_segmented_mla_decode()
183 + except Exception: # noqa: BLE001
184 + return False
185 + return True
186 +
187 +
188 +def _segmented_dcp_verify_supported(dcp_world_size: int, cp_interleave: int) -> bool:
189 + """Whether this configuration can serve DCP verify on segmented MLA.
190 +
191 + Configuration only -- whether a given batch takes the route additionally
192 + depends on its query length. Round-robin interleaving other than 1 is
193 + excluded because the per-row causal window has not been validated there.
194 + """
195 + return (
196 + dcp_world_size > 1 and cp_interleave == 1 and _segmented_mla_decode_supported()
197 + )
198 +
199 +
200 def _aiter_mla_small_head_mode() -> str:
201 """Small-head (<16) MLA decode kernel selection.
202
@@ -200,6 +260,31 @@ class AiterMLABackend(MLACommonBackend):
260 return AiterMLAMetadataBuilder
261
262
263 +@dataclass
264 +class AiterMLADCPVerifyMetadata:
265 + """One paged-KV row per verify token, for segmented DCP verification.
266 +
267 + These are produced together or not at all, so they travel as one value: its
268 + presence on the decode metadata *is* the routing decision the builder made,
269 + and the impl does not re-derive it.
270 + """
271 +
272 + # Local KV length of each verify row, already causally bounded.
273 + row_lens: torch.Tensor
274 + # Subpage IDs each row reads, shape [num_rows, max_local_pages].
275 + block_table: torch.Tensor
276 + # One query per row, so this is an arange.
277 + qo_indptr: torch.Tensor
278 + # Subpage size the block table is expressed in, and the segmented kernel's
279 + # TILE_SIZE. Carried so the impl reinterprets the cache exactly the way the
280 + # block table was built, instead of re-deriving it from the cache shape.
281 + page_size: int
282 + # Kernel-visible KV bound. Whenever full graphs are enabled this is the
283 + # configuration's maximum for every batch, not just during capture, so the
284 + # page table keeps one shape across replays.
285 + max_kv_seq_len: int
286 +
287 +
288 @dataclass
289 class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
290 # The indptr of the paged kv cache, shape: [batch_size + 1]
@@ -217,8 +302,12 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
302 max_qo_len: int | None = None
303 # Minimum KV length used by Gluon to choose a safe split count.
304 min_kv_seq_len: int = 1
305 + # Set exactly when this batch routes to segmented DCP verification.
306 + dcp_verify: AiterMLADCPVerifyMetadata | None = None
307 # Small-head decode uses Gluon (avoids padding to 16).
308 use_gluon_decode: bool = False
309 + # Small-head non-DCP multi-token verify uses Gluon's native MTP entry.
310 + use_gluon_verify: bool = False
311 # Whether persistent MLA metadata was computed
312 has_persistent_metadata: bool = False
313
@@ -305,13 +394,27 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
394 vllm_config: VllmConfig,
395 device: torch.device,
396 ):
397 + parallel_config = vllm_config.parallel_config
398 + supports_segmented_dcp_verify = _segmented_dcp_verify_supported(
399 + parallel_config.decode_context_parallel_size,
400 + parallel_config.cp_kv_cache_interleave_size,
401 + )
402 super().__init__(
309 - kv_cache_spec, layer_names, vllm_config, device, AiterMLAMetadata
403 + kv_cache_spec,
404 + layer_names,
405 + vllm_config,
406 + device,
407 + AiterMLAMetadata,
408 + supports_dcp_with_varlen=supports_segmented_dcp_verify,
409 )
410 + self._supports_segmented_dcp_verify = supports_segmented_dcp_verify
411
412 self.compilation_config = vllm_config.compilation_config
413 self.decode_attn_out_dtype = vllm_config.model_config.dtype
414
415 + # Needed to place a verify row's causal window on this rank's KV shard.
416 + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
417 +
418 # reorder_batch_threshold is the largest query length decode can be
419 # handed, and already accounts for the drafting scheme. A method-name
420 # whitelist sizes unlisted drafters for qlen=1, which closes the
@@ -323,6 +426,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
426 # (e.g. 16 with Eagle3), we expand block-level indices into per-token
427 # flat indices since the aiter kernel always uses page_size=1 internally.
428 self.kernel_block_size = kv_cache_spec.block_size
429 + self._segmented_page_size = _segmented_mla_page_size(self.kernel_block_size)
430
431 # In the flat view (.view(-1,1,1,H)), each token is its own page,
432 # so max_num_pages_per_req = max_model_len regardless of
@@ -351,10 +455,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
455
456 from aiter import dtypes, get_mla_metadata_info_v1
457
458 + # Decode kernels consume the DCP-gathered query heads.
459 + self._decode_num_heads = self.num_heads * self.dcp_world_size
460 # Keep metadata sizing consistent with the padded tensor shape passed
355 - # to mla_decode_fwd.
461 + # to mla_decode_fwd, including native 24-head AITER builds.
462 self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
357 - self.num_heads
463 + self._decode_num_heads
464 )
465 kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
466 if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
@@ -437,6 +543,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
543 device,
544 )
545
546 + # Persistent buffers for segmented DCP verification. Captured graphs
547 + # require stable addresses while row lengths vary between replays.
548 + self._dcp_verify_buffers: AiterMLADCPVerifyMetadata | None = None
549 + self._graph_seq_lens: torch.Tensor | None = None
550 if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
551 self.paged_kv_indptr = torch.zeros(
552 max_num_reqs + 1, dtype=torch.int32, device=device
@@ -446,6 +556,39 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
556 max_num_reqs + 1, dtype=torch.int32, device=device
557 )
558
559 + # Full graphs require a stable address after uniform-MTP padding.
560 + self._graph_seq_lens = torch.zeros(
561 + max_num_reqs, dtype=torch.int32, device=device
562 + )
563 +
564 + if self._supports_segmented_dcp_verify and self._mtp_decode_qlen > 1:
565 + # A DCP rank's shard of the longest sequence bounds every verify
566 + # row, and full graphs need that bound to be constant.
567 + num_dcp_partitions = (
568 + self.dcp_world_size * self.cp_kv_cache_interleave_size
569 + )
570 + graph_max_kv_seq_len = (
571 + cdiv(vllm_config.model_config.max_model_len, num_dcp_partitions)
572 + * self.cp_kv_cache_interleave_size
573 + )
574 + max_verify_rows = max_num_reqs * self._mtp_decode_qlen
575 + max_local_pages = cdiv(graph_max_kv_seq_len, self._segmented_page_size)
576 + self._dcp_verify_buffers = AiterMLADCPVerifyMetadata(
577 + row_lens=torch.zeros(
578 + max_verify_rows, dtype=torch.int32, device=device
579 + ),
580 + block_table=torch.zeros(
581 + (max_verify_rows, max_local_pages),
582 + dtype=torch.int32,
583 + device=device,
584 + ),
585 + qo_indptr=torch.arange(
586 + max_verify_rows + 1, dtype=torch.int32, device=device
587 + ),
588 + page_size=self._segmented_page_size,
589 + max_kv_seq_len=graph_max_kv_seq_len,
590 + )
591 +
592 def _init_fp8_prefill_ps_buffers(
593 self,
594 max_num_reqs: int,
@@ -667,6 +810,99 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
810 metadata.fp8_prefill_max_q_len = prefill.max_query_len
811 metadata.fp8_prefill_num_partial_tiles = num_partial_tiles
812
813 + def _build_dcp_verify_row_view(
814 + self,
815 + qlen: int,
816 + block_table: torch.Tensor,
817 + dcp_tot_seq_lens: torch.Tensor,
818 + ) -> AiterMLADCPVerifyMetadata:
819 + """Build one paged-KV row per verify token for segmented DCP verification.
820 +
821 + Every row is a single query (``qo_indptr`` is an arange), so the whole
822 + causal structure is carried by ``dcp_local_verify_row_lens`` and the
823 + kernel applies no tail of its own.
824 + """
825 + assert self.dcp_world_size > 1
826 + num_reqs = dcp_tot_seq_lens.numel()
827 + row_lens = AiterMLAHelper.dcp_local_verify_row_lens(
828 + dcp_tot_seq_lens,
829 + qlen,
830 + self.dcp_world_size,
831 + self.dcp_rank,
832 + self.cp_kv_cache_interleave_size,
833 + )
834 + num_rows = row_lens.numel()
835 + buffers = self._dcp_verify_buffers
836 + # Persistent buffers exist exactly when full graphs are on, and they
837 + # already carry the static bound they were sized for.
838 + max_kv_seq_len = (
839 + buffers.max_kv_seq_len
840 + if buffers is not None
841 + else max(1, int(row_lens.max().item()))
842 + )
843 + page_size = self._segmented_page_size
844 + max_local_pages = cdiv(max_kv_seq_len, page_size)
845 + if buffers is not None:
846 + row_block_table = buffers.block_table[:num_rows]
847 + buffers.row_lens[:num_rows].copy_(row_lens, non_blocking=True)
848 + row_lens = buffers.row_lens[:num_rows]
849 + qo_indptr = buffers.qo_indptr[: num_rows + 1]
850 + else:
851 + row_block_table = torch.empty(
852 + (num_rows, max_local_pages),
853 + dtype=torch.int32,
854 + device=block_table.device,
855 + )
856 + qo_indptr = torch.arange(
857 + num_rows + 1,
858 + dtype=torch.int32,
859 + device=block_table.device,
860 + )
861 + self._fill_dcp_verify_page_table(row_block_table, block_table, num_reqs, qlen)
862 + return AiterMLADCPVerifyMetadata(
863 + row_lens=row_lens,
864 + block_table=row_block_table,
865 + qo_indptr=qo_indptr,
866 + page_size=page_size,
867 + max_kv_seq_len=max_kv_seq_len,
868 + )
869 +
870 + def _fill_dcp_verify_page_table(
871 + self,
872 + row_block_table: torch.Tensor,
873 + block_table: torch.Tensor,
874 + num_reqs: int,
875 + qlen: int,
876 + ) -> None:
877 + """Expand each request's blocks into the subpages one verify row reads.
878 +
879 + Every row of a request shares the request's shard, so the page list is
880 + built once per request and repeated; only ``row_lens`` distinguishes the
881 + rows. A DCP rank holds ``1/dcp_world_size`` of the sequence, so the
882 + request's block table is always wider than the pages a row can reach.
883 + """
884 + pages_per_block = self.kernel_block_size // self._segmented_page_size
885 + max_local_pages = row_block_table.shape[1]
886 + max_local_blocks = cdiv(max_local_pages, pages_per_block)
887 + assert max_local_blocks <= block_table.shape[1], (
888 + f"DCP verify needs {max_local_blocks} blocks per request but the "
889 + f"block table only has {block_table.shape[1]}"
890 + )
891 + subpage_offsets = torch.arange(
892 + pages_per_block,
893 + dtype=block_table.dtype,
894 + device=block_table.device,
895 + )
896 + per_req_page_table = (
897 + block_table[:num_reqs, :max_local_blocks, None] * pages_per_block
898 + + subpage_offsets
899 + ).flatten(1)[:, :max_local_pages]
900 + # Broadcast the request's page list across its rows instead of
901 + # materializing a repeat_interleave copy.
902 + row_block_table.unflatten(0, (num_reqs, qlen)).copy_(
903 + per_req_page_table.unsqueeze(1)
904 + )
905 +
906 def _build_decode(
907 self,
908 block_table_tensor: torch.Tensor,
@@ -701,6 +937,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
937 seq_lens_for_kernel.new_full((), max_qo_len),
938 )
939
940 + if self._graph_seq_lens is not None:
941 + self._graph_seq_lens[:num_kernel_reqs].copy_(
942 + seq_lens_for_kernel, non_blocking=True
943 + )
944 + seq_lens_for_kernel = self._graph_seq_lens[:num_kernel_reqs]
945 +
946 # The aiter kernel always operates with page_size=1 (the wrapper
947 # flattens kv_buffer). last_page_len is always 1.
948 paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs]
@@ -713,26 +955,43 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
955 ]
956 )
957 use_gluon_decode = AiterMLAHelper.use_gluon_decode(
716 - self.num_heads, int(max_qo_len), self._kv_cache_dtype_str
958 + self._decode_num_heads,
959 + int(max_qo_len),
960 + self._kv_cache_dtype_str,
961 )
718 -
719 - if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
720 - self.paged_kv_indices.fill_(-1)
721 -
722 - # Expand block_table entries into per-token flat indices.
723 - # When kernel_block_size=1, this degrades to a direct copy (identical
724 - # to the original _copy_page_indices_kernel).
725 - # When kernel_block_size=K>1, block_table entry b covering K tokens
726 - # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
727 - _expand_page_indices_kernel[(num_reqs,)](
728 - self.paged_kv_indices,
729 - block_table_tensor,
730 - block_table_tensor.stride(0),
731 - paged_kv_indptr,
732 - KERNEL_BLOCK_SIZE=self.kernel_block_size,
733 - BLOCK_SIZE=1024,
962 + use_gluon_verify = AiterMLAHelper.use_gluon_verify(
963 + self._decode_num_heads,
964 + int(max_qo_len),
965 + self._kv_cache_dtype_str,
966 + self.dcp_world_size,
967 )
735 - paged_kv_indices = self.paged_kv_indices
968 + use_segmented_dcp_verify = (
969 + self._supports_segmented_dcp_verify and max_qo_len > 1
970 + )
971 +
972 + # Segmented DCP verify carries its own per-row subpage table, so the
973 + # flat per-token view is dead work for it. Leave the buffer alone and
974 + # hand the metadata None, so a future reader cannot pick up whatever
975 + # the previous batch left behind.
976 + paged_kv_indices = None
977 + if not use_segmented_dcp_verify:
978 + if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
979 + self.paged_kv_indices.fill_(-1)
980 +
981 + # Expand block_table entries into per-token flat indices.
982 + # When kernel_block_size=1, this degrades to a direct copy (identical
983 + # to the original _copy_page_indices_kernel).
984 + # When kernel_block_size=K>1, block_table entry b covering K tokens
985 + # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
986 + _expand_page_indices_kernel[(num_reqs,)](
987 + self.paged_kv_indices,
988 + block_table_tensor,
989 + block_table_tensor.stride(0),
990 + paged_kv_indptr,
991 + KERNEL_BLOCK_SIZE=self.kernel_block_size,
992 + BLOCK_SIZE=1024,
993 + )
994 + paged_kv_indices = self.paged_kv_indices
995
996 if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
997 self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_(
@@ -745,12 +1004,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1004 # (set above), so no copy needed - buffer is always 1s.
1005
1006 if pad_uniform_mtp:
748 - qo_indptr_src = torch.arange(
749 - 0,
750 - (num_kernel_reqs + 1) * max_qo_len,
751 - step=max_qo_len,
752 - dtype=torch.int32,
753 - device=device,
1007 + qo_indptr_src = AiterMLAHelper.qo_indptr_for_uniform_qlen(
1008 + num_kernel_reqs, int(max_qo_len), device
1009 )
1010 else:
1011 qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs]
@@ -762,42 +1017,32 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1017
1018 else:
1019 if max_qo_len == 1:
765 - qo_indptr = torch.arange(
766 - 0,
767 - num_kernel_reqs + 1,
768 - step=1,
769 - dtype=torch.int32,
770 - device=device,
1020 + qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen(
1021 + num_kernel_reqs, 1, device
1022 )
1023 else:
1024 if pad_uniform_mtp:
774 - qo_indptr = torch.arange(
775 - 0,
776 - (num_kernel_reqs + 1) * max_qo_len,
777 - step=max_qo_len,
778 - dtype=torch.int32,
779 - device=device,
1025 + qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen(
1026 + num_kernel_reqs, int(max_qo_len), device
1027 )
1028 else:
1029 qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
1030
1031 + has_persistent_metadata = False
1032 # Only the asm decode consumes the schedule, so gate on the routing
1033 # rather than on num_heads >= 16, which denies it to a padded rank
786 - # running the same asm kernels. The two predicates are disjoint --
787 - # decode is qlen==1, verify is qlen>1 -- and cover both Gluon entries.
788 - has_persistent_metadata = False
1034 + # running the same asm kernels. The predicates are disjoint -- decode
1035 + # is qlen==1, verify is qlen>1 -- and cover both Gluon entries plus the
1036 + # segmented DCP verify.
1037 use_persistent_metadata = (
790 - not AiterMLAHelper.use_gluon_decode(
791 - self.num_heads, max_qo_len, self._kv_cache_dtype_str
792 - )
793 - and not AiterMLAHelper.use_gluon_verify(
794 - self.num_heads, max_qo_len, self._kv_cache_dtype_str
795 - )
1038 + not use_gluon_decode
1039 + and not use_gluon_verify
1040 + and not use_segmented_dcp_verify
1041 # A padded rank has no bf16 persistent kernel past qlen 4 where the
1042 # gfx950 fold is absent; the non-persistent entry covers it. fp8
1043 # keeps the schedule -- its fold rejects non-persistent outright.
1044 and (
800 - self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
1045 + self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
1046 or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
1047 or is_quantized_kv_cache(self._kv_cache_dtype_str)
1048 )
@@ -840,11 +1085,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1085 # active requests, not cudagraph padding rows pinned to max_qo_len.
1086 # The .item() below runs in the builder, outside the captured region,
1087 # so it does not abort HIP graph capture the way the per-layer syncs
843 - # in forward_mqa did.
1088 + # in forward_mqa did. DCP verify never reaches here -- it carries its
1089 + # own per-row lengths instead of a single split bound.
1090 min_kv_seq_len = 1
845 - if AiterMLAHelper.use_gluon_verify(
846 - self.num_heads, int(max_qo_len), self._kv_cache_dtype_str
847 - ):
1091 + if use_gluon_verify:
1092 per_req_len = paged_kv_indptr[1:] - paged_kv_indptr[:-1]
1093 if pad_uniform_mtp:
1094 active = qo_lens_device > 0
@@ -853,6 +1097,15 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1097 else:
1098 min_kv_seq_len = int(per_req_len.min().item())
1099
1100 + dcp_verify = None
1101 + if use_segmented_dcp_verify:
1102 + assert dcp_tot_seq_lens_device is not None
1103 + dcp_verify = self._build_dcp_verify_row_view(
1104 + int(max_qo_len),
1105 + block_table_tensor,
1106 + dcp_tot_seq_lens_device,
1107 + )
1108 +
1109 attn_metadata = AiterMLADecodeMetadata(
1110 block_table=block_table_tensor,
1111 seq_lens=seq_lens_for_kernel,
@@ -863,7 +1116,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1116 dcp_tot_seq_lens=dcp_tot_seq_lens_device,
1117 max_qo_len=max_qo_len,
1118 min_kv_seq_len=min_kv_seq_len,
1119 + dcp_verify=dcp_verify,
1120 use_gluon_decode=use_gluon_decode,
1121 + use_gluon_verify=use_gluon_verify,
1122 attn_out_dtype=self.decode_attn_out_dtype,
1123 has_persistent_metadata=has_persistent_metadata,
1124 )
@@ -966,6 +1221,22 @@ class AiterMLAHelper:
1221 _ASM_PADDED_MAX_PS_QLEN: Final = 4
1222 _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
1223
1224 + @staticmethod
1225 + def qo_indptr_for_uniform_qlen(
1226 + num_reqs: int,
1227 + qlen: int,
1228 + device: torch.device,
1229 + dtype: torch.dtype = torch.int32,
1230 + ) -> torch.Tensor:
1231 + """Build ``[0, qlen, 2*qlen, ..., num_reqs*qlen]``."""
1232 + return torch.arange(
1233 + 0,
1234 + (num_reqs + 1) * qlen,
1235 + step=qlen,
1236 + dtype=dtype,
1237 + device=device,
1238 + )
1239 +
1240 @staticmethod
1241 def check_num_heads_validity(num_heads: int):
1242 assert AiterMLAHelper.is_valid_num_heads(num_heads), (
@@ -1042,14 +1313,22 @@ class AiterMLAHelper:
1313
1314 @staticmethod
1315 def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:
1316 + return AiterMLAHelper._get_mla_unpadded_heads(num_heads, o)
1317 +
1318 + @staticmethod
1319 + def _get_mla_unpadded_heads(num_heads: int, tensor: torch.Tensor) -> torch.Tensor:
1320 m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
1321 if num_heads == m:
1047 - return o
1322 + return tensor
1323 if m % num_heads == 0:
1049 - return o[:, :: m // num_heads, :]
1324 + return tensor[:, :: m // num_heads, ...]
1325 # Undo the tile-padding from get_mla_padded_q: the real heads are the
1326 # first num_heads.
1052 - return o[:, :num_heads, :]
1327 + return tensor[:, :num_heads, ...]
1328 +
1329 + @staticmethod
1330 + def get_mla_unpadded_lse(num_heads: int, lse: torch.Tensor) -> torch.Tensor:
1331 + return AiterMLAHelper._get_mla_unpadded_heads(num_heads, lse)
1332
1333 @staticmethod
1334 def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
@@ -1074,7 +1353,12 @@ class AiterMLAHelper:
1353 return m % num_heads == 0 and gluon_supported
1354
1355 @staticmethod
1077 - def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
1356 + def use_gluon_verify(
1357 + num_heads: int,
1358 + max_qo_len: int,
1359 + kv_cache_dtype: str,
1360 + dcp_world_size: int = 1,
1361 + ) -> bool:
1362 """Whether a small-head multi-token verify uses native Gluon MTP.
1363
1364 bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through
@@ -1083,16 +1367,67 @@ class AiterMLAHelper:
1367 come here: the MTP path hands Gluon the batch size its fp8 regime
1368 asserts against. A predicate rather than inline in forward_mqa so the
1369 builder sees the same answer the impl acts on.
1370 +
1371 + DCP verify is excluded: its per-row causal windows are served by the
1372 + segmented path, which Gluon's MTP entry cannot express.
1373 """
1087 - if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1:
1374 + if max_qo_len <= 1 or dcp_world_size > 1:
1375 return False
1376 if is_quantized_kv_cache(kv_cache_dtype):
1377 return False
1378 + if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS:
1379 + return False
1380 + if not _gluon_mla_decode_supported():
1381 + return False
1382 # Same arch and mode gating as use_gluon_decode.
1092 - return _aiter_mla_small_head_mode() != "asm" and _gluon_mla_decode_supported()
1383 + return _aiter_mla_small_head_mode() != "asm"
1384 +
1385 + @staticmethod
1386 + def dcp_local_verify_row_lens(
1387 + tot_seq_lens: torch.Tensor,
1388 + qlen: int,
1389 + dcp_world_size: int,
1390 + dcp_rank: int,
1391 + cp_interleave: int,
1392 + ) -> torch.Tensor:
1393 + """Local KV length of every verify row, in global-causal order.
1394 +
1395 + Row ``t`` of a qlen-token verify attends global positions
1396 + ``[0, seq_len - qlen + t]``, its own token included, so its local length
1397 + is the round-robin count evaluated at that bound. Truncating causally in
1398 + global coordinates has to happen before the shard mapping: the tokens a
1399 + row drops sit on ``qlen - 1 - t`` different ranks, while subtracting the
1400 + offset from a request's local length takes one from every rank.
1401 + """
1402 + offsets = torch.arange(
1403 + 1,
1404 + qlen + 1,
1405 + device=tot_seq_lens.device,
1406 + dtype=tot_seq_lens.dtype,
1407 + )
1408 + visible = (tot_seq_lens.unsqueeze(1) - qlen + offsets).clamp_(min=0)
1409 + return get_dcp_local_seq_lens(
1410 + visible,
1411 + dcp_world_size,
1412 + dcp_rank,
1413 + cp_interleave,
1414 + ).flatten()
1415
1416
1417 class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1418 + # DCP decode paths return natural-log softmax LSE for the cross-rank merge.
1419 + can_return_lse_for_decode: bool = True
1420 + # Measured on gfx950: aiter.mla.mla_decode_fwd(return_lse=True) matches
1421 + # logsumexp to fp32 exactly, and merge_mla_segments_triton converts AITER's
1422 + # base-2 segment statistics with LOGE2. Stated rather than inherited because
1423 + # the DCP combine silently corrupts the softmax denominator if it disagrees.
1424 + lse_base_on_e: bool = True
1425 +
1426 + @property
1427 + def _decode_num_heads(self) -> int:
1428 + """Return the query-head count after DCP gathering."""
1429 + return self.num_heads * self.dcp_world_size
1430 +
1431 def __init__(
1432 self,
1433 num_heads: int,
@@ -1122,6 +1457,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1457 **mla_args,
1458 )
1459 AiterMLAHelper.check_num_heads_validity(num_heads)
1460 + AiterMLAHelper.check_num_heads_validity(self._decode_num_heads)
1461
1462 unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
1463 if any(unsupported_features):
@@ -1363,6 +1699,59 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1699
1700 self._mla_fp8_prefill_attn(q, k, v, attn_metadata, output)
1701
1702 + def _forward_segmented_dcp_verify(
1703 + self,
1704 + q_nope: torch.Tensor,
1705 + q_pe: torch.Tensor,
1706 + verify: AiterMLADCPVerifyMetadata,
1707 + kv_c_and_k_pe_cache: torch.Tensor,
1708 + layer: AttentionLayer,
1709 + out_dtype: torch.dtype,
1710 + ) -> tuple[torch.Tensor, torch.Tensor]:
1711 + """Run segmented attention over this rank's shard of every verify row.
1712 +
1713 + Each row's length already covers the current tokens this rank holds, so
1714 + the cross-rank LSE merge in the MLA layer completes the causal block.
1715 + """
1716 + q_mla = torch.cat([q_nope, q_pe], dim=-1)
1717 + # skip_reduce with NUM_SEGMENTS>1 returns the partials and never writes
1718 + # `out`. Pass None rather than aliasing q_mla: if that ever stops
1719 + # holding, this fails loudly instead of scribbling over the query.
1720 + segment_partials = _get_segmented_mla_decode()(
1721 + q_mla,
1722 + kv_c_and_k_pe_cache.view(
1723 + -1,
1724 + verify.page_size,
1725 + 1,
1726 + kv_c_and_k_pe_cache.shape[-1],
1727 + ),
1728 + None,
1729 + verify.qo_indptr,
1730 + verify.row_lens,
1731 + verify.max_kv_seq_len,
1732 + verify.block_table,
1733 + self.scale,
1734 + self.kv_lora_rank,
1735 + self.qk_rope_head_dim,
1736 + causal=True,
1737 + q_descale=None,
1738 + kv_descale=layer._k_scale,
1739 + skip_reduce=True,
1740 + )
1741 + assert isinstance(segment_partials, tuple) and len(segment_partials) == 3, (
1742 + "AITER segmented MLA verify must return segment partials "
1743 + "when skip_reduce=True."
1744 + )
1745 + segm_output, segm_max, segm_expsum = segment_partials
1746 + return merge_mla_segments_triton(
1747 + segm_output,
1748 + segm_max,
1749 + segm_expsum,
1750 + verify.row_lens,
1751 + verify.page_size,
1752 + out_dtype,
1753 + )
1754 +
1755 def forward_mqa(
1756 self,
1757 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -1375,9 +1764,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1764
1765 decode = attn_metadata.decode
1766 assert decode.max_qo_len is not None
1378 - assert decode.paged_kv_indptr is not None
1379 - assert decode.paged_kv_indices is not None
1767 if decode.use_gluon_decode:
1768 + assert decode.paged_kv_indptr is not None
1769 + assert decode.paged_kv_indices is not None
1770 if type(q) is tuple:
1771 q_nope, q_pe = q
1772 else:
@@ -1416,12 +1805,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1805 # The block is causal -- the target is checking draft tokens, so
1806 # position t must not see t+1 -- and the kernel bounds each query
1807 # position's scores itself.
1419 - # Arch, mode and dtype gating all live in use_gluon_verify, so that the
1420 - # builder -- which has to know whether the asm decode will run -- sees
1421 - # the same answer as this branch.
1422 - if AiterMLAHelper.use_gluon_verify(
1423 - self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype
1424 - ):
1808 + # Arch, mode and dtype gating all live in use_gluon_verify, which the
1809 + # builder evaluates and records on the metadata, so this branch acts on
1810 + # the same answer the builder used to decide whether the asm decode
1811 + # would run. DCP verify is routed to the segmented path below instead.
1812 + if decode.use_gluon_verify:
1813 if type(q) is tuple:
1814 q_nope, q_pe = q
1815 else:
@@ -1473,14 +1861,48 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1861 )
1862 return o, None
1863
1864 + verify = decode.dcp_verify
1865 + if verify is not None:
1866 + if type(q) is tuple:
1867 + q_nope, q_pe = q
1868 + else:
1869 + q_nope, q_pe = torch.split(
1870 + q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
1871 + )
1872 + if (
1873 + is_quantized_kv_cache(self.kv_cache_dtype)
1874 + and q_nope.dtype != torch.bfloat16
1875 + ):
1876 + q_nope = q_nope.to(torch.bfloat16) * layer._q_scale
1877 + q_pe = q_pe.to(torch.bfloat16) * layer._q_scale
1878 + return self._forward_segmented_dcp_verify(
1879 + q_nope,
1880 + q_pe,
1881 + verify,
1882 + kv_c_and_k_pe_cache,
1883 + layer,
1884 + decode.attn_out_dtype,
1885 + )
1886 +
1887 + if self.dcp_world_size > 1 and int(decode.max_qo_len) > 1:
1888 + raise RuntimeError(
1889 + "ROCM_AITER_MLA DCP multi-token verify requires segmented MLA."
1890 + )
1891 +
1892 if type(q) is tuple:
1893 q = torch.cat(q, dim=-1)
1894
1895 assert isinstance(q, torch.Tensor)
1896 + assert decode.paged_kv_indptr is not None
1897 + assert decode.paged_kv_indices is not None
1898 B = q.shape[0]
1899
1482 - mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
1483 - mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
1900 + assert q.shape[1] == self._decode_num_heads, (
1901 + "ROCM_AITER_MLA decode expected the DCP-gathered query head count "
1902 + f"{self._decode_num_heads}, got {q.shape[1]}"
1903 + )
1904 + mla_padded_q = AiterMLAHelper.get_mla_padded_q(self._decode_num_heads, q)
1905 + mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads)
1906 o = torch.empty(
1907 B,
1908 mla_num_heads,
@@ -1512,17 +1934,43 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1934 reduce_partial_map=attn_metadata.reduce_partial_map,
1935 )
1936
1515 - rocm_aiter_ops.mla_decode_fwd(
1516 - mla_padded_q,
1517 - kv_buffer,
1518 - o,
1519 - self.scale,
1520 - decode.qo_indptr,
1521 - decode.max_qo_len,
1522 - decode.paged_kv_indptr,
1523 - decode.paged_kv_indices,
1524 - decode.paged_kv_last_page_len,
1525 - **mla_kwargs,
1526 - )
1937 + lse = None
1938 + if self.dcp_world_size > 1:
1939 + # The vLLM custom-op wrapper exposes only the in-place output and
1940 + # drops aiter's final LSE, which the cross-shard merge needs, so go
1941 + # through aiter's native entry point on the DCP path.
1942 + _, lse = _get_aiter_mla_decode()(
1943 + mla_padded_q,
1944 + kv_buffer.view(-1, 1, 1, mla_padded_q.shape[-1]),
1945 + o,
1946 + decode.qo_indptr,
1947 + decode.paged_kv_indptr,
1948 + decode.paged_kv_indices,
1949 + decode.paged_kv_last_page_len,
1950 + decode.max_qo_len,
1951 + sm_scale=self.scale,
1952 + return_lse=True,
1953 + **mla_kwargs,
1954 + )
1955 + assert lse is not None, (
1956 + "aiter mla_decode_fwd(return_lse=True) returned no LSE; upgrade "
1957 + "aiter to a build with decode LSE support."
1958 + )
1959 + else:
1960 + rocm_aiter_ops.mla_decode_fwd(
1961 + mla_padded_q,
1962 + kv_buffer,
1963 + o,
1964 + self.scale,
1965 + decode.qo_indptr,
1966 + decode.max_qo_len,
1967 + decode.paged_kv_indptr,
1968 + decode.paged_kv_indices,
1969 + decode.paged_kv_last_page_len,
1970 + **mla_kwargs,
1971 + )
1972
1528 - return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, o), None
1973 + output = AiterMLAHelper.get_mla_unpadded_o(self._decode_num_heads, o)
1974 + if lse is not None:
1975 + lse = AiterMLAHelper.get_mla_unpadded_lse(self._decode_num_heads, lse)
1976 + return output, lse
vllm/v1/attention/ops/rocm_aiter_mla_merge.py new
+136
@@ -0,0 +1,136 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Merge AITER segmented MLA split-K partials into output plus natural-log LSE.
4 +
5 +This mirrors AITER's own segment reduction, including its
6 +``tiles_per_segment = cdiv(seq_len, NUM_SEGMENTS * TILE_SIZE)`` partitioning, so
7 +it has to move together with the ``skip_reduce=True`` call in the AITER MLA
8 +backend. It is a rank-local split-K merge and unrelated to any collective
9 +reduce; the natural-log LSE it returns is what the cross-rank DCP merge
10 +consumes.
11 +"""
12 +
13 +import torch
14 +
15 +from vllm.triton_utils import LOGE2, tl, triton
16 +
17 +
18 +@triton.jit
19 +def _merge_mla_segments_kernel(
20 + out_ptr,
21 + lse_ptr,
22 + segm_output_ptr,
23 + segm_max_ptr,
24 + segm_expsum_ptr,
25 + seq_lens_ptr,
26 + num_query_heads: tl.constexpr,
27 + out_stride0: tl.int64,
28 + out_stride1: tl.int64,
29 + lse_stride0: tl.int64,
30 + TILE_SIZE: tl.constexpr,
31 + KV_LORA_RANK: tl.constexpr,
32 + NUM_SEGMENTS_PER_SEQ: tl.constexpr,
33 + LOGE2: tl.constexpr,
34 +):
35 + token_idx = tl.program_id(0)
36 + head_idx = tl.program_id(1)
37 +
38 + seq_len = tl.load(seq_lens_ptr + token_idx)
39 + tiles_per_segment = tl.maximum(
40 + 1, tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
41 + )
42 + active_segments = tl.where(
43 + seq_len > 0,
44 + tl.cdiv(seq_len, tiles_per_segment * TILE_SIZE),
45 + 0,
46 + )
47 + segment_offsets = tl.arange(0, NUM_SEGMENTS_PER_SEQ)
48 + segment_mask = segment_offsets < active_segments
49 +
50 + segment_base = (
51 + token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
52 + + head_idx * NUM_SEGMENTS_PER_SEQ
53 + )
54 + segment_indices = segment_base + segment_offsets
55 + segment_max = tl.load(
56 + segm_max_ptr + segment_indices,
57 + mask=segment_mask,
58 + other=float("-inf"),
59 + )
60 + overall_max = tl.max(segment_max)
61 +
62 + segment_expsum = tl.load(
63 + segm_expsum_ptr + segment_indices,
64 + mask=segment_mask,
65 + other=0.0,
66 + )
67 + segment_weight = tl.where(
68 + segment_mask,
69 + tl.exp2(segment_max - overall_max),
70 + 0.0,
71 + )
72 + overall_expsum = tl.sum(segment_expsum * segment_weight)
73 +
74 + output_offsets = (
75 + token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK)
76 + + head_idx * (NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK)
77 + + segment_offsets[:, None] * KV_LORA_RANK
78 + + tl.arange(0, KV_LORA_RANK)[None, :]
79 + )
80 + segment_output = tl.load(
81 + segm_output_ptr + output_offsets,
82 + mask=segment_mask[:, None],
83 + other=0.0,
84 + )
85 + accumulator = tl.sum(segment_output * segment_weight[:, None], axis=0)
86 + accumulator = tl.where(overall_expsum == 0.0, 0.0, accumulator / overall_expsum)
87 + lse = tl.where(
88 + overall_expsum == 0.0,
89 + float("-inf"),
90 + (overall_max + tl.log2(overall_expsum)) * LOGE2,
91 + )
92 +
93 + output_indices = (
94 + token_idx * out_stride0 + head_idx * out_stride1 + tl.arange(0, KV_LORA_RANK)
95 + )
96 + tl.store(out_ptr + output_indices, accumulator)
97 + tl.store(lse_ptr + token_idx * lse_stride0 + head_idx, lse)
98 +
99 +
100 +def merge_mla_segments_triton(
101 + segm_output: torch.Tensor,
102 + segm_max: torch.Tensor,
103 + segm_expsum: torch.Tensor,
104 + seq_lens: torch.Tensor,
105 + tile_size: int,
106 + out_dtype: torch.dtype,
107 +) -> tuple[torch.Tensor, torch.Tensor]:
108 + """Merge AITER base-2 segment partials into output and natural-log LSE."""
109 + num_tokens, num_heads, num_segments, kv_lora_rank = segm_output.shape
110 + output = torch.empty(
111 + (num_tokens, num_heads, kv_lora_rank),
112 + dtype=out_dtype,
113 + device=segm_output.device,
114 + )
115 + lse = torch.empty(
116 + (num_tokens, num_heads),
117 + dtype=torch.float32,
118 + device=segm_output.device,
119 + )
120 + _merge_mla_segments_kernel[(num_tokens, num_heads)](
121 + output,
122 + lse,
123 + segm_output,
124 + segm_max,
125 + segm_expsum,
126 + seq_lens,
127 + num_query_heads=num_heads,
128 + out_stride0=output.stride(0),
129 + out_stride1=output.stride(1),
130 + lse_stride0=lse.stride(0),
131 + TILE_SIZE=tile_size,
132 + KV_LORA_RANK=kv_lora_rank,
133 + NUM_SEGMENTS_PER_SEQ=num_segments,
134 + LOGE2=LOGE2,
135 + )
136 + return output, lse