diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index e0a39cbae..a8bf5668d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -51,6 +51,12 @@ _deferred_finalize_enabled: contextvars.ContextVar[bool] = contextvars.ContextVa _TRTLLM_MOE_PDL_MAX_TOKENS = envs.SGLANG_TRTLLM_MOE_PDL_MAX_TOKENS.get() +def trtllm_moe_enable_pdl(num_tokens: int) -> bool: + from sglang.kernels.jit.utils import is_arch_support_pdl + + return is_arch_support_pdl() and num_tokens <= _TRTLLM_MOE_PDL_MAX_TOKENS + + @dataclass class FlashInferTrtllmDeferredFinalizeOutput: gemm2_out: torch.Tensor @@ -74,7 +80,6 @@ def finalize_flashinfer_trtllm_deferred_output( deferred_output: FlashInferTrtllmDeferredFinalizeOutput, shared_output: torch.Tensor, ) -> torch.Tensor: - from sglang.kernels.jit.utils import is_arch_support_pdl from sglang.kernels.ops.moe.moe_finalize_fuse_shared import moe_finalize_fuse_shared return moe_finalize_fuse_shared( @@ -83,7 +88,7 @@ def finalize_flashinfer_trtllm_deferred_output( deferred_output.expert_weights, shared_output, deferred_output.top_k, - enable_pdl=is_arch_support_pdl(), + enable_pdl=trtllm_moe_enable_pdl(deferred_output.expert_weights.shape[0]), ) @@ -1087,7 +1092,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( activation_type=activation_type, tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), output=symm_output, - enable_pdl=hs_fp4.shape[0] <= _TRTLLM_MOE_PDL_MAX_TOKENS, + enable_pdl=trtllm_moe_enable_pdl(hs_fp4.shape[0]), )[0] else: assert TopKOutputChecker.format_is_bypassed(topk_output) @@ -1131,7 +1136,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( do_finalize=not defer_finalize, activation_type=activation_type, tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), - enable_pdl=hs_fp4.shape[0] <= _TRTLLM_MOE_PDL_MAX_TOKENS, + enable_pdl=trtllm_moe_enable_pdl(hs_fp4.shape[0]), ) if not defer_finalize: moe_kwargs["output"] = symm_output diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 4ad7bfc39..b3136cff4 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -1434,6 +1434,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): if self._fi_kernel == "cutlass_sm120": return self._apply_sm120_cutlass(layer, dispatch_output) if self.use_flashinfer: + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + trtllm_moe_enable_pdl, + ) + # When bf16 mode is enabled, we don't need to quantize the input, # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations, # which can theoretically improve performance @@ -1591,6 +1595,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), output=symm_output, do_finalize=not defer_finalize, + enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), ) if defer_finalize: from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( @@ -1650,6 +1655,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): local_num_experts=layer.num_local_experts, tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), output=symm_output, + enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), ) return StandardCombineInput(hidden_states=symm_output) @@ -1682,6 +1688,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): True, # do finalize tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), output=symm_output, + enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), )[0] return StandardCombineInput(hidden_states=trtllm_gen_output) if _use_aiter: diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index dbdcc713b..f4ec27acc 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -319,6 +319,10 @@ class Mxfp4FlashinferTrtllmMoEMethod: else: raise NotImplementedError(f"Unsupported mxfp4 moe precision: {precision}") + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + trtllm_moe_enable_pdl, + ) + with use_symmetric_memory( get_tp_group(), disabled=not is_allocation_symmetric() ): @@ -361,6 +365,7 @@ class Mxfp4FlashinferTrtllmMoEMethod: do_finalize=True, tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), output=symm_output, + enable_pdl=trtllm_moe_enable_pdl(num_tokens), )[0] return StandardCombineInput(hidden_states=output)