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
[BugFix] Fix and optimize max_num_blocks_per_req calculation for MambaSpec (#34440)
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Harry Huang committed
Feb 13, 2026 at 16:13 UTC
7a8a46ddcb05ba754e1f0f3f428ebbeb572d0f02
1 file changed
+6
-11
vllm/v1/worker/gpu_model_runner.py
+6
-11
@@ -5698,28 +5698,23 @@ class GPUModelRunner(
5698
kv_cache_config: The KV cache configuration.
5699
kernel_block_sizes: The kernel block sizes for each KV cache group.
5700
"""
5701
- block_sizes = [
5702
- kv_cache_group.kv_cache_spec.block_size
5703
- for kv_cache_group in kv_cache_config.kv_cache_groups
5704
- if not isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec)
5705
- ]
5701
+ block_sizes = []
5702
max_num_blocks = []
5703
max_model_len = max(self.max_model_len, self.max_encoder_len)
5708
- for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
5704
+ for kv_cache_group in kv_cache_config.kv_cache_groups:
5705
if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):
5706
continue
5707
+ block_size = kv_cache_group.kv_cache_spec.block_size
5708
+ block_sizes.append(block_size)
5709
max_num_blocks_per_req = cdiv(
5712
- max_model_len, block_sizes[i] * get_total_cp_world_size()
5710
+ max_model_len, block_size * get_total_cp_world_size()
5711
)
5712
if isinstance(kv_cache_group.kv_cache_spec, MambaSpec):
5715
- mamba_blocks_per_req = (
5713
+ max_num_blocks_per_req = (
5714
max_num_blocks_per_req
5715
if self.cache_config.enable_prefix_caching
5716
else 1
5717
) + kv_cache_group.kv_cache_spec.num_speculative_blocks
5720
- max_num_blocks_per_req = max(
5721
- max_num_blocks_per_req, mamba_blocks_per_req
5722
- )
5718
max_num_blocks.append(max_num_blocks_per_req)
5719
5720
if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [