diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py new file mode 100644 index 0000000000..844bee0f71 --- /dev/null +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm AITER MLA tests", allow_module_level=True) + +from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402 +from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402 + AiterMLAMetadataBuilder, +) + + +class _NoOpTritonKernel: + def __getitem__(self, grid): + self.grid = grid + return self + + def __call__(self, *args, **kwargs): + pass + + +class _ExpandPageIndicesKernel: + def __getitem__(self, grid): + self.grid = grid + return self + + def __call__( + self, + page_indices, + block_table_tensor, + stride, + paged_kv_indptr, + seq_lens_for_kernel, + *, + KERNEL_BLOCK_SIZE, + BLOCK_SIZE, + ): + self.kernel_block_size = KERNEL_BLOCK_SIZE + for req_idx in range(self.grid[0]): + out_start = int(paged_kv_indptr[req_idx].item()) + seq_len = int(seq_lens_for_kernel[req_idx].item()) + for token_idx in range(seq_len): + block_id = int( + block_table_tensor[req_idx, token_idx // KERNEL_BLOCK_SIZE].item() + ) + page_indices[out_start + token_idx] = ( + block_id * KERNEL_BLOCK_SIZE + token_idx % KERNEL_BLOCK_SIZE + ) + + +def _builder( + *, + mtp_decode_qlen: int, + has_full_cudagraphs: bool = False, + kernel_block_size: int = 1, + max_decode_rows: int = 32, + num_heads: int = 16, +): + return SimpleNamespace( + device=torch.device("cpu"), + num_heads=num_heads, + paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32), + paged_kv_indices=torch.empty(1024, dtype=torch.int32), + paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32), + qo_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32), + compilation_config=SimpleNamespace( + cudagraph_mode=SimpleNamespace( + has_full_cudagraphs=lambda: has_full_cudagraphs + ) + ), + _mtp_decode_qlen=mtp_decode_qlen, + _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len), + _use_persistent_metadata=False, + kernel_block_size=kernel_block_size, + _num_attention_heads=16, + _mla_work_meta_data=torch.empty(1, dtype=torch.int32), + _mla_work_info_set=torch.empty(1, dtype=torch.int32), + _mla_work_indptr=torch.empty(1, dtype=torch.int32), + _mla_reduce_indptr=torch.empty(1, dtype=torch.int32), + _mla_reduce_final_map=torch.empty(1, dtype=torch.int32), + _mla_reduce_partial_map=torch.empty(1, dtype=torch.int32), + _mla_q_dtype=torch.bfloat16, + _mla_kv_dtype=torch.bfloat16, + decode_attn_out_dtype=torch.bfloat16, + ) + + +def test_backend_declares_uniform_batch_support(): + # UNIFORM/UNIFORM_BATCH is unconditional: MTP yields uniform qlen>1 and + # non-MTP yields qlen==1, both uniform batches. + assert ( + AiterMLAMetadataBuilder.query_len_support + == rocm_aiter_mla.QueryLenSupport.UNIFORM + ) + assert ( + AiterMLAMetadataBuilder._cudagraph_support + == rocm_aiter_mla.AttentionCGSupport.UNIFORM_BATCH + ) + + +@pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128]) +def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads): + """Aiter init passes real MTP qlen/dtypes to rocm_aiter_mla.py:349-355. + + Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the + metadata, covering the fp8 nhead=32 (TP4) fold path. + """ + + dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16") + info_calls = [] + + def get_mla_metadata_info_v1( + max_batch_size, + max_qo_len, + num_attention_heads, + q_dtype, + kv_dtype, + *, + is_sparse, + fast_mode, + ): + info_calls.append( + { + "max_batch_size": max_batch_size, + "max_qo_len": max_qo_len, + "num_attention_heads": num_attention_heads, + "q_dtype": q_dtype, + "kv_dtype": kv_dtype, + "is_sparse": is_sparse, + "fast_mode": fast_mode, + } + ) + return tuple((1, torch.int32) for _ in range(6)) + + def init_common_builder(self, *args, **kwargs): + self.num_heads = num_heads + + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace( + dtypes=dtypes, + get_mla_metadata_info_v1=get_mla_metadata_info_v1, + ), + ) + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, + "__init__", + init_common_builder, + ) + monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False) + + config = SimpleNamespace( + speculative_config=SimpleNamespace( + method="deepseek_mtp", + num_speculative_tokens=3, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=8), + model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16), + 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) + ), + ) + builder = AiterMLAMetadataBuilder( + kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16), + layer_names=["layer.0"], + vllm_config=config, + device=torch.device("cpu"), + ) + + assert info_calls == [ + { + "max_batch_size": config.scheduler_config.max_num_seqs, + "max_qo_len": config.speculative_config.num_speculative_tokens + 1, + "num_attention_heads": max(16, num_heads), + "q_dtype": dtypes.fp8, + "kv_dtype": dtypes.fp8, + "is_sparse": False, + "fast_mode": True, + } + ] + assert builder._mla_q_dtype == dtypes.fp8 + assert builder._mla_kv_dtype == dtypes.fp8 + + +def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch): + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + + metadata = AiterMLAMetadataBuilder._build_decode( + _builder(mtp_decode_qlen=4), + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([7, 5], dtype=torch.int32), + max_seq_len=7, + query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, 4, 8], dtype=torch.int32), + num_decode_tokens=8, + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == 4 + assert torch.equal(metadata.seq_lens, torch.tensor([7, 5], dtype=torch.int32)) + assert torch.equal(metadata.qo_indptr, torch.tensor([0, 4, 8], dtype=torch.int32)) + assert metadata.has_persistent_metadata + assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == 4 + assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == 4 + + +def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr( + monkeypatch, +): + """Full-CG zero-qo rows follow rocm_aiter_mla.py:608-657,717-759.""" + + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + + mtp_qlen = 4 + seq_lens = torch.tensor([7, 0], dtype=torch.int32) + qo_lens = torch.tensor([mtp_qlen, 0], dtype=torch.int32) + expected_seq_lens = torch.where(qo_lens > 0, seq_lens, mtp_qlen) + expected_paged_kv_indptr = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + expected_seq_lens.cumsum(dim=0, dtype=torch.int32), + ] + ) + expected_qo_indptr = torch.arange( + 0, + (seq_lens.numel() + 1) * mtp_qlen, + step=mtp_qlen, + dtype=torch.int32, + ) + + builder = _builder( + mtp_decode_qlen=mtp_qlen, + has_full_cudagraphs=True, + max_decode_rows=4, + ) + metadata = AiterMLAMetadataBuilder._build_decode( + builder, + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=seq_lens, + max_seq_len=int(seq_lens.max().item()), + query_start_loc_cpu=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32), + num_decode_tokens=seq_lens.numel() * mtp_qlen, + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == mtp_qlen + assert torch.equal(metadata.seq_lens, expected_seq_lens) + assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr) + assert torch.equal(metadata.qo_indptr, expected_qo_indptr) + assert torch.all( + builder.paged_kv_indptr[expected_paged_kv_indptr.numel() :] + == expected_paged_kv_indptr[-1] + ) + assert torch.all( + builder.qo_indptr[expected_qo_indptr.numel() :] == expected_qo_indptr[-1] + ) + assert metadata.has_persistent_metadata + assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == mtp_qlen + assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == mtp_qlen + + +def test_decode_expands_kernel_block_page_indices(monkeypatch): + """kernel_block_size>1 expands b -> b*K+offset at rocm_aiter_mla.py:696-704.""" + + expand_kernel = _ExpandPageIndicesKernel() + monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand_kernel) + # qlen==1 now takes the persistent-metadata path, which imports + # get_mla_metadata_v1 from aiter; mock it so this test needs no real kernel. + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()), + ) + + kernel_block_size = 2 + block_table = torch.tensor( + [ + [10, 11, 99], + [20, 21, 22], + ], + dtype=torch.int32, + ) + seq_lens = torch.tensor([3, 5], dtype=torch.int32) + expected_paged_kv_indptr = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + seq_lens.cumsum(dim=0, dtype=torch.int32), + ] + ) + expected_indices = torch.tensor( + [ + int(block_table[req_idx, token_idx // kernel_block_size].item()) + * kernel_block_size + + token_idx % kernel_block_size + for req_idx, seq_len in enumerate(seq_lens.tolist()) + for token_idx in range(seq_len) + ], + dtype=torch.int32, + ) + + metadata = AiterMLAMetadataBuilder._build_decode( + _builder( + mtp_decode_qlen=1, + kernel_block_size=kernel_block_size, + ), + block_table_tensor=block_table, + seq_lens_device=seq_lens, + max_seq_len=int(seq_lens.max().item()), + query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, 1, 2], dtype=torch.int32), + num_decode_tokens=seq_lens.numel(), + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == 1 + assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr) + assert torch.equal( + metadata.paged_kv_indices[: expected_indices.numel()], + expected_indices, + ) + assert expand_kernel.grid == (seq_lens.numel(),) + assert expand_kernel.kernel_block_size == kernel_block_size + + +@pytest.mark.parametrize( + "mtp_decode_qlen, qo_len, num_heads, expect_persistent", + [ + (1, 1, 16, True), # non-MTP decode + (4, 2, 16, True), # MTP deployment, in-range step + (4, 4, 16, True), # MTP deployment, full-qlen verification step + (2, 4, 16, False), # step demand exceeds provisioned K -> fallback + (1, 1, 8, False), # small head count -> Gluon decode owns qlen==1 + (4, 4, 8, False), # small head count -> Gluon flatten owns qlen>1 + ], +) +def test_persistent_metadata_gate( + monkeypatch, mtp_decode_qlen, qo_len, num_heads, expect_persistent +): + """Persistent metadata is passed iff num_heads >= 16 and 1 <= max_qo_len <= K. + + K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets + the pre-built schedule only when its qlen fits those buffers, otherwise it + falls back to the kernel computing its own. qlen==1 (non-MTP) must stay + in-range -- dropping it is the regression this guards. Fewer than 16 heads + is served by the Gluon decode paths, which never read this schedule. + """ + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + + # Uniform, non-CUDA-graph batch: every request has exactly qo_len tokens, so + # num_decode_tokens == sum(qo_len) and no dummy-row padding kicks in. + num_reqs = 2 + query_start_loc = torch.arange( + 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32 + ) + metadata = AiterMLAMetadataBuilder._build_decode( + _builder(mtp_decode_qlen=mtp_decode_qlen, num_heads=num_heads), + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([8, 8], dtype=torch.int32), + max_seq_len=8, + query_start_loc_cpu=query_start_loc, + query_start_loc_device=query_start_loc, + num_decode_tokens=num_reqs * qo_len, + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == qo_len + assert metadata.has_persistent_metadata is expect_persistent + assert get_mla_metadata_v1.called is expect_persistent + if expect_persistent: + assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == qo_len + assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == qo_len diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 4c7a90ed02..24c1c16540 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -132,7 +132,7 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): min_kv_seq_len: int = 1 # Small-head decode uses Gluon (avoids padding to 16). use_gluon_decode: bool = False - # Whether persistent MLA metadata was computed (only for qseqlen=1) + # Whether persistent MLA metadata was computed has_persistent_metadata: bool = False @@ -171,6 +171,46 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM + @staticmethod + def _uniform_padded_mtp_qo_len( + qo_len: torch.Tensor, + max_qo_len: int, + num_decode_tokens: int, + ) -> int: + num_reqs = qo_len.numel() + if num_reqs == 0 or num_decode_tokens <= 0: + return 0 + + # Full-CG pads q to a captured token count while leaving + # query_start_loc flat for dummy requests. Only synthesize dummy rows + # when every padded request maps to the same qlen and the q buffer has + # exactly that many rows. + if num_decode_tokens <= int(qo_len.sum().item()): + return 0 + if num_decode_tokens % num_reqs != 0: + return 0 + + uniform_qo_len = num_decode_tokens // num_reqs + if uniform_qo_len <= 1: + return 0 + + positive_qo_len = qo_len[qo_len > 0] + if positive_qo_len.numel() == qo_len.numel(): + return 0 + if positive_qo_len.numel() > 0: + if max_qo_len != uniform_qo_len: + return 0 + if not torch.all(positive_qo_len == uniform_qo_len): + return 0 + + zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten() + if zero_positions.numel() > 0: + first_zero = int(zero_positions[0].item()) + if torch.any(qo_len[first_zero:] > 0): + return 0 + + return uniform_qo_len + def __init__( self, kv_cache_spec: AttentionSpec, @@ -185,6 +225,18 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): self.compilation_config = vllm_config.compilation_config self.decode_attn_out_dtype = vllm_config.model_config.dtype + # MTP/deepseek_mtp verification runs decode with qlen = num_spec + 1; + # any other config (including no spec) stays at single-token decode. + speculative_config = vllm_config.speculative_config + if ( + speculative_config is not None + and speculative_config.method in ("mtp", "deepseek_mtp") + and speculative_config.num_speculative_tokens is not None + ): + self._mtp_decode_qlen = int(speculative_config.num_speculative_tokens) + 1 + else: + self._mtp_decode_qlen = 1 + # Store the kernel block size from the spec. When kernel_block_size=1 # (no spec-dec), behavior is identical to the original. When > 1 # (e.g. 16 with Eagle3), we expand block-level indices into per-token @@ -222,14 +274,20 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent # with the actual tensor shape passed to mla_decode_fwd. self._num_attention_heads = max(16, self.num_heads) - q_dtype = self.decode_attn_out_dtype kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): kv_cache_dtype_str = "fp8" - q_dtype = dtypes.fp8 + kv_dtype = dtypes.fp8 else: - kv_cache_dtype_str = "bf16" - kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16) + kv_dtype = { + torch.float16: dtypes.fp16, + torch.bfloat16: dtypes.bf16, + }[kv_cache_spec.dtype] + # MLAAttention quantizes decode Q to FP8 before calling this backend + # whenever the KV cache is FP8 and supports_quant_query_input is true. + q_dtype = ( + dtypes.fp8 if kv_cache_dtype_str == "fp8" else self.decode_attn_out_dtype + ) # Persist for get_mla_metadata_v1 (decode build): omitting these causes # wrong split/reduce metadata for the gfx950 fp8 nhead=32 fold path. self._mla_q_dtype = q_dtype @@ -243,7 +301,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): (reduce_partial_map_size, reduce_partial_map_type), ) = get_mla_metadata_info_v1( max_num_reqs, - 1, + self._mtp_decode_qlen, self._num_attention_heads, q_dtype, kv_dtype, @@ -488,20 +546,39 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): ) -> AiterMLADecodeMetadata: device = self.device num_reqs = seq_lens_device.size(0) + qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + max_qo_len = qo_len.max().item() + padded_mtp_qo_len = self._uniform_padded_mtp_qo_len( + qo_len, max_qo_len, num_decode_tokens + ) + if padded_mtp_qo_len > 0: + max_qo_len = padded_mtp_qo_len + pad_uniform_mtp = padded_mtp_qo_len > 0 + + seq_lens_for_kernel = seq_lens_device + num_kernel_reqs = num_reqs + if pad_uniform_mtp: + qo_lens_device = ( + query_start_loc_device[1 : num_reqs + 1] + - query_start_loc_device[:num_reqs] + ).to(torch.int32) + seq_lens_for_kernel = torch.where( + qo_lens_device > 0, + seq_lens_for_kernel, + seq_lens_for_kernel.new_full((), max_qo_len), + ) # 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_reqs] + paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs] # indptr: cumsum of seq_lens (one page per token in the flat view) paged_kv_indptr = torch.cat( [ - torch.zeros(1, dtype=seq_lens_device.dtype, device=device), - seq_lens_device.cumsum(dim=0, dtype=torch.int32), + torch.zeros(1, dtype=torch.int32, device=device), + seq_lens_for_kernel.cumsum(dim=0, dtype=torch.int32), ] ) - qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - max_qo_len = qo_len.max().item() use_gluon_decode = AiterMLAHelper.use_gluon_decode( self.num_heads, int(max_qo_len) ) @@ -519,44 +596,75 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): block_table_tensor, block_table_tensor.stride(0), paged_kv_indptr, - seq_lens_device, + seq_lens_for_kernel, 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_reqs].copy_( + self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_( paged_kv_indptr, non_blocking=True ) - self.paged_kv_indptr[1 + num_reqs :].fill_(paged_kv_indptr[-1]) - paged_kv_indptr = self.paged_kv_indptr[: 1 + num_reqs] + self.paged_kv_indptr[1 + num_kernel_reqs :].fill_(paged_kv_indptr[-1]) + paged_kv_indptr = self.paged_kv_indptr[: 1 + num_kernel_reqs] # paged_kv_last_page_len already uses the pre-initialized buffer slice # (set above), so no copy needed - buffer is always 1s. - self.qo_indptr[: 1 + num_reqs].copy_( - query_start_loc_device, non_blocking=True + 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, + ) + else: + qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs] + self.qo_indptr[: 1 + num_kernel_reqs].copy_( + qo_indptr_src, non_blocking=True ) - self.qo_indptr[1 + num_reqs :] = query_start_loc_device[-1] - qo_indptr = self.qo_indptr[: 1 + num_reqs] + self.qo_indptr[1 + num_kernel_reqs :] = qo_indptr_src[-1] + qo_indptr = self.qo_indptr[: 1 + num_kernel_reqs] else: - qo_indptr = torch.arange( - 0, num_reqs + 1, step=1, dtype=torch.int32, device=device - ) + if max_qo_len == 1: + qo_indptr = torch.arange( + 0, + num_kernel_reqs + 1, + step=1, + dtype=torch.int32, + device=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, + ) + else: + qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs] - # The aiter MLA ASM kernel only supports qseqlen=1 (single-token - # decode). With speculative decoding, the verification step has - # qseqlen > 1 (e.g. 8 for spec7). get_mla_metadata_v1 calls - # get_heuristic_kernel_mla which fails for qseqlen > 1. - # We track whether persistent metadata was successfully computed - # so forward_mqa can skip passing it (falling back to the kernel - # computing its own metadata internally, like v0.18.0). + # Pass persistent metadata for every uniform decode we sized buffers for + # (normal qlen==1 through MTP verification qlen==K): the fp8 nhead=32 fold + # path breaks without it. qlen>K falls back to kernel-internal metadata. + # Small-head (<16) decode takes the Gluon paths and never consumes it. has_persistent_metadata = False - if max_qo_len == 1 and not use_gluon_decode: + use_persistent_metadata = ( + self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS + and max_qo_len >= 1 + and max_qo_len <= self._mtp_decode_qlen + ) + if use_persistent_metadata: from aiter import get_mla_metadata_v1 + uni_qo_len = ( + max_qo_len if pad_uniform_mtp or torch.all(qo_len == max_qo_len) else -1 + ) get_mla_metadata_v1( qo_indptr, paged_kv_indptr, @@ -573,7 +681,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): page_size=1, kv_granularity=16, max_seqlen_qo=max_qo_len, - uni_seqlen_qo=max_qo_len, + uni_seqlen_qo=uni_qo_len, fast_mode=True, dtype_q=self._mla_q_dtype, dtype_kv=self._mla_kv_dtype, @@ -582,7 +690,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): attn_metadata = AiterMLADecodeMetadata( block_table=block_table_tensor, - seq_lens=seq_lens_device, + seq_lens=seq_lens_for_kernel, paged_kv_indptr=paged_kv_indptr, paged_kv_indices=paged_kv_indices, paged_kv_last_page_len=paged_kv_last_page_len, @@ -1076,13 +1184,16 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): dtype=attn_metadata.decode.attn_out_dtype, device=q.device, ) + if decode.max_qo_len > 1 and not decode.has_persistent_metadata: + # MTP verification can call the AITER MLA decode kernel with + # qlen > 1. If that path is running without persistent metadata, + # zero-fill so unwritten lanes cannot leak into logits. + o.zero_() kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) # Build kwargs for mla_decode_fwd. Pass persistent metadata only - # when it was successfully computed (qseqlen=1 decode steps). - # For multi-token verification steps (spec-dec), the kernel falls - # back to computing metadata internally. + # when it was successfully computed. mla_kwargs = dict( q_scale=layer._q_scale, kv_scale=layer._k_scale,