Skip to content

[1/5] refactor(ops): let the sparse-MoE decode codecs own their matrix addressing - #286

Open
MichaelDementii wants to merge 1 commit into
Neroued:masterfrom
MichaelDementii:prep/moe-decode-codec-contract
Open

MichaelDementii wants to merge 1 commit into
Neroued:masterfrom
MichaelDementii:prep/moe-decode-codec-contract

Conversation

@MichaelDementii

Copy link
Copy Markdown
Contributor

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. Q8Codec was hardcoded in the shared
branch 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. One
behavioural detail worth naming: the shared expert's high plane is now passed rather than null.
Q8Codec discards 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*c is re-fused by the compiler, and a reversed accumulation order is absorbed by the BF16
output 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_eight through Q4 and load_one through the Q8 shared expert, D4 reaches load_eight
through Q5/Q6 and load_pair through Q8.

ctest 120 of 120 at this commit, pytest unchanged at 47 — this commit touches no Python.
clang-format --dry-run (21.1.8): 0 notes on this branch and 0 on master. RTX 5090 sm_120a, driver 616.64, CUDA 13.1, Release, WSL2.

Checks not run

  • Nothing measured, and in particular no register, shared-memory or spill comparison for the four
    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.
  • ctest was run on this commit, not on its parent from the same build directory.
  • One card, one architecture.

🤖 Generated with Claude Code

…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.
@MichaelDementii

Copy link
Copy Markdown
Contributor Author

This one is a pure refactor: the decode dot product stops asking each codec how many values a lane
owns at run time and takes it from a kSingleValuePerLane trait, with a trap in the final else so
a codec that fits neither ownership shape fails to compile rather than silently dropping work. The
three shipped decode profiles are bit-identical over the witness; the trait is the only new contract
surface, so that is where a second pair of eyes is worth most.

@codex review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants