Skip to content

fix: size MegaMoE full-pool MXFP8FP4 pull buffer #82

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
shyeh25 wants to merge 1 commit into from
Closed

fix: size MegaMoE full-pool MXFP8FP4 pull buffer #82

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions csrc/jit_kernels/heuristics/mega_moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,7 @@ static MegaMoEConfig get_mega_moe_config(
const int& hidden, const int& intermediate_hidden,
const int& num_ring_tokens,
const int& num_sf_ring_tokens,
const int& num_shared_experts,
const MmaKind& mma_kind) {

// Block config
Expand All @@ -239,10 +240,21 @@ static MegaMoEConfig get_mega_moe_config(
const int num_dispatch_threads = 128;
const int num_non_epilogue_threads = 128;

// Pull: divide token bytes by 2 until <= kPullThreshold
// The MXFP8FP4 full-pool path issues one-shot TMA transfers for an entire
// token. Its per-warp scratch slice must therefore hold the complete token.
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

@nvpohanh nvpohanh Sep 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: Why doesn't other kinds of mma_kind have the same issue?

@shyeh25 shyeh25 Sep 4, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@nvpohanh nvpohanh Sep 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up question: where does the hidden=7168 assumption come from? These kernels support many models with many kinds of hidden size.

@shyeh25 shyeh25 Sep 4, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

num_ring_tokens >= num_max_pool_tokens and
num_shared_experts == 0 and

@nvpohanh nvpohanh Sep 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: Why num_shared_experts > 0 won't have this issue?

@shyeh25 shyeh25 Sep 4, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

block_k == block_n and

@nvpohanh nvpohanh Sep 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: what is this checking? what if block_k != block_n?

@shyeh25 shyeh25 Sep 4, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

intermediate_hidden / block_k <= 32;

// Reusable-ring paths can split a token into smaller pull chunks.
constexpr int kPullThreshold = 4096;
int num_bytes_per_pull = hidden * get_element_bits(mma_kind) / 8;
while (num_bytes_per_pull > kPullThreshold) {
while (not use_full_pool_fp8_fp4_path and num_bytes_per_pull > kPullThreshold) {
DG_HOST_ASSERT(num_bytes_per_pull % 2 == 0);
num_bytes_per_pull /= 2;
}
Expand Down
2 changes: 1 addition & 1 deletion csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ static void sm100_bf16_mega_moe(
const auto config = get_mega_moe_config(
num_ranks, num_experts, num_experts_per_rank,
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden,
num_ring_tokens, 0, MmaKind::BF16);
num_ring_tokens, 0, num_shared_experts, MmaKind::BF16);

// Make tensormap
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
Expand Down
1 change: 1 addition & 0 deletions csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,7 @@ static void sm100_fp8_fp4_mega_moe(
num_ranks, num_experts, num_experts_per_rank,
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden,
num_ring_tokens, num_sf_ring_tokens,
num_shared_experts,
mma_kind);

// Make tensormap
Expand Down
3 changes: 3 additions & 0 deletions deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -582,6 +582,9 @@ sm100_fp8_fp4_mega_moe_impl(void* y,

// Hidden bytes are divided into chunks
constexpr uint32_t kNumHiddenTokenBytes = kHidden * kNumElemBits / 8;
DG_STATIC_ASSERT(
not kUseFullPoolFP8FP4Path or kNumBytesPerPull == kNumHiddenTokenBytes,
"The full-pool MXFP8FP4 path requires a complete-token pull buffer");
constexpr uint32_t kNumChunks = kNumHiddenTokenBytes / kNumBytesPerPull;
DG_STATIC_ASSERT(kNumChunks * kNumBytesPerPull == kNumHiddenTokenBytes, "kNumBytesPerPull must divide the token bytes");
const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert;
Expand Down