diff --git a/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py index dab49c33fa..ae66f98a88 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py @@ -58,6 +58,13 @@ PAGE_SIZE = 1 MAX_MODEL_LEN = 1024 MAX_NUM_SEQS = 8 +# The builder resolves the verify routing and records it on the metadata, so +# the gfx950 feature probe has to be forced for the build as well as for the +# forward; ROCm CI also runs gfx942, where the real probe is False. +_GLUON_SUPPORTED_TARGET = ( + "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported" +) + def _seq_lens() -> list[int]: return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS @@ -99,6 +106,10 @@ def _run_verify_block(): vllm_config.speculative_config = SpeculativeConfig( method="ngram", num_speculative_tokens=QLEN - 1 ) + vllm_config.model_config.get_num_attention_heads = types.MethodType( + lambda self, parallel_config, arch_config=None: NUM_QUERY_HEADS, + vllm_config.model_config, + ) spec = MLAAttentionSpec( block_size=PAGE_SIZE, @@ -138,7 +149,10 @@ def _run_verify_block(): captured["page_table"] = kwargs["page_table"].detach().clone() captured["seq_info"] = kwargs["seq_info"].detach().clone() - with set_current_vllm_config(vllm_config): + with ( + patch(_GLUON_SUPPORTED_TARGET, lambda: True), + set_current_vllm_config(vllm_config), + ): builder = builder_cls(spec, [layer_name], vllm_config, device) common_attn_metadata = create_common_attn_metadata( batch_spec, PAGE_SIZE, device, arange_block_indices=True @@ -182,10 +196,7 @@ def _run_verify_block(): # The Gluon kernel only reads the metadata this test is about, so a spy in # its place keeps the assertions independent of the AITER build. with ( - patch( - "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported", - lambda: True, - ), + patch(_GLUON_SUPPORTED_TARGET, lambda: True), patch( "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon", lambda: spy, @@ -211,6 +222,10 @@ def test_verify_mtp_uses_native_4d_gluon_entry(): f"expected a {QLEN}-token verify block, got max_qo_len=" f"{decode.max_qo_len}; the MTP path under test was not reached" ) + assert decode.use_gluon_verify, ( + "the builder owns verify routing, so the metadata -- not the impl -- " + "has to select the Gluon MTP entry for a small-head bf16 verify block" + ) assert captured, "forward_mqa did not reach the Gluon kernel" seq_lens = _seq_lens() diff --git a/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py index 8a284e320a..f9130460ef 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py @@ -324,7 +324,9 @@ def test_h12_aiter_mla_decode_matches_reference(): impl = object.__new__(AiterMLAImpl) impl.num_heads = NUM_HEADS + impl.dcp_world_size = 1 impl.kv_lora_rank = KV_LORA_RANK + impl.qk_rope_head_dim = QK_ROPE_HEAD_DIM impl.scale = SCALE impl.kv_cache_dtype = "auto" diff --git a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py index 31e16c6250..06684302cf 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py +++ b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py @@ -8,6 +8,10 @@ batch_size == 1, so it cannot serve a decode batch. Every fp8 shape therefore has to land on the asm kernels, which ship real fp8 variants for gqa=16. These tests pin that down at the predicates, since the failure it prevents is either a batch assertion or -- worse -- a silently wrong result. + +Gluon is still the route for supported non-DCP decode shapes. DCP multi-token +verification uses segmented MLA instead, so it depends on neither Gluon's +query-head nor its long-context pipeline limits. """ import pytest @@ -41,8 +45,8 @@ def gluon_available(monkeypatch): @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES) @pytest.mark.parametrize("num_heads", [1, 2, 4, 5, 6, 8, 12, 16, 32, 128]) @pytest.mark.parametrize("max_qo_len", [1, 2, 4, 5, 8, 15]) -def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len): - """No fp8 shape reaches either Gluon entry point. +def test_non_dcp_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len): + """No non-DCP fp8 shape reaches either Gluon entry point. The head count is deliberately swept across divisors of 16 as well as 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): def test_fp8_never_routes_to_gluon_under_any_mode( monkeypatch, kv_cache_dtype, num_heads, mode ): - """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force an fp8 cache onto Gluon. + """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force a non-DCP fp8 cache onto Gluon. The dtype guard deliberately precedes the mode knob: honouring an explicit "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 assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) +@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES) +@pytest.mark.parametrize("num_heads", [8, 16, 32, 96]) +@pytest.mark.parametrize("max_qo_len", [2, 3, 8]) +def test_dcp_multitoken_verify_never_uses_gluon( + gluon_available, kv_cache_dtype, num_heads, max_qo_len +): + assert not AiterMLAHelper.use_gluon_verify( + num_heads, max_qo_len, kv_cache_dtype, dcp_world_size=8 + ) + + +def test_segmented_dcp_verify_does_not_depend_on_gluon(monkeypatch): + monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False) + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) + + assert rocm_aiter_mla._segmented_dcp_verify_supported(8, 1) + # Round-robin interleaving other than 1 is not served by this route. + assert not rocm_aiter_mla._segmented_dcp_verify_supported(8, 4) + + @pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) @pytest.mark.parametrize("num_heads", [1, 2, 4, 8]) def test_unquantized_divisor_heads_keep_gluon_decode( diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 604e7b84f7..112cdff324 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math import sys from types import SimpleNamespace from unittest import mock @@ -16,8 +17,12 @@ if not current_platform.is_rocm(): from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402 from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402 AiterMLAHelper, + AiterMLAImpl, AiterMLAMetadataBuilder, ) +from vllm.v1.attention.ops.rocm_aiter_mla_merge import ( # noqa: E402 + merge_mla_segments_triton, +) class _NoOpTritonKernel: @@ -29,6 +34,19 @@ class _NoOpTritonKernel: pass +def _lse_combine_natural( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, + out_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + normalizer = torch.logaddexp(lse_a, lse_b) + out = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1) + out += out_b * torch.exp(lse_b - normalizer).unsqueeze(-1) + return out.to(out_dtype), normalizer + + class _ExpandPageIndicesKernel: def __getitem__(self, grid): self.grid = grid @@ -67,10 +85,27 @@ def _builder( max_decode_rows: int = 32, num_heads: int = 16, kv_cache_dtype: str = "auto", + dcp_world_size: int = 1, + dcp_rank: int = 0, ): - return SimpleNamespace( + stub = SimpleNamespace( device=torch.device("cpu"), num_heads=num_heads, + # DCP gathers every rank's head shard before decode, and the routing + # predicates read the gathered count. + _decode_num_heads=num_heads * dcp_world_size, + dcp_world_size=dcp_world_size, + dcp_rank=dcp_rank, + cp_kv_cache_interleave_size=1, + # Mirrors the production constructor: configuration decides whether the + # segmented route is available, query length decides per batch. + _supports_segmented_dcp_verify=rocm_aiter_mla._segmented_dcp_verify_supported( + dcp_world_size, 1 + ), + # Derived once in the real constructor, so derive it once here too. + _segmented_page_size=rocm_aiter_mla._segmented_mla_page_size(kernel_block_size), + _dcp_verify_buffers=None, + _graph_seq_lens=None, _kv_cache_dtype_str=kv_cache_dtype, paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32), paged_kv_indices=torch.empty(1024, dtype=torch.int32), @@ -96,6 +131,15 @@ def _builder( _mla_kv_dtype=torch.bfloat16, decode_attn_out_dtype=torch.bfloat16, ) + # Bound to the stub rather than faked: the verify flatten's per-row view is + # part of what _build_decode is being tested for. + stub._build_dcp_verify_row_view = ( + AiterMLAMetadataBuilder._build_dcp_verify_row_view.__get__(stub) + ) + stub._fill_dcp_verify_page_table = ( + AiterMLAMetadataBuilder._fill_dcp_verify_page_table.__get__(stub) + ) + return stub def test_backend_declares_uniform_batch_support(): @@ -111,6 +155,324 @@ def test_backend_declares_uniform_batch_support(): ) +def test_dcp_verify_row_view_is_causal_per_row(): + """Row t must cover this rank's share of positions [0, seq_len - qlen + t]. + + Regression guard: giving every row of a request the committed prefix only + drops the whole verify block, including each row's own token, since nothing + else adds it back on the target path. + """ + qlen = 4 + builder = _builder(mtp_decode_qlen=qlen, dcp_world_size=2, kernel_block_size=2) + block_table = torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32) + + view = builder._build_dcp_verify_row_view( + qlen, + block_table, + torch.tensor([10, 12], dtype=torch.int32), + ) + + # seq_len 10 -> rows see 7, 8, 9, 10 global positions, of which rank 0 of 2 + # holds ceil(n / 2); seq_len 12 -> 9, 10, 11, 12. + assert view.row_lens.tolist() == [4, 4, 5, 5, 5, 5, 6, 6] + assert view.qo_indptr.tolist() == list(range(2 * qlen + 1)) + # One query per row, so the page table is the request's ascending shard and + # only the row length selects the causal prefix. + assert view.block_table.tolist() == [[7, 8, 9]] * qlen + [[10, 11, 12]] * qlen + assert view.max_kv_seq_len == 6 + + +def test_dcp_verify_rows_sum_to_the_global_window_across_ranks(): + """Every rank's rows together must cover each row's window exactly once.""" + qlen = 4 + dcp_world_size = 2 + seq_lens = torch.tensor([10, 12], dtype=torch.int32) + per_rank = [] + for dcp_rank in range(dcp_world_size): + builder = _builder( + mtp_decode_qlen=qlen, + dcp_world_size=dcp_world_size, + dcp_rank=dcp_rank, + kernel_block_size=2, + ) + view = builder._build_dcp_verify_row_view( + qlen, + torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32), + seq_lens, + ) + per_rank.append(view.row_lens.tolist()) + + expected = [ + int(seq_len) - qlen + t + 1 + for seq_len in seq_lens.tolist() + for t in range(qlen) + ] + assert [sum(lens) for lens in zip(*per_rank)] == expected + + +def test_dcp_verify_row_view_uses_static_graph_bound(): + """Under full graphs the bound comes from the buffers, not the batch.""" + qlen = 3 + builder = _builder( + mtp_decode_qlen=qlen, + dcp_world_size=8, + has_full_cudagraphs=True, + kernel_block_size=1536, + ) + builder._dcp_verify_buffers = rocm_aiter_mla.AiterMLADCPVerifyMetadata( + row_lens=torch.zeros(qlen, dtype=torch.int32), + block_table=torch.zeros((qlen, 24), dtype=torch.int32), + qo_indptr=torch.zeros(qlen + 1, dtype=torch.int32), + page_size=128, + max_kv_seq_len=3072, + ) + + view = builder._build_dcp_verify_row_view( + qlen, + torch.arange(2, dtype=torch.int32).view(1, 2), + torch.tensor([1000], dtype=torch.int32), + ) + + buffers = builder._dcp_verify_buffers + assert view.row_lens.tolist() == [125] * qlen + assert view.block_table.data_ptr() == buffers.block_table.data_ptr() + assert view.qo_indptr.data_ptr() == buffers.qo_indptr.data_ptr() + assert view.block_table[0].tolist() == list(range(24)) + assert view.max_kv_seq_len == 3072 + + +def test_dcp_fp8_verify_build_uses_segmented(monkeypatch): + qlen = 4 + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) + + metadata = AiterMLAMetadataBuilder._build_decode( + _builder( + mtp_decode_qlen=qlen, + dcp_world_size=2, + kv_cache_dtype="fp8", + kernel_block_size=2, + ), + block_table_tensor=torch.tensor([[0, 1, 2], [10, 11, 12]], dtype=torch.int32), + seq_lens_device=torch.tensor([5, 6], dtype=torch.int32), + max_seq_len=6, + query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + num_decode_tokens=2 * qlen, + dcp_tot_seq_lens_device=torch.tensor([10, 12], dtype=torch.int32), + ) + + assert metadata.dcp_verify is not None + + +def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch): + """Single-token DCP decode must come back with an LSE the merge can use. + + This is the path every DCP step takes outside verification. The vLLM custom + op drops aiter's LSE, so the impl calls aiter directly; without an LSE the + MLA layer's cross-rank combine asserts. The head count is deliberately a + non-divisor of 16 so both the output and the LSE have to be unpadded back + to the gathered head count. + """ + num_heads, dcp_world_size = 6, 2 + decode_heads = num_heads * dcp_world_size # 12, padded to 16 by aiter + num_tokens, head_dim = 3, 576 + captured = {} + + def fake_aiter_decode(q, kv_buffer, out, *args, **kwargs): + captured["q_heads"] = q.shape[1] + captured["return_lse"] = kwargs.get("return_lse") + return None, torch.zeros(num_tokens, q.shape[1]) + + monkeypatch.setattr( + rocm_aiter_mla, "_get_aiter_mla_decode", lambda: fake_aiter_decode + ) + + impl = object.__new__(AiterMLAImpl) + impl.num_heads = num_heads + impl.dcp_world_size = dcp_world_size + impl.kv_cache_dtype = "auto" + impl.kv_lora_rank = 512 + impl.qk_rope_head_dim = 64 + impl.scale = head_dim**-0.5 + decode = SimpleNamespace( + max_qo_len=1, + qo_indptr=torch.arange(num_tokens + 1, dtype=torch.int32), + paged_kv_indptr=torch.zeros(num_tokens + 1, dtype=torch.int32), + paged_kv_indices=torch.zeros(1, dtype=torch.int32), + paged_kv_last_page_len=torch.ones(num_tokens, dtype=torch.int32), + use_gluon_decode=False, + use_gluon_verify=False, + dcp_verify=None, + has_persistent_metadata=False, + attn_out_dtype=torch.bfloat16, + ) + attn_metadata = SimpleNamespace(decode=decode, causal=True, work_meta_data=None) + layer = SimpleNamespace(_q_scale=torch.tensor(1.0), _k_scale=torch.tensor(1.0)) + + q = torch.zeros(num_tokens, decode_heads, head_dim, dtype=torch.bfloat16) + output, lse = impl.forward_mqa(q, torch.zeros(1, 1, head_dim), attn_metadata, layer) + + assert captured["return_lse"] is True + # aiter is handed the padded head count, but both results come back at the + # gathered count the cross-rank merge expects. + assert captured["q_heads"] == 16 + assert lse is not None + assert output.shape[1] == decode_heads + assert lse.shape == (num_tokens, decode_heads) + + +def test_verify_partial_attention_merge(): + device = torch.device("cuda") + out_a = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]], device=device) + out_b = torch.tensor([[[5.0, 6.0]], [[7.0, 8.0]]], device=device) + lse_a = torch.tensor([[0.0], [1.0]], device=device) + lse_b = torch.tensor([[1.0], [0.0]], device=device) + + out, lse = _lse_combine_natural(out_a, lse_a, out_b, lse_b, torch.float32) + normalizer = torch.logaddexp(lse_a, lse_b) + expected = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1) + out_b * torch.exp( + lse_b - normalizer + ).unsqueeze(-1) + + torch.testing.assert_close(out, expected) + torch.testing.assert_close(lse, normalizer) + + +def test_segmented_verify_reduce_returns_natural_lse_and_masks_empty_rows(): + segment_output = torch.tensor( + [1.0, 3.0, 99.0, 99.0], + device="cuda", + ).reshape(2, 1, 2, 1) + segment_max = torch.tensor( + [[[0.0, 1.0]], [[99.0, 99.0]]], + device="cuda", + ) + segment_expsum = torch.ones_like(segment_max) + + output, lse = merge_mla_segments_triton( + segment_output, + segment_max, + segment_expsum, + torch.tensor([2, 0], dtype=torch.int32, device="cuda"), + tile_size=1, + out_dtype=torch.float32, + ) + + torch.testing.assert_close(output[0], torch.tensor([[7.0 / 3]], device="cuda")) + assert output[1].item() == 0 + torch.testing.assert_close(lse[0], torch.tensor([math.log(3.0)], device="cuda")) + assert lse[1].item() == float("-inf") + + +def test_segmented_dcp_verify_matches_causal_attention(monkeypatch): + """Exercise the target qlen>1 path and merge two rank-local results.""" + monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) + + torch.manual_seed(0) + device = torch.device("cuda") + dcp_world_size = 2 + qlen = 3 + num_heads = 64 + kv_lora_rank = 512 + rope_dim = 64 + head_dim = kv_lora_rank + rope_dim + global_seq_len = 259 + block_size = 128 + kv_scale = 0.02 + sm_scale = head_dim**-0.5 + + q_nope = torch.randn( + qlen, num_heads, kv_lora_rank, dtype=torch.bfloat16, device=device + ) + q_pe = torch.randn(qlen, num_heads, rope_dim, dtype=torch.bfloat16, device=device) + kv_source = torch.randn( + global_seq_len, head_dim, dtype=torch.float32, device=device + ) + kv_fp8 = (kv_source / kv_scale).to(torch.float8_e4m3fn) + kv_dequant = kv_fp8.float() * kv_scale + + partials = [] + for dcp_rank in range(dcp_world_size): + local_kv = kv_fp8[dcp_rank::dcp_world_size] + num_blocks = math.ceil(local_kv.shape[0] / block_size) + kv_cache = torch.zeros( + num_blocks, + block_size, + head_dim, + dtype=torch.float8_e4m3fn, + device=device, + ) + kv_cache.view(-1, head_dim)[: local_kv.shape[0]].copy_(local_kv) + + builder = _builder( + mtp_decode_qlen=qlen, + kernel_block_size=block_size, + num_heads=num_heads // dcp_world_size, + kv_cache_dtype="fp8", + dcp_world_size=dcp_world_size, + dcp_rank=dcp_rank, + ) + block_table = torch.arange( + num_blocks, dtype=torch.int32, device=device + ).unsqueeze(0) + view = builder._build_dcp_verify_row_view( + qlen, + block_table, + torch.tensor([global_seq_len], dtype=torch.int32, device=device), + ) + decode = SimpleNamespace( + max_qo_len=qlen, + paged_kv_indptr=torch.tensor( + [0, local_kv.shape[0]], dtype=torch.int32, device=device + ), + paged_kv_indices=torch.arange( + local_kv.shape[0], dtype=torch.int32, device=device + ), + use_gluon_decode=False, + use_gluon_verify=False, + dcp_verify=view, + attn_out_dtype=torch.bfloat16, + ) + attn_metadata = SimpleNamespace(decode=decode, causal=True) + impl = object.__new__(AiterMLAImpl) + impl.num_heads = num_heads // dcp_world_size + impl.dcp_world_size = dcp_world_size + impl.kv_cache_dtype = "fp8" + impl.kv_lora_rank = kv_lora_rank + impl.qk_rope_head_dim = rope_dim + impl.scale = sm_scale + layer = SimpleNamespace(_k_scale=torch.tensor(kv_scale, device=device)) + + partials.append( + impl.forward_mqa((q_nope, q_pe), kv_cache, attn_metadata, layer) + ) + + output, lse = partials[0] + assert lse is not None + for rank_output, rank_lse in partials[1:]: + assert rank_lse is not None + output, lse = _lse_combine_natural( + output.float(), lse, rank_output.float(), rank_lse, torch.float32 + ) + + reference = torch.empty_like(output) + q_nope_fp32 = q_nope.float() + q_pe_fp32 = q_pe.float() + for query_pos in range(qlen): + visible = global_seq_len - qlen + query_pos + 1 + keys = kv_dequant[:visible] + scores = torch.einsum( + "hd,nd->hn", q_nope_fp32[query_pos], keys[:, :kv_lora_rank] + ) + scores += torch.einsum( + "hd,nd->hn", q_pe_fp32[query_pos], keys[:, kv_lora_rank:] + ) + probs = torch.softmax(scores * sm_scale, dim=-1) + reference[query_pos] = probs @ keys[:, :kv_lora_rank] + + torch.testing.assert_close(output, reference, rtol=3e-2, atol=3e-2) + + @pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128]) @pytest.mark.parametrize( "spec_method, parallel_drafting", @@ -120,7 +482,7 @@ def test_backend_declares_uniform_batch_support(): # and a parallel one, so the threshold is 1 + 2 * num_spec rather than # 1 + num_spec. Sizing the metadata off a method name instead leaves # these at qlen=1 while the router still admits the full range. - ("dspark", True), + ("custom", True), ("eagle", False), ], ) @@ -161,6 +523,8 @@ def test_mtp_builder_init_sizes_native_fp8_metadata( def init_common_builder(self, *args, **kwargs): self.num_heads = num_heads + self.dcp_world_size = 1 + self.cp_kv_cache_interleave_size = 1 # Mirror what _init_reorder_batch_threshold would have left behind: the # metadata is sized from the routing threshold, so a stub that skips it # 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( num_speculative_tokens=3, parallel_drafting=parallel_drafting, ), - parallel_config=SimpleNamespace(tensor_parallel_size=8), - model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16), + parallel_config=SimpleNamespace( + tensor_parallel_size=8, + decode_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + model_config=SimpleNamespace( + max_model_len=16, + dtype=torch.bfloat16, + get_num_attention_heads=lambda parallel_config: num_heads, + ), scheduler_config=SimpleNamespace(max_num_seqs=2), cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"), compilation_config=SimpleNamespace( - cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False) + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False), + # Empty: the per-layer head-count probe finds no attention layer and + # falls back to the model config above. + static_forward_context={}, ), ) builder = AiterMLAMetadataBuilder( diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9072c71426..280dec2d31 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -900,23 +900,17 @@ class RocmPlatform(Platform): compilation_config = vllm_config.compilation_config parallel_config = vllm_config.parallel_config - if compilation_config.cudagraph_mode.has_full_cudagraphs(): - # decode context parallel does not support full cudagraphs - if parallel_config.decode_context_parallel_size > 1: - logger.warning_once( - "Decode context parallel (DCP) is enabled, which is " - "incompatible with full CUDA graphs. " - "Overriding cudagraph_mode to PIECEWISE." - ) - compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE + if ( + compilation_config.cudagraph_mode.has_full_cudagraphs() + and parallel_config.prefill_context_parallel_size > 1 + ): # prefill context parallel do not support full cudagraphs - elif parallel_config.prefill_context_parallel_size > 1: - logger.warning_once( - "Prefill context parallel (PCP) is enabled, which is " - "incompatible with full CUDA graphs. " - "Overriding cudagraph_mode to PIECEWISE." - ) - compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE + logger.warning_once( + "Prefill context parallel (PCP) is enabled, which is " + "incompatible with full CUDA graphs. " + "Overriding cudagraph_mode to PIECEWISE." + ) + compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE if parallel_config.worker_cls == "auto": parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker" diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 7b1263dedd..ddf8bbd54e 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -11,6 +11,7 @@ import torch from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.config.cache import CacheDType +from vllm.distributed import get_dcp_group from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( MLACommonBackend, @@ -21,17 +22,39 @@ from vllm.model_executor.layers.attention.mla_attention import ( QueryLenSupport, ) from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor from vllm.v1.attention.backend import ( AttentionCGSupport, AttentionLayer, CommonAttentionMetadata, MultipleOf, ) +from vllm.v1.attention.backends.utils import ( + get_dcp_local_seq_lens, +) +from vllm.v1.attention.ops.rocm_aiter_mla_merge import ( + merge_mla_segments_triton, +) from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache logger = init_logger(__name__) +def _segmented_mla_page_size(block_size: int) -> int: + """Largest supported power-of-two subpage dividing a physical KV block. + + The subpage size becomes the segmented kernel's ``TILE_SIZE``: the cache is + reinterpreted as ``(-1, page_size, 1, head_dim)`` before the call. It must + divide the physical block and be a power of two, since the kernel walks a + tile with ``tl.arange(0, TILE_SIZE)``. 128 is the largest tile that tiling + supports -- a tile holds ``TILE_SIZE x kv_lora_rank`` keys and its scores + are ``BLOCK_M x TILE_SIZE`` -- so a larger manager block is split into + several subpages instead of widening the tile. + """ + assert block_size > 0 + return min(128, largest_power_of_2_divisor(block_size)) + + @functools.lru_cache(maxsize=1) def _get_mla_gluon(): """Load the small-head Gluon MLA entry point.""" @@ -137,6 +160,43 @@ def _gluon_mla_decode_supported() -> bool: return on_gfx950() +@functools.lru_cache(maxsize=1) +def _get_segmented_mla_decode(): + """Load AITER's segmented MLA decode with unreduced partial output.""" + from aiter.ops.triton.attention.mla import mla_decode_fwd + + return mla_decode_fwd + + +@functools.lru_cache(maxsize=1) +def _get_aiter_mla_decode(): + from aiter.mla import mla_decode_fwd + + return mla_decode_fwd + + +@functools.lru_cache(maxsize=1) +def _segmented_mla_decode_supported() -> bool: + """Whether AITER exposes the segmented MLA decode used by DCP verify.""" + try: + _get_segmented_mla_decode() + except Exception: # noqa: BLE001 + return False + return True + + +def _segmented_dcp_verify_supported(dcp_world_size: int, cp_interleave: int) -> bool: + """Whether this configuration can serve DCP verify on segmented MLA. + + Configuration only -- whether a given batch takes the route additionally + depends on its query length. Round-robin interleaving other than 1 is + excluded because the per-row causal window has not been validated there. + """ + return ( + dcp_world_size > 1 and cp_interleave == 1 and _segmented_mla_decode_supported() + ) + + def _aiter_mla_small_head_mode() -> str: """Small-head (<16) MLA decode kernel selection. @@ -200,6 +260,31 @@ class AiterMLABackend(MLACommonBackend): return AiterMLAMetadataBuilder +@dataclass +class AiterMLADCPVerifyMetadata: + """One paged-KV row per verify token, for segmented DCP verification. + + These are produced together or not at all, so they travel as one value: its + presence on the decode metadata *is* the routing decision the builder made, + and the impl does not re-derive it. + """ + + # Local KV length of each verify row, already causally bounded. + row_lens: torch.Tensor + # Subpage IDs each row reads, shape [num_rows, max_local_pages]. + block_table: torch.Tensor + # One query per row, so this is an arange. + qo_indptr: torch.Tensor + # Subpage size the block table is expressed in, and the segmented kernel's + # TILE_SIZE. Carried so the impl reinterprets the cache exactly the way the + # block table was built, instead of re-deriving it from the cache shape. + page_size: int + # Kernel-visible KV bound. Whenever full graphs are enabled this is the + # configuration's maximum for every batch, not just during capture, so the + # page table keeps one shape across replays. + max_kv_seq_len: int + + @dataclass class AiterMLADecodeMetadata(MLACommonDecodeMetadata): # The indptr of the paged kv cache, shape: [batch_size + 1] @@ -217,8 +302,12 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): max_qo_len: int | None = None # Minimum KV length used by Gluon to choose a safe split count. min_kv_seq_len: int = 1 + # Set exactly when this batch routes to segmented DCP verification. + dcp_verify: AiterMLADCPVerifyMetadata | None = None # Small-head decode uses Gluon (avoids padding to 16). use_gluon_decode: bool = False + # Small-head non-DCP multi-token verify uses Gluon's native MTP entry. + use_gluon_verify: bool = False # Whether persistent MLA metadata was computed has_persistent_metadata: bool = False @@ -305,13 +394,27 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): vllm_config: VllmConfig, device: torch.device, ): - super().__init__( - kv_cache_spec, layer_names, vllm_config, device, AiterMLAMetadata + parallel_config = vllm_config.parallel_config + supports_segmented_dcp_verify = _segmented_dcp_verify_supported( + parallel_config.decode_context_parallel_size, + parallel_config.cp_kv_cache_interleave_size, ) + super().__init__( + kv_cache_spec, + layer_names, + vllm_config, + device, + AiterMLAMetadata, + supports_dcp_with_varlen=supports_segmented_dcp_verify, + ) + self._supports_segmented_dcp_verify = supports_segmented_dcp_verify self.compilation_config = vllm_config.compilation_config self.decode_attn_out_dtype = vllm_config.model_config.dtype + # Needed to place a verify row's causal window on this rank's KV shard. + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 + # reorder_batch_threshold is the largest query length decode can be # handed, and already accounts for the drafting scheme. A method-name # whitelist sizes unlisted drafters for qlen=1, which closes the @@ -323,6 +426,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # (e.g. 16 with Eagle3), we expand block-level indices into per-token # flat indices since the aiter kernel always uses page_size=1 internally. self.kernel_block_size = kv_cache_spec.block_size + self._segmented_page_size = _segmented_mla_page_size(self.kernel_block_size) # In the flat view (.view(-1,1,1,H)), each token is its own page, # so max_num_pages_per_req = max_model_len regardless of @@ -351,10 +455,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): from aiter import dtypes, get_mla_metadata_info_v1 + # Decode kernels consume the DCP-gathered query heads. + self._decode_num_heads = self.num_heads * self.dcp_world_size # Keep metadata sizing consistent with the padded tensor shape passed - # to mla_decode_fwd. + # to mla_decode_fwd, including native 24-head AITER builds. self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads( - self.num_heads + self._decode_num_heads ) kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): @@ -437,6 +543,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): device, ) + # Persistent buffers for segmented DCP verification. Captured graphs + # require stable addresses while row lengths vary between replays. + self._dcp_verify_buffers: AiterMLADCPVerifyMetadata | None = None + self._graph_seq_lens: torch.Tensor | None = None if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): self.paged_kv_indptr = torch.zeros( max_num_reqs + 1, dtype=torch.int32, device=device @@ -446,6 +556,39 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): max_num_reqs + 1, dtype=torch.int32, device=device ) + # Full graphs require a stable address after uniform-MTP padding. + self._graph_seq_lens = torch.zeros( + max_num_reqs, dtype=torch.int32, device=device + ) + + if self._supports_segmented_dcp_verify and self._mtp_decode_qlen > 1: + # A DCP rank's shard of the longest sequence bounds every verify + # row, and full graphs need that bound to be constant. + num_dcp_partitions = ( + self.dcp_world_size * self.cp_kv_cache_interleave_size + ) + graph_max_kv_seq_len = ( + cdiv(vllm_config.model_config.max_model_len, num_dcp_partitions) + * self.cp_kv_cache_interleave_size + ) + max_verify_rows = max_num_reqs * self._mtp_decode_qlen + max_local_pages = cdiv(graph_max_kv_seq_len, self._segmented_page_size) + self._dcp_verify_buffers = AiterMLADCPVerifyMetadata( + row_lens=torch.zeros( + max_verify_rows, dtype=torch.int32, device=device + ), + block_table=torch.zeros( + (max_verify_rows, max_local_pages), + dtype=torch.int32, + device=device, + ), + qo_indptr=torch.arange( + max_verify_rows + 1, dtype=torch.int32, device=device + ), + page_size=self._segmented_page_size, + max_kv_seq_len=graph_max_kv_seq_len, + ) + def _init_fp8_prefill_ps_buffers( self, max_num_reqs: int, @@ -667,6 +810,99 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): metadata.fp8_prefill_max_q_len = prefill.max_query_len metadata.fp8_prefill_num_partial_tiles = num_partial_tiles + def _build_dcp_verify_row_view( + self, + qlen: int, + block_table: torch.Tensor, + dcp_tot_seq_lens: torch.Tensor, + ) -> AiterMLADCPVerifyMetadata: + """Build one paged-KV row per verify token for segmented DCP verification. + + Every row is a single query (``qo_indptr`` is an arange), so the whole + causal structure is carried by ``dcp_local_verify_row_lens`` and the + kernel applies no tail of its own. + """ + assert self.dcp_world_size > 1 + num_reqs = dcp_tot_seq_lens.numel() + row_lens = AiterMLAHelper.dcp_local_verify_row_lens( + dcp_tot_seq_lens, + qlen, + self.dcp_world_size, + self.dcp_rank, + self.cp_kv_cache_interleave_size, + ) + num_rows = row_lens.numel() + buffers = self._dcp_verify_buffers + # Persistent buffers exist exactly when full graphs are on, and they + # already carry the static bound they were sized for. + max_kv_seq_len = ( + buffers.max_kv_seq_len + if buffers is not None + else max(1, int(row_lens.max().item())) + ) + page_size = self._segmented_page_size + max_local_pages = cdiv(max_kv_seq_len, page_size) + if buffers is not None: + row_block_table = buffers.block_table[:num_rows] + buffers.row_lens[:num_rows].copy_(row_lens, non_blocking=True) + row_lens = buffers.row_lens[:num_rows] + qo_indptr = buffers.qo_indptr[: num_rows + 1] + else: + row_block_table = torch.empty( + (num_rows, max_local_pages), + dtype=torch.int32, + device=block_table.device, + ) + qo_indptr = torch.arange( + num_rows + 1, + dtype=torch.int32, + device=block_table.device, + ) + self._fill_dcp_verify_page_table(row_block_table, block_table, num_reqs, qlen) + return AiterMLADCPVerifyMetadata( + row_lens=row_lens, + block_table=row_block_table, + qo_indptr=qo_indptr, + page_size=page_size, + max_kv_seq_len=max_kv_seq_len, + ) + + def _fill_dcp_verify_page_table( + self, + row_block_table: torch.Tensor, + block_table: torch.Tensor, + num_reqs: int, + qlen: int, + ) -> None: + """Expand each request's blocks into the subpages one verify row reads. + + Every row of a request shares the request's shard, so the page list is + built once per request and repeated; only ``row_lens`` distinguishes the + rows. A DCP rank holds ``1/dcp_world_size`` of the sequence, so the + request's block table is always wider than the pages a row can reach. + """ + pages_per_block = self.kernel_block_size // self._segmented_page_size + max_local_pages = row_block_table.shape[1] + max_local_blocks = cdiv(max_local_pages, pages_per_block) + assert max_local_blocks <= block_table.shape[1], ( + f"DCP verify needs {max_local_blocks} blocks per request but the " + f"block table only has {block_table.shape[1]}" + ) + subpage_offsets = torch.arange( + pages_per_block, + dtype=block_table.dtype, + device=block_table.device, + ) + per_req_page_table = ( + block_table[:num_reqs, :max_local_blocks, None] * pages_per_block + + subpage_offsets + ).flatten(1)[:, :max_local_pages] + # Broadcast the request's page list across its rows instead of + # materializing a repeat_interleave copy. + row_block_table.unflatten(0, (num_reqs, qlen)).copy_( + per_req_page_table.unsqueeze(1) + ) + def _build_decode( self, block_table_tensor: torch.Tensor, @@ -701,6 +937,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): seq_lens_for_kernel.new_full((), max_qo_len), ) + if self._graph_seq_lens is not None: + self._graph_seq_lens[:num_kernel_reqs].copy_( + seq_lens_for_kernel, non_blocking=True + ) + seq_lens_for_kernel = self._graph_seq_lens[:num_kernel_reqs] + # The aiter kernel always operates with page_size=1 (the wrapper # flattens kv_buffer). last_page_len is always 1. paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs] @@ -713,26 +955,43 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): ] ) use_gluon_decode = AiterMLAHelper.use_gluon_decode( - self.num_heads, int(max_qo_len), self._kv_cache_dtype_str + self._decode_num_heads, + int(max_qo_len), + self._kv_cache_dtype_str, + ) + use_gluon_verify = AiterMLAHelper.use_gluon_verify( + self._decode_num_heads, + int(max_qo_len), + self._kv_cache_dtype_str, + self.dcp_world_size, + ) + use_segmented_dcp_verify = ( + self._supports_segmented_dcp_verify and max_qo_len > 1 ) - if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): - self.paged_kv_indices.fill_(-1) + # Segmented DCP verify carries its own per-row subpage table, so the + # flat per-token view is dead work for it. Leave the buffer alone and + # hand the metadata None, so a future reader cannot pick up whatever + # the previous batch left behind. + paged_kv_indices = None + if not use_segmented_dcp_verify: + if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): + self.paged_kv_indices.fill_(-1) - # Expand block_table entries into per-token flat indices. - # When kernel_block_size=1, this degrades to a direct copy (identical - # to the original _copy_page_indices_kernel). - # When kernel_block_size=K>1, block_table entry b covering K tokens - # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). - _expand_page_indices_kernel[(num_reqs,)]( - self.paged_kv_indices, - block_table_tensor, - block_table_tensor.stride(0), - paged_kv_indptr, - KERNEL_BLOCK_SIZE=self.kernel_block_size, - BLOCK_SIZE=1024, - ) - paged_kv_indices = self.paged_kv_indices + # Expand block_table entries into per-token flat indices. + # When kernel_block_size=1, this degrades to a direct copy (identical + # to the original _copy_page_indices_kernel). + # When kernel_block_size=K>1, block_table entry b covering K tokens + # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). + _expand_page_indices_kernel[(num_reqs,)]( + self.paged_kv_indices, + block_table_tensor, + block_table_tensor.stride(0), + paged_kv_indptr, + KERNEL_BLOCK_SIZE=self.kernel_block_size, + BLOCK_SIZE=1024, + ) + paged_kv_indices = self.paged_kv_indices if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_( @@ -745,12 +1004,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # (set above), so no copy needed - buffer is always 1s. if pad_uniform_mtp: - qo_indptr_src = torch.arange( - 0, - (num_kernel_reqs + 1) * max_qo_len, - step=max_qo_len, - dtype=torch.int32, - device=device, + qo_indptr_src = AiterMLAHelper.qo_indptr_for_uniform_qlen( + num_kernel_reqs, int(max_qo_len), device ) else: qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs] @@ -762,42 +1017,32 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): else: if max_qo_len == 1: - qo_indptr = torch.arange( - 0, - num_kernel_reqs + 1, - step=1, - dtype=torch.int32, - device=device, + qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen( + num_kernel_reqs, 1, device ) else: if pad_uniform_mtp: - qo_indptr = torch.arange( - 0, - (num_kernel_reqs + 1) * max_qo_len, - step=max_qo_len, - dtype=torch.int32, - device=device, + qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen( + num_kernel_reqs, int(max_qo_len), device ) else: qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs] + has_persistent_metadata = False # Only the asm decode consumes the schedule, so gate on the routing # rather than on num_heads >= 16, which denies it to a padded rank - # running the same asm kernels. The two predicates are disjoint -- - # decode is qlen==1, verify is qlen>1 -- and cover both Gluon entries. - has_persistent_metadata = False + # running the same asm kernels. The predicates are disjoint -- decode + # is qlen==1, verify is qlen>1 -- and cover both Gluon entries plus the + # segmented DCP verify. use_persistent_metadata = ( - not AiterMLAHelper.use_gluon_decode( - self.num_heads, max_qo_len, self._kv_cache_dtype_str - ) - and not AiterMLAHelper.use_gluon_verify( - self.num_heads, max_qo_len, self._kv_cache_dtype_str - ) + not use_gluon_decode + and not use_gluon_verify + and not use_segmented_dcp_verify # A padded rank has no bf16 persistent kernel past qlen 4 where the # gfx950 fold is absent; the non-persistent entry covers it. fp8 # keeps the schedule -- its fold rejects non-persistent outright. and ( - self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS + self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN or is_quantized_kv_cache(self._kv_cache_dtype_str) ) @@ -840,11 +1085,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # active requests, not cudagraph padding rows pinned to max_qo_len. # The .item() below runs in the builder, outside the captured region, # so it does not abort HIP graph capture the way the per-layer syncs - # in forward_mqa did. + # in forward_mqa did. DCP verify never reaches here -- it carries its + # own per-row lengths instead of a single split bound. min_kv_seq_len = 1 - if AiterMLAHelper.use_gluon_verify( - self.num_heads, int(max_qo_len), self._kv_cache_dtype_str - ): + if use_gluon_verify: per_req_len = paged_kv_indptr[1:] - paged_kv_indptr[:-1] if pad_uniform_mtp: active = qo_lens_device > 0 @@ -853,6 +1097,15 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): else: min_kv_seq_len = int(per_req_len.min().item()) + dcp_verify = None + if use_segmented_dcp_verify: + assert dcp_tot_seq_lens_device is not None + dcp_verify = self._build_dcp_verify_row_view( + int(max_qo_len), + block_table_tensor, + dcp_tot_seq_lens_device, + ) + attn_metadata = AiterMLADecodeMetadata( block_table=block_table_tensor, seq_lens=seq_lens_for_kernel, @@ -863,7 +1116,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): dcp_tot_seq_lens=dcp_tot_seq_lens_device, max_qo_len=max_qo_len, min_kv_seq_len=min_kv_seq_len, + dcp_verify=dcp_verify, use_gluon_decode=use_gluon_decode, + use_gluon_verify=use_gluon_verify, attn_out_dtype=self.decode_attn_out_dtype, has_persistent_metadata=has_persistent_metadata, ) @@ -966,6 +1221,22 @@ class AiterMLAHelper: _ASM_PADDED_MAX_PS_QLEN: Final = 4 _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = () + @staticmethod + def qo_indptr_for_uniform_qlen( + num_reqs: int, + qlen: int, + device: torch.device, + dtype: torch.dtype = torch.int32, + ) -> torch.Tensor: + """Build ``[0, qlen, 2*qlen, ..., num_reqs*qlen]``.""" + return torch.arange( + 0, + (num_reqs + 1) * qlen, + step=qlen, + dtype=dtype, + device=device, + ) + @staticmethod def check_num_heads_validity(num_heads: int): assert AiterMLAHelper.is_valid_num_heads(num_heads), ( @@ -1042,14 +1313,22 @@ class AiterMLAHelper: @staticmethod def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor: + return AiterMLAHelper._get_mla_unpadded_heads(num_heads, o) + + @staticmethod + def _get_mla_unpadded_heads(num_heads: int, tensor: torch.Tensor) -> torch.Tensor: m = AiterMLAHelper.get_actual_mla_num_heads(num_heads) if num_heads == m: - return o + return tensor if m % num_heads == 0: - return o[:, :: m // num_heads, :] + return tensor[:, :: m // num_heads, ...] # Undo the tile-padding from get_mla_padded_q: the real heads are the # first num_heads. - return o[:, :num_heads, :] + return tensor[:, :num_heads, ...] + + @staticmethod + def get_mla_unpadded_lse(num_heads: int, lse: torch.Tensor) -> torch.Tensor: + return AiterMLAHelper._get_mla_unpadded_heads(num_heads, lse) @staticmethod def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool: @@ -1074,7 +1353,12 @@ class AiterMLAHelper: return m % num_heads == 0 and gluon_supported @staticmethod - def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool: + def use_gluon_verify( + num_heads: int, + max_qo_len: int, + kv_cache_dtype: str, + dcp_world_size: int = 1, + ) -> bool: """Whether a small-head multi-token verify uses native Gluon MTP. bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through @@ -1083,16 +1367,67 @@ class AiterMLAHelper: come here: the MTP path hands Gluon the batch size its fp8 regime asserts against. A predicate rather than inline in forward_mqa so the builder sees the same answer the impl acts on. + + DCP verify is excluded: its per-row causal windows are served by the + segmented path, which Gluon's MTP entry cannot express. """ - if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1: + if max_qo_len <= 1 or dcp_world_size > 1: return False if is_quantized_kv_cache(kv_cache_dtype): return False + if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS: + return False + if not _gluon_mla_decode_supported(): + return False # Same arch and mode gating as use_gluon_decode. - return _aiter_mla_small_head_mode() != "asm" and _gluon_mla_decode_supported() + return _aiter_mla_small_head_mode() != "asm" + + @staticmethod + def dcp_local_verify_row_lens( + tot_seq_lens: torch.Tensor, + qlen: int, + dcp_world_size: int, + dcp_rank: int, + cp_interleave: int, + ) -> torch.Tensor: + """Local KV length of every verify row, in global-causal order. + + Row ``t`` of a qlen-token verify attends global positions + ``[0, seq_len - qlen + t]``, its own token included, so its local length + is the round-robin count evaluated at that bound. Truncating causally in + global coordinates has to happen before the shard mapping: the tokens a + row drops sit on ``qlen - 1 - t`` different ranks, while subtracting the + offset from a request's local length takes one from every rank. + """ + offsets = torch.arange( + 1, + qlen + 1, + device=tot_seq_lens.device, + dtype=tot_seq_lens.dtype, + ) + visible = (tot_seq_lens.unsqueeze(1) - qlen + offsets).clamp_(min=0) + return get_dcp_local_seq_lens( + visible, + dcp_world_size, + dcp_rank, + cp_interleave, + ).flatten() class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): + # DCP decode paths return natural-log softmax LSE for the cross-rank merge. + can_return_lse_for_decode: bool = True + # Measured on gfx950: aiter.mla.mla_decode_fwd(return_lse=True) matches + # logsumexp to fp32 exactly, and merge_mla_segments_triton converts AITER's + # base-2 segment statistics with LOGE2. Stated rather than inherited because + # the DCP combine silently corrupts the softmax denominator if it disagrees. + lse_base_on_e: bool = True + + @property + def _decode_num_heads(self) -> int: + """Return the query-head count after DCP gathering.""" + return self.num_heads * self.dcp_world_size + def __init__( self, num_heads: int, @@ -1122,6 +1457,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): **mla_args, ) AiterMLAHelper.check_num_heads_validity(num_heads) + AiterMLAHelper.check_num_heads_validity(self._decode_num_heads) unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] if any(unsupported_features): @@ -1363,6 +1699,59 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): self._mla_fp8_prefill_attn(q, k, v, attn_metadata, output) + def _forward_segmented_dcp_verify( + self, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + verify: AiterMLADCPVerifyMetadata, + kv_c_and_k_pe_cache: torch.Tensor, + layer: AttentionLayer, + out_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run segmented attention over this rank's shard of every verify row. + + Each row's length already covers the current tokens this rank holds, so + the cross-rank LSE merge in the MLA layer completes the causal block. + """ + q_mla = torch.cat([q_nope, q_pe], dim=-1) + # skip_reduce with NUM_SEGMENTS>1 returns the partials and never writes + # `out`. Pass None rather than aliasing q_mla: if that ever stops + # holding, this fails loudly instead of scribbling over the query. + segment_partials = _get_segmented_mla_decode()( + q_mla, + kv_c_and_k_pe_cache.view( + -1, + verify.page_size, + 1, + kv_c_and_k_pe_cache.shape[-1], + ), + None, + verify.qo_indptr, + verify.row_lens, + verify.max_kv_seq_len, + verify.block_table, + self.scale, + self.kv_lora_rank, + self.qk_rope_head_dim, + causal=True, + q_descale=None, + kv_descale=layer._k_scale, + skip_reduce=True, + ) + assert isinstance(segment_partials, tuple) and len(segment_partials) == 3, ( + "AITER segmented MLA verify must return segment partials " + "when skip_reduce=True." + ) + segm_output, segm_max, segm_expsum = segment_partials + return merge_mla_segments_triton( + segm_output, + segm_max, + segm_expsum, + verify.row_lens, + verify.page_size, + out_dtype, + ) + def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -1375,9 +1764,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): decode = attn_metadata.decode assert decode.max_qo_len is not None - assert decode.paged_kv_indptr is not None - assert decode.paged_kv_indices is not None if decode.use_gluon_decode: + assert decode.paged_kv_indptr is not None + assert decode.paged_kv_indices is not None if type(q) is tuple: q_nope, q_pe = q else: @@ -1416,12 +1805,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): # The block is causal -- the target is checking draft tokens, so # position t must not see t+1 -- and the kernel bounds each query # position's scores itself. - # Arch, mode and dtype gating all live in use_gluon_verify, so that the - # builder -- which has to know whether the asm decode will run -- sees - # the same answer as this branch. - if AiterMLAHelper.use_gluon_verify( - self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype - ): + # Arch, mode and dtype gating all live in use_gluon_verify, which the + # builder evaluates and records on the metadata, so this branch acts on + # the same answer the builder used to decide whether the asm decode + # would run. DCP verify is routed to the segmented path below instead. + if decode.use_gluon_verify: if type(q) is tuple: q_nope, q_pe = q else: @@ -1473,14 +1861,48 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): ) return o, None + verify = decode.dcp_verify + if verify is not None: + if type(q) is tuple: + q_nope, q_pe = q + else: + q_nope, q_pe = torch.split( + q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + if ( + is_quantized_kv_cache(self.kv_cache_dtype) + and q_nope.dtype != torch.bfloat16 + ): + q_nope = q_nope.to(torch.bfloat16) * layer._q_scale + q_pe = q_pe.to(torch.bfloat16) * layer._q_scale + return self._forward_segmented_dcp_verify( + q_nope, + q_pe, + verify, + kv_c_and_k_pe_cache, + layer, + decode.attn_out_dtype, + ) + + if self.dcp_world_size > 1 and int(decode.max_qo_len) > 1: + raise RuntimeError( + "ROCM_AITER_MLA DCP multi-token verify requires segmented MLA." + ) + if type(q) is tuple: q = torch.cat(q, dim=-1) assert isinstance(q, torch.Tensor) + assert decode.paged_kv_indptr is not None + assert decode.paged_kv_indices is not None B = q.shape[0] - mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q) - mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads) + assert q.shape[1] == self._decode_num_heads, ( + "ROCM_AITER_MLA decode expected the DCP-gathered query head count " + f"{self._decode_num_heads}, got {q.shape[1]}" + ) + mla_padded_q = AiterMLAHelper.get_mla_padded_q(self._decode_num_heads, q) + mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads) o = torch.empty( B, mla_num_heads, @@ -1512,17 +1934,43 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): reduce_partial_map=attn_metadata.reduce_partial_map, ) - rocm_aiter_ops.mla_decode_fwd( - mla_padded_q, - kv_buffer, - o, - self.scale, - decode.qo_indptr, - decode.max_qo_len, - decode.paged_kv_indptr, - decode.paged_kv_indices, - decode.paged_kv_last_page_len, - **mla_kwargs, - ) + lse = None + if self.dcp_world_size > 1: + # The vLLM custom-op wrapper exposes only the in-place output and + # drops aiter's final LSE, which the cross-shard merge needs, so go + # through aiter's native entry point on the DCP path. + _, lse = _get_aiter_mla_decode()( + mla_padded_q, + kv_buffer.view(-1, 1, 1, mla_padded_q.shape[-1]), + o, + decode.qo_indptr, + decode.paged_kv_indptr, + decode.paged_kv_indices, + decode.paged_kv_last_page_len, + decode.max_qo_len, + sm_scale=self.scale, + return_lse=True, + **mla_kwargs, + ) + assert lse is not None, ( + "aiter mla_decode_fwd(return_lse=True) returned no LSE; upgrade " + "aiter to a build with decode LSE support." + ) + else: + rocm_aiter_ops.mla_decode_fwd( + mla_padded_q, + kv_buffer, + o, + self.scale, + decode.qo_indptr, + decode.max_qo_len, + decode.paged_kv_indptr, + decode.paged_kv_indices, + decode.paged_kv_last_page_len, + **mla_kwargs, + ) - return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, o), None + output = AiterMLAHelper.get_mla_unpadded_o(self._decode_num_heads, o) + if lse is not None: + lse = AiterMLAHelper.get_mla_unpadded_lse(self._decode_num_heads, lse) + return output, lse diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_merge.py b/vllm/v1/attention/ops/rocm_aiter_mla_merge.py new file mode 100644 index 0000000000..22774464f8 --- /dev/null +++ b/vllm/v1/attention/ops/rocm_aiter_mla_merge.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Merge AITER segmented MLA split-K partials into output plus natural-log LSE. + +This mirrors AITER's own segment reduction, including its +``tiles_per_segment = cdiv(seq_len, NUM_SEGMENTS * TILE_SIZE)`` partitioning, so +it has to move together with the ``skip_reduce=True`` call in the AITER MLA +backend. It is a rank-local split-K merge and unrelated to any collective +reduce; the natural-log LSE it returns is what the cross-rank DCP merge +consumes. +""" + +import torch + +from vllm.triton_utils import LOGE2, tl, triton + + +@triton.jit +def _merge_mla_segments_kernel( + out_ptr, + lse_ptr, + segm_output_ptr, + segm_max_ptr, + segm_expsum_ptr, + seq_lens_ptr, + num_query_heads: tl.constexpr, + out_stride0: tl.int64, + out_stride1: tl.int64, + lse_stride0: tl.int64, + TILE_SIZE: tl.constexpr, + KV_LORA_RANK: tl.constexpr, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, + LOGE2: tl.constexpr, +): + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + + seq_len = tl.load(seq_lens_ptr + token_idx) + tiles_per_segment = tl.maximum( + 1, tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + ) + active_segments = tl.where( + seq_len > 0, + tl.cdiv(seq_len, tiles_per_segment * TILE_SIZE), + 0, + ) + segment_offsets = tl.arange(0, NUM_SEGMENTS_PER_SEQ) + segment_mask = segment_offsets < active_segments + + segment_base = ( + token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + head_idx * NUM_SEGMENTS_PER_SEQ + ) + segment_indices = segment_base + segment_offsets + segment_max = tl.load( + segm_max_ptr + segment_indices, + mask=segment_mask, + other=float("-inf"), + ) + overall_max = tl.max(segment_max) + + segment_expsum = tl.load( + segm_expsum_ptr + segment_indices, + mask=segment_mask, + other=0.0, + ) + segment_weight = tl.where( + segment_mask, + tl.exp2(segment_max - overall_max), + 0.0, + ) + overall_expsum = tl.sum(segment_expsum * segment_weight) + + output_offsets = ( + token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK) + + head_idx * (NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK) + + segment_offsets[:, None] * KV_LORA_RANK + + tl.arange(0, KV_LORA_RANK)[None, :] + ) + segment_output = tl.load( + segm_output_ptr + output_offsets, + mask=segment_mask[:, None], + other=0.0, + ) + accumulator = tl.sum(segment_output * segment_weight[:, None], axis=0) + accumulator = tl.where(overall_expsum == 0.0, 0.0, accumulator / overall_expsum) + lse = tl.where( + overall_expsum == 0.0, + float("-inf"), + (overall_max + tl.log2(overall_expsum)) * LOGE2, + ) + + output_indices = ( + token_idx * out_stride0 + head_idx * out_stride1 + tl.arange(0, KV_LORA_RANK) + ) + tl.store(out_ptr + output_indices, accumulator) + tl.store(lse_ptr + token_idx * lse_stride0 + head_idx, lse) + + +def merge_mla_segments_triton( + segm_output: torch.Tensor, + segm_max: torch.Tensor, + segm_expsum: torch.Tensor, + seq_lens: torch.Tensor, + tile_size: int, + out_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Merge AITER base-2 segment partials into output and natural-log LSE.""" + num_tokens, num_heads, num_segments, kv_lora_rank = segm_output.shape + output = torch.empty( + (num_tokens, num_heads, kv_lora_rank), + dtype=out_dtype, + device=segm_output.device, + ) + lse = torch.empty( + (num_tokens, num_heads), + dtype=torch.float32, + device=segm_output.device, + ) + _merge_mla_segments_kernel[(num_tokens, num_heads)]( + output, + lse, + segm_output, + segm_max, + segm_expsum, + seq_lens, + num_query_heads=num_heads, + out_stride0=output.stride(0), + out_stride1=output.stride(1), + lse_stride0=lse.stride(0), + TILE_SIZE=tile_size, + KV_LORA_RANK=kv_lora_rank, + NUM_SEGMENTS_PER_SEQ=num_segments, + LOGE2=LOGE2, + ) + return output, lse