From 111c739c9cbeaa22ee17b4627c5c604d22c3769a Mon Sep 17 00:00:00 2001 From: Jimmy Shong Date: Sun, 6 Sep 2026 23:35:36 -0700 Subject: [PATCH 1/2] [Kernel] moe_fused_gate: wait for PDL dependency before loading the bias The Triton router kernel prefetched `bias` before `gdc_wait()` and loaded `scores` after it. The bias is not an immutable input: `fused_topk` passes a buffer written by the immediately preceding kernel (a dtype cast of the correction bias today; a fresh `torch.zeros` before #36811). Under programmatic dependent launch the router can start once that kernel's blocks exit, before its stores are visible, so the pre-wait load could read the buffer's previous contents. On DGX Spark (GB10) with Qwen3.8-Flash-Next NVFP4 + NEXTN this surfaced as batch-wide NaN routing weights and an output collapse to token 0 (#37111). Move the wait above the bias load so every dependent load sits after it. Co-Authored-By: Claude Fable 5.1 --- python/sglang/kernels/ops/moe/moe_fused_gate.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index b60ecce2ec54..f5a89a5aeea7 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -127,17 +127,22 @@ def _router_triton_kernel( mask_m = offs_m < M mask_n = offs_n < N - # Prefetch a real bias before the PDL wait. Plain softmax routing has no - # bias, so keep the zero value in registers rather than materializing and - # clearing a device tensor for every routing call. + # With programmatic dependent launch this grid may start before the + # preceding kernel's stores are visible, so every load that can depend on + # prior work must come after the wait. The bias is such an input: callers + # pass buffers produced right before this launch (a dtype cast of the + # correction bias, or historically a fresh torch.zeros), and a load placed + # before the wait read the buffer's previous contents. + if USE_PDL: + tl.extra.cuda.gdc_wait() + + # Plain softmax routing has no bias, so keep the zero value in registers + # rather than materializing and clearing a device tensor per call. if HAS_BIAS: bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) else: bias = tl.zeros([BLOCK_N], dtype=tl.float32) - if USE_PDL: - tl.extra.cuda.gdc_wait() - row_ptr = scores_ptr + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn mask2d = mask_m[:, None] & mask_n[None, :] scores = tl.load(row_ptr, mask=mask2d, other=0.0).to( From 522140095486af1ee5af3720e5a0c923b324b261 Mon Sep 17 00:00:00 2001 From: Jimmy Shong Date: Mon, 7 Sep 2026 00:34:34 -0700 Subject: [PATCH 2/2] [Kernel] route_radix: wait for PDL dependency before loading the bias Same ordering defect as in the Triton router: route_radix_block prefetched the bias before PDLWaitPrimary(), assuming a frozen weight, but the public moe_fused_gate dispatches covered sigmoid inputs here with a bias that can be a fresh fp32 cast or a fresh torch.zeros. Move the wait above the load. Also shorten the Triton kernel comment to state only what is established. Co-Authored-By: Claude Fable 5.1 --- python/sglang/kernels/jit/csrc/moe/route_radix.cuh | 5 +++-- python/sglang/kernels/ops/moe/moe_fused_gate.py | 9 +++------ 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/moe/route_radix.cuh b/python/sglang/kernels/jit/csrc/moe/route_radix.cuh index 0c39f604794c..492d1b980646 100644 --- a/python/sglang/kernels/jit/csrc/moe/route_radix.cuh +++ b/python/sglang/kernels/jit/csrc/moe/route_radix.cuh @@ -142,9 +142,10 @@ SGL_DEVICE void route_radix_block(const RouteRadixParams& params, typename Large // radix math below is fp32 either way — only the load width differs. AlignedVector, kVecSize / 2> scores_vec; - // prefetch bias (frozen weight) before the PDL wait - bias_vec.load(params.bias, tx); + // Bias may be produced by a preceding cast or fill kernel (the caller + // does not guarantee a frozen weight), so wait before loading either input. PDLWaitPrimary(); + bias_vec.load(params.bias, tx); scores_vec.load(scores, tx); #pragma unroll diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index f5a89a5aeea7..d17227e7b9b7 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -127,12 +127,9 @@ def _router_triton_kernel( mask_m = offs_m < M mask_n = offs_n < N - # With programmatic dependent launch this grid may start before the - # preceding kernel's stores are visible, so every load that can depend on - # prior work must come after the wait. The bias is such an input: callers - # pass buffers produced right before this launch (a dtype cast of the - # correction bias, or historically a fresh torch.zeros), and a load placed - # before the wait read the buffer's previous contents. + # PDL may start this grid before prior kernel stores are visible. Bias can + # be produced by a preceding cast or fill kernel, so wait before loading + # either bias or scores. if USE_PDL: tl.extra.cuda.gdc_wait()