From 16e8eae66d4f95d739b60728046e48a92ac3b0b7 Mon Sep 17 00:00:00 2001 From: shyeh25 <206795756+shyeh25@users.noreply.github.com> Date: Thu, 3 Sep 2026 19:35:26 -0700 Subject: [PATCH] fix: size MegaMoE full-pool MXFP8FP4 pull buffer 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. --- csrc/jit_kernels/heuristics/mega_moe.hpp | 16 ++++++++++++++-- csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp | 2 +- .../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp | 1 + .../deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh | 3 +++ 4 files changed, 19 insertions(+), 3 deletions(-) diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index edf70ba836..2486e4a007 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -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 @@ -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 + num_ring_tokens >= num_max_pool_tokens and + num_shared_experts == 0 and + block_k == block_n and + 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; } diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp index 273874cfd9..ece8b1a040 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp @@ -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, diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index 1b14558121..a8188a0174 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -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 diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh index 8c5d5976b1..0fe28c91f1 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh @@ -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;