[1/5] refactor(ops): let the sparse-MoE decode codecs own their matrix addressing - #286
Open
MichaelDementii wants to merge 1 commit into
Open
MichaelDementii wants to merge 1 commit into
MichaelDementii wants to merge 1 commit into
Conversation
…essing The four decode kernels addressed a matrix from the call site: a flat group index built from the row, and three plane pointers picked by hand. That holds while every codec keeps one scale per group per row. It does not generalise - a codec whose scale plane is swizzled needs the row to form an address, and the call site cannot compute it without already knowing the codec. A codec now receives the matrix as a plane set plus a row, with the column count as a template argument. The shared expert becomes a codec parameter in the same move: every registered profile stores it as Q8 today, but that is a property of the profiles rather than of the kernels, and the kernels no longer assert it. No behaviour change. The arithmetic, the accumulation order and the launch geometry are untouched for all four registered codecs.
Contributor
Author
|
This one is a pure refactor: the decode dot product stops asking each codec how many values a lane @codex review |
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
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.
Related Issue: #283
First of five. No NVFP4 in it — this is the prerequisite the rest stands on. No performance claim.
What changes
A codec now receives the matrix as a plane set plus a row, with the column count as a template
argument, and addresses itself. Before, the call site built a flat group index from the row and
picked the plane pointers by hand, which holds only while every codec keeps one scale per group per
row: a swizzled scale plane needs the row to form an address, and the call site cannot compute it
without already knowing the codec. Both lane ownerships move, packed and scalar.
The shared expert becomes a codec parameter in the same move.
Q8Codecwas hardcoded in the sharedbranch of all four kernel bodies; every registered profile does store it as Q8, but that is a
property of the profiles, not of the kernels. This is the compile-time specialization point you
raised on #99, applied here.
Both are prerequisites, not cleanup: a profile whose scale plane is swizzled and whose shared expert
is not Q8 cannot be added without them.
Contract
Internal to
src/ops/sparse_moe/decode; no public signature changes, so no caller moves. Onebehavioural detail worth naming: the shared expert's
highplane is now passed rather than null.Q8Codecdiscards it, so nothing reads it, but it is a change in what is handed over.Verification
This commit is on the live decode path of every current MoE profile, and the suite's criterion is a
tolerance, so it cannot separate a refactor that moved a bit from one that did not.
The gate is a witness: an FNV-1a 64 over the decoded output of every operator case, read from two
binaries — this commit and its parent. Over the suite's sixteen cases on Q4+Q5, Q4+Q6 and Q8+Q8 they
are identical on all sixteen. The witness is a scratch patch applied to both arms for the gate
and is not part of this commit; nothing in the diff reads the environment.
Strength control on each of the two decode kernels: with the D3 pair made to read its first row
twice the witness moves on 6 of 16 cases, and with the D4 block's row index pinned to its base it
moves on 2 of 16; both fail the test. Two weaker controls moved nothing and bound what the gate
covers —
a + b*cis re-fused by the compiler, and a reversed accumulation order is absorbed by the BF16output the hash is taken over. It resolves addressing and scaling, not FP32 reassociation.
The three pairs are every profile the engine registers, and between them they cover every path: D3
reaches
load_eightthrough Q4 andload_onethrough the Q8 shared expert, D4 reachesload_eightthrough Q5/Q6 and
load_pairthrough Q8.ctest120 of 120 at this commit,pytestunchanged at 47 — this commit touches no Python.clang-format --dry-run(21.1.8): 0 notes on this branch and 0 on master. RTX 5090sm_120a, driver 616.64, CUDA 13.1, Release, WSL2.Checks not run
decode kernels, which gain a template parameter each. The change is meant to be output-identical
and the witness is the evidence that it is.
ctestwas run on this commit, not on its parent from the same build directory.🤖 Generated with Claude Code