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 10, 06:30 UTC
[Bugfix] Fix stale SSM state for new Mamba requests scheduled as decode (#32118)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
Asaf Joseph Gardin committed
Jan 12, 2026 at 19:02 UTC
8fb2c135be35d4d07b3ef25e36c51733345a01eb
2 files changed
+24
-3
tests/v1/attention/test_batch_reordering.py
+21
@@ -98,6 +98,27 @@ REORDER_TEST_CASES = {
98
expected_order=[0, 1, 6, 8, 4, 3, 2, 7, 5],
99
expected_modified=True,
100
),
101
+ "new_request_single_token_prefill": ReorderTestCase(
102
+ requests=[
103
+ (100, 0),
104
+ (1, 0), # New request with only 1 token (STILL prefill)
105
+ (50, 100),
106
+ (1, 10),
107
+ ],
108
+ # Only index 3 is a true decode (has num_computed_tokens > 0)
109
+ expected_order=[3, 2, 0, 1],
110
+ expected_modified=True,
111
+ ),
112
+ "multiple_new_requests_single_token_prefill": ReorderTestCase(
113
+ requests=[
114
+ (1, 0), # New prefill (1 token, no computed)
115
+ (1, 0), # New prefill (1 token, no computed)
116
+ (1, 50),
117
+ (200, 0),
118
+ ],
119
+ expected_order=[2, 1, 0, 3],
120
+ expected_modified=True,
121
+ ),
122
}
123
124
vllm/v1/attention/backends/utils.py
+3
-3
@@ -1040,9 +1040,9 @@ def reorder_batch_to_split_decodes_and_prefills(
1040
num_scheduled_tokens_np = np.array(num_scheduled_tokens)
1041
num_computed_tokens_np = input_batch.num_computed_tokens_cpu[:num_reqs]
1042
1043
- is_decode = num_scheduled_tokens_np <= decode_threshold
1044
- is_extend = (~is_decode) & (num_computed_tokens_np > 0)
1045
- is_prefill = (~is_decode) & (num_computed_tokens_np == 0)
1043
+ is_prefill = num_computed_tokens_np == 0
1044
+ is_decode = (num_scheduled_tokens_np <= decode_threshold) & (~is_prefill)
1045
+ is_extend = (num_scheduled_tokens_np > decode_threshold) & (~is_prefill)
1046
1047
# Desired order: decode → extend → prefill
1048
req_regions = np.zeros(is_decode.shape, dtype=np.int32) # 0 = decode by default