[Kernel] Wait for the PDL dependency before loading the router bias (Triton + radix)

[Jiminator]

Targets qwen4-main-squashed, on top of #38308 (the cherry-pick of #36811, merged).

Motivation

Two PDL-launched routing kernels load their bias before the programmatic-dependent-launch wait and load scores after it:

• python/sglang/kernels/ops/moe/moe_fused_gate.py::_router_triton_kernel (Triton, launch_pdl=True on compute capability >= 9)
• python/sglang/kernels/jit/csrc/moe/route_radix.cuh::route_radix_block (the radix fast path that moe_fused_gate dispatches covered sigmoid inputs to; also reused by route_quant_fused.cuh)

Under PDL a grid may start once the preceding kernel's blocks have exited, before that kernel's stores are visible; the wait (griddepcontrol.wait) is what establishes completion and visibility. A load placed before it must not depend on prior work. Both kernels assumed the bias is a frozen weight, but the callers do not guarantee that: fused_topk's sigmoid path passes correction_bias.to(torch.float32) (a fresh tensor when the parameter is fp16/bf16) or a fresh torch.zeros when there is no correction bias (topk.py:1047-1054), the grouped and ungrouped wrappers cast as well, and before #36811 the softmax path passed a fresh torch.zeros on every call. The compiled sm_121 Triton kernel confirms the pre-fix order: four LDG.E.128 bias loads precede ACQBULK, the score loads follow.

This is the mechanism behind #37111 (Qwen3.8-Flash-Next NVFP4 with NEXTN on DGX Spark, output collapsing to repeated !), confirmed by capturing the values the kernel consumed during an organic collapse. With an instrumented copy of the unmodified kernel that additionally stores, per row, the bias it loaded before the wait, the scores it loaded after the wait, and a volatile re-read of the bias after the wait, an onset step (32 rows, 23 with NaN top-k weights) showed for every bad row: the bias loaded before the wait contained NaN (256 of 512 elements in most rows; in others the first 64 were zero and the next 64 to 128 NaN), the bias re-read after the wait was all zero, and the scores loaded matched the settled logits exactly. Every clean row had seen an all-zero bias before the wait. Upstream of the router the block input and the gate logits were finite; downstream, the expert kernel produced NaN in exactly those rows and the batch collapsed. #36811 already removed the fresh-zero trigger from plain softmax routing; dynamic-bias paths still need correctly ordered loads.

Modifications

Move the wait above the bias load in both kernels so every dependent load sits after it. No change to routing semantics, dispatch, launch geometry, or the outgoing gdc_launch_dependents(); the HAS_BIAS=False register-zero path is unaffected.

Accuracy / performance

On 1x DGX Spark (GB10), RadixArk/Qwen3.8-Flash-Next-NVFP4, NEXTN 3/1/4, radix cache off, 8 running requests, GSM8K 1319 (chat, greedy, 8192 max tokens) plus a churn client, on the pre-#36811 build 9b2aee2283 where the collapse reproduced within 2 to 25 minutes in most runs: with only the Triton reorder applied (PDL on, zero-bias still allocated) the full set completed at 0.967 with 0 invalid answers, 0 NaN events, and 0 routing mismatches against a reference top-k recomputed after the kernel on the same stream. Two workarounds that also remove the dependency (PDL off for the kernel; SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK=0) each completed the same set clean (0.970 / 0.971). Author-reported; the radix change was not exercised by these runs (softmax routing does not dispatch to radix).

Cost: the bias load no longer overlaps the tail of the previous kernel, so a program pays that load's latency after the wait instead of hiding it. Per launch, not per row (rows run concurrently); zero on the softmax path with #36811 since no bias is loaded. Not measured here.

Follow-ups (not in this PR)

• A deterministic regression check: compile a bias-present specialization and assert every ld.global in the PTX comes after griddepcontrol.wait.
• moe_fused_gate.py grouped routing removes all tied group maxima when computing the second maximum (vals2 = tl.where(vals >= top1, ...)), so tied top-1 values yield the wrong group score. Pre-existing, unrelated to this change.
• kernels/jit/include/sgl_kernel/utils.cuh:170-172 describes the PDL trigger as covering writes issued before it; it is a launch hint, and consumers must wait.

Checklist

• CI kernel tests (test/registered/kernels/ops/moe/test_moe_fused_gate.py, test/registered/kernels/ops/kimi_k3/test_compute.py)

---

CI States

Latest PR Test (Base): ❌ Run #34108740533
Latest PR Test (Extra): ❌ Run #34108740503
Latest PR Test (AMD ROCm 7.2): ❌ Run #34108740687

[BBuf]

LGTM.