From 61897f9b5aa1ce5524cfdcaaa03c8955a8ab645a Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Mon, 17 Aug 2026 18:01:33 +0800 Subject: [PATCH 1/3] Add SM90 (Hopper) MegaMoE support Ports upstream PR sgl-project/DeepGEMM#36 ("Sm90 mega moe on sgl dev"), which added a Hopper FP8xFP8 fused MoE GEMM kernel on the old `dev` branch layout (tvm_ffi_api.cpp / sgl_deep_gemm), onto nv_dev's current layout (csrc/apis/*.hpp, csrc/python_api.cpp, deep_gemm/mega). nv_dev already had MegaMoE, but only for SM100 (Blackwell), built on a ring-buffer / cluster-pair-interleaved persistent-grid scheduler and a UE8M0/FP4-oriented workspace layout. The SM90 kernel in PR #36 instead uses a simpler pool-based workspace and per-expert-wave scheduler with float (non-UE8M0) scale factors, FP8-only weights (no FP4), and no 2-CTA clusters or shared-expert support -- these are not compatible data layouts, so the SM90 path is added as a fully parallel path alongside (not replacing) the SM100 one. New files (ported from the PR's `.cuh`/`.hpp`, adapted to nv_dev's current naming/dispatch conventions and symbol-renamed to avoid colliding with SM100's `layout::Workspace` / `sched::MegaMoEScheduler` / `sched::BlockPhase`): - deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh (`layout::MegaMoESM90Workspace`, pool-based; reuses the existing shared `layout::Data`/`layout::Buffer`/`layout::TokenSrcMetadata` and `layout::get_num_max_pool_tokens` from layout/mega_moe.cuh) - deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh (`sched::MegaMoESM90Scheduler` / `sched::MegaMoESM90BlockPhase`) - deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh (the ported kernel, ~1935 lines, WGMMA/TMA-based Hopper impl) - csrc/jit_kernels/heuristics/sm90_mega_moe.hpp (`MegaMoESM90Config` and block/pipeline/wave heuristics, mirroring csrc/jit_kernels/heuristics/mega_moe.hpp's structure) - csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp (JIT host runtime, mirroring csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp) - tests/test_mega_moe_hopper.py (correctness test with an SM90 capability guard; skips cleanly on non-Hopper GPUs) Modified files: - csrc/apis/mega.hpp: adds `get_symm_buffer_size_for_sm90_mega_moe` and `fp8_mega_moe_sm90`, registered via the existing `deep_gemm::mega::register_apis` (already wired into csrc/python_api.cpp, so no python_api.cpp changes were needed). - deep_gemm/mega/__init__.py + deep_gemm/__init__.py: expose `Sm90SymmBuffer`, `get_symm_buffer_for_sm90_mega_moe`, `transform_weights_for_mega_moe_sm90`, `fp8_mega_moe_sm90`, following the existing SM100 exposure pattern. - deep_gemm/include/deep_gemm/comm/barrier.cuh: generalizes `grid_sync`/`nvlink_barrier` to a templated workspace type (instead of hard-coding `layout::Workspace`) so the SM90 kernel can reuse them with `layout::MegaMoESM90Workspace`; existing SM100 call sites are unaffected (the type is still deduced from the argument). Also ports the PR's ARCH 900-1000 guarded trap-instead -of-printf change to the NVLink barrier timeout path. Not ported: shared-expert support and the `situ` activation (the PR does not implement either for SM90). No GPU/CUDA toolchain is available in this environment, so this is a best-effort, close-reading port verified via `python3 -m py_compile`, brace/paren balance checks on all new/modified C++/CUDA files, and confirming `import deep_gemm` fails identically (missing compiled `_C` extension) before and after these changes -- i.e. no regression introduced. It has not been compiled or run on real Hopper hardware. Co-Authored-By: Claude Sonnet 5 --- csrc/apis/mega.hpp | 208 ++ csrc/jit_kernels/heuristics/sm90_mega_moe.hpp | 236 ++ csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp | 307 +++ deep_gemm/__init__.py | 5 + deep_gemm/include/deep_gemm/comm/barrier.cuh | 23 +- .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 1936 +++++++++++++++++ .../deep_gemm/layout/sm90_mega_moe.cuh | 179 ++ .../deep_gemm/scheduler/sm90_mega_moe.cuh | 235 ++ deep_gemm/mega/__init__.py | 111 + tests/test_mega_moe_hopper.py | 172 ++ 10 files changed, 3408 insertions(+), 4 deletions(-) create mode 100644 csrc/jit_kernels/heuristics/sm90_mega_moe.hpp create mode 100644 csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh create mode 100644 deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh create mode 100644 deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh create mode 100644 tests/test_mega_moe_hopper.py diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 93a9138ce4..bd4c8f36c5 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -15,6 +15,7 @@ #include "../jit_kernels/impls/sm100_bf16_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp" +#include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp" namespace deep_gemm::mega { @@ -514,6 +515,211 @@ static void fp4_fp4_mega_moe( sym_buffer.zero_(); } + +// ============================================================================ +// SM90 (Hopper) FP8 MegaMoE +// ---------------------------------------------------------------------------- +// Ported from upstream PR https://github.com/sgl-project/DeepGEMM/pull/36 +// (`csrc/apis/sm90_mega.hpp` on the old `dev` branch). +// +// Unlike `fp8_fp4_mega_moe` (SM100, FP4 weights, UE8M0 SF, ring-buffer workspace, optional +// shared experts), this path is FP8-only (both activations and weights), uses float SF at +// per-128 (L1)/per-64 (L2) K granularity, and uses a simpler pool-based workspace/scheduler +// (`layout::MegaMoESM90Workspace` / `sched::MegaMoESM90Scheduler`). It does not support shared +// experts. The symmetric buffer layout is therefore also different and is *not* +// interchangeable with `get_symm_buffer_size_for_mega_moe`'s buffer. +// ============================================================================ + +static std::tuple(const torch::Tensor&)>> +get_symm_buffer_size_for_sm90_mega_moe( + const int& num_ranks, const int& num_experts, + const int& num_max_tokens_per_rank, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const bool& use_fp8_dispatch, const std::string& activation) { + DG_HOST_ASSERT(num_experts % num_ranks == 0); + DG_HOST_ASSERT(use_fp8_dispatch); + DG_HOST_ASSERT(activation == "swiglu"); + + const auto workspace = layout::MegaMoESM90Workspace(nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk); + + const auto fp8_token_layout = layout::Data(hidden); + const auto bf16_token_layout = layout::Data(hidden * 2); + const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden); + const auto fp8_sf_layout = layout::Data(hidden / 32); + const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 16); + const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); + const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false); + const auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + + const auto input_token_buffer = layout::Buffer( + fp8_token_layout, 1, num_max_tokens_per_rank, + workspace.get_end_ptr()); + const auto input_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, num_max_tokens_per_rank, + input_token_buffer.get_end_ptr()); + const auto input_topk_idx_buffer = layout::Buffer( + input_topk_idx_layout, 1, num_max_tokens_per_rank, + input_sf_buffer.get_end_ptr()); + const auto input_topk_weights_buffer = layout::Buffer( + input_topk_weights_layout, 1, num_max_tokens_per_rank, + input_topk_idx_buffer.get_end_ptr()); + + const auto num_max_pool_tokens = static_cast(workspace.num_max_pool_tokens); + int num_max_padded_sf_pool_tokens = 0; + for (int block_m: layout::kCandidateBlockM) { + num_max_padded_sf_pool_tokens = std::max( + num_max_padded_sf_pool_tokens, + layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m) + ); + } + + const auto l1_token_buffer = layout::Buffer( + fp8_token_layout, 1, num_max_pool_tokens, + input_topk_weights_buffer.get_end_ptr()); + const auto l1_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, num_max_padded_sf_pool_tokens, + l1_token_buffer.get_end_ptr()); + const auto l1_topk_weights_buffer = layout::Buffer( + l1_topk_weights_layout, 1, num_max_pool_tokens, + l1_sf_buffer.get_end_ptr()); + + const auto l2_token_buffer = layout::Buffer( + fp8_intermediate_token_layout, 1, num_max_pool_tokens, + l1_topk_weights_buffer.get_end_ptr()); + const auto l2_sf_buffer = layout::Buffer( + fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens, + l2_token_buffer.get_end_ptr()); + + const auto combine_token_buffer = layout::Buffer( + bf16_token_layout, num_topk, num_max_tokens_per_rank, + l2_sf_buffer.get_end_ptr()); + + DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); + + auto slice_input_buffers = [=](const torch::Tensor& buffer) { + auto x = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_token_buffer.base)), + {num_max_tokens_per_rank, hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto x_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_sf_buffer.base)), + {num_max_tokens_per_rank, hidden / 128}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto topk_idx = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_topk_idx_buffer.base)), + {num_max_tokens_per_rank, num_topk}, + torch::TensorOptions().dtype(torch::kInt64).device(buffer.device())); + auto topk_weights = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_topk_weights_buffer.base)), + {num_max_tokens_per_rank, num_topk}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto l1_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_token_buffer.base)), + {num_max_pool_tokens, hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto l1_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_sf_buffer.base)), + {num_max_padded_sf_pool_tokens, hidden / 128}, + {1, num_max_padded_sf_pool_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto l2_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_token_buffer.base)), + {num_max_pool_tokens, intermediate_hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto l2_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_sf_buffer.base)), + {num_max_padded_sf_pool_tokens, intermediate_hidden / 64}, + {1, num_max_padded_sf_pool_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); + }; + return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; +} + +static void fp8_mega_moe_sm90( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, const int& rank_idx, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + const std::tuple& recipe, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math +) { + const auto [l1_weights, l1_weights_sf] = l1_weights_tuple; + const auto [l2_weights, l2_weights_sf] = l2_weights_tuple; + + const auto arch_major = device_runtime->get_arch_major(); + DG_HOST_ASSERT(arch_major == 9); + + const auto num_tokens = static_cast(y.size(0)); + const auto [rm, rn, rk] = recipe; + DG_HOST_ASSERT(rm == 128 and rn == 128 and rk == 128); + DG_HOST_ASSERT(activation == "swiglu"); + + const auto activation_clamp = + activation_clamp_opt.value_or(std::numeric_limits::infinity()); + DG_HOST_ASSERT(activation_clamp >= 0); + + DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(l1_weights.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(l2_weights.scalar_type() == torch::kFloat8_e4m3fn); + const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = get_shape<3>(l1_weights); + const auto [num_experts_per_rank_, hidden_, intermediate_hidden] = get_shape<3>(l2_weights); + DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank); + DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_); + DG_HOST_ASSERT(hidden == hidden_); + DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); + DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); + DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); + DG_HOST_ASSERT(intermediate_hidden / 64 <= 64); + + constexpr int kGranMN = 128, kGranK = 128; + check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, + num_experts_per_rank, false, true, torch::kFloat); + check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, + num_experts_per_rank, false, true, torch::kFloat); + + if (cumulative_local_expert_recv_stats.has_value()) { + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous()); + } + + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts_ = num_experts_per_rank * num_ranks; + const auto [num_required_bytes, slice] = get_symm_buffer_size_for_sm90_mega_moe( + num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + true, activation); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(num_experts == num_experts_); + + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + + sm90_fp8_mega_moe(y, + l1_acts, l1_acts_sf, + l2_acts, l2_acts_sf, + l1_weights, l2_weights, + l1_weights_sf, l2_weights_sf, + cumulative_local_expert_recv_stats, + sym_buffer_ptrs, + rank_idx, num_max_tokens_per_rank, + num_experts_per_rank, + num_tokens, num_topk, + hidden, intermediate_hidden, + activation_clamp, fast_math); + + if (get_env("DG_COMM_KERNEL_DEBUG")) + sym_buffer.zero_(); +} + static void bf16_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_weights, @@ -630,6 +836,8 @@ static void register_apis(pybind11::module_& m) { m.def("fp8_fp4_mega_moe", &fp8_fp4_mega_moe); m.def("fp4_fp4_mega_moe", &fp4_fp4_mega_moe); m.def("bf16_mega_moe", &bf16_mega_moe); + m.def("get_symm_buffer_size_for_sm90_mega_moe", &get_symm_buffer_size_for_sm90_mega_moe); + m.def("fp8_mega_moe_sm90", &fp8_mega_moe_sm90); #endif } diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp new file mode 100644 index 0000000000..74ba1dec4e --- /dev/null +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -0,0 +1,236 @@ +#pragma once + +#include +#include + +#include +#include + +#include "../../utils/exception.hpp" +#include "../../utils/math.hpp" +#include "../../utils/system.hpp" +#include "sm90.hpp" + +namespace deep_gemm { + +// ============================================================================ +// SM90 (Hopper) MegaMoE configuration +// ---------------------------------------------------------------------------- +// Ported from upstream PR https://github.com/sgl-project/DeepGEMM/pull/36. +// SM90 differs from SM100 in: +// - No tensor memory (TMEM): WGMMA accumulators live in registers. +// - No FP4: weights are FP8 e4m3 with per-128 channel float scales. +// - No 2-CTA cluster MMA: this kernel always runs with cluster_size = 1. +// - Activation SF is float, not UE8M0 int: L1 input uses per-128 K and the +// fused L1 epilogue writes L2 activation SF at per-64 K granularity. +// - Scheduling is a simple pool/wave scheduler (`sched::MegaMoESM90Scheduler`), +// not the SM100 ring-buffer/cluster-pair persistent-grid scheduler. +// The kernel implementation is in `deep_gemm/impls/sm90_fp8_mega_moe.cuh`. +// ============================================================================ + +struct MegaMoESM90Config { + int block_m, block_n, block_k; + int cluster_size; + int num_max_pool_tokens; + int num_padded_sf_pool_tokens; + int swizzle_acts_mode, swizzle_weights_mode; + int num_experts_per_wave; + int num_stages, smem_size; + int num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads; + + friend std::ostream& operator << (std::ostream& os, const MegaMoESM90Config& config) { + os << "MegaMoESM90Config(" + << "block_m=" << config.block_m << ", block_n=" << config.block_n << ", block_k=" << config.block_k + << ", cluster_size=" << config.cluster_size + << ", num_max_pool_tokens=" << config.num_max_pool_tokens + << ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens + << ", swizzle_acts_mode=" << config.swizzle_acts_mode << ", swizzle_weights_mode=" << config.swizzle_weights_mode + << ", num_experts_per_wave=" << config.num_experts_per_wave + << ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size + << ", num_dispatch_threads=" << config.num_dispatch_threads + << ", num_non_epilogue_threads=" << config.num_non_epilogue_threads + << ", num_epilogue_threads=" << config.num_epilogue_threads << ")"; + return os; + } +}; + +static std::tuple get_block_config_for_mega_moe_sm90( + const int& num_ranks, const int& num_experts, + const int& num_max_tokens_per_rank, const int& num_topk, + const int& num_tokens) { + const float expected_tokens_per_expert = + static_cast(num_tokens) * num_ranks * num_topk / num_experts; + const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; + if (auto_split_mn) + return {128, 512}; + + const int block_m = 64; + const int num_epilogue_warpgroups = 2; + + DG_HOST_ASSERT(std::any_of( + layout::kCandidateBlockM, layout::kCandidateBlockM + layout::kNumCandidateBlockMs, + [=](const auto& candidate) { return candidate == block_m; }) + ); + return {block_m, num_epilogue_warpgroups * 128}; +} + +// Base (non-arch-specific) heuristic for how many local experts to batch into a single +// scheduling "wave". Ported as-is from the old `dev` branch's shared +// `get_num_experts_per_wave_for_mega_moe` (pre-dates `nv_dev`'s ring-buffer scheduler, which has +// no equivalent wave concept), kept private to the SM90 path since nothing else references it. +static int get_num_experts_per_wave_for_mega_moe( + const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, + const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms) { + float expected_tokens_per_expert = static_cast(num_tokens) * num_topk / num_experts_per_rank; + if (expected_tokens_per_expert < 1) { + // Most experts don't have tokens, calculate all experts at once + return num_experts_per_rank; + } + + // Reduce per-expert block count by this factor since uneven routing leaves some experts with fewer tokens + constexpr int kImbalanceFactor = 2; + + // Count L1 blocks per expert assuming tokens are evenly spread across experts + const int num_m_blocks = ceil_div(static_cast(std::ceil(expected_tokens_per_expert)), block_m); + const int num_n_blocks = (2 * intermediate_hidden) / block_n; + const int num_l1_blocks_per_expert = num_m_blocks * num_n_blocks; + + // Pick the smallest value whose total blocks (after imbalance reduction) can keep all SMs busy + int num_experts_per_wave = num_l1_blocks_per_expert > 0 + ? ceil_div(kImbalanceFactor * num_sms, num_l1_blocks_per_expert) : 1; + num_experts_per_wave = std::min(num_experts_per_wave, num_experts_per_rank); + + // Round up to the nearest divisor of num_experts_per_rank so every wave processes the same count + while (num_experts_per_wave < num_experts_per_rank and num_experts_per_rank % num_experts_per_wave != 0) + ++ num_experts_per_wave; + + return num_experts_per_wave; +} + +static int get_num_experts_per_wave_for_mega_moe_sm90( + const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, + const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms) { + const float expected_tokens_per_expert = + static_cast(num_tokens) * num_topk / num_experts_per_rank; + if (expected_tokens_per_expert < 1.0f or expected_tokens_per_expert > 4.0f) + return num_experts_per_rank; + + if (block_m == 64 and intermediate_hidden >= 3072) { + const int num_n_blocks_per_expert = (2 * intermediate_hidden) / block_n; + const int single_wave_blocks = + num_experts_per_rank * num_n_blocks_per_expert; + if (single_wave_blocks >= 4 * num_sms) + return num_experts_per_rank; + } + return get_num_experts_per_wave_for_mega_moe( + num_experts_per_rank, num_tokens, num_topk, + intermediate_hidden, block_m, block_n, num_sms); +} + +static std::pair get_pipeline_config_for_mega_moe_sm90( + const int& smem_capacity, + const int& num_experts, const int& hidden, + const int& block_m, const int& block_n, const int& block_k, + const int& num_dispatch_warps, const int& num_epilogue_warps) { + constexpr int kSmemAlignment = 1024; + + const int smem_expert_count_size = align( + num_experts * static_cast(sizeof(uint32_t)), kSmemAlignment); + const int smem_send_buffers_size = align( + static_cast(layout::Buffer(layout::Data(hidden), num_dispatch_warps, 1).get_num_bytes()), + kSmemAlignment); + const int smem_dispatch_size = smem_expert_count_size + smem_send_buffers_size; + + const int smem_cd_l1 = block_m * (block_n / 2); + const int smem_cd_l2 = block_m * block_n * static_cast(sizeof(nv_bfloat16)); + const int smem_cd = align(std::max(smem_cd_l1, smem_cd_l2), kSmemAlignment); + + const int smem_sfa_per_stage = align(2 * block_m * static_cast(sizeof(float)), 128); + const int smem_sfb_per_stage = 0; + const int smem_per_stage = block_m * block_k + block_n * block_k + + smem_sfa_per_stage + smem_sfb_per_stage; + + const int smem_barriers_fixed = (num_dispatch_warps + 2 * num_epilogue_warps) * 8; + const int smem_barriers_per_stage = 2 * 8; + const int smem_fixed = smem_dispatch_size + smem_cd + smem_barriers_fixed; + + const int num_stages = (smem_capacity - smem_fixed) / + (smem_per_stage + smem_barriers_per_stage); + DG_HOST_ASSERT(num_stages >= 2); + const int smem_size = smem_fixed + num_stages * (smem_per_stage + smem_barriers_per_stage); + DG_HOST_ASSERT(smem_size <= smem_capacity); + return {num_stages, smem_size}; +} + +static MegaMoESM90Config get_mega_moe_config_sm90( + const int& num_ranks, const int& num_experts, const int& num_experts_per_rank, + 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_padded_sf_pool_tokens) { + const auto [block_m, num_epilogue_threads] = get_block_config_for_mega_moe_sm90( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens); + const float expected_tokens_per_expert = + static_cast(num_tokens) * num_ranks * num_topk / num_experts; + const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; + const bool decode_split_n_path = + block_m == 64 and num_epilogue_threads == 256; + const bool decode_use_block_n_256 = + decode_split_n_path and intermediate_hidden >= 3072 and + expected_tokens_per_expert >= 0.25f and + (2 * intermediate_hidden) % 256 == 0; + const int block_n = auto_split_mn ? 256 + : (decode_use_block_n_256 ? 256 : 128); + const int block_k = 128; + const int cluster_size = 1; + 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 int swizzle_acts_mode = 128; + const int swizzle_weights_mode = 128; + + const int num_sms = device_runtime->get_num_sms(); + const int num_experts_per_wave = get_num_experts_per_wave_for_mega_moe_sm90( + num_experts_per_rank, num_tokens, num_topk, + intermediate_hidden, block_m, block_n, num_sms); + + const bool reduce_decode_threads = num_epilogue_threads == 128; + const bool decode_split_n = + block_m == 64 and num_epilogue_threads == 256; + const bool shrink_non_epilogue = reduce_decode_threads or decode_split_n; + const int num_dispatch_threads = + (num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128; + const bool split_sfa_loader_warp = false; + const int num_non_epilogue_threads = + split_sfa_loader_warp ? 128 : + ((num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128); + DG_HOST_ASSERT((num_dispatch_threads + num_non_epilogue_threads) % 128 == 0); + + const auto [num_stages, smem_size] = get_pipeline_config_for_mega_moe_sm90( + SM90ArchSpec::smem_capacity, + num_experts, hidden, + block_m, block_n, block_k, + num_dispatch_threads / 32, num_epilogue_threads / 32); + + const auto config = MegaMoESM90Config { + block_m, block_n, block_k, + cluster_size, + num_max_pool_tokens, num_padded_sf_pool_tokens, + swizzle_acts_mode, swizzle_weights_mode, + num_experts_per_wave, + num_stages, smem_size, + num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads + }; + + if (get_env("DG_JIT_DEBUG") or get_env("DG_PRINT_CONFIGS")) { + const auto key = fmt::format( + "MegaMoESM90Config(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={})", + num_ranks, num_experts, hidden, intermediate_hidden, num_max_tokens_per_rank, num_tokens, num_topk); + static std::unordered_set printed; + if (printed.count(key) == 0) { + std::cout << key << ": " << config << std::endl; + printed.insert(key); + } + } + return config; +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp new file mode 100644 index 0000000000..8416196a6d --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -0,0 +1,307 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "runtime_utils.hpp" + +#include +#include +#include + +#include "../heuristics/sm90_mega_moe.hpp" + +namespace deep_gemm { + +// ============================================================================ +// SM90 (Hopper) FP8 MegaMoE host runtime +// ---------------------------------------------------------------------------- +// Ported from upstream PR https://github.com/sgl-project/DeepGEMM/pull/36. This is the SM90 +// counterpart of `SM100FP8FP4MegaMoERuntime`. The kernel itself lives in +// `deep_gemm/impls/sm90_fp8_mega_moe.cuh`. +// +// Differences from the SM100 path: +// * Activations and weights are both FP8 (e4m3); no FP4. +// * Activation/weight scale factors (SF) are float, not UE8M0 int + per-32 UTCCP layout. L1 +// activation SF and weight SF are per-128 K; the fused L1 epilogue writes L2 activation SF at +// per-64 K granularity. +// * No tensor memory: WGMMA accumulators are register-resident. +// * No 2-CTA cluster: `cluster_size` is always 1. +// * No shared-expert support (the PR this was ported from does not implement it for SM90). +// ============================================================================ + +class SM90FP8MegaMoERuntime final : public LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_max_tokens_per_rank; + int hidden, intermediate_hidden; + int num_experts, num_topk; + int num_ranks; + float activation_clamp; + bool fast_math; + int epilogue_registers; + bool reuse_accum_as_final; + bool l2_arrival_counter; + bool l2_epilogue_requires_full_sync; + bool split_phase_hot_path; + MegaMoESM90Config config; + + // Runtime arguments + void* y; + int* cumulative_local_expert_recv_stats; + int num_tokens; + layout::SymBuffer<> sym_buffer_ptrs; + + // Tensormaps for activations and weights. Weight scale factors use block (128, 128) + // quantization and are loaded by the math warpgroup directly from global memory (no TMA + // descriptor required). + CUtensorMap tensor_map_l1_acts; + CUtensorMap tensor_map_l1_acts_sf; + CUtensorMap tensor_map_l1_weights; + const float* l1_weights_sf; + CUtensorMap tensor_map_l1_output; + CUtensorMap tensor_map_l2_acts; + CUtensorMap tensor_map_l2_acts_sf; + CUtensorMap tensor_map_l2_weights; + const float* l2_weights_sf; + + // Launch configs + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_mega_moe_impl< + {}, + {}, {}, + {}, {}, + {}, + {}, {}, {}, + {}, + {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {} + >); +}}; +)", + args.num_max_tokens_per_rank, + args.hidden, args.intermediate_hidden, + args.num_experts, args.num_topk, + args.config.num_experts_per_wave, + args.config.block_m, args.config.block_n, args.config.block_k, + args.config.num_max_pool_tokens, + args.config.num_padded_sf_pool_tokens, + args.config.num_stages, + args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads, + args.launch_args.grid_dim.first, args.num_ranks, + to_string(args.activation_clamp), + args.fast_math ? "true" : "false", + args.epilogue_registers, + args.reuse_accum_as_final ? "true" : "false", + args.l2_arrival_counter ? "true" : "false", + args.l2_epilogue_requires_full_sync ? "true" : "false", + args.split_phase_hot_path ? "true" : "false"); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.y, + args.cumulative_local_expert_recv_stats, + args.num_tokens, + args.sym_buffer_ptrs, + args.tensor_map_l1_acts, + args.tensor_map_l1_acts_sf, + args.tensor_map_l1_weights, + args.l1_weights_sf, + args.tensor_map_l1_output, + args.tensor_map_l2_acts, + args.tensor_map_l2_acts_sf, + args.tensor_map_l2_weights, + args.l2_weights_sf + )); + } +}; + +static void sm90_fp8_mega_moe( + const torch::Tensor& y, + const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf, + const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf, + const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, + const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf, + const std::optional cumulative_local_expert_recv_stats, + const std::vector& sym_buffer_ptrs, + const int& rank_idx, const int& num_max_tokens_per_rank, + const int& num_experts_per_rank, + const int& num_tokens, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const float& activation_clamp, + const bool& fast_math +) { + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts = num_experts_per_rank * num_ranks; + const auto num_padded_sf_pool_tokens = static_cast(l1_acts_sf.size(0)); + + // Heuristics + const auto config = get_mega_moe_config_sm90( + num_ranks, num_experts, num_experts_per_rank, + num_max_tokens_per_rank, num_tokens, num_topk, + hidden, intermediate_hidden, num_padded_sf_pool_tokens); + const int default_epilogue_registers = + config.num_epilogue_threads == 512 ? 112 : 0; + const int epilogue_registers = default_epilogue_registers; + if (epilogue_registers > 0) { + const int dispatch_registers = + config.num_epilogue_threads == 512 ? 32 : 48; + const int non_epilogue_registers = + config.num_epilogue_threads == 512 ? 24 : 40; + DG_HOST_ASSERT(dispatch_registers * config.num_dispatch_threads + + non_epilogue_registers * config.num_non_epilogue_threads + + epilogue_registers * config.num_epilogue_threads <= 64512); + } + const bool reuse_accum_as_final = config.block_m == 128; + const bool default_split_mn_barrier_opt = + config.block_m == 128 and config.block_n == 256 and + config.num_epilogue_threads == 512; + const bool split_phase_hot_path = + config.block_m == 128 and config.block_n == 256 and hidden >= 7168; + const bool decode_split_n_path = + config.block_m == 64 and config.num_epilogue_threads == 256; + const bool decode_split_n_bn256 = + decode_split_n_path and config.block_n == 256; + const bool decode_l2_counter = + decode_split_n_bn256 and num_tokens >= 4 and num_tokens <= 128; + const bool l2_arrival_counter = + default_split_mn_barrier_opt or decode_l2_counter; + const bool l2_epilogue_requires_full_sync = + not l2_arrival_counter; + + // Tensormap construction + // Acts/weights: standard 2D TMA descriptors (FP8 K-major). + // Activation SF: per-128 channel float for L1, per-64 for L2 (MN-major, no swizzle). + // Weight SF: block (128, 128) raw float pointer (no TMA descriptor). + constexpr int kGranK = 128; + constexpr int kL2ActsSFGranK = 64; + const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, + hidden, config.num_max_pool_tokens, + config.block_k, config.block_m, + static_cast(l1_acts.stride(-2)), + config.swizzle_acts_mode); + const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, + config.num_padded_sf_pool_tokens, hidden, + config.block_m, kGranK, + 1, 0); + const int weight_tma_block_n = config.block_n > 256 ? 256 : config.block_n; + const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights, + hidden, num_experts_per_rank * intermediate_hidden * 2, + config.block_k, weight_tma_block_n, + static_cast(l1_weights.stride(-2)), + config.swizzle_weights_mode); + // L1 output (post-SwiGLU FP8): N is halved. The correctness path stages this tile in plain + // row-major SMEM before the TMA store. Later L2 TMA loads may still swizzle from this + // row-major global buffer into their own SMEM tile. + // The usual TMA store is issued per warpgroup, each writing a `WG_BLOCK_M` row tile from its + // own SMEM offset. The m64n128 2-WG split-N decode path is different: both warpgroups stage + // one joint 64-column L1-output tile and a single warpgroup issues the combined store, so the + // descriptor must cover the full block_m x (block_n / 2) tile. + const int num_epilogue_warpgroups_h = config.num_epilogue_threads / 128; + const bool split_n_warpgroups = + config.block_m == 64 and num_epilogue_warpgroups_h > 1 and + config.block_n % num_epilogue_warpgroups_h == 0 and + (config.block_n / num_epilogue_warpgroups_h == 64 or + config.block_n / num_epilogue_warpgroups_h == 128); + const bool split_mn_warpgroups = + config.block_m == 128 and config.block_n == 256 and num_epilogue_warpgroups_h == 4; + const int wg_split_m = split_n_warpgroups ? 1 : + (split_mn_warpgroups ? 2 : num_epilogue_warpgroups_h); + const int wg_split_n = split_n_warpgroups ? num_epilogue_warpgroups_h : + (split_mn_warpgroups ? 2 : 1); + DG_HOST_ASSERT(wg_split_m * wg_split_n == num_epilogue_warpgroups_h); + const int wg_block_m = config.block_m / wg_split_m; + const int wg_block_n = config.block_n / wg_split_n; + const int wg_l1_out_block_n = wg_block_n / 2; + const bool split_n_shares_sf = + split_n_warpgroups and wg_l1_out_block_n < kL2ActsSFGranK; + const int l1_output_swizzle_mode = 0; + const int l1_output_box_n = + split_n_shares_sf ? config.block_n / 2 : wg_l1_out_block_n; + const int l1_output_box_m = + split_n_shares_sf ? config.block_m : wg_block_m; + const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_max_pool_tokens, + l1_output_box_n, l1_output_box_m, + static_cast(l2_acts.stride(-2)), + l1_output_swizzle_mode); + const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_max_pool_tokens, + config.block_k, config.block_m, + static_cast(l2_acts.stride(-2)), + config.swizzle_acts_mode); + const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, + config.num_padded_sf_pool_tokens, intermediate_hidden, + config.block_m, kL2ActsSFGranK, + 1, 0); + const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights, + intermediate_hidden, num_experts_per_rank * hidden, + config.block_k, weight_tma_block_n, + static_cast(l2_weights.stride(-2)), + config.swizzle_weights_mode); + + // Stats can be optional + int* cumulative_local_expert_recv_stats_ptr = nullptr; + if (cumulative_local_expert_recv_stats.has_value()) + cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); + + // Launch + const auto num_sms = device_runtime->get_num_sms(); + const SM90FP8MegaMoERuntime::Args args = { + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .hidden = hidden, .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, .num_topk = num_topk, + .num_ranks = num_ranks, + .activation_clamp = activation_clamp, + .fast_math = fast_math, + .epilogue_registers = epilogue_registers, + .reuse_accum_as_final = reuse_accum_as_final, + .l2_arrival_counter = l2_arrival_counter, + .l2_epilogue_requires_full_sync = l2_epilogue_requires_full_sync, + .split_phase_hot_path = split_phase_hot_path, + .config = config, + .y = y.data_ptr(), + .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, + .num_tokens = num_tokens, + .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + .tensor_map_l1_acts = tensor_map_l1_acts, + .tensor_map_l1_acts_sf = tensor_map_l1_acts_sf, + .tensor_map_l1_weights = tensor_map_l1_weights, + .l1_weights_sf = l1_weights_sf.data_ptr(), + .tensor_map_l1_output = tensor_map_l1_output, + .tensor_map_l2_acts = tensor_map_l2_acts, + .tensor_map_l2_acts_sf = tensor_map_l2_acts_sf, + .tensor_map_l2_weights = tensor_map_l2_weights, + .l2_weights_sf = l2_weights_sf.data_ptr(), + .launch_args = LaunchArgs(num_sms, config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, + config.smem_size, config.cluster_size) + }; + const auto code = SM90FP8MegaMoERuntime::generate(args); + const auto runtime = compiler->build("sm90_fp8_mega_moe", code); + SM90FP8MegaMoERuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 39bdcc0d53..033930e2c8 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -90,6 +90,11 @@ fp8_fp4_mega_moe, fp4_fp4_mega_moe, bf16_mega_moe, + # SM90 (Hopper) MegaMoE + Sm90SymmBuffer, + get_symm_buffer_for_sm90_mega_moe, + transform_weights_for_mega_moe_sm90, + fp8_mega_moe_sm90, ) # Some utils diff --git a/deep_gemm/include/deep_gemm/comm/barrier.cuh b/deep_gemm/include/deep_gemm/comm/barrier.cuh index 50b894e937..c1672cda66 100644 --- a/deep_gemm/include/deep_gemm/comm/barrier.cuh +++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh @@ -18,8 +18,14 @@ CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { cute::cluster_wait(); } -template -CUTLASS_DEVICE void grid_sync(const layout::Workspace& workspace, +// NOTES: `WorkspaceT` is templated (rather than hard-coded to `layout::Workspace`) so that +// architectures with a different workspace layout (e.g. SM90 MegaMoE's pool-based +// `layout::MegaMoESM90Workspace`) can reuse the same grid/NVLink barrier logic, as long as they +// expose the same `get_grid_sync_count_ptr`/`get_nvl_barrier_counter_ptr`/`get_nvl_barrier_signal_ptr` +// accessors. This is purely a signature generalization; existing SM100 call sites are unaffected +// since `layout::Workspace` is still deduced automatically. +template +CUTLASS_DEVICE void grid_sync(const WorkspaceT& workspace, const uint32_t& sm_idx, const uint32_t& thread_idx, const sync_scope_t& sync_scope) { // NOTES: the implementation idea is from `cooperative_groups::this_grid().sync()` @@ -43,8 +49,8 @@ CUTLASS_DEVICE void grid_sync(const layout::Workspace& workspace, sync_scope(); } -template -CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, +template +CUTLASS_DEVICE void nvlink_barrier(const WorkspaceT& workspace, const layout::SymBuffer& sym_buffer, const uint32_t& sm_idx, const uint32_t& thread_idx, const sync_scope_t& sync_scope, @@ -75,9 +81,18 @@ CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, const auto start_clock = clock64(); while (ptx::ld_acq_sys(signal_ptr) != target) { if (clock64() - start_clock >= kNumTimeoutCycles) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1000) && \ + !(defined(DG_NVLINK_BARRIER_VERBOSE_TIMEOUT) && DG_NVLINK_BARRIER_VERBOSE_TIMEOUT) + // NOTES: on SM90, a bare `trap` is used instead of the verbose `printf` + + // assert path below, mirroring upstream PR #36. This avoids pulling in extra + // device-side formatting code on the Hopper MegaMoE hot path; define + // `DG_NVLINK_BARRIER_VERBOSE_TIMEOUT=1` to opt back into the verbose message. + DG_TRAP_ONLY_DEVICE_ASSERT(false); +#else printf("DeepGEMM NVLink barrier timeout: rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); DG_DEVICE_ASSERT(false and "NVLink barrier timeout"); +#endif } } } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh new file mode 100644 index 0000000000..b5ac5eab39 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -0,0 +1,1936 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#define __CLION_IDE__ + +namespace deep_gemm { + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_clamp_gate(float x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + return x; +} + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_clamp_up(float x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + return x; +} + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_silu(float x) { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; +} + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_swiglu(float g, float u) { + g = sm90_fp8_mega_moe_clamp_gate(g); + u = sm90_fp8_mega_moe_clamp_up(u); + return sm90_fp8_mega_moe_silu(g) * u; +} + +__forceinline__ __device__ void sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( + const float2& amax, float2& sf, float2& sf_inv) { + constexpr float kScale = 1.0f / 448.0f; + const auto scaled = make_float2(__fmul_rn(amax.x, kScale), __fmul_rn(amax.y, kScale)); + const auto exp_x = math::fast_log2_ceil(scaled.x); + const auto exp_y = math::fast_log2_ceil(scaled.y); + sf.x = math::fast_pow2(exp_x), sf_inv.x = math::fast_pow2(-exp_x); + sf.y = math::fast_pow2(exp_y), sf_inv.y = math::fast_pow2(-exp_y); +} + +template +CUTLASS_DEVICE void sm90_fp8_mega_moe_for_each_block_split( + sched::MegaMoESM90Scheduler& scheduler, + L1Func&& l1_func, L2Func&& l2_func) { + scheduler.fetch_expert_recv_count(); + scheduler.set_expert_idx(0); + + while (true) { + CUTE_TIE_DECL(scheduler.get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx); + if (block_phase == sched::MegaMoESM90BlockPhase::None) + break; + + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + l1_func(current_local_expert_idx, kNumL1BlockKs, m_block_idx, n_block_idx); + } else { + l2_func(current_local_expert_idx, kNumL2BlockKs, m_block_idx, n_block_idx); + } + } +} + +// ============================================================================ +// SM90 (Hopper) FP8 MegaMoE — full implementation +// ---------------------------------------------------------------------------- +// Pipeline (cluster=1, no TMA multicast): +// * Dispatch warps: pull tokens (FP8) and SF (per-128 channel float) from +// remote ranks via NVLink into the local L1 pool. +// * GEMM TMA-load warps (1 for A+SFA, 1 for B+SFB) feed the pipeline stages. +// * Math warpgroups (totalling kNumEpilogueThreads) consume each +// stage with WGMMA, accumulate into registers, then run the epilogue: +// - L1 (Linear1): SwiGLU with gate/up granularity-8 interleaved layout, +// per-row amax over each output-SF group, FP8 e4m3 quantize, STSM into +// SMEM, TMA store to local L1 output buffer. +// The per-row SF is written as a *float* into the L2-acts SF buffer at +// per-64 K granularity (one SF per L1 N block), so each block is fully +// self-contained and no cross-CTA amax synchronisation is needed. +// - L2 (Linear2): BF16 cast of the GEMM output, STSM into SMEM, then +// NVLink scatter to remote combine buffers. +// * After all GEMM blocks, the math warps run the COMBINE step (top-k +// reduction in BF16) — ported verbatim from the SM100 kernel. +// ============================================================================ + +template < + uint32_t kNumMaxTokensPerRank, + uint32_t kHidden, uint32_t kIntermediateHidden, + uint32_t kNumExperts, uint32_t kNumTopk, + uint32_t kNumExpertsPerWave, + uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, + uint32_t kNumMaxPoolTokens, + uint32_t kNumPaddedSFPoolTokens, + uint32_t kNumStages, + uint32_t kNumDispatchThreads, uint32_t kNumNonEpilogueThreads, + uint32_t kNumEpilogueThreads, + uint32_t kNumSMs, uint32_t kNumRanks, + float kActivationClamp, + bool kFastMath, + uint32_t kEpilogueRegisterBudget, + bool kReuseAccumAsFinal, + bool kL2ArrivalCounter, + bool kL2EpilogueRequiresFullSync, + bool kSplitPhaseHotPath, + uint32_t L1_SHAPE_N = kIntermediateHidden * 2, + uint32_t L1_SHAPE_K = kHidden, + uint32_t L2_SHAPE_N = kHidden, + uint32_t L2_SHAPE_K = kIntermediateHidden, + uint32_t kNumDispatchWarps = kNumDispatchThreads / 32, + uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32, + uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32, + uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4, + uint32_t kNumThreads = kNumDispatchThreads + kNumNonEpilogueThreads + kNumEpilogueThreads, + uint32_t kNumTokensPerWarp = 32 / kNumTopk, + uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks +> +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void +sm90_fp8_mega_moe_impl(void* y, + int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const float* __restrict__ l1_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const float* __restrict__ l2_weights_sf) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + // ===================================================================== + // Template checks + // ===================================================================== + DG_STATIC_ASSERT(kNumDispatchThreads >= 64 and kNumDispatchThreads % 64 == 0, + "Invalid number of dispatch threads"); + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 64 or kNumNonEpilogueThreads == 128, + "Invalid number of GEMM TMA warps"); + DG_STATIC_ASSERT((kNumDispatchThreads + kNumNonEpilogueThreads) % 128 == 0, + "Math warpgroup start must be 128-thread aligned"); + DG_STATIC_ASSERT(kNumEpilogueThreads % 128 == 0, "Invalid number of math/epilogue threads"); + DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + DG_STATIC_ASSERT(BLOCK_M % 64 == 0, "BLOCK_M must be a multiple of WGMMA::M (64)"); + DG_STATIC_ASSERT(BLOCK_N == 128 or BLOCK_N == 256 or BLOCK_N == 512, + "SM90 MegaMoE supports CTA BLOCK_N=128/256/512"); + DG_STATIC_ASSERT(BLOCK_K == 128, "BLOCK_K is fixed to 128 (per-128 SF)"); + + // ===================================================================== + // Thread / warp identification + // ===================================================================== + const uint32_t sm_idx = blockIdx.x; + const uint32_t thread_idx = threadIdx.x; + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + + // Prefetch all TMA descriptors at the very beginning + if (warp_idx == 0 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); + } + + // ===================================================================== + // Workspaces and symmetric buffer slicing (mirror SM100 layout, except SF + // for L2 activations uses per-64 K granularity) + // ===================================================================== + const auto workspace = layout::MegaMoESM90Workspace( + sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk); + + constexpr auto fp8_token_layout = layout::Data(kHidden); + constexpr auto bf16_token_layout = layout::Data(kHidden * sizeof(nv_bfloat16)); + constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden); + // Per-128 K float SF: 4 bytes per per-128 group => `kHidden / 32` bytes/token (same as SM100 packing) + constexpr auto fp8_sf_layout = layout::Data(kHidden / 32); + // Per-64 K float SF (SM90 only): 4 bytes per per-64 group => `kIntermediateHidden / 16` bytes/token + constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 16); + constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * sizeof(int64_t), false); + constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * sizeof(float), false); + constexpr auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + + // Registered input area + const auto input_token_buffer = layout::Buffer(fp8_token_layout, 1, kNumMaxTokensPerRank, workspace.get_end_ptr()); + const auto input_sf_buffer = layout::Buffer(fp8_sf_layout, 1, kNumMaxTokensPerRank, input_token_buffer.get_end_ptr()); + const auto input_topk_idx_buffer = layout::Buffer(input_topk_idx_layout, 1, kNumMaxTokensPerRank, input_sf_buffer.get_end_ptr()); + const auto input_topk_weights_buffer = layout::Buffer(input_topk_weights_layout, 1, kNumMaxTokensPerRank, input_topk_idx_buffer.get_end_ptr()); + + // L1 input area + const auto l1_token_buffer = layout::Buffer(fp8_token_layout, 1, kNumMaxPoolTokens, input_topk_weights_buffer.get_end_ptr()); + const auto l1_sf_buffer = layout::Buffer(fp8_sf_layout, 1, kNumPaddedSFPoolTokens, l1_token_buffer.get_end_ptr()); + const auto l1_topk_weights_buffer = layout::Buffer(l1_topk_weights_layout, 1, kNumMaxPoolTokens, l1_sf_buffer.get_end_ptr()); + + // L2 input area + const auto l2_token_buffer = layout::Buffer(fp8_intermediate_token_layout, 1, kNumMaxPoolTokens, l1_topk_weights_buffer.get_end_ptr()); + const auto l2_sf_buffer = layout::Buffer(fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens, l2_token_buffer.get_end_ptr()); + + // Combine input area + const auto combine_token_buffer = layout::Buffer(bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, l2_sf_buffer.get_end_ptr()); + + // ===================================================================== + // GEMM data types and shape constants + // ===================================================================== + using a_dtype_t = cutlass::float_e4m3_t; + using b_dtype_t = cutlass::float_e4m3_t; + constexpr bool kSplitNWarpgroups = + BLOCK_M == 64 and kNumEpilogueWarpgroups > 1 and + BLOCK_N % kNumEpilogueWarpgroups == 0 and + ((BLOCK_N / kNumEpilogueWarpgroups == 64) or (BLOCK_N / kNumEpilogueWarpgroups == 128)); + constexpr bool kSplitMNWarpgroups = + BLOCK_M == 128 and BLOCK_N == 256 and kNumEpilogueWarpgroups == 4; + constexpr uint32_t kWarpgroupSplitM = kSplitNWarpgroups ? 1 : + (kSplitMNWarpgroups ? 2 : kNumEpilogueWarpgroups); + constexpr uint32_t kWarpgroupSplitN = kSplitNWarpgroups ? kNumEpilogueWarpgroups : + (kSplitMNWarpgroups ? 2 : 1); + constexpr uint32_t WG_BLOCK_M = BLOCK_M / kWarpgroupSplitM; + constexpr uint32_t WG_BLOCK_N = BLOCK_N / kWarpgroupSplitN; + constexpr uint32_t kNumCombineWarps = kNumEpilogueWarps; + using L1WGMMA = typename mma::sm90::FP8MMASelector::type; // M=64, N=WG_BLOCK_N, K=32 + using L2WGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kL1OutputArrivalParts = 1; + static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, + "Unexpected WGMMA shape"); + DG_STATIC_ASSERT(kWarpgroupSplitM * kWarpgroupSplitN == kNumEpilogueWarpgroups, + "Invalid warpgroup split"); + DG_STATIC_ASSERT(WG_BLOCK_M == L1WGMMA::M, + "Each warpgroup must run exactly one WGMMA-M tile"); + DG_STATIC_ASSERT(kNumCombineWarps <= kNumEpilogueWarps, + "Combine warp count must fit in epilogue warps"); + + // Cluster=1 -> no multicast, A/B are loaded full-sized + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M; + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; + constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; // post-SwiGLU + constexpr uint32_t WG_L1_OUT_BLOCK_N = WG_BLOCK_N / 2; + // When WG_L1_OUT_BLOCK_N < 64 the two N-split warpgroups jointly own a + // single per-64 L2-acts SF group, so they must publish ONE shared SF slot + // (k_sf_idx == n_block_idx) instead of one per warpgroup. The amax that + // feeds that shared SF must be reduced across both warpgroups. + constexpr bool kSplitNSharesSF = kSplitNWarpgroups and (WG_L1_OUT_BLOCK_N < 64); + constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); // 128 + constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t); // 128 + constexpr uint32_t kSwizzleCDMode = 128; + constexpr uint32_t kGranK = 128; // L1 acts SF, weights SF + constexpr uint32_t kL2ActsSFGranK = 64; // L2 acts SF (per-64 K, SM90 only) + + // ===================================================================== + // Shared memory layout + // ===================================================================== + constexpr uint32_t kSharedMemoryAlignment = 1024; + extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); + constexpr uint32_t SMEM_SEND_BUFFER_SIZE = + math::constexpr_align(fp8_token_layout.get_num_bytes() * kNumDispatchWarps, kSharedMemoryAlignment); + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); + // SFA per-stage must be sized for the larger of L1 (BLOCK_M floats) and L2 (2*BLOCK_M floats per-64). + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = + math::constexpr_align(2 * BLOCK_M * sizeof(float), 128u); + // Block (128, 128) weight SF is loaded directly from global by the math + // warpgroup, so no SMEM is needed. + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = 0; + + // CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte) and + // L2 BF16 (BLOCK_M * BLOCK_N * 2 bytes). Split-M warpgroups own disjoint + // row slices; shared-SF split-N warpgroups stage disjoint column slices + // into one CTA tile. + constexpr uint32_t SMEM_CD_L1_SIZE = BLOCK_M * L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t); + constexpr uint32_t SMEM_CD_L2_SIZE = BLOCK_M * BLOCK_N * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_CD_SIZE = math::constexpr_align( + SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE, kSharedMemoryAlignment); + + constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = + SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + + // SMEM pointers + auto smem_expert_count = reinterpret_cast(smem_buffer); + const auto smem_send_buffers = layout::Buffer( + fp8_token_layout, kNumDispatchWarps, 1, + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE)); + + auto smem_gemm_base = math::advance_ptr( + smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE); + + // CD output is shared by L1 (FP8) and L2 (BF16); reinterpret-cast as needed. + auto smem_cd_l1 = reinterpret_cast(smem_gemm_base); + auto smem_cd_l2 = reinterpret_cast(smem_gemm_base); + + auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); + }); + auto sf_start_ptr = math::advance_ptr(smem_gemm_base, + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); + auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); + }); + + // Barriers live after SF (SFB is loaded directly from global, no SMEM) + auto barrier_start_ptr = reinterpret_cast( + sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE); + auto dispatch_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + i; }); + auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + i; }); + auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages + i; }); + auto combine_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + i; }); + + // ===================================================================== + // Initialization + // ===================================================================== + if (warp_idx == 0) { + // Clean expert-count shared memory + #pragma unroll + for (uint32_t i = lane_idx; i < kNumExperts; i += 32) + ptx::st_shared(smem_expert_count + i, 0u); + } else if (warp_idx == 1) { + // Init dispatch m-barriers + #pragma unroll + for (uint32_t i = lane_idx; i < kNumDispatchWarps; i += 32) + dispatch_barriers[i]->init(1); + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 2) { + // Init GEMM full/empty barriers and combine barriers + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + // Two producer warps (A+SFA loader, B+SFB loader) each call + // `arrive_and_expect_tx` per stage, so init count must be 2. + full_barriers[i]->init(2); + // Each math warp arrives once per stage release. + empty_barriers[i]->init(kNumEpilogueWarps); + } + #pragma unroll + for (uint32_t i = 0; i < kNumCombineWarps * 2; ++ i) + combine_barriers[i]->init(1); + } + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); + + // ===================================================================== + // Scheduler (cluster=1) + // ===================================================================== + auto scheduler = sched::MegaMoESM90Scheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, kNumExpertsPerWave, + kNumSMs, kNumRanks>(workspace); + + // Pipeline state shared by TMA loaders and math warpgroups + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + // Intra-SM barrier indices (mirroring SM100) + constexpr uint32_t kDispatchBarrierIdx = 0; + constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; + constexpr uint32_t kEpilogueFullBarrierIdx = 2; + constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; + + // Cross-rank NVLink barrier tags + constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; + constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + + // Register reconfiguration counts (chosen to fit in 64512 reg budget). + // For the 256-epilogue-thread split-N decode path: + // 64*48 + 64*40 + 256*168 = 48640 <= 64512. + // For the 512-epilogue-thread split-MN path, trim dispatch and loader roles + // so launch bounds still leave enough WGMMA registers. + // Reduced-thread decode (kNumThreads<=256) raises the launch-bounds + // register ceiling to 65536/256=256; grant the epilogue warpgroup the full + // 256 so the accumulator double-buffer fits without spilling. + // 64*48 + 64*40 + 128*256 = 38400 <= 64512. + constexpr uint32_t kNumEpilogueRegisters = + kEpilogueRegisterBudget == 0 ? + (kNumEpilogueThreads == 512 ? 112 : + (kNumEpilogueThreads == 256 ? 168 : + (kNumThreads <= 256u ? 256 : 208))) : + kEpilogueRegisterBudget; + // The 512-epilogue-thread path has only 3584 registers of headroom at + // epilogue=112. Raising epilogue to 120 is not viable without changing + // the role topology: dispatch=24 stalls the split-MN path, while + // non-epilogue=16 is below ptxas' setmaxnreg.dec legal minimum. + constexpr uint32_t kNumDispatchRegisters = + kNumEpilogueThreads == 512 ? 32 : 48; + constexpr uint32_t kNumNonEpilogueRegisters = + kNumEpilogueThreads == 512 ? 24 : 40; + DG_STATIC_ASSERT(kNumDispatchRegisters * kNumDispatchThreads + + kNumNonEpilogueRegisters * kNumNonEpilogueThreads + + kNumEpilogueRegisters * kNumEpilogueThreads <= 64512, + "Too many registers"); + + constexpr uint32_t kDispatchGridSyncIndex = 0; + constexpr uint32_t kEpilogueGridSyncIndex = 1; + + // ===================================================================== + // ROLE 1: DISPATCH WARPS + // Mirrors SM100 dispatch with two changes: + // * SF is per-128 channel float (no UTCCP transpose). We store the + // remote per-token SF directly into the local L1 SF buffer in + // MN-major layout: `local_sf[k_chunk * num_padded_sf_pool_tokens + token_idx]`. + // * The "token_idx_in_expert" → SF token index is now the simple + // per-block linear mapping (no 4×32 transpose). + // ===================================================================== + if (warp_idx < kNumDispatchWarps) { + cutlass::arch::warpgroup_reg_dealloc(); + + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + #pragma unroll + for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; + i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) { + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) + process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + }; + + // Count tokens per expert + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + atomicAdd_block(smem_expert_count + expert_idx, 1); + }); + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Stake out per-expert SM offsets via global atomic + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); + smem_expert_count[i] = static_cast( + ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Write source token-topk indices to remote ranks + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + const auto dst_rank_idx = expert_idx / kNumExpertsPerRank; + const auto dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1); + const auto dst_ptr = workspace.get_src_token_topk_idx_ptr( + expert_idx % kNumExpertsPerRank, sym_buffer.rank_idx, dst_slot_idx); + *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; + }); + + comm::grid_sync( + workspace, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + + if (sm_idx == 0) { + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + false, true); + + // Sync with epilogue warps before pulling tokens. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Token / SF pull loop + uint32_t pull_mbarrier_phase = 0; + const auto pull_buffer = smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0); + const auto pull_mbarrier = dispatch_barriers[warp_idx]; + + scheduler.fetch_expert_recv_count(); + + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalWarps) { + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++ current_expert_idx >= kNumExpertsPerRank) + break; + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + if (current_expert_idx >= kNumExpertsPerRank) + break; + + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t j = i * 32 + lane_idx; + stored_rank_count[i] = j < kNumRanks ? + static_cast(*workspace.get_expert_recv_count_ptr(j, current_expert_idx)) : 0; + } + } + + // Round-robin rank selection (identical to SM100) + uint32_t current_rank_in_expert_idx; + uint32_t remaining[kNumRanksPerLane]; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] = stored_rank_count[i]; + uint32_t offset = 0; + uint32_t token_idx_in_expert = token_idx - expert_start_idx; + uint32_t slot_idx = token_idx_in_expert; + uint32_t token_idx_in_rank; + while (true) { + uint32_t num_actives_in_lane = 0; + uint32_t min_in_lane = 0xffffffff; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + num_actives_in_lane += remaining[i] > 0; + if (remaining[i] > 0) + min_in_lane = cute::min(min_in_lane, remaining[i]); + } + const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane); + const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane); + + const uint32_t num_round_tokens = length * num_active_ranks; + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = slot_idx % num_active_ranks; + uint32_t num_seen_ranks = 0; + current_rank_in_expert_idx = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0); + const uint32_t num_active_lanes = __popc(mask); + if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes) + current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); + num_seen_ranks += num_active_lanes; + } + token_idx_in_rank = offset + (slot_idx / num_active_ranks); + break; + } + slot_idx -= num_round_tokens; + offset += length; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] -= cute::min(remaining[i], length); + } + + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + + // Pull token data. Overlap a remote TMA load with SF copy and + // then use TMA store to materialize the local L1 input. + if (cute::elect_one_sync()) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx), + pull_mbarrier, kHidden); + } + __syncwarp(); + + // Copy SF: per-128 K floats, written linearly (no UTCCP transpose). + constexpr uint32_t kNumSFFloats = kHidden / 128; + DG_STATIC_ASSERT(kNumSFFloats > 0 and kHidden % 128 == 0, "Invalid SF"); + const auto remote_sf_ptr = sym_buffer.map( + input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + const uint32_t sf_pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + #pragma unroll + for (uint32_t i = 0; i < math::constexpr_ceil_div(kNumSFFloats, 32u); ++ i) { + const uint32_t j = i * 32 + lane_idx; + if (j < kNumSFFloats) + local_sf_ptr[j * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[j]; + } + __syncwarp(); + + if (cute::elect_one_sync()) { + const auto weight = *sym_buffer.map( + input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; + } + __syncwarp(); + + if (cute::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kHidden); + ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); + + ptx::tma_store_1d( + l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(), + pull_buffer.get_base_ptr(), pull_buffer.get_num_bytes()); + + *workspace.get_token_src_metadata_ptr(pool_token_idx) = + {current_rank_in_expert_idx, src_token_idx, src_topk_idx}; + + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + ptx::red_add_rel( + workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1); + } + __syncwarp(); + } + + // Cleanup workspace, overlapping with combine. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + DG_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + } else { + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + const auto num_recv_tokens = static_cast( + *workspace.get_expert_recv_count_sum_ptr(i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + + expert_pool_block_offset = scheduler.get_pool_block_offset(i); + + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + DG_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); + if (warp_idx == 0) { + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + } else if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + + for (uint32_t j = thread_idx; j < kNumRanks; j += kNumDispatchThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + __syncwarp(); + + for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { + *workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + j) = 0; + *workspace.get_l2_arrival_mask_ptr(expert_pool_block_offset + j) = 0; + } + __syncwarp(); + } + } + + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + true, false); + + // ===================================================================== + // ROLE 2: GEMM TMA LOAD warps (load A+SFA, B+SFB) + // Warps inside `kNumNonEpilogueThreads`: warp 0 loads A + SFA, + // warp 1 loads B. + // ===================================================================== + } else if (warp_idx == kNumDispatchWarps) { + cutlass::arch::warpgroup_reg_dealloc(); + + auto process_a_sfa_block = [&](const auto& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const auto tensor_map_a_ptr = block_phase == sched::MegaMoESM90BlockPhase::Linear2 + ? &tensor_map_l2_acts : &tensor_map_l1_acts; + const auto tensor_map_sfa_ptr = block_phase == sched::MegaMoESM90BlockPhase::Linear2 + ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; + + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + + // Wait for the pool to be ready + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); + const auto expected = scheduler.template get_valid_m(); + while (ptx::ld_acq(ptr) != expected); + } else { + constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; + if constexpr (kL2ArrivalCounter) { + const auto ptr = reinterpret_cast( + workspace.get_l2_arrival_mask_ptr(pool_block_idx)); + const uint32_t active_m_wgs = math::ceil_div( + scheduler.template get_valid_m(), WG_BLOCK_M); + const uint32_t expected = + kNumL1BlockNs * active_m_wgs * kWarpgroupSplitN * kL1OutputArrivalParts; + while (ptx::ld_acq(ptr) != expected); + } else { + const auto ptr = workspace.get_l2_arrival_mask_ptr(pool_block_idx); + // Each L1 N block sets one bit; total bits = L1_SHAPE_N / BLOCK_N. + const uint64_t expected = (kNumL1BlockNs >= 64) + ? ~0ull : ((1ull << kNumL1BlockNs) - 1ull); + while (ptx::ld_acq_gpu(ptr) != expected); + } + } + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // TMA load A + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], + k_idx, m_idx, 1); + + // TMA load SFA + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + // L1 SFA per-128: load (BLOCK_M, 1) at K=k_block_idx + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], + m_idx, k_block_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + BLOCK_M * sizeof(float)); + } else { + // L2 SFA per-64: descriptor box is (block_mn, 1) (see make_tma_sf_desc), + // so we must issue two single-group TMAs and place them at smem offsets + // 0 and BLOCK_M to match math's load offsets (`+ 0 * BLOCK_M` / `+ 1 * BLOCK_M`). + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], + m_idx, k_block_idx * 2, 1); + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx] + BLOCK_M, + m_idx, k_block_idx * 2 + 1, 1); + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + 2 * BLOCK_M * sizeof(float)); + } + } + __syncwarp(); + } + }; + + if constexpr (kSplitPhaseHotPath) { + sm90_fp8_mega_moe_for_each_block_split( + scheduler, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_a_sfa_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_a_sfa_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } else { + scheduler.for_each_block([&](const sched::MegaMoESM90BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_a_sfa_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } + + } else if (warp_idx == kNumDispatchWarps + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + + scheduler.for_each_block([&](const sched::MegaMoESM90BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const auto tensor_map_b_ptr = + block_phase == sched::MegaMoESM90BlockPhase::Linear2 ? &tensor_map_l2_weights : &tensor_map_l1_weights; + + const uint32_t shape_n = block_phase == sched::MegaMoESM90BlockPhase::Linear2 ? L2_SHAPE_N : L1_SHAPE_N; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + const uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // TMA load B (weight SF is now loaded directly by math warps from global) + if constexpr (LOAD_BLOCK_N <= 256) { + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], smem_b[stage_idx], + k_idx, n_idx, 1); + } else { + DG_STATIC_ASSERT(LOAD_BLOCK_N % 256 == 0, + "Large B tiles are loaded as 256-column TMA slices"); + #pragma unroll + for (uint32_t b_slice_idx = 0; b_slice_idx < LOAD_BLOCK_N / 256; ++ b_slice_idx) { + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], + smem_b[stage_idx] + b_slice_idx * 256 * BLOCK_K, + k_idx, n_idx + b_slice_idx * 256, 1); + } + } + + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_B_SIZE_PER_STAGE); + } + __syncwarp(); + } + }); + + } else if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // Idle non-epilogue warps (kNumDispatchWarps+2, +3). They must still + // participate in the warpgroup-collective `setmaxnreg.dec.sync.aligned` + // so that the math warpgroup's `warpgroup_reg_alloc` can succeed. + cutlass::arch::warpgroup_reg_dealloc(); + + } else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // ===================================================================== + // ROLE 3: MATH WARPGROUPS (WGMMA + epilogue + combine) + // ===================================================================== + cutlass::arch::warpgroup_reg_alloc(); + + const uint32_t epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const uint32_t epilogue_wg_idx = epilogue_warp_idx / 4; + const uint32_t epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const uint32_t warp_idx_in_wg = epilogue_warp_idx % 4; + + // WGMMA-output register layout helpers + const uint32_t row_idx = lane_idx / 4; + const uint32_t col_idx = lane_idx % 4; + const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + + // When the two N-split warpgroups share a single per-64 SF group they + // also stage into ONE shared row-major L1-output tile (stride + // L1_OUT_BLOCK_N), each writing its own WG_L1_OUT_BLOCK_N-column half, + // so a single combined TMA store matches the host descriptor box. + constexpr uint32_t WG_SMEM_CD_L1_STRIDE_N = + kSplitNSharesSF ? L1_OUT_BLOCK_N : WG_L1_OUT_BLOCK_N; + constexpr uint32_t WG_SMEM_CD_L2_STRIDE_N = WG_BLOCK_N; + + // Sync with dispatch in the full communication path. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + auto process_math_block = [&](const auto& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const uint32_t valid_m = scheduler.template get_valid_m(); + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t n_idx = n_block_idx * BLOCK_N; + const uint32_t epilogue_wg_m_idx = epilogue_wg_idx / kWarpgroupSplitN; + const uint32_t epilogue_wg_n_idx = epilogue_wg_idx - epilogue_wg_m_idx * kWarpgroupSplitN; + const uint32_t wg_n_offset = epilogue_wg_n_idx * WG_BLOCK_N; + const uint32_t wg_l1_out_n_offset = epilogue_wg_n_idx * WG_L1_OUT_BLOCK_N; + const uint32_t row_base = epilogue_wg_m_idx * WG_BLOCK_M; + const uint32_t row_offset_r0 = row_base + r_0; + const uint32_t row_offset_r1 = row_base + r_1; + const uint32_t sf_n_block_idx = kSplitNSharesSF ? n_block_idx + : (n_block_idx * kWarpgroupSplitN + epilogue_wg_n_idx); + const uint32_t smem_a_wg_offset = epilogue_wg_m_idx * WG_BLOCK_M * BLOCK_K; + const uint32_t smem_b_wg_offset = epilogue_wg_n_idx * WG_BLOCK_N * BLOCK_K; + // In the shared-tile case the WG stages into the joint L1-output tile + // at its own column offset (row stride L1_OUT_BLOCK_N); otherwise each + // WG owns a disjoint contiguous WG_BLOCK_M x WG_L1_OUT_BLOCK_N slice. + const uint32_t smem_cd_l1_wg_offset = kSplitNSharesSF ? wg_l1_out_n_offset + : (epilogue_wg_idx * WG_BLOCK_M * WG_L1_OUT_BLOCK_N); + const uint32_t smem_cd_l2_wg_offset = epilogue_wg_idx * WG_BLOCK_M * WG_BLOCK_N; + const bool valid_r0 = row_offset_r0 < valid_m; + const bool valid_r1 = row_offset_r1 < valid_m; + + // ---------------- GEMM ---------------- + using WGMMA = L1WGMMA; + constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; + float final_accum[kAccumPerThread] = {}; + + if constexpr (kReuseAccumAsFinal) { + auto prescale_l1_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& gate_sf, const float& up_sf) { + const float inv_s0_gate = kFastMath ? math::fast_rcp(scale_a_0 * gate_sf) : 1.0f / (scale_a_0 * gate_sf); + const float inv_s1_gate = kFastMath ? math::fast_rcp(scale_a_1 * gate_sf) : 1.0f / (scale_a_1 * gate_sf); + const float inv_s0_up = kFastMath ? math::fast_rcp(scale_a_0 * up_sf) : 1.0f / (scale_a_0 * up_sf); + const float inv_s1_up = kFastMath ? math::fast_rcp(scale_a_1 * up_sf) : 1.0f / (scale_a_1 * up_sf); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float inv_s0 = (i & 1u) ? inv_s0_up : inv_s0_gate; + const float inv_s1 = (i & 1u) ? inv_s1_up : inv_s1_gate; + final_accum[i*4+0] *= inv_s0; + final_accum[i*4+1] *= inv_s0; + final_accum[i*4+2] *= inv_s1; + final_accum[i*4+3] *= inv_s1; + } + }; + auto postscale_l1_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& gate_sf, const float& up_sf) { + const float s0_gate = scale_a_0 * gate_sf; + const float s1_gate = scale_a_1 * gate_sf; + const float s0_up = scale_a_0 * up_sf; + const float s1_up = scale_a_1 * up_sf; + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float s0 = (i & 1u) ? s0_up : s0_gate; + const float s1 = (i & 1u) ? s1_up : s1_gate; + final_accum[i*4+0] *= s0; + final_accum[i*4+1] *= s0; + final_accum[i*4+2] *= s1; + final_accum[i*4+3] *= s1; + } + }; + auto prescale_l2_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& l2_sf) { + const float inv_s0 = kFastMath ? math::fast_rcp(scale_a_0 * l2_sf) : 1.0f / (scale_a_0 * l2_sf); + const float inv_s1 = kFastMath ? math::fast_rcp(scale_a_1 * l2_sf) : 1.0f / (scale_a_1 * l2_sf); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= inv_s0; + final_accum[i*4+1] *= inv_s0; + final_accum[i*4+2] *= inv_s1; + final_accum[i*4+3] *= inv_s1; + } + }; + auto postscale_l2_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& l2_sf) { + const float s0 = scale_a_0 * l2_sf; + const float s1 = scale_a_1 * l2_sf; + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= s0; + final_accum[i*4+1] *= s0; + final_accum[i*4+2] *= s1; + final_accum[i*4+3] *= s1; + } + }; + auto rescale_l1_final = [&](const float& prev_scale_a_0, const float& prev_scale_a_1, + const float& prev_gate_sf, const float& prev_up_sf, + const float& scale_a_0, const float& scale_a_1, + const float& gate_sf, const float& up_sf) { + const float r0_gate = (prev_scale_a_0 * prev_gate_sf) * + (kFastMath ? math::fast_rcp(scale_a_0 * gate_sf) : 1.0f / (scale_a_0 * gate_sf)); + const float r1_gate = (prev_scale_a_1 * prev_gate_sf) * + (kFastMath ? math::fast_rcp(scale_a_1 * gate_sf) : 1.0f / (scale_a_1 * gate_sf)); + const float r0_up = (prev_scale_a_0 * prev_up_sf) * + (kFastMath ? math::fast_rcp(scale_a_0 * up_sf) : 1.0f / (scale_a_0 * up_sf)); + const float r1_up = (prev_scale_a_1 * prev_up_sf) * + (kFastMath ? math::fast_rcp(scale_a_1 * up_sf) : 1.0f / (scale_a_1 * up_sf)); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float r0 = (i & 1u) ? r0_up : r0_gate; + const float r1 = (i & 1u) ? r1_up : r1_gate; + final_accum[i*4+0] *= r0; + final_accum[i*4+1] *= r0; + final_accum[i*4+2] *= r1; + final_accum[i*4+3] *= r1; + } + }; + auto rescale_l2_final = [&](const float& prev_scale_a_0, const float& prev_scale_a_1, + const float& prev_l2_sf, + const float& scale_a_0, const float& scale_a_1, + const float& l2_sf) { + const float r0 = (prev_scale_a_0 * prev_l2_sf) * + (kFastMath ? math::fast_rcp(scale_a_0 * l2_sf) : 1.0f / (scale_a_0 * l2_sf)); + const float r1 = (prev_scale_a_1 * prev_l2_sf) * + (kFastMath ? math::fast_rcp(scale_a_1 * l2_sf) : 1.0f / (scale_a_1 * l2_sf)); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= r0; + final_accum[i*4+1] *= r0; + final_accum[i*4+2] *= r1; + final_accum[i*4+3] *= r1; + } + }; + auto rescale_l2_act_final = [&](const float& prev_scale_a_0, const float& prev_scale_a_1, + const float& scale_a_0, const float& scale_a_1) { + const float r0 = prev_scale_a_0 * (kFastMath ? math::fast_rcp(scale_a_0) : 1.0f / scale_a_0); + const float r1 = prev_scale_a_1 * (kFastMath ? math::fast_rcp(scale_a_1) : 1.0f / scale_a_1); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= r0; + final_accum[i*4+1] *= r0; + final_accum[i*4+2] *= r1; + final_accum[i*4+3] *= r1; + } + }; + + if constexpr (kHidden >= 7168) { + float prev_scale_a_0 = 1.0f, prev_scale_a_1 = 1.0f; + float prev_gate_sf = 1.0f, prev_up_sf = 1.0f, prev_l2_sf = 1.0f; + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + + float scale_a_0_lo, scale_a_1_lo; + float scale_a_0_hi, scale_a_1_hi; + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + + constexpr uint32_t kL1SFKBlocks = kHidden / 128; + constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; + constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; + float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + const uint32_t gate_n = sf_n_block_idx / 2u; + const uint32_t up_n = kL1SFGateBlks + gate_n; + const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kL1SFKBlocks); + up_sf = __ldg(base + up_n * kL1SFKBlocks); + } else { + l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert + + sf_n_block_idx * kL2SFKBlocks + k_block_idx); + } + + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + if (k_block_idx != 0) + rescale_l1_final(prev_scale_a_0, prev_scale_a_1, + prev_gate_sf, prev_up_sf, + scale_a_0_lo, scale_a_1_lo, + gate_sf, up_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + prev_scale_a_0 = scale_a_0_lo; + prev_scale_a_1 = scale_a_1_lo; + prev_gate_sf = gate_sf; + prev_up_sf = up_sf; + } else { + if (k_block_idx != 0) + rescale_l2_final(prev_scale_a_0, prev_scale_a_1, prev_l2_sf, + scale_a_0_lo, scale_a_1_lo, l2_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + rescale_l2_act_final(scale_a_0_lo, scale_a_1_lo, + scale_a_0_hi, scale_a_1_hi); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + prev_scale_a_0 = scale_a_0_hi; + prev_scale_a_1 = scale_a_1_hi; + prev_l2_sf = l2_sf; + } + } + + if (num_k_blocks != 0) { + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + postscale_l1_final(prev_scale_a_0, prev_scale_a_1, + prev_gate_sf, prev_up_sf); + } else { + postscale_l2_final(prev_scale_a_0, prev_scale_a_1, prev_l2_sf); + } + } + } else { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + + float scale_a_0_lo, scale_a_1_lo; + float scale_a_0_hi, scale_a_1_hi; + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + + constexpr uint32_t kL1SFKBlocks = kHidden / 128; + constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; + constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; + float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + const uint32_t gate_n = sf_n_block_idx / 2u; + const uint32_t up_n = kL1SFGateBlks + gate_n; + const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kL1SFKBlocks); + up_sf = __ldg(base + up_n * kL1SFKBlocks); + } else { + l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert + + sf_n_block_idx * kL2SFKBlocks + k_block_idx); + } + + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + if (k_block_idx != 0) + prescale_l1_final(scale_a_0_lo, scale_a_1_lo, gate_sf, up_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + postscale_l1_final(scale_a_0_lo, scale_a_1_lo, gate_sf, up_sf); + } else { + if (k_block_idx != 0) + prescale_l2_final(scale_a_0_lo, scale_a_1_lo, l2_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + postscale_l2_final(scale_a_0_lo, scale_a_1_lo, l2_sf); + prescale_l2_final(scale_a_0_hi, scale_a_1_hi, l2_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + postscale_l2_final(scale_a_0_hi, scale_a_1_hi, l2_sf); + } + } + } + } else { + float accum[kAccumPerThread]; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + + // Read SF (must precede warpgroup_arrive) + float scale_a_0_lo, scale_a_1_lo; + float scale_a_0_hi, scale_a_1_hi; // Only used in L2 (per-64 K) + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + // L2: SFA layout is (K=2, M=BLOCK_M) MN-major; first half SF at offset 0, second at BLOCK_M + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + + // ----- Block (128, 128) weight SF (loaded directly from global) ----- + // L1 weight SF shape: (E, 2*IH/128, H/128) MN-major. The N axis is + // [gate(IH/128), up(IH/128)]; with the gate/up gran-8 interleave on + // the FP8 weight, each logical 128-wide N tile covers 64 rows of gate + // plus 64 rows of up taken from the same original 128-row block, so: + // gate_sf_n = sf_n_block_idx / 2 + // up_sf_n = (IH/128) + sf_n_block_idx / 2 + // + // L2 weight SF shape: (E, H/128, IH/128) MN-major. One scalar per + // logical 128x128 weight-SF tile, broadcast across the matching + // WGMMA accumulators. + // + // Load the weight scale after the barrier from all WG threads. + // This keeps scale loads close to their WGMMA use and lets the + // read-only cache coalesce the same-address accesses. + constexpr uint32_t kL1SFKBlocks = kHidden / 128; + constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; + constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; + float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + const uint32_t gate_n = sf_n_block_idx / 2u; + const uint32_t up_n = kL1SFGateBlks + gate_n; + const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kL1SFKBlocks); + up_sf = __ldg(base + up_n * kL1SFKBlocks); + } else { + l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert + + sf_n_block_idx * kL2SFKBlocks + k_block_idx); + } + + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + // Single per-128 K-block WGMMA group + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + // L1: gate/up alternate at gran=8 along N; each `i` block of 8 + // cols belongs entirely to one of {gate, up}, so .x and .y + // share the same scalar. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float sb = (i & 1u) ? up_sf : gate_sf; + final_accum[i*4+0] += scale_a_0_lo * sb * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_lo * sb * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_lo * sb * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_lo * sb * accum[i*4+3]; + } + } else { + // L2: split BLOCK_K=128 into two halves (per-64 SFA), each 2 WGMMAs. + // First half: K=0..63, SFA = scale_a_*_lo + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + // L2 first half: single scalar `l2_sf` broadcast across N. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] += scale_a_0_lo * l2_sf * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_lo * l2_sf * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_lo * l2_sf * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_lo * l2_sf * accum[i*4+3]; + } + + // Second half: K=64..127, SFA = scale_a_*_hi + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + // L2 second half: same broadcast scalar `l2_sf`. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] += scale_a_0_hi * l2_sf * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_hi * l2_sf * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_hi * l2_sf * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_hi * l2_sf * accum[i*4+3]; + } + } + } + } + + // Skip epilogue when block is past valid M (still must release via empty) + if (row_base >= valid_m) { + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + if constexpr (not kL2ArrivalCounter) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } else { + if constexpr (kL2EpilogueRequiresFullSync) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + return; + } + + if (block_phase == sched::MegaMoESM90BlockPhase::Linear1) { + + // ---------------- L1 EPILOGUE: activation + FP8 quantize + TMA store ---------------- + // Layout in `final_accum`: + // kAccumPerThread/4 chunks, each chunk = 4 floats per thread = + // (r0c0, r0c1, r1c0, r1c1). + // Gate and up chunks alternate; pair `p` uses chunks 2p and 2p+1. + // + // For each pair we produce 4 post-SwiGLU floats per thread, mapped to + // output cols (p*8 + col_idx*2 + {0,1}) for both r0 and r1. + + constexpr uint32_t kNumPairs = kAccumPerThread / 8; + float sf_r0, sf_inv_r0; + float sf_r1, sf_inv_r1; + + float swiglu_r0[kNumPairs][2]; + float swiglu_r1[kNumPairs][2]; + float amax_r0 = 0.0f, amax_r1 = 0.0f; + + auto clamp_gate = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + }; + auto clamp_up = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + }; + auto silu = [](float x) -> float { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; + }; + + #pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++ p) { + const uint32_t gate = 2 * p, up = 2 * p + 1; + + float g_r0_c0 = final_accum[gate*4 + 0]; + float g_r0_c1 = final_accum[gate*4 + 1]; + float g_r1_c0 = final_accum[gate*4 + 2]; + float g_r1_c1 = final_accum[gate*4 + 3]; + float u_r0_c0 = final_accum[up*4 + 0]; + float u_r0_c1 = final_accum[up*4 + 1]; + float u_r1_c0 = final_accum[up*4 + 2]; + float u_r1_c1 = final_accum[up*4 + 3]; + clamp_gate(g_r0_c0); + clamp_gate(g_r0_c1); + clamp_gate(g_r1_c0); + clamp_gate(g_r1_c1); + clamp_up(u_r0_c0); + clamp_up(u_r0_c1); + clamp_up(u_r1_c0); + clamp_up(u_r1_c1); + + if (valid_r0) { + swiglu_r0[p][0] = silu(g_r0_c0) * u_r0_c0; + swiglu_r0[p][1] = silu(g_r0_c1) * u_r0_c1; + amax_r0 = cute::max(amax_r0, cute::max(cute::abs(swiglu_r0[p][0]), cute::abs(swiglu_r0[p][1]))); + } else { + swiglu_r0[p][0] = 0.0f; + swiglu_r0[p][1] = 0.0f; + } + if (valid_r1) { + swiglu_r1[p][0] = silu(g_r1_c0) * u_r1_c0; + swiglu_r1[p][1] = silu(g_r1_c1) * u_r1_c1; + amax_r1 = cute::max(amax_r1, cute::max(cute::abs(swiglu_r1[p][0]), cute::abs(swiglu_r1[p][1]))); + } else { + swiglu_r1[p][0] = 0.0f; + swiglu_r1[p][1] = 0.0f; + } + } + + // Apply token weight: SwiGLU * topk_weight (single load per row) + const float weight_r0 = valid_r0 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r0) + .get_base_ptr() : 0.0f; + const float weight_r1 = valid_r1 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r1) + .get_base_ptr() : 0.0f; + #pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++ p) { + swiglu_r0[p][0] *= weight_r0; + swiglu_r0[p][1] *= weight_r0; + swiglu_r1[p][0] *= weight_r1; + swiglu_r1[p][1] *= weight_r1; + } + + amax_r0 *= cute::abs(weight_r0); + amax_r1 *= cute::abs(weight_r1); + + // Reduce amax across the 4 col-lanes that share the same row. In the + // SM90 WGMMA output layout, lanes with the same `lane_idx >> 2` and + // different `lane_idx & 3` partition the WG-owned output columns for + // the same r_0/r_1, so this is an INTRA-group reduction + // (`warp_reduce<4, false>`). Using `<4, true>` would instead merge + // amax across 8 different rows -- giving wrong per-row SF. + amax_r0 = math::warp_reduce<4, false>(amax_r0, math::ReduceMax()); + amax_r1 = math::warp_reduce<4, false>(amax_r1, math::ReduceMax()); + + // Phase 2: cross-WG amax. When two N-split warpgroups share one + // per-64 SF group, each WG so far only saw its own + // WG_L1_OUT_BLOCK_N columns; the true per-row amax spans both + // halves. Reduce across both warpgroups through a small smem + // scratch (carved from the upper, currently-unused half of the + // CD staging region) so BOTH WGs quantize with the SAME SF. + if constexpr (kSplitNSharesSF) { + float* amax_scratch = reinterpret_cast( + reinterpret_cast(smem_cd_l1) + SMEM_CD_SIZE / 2); + #pragma unroll + for (uint32_t i = epilogue_thread_idx; i < BLOCK_M; i += kNumEpilogueThreads) + amax_scratch[i] = 0.0f; + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (col_idx == 0) { + atomicMax(reinterpret_cast(&amax_scratch[r_0]), __float_as_uint(amax_r0)); + atomicMax(reinterpret_cast(&amax_scratch[r_1]), __float_as_uint(amax_r1)); + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + amax_r0 = amax_scratch[r_0]; + amax_r1 = amax_scratch[r_1]; + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + + // Compute SF and inverse SF for each row + float2 amax_pair = {amax_r0, amax_r1}; + float2 sf_pair, sf_inv_pair; + sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv(amax_pair, sf_pair, sf_inv_pair); + sf_r0 = sf_pair.x; sf_inv_r0 = sf_inv_pair.x; + sf_r1 = sf_pair.y; sf_inv_r1 = sf_inv_pair.y; + + // Quantize and write to the shared-memory staging tile. + auto* smem_cd_l1_wg = smem_cd_l1 + smem_cd_l1_wg_offset; + DG_STATIC_ASSERT(kNumPairs % 2 == 0, "L1 staging stores two 8-byte chunks at once"); + #pragma unroll + for (uint32_t p_base = 0; p_base < kNumPairs; p_base += 2) { + uint16_t r0_bits[2], r1_bits[2]; + #pragma unroll + for (uint32_t q = 0; q < 2; ++ q) { + const uint32_t p = p_base + q; + const float v00 = swiglu_r0[p][0] * sf_inv_r0; + const float v01 = swiglu_r0[p][1] * sf_inv_r0; + const float v10 = swiglu_r1[p][0] * sf_inv_r1; + const float v11 = swiglu_r1[p][1] * sf_inv_r1; + + const __nv_fp8x2_e4m3 r0_pair(make_float2(v00, v01)); + const __nv_fp8x2_e4m3 r1_pair(make_float2(v10, v11)); + r0_bits[q] = valid_r0 ? r0_pair.__x : 0u; + r1_bits[q] = valid_r1 ? r1_pair.__x : 0u; + } + + #pragma unroll + for (uint32_t q = 0; q < 2; ++ q) { + const uint32_t p = p_base + q; + const uint32_t col = p * 8 + col_idx * 2; + auto* p0 = reinterpret_cast( + smem_cd_l1_wg + r_0 * WG_SMEM_CD_L1_STRIDE_N + col); + auto* p1 = reinterpret_cast( + smem_cd_l1_wg + r_1 * WG_SMEM_CD_L1_STRIDE_N + col); + if (valid_r0) + *p0 = r0_bits[q]; + if (valid_r1) + *p1 = r1_bits[q]; + } + } + + // Write SF as float at `[token, n_block_idx]` in L2 acts SF buffer (per-64 layout). + // Each row is contributed by lanes col_idx in {0..3}; only col_idx == 0 writes. + // In the shared-SF split both warpgroups own the same per-64 group and rows, so + // only the first N-split warpgroup publishes the SF slot to avoid a write race. + if (col_idx == 0 and (not kSplitNSharesSF or epilogue_wg_n_idx == 0)) { + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + // SF buffer is (kNumPaddedSFPoolTokens x kIntermediateHidden/64), MN-major: + // addr[k_idx * num_padded_sf_pool_tokens + token_idx] + const uint32_t token_r0 = pool_block_idx * BLOCK_M + row_offset_r0; + const uint32_t token_r1 = pool_block_idx * BLOCK_M + row_offset_r1; + const uint32_t k_sf_idx = sf_n_block_idx; // one per-64 post-SwiGLU group + if (valid_r0) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0; + if (valid_r1) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1; + } + + // Sync the warpgroup before TMA store. In the shared-tile split + // both N-split warpgroups must finish writing their halves of the + // joint L1-output tile, so sync across all epilogue threads. + if constexpr (kSplitNSharesSF) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + else + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Issue TMA store of the entire tile. Padding rows beyond + // `valid_m` are written with stale/garbage FP8 to the L1-output + // pool buffer, but they are never consumed downstream: the L2 + // GEMM tile loads them, but its NVLink-scatter epilogue is + // gated by `m_idx_in_block >= valid_m`, and stale SF in the + // padding rows can produce NaN accumulators that simply stay + // in registers (only valid rows are converted to BF16 and + // STSM'd into smem). Using TMA for partial tiles is a large + // win for low-batch / decode where every tile is partial. + if constexpr (kSplitNSharesSF) { + // One combined store of the joint L1_OUT_BLOCK_N tile, issued + // by the first N-split warpgroup once both halves are staged. + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_l1_output, + smem_cd_l1, + out_n_idx, + m_idx + row_base); + cute::tma_store_arrive(); + } + } else { + if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N + wg_l1_out_n_offset; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_l1_output, + smem_cd_l1 + smem_cd_l1_wg_offset, + out_n_idx, + m_idx + row_base); + cute::tma_store_arrive(); + } + } + __syncwarp(); + ptx::tma_store_wait<0>(); + + // Notify L2 that this L1 output (and SF) is ready. Counter mode lets + // independent WG tiles publish arrivals without the CTA-wide barrier + // needed before the single bit-mask update. + if constexpr (kL2ArrivalCounter) { + if constexpr (kSplitNSharesSF) { + // The combined tile counts for both N-split warpgroups; the + // storing warpgroup publishes all kWarpgroupSplitN arrivals + // after its TMA store has drained. + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), + kWarpgroupSplitN); + } + } else if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), 1); + } + } else { + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_or_rel_gpu( + workspace.get_l2_arrival_mask_ptr(pool_block_idx), + 1ull << n_block_idx); + } + } + __syncwarp(); + // In the shared-tile split only the first warpgroup issues and + // drains the combined TMA store; gate the other warpgroup so it + // cannot overwrite the joint smem tile in the next block until + // that store has drained. + if constexpr (kSplitNSharesSF) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } else { + // ---------------- L2 EPILOGUE: BF16 cast + NVLink scatter ---------------- + constexpr uint32_t kNumRowsPerWarp = WG_BLOCK_M / 8; + + const uint32_t row_in_warp_block = lane_idx / 16; // 0 or 1 + const uint32_t lane_in_row = lane_idx % 16; + const uint32_t cols_per_lane = WG_BLOCK_N / 16; + + // STSM into smem_cd_l2 (BF16). Reuse SM100 column-swizzle layout. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 8; ++ i) { + // Each i consumes 8 floats (one 16x256b chunk in SM100 terms). + // For SM90 WGMMA layout, 8 floats per i correspond to 2 chunks of 4 floats: + // final_accum[i*8 + (0..3)] = chunk 2i: (r0c0, r0c1, r1c0, r1c1) + // final_accum[i*8 + (4..7)] = chunk 2i+1: same shape + const uint32_t chunk_lo = 2 * i, chunk_hi = 2 * i + 1; + + auto write_pair = [&](uint32_t row, uint32_t col, uint32_t packed) { + auto smem_ptr = smem_cd_l2 + + smem_cd_l2_wg_offset + + row * WG_BLOCK_N + + col; + // BF16 STS: 2 bf16 elements + *reinterpret_cast(smem_ptr) = packed; + }; + if (valid_r0) { + const uint32_t r0_lo = math::cast_into_bf16_and_pack( + final_accum[chunk_lo*4 + 0], final_accum[chunk_lo*4 + 1]); + const uint32_t r0_hi = math::cast_into_bf16_and_pack( + final_accum[chunk_hi*4 + 0], final_accum[chunk_hi*4 + 1]); + write_pair(r_0, chunk_lo * 8 + col_idx * 2, r0_lo); + write_pair(r_0, chunk_hi * 8 + col_idx * 2, r0_hi); + } + if (valid_r1) { + const uint32_t r1_lo = math::cast_into_bf16_and_pack( + final_accum[chunk_lo*4 + 2], final_accum[chunk_lo*4 + 3]); + const uint32_t r1_hi = math::cast_into_bf16_and_pack( + final_accum[chunk_hi*4 + 2], final_accum[chunk_hi*4 + 3]); + write_pair(r_1, chunk_lo * 8 + col_idx * 2, r1_lo); + write_pair(r_1, chunk_hi * 8 + col_idx * 2, r1_hi); + } + } + + // Each warp writes and then scatters only its own 16-row + // slice, so a warp-level fence is enough before reading + // back from shared memory. + __syncwarp(); + + // Scatter to remote ranks via NVLink (one row per warp-pair) + // Each warpgroup-warp covers 8 unique rows x 2 (r_0 + r_1 doubled by warps) + // Lane group of 16 within a warp -> 1 row. + // Each lane copies `cols_per_lane` BF16 (= cols_per_lane*2 bytes) as one + // vector. WG_BLOCK_N=128 -> 8 BF16 = uint4; WG_BLOCK_N=64 -> 4 BF16 = uint2. + using ScatterVec = std::conditional_t<(WG_BLOCK_N <= 64), uint2, uint4>; + DG_STATIC_ASSERT(cols_per_lane * sizeof(nv_bfloat16) == sizeof(ScatterVec), + "Scatter vector width must match cols_per_lane"); + #pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++ j) { + const uint32_t row_in_wg = warp_idx_in_wg * 16 + j * 2 + row_in_warp_block; + const uint32_t m_idx_in_block = row_base + row_in_wg; + if (m_idx_in_block >= valid_m) break; + + // Read cols_per_lane BF16 (= one ScatterVec) from smem + auto smem_ptr = smem_cd_l2 + + smem_cd_l2_wg_offset + + row_in_wg * WG_BLOCK_N + + lane_in_row * cols_per_lane; + const auto packed = *reinterpret_cast(smem_ptr); + + const auto src_metadata = *workspace.get_token_src_metadata_ptr(m_idx + m_idx_in_block); + const uint32_t dst_rank_idx = src_metadata.rank_idx; + const uint32_t dst_token_idx = src_metadata.token_idx; + const uint32_t dst_topk_idx = src_metadata.topk_idx; + const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) + .get_data_buffer(dst_token_idx); + auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), + (n_idx + wg_n_offset) * sizeof(nv_bfloat16) + lane_in_row * sizeof(ScatterVec)); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } + + if constexpr (kL2EpilogueRequiresFullSync) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + }; + + if constexpr (kSplitPhaseHotPath) { + sm90_fp8_mega_moe_for_each_block_split( + scheduler, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_math_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_math_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } else { + scheduler.for_each_block([&](const sched::MegaMoESM90BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_math_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } + + // ---------------- COMBINE ---------------- + // NVLink barrier first: signals remote ranks that this rank's GEMM + // outputs (NVLink scatter targets) are fully written. + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + + // Sync with dispatch (paired with dispatch's pre-cleanup sync) so that + // dispatch may now safely clean workspace state. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + if (epilogue_warp_idx >= kNumCombineWarps) + return; + + constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); + + constexpr uint32_t kNumChunkSlots = 3; + constexpr uint32_t kNumMaxRegistersForBuffer = 128; + constexpr uint32_t kDefaultNumChunks = + (kNumChunkSlots * kNumCombineWarps * kNumHiddenBytes <= SMEM_BEFORE_BARRIER_SIZE + and kHidden <= 32 * kNumMaxRegistersForBuffer) ? 1 : 2; + // Flash-style hidden=7168 is 7 * 1024. Splitting combine into 7 chunks + // keeps each lane's BF16 reduce accumulator much smaller without + // violating the 32-lane uint4 mapping. + constexpr uint32_t kSplitMNNumChunks = (kHidden % 7 == 0) ? 7 : (kHidden >= 1024 ? 4 : 1); + constexpr uint32_t kNumChunks = kSplitMNWarpgroups ? kSplitMNNumChunks : kDefaultNumChunks; + constexpr uint32_t kNumChunkBytes = kNumHiddenBytes / kNumChunks; + constexpr uint32_t kNumChunkUint4 = kNumChunkBytes / sizeof(uint4); + constexpr uint32_t kNumUint4PerLane = kNumChunkUint4 / 32; + DG_STATIC_ASSERT(kHidden % kNumChunks == 0, "Hidden must be divisible by number of chunks"); + DG_STATIC_ASSERT(kNumChunkSlots * kNumCombineWarps * kNumHiddenBytes / kNumChunks <= SMEM_BEFORE_BARRIER_SIZE, "Hidden is too large"); + DG_STATIC_ASSERT(kNumChunkBytes % 16 == 0, "Combine chunk must be TMA-aligned (16 bytes)"); + DG_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, "Combine chunk must be divisible by 16 bytes"); + DG_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, "Combine chunk must be a multiple of 32 16-byte elements"); + DG_STATIC_ASSERT(kNumTopk <= 32, "Top-k must fit in a single warp"); + + DG_TRAP_ONLY_DEVICE_ASSERT(kNumChunkSlots * kNumCombineWarps * kNumChunkBytes <= static_cast( + reinterpret_cast(barrier_start_ptr) - smem_buffer)); + + const auto combine_load_buffer = utils::PatternVisitor([&](const uint32_t& i) { + return math::advance_ptr(smem_buffer, (epilogue_warp_idx + i * kNumCombineWarps) * kNumChunkBytes); + }); + const auto combine_store_buffer = math::advance_ptr( + smem_buffer, (epilogue_warp_idx + kNumCombineWarps * 2) * kNumChunkBytes); + + auto combine_load_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return combine_barriers[i + epilogue_warp_idx * 2]; + }); + + uint32_t combine_phase = 0; + uint32_t load_stage_idx = 0; + for (uint32_t token_idx = sm_idx * kNumCombineWarps + epilogue_warp_idx; + token_idx < num_tokens; + token_idx += kNumSMs * kNumCombineWarps) { + const int stored_topk_slot_idx = lane_idx < kNumTopk ? + static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; + const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); + + for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { + const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; + + uint32_t mask = total_mask; + const auto move_mask_and_load = [&](const uint32_t& i) { + if (mask) { + const uint32_t slot_idx = __ffs(mask) - 1; + mask ^= 1 << slot_idx; + if (cute::elect_one_sync()) { + const auto src_ptr = math::advance_ptr( + combine_token_buffer.get_rank_buffer(slot_idx) + .get_data_buffer(token_idx).get_base_ptr(), + chunk_byte_offset); + ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], kNumChunkBytes); + ptx::mbarrier_arrive_and_set_tx(combine_load_barriers[i], kNumChunkBytes); + } + __syncwarp(); + return true; + } + return false; + }; + + bool do_reduce = move_mask_and_load(load_stage_idx); + + float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {}; + while (do_reduce) { + do_reduce = move_mask_and_load(load_stage_idx ^ 1); + combine_load_barriers[load_stage_idx]->wait(combine_phase); + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; + const auto bf16_values = reinterpret_cast(&uint4_values); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); + } + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } + + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); + + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); + } + __syncwarp(); + + if (cute::elect_one_sync()) { + cute::tma_store_fence(); + ptx::tma_store_1d( + math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + chunk_byte_offset), + combine_store_buffer, kNumChunkBytes); + cute::tma_store_arrive(); + } + __syncwarp(); + } + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_90"); +#endif +} + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh b/deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh new file mode 100644 index 0000000000..1836fac953 --- /dev/null +++ b/deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh @@ -0,0 +1,179 @@ +#pragma once + +#include + +#include +#include +#include + +// ============================================================================ +// SM90 (Hopper) MegaMoE workspace layout +// ---------------------------------------------------------------------------- +// Ported from upstream PR https://github.com/sgl-project/DeepGEMM/pull/36 +// (`sched::MegaMoEScheduler` / `layout::Workspace` on the old `dev` branch). +// +// This is a *pool*-based workspace: all local experts share one contiguous +// token pool sized by `layout::get_num_max_pool_tokens` (already shared with +// the SM100 path), with no ring-buffer bookkeeping. This is structurally +// different from `nv_dev`'s current `layout::Workspace` (which is a +// *ring*-buffer workspace built for the SM100 persistent-grid, cluster-pair +// interleaved L1/L2 scheduler). Since the two workspace layouts are not +// binary-compatible (different byte layout, different constructor, different +// arrival-tracking scheme), this SM90 pool workspace is defined under its own +// name (`MegaMoESM90Workspace`) to avoid colliding with the SM100 +// `layout::Workspace` struct, rather than replacing/aliasing it. +// +// `layout::Data`, `layout::Buffer`, and `layout::TokenSrcMetadata` (all in +// `deep_gemm/layout/mega_moe.cuh`) are structurally unchanged between the two +// designs, so they are reused directly by the SM90 kernel instead of being +// redefined here. +// ============================================================================ + +namespace deep_gemm::layout { + +// SF pool capacity: all experts share a contiguous SF region, sized by pool blocks x SF_BLOCK_M. +// NOTES: this is the SM90 (pool-based) counterpart of `get_num_sf_ring_tokens` (ring-based, used +// by SM100). `get_num_max_pool_tokens` itself is already shared between both paths. +template +CUTLASS_HOST_DEVICE constexpr T get_num_padded_sf_pool_tokens(T num_max_pool_tokens, T block_m) { + return (num_max_pool_tokens / block_m) * math::constexpr_align(block_m, static_cast(128)); +} + +struct MegaMoESM90Workspace { + void* base; + uint32_t num_ranks, num_experts; + uint32_t num_experts_per_rank; + uint32_t num_max_tokens_per_rank; + uint32_t num_max_recv_tokens_per_expert; + + // Pool capacity: all local experts share a contiguous token pool + uint32_t num_max_pool_tokens; + uint32_t num_max_pool_blocks; + + // For both grid barrier and NVLink barrier + static constexpr uint64_t kNumBarrierSignalBytes = 32; + + MegaMoESM90Workspace() = default; + + CUTLASS_HOST_DEVICE + MegaMoESM90Workspace(void* base, + const uint32_t& num_ranks, + const uint32_t& num_experts, + const uint32_t& num_max_tokens_per_rank, + const uint32_t& num_topk): + base(base), + num_ranks(num_ranks), num_experts(num_experts), + num_max_tokens_per_rank(num_max_tokens_per_rank) { + num_experts_per_rank = num_experts / num_ranks; + num_max_recv_tokens_per_expert = num_ranks * num_max_tokens_per_rank; + num_max_pool_tokens = get_num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); + num_max_pool_blocks = num_max_pool_tokens / kMinCandidateBlockM; + } + + CUTLASS_HOST_DEVICE + uint64_t get_num_bytes() const { + uint64_t num_bytes = 0; + + // Barrier + num_bytes += kNumBarrierSignalBytes; + + // Expert send/recv count + num_bytes += num_experts * sizeof(uint64_t) * 2; + + // Expert recv count sum + num_bytes += num_experts_per_rank * sizeof(uint64_t); + + // L1 arrival count (padded to even entry count for `uint64_t` alignment of L2 mask) + num_bytes += math::align(num_max_pool_blocks, 2u) * sizeof(uint32_t); + + // L2 block arrival mask + num_bytes += num_max_pool_blocks * sizeof(uint64_t); + + // Dispatch pulling source token-topk + num_bytes += num_experts_per_rank * num_ranks * num_max_recv_tokens_per_expert * sizeof(int); + + // Combine push source indices + num_bytes += num_max_pool_tokens * sizeof(TokenSrcMetadata); + + // Align to TMA descriptor requirements + num_bytes = math::align(num_bytes, 16); + return num_bytes; + } + + CUTLASS_HOST_DEVICE + void* get_end_ptr() const { + return math::advance_ptr(base, get_num_bytes()); + } + + // Grid sync counters: `kNumBarrierSignalBytes` layout + // [ 0..15]: 4 x `uint32_t` grid sync counters + // [16..20]: `uint32_t` NVLink barrier counter + // [20..27]: 2 x `int` NVLink barrier signals (phase 0 and 1) + static constexpr uint32_t kNumMaxGridSyncCounters = 4; + + template + CUTLASS_DEVICE + uint32_t* get_grid_sync_count_ptr() const { + DG_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds"); + return static_cast(base) + kIndex; + } + + CUTLASS_DEVICE + uint32_t* get_nvl_barrier_counter_ptr() const { + return static_cast(base) + kNumMaxGridSyncCounters; + } + + CUTLASS_DEVICE + int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const { + // NOTES: the signal is signed, as we may minus + return math::advance_ptr(base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int)); + } + + CUTLASS_DEVICE + uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const { + return math::advance_ptr(base, kNumBarrierSignalBytes) + expert_idx; + } + + CUTLASS_DEVICE + uint64_t* get_expert_recv_count_ptr( + const uint32_t& rank_idx = 0, const uint32_t& expert_idx = 0) const { + return get_expert_send_count_ptr(num_experts) + rank_idx * num_experts_per_rank + expert_idx; + } + + CUTLASS_DEVICE + uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx = 0) const { + return get_expert_send_count_ptr(num_experts * 2) + expert_idx; + } + + CUTLASS_DEVICE + uint32_t* get_l1_arrival_count_ptr(const uint32_t& pool_block_idx = 0) const { + const auto base = get_expert_recv_count_sum_ptr(num_experts_per_rank); + return reinterpret_cast(base) + pool_block_idx; + } + + CUTLASS_DEVICE + uint64_t* get_l2_arrival_mask_ptr(const uint32_t& pool_block_idx = 0) const { + // Pad L1 entry count to even so that the `l2_arrival_mask` is 8-byte aligned + const auto base = get_l1_arrival_count_ptr(math::align(num_max_pool_blocks, 2u)); + return reinterpret_cast(base) + pool_block_idx; + } + + // For dispatch pulling + CUTLASS_DEVICE + uint32_t* get_src_token_topk_idx_ptr( + const uint32_t& expert_idx = 0, const uint32_t& rank_idx = 0, const uint32_t& token_idx = 0) const { + const auto base = get_l2_arrival_mask_ptr(num_max_pool_blocks); + return reinterpret_cast(base) + + expert_idx * (num_ranks * num_max_recv_tokens_per_expert) + + rank_idx * num_max_recv_tokens_per_expert + token_idx; + } + + // For combine usages + CUTLASS_DEVICE + TokenSrcMetadata* get_token_src_metadata_ptr(const uint32_t& pool_token_idx = 0) const { + const auto base = reinterpret_cast(get_src_token_topk_idx_ptr(num_experts_per_rank)); + return base + pool_token_idx; + } +}; + +} // namespace deep_gemm::layout diff --git a/deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh b/deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh new file mode 100644 index 0000000000..303ee31936 --- /dev/null +++ b/deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh @@ -0,0 +1,235 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +// ============================================================================ +// SM90 (Hopper) MegaMoE scheduler +// ---------------------------------------------------------------------------- +// Ported from upstream PR https://github.com/sgl-project/DeepGEMM/pull/36 +// (`sched::MegaMoEScheduler` / `sched::BlockPhase` on the old `dev` branch). +// +// This scheduler walks a *wave* of local experts at a time (grouped by +// `kNumExpertsPerWave`), alternating between all L1 (Linear1) blocks and all +// L2 (Linear2) blocks of the wave, over a persistent grid of `kNumSMs` +// single-CTA "tiles" (no 2-CTA cluster pairing, unlike SM100's +// `sched::MegaMoEScheduler`). It is intentionally simpler than SM100's +// ring-buffered, cluster-pair-interleaved scheduler in `mega_moe.cuh`, which +// is why it is kept as a separate class (`MegaMoESM90Scheduler`) and enum +// (`MegaMoESM90BlockPhase`) rather than reusing/extending the SM100 one: the +// two have incompatible state machines and workspace layouts +// (`layout::MegaMoESM90Workspace` vs. `layout::Workspace`). +// ============================================================================ + +namespace deep_gemm::sched { + +// Computation phase for the current block (SM90 MegaMoE only; do not confuse with the +// unrelated SM100 `sched::BlockPhase`, which additionally has `SharedLinear1`/`SharedLinear2`). +enum class MegaMoESM90BlockPhase { + None = 0, + Linear1 = 1, + Linear2 = 2 +}; + +template +struct MegaMoESM90Scheduler { + DG_STATIC_ASSERT(L1_SHAPE_N % BLOCK_N == 0, "Invalid shape"); + DG_STATIC_ASSERT(L2_SHAPE_N % BLOCK_N == 0, "Invalid shape"); + DG_STATIC_ASSERT(L1_SHAPE_K % BLOCK_K == 0, "Invalid shape"); + DG_STATIC_ASSERT(L2_SHAPE_K % BLOCK_K == 0, "Invalid shape"); + DG_STATIC_ASSERT(kNumExpertsPerRank % kNumExpertsPerWave == 0, "Invalid wave config"); + + // Arrival counts + const layout::MegaMoESM90Workspace& workspace; + + // Scheduler state + MegaMoESM90BlockPhase next_phase = MegaMoESM90BlockPhase::Linear1; + + // Current expert and block indices + uint32_t current_local_expert_idx = 0; + uint32_t current_num_tokens = 0; + uint32_t current_pool_block_offset = 0; + uint32_t block_idx = 0; + uint32_t m_block_idx = 0; + uint32_t n_block_idx = 0; + + // Pre-cached per-expert token counts (filled during `for_each_block` init) + // Layout: `stored_num_tokens_per_expert[i]` holds expert (i * 32 + lane_idx)'s count + uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {}; + + CUTLASS_DEVICE explicit MegaMoESM90Scheduler(const layout::MegaMoESM90Workspace& workspace): workspace(workspace) { + block_idx = blockIdx.x; + } + + CUTLASS_DEVICE uint32_t get_wave_expert_end_idx() const { + return math::align(current_local_expert_idx + 1, kNumExpertsPerWave); + } + + CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const { + uint32_t valid_value; + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ? + stored_num_tokens_per_expert[i] : valid_value; + } + return ptx::exchange(valid_value, expert_idx % 32); + } + + // Get pool block offset for a given expert index from a per-lane token count array + CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) { + uint32_t num_blocks = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + if (i * 32 + ptx::get_lane_idx() < expert_idx) + num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M); + } + return __reduce_add_sync(0xffffffff, num_blocks); + } + + CUTLASS_DEVICE void advance_expert_idx() { + current_pool_block_offset += get_current_num_m_blocks(); + current_local_expert_idx += 1; + current_num_tokens = get_num_tokens(current_local_expert_idx); + } + + CUTLASS_DEVICE void set_expert_idx(const uint32_t& expert_idx) { + current_local_expert_idx = expert_idx; + current_num_tokens = get_num_tokens(expert_idx); + current_pool_block_offset = get_pool_block_offset(expert_idx); + } + + CUTLASS_DEVICE uint32_t get_current_pool_block_offset() const { + return current_pool_block_offset; + } + + CUTLASS_DEVICE uint32_t get_current_num_m_blocks() const { + return math::ceil_div(current_num_tokens, BLOCK_M); + } + + template + CUTLASS_DEVICE uint32_t get_valid_m() const { + const auto m = cute::min(current_num_tokens - m_block_idx * BLOCK_M, BLOCK_M); + return kDoUMMAAligned ? math::align(m, 16u) : m; + } + + CUTLASS_DEVICE bool fetch_next_l1_block() { + const auto wave_end_expert_idx = get_wave_expert_end_idx(); + while (current_local_expert_idx < wave_end_expert_idx) { + const auto num_m_blocks = get_current_num_m_blocks(); + m_block_idx = block_idx / kNumL1BlockNs; + if (m_block_idx < num_m_blocks) + return true; + + // Current expert is fully assigned, move to the next + block_idx -= num_m_blocks * kNumL1BlockNs; + advance_expert_idx(); + } + return false; + } + + CUTLASS_DEVICE bool fetch_next_l2_block() { + const auto wave_end_expert_idx = get_wave_expert_end_idx(); + while (current_local_expert_idx < wave_end_expert_idx) { + const auto num_m_blocks = get_current_num_m_blocks(); + if (block_idx < num_m_blocks * kNumL2BlockNs) { + m_block_idx = block_idx / kNumL2BlockNs; + return true; + } + + // Current expert is fully assigned, move to the next + block_idx -= num_m_blocks * kNumL2BlockNs; + advance_expert_idx(); + } + return false; + } + + // Core state machine: assigns the next block + CUTLASS_DEVICE cute::tuple get_next_block() { + while (true) { + if (current_local_expert_idx >= kNumExpertsPerRank) + break; + + if (next_phase == MegaMoESM90BlockPhase::Linear1) { + if (fetch_next_l1_block()) { + // Found a new L1 block + n_block_idx = block_idx - m_block_idx * kNumL1BlockNs; + // Jump to next block + block_idx += kNumSMs; + return {MegaMoESM90BlockPhase::Linear1, current_local_expert_idx, m_block_idx, n_block_idx}; + } else { + // L1 for the current wave is complete, transition to L2 + next_phase = MegaMoESM90BlockPhase::Linear2; + set_expert_idx(math::align(current_local_expert_idx - 1, kNumExpertsPerWave)); + } + } else { + if (fetch_next_l2_block()) { + // Found a new L2 block + n_block_idx = block_idx - m_block_idx * kNumL2BlockNs; + // Jump to next block + block_idx += kNumSMs; + return {MegaMoESM90BlockPhase::Linear2, current_local_expert_idx, m_block_idx, n_block_idx}; + } else { + // Move to L1 of the next wave + next_phase = MegaMoESM90BlockPhase::Linear1; + } + } + } + + // All waves and experts are fully processed + return {MegaMoESM90BlockPhase::None, 0, 0, 0}; + } + + CUTLASS_DEVICE void fetch_expert_recv_count() { + // NOTES: each lane caches experts at indices (i * 32 + lane_idx) + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + const auto expert_idx = i * 32 + ptx::get_lane_idx(); + uint64_t value = 0; + if (expert_idx < kNumExpertsPerRank) { + do { + value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx)); + } while (static_cast(value >> 32) != kNumSMs * kNumRanks); + } + stored_num_tokens_per_expert[i] = static_cast(value); + } + __syncwarp(); + } + + template + CUTLASS_DEVICE void for_each_block(Func&& func) { + // Wait for all expert counters to be finalized + fetch_expert_recv_count(); + + // Initialize current expert with 0 + set_expert_idx(0); + + // Iterate over all blocks + // TODO: add swizzle within expert waves for better L2 cache utilization + while (true) { + CUTE_TIE_DECL(get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx); + if (block_phase == MegaMoESM90BlockPhase::None) + break; + + func(block_phase, current_local_expert_idx, + block_phase == MegaMoESM90BlockPhase::Linear2 ? kNumL2BlockKs : kNumL1BlockKs, + m_block_idx, n_block_idx); + } + } +}; + +} // namespace deep_gemm::sched diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index a3ffba5c32..c75a053941 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -98,6 +98,74 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, ) +class Sm90SymmBuffer: + """Symmetric buffer for the SM90 (Hopper) FP8 MegaMoE path (`fp8_mega_moe_sm90`). + + Unlike `SymmBuffer` (SM100), this always uses FP8xFP8 with float scale factors, a pool-based + (non-ring) workspace, and does not support shared experts or the `situ` activation. + """ + + def __init__(self, group: dist.ProcessGroup, + num_experts: int, + num_max_tokens_per_rank: int, num_topk: int, + hidden: int, intermediate_hidden: int, + activation: str = 'swiglu'): + assert activation == 'swiglu', f'SM90 MegaMoE only supports `swiglu`, got activation={activation!r}' + self.group = group + self.num_experts = num_experts + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_topk = num_topk + self.hidden = hidden + self.intermediate_hidden = intermediate_hidden + + # Allocate a symmetric buffer + num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe( + group.size(), num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + True, activation + ) + allocator = torch if group.size() == 1 else symm_mem + self.buffer = allocator.empty(num_bytes, dtype=torch.int8, device='cuda') + self.handle = ( + types.SimpleNamespace(buffer_ptrs=[self.buffer.data_ptr()]) + if group.size() == 1 + else symm_mem.rendezvous(self.buffer, group=group) + ) + self.buffer.zero_() + self.group.barrier() + torch.cuda.synchronize() + + # Create input buffer views + (self.x, self.x_sf, + self.topk_idx, self.topk_weights, + self.l1_acts, self.l1_acts_sf, + self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer) + + def destroy(self): + self.handle = None + self.buffer = None + self.group = None + self.x = None + self.x_sf = None + + +def get_symm_buffer_for_sm90_mega_moe(group: dist.ProcessGroup, + num_experts: int, + num_max_tokens_per_rank: int, num_topk: int, + hidden: int, intermediate_hidden: int, + activation: str = 'swiglu') -> Sm90SymmBuffer: + # Align token count (SM90 MegaMoE shares the same token alignment requirement as SM100) + num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) + + return Sm90SymmBuffer( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + activation=activation + ) + + def _interleave_weights(t: torch.Tensor, gran: int = 8) -> torch.Tensor: # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up] # Unsqueeze for 2D @@ -157,6 +225,21 @@ def transform_weights_for_mega_moe( +def transform_weights_for_mega_moe_sm90( + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor] +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + """Weight preprocessing for `fp8_mega_moe_sm90`. + + Unlike `transform_weights_for_mega_moe` (SM100), the SM90 kernel consumes weight SF as plain + float in block-(128, 128) layout (no UE8M0 packing, no UTCCP transpose), so only the L1 + gate/up interleave is needed; L2 weights and both SF tensors pass through unchanged. + """ + l1_fp8, l1_sf = l1_weights + l1_transformed = (_interleave_weights(l1_fp8), l1_sf) + return l1_transformed, l2_weights + + def fp8_fp4_mega_moe(y: torch.Tensor, l1_weights: Tuple[torch.Tensor, torch.Tensor], l2_weights: Tuple[torch.Tensor, torch.Tensor], @@ -268,3 +351,31 @@ def bf16_mega_moe(y: torch.Tensor, activation, activation_clamp, fast_math ) + + +def fp8_mega_moe_sm90(y: torch.Tensor, + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], + sym_buffer: Sm90SymmBuffer, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + recipe: Tuple[int, int, int] = (128, 128, 128), + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True): + """SM90 (Hopper) FP8 MegaMoE fused GEMM, ported from upstream PR #36. + + Unlike `fp8_fp4_mega_moe` (SM100), both L1/L2 weights are FP8 e4m3 with block-(128, 128) + float scale factors (not FP4/UE8M0), and there is no shared-expert or `situ` support. + """ + _C.fp8_mega_moe_sm90( + y, + l1_weights, l2_weights, + cumulative_local_expert_recv_stats, + sym_buffer.buffer, + sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), + sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_experts, sym_buffer.num_topk, + recipe, + activation, activation_clamp, + fast_math + ) diff --git a/tests/test_mega_moe_hopper.py b/tests/test_mega_moe_hopper.py new file mode 100644 index 0000000000..cb6a142088 --- /dev/null +++ b/tests/test_mega_moe_hopper.py @@ -0,0 +1,172 @@ +"""SM90 (Hopper) MegaMoE correctness test. + +Ported from upstream PR https://github.com/sgl-project/DeepGEMM/pull/36 +(`tests/test_mega_moe_hopper.py` on the old `dev` branch) and adapted to follow the structure and +test harness (`init_dist`/`dist_print`, multiprocessing-spawn driver) of `tests/test_mega_moe.py` +(the SM100 FP4 MegaMoE test in this repo), and to call the SM90 API surface added in +`deep_gemm.mega` (`Sm90SymmBuffer`, `get_symm_buffer_for_sm90_mega_moe`, +`transform_weights_for_mega_moe_sm90`, `fp8_mega_moe_sm90`). + +Unlike the SM100 path, `fp8_mega_moe_sm90`: + * is FP8xFP8-only (no FP4 weights, no `bf16xbf16` variant here); + * uses plain float scale factors at block-(128, 128)/(token, 128) granularity (no UE8M0 packing); + * has no shared-expert support; + * runs on a simpler pool-based scheduler (no ring buffer). + +This test is skipped cleanly (prints a message and returns) on any non-Hopper (SM90) GPU. +""" + +import argparse +import random +from typing import Tuple + +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.utils import per_block_cast_to_fp8, per_token_cast_to_fp8 +from deep_gemm.utils.dist import dist_print, init_dist +from deep_gemm.testing import calc_diff, get_arch_major + + +def _quantize_weights_block_128_128(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a `(num_groups, n, k)` BF16 weight tensor into FP8 e4m3 with block-(128, 128) + float scale factors, matching the layout `fp8_mega_moe_sm90` expects for L1/L2 weights.""" + num_groups, n, k = bf16_weights.shape + w = torch.empty((num_groups, n, k), device='cuda', dtype=torch.float8_e4m3fn) + w_sf = torch.empty((num_groups, n // 128, k // 128), device='cuda', dtype=torch.float) + for i in range(num_groups): + w[i], w_sf[i] = per_block_cast_to_fp8(bf16_weights[i], use_ue8m0=False, gran_k=128) + return w, w_sf + + +def _reference_mega_moe( + x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, + l1_weights: torch.Tensor, l2_weights: torch.Tensor, + num_experts_per_rank: int, activation_clamp: float, +) -> torch.Tensor: + """Naive reference: dense routed SwiGLU MLP in FP32, evaluated per (token, topk) pair. + NOTES: this assumes a single rank (EP=1), matching how this test invokes the kernel.""" + num_tokens, hidden = x.shape + num_topk = topk_idx.shape[1] + y = torch.zeros((num_tokens, hidden), dtype=torch.float32, device='cuda') + x_f32 = x.float() + l1_f32 = l1_weights.float() + l2_f32 = l2_weights.float() + for t in range(num_tokens): + for k in range(num_topk): + e = topk_idx[t, k].item() + if e < 0 or e >= num_experts_per_rank: + continue + w = topk_weights[t, k].item() + gate_up = x_f32[t] @ l1_f32[e].T + intermediate_hidden = gate_up.shape[0] // 2 + gate, up = gate_up[:intermediate_hidden], gate_up[intermediate_hidden:] + gate = gate.clamp(max=activation_clamp) + up = up.clamp(min=-activation_clamp, max=activation_clamp) + act = torch.nn.functional.silu(gate) * up + out = act @ l2_f32[e].T + y[t] += w * out + return y.to(torch.bfloat16) + + +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(rank_idx) + random.seed(rank_idx) + + if not torch.cuda.is_available() or get_arch_major() != 9: + dist_print(f'Skipping SM90 MegaMoE test: requires a Hopper (SM90) GPU, ' + f'got arch major {get_arch_major() if torch.cuda.is_available() else "N/A"}', + once_in_node=True) + dist.barrier() + dist.destroy_process_group() + return + + num_max_tokens_per_rank = args.num_max_tokens_per_rank + num_tokens = args.num_tokens if args.num_tokens > 0 else num_max_tokens_per_rank + num_experts, num_topk = args.num_experts, args.num_topk + num_experts_per_rank = num_experts // num_ranks + hidden, intermediate_hidden = args.hidden, args.intermediate_hidden + activation_clamp = args.activation_clamp + assert num_tokens <= num_max_tokens_per_rank + + # Allocate the SM90 symmetric buffer + buffer = deep_gemm.get_symm_buffer_for_sm90_mega_moe( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + ) + + # Random inputs + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1_weights_bf16 = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), dtype=torch.bfloat16, device='cuda') + l2_weights_bf16 = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), dtype=torch.bfloat16, device='cuda') + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + # Keep only experts routed to this rank, and remap to local expert indices for the reference + local_mask = (topk_idx // num_experts_per_rank) == rank_idx + local_topk_idx = torch.where(local_mask, topk_idx % num_experts_per_rank, -1) + + # Quantize activations (per-token, per-128 K, float SF) and weights (block-128x128, float SF) + x_fp8, x_sf = per_token_cast_to_fp8(x, use_ue8m0=False, gran_k=128) + l1_weights_fp8, l1_weights_sf = _quantize_weights_block_128_128(l1_weights_bf16) + l2_weights_fp8, l2_weights_sf = _quantize_weights_block_128_128(l2_weights_bf16) + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe_sm90( + (l1_weights_fp8, l1_weights_sf), (l2_weights_fp8, l2_weights_sf)) + + # Copy inputs into the symmetric buffer + buffer.x[:num_tokens].copy_(x_fp8) + buffer.x_sf[:num_tokens].copy_(x_sf) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + if num_tokens < num_max_tokens_per_rank: + buffer.x[num_tokens:].zero_() + buffer.topk_idx[num_tokens:].fill_(-1) + buffer.topk_weights[num_tokens:].zero_() + + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_mega_moe_sm90( + y, transformed_l1, transformed_l2, buffer, + activation_clamp=activation_clamp, + ) + + if num_ranks == 1: + y_ref = _reference_mega_moe( + x, local_topk_idx, topk_weights, + l1_weights_bf16, l2_weights_bf16, + num_experts_per_rank, activation_clamp) + diff = calc_diff(y.float(), y_ref.float()) + dist_print(f' > EP {rank_idx:2}/{num_ranks} | correctness diff: {diff:.6f}', once_in_node=True) + assert diff < 0.05, f'MegaMoE SM90 correctness check failed: diff={diff}' + else: + # Cross-rank combine correctness needs a distributed reference; only run the smoke check + # (kernel executes without error / NaNs) for multi-rank configurations. + assert torch.isfinite(y.float()).all(), 'MegaMoE SM90 produced non-finite outputs' + dist_print(f' > EP {rank_idx:2}/{num_ranks} | smoke test passed (finite outputs)', once_in_node=True) + + dist.barrier() + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='SM90 (Hopper) MegaMoE correctness test') + parser.add_argument('--num-processes', type=int, default=1, help='Number of processes to spawn (default: 1)') + parser.add_argument('--num-max-tokens-per-rank', type=int, default=512, help='Number of maximum tokens per rank') + parser.add_argument('--num-tokens', type=int, default=0, help='Number of tokens per rank (follow max if 0)') + parser.add_argument('--hidden', type=int, default=2048, help='Hidden size') + parser.add_argument('--intermediate-hidden', type=int, default=1024, help='Intermediate hidden size') + parser.add_argument('--activation-clamp', type=float, default=10, help='Clamp value for activation') + parser.add_argument('--num-experts', type=int, default=8, help='Number of experts') + parser.add_argument('--num-topk', type=int, default=2, help='Number of expert selections') + args = parser.parse_args() + + if args.num_processes == 1: + test(0, 1, args) + else: + torch.multiprocessing.spawn( + test, args=(args.num_processes, args), nprocs=args.num_processes + ) From ae4228ad7acb34f39f3daa7a0633333d044b9921 Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Tue, 18 Aug 2026 11:13:46 +0800 Subject: [PATCH 2/3] Fix missing 'template' disambiguator in grid_sync dependent-type call workspace.get_grid_sync_count_ptr() failed NVCC compilation because workspace is a dependent-type parameter (WorkspaceT&); without the 'template' keyword the '<' is parsed as less-than. --- deep_gemm/include/deep_gemm/comm/barrier.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deep_gemm/include/deep_gemm/comm/barrier.cuh b/deep_gemm/include/deep_gemm/comm/barrier.cuh index c1672cda66..39fe50754b 100644 --- a/deep_gemm/include/deep_gemm/comm/barrier.cuh +++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh @@ -32,7 +32,7 @@ CUTLASS_DEVICE void grid_sync(const WorkspaceT& workspace, static constexpr uint32_t kFinishSumTag = 0x80000000u; sync_scope(); if (thread_idx == 0) { - const auto count_ptr = workspace.get_grid_sync_count_ptr(); + const auto count_ptr = workspace.template get_grid_sync_count_ptr(); const auto old_value = ptx::atomic_add_rel( count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1); uint32_t new_value; From 11885a0e83738b68f226a80bffe1efdc83dc5364 Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Mon, 31 Aug 2026 08:21:58 +0000 Subject: [PATCH 3/3] add hidden size check to 256 Signed-off-by: rongfu.leng --- csrc/apis/mega.hpp | 23 ++++++++++++++++++++--- csrc/utils/layout.hpp | 7 ++++--- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index bd4c8f36c5..00d8df0fdb 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -539,6 +539,11 @@ get_symm_buffer_size_for_sm90_mega_moe( DG_HOST_ASSERT(num_experts % num_ranks == 0); DG_HOST_ASSERT(use_fp8_dispatch); DG_HOST_ASSERT(activation == "swiglu"); + // `get_mega_moe_config_sm90` may pick `num_dispatch_threads == 64`, and the SM90 + // `nvlink_barrier` only has one signaling thread per rank (`thread_idx < kNumRanks`), + // so more than 64 ranks either fails the kernel's `kNumRanks <= kNumThreads` + // static_assert at JIT time or, if that were relaxed, would hang on missed signals. + DG_HOST_ASSERT(num_ranks <= 64); const auto workspace = layout::MegaMoESM90Workspace(nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk); @@ -565,8 +570,13 @@ get_symm_buffer_size_for_sm90_mega_moe( input_topk_idx_buffer.get_end_ptr()); const auto num_max_pool_tokens = static_cast(workspace.num_max_pool_tokens); + // Unlike SM100 (which can select any of `layout::kCandidateBlockM`), SM90's + // `get_block_config_for_mega_moe_sm90` only ever picks block_m in {64, 128}. Sizing the + // SF pool against the full shared candidate set (which includes block_m=8) would + // over-allocate the SF pool by ~8x. + constexpr int kSm90CandidateBlockM[] = {64, 128}; int num_max_padded_sf_pool_tokens = 0; - for (int block_m: layout::kCandidateBlockM) { + for (int block_m: kSm90CandidateBlockM) { num_max_padded_sf_pool_tokens = std::max( num_max_padded_sf_pool_tokens, layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m) @@ -594,7 +604,9 @@ get_symm_buffer_size_for_sm90_mega_moe( bf16_token_layout, num_topk, num_max_tokens_per_rank, l2_sf_buffer.get_end_ptr()); - DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); + // Kept in sync with the stricter check in `fp8_mega_moe_sm90` (see comment there): + // hidden must be a multiple of 256 for the scheduler's BLOCK_N=256 case to compile. + DG_HOST_ASSERT(hidden % 256 == 0 and intermediate_hidden % 128 == 0); auto slice_input_buffers = [=](const torch::Tensor& buffer) { auto x = torch::from_blob( @@ -676,7 +688,12 @@ static void fp8_mega_moe_sm90( DG_HOST_ASSERT(hidden == hidden_); DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); - DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); + // `get_mega_moe_config_sm90` may pick BLOCK_N=256 for either the L1 (2 * intermediate_hidden) + // or L2 (hidden) GEMM depending on the runtime token distribution, and the scheduler + // requires L1_SHAPE_N/L2_SHAPE_N to be an exact multiple of BLOCK_N or the JIT compile + // fails. Require hidden % 256 == 0 so this holds regardless of which BLOCK_N is chosen; + // intermediate_hidden % 128 == 0 already implies (2 * intermediate_hidden) % 256 == 0. + DG_HOST_ASSERT(hidden % 256 == 0 and intermediate_hidden % 128 == 0); DG_HOST_ASSERT(intermediate_hidden / 64 <= 64); constexpr int kGranMN = 128, kGranK = 128; diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 07a81c4e37..bd0d0eb9c4 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -116,12 +116,13 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf, DG_HOST_ASSERT(sf.stride(-1) == get_tma_aligned_size(mn, sf.element_size())); } - // SM90 SFB must be contiguous, or contiguous after transposing the last two dimensions + // SM90 SFB must be K-major contiguous: the kernel indexes it with a fixed + // `expert * per_expert + n_block * k_blocks + k_block` formula that ignores tensor + // stride, so an MN-major (transposed) input would silently read the wrong scale. if (sm90_sfb_check) { if (num_groups.has_value()) DG_HOST_ASSERT(sf.stride(-3) == sf.size(-2) * sf.size(-1)); - DG_HOST_ASSERT((sf.stride(-1) == 1 and sf.stride(-2) == sf.size(-1)) or - (sf.stride(-1) == sf.size(-2) and sf.stride(-2) == 1)); + DG_HOST_ASSERT(sf.stride(-1) == 1 and sf.stride(-2) == sf.size(-1)); } return sf; }