[Bugfix] Fix Mamba all-mode CPU offload boundary alignment (#51100)
Signed-off-by: qianxu <admin@picasso-lab.com> Co-authored-by: qianxu <admin@picasso-lab.com>
Qianxu Wang committed
Aug 6, 2026 at 21:37 UTC
c810e5ee9976ad86b81d1277b53e76d0ee639414
2 files changed
+21
-13
tests/v1/kv_connector/unit/test_offloading_connector.py
+14
-9
@@ -542,10 +542,13 @@ def test_fs_tiering_offloading(tmp_path) -> None:
542
[
543
# ("Qwen/Qwen3.6-35B-A3B", 1056, 2),
544
# ("tiiuae/falcon-mamba-7b", 16, 1),
545
- ("state-spaces/mamba-1.4b-hf", 16, 1)
545
+ ("state-spaces/mamba-1.4b-hf", 16, 1),
546
],
547
)
548
-def test_mamba_align_cpu_offload(model: str, block_size: int, tp_size: int):
548
+@pytest.mark.parametrize("mamba_cache_mode", ["align", "all"])
549
+def test_mamba_cpu_offload_boundary(
550
+ model: str, block_size: int, tp_size: int, mamba_cache_mode: str
551
+):
552
kv_transfer_config = KVTransferConfig(
553
kv_connector="OffloadingConnector",
554
kv_role="kv_both",
@@ -562,7 +565,11 @@ def test_mamba_align_cpu_offload(model: str, block_size: int, tp_size: int):
565
kv_transfer_config=kv_transfer_config,
566
language_model_only=True,
567
enable_prefix_caching=True,
565
- mamba_cache_mode="align",
568
+ mamba_cache_mode=mamba_cache_mode,
569
+ # Use lossless state storage so this exact-equality regression isolates
570
+ # offload boundary selection from low-precision Mamba checkpointing.
571
+ mamba_cache_dtype="float32",
572
+ mamba_ssm_cache_dtype="float32",
573
disable_hybrid_kv_cache_manager=False,
574
)
575
@@ -599,12 +606,10 @@ def test_mamba_align_cpu_offload(model: str, block_size: int, tp_size: int):
606
)
607
608
try:
602
- # Mamba has only a single state. The CPU cache stores are triggered
603
- # at offload block boundaries. When the prompt is exactly at the boundary,
604
- # The CPU offload should not load the cached block.
605
- # This is because we'd use that state to recompute the last token. This
606
- # does not work for mamba as there is only one KV value and that is for
607
- # for the token at the boundary.
609
+ # A cached Mamba entry is a point state. CPU cache stores are triggered
610
+ # at offload block boundaries. When the prompt is exactly at a boundary,
611
+ # CPU offload must not load that boundary's cached state because it already
612
+ # includes the last token that must be recomputed.
613
# This is fine for other attention types as we have all the necessary
614
# token KV values in the hit blocks.
615
prompt = TokensPrompt(prompt_token_ids=initial_ids)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+7
-4
@@ -146,14 +146,17 @@ def resolve_mamba_align_size(
146
"""Scan all KV cache groups in *spec* and return the single mamba alignment
147
size, or None if no group requires mamba alignment.
148
149
- For MambaSpec groups in "align" cache mode the hit window must be rounded
150
- down to a multiple of the offloaded chunk size. Asserts that all such
151
- groups agree on the same value.
149
+ For MambaSpec groups in "align" or "all" cache mode the hit window must be
150
+ rounded down to a multiple of the offloaded chunk size. Asserts that all
151
+ such groups agree on the same value.
152
"""
153
mamba_align_size: int | None = None
154
for idx, tokens_per_block in enumerate(spec.tokens_per_block):
155
kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
156
- if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align":
156
+ if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode in (
157
+ "align",
158
+ "all",
159
+ ):
160
tokens_per_chunk = tokens_per_block * spec.blocks_per_chunk
161
assert mamba_align_size is None or mamba_align_size == tokens_per_chunk
162
mamba_align_size = tokens_per_chunk