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
[BugFix] Fix regression caused by mamba state dtype PR (#22998)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Aug 16, 2025 at 00:55 UTC
f5d412bafbd9d4700ff57cb6a2d5220cf2b7637e
2 files changed
+12
-4
vllm/model_executor/models/phi4flash.py
+6
-2
@@ -650,8 +650,12 @@ class Phi4FlashForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):
650
num_mamba_layers = self.config.num_hidden_layers \
651
// 2 // self.config.mb_per_layer + 1
652
self.mamba_cache = MambaCacheManager(
653
- self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,
654
- *self._get_mamba_cache_shape())
653
+ self.vllm_config,
654
+ num_mamba_layers,
655
+ *self._get_mamba_cache_shape(),
656
+ self.lm_head.weight.dtype,
657
+ self.lm_head.weight.dtype,
658
+ )
659
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
660
661
attn_metadata = get_forward_context().attn_metadata
vllm/model_executor/models/plamo2.py
+6
-2
@@ -767,8 +767,12 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, SupportsPP,
767
self.vllm_config.parallel_config, LayerBlockType.mamba)
768
769
self.mamba_cache = MambaCacheManager(
770
- self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,
771
- *self._get_mamba_cache_shape())
770
+ self.vllm_config,
771
+ num_mamba_layers,
772
+ *self._get_mamba_cache_shape(),
773
+ self.lm_head.weight.dtype,
774
+ self.lm_head.weight.dtype,
775
+ )
776
777
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
778