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
[Hybrid] Fix and optimize block-aligned splitting in mamba cache align mode (#33706)
[Hybrid] Fix and optimize block-aligned splitting in mamba cache align mode (#33706) Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Harry Huang committed
Feb 13, 2026 at 10:21 UTC
bf37812ca77acf7f00c7761bdb0cf257d0e391a3
1 file changed
+14
-11
vllm/v1/core/sched/scheduler.py
+14
-11
@@ -281,27 +281,30 @@ class Scheduler(SchedulerInterface):
281
assert num_external_computed_tokens == 0, (
282
"External KV connector is not verified yet"
283
)
284
- # TODO: need check for resume requests
285
- if request.num_output_tokens == 0: # prefill
284
+ num_computed_tokens = (
285
+ request.num_computed_tokens
286
+ + num_new_local_computed_tokens
287
+ + num_external_computed_tokens
288
+ )
289
+ # Perform block-aligned splitting at prefill phase, including:
290
+ # * non-resumed requests: num_computed_tokens < num_prompt_tokens + 0
291
+ # * resumed requests: num_computed_tokens < (
292
+ # num_prompt_tokens + num_output_tokens
293
+ # )
294
+ # NOTE: Use `request.num_tokens - 1` to bypass normal decoding.
295
+ if num_computed_tokens < max(request.num_prompt_tokens, request.num_tokens - 1):
296
# To enable block-aligned caching of the Mamba state, `num_new_tokens`
297
# must be a multiple of `block_size`.
298
# As an exception, if `num_new_tokens` is less than `block_size`, the
299
# state is simply not cached, requiring no special handling.
300
# Additionally, when Eagle mode is enabled, FullAttn prunes the last
301
# matching block. To prevent this from causing a Mamba cache miss, the
292
- # last chunk must be larger than `block_size`.
302
+ # last chunk must be not smaller than `block_size`.
303
block_size = self.cache_config.block_size
294
- last_cache_position = (
295
- request.num_prompt_tokens - request.num_prompt_tokens % block_size
296
- )
304
+ last_cache_position = request.num_tokens - request.num_tokens % block_size
305
# eagle prune
306
if self.use_eagle:
307
last_cache_position = max(last_cache_position - block_size, 0)
300
- num_computed_tokens = (
301
- request.num_computed_tokens
302
- + num_new_local_computed_tokens
303
- + num_external_computed_tokens
304
- )
308
num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens
309
if num_computed_tokens_after_sched < last_cache_position:
310
# align to block_size