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 14, 05:30 UTC
[Model][gpt-oss] Support DP+EP for GPT-OSS with FlashInfer trtllm-gen MoE (#23819)
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
Po-Han Huang (NVIDIA) committed
Aug 28, 2025 at 21:56 UTC
95089607fa307c5facfb9706ea919292fb56e78c
3 files changed
+14
-15
vllm/model_executor/layers/fused_moe/config.py
+8
-7
@@ -190,12 +190,6 @@ class FusedMoEParallelConfig:
190
return (self.use_all2all_kernels
191
and envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency")
192
193
- @property
194
- def use_flashinfer_cutlass_kernels(self):
195
- return (envs.VLLM_USE_FLASHINFER_MOE_FP4
196
- and has_flashinfer_cutlass_fused_moe()
197
- and envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput")
198
-
193
@staticmethod
194
def make(tp_size_: int, dp_size_: int,
195
vllm_parallel_config: ParallelConfig) -> "FusedMoEParallelConfig":
@@ -404,7 +398,14 @@ class FusedMoEConfig:
398
399
@property
400
def use_flashinfer_cutlass_kernels(self):
407
- return self.moe_parallel_config.use_flashinfer_cutlass_kernels
401
+ """
402
+ Whether to use FlashInfer cutlass kernels for NVFP4 MoE.
403
+ """
404
+ return (self.quant_config is not None
405
+ and self.quant_config.quant_dtype == "nvfp4"
406
+ and envs.VLLM_USE_FLASHINFER_MOE_FP4
407
+ and has_flashinfer_cutlass_fused_moe()
408
+ and envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput")
409
410
@staticmethod
411
def make(
vllm/model_executor/layers/fused_moe/layer.py
+4
-4
@@ -920,7 +920,7 @@ class FusedMoE(CustomOp):
920
self.batched_router_logits: Optional[torch.Tensor] = None
921
if (self.moe_parallel_config.use_pplx_kernels
922
or self.moe_parallel_config.use_deepep_ll_kernels
923
- or self.moe_parallel_config.use_flashinfer_cutlass_kernels):
923
+ or self.moe_config.use_flashinfer_cutlass_kernels):
924
self.batched_hidden_states = torch.zeros(
925
(moe.max_num_tokens, self.hidden_size),
926
dtype=moe.in_dtype,
@@ -974,7 +974,7 @@ class FusedMoE(CustomOp):
974
975
@property
976
def use_flashinfer_cutlass_kernels(self):
977
- return self.moe_parallel_config.use_flashinfer_cutlass_kernels
977
+ return self.moe_config.use_flashinfer_cutlass_kernels
978
979
def update_expert_map(self):
980
# ep_size and ep_rank should already be updated
@@ -1665,7 +1665,7 @@ class FusedMoE(CustomOp):
1665
# only when data parallelism (DP) is enabled.
1666
use_flashinfer_cutlass_kernels = (
1667
self.dp_size > 1
1668
- and self.moe_parallel_config.use_flashinfer_cutlass_kernels)
1668
+ and self.moe_config.use_flashinfer_cutlass_kernels)
1669
if (self.moe_parallel_config.use_pplx_kernels
1670
or self.moe_parallel_config.use_deepep_ll_kernels
1671
or use_flashinfer_cutlass_kernels):
@@ -1674,7 +1674,7 @@ class FusedMoE(CustomOp):
1674
do_naive_dispatch_combine: bool = (
1675
self.dp_size > 1
1676
and not self.moe_parallel_config.use_deepep_ht_kernels
1677
- and not self.moe_parallel_config.use_flashinfer_cutlass_kernels)
1677
+ and not self.moe_config.use_flashinfer_cutlass_kernels)
1678
if do_naive_dispatch_combine:
1679
hidden_states, router_logits = get_ep_group().dispatch(
1680
hidden_states, router_logits)
vllm/model_executor/layers/quantization/mxfp4.py
+2
-4
@@ -623,8 +623,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
623
624
if should_use_flashinfer_mxfp4():
625
from flashinfer import mxfp8_quantize, trtllm_fp4_block_scale_moe
626
- assert not self.moe.use_ep, (
627
- "EP is not supported for flashinfer mxfp4 moe backend yet.")
626
if _should_use_flashinfer_mxfp4_bf16():
627
assert x.dtype == torch.bfloat16
628
x_quant = x
@@ -650,12 +648,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
648
None, # output1_scale_scalar
649
None, # output1_scale_gate_scalar
650
None, # output2_scale_scalar
653
- self.num_experts,
651
+ global_num_experts,
652
top_k,
653
None, # n_group
654
None, # topk_group
655
self.intermediate_size, # padded to multiple of 256
658
- 0, # local_expert_offset
656
+ layer.ep_rank * layer.local_num_experts, # local_expert_offset
657
self.num_experts, # local num experts
658
None,
659
self._get_tile_tokens_dim(x, top_k),