[ROCm][Kimi-K3] Support DSpark adaptive verification on MI355X (#59824)

Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

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

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,
         )

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,
+    )

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

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

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