spec(GDN-MOE-BF16-OUT): the GDN recurrence output and the z gate are f32 on every MoE checkpoint, and both references keep them bf16 - #1173
Merged
Conversation
…f32 on every MoE checkpoint, and vLLM keeps both bf16 (#1168, #1169, #1170) `GdnOutDType(bool dense_model)` resolves bf16 for a dense checkpoint and f32 for a MoE one, and all three call sites pass `cfg.num_experts == 0`. So on the MoE arms the recurrence output `dcore`, the `z` gate and the gated-RMSNorm weight that must match them all stay f32, which doubles four passes over the two largest per-layer GDN activations. vLLM allocates `core_attn_out` at `hidden_states.dtype` and slices `z` out of the bf16 in_proj output, with no dense-versus-MoE branch anywhere. So this is a deviation from the PRIMARY oracle rather than an SGLang-only difference, and it is the exact shape AGENTS.md describes when it says a token gate cannot detect a dtype that is too wide: the 35B gate is 315/315 BECAUSE f32 is the more precise error. The tree already derived the whole finding in the comment above the function and deferred it to "its later independently gated campaign". That campaign was never opened. This spec opens it, and it is spec-only: no product code, no measurement. Two findings changed the shape of the row while writing it, so both are filed with an owner rather than carried as assumptions. Removing the redundant `dense_model` term from the packed-decode eligibility does NOT reach packed decode on a MoE model, because only the dense loader builds the merged `in_proj_ba` owner (#1169). And all four GDN Triton AOT arms are pinned to 48 or 32 linear V-heads, so the 2.4T's 128 misses every one of them and runs the hand kernels the vendoring exists to replace (#1170). Neither is fixed here; both are listed under `## Owed`. The gap was verified against the four unmerged local dtype branches before the spec was written. None narrows this default. Two collide with it, and the spec records both: `row/REFACTOR-DTYPE-CONSISTENCY` relocates the function verbatim, and `row/PERF-GDN-BF16-CHAIN` uses the MoE f32 default as an eligibility bound that this change removes. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5-1m [Claude Code]
Spec and issue-index rows for #1168: GDN core/z stay f32 on every MoE checkpoint while both references keep them at the bf16 model dtype. Adds #1169 (packed GDN decode has a second, loader-side blocker) and #1170 (the Triton AOT arms are pinned to 48 or 32 linear V-heads) as owed debt. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude:Opus-5 [Claude Code]
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.
The GDN recurrence output
dcoreand thezgate are held in f32 on every MoEcheckpoint, while both reference engines keep them at the bf16 model dtype.
GdnOutDTyperesolvesbf16 = dense_modeland every call site passescfg.num_experts == 0, so the dense arm is already bf16 and only the MoE armpays. The cost is the largest per-layer GDN activation carried at double width
twice: written by the recurrence, read by the gated RMSNorm.
vLLM keeps both at the model dtype, and so does SGLang. A token gate cannot see
this: the tokens match and the goldens pass while the path moves twice the
bytes, which is the dtype-polarity case the protocol calls out by name.
The row is spec-only. No product code changes here, and every gate is PENDING.
Two findings surfaced while verifying the premise, and both are filed rather
than folded in. Narrowing the dtype does not on its own make packed GDN decode
reachable on MoE: the eligibility also requires
has_packed_ba, andin_proj_bais written at exactly one site, the dense loader, so the MoE andGGUF loaders leave it split. That is a second, loader-side blocker and it is
tracked separately. The Triton AOT arms are pinned to 48 or 32 linear V-heads,
so the 2.4T checkpoint misses all four of them and runs the hand kernels; the
27B is unaffected because it carries 48. Both are recorded as owed debt with a
named owner rather than fixed in this flow.
Landing this also loosens an assumption on an unmerged branch: once the MoE arm
can be bf16,
PERF-GDN-BF16-CHAIN's third eligibility term stops being boundedby model shape and is bounded only by its default-OFF toggle. No behaviour moves
in either merge order, so it is a review obligation, logged under
## Owed.Refs #1168, #1169, #1170, #521.
FOLLOWING_AGENTS_PROTOCOL
Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Claude:Opus-5 [Claude Code]