fix: size MegaMoE full-pool MXFP8FP4 pull buffer - #82
Conversation
Keep a complete-token pull buffer only when the SM100 MXFP8FP4 kernel selects the one-shot full-pool path. Preserve chunked pulls for all other paths and assert the full-token invariant at compile time.
| const int num_max_pool_tokens = layout::get_num_max_pool_tokens( | ||
| num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); | ||
| const bool use_full_pool_fp8_fp4_path = | ||
| mma_kind == MmaKind::MXFP8FP4 and |
There was a problem hiding this comment.
Question: Why doesn't other kinds of mma_kind have the same issue?
There was a problem hiding this comment.
BF16 does not have this issue because it uses a separate kernel whose dispatch pull always transfers the token in kNumBytesPerPull chunks.
MXFP4/NVFP4 share this kernel, but their activations are packed FP4. For the affected hidden=7168 shape, a complete token is 3,584 bytes, already below the 4 KiB threshold, so the host heuristic does not split it.
MXFP8FP4 uses FP8 activations, making the complete token 7,168 bytes for this shape. The host reduced the scratch slice to 3,584 bytes while the full-pool branch still issued a one-shot 7,168-byte TMA transfer, causing the overflow.
The added static assertion also guards against a future configuration producing the same host/device size mismatch.
There was a problem hiding this comment.
Follow-up question: where does the hidden=7168 assumption come from? These kernels support many models with many kinds of hidden size.
There was a problem hiding this comment.
Good question. The hidden=7168 value comes from the DeepSeek-V4 workload that reproduced the issue, but the safety of the other MMA kinds does not actually depend on that hidden size.
For MXFP4 and NVFP4, get_element_bits(mma_kind) is 4. The heuristic initially selects a block_k of either 128 or 256, then scales it with:
block_k = block_k * 8 / get_element_bits(mma_kind);
Therefore, MXFP4/NVFP4 always use a final block_k of 256 or 512, while block_n is fixed at 128. They can never satisfy the existing block_k == block_n condition in kUseFullPoolFP8FP4Path, so they always use the chunked pull path regardless of hidden size.
BF16 uses a separate kernel whose dispatch pull is also always chunked. MXFP8FP4 has 8-bit activations, so its final block_k can remain 128 and satisfy block_k == block_n, allowing it to enter the one-shot full-pool path.
My previous answer focused too much on the reproduced hidden=7168 shape; the block-shape condition is the general reason the other MMA kinds are unaffected.
| const bool use_full_pool_fp8_fp4_path = | ||
| mma_kind == MmaKind::MXFP8FP4 and | ||
| num_ring_tokens >= num_max_pool_tokens and | ||
| num_shared_experts == 0 and |
There was a problem hiding this comment.
Question: Why num_shared_experts > 0 won't have this issue?
There was a problem hiding this comment.
When num_shared_experts > 0, kHasShared is true, so the existing device-side predicate disables kUseFullPoolFP8FP4Path.
The kernel therefore uses the general chunked pull path, where every TMA transfer is limited to kNumBytesPerPull, so a reduced scratch slice is safe. Shared-expert phases also use separate buffers and synchronization counters; the current full-pool arrival protocol is only enabled for the routed-only case.
This host-side check simply mirrors that existing device-side eligibility condition.
| mma_kind == MmaKind::MXFP8FP4 and | ||
| num_ring_tokens >= num_max_pool_tokens and | ||
| num_shared_experts == 0 and | ||
| block_k == block_n and |
There was a problem hiding this comment.
Question: what is this checking? what if block_k != block_n?
There was a problem hiding this comment.
The full-pool protocol uses an arrival bitmask to connect L1 output tiles, produced along N in BLOCK_N units, with L2 input tiles, consumed along K in BLOCK_K units. Its indexing assumes a one-to-one mapping between those tiles, which requires BLOCK_K == BLOCK_N.
If BLOCK_K != BLOCK_N, the existing device predicate makes kUseFullPoolFP8FP4Path false, and the kernel falls back to the general counter-based/chunked path. Therefore, the complete-token scratch buffer is not required in that case.
This condition also mirrors the static assertion already present in the device full-pool branch.
|
@Fridge003 could you review this? Thanks! |
2 participants
Summary
kernel selects its one-shot full-pool path.
non-full-pool paths.
compile-time assertion at the transfer site to prevent future host/device
drift.
Related issue: sgl-project/sglang#37559
Root cause
The v0.1.7 host heuristic repeatedly halves
num_bytes_per_pulluntil it is atmost 4 KiB. For DeepSeek-V4's 7,168-byte FP8 activation token, this produces a
3,584-byte per-warp scratch slice.
The retained MXFP8FP4 full-pool fast path does not consume that token in chunks:
it issues a one-shot TMA transfer for all 7,168 bytes. The transfer therefore
writes beyond the 3,584-byte slice and corrupts adjacent shared memory,
eventually surfacing as a CUDA illegal-memory-access failure under sustained
high-concurrency serving.
Validation
All GPU testing used Blackwell B300/GB300 systems.
Correctness and memory safety
B300 GPUs.
additional concurrency-4,096 stress run also passed.
sgl_deep_gemm/run_tests.shattention andMegaMoE tests passed; documented architecture/upstream exclusions remained
skips.
match.
normalized per-sample results.
No functional or accuracy regression was detected in the validation above.