A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 08:45 UTC
[Bug Fix] Fix MambaManager.cache_blocks() crash on null blocks in align mode (#34418)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
Feb 13, 2026 at 16:13 UTC
4137c5dfa7c0de6c0ff74ad3774224b6b3280349
2 files changed
+48
tests/v1/core/test_prefix_caching.py
+46
@@ -744,6 +744,12 @@ def _make_hybrid_kv_cache_config(
744
shapes=(1, 1),
745
dtypes=(torch.float32,),
746
),
747
+ "mamba_align": lambda: MambaSpec(
748
+ block_size=block_size,
749
+ shapes=(1, 1),
750
+ dtypes=(torch.float32,),
751
+ mamba_cache_mode="align",
752
+ ),
753
}
754
755
kv_cache_groups = [
@@ -962,6 +968,46 @@ def test_prefill_hybrid_model_combinations_eagle(
968
manager.free(req1)
969
970
971
+def test_prefill_hybrid_model_mamba_align():
972
+ """Test that MambaManager.cache_blocks() handles null blocks in align mode.
973
+
974
+ Regression test for https://github.com/vllm-project/vllm/issues/34361.
975
+ In mamba_cache_mode="align", allocate_new_blocks() pads req_to_blocks with
976
+ null blocks. cache_full_blocks() correctly skips them, but
977
+ MambaManager.cache_blocks() must also skip null blocks when tracking
978
+ cached_blocks_this_step.
979
+ """
980
+ block_size = 16
981
+ num_blocks = 30
982
+
983
+ kv_cache_config = _make_hybrid_kv_cache_config(
984
+ block_size, num_blocks, ["full", "mamba_align"]
985
+ )
986
+ manager = KVCacheManager(
987
+ kv_cache_config,
988
+ max_model_len=8192,
989
+ enable_caching=True,
990
+ hash_block_size=block_size,
991
+ )
992
+
993
+ hash_fn = sha256
994
+
995
+ # 3 full blocks (48 tokens) + 7 partial tokens = 55 tokens total
996
+ all_token_ids = [i for i in range(3) for _ in range(block_size)] + [3] * 7
997
+
998
+ # First request: allocate_slots should not crash with the assertion error
999
+ # in MambaManager.cache_blocks() when null blocks are present.
1000
+ req0 = make_request("0", all_token_ids, block_size, hash_fn)
1001
+ computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
1002
+ assert num_computed_tokens == 0
1003
+
1004
+ blocks = manager.allocate_slots(req0, 55, num_computed_tokens, computed_blocks)
1005
+ assert blocks is not None
1006
+ assert len(blocks.get_block_ids()) == 2 # full_attn + mamba groups
1007
+
1008
+ manager.free(req0)
1009
+
1010
+
1011
def test_prefill_plp():
1012
"""Test prefill with APC and some prompt logprobs (plp) requests.
1013
vllm/v1/core/single_type_kv_cache_manager.py
+2
@@ -1000,6 +1000,8 @@ class MambaManager(SingleTypeKVCacheManager):
1000
for block in self.req_to_blocks[request.request_id][
1001
num_cached_blocks_before:num_cached_blocks_after
1002
]:
1003
+ if block.is_null:
1004
+ continue
1005
assert block.block_hash is not None
1006
self.cached_blocks_this_step.add(block.block_hash)
1007