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

[Jiminator]

Motivation

PDL-launched routers can read a stale bias produced by a preceding cast or fill kernel. Both the Triton router and radix router currently load bias before waiting for the PDL dependency, although scores are loaded after the wait. The sigmoid/grouped callers can supply freshly cast bias tensors, so this is not limited to Qwen.

Upstreams #38290 onto main after #37500. Related reports: #37111 and #37052.

Modifications

Move the existing dependency wait ahead of both input loads in the Triton and radix kernels. No additional allocation, launch, load, or wait. The no-bias path continues to use register zeros.

Accuracy Tests

September 11, 2026: current-main combined deployment validation

Tested #38568 + #38569 + #38570 together on main@165d8dd17736fa4938f2e73018f26835cd6e0778, combined merge 69d38ccbd8a03646cae9168c971b3783883c2903. PR heads: #38568 626f614f7bcd5dc531de5c268a694ff1a9f31f28, #38569 825663897b1fd9ee5c341eeb113b795e4ecfc3f8, #38570 b14d2c55f37f8d3dc2b5fe21175b56fbe62771e1. All merged cleanly; no additional source fixes were needed.

Native host installation on 1x DGX Spark / GB10, TP=1, rebuilt from the merged checkout without a patched Docker image. Model: nvidia/Qwen3.8-Flash-Next-NVFP4, checkpoint revision fab0aecb760cec45227f6656abcaafa11abca87a. GSM8K first 200 test questions, greedy chat API, thinking disabled, 8,192-token generation limit:

Configuration	Accuracy	Request errors / invalid / truncated	Ready	Aggregate output
MTP, concurrency 8	195/200 (97.5%)	0 / 0 / 0	763.26 s	92.87 tokens/s
Non-MTP, concurrency 24	194/200 (97.0%)	0 / 0 / 0	633.17 s	89.53 tokens/s

Both smoke responses were correct and CUDA graphs were active. MTP used the expected FP8 Triton expert fallback and reported average draft acceptance 3.52925. No serving NaN/OOM/CUDA-error/exception reports appeared before intentional shutdown; shutdown cancellation diagnostics were recorded separately.

Serving flags: TP=1, FlashInfer CUTLASS MoE/FP4 backends, page size 64, chunked prefill 4096, requested context 262144, static memory fraction 0.85, and file-backed PLE with a fresh local-NVMe file per boot. MTP: NEXTN, 3 steps, top-k 1, 4 draft tokens, 8 maximum requests, 40 Mamba slots. Non-MTP: extra_buffer_lazy, 24 maximum requests, 96 Mamba slots.

The isolated host environment was updated from current main and native extensions rebuilt (PyTorch 2.13.0+cu130, FlashInfer 0.6.18, Triton 3.7.1, sglang-kernel 0.4.6.post1, Transformers 5.12.1). Runs reused previously compiled host kernels with MAX_JOBS=2; startup timings exclude cold kernel compilation. Throughput includes prefill and the lower-concurrency tail and is not an isolated measurement of this PR's speed impact.

Scope: combined TP=1 deployment validation, not independent deterministic reproduction of the PDL race, TP>1 prefetch validation, or the full 1,319-question benchmark. Effective max input was 180,282 tokens with MTP and 262,138 without MTP; no full-length 262K prompt was tested. Both servers were stopped and their generated PLE files removed after evaluation.

Latest-main necessity check: source-reviewed main@45715e7f20629d1813ff8f543a23bb437128db2e after the deployment run. Both radix and Triton routers still load bias before their PDL dependency wait. The combined PRs merge cleanly onto this newer revision. The accuracy numbers above remain measurements of 165d8dd177, not a GPU rerun on this newer base.

Earlier validation

• On this port, local GB10: targeted test_moe_fused_gate.py reference comparisons, no-bias/zero-bias equivalence, and non-fp32 bias tests: 15 passed. These check routing semantics, not deterministic reproduction of the race.
• All pre-commit hooks for changed files pass.
• Prior validation recorded for [Kernel] Wait for the PDL dependency before loading the router bias (Triton + radix) #38290 (not rerun as a full-model evaluation on this port): the unfixed sm121 cubin had four bias LDG.E.128 instructions before ACQBULK; fixed PTX had griddepcontrol.wait before all six global loads. During an instrumented collapse, all 23 bad rows in a 32-row target-verify step saw NaNs in the pre-wait bias and zeros on a post-wait reread, with settled scores unchanged.
• Prior full GSM8K run on 1x DGX Spark, 1,319 questions, greedy chat API, 8,192-token budget, NEXTN 3/1/4, radix cache disabled, 8 running requests plus churn: 0.967 accuracy, 0 invalid answers, 0 NaN events, 0 routing mismatches, with PDL enabled and only the wait reorder.

Combined deployment validation with #38569 and #38570 also passed on a fresh-host GB10: 200-question GSM8K samples scored 98.0% (MTP) / 97.0% (non-MTP) on base 52fecfdf09, with no invalid/capped answers or request errors. Both commands also booted and served 24-request sanity samples on newer base db272201a2. Full environment and result details are in #38569. These combined model runs do not deterministically reproduce the bias-carrying PDL race.

Speed Tests and Profiling

No new performance measurement on this port. The bias load no longer overlaps the preceding kernel's tail; the exposed load latency is per launch, not multiplied by row count. The plain-softmax no-bias path has no bias load. Performance on bias-carrying paths remains to be measured.

Checklist

• Format changed code with pre-commit.
• Run targeted existing correctness tests.
• Run full-model accuracy and speed benchmarks on the rebased branch.

---

CI States

Latest PR Test (Base): ❌ Run #34592943620
Latest PR Test (Extra): ❌ Run #34592943490
Latest PR Test (AMD ROCm 10): ❌ Run #34592943551

[Jiminator]

Superseded by combined draft PR #39126, which preserves all three component changes and includes the DGX Spark MTP/non-MTP accuracy and throughput results. Closing this separate PR in favor of the combined review.