[qwen4-main-squashed] Cherry-pick #36811: avoid zero-bias allocation in fused softmax routing (fixes GB10 NEXTN collapse) - #38308
Merged
Jiminator merged 1 commit intoSep 7, 2026
Conversation
(cherry picked from commit 221a627)
Jiminator
requested review from
BBuf,
DarkSharpness,
Edwardf0t1,
Fridge003,
HaiShaw,
HydraQYH,
Ying1123,
celve,
ch-wan,
ispobock,
merrymercy and
yuan-luo
as code owners
September 7, 2026 09:14
This was referenced Sep 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cherry-pick of #36811 (main commit 221a627) onto
qwen4-main-squashed. One trivial conflict resolved: this branch still had theassert bias.dtype == torch.float32line that #36811 deletes.Why
On DGX Spark (GB10) every NEXTN cell of Qwen3.8-Flash-Next-NVFP4 built from this branch sporadically collapses to repeated
!(NaN logits, #37111). Root cause:fused_topkpasses a freshtorch.zerosbias to the Triton router kernel, which loads it before its PDL wait, so under programmatic dependent launch it can read the buffer's previous bytes; on GB10 the MTP draft leaves NaN there. An instrumented run captured the kernel consuming NaN bias before the wait, zeros after it, with correct logits. #36811 removes the zero-bias allocation for softmax routing, so there is no buffer left to read early.Validation (1x DGX Spark, RadixArk/Qwen3.8-Flash-Next-NVFP4, NEXTN 3/1/4, radix cache off, 8 running requests, GSM8K 1319 + churn client, the shape that collapsed within 2 to 25 min in most runs on this branch tip)
With only this cherry-pick applied: 0.971 accuracy, 0 invalid, 1319/1319 stopped normally, 0 NaN events, 0 routing mismatches against a settled reference recomputed after the kernel. The unpatched branch collapsed twice in the same hours on the same shape.
The kernel-side hardening (wait before the bias load, Triton + radix) is proposed separately against main in #38290 and is not required for this branch's cells.
CI States
Latest PR Test (Base): ❌ Run #34104919156
Latest PR Test (Extra): ❌ Run #34104918908
Latest PR Test (AMD ROCm 7.2): ❌ Run #34104919185