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 15, 06:00 UTC
[Model] Fix minimax model cache & lm_head precision (#19592)
Signed-off-by: qingjun <qingjun@minimaxi.com>
qscqesze committed
Jun 13, 2025 at 20:08 UTC
a24cb91600bdfcafd4c18b6647e9184a2b47fcb4
1 file changed
+3
-3
vllm/model_executor/models/minimax_text_01.py
+3
-3
@@ -856,7 +856,7 @@ class MiniMaxText01Model(nn.Module):
856
self._dtype = _dummy.dtype
857
del _dummy
858
859
- self.minimax_cache = MinimaxCacheManager(dtype=self._dtype,
859
+ self.minimax_cache = MinimaxCacheManager(dtype=torch.float32,
860
cache_shape=self.cache_shape)
861
862
rope_theta = getattr(config, "rope_theta", 10000)
@@ -1021,7 +1021,7 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid,
1021
1022
else:
1023
self.lm_head = PPMissingLayer()
1024
-
1024
+ self.lm_head.float()
1025
flash_layer_count = sum(1 for attn_type in self.config.attn_type_list
1026
if attn_type == 1)
1027
self.kv_cache = [torch.tensor([]) for _ in range(flash_layer_count)]
@@ -1054,7 +1054,7 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid,
1054
1055
def compute_logits(self, hidden_states: torch.Tensor,
1056
sampling_metadata: SamplingMetadata) -> torch.Tensor:
1057
- logits = self.logits_processor(self.lm_head, hidden_states,
1057
+ logits = self.logits_processor(self.lm_head, hidden_states.float(),
1058
sampling_metadata)
1059
1060
return logits