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 4, 21:15 UTC
Default to 'align' mamba cache mode for Mamba-based models when speculative decoding is enabled (#40454)
Signed-off-by: Roi Koren <roik@nvidia.com>
roikoren755 committed
Apr 21, 2026 at 17:51 UTC
f819265a4ab0187181575c02174ec4a2f91d9220
1 file changed
+20
-9
vllm/model_executor/models/config.py
+20
-9
@@ -325,15 +325,26 @@ class MambaModelConfig(VerifyAndUpdateConfig):
325
326
if cache_config.enable_prefix_caching:
327
if cache_config.mamba_cache_mode == "none":
328
- cache_config.mamba_cache_mode = (
329
- "all" if model_config.supports_mamba_prefix_caching else "align"
330
- )
331
- logger.warning(
332
- "Mamba cache mode is set to '%s' for %s by default "
333
- "when prefix caching is enabled",
334
- cache_config.mamba_cache_mode,
335
- model_config.architecture,
336
- )
328
+ if (
329
+ model_config.supports_mamba_prefix_caching
330
+ and vllm_config.speculative_config is not None
331
+ ):
332
+ cache_config.mamba_cache_mode = "align"
333
+ logger.warning(
334
+ "Mamba cache mode is set to 'align' for %s by default "
335
+ "when prefix caching and speculative decoding are enabled",
336
+ model_config.architecture,
337
+ )
338
+ else:
339
+ cache_config.mamba_cache_mode = (
340
+ "all" if model_config.supports_mamba_prefix_caching else "align"
341
+ )
342
+ logger.warning(
343
+ "Mamba cache mode is set to '%s' for %s by default "
344
+ "when prefix caching is enabled",
345
+ cache_config.mamba_cache_mode,
346
+ model_config.architecture,
347
+ )
348
if (
349
cache_config.mamba_cache_mode == "all"
350
and not model_config.supports_mamba_prefix_caching