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, 05:00 UTC
[NemotronH] Do not force router to run in fp32 (#34582)
Signed-off-by: Roi Koren <roik@nvidia.com>
roikoren755 committed
Feb 16, 2026 at 20:15 UTC
3b30e6150777de549b11f67dde3ecc0d3b1f3f50
2 files changed
+5
-4
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+4
@@ -309,6 +309,10 @@ def fi_trtllm_fp8_per_tensor_moe(
309
310
from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe
311
312
+ # The DeepSeekV3 routing method requires float32 router logits.
313
+ if routing_method_type == RoutingMethodType.DeepSeekV3:
314
+ routing_logits = routing_logits.to(torch.float32)
315
+
316
return flashinfer_trtllm_fp8_per_tensor_scale_moe(
317
routing_logits=routing_logits,
318
routing_bias=routing_bias,
vllm/model_executor/models/nemotron_h.py
+1
-4
@@ -148,12 +148,10 @@ class NemotronHMoE(nn.Module):
148
149
self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
150
151
- router_logits_dtype = torch.float32
151
self.gate = ReplicatedLinear(
152
config.hidden_size,
153
config.n_routed_experts,
154
bias=False,
156
- params_dtype=router_logits_dtype,
155
quant_config=None,
156
prefix=f"{prefix}.gate",
157
)
@@ -232,7 +230,6 @@ class NemotronHMoE(nn.Module):
230
enable_eplb=self.enable_eplb,
231
num_redundant_experts=self.n_redundant_experts,
232
is_sequence_parallel=self.is_sequence_parallel,
235
- router_logits_dtype=router_logits_dtype,
233
routed_input_transform=self.fc1_latent_proj,
234
)
235
@@ -244,7 +241,7 @@ class NemotronHMoE(nn.Module):
241
hidden_states = sequence_parallel_chunk(hidden_states)
242
243
# router_logits: (num_tokens, n_experts)
247
- router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
244
+ router_logits, _ = self.gate(hidden_states)
245
246
# SharedFusedMoE handles:
247
# - shared experts (with original hidden_states)