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 2, 17:01 UTC
[Bugfix] Pass `routed_scaling_factor` to FlashInfer TRTLLM BF16 MoE (#43769)
Thien Tran committed
May 28, 2026 at 12:29 UTC
e54eff769dea26f00b122be89ebaa870875da03e
2 files changed
+6
-2
tests/kernels/moe/test_moe.py
+5
-2
@@ -1656,7 +1656,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
1656
layer.routing_method_type = RoutingMethodType.Renormalize
1657
layer.expert_map = None
1658
layer.apply_router_weight_on_input = False
1659
- layer.routed_scaling_factor = None
1659
+ layer.routed_scaling_factor = 2.446
1660
layer.shared_experts = None
1661
layer._expert_routing_tables = lambda: None
1662
@@ -1678,7 +1678,10 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
1678
# Compute torch baseline
1679
w1_original = w1.clone()
1680
w2_original = w2.clone()
1681
- baseline_output = torch_moe(a, w1_original, w2_original, router_logits, topk)
1681
+ baseline_output = (
1682
+ torch_moe(a, w1_original, w2_original, router_logits, topk)
1683
+ * layer.routed_scaling_factor
1684
+ )
1685
1686
close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85)
1687
assert close.float().mean() > 0.925
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+1
@@ -137,5 +137,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
137
intermediate_size=self.intermediate_size_per_partition,
138
local_expert_offset=self.ep_rank * self.local_num_experts,
139
local_num_experts=self.local_num_experts,
140
+ routed_scaling_factor=routed_scaling_factor,
141
routing_method_type=self.routing_method_type,
142
)