Add an explicit RMSNorm kernel-family contract - #19
Closed
kiddyboots216 wants to merge 1 commit into
Closed
Conversation
Broly Security ScanNote ✅ Clean scan Note Re-scan this PR anytime with
|
Two batch-invariant rmsnorm kernel families coexist (looped tl.sum + 1.0/tl.sqrt vs mean_dim + tl.rsqrt) and disagree at 1 ulp on rare bf16 boundary values, so an implicit family flip against serving is a silent K3 seed (the 2026-07-04 norm-seed incident: 2.99e-5 from five such seeds). Make the family an explicit declaration instead of a consequence of how the call site is written. - batch_invariant_ops: RMSNormFamily constants plus the bi_rms_norm / bi_fused_add_rms_norm funnels. The aten::rms_norm interpose, the fused autograd wrappers, and the family-1 trunk-contract wrapper route through them. Kernels untouched. - The Qwen3.5 zero-centered twin registers as family-1 with a zero_centered fold in the funnel (same reduction tree, fp32 1+weight fold) rather than becoming a third family. - RMSNorm: family declaration at construction or per call; violations raise; undeclared calls in a parity lane warn once and raise under XORL_RMSNORM_REQUIRE_FAMILY=1. Qwen3, Qwen3-MoE and the shared attention qk-norms declare their site families. - tests: family contract guards plus cross-engine bitwise gates against SGLang's dispatched kernels per site-class, on adversarial shapes including the [4096, 128] 1-ulp seed shape. batch_invariant_ops.py is new here and carries only the surface this contract needs: the shared batch-invariant kernels and mode plumbing, the rmsnorm kernels, and the family funnels. The shape-keyed GEMM config table, the fused LM-head logprob, and the MoE router GEMM are separate contracts and are not included.
kiddyboots216
force-pushed
the
pr/k3-rmsnorm-family-contract
branch
from
July 29, 2026 03:45
cc6e633 to
5f718cb
Compare
qywu
self-requested a review
July 30, 2026 03:38
qywu
approved these changes
Jul 30, 2026
Contributor
Author
|
Superseded by #22, which consolidates the explicit RMSNorm-family dispatch with the shared numerical contract on the repaired Foundation branch. This standalone copy should not be merged independently. |
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.
Which RMSNorm kernel runs was decided by whether a call site passed a residual argument: with one, the
mean_dimreduction tree; without, the looped kernel. The two sum the hidden dimension in a different order, so rewriting a call site could move a layer between kernels without changing any declaration, leaving trainer and sampler on different reduction orders for the same layer.A site now declares its family explicitly. Dispatch is bit-identical to the
force_sglang_residualexpressions the call sites used to carry, in every combination of rmsnorm and batch-invariant mode.Sampler half:
togethercomputer/xorl-sglangbranchpr/k3-rmsnorm-family-contract.Self-contained by design — carries its own
batch_invariant_ops.pyand conflicts with the sibling contract PRs.