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 13, 02:00 UTC
[Performance][Fix] update nvfp4 code to support renorm routing (#28569)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
jiahanc committed
Nov 16, 2025 at 18:02 UTC
561253b37faadaafe68168ea32d8d8157621a6b4
2 files changed
+15
-8
vllm/model_executor/layers/quantization/modelopt.py
+11
-7
@@ -15,6 +15,7 @@ from vllm.logger import init_logger
15
from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEConfig,
17
FusedMoEQuantConfig,
18
+ RoutingMethodType,
19
fp8_w8a8_moe_quant_config,
20
nvfp4_moe_quant_config,
21
)
@@ -1657,16 +1658,19 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1658
use_llama4_routing = (
1659
custom_routing_function is Llama4MoE.custom_routing_function
1660
)
1660
- routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3
1661
+ routing_method_type = layer.routing_method_type
1662
if use_llama4_routing:
1662
- routing_method_type = flashinfer.RoutingMethodType.Llama4
1663
+ routing_method_type = RoutingMethodType.Llama4
1664
+ router_logits = (
1665
+ router_logits.to(torch.float32)
1666
+ if routing_method_type == RoutingMethodType.DeepSeekV3
1667
+ else router_logits
1668
+ )
1669
routing_bias = e_score_correction_bias
1670
if routing_bias is not None:
1671
routing_bias = routing_bias.to(torch.bfloat16)
1672
out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
1667
- routing_logits=router_logits
1668
- if use_llama4_routing
1669
- else router_logits.to(torch.float32),
1673
+ routing_logits=router_logits,
1674
routing_bias=routing_bias,
1675
hidden_states=hidden_states_fp4,
1676
hidden_states_scale=hidden_states_scale_linear_fp4.view(
@@ -1690,8 +1694,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1694
output2_scale_scalar=layer.g2_alphas.data,
1695
num_experts=global_num_experts,
1696
top_k=top_k,
1693
- n_group=num_expert_group if num_expert_group is not None else 0,
1694
- topk_group=topk_group if topk_group is not None else 0,
1697
+ n_group=num_expert_group,
1698
+ topk_group=topk_group,
1699
intermediate_size=layer.intermediate_size_per_partition,
1700
local_expert_offset=layer.ep_rank * layer.local_num_experts,
1701
local_num_experts=layer.local_num_experts,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+4
-1
@@ -291,5 +291,8 @@ def get_flashinfer_moe_backend() -> FlashinferMoeBackend:
291
292
def is_flashinfer_supporting_global_sf(backend: FlashinferMoeBackend | None) -> bool:
293
# TODO(shuw@nvidia): Update when new backends are added.
294
- backends_supporting_global_sf = (FlashinferMoeBackend.CUTLASS,)
294
+ backends_supporting_global_sf = (
295
+ FlashinferMoeBackend.CUTLASS,
296
+ FlashinferMoeBackend.TENSORRT_LLM,
297
+ )
298
return backend in backends_supporting_global_sf