Skip to content

Match the MegaMoE dispatch pull to the kernel's full-pool predicate - #85

Open
mmangkad wants to merge 2 commits into
sgl-project:devfrom
mmangkad-dev:mmangkad/mxfp8fp4-full-pool-pull
Open

Match the MegaMoE dispatch pull to the kernel's full-pool predicate#85
mmangkad wants to merge 2 commits into
sgl-project:devfrom
mmangkad-dev:mmangkad/mxfp8fp4-full-pool-pull

Conversation

@mmangkad

@mmangkad mmangkad commented Sep 10, 2026

Copy link
Copy Markdown

Summary

Fixes a regression from #80, shipped in v0.1.7.

num_bytes_per_pull sizes 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:

sglang commit sgl-deep-gemm result
20a491d1d3 0.1.5.post3 / 0.1.6 pass
20a491d1d3 0.1.7 NaN
dc2843801d 0.1.6 pass
dc2843801d 0.1.7 NaN

Restoring 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, and L2_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 threads num_shared_experts into get_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=128 gets the whole 7168-byte token, block_k=256 stays chunked at 3584. The MegaMoE test family passes, including #80's own test_mega_moe_nvfp4_alphas.

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.

1 participant