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 8, 01:15 UTC
Revert "[NemotronH] Do not force router to run in fp32 (#34582)" (#34808)
Signed-off-by: Roi Koren <roik@nvidia.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
roikoren755 committed
Feb 19, 2026 at 19:47 UTC
3eff45d793daa976a21d0df5954cf6cc6723335f
1 file changed
+4
-1
vllm/model_executor/models/nemotron_h.py
+4
-1
@@ -148,10 +148,12 @@ class NemotronHMoE(nn.Module):
148
149
self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
150
151
+ router_logits_dtype = torch.float32
152
self.gate = ReplicatedLinear(
153
config.hidden_size,
154
config.n_routed_experts,
155
bias=False,
156
+ params_dtype=router_logits_dtype,
157
quant_config=None,
158
prefix=f"{prefix}.gate",
159
)
@@ -230,6 +232,7 @@ class NemotronHMoE(nn.Module):
232
enable_eplb=self.enable_eplb,
233
num_redundant_experts=self.n_redundant_experts,
234
is_sequence_parallel=self.is_sequence_parallel,
235
+ router_logits_dtype=router_logits_dtype,
236
routed_input_transform=self.fc1_latent_proj,
237
)
238
@@ -241,7 +244,7 @@ class NemotronHMoE(nn.Module):
244
hidden_states = sequence_parallel_chunk(hidden_states)
245
246
# router_logits: (num_tokens, n_experts)
244
- router_logits, _ = self.gate(hidden_states)
247
+ router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
248
249
# SharedFusedMoE handles:
250
# - shared experts (with original hidden_states)