Default to 'align' mamba cache mode for Mamba-based models when speculative...

Default to 'align' mamba cache mode for Mamba-based models when speculative decoding is enabled (#40454)
Signed-off-by: Roi Koren <roik@nvidia.com>

vllm/model_executor/models/config.py

@@ -325,6 +325,17 @@ class MambaModelConfig(VerifyAndUpdateConfig):
 
         if cache_config.enable_prefix_caching:
             if cache_config.mamba_cache_mode == "none":
+                if (
+                    model_config.supports_mamba_prefix_caching
+                    and vllm_config.speculative_config is not None
+                ):
+                    cache_config.mamba_cache_mode = "align"
+                    logger.warning(
+                        "Mamba cache mode is set to 'align' for %s by default "
+                        "when prefix caching and speculative decoding are enabled",
+                        model_config.architecture,
+                    )
+                else:
                     cache_config.mamba_cache_mode = (
                         "all" if model_config.supports_mamba_prefix_caching else "align"
                     )