Match the MegaMoE dispatch pull to the kernel's full-pool predicate - #85
Open
mmangkad wants to merge 2 commits into
Open
Match the MegaMoE dispatch pull to the kernel's full-pool predicate#85mmangkad wants to merge 2 commits into
mmangkad wants to merge 2 commits into
Conversation
mmangkad
force-pushed
the
mmangkad/mxfp8fp4-full-pool-pull
branch
from
September 10, 2026 09:05
cd4ec3d to
cea487b
Compare
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.
Summary
Fixes a regression from #80, shipped in v0.1.7.
num_bytes_per_pullsizes the dispatch pull buffer on the host, and the kernel either TMAs a whole token into that buffer or walks it in chunks. The two have to agree. v0.1.6 kept them in sync with a host-side escape that skipped chunking whenever the kernel took its full-pool path; #80 deleted the escape and left the halving loop unconditional.What follows is an overrun. DeepSeek-V4-Pro FP4 (
hidden=7168, EP4) takes the fast path, the host hands it a 3584-byte buffer, the kernel writes 7168 bytes into it, and the activations come back NaN. That is what killed sglang's GB300 nightly. Moving only the wheel pins it down:20a491d1d320a491d1d3dc2843801ddc2843801dRestoring v0.1.6's escape verbatim would not be enough, because it keyed on the ring alone while the kernel also demands no shared experts,
BLOCK_K == BLOCK_N, andL2_SHAPE_K / BLOCK_K <= 32. Exempt more than the kernel does and the dispatch buffers grow, the pipeline drops a stage, and a shared-expert shape stops compiling outright with"Hidden is too large". So this mirrors the kernel's predicate term for term and threadsnum_shared_expertsintoget_mega_moe_config().BF16 stays chunked — different kernel, no unchunked branch. MXFP4 and NVFP4 are covered, since #80 routes them through this same kernel and its predicate ignores the MMA kind.
Verified on 4x GB300: gsm8k 0.96499 against a 0.935 floor, zero NaN, ~1012 tok/s against 1013 on v0.1.6. Pull size now tracks the kernel per config —
block_k=128gets the whole 7168-byte token,block_k=256stays chunked at 3584. The MegaMoE test family passes, including #80's owntest_mega_moe_nvfp4_alphas.