[Mamba] Declare prefill checkpoint alignment once per layer

tests/models/kimi_k3/test_amd_kda_checkpoint.py

@@ -19,7 +19,6 @@
     MambaPrefillCheckpointExporter,
     MambaPrefillCheckpointMetadata,
 )
-from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
 from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_fn
 from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention
 from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDAMetadataBuilder
@@ -474,10 +473,12 @@ def test_checkpoint_metadata_skips_unaligned_offset() -> None:
 
 
 def _spec_for(monkeypatch, **attrs) -> MambaSpec:
-    """Run the layer's `get_kv_cache_spec` over a stubbed base spec."""
-    base = MambaSpec(block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,))
+    """Run the layer's `get_kv_cache_spec` over stubbed state shapes."""
     monkeypatch.setattr(
-        GatedDeltaNetAttention, "get_kv_cache_spec", lambda self, cfg: base
+        KimiK3DeltaAttention, "get_state_shape", lambda self: ((16, 64),)
+    )
+    monkeypatch.setattr(
+        KimiK3DeltaAttention, "get_state_dtype", lambda self: (torch.float16,)
     )
     layer = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
     for name, value in dict(
@@ -489,7 +490,10 @@ def _spec_for(monkeypatch, **attrs) -> MambaSpec:
         setattr(layer, name, attrs.get(name, value))
     vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B")
     vllm_config.cache_config.mamba_cache_mode = "align"
-    return layer.get_kv_cache_spec(vllm_config)
+    vllm_config.cache_config.mamba_block_size = BLOCK_SIZE
+    spec = layer.get_kv_cache_spec(vllm_config)
+    assert isinstance(spec, MambaSpec)
+    return spec
 
 
 def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None:

tests/v1/attention/test_attention_backends_selection.py

@@ -35,6 +35,7 @@ def test_replayssm_does_not_reserve_speculative_state_blocks(
         get_state_dtype=lambda: (torch.float32,),
         mamba_type=MambaAttentionBackendEnum.MAMBA2,
         is_kv_cache_tp_replicated=False,
+        get_prefill_checkpoint_alignment=lambda: None,
     )
     vllm_config = SimpleNamespace(
         cache_config=SimpleNamespace(
@@ -43,6 +44,7 @@ def test_replayssm_does_not_reserve_speculative_state_blocks(
             mamba_cache_mode="none",
             use_replayssm=use_replayssm,
         ),
+        compilation_config=SimpleNamespace(static_forward_context={}),
         num_speculative_tokens=3,
     )
 

tests/v1/attention/test_mamba2_prefill_checkpoints.py

@@ -2,17 +2,22 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Internal prefill checkpoints for Mamba2 align-mode prefix caching."""
 
+from types import SimpleNamespace
+
+import pytest
 import torch
 
 from tests.v1.attention.utils import (
     BatchSpec,
     create_common_attn_metadata,
     create_vllm_config,
 )
+from vllm.model_executor.layers.mamba.abstract import MambaBase
 from vllm.v1.attention.backends.mamba2_attn import (
     Mamba2AttentionMetadataBuilder,
 )
 from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.kv_cache_interface import MambaSpec
 
@@ -166,3 +171,58 @@ def test_builder_emits_nothing_outside_align_mode():
     meta = _build(builder, seq_lens=[900], query_lens=[900])
 
     assert meta.checkpoint_chunk_idx is None
+
+
+class _CheckpointingMambaLayer(MambaBase):
+    mamba_type = MambaAttentionBackendEnum.MAMBA2
+
+    def __init__(self, prefill_checkpoint_alignment: int | None):
+        self.prefill_checkpoint_alignment = prefill_checkpoint_alignment
+
+    def get_prefill_checkpoint_alignment(self) -> int | None:
+        return self.prefill_checkpoint_alignment
+
+    def get_state_shape(self):
+        return ((4, 8),)
+
+    def get_state_dtype(self):
+        return (torch.float32,)
+
+
+def _first_layer_spec(alignments: list[int | None]) -> MambaSpec:
+    layers = {
+        f"layer.{i}": _CheckpointingMambaLayer(alignment)
+        for i, alignment in enumerate(alignments)
+    }
+    vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            mamba_block_size=64,
+            mamba_page_size_padded=None,
+            mamba_cache_mode="align",
+            use_replayssm=False,
+        ),
+        compilation_config=SimpleNamespace(static_forward_context=layers),
+        num_speculative_tokens=0,
+    )
+    spec = layers["layer.0"].get_kv_cache_spec(vllm_config)
+    assert isinstance(spec, MambaSpec)
+    return spec
+
+
+@pytest.mark.parametrize("alignment", [1, 16, None])
+def test_layer_checkpoint_alignment_reaches_the_kv_cache_spec(alignment):
+    """The scheduler learns a layer's checkpoint grid only from its spec, and a
+    checkpoint block reserved without one would cache a state never written."""
+    spec = _first_layer_spec([alignment])
+
+    assert spec.num_prefill_checkpoint_blocks == int(alignment is not None)
+    assert spec.prefill_checkpoint_alignment == alignment
+
+
+def test_no_checkpoint_block_unless_every_mamba_layer_declares_one():
+    """If any Mamba layer cannot export checkpoints, the scheduler splits the
+    prefill, so a checkpoint block would never be written."""
+    spec = _first_layer_spec([1, None])
+
+    assert spec.num_prefill_checkpoint_blocks == 0
+    assert spec.prefill_checkpoint_alignment is None

vllm/model_executor/layers/mamba/abstract.py

@@ -64,15 +64,25 @@ def get_state_dtype(self) -> tuple[torch.dtype, ...]:
 
     def get_prefill_checkpoint_alignment(self) -> int | None:
         """Token alignment of the checkpoint offsets this layer can export from
-        inside a prefill, if any; in align mode, the platform makes Mamba blocks
-        multiples of it.
+        inside a prefill, or None if it cannot.
+
+        Only read in align mode, where the platform makes Mamba blocks multiples
+        of it. Unless every Mamba layer declares one, the scheduler splits the
+        prefill instead.
         """
         return None
 
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
         mamba_block_size = vllm_config.cache_config.mamba_block_size
         assert mamba_block_size is not None
         page_size_padded = vllm_config.cache_config.mamba_page_size_padded
+        # The scheduler takes internal checkpoints only when every Mamba layer
+        # can export them.
+        checkpoint_alignment = (
+            self.get_prefill_checkpoint_alignment()
+            if get_mamba_prefill_checkpoint_alignment(vllm_config) is not None
+            else None
+        )
         return MambaSpec(
             shapes=tuple(self.get_state_shape()),
             dtypes=self.get_state_dtype(),
@@ -88,6 +98,8 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
                 if vllm_config.cache_config.use_replayssm
                 else vllm_config.num_speculative_tokens
             ),
+            num_prefill_checkpoint_blocks=int(checkpoint_alignment is not None),
+            prefill_checkpoint_alignment=checkpoint_alignment,
         )
 
     def get_attn_backend(self) -> type[AttentionBackend]:
@@ -117,8 +129,8 @@ def get_mamba_prefill_checkpoint_grid(vllm_config: VllmConfig) -> int:
     """Token grid every Mamba layer keeps its prefill checkpoint offsets on.
 
     Prefix-cache boundaries, made multiples of the layers' checkpoint alignment
-    (the platform already keeps blocks on it). The scheduler checks offsets
-    against one alignment, so every Mamba spec must use this one.
+    (the platform already keeps blocks on it), since the scheduler checks
+    offsets against the LCM of every Mamba spec's alignment.
     """
     alignment = get_mamba_prefill_checkpoint_alignment(vllm_config)
     assert alignment is not None

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -6,7 +6,6 @@
 import itertools
 import os
 from collections.abc import Callable
-from dataclasses import replace
 from typing import Literal
 
 import torch
@@ -84,7 +83,6 @@
     FLASHINFER_GDN_CHECKPOINT_ALIGNMENT,
     GDNAttentionMetadata,
 )
-from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
 from vllm.v1.worker.workspace import current_workspace_manager
 
 # Optional ROCm AITER Triton kernels for the GDN decode path.
@@ -736,32 +734,17 @@ def __init__(
     def get_prefill_checkpoint_alignment(self) -> int | None:
         return self._checkpoint_alignment
 
-    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
-        spec = super().get_kv_cache_spec(vllm_config)
-        if not self._exports_checkpoints or not isinstance(spec, MambaSpec):
-            return spec
-        # FlashInfer emits states at one stride for the whole batch, the gcd of
-        # the offsets; keeping every offset on the prefix-cache grid keeps that
-        # stride a multiple of the grid, which bounds the states emitted.
-        return replace(
-            spec,
-            num_prefill_checkpoint_blocks=1,
-            prefill_checkpoint_alignment=get_mamba_prefill_checkpoint_grid(vllm_config),
-        )
-
     @functools.cached_property
     def _exports_checkpoints(self) -> bool:
         # Known only once every layer is built: a Mamba layer that cannot
-        # checkpoint (e.g. Qwen4Exp's PLE short conv) keeps the scheduler
-        # splitting the prefill instead.
-        return (
-            self._checkpoint_alignment is not None
-            and get_mamba_prefill_checkpoint_alignment(self.vllm_config) is not None
-        )
+        # checkpoint keeps the scheduler splitting the prefill instead.
+        return get_mamba_prefill_checkpoint_alignment(self.vllm_config) is not None
 
     @functools.cached_property
     def _checkpoint_buffer_spec(self) -> tuple[tuple[int, ...], torch.dtype]:
-        # Offsets on the grid emit at most one state per grid step of a batch.
+        # FlashInfer emits states at the gcd of a batch's offsets. Offsets are
+        # prefix-cache boundaries on every layer's alignment, i.e. on the grid,
+        # so a batch emits at most one state per grid step.
         # Read lazily: the block size is final only after the model is built.
         return (
             (
@@ -1839,6 +1822,7 @@ def _forward_core(
                 (state_checkpoints,) = current_workspace_manager().get_simultaneous(
                     self._checkpoint_buffer_spec
                 )
+                assert plan.num_checkpoints <= state_checkpoints.shape[0]
                 state_checkpoints = state_checkpoints[: plan.num_checkpoints]
                 (
                     core_attn_out_non_spec,

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -4,13 +4,12 @@
 
 import itertools
 from collections.abc import Sequence
-from dataclasses import replace
 
 import torch
 from torch import nn
 
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
-from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
+from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
 from vllm.config.mamba import MambaBackendEnum
 from vllm.distributed import (
     divide,
@@ -70,7 +69,7 @@
 from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
-from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec, MambaSpec
+from vllm.v1.kv_cache_interface import KVCacheGroupSpec
 
 logger = init_logger(__name__)
 
@@ -1105,17 +1104,10 @@ def conv_ssm_forward(
                     is_blackwell=self.is_blackwell,
                 )
 
-    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
-        spec = super().get_kv_cache_spec(vllm_config)
-        if not isinstance(spec, MambaSpec):
-            return spec
-        # The block-boundary state can be exported from inside one forward
-        # pass, so the scheduler need not split the prefill; every reader
-        # acts on this only in align mode. The metadata builder puts a chunk
-        # end exactly on the checkpoint, so any token position works.
-        return replace(
-            spec, num_prefill_checkpoint_blocks=1, prefill_checkpoint_alignment=1
-        )
+    def get_prefill_checkpoint_alignment(self) -> int:
+        # The metadata builder puts a chunk end exactly on the checkpoint, so
+        # any token position works.
+        return 1
 
     def get_state_dtype(self) -> tuple[torch.dtype, ...]:
         assert self.model_config is not None

vllm/models/glm5next/common/kda.py

@@ -2,8 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """GLM-5.3-Flash KDA layer with separate convolutions and a bounded safe gate."""
 
-from dataclasses import replace
-
 import torch
 from torch import nn
 from transformers import Glm5NextTextConfig
@@ -49,7 +47,6 @@
 from vllm.platforms import current_platform
 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
-from vllm.v1.kv_cache_interface import MambaSpec
 from vllm.v1.worker.workspace import current_workspace_manager
 
 if current_platform.is_rocm():
@@ -198,15 +195,8 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
     num_heads: int
     conv_size: int
 
-    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
-        spec = super().get_kv_cache_spec(vllm_config)
-        assert isinstance(spec, MambaSpec)
-        alignment = kda_prefill_checkpoint_alignment(self.kda_prefill_backend)
-        return replace(
-            spec,
-            num_prefill_checkpoint_blocks=int(alignment is not None),
-            prefill_checkpoint_alignment=alignment,
-        )
+    def get_prefill_checkpoint_alignment(self) -> int | None:
+        return kda_prefill_checkpoint_alignment(self.kda_prefill_backend)
 
     def get_state_dtype(
         self,

vllm/models/kimi_k3/amd/kda.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from dataclasses import replace
-
 import torch
 from einops import rearrange
 from torch import nn
@@ -68,7 +66,6 @@
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
 from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
-from vllm.v1.kv_cache_interface import MambaSpec
 
 logger = init_logger(__name__)
 
@@ -95,24 +92,14 @@ def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
             num_spec=self.num_spec,
         )
 
-    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
-        spec = super().get_kv_cache_spec(vllm_config)
-        assert isinstance(spec, MambaSpec)
+    def get_prefill_checkpoint_alignment(self) -> int | None:
         # Only the fused path exports the checkpoint snapshots
         can_checkpoint = (
             self.use_prefill_checkpoint
             and self.kda_prefill_backend == "fused"
             and self.use_safe_gate
         )
-        if can_checkpoint and vllm_config.cache_config.mamba_cache_mode == "align":
-            logger.info_once("Kimi-K3 KDA prefill checkpoint enabled.")
-        return replace(
-            spec,
-            num_prefill_checkpoint_blocks=int(can_checkpoint),
-            prefill_checkpoint_alignment=(
-                KDA_CHECKPOINT_ALIGNMENT if can_checkpoint else None
-            ),
-        )
+        return KDA_CHECKPOINT_ALIGNMENT if can_checkpoint else None
 
     def __init__(
         self,
@@ -280,6 +267,11 @@ def __init__(
             if isinstance(additional_config, dict)
             else True
         )
+        if (
+            self.get_prefill_checkpoint_alignment() is not None
+            and vllm_config.cache_config.mamba_cache_mode == "align"
+        ):
+            logger.info_once("Kimi-K3 KDA prefill checkpoint enabled.")
         self._checkpoint_exporter = KimiK3ROCmKDAPrefillCheckpointExporter(
             state_len=self.conv_size - 1
         )