From 69fdbcae1a62c0ac8ca8498a3e763e681832ff90 Mon Sep 17 00:00:00 2001 From: Mohammad Angkad Date: Sun, 6 Sep 2026 07:05:47 +0000 Subject: [PATCH 1/2] [SM100] Pull whole MXFP8FP4 tokens when the ring spans the full pool --- csrc/jit_kernels/heuristics/mega_moe.hpp | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index edf70ba836..8bb170e840 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -239,10 +239,17 @@ static MegaMoEConfig get_mega_moe_config( const int num_dispatch_threads = 128; const int num_non_epilogue_threads = 128; + // A ring that already spans the whole pool pulls each MXFP8FP4 token whole; + // splitting it desynchronises the full-pool dispatch and yields NaN activations. + 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; + // Pull: divide token bytes by 2 until <= kPullThreshold 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; } From cea487bb4a32bd82a4056e2db78c85a687048c03 Mon Sep 17 00:00:00 2001 From: Mohammad Angkad Date: Sun, 6 Sep 2026 07:27:27 +0000 Subject: [PATCH 2/2] [SM100] Match the MegaMoE pull exemption to the kernel's full-pool predicate --- csrc/jit_kernels/heuristics/mega_moe.hpp | 26 ++++++++++++++++--- .../jit_kernels/impls/sm100_bf16_mega_moe.hpp | 2 +- .../impls/sm100_fp8_fp4_mega_moe.hpp | 2 +- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index 8bb170e840..361b2101a5 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -104,6 +104,21 @@ static bool is_nvfp4_mma_kind(const MmaKind& mma_kind) { return mma_kind == MmaKind::NVFP4; } +// Mirror of `kUseFullPoolFP8FP4Path` in `sm100_fp8_fp4_mega_moe.cuh`. On that +// branch the kernel pulls a whole token with a single TMA, so the pull buffer +// has to be sized for the whole token; every other case keeps chunked pulls, +// which also keeps the dispatch buffers small enough to hold the pipeline +// depth. BF16 runs `sm100_bf16_mega_moe.cuh`, which always chunks. +// Keep this predicate in step with the kernel's. +static bool uses_full_pool_fp8_fp4_path( + const MmaKind& mma_kind, const int& num_shared_experts, + const int& num_ring_tokens, const int& num_max_pool_tokens, + const int& block_n, const int& block_k, const int& intermediate_hidden) { + return mma_kind != MmaKind::BF16 and num_shared_experts == 0 and + num_ring_tokens >= num_max_pool_tokens and block_k == block_n and + intermediate_hidden / block_k <= 32; +} + static std::tuple get_block_config_for_mega_moe( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, @@ -215,6 +230,7 @@ static std::pair get_pipeline_config_for_mega_moe( static MegaMoEConfig get_mega_moe_config( const int& num_ranks, const int& num_experts, const int& num_experts_per_rank, + const int& num_shared_experts, const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, const int& num_ring_tokens, @@ -239,12 +255,14 @@ static MegaMoEConfig get_mega_moe_config( const int num_dispatch_threads = 128; const int num_non_epilogue_threads = 128; - // A ring that already spans the whole pool pulls each MXFP8FP4 token whole; - // splitting it desynchronises the full-pool dispatch and yields NaN activations. + // The kernel's full-pool branch pulls a whole token in one TMA into a buffer + // this size, so it must cover the token; splitting it overruns that buffer + // and leaves NaN activations behind. 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; + const bool use_full_pool_fp8_fp4_path = uses_full_pool_fp8_fp4_path( + mma_kind, num_shared_experts, num_ring_tokens, num_max_pool_tokens, + block_n, block_k, intermediate_hidden); // Pull: divide token bytes by 2 until <= kPullThreshold constexpr int kPullThreshold = 4096; diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp index 273874cfd9..3cff677289 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp @@ -131,7 +131,7 @@ static void sm100_bf16_mega_moe( // Heuristics const auto config = get_mega_moe_config( - num_ranks, num_experts, num_experts_per_rank, + num_ranks, num_experts, num_experts_per_rank, num_shared_experts, num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_ring_tokens, 0, MmaKind::BF16); 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..db3683de3e 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -188,7 +188,7 @@ static void sm100_fp8_fp4_mega_moe( // Heuristics const auto config = get_mega_moe_config( - num_ranks, num_experts, num_experts_per_rank, + num_ranks, num_experts, num_experts_per_rank, num_shared_experts, num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_ring_tokens, num_sf_ring_tokens, mma_kind);