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 9, 10:00 UTC
Bugfix: Pass router logits dtype in nemotron shared experts (#32669)
Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com>
amirkl94 committed
Jan 29, 2026 at 01:36 UTC
e01ff5c070f47a0562711e61f5f885d4ccf241c9
1 file changed
+3
-1
vllm/model_executor/models/nemotron_h.py
+3
-1
@@ -145,11 +145,12 @@ class NemotronHMoE(nn.Module):
145
146
self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
147
148
+ router_logits_dtype = torch.float32
149
self.gate = ReplicatedLinear(
150
config.hidden_size,
151
config.n_routed_experts,
152
bias=False,
152
- params_dtype=torch.float32,
153
+ params_dtype=router_logits_dtype,
154
quant_config=None,
155
prefix=f"{prefix}.gate",
156
)
@@ -209,6 +210,7 @@ class NemotronHMoE(nn.Module):
210
enable_eplb=self.enable_eplb,
211
num_redundant_experts=self.n_redundant_experts,
212
is_sequence_parallel=self.is_sequence_parallel,
213
+ router_logits_dtype=router_logits_dtype,
214
)
215
216
if self.use_latent_moe: