A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix]: Fix assertion in MambaManager.allocate_slots() (#39562)
Signed-off-by: Holworth <kangqihan17@mails.ucas.ac.cn> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
KANG Qihan committed
Jun 8, 2026 at 12:34 UTC
303916e93d66da301231c9aee80489951a5cd8f6
2 files changed
+55
-5
tests/v1/core/test_prefix_caching.py
+52
@@ -1023,6 +1023,58 @@ def test_prefill_hybrid_model_mamba_align():
1023
manager.free(req0)
1024
1025
1026
+def test_hybrid_model_mamba_align_with_dynamic_draft_tokens():
1027
+ """Regression test for https://github.com/vllm-project/vllm/issues/39271.
1028
+
1029
+ With suffix decoding enabled, the number of proposed draft token may
1030
+ change dynamically each round, causing the MambaManager to crash during
1031
+ allocate_slots() as it originally assumes the `num_blocks` to increase.
1032
+ """
1033
+ block_size = 16
1034
+ num_blocks = 30
1035
+
1036
+ kv_cache_config = _make_hybrid_kv_cache_config(
1037
+ block_size, num_blocks, ["full", "mamba_align"]
1038
+ )
1039
+ manager = KVCacheManager(
1040
+ kv_cache_config,
1041
+ max_model_len=8192,
1042
+ enable_caching=True,
1043
+ hash_block_size=block_size,
1044
+ scheduler_block_size=block_size,
1045
+ )
1046
+
1047
+ # the default hash function is sha256
1048
+ hash_fn = sha256
1049
+
1050
+ all_token_ids = [i for i in range(3) for _ in range(block_size)] + [3] * 7
1051
+ req0 = make_request("0", all_token_ids, block_size, hash_fn)
1052
+ computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
1053
+ assert num_computed_tokens == 0
1054
+ blocks = manager.allocate_slots(
1055
+ req0, len(all_token_ids), num_computed_tokens, computed_blocks
1056
+ )
1057
+ assert blocks is not None
1058
+
1059
+ # prefill forward finished
1060
+ req0.append_output_token_ids([1])
1061
+ req0.num_computed_tokens = len(all_token_ids)
1062
+
1063
+ # Round1: propose 16 draft tokens, accept only one
1064
+ req0.spec_token_ids = [4] * 16
1065
+ blocks = manager.allocate_slots(req0, num_new_tokens=16, num_new_computed_tokens=0)
1066
+ assert blocks is not None
1067
+ req0.append_output_token_ids([4])
1068
+ req0.num_computed_tokens += 1
1069
+
1070
+ # Round2: propose only one token, allocate should not crash
1071
+ req0.spec_token_ids = [5] * 1
1072
+ blocks = manager.allocate_slots(req0, num_new_tokens=1, num_new_computed_tokens=0)
1073
+ assert blocks is not None and all(len(group) == 0 for group in blocks.blocks)
1074
+
1075
+ manager.free(req0)
1076
+
1077
+
1078
def test_prefill_plp():
1079
"""Test prefill with APC and some prompt logprobs (plp) requests.
1080
vllm/v1/core/single_type_kv_cache_manager.py
+3
-5
@@ -1143,13 +1143,11 @@ class MambaManager(SingleTypeKVCacheManager):
1143
num_required_blocks = (
1144
cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
1145
)
1146
- if num_required_blocks == len(req_blocks):
1146
+ # `num_required_blocks` might be less than `len(req_blocks)` if blocks are
1147
+ # over-allocated at last round.
1148
+ if num_required_blocks <= len(req_blocks):
1149
return []
1150
else:
1149
- assert num_required_blocks > len(req_blocks), (
1150
- "num_required_blocks "
1151
- f"{num_required_blocks} < len(req_blocks) {len(req_blocks)}"
1152
- )
1151
prev_block_len = len(req_blocks)
1152
blocks_allocated = request_id in self._allocated_block_reqs
1153
# Record the last state block