[MoE] Route every trtllm-gen MoE call site through one PDL guard (#34789)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

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:

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)