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 14, 05:30 UTC
[Hybrid] Add mamba_block_size to Engine Args (#27289)
Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>
Asaf Joseph Gardin committed
Oct 28, 2025 at 14:54 UTC
05181cc57f7939cb6ad44348ee7c99f4a0cee900
3 files changed
+25
-8
vllm/config/cache.py
+13
-3
@@ -5,7 +5,7 @@ import hashlib
5
from dataclasses import field
6
from typing import TYPE_CHECKING, Any, Literal
7
8
-from pydantic import Field, SkipValidation, field_validator
8
+from pydantic import Field, SkipValidation, field_validator, model_validator
9
from pydantic.dataclasses import dataclass
10
11
from vllm.config.utils import config
@@ -90,8 +90,10 @@ class CacheConfig:
90
mamba_page_size_padded: int | None = None
91
""" Optional override for mamba page size; used by hybrid mamba/attention
92
models to ensure exact alignment with attention page size."""
93
- mamba_block_size: int | None = None
94
- """Size of a contiguous cache block in number of tokens for mamba cache."""
93
+ mamba_block_size: int | None = Field(default=None, gt=0)
94
+ """Size of a contiguous cache block in number of tokens for mamba cache.
95
+ Can be set only when prefix caching is enabled.
96
+ Value must be a multiple of 8 to align with causal_conv1d kernel."""
97
mamba_cache_dtype: MambaDType = "auto"
98
"""The data type to use for the Mamba cache (both the conv as well as the
99
ssm state). If set to 'auto', the data type will be inferred from the model
@@ -183,3 +185,11 @@ class CacheConfig:
185
raise ValueError("Too large swap space. " + msg)
186
elif cpu_memory_usage > 0.4 * total_cpu_memory:
187
logger.warning("Possibly too large swap space. %s", msg)
188
+
189
+ @model_validator(mode="after")
190
+ def validate_mamba_block_size(self) -> "CacheConfig":
191
+ if self.mamba_block_size is not None and not self.enable_prefix_caching:
192
+ raise ValueError(
193
+ "--mamba-block-size can only be set with --enable-prefix-caching"
194
+ )
195
+ return self
vllm/engine/arg_utils.py
+5
@@ -535,6 +535,7 @@ class EngineArgs:
535
calculate_kv_scales: bool = CacheConfig.calculate_kv_scales
536
mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype
537
mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype
538
+ mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
539
540
additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config")
541
@@ -893,6 +894,9 @@ class EngineArgs:
894
cache_group.add_argument(
895
"--mamba-ssm-cache-dtype", **cache_kwargs["mamba_ssm_cache_dtype"]
896
)
897
+ cache_group.add_argument(
898
+ "--mamba-block-size", **cache_kwargs["mamba_block_size"]
899
+ )
900
901
# Multimodal related configs
902
multimodal_kwargs = get_kwargs(MultiModalConfig)
@@ -1390,6 +1394,7 @@ class EngineArgs:
1394
kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
1395
mamba_cache_dtype=self.mamba_cache_dtype,
1396
mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,
1397
+ mamba_block_size=self.mamba_block_size,
1398
)
1399
1400
ray_runtime_env = None
vllm/model_executor/models/config.py
+7
-5
@@ -291,9 +291,8 @@ class MambaModelConfig(VerifyAndUpdateConfig):
291
model_config = vllm_config.model_config
292
cache_config = vllm_config.cache_config
293
294
- # Set mamba block size to max_model_len (this may get
295
- # override by prefix caching logic later)
296
- cache_config.mamba_block_size = model_config.max_model_len
294
+ if cache_config.mamba_block_size is None:
295
+ cache_config.mamba_block_size = model_config.max_model_len
296
297
if cache_config.enable_prefix_caching:
298
if model_config.supports_mamba_prefix_caching:
@@ -333,6 +332,8 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
332
if not envs.VLLM_USE_V1:
333
return
334
335
+ # Save the user input before it gets modified by MambaModelConfig
336
+ mamba_block_size = vllm_config.cache_config.mamba_block_size
337
# Enable FULL_AND_PIECEWISE by default
338
MambaModelConfig.verify_and_update_config(vllm_config)
339
@@ -386,7 +387,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
387
# With prefix caching, select attention block size to
388
# optimize for mamba kernel performance
389
389
- # mamba SSD kernel uses a chunk_size, e.g. 256
390
+ # Mamba2 SSD kernel uses a chunk_size, e.g. 256
391
# Align the block to the kernel: use lowest multiple of chunk_size
392
# of attention tokens that would fit mamba_page_size:
393
# e.g. for mamba page size = 788kB
@@ -404,7 +405,8 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
405
def lcm(a, b):
406
return a * b // gcd(a, b)
407
407
- base_chunk_size = model_config.get_mamba_chunk_size()
408
+ base_chunk_size = mamba_block_size or model_config.get_mamba_chunk_size()
409
+
410
attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
411
412
chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)