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, 23:00 UTC
[BugFix] Fix TRT-LLM NVFP4 DP/EP (#32349)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
jiahanc committed
Jan 19, 2026 at 11:32 UTC
73503317188e4b6430ab2857d98546b5b432cb43
4 files changed
+21
-11
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml
new
+8
@@ -0,0 +1,8 @@
1
+model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
6
+env:
7
+ VLLM_USE_FLASHINFER_MOE_FP4: "1"
8
+ VLLM_FLASHINFER_MOE_BACKEND: "latency"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
+1
@@ -1,6 +1,7 @@
1
Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml
2
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml
3
Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml
4
+Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml
5
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml
6
Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ht.yaml
7
Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml
vllm/model_executor/layers/fused_moe/layer.py
+1
-8
@@ -53,7 +53,6 @@ from vllm.model_executor.layers.quantization.base_config import (
53
QuantizationConfig,
54
)
55
from vllm.platforms import current_platform
56
-from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
56
from vllm.utils.math_utils import cdiv, round_up
57
from vllm.utils.torch_utils import (
58
aux_stream,
@@ -1761,17 +1760,11 @@ class FusedMoE(CustomOp):
1760
with sp_ctx:
1761
extra_tensors = None
1762
if do_naive_dispatch_combine:
1764
- # Avoid circular import
1765
- from vllm.model_executor.layers.quantization.modelopt import (
1766
- ModelOptNvFp4FusedMoE,
1767
- )
1768
-
1763
post_quant_allgather = (
1764
self.quant_method is not None
1765
and self.dp_size > 1
1766
and self.use_ep
1773
- and isinstance(self.quant_method, ModelOptNvFp4FusedMoE)
1774
- and has_flashinfer_trtllm_fused_moe()
1767
+ and getattr(self.quant_method, "do_post_quant_allgather", False)
1768
)
1769
if post_quant_allgather:
1770
hidden_states_to_dispatch, extra_tensors = (
vllm/model_executor/layers/quantization/modelopt.py
+11
-3
@@ -1564,6 +1564,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1564
moe_config=self.moe,
1565
)
1566
1567
+ @property
1568
+ def do_post_quant_allgather(self):
1569
+ return self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM
1570
+
1571
def prepare_dp_allgather_tensor(
1572
self,
1573
layer: FusedMoE,
@@ -1571,13 +1575,17 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1575
router_logits: torch.Tensor,
1576
) -> tuple[torch.Tensor, list[torch.Tensor]]:
1577
"""Optionally prepare extra tensors to carry through DP allgather/EP."""
1578
+ if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM:
1579
+ raise RuntimeError(
1580
+ "prepare_dp_allgather_tensor is only supported for "
1581
+ "FlashInfer TRTLLM NVFP4 MoE backend."
1582
+ )
1583
+
1584
import flashinfer
1585
1576
- assert self.moe_quant_config is not None
1577
- a1_gscale = self.moe_quant_config.a1_gscale
1586
hidden_states_fp4, hidden_states_sf = flashinfer.fp4_quantize(
1587
hidden_states,
1580
- a1_gscale,
1588
+ layer.a1_gscale,
1589
is_sf_swizzled_layout=False,
1590
)
1591
extra_tensors: list[torch.Tensor] = [hidden_states_sf]