[4/5] feat(engine): accept NVFP4 sparse-MoE expert banks - #289
Open
MichaelDementii wants to merge 4 commits into
Open
MichaelDementii wants to merge 4 commits into
MichaelDementii wants to merge 4 commits 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.
…d tensor A tensor object holds one NVFP4 weight divisor, which is right for a plane quantised as one matrix and wrong for a plane assembled from several that were quantised apart. Published NVFP4 checkpoints of MoE models are the second case: each expert matrix carries its own weight_global_scale, so the converter has to refuse the pack. The tensor object gains an optional `divisors`: the number of separately quantised source matrices stacked into the plane, default 1. The payload keeps that many FP32 words immediately after the scale plane, in stacking order, so payload_bytes = weight_divisor_offset + 4 * divisors, and row r of an N-row plane uses word r / (N / divisors). A divisor's share must be a whole number of 128-row scale tiles, so a boundary never splits one. Absent means one word and the present layout exactly, so every existing artifact stays valid and no reader changes behaviour. The count is decided where the parent's shape is chosen, in the recipe plan, rather than being discovered at encode time; the converter refuses it outside NVFP4, and `validate_nvfp4_weight` refuses a plane whose divisor count does not divide its rows.
The operation admitted expert banks as Q4+Q5/Q6 or Q8+Q8 and nothing else, so an artifact whose experts are NVFP4 had no route at any token count. This registers the profile across all three: decode, small-token and prefill. Decode reads NVFP4 against represented BF16 activations. The kernels are CUDA-core dot products, so four-bit activations buy no arithmetic there and would cost a quantiser of their own; keeping the activation represented puts rel_l2 at T=1 at 1.60e-3 instead of 1.70e-1. The D4 block takes two hidden rows, because the scale plane puts rows r and r+1 in one 32-byte sector. Prefill quantises the chunk once and runs W4A4 over the existing work list, and the gate/up epilogue writes the SwiGLU intermediate already encoded, so `down` needs no separate quantiser. The per-source divisor is applied in that epilogue rather than inside the MMA loop, by a shift where the stride is a power of two. The frontier between the small-token and prefill routes is 13 for this profile, swept cold over [2,46] against its own routes rather than inherited from the groupwise ones, and the small-token schedule is chosen by the same sweep. The staging buffers the route needs are charged to this profile alone. Three launches written for it - the small-token router, the fused selection/scan/index block and a narrower route-job bound - are gated on it for the same reason: they would otherwise change what a groupwise call runs on the ragged tail of a sliced chunk. The wrapper also requires the routed and shared banks to agree on their activation divisor, because the route encodes the chunk once and both gate/up GEMMs read that one plane. The operator test walks the profile at T = 1, 2, 12, 13, 64, 768, 4097, with a separate tolerance either side of the frontier, and its fixture stacks gate and up under divisors that differ, as the published checkpoint does.
`prepare_sparse_moe_weights` enumerates the admissible bank formats and had no NVFP4 arm, so an artifact whose experts carry the format could not be bound even once the operation could run it. The profile takes all four expert matrices as NVFP4 - they are inseparable, because the gate/up epilogue emits an encoded intermediate that `down` reads. It requires AllowA4 on every expert input, since the prefill route quantises the hidden state to four bits, and every other NVFP4 consumer refuses an input without it.
Contributor
Author
|
Small one: the engine's expert bank learns to admit NVFP4 planes. The question worth asking is @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
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.
Related Issue: #283
Fourth of five. Thirty lines, and the point at which the profile becomes reachable. It needs the
third. No performance claim.
What changes
prepare_sparse_moe_weightsenumerates the admissible bank formats and had no NVFP4 arm, so anartifact whose experts carry the format could not be bound even once the operation could run it.
The profile takes all four expert matrices as NVFP4. They are inseparable: the gate/up epilogue emits
an encoded intermediate that
downreads, so a mixed bank has no route.It also requires
AllowA4on every expert input. The prefill route quantises the hidden state to fourbits, which is exactly what that permission grants; every other NVFP4 consumer refuses an input
without it, and this one would have quantised the activation of an input whose policy says it must
not, and said nothing.
Contract
A user who does not grant
AllowA4now gets a refusal at bind time naming the reason, where beforethe bank was rejected as an unsupported format. No existing profile's admissibility changes.
Verification
ctest120 of 120.End to end: the published NVFP4 artifact loads and serves with this commit and is refused without it.
The serving numbers for the whole series are on the Issue.
Checks not run
AllowA4refusal. The only in-tree callerof
prepare_sparse_moe_weightsoutside production is the real-artifact loading test, which needsan artifact this suite does not carry and therefore skips. The evidence is that the artifact binds
with this commit and not without — weaker than a negative case.
🤖 Generated with Claude Code