From e22519ef629df4a7c43957ad33e74a1ec019fb3a Mon Sep 17 00:00:00 2001 From: zjy0516 Date: Tue, 13 Oct 2026 17:56:00 +0000 Subject: [PATCH 1/2] [Mamba] Declare prefill checkpoint alignment once per layer Co-Authored-By: Claude Opus 5.5 Signed-off-by: zjy0516 --- .../models/kimi_k3/test_amd_kda_checkpoint.py | 14 ++++--- .../test_mamba2_prefill_checkpoints.py | 42 +++++++++++++++++++ vllm/model_executor/layers/mamba/abstract.py | 12 ++++++ .../layers/mamba/mamba_mixer2.py | 20 +++------ vllm/models/glm5next/common/kda.py | 14 +------ vllm/models/kimi_k3/amd/kda.py | 22 ++++------ vllm/models/kimi_k3/nvidia/kda.py | 13 +----- 7 files changed, 80 insertions(+), 57 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py index 6f3129a224e5..b2506ee9e8fd 100644 --- a/tests/models/kimi_k3/test_amd_kda_checkpoint.py +++ b/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: diff --git a/tests/v1/attention/test_mamba2_prefill_checkpoints.py b/tests/v1/attention/test_mamba2_prefill_checkpoints.py index a3968b5a0bc8..cd0aebcd7d79 100644 --- a/tests/v1/attention/test_mamba2_prefill_checkpoints.py +++ b/tests/v1/attention/test_mamba2_prefill_checkpoints.py @@ -2,6 +2,9 @@ # 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 ( @@ -9,10 +12,12 @@ 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,40 @@ 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,) + + +@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.""" + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + mamba_block_size=64, + mamba_page_size_padded=None, + mamba_cache_mode="align", + use_replayssm=False, + ), + num_speculative_tokens=0, + ) + + spec = _CheckpointingMambaLayer(alignment).get_kv_cache_spec(vllm_config) + + assert isinstance(spec, MambaSpec) + assert spec.num_prefill_checkpoint_blocks == int(alignment is not None) + assert spec.prefill_checkpoint_alignment == alignment diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 5322f899dd63..0b476203e094 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -62,10 +62,20 @@ def is_kv_cache_tp_replicated(self) -> bool: def get_state_dtype(self) -> tuple[torch.dtype, ...]: pass + def get_prefill_checkpoint_alignment(self) -> int | None: + """Token alignment of the checkpoint offsets this layer can export from + inside a prefill, or None if it cannot. + + Only read in align mode, where the scheduler declines checkpoints off + this grid and 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 + checkpoint_alignment = self.get_prefill_checkpoint_alignment() return MambaSpec( shapes=tuple(self.get_state_shape()), dtypes=self.get_state_dtype(), @@ -81,6 +91,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]: diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index e310aaf40f12..f8ade4115524 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -3,13 +3,12 @@ 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, @@ -69,7 +68,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__) @@ -1093,17 +1092,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 diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 68d405b40cfc..57b344e2d05c 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/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 @@ -46,7 +44,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(): @@ -165,15 +162,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, diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 409cf6b1bf90..5fed5544db3f 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/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 @@ -66,7 +64,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__) @@ -91,22 +88,12 @@ 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.use_fused_chunk 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, @@ -261,6 +248,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 ) diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 26fe7cb5b457..b3a11b5430ab 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from dataclasses import replace import torch from einops import rearrange @@ -62,7 +61,6 @@ has_flashinfer_recurrent_kda, ) from vllm.v1.attention.backend import AttentionBackend -from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) @@ -826,15 +824,8 @@ def __init__( raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - 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 forward( self, From b81127a9e9a5e904b6af4ac8f2bb5d8bb2f10269 Mon Sep 17 00:00:00 2001 From: zjy0516 Date: Wed, 28 Oct 2026 09:20:33 +0000 Subject: [PATCH 2/2] [Mamba] Derive GDN and PLE checkpoint fields in the base, LCM in scheduler Co-Authored-By: Claude Opus 5.5 Signed-off-by: zjy0516 --- .../test_attention_backends_selection.py | 2 ++ .../test_mamba2_prefill_checkpoints.py | 4 +-- vllm/model_executor/layers/mamba/abstract.py | 4 +-- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 28 ++++--------------- vllm/models/qwen4_exp/nvidia/ple_layer.py | 20 +++---------- vllm/v1/core/sched/scheduler.py | 20 +++++++------ 6 files changed, 27 insertions(+), 51 deletions(-) diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index cf6a02ea7beb..974efa1f4e3f 100644 --- a/tests/v1/attention/test_attention_backends_selection.py +++ b/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, ) diff --git a/tests/v1/attention/test_mamba2_prefill_checkpoints.py b/tests/v1/attention/test_mamba2_prefill_checkpoints.py index c449e9d9dc24..84211bc27906 100644 --- a/tests/v1/attention/test_mamba2_prefill_checkpoints.py +++ b/tests/v1/attention/test_mamba2_prefill_checkpoints.py @@ -177,10 +177,10 @@ class _CheckpointingMambaLayer(MambaBase): mamba_type = MambaAttentionBackendEnum.MAMBA2 def __init__(self, prefill_checkpoint_alignment: int | None): - self._prefill_checkpoint_alignment = prefill_checkpoint_alignment + self.prefill_checkpoint_alignment = prefill_checkpoint_alignment def get_prefill_checkpoint_alignment(self) -> int | None: - return self._prefill_checkpoint_alignment + return self.prefill_checkpoint_alignment def get_state_shape(self): return ((4, 8),) diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 3f38f6f72ba8..be0fe9ef924e 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -129,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 diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index b712268a9460..eabe402af25e 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/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, diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 92ec76495f4c..1722afa7725c 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -3,7 +3,6 @@ """GPU-resident Qwen4Exp position-learning enhancement layers.""" from collections.abc import Sequence -from dataclasses import replace import torch from torch import nn @@ -15,7 +14,6 @@ from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.model_executor.layers.mamba.abstract import ( MambaBase, - get_mamba_prefill_checkpoint_alignment, get_mamba_prefill_checkpoint_grid, ) from vllm.model_executor.layers.mamba.checkpoint import ( @@ -204,20 +202,10 @@ def get_prefill_checkpoint_alignment(self) -> int | None: def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: spec = super().get_kv_cache_spec(vllm_config) - # The scheduler takes internal checkpoints only when every Mamba layer - # can export them. - if ( - self.get_prefill_checkpoint_alignment() is None - or get_mamba_prefill_checkpoint_alignment(vllm_config) is None - or not isinstance(spec, MambaSpec) - ): - return spec - grid = get_mamba_prefill_checkpoint_grid(vllm_config) - # Offsets are multiples of the grid, so the window holds only inputs. - assert grid >= self.conv_state_len - return replace( - spec, num_prefill_checkpoint_blocks=1, prefill_checkpoint_alignment=grid - ) + if isinstance(spec, MambaSpec) and spec.num_prefill_checkpoint_blocks: + # Offsets lie on the grid, so the window holds only inputs. + assert get_mamba_prefill_checkpoint_grid(vllm_config) >= self.conv_state_len + return spec def _short_conv_dilated_dispatch( self, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 5d97ef2dd91e..bf33061ff21d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -5,6 +5,7 @@ from collections import defaultdict, deque from collections.abc import Iterable from dataclasses import replace +from math import lcm from typing import Any from vllm.compilation.cuda_graph import CUDAGraphStat @@ -425,15 +426,16 @@ def __init__( self.need_mamba_block_aligned_split = ( self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" ) - # TODO: Support models with multiple Mamba specs that require different - # prefill checkpoint alignments instead of selecting the first one. - self.mamba_prefill_checkpoint_alignment = next( - ( - group.kv_cache_spec.prefill_checkpoint_alignment - for group in kv_cache_config.kv_cache_groups - if isinstance(group.kv_cache_spec, MambaSpec) - ), - None, + mamba_checkpoint_alignments = [ + group.kv_cache_spec.prefill_checkpoint_alignment + for group in kv_cache_config.kv_cache_groups + if isinstance(group.kv_cache_spec, MambaSpec) + ] + # A checkpoint offset must lie on every Mamba group's alignment. + self.mamba_prefill_checkpoint_alignment = ( + lcm(*filter(None, mamba_checkpoint_alignments)) + if mamba_checkpoint_alignments and None not in mamba_checkpoint_alignments + else None ) self.mamba_has_prefill_checkpoint_blocks = self.has_mamba_layers and all( not isinstance(group.kv_cache_spec, MambaSpec)