diff --git a/tests/models/kimi_k3/test_amd_kda_adaptive_verification.py b/tests/models/kimi_k3/test_amd_kda_adaptive_verification.py new file mode 100644 index 0000000000..9cb4307e56 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_kda_adaptive_verification.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm Kimi-K3 KDA metadata under adaptive verification. + +Adaptive verification splits the draft budget evenly on the host, then trims +each verify request on device. The KDA spec path must therefore plan from the +device offsets, never from host lengths that merely look uniform. +""" + +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import SpeculativeConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.models.kimi_k3.amd.kda_metadata import ( + KimiK3ROCmKDABackend, + KimiK3ROCmKDAMetadataBuilder, +) +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") +NUM_SPEC = 3 + + +def _config(adaptive: bool): + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + ) + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=NUM_SPEC, + enable_adaptive_verification=adaptive, + ) + vllm_config.cache_config.mamba_cache_mode = "none" + return vllm_config + + +def _spec() -> MambaSpec: + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=NUM_SPEC, + ) + + +def _builder(adaptive: bool) -> KimiK3ROCmKDAMetadataBuilder: + return KimiK3ROCmKDAMetadataBuilder( + kv_cache_spec=_spec(), + layer_names=["layer.0"], + vllm_config=_config(adaptive), + device=DEVICE, + ) + + +def _trimmed_verify_batch(builder: KimiK3ROCmKDAMetadataBuilder): + """Host sees [3, 3]; the device was trimmed to [4, 2] (same total).""" + batch = BatchSpec(seq_lens=[64, 64], query_lens=[4, 2]) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE, arange_block_indices=True + ) + common = common.replace( + query_start_loc_cpu=torch.tensor([0, 3, 6], dtype=torch.int32) + ) + accepted = torch.ones(2, dtype=torch.int32, device=DEVICE) + num_decode_draft_tokens_cpu = torch.tensor([NUM_SPEC, NUM_SPEC], dtype=torch.int32) + return builder.build(0, common, accepted, num_decode_draft_tokens_cpu) + + +def test_adaptive_verification_plans_from_device_offsets(): + """Equal host lengths would select the fixed-length recurrent kernel, + which places sequence i at i * L and ignores the device cu_seqlens.""" + assert KimiK3ROCmKDABackend.supports_device_cpu_query_lens_mismatch() + bound = KimiK3ROCmKDAMetadataBuilder.get_varlen_cudagraph_max_query_len + assert bound(_config(adaptive=True), _spec()) == NUM_SPEC + 1 + assert bound(_config(adaptive=False), _spec()) is None + + meta = _trimmed_verify_batch(_builder(adaptive=True)) + assert meta.num_spec_decodes == 2 + assert meta.uniform_spec_sequence_length is None + assert torch.equal( + meta.spec_query_start_loc.cpu(), torch.tensor([0, 4, 6], dtype=torch.int32) + ) + # Without adaptive verification the host lengths are exact and kept. + meta = _trimmed_verify_batch(_builder(adaptive=False)) + assert meta.uniform_spec_sequence_length == 3 diff --git a/tests/v1/attention/test_expand_page_indices.py b/tests/v1/attention/test_expand_page_indices.py index 73163fe13d..1fc078cf56 100644 --- a/tests/v1/attention/test_expand_page_indices.py +++ b/tests/v1/attention/test_expand_page_indices.py @@ -51,6 +51,7 @@ def _run(block_table, cu_num_tokens, kernel_block_size, out_len, device): block_table.stride(0), block_table.stride(1), cu_num_tokens, + cu_num_tokens[1:] - cu_num_tokens[:-1], KERNEL_BLOCK_SIZE=kernel_block_size, BLOCK_SIZE=1024, ) @@ -109,6 +110,7 @@ def test_ragged_batch_leaves_other_rows_untouched(): bt.stride(0), bt.stride(1), cu, + cu[1:] - cu[:-1], KERNEL_BLOCK_SIZE=kbs, BLOCK_SIZE=1024, ) @@ -153,6 +155,7 @@ def test_over_provisioned_chunks_do_not_change_result(): bt.stride(0), bt.stride(1), cu, + cu[1:] - cu[:-1], KERNEL_BLOCK_SIZE=kbs, BLOCK_SIZE=1024, ) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index b696e35b77..891aa8880d 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -418,6 +418,11 @@ for backend in BACKENDS_TO_TEST: if query_len_support != QueryLenSupport.SINGLE_ONLY: SPEC_DECODE_BACKENDS.append(backend) +# Backends whose decode pipeline replays device-trimmed verify batches. +ADAPTIVE_VARLEN_BACKENDS = [ + b for b in BACKENDS_TO_TEST if b == AttentionBackendEnum.ROCM_AITER_MLA +] + BACKEND_BLOCK_SIZES = {} for backend in BACKENDS_TO_TEST: supported_sizes = backend.get_class().get_supported_kernel_block_sizes() @@ -470,6 +475,10 @@ BATCH_SPECS = { "spec_decode_small": BatchSpec( seq_lens=[128, 256, 512, 1024], query_lens=[4, 4, 4, 4] ), + # Adaptive verification: verify requests trimmed to 1..k+1 tokens on device. + "spec_decode_ragged_small": BatchSpec( + seq_lens=[128, 256, 512, 1024], query_lens=[8, 3, 1, 5] + ), "spec_decode_medium": BatchSpec( seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[8, 8, 8, 8, 8, 8] ), @@ -1581,6 +1590,34 @@ def run_attention_backend( return output +def _with_even_cpu_query_split( + common_attn_metadata: CommonAttentionMetadata, +) -> CommonAttentionMetadata: + """Host query lengths as adaptive verification leaves them. + + The device keeps the trimmed per-request lengths; the host sees the same + total spread evenly, the first requests taking the remainder + (AdaptiveVerificationManager.compact_batch). + """ + num_reqs = common_attn_metadata.num_reqs + total = int(common_attn_metadata.query_start_loc_cpu[num_reqs]) + cpu_lens = torch.full((num_reqs,), total // num_reqs, dtype=torch.int32) + cpu_lens[: total % num_reqs] += 1 + query_start_loc_cpu = torch.zeros_like(common_attn_metadata.query_start_loc_cpu) + torch.cumsum(cpu_lens, dim=0, out=query_start_loc_cpu[1 : num_reqs + 1]) + return common_attn_metadata.replace( + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=int( + ( + common_attn_metadata.query_start_loc[1:] + - common_attn_metadata.query_start_loc[:-1] + ) + .max() + .item() + ), + ) + + def _run_backend_correctness( default_vllm_config, dist_init, @@ -1595,6 +1632,7 @@ def _run_backend_correctness( qk_nope_head_dim: int, v_head_dim: int, chunked_prefill_workspace_size: int | None = None, + adaptive_verification: bool = False, ): """Test that all backends produce similar outputs to a reference implementation using torch.nn.functional.scaled_dot_product_attention. @@ -1626,6 +1664,10 @@ def _run_backend_correctness( ) and AttentionBackendEnum.CUTLASS_MLA in backends_to_test: # CUTLASS_MLA does not support non-1 Q/K scales backends_to_test.remove(AttentionBackendEnum.CUTLASS_MLA) + if adaptive_verification: + backends_to_test = [ + b for b in backends_to_test if b in ADAPTIVE_VARLEN_BACKENDS + ] if not backends_to_test: pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}") @@ -1675,13 +1717,16 @@ def _run_backend_correctness( if is_spec_decode_test: from vllm.config import SpeculativeConfig - # Get the query length from the batch spec (they should all be uniform) - query_len = batch_spec.query_lens[0] + # Get the query length from the batch spec (they should all be uniform, + # except under adaptive verification, where the longest is k+1) + query_len = max(batch_spec.query_lens) # Set num_speculative_tokens to query_len - 1 # (since threshold is 1 + num_spec_tokens) # Use ngram method which doesn't require a draft model vllm_config.speculative_config = SpeculativeConfig( - method="ngram", num_speculative_tokens=query_len - 1 + method="ngram", + num_speculative_tokens=query_len - 1, + enable_adaptive_verification=adaptive_verification, ) device = torch.device(f"{DEVICE_TYPE}:0") @@ -1951,6 +1996,13 @@ def _run_backend_correctness( ) kv_cache_per_block_size[block_size] = kv_cache + if adaptive_verification: + # Only after the cache is filled: prepopulation reads the context + # lengths off the host query lengths. + metadata_per_block_size[block_size] = _with_even_cpu_query_split( + common_attn_metadata + ) + # 4. Run vLLM backends and compare rtol = 1e-2 atol = { @@ -2136,3 +2188,29 @@ def test_chunked_context_backend_correctness( v_head_dim=v_head_dim, chunked_prefill_workspace_size=1024, ) + + +@pytest.mark.skipif( + not ADAPTIVE_VARLEN_BACKENDS, reason="No adaptive varlen decode backend" +) +def test_adaptive_varlen_decode_correctness( + default_vllm_config, dist_init, workspace_init +): + """Device query lengths 1..k+1 against an even host split, as adaptive + verification hands them to the target model's attention. TP 8 gives 16 + heads, the padded Kimi-K3 count.""" + _run_backend_correctness( + default_vllm_config, + dist_init, + workspace_init, + "spec_decode_ragged_small", + "deepseek-ai/DeepSeek-R1", + 8, + "fp8", + 1.0, + 1.0, + None, + 128, + 128, + adaptive_verification=True, + ) 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 187bbea4bf..fdc89f4bbe 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -63,9 +63,11 @@ class _ExpandPageIndicesKernel: stride_0, stride_1, paged_kv_indptr, + num_real_tokens, *, KERNEL_BLOCK_SIZE, BLOCK_SIZE, + CLAMP_TAIL=False, ): # Both block-table strides are passed explicitly: the kernel no longer # assumes the column stride is 1. This stand-in indexes the tensor @@ -145,6 +147,7 @@ def _builder( ) ), _mtp_decode_qlen=mtp_decode_qlen, + _varlen_decode_qlen=0, _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len), _use_persistent_metadata=False, kernel_block_size=kernel_block_size, @@ -500,6 +503,7 @@ def test_non_causal_dcp_block_bypasses_cprr_in_forward(monkeypatch): impl.qk_rope_head_dim = 64 impl.scale = 576**-0.5 decode = SimpleNamespace( + varlen_q_rows=None, max_qo_len=4, qo_indptr=torch.tensor([0, 4], dtype=torch.int32), paged_kv_indptr=torch.tensor([0, 1], dtype=torch.int32), @@ -563,6 +567,7 @@ def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch): impl.qk_rope_head_dim = 64 impl.scale = head_dim**-0.5 decode = SimpleNamespace( + varlen_q_rows=None, 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), @@ -787,6 +792,7 @@ def test_segmented_dcp_verify_matches_causal_attention(monkeypatch): torch.tensor([global_seq_len], dtype=torch.int32, device=device), ) decode = SimpleNamespace( + varlen_q_rows=None, max_qo_len=qlen, paged_kv_indptr=torch.tensor( [0, local_kv.shape[0]], dtype=torch.int32, device=device @@ -1452,3 +1458,29 @@ def test_a_two_token_bf16_block_is_allowed(monkeypatch): monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8 ) assert metadata.has_persistent_metadata + + +def _adaptive_config( + *, adaptive: bool = True, dcp: int = 1, num_speculative_tokens: int = 7 +) -> SimpleNamespace: + return SimpleNamespace( + speculative_config=SimpleNamespace( + num_speculative_tokens=num_speculative_tokens, + enable_adaptive_verification=adaptive, + parallel_drafting=False, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp), + use_v2_model_runner=True, + ) + + +def test_varlen_cudagraph_bound_follows_adaptive_verification(): + bound = AiterMLAMetadataBuilder.get_varlen_cudagraph_max_query_len + assert bound(_adaptive_config(), SimpleNamespace()) == 8 + assert bound(_adaptive_config(adaptive=False), SimpleNamespace()) is None + # DCP routes plan per-row windows from host lengths. + assert bound(_adaptive_config(dcp=2), SimpleNamespace()) is None + # Kimi-K3 DSpark draft layers share the target's group; causality is + # decided per batch, so the merged non-causal spec keeps the bound. + shared_draft_group = SimpleNamespace(non_causal_multi_token_decode=True) + assert bound(_adaptive_config(), shared_draft_group) == 8 diff --git a/vllm/models/kimi_k3/amd/kda_metadata.py b/vllm/models/kimi_k3/amd/kda_metadata.py index dcc03d2814..90887c0491 100644 --- a/vllm/models/kimi_k3/amd/kda_metadata.py +++ b/vllm/models/kimi_k3/amd/kda_metadata.py @@ -11,6 +11,7 @@ from dataclasses import replace import torch +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.mamba.checkpoint import ( MambaPrefillCheckpointMetadata, @@ -18,12 +19,14 @@ from vllm.model_executor.layers.mamba.checkpoint import ( from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 +from vllm.v1.attention.backend import max_decode_query_len from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionBackend, GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) from vllm.v1.attention.backends.utils import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import KVCacheSpec logger = init_logger(__name__) @@ -92,7 +95,24 @@ def prepare_chunk_metadata_device( return chunk_indices, chunk_offsets +def _adaptive_verification_enabled(vllm_config: VllmConfig) -> bool: + speculative_config = vllm_config.speculative_config + return speculative_config is not None and bool( + getattr(speculative_config, "enable_adaptive_verification", False) + ) + + class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder): + @classmethod + def get_varlen_cudagraph_max_query_len( + cls: type["KimiK3ROCmKDAMetadataBuilder"], + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> int | None: + if not _adaptive_verification_enabled(vllm_config): + return None + return max_decode_query_len(vllm_config) + def build( # type: ignore[override] self, common_prefix_len: int, @@ -108,6 +128,11 @@ class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder): num_decode_draft_tokens_cpu, fast_build, ) + if ( + attn_metadata.uniform_spec_sequence_length is not None + and _adaptive_verification_enabled(self.vllm_config) + ): + attn_metadata = replace(attn_metadata, uniform_spec_sequence_length=None) if attn_metadata.checkpoint is None: return attn_metadata # The shared builder keys checkpoint rows on every request; the ROCm @@ -163,3 +188,7 @@ class KimiK3ROCmKDABackend(GDNAttentionBackend): @staticmethod def get_builder_cls() -> type[KimiK3ROCmKDAMetadataBuilder]: return KimiK3ROCmKDAMetadataBuilder + + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + return True diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index e485f5462a..bb29e55826 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -35,6 +35,7 @@ from vllm.v1.attention.backend import ( AttentionLayer, CommonAttentionMetadata, MultipleOf, + max_decode_query_len, ) from vllm.v1.attention.backends.mla.triton_mla import _compute_num_kv_splits from vllm.v1.attention.backends.utils import ( @@ -49,7 +50,11 @@ from vllm.v1.attention.ops.rocm_aiter_mla_prefill import ( gather_compressed_context, ) from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd -from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheSpec, + is_quantized_kv_cache, +) from vllm.v1.worker.workspace import ( current_workspace_manager, is_workspace_manager_initialized, @@ -457,6 +462,16 @@ def _aiter_mla_small_head_mode() -> str: return mode +def _adaptive_varlen_verify_enabled(vllm_config: VllmConfig) -> bool: + speculative_config = vllm_config.speculative_config + return ( + speculative_config is not None + and bool(getattr(speculative_config, "enable_adaptive_verification", False)) + and vllm_config.parallel_config.decode_context_parallel_size == 1 + and max_decode_query_len(vllm_config) > 1 + ) + + class AiterMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -620,6 +635,8 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): use_gluon_verify: bool = False # Whether persistent MLA metadata was computed has_persistent_metadata: bool = False + varlen_q_rows: torch.Tensor | None = None + varlen_o_rows: torch.Tensor | None = None @dataclass @@ -720,6 +737,21 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): return uniform_qo_len + @classmethod + def get_varlen_cudagraph_max_query_len( + cls: type["AiterMLAMetadataBuilder"], + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> int | None: + if not _adaptive_varlen_verify_enabled(vllm_config): + return None + if ( + cls.get_cudagraph_support(vllm_config, kv_cache_spec) + != AttentionCGSupport.UNIFORM_BATCH + ): + return None + return max_decode_query_len(vllm_config) + def __init__( self, kv_cache_spec: AttentionSpec, @@ -823,6 +855,26 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # persistent gate below and makes aiter raise a KeyError mid-run. self._mtp_decode_qlen = self.reorder_batch_threshold or 1 + self._varlen_decode_qlen = 0 + if _adaptive_varlen_verify_enabled(vllm_config): + assert self._mtp_decode_qlen == max_decode_query_len(vllm_config), ( + self._mtp_decode_qlen, + max_decode_query_len(vllm_config), + ) + self._varlen_decode_qlen = self._mtp_decode_qlen + self.query_len_support = QueryLenSupport.VARLEN # type: ignore[misc] + max_rows = vllm_config.scheduler_config.max_num_seqs * ( + self._varlen_decode_qlen + ) + row_ids = torch.arange(max_rows, dtype=torch.int64, device=device) + self._varlen_row_req = row_ids // self._varlen_decode_qlen + self._varlen_row_pos = row_ids % self._varlen_decode_qlen + self._varlen_row_ids = row_ids + self._varlen_q_rows = torch.zeros_like(row_ids) + self._varlen_o_rows = torch.zeros( + max_rows + 1, dtype=torch.int64, device=device + ) + # 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 @@ -850,6 +902,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): max_num_pages_per_req = vllm_config.model_config.max_model_len max_num_reqs = vllm_config.scheduler_config.max_num_seqs max_num_pages = max_num_reqs * max_num_pages_per_req + speculative_config = vllm_config.speculative_config + if _adaptive_varlen_verify_enabled(vllm_config): + assert speculative_config is not None + max_num_pages += max_num_reqs * ( + speculative_config.num_speculative_tokens or 0 + ) # Preparing persistent buffers # TODO: we can disambiguate between decode and mixed-prefill decode here @@ -1437,13 +1495,18 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): 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 + varlen_decode = self._varlen_decode_qlen > 0 and causal + if varlen_decode: + max_qo_len = self._varlen_decode_qlen + pad_uniform_mtp = True + else: + 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 # the global lengths need the same dummy-row @@ -1477,6 +1540,13 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): g_tot_seq_lens.new_full((), max_qo_len * self.dcp_world_size), ) + varlen_q_rows = varlen_o_rows = None + if varlen_decode: + seq_lens_for_kernel = seq_lens_device + (max_qo_len - qo_lens_device) + varlen_q_rows, varlen_o_rows = self._fill_varlen_rows( + query_start_loc_device, qo_lens_device, num_reqs + ) + if self._graph_seq_lens is not None: self._graph_seq_lens[:num_kernel_reqs].copy_( seq_lens_for_kernel, non_blocking=True @@ -1537,6 +1607,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # device synchronisation. Programs whose chunk lies past # num_tokens mask out entirely. max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size + if varlen_decode: + max_tokens_per_req += max_qo_len - 1 num_chunks = max(1, cdiv(max_tokens_per_req, 1024)) _expand_page_indices_kernel[(num_reqs, num_chunks)]( self.paged_kv_indices, @@ -1544,8 +1616,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): block_table_tensor.stride(0), block_table_tensor.stride(1), paged_kv_indptr, + seq_lens_device, KERNEL_BLOCK_SIZE=self.kernel_block_size, BLOCK_SIZE=1024, + CLAMP_TAIL=varlen_decode, ) paged_kv_indices = self.paged_kv_indices @@ -1740,10 +1814,38 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): use_gluon_verify=use_gluon_verify, attn_out_dtype=self.decode_attn_out_dtype, has_persistent_metadata=has_persistent_metadata, + varlen_q_rows=varlen_q_rows, + varlen_o_rows=varlen_o_rows, ) return attn_metadata + def _fill_varlen_rows( + self, + query_start_loc_device: torch.Tensor, + qo_lens_device: torch.Tensor, + num_reqs: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + width = self._varlen_decode_qlen + num_rows = num_reqs * width + row_req = self._varlen_row_req[:num_rows] + row_pos = self._varlen_row_pos[:num_rows] + row_qlen = qo_lens_device.to(torch.int64).index_select(0, row_req) + row_start = query_start_loc_device.to(torch.int64).index_select(0, row_req) + is_real = row_pos < row_qlen + packed_row = row_start + row_pos + q_rows = self._varlen_q_rows[:num_rows] + q_rows.copy_(torch.where(is_real, packed_row, 0)) + o_rows = self._varlen_o_rows + o_rows.zero_() + dump_slot = o_rows.size(0) - 1 + o_rows.scatter_( + 0, + torch.where(is_real, packed_row, dump_slot), + self._varlen_row_ids[:num_rows], + ) + return q_rows, o_rows + def build( self, common_prefix_len: int, @@ -1795,8 +1897,10 @@ def _expand_page_indices_kernel( block_table_stride_0, block_table_stride_1, cu_num_tokens, + num_real_tokens, KERNEL_BLOCK_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, + CLAMP_TAIL: tl.constexpr = False, ): """Expand block table entries into per-token flat page indices. @@ -1810,6 +1914,10 @@ def _expand_page_indices_kernel( When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens) is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). + + With CLAMP_TAIL, entries past num_real_tokens[req] (padding a request's + page list out to its kernel length) repeat the request's first page, so + they never index past its block-table row. """ # One program per (request, token-chunk). Parallelising over requests # alone degenerates at low concurrency: with num_reqs == 1 a single @@ -1819,6 +1927,8 @@ def _expand_page_indices_kernel( row_ptr = block_table + req_idx * block_table_stride_0 start_idx = tl.load(cu_num_tokens + req_idx) num_tokens = tl.load(cu_num_tokens + req_idx + 1) - start_idx + if CLAMP_TAIL: + real_tokens = tl.load(num_real_tokens + req_idx) # The grid is sized from the block table width, an upper bound over all # requests, so a ragged batch launches chunks past a short request's end. @@ -1830,11 +1940,14 @@ def _expand_page_indices_kernel( token_offsets = chunk_start + tl.arange(0, BLOCK_SIZE) mask = token_offsets < num_tokens + source_offsets = token_offsets + if CLAMP_TAIL: + source_offsets = tl.where(token_offsets < real_tokens, token_offsets, 0) # Which block in the block table does this token belong to? - block_idx = token_offsets // KERNEL_BLOCK_SIZE + block_idx = source_offsets // KERNEL_BLOCK_SIZE # Offset within that block - offset_in_block = token_offsets % KERNEL_BLOCK_SIZE + offset_in_block = source_offsets % KERNEL_BLOCK_SIZE # Load the block ID from the block table # Both strides are taken from the caller: the block table is a view owned @@ -2600,6 +2713,28 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): kv_c_and_k_pe_cache: torch.Tensor, attn_metadata: AiterMLAMetadata, layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + decode = attn_metadata.decode + if decode is None or decode.varlen_q_rows is None: + return self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer) + assert decode.varlen_o_rows is not None + q_rows = decode.varlen_q_rows + if isinstance(q, tuple): + num_tokens = q[0].shape[0] + q = (q[0].index_select(0, q_rows), q[1].index_select(0, q_rows)) + else: + num_tokens = q.shape[0] + q = q.index_select(0, q_rows) + o, lse = self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer) + assert lse is None + return o.index_select(0, decode.varlen_o_rows[:num_tokens]), None + + def _forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: AiterMLAMetadata, + layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: assert kv_c_and_k_pe_cache.numel() > 0 assert attn_metadata.decode is not None