[MoE] Route every trtllm-gen MoE call site through one PDL guard

[kpham-sgl]

Motivation

PDL on the trtllm-gen MoE can leave its grid-dependency wait unreleased when another stream overlaps the launch. The stalled rank never reaches its next collective, so the entire TP group hangs. 6a1d1f4422 (landed in #31681) identified this and capped PDL by token count via SGLANG_TRTLLM_MOE_PDL_MAX_TOKENS — but only at the two fp4 call sites in moe_runner/flashinfer_trtllm.py.

The mxfp4 call sites never passed enable_pdl at all. FlashInfer then defaults it to device_support_pdl() — on for SM90+, with no ceiling. The env var cannot reach those paths at any setting.

That includes the Kimi-K3 SiTU routed path, where a TP8 SM103 engine hangs during long-context prefill at the default --chunked-prefill-size 16384, twice the existing 8192 threshold.

Evidence

From an 8-rank CUDA coredump taken by the watchdog while the ranks were still wedged:

• All 5,296 warps of the finalize kernel (moe::dev::finalize::finalizeKernelVecLoad<KernelParams<bfloat16_t,bfloat16_t,4,true>> — trailing true is UsePdl) sit at one PC: +5200, the ACQBULK that implements cudaGridDependencySynchronize(). Zero divergence across the whole grid.
• Its producer GEMM has 26 of 82,236 CTAs resident. All 148 SMs are at their block limit: 122 SMs × 5 finalize CTAs + 26 SMs × (2 finalize + 1 GEMM) = 662 finalize CTAs, exactly as observed.
• The producer runs 2-CTA clusters (c2x1x1); every cluster pair occupies consecutive SM ids, so placement needs a co-schedulable SM pair that never appears.
• The other seven ranks spin in all_reduce_pull_norm_kernel on an unbounded semaphore loop with no timeout, waiting for the eighth participant.
• The finalize grid is launched with numBlocks = data.numTokens (trtllm_fused_moe_dev_kernel.cu:1034), one block per token.

Changes

Add trtllm_moe_enable_pdl(num_tokens) next to the threshold it reads, and route every call site through it so a new one cannot silently inherit PDL-on:

• mxfp4.py — SiTU routed (the path that hangs), SiTU bypassed-topk, and block-scale (GPT-OSS)
• mxfp4_flashinfer_trtllm_moe.py — fp4 routed
• the deferred-finalize path, which fuses the shared-expert add into sglang's own moe_finalize_fuse_shared and gated PDL on is_arch_support_pdl() alone. Kimi-K3 takes this branch, and its PDL consumer is a different kernel from FlashInfer's in-op finalize, so threading enable_pdl into the trtllm call would not have covered it.
• the two already-capped sites, now sharing the one definition

trtllm_bf16_moe is deliberately untouched: no call site passes enable_pdl today, so its signature is unverified from this tree.

Threshold value is out of scope

The coredump suggests 8192 may be too high: the finalize kernel's SM footprint is clamped by occupancy at ~740 resident blocks (148 SMs × 5) rather than scaling with token count, so an 8192-token grid floods the machine as thoroughly as the 16384-token one that hangs. That would leave 740–8192 nominally capped but still in the flood regime.

It is not settled enough to move the constant. tune_max_num_tokens=next_power_of_2(num_tokens) may select a different GEMM tactic at 8192 than at 16384, which would account for the soak result behind the current default without the occupancy model holding.

This PR restores the knob's reach. Retuning it belongs in a separate change, after a tactic comparison at both sizes.

Testing

[TODO]

🤖 Generated with Claude Code

[kpham-sgl]

/rerun-test registered/unit/layers/quantization/test_nvfp4_moe_backends.py registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py

[github-actions]

Results for /rerun-test registered/unit/layers/quantization/test_nvfp4_moe_backends.py registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py:

🚀 4-gpu-b200 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_nvfp4_moe_backends.py
cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py
cd test/ && python3 registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py