From ec72c60ca81f4ef0ea8531926e7a61bcd2a77ae5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E6=A5=BC?= Date: Thu, 20 Aug 2026 20:23:42 +0800 Subject: [PATCH 1/3] feat(sm90): add FP8 x MXFP4 MegaMoE support --- csrc/apis/sm90_mega.hpp | 435 ++ csrc/jit/handle.hpp | 64 +- csrc/jit/kernel_runtime.hpp | 34 +- csrc/jit_kernels/heuristics/sm90_mega_moe.hpp | 138 + csrc/jit_kernels/impls/runtime_utils.hpp | 9 +- csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp | 464 ++ csrc/python_api.cpp | 2 + deep_gemm/__init__.py | 5 + .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 3791 +++++++++++++++++ .../include/deep_gemm/layout/mega_moe.cuh | 110 +- .../deep_gemm/layout/sm90_mega_moe.cuh | 20 + deep_gemm/include/deep_gemm/ptx/ld_st.cuh | 14 + .../include/deep_gemm/scheduler/mega_moe.cuh | 222 +- deep_gemm/mega/__init__.py | 287 ++ deep_gemm/mega/mxfp4.py | 471 ++ tests/bench_mega_moe_sm90.py | 597 +++ tests/test_mega_moe_mxfp4_preprocess.py | 524 +++ tests/test_mega_moe_sm90.py | 1062 +++++ 18 files changed, 8186 insertions(+), 63 deletions(-) create mode 100644 csrc/apis/sm90_mega.hpp 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/mega/mxfp4.py create mode 100644 tests/bench_mega_moe_sm90.py create mode 100644 tests/test_mega_moe_mxfp4_preprocess.py create mode 100644 tests/test_mega_moe_sm90.py diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp new file mode 100644 index 0000000000..b70bf06423 --- /dev/null +++ b/csrc/apis/sm90_mega.hpp @@ -0,0 +1,435 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#if DG_TENSORMAP_COMPATIBLE +#include "../jit/compiler.hpp" +#endif +#include "../jit/device_runtime.hpp" +#include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp" +#include "../utils/layout.hpp" +#include "../utils/system.hpp" + +namespace deep_gemm::mega { + +static constexpr int kSM90MegaMoEBlockM = MegaMoESM90Config::block_m; +static constexpr int kSM90MegaMoETokenAlignment = 128; +static constexpr int kSM90MegaMoEBlockN = MegaMoESM90Config::block_n; +static constexpr int kSM90MegaMoEWorkerCTAsPerSM = 2; +static constexpr int kSM90MegaMoECTAsPerTask = 1; +// The dispatch-side NVLink barrier assigns one signaling thread per rank and +// the compact MXFP4 frontend has 64 dispatch threads. +static constexpr int kSM90MegaMoEMaxRanks = 64; + +using SM90MegaMoEBufferViews = std::tuple< + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>; +using SM90MegaMoEBufferSlicer = std::function; + +static bool is_valid_hidden_for_sm90_mega_moe(const int hidden) { + // The combine epilogue uses four chunks above 8192 elements, with each + // chunk vectorized in groups of 256 BF16 values. + return hidden % 512 == 0 and (hidden <= 8192 or hidden % 1024 == 0); +} + +static int get_token_alignment_for_sm90_mega_moe() { + return kSM90MegaMoETokenAlignment; +} + +static int get_num_ring_tokens_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 auto num_experts_per_rank = num_experts / num_ranks; + // Only physical activation/SF/top-k-weight slots are ring-sized; Workspace + // retains the full logical-pool source metadata. + const auto num_worker_ctas = + kSM90MegaMoEWorkerCTAsPerSM * device_runtime->get_num_sms(); + const auto num_active_topk = std::min(num_topk, num_experts_per_rank); + const auto num_max_routed_tokens = + num_max_tokens_per_rank * num_ranks * num_active_topk; + const auto num_pool_blocks = + math::ceil_div(num_max_routed_tokens, kSM90MegaMoEBlockM) + + num_experts_per_rank; + const auto num_live_pool_blocks = sched::get_num_max_live_pool_blocks( + num_pool_blocks, num_worker_ctas, hidden, intermediate_hidden, + kSM90MegaMoEBlockN, kSM90MegaMoECTAsPerTask); + const auto num_ring_tokens = + num_live_pool_blocks * kSM90MegaMoEBlockM; + return math::align(num_ring_tokens, kSM90MegaMoETokenAlignment); +} + +static std::tuple +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 std::string& mma_type, const std::string& activation, + const int& num_shared_experts = 0) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + DG_HOST_ASSERT(num_ranks > 0 and num_ranks <= kSM90MegaMoEMaxRanks); + DG_HOST_ASSERT(num_experts > 0); + DG_HOST_ASSERT(num_experts % num_ranks == 0); + DG_HOST_ASSERT(num_max_tokens_per_rank > 0); + DG_HOST_ASSERT(num_max_tokens_per_rank % kSM90MegaMoETokenAlignment == 0); + DG_HOST_ASSERT(num_topk > 0 and num_topk <= std::min(num_experts, 32)); + DG_HOST_ASSERT(num_shared_experts >= 0); + DG_HOST_ASSERT(num_topk + (num_shared_experts > 0 ? 1 : 0) <= 32); + DG_HOST_ASSERT(hidden > 0 and intermediate_hidden > 0); + // Input/L1 K128 FP32 SF rows occupy H/32 bytes and must retain 16-byte + // TMA alignment. L2 K64 FP32 SF rows occupy I/16 bytes. + DG_HOST_ASSERT(is_valid_hidden_for_sm90_mega_moe(hidden) and + intermediate_hidden % 256 == 0); + // Keep the buffer API weight-format aware even though only MXFP4 is + // implemented today. Future FP8 and INT4 backends can select their layout + // here without introducing another public allocator. + DG_HOST_ASSERT(mma_type == "fp8xmxfp4"); + DG_HOST_ASSERT(activation == "swiglu"); + + const auto num_ring_tokens = get_num_ring_tokens_for_sm90_mega_moe( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden); + const auto num_sf_ring_tokens = layout::get_num_sf_ring_tokens( + num_ring_tokens, kSM90MegaMoEBlockM); + const auto scale_layout_spec = + layout::ScaleLayoutSpec::sm90_fp32_k128_k64(hidden, intermediate_hidden); + const auto mega_buffer = layout::MegaMoEBuffer( + nullptr, hidden, intermediate_hidden, + num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, num_ring_tokens, num_sf_ring_tokens, + /*with_sf=*/ true, num_shared_experts, scale_layout_spec, + kSM90MegaMoEBlockM); + const auto shared_intermediate_hidden = intermediate_hidden * num_shared_experts; + const auto num_max_shared_sf_tokens = + layout::get_num_shared_sf_tokens( + num_max_tokens_per_rank, kSM90MegaMoEBlockM); + + // Slice function follows main's twelve-view ABI. Shared weights remain + // FP8+FP32 on SM90; only routed weights use packed MXFP4+UE8M0. + // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major + auto slice_input_buffers = [=](const torch::Tensor& buffer) { + auto x = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.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(mega_buffer.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(mega_buffer.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(mega_buffer.input_topk_weights_buffer.base)), + {num_max_tokens_per_rank, num_topk}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + + auto shared_l1_acts = x; + auto shared_l1_acts_sf = num_shared_experts > 0 ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l1_sf_buffer.base)), + {num_max_shared_sf_tokens, hidden / 128}, + {1, num_max_shared_sf_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())) : torch::Tensor(); + auto shared_l2_acts = num_shared_experts > 0 ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l2_token_buffer.base)), + {num_max_tokens_per_rank, shared_intermediate_hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())) : torch::Tensor(); + auto shared_l2_acts_sf = num_shared_experts > 0 ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l2_sf_buffer.base)), + {num_max_shared_sf_tokens, shared_intermediate_hidden / 64}, + {1, num_max_shared_sf_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())) : torch::Tensor(); + + auto l1_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l1_token_buffer.base)), + {num_ring_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(mega_buffer.l1_sf_buffer.base)), + {num_sf_ring_tokens, hidden / 128}, + {1, num_sf_ring_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto l2_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l2_token_buffer.base)), + {num_ring_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(mega_buffer.l2_sf_buffer.base)), + {num_sf_ring_tokens, intermediate_hidden / 64}, + {1, num_sf_ring_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + return std::make_tuple( + x, x_sf, topk_idx, topk_weights, + shared_l1_acts, shared_l1_acts_sf, + shared_l2_acts, shared_l2_acts_sf, + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); + }; + return {mega_buffer.get_num_bytes(), slice_input_buffers}; +} + +// SM90 (Hopper) FP8-activation MegaMoE entry point. +// +// Shared validation and dispatch for preprocessed packed MXFP4 weights with +// K32 relative UE8M0 scales plus one FP32 secondary scale per expert. +// Top-level routing is the caller's responsibility (see +// `deep_gemm/mega/__init__.py`). +static void run_sm90_fp8_mxfp4_mega_moe( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional>& shared_l1_weights_tuple_opt, + const std::optional>& shared_l2_weights_tuple_opt, + 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_mxfp4_weights_sf, l1_mxfp4_secondary] = l1_weights_tuple; + const auto [l2_weights, l2_mxfp4_weights_sf, l2_mxfp4_secondary] = l2_weights_tuple; + torch::Tensor shared_l1_weights, shared_l1_weights_sf; + torch::Tensor shared_l2_weights, shared_l2_weights_sf; + DG_HOST_ASSERT( + shared_l1_weights_tuple_opt.has_value() == + shared_l2_weights_tuple_opt.has_value()); + + // Architecture check + const auto arch_major = device_runtime->get_arch_major(); + DG_HOST_ASSERT(arch_major == 9); + + // MXFP4 weights use K32 relative UE8M0 SF plus an FP32 per-expert + // secondary. Activations use per-token per-128-K float SF. + DG_HOST_ASSERT(y.is_cuda()); + DG_HOST_ASSERT(y.dim() == 2); + DG_HOST_ASSERT(y.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(y.is_contiguous()); + const auto num_tokens = static_cast(y.size(0)); + const auto [rm, rn, rk] = recipe; + DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32); + DG_HOST_ASSERT(activation == "swiglu"); + + // Activation checks + const auto activation_clamp = + activation_clamp_opt.value_or(std::numeric_limits::infinity()); + DG_HOST_ASSERT(activation_clamp >= 0); + + // Tensor checks: preprocessed Humming MXFP4 uses int8 packed E2M1 + // [E, N, K/2], an opaque contiguous uint8 relative-UE8M0 payload with + // public shape [E, N, K/32], and the FP32 per-expert secondary below. + 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() == kPackedFP4); + DG_HOST_ASSERT(l2_weights.scalar_type() == kPackedFP4); + const auto [num_experts_per_rank, intermediate_hidden_2, l1_stored_k] = + get_shape<3>(l1_weights); + const auto [num_experts_per_rank_, hidden_, l2_stored_k] = + get_shape<3>(l2_weights); + const int hidden = static_cast(l1_stored_k) * 2; + const int intermediate_hidden = static_cast(l2_stored_k) * 2; + DG_HOST_ASSERT(y.size(1) == hidden); + DG_HOST_ASSERT(num_tokens >= 0); + 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(l1_mxfp4_weights_sf.is_cuda() and l2_mxfp4_weights_sf.is_cuda()); + DG_HOST_ASSERT(l1_weights.device() == y.device() and + l2_weights.device() == y.device() and + l1_mxfp4_weights_sf.device() == y.device() and + l2_mxfp4_weights_sf.device() == y.device()); + + // Keep host validation aligned with the default TMA-aligned Data layouts + // reconstructed by the generated SM90 kernel. + DG_HOST_ASSERT(is_valid_hidden_for_sm90_mega_moe(hidden) and + intermediate_hidden % 256 == 0); + + // Check the public weight-SF shape. SF is producer-warp loaded rather than + // TMA-loaded; preprocessing owns its hidden-dependent physical ordering. + DG_HOST_ASSERT(l1_mxfp4_weights_sf.scalar_type() == torch::kUInt8 and + l2_mxfp4_weights_sf.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT(l1_mxfp4_weights_sf.is_contiguous() and + l2_mxfp4_weights_sf.is_contiguous()); + DG_HOST_ASSERT(l1_mxfp4_weights_sf.dim() == 3 and + l1_mxfp4_weights_sf.size(0) == num_experts_per_rank and + l1_mxfp4_weights_sf.size(1) == intermediate_hidden * 2 and + l1_mxfp4_weights_sf.size(2) == hidden / 32); + DG_HOST_ASSERT(l2_mxfp4_weights_sf.dim() == 3 and + l2_mxfp4_weights_sf.size(0) == num_experts_per_rank and + l2_mxfp4_weights_sf.size(1) == hidden and + l2_mxfp4_weights_sf.size(2) == intermediate_hidden / 32); + DG_HOST_ASSERT(l1_mxfp4_secondary.is_cuda() and + l2_mxfp4_secondary.is_cuda()); + DG_HOST_ASSERT(l1_mxfp4_secondary.device() == y.device() and + l2_mxfp4_secondary.device() == y.device()); + DG_HOST_ASSERT(l1_mxfp4_secondary.scalar_type() == torch::kFloat and + l2_mxfp4_secondary.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(l1_mxfp4_secondary.is_contiguous() and + l2_mxfp4_secondary.is_contiguous()); + DG_HOST_ASSERT(l1_mxfp4_secondary.dim() == 1 and + l1_mxfp4_secondary.size(0) == num_experts_per_rank); + DG_HOST_ASSERT(l2_mxfp4_secondary.dim() == 1 and + l2_mxfp4_secondary.size(0) == num_experts_per_rank); + + // Match main's mixed-precision shared-expert contract: routed weights are + // MXFP4, while one or more shared experts are concatenated into a single + // wide FP8 FFN. Its output occupies one extra combine slot and has no + // router weight. SharedLinear1/SharedLinear2 are generated by the same + // persistent runtime that schedules routed live-ring tasks. + int num_shared_experts = 0; + if (shared_l1_weights_tuple_opt.has_value()) { + std::tie(shared_l1_weights, shared_l1_weights_sf) = + shared_l1_weights_tuple_opt.value(); + std::tie(shared_l2_weights, shared_l2_weights_sf) = + shared_l2_weights_tuple_opt.value(); + DG_HOST_ASSERT(shared_l1_weights.dim() == 2 and shared_l2_weights.dim() == 2); + DG_HOST_ASSERT(shared_l1_weights.scalar_type() == torch::kFloat8_e4m3fn and + shared_l2_weights.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(shared_l1_weights.is_contiguous() and + shared_l2_weights.is_contiguous()); + DG_HOST_ASSERT(get_major_type_ab(shared_l1_weights) == cute::UMMA::Major::K and + get_major_type_ab(shared_l2_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(shared_l1_weights.device() == y.device() and + shared_l2_weights.device() == y.device()); + + const auto shared_intermediate_hidden = + static_cast(shared_l2_weights.size(1)); + DG_HOST_ASSERT(shared_intermediate_hidden > 0 and + shared_intermediate_hidden % intermediate_hidden == 0); + num_shared_experts = shared_intermediate_hidden / intermediate_hidden; + DG_HOST_ASSERT(shared_l1_weights.size(0) == shared_intermediate_hidden * 2 and + shared_l1_weights.size(1) == hidden); + DG_HOST_ASSERT(shared_l2_weights.size(0) == hidden); + + // FP32 block scales remain in natural C-contiguous layout. Unlike + // main's SM100 UE8M0 scales, these are neither packed nor UTCCP-swizzled. + DG_HOST_ASSERT(shared_l1_weights_sf.is_cuda() and + shared_l2_weights_sf.is_cuda()); + DG_HOST_ASSERT(shared_l1_weights_sf.device() == y.device() and + shared_l2_weights_sf.device() == y.device()); + DG_HOST_ASSERT(shared_l1_weights_sf.scalar_type() == torch::kFloat and + shared_l2_weights_sf.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(shared_l1_weights_sf.is_contiguous() and + shared_l2_weights_sf.is_contiguous()); + DG_HOST_ASSERT(shared_l1_weights_sf.dim() == 2 and + shared_l1_weights_sf.size(0) == shared_intermediate_hidden * 2 / 128 and + shared_l1_weights_sf.size(1) == hidden / 128); + DG_HOST_ASSERT(shared_l2_weights_sf.dim() == 2 and + shared_l2_weights_sf.size(0) == hidden / 128 and + shared_l2_weights_sf.size(1) == shared_intermediate_hidden / 128); + DG_HOST_ASSERT(num_topk + 1 <= 32); + } + + // Check stats counter + if (cumulative_local_expert_recv_stats.has_value()) { + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_cuda()); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->device() == y.device()); + 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()); + } + + // Check buffer bytes + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + DG_HOST_ASSERT(num_ranks > 0 and num_ranks <= kSM90MegaMoEMaxRanks); + DG_HOST_ASSERT(rank_idx >= 0 and rank_idx < num_ranks); + DG_HOST_ASSERT(sym_buffer.is_cuda()); + DG_HOST_ASSERT(sym_buffer.device() == y.device()); + DG_HOST_ASSERT(sym_buffer.scalar_type() == torch::kChar); + DG_HOST_ASSERT(sym_buffer.is_contiguous()); + DG_HOST_ASSERT(num_max_tokens_per_rank > 0); + DG_HOST_ASSERT(num_max_tokens_per_rank % kSM90MegaMoETokenAlignment == 0); + DG_HOST_ASSERT(num_experts > 0 and num_topk > 0 and + num_topk <= std::min(num_experts, 32)); + const auto num_experts_ = num_experts_per_rank * num_ranks; + DG_HOST_ASSERT(num_experts == num_experts_); + 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, + "fp8xmxfp4", activation, num_shared_experts); + // The live-ring offsets depend on H/I/S and the active worker-SM count. + // Require the exact allocation contract so a buffer created for another + // shape, shared-expert count, or `set_num_sms` value cannot be re-sliced + // with silently shifted routed regions. + DG_HOST_ASSERT(sym_buffer.nbytes() == static_cast(num_required_bytes)); + + // Already registered tensors + const auto [x, x_sf, topk_idx, topk_weights, + shared_l1_acts, shared_l1_acts_sf, + shared_l2_acts, shared_l2_acts_sf, + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + + sm90_fp8_mxfp4_mega_moe(y, + l1_acts, l1_acts_sf, + l2_acts, l2_acts_sf, + shared_l1_acts, shared_l1_acts_sf, + shared_l2_acts, shared_l2_acts_sf, + l1_weights, l2_weights, + l1_mxfp4_weights_sf, l2_mxfp4_weights_sf, + shared_l1_weights, shared_l2_weights, + shared_l1_weights_sf, shared_l2_weights_sf, + cumulative_local_expert_recv_stats, + sym_buffer_ptrs, + rank_idx, num_max_tokens_per_rank, + num_experts_per_rank, + num_shared_experts, + num_tokens, num_topk, + hidden, intermediate_hidden, + activation_clamp, fast_math, + l1_mxfp4_secondary, l2_mxfp4_secondary); + + if (get_env("DG_COMM_KERNEL_DEBUG")) + sym_buffer.zero_(); +} + +static void fp8_mxfp4_mega_moe( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional>& shared_l1_weights_tuple_opt, + const std::optional>& shared_l2_weights_tuple_opt, + 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) { + run_sm90_fp8_mxfp4_mega_moe( + y, l1_weights_tuple, l2_weights_tuple, + shared_l1_weights_tuple_opt, shared_l2_weights_tuple_opt, + cumulative_local_expert_recv_stats, + sym_buffer, sym_buffer_ptrs, rank_idx, + num_max_tokens_per_rank, num_experts, num_topk, + recipe, activation, activation_clamp_opt, fast_math); +} + +static void register_sm90_apis(pybind11::module_& m) { +#if DG_TENSORMAP_COMPATIBLE + m.def("get_token_alignment_for_sm90_mega_moe", &get_token_alignment_for_sm90_mega_moe); + m.def("get_symm_buffer_size_for_sm90_mega_moe", &get_symm_buffer_size_for_sm90_mega_moe); + m.def("fp8_mxfp4_mega_moe", &fp8_mxfp4_mega_moe); +#endif +} + +} // namespace deep_gemm::mega diff --git a/csrc/jit/handle.hpp b/csrc/jit/handle.hpp index be3bc31c07..868a4382a2 100644 --- a/csrc/jit/handle.hpp +++ b/csrc/jit/handle.hpp @@ -36,6 +36,7 @@ static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \ DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorName); DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorString); DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute); +DECL_LAZY_CUDA_DRIVER_FUNCTION(cuOccupancyMaxActiveBlocksPerMultiprocessor); DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad); DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload); DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction); @@ -74,7 +75,9 @@ static void unload_library(const LibraryHandle& library) { static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, const cudaStream_t& stream, const int& smem_size, - const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) { + const dim3& grid_dim, const dim3& block_dim, + const int& cluster_dim, const bool& enable_pdl, + const bool& cooperative) { if (smem_size > 0) DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); @@ -85,8 +88,9 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, config.stream = stream; // Create attributes - // NOTES: must use `static` or the `attr` will be deconstructed - static LaunchAttrHandle attrs[2]; + // The launch config only borrows this array until the immediate launch. + // Thread-local storage preserves that lifetime without cross-thread races. + static thread_local LaunchAttrHandle attrs[3]; config.numAttrs = 0; config.attrs = attrs; @@ -104,9 +108,32 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, attr.val.programmaticStreamSerializationAllowed = 1; } + // Cooperative launch makes the all-CTA residency contract explicit for + // kernels that use a software grid barrier. + if (cooperative) { + auto& attr = attrs[config.numAttrs ++]; + attr.id = cudaLaunchAttributeCooperative; + attr.val.cooperative = 1; + } + return config; } +static void prefer_max_shared_memory_carveout(const KernelHandle& kernel) { + DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, + cudaSharedmemCarveoutMaxShared)); +} + +static int get_max_active_blocks_per_sm( + const KernelHandle& kernel, const int num_threads, const int smem_size) { + int num_blocks = 0; + DG_CUDA_RUNTIME_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &num_blocks, reinterpret_cast(kernel), + num_threads, smem_size)); + return num_blocks; +} + template static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) { void *ptr_args[] = { &args... }; @@ -173,7 +200,9 @@ static void unload_library(const LibraryHandle& library) { static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, const cudaStream_t& stream, const int& smem_size, - const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) { + const dim3& grid_dim, const dim3& block_dim, + const int& cluster_dim, const bool& enable_pdl, + const bool& cooperative) { if (smem_size > 0) DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size)); @@ -188,8 +217,9 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, config.hStream = stream; // Create attributes - // NOTES: must use `static` or the `attr` will be deconstructed - static LaunchAttrHandle attrs[2]; + // The launch config only borrows this array until the immediate launch. + // Thread-local storage preserves that lifetime without cross-thread races. + static thread_local LaunchAttrHandle attrs[3]; config.numAttrs = 0; config.attrs = attrs; @@ -209,9 +239,31 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, attr.value.programmaticStreamSerializationAllowed = 1; } + // Cooperative launch makes the all-CTA residency contract explicit for + // kernels that use a software grid barrier. + if (cooperative) { + auto& attr = attrs[config.numAttrs ++]; + attr.id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE; + attr.value.cooperative = 1; + } + return config; } +static void prefer_max_shared_memory_carveout(const KernelHandle& kernel) { + DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute( + kernel, CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, + CU_SHAREDMEM_CARVEOUT_MAX_SHARED)); +} + +static int get_max_active_blocks_per_sm( + const KernelHandle& kernel, const int num_threads, const int smem_size) { + int num_blocks = 0; + DG_CUDA_DRIVER_CHECK(lazy_cuOccupancyMaxActiveBlocksPerMultiprocessor( + &num_blocks, kernel, num_threads, smem_size)); + return num_blocks; +} + template static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) { void *ptr_args[] = { &args... }; diff --git a/csrc/jit/kernel_runtime.hpp b/csrc/jit/kernel_runtime.hpp index 40597fb448..d5e7a19084 100644 --- a/csrc/jit/kernel_runtime.hpp +++ b/csrc/jit/kernel_runtime.hpp @@ -17,12 +17,19 @@ struct LaunchArgs { int smem_size; int cluster_dim; bool enable_pdl; - - LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true): - grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {} - - LaunchArgs(const std::pair& grid_dim, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true): - grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {} + bool cooperative; + + LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, + const int& cluster_dim = 1, const bool& enable_pdl = true, + const bool& cooperative = false): + grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), + cluster_dim(cluster_dim), enable_pdl(enable_pdl), cooperative(cooperative) {} + + LaunchArgs(const std::pair& grid_dim, const int& num_threads, const int& smem_size = 0, + const int& cluster_dim = 1, const bool& enable_pdl = true, + const bool& cooperative = false): + grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), + cluster_dim(cluster_dim), enable_pdl(enable_pdl), cooperative(cooperative) {} }; class KernelRuntime final { @@ -141,22 +148,25 @@ class LaunchRuntime { const auto stream = at::cuda::getCurrentCUDAStream(); LaunchArgs launch_args = args.launch_args; - // Allow runtime override from Python. - // NOTES: the default is enabled. - launch_args.enable_pdl = device_runtime->get_pdl(); + // A caller may disable PDL for grid-barrier or residency constraints. + // The global runtime switch can only disable PDL further, never force it on. + launch_args.enable_pdl = + launch_args.enable_pdl and device_runtime->get_pdl(); const dim3 grid_dim = {static_cast(launch_args.grid_dim.first), static_cast(launch_args.grid_dim.second), 1}; const dim3 block_dim = {static_cast(launch_args.num_threads), 1, 1}; auto config = construct_launch_config(kernel, stream, launch_args.smem_size, - grid_dim, block_dim, launch_args.cluster_dim, launch_args.enable_pdl); + grid_dim, block_dim, launch_args.cluster_dim, + launch_args.enable_pdl, launch_args.cooperative); // Launch in the derived class if (get_env("DG_JIT_DEBUG")) { - printf("Launch kernel with {%d, %d} x %d, shared memory: %d bytes, cluster: %d, pdl: %d, stream: %ld\n", + printf("Launch kernel with {%d, %d} x %d, shared memory: %d bytes, cluster: %d, pdl: %d, cooperative: %d, stream: %ld\n", launch_args.grid_dim.first, launch_args.grid_dim.second, launch_args.num_threads, - launch_args.smem_size, launch_args.cluster_dim, launch_args.enable_pdl, stream.id()); + launch_args.smem_size, launch_args.cluster_dim, launch_args.enable_pdl, + launch_args.cooperative, stream.id()); } Derived::launch_impl(kernel, config, args); } 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..1a3ece1cc2 --- /dev/null +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -0,0 +1,138 @@ +#pragma once + +#include +#include + +#include "../../utils/exception.hpp" + +#include +#include +#include + +#include "../../utils/math.hpp" +#include "sm90.hpp" + +namespace deep_gemm { + +// The retained SM90 path is the compact Humming FP8 x MXFP4 persistent kernel. +// Its block/thread schedule is fixed so the host sizing logic and the kernel's +// two-CTA occupancy contract cannot drift independently. +struct MegaMoESM90Config { + using Schedule = layout::Sm90HummingMoeSchedule; + static constexpr int block_m = static_cast(Schedule::block_m); + static constexpr int block_n = static_cast(Schedule::block_n); + static constexpr int block_k = static_cast(Schedule::block_k); + static constexpr int num_stages = static_cast(Schedule::num_stages); + static constexpr int num_dispatch_threads = + static_cast(Schedule::num_dispatch_threads); + static constexpr int num_non_epilogue_threads = + static_cast(Schedule::num_non_epilogue_threads); + static constexpr int num_epilogue_threads = + static_cast(Schedule::num_epilogue_threads); + + int num_sms; + int smem_size; +}; + +constexpr int kSM90MoeMaxLatencyOverlapTokens = 4096; + +struct Sm90MoeHeuristicInput { + int launch_num_sms; + + int num_experts; + int num_tokens; + int hidden, intermediate_hidden; +}; + +static int get_mxfp4_pipeline_smem_size_for_mega_moe_sm90( + const int smem_capacity, + const int num_experts, + const int hidden, + const bool double_buffer_mxfp4_expanded_b) { + constexpr int kSmemAlignment = 1024; + constexpr int block_m = MegaMoESM90Config::block_m; + constexpr int block_n = MegaMoESM90Config::block_n; + constexpr int block_k = MegaMoESM90Config::block_k; + constexpr int num_stages = MegaMoESM90Config::num_stages; + constexpr int num_dispatch_warps = + MegaMoESM90Config::num_dispatch_threads / 32; + constexpr int num_epilogue_warps = + MegaMoESM90Config::num_epilogue_threads / 32; + + 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; + + constexpr int smem_cd_l1 = + block_m * (block_n / 2); + constexpr 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); + + constexpr int smem_sfa_half_stride_bytes = + ((block_m * static_cast(sizeof(float)) + 127) / 128) * 128; + constexpr int smem_sfa_per_stage = + (block_k / 64) * smem_sfa_half_stride_bytes; + constexpr int smem_sfb_per_stage = block_n * (block_k / 32); + constexpr int smem_packed_b_per_stage = block_n * block_k / 2; + constexpr int smem_a_per_stage = block_m * block_k; + constexpr int smem_barriers_per_stage = 2 * 8; + constexpr int smem_per_stage = + smem_a_per_stage + smem_packed_b_per_stage + + smem_sfa_per_stage + smem_sfb_per_stage + + smem_barriers_per_stage; + + const int smem_expanded_b_scratch = + block_n * block_k * (double_buffer_mxfp4_expanded_b ? 2 : 1); + constexpr int smem_barriers_fixed = + (num_dispatch_warps + 2 * num_epilogue_warps) * 8; + const int smem_fixed = smem_dispatch_size + smem_cd + + smem_expanded_b_scratch + smem_barriers_fixed; + const int smem_size = smem_fixed + num_stages * smem_per_stage; + return smem_size <= smem_capacity ? smem_size : 0; +} + +// Compact Hopper MXFP4 schedule. One math warpgroup owns fixed expanded-B +// scratch. Flash reserves a second ping-pong tile for both latency and +// throughput workloads so the next packed-B stage can be decoded while the +// current WGMMA group is in flight. A, packed-B, SFA, and SFB use a three-stage +// producer pipeline. Two logical worker CTAs are launched per physical H20 SM; +// the exact-kernel occupancy check at launch is the hard safety gate for grid +// barriers. +static MegaMoESM90Config select_mxfp4_mega_moe_sm90( + const Sm90MoeHeuristicInput& input) { + constexpr int block_m = MegaMoESM90Config::block_m; + constexpr int block_n = MegaMoESM90Config::block_n; + constexpr int block_k = MegaMoESM90Config::block_k; + + DG_HOST_ASSERT((2 * input.intermediate_hidden) % block_n == 0 and + input.hidden % block_n == 0); + DG_HOST_ASSERT(input.hidden % block_k == 0 and + input.intermediate_hidden % block_k == 0); + + const int num_worker_ctas = 2 * input.launch_num_sms; + const bool double_buffer_mxfp4_expanded_b = + input.hidden == 4096; + const int smem_size = get_mxfp4_pipeline_smem_size_for_mega_moe_sm90( + SM90ArchSpec::smem_capacity, + input.num_experts, + input.hidden, + double_buffer_mxfp4_expanded_b); + // `smem_capacity` is the opt-in per-block limit. Reserve the remaining + // 1 KiB/CTA implementation overhead in the two-CTA static precheck; the + // exact JIT kernel still goes through the runtime occupancy hard gate. + DG_HOST_ASSERT(smem_size > 0 and + 2 * smem_size <= SM90ArchSpec::smem_capacity - 1024); + return { + num_worker_ctas, + smem_size, + }; +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index 2e617d11a5..005c3b57f0 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -110,18 +110,20 @@ static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const in } } +// `force_uint8` is reserved for staging packed SM90 MXFP4 int8 payloads as raw bytes. static CUtensorMap make_tma_2d_desc(const torch::Tensor& t, int gmem_inner_dim, int gmem_outer_dim, int smem_inner_dim, int smem_outer_dim, const int& gmem_outer_stride, const int& swizzle_mode, const int& swizzle_base = 0, const bool& allow_tf32 = false, - const bool& fp4_unpacked_smem = true) { + const bool& fp4_unpacked_smem = true, + const bool& force_uint8 = false) { const auto elem_size = static_cast(t.element_size()); if (swizzle_mode != 0) smem_inner_dim = swizzle_mode / elem_size; - if (t.scalar_type() == kPackedFP4) { + if (t.scalar_type() == kPackedFP4 and not force_uint8) { // Inner dim must be a multiple of 64B for .b4x16_p64 DG_HOST_ASSERT(not fp4_unpacked_smem or gmem_inner_dim % 128 == 0); @@ -142,7 +144,8 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t, reinterpret_cast(t.data_ptr())); } DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled( - &tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem), + &tensor_map, force_uint8 ? CU_TENSOR_MAP_DATA_TYPE_UINT8 : + aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem), 2, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base), CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); 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..507ed4e132 --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -0,0 +1,464 @@ +#pragma once + +#include +#include + +#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-activation, MXFP4-weight MegaMoE host runtime +// ---------------------------------------------------------------------------- +// This is the SM90 counterpart of `SM100FP8FP4MegaMoERuntime`. The kernel +// itself lives in `deep_gemm/impls/sm90_fp8_mega_moe.cuh` and uses the same +// dispatch/combine contract with an SM90 FP8 TMA/WGMMA implementation. +// +// Differences from SM100 path: +// * Routed activations are FP8 (e4m3). Routed weights are packed MXFP4 and +// expanded to FP8 in shared memory before Hopper WGMMA. +// * Routed MXFP4 uses relative/rebased UE8M0 bytes at K32 granularity plus +// one FP32 secondary scale per expert. Shared experts retain FP8 weights +// with per-128-channel FP32 scales; neither uses the SM100 UTCCP layout. +// * No tensor memory: WGMMA accumulators are register-resident. +// * One CTA processes each work item; there is no cluster multicast or 2-CTA UMMA. +// ============================================================================ + +class SM90FP8MXFP4MegaMoERuntime final : public LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_max_tokens_per_rank; + int hidden, intermediate_hidden; + int num_experts, num_shared_experts, num_topk; + int num_ranks; + int num_ring_tokens, num_sf_ring_tokens; + float activation_clamp; + bool fast_math; + MegaMoESM90Config config; + + // Runtime arguments. num_tokens also selects compile-time MXFP4 + // scale-overlap and L2 C/D swizzle buckets during generated-source + // construction. + void* y; + int* cumulative_local_expert_recv_stats; + int num_tokens; + layout::SymBuffer<> sym_buffer_ptrs; + + // Tensormaps for activations and weights. The B producer can stage + // MXFP4 relative-scale bytes with ordinary global/shared instructions + // because each row exposes only a four-byte TMA box, illegal on SM90. + CUtensorMap tensor_map_l1_acts; + CUtensorMap tensor_map_l1_acts_sf; + CUtensorMap tensor_map_l1_weights; + const float* l1_mxfp4_secondary; + const uint8_t* l1_mxfp4_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_mxfp4_secondary; + const uint8_t* l2_mxfp4_weights_sf; + CUtensorMap tensor_map_shared_l1_acts; + CUtensorMap tensor_map_shared_l1_acts_sf; + CUtensorMap tensor_map_shared_l1_weights; + const float* shared_l1_weights_sf; + CUtensorMap tensor_map_shared_l1_output; + CUtensorMap tensor_map_shared_l2_acts; + CUtensorMap tensor_map_shared_l2_acts_sf; + CUtensorMap tensor_map_shared_l2_weights; + const float* shared_l2_weights_sf; + + // Launch configs + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + const bool overlap_mxfp4_scale_path = + args.hidden >= 4096 or + args.num_tokens <= kSM90MoeMaxLatencyOverlapTokens; + const bool use_prmt_mxfp4_exponent = + (args.hidden == 4096 and args.num_tokens < 1024) or + (args.hidden == 7168 and + (args.num_tokens == 8 or + args.num_tokens == 16 or + args.num_tokens == 32 or + args.num_tokens == 64)); + const bool use_incremental_mxfp4_descriptor = + args.hidden == 4096 and + args.num_tokens > kSM90MoeMaxLatencyOverlapTokens; + // The mature swap path's vector scale staging, packed promotion, and + // bank-permuted decoder move Flash M128 below the regular crossover. + const bool small_m_swap_ab = + args.num_shared_experts == 0 and + ((args.hidden == 4096 and args.num_tokens <= 128) or + (args.hidden == 7168 and args.num_tokens <= 128)); + const uint32_t max_swap_ab_tokens = + args.num_tokens <= 8 ? 8 : + args.num_tokens <= 16 ? 16 : + args.num_tokens <= 32 ? 32 : 64; + const bool packed_bf16_swap_epilogue = + args.num_shared_experts == 0 and + ((args.hidden == 4096 and + (args.num_tokens == 16 or args.num_tokens == 32)) or + (args.hidden == 7168 and + (args.num_tokens == 16 or args.num_tokens == 32))); + const bool bank_permute_mxfp4_pair_loads = + args.hidden == 4096 or + (args.hidden == 7168 and + ((small_m_swap_ab and + (args.num_tokens == 8 or + args.num_tokens == 32 or + args.num_tokens == 64 or + args.num_tokens == 128)) or + (not small_m_swap_ab and args.num_tokens >= 512))); + constexpr int kL2CDSwizzleMinTokens = 1024; + const bool swizzle_l2_cd = + args.num_tokens >= kL2CDSwizzleMinTokens; + const bool sparse_dispatch_completion = + args.hidden == 4096 and + (args.num_tokens == 32 or args.num_tokens == 1024); + return fmt::format(R"( +{} +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_mxfp4_mega_moe_persistent_impl< + {}, + {}, {}, + {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, {}, {} + >); +}}; +)", + sparse_dispatch_completion ? + "#define DG_SM90_SPARSE_DISPATCH_COMPLETION 1" : "", + args.num_max_tokens_per_rank, + args.hidden, args.intermediate_hidden, + args.num_experts, args.num_topk, + args.config.num_sms, args.num_ranks, + to_string(args.activation_clamp), + args.fast_math ? "true" : "false", + small_m_swap_ab ? "true" : "false", + max_swap_ab_tokens, + packed_bf16_swap_epilogue ? "true" : "false", + swizzle_l2_cd ? "true" : "false", + overlap_mxfp4_scale_path ? "true" : "false", + use_prmt_mxfp4_exponent ? "true" : "false", + use_incremental_mxfp4_descriptor ? "true" : "false", + bank_permute_mxfp4_pair_loads ? "true" : "false", + args.num_ring_tokens, args.num_sf_ring_tokens, args.num_shared_experts); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + const int num_threads = + args.config.num_dispatch_threads + + args.config.num_non_epilogue_threads + + args.config.num_epilogue_threads; + DG_HOST_ASSERT(num_threads == 256 and args.config.num_stages == 3 and + args.config.num_sms == 2 * device_runtime->get_num_sms() and + args.launch_args.grid_dim.first == args.config.num_sms); + DG_HOST_ASSERT(device_runtime->get_prop()->cooperativeLaunch and + "SM90 MXFP4 MegaMoE requires cooperative launch support"); + + // The JIT cache owns every loaded handle for process lifetime, so the + // exact-kernel residency preflight is stable after its first successful + // launch. Avoid repeating carveout and occupancy driver calls on every + // decode forward while remaining host-thread safe. + { + static std::mutex preflight_mutex; + static std::unordered_set preflighted_kernels; + const std::lock_guard lock(preflight_mutex); + if (preflighted_kernels.count(kernel) == 0) { + prefer_max_shared_memory_carveout(kernel); + const int max_active_blocks = get_max_active_blocks_per_sm( + kernel, num_threads, args.config.smem_size); + DG_HOST_ASSERT(max_active_blocks >= 2 and + "MXFP4 logical 2x-SM grid requires two resident CTAs per physical SM"); + DG_HOST_ASSERT(args.launch_args.grid_dim.first <= + max_active_blocks * device_runtime->get_prop()->multiProcessorCount and + "SM90 MXFP4 MegaMoE cooperative grid is too large for this device"); + preflighted_kernels.insert(kernel); + } + } + 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_mxfp4_secondary, + args.l1_mxfp4_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_mxfp4_secondary, + args.l2_mxfp4_weights_sf, + args.tensor_map_shared_l1_acts, + args.tensor_map_shared_l1_acts_sf, + args.tensor_map_shared_l1_weights, + args.shared_l1_weights_sf, + args.tensor_map_shared_l1_output, + args.tensor_map_shared_l2_acts, + args.tensor_map_shared_l2_acts_sf, + args.tensor_map_shared_l2_weights, + args.shared_l2_weights_sf + )); + } +}; + +static void sm90_fp8_mxfp4_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& shared_l1_acts, const torch::Tensor& shared_l1_acts_sf, + const torch::Tensor& shared_l2_acts, const torch::Tensor& shared_l2_acts_sf, + const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, + const torch::Tensor& l1_mxfp4_weights_sf, const torch::Tensor& l2_mxfp4_weights_sf, + const torch::Tensor& shared_l1_weights, const torch::Tensor& shared_l2_weights, + const torch::Tensor& shared_l1_weights_sf, const torch::Tensor& shared_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_shared_experts, + const int& num_tokens, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const float& activation_clamp, + const bool& fast_math, + const torch::Tensor& l1_mxfp4_secondary, + const torch::Tensor& l2_mxfp4_secondary +) { + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts = num_experts_per_rank * num_ranks; + const auto num_ring_tokens = static_cast(l1_acts.size(0)); + const auto num_sf_ring_tokens = static_cast(l1_acts_sf.size(0)); + const auto shared_intermediate_hidden = + intermediate_hidden * num_shared_experts; + DG_HOST_ASSERT(num_shared_experts >= 0); + DG_HOST_ASSERT(num_shared_experts == 0 or + (shared_l1_acts.defined() and shared_l1_acts_sf.defined() and + shared_l2_acts.defined() and shared_l2_acts_sf.defined() and + shared_l1_weights.defined() and shared_l2_weights.defined() and + shared_l1_weights_sf.defined() and shared_l2_weights_sf.defined())); + + // Resolve the production MXFP4 persistent launch config once. + const int num_sms = device_runtime->get_num_sms(); + const Sm90MoeHeuristicInput heuristic_input { + num_sms, + num_experts, + num_tokens, + hidden, intermediate_hidden + }; + const auto config = select_mxfp4_mega_moe_sm90(heuristic_input); + + // 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). + // Routed weight SF: K32 relative UE8M0 bytes plus a per-expert FP32 + // secondary, both passed as ordinary pointers. Shared-expert weight SF is + // a block-(128, 128) FP32 pointer. + constexpr int kGranK = 128; + constexpr int kL2ActsSFGranK = 64; + // Keep each tensor-map box within Hopper's TMA limits if the schedule is + // widened later; the current Humming schedule resolves both to 128. + const int tma_block_k = std::min(config.block_k, kGranK); + const int tma_block_n = std::min(config.block_n, 256); + const int pool_tokens = num_ring_tokens; + const int sf_stride_tokens = num_sf_ring_tokens; + const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, + hidden, pool_tokens, + tma_block_k, config.block_m, + static_cast(l1_acts.stride(-2)), + 128); + const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, + sf_stride_tokens, hidden, + config.block_m, kGranK, + 1, 0); + const auto tensor_map_l1_weights = make_tma_2d_desc( + l1_weights, + hidden / 2, + num_experts_per_rank * intermediate_hidden * 2, + tma_block_k / 2, + tma_block_n, + static_cast(l1_weights.stride(-2)), + 64, 0, false, true, true); + // L1 output (post-SwiGLU FP8): N is halved. The SM90 epilogue writes this + // staging tile to SMEM as plain row-major bytes, so the TMA store descriptor + // must use no shared-memory swizzle. Later L2 TMA loads may still swizzle + // from this row-major global buffer into their own SMEM tile. + // One warpgroup produces and stores the complete post-SwiGLU tile. + constexpr int l1_output_box_n = MegaMoESM90Config::block_n / 2; + constexpr int l1_output_box_m = MegaMoESM90Config::block_m; + const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, + intermediate_hidden, pool_tokens, + l1_output_box_n, l1_output_box_m, + static_cast(l2_acts.stride(-2)), + 0); + const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts, + intermediate_hidden, pool_tokens, + tma_block_k, config.block_m, + static_cast(l2_acts.stride(-2)), + 128); + const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, + sf_stride_tokens, intermediate_hidden, + config.block_m, kL2ActsSFGranK, + 1, 0); + const auto tensor_map_l2_weights = make_tma_2d_desc( + l2_weights, + intermediate_hidden / 2, + num_experts_per_rank * hidden, + tma_block_k / 2, + tma_block_n, + static_cast(l2_weights.stride(-2)), + 64, 0, false, true, true); + + // Shared experts retain Hopper's native FP8 K-major weight path. Their + // FP32 block-(128, 128) weight scales stay in natural row-major storage + // and are passed as ordinary pointers rather than TMA descriptors. + const auto tensor_map_shared_l1_acts = num_shared_experts > 0 ? + make_tma_2d_desc( + shared_l1_acts, + hidden, num_max_tokens_per_rank, + tma_block_k, config.block_m, + static_cast(shared_l1_acts.stride(-2)), + 128) : tensor_map_l1_acts; + const auto tensor_map_shared_l1_acts_sf = num_shared_experts > 0 ? + make_tma_sf_desc( + cute::UMMA::Major::MN, shared_l1_acts_sf, + static_cast(shared_l1_acts_sf.size(0)), hidden, + config.block_m, kGranK, + 1, 0) : tensor_map_l1_acts_sf; + const auto tensor_map_shared_l1_weights = num_shared_experts > 0 ? + make_tma_2d_desc( + shared_l1_weights, + hidden, shared_intermediate_hidden * 2, + tma_block_k, tma_block_n, + static_cast(shared_l1_weights.stride(-2)), + 128) : tensor_map_l1_weights; + const auto tensor_map_shared_l1_output = num_shared_experts > 0 ? + make_tma_2d_desc( + shared_l2_acts, + shared_intermediate_hidden, num_max_tokens_per_rank, + l1_output_box_n, l1_output_box_m, + static_cast(shared_l2_acts.stride(-2)), + 0) : tensor_map_l1_output; + const auto tensor_map_shared_l2_acts = num_shared_experts > 0 ? + make_tma_2d_desc( + shared_l2_acts, + shared_intermediate_hidden, num_max_tokens_per_rank, + tma_block_k, config.block_m, + static_cast(shared_l2_acts.stride(-2)), + 128) : tensor_map_l2_acts; + const auto tensor_map_shared_l2_acts_sf = num_shared_experts > 0 ? + make_tma_sf_desc( + cute::UMMA::Major::MN, shared_l2_acts_sf, + static_cast(shared_l2_acts_sf.size(0)), + shared_intermediate_hidden, + config.block_m, kL2ActsSFGranK, + 1, 0) : tensor_map_l2_acts_sf; + const auto tensor_map_shared_l2_weights = num_shared_experts > 0 ? + make_tma_2d_desc( + shared_l2_weights, + shared_intermediate_hidden, hidden, + tma_block_k, tma_block_n, + static_cast(shared_l2_weights.stride(-2)), + 128) : tensor_map_l2_weights; + + // 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 + auto persistent_config = config; + // The unified sizing already covers both logical phases. Add only the + // persistent scheduler mailboxes/barriers here; the physical SF ring is the + // descriptor stride used after wraparound. + persistent_config.smem_size += + (4 + (num_shared_experts > 0 ? 2 : 0)) * + static_cast(sizeof(cutlass::arch::ClusterTransactionBarrier)) + + 2 * static_cast(sizeof(sched::TaskInfo)); + + const SM90FP8MXFP4MegaMoERuntime::Args args = { + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .hidden = hidden, .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, + .num_shared_experts = num_shared_experts, + .num_topk = num_topk, + .num_ranks = num_ranks, + .num_ring_tokens = num_ring_tokens, + .num_sf_ring_tokens = num_sf_ring_tokens, + .activation_clamp = activation_clamp, + .fast_math = fast_math, + .config = persistent_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_mxfp4_secondary = l1_mxfp4_secondary.data_ptr(), + .l1_mxfp4_weights_sf = l1_mxfp4_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_mxfp4_secondary = l2_mxfp4_secondary.data_ptr(), + .l2_mxfp4_weights_sf = l2_mxfp4_weights_sf.data_ptr(), + .tensor_map_shared_l1_acts = tensor_map_shared_l1_acts, + .tensor_map_shared_l1_acts_sf = tensor_map_shared_l1_acts_sf, + .tensor_map_shared_l1_weights = tensor_map_shared_l1_weights, + .shared_l1_weights_sf = num_shared_experts > 0 ? + shared_l1_weights_sf.data_ptr() : nullptr, + .tensor_map_shared_l1_output = tensor_map_shared_l1_output, + .tensor_map_shared_l2_acts = tensor_map_shared_l2_acts, + .tensor_map_shared_l2_acts_sf = tensor_map_shared_l2_acts_sf, + .tensor_map_shared_l2_weights = tensor_map_shared_l2_weights, + .shared_l2_weights_sf = num_shared_experts > 0 ? + shared_l2_weights_sf.data_ptr() : nullptr, + .launch_args = LaunchArgs( + config.num_sms, + persistent_config.num_dispatch_threads + + persistent_config.num_non_epilogue_threads + + persistent_config.num_epilogue_threads, + persistent_config.smem_size, 1, false, true) + }; + // A live-block ring requires L1 production and L2 consumption to overlap + // in one cooperative grid. + const auto code = SM90FP8MXFP4MegaMoERuntime::generate(args); + const auto runtime = compiler->build( + "sm90_fp8_mxfp4_mega_moe_persistent_impl", + code); + SM90FP8MXFP4MegaMoERuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index a966afe1ed..55c0fa2b33 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -7,6 +7,7 @@ #include "apis/gemm.hpp" #include "apis/layout.hpp" #include "apis/mega.hpp" +#include "apis/sm90_mega.hpp" #include "apis/runtime.hpp" #ifndef TORCH_EXTENSION_NAME @@ -24,5 +25,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { deep_gemm::gemm::register_apis(m); deep_gemm::layout::register_apis(m); deep_gemm::mega::register_apis(m); + deep_gemm::mega::register_sm90_apis(m); deep_gemm::runtime::register_apis(m); } diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index f5e52b2d12..a75c7817fc 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -85,9 +85,14 @@ # Mega kernels from .mega import ( SymmBuffer, + SM90SymmBuffer, get_symm_buffer_for_mega_moe, + get_symm_buffer_for_sm90_mega_moe, transform_weights_for_mega_moe, + transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90, + transform_shared_weights_for_fp8_mega_moe_sm90, fp8_fp4_mega_moe, + fp8_mega_moe, bf16_mega_moe, ) 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..3eb3804f42 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -0,0 +1,3791 @@ +#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 +#include +#include +#include +#include +namespace deep_gemm { + +// ============================================================================ +// SM90 (Hopper) FP8 x MXFP4 MegaMoE — persistent 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 (1 or 2, 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 the 64 post-SwiGLU columns of this block, 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. +// ============================================================================ + +// The model-load preprocessing keeps every magnitude nibble in place but moves +// the eight signs in each packed word to [s0,s4,s1,s5,s2,s6,s3,s7]. The low +// nibble signs of the four bytes therefore belong to outputs 0..3, and the +// high nibble signs belong to outputs 4..7. This removes the two sign-gather +// PRMTs from every decode without changing its byte size. +CUTLASS_DEVICE uint2 sm90_mxfp4_e4m3_lookup( + const uint32_t exponent_offset) { + uint2 result; + asm volatile( + "{\n\t" + "mad.lo.u32 %0, %2, 0x08080800, 0x0c080000;\n\t" + "mad.lo.u32 %1, %2, 0x08080808, 0x1c181410;\n\t" + "}" + : "=r"(result.x), "=r"(result.y) + : "r"(exponent_offset)); + return result; +} + +CUTLASS_DEVICE uint2 sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + const uint32_t packed, const uint2 lookup) { + uint2 result; + asm volatile( + "{\n\t" + ".reg .b32 temp0, temp1, out_lo;\n\t" + "shl.b32 out_lo, %2, 4;\n\t" + "and.b32 temp0, %2, 0x77777777;\n\t" + "prmt.b32 temp1, %3, %4, temp0;\n\t" + "lop3.b32 out_lo, out_lo, 0x80808080, temp1, 0xea;\n\t" + "shr.u32 temp0, temp0, 16;\n\t" + "prmt.b32 temp1, %3, %4, temp0;\n\t" + "lop3.b32 %1, %2, 0x80808080, temp1, 0xea;\n\t" + "mov.b32 %0, out_lo;\n\t" + "}" + : "=r"(result.x), "=r"(result.y) + : "r"(packed), "r"(lookup.x), "r"(lookup.y)); + return result; +} + +CUTLASS_DEVICE uint2 sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + const uint32_t packed, const uint32_t exponent_offset) { + uint2 result; + asm volatile( + "{\n\t" + ".reg .b32 lookup_lo, lookup_hi, temp0, temp1, out_lo;\n\t" + "mad.lo.u32 lookup_lo, %3, 0x08080800, 0x0c080000;\n\t" + "mad.lo.u32 lookup_hi, %3, 0x08080808, 0x1c181410;\n\t" + "shl.b32 out_lo, %2, 4;\n\t" + "and.b32 temp0, %2, 0x77777777;\n\t" + "prmt.b32 temp1, lookup_lo, lookup_hi, temp0;\n\t" + "lop3.b32 out_lo, out_lo, 0x80808080, temp1, 0xea;\n\t" + "shr.u32 temp0, temp0, 16;\n\t" + "prmt.b32 temp1, lookup_lo, lookup_hi, temp0;\n\t" + "lop3.b32 %1, %2, 0x80808080, temp1, 0xea;\n\t" + "mov.b32 %0, out_lo;\n\t" + "}" + : "=r"(result.x), "=r"(result.y) + : "r"(packed), "r"(exponent_offset)); + return result; +} + +CUTLASS_DEVICE uint4 sm90_mxfp4_reordered_signs_e2m1x16_to_e4m3x16_bits( + const uint32_t packed0, + const uint32_t packed1, + const uint32_t exponent_offset) { + uint4 result; + asm volatile( + "{\n\t" + ".reg .b32 lookup_lo, lookup_hi, temp0, temp1, out_lo;\n\t" + "mad.lo.u32 lookup_lo, %6, 0x08080800, 0x0c080000;\n\t" + "mad.lo.u32 lookup_hi, %6, 0x08080808, 0x1c181410;\n\t" + "shl.b32 out_lo, %4, 4;\n\t" + "and.b32 temp0, %4, 0x77777777;\n\t" + "prmt.b32 temp1, lookup_lo, lookup_hi, temp0;\n\t" + "lop3.b32 out_lo, out_lo, 0x80808080, temp1, 0xea;\n\t" + "shr.u32 temp0, temp0, 16;\n\t" + "prmt.b32 temp1, lookup_lo, lookup_hi, temp0;\n\t" + "lop3.b32 %1, %4, 0x80808080, temp1, 0xea;\n\t" + "mov.b32 %0, out_lo;\n\t" + "shl.b32 out_lo, %5, 4;\n\t" + "and.b32 temp0, %5, 0x77777777;\n\t" + "prmt.b32 temp1, lookup_lo, lookup_hi, temp0;\n\t" + "lop3.b32 out_lo, out_lo, 0x80808080, temp1, 0xea;\n\t" + "shr.u32 temp0, temp0, 16;\n\t" + "prmt.b32 temp1, lookup_lo, lookup_hi, temp0;\n\t" + "lop3.b32 %3, %5, 0x80808080, temp1, 0xea;\n\t" + "mov.b32 %2, out_lo;\n\t" + "}" + : "=r"(result.x), "=r"(result.y), "=r"(result.z), "=r"(result.w) + : "r"(packed0), "r"(packed1), "r"(exponent_offset)); + return result; +} + +__forceinline__ __device__ uint32_t sm90_extract_u8_prmt( + const uint32_t packed, const uint32_t byte_idx) { + uint32_t result = 0; + switch (byte_idx) { + case 0: + asm volatile("prmt.b32 %0, %1, 0, 0x4440;" + : "=r"(result) : "r"(packed)); + break; + case 1: + asm volatile("prmt.b32 %0, %1, 0, 0x4441;" + : "=r"(result) : "r"(packed)); + break; + case 2: + asm volatile("prmt.b32 %0, %1, 0, 0x4442;" + : "=r"(result) : "r"(packed)); + break; + default: + asm volatile("prmt.b32 %0, %1, 0, 0x4443;" + : "=r"(result) : "r"(packed)); + break; + } + return result; +} + +__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); +} + +// Keep SM90 MegaMoE grid barriers on a trap-only timeout path. The shared +// SM100 barrier intentionally retains nv_dev's diagnostic printf, but compiling +// that printf into this register-heavy Hopper kernel creates a per-thread stack +// frame and can invalidate the two-CTA occupancy contract. +template +CUTLASS_DEVICE void sm90_grid_sync( + const layout::Workspace& workspace, + const uint32_t& sm_idx, const uint32_t& thread_idx, + const sync_scope_t& sync_scope) { + static constexpr uint32_t kFinishSumTag = 0x80000000u; + sync_scope(); + if (thread_idx == 0) { + const auto count_ptr = workspace.get_grid_sync_count_ptr(); + const auto old_value = ptx::atomic_add_rel( + count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1); + const auto start_clock = clock64(); + uint32_t new_value; + do { + new_value = ptx::ld_acq(count_ptr); + if (clock64() - start_clock >= comm::kNumTimeoutCycles) + DG_TRAP_ONLY_DEVICE_ASSERT(false and "SM90 grid sync timeout"); + } while (((new_value ^ old_value) & kFinishSumTag) == 0); + } + sync_scope(); +} + +// SM90-local NVLink barrier using the trap-only grid barrier above. The NVLink +// rendezvous keeps the long 300-second timeout because all participating +// ranks may compile or enter the first launch at different times. +template +CUTLASS_DEVICE void sm90_nvlink_barrier( + const layout::Workspace& workspace, + const layout::SymBuffer& sym_buffer, + const uint32_t& sm_idx, const uint32_t& thread_idx, + const sync_scope_t& sync_scope, + const bool& sync_prologue = true, + const bool& sync_epilogue = true) { + DG_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads"); + + if (sync_prologue) + sm90_grid_sync( + workspace, sm_idx, thread_idx, sync_scope); + + if (sm_idx == 0) { + auto* counter_ptr = workspace.get_nvl_barrier_counter_ptr(); + const auto status = (*counter_ptr) & 3; + const auto signal_phase = status & 1, signal_sign = status >> 1; + auto* signal_ptr = workspace.get_nvl_barrier_signal_ptr(signal_phase); + + if (thread_idx < kNumRanks) + ptx::red_add_rel_sys( + sym_buffer.map(signal_ptr, thread_idx), signal_sign ? -1 : 1); + sync_scope(); + + constexpr int64_t kNumTimeoutCycles = 300ll * 2000000000ll; + if (thread_idx == 0) { + ptx::red_add(counter_ptr, 1); + const int target = signal_sign ? 0 : static_cast(kNumRanks); + const auto start_clock = clock64(); + while (ptx::ld_acq_sys(signal_ptr) != target) { + if (clock64() - start_clock >= kNumTimeoutCycles) + DG_TRAP_ONLY_DEVICE_ASSERT(false and "NVLink barrier timeout"); + } + } + } + + if (sync_epilogue) + sm90_grid_sync( + workspace, sm_idx, thread_idx, sync_scope); +} + +#define DG_SM90_FP8_MOE_TEMPLATE_PARAMS \ + uint32_t kNumMaxTokensPerRank, \ + uint32_t kHidden, uint32_t kIntermediateHidden, \ + uint32_t kNumExperts, uint32_t kNumTopk, \ + uint32_t kNumSMs, uint32_t kNumRanks, \ + float kActivationClamp, \ + bool kFastMath, \ + bool kSmallMSwapAB, \ + uint32_t kMaxSwapABTokens, \ + bool kPackedBF16SwapEpilogue, \ + bool kSwizzleL2CD, \ + bool kOverlapMXFP4ScalePath, \ + bool kUsePRMTMXFP4Exponent, \ + bool kUseIncrementalMXFP4Descriptor, \ + bool kBankPermuteMXFP4PairLoads, \ + uint32_t kNumRingTokens, \ + uint32_t kNumSFRingTokens, \ + uint32_t kNumSharedExperts + +#define DG_SM90_FP8_MOE_KERNEL_ARGS_DECL \ + 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_mxfp4_secondary, \ + const uint8_t* __restrict__ l1_mxfp4_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_mxfp4_secondary, \ + const uint8_t* __restrict__ l2_mxfp4_weights_sf, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_acts, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_acts_sf, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_weights, \ + const float* __restrict__ shared_l1_weights_sf, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_output, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l2_acts, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l2_acts_sf, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l2_weights, \ + const float* __restrict__ shared_l2_weights_sf + +#define DG_SM90_FP8_MOE_CORE_ARGS_DECL \ + void* y, \ + int* cumulative_local_expert_recv_stats, \ + const uint32_t num_tokens, \ + const layout::SymBuffer& sym_buffer, \ + const cute::TmaDescriptor& tensor_map_l1_acts, \ + const cute::TmaDescriptor& tensor_map_l1_acts_sf, \ + const cute::TmaDescriptor& tensor_map_l1_weights, \ + const float* __restrict__ l1_mxfp4_secondary, \ + const uint8_t* __restrict__ l1_mxfp4_weights_sf, \ + const cute::TmaDescriptor& tensor_map_l1_output, \ + const cute::TmaDescriptor& tensor_map_l2_acts, \ + const cute::TmaDescriptor& tensor_map_l2_acts_sf, \ + const cute::TmaDescriptor& tensor_map_l2_weights, \ + const float* __restrict__ l2_mxfp4_secondary, \ + const uint8_t* __restrict__ l2_mxfp4_weights_sf, \ + const cute::TmaDescriptor& tensor_map_shared_l1_acts, \ + const cute::TmaDescriptor& tensor_map_shared_l1_acts_sf, \ + const cute::TmaDescriptor& tensor_map_shared_l1_weights, \ + const float* __restrict__ shared_l1_weights_sf, \ + const cute::TmaDescriptor& tensor_map_shared_l1_output, \ + const cute::TmaDescriptor& tensor_map_shared_l2_acts, \ + const cute::TmaDescriptor& tensor_map_shared_l2_acts_sf, \ + const cute::TmaDescriptor& tensor_map_shared_l2_weights, \ + const float* __restrict__ shared_l2_weights_sf + +#define DG_SM90_FP8_MOE_KERNEL_ARGS \ + y, cumulative_local_expert_recv_stats, num_tokens, sym_buffer, \ + tensor_map_l1_acts, tensor_map_l1_acts_sf, tensor_map_l1_weights, \ + l1_mxfp4_secondary, l1_mxfp4_weights_sf, tensor_map_l1_output, tensor_map_l2_acts, \ + tensor_map_l2_acts_sf, tensor_map_l2_weights, l2_mxfp4_secondary, \ + l2_mxfp4_weights_sf, tensor_map_shared_l1_acts, \ + tensor_map_shared_l1_acts_sf, tensor_map_shared_l1_weights, \ + shared_l1_weights_sf, tensor_map_shared_l1_output, \ + tensor_map_shared_l2_acts, tensor_map_shared_l2_acts_sf, \ + tensor_map_shared_l2_weights, shared_l2_weights_sf + +#define DG_SM90_FP8_MOE_CORE_TEMPLATE_ARGS \ + kNumMaxTokensPerRank, kHidden, kIntermediateHidden, kNumExperts, kNumTopk, \ + kNumSMs, kNumRanks, \ + kActivationClamp, kFastMath, kSmallMSwapAB, \ + kMaxSwapABTokens, \ + kPackedBF16SwapEpilogue, kSwizzleL2CD, \ + kOverlapMXFP4ScalePath, \ + kUsePRMTMXFP4Exponent, \ + kUseIncrementalMXFP4Descriptor, \ + kBankPermuteMXFP4PairLoads, \ + kNumRingTokens, kNumSFRingTokens, kNumSharedExperts + +template +CUTLASS_DEVICE void +sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Schedule = layout::Sm90HummingMoeSchedule; + constexpr uint32_t BLOCK_M = Schedule::block_m; + constexpr uint32_t BLOCK_N = Schedule::block_n; + constexpr uint32_t BLOCK_K = Schedule::block_k; + constexpr uint32_t kNumStages = Schedule::num_stages; + constexpr uint32_t kNumDispatchThreads = Schedule::num_dispatch_threads; + constexpr uint32_t kNumNonEpilogueThreads = + Schedule::num_non_epilogue_threads; + constexpr uint32_t kNumEpilogueThreads = Schedule::num_epilogue_threads; + constexpr uint32_t L1_SHAPE_N = kIntermediateHidden * 2; + constexpr uint32_t L1_SHAPE_K = kHidden; + constexpr uint32_t L2_SHAPE_N = kHidden; + constexpr uint32_t L2_SHAPE_K = kIntermediateHidden; + constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32; + constexpr uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32; + constexpr uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32; + constexpr uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4; + constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk; + constexpr uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks; + constexpr uint32_t kNumRingBlocks = kNumRingTokens / BLOCK_M; + constexpr bool kHasSharedExperts = kNumSharedExperts > 0; + DG_STATIC_ASSERT(not kSmallMSwapAB or + ((kHidden == 4096 or kHidden == 7168) and + not kHasSharedExperts), + "Small-M swap-AB is routed-only and model-specific"); + DG_STATIC_ASSERT(not kPackedBF16SwapEpilogue or + (kSmallMSwapAB and + (kHidden == 4096 or kHidden == 7168) and + not kHasSharedExperts), + "Packed-BF16 swap epilogue is routed small-M only"); + DG_STATIC_ASSERT( + kMaxSwapABTokens == 8 or kMaxSwapABTokens == 16 or + kMaxSwapABTokens == 32 or kMaxSwapABTokens == 64, + "Swap-AB token bound must select a supported WGMMA bucket"); + + // ===================================================================== + // Template checks + // ===================================================================== + DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + DG_STATIC_ASSERT(kNumRingTokens > 0 and kNumRingTokens % BLOCK_M == 0, + "Ring tokens must contain complete M blocks"); + DG_STATIC_ASSERT(kNumSFRingTokens >= kNumRingTokens, + "SM90 SF ring must cover every physical ring token"); + DG_STATIC_ASSERT(kNumTopk + (kHasSharedExperts ? 1u : 0u) <= 32, + "Top-k plus shared contribution must fit one warp"); + DG_STATIC_ASSERT(kHidden % BLOCK_K == 0 and kIntermediateHidden % BLOCK_K == 0, + "GEMM K dimensions must be divisible by BLOCK_K"); + // ===================================================================== + // 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(); + + // The persistent kernel executes both logical phases. + 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); + if constexpr (kHasSharedExperts) { + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_shared_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_shared_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_shared_l2_weights); + } + } + + // ===================================================================== + // Workspaces and symmetric buffer slicing (mirror SM100 layout, except SF + // for L2 activations uses per-64 K granularity) + // ===================================================================== + const auto buffer = layout::MegaMoEBuffer( + sym_buffer.get_base_ptr(), kHidden, kIntermediateHidden, + kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk, + kNumRingTokens, kNumSFRingTokens, /*with_sf=*/ true, + kNumSharedExperts, + layout::ScaleLayoutSpec::sm90_fp32_k128_k64( + kHidden, kIntermediateHidden), + BLOCK_M); + const auto workspace = buffer.workspace; + + constexpr auto fp8_token_layout = layout::Data(kHidden); + const auto input_token_buffer = buffer.input_token_buffer; + const auto input_sf_buffer = buffer.input_sf_buffer; + const auto input_topk_idx_buffer = buffer.input_topk_idx_buffer; + const auto input_topk_weights_buffer = buffer.input_topk_weights_buffer; + const auto shared_l1_token_buffer = buffer.shared_l1_token_buffer; + const auto shared_l1_sf_buffer = buffer.shared_l1_sf_buffer; + const auto shared_l2_token_buffer = buffer.shared_l2_token_buffer; + const auto shared_l2_sf_buffer = buffer.shared_l2_sf_buffer; + const auto l1_token_buffer = buffer.l1_token_buffer; + const auto l1_sf_buffer = buffer.l1_sf_buffer; + const auto l1_topk_weights_buffer = buffer.l1_topk_weights_buffer; + const auto l2_token_buffer = buffer.l2_token_buffer; + const auto l2_sf_buffer = buffer.l2_sf_buffer; + + // Combine input area stores BF16 L2 contributions for the final reduction. + constexpr uint32_t kCombineElementBytes = sizeof(nv_bfloat16); + const auto combine_token_buffer = buffer.combine_token_buffer; + + // ===================================================================== + // GEMM data types and shape constants + // ===================================================================== + using a_dtype_t = cutlass::float_e4m3_t; + using b_dtype_t = cutlass::float_e4m3_t; + using task_info_t = sched::TaskInfo; + DG_STATIC_ASSERT(sizeof(Barrier) == sizeof(uint64_t), + "Host SMEM accounting assumes eight-byte barriers"); + DG_STATIC_ASSERT(sizeof(task_info_t) == 32, + "Host SMEM accounting assumes 32-byte TaskInfo payloads"); + constexpr uint32_t WG_BLOCK_M = BLOCK_M; + constexpr uint32_t WG_BLOCK_N = BLOCK_N; + constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; // post-SwiGLU tile N + constexpr uint32_t WG_L1_OUT_BLOCK_N = WG_BLOCK_N / 2; // post-SwiGLU per-WG N + // The JIT knows the global token count. No routed expert can own more + // tokens than that, so small-M epilogues only need to materialize chunks + // up to this compile-time bound. The math mainloop still selects N8/N16/ + // N32/N64 from each expert's actual valid_m. + constexpr uint32_t kSwapABTokenChunks = kMaxSwapABTokens / 8; + constexpr uint32_t kSwapABWeightHalves = BLOCK_N / 64; + constexpr uint32_t kSwapABHalfAccumPerThread = 64 * 64 / 128; + DG_STATIC_ASSERT(BLOCK_M == 64 and BLOCK_N == 128 and + kSwapABWeightHalves == 2, + "Small-M swap-AB requires the fixed Humming tile"); + // Two-CTA MXFP4 is compiled at 128 registers/thread. Keeping both the + // 64-value WGMMA fragment and a 64-value FP32 persistent sum live creates + // a short local-memory frame. Retain the cross-promotion sum as 32 packed + // BF16 pairs and expand it only after the mainloop has released the WGMMA + // fragment. Strict mode multiplies and accumulates in FP32 and rounds only + // the persistent storage. Fast math also rounds each promotion's scale and + // WGMMA fragment to BF16 so the packed pairs can be updated with HFMA2. + using L1WGMMA = typename mma::sm90::FP8MMASelector::type; + static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, + "Unexpected WGMMA shape"); + + // SM90 MegaMoE uses one CTA per work item; A and B are CTA-local. + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M; + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; + constexpr uint32_t kSwizzleAMode = 128; + constexpr uint32_t kSwizzleBMode = 128; + constexpr uint32_t kGranK = 128; // L1 acts SF, weights SF + constexpr uint32_t kL2ActsSFGranK = 64; // L2 acts SF (per-64 K, SM90 only) + constexpr uint32_t kTMATileK = 128; + constexpr uint32_t kNumTMATilesPerStage = BLOCK_K / kTMATileK; + constexpr uint32_t kTMATileN = BLOCK_N > 256 ? 256 : BLOCK_N; + constexpr uint32_t kNumTMANTilesPerStage = BLOCK_N / kTMATileN; + + // ===================================================================== + // Shared memory layout + // ===================================================================== + constexpr uint32_t kSharedMemoryAlignment = 1024; + extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[]; + + // Combine reuses the pre-barrier SMEM region, including dispatch scratch. + 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); + // Flash reserves a second expanded tile so the math WG can decode the + // next packed-B stage while the current phase's final WGMMA group is in + // flight. Pro cannot reserve another 16 KiB without breaking the two-CTA + // occupancy contract, so it aliases the second mainloop-only tile with the + // C/D region, whose lifetime starts after the routed GEMM loop. Shared FP8 + // tasks keep their original expanded/packed slots because their independent + // B producer may prefetch across the math warpgroup's epilogue. + constexpr bool kDoubleBufferedMXFP4ExpandedBStorage = + kHidden == 4096 and + kOverlapMXFP4ScalePath; + constexpr bool kAliasMXFP4ExpandedBWithCD = + kHidden > 4096 and + kOverlapMXFP4ScalePath; + constexpr bool kPipelineMXFP4ExpandedB = + kDoubleBufferedMXFP4ExpandedBStorage or + kAliasMXFP4ExpandedBWithCD; + constexpr uint32_t SMEM_B_STORAGE_SIZE = + (kDoubleBufferedMXFP4ExpandedBStorage ? 2u : 1u) * + SMEM_B_SIZE_PER_STAGE; + // Packed MXFP4 is TMA-loaded with B64 swizzle into a temporary half-sized + // region. The packed row is exactly 64 bytes at BK128, so B64 spreads the + // one-word-per-row decoder loads across banks. The math warpgroup expands + // it into the normal B128-swizzled FP8 B tile before issuing WGMMA. + constexpr uint32_t SMEM_B_PACKED_SIZE_PER_STAGE = + LOAD_BLOCK_N * BLOCK_K / 2; + // SFA holds one aligned BLOCK_M-float vector per 64 channels. L1 uses every + // other vector (per-128); L2 uses all of them (per-64). + constexpr uint32_t kL2SFAHalfStride = + math::constexpr_align(BLOCK_M * sizeof(float), 128u) / sizeof(float); + constexpr uint32_t kNumL2SFAKGroups = BLOCK_K / kL2ActsSFGranK; + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = + kNumL2SFAKGroups * kL2SFAHalfStride * sizeof(float); + // MXFP4 relative scales contain one UE8M0 byte per logical (N, K32) + // group. Preprocessing stores K128 words contiguously across N for the + // supported DSV4 hidden sizes; larger shapes retain natural row-major + // storage. The B + // producer stages one word per output row alongside each packed-B tile. + constexpr uint32_t kMXFP4WeightGranK = 32; + constexpr uint32_t kMXFP4CoalescedScaleMaxHidden = 8192; + constexpr uint32_t kNumMXFP4SFBKGroups = BLOCK_K / kMXFP4WeightGranK; + constexpr uint32_t kL1MXFP4WeightSFStrideK = kHidden / kMXFP4WeightGranK; + constexpr uint32_t kL2MXFP4WeightSFStrideK = + kIntermediateHidden / kMXFP4WeightGranK; + constexpr uint32_t kL1MXFP4WeightSFPerExpert = + (kIntermediateHidden * 2) * kL1MXFP4WeightSFStrideK; + constexpr uint32_t kL2MXFP4WeightSFPerExpert = + kHidden * kL2MXFP4WeightSFStrideK; + // The unified persistent specialization stages routed weight SF for both + // logical phases. One shared-memory layout must cover whichever task the + // dynamic scheduler publishes next. + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = + BLOCK_N * kNumMXFP4SFBKGroups * sizeof(uint8_t); + constexpr uint32_t SMEM_SFB_STORAGE_SIZE = + kNumStages * SMEM_SFB_SIZE_PER_STAGE; + // CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte * num_wg) and + // L2 BF16 contribution. + constexpr uint32_t SMEM_CD_L1_SIZE = + kNumEpilogueWarpgroups * WG_BLOCK_M * WG_L1_OUT_BLOCK_N * + sizeof(cutlass::float_e4m3_t); + constexpr uint32_t SMEM_CD_L2_SIZE = + kNumEpilogueWarpgroups * WG_BLOCK_M * WG_BLOCK_N * + kCombineElementBytes; + constexpr uint32_t SMEM_CD_OUTPUT_BASE_SIZE = + SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE; + constexpr uint32_t SMEM_CD_OUTPUT_SIZE = math::constexpr_align( + SMEM_CD_OUTPUT_BASE_SIZE, kSharedMemoryAlignment); + constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_OUTPUT_SIZE; + DG_STATIC_ASSERT(not kAliasMXFP4ExpandedBWithCD or + SMEM_CD_SIZE >= SMEM_B_SIZE_PER_STAGE, + "C/D alias must hold one expanded MXFP4 B tile"); + + 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_STORAGE_SIZE + + kNumStages * SMEM_B_PACKED_SIZE_PER_STAGE; + + constexpr uint32_t kCombineInputHiddenBytes = kHidden * kCombineElementBytes; + constexpr uint32_t kCombineOutputHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kCombineMaxRegistersForBuffer = 128; + constexpr bool kCombineOneChunkFits = + kNumEpilogueWarps * (2 * kCombineInputHiddenBytes + kCombineOutputHiddenBytes) <= + SMEM_BEFORE_BARRIER_SIZE and + kHidden <= 32 * kCombineMaxRegistersForBuffer; + constexpr bool kCombineTwoChunksFits = + kHidden % 2 == 0 and + kNumEpilogueWarps * + (2 * (kCombineInputHiddenBytes / 2) + kCombineOutputHiddenBytes / 2) <= + SMEM_BEFORE_BARRIER_SIZE and + kHidden <= 2 * 32 * kCombineMaxRegistersForBuffer; + constexpr uint32_t kCombineNumChunks = kCombineOneChunkFits ? 1 : + (kCombineTwoChunksFits ? 2 : 4); + constexpr uint32_t kCombineChunkElems = kHidden / kCombineNumChunks; + constexpr uint32_t kCombineInputChunkBytes = + kCombineInputHiddenBytes / kCombineNumChunks; + constexpr uint32_t kCombineOutputChunkBytes = + kCombineOutputHiddenBytes / kCombineNumChunks; + constexpr uint32_t SMEM_COMBINE_ALIAS_SIZE = kNumEpilogueWarps * + (2 * kCombineInputChunkBytes + kCombineOutputChunkBytes); + DG_STATIC_ASSERT(kHidden % kCombineNumChunks == 0, "Hidden must be divisible by number of combine chunks"); + DG_STATIC_ASSERT(SMEM_COMBINE_ALIAS_SIZE <= SMEM_BEFORE_BARRIER_SIZE, + "Combine SMEM alias exceeds the pre-barrier scratch region"); + + // 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); + + auto smem_cd_base = smem_gemm_base; + // CD output is shared by L1 and L2; reinterpret-cast as needed. + auto smem_cd_l1 = reinterpret_cast(smem_cd_base); + + constexpr uint32_t SMEM_A_OFFSET = SMEM_CD_SIZE; + constexpr uint32_t SMEM_B_OFFSET = + SMEM_A_OFFSET + kNumStages * SMEM_A_SIZE_PER_STAGE; + constexpr uint32_t SMEM_B_PACKED_OFFSET = + SMEM_B_OFFSET + SMEM_B_STORAGE_SIZE; + constexpr uint32_t SMEM_SFA_OFFSET = + SMEM_B_PACKED_OFFSET + kNumStages * SMEM_B_PACKED_SIZE_PER_STAGE; + constexpr uint32_t SMEM_BARRIER_OFFSET = + SMEM_SFA_OFFSET + kNumStages * SMEM_SFA_SIZE_PER_STAGE + + SMEM_SFB_STORAGE_SIZE; + DG_STATIC_ASSERT(SMEM_A_SIZE_PER_STAGE == 8192 and + SMEM_B_STORAGE_SIZE == + (kDoubleBufferedMXFP4ExpandedBStorage ? + 32768u : 16384u) and + SMEM_B_PACKED_SIZE_PER_STAGE == 8192 and + SMEM_SFA_SIZE_PER_STAGE == 512 and + SMEM_SFB_SIZE_PER_STAGE == 512 and + SMEM_SFB_STORAGE_SIZE == 1536, + "Unexpected compact MXFP4 shared-memory tile sizes"); + DG_STATIC_ASSERT(SMEM_A_OFFSET % 128 == 0 and + SMEM_B_OFFSET % 128 == 0 and + SMEM_B_PACKED_OFFSET % 128 == 0 and + SMEM_SFA_OFFSET % 128 == 0 and + SMEM_BARRIER_OFFSET % alignof(Barrier) == 0, + "SM90 MegaMoE shared-memory regions must be 128-byte aligned"); + + auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr( + smem_gemm_base, SMEM_A_OFFSET + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b_expanded_base = math::advance_ptr( + smem_gemm_base, SMEM_B_OFFSET); + auto smem_b_expanded = utils::PatternVisitor([=](const uint32_t& i) { + DG_DEVICE_ASSERT(i < (kPipelineMXFP4ExpandedB ? 2u : 1u)); + if (i == 0) + return smem_b_expanded_base; + if constexpr (kDoubleBufferedMXFP4ExpandedBStorage) + return smem_b_expanded_base + SMEM_B_SIZE_PER_STAGE; + return reinterpret_cast(smem_cd_base); + }); + auto smem_b = utils::PatternVisitor([=](const uint32_t&) { + return smem_b_expanded_base; + }); + auto smem_b_packed = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr( + smem_gemm_base, SMEM_B_PACKED_OFFSET + + i * SMEM_B_PACKED_SIZE_PER_STAGE); + }); + // Shared experts keep FP8 weights. Reuse two disjoint physical B slots + // from the routed MXFP4 expanded/packed regions and protect the alias with + // its own two-stage consumer barrier. This preserves two CTAs/SM without + // reserving three additional 16-KiB FP8 tiles. + constexpr uint32_t kNumSharedBStages = 2; + auto smem_shared_b = utils::PatternVisitor([=](const uint32_t& i) { + DG_DEVICE_ASSERT(i < kNumSharedBStages); + if (i == 0) + return smem_b_expanded_base; + if constexpr (kDoubleBufferedMXFP4ExpandedBStorage) + return smem_b_expanded[1]; + return reinterpret_cast(smem_b_packed[0]); + }); + auto sf_start_ptr = math::advance_ptr(smem_gemm_base, + SMEM_SFA_OFFSET); + auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); + }); + auto sfb_start_ptr = sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE; + auto smem_sfb = utils::PatternVisitor([=](const uint32_t& i) { + return sfb_start_ptr + i * SMEM_SFB_SIZE_PER_STAGE; + }); + // Barriers live after the activation- and weight-SF stages. + auto barrier_start_ptr = reinterpret_cast( + math::advance_ptr(smem_gemm_base, SMEM_BARRIER_OFFSET)); + 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; }); + constexpr uint32_t kNumBaseBarriers = + kNumDispatchWarps + kNumStages * 2 + kNumEpilogueWarps * 2; + constexpr uint32_t kNumScheduleStages = 2; + constexpr uint32_t kTaskInfoSmemOffset = SMEM_BARRIER_OFFSET + + (kNumBaseBarriers + + (kHasSharedExperts ? kNumSharedBStages : 0u) + + kNumScheduleStages * 2) * sizeof(Barrier); + DG_STATIC_ASSERT(kTaskInfoSmemOffset % alignof(task_info_t) == 0, + "TaskInfo mailboxes must remain 16-byte aligned"); + auto shared_b_empty_barriers = barrier_start_ptr + kNumBaseBarriers; + auto task_info_full_barriers = + shared_b_empty_barriers + (kHasSharedExperts ? kNumSharedBStages : 0u); + auto task_info_empty_barriers = + task_info_full_barriers + kNumScheduleStages; + auto task_infos = reinterpret_cast( + task_info_empty_barriers + kNumScheduleStages); + + // ===================================================================== + // 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) { + // Producer arrivals: A(+SFA) + B. + full_barriers[i]->init(2); + empty_barriers[i]->init(kNumEpilogueWarps); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++ i) + combine_barriers[i]->init(1); + if constexpr (kHasSharedExperts) { + #pragma unroll + for (uint32_t i = 0; i < kNumSharedBStages; ++ i) + shared_b_empty_barriers[i].init(kNumEpilogueWarps); + } + #pragma unroll + for (uint32_t i = 0; i < kNumScheduleStages; ++ i) { + task_info_full_barriers[i].init(1); + // A-loader and math threads both consume the CTA-local + // mailbox. The B-loader may reuse the slot only after all + // of them have copied the payload into registers. + task_info_empty_barriers[i].init( + kNumEpilogueThreads + 32); + } + } + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); + + // ===================================================================== + // Scheduler (cluster=1) + // ===================================================================== + using scheduler_t = sched::MegaMoEScheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, kNumSMs, kNumRanks, + kNumRingBlocks, kNumSharedExperts, 1, + kSmallMSwapAB and + (kMaxSwapABTokens == 8 or kMaxSwapABTokens == 64)>; + auto scheduler = scheduler_t( + workspace, task_info_full_barriers, + task_info_empty_barriers, task_infos); + + // 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; + }; + uint32_t shared_b_stage_idx = 0, shared_b_phase = 0; + auto advance_shared_b_pipeline = [&]() { + shared_b_stage_idx ^= 1; + shared_b_phase ^= shared_b_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; + constexpr uint32_t kGemmPhaseBoundaryBarrierIdx = + kEpilogueWGBarrierStartIdx + kNumEpilogueWarpgroups; + constexpr uint32_t kExpertCountCacheBarrierIdx = + kGemmPhaseBoundaryBarrierIdx + 1; + constexpr bool kCacheProExpertCounts = + not kHasSharedExperts and kHidden == 7168 and + ((kSmallMSwapAB and kMaxSwapABTokens == 8) or + (not kSmallMSwapAB and kBankPermuteMXFP4PairLoads and + not kSwizzleL2CD)); + + // Cross-rank NVLink barrier tags + constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; + constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + + // Register reconfiguration for the fixed 64+64+128-thread CTA consumes + // half of the SM register file, preserving the two-CTA occupancy contract. + constexpr uint32_t kNumDispatchRegisters = 48; + constexpr uint32_t kNumNonEpilogueRegisters = 48; + constexpr uint32_t kNumEpilogueRegisters = 208; + constexpr uint32_t kCTARegisterBudget = + kNumDispatchRegisters * kNumDispatchThreads + + kNumNonEpilogueRegisters * kNumNonEpilogueThreads + + kNumEpilogueRegisters * kNumEpilogueThreads; + DG_STATIC_ASSERT(kCTARegisterBudget == 32768, + "Two-CTA MXFP4 must use half the SM register file"); + + constexpr uint32_t kDispatchGridSyncIndex = 0; + constexpr uint32_t kEpilogueGridSyncIndex = 1; + + int previous_task_kind = -1; + const auto sync_task_storage_alias = [&](const bool is_shared_task) { + if constexpr (kHasSharedExperts) { + const int task_kind = is_shared_task ? 1 : 0; + if (previous_task_kind >= 0 and previous_task_kind != task_kind) { + ptx::sync_unaligned( + kNumNonEpilogueThreads + kNumEpilogueThreads, + kGemmPhaseBoundaryBarrierIdx); + } + previous_task_kind = task_kind; + } + }; + + const auto invoke_persistent_task = [&](const task_info_t& task_info, + auto&& func) { + const uint32_t num_k_blocks = + math::ceil_div(task_info.shape_k, BLOCK_K); + const uint32_t n_block_idx = + scheduler_t::get_n_block_idx(task_info); + if (task_info.block_phase == sched::BlockPhase::Linear1) { + func(std::integral_constant< + sched::BlockPhase, sched::BlockPhase::Linear1>{}, + task_info.local_expert_idx, num_k_blocks, + task_info.m_block_idx, n_block_idx, + task_info.pool_block_idx, task_info.valid_m); + } else if (task_info.block_phase == sched::BlockPhase::Linear2) { + func(std::integral_constant< + sched::BlockPhase, sched::BlockPhase::Linear2>{}, + task_info.local_expert_idx, num_k_blocks, + task_info.m_block_idx, n_block_idx, + task_info.pool_block_idx, task_info.valid_m); + } else if constexpr (kHasSharedExperts) { + if (task_info.block_phase == sched::BlockPhase::SharedLinear1) { + func(std::integral_constant< + sched::BlockPhase, sched::BlockPhase::SharedLinear1>{}, + task_info.local_expert_idx, num_k_blocks, + task_info.m_block_idx, n_block_idx, + task_info.pool_block_idx, task_info.valid_m); + } else { + func(std::integral_constant< + sched::BlockPhase, sched::BlockPhase::SharedLinear2>{}, + task_info.local_expert_idx, num_k_blocks, + task_info.m_block_idx, n_block_idx, + task_info.pool_block_idx, task_info.valid_m); + } + } + }; + + const auto for_each_selected_block = [&](auto&& func) { + task_info_t task_info; + while (scheduler.get_next_task(task_info)) + invoke_persistent_task(task_info, func); + }; + + const auto produce_selected_blocks = [&](auto&& func) { + const auto process_task = [&](const task_info_t& task_info) { + invoke_persistent_task(task_info, func); + }; + if constexpr (kCacheProExpertCounts) + scheduler.template mainloop_with_task( + num_tokens, process_task); + else + scheduler.mainloop_with_task(num_tokens, process_task); + }; + + const auto cleanup_workspace = [&]() { + DG_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { + for (uint32_t i = thread_idx; i < kNumExperts; + i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + for (uint32_t i = thread_idx; i < workspace.num_ring_blocks; + i += kNumDispatchThreads) { + *workspace.get_l1_full_count_ptr(i) = 0; + *workspace.get_l1_empty_count_ptr(i) = 0; + *workspace.get_l2_full_count_ptr(i) = 0; + *workspace.get_l2_empty_count_ptr(i) = 0; + } + for (uint32_t i = thread_idx; + i < workspace.num_shared_l2_pool_blocks; + i += kNumDispatchThreads) + *workspace.get_shared_l2_full_count_ptr(i) = 0; + if (thread_idx == 0) { + *workspace.get_l1_task_count_ptr() = 0; + *workspace.get_l2_task_count_ptr() = 0; + *workspace.get_shared_l1_task_count_ptr() = 0; + *workspace.get_shared_l2_task_count_ptr() = 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)); + ptx::sync_aligned( + kNumDispatchThreads, kDispatchBarrierIdx); + if (warp_idx == 0) { + if (lane_idx == 0) + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + for (uint32_t rank = lane_idx; rank < kNumRanks; + rank += 32) + *workspace.get_expert_recv_count_ptr(rank, 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(); + } + } + }; + + // The compact frontend's dispatch and TMA warps share one warpgroup. + // `setmaxnreg` is warpgroup-collective, so all four warps must execute this + // single dynamic instruction rather than equivalent role-local call sites. + if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) + cutlass::arch::warpgroup_reg_dealloc(); + else + cutlass::arch::warpgroup_reg_alloc(); + + // ===================================================================== + // 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) { + 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) { + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + const int 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 uint32_t local_count = smem_expert_count[i]; +#if defined(DG_SM90_SPARSE_DISPATCH_COMPLETION) + if (local_count != 0) { +#endif + const uint64_t send_value = + (1ull << 32) | static_cast(local_count); + smem_expert_count[i] = static_cast( + ptx::atomic_add( + workspace.get_expert_send_count_ptr(i), send_value)); +#if defined(DG_SM90_SPARSE_DISPATCH_COMPLETION) + } +#endif + } + 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; + }); + + sm90_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 raw_expert_status = + *workspace.get_expert_send_count_ptr(i); +#if defined(DG_SM90_SPARSE_DISPATCH_COMPLETION) + // The grid barrier proves every CTA has completed its token + // writes, so publish the equivalent all-CTA completion count. + const uint64_t expert_status = + (static_cast(kNumSMs) << 32) | + static_cast(raw_expert_status); +#else + const uint64_t expert_status = raw_expert_status; +#endif + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; +#if not defined(DG_SM90_SPARSE_DISPATCH_COMPLETION) + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), + expert_status); +#endif + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + +#if defined(DG_SM90_SPARSE_DISPATCH_COMPLETION) + sm90_nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + false, false); + + if (sm_idx == 0) { + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + for (uint32_t local_expert_idx = thread_idx; + local_expert_idx < kNumExpertsPerRank; + local_expert_idx += kNumDispatchThreads) { + uint32_t num_recv_tokens = 0; + #pragma unroll + for (uint32_t rank_idx = 0; rank_idx < kNumRanks; ++ rank_idx) + num_recv_tokens += static_cast( + *workspace.get_expert_recv_count_ptr( + rank_idx, local_expert_idx)); + *workspace.get_expert_recv_count_sum_ptr(local_expert_idx) = + (static_cast(kNumSMs * kNumRanks) << 32) | + num_recv_tokens; + } + } + sm90_grid_sync( + workspace, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }); +#else + sm90_nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + false, true); +#endif + + // Selected Pro buckets make both dispatch warps and the B-loader + // scheduler independently poll and load the same completed expert + // totals. Publish one warp's snapshot through the dispatch scratch + // after the NVLink rendezvous, then let all three consumers initialize + // their private scheduler state from shared memory. + if constexpr (kCacheProExpertCounts) { + if (warp_idx == 0) + scheduler.cache_expert_recv_count(smem_expert_count); + ptx::sync_unaligned( + kNumDispatchThreads + 32, kExpertCountCacheBarrierIdx); + scheduler.fetch_cached_expert_recv_count(smem_expert_count); + } + + // Shared L1 does not depend on routed dispatch. Let dispatch pull + // routed tokens while the math warpgroup computes shared L1 instead + // of serializing both paths at the frontend barrier. + if constexpr (not kHasSharedExperts) + 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]; + + if constexpr (not kCacheProExpertCounts) + 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; + } + } + + uint32_t current_rank_in_expert_idx = 0; + uint32_t token_idx_in_rank = 0; + const uint32_t token_idx_in_expert = token_idx - expert_start_idx; + + // At small routed Flash M, most experts receive at most one route + // from each source rank. Select the source directly from the + // nonempty mask in that common case; preserve the general + // round-robin path for duplicate routes from any rank. + bool used_single_slot_rank_selection = false; + if constexpr (kSmallMSwapAB and kHidden == 4096 and + (kMaxSwapABTokens == 8 or kMaxSwapABTokens == 16 or + kMaxSwapABTokens == 32) and + kNumRanks <= 32) { + const uint32_t multi_slot_rank_mask = __ballot_sync( + 0xffffffff, stored_rank_count[0] > 1); + if (multi_slot_rank_mask == 0) { + const uint32_t nonempty_rank_mask = __ballot_sync( + 0xffffffff, stored_rank_count[0] != 0); + current_rank_in_expert_idx = __fns( + nonempty_rank_mask, 0, token_idx_in_expert + 1); + used_single_slot_rank_selection = true; + } + } + + // General round-robin rank selection (identical to SM100). + 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 slot_idx = token_idx_in_expert; + while (not used_single_slot_rank_selection) { + 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; + const uint32_t pool_block_idx = pool_token_idx / BLOCK_M; + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t physical_token_idx = + ring_block_idx * BLOCK_M + token_idx_in_expert % BLOCK_M; + + // Do not overwrite a live L1 slot from the previous generation. + constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; + const uint32_t empty_target = + (pool_block_idx / kNumRingBlocks) * kNumL1BlockNs; + if (empty_target > 0) { + const auto empty_ptr = + workspace.get_l1_empty_count_ptr(ring_block_idx); + if constexpr (kHidden == 4096 and + kUseIncrementalMXFP4Descriptor) { + // Poll the monotonic counter without invalidating L1 on + // every miss, then acquire once before reusing the slot. + while (ptx::ld_relaxed_gpu(empty_ptr) < empty_target) {} + while (ptx::ld_acq(empty_ptr) < empty_target) {} + } else { + while (ptx::ld_acq(empty_ptr) < empty_target) { + // For one-block Pro dispatch, avoid hammering the + // counter while the wider GEMM tile retires the + // previous slot. + if constexpr (kNumRanks > 1 and kHidden > 4096) { + if (num_tokens <= BLOCK_M) + __nanosleep(64); + } + // Flash M128 wraps padded expert blocks onto a second + // ring generation. Briefly back off contending + // dispatch lanes. + if constexpr (kNumRanks > 1 and kHidden == 4096) { + if (num_tokens == 2 * BLOCK_M) + __nanosleep(16); + } + } + } + } + + // TMA pull token data into SMEM + 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(); + #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 * kNumSFRingTokens + physical_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(physical_token_idx) + .get_base_ptr() = weight; + + 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(physical_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>(); + const bool is_last_token = token_idx == expert_end_idx - 1; + const uint32_t token_idx_in_block = + token_idx_in_expert % BLOCK_M; + ptx::red_add_rel( + workspace.get_l1_full_count_ptr(ring_block_idx), + is_last_token ? BLOCK_M - token_idx_in_block : 1u); + } + __syncwarp(); + } + + // Pair with the epilogue after all L2 writes and combine loads are + // globally visible, then clean every generation/task counter. + ptx::sync_unaligned( + kNumDispatchThreads + kNumEpilogueThreads, + kDispatchWithEpilogueBarrierIdx); + cleanup_workspace(); + sm90_nvlink_barrier< + kNumRanks, kNumSMs, kNumDispatchThreads, + kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag>( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { + ptx::sync_aligned( + kNumDispatchThreads, kDispatchBarrierIdx); + }, + true, false); + return; + + // ===================================================================== + // ROLE 2: two GEMM TMA producer warps, one for A+SFA and one for B+SFB. + // ===================================================================== + } else if (warp_idx == kNumDispatchWarps) { + for_each_selected_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& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t< + std::remove_reference_t>; + constexpr bool is_linear1_phase = + BlockPhaseTag::value == sched::BlockPhase::Linear1 or + BlockPhaseTag::value == sched::BlockPhase::SharedLinear1; + constexpr bool is_shared_phase = + BlockPhaseTag::value == sched::BlockPhase::SharedLinear1 or + BlockPhaseTag::value == sched::BlockPhase::SharedLinear2; + sync_task_storage_alias(is_shared_phase); + scheduler.release_task_info(); + const auto tensor_map_a_ptr = is_shared_phase ? + (is_linear1_phase ? &tensor_map_shared_l1_acts : + &tensor_map_shared_l2_acts) : + (is_linear1_phase ? &tensor_map_l1_acts : + &tensor_map_l2_acts); + const auto tensor_map_sfa_ptr = is_shared_phase ? + (is_linear1_phase ? &tensor_map_shared_l1_acts_sf : + &tensor_map_shared_l2_acts_sf) : + (is_linear1_phase ? &tensor_map_l1_acts_sf : + &tensor_map_l2_acts_sf); + + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t block_idx = is_shared_phase ? pool_block_idx : + ring_block_idx; + const bool has_valid_m = valid_m > 0; + + // Wait for the pool to be ready. Cluster peers can be dummy CTAs for + // the tail M unit when an expert has an odd number of M blocks. + if (has_valid_m) { + if constexpr (BlockPhaseTag::value == sched::BlockPhase::Linear1) { + const auto ptr = + workspace.get_l1_full_count_ptr(ring_block_idx); + const auto expected = BLOCK_M * + (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != expected) {} + } else if constexpr (BlockPhaseTag::value == sched::BlockPhase::Linear2) { + const auto ptr = + workspace.get_l2_full_count_ptr(ring_block_idx); + const auto expected = (L1_SHAPE_N / BLOCK_N) * + (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != expected) {} + } else if constexpr ( + BlockPhaseTag::value == sched::BlockPhase::SharedLinear2) { + const auto ptr = workspace.get_shared_l2_full_count_ptr( + pool_block_idx); + constexpr uint32_t kNumSharedL1BlockNs = + (kIntermediateHidden * kNumSharedExperts * 2) / + BLOCK_N; + while (ptx::ld_acq(ptr) != kNumSharedL1BlockNs) {} + } + } + 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()) { + if (has_valid_m) { + const uint32_t m_idx = block_idx * BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + #pragma unroll + for (uint32_t k_tile = 0; + k_tile < kNumTMATilesPerStage; ++ k_tile) { + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], + smem_a[stage_idx] + + k_tile * LOAD_BLOCK_M * kTMATileK, + k_idx + k_tile * kTMATileK, m_idx, 1); + } + + // TMA load SFA with A on the same producer warp. + if (is_linear1_phase) { + // L1 SFA per-128: one vector per BK128 plane. + #pragma unroll + for (uint32_t sf_group = 0; + sf_group < BLOCK_K / kGranK; ++ sf_group) { + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx] + + sf_group * kL2SFAHalfStride, + m_idx, + k_block_idx * (BLOCK_K / kGranK) + sf_group, + 1); + } + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + + (BLOCK_K / kGranK) * BLOCK_M * sizeof(float)); + } else { + // L2 SFA per-64: one TMA per scale group. + #pragma unroll + for (uint32_t sf_group = 0; + sf_group < kNumL2SFAKGroups; ++ sf_group) { + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx] + + sf_group * kL2SFAHalfStride, + m_idx, + k_block_idx * kNumL2SFAKGroups + sf_group, + 1); + } + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + + kNumL2SFAKGroups * BLOCK_M * sizeof(float)); + } + } else { + full_barriers[stage_idx]->arrive(); + } + } + __syncwarp(); + } + }); + + } else if (warp_idx == kNumDispatchWarps + 1) { + if constexpr (kCacheProExpertCounts) { + ptx::sync_unaligned( + kNumDispatchThreads + 32, kExpertCountCacheBarrierIdx); + scheduler.fetch_cached_expert_recv_count(smem_expert_count); + } + const auto load_b_task = [&](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& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t< + std::remove_reference_t>; + constexpr bool is_linear1_phase = + BlockPhaseTag::value == sched::BlockPhase::Linear1 or + BlockPhaseTag::value == sched::BlockPhase::SharedLinear1; + constexpr bool is_shared_phase = + BlockPhaseTag::value == sched::BlockPhase::SharedLinear1 or + BlockPhaseTag::value == sched::BlockPhase::SharedLinear2; + sync_task_storage_alias(is_shared_phase); + constexpr bool use_mxfp4_task = not is_shared_phase; + const auto tensor_map_b_ptr = is_shared_phase ? + (is_linear1_phase ? &tensor_map_shared_l1_weights : + &tensor_map_shared_l2_weights) : + (is_linear1_phase ? &tensor_map_l1_weights : + &tensor_map_l2_weights); + + constexpr uint32_t kSharedL1ShapeN = + kIntermediateHidden * kNumSharedExperts * 2; + constexpr uint32_t kSharedL2ShapeN = kHidden; + const uint32_t shape_n = is_shared_phase ? + (is_linear1_phase ? kSharedL1ShapeN : kSharedL2ShapeN) : + (is_linear1_phase ? L1_SHAPE_N : L2_SHAPE_N); + + constexpr bool kCoalescedMXFP4WeightSF = + kHidden <= kMXFP4CoalescedScaleMaxHidden; + const uint32_t weight_sf_stride_k = is_linear1_phase ? + kL1MXFP4WeightSFStrideK : kL2MXFP4WeightSFStrideK; + const uint32_t weight_sf_per_expert = is_linear1_phase ? + kL1MXFP4WeightSFPerExpert : kL2MXFP4WeightSFPerExpert; + + 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 constexpr (is_shared_phase) + shared_b_empty_barriers[shared_b_stage_idx].wait( + shared_b_phase ^ 1); + + const bool elected = cute::elect_one_sync(); + const uint32_t local_n_idx = n_block_idx * BLOCK_N; + if (elected) { + const uint32_t n_idx = is_shared_phase ? local_n_idx : + local_expert_idx * shape_n + local_n_idx; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + if constexpr (use_mxfp4_task) { + // Load the packed storage as ordinary bytes. Using a + // UINT8 descriptor over K/2 bytes preserves the + // existing SM90 CUDA baseline and avoids relying on + // newer packed-FP4 TMA layout semantics. + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], + smem_b_packed[stage_idx], k_idx / 2, n_idx, 1); + } else { + #pragma unroll + for (uint32_t k_tile = 0; + k_tile < kNumTMATilesPerStage; ++ k_tile) { + #pragma unroll + for (uint32_t n_tile = 0; + n_tile < kNumTMANTilesPerStage; ++ n_tile) { + tma::copy( + tensor_map_b_ptr, + full_barriers[stage_idx], + (is_shared_phase ? + smem_shared_b[shared_b_stage_idx] : + smem_b[stage_idx]) + + k_tile * LOAD_BLOCK_N * kTMATileK + + n_tile * kTMATileN * kTMATileK, + k_idx + k_tile * kTMATileK, + n_idx + n_tile * kTMATileN, + 1); + } + } + } + } + + if constexpr (use_mxfp4_task) { + // A natural-layout TMA box would have only four contiguous + // bytes per row, below Hopper's 16-byte requirement. Use + // the B producer warp to issue four strided row groups and + // publish them before its transaction-barrier arrival. + const auto* weight_sf_base = (is_linear1_phase ? + l1_mxfp4_weights_sf : l2_mxfp4_weights_sf) + + local_expert_idx * weight_sf_per_expert; + if constexpr (kCoalescedMXFP4WeightSF) { + // The preprocessed layout makes all 128 K128 scale + // words contiguous in N. Coalesce the producer's four + // scalar transfers into one 16-byte transaction per + // lane for every routed bucket. + const auto* weight_sf_words = + reinterpret_cast(weight_sf_base) + + k_block_idx * shape_n + local_n_idx; + const uint4 scale_words = __ldg( + reinterpret_cast(weight_sf_words) + + lane_idx); + ptx::st_shared( + reinterpret_cast(smem_sfb[stage_idx]) + + lane_idx, + scale_words.x, scale_words.y, + scale_words.z, scale_words.w); + } else { + #pragma unroll + for (uint32_t local_n = lane_idx; local_n < BLOCK_N; + local_n += 32) { + uint32_t scale_word; + if constexpr (kCoalescedMXFP4WeightSF) { + const auto* weight_sf_words = + reinterpret_cast( + weight_sf_base) + + k_block_idx * shape_n; + scale_word = __ldg( + weight_sf_words + local_n_idx + local_n); + } else { + scale_word = __ldg( + reinterpret_cast( + weight_sf_base + + (local_n_idx + local_n) * + weight_sf_stride_k + + k_block_idx * + kNumMXFP4SFBKGroups)); + } + ptx::st_shared( + reinterpret_cast( + smem_sfb[stage_idx]) + local_n, + scale_word); + } + } + __syncwarp(); + } + + if (elected) { + full_barriers[stage_idx]->arrive_and_expect_tx( + use_mxfp4_task ? SMEM_B_PACKED_SIZE_PER_STAGE : + SMEM_B_SIZE_PER_STAGE); + } + __syncwarp(); + if constexpr (is_shared_phase) + advance_shared_b_pipeline(); + } + }; + produce_selected_blocks(load_b_task); + + } else { + // ===================================================================== + // ROLE 3: MATH WARPGROUPS (WGMMA + epilogue + combine) + // ===================================================================== + const uint32_t epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const uint32_t epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const uint32_t warp_idx_in_wg = epilogue_warp_idx; + + const auto arrive_empty_barrier = [&](const uint32_t& s) { + if (lane_idx == 0) + empty_barriers[s]->arrive(); + }; + + // WGMMA-output register layout helpers + constexpr uint32_t WG_SMEM_CD_L1_STRIDE_N = WG_L1_OUT_BLOCK_N; + DG_STATIC_ASSERT(WG_BLOCK_M == L1WGMMA::M and + WG_BLOCK_N == L1WGMMA::N, + "The Humming warpgroup owns one complete WGMMA tile"); + + // With shared experts, dispatch and shared L1 are independent and can + // start concurrently. Their end-of-kernel rendezvous remains paired. + if constexpr (not kHasSharedExperts) + ptx::sync_unaligned( + kNumDispatchThreads + kNumEpilogueThreads, + kDispatchWithEpilogueBarrierIdx); + + for_each_selected_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& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t< + std::remove_reference_t>; + constexpr bool is_linear1_phase = + BlockPhaseTag::value == sched::BlockPhase::Linear1 or + BlockPhaseTag::value == sched::BlockPhase::SharedLinear1; + constexpr bool is_shared_phase = + BlockPhaseTag::value == sched::BlockPhase::SharedLinear1 or + BlockPhaseTag::value == sched::BlockPhase::SharedLinear2; + sync_task_storage_alias(is_shared_phase); + scheduler.release_task_info(); + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t block_idx = is_shared_phase ? pool_block_idx : + ring_block_idx; + const uint32_t m_idx = block_idx * BLOCK_M; + const uint32_t pool_m_idx = pool_block_idx * BLOCK_M; + const uint32_t n_idx = n_block_idx * BLOCK_N; + const auto arrive_task_empty_barrier = [&](const uint32_t& s) { + arrive_empty_barrier(s); + if constexpr (is_shared_phase) { + if (lane_idx == 0) + shared_b_empty_barriers[shared_b_stage_idx].arrive(); + advance_shared_b_pipeline(); + } + }; + // ---------------- GEMM ---------------- + using WGMMA = L1WGMMA; + constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; // 64 for M=64,N=128 + float final_accum[ + kPackedBF16SwapEpilogue ? 1 : kAccumPerThread]; + if constexpr (is_shared_phase) { + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) + final_accum[i] = 0.0f; + } + // Swap-AB consumes one N64 weight half at a time. Reuse the same + // fragment for the second half after promotion instead of keeping + // both halves live across WGMMA. The regular orientation still + // needs the complete M64xN128 fragment. + constexpr bool kReuseSwapABFragment = kSmallMSwapAB; + constexpr uint32_t kAccumStorage = kReuseSwapABFragment ? + kSwapABHalfAccumPerThread : kAccumPerThread; + float accum[kAccumStorage]; + nv_bfloat162 mxfp4_final_bf16[kAccumPerThread / 2]; + + const auto run_mxfp4_gemm_loop = [&]() { + { + constexpr uint32_t kWeightGranK = kMXFP4WeightGranK; + constexpr uint32_t kWGThreads = 128; + const uint32_t wg_thread_idx = warp_idx_in_wg * 32 + lane_idx; + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 2; ++ i) + mxfp4_final_bf16[i] = + __float2bfloat162_rn(0.0f); + + const auto prepare_stage_weights = [&]( + const uint32_t pipeline_stage, + const uint32_t expanded_slot) { + const auto* packed = smem_b_packed[pipeline_stage]; + auto* expanded = reinterpret_cast( + smem_b_expanded[expanded_slot]); + + const uint32_t local_n = wg_thread_idx; + const uint32_t scale_word = ptx::ld_shared( + reinterpret_cast( + smem_sfb[pipeline_stage]) + local_n); + { + DG_STATIC_ASSERT(WG_BLOCK_N == kWGThreads, + "MXFP4 assigns one N row per WG thread"); + constexpr uint32_t kPackedWordsPerK32 = 32 / 8; + constexpr uint32_t kRowsPerDecodeGroup = 8; + DG_STATIC_ASSERT( + kPackedWordsPerK32 == 4 and + 32 % kRowsPerDecodeGroup == 0, + "MXFP4 warp mapping requires 8 rows x 4 words"); + + constexpr bool kPairPackedWords = + kOverlapMXFP4ScalePath; + if constexpr (kPairPackedWords) { + // Routed small-M gives each lane one complete + // packed row. Both adjacent word pairs share + // the row's exponent lookup, and the lane owns + // the scale directly instead of shuffling it + // twice from the two-lane-per-row mapping. + constexpr bool kFullRowPairDecode = + kSmallMSwapAB; + if constexpr (kFullRowPairDecode) { + const uint32_t decoded_local_n = local_n; + const uint32_t packed_row_base = + decoded_local_n * (BLOCK_K / 2); + const uint32_t packed_row_xor = + cute::Swizzle<2, 4, 3>::apply( + packed_row_base) ^ packed_row_base; + const uint32_t first_pair_in_k32 = + __byte_perm( + 0x00020200u, 0u, lane_idx >> 3); + uint2 packed_current[2]; + #pragma unroll + for (uint32_t pair_idx = 0; + pair_idx < 2; ++ pair_idx) { + const uint32_t pair_in_k32 = + first_pair_in_k32 ^ (pair_idx * 2u); + const uint32_t packed_byte_offset = + packed_row_base + + ((pair_in_k32 * sizeof(uint32_t)) ^ + packed_row_xor); + packed_current[pair_idx] = + ptx::ld_shared( + reinterpret_cast( + packed + + packed_byte_offset)); + } + #pragma unroll + for (uint32_t k32_idx = 0; + k32_idx < kNumMXFP4SFBKGroups; + ++ k32_idx) { + uint2 packed_next[2] = { + make_uint2(0u, 0u), + make_uint2(0u, 0u)}; + if (k32_idx + 1 < + kNumMXFP4SFBKGroups) { + #pragma unroll + for (uint32_t pair_idx = 0; + pair_idx < 2; ++ pair_idx) { + const uint32_t pair_in_k32 = + first_pair_in_k32 ^ + (pair_idx * 2u); + const uint32_t next_packed_pair = + (k32_idx + 1) * + kPackedWordsPerK32 + + pair_in_k32; + const uint32_t + next_packed_byte_offset = + packed_row_base + + ((next_packed_pair * + sizeof(uint32_t)) ^ + packed_row_xor); + packed_next[pair_idx] = + ptx::ld_shared( + reinterpret_cast< + const uint2*>( + packed + + next_packed_byte_offset)); + } + } + const uint32_t exponent_offset = + sm90_extract_u8_prmt( + scale_word, k32_idx); + const uint2 lookup = + sm90_mxfp4_e4m3_lookup( + exponent_offset); + #pragma unroll + for (uint32_t pair_idx = 0; + pair_idx < 2; ++ pair_idx) { + const uint2 decoded0 = + sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + packed_current[pair_idx].x, + lookup); + const uint2 decoded1 = + sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + packed_current[pair_idx].y, + lookup); + const uint32_t pair_in_k32 = + first_pair_in_k32 ^ + (pair_idx * 2u); + const uint32_t packed_k_pair = + k32_idx * + kPackedWordsPerK32 + + pair_in_k32; + const uint32_t flat = + decoded_local_n * BLOCK_K + + packed_k_pair * 8u; + const uint32_t swizzled = + cute::Swizzle<3, 4, 3>::apply( + flat); + ptx::st_shared( + expanded + swizzled, + decoded0.x, decoded0.y, + decoded1.x, decoded1.y); + packed_current[pair_idx] = + packed_next[pair_idx]; + } + } + } else { + constexpr uint32_t kPairRowsPerDecodeGroup = 16; + const uint32_t pair_row_in_decode_group = + lane_idx % 16; + // For the conflict-heavy Flash and selected + // Pro swap-AB/regular buckets, + // alternate the two adjacent word pairs across + // each half warp's lower and upper eight rows. + // Under the packed B64 swizzle this covers every + // bank once per LDS.64 wave instead of aliasing + // rows r and r+8 onto the same bank pair. The + // full warp still owns the identical 16-row x + // 4-word address set. + constexpr bool kBankPermutedPairLoads = + kBankPermuteMXFP4PairLoads; + constexpr bool kPrmtPairLoads = + kSmallMSwapAB and + ((kHidden == 4096 and + (kMaxSwapABTokens == 8 or + kMaxSwapABTokens == 16 or + kMaxSwapABTokens == 32 or + kMaxSwapABTokens == 64)) or + (kHidden == 7168 and + (kMaxSwapABTokens == 8 or + kMaxSwapABTokens == 32 or + kMaxSwapABTokens == 64))); + const uint32_t packed_k_pair_in_k32 = + kBankPermutedPairLoads ? + (kPrmtPairLoads ? + __byte_perm( + 0x00020200u, 0u, + lane_idx >> 3) : + (((lane_idx >> 3) ^ + (lane_idx >> 4)) & 1u) * 2u) : + (lane_idx / 16) * 2; + #pragma unroll + for (uint32_t row_group = 0; + row_group < + 32 / kPairRowsPerDecodeGroup; + ++ row_group) { + const uint32_t row_in_warp = + row_group * kPairRowsPerDecodeGroup + + pair_row_in_decode_group; + const uint32_t decoded_local_n = + warp_idx_in_wg * 32 + row_in_warp; + const uint32_t decoded_scale_word = + __shfl_sync( + 0xffffffffu, scale_word, row_in_warp); + const uint32_t packed_row_base = + decoded_local_n * (BLOCK_K / 2); + const uint32_t packed_row_xor = + cute::Swizzle<2, 4, 3>::apply( + packed_row_base) ^ packed_row_base; + const uint32_t first_packed_byte_offset = + packed_row_base + + ((packed_k_pair_in_k32 * + sizeof(uint32_t)) ^ packed_row_xor); + uint2 packed_current = ptx::ld_shared( + reinterpret_cast( + packed + first_packed_byte_offset)); + #pragma unroll + for (uint32_t k32_idx = 0; + k32_idx < kNumMXFP4SFBKGroups; + ++ k32_idx) { + const uint32_t packed_k_pair = + k32_idx * kPackedWordsPerK32 + + packed_k_pair_in_k32; + uint2 packed_next{0u, 0u}; + if (k32_idx + 1 < + kNumMXFP4SFBKGroups) { + const uint32_t next_packed_k_pair = + packed_k_pair + + kPackedWordsPerK32; + const uint32_t next_packed_byte_offset = + packed_row_base + + ((next_packed_k_pair * + sizeof(uint32_t)) ^ + packed_row_xor); + packed_next = ptx::ld_shared( + reinterpret_cast( + packed + + next_packed_byte_offset)); + } + const uint32_t exponent_offset = [&]() { + if constexpr (kUsePRMTMXFP4Exponent) { + return sm90_extract_u8_prmt( + decoded_scale_word, k32_idx); + } else { + return (decoded_scale_word >> + (k32_idx * 8u)) & 0xffu; + } + }(); + uint4 decoded; + if constexpr (kHidden == 4096) { + decoded = + sm90_mxfp4_reordered_signs_e2m1x16_to_e4m3x16_bits( + packed_current.x, + packed_current.y, + exponent_offset); + } else { + const uint2 lookup = + sm90_mxfp4_e4m3_lookup( + exponent_offset); + const uint2 decoded0 = + sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + packed_current.x, lookup); + const uint2 decoded1 = + sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + packed_current.y, lookup); + decoded = make_uint4( + decoded0.x, decoded0.y, + decoded1.x, decoded1.y); + } + const uint32_t logical_k0 = + packed_k_pair * 8; + const uint32_t flat0 = + decoded_local_n * BLOCK_K + + logical_k0; + const uint32_t swizzled0 = + cute::Swizzle<3, 4, 3>::apply(flat0); + ptx::st_shared( + expanded + swizzled0, + decoded.x, decoded.y, + decoded.z, decoded.w); + packed_current = packed_next; + } + } + } + } else { + // Each half warp covers the same eight rows and + // two disjoint words. This distributes B64 loads + // over all 32 banks and also gives each half-warp + // B128 STS.64 transaction one bank-word per bank. + const uint32_t lane_in_half_warp = lane_idx % 16; + const uint32_t row_in_decode_group = + lane_in_half_warp / 2; + const uint32_t packed_k_in_k32 = + (lane_idx / 16) * 2 + lane_in_half_warp % 2; + #pragma unroll + for (uint32_t row_group = 0; + row_group < 32 / kRowsPerDecodeGroup; ++ row_group) { + const uint32_t row_in_warp = + row_group * kRowsPerDecodeGroup + + row_in_decode_group; + const uint32_t decoded_local_n = + warp_idx_in_wg * 32 + row_in_warp; + const uint32_t decoded_scale_word = __shfl_sync( + 0xffffffffu, scale_word, row_in_warp); + const uint32_t packed_row_base = + decoded_local_n * (BLOCK_K / 2); + const uint32_t packed_row_xor = + cute::Swizzle<2, 4, 3>::apply(packed_row_base) ^ + packed_row_base; + + // Reuse the validated two-word LDS lookahead + // in every overlap-enabled routed phase. + // Decoder temporaries die before the first + // QGMMA, so they do not cross accumulator + // lifetime like rejected next-stage overlap. + constexpr bool kPipelinePackedLDS = + kOverlapMXFP4ScalePath; + if constexpr (kPipelinePackedLDS) { + const uint32_t first_packed_byte_offset = + packed_row_base + + ((packed_k_in_k32 * sizeof(uint32_t)) ^ + packed_row_xor); + uint32_t packed_current = ptx::ld_shared( + reinterpret_cast( + packed + first_packed_byte_offset)); + #pragma unroll + for (uint32_t k32_idx = 0; + k32_idx < kNumMXFP4SFBKGroups; + ++ k32_idx) { + const uint32_t packed_k = + k32_idx * kPackedWordsPerK32 + + packed_k_in_k32; + uint32_t packed_next = 0; + if (k32_idx + 1 < + kNumMXFP4SFBKGroups) { + const uint32_t next_packed_k = + packed_k + kPackedWordsPerK32; + const uint32_t next_packed_byte_offset = + packed_row_base + + ((next_packed_k * sizeof(uint32_t)) ^ + packed_row_xor); + packed_next = ptx::ld_shared( + reinterpret_cast( + packed + + next_packed_byte_offset)); + } + const uint32_t exponent_offset = [&]() { + if constexpr (kUsePRMTMXFP4Exponent) { + return sm90_extract_u8_prmt( + decoded_scale_word, k32_idx); + } else { + return (decoded_scale_word >> + (k32_idx * 8u)) & 0xffu; + } + }(); + const uint2 decoded = + sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + packed_current, + exponent_offset); + const uint32_t logical_n = + decoded_local_n; + const uint32_t logical_k0 = + packed_k * 8; + const uint32_t flat0 = + logical_n * BLOCK_K + logical_k0; + const uint32_t swizzled0 = + cute::Swizzle<3, 4, 3>::apply(flat0); + ptx::st_shared( + expanded + swizzled0, + decoded.x, decoded.y); + packed_current = packed_next; + } + } else { + #pragma unroll + for (uint32_t k32_idx = 0; + k32_idx < kNumMXFP4SFBKGroups; + ++ k32_idx) { + const uint32_t packed_k = + k32_idx * kPackedWordsPerK32 + + packed_k_in_k32; + const uint32_t exponent_offset = [&]() { + if constexpr (kUsePRMTMXFP4Exponent) { + return sm90_extract_u8_prmt( + decoded_scale_word, k32_idx); + } else { + return (decoded_scale_word >> + (k32_idx * 8u)) & 0xffu; + } + }(); + const uint32_t packed_byte_offset = + packed_row_base + + ((packed_k * sizeof(uint32_t)) ^ + packed_row_xor); + const uint2 decoded = + sm90_mxfp4_reordered_signs_e2m1x8_to_e4m3x8_bits( + ptx::ld_shared( + reinterpret_cast( + packed + + packed_byte_offset)), + exponent_offset); + const uint32_t logical_n = + decoded_local_n; + const uint32_t logical_k0 = packed_k * 8; + const uint32_t flat0 = + logical_n * BLOCK_K + logical_k0; + const uint32_t swizzled0 = + cute::Swizzle<3, 4, 3>::apply(flat0); + ptx::st_shared( + expanded + swizzled0, + decoded.x, decoded.y); + } + } + } + } + } + // Generic shared stores are not ordered with WGMMA's + // async proxy by a named barrier alone. Publish every + // decoded E4M3 byte before the warpgroup consumes it. + cutlass::arch::fence_view_async_shared(); + ptx::sync_aligned( + kWGThreads, kEpilogueWGBarrierStartIdx); + }; + + // Producer-staged SF uses source-order release for short + // phase-K loops. + // For longer loops ptxas may still advance the loop-bottom + // arrival once the final stage read and WGMMA wait finish; + // no machine-level order relative to register-only HFMA2 + // is required for correctness. + constexpr uint32_t kPhaseKBlocks = is_linear1_phase ? + kHidden / BLOCK_K : kIntermediateHidden / BLOCK_K; + constexpr bool kEarlyReleaseMXFP4Stage = + kPhaseKBlocks <= 16; + const auto issue_mxfp4_wgmma = [&]( + const uint32_t pipeline_stage, + const uint32_t expanded_slot) { + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + if constexpr (kUseIncrementalMXFP4Descriptor) { + const auto desc_a_base = mma::sm90::make_smem_desc( + smem_a[pipeline_stage] + + kStartK32 * kWeightGranK, 1); + const auto desc_b_base = mma::sm90::make_smem_desc( + smem_b_expanded[expanded_slot] + + kStartK32 * kWeightGranK, 1); + #pragma unroll + for (uint32_t k = 0; k < kNumWGMMAs; ++ k) { + // K32 advances the descriptor's 16-byte start + // address field by two without changing layout. + const cute::GmmaDescriptor desc_a( + desc_a_base.desc_ + k * 2u); + const cute::GmmaDescriptor desc_b( + desc_b_base.desc_ + k * 2u); + WGMMA::wgmma( + desc_a, desc_b, accum, + k != 0); + } + } else { + #pragma unroll + for (uint32_t k = 0; k < kNumWGMMAs; ++ k) { + const uint32_t k32_idx = kStartK32 + k; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[pipeline_stage] + + k32_idx * kWeightGranK, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b_expanded[expanded_slot] + + k32_idx * kWeightGranK, 1); + WGMMA::wgmma( + desc_a, desc_b, accum, + k != 0); + } + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + if constexpr (not kOverlapMXFP4ScalePath) + ptx::warpgroup_wait<0>(); + }; + + const auto promote_mxfp4 = [&]( + const uint32_t pipeline_stage, + const uint32_t activation_sf_group, + const float secondary) { + // Rematerialize the cheap row offset at the SFA load. + // Keeping it live across decode/WGMMA spills it to the + // local stack once per stage on the 128-register build. + const uint32_t stage_row_offset_r0 = + warp_idx_in_wg * 16 + lane_idx / 4; + const float scale_a_0 = ptx::ld_shared( + smem_sfa[pipeline_stage] + + activation_sf_group * kL2SFAHalfStride + + stage_row_offset_r0); + const float scale_a_1 = ptx::ld_shared( + smem_sfa[pipeline_stage] + + activation_sf_group * kL2SFAHalfStride + + stage_row_offset_r0 + 8); + // Applying the E2M1-to-E4M3 bias correction after + // scale_a * secondary can underflow for valid UE8M0 + // codes 0/1. Apply x64 to the secondary first except + // at the high endpoint where that product would + // overflow; the branch is uniform for the expert. + constexpr float kMaxSecondaryBeforeX64 = 0x1p121f; + const bool compensate_secondary = + secondary <= kMaxSecondaryBeforeX64; + const float compensated_secondary = + compensate_secondary ? secondary * 64.0f : secondary; + const float compensated_scale_a_0 = + compensate_secondary ? scale_a_0 : scale_a_0 * 64.0f; + const float compensated_scale_a_1 = + compensate_secondary ? scale_a_1 : scale_a_1 * 64.0f; + const float combined_scale_0 = + compensated_scale_a_0 * compensated_secondary; + const float combined_scale_1 = + compensated_scale_a_1 * compensated_secondary; + // The scale path is independent of the WGMMA result. + // Keep the group in flight while loading SFA and + // preparing the two promotion multipliers, then wait + // immediately before the first accumulator access. + if constexpr (kFastMath) { + const nv_bfloat162 combined_scale_bf16_0 = + __float2bfloat162_rn(combined_scale_0); + const nv_bfloat162 combined_scale_bf16_1 = + __float2bfloat162_rn(combined_scale_1); + if constexpr (kOverlapMXFP4ScalePath) + ptx::warpgroup_wait<0>(); + // The final wait ends every shared-memory access + // to this stage. Release it before register-only + // accumulator promotion so the producer can start + // refilling the stage in parallel. + if constexpr (kReleaseStage) { + if constexpr (kOverlapMXFP4ScalePath) + arrive_task_empty_barrier(pipeline_stage); + } + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + mxfp4_final_bf16[i * 2] = __hfma2( + combined_scale_bf16_0, + __floats2bfloat162_rn( + accum[i * 4], accum[i * 4 + 1]), + mxfp4_final_bf16[i * 2]); + mxfp4_final_bf16[i * 2 + 1] = __hfma2( + combined_scale_bf16_1, + __floats2bfloat162_rn( + accum[i * 4 + 2], accum[i * 4 + 3]), + mxfp4_final_bf16[i * 2 + 1]); + } + } else { + if constexpr (kOverlapMXFP4ScalePath) + ptx::warpgroup_wait<0>(); + if constexpr (kReleaseStage) { + if constexpr (kOverlapMXFP4ScalePath) + arrive_task_empty_barrier(pipeline_stage); + } + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float2 persistent_0 = + __bfloat1622float2(mxfp4_final_bf16[i * 2]); + const float2 persistent_1 = + __bfloat1622float2(mxfp4_final_bf16[i * 2 + 1]); + mxfp4_final_bf16[i * 2] = + __floats2bfloat162_rn( + fmaf(combined_scale_0, + accum[i * 4], persistent_0.x), + fmaf(combined_scale_0, + accum[i * 4 + 1], persistent_0.y)); + mxfp4_final_bf16[i * 2 + 1] = + __floats2bfloat162_rn( + fmaf(combined_scale_1, + accum[i * 4 + 2], persistent_1.x), + fmaf(combined_scale_1, + accum[i * 4 + 3], persistent_1.y)); + } + } + }; + + const float mxfp4_secondary = __ldg( + (is_linear1_phase ? l1_mxfp4_secondary : l2_mxfp4_secondary) + + local_expert_idx); + if constexpr (kSmallMSwapAB) { + // The regular M64xN128 orientation wastes nearly all + // tensor-core work when each routed expert owns only a + // handful of tokens. Use each N64 weight half as the + // WGMMA M operand and bucket valid tokens along N. + // Unlike the retired per-tensor path, every K128/K64 + // group is promoted with the token's staged activation + // scale before the temporary accumulator is reused. + auto run_swap_ab = [&]() { + using SwapWGMMA = typename + mma::sm90::FP8MMASelector::type; + constexpr uint32_t kSwapAccum = + SwapWGMMA::kNumAccum; + DG_STATIC_ASSERT( + kSwapAccum <= kSwapABHalfAccumPerThread, + "Invalid swap-AB accumulator bucket"); + // Selected compact Pro buckets can keep two + // bucket-sized fragments inside the existing + // 32-float allocation. Commit them as separate + // WGMMA groups so promoting half 0 can overlap + // execution of half 1. Packed-BF16 buckets and N64 + // keep reusing one fragment sequentially. + constexpr bool kPipelineWeightHalves = + kReuseSwapABFragment and + not kPackedBF16SwapEpilogue and + 2 * kSwapAccum <= kAccumStorage; + constexpr uint32_t kWeightHalfAccumStride = + kPipelineWeightHalves ? kSwapAccum : + kSwapABHalfAccumPerThread; + DG_STATIC_ASSERT( + not kPipelineWeightHalves or + 2 * kWeightHalfAccumStride <= kAccumStorage, + "Pipelined swap-AB fragments exceed storage"); + const uint32_t swap_col_idx = lane_idx % 4; + + const auto issue_swap_wgmma = [&]< + uint32_t kFirstWeightHalf, + uint32_t kNumWeightHalves, + uint32_t kStartK32, + uint32_t kNumWGMMAs>( + const uint32_t pipeline_stage, + const uint32_t expanded_slot) { + DG_STATIC_ASSERT( + kNumWeightHalves > 0 and + kFirstWeightHalf + kNumWeightHalves <= + kSwapABWeightHalves, + "Invalid swap-AB weight-half range"); + #pragma unroll + for (uint32_t half_idx = 0; + half_idx < kNumWeightHalves; ++ half_idx) { + const uint32_t weight_half = + kFirstWeightHalf + half_idx; + auto* wgmma_accum = accum + + (kNumWeightHalves == 1 and + not kPipelineWeightHalves ? + 0u : + weight_half * + kWeightHalfAccumStride); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand( + wgmma_accum[i]); + } + ptx::warpgroup_arrive(); + const auto desc_b_base = + mma::sm90::make_smem_desc( + smem_a[pipeline_stage], 1); + #pragma unroll + for (uint32_t half_idx = 0; + half_idx < kNumWeightHalves; ++ half_idx) { + const uint32_t weight_half = + kFirstWeightHalf + half_idx; + auto* wgmma_accum = accum + + (kNumWeightHalves == 1 and + not kPipelineWeightHalves ? + 0u : + weight_half * + kWeightHalfAccumStride); + const auto desc_a_base = + mma::sm90::make_smem_desc( + smem_b_expanded[expanded_slot] + + weight_half * 64u * BLOCK_K, + 1); + #pragma unroll + for (uint32_t k = 0; + k < kNumWGMMAs; ++ k) { + const uint32_t k32_idx = kStartK32 + k; + const cute::GmmaDescriptor desc_a( + desc_a_base.desc_ + k32_idx * 2u); + const cute::GmmaDescriptor desc_b( + desc_b_base.desc_ + k32_idx * 2u); + SwapWGMMA::wgmma( + desc_a, desc_b, wgmma_accum, + k != 0); + } + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t half_idx = 0; + half_idx < kNumWeightHalves; ++ half_idx) { + const uint32_t weight_half = + kFirstWeightHalf + half_idx; + auto* wgmma_accum = accum + + (kNumWeightHalves == 1 and + not kPipelineWeightHalves ? + 0u : + weight_half * + kWeightHalfAccumStride); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand( + wgmma_accum[i]); + } + }; + + const auto promote_swap_ab = [&]< + uint32_t kFirstWeightHalf, + uint32_t kNumWeightHalves, + uint32_t kWaitGroups, + bool kReleaseStage>( + const uint32_t pipeline_stage, + const uint32_t activation_sf_group) { + DG_STATIC_ASSERT( + kNumWeightHalves > 0 and + kFirstWeightHalf + kNumWeightHalves <= + kSwapABWeightHalves, + "Invalid swap-AB weight-half range"); + constexpr float kMaxSecondaryBeforeX64 = + 0x1p121f; + const bool compensate_secondary = + mxfp4_secondary <= + kMaxSecondaryBeforeX64; + const float compensated_secondary = + compensate_secondary ? + mxfp4_secondary * 64.0f : + mxfp4_secondary; + ptx::warpgroup_wait(); + + #pragma unroll + for (uint32_t chunk = 0; + chunk < kSwapAccum / 4; ++ chunk) { + const uint32_t token_0 = + chunk * 8 + swap_col_idx * 2; + const uint32_t token_1 = token_0 + 1; + const float scale_a_0 = + token_0 < valid_m ? + ptx::ld_shared( + smem_sfa[pipeline_stage] + + activation_sf_group * + kL2SFAHalfStride + + token_0) : + 0.0f; + const float scale_a_1 = + token_1 < valid_m ? + ptx::ld_shared( + smem_sfa[pipeline_stage] + + activation_sf_group * + kL2SFAHalfStride + + token_1) : + 0.0f; + const float combined_scale_0 = + (compensate_secondary ? + scale_a_0 : scale_a_0 * 64.0f) * + compensated_secondary; + const float combined_scale_1 = + (compensate_secondary ? + scale_a_1 : scale_a_1 * 64.0f) * + compensated_secondary; + #pragma unroll + for (uint32_t half_idx = 0; + half_idx < kNumWeightHalves; + ++ half_idx) { + const uint32_t weight_half = + kFirstWeightHalf + half_idx; + const uint32_t accum_offset = + (kNumWeightHalves == 1 and + not kPipelineWeightHalves ? + 0u : + weight_half * + kWeightHalfAccumStride) + + chunk * 4; + const uint32_t pair_offset = + weight_half * + (kSwapABHalfAccumPerThread / 2) + + chunk * 2; + // Matched cold-L2 A/B selects the + // endpoint buckets for both routed + // DSV4 models and Flash M16/M32's joint + // PRMT/HFMA2 plus packed-epilogue path. + // Pro M16/M32 retain scalar FP32 + // promotion. + if constexpr ( + kFastMath and + (kHidden == 4096 or + kHidden == 7168) and + kSmallMSwapAB and + (kMaxSwapABTokens == 8 or + kMaxSwapABTokens == 16 or + (kHidden == 4096 and + kMaxSwapABTokens == 32 and + kPackedBF16SwapEpilogue) or + kMaxSwapABTokens == 64)) { + const nv_bfloat162 scale_pair = + __floats2bfloat162_rn( + combined_scale_0, + combined_scale_1); + mxfp4_final_bf16[pair_offset] = + __hfma2( + scale_pair, + __floats2bfloat162_rn( + accum[accum_offset], + accum[accum_offset + 1]), + mxfp4_final_bf16[ + pair_offset]); + mxfp4_final_bf16[pair_offset + 1] = + __hfma2( + scale_pair, + __floats2bfloat162_rn( + accum[accum_offset + 2], + accum[accum_offset + 3]), + mxfp4_final_bf16[ + pair_offset + 1]); + } else { + const float2 persistent_0 = + __bfloat1622float2( + mxfp4_final_bf16[ + pair_offset]); + const float2 persistent_1 = + __bfloat1622float2( + mxfp4_final_bf16[ + pair_offset + 1]); + mxfp4_final_bf16[pair_offset] = + __floats2bfloat162_rn( + fmaf(combined_scale_0, + accum[accum_offset], + persistent_0.x), + fmaf(combined_scale_1, + accum[accum_offset + 1], + persistent_0.y)); + mxfp4_final_bf16[pair_offset + 1] = + __floats2bfloat162_rn( + fmaf(combined_scale_0, + accum[accum_offset + 2], + persistent_1.x), + fmaf(combined_scale_1, + accum[accum_offset + 3], + persistent_1.y)); + } + } + } + // SFA belongs to the same producer stage as A. + // Release it only after every token scale has + // been consumed; otherwise the loader can + // overwrite a later chunk during promotion. + if constexpr (kReleaseStage) + arrive_task_empty_barrier(pipeline_stage); + }; + + for (uint32_t k_block_idx = 0; + k_block_idx < num_k_blocks; + advance_pipeline(k_block_idx)) { + const uint32_t expanded_slot = + kPipelineMXFP4ExpandedB ? + (k_block_idx & 1u) : 0u; + if constexpr (kPipelineMXFP4ExpandedB) { + if (k_block_idx == 0) { + full_barriers[stage_idx]->wait(phase); + prepare_stage_weights( + stage_idx, expanded_slot); + } + } else { + full_barriers[stage_idx]->wait(phase); + prepare_stage_weights( + stage_idx, expanded_slot); + } + + if constexpr (kPipelineWeightHalves) { + if constexpr (is_linear1_phase) { + issue_swap_wgmma.template operator()< + 0, 1, 0, 4>( + stage_idx, expanded_slot); + issue_swap_wgmma.template operator()< + 1, 1, 0, 4>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 0, 1, 1, false>(stage_idx, 0); + } else { + issue_swap_wgmma.template operator()< + 0, 1, 0, 2>( + stage_idx, expanded_slot); + issue_swap_wgmma.template operator()< + 1, 1, 0, 2>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 0, 1, 1, false>(stage_idx, 0); + promote_swap_ab.template operator()< + 1, 1, 0, false>(stage_idx, 0); + issue_swap_wgmma.template operator()< + 0, 1, 2, 2>( + stage_idx, expanded_slot); + issue_swap_wgmma.template operator()< + 1, 1, 2, 2>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 0, 1, 1, false>(stage_idx, 1); + } + } else if constexpr (kReuseSwapABFragment) { + if constexpr (is_linear1_phase) { + issue_swap_wgmma.template operator()< + 0, 1, 0, 4>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 0, 1, 0, false>(stage_idx, 0); + issue_swap_wgmma.template operator()< + 1, 1, 0, 4>( + stage_idx, expanded_slot); + } else { + issue_swap_wgmma.template operator()< + 0, 1, 0, 2>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 0, 1, 0, false>(stage_idx, 0); + issue_swap_wgmma.template operator()< + 1, 1, 0, 2>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 1, 1, 0, false>(stage_idx, 0); + issue_swap_wgmma.template operator()< + 0, 1, 2, 2>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 0, 1, 0, false>(stage_idx, 1); + issue_swap_wgmma.template operator()< + 1, 1, 2, 2>( + stage_idx, expanded_slot); + } + } else { + if constexpr (is_linear1_phase) { + issue_swap_wgmma.template operator()< + 0, 2, 0, 4>( + stage_idx, expanded_slot); + } else { + issue_swap_wgmma.template operator()< + 0, 2, 0, 2>( + stage_idx, expanded_slot); + promote_swap_ab.template operator()< + 0, 2, 0, false>(stage_idx, 0); + issue_swap_wgmma.template operator()< + 0, 2, 2, 2>( + stage_idx, expanded_slot); + } + } + + if constexpr (kPipelineMXFP4ExpandedB) { + if (k_block_idx + 1 < num_k_blocks) { + const uint32_t next_stage = + stage_idx == kNumStages - 1 ? + 0u : stage_idx + 1u; + const uint32_t next_phase = + phase ^ (next_stage == 0u); + full_barriers[next_stage]->wait( + next_phase); + prepare_stage_weights( + next_stage, + expanded_slot ^ 1u); + } + } + if constexpr (kPipelineWeightHalves) { + promote_swap_ab.template operator()< + 1, 1, 0, true>( + stage_idx, + is_linear1_phase ? 0u : 1u); + } else if constexpr (kReuseSwapABFragment) { + promote_swap_ab.template operator()< + 1, 1, 0, true>( + stage_idx, + is_linear1_phase ? 0u : 1u); + } else { + promote_swap_ab.template operator()< + 0, 2, 0, true>( + stage_idx, + is_linear1_phase ? 0u : 1u); + } + } + }; + + const uint32_t n_swap = + math::ceil_div(valid_m, 8u) * 8u; + if (n_swap <= 8) + run_swap_ab.template operator()<8>(); + else if (n_swap <= 16) + run_swap_ab.template operator()<16>(); + else if (n_swap <= 32) + run_swap_ab.template operator()<32>(); + else + run_swap_ab.template operator()<64>(); + + if constexpr (not kPackedBF16SwapEpilogue) { + #pragma unroll + for (uint32_t i = 0; + i < kAccumPerThread / 2; ++ i) { + const float2 pair = + __bfloat1622float2( + mxfp4_final_bf16[i]); + final_accum[i * 2] = pair.x; + final_accum[i * 2 + 1] = pair.y; + } + } + } else { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; + advance_pipeline(k_block_idx)) { + const uint32_t expanded_slot = + kPipelineMXFP4ExpandedB ? + (k_block_idx & 1u) : 0u; + if constexpr (kPipelineMXFP4ExpandedB) { + // Stage zero is the prologue. Later expanded tiles + // were decoded during the previous WGMMA flight. + if (k_block_idx == 0) { + full_barriers[stage_idx]->wait(phase); + prepare_stage_weights(stage_idx, expanded_slot); + } + } else { + full_barriers[stage_idx]->wait(phase); + prepare_stage_weights(stage_idx, expanded_slot); + } + + if constexpr (is_linear1_phase) { + issue_mxfp4_wgmma.template operator()<0, 4>( + stage_idx, expanded_slot); + if constexpr (kPipelineMXFP4ExpandedB) { + // Decode the next packed stage into the + // other expanded slot while this WGMMA + // group is still in flight. + if (k_block_idx + 1 < num_k_blocks) { + const uint32_t next_stage = + stage_idx == kNumStages - 1 ? + 0u : stage_idx + 1u; + const uint32_t next_phase = + phase ^ (next_stage == 0u); + full_barriers[next_stage]->wait( + next_phase); + prepare_stage_weights( + next_stage, expanded_slot ^ 1u); + } + } + promote_mxfp4.template operator()< + kEarlyReleaseMXFP4Stage>( + stage_idx, 0, mxfp4_secondary); + } else { + issue_mxfp4_wgmma.template operator()<0, 2>( + stage_idx, expanded_slot); + promote_mxfp4.template operator()( + stage_idx, 0, mxfp4_secondary); + issue_mxfp4_wgmma.template operator()<2, 2>( + stage_idx, expanded_slot); + if constexpr (kPipelineMXFP4ExpandedB) { + // L2's first K64 group must be promoted + // before the second group can reuse the + // fragment. Hide the next packed-stage + // decode under the final K64 WGMMA flight. + if (k_block_idx + 1 < num_k_blocks) { + const uint32_t next_stage = + stage_idx == kNumStages - 1 ? + 0u : stage_idx + 1u; + const uint32_t next_phase = + phase ^ (next_stage == 0u); + full_barriers[next_stage]->wait( + next_phase); + prepare_stage_weights( + next_stage, expanded_slot ^ 1u); + } + } + promote_mxfp4.template operator()< + kEarlyReleaseMXFP4Stage>( + stage_idx, 1, mxfp4_secondary); + } + if constexpr (kEarlyReleaseMXFP4Stage) { + if constexpr (not kOverlapMXFP4ScalePath) + arrive_task_empty_barrier(stage_idx); + } else { + arrive_task_empty_barrier(stage_idx); + } + } + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 2; ++ i) { + const float2 pair = + __bfloat1622float2(mxfp4_final_bf16[i]); + final_accum[i * 2] = pair.x; + final_accum[i * 2 + 1] = pair.y; + } + } + } + }; + + const auto run_shared_fp8_gemm_loop = [&]() { + constexpr uint32_t kL1SFKBlocks = kHidden / 128; + constexpr uint32_t kTaskIntermediateHidden = is_shared_phase ? + kIntermediateHidden * kNumSharedExperts : + kIntermediateHidden; + constexpr uint32_t kL2SFKBlocks = + kTaskIntermediateHidden / 128; + constexpr uint32_t kL1SFGateBlks = + kTaskIntermediateHidden / 128; + constexpr uint32_t kL1SFPerExpert = + (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; + constexpr uint32_t kL2SFPerExpert = + (kHidden / 128) * kL2SFKBlocks; + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; + advance_pipeline(k_block_idx)) { + float gate_sf = 0.0f, up_sf = 0.0f; + float l2_sf_lo = 0.0f, l2_sf_hi = 0.0f; + full_barriers[stage_idx]->wait(phase); + const auto task_smem_b = is_shared_phase ? + smem_shared_b[shared_b_stage_idx] : smem_b[stage_idx]; + + // 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) + const uint32_t stage_row_offset_r0 = + warp_idx_in_wg * 16 + lane_idx / 4; + if (is_linear1_phase) { + scale_a_0_lo = ptx::ld_shared( + smem_sfa[stage_idx] + stage_row_offset_r0); + scale_a_1_lo = ptx::ld_shared( + smem_sfa[stage_idx] + stage_row_offset_r0 + 8); + } 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] + stage_row_offset_r0); + scale_a_1_lo = ptx::ld_shared( + smem_sfa[stage_idx] + stage_row_offset_r0 + 8); + scale_a_0_hi = ptx::ld_shared( + smem_sfa[stage_idx] + kL2SFAHalfStride + + stage_row_offset_r0); + scale_a_1_hi = ptx::ld_shared( + smem_sfa[stage_idx] + kL2SFAHalfStride + + stage_row_offset_r0 + 8); + } + + // ----- 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 BLOCK_N=128 tile covers 64 rows of gate plus + // 64 rows of up taken from the same original 128-row block, so: + // gate_sf_n = n_block_idx / 2 + // up_sf_n = (IH/128) + n_block_idx / 2 + // + // L2 weight SF shape: (E, H/128, IH/128) MN-major. One scalar per + // (BLOCK_N, BLOCK_K) tile, broadcast across all WGMMA accumulators. + // + if (is_linear1_phase) { + const uint32_t gate_n = n_block_idx * BLOCK_N / 256u; + const uint32_t up_n = kL1SFGateBlks + gate_n; + const float* base = (is_shared_phase ? + shared_l1_weights_sf : l1_mxfp4_secondary) + + (is_shared_phase ? 0u : + local_expert_idx * kL1SFPerExpert) + + k_block_idx; + gate_sf = __ldg(base + gate_n * kL1SFKBlocks); + up_sf = __ldg(base + up_n * kL1SFKBlocks); + } else { + const uint32_t sf_n = n_block_idx * BLOCK_N / 128u; + const float* base = (is_shared_phase ? + shared_l2_weights_sf : l2_mxfp4_secondary) + + (is_shared_phase ? 0u : + local_expert_idx * kL2SFPerExpert) + + k_block_idx; + l2_sf_lo = __ldg(base + sf_n * kL2SFKBlocks); + l2_sf_hi = l2_sf_lo; + } + + if (is_linear1_phase) { + // 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] + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + task_smem_b + 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>(); + + arrive_task_empty_barrier(stage_idx); + + // 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] + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + task_smem_b + 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 weight SF is per 128 output columns; M64N256 spans two SF groups. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float l2_sf = (i < 16u) ? l2_sf_lo : l2_sf_hi; + 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] + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + task_smem_b + 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>(); + + arrive_task_empty_barrier(stage_idx); + + // L2 second half: same SFA half, still choose weight SF by N chunk. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float l2_sf = (i < 16u) ? l2_sf_lo : l2_sf_hi; + 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]; + } + } + } + }; + + if constexpr (not is_shared_phase) + run_mxfp4_gemm_loop(); + else + run_shared_fp8_gemm_loop(); + + if constexpr (BlockPhaseTag::value == sched::BlockPhase::Linear1) { + // L1 may only overwrite an intermediate slot after every L2 N + // task from the previous generation has consumed it. + const auto empty_ptr = + workspace.get_l2_empty_count_ptr(ring_block_idx); + const uint32_t empty_target = (L2_SHAPE_N / BLOCK_N) * + (pool_block_idx / kNumRingBlocks); + while (ptx::ld_acq(empty_ptr) != empty_target) {} + } + if constexpr (BlockPhaseTag::value == sched::BlockPhase::Linear2) { + // Every math thread has completed all A loads before this + // barrier; one thread releases the physical L2 slot. + ptx::sync_aligned( + kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) + ptx::red_add( + workspace.get_l2_empty_count_ptr(ring_block_idx), 1u); + __syncwarp(); + } + + // Skip epilogue when block is past valid M (still must release via empty). + if (valid_m == 0) { + ptx::sync_aligned( + kNumEpilogueThreads, kEpilogueFullBarrierIdx); + return; + } + + 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; + const uint32_t row_offset_r0 = r_0; + const uint32_t row_offset_r1 = r_1; + const bool valid_r0 = row_offset_r0 < valid_m; + const bool valid_r1 = row_offset_r1 < valid_m; + + if (is_linear1_phase) { + auto* smem_cd_l1_wg = smem_cd_l1; + if constexpr (kSmallMSwapAB and not is_shared_phase) { + // Swap-AB maps tokens across WGMMA N and output channels + // across lanes/warps. Derive the dynamic per-token K64 + // output scale with a two-level warpgroup reduction, then + // write the same row-major FP8/SF contract consumed by L2. + float swap_swiglu[ + kSwapABWeightHalves][kSwapABTokenChunks][2]; + auto 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; + }; + 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); + }; + const uint32_t num_swap_token_chunks = + math::ceil_div(valid_m, 8u); + // Pipeline SFA is reusable by the producer immediately + // after the math loop releases a stage. Use C/D storage, + // which is epilogue-exclusive, for the cross-warp amax. + auto* swap_scale_scratch = + reinterpret_cast(smem_cd_base); + + #pragma unroll + for (uint32_t chunk = 0; + chunk < kSwapABTokenChunks; ++ chunk) { + const uint32_t token_0 = + chunk * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + const bool active_chunk = + chunk < num_swap_token_chunks; + const float weight_0 = + active_chunk and token_0 < valid_m ? + *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_0) + .template get_base_ptr() : + 0.0f; + const float weight_1 = + active_chunk and token_1 < valid_m ? + *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_1) + .template get_base_ptr() : + 0.0f; + #pragma unroll + for (uint32_t half = 0; + half < kSwapABWeightHalves; ++ half) { + const uint32_t accum_offset = + half * kSwapABHalfAccumPerThread + + chunk * 4; + float gate_0, gate_1, up_0, up_1; + if constexpr (kPackedBF16SwapEpilogue) { + const float2 gate_pair = + __bfloat1622float2( + mxfp4_final_bf16[ + accum_offset / 2]); + const float2 up_pair = + __bfloat1622float2( + mxfp4_final_bf16[ + accum_offset / 2 + 1]); + gate_0 = gate_pair.x; + gate_1 = gate_pair.y; + up_0 = up_pair.x; + up_1 = up_pair.y; + } else { + gate_0 = final_accum[accum_offset]; + gate_1 = final_accum[accum_offset + 1]; + up_0 = final_accum[accum_offset + 2]; + up_1 = final_accum[accum_offset + 3]; + } + clamp_gate(gate_0); + clamp_gate(gate_1); + clamp_up(up_0); + clamp_up(up_1); + swap_swiglu[half][chunk][0] = + silu(gate_0) * up_0 * weight_0; + swap_swiglu[half][chunk][1] = + silu(gate_1) * up_1 * weight_1; + } + + float partial_0 = cute::max( + cute::abs(swap_swiglu[0][chunk][0]), + cute::abs(swap_swiglu[1][chunk][0])); + float partial_1 = cute::max( + cute::abs(swap_swiglu[0][chunk][1]), + cute::abs(swap_swiglu[1][chunk][1])); + #pragma unroll + for (uint32_t delta = 4; delta <= 16; delta *= 2) { + partial_0 = cute::max( + partial_0, + __shfl_xor_sync( + 0xffffffffu, partial_0, delta)); + partial_1 = cute::max( + partial_1, + __shfl_xor_sync( + 0xffffffffu, partial_1, delta)); + } + if (row_idx == 0 and active_chunk) { + swap_scale_scratch[ + warp_idx_in_wg * BLOCK_M + token_0] = + partial_0; + swap_scale_scratch[ + warp_idx_in_wg * BLOCK_M + token_1] = + partial_1; + } + } + ptx::sync_aligned( + kNumEpilogueThreads, + kEpilogueWGBarrierStartIdx); + + if (warp_idx_in_wg == 0) { + #pragma unroll + for (uint32_t half = 0; half < 2; ++ half) { + const uint32_t token = lane_idx + half * 32; + if (token < valid_m) { + float amax = 0.0f; + #pragma unroll + for (uint32_t warp = 0; + warp < kNumEpilogueWarps; ++ warp) + amax = cute::max( + amax, + swap_scale_scratch[ + warp * BLOCK_M + token]); + float2 amax_pair = {amax, amax}; + float2 sf_pair, sf_inv_pair; + sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( + amax_pair, sf_pair, sf_inv_pair); + swap_scale_scratch[token] = sf_pair.x; + swap_scale_scratch[BLOCK_M + token] = + sf_inv_pair.x; + } + } + } + ptx::sync_aligned( + kNumEpilogueThreads, + kEpilogueWGBarrierStartIdx); + + float swap_sf_inv[kSwapABTokenChunks][2]; + #pragma unroll + for (uint32_t chunk = 0; + chunk < kSwapABTokenChunks; ++ chunk) { + const uint32_t token_0 = + chunk * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + swap_sf_inv[chunk][0] = token_0 < valid_m ? + swap_scale_scratch[BLOCK_M + token_0] : 0.0f; + swap_sf_inv[chunk][1] = token_1 < valid_m ? + swap_scale_scratch[BLOCK_M + token_1] : 0.0f; + } + if (warp_idx_in_wg == 0) { + auto sf_base_ptr = + l2_sf_buffer.get_base_ptr(); + const uint32_t base_k_sf_idx = n_block_idx; + #pragma unroll + for (uint32_t half = 0; half < 2; ++ half) { + const uint32_t token = lane_idx + half * 32; + if (token < valid_m) + sf_base_ptr[ + base_k_sf_idx * kNumSFRingTokens + + m_idx + token] = + swap_scale_scratch[token]; + } + } + // Every thread must retain its inverse scales before the + // row-major FP8 stores overwrite C/D scratch. + ptx::sync_aligned( + kNumEpilogueThreads, + kEpilogueWGBarrierStartIdx); + + #pragma unroll + for (uint32_t chunk = 0; + chunk < kSwapABTokenChunks; ++ chunk) { + if (chunk < num_swap_token_chunks) { + const uint32_t token_0 = + chunk * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + #pragma unroll + for (uint32_t half = 0; + half < kSwapABWeightHalves; ++ half) { + const uint32_t out_col = + half * 32u + + warp_idx_in_wg * 8 + row_idx; + if (token_0 < valid_m) { + const __nv_fp8_e4m3 q( + swap_swiglu[half][chunk][0] * + swap_sf_inv[chunk][0]); + reinterpret_cast( + smem_cd_l1_wg)[ + token_0 * + WG_SMEM_CD_L1_STRIDE_N + + out_col] = + *reinterpret_cast(&q); + } + if (token_1 < valid_m) { + const __nv_fp8_e4m3 q( + swap_swiglu[half][chunk][1] * + swap_sf_inv[chunk][1]); + reinterpret_cast( + smem_cd_l1_wg)[ + token_1 * + WG_SMEM_CD_L1_STRIDE_N + + out_col] = + *reinterpret_cast(&q); + } + } + } + } + + } else { + // ---------------- L1 EPILOGUE: SwiGLU + FP8 quantize + TMA store ---------------- + // Layout in `final_accum`: + // 16 chunks of 8 N-cols, each chunk = 4 floats per thread = (r0c0, r0c1, r1c0, r1c1). + // Gate chunks: even (0, 2, ..., 14). Up chunks: odd (1, 3, ..., 15). + // Pair `p` ∈ [0, 8): gate chunk = 2p, up chunk = 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; + DG_STATIC_ASSERT(WG_L1_OUT_BLOCK_N % 64 == 0, + "Each L1 consumer must cover complete 64-column SF groups"); + constexpr uint32_t kNumSFGroups = WG_L1_OUT_BLOCK_N / 64; + float swiglu_r0[kNumPairs][2]; + float swiglu_r1[kNumPairs][2]; + + // Per-row amax, one scale for each 64-col L1 output group. + float amax_r0[kNumSFGroups] = {}; + float amax_r1[kNumSFGroups] = {}; + + // Compute SwiGLU + per-group amax. + #pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++ p) { + const uint32_t gate = 2 * p, up = 2 * p + 1; + const uint32_t sf_group = p / 8; + + 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); + }; + float g_r0_c0 = final_accum[gate*4 + 0]; clamp_gate(g_r0_c0); + float g_r0_c1 = final_accum[gate*4 + 1]; clamp_gate(g_r0_c1); + float g_r1_c0 = final_accum[gate*4 + 2]; clamp_gate(g_r1_c0); + float g_r1_c1 = final_accum[gate*4 + 3]; clamp_gate(g_r1_c1); + float u_r0_c0 = final_accum[up*4 + 0]; clamp_up(u_r0_c0); + float u_r0_c1 = final_accum[up*4 + 1]; clamp_up(u_r0_c1); + float u_r1_c0 = final_accum[up*4 + 2]; clamp_up(u_r1_c0); + float u_r1_c1 = final_accum[up*4 + 3]; clamp_up(u_r1_c1); + + auto 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; + }; + + 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[sf_group] = cute::max( + amax_r0[sf_group], + 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[sf_group] = cute::max( + amax_r1[sf_group], + 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; + } + } + + + float weight_r0 = 0.0f, weight_r1 = 0.0f; + if constexpr (kNumMaxTokensPerRank <= 1024) { + const int topk_weight_src_lane = static_cast(lane_idx - col_idx); + if (col_idx == 0) { + weight_r0 = is_shared_phase ? 1.0f : + (valid_r0 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r0) + .template get_base_ptr() : 0.0f); + weight_r1 = is_shared_phase ? 1.0f : + (valid_r1 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r1) + .template get_base_ptr() : 0.0f); + } + weight_r0 = __shfl_sync(0xffffffff, weight_r0, topk_weight_src_lane); + weight_r1 = __shfl_sync(0xffffffff, weight_r1, topk_weight_src_lane); + } else { + weight_r0 = is_shared_phase ? 1.0f : + (valid_r0 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r0) + .template get_base_ptr() : 0.0f); + weight_r1 = is_shared_phase ? 1.0f : + (valid_r1 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r1) + .template 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; + } + #pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++ g) { + amax_r0[g] *= cute::abs(weight_r0); + amax_r1[g] *= cute::abs(weight_r1); + } + #pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++ g) { + amax_r0[g] = math::warp_reduce<4, false>(amax_r0[g], math::ReduceMax()); + amax_r1[g] = math::warp_reduce<4, false>(amax_r1[g], math::ReduceMax()); + } + + float sf_r0[kNumSFGroups], sf_inv_r0[kNumSFGroups]; + float sf_r1[kNumSFGroups], sf_inv_r1[kNumSFGroups]; + #pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++ g) { + float2 amax_pair = {amax_r0[g], amax_r1[g]}; + 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[g] = sf_pair.x; sf_inv_r0[g] = sf_inv_pair.x; + sf_r1[g] = sf_pair.y; sf_inv_r1[g] = sf_inv_pair.y; + } + + // Quantize and write to smem_cd_l1 (row-major, no swizzle). + #pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++ p) { + const uint32_t sf_group = p / 8; + const float v00 = swiglu_r0[p][0] * sf_inv_r0[sf_group]; + const float v01 = swiglu_r0[p][1] * sf_inv_r0[sf_group]; + const float v10 = swiglu_r1[p][0] * sf_inv_r1[sf_group]; + const float v11 = swiglu_r1[p][1] * sf_inv_r1[sf_group]; + + const __nv_fp8x2_e4m3 r0_pair(make_float2(v00, v01)); + const __nv_fp8x2_e4m3 r1_pair(make_float2(v10, v11)); + + 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_pair.__x; + if (valid_r1) + *p1 = r1_pair.__x; + } + + // Write L2-activation SF as float, one value per 64 output columns. + if (col_idx == 0) { + auto sf_base_ptr = is_shared_phase ? + shared_l2_sf_buffer.get_base_ptr() : + l2_sf_buffer.get_base_ptr(); + constexpr uint32_t kSharedSFStride = + layout::get_num_shared_sf_tokens( + kNumMaxTokensPerRank, BLOCK_M); + const uint32_t sf_stride = is_shared_phase ? + kSharedSFStride : kNumSFRingTokens; + const uint32_t token_r0 = m_idx + row_offset_r0; + const uint32_t token_r1 = m_idx + row_offset_r1; + const uint32_t base_k_sf_idx = + n_block_idx * L1_OUT_BLOCK_N / 64u; + #pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++ g) { + if (valid_r0) + sf_base_ptr[(base_k_sf_idx + g) * sf_stride + token_r0] = sf_r0[g]; + if (valid_r1) + sf_base_ptr[(base_k_sf_idx + g) * sf_stride + token_r1] = sf_r1[g]; + } + } + } + + // 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. + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx); + if (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( + is_shared_phase ? &tensor_map_shared_l1_output : + &tensor_map_l1_output, + smem_cd_l1_wg, + out_n_idx, + m_idx); + cute::tma_store_arrive(); + } + __syncwarp(); + + // Publish L1 only after every ordinary SF store and every + // asynchronous TMA output store is globally visible. One + // completion is contributed by each L1 N task. + ptx::tma_store_wait<0>(); + ptx::sync_aligned( + kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + if constexpr (is_shared_phase) { + ptx::red_add_rel( + workspace.get_shared_l2_full_count_ptr( + pool_block_idx), + 1u); + } else { + ptx::red_add_rel( + workspace.get_l2_full_count_ptr(ring_block_idx), + 1u); + ptx::red_add( + workspace.get_l1_empty_count_ptr(ring_block_idx), + 1u); + } + } + __syncwarp(); + } else { + // ---------------- L2 EPILOGUE: BF16 cast + NVLink scatter ---------------- + constexpr uint32_t kNumRowsPerWarp = WG_BLOCK_M / 8; + + const auto get_l2_cd_byte_idx = [](const uint32_t byte_idx) { + if constexpr (kSwizzleL2CD) + return cute::Swizzle<3, 4, 3>::apply(byte_idx); + return byte_idx; + }; + auto store_l2_pair = [&](const uint32_t& elem_idx, + float value0, float value1) { + const uint32_t storage_byte_idx = get_l2_cd_byte_idx( + elem_idx * sizeof(nv_bfloat16)); + *reinterpret_cast( + static_cast(smem_cd_base) + + storage_byte_idx) = + math::cast_into_bf16_and_pack(value0, value1); + }; + auto store_l2_scalar = [&](const uint32_t& elem_idx, + float value) { + const uint32_t storage_byte_idx = get_l2_cd_byte_idx( + elem_idx * sizeof(nv_bfloat16)); + *reinterpret_cast( + static_cast(smem_cd_base) + + storage_byte_idx) = __float2bfloat16_rn(value); + }; + if constexpr (kSmallMSwapAB and not is_shared_phase) { + const uint32_t num_swap_token_chunks = + math::ceil_div(valid_m, 8u); + #pragma unroll + for (uint32_t chunk = 0; + chunk < kSwapABTokenChunks; ++ chunk) { + if (chunk < num_swap_token_chunks) { + const uint32_t token_0 = + chunk * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + #pragma unroll + for (uint32_t half = 0; + half < kSwapABWeightHalves; ++ half) { + const uint32_t accum_offset = + half * kSwapABHalfAccumPerThread + + chunk * 4; + const uint32_t col_offset = half * 64u; + if constexpr ( + kPackedBF16SwapEpilogue) { + const float2 gate_pair = + __bfloat1622float2( + mxfp4_final_bf16[ + accum_offset / 2]); + const float2 up_pair = + __bfloat1622float2( + mxfp4_final_bf16[ + accum_offset / 2 + 1]); + if (token_0 < valid_m) { + store_l2_scalar( + token_0 * WG_BLOCK_N + + col_offset + r_0, + gate_pair.x); + store_l2_scalar( + token_0 * WG_BLOCK_N + + col_offset + r_1, + up_pair.x); + } + if (token_1 < valid_m) { + store_l2_scalar( + token_1 * WG_BLOCK_N + + col_offset + r_0, + gate_pair.y); + store_l2_scalar( + token_1 * WG_BLOCK_N + + col_offset + r_1, + up_pair.y); + } + } else { + if (token_0 < valid_m) { + store_l2_scalar( + token_0 * WG_BLOCK_N + + col_offset + r_0, + final_accum[ + accum_offset]); + store_l2_scalar( + token_0 * WG_BLOCK_N + + col_offset + r_1, + final_accum[ + accum_offset + 2]); + } + if (token_1 < valid_m) { + store_l2_scalar( + token_1 * WG_BLOCK_N + + col_offset + r_0, + final_accum[ + accum_offset + 1]); + store_l2_scalar( + token_1 * WG_BLOCK_N + + col_offset + r_1, + final_accum[ + accum_offset + 3]); + } + } + } + } + } + } else { + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 8; ++ i) { + const uint32_t chunk_lo = 2 * i, chunk_hi = 2 * i + 1; + auto write_pair = [&](const uint32_t row, + const uint32_t col, + const float value0, + const float value1) { + const uint32_t elem_idx = + row * WG_BLOCK_N + col; + store_l2_pair(elem_idx, value0, value1); + }; + if (valid_r0) { + write_pair(r_0, chunk_lo * 8 + col_idx * 2, + final_accum[chunk_lo * 4 + 0], + final_accum[chunk_lo * 4 + 1]); + write_pair(r_0, chunk_hi * 8 + col_idx * 2, + final_accum[chunk_hi * 4 + 0], + final_accum[chunk_hi * 4 + 1]); + } + if (valid_r1) { + write_pair(r_1, chunk_lo * 8 + col_idx * 2, + final_accum[chunk_lo * 4 + 2], + final_accum[chunk_lo * 4 + 3]); + write_pair(r_1, chunk_hi * 8 + col_idx * 2, + final_accum[chunk_hi * 4 + 2], + final_accum[chunk_hi * 4 + 3]); + } + } + } + + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx); + + // Scatter to remote ranks via NVLink (one row per warp-pair) + // Each warpgroup-warp covers 8 unique rows × 2 (r_0 + r_1 doubled by warps) + // Lane group of 16 within a warp → 1 row. + 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; + #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_in_wg; + if (m_idx_in_block >= valid_m) break; + + uint32_t dst_rank_idx, dst_token_idx, dst_topk_idx; + if constexpr (is_shared_phase) { + dst_rank_idx = sym_buffer.rank_idx; + dst_token_idx = pool_m_idx + m_idx_in_block; + dst_topk_idx = kNumTopk; + } else { + const auto src_metadata = + *workspace.get_token_src_metadata_ptr( + pool_m_idx + m_idx_in_block); + dst_rank_idx = src_metadata.rank_idx; + dst_token_idx = src_metadata.token_idx; + dst_topk_idx = src_metadata.topk_idx; + } + + const uint32_t smem_elem_idx = + row_in_wg * WG_BLOCK_N + + lane_in_row * cols_per_lane; + constexpr uint32_t kScatterBytesPerLane = + (WG_BLOCK_N / 16) * kCombineElementBytes; + DG_STATIC_ASSERT( + kScatterBytesPerLane == 4 or + kScatterBytesPerLane == 8 or + kScatterBytesPerLane == 16 or + kScatterBytesPerLane == 32, + "Unexpected L2 scatter width"); + // B128 swizzle preserves each aligned 16-byte segment, + // so the vectorized scatter load stays contiguous while + // epilogue stores spread row groups across SMEM banks. + const uint32_t storage_smem_byte_idx = get_l2_cd_byte_idx( + smem_elem_idx * kCombineElementBytes); + auto smem_ptr = math::advance_ptr( + smem_cd_base, storage_smem_byte_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 * kCombineElementBytes + + lane_in_row * kScatterBytesPerLane); + auto mapped_dst_ptr = sym_buffer.map(dst_ptr, dst_rank_idx); + + if constexpr (kScatterBytesPerLane == 32) { + const auto packed0 = + *reinterpret_cast(smem_ptr); + const auto packed1 = + *(reinterpret_cast(smem_ptr) + 1); + reinterpret_cast(mapped_dst_ptr)[0] = packed0; + reinterpret_cast(mapped_dst_ptr)[1] = packed1; + } else if constexpr (kScatterBytesPerLane == 16) { + const auto packed = + *reinterpret_cast(smem_ptr); + *reinterpret_cast(mapped_dst_ptr) = packed; + } else if constexpr (kScatterBytesPerLane == 8) { + const auto packed = + *reinterpret_cast(smem_ptr); + *reinterpret_cast(mapped_dst_ptr) = packed; + } else { + // The width assertion above leaves only the 4-byte case. + *reinterpret_cast(mapped_dst_ptr) = + *reinterpret_cast(smem_ptr); + } + } + + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + }); + + // ---------------- COMBINE ---------------- + // NVLink barrier first: signals remote ranks that this rank's GEMM + // outputs (NVLink scatter targets) are fully written. + sm90_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); + + constexpr uint32_t kNumChunks = kCombineNumChunks; + constexpr uint32_t kInputChunkBytes = kCombineInputChunkBytes; + constexpr uint32_t kOutputChunkBytes = kCombineOutputChunkBytes; + constexpr uint32_t kElemsPerVector = 8; + constexpr uint32_t kInputVectorBytes = kElemsPerVector * kCombineElementBytes; + constexpr uint32_t kNumVectorsPerLane = + kCombineChunkElems / (32 * kElemsPerVector); + constexpr uint32_t kNumBF16PairsPerVector = kElemsPerVector / 2; + DG_STATIC_ASSERT(kInputChunkBytes % 16 == 0, + "Combine input chunk must be TMA-aligned"); + DG_STATIC_ASSERT(kOutputChunkBytes % 16 == 0, + "Combine output chunk must be TMA-aligned"); + DG_STATIC_ASSERT(kCombineChunkElems % (32 * kElemsPerVector) == 0, + "Combine chunk must distribute evenly across a warp"); + DG_STATIC_ASSERT( + kNumTopk + (kHasSharedExperts ? 1u : 0u) <= 32, + "Top-k plus shared expert must fit in a single warp"); + + const auto combine_load_buffer = utils::PatternVisitor([&](const uint32_t& i) { + return math::advance_ptr( + smem_buffer, + (epilogue_warp_idx + i * kNumEpilogueWarps) * kInputChunkBytes); + }); + const auto combine_store_buffer = math::advance_ptr(smem_buffer, + 2 * kNumEpilogueWarps * kInputChunkBytes + + epilogue_warp_idx * kOutputChunkBytes); + + 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 * kNumEpilogueWarps + epilogue_warp_idx; + token_idx < num_tokens; + token_idx += kNumSMs * kNumEpilogueWarps) { + const int stored_topk_slot_idx = lane_idx < kNumTopk ? + static_cast(__ldg( + input_topk_idx_buffer.get_base_ptr() + + token_idx * kNumTopk + lane_idx)) : + (kHasSharedExperts and lane_idx == kNumTopk ? + static_cast(kNumTopk) : -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 input_chunk_byte_offset = chunk * kInputChunkBytes; + const uint32_t output_chunk_byte_offset = chunk * kOutputChunkBytes; + + 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(), + input_chunk_byte_offset); + ptx::tma_load_1d( + combine_load_buffer[i], src_ptr, + combine_load_barriers[i], kInputChunkBytes); + ptx::mbarrier_arrive_and_set_tx( + combine_load_barriers[i], kInputChunkBytes); + } + __syncwarp(); + return true; + } + return false; + }; + + bool do_reduce = move_mask_and_load(load_stage_idx); + + // Fast mode rounds after each routed contribution so the combine + // state stays in packed BF16 registers; strict mode retains FP32. + using combine_accum_t = std::conditional_t< + kFastMath, nv_bfloat162, float2>; + combine_accum_t reduced[ + kNumVectorsPerLane * kNumBF16PairsPerVector] = {}; + 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 < kNumVectorsPerLane; ++ j) { + const uint32_t vector_idx = j * 32 + lane_idx; + const auto packed = *reinterpret_cast( + combine_load_buffer[load_stage_idx] + + vector_idx * kInputVectorBytes); + const auto bf16_values = + reinterpret_cast(&packed); + #pragma unroll + for (uint32_t l = 0; l < kNumBF16PairsPerVector; ++ l) { + const uint32_t accum_idx = + j * kNumBF16PairsPerVector + l; + if constexpr (kFastMath) { + reduced[accum_idx] = __hadd2( + reduced[accum_idx], bf16_values[l]); + } else { + ptx::accumulate( + reduced[accum_idx], bf16_values[l]); + } + } + } + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } + + #pragma unroll + for (uint32_t j = 0; j < kNumVectorsPerLane; ++ j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); + #pragma unroll + for (uint32_t l = 0; l < kNumBF16PairsPerVector; ++ l) { + const uint32_t accum_idx = + j * kNumBF16PairsPerVector + l; + if constexpr (kFastMath) { + casted_bf16[l] = reduced[accum_idx]; + } else { + casted_bf16[l] = __float22bfloat162_rn( + reduced[accum_idx]); + } + } + + 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) * kCombineOutputHiddenBytes + + output_chunk_byte_offset), + combine_store_buffer, kOutputChunkBytes); + 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 +} + +template +CUTLASS_GLOBAL __launch_bounds__( + layout::Sm90HummingMoeSchedule::num_dispatch_threads + + layout::Sm90HummingMoeSchedule::num_non_epilogue_threads + + layout::Sm90HummingMoeSchedule::num_epilogue_threads, + 2) void +sm90_fp8_mxfp4_mega_moe_persistent_impl( + DG_SM90_FP8_MOE_KERNEL_ARGS_DECL) { + sm90_fp8_mega_moe_core< + DG_SM90_FP8_MOE_CORE_TEMPLATE_ARGS>( + DG_SM90_FP8_MOE_KERNEL_ARGS); +} + +#undef DG_SM90_FP8_MOE_TEMPLATE_PARAMS +#undef DG_SM90_FP8_MOE_KERNEL_ARGS_DECL +#undef DG_SM90_FP8_MOE_CORE_ARGS_DECL +#undef DG_SM90_FP8_MOE_KERNEL_ARGS +#undef DG_SM90_FP8_MOE_CORE_TEMPLATE_ARGS + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh index 410c0b6a2a..0d205f8470 100644 --- a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh @@ -30,10 +30,27 @@ CUTLASS_HOST_DEVICE constexpr T get_num_sf_ring_tokens(T num_ring_tokens, T bloc return (num_ring_tokens / block_m) * math::constexpr_align(block_m, static_cast(128)); } -// Shared L2 input SF capacity: worst-case aligned SF pages over all candidate BLOCK_M. +// SM90's split L1/L2 MegaMoE backend uses the complete routed-token pool +// instead of the SM100 live-block ring. The SF padding rule is otherwise +// identical. Keep the legacy name as a narrow compatibility entry point so +// the SM90 kernel does not need to reuse the SM100 scheduling vocabulary. +template +CUTLASS_HOST_DEVICE constexpr T get_num_padded_sf_pool_tokens(T num_max_pool_tokens, T block_m) { + return get_num_sf_ring_tokens(num_max_pool_tokens, block_m); +} + +// Shared L2 input SF capacity for a fixed schedule block size. +template +CUTLASS_HOST_DEVICE constexpr T get_num_shared_sf_tokens( + const T& num_max_tokens_per_rank, const T& block_m) { + return math::constexpr_ceil_div(num_max_tokens_per_rank, block_m) * 128; +} + +// Backward-compatible worst case over all candidate BLOCK_M values. template CUTLASS_HOST_DEVICE constexpr T get_num_max_shared_sf_tokens(const T& num_max_tokens_per_rank) { - return math::constexpr_ceil_div(num_max_tokens_per_rank, kMinCandidateBlockM) * 128; + return get_num_shared_sf_tokens( + num_max_tokens_per_rank, static_cast(kMinCandidateBlockM)); } // Per-token source metadata for combine write-back @@ -82,6 +99,26 @@ struct Workspace { num_shared_l2_pool_blocks = math::ceil_div(num_max_tokens_per_rank, kMinCandidateBlockM); } + // Full-pool compatibility constructor for the SM90 split-kernel backend. + // Current SM100 callers always use the six-argument constructor above. + CUTLASS_HOST_DEVICE + Workspace(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): + Workspace( + base, + num_ranks, + num_experts, + num_max_tokens_per_rank, + num_topk, + get_num_max_pool_tokens( + num_ranks, + num_max_tokens_per_rank, + num_topk, + num_experts / num_ranks)) {} + CUTLASS_HOST_DEVICE uint64_t get_num_bytes() const { uint64_t num_bytes = 0; @@ -197,6 +234,14 @@ struct Workspace { return reinterpret_cast(base) + ring_block_idx; } + // SM90 publishes one arrival count per full-pool block. With the + // five-argument constructor `num_ring_tokens == num_max_pool_tokens`, so + // the existing L1 full-count region is exactly the required storage. + CUTLASS_DEVICE + uint32_t* get_l1_arrival_count_ptr(const uint32_t& pool_block_idx = 0) const { + return get_l1_full_count_ptr(pool_block_idx); + } + CUTLASS_DEVICE uint32_t* get_l1_empty_count_ptr(const uint32_t& ring_block_idx = 0) const { const auto base = get_l1_full_count_ptr(num_ring_blocks); @@ -328,6 +373,36 @@ struct Buffer { } }; +// Per-token activation scale-factor storage. The default matches the SM100 +// packed-UE8M0 ABI (one 32-bit word per 128 K elements). Hopper keeps FP32 +// activation scales and uses K64 for the L2 input, so it overrides only the +// intermediate row size while reusing the rest of MegaMoEBuffer. +struct ScaleLayoutSpec { + uint32_t input_sf_bytes_per_token; + uint32_t intermediate_sf_bytes_per_token; + + CUTLASS_HOST_DEVICE + constexpr ScaleLayoutSpec( + const uint32_t& input_sf_bytes_per_token = 0, + const uint32_t& intermediate_sf_bytes_per_token = 0): + input_sf_bytes_per_token(input_sf_bytes_per_token), + intermediate_sf_bytes_per_token(intermediate_sf_bytes_per_token) {} + + CUTLASS_HOST_DEVICE + static constexpr ScaleLayoutSpec sm100( + const uint32_t& hidden, + const uint32_t& intermediate_hidden) { + return {hidden / 32, intermediate_hidden / 32}; + } + + CUTLASS_HOST_DEVICE + static constexpr ScaleLayoutSpec sm90_fp32_k128_k64( + const uint32_t& hidden, + const uint32_t& intermediate_hidden) { + return {hidden / 32, intermediate_hidden / 16}; + } +}; + struct MegaMoEBuffer { Workspace workspace; @@ -337,7 +412,7 @@ struct MegaMoEBuffer { input_topk_idx_buffer, input_topk_weights_buffer; - // Routed expert ring buffers + // Shared expert buffers // NOTE: shared L1 tokens reuse `input_token_buffer`. Buffer shared_l1_token_buffer, shared_l1_sf_buffer, shared_l2_token_buffer, shared_l2_sf_buffer; @@ -361,14 +436,32 @@ struct MegaMoEBuffer { const uint32_t& num_ring_tokens, const uint32_t& num_sf_ring_tokens, const bool& with_sf, - const uint32_t& num_shared_experts = 0) { + const uint32_t& num_shared_experts = 0, + const ScaleLayoutSpec& scale_layout_spec = ScaleLayoutSpec(), + const uint32_t shared_sf_block_m = kMinCandidateBlockM) { // Workspace workspace = Workspace(base, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_ring_tokens); // Shared const auto shared_intermediate_hidden = intermediate_hidden * num_shared_experts; - const auto num_max_shared_sf_tokens = with_sf ? get_num_max_shared_sf_tokens(num_max_tokens_per_rank) : 0u; + const auto num_max_shared_sf_tokens = with_sf ? get_num_shared_sf_tokens( + num_max_tokens_per_rank, shared_sf_block_m) : 0u; + + // A zero-initialized spec is the backward-compatible SM100 default. + // Keeping the scale row sizes explicit prevents an SM90 K64 L2 scale + // buffer from being silently allocated with the SM100 K128 byte span. + const auto default_scale_layout = ScaleLayoutSpec::sm100(hidden, intermediate_hidden); + const auto input_sf_bytes_per_token = with_sf ? + (scale_layout_spec.input_sf_bytes_per_token == 0 ? + default_scale_layout.input_sf_bytes_per_token : + scale_layout_spec.input_sf_bytes_per_token) : 0u; + const auto intermediate_sf_bytes_per_token = with_sf ? + (scale_layout_spec.intermediate_sf_bytes_per_token == 0 ? + default_scale_layout.intermediate_sf_bytes_per_token : + scale_layout_spec.intermediate_sf_bytes_per_token) : 0u; + DG_UNIFIED_ASSERT(not with_sf or input_sf_bytes_per_token % 16 == 0); + DG_UNIFIED_ASSERT(not with_sf or intermediate_sf_bytes_per_token % 16 == 0); // Layouts const uint32_t num_mma_elem_bytes = with_sf ? 1 : 2; @@ -376,9 +469,10 @@ struct MegaMoEBuffer { const auto bf16_token_layout = layout::Data(hidden * 2); const auto intermediate_token_layout = layout::Data(intermediate_hidden * num_mma_elem_bytes); const auto shared_intermediate_token_layout = layout::Data(shared_intermediate_hidden * num_mma_elem_bytes); - const auto input_sf_layout = layout::Data(with_sf ? hidden / 32 : 0); - const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / 32 : 0); - const auto shared_intermediate_sf_layout = layout::Data(with_sf ? shared_intermediate_hidden / 32 : 0); + const auto input_sf_layout = layout::Data(input_sf_bytes_per_token); + const auto intermediate_sf_layout = layout::Data(intermediate_sf_bytes_per_token); + const auto shared_intermediate_sf_layout = layout::Data( + intermediate_sf_bytes_per_token * num_shared_experts); 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); 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..85b7872c79 --- /dev/null +++ b/deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh @@ -0,0 +1,20 @@ +#pragma once + +#include + +namespace deep_gemm::layout { + +// Single production schedule for the Hopper Humming FP8 x MXFP4 kernel. +// Host shared-memory sizing, TMA descriptors, JIT instantiation, and device +// launch bounds all consume this contract. +struct Sm90HummingMoeSchedule { + static constexpr uint32_t block_m = 64; + static constexpr uint32_t block_n = 128; + static constexpr uint32_t block_k = 128; + static constexpr uint32_t num_stages = 3; + static constexpr uint32_t num_dispatch_threads = 64; + static constexpr uint32_t num_non_epilogue_threads = 64; + static constexpr uint32_t num_epilogue_threads = 128; +}; + +} // namespace deep_gemm::layout diff --git a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh index 806a4c2e20..cf0f103106 100644 --- a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh +++ b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh @@ -105,6 +105,14 @@ CUTLASS_DEVICE uint32_t ld_shared(const uint32_t* ptr) { return ret; } +CUTLASS_DEVICE uint2 ld_shared(const uint2* ptr) { + uint2 ret; + asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];" + : "=r"(ret.x), "=r"(ret.y) + : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + CUTLASS_DEVICE float2 ld_shared(const float2* ptr) { float2 ret; asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];" : "=f"(ret.x), "=f"(ret.y) : "l"(__cvta_generic_to_shared(ptr))); @@ -205,6 +213,12 @@ CUTLASS_DEVICE uint32_t ld_acq(const uint32_t* ptr) { return ret; } +CUTLASS_DEVICE uint32_t ld_relaxed_gpu(const uint32_t* ptr) { + uint32_t ret; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + CUTLASS_DEVICE uint64_t ld_acq_sys(const uint64_t* ptr) { uint64_t ret; asm volatile("ld.acquire.sys.global.b64 %0, [%1];" : "=l"(ret) : "l"(ptr)); diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh index ba0e11e8f1..f09e920096 100644 --- a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh @@ -49,15 +49,18 @@ CUTLASS_HOST_DEVICE constexpr int get_num_max_live_pool_blocks( const int& num_total_m_blocks, const int& num_sms, const int& hidden, - const int& intermediate_hidden) { - constexpr int kMegaMoEBlockN = 128; - constexpr int kNumCTAsPerCluster = 2; - - DG_UNIFIED_ASSERT((intermediate_hidden * 2) % (kNumCTAsPerCluster * kMegaMoEBlockN) == 0); - DG_UNIFIED_ASSERT(hidden % (kNumCTAsPerCluster * kMegaMoEBlockN) == 0); - const int num_clusters = num_sms / kNumCTAsPerCluster; - const int num_l1_n_clusters = intermediate_hidden * 2 / (kNumCTAsPerCluster * kMegaMoEBlockN); - const int num_l2_n_clusters = hidden / (kNumCTAsPerCluster * kMegaMoEBlockN); + const int& intermediate_hidden, + const int& block_n = 128, + const int& num_ctas_per_task = 2) { + + DG_UNIFIED_ASSERT(block_n > 0); + DG_UNIFIED_ASSERT(num_ctas_per_task > 0); + DG_UNIFIED_ASSERT(num_sms % num_ctas_per_task == 0); + DG_UNIFIED_ASSERT((intermediate_hidden * 2) % (num_ctas_per_task * block_n) == 0); + DG_UNIFIED_ASSERT(hidden % (num_ctas_per_task * block_n) == 0); + const int num_clusters = num_sms / num_ctas_per_task; + const int num_l1_n_clusters = intermediate_hidden * 2 / (num_ctas_per_task * block_n); + const int num_l2_n_clusters = hidden / (num_ctas_per_task * block_n); const int num_l1_clusters = num_total_m_blocks * num_l1_n_clusters; const int num_l1_waves = math::constexpr_ceil_div(num_l1_clusters, num_clusters); const int num_min_l1_warmup_waves = get_num_l1_warmup_waves( @@ -144,11 +147,13 @@ template + uint32_t kNumL1Clusters = kNumL1BlockNs / kNumCTAsPerTask, + uint32_t kNumL2Clusters = kNumL2BlockNs / kNumCTAsPerTask> struct MegaMoEScheduler { static constexpr bool kHasShared = kNumSharedExperts > 0; static constexpr uint32_t SHARED_L1_SHAPE_N = L1_SHAPE_N * kNumSharedExperts; @@ -157,18 +162,23 @@ struct MegaMoEScheduler { static constexpr uint32_t SHARED_L2_SHAPE_K = L2_SHAPE_K * kNumSharedExperts; using task_info_t = TaskInfo; - DG_STATIC_ASSERT(L1_SHAPE_N % (BLOCK_N * 2) == 0, "Invalid shape"); - DG_STATIC_ASSERT(L2_SHAPE_N % (BLOCK_N * 2) == 0, "Invalid shape"); + DG_STATIC_ASSERT(kNumCTAsPerTask == 1 or kNumCTAsPerTask == 2, + "MegaMoE task groups support one or two CTAs"); + DG_STATIC_ASSERT(L1_SHAPE_N % (BLOCK_N * kNumCTAsPerTask) == 0, "Invalid shape"); + DG_STATIC_ASSERT(L2_SHAPE_N % (BLOCK_N * kNumCTAsPerTask) == 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(SHARED_L1_SHAPE_N % (BLOCK_N * 2) == 0, "Invalid shared shape"); - DG_STATIC_ASSERT(SHARED_L2_SHAPE_N % (BLOCK_N * 2) == 0, "Invalid shared shape"); + DG_STATIC_ASSERT(SHARED_L1_SHAPE_N % (BLOCK_N * kNumCTAsPerTask) == 0, + "Invalid shared shape"); + DG_STATIC_ASSERT(SHARED_L2_SHAPE_N % (BLOCK_N * kNumCTAsPerTask) == 0, + "Invalid shared shape"); DG_STATIC_ASSERT(SHARED_L1_SHAPE_K % BLOCK_K == 0, "Invalid shared shape"); DG_STATIC_ASSERT(SHARED_L2_SHAPE_K % BLOCK_K == 0, "Invalid shared shape"); - // NOTES: N block counts must be even so that 2 adjacent CTAs in a cluster - // always land on the same m_block_idx with n_block_idx differing by 1 - DG_STATIC_ASSERT(kNumSMs % 2 == 0, "Number of SMs must be even for 2-CTA cluster"); + // SM100 groups two adjacent cluster CTAs; SM90 can publish the same task + // payload to one CTA-local consumer group. + DG_STATIC_ASSERT(kNumSMs % kNumCTAsPerTask == 0, + "Number of workers must divide into complete task groups"); DG_STATIC_ASSERT(kNumRingBlocks > 0, "Invalid ring buffer config"); // Workspace @@ -212,6 +222,7 @@ struct MegaMoEScheduler { CUTLASS_DEVICE bool get_next_task(task_info_t& task_info) { task_info_full_barriers[sched_stage_idx].wait(sched_phase); + asm volatile("" ::: "memory"); task_info = task_infos[sched_stage_idx]; advance_sched_pipeline(); return task_info.is_valid(); @@ -246,8 +257,19 @@ struct MegaMoEScheduler { return get_pool_block_offset(kNumExpertsPerRank); } - CUTLASS_DEVICE void fetch_expert_recv_count() { - // NOTES: each lane caches experts at indices (i * 32 + lane_idx) + CUTLASS_DEVICE void finalize_expert_recv_count() { + __syncwarp(); + + num_total_m_blocks = get_num_total_pool_blocks(); + const uint32_t num_total_l1_tasks = num_total_m_blocks * kNumL1Clusters; + const uint32_t num_task_groups = kNumSMs / kNumCTAsPerTask; + const uint32_t num_total_l1_waves = math::ceil_div(num_total_l1_tasks, num_task_groups); + const uint32_t min_l1_warmup_waves = get_num_l1_warmup_waves( + num_total_m_blocks, num_task_groups, kNumL1Clusters, kNumL2Clusters); + num_sched_l1_waves = cute::min(min_l1_warmup_waves, num_total_l1_waves); + } + + CUTLASS_DEVICE void cache_expert_recv_count(uint32_t* smem_counts) const { #pragma unroll for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { const auto expert_idx = i * 32 + ptx::get_lane_idx(); @@ -256,17 +278,40 @@ struct MegaMoEScheduler { do { value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx)); } while (static_cast(value >> 32) != kNumSMs * kNumRanks); + ptx::st_shared(smem_counts + expert_idx, + static_cast(value)); + } + } + } + + 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(); + finalize_expert_recv_count(); + } - num_total_m_blocks = get_num_total_pool_blocks(); - const uint32_t num_total_l1_tasks = num_total_m_blocks * kNumL1Clusters; - const uint32_t num_total_l1_waves = math::ceil_div(num_total_l1_tasks, kNumSMs / 2); - const uint32_t min_l1_warmup_waves = get_num_l1_warmup_waves( - num_total_m_blocks, kNumSMs / 2, kNumL1Clusters, kNumL2Clusters); - num_sched_l1_waves = cute::min(min_l1_warmup_waves, num_total_l1_waves); + CUTLASS_DEVICE void fetch_cached_expert_recv_count( + const uint32_t* smem_counts) { + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + const auto expert_idx = i * 32 + ptx::get_lane_idx(); + stored_num_tokens_per_expert[i] = + expert_idx < kNumExpertsPerRank ? + ptx::ld_shared(smem_counts + expert_idx) : 0u; + } + finalize_expert_recv_count(); } CUTLASS_DEVICE task_info_t create_task(const BlockPhase& block_phase, @@ -279,6 +324,30 @@ struct MegaMoEScheduler { const uint32_t n_cluster_idx = task_idx % num_clusters; task_info_t result(block_phase, 0, 0, n_cluster_idx, m_block_idx, 0, shape_n, shape_k); + // DSV4 Flash places one local expert on each lane. At small M, an + // expert normally owns at most one M64 block, so the pool-block index + // is just its ordinal in the nonempty-lane mask. Avoid rebuilding the + // same warp prefix sum for every L1/L2 N tile. A skewed expert with + // more than BLOCK_M tokens takes the general multi-block path below. + if constexpr (kSingleBlockTaskFastPath and + kNumExpertsPerLane == 1) { + const uint32_t expert_idx = lane_idx; + const uint32_t num_tokens = stored_num_tokens_per_expert[0]; + const uint32_t multi_block_mask = __ballot_sync( + 0xffffffff, expert_idx < kNumExpertsPerRank and + num_tokens > BLOCK_M); + if (multi_block_mask == 0) { + const uint32_t owner_mask = __ballot_sync( + 0xffffffff, expert_idx < kNumExpertsPerRank and + num_tokens != 0); + const uint32_t owner_lane_idx = __fns( + owner_mask, 0, m_block_idx + 1); + result.local_expert_idx = owner_lane_idx; + result.valid_m = ptx::exchange(num_tokens, owner_lane_idx); + return result; + } + } + uint32_t block_offset = 0; #pragma unroll for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { @@ -350,20 +419,33 @@ struct MegaMoEScheduler { } CUTLASS_DEVICE void publish_task(const task_info_t& task_info, const uint32_t& lane_idx) { - if (lane_idx < 2) { - task_info_full_barriers[sched_stage_idx].arrive_and_expect_tx(sizeof(task_info_t), lane_idx); - ptx::st_async_cluster( - task_infos + sched_stage_idx, task_info, - lane_idx, task_info_full_barriers[sched_stage_idx] - ); + if constexpr (kNumCTAsPerTask == 1) { + if (cute::elect_one_sync()) { + task_infos[sched_stage_idx] = task_info; + __threadfence_block(); + task_info_full_barriers[sched_stage_idx].arrive(); + } + } else { + if (lane_idx < kNumCTAsPerTask) { + task_info_full_barriers[sched_stage_idx].arrive_and_expect_tx( + sizeof(task_info_t), lane_idx); + ptx::st_async_cluster( + task_infos + sched_stage_idx, task_info, + lane_idx, task_info_full_barriers[sched_stage_idx]); + } } __syncwarp(); advance_sched_pipeline(); } + CUTLASS_DEVICE static uint32_t get_n_block_idx( + const task_info_t& task_info, const uint32_t& cta_rank = 0) { + return task_info.n_cluster_idx * kNumCTAsPerTask + cta_rank; + } + template CUTLASS_DEVICE void shared_mainloop(const uint32_t& num_tokens, const uint32_t& lane_idx, const uint32_t* task_count_ptr) { - constexpr uint32_t kNumNClusters = kShapeN / BLOCK_N / 2; + constexpr uint32_t kNumNClusters = kShapeN / BLOCK_N / kNumCTAsPerTask; const uint32_t num_m_blocks = math::ceil_div(num_tokens, BLOCK_M); const uint32_t num_tasks = num_m_blocks * kNumNClusters; while (true) { @@ -412,6 +494,78 @@ struct MegaMoEScheduler { task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1); publish_task(task_info_t(BlockPhase::None, 0, 0, 0, 0, 0, 0, 0), lane_idx); } + + // SM90 reuses the weight-loader warp as the task producer. It must publish + // each payload before issuing the matching B load so the A loader and math + // consumers observe exactly the same dynamic schedule. SM100 keeps using + // `mainloop()` above from its dedicated scheduler warp. + template + CUTLASS_DEVICE void shared_mainloop_with_task( + const uint32_t& num_tokens, + const uint32_t& lane_idx, + const uint32_t* task_count_ptr, + Func&& process_task) { + constexpr uint32_t kNumNClusters = + kShapeN / BLOCK_N / kNumCTAsPerTask; + const uint32_t num_m_blocks = math::ceil_div(num_tokens, BLOCK_M); + const uint32_t num_tasks = num_m_blocks * kNumNClusters; + while (true) { + task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1); + const uint32_t task_idx = get_next_task_idx(task_count_ptr); + if (task_idx >= num_tasks) + break; + const uint32_t m_block_idx = task_idx / kNumNClusters; + const uint32_t n_cluster_idx = task_idx % kNumNClusters; + const uint32_t valid_m = + cute::min(num_tokens - m_block_idx * BLOCK_M, BLOCK_M); + const task_info_t task_info( + kBlockPhase, 0, m_block_idx, n_cluster_idx, m_block_idx, + valid_m, kShapeN, kShapeK); + publish_task(task_info, lane_idx); + process_task(task_info); + } + } + + template + CUTLASS_DEVICE void mainloop_with_task( + const uint32_t& num_tokens, Func&& process_task) { + const auto lane_idx = ptx::get_lane_idx(); + + if constexpr (kHasShared) { + shared_mainloop_with_task< + BlockPhase::SharedLinear1, + SHARED_L1_SHAPE_N, + SHARED_L1_SHAPE_K>( + num_tokens, lane_idx, + workspace.get_shared_l1_task_count_ptr(), process_task); + } + + if constexpr (kFetchExpertRecvCount) + fetch_expert_recv_count(); + task_info_t task_info; + do { + task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1); + task_info = get_next_task(); + if (task_info.is_valid()) { + publish_task(task_info, lane_idx); + process_task(task_info); + } + } while (task_info.is_valid()); + + if constexpr (kHasShared) { + shared_mainloop_with_task< + BlockPhase::SharedLinear2, + SHARED_L2_SHAPE_N, + SHARED_L2_SHAPE_K>( + num_tokens, lane_idx, + workspace.get_shared_l2_task_count_ptr(), process_task); + } + + task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1); + publish_task( + task_info_t(BlockPhase::None, 0, 0, 0, 0, 0, 0, 0), lane_idx); + } }; #endif diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 2314564e22..5582b5c0eb 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -3,6 +3,14 @@ import warnings from typing import Tuple, Optional, Union from ..utils.math import align +from .mxfp4 import ( + _normalize_mxfp4_ue8m0, + _process_mxfp4_e8m0, + _reorder_mxfp4_sign_bits_for_sm90, + _restore_mxfp4_sign_bits_from_sm90, + transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90, + _validate_processed_mxfp4_kernel_weights, +) # noinspection PyBroadException try: @@ -15,6 +23,11 @@ from .. import _C +_SM90_MEGA_MOE_OP_NAMES = { + 'fp8xmxfp4': 'fp8_mxfp4_mega_moe', +} + + class SymmBuffer: def __init__(self, group: dist.ProcessGroup, num_experts: int, @@ -66,6 +79,76 @@ def destroy(self): self.x_sf = None +class SM90SymmBuffer: + """Symmetric buffer for the SM90 persistent MegaMoE backend. + + The Hopper backend follows the common twelve-view live-ring/shared-expert + ABI while retaining its architecture-specific FP32 K128 input scales and + FP32 K64 intermediate scales. + """ + def __init__(self, group: dist.ProcessGroup, + num_experts: int, + num_max_tokens_per_rank: int, num_topk: int, + hidden: int, intermediate_hidden: int, + mma_type: str = 'fp8xmxfp4', + activation: str = 'swiglu', + num_shared_experts: int = 0): + if num_shared_experts < 0: + raise ValueError('num_shared_experts must be non-negative') + if mma_type not in _SM90_MEGA_MOE_OP_NAMES: + raise ValueError( + f'Unsupported SM90 MegaMoE mma_type `{mma_type}`; ' + f'supported values: {tuple(_SM90_MEGA_MOE_OP_NAMES)}') + if activation != 'swiglu': + raise ValueError( + f'Only `swiglu` activation is supported, got `{activation}`') + + 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 + self.num_shared_experts = num_shared_experts + self.mma_type = mma_type + + 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, + mma_type, activation, num_shared_experts, + ) + 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() + + (self.x, self.x_sf, + self.topk_idx, self.topk_weights, + self.shared_l1_acts, self.shared_l1_acts_sf, + self.shared_l2_acts, self.shared_l2_acts_sf, + 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 + for name in ( + 'x', 'x_sf', 'topk_idx', 'topk_weights', + 'shared_l1_acts', 'shared_l1_acts_sf', + 'shared_l2_acts', 'shared_l2_acts_sf', + 'l1_acts', 'l1_acts_sf', 'l2_acts', 'l2_acts_sf', + ): + setattr(self, name, None) + self.buffer = None + self.group = None + + def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, num_experts: int, num_max_tokens_per_rank: int, num_topk: int, @@ -94,6 +177,30 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, ) +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, + mma_type: str = 'fp8xmxfp4', + activation: str = 'swiglu', + num_shared_experts: int = 0) -> SM90SymmBuffer: + """Allocate the twelve-view SM90 live-ring MegaMoE symmetric buffer. + + ``num_experts`` is global. Each rank supplies ``num_experts / group.size()`` + weights for its contiguous global expert-ID interval. + """ + num_max_tokens_per_rank = align( + num_max_tokens_per_rank, _C.get_token_alignment_for_sm90_mega_moe()) + return SM90SymmBuffer( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + mma_type=mma_type, + activation=activation, + num_shared_experts=num_shared_experts, + ) + + 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 @@ -149,6 +256,63 @@ def transform_weights_for_mega_moe( return l1_transformed, l2_transformed +def transform_shared_weights_for_fp8_mega_moe_sm90( + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], + activation: str = 'swiglu', +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], + Tuple[torch.Tensor, torch.Tensor]]: + """Prepare replicated SM90 shared-expert FP8 weights. + + Input scales are natural row-major FP32 block-(128, 128) tensors. Only the + L1 FP8 rows are gate/up-interleaved at granularity 8; unlike the SM100 + helper, neither scale tensor is UTCCP-transposed or interleaved. + """ + if activation != 'swiglu': + raise ValueError( + f'Only `swiglu` activation is supported, got `{activation}`') + if not (isinstance(l1_weights, tuple) and len(l1_weights) == 2 and + isinstance(l2_weights, tuple) and len(l2_weights) == 2): + raise TypeError('shared SM90 weights must be FP8/FP32-scale pairs') + l1_weight, l1_scale = l1_weights + l2_weight, l2_scale = l2_weights + if not all(isinstance(tensor, torch.Tensor) for tensor in ( + l1_weight, l1_scale, l2_weight, l2_scale)): + raise TypeError('shared SM90 weights and scales must be torch.Tensor objects') + if l1_weight.dtype != torch.float8_e4m3fn or \ + l2_weight.dtype != torch.float8_e4m3fn: + raise TypeError('shared SM90 weights must use torch.float8_e4m3fn') + if l1_scale.dtype != torch.float32 or l2_scale.dtype != torch.float32: + raise TypeError('shared SM90 weight scales must use torch.float32') + if not (l1_weight.device == l1_scale.device == + l2_weight.device == l2_scale.device): + raise ValueError('shared SM90 weights and scales must be on the same device') + if l1_weight.dim() != 2 or l2_weight.dim() != 2: + raise ValueError('shared SM90 weights must be two-dimensional') + shared_intermediate_hidden = l2_weight.size(1) + hidden = l2_weight.size(0) + if tuple(l1_weight.shape) != (2 * shared_intermediate_hidden, hidden): + raise ValueError( + 'shared L1/L2 weight shapes must be [2*S*I,H] and [H,S*I]') + if hidden % 128 != 0 or shared_intermediate_hidden % 128 != 0: + raise ValueError('shared SM90 H and S*I must be divisible by 128') + if tuple(l1_scale.shape) != ( + 2 * shared_intermediate_hidden // 128, hidden // 128) or \ + tuple(l2_scale.shape) != ( + hidden // 128, shared_intermediate_hidden // 128): + raise ValueError( + 'shared SM90 scales must use natural block-(128,128) shapes') + if not all(tensor.is_contiguous() for tensor in ( + l1_weight, l1_scale, l2_weight, l2_scale)): + raise ValueError('shared SM90 weights and scales must be contiguous') + if not all( + bool(torch.isfinite(scale).all().item()) and + bool((scale > 0).all().item()) + for scale in (l1_scale, l2_scale)): + raise ValueError('shared SM90 weight scales must contain finite positive values') + return (_interleave_weights(l1_weight), l1_scale), (l2_weight, l2_scale) + + def fp8_fp4_mega_moe(y: torch.Tensor, l1_weights: Tuple[torch.Tensor, torch.Tensor], @@ -175,6 +339,129 @@ def fp8_fp4_mega_moe(y: torch.Tensor, fast_math ) + +def _validate_sm90_mxfp4_routed_weights( + l1_weights: Tuple[torch.Tensor, ...], + l2_weights: Tuple[torch.Tensor, ...], + sym_buffer: SM90SymmBuffer, +) -> None: + _validate_processed_mxfp4_kernel_weights(l1_weights, l2_weights) + num_ranks = sym_buffer.group.size() + if sym_buffer.num_experts % num_ranks != 0: + raise ValueError('global num_experts must be divisible by the number of ranks') + expected_local_experts = sym_buffer.num_experts // num_ranks + if l1_weights[0].size(0) != expected_local_experts: + raise ValueError( + 'SM90 MXFP4 weights must contain the contiguous local expert shard: ' + f'expected E_local={expected_local_experts}, got {l1_weights[0].size(0)}') + expected_l1_shape = ( + expected_local_experts, + 2 * sym_buffer.intermediate_hidden, + sym_buffer.hidden // 2, + ) + expected_l2_shape = ( + expected_local_experts, + sym_buffer.hidden, + sym_buffer.intermediate_hidden // 2, + ) + if tuple(l1_weights[0].shape) != expected_l1_shape or \ + tuple(l2_weights[0].shape) != expected_l2_shape: + raise ValueError( + 'SM90 routed MXFP4 weights do not match the symmetric-buffer ' + f'H/I contract: expected {expected_l1_shape} and ' + f'{expected_l2_shape}, got {tuple(l1_weights[0].shape)} and ' + f'{tuple(l2_weights[0].shape)}') + + +_SM90_MEGA_MOE_WEIGHT_VALIDATORS = { + 'fp8xmxfp4': _validate_sm90_mxfp4_routed_weights, +} + + +def fp8_mega_moe(y: torch.Tensor, + l1_weights: Tuple[torch.Tensor, ...], + l2_weights: Tuple[torch.Tensor, ...], + sym_buffer: SM90SymmBuffer, + shared_l1_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l2_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + recipe: Tuple[int, int, int] = (1, 1, 32), + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True): + """Run an SM90 FP8-activation MegaMoE kernel selected by ``mma_type``. + + ``sym_buffer.mma_type`` selects the routed-weight backend. The current + ``fp8xmxfp4`` backend requires processed triples + ``(processed_e2m1, relative_ue8m0, weight_scale_2)`` returned by + :func:`transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90`. Future + FP8/FP8 and FP8/INT4 backends can add dispatch branches without changing + the public launch API. + + When shared experts are enabled, prepare their FP8/FP32 pairs with + :func:`transform_shared_weights_for_fp8_mega_moe_sm90` and copy the + input K128 FP32 scales into ``sym_buffer.shared_l1_acts_sf`` before launch. + ``shared_l1_acts`` itself aliases ``sym_buffer.x``. + """ + if not isinstance(sym_buffer, SM90SymmBuffer): + raise TypeError( + 'fp8_mega_moe requires an SM90SymmBuffer allocated by ' + 'get_symm_buffer_for_sm90_mega_moe') + try: + op_name = _SM90_MEGA_MOE_OP_NAMES[sym_buffer.mma_type] + validate_routed_weights = _SM90_MEGA_MOE_WEIGHT_VALIDATORS[ + sym_buffer.mma_type] + except KeyError as exception: + raise ValueError( + f'Unsupported SM90 MegaMoE mma_type `{sym_buffer.mma_type}`') from exception + if (shared_l1_weights is None) != (shared_l2_weights is None): + raise ValueError( + 'shared_l1_weights and shared_l2_weights must be provided together') + num_shared_experts = getattr(sym_buffer, 'num_shared_experts', 0) + if num_shared_experts == 0 and shared_l1_weights is not None: + raise ValueError( + 'shared weights require an SM90SymmBuffer allocated with ' + 'num_shared_experts > 0') + if num_shared_experts > 0 and shared_l1_weights is None: + raise ValueError( + 'an SM90SymmBuffer with shared experts requires both shared weight tuples') + validate_routed_weights(l1_weights, l2_weights, sym_buffer) + if num_shared_experts > 0: + shared_intermediate_hidden = ( + num_shared_experts * sym_buffer.intermediate_hidden) + expected_shared_l1_shape = ( + 2 * shared_intermediate_hidden, sym_buffer.hidden) + expected_shared_l2_shape = ( + sym_buffer.hidden, shared_intermediate_hidden) + if tuple(shared_l1_weights[0].shape) != expected_shared_l1_shape or \ + tuple(shared_l2_weights[0].shape) != expected_shared_l2_shape: + raise ValueError( + 'SM90 shared FP8 weights do not match the symmetric-buffer ' + f'shared-expert contract: expected {expected_shared_l1_shape} ' + f'and {expected_shared_l2_shape}, got ' + f'{tuple(shared_l1_weights[0].shape)} and ' + f'{tuple(shared_l2_weights[0].shape)}') + try: + op = getattr(_C, op_name) + except AttributeError as exception: + raise RuntimeError( + f'DeepGEMM was built without the SM90 MegaMoE binding `{op_name}` ' + f'required by `fp8_mega_moe(mma_type="{sym_buffer.mma_type}")`; ' + 'rebuild the extension after enabling the SM90 backend') from exception + op( + y, + l1_weights, l2_weights, + shared_l1_weights, shared_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 + ) + def bf16_mega_moe(y: torch.Tensor, l1_weights: torch.Tensor, l2_weights: torch.Tensor, diff --git a/deep_gemm/mega/mxfp4.py b/deep_gemm/mega/mxfp4.py new file mode 100644 index 0000000000..cbb0f385ed --- /dev/null +++ b/deep_gemm/mega/mxfp4.py @@ -0,0 +1,471 @@ +"""SM90 weight preparation for Humming-compatible MXFP4 MegaMoE. + +The public transforms in this module intentionally do not depend on CUDA or +the DeepGEMM extension. Weight conversion is a model-load-time operation and +can therefore be contract-tested on CPU. +""" + +from typing import Optional, Tuple + +import torch + + +MXFP4CheckpointWeights = Tuple[torch.Tensor, torch.Tensor] +_MXFP4Payload = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] +_SM90_MXFP4_COALESCED_SCALE_MAX_HIDDEN = 8192 + + +def _is_valid_sm90_mxfp4_hidden_size(hidden: int) -> bool: + """Match the vectorized SM90 combine-chunk divisibility contract.""" + return hidden % 512 == 0 and (hidden <= 8192 or hidden % 1024 == 0) + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +class MXFP4ProcessedWeights(tuple): + """Validated tuple payload produced by the SM90 MXFP4 transform. + + A tuple subclass preserves the public/PyBind tuple ABI while distinguishing + transformed relative-UE8M0 payloads from raw checkpoint triplets. Semantic + validation happens once at model-load or deserialization time instead of + scanning the full scale tensor before every kernel launch. + """ + + __slots__ = () + + def __new__(cls, values: _MXFP4Payload) -> 'MXFP4ProcessedWeights': + if not isinstance(values, tuple) or len(values) != 3: + raise TypeError( + 'SM90 processed MXFP4 weights require a three-tensor tuple') + weight, relative_sf, secondary = values + if not all(isinstance(tensor, torch.Tensor) for tensor in values): + raise TypeError('SM90 processed MXFP4 entries must be torch.Tensor objects') + if relative_sf.dtype != torch.uint8: + raise TypeError('SM90 relative UE8M0 scales must have dtype uint8') + _require( + bool(((relative_sf >= 1) & (relative_sf <= 12)).all().item()), + 'SM90 processed MXFP4 requires relative UE8M0 values in [1, 12]') + if secondary.dtype != torch.float32: + raise TypeError('SM90 MXFP4 secondary scale must have dtype float32') + _require( + bool(torch.isfinite(secondary).all().item()) and + bool((secondary > 0).all().item()), + 'SM90 MXFP4 secondary scale must contain finite positive values') + return super().__new__(cls, (weight, relative_sf, secondary)) + + def __getnewargs__(self) -> Tuple[_MXFP4Payload]: + return (tuple(self),) + + +def _normalize_mxfp4_ue8m0(sf: torch.Tensor) -> torch.Tensor: + """Return natural-layout UE8M0 exponent bytes. + + Checkpoints normally expose UE8M0 as ``uint8`` or + ``float8_e8m0fnu``. FP32 powers of two are also accepted because + DeepGEMM's quantization utilities use that representation. UE8M0 code 0 + is the finite endpoint ``2**-127``; zero, negative values, non-powers of + two, and infinities are rejected. + """ + if not isinstance(sf, torch.Tensor): + raise TypeError('MXFP4 scale must be a torch.Tensor') + if sf.dtype == torch.uint8: + return sf.contiguous() + + e8m0_dtype = getattr(torch, 'float8_e8m0fnu', None) + if e8m0_dtype is not None and sf.dtype == e8m0_dtype: + return sf.contiguous().view(torch.uint8) + + if sf.dtype != torch.float32: + raise TypeError( + 'MXFP4 scale must have dtype uint8, float8_e8m0fnu, or float32, ' + f'got {sf.dtype}') + + bits = sf.contiguous().view(torch.int32) + exponent = (bits >> 23) & 0xff + mantissa = bits & ((1 << 23) - 1) + is_min_subnormal = bits == (1 << 22) # UE8M0 code 0: 2**-127. + is_normal_power_of_two = ( + ((bits >> 31) == 0) & + (mantissa == 0) & + (exponent >= 1) & + (exponent <= 254) + ) + _require( + bool((is_min_subnormal | is_normal_power_of_two).all().item()), + 'FP32 MXFP4 scales must be finite positive UE8M0 powers of two') + return torch.where( + is_min_subnormal, torch.zeros_like(exponent), exponent).to(torch.uint8) + + +def _normalize_mxfp4_packed_weight(weight: torch.Tensor) -> torch.Tensor: + """Return packed E2M1 bytes using DeepGEMM's signed-byte tensor view.""" + if not isinstance(weight, torch.Tensor): + raise TypeError('MXFP4 weight must be a torch.Tensor') + if weight.dtype not in (torch.uint8, torch.int8): + raise TypeError( + f'packed MXFP4 weight must have dtype uint8 or int8, got {weight.dtype}') + weight = weight.contiguous() + return weight if weight.dtype == torch.int8 else weight.view(torch.int8) + + +def _validate_checkpoint_mxfp4_weights( + weights: MXFP4CheckpointWeights, + name: str, +) -> MXFP4CheckpointWeights: + if not isinstance(weights, tuple) or len(weights) != 2: + raise TypeError(f'{name} must be a (packed_weight, ue8m0_scale) tuple') + weight = _normalize_mxfp4_packed_weight(weights[0]) + sf = _normalize_mxfp4_ue8m0(weights[1]) + _require(weight.dim() == 3, f'{name} weight must have shape [E, N, K/2]') + _require(sf.dim() == 3, f'{name} scale must have shape [E, N, K/32]') + _require(weight.device == sf.device, f'{name} weight and scale must be on the same device') + _require(weight.size(0) > 0 and weight.size(1) > 0 and weight.size(2) > 0, + f'{name} dimensions must be nonzero') + _require(weight.shape[:2] == sf.shape[:2], + f'{name} weight and scale must have matching [E, N] dimensions') + _require(weight.size(2) == sf.size(2) * 16, + f'{name} expects packed K/2 weights and K/32 scales') + return weight, sf + + +def _validate_mxfp4_layer_pair( + l1_weights: MXFP4CheckpointWeights, + l2_weights: MXFP4CheckpointWeights, +) -> None: + l1_w, _ = l1_weights + l2_w, _ = l2_weights + _require(l1_w.device == l2_w.device, + 'L1 and L2 MXFP4 weights must be on the same device') + _require(l1_w.size(0) == l2_w.size(0), + 'L1 and L2 MXFP4 weights must contain the same number of experts') + _require(l1_w.size(1) % 16 == 0, + 'L1 gate/up rows must be divisible by 16 for granularity-8 interleave') + # L1 is [E, 2I, H/2] and L2 is [E, H, I/2]. + _require(l1_w.size(1) == l2_w.size(2) * 4, + 'incompatible L1/L2 shapes: expected L1 N == 4 * L2 packed-K') + _require(l2_w.size(1) == l1_w.size(2) * 2, + 'incompatible L1/L2 shapes: expected L2 N == 2 * L1 packed-K') + hidden = l2_w.size(1) + intermediate_hidden = l1_w.size(1) // 2 + _require( + _is_valid_sm90_mxfp4_hidden_size(hidden), + 'SM90 MXFP4 MegaMoE hidden size must be divisible by 512 and, ' + 'when greater than 8192, divisible by 1024') + _require(intermediate_hidden % 256 == 0, + 'SM90 MXFP4 MegaMoE intermediate hidden size must be divisible by 256') + + +def _interleave_mxfp4_rows(tensor: torch.Tensor, granularity: int = 8) -> torch.Tensor: + """Convert ``[gate | up]`` rows to granularity-8 gate/up interleave.""" + _require(tensor.dim() == 3, 'MXFP4 expert tensor must be three-dimensional') + num_experts, num_rows, *tail = tensor.shape + _require(num_rows % (2 * granularity) == 0, + f'gate/up row count must be divisible by {2 * granularity}') + half = num_rows // 2 + gate = tensor[:, :half].reshape( + num_experts, half // granularity, granularity, *tail) + up = tensor[:, half:].reshape( + num_experts, half // granularity, granularity, *tail) + result = torch.empty_like(tensor) + result.view( + num_experts, half // granularity, 2, granularity, *tail + ).copy_(torch.stack((gate, up), dim=2)) + return result + + +def _transpose_mxfp4_scales_for_sm90(tensor: torch.Tensor) -> torch.Tensor: + """Store K128 scale words contiguously across N for coalesced SM90 loads. + + The returned tensor deliberately retains the public ``[E, N, K/32]`` + shape. Its opaque processed payload is physically ordered as + ``[E, K/128, N, 4]`` and consumed only by the SM90 Humming kernel. + """ + _require(tensor.dim() == 3, + 'MXFP4 scale tensor must be three-dimensional') + num_experts, num_rows, num_k32_groups = tensor.shape + _require(num_k32_groups % 4 == 0, + 'SM90 MXFP4 scale K/32 dimension must be divisible by 4') + return tensor.reshape( + num_experts, num_rows, num_k32_groups // 4, 4 + ).permute(0, 2, 1, 3).contiguous().view(tensor.shape) + + +def _uses_coalesced_mxfp4_scales_sm90(hidden: int) -> bool: + """Match the compile-time scale-load specialization in the SM90 kernel.""" + return hidden <= _SM90_MXFP4_COALESCED_SCALE_MAX_HIDDEN + + +def _restore_mxfp4_scales_from_sm90(tensor: torch.Tensor) -> torch.Tensor: + """Restore an opaque SM90 scale payload to logical ``[E, N, K/32]``.""" + _require(tensor.dim() == 3, + 'MXFP4 scale tensor must be three-dimensional') + num_experts, num_rows, num_k32_groups = tensor.shape + _require(num_k32_groups % 4 == 0, + 'SM90 MXFP4 scale K/32 dimension must be divisible by 4') + return tensor.view( + num_experts, num_k32_groups // 4, num_rows, 4 + ).permute(0, 2, 1, 3).contiguous().view(tensor.shape) + + +def _reorder_mxfp4_sign_bits_for_sm90(weight: torch.Tensor) -> torch.Tensor: + """Reorder sign bits in each packed 32-bit word for the SM90 decoder. + + Magnitude nibbles remain in checkpoint order. Sign bits change from + ``[s0,s1,s2,s3,s4,s5,s6,s7]`` to + ``[s0,s4,s1,s5,s2,s6,s3,s7]`` so the device decoder can form the two FP8 + output words without additional sign permutations. + """ + if weight.dtype not in (torch.uint8, torch.int8): + raise TypeError('SM90 MXFP4 sign reordering requires uint8 or int8 bytes') + _require(weight.size(-1) % 4 == 0, + 'SM90 MXFP4 sign reordering requires K/2 bytes divisible by 4') + original_dtype = weight.dtype + source = weight.contiguous().view(torch.uint8).reshape( + *weight.shape[:-1], -1, 4) + reordered = source & 0x77 + reordered[..., 0] |= (source[..., 0] & 0x08) | ((source[..., 2] & 0x08) << 4) + reordered[..., 1] |= ((source[..., 0] & 0x80) >> 4) | (source[..., 2] & 0x80) + reordered[..., 2] |= (source[..., 1] & 0x08) | ((source[..., 3] & 0x08) << 4) + reordered[..., 3] |= ((source[..., 1] & 0x80) >> 4) | (source[..., 3] & 0x80) + return reordered.reshape(weight.shape).view(original_dtype) + + +def _restore_mxfp4_sign_bits_from_sm90(weight: torch.Tensor) -> torch.Tensor: + """Restore standard consecutive-nibble signs from the SM90 layout.""" + if weight.dtype not in (torch.uint8, torch.int8): + raise TypeError('SM90 MXFP4 sign restoration requires uint8 or int8 bytes') + _require(weight.size(-1) % 4 == 0, + 'SM90 MXFP4 sign restoration requires K/2 bytes divisible by 4') + original_dtype = weight.dtype + source = weight.contiguous().view(torch.uint8).reshape( + *weight.shape[:-1], -1, 4) + restored = source & 0x77 + restored[..., 0] |= (source[..., 0] & 0x08) | ((source[..., 1] & 0x08) << 4) + restored[..., 1] |= (source[..., 2] & 0x08) | ((source[..., 3] & 0x08) << 4) + restored[..., 2] |= ((source[..., 0] & 0x80) >> 4) | (source[..., 1] & 0x80) + restored[..., 3] |= ((source[..., 2] & 0x80) >> 4) | (source[..., 3] & 0x80) + return restored.reshape(weight.shape).view(original_dtype) + + +def _process_mxfp4_e8m0( + weight: torch.Tensor, + sf: torch.Tensor, + interleave_rows: bool = False, +) -> _MXFP4Payload: + """Convert raw UE8M0 scales to bounded per-expert exponent offsets. + + The retained relative exponent range is 11, matching Humming's fused-E8M0 + weight contract. Groups below that range are requantized in packed E2M1; + the returned FP32 secondary scale restores the common expert exponent in + the MegaMoE accumulation path. + """ + weight, raw_sf = _validate_checkpoint_mxfp4_weights((weight, sf), 'MXFP4') + _require(not bool((raw_sf == 255).any().item()), + 'processed MXFP4 does not accept UE8M0 NaN code 255') + if interleave_rows: + _require(weight.size(1) % 16 == 0, + 'processed L1 rows must be divisible by 16') + + sf_i16 = raw_sf.to(torch.int16) + max_exp = sf_i16.flatten(1).amax(dim=1) + min_exp = sf_i16.flatten(1).amin(dim=1) + base_exp = max_exp - torch.minimum( + max_exp - min_exp, torch.full_like(max_exp, 11)) + + base_view = base_exp.view(-1, 1, 1) + clamped = torch.maximum(sf_i16, base_view) + delta = clamped - sf_i16 + # Humming's unsigned exponent construction wraps for delta >= 128. + _require(not bool((delta >= 128).any().item()), + 'processed MXFP4 exponent delta must be less than 128') + offsets = (clamped - base_view + 1).to(torch.uint8) + secondary = torch.exp2(base_exp.float() - 128.0).contiguous() + + # Exact E2M1 requantization lookup for exponent deltas 0..5. Larger + # accepted deltas map all finite magnitudes to zero and reuse row 5. + nibble_lut = torch.tensor([ + 0, 1, 2, 3, 4, 5, 6, 7, 0, 9, 10, 11, 12, 13, 14, 15, + 0, 1, 1, 2, 2, 3, 4, 5, 0, 9, 9, 10, 10, 11, 12, 13, + 0, 0, 1, 1, 1, 2, 2, 3, 0, 8, 9, 9, 9, 10, 10, 11, + 0, 0, 0, 0, 1, 1, 1, 2, 0, 8, 8, 8, 9, 9, 9, 10, + 0, 0, 0, 0, 0, 0, 1, 1, 0, 8, 8, 8, 8, 8, 9, 9, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, + ], dtype=torch.uint8, device=weight.device).view(6, 16) + packed_values = torch.arange(256, dtype=torch.long, device=weight.device) + packed_lut = ( + nibble_lut[:, packed_values & 0x0f] | + (nibble_lut[:, packed_values >> 4] << 4) + ).reshape(-1) + + packed = weight.view(torch.uint8) + rewritten = torch.empty_like(packed) + if interleave_rows: + half_rows = weight.size(1) // 2 + rewritten_view = rewritten.view( + weight.size(0), half_rows // 8, 2, 8, weight.size(2)) + + num_k_groups = raw_sf.size(2) + num_rows_per_chunk = 1024 + for expert_idx in range(weight.size(0)): + row_regions = ((0, weight.size(1), -1),) if not interleave_rows else ( + (0, half_rows, 0), (half_rows, weight.size(1), 1)) + for region_start, region_end, gate_up_idx in row_regions: + for row_start in range(region_start, region_end, num_rows_per_chunk): + row_end = min(row_start + num_rows_per_chunk, region_end) + packed_chunk = packed[expert_idx, row_start:row_end].view( + row_end - row_start, num_k_groups, 16) + delta_chunk = delta[ + expert_idx, row_start:row_end].clamp_max(5).unsqueeze(-1) + lut_idx = delta_chunk.to(torch.long) * 256 + packed_chunk.to(torch.long) + rewritten_chunk = packed_lut[lut_idx].reshape(row_end - row_start, -1) + rewritten_chunk = _reorder_mxfp4_sign_bits_for_sm90(rewritten_chunk) + if interleave_rows: + relative_start = row_start - region_start + relative_end = row_end - region_start + rewritten_view[ + expert_idx, + relative_start // 8:relative_end // 8, + gate_up_idx, + ].copy_(rewritten_chunk.view(-1, 8, weight.size(2))) + else: + rewritten[expert_idx, row_start:row_end].copy_(rewritten_chunk) + + return rewritten.view(torch.int8), offsets.contiguous(), secondary + + +def _apply_weight_scale_2( + scale: torch.Tensor, + secondary: torch.Tensor, + name: str, +) -> torch.Tensor: + if not isinstance(scale, torch.Tensor): + raise TypeError(f'{name}_weight_scale_2 must be a torch.Tensor') + if scale.dtype != torch.float32: + raise TypeError(f'{name}_weight_scale_2 must have dtype float32') + _require(scale.dim() == 1 and scale.numel() == secondary.numel(), + f'{name}_weight_scale_2 must have shape [E]') + _require(scale.device == secondary.device, + f'{name}_weight_scale_2 must be on the weight device') + _require(bool(torch.isfinite(scale).all().item()), + f'{name}_weight_scale_2 must contain only finite values') + _require(bool((scale > 0).all().item()), + f'{name}_weight_scale_2 must contain only positive values') + result = (secondary * scale).contiguous() + _require( + bool(torch.isfinite(result).all().item()) and + bool((result > 0).all().item()), + f'{name}_weight_scale_2 result must contain finite positive values') + return result + + +def transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + l1_weights: MXFP4CheckpointWeights, + l2_weights: MXFP4CheckpointWeights, + l1_weight_scale_2: Optional[torch.Tensor] = None, + l2_weight_scale_2: Optional[torch.Tensor] = None, +) -> Tuple[MXFP4ProcessedWeights, MXFP4ProcessedWeights]: + """Prepare Humming-compatible MXFP4 triples for Hopper MegaMoE. + + Each result is ``(processed_e2m1, relative_ue8m0, weight_scale_2)``. + ``weight_scale_2`` is FP32 ``[E]`` and includes the optional Humming + checkpoint secondary scale. For hidden sizes up to 8192, the scale tensor + keeps its public ``[E, N, K/32]`` shape but stores an opaque physical + ``[E, K/128, N, 4]`` payload for coalesced producer-warp loads. Larger + hidden sizes retain natural row-major storage. This is the only routed + MXFP4 weight contract accepted by the SM90 runtime. + """ + l1 = _validate_checkpoint_mxfp4_weights(l1_weights, 'L1 MXFP4') + l2 = _validate_checkpoint_mxfp4_weights(l2_weights, 'L2 MXFP4') + _validate_mxfp4_layer_pair(l1, l2) + + _require((l1_weight_scale_2 is None) == (l2_weight_scale_2 is None), + 'L1 and L2 weight_scale_2 must either both be provided or both be omitted') + + l1_w, l1_sf, l1_secondary = _process_mxfp4_e8m0( + *l1, interleave_rows=True) + l2_w, l2_sf, l2_secondary = _process_mxfp4_e8m0(*l2) + if l1_weight_scale_2 is not None: + l1_secondary = _apply_weight_scale_2( + l1_weight_scale_2, l1_secondary, 'l1') + l2_secondary = _apply_weight_scale_2( + l2_weight_scale_2, l2_secondary, 'l2') + + hidden = l2_w.size(1) + l1_sf = _interleave_mxfp4_rows(l1_sf) + if _uses_coalesced_mxfp4_scales_sm90(hidden): + l1_sf = _transpose_mxfp4_scales_for_sm90(l1_sf) + l2_sf = _transpose_mxfp4_scales_for_sm90(l2_sf) + + return MXFP4ProcessedWeights(( + l1_w, + l1_sf, + l1_secondary, + )), MXFP4ProcessedWeights(( + l2_w.contiguous(), + l2_sf.contiguous(), + l2_secondary, + )) + + +def _validate_processed_mxfp4_kernel_weights( + l1_weights: MXFP4ProcessedWeights, + l2_weights: MXFP4ProcessedWeights, +) -> None: + """Validate transformed Humming triples for the SM90 wrapper.""" + if not isinstance(l1_weights, tuple) or not isinstance(l2_weights, tuple): + raise TypeError('SM90 MXFP4 weights must be tuples') + if len(l1_weights) != 3 or len(l2_weights) != 3: + raise ValueError( + 'L1/L2 SM90 MXFP4 weights must be Humming processed triples; ' + 'call transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90 first') + + def validate_payload( + weights: MXFP4ProcessedWeights, + name: str, + ) -> Tuple[torch.Tensor, torch.Tensor]: + weight, sf = weights[:2] + if not isinstance(weight, torch.Tensor) or not isinstance(sf, torch.Tensor): + raise TypeError(f'{name} payload entries must be torch.Tensor objects') + if weight.dtype != torch.int8: + raise TypeError( + f'{name} packed weight must be transformed to int8 before launch') + if sf.dtype != torch.uint8: + raise TypeError( + f'{name} scale must be transformed to uint8 before launch') + _require(weight.is_contiguous() and sf.is_contiguous(), + f'{name} weight and scale must be contiguous before launch') + _require(weight.dim() == 3 and sf.dim() == 3, + f'{name} must have [E, N, K/2] weight and [E, N, K/32] scale') + _require(weight.device == sf.device, + f'{name} weight and scale must be on the same device') + _require(weight.size(0) > 0 and weight.size(1) > 0 and weight.size(2) > 0, + f'{name} dimensions must be nonzero') + _require(weight.shape[:2] == sf.shape[:2] and weight.size(2) == sf.size(2) * 16, + f'{name} expects packed K/2 weights and K/32 scales') + return weight, sf + + l1 = validate_payload(l1_weights, 'transformed L1 MXFP4') + l2 = validate_payload(l2_weights, 'transformed L2 MXFP4') + _validate_mxfp4_layer_pair(l1, l2) + for weights, payload, name in ( + (l1_weights, l1, 'l1'), + (l2_weights, l2, 'l2'), + ): + scale = weights[2] + if not isinstance(scale, torch.Tensor): + raise TypeError(f'{name}_weight_scale_2 must be a torch.Tensor') + if scale.dtype != torch.float32: + raise TypeError(f'{name}_weight_scale_2 must have dtype float32') + _require(scale.dim() == 1 and scale.numel() == payload[0].size(0), + f'{name}_weight_scale_2 must have shape [E]') + _require(scale.device == payload[0].device, + f'{name}_weight_scale_2 must be on the weight device') + _require(scale.is_contiguous(), + f'{name}_weight_scale_2 must be contiguous before launch') + if not isinstance(l1_weights, MXFP4ProcessedWeights) or \ + not isinstance(l2_weights, MXFP4ProcessedWeights): + raise ValueError( + 'L1/L2 SM90 MXFP4 weights must be returned by the SM90 transform') diff --git a/tests/bench_mega_moe_sm90.py b/tests/bench_mega_moe_sm90.py new file mode 100644 index 0000000000..f2614a85d4 --- /dev/null +++ b/tests/bench_mega_moe_sm90.py @@ -0,0 +1,597 @@ +"""Benchmark the SM90 Humming FP8 x MXFP4 persistent MegaMoE kernel. + +The benchmark measures the CUDA kernel selected by ``bench_kineto``. Weight +quantization and preprocessing happen before timing, and input-buffer copies +are not included in the reported kernel duration. Distributed results report +both rank-0 time and the maximum rank-local average. + +DeepGEMM imports are intentionally delayed until the CUDA worker starts so +``--help`` and ``py_compile`` remain usable on non-CUDA hosts. +""" + +import argparse +import json +import math +import os +import random +import statistics +import sys +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.distributed as dist + + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + + +MODEL_CONFIGS: Dict[str, Dict[str, int]] = { + "flash": { + "hidden": 4096, + "intermediate_hidden": 2048, + "num_experts": 256, + "num_topk": 6, + }, + "pro": { + "hidden": 7168, + "intermediate_hidden": 3072, + "num_experts": 384, + "num_topk": 6, + }, +} + +DEFAULT_BATCHES = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] +KERNEL_NAME = "sm90_fp8_mxfp4_mega_moe_persistent_impl" + + +class _SingleProcessGroup: + """Minimal rank-0 group for profiler runs that should avoid NCCL.""" + + @staticmethod + def size() -> int: + return 1 + + @staticmethod + def rank() -> int: + return 0 + + @staticmethod + def barrier() -> None: + torch.cuda.synchronize() + + +def _stable_seed(name: str) -> int: + return ( + sum((index + 1) * ord(character) for index, character in enumerate(name)) + & 0x7FFFFFFF + ) + + +def _barrier(group: Any) -> None: + if isinstance(group, _SingleProcessGroup): + group.barrier() + else: + dist.barrier(group=group) + + +def _quantize_grouped_mxfp4( + weight: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Create packed E2M1 ``[E,N,K/2]`` and K32 UE8M0 FP32 scales.""" + from deep_gemm.utils import per_token_cast_to_fp4 + + num_experts, n, k = weight.shape + assert k % 32 == 0 + packed = torch.empty( + (num_experts, n, k // 2), dtype=torch.int8, device=weight.device + ) + scale = torch.empty( + (num_experts, n, k // 32), + dtype=torch.float32, + device=weight.device, + ) + for expert_idx in range(num_experts): + packed[expert_idx], scale[expert_idx] = per_token_cast_to_fp4( + weight[expert_idx], use_ue8m0=True, gran_k=32 + ) + return packed, scale + + +def _quantize_block_fp8( + weight: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a shared-expert weight with FP32 block-(128,128) scales.""" + from deep_gemm.utils import per_block_cast_to_fp8 + + fp8, scale = per_block_cast_to_fp8(weight, use_ue8m0=False, gran_k=128) + assert scale.dtype == torch.float32 + return fp8, scale.contiguous() + + +def _copy_shared_l1_scale(dst: torch.Tensor, src: torch.Tensor) -> None: + """Copy K128 input scales into the shared L1 column-major view.""" + num_tokens, num_scale_columns = src.shape + assert dst.size(0) >= num_tokens and dst.size(1) == num_scale_columns + dst.zero_() + dst[:num_tokens].copy_(src) + + +def _emit_json(prefix: str, payload: Dict[str, Any]) -> None: + print(f"{prefix} {json.dumps(payload, sort_keys=True)}", flush=True) + + +def _cache_mode(flush_l2: int) -> str: + return "cold_l2" if flush_l2 else "no_explicit_l2_flush" + + +def _benchmark_case( + args: argparse.Namespace, + model_name: str, + num_tokens: int, + rank_idx: int, + num_ranks: int, + group: Any, +) -> None: + import deep_gemm + from deep_gemm.testing import bench_kineto + from deep_gemm.utils import per_token_cast_to_fp8 + + model = MODEL_CONFIGS[model_name] + hidden = model["hidden"] + intermediate_hidden = model["intermediate_hidden"] + num_experts = args.num_experts_override or model["num_experts"] + num_topk = model["num_topk"] + num_local_experts = num_experts // num_ranks + num_shared_experts = args.num_shared_experts + + if num_experts % num_ranks != 0: + raise ValueError( + f"num_experts={num_experts} must be divisible by num_ranks={num_ranks}" + ) + if num_ranks > 64: + raise ValueError("SM90 Humming MegaMoE supports at most 64 ranks") + if num_topk > num_experts: + raise ValueError( + f"num_topk={num_topk} must not exceed num_experts={num_experts}" + ) + if num_topk + int(num_shared_experts > 0) > 32: + raise ValueError("num_topk plus the shared-expert contribution must be <= 32") + if num_tokens > args.num_max_tokens_per_rank: + raise ValueError( + f"num_tokens={num_tokens} exceeds capacity {args.num_max_tokens_per_rank}" + ) + + case_seed = ( + args.seed + + rank_idx * 1000003 + + _stable_seed(f"mxfp4:{model_name}:{num_tokens}:shared={num_shared_experts}") + ) + torch.manual_seed(case_seed) + random.seed(case_seed) + + buffer = deep_gemm.get_symm_buffer_for_sm90_mega_moe( + group, + num_experts, + args.num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + num_shared_experts=num_shared_experts, + ) + + try: + x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + x_fp8, x_scale = per_token_cast_to_fp8( + x_bf16, + use_ue8m0=False, + gran_k=128, + use_packed_ue8m0=False, + ) + del x_bf16 + + l1_bf16 = ( + torch.randn( + (num_local_experts, 2 * intermediate_hidden, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + l1_quantized = _quantize_grouped_mxfp4(l1_bf16) + del l1_bf16 + + l2_bf16 = ( + torch.randn( + (num_local_experts, hidden, intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + l2_quantized = _quantize_grouped_mxfp4(l2_bf16) + del l2_bf16 + + transformed_l1, transformed_l2 = ( + deep_gemm.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + l1_quantized, l2_quantized + ) + ) + del l1_quantized, l2_quantized + + scores = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + topk_weights, topk_idx = torch.topk( + scores, num_topk, dim=-1, largest=True, sorted=False + ) + del scores + if args.masked_ratio > 0: + mask = torch.rand_like(topk_idx, dtype=torch.float32) < args.masked_ratio + topk_idx.masked_fill_(mask, -1) + topk_weights.masked_fill_(mask, 0.0) + + transformed_shared_l1: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + transformed_shared_l2: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + if num_shared_experts > 0: + shared_generator = torch.Generator(device="cuda") + shared_generator.manual_seed( + args.seed + _stable_seed(f"shared:{model_name}:{num_shared_experts}") + ) + shared_intermediate_hidden = num_shared_experts * intermediate_hidden + shared_l1_bf16 = ( + torch.randn( + (2 * shared_intermediate_hidden, hidden), + dtype=torch.bfloat16, + device="cuda", + generator=shared_generator, + ) + * 0.05 + ) + shared_l1_quantized = _quantize_block_fp8(shared_l1_bf16) + del shared_l1_bf16 + shared_l2_bf16 = ( + torch.randn( + (hidden, shared_intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + generator=shared_generator, + ) + * 0.05 + ) + shared_l2_quantized = _quantize_block_fp8(shared_l2_bf16) + del shared_l2_bf16 + transformed_shared_l1, transformed_shared_l2 = ( + deep_gemm.transform_shared_weights_for_fp8_mega_moe_sm90( + shared_l1_quantized, shared_l2_quantized + ) + ) + del shared_l1_quantized, shared_l2_quantized + + output = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + cumulative_recv_stats = torch.zeros( + num_local_experts, dtype=torch.int32, device="cuda" + ) + + if num_shared_experts > 0: + assert buffer.shared_l1_acts_sf is not None + _copy_shared_l1_scale(buffer.shared_l1_acts_sf, x_scale) + + def run_kernel() -> torch.Tensor: + buffer.x[:num_tokens].copy_(x_fp8) + buffer.x_sf[:num_tokens].copy_(x_scale) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + deep_gemm.fp8_mega_moe( + output, + transformed_l1, + transformed_l2, + buffer, + shared_l1_weights=transformed_shared_l1, + shared_l2_weights=transformed_shared_l2, + cumulative_local_expert_recv_stats=cumulative_recv_stats, + recipe=(1, 1, 32), + activation="swiglu", + activation_clamp=( + args.activation_clamp + if math.isfinite(args.activation_clamp) + else None + ), + fast_math=bool(args.fast_math), + ) + return output + + case_metadata = { + "implementation": args.implementation, + "kernel": KERNEL_NAME, + "model": model_name, + "m": num_tokens, + "num_ranks": num_ranks, + "hidden": hidden, + "intermediate_hidden": intermediate_hidden, + "num_experts": num_experts, + "num_topk": num_topk, + "num_shared_experts": num_shared_experts, + "num_max_tokens_per_rank": buffer.num_max_tokens_per_rank, + "fast_math": args.fast_math, + "activation_clamp": args.activation_clamp, + "masked_ratio": args.masked_ratio, + "seed": args.seed, + } + + if args.profile_only: + if rank_idx == 0: + _emit_json("PROFILE_CASE_JSON", case_metadata) + use_ncu_range = bool(int(os.environ.get("DG_NCU_RANGE", "0"))) + if use_ncu_range: + torch.cuda.cudart().cudaProfilerStart() + run_kernel() + torch.cuda.synchronize() + _barrier(group) + if use_ncu_range: + torch.cuda.cudart().cudaProfilerStop() + return + + for _ in range(args.num_warmups): + run_kernel() + torch.cuda.synchronize() + _barrier(group) + + local_nonfinite = torch.tensor( + [int(not bool(torch.isfinite(output).all().item()))], + dtype=torch.int32, + device="cuda", + ) + if num_ranks > 1: + dist.all_reduce(local_nonfinite, op=dist.ReduceOp.MAX, group=group) + if local_nonfinite.item() != 0: + raise RuntimeError("warmup produced non-finite output") + + repeats = args.repeats + if repeats is None: + repeats = args.small_repeats if num_tokens <= 128 else args.large_repeats + + rank0_observations = [] + max_rank_observations = [] + for repeat in range(repeats): + kernel_time = bench_kineto( + run_kernel, + KERNEL_NAME, + barrier=lambda: _barrier(group), + num_tests=args.num_tests, + suppress_kineto_output=True, + flush_l2=bool(args.flush_l2), + ) + if not math.isfinite(kernel_time) or kernel_time <= 0: + raise RuntimeError( + f"failed to find a positive duration for {KERNEL_NAME}; " + f"got {kernel_time}" + ) + + max_rank_time = torch.tensor( + kernel_time, dtype=torch.float64, device="cuda" + ) + if num_ranks > 1: + dist.all_reduce(max_rank_time, op=dist.ReduceOp.MAX, group=group) + + rank0_observations.append(kernel_time) + max_rank_observations.append(max_rank_time.item()) + if rank_idx == 0: + _emit_json( + "BENCH_OBS_JSON", + { + **case_metadata, + "repeat": repeat, + "rank0_us": kernel_time * 1e6, + "persistent_rank0_us": kernel_time * 1e6, + "max_rank_us": max_rank_time.item() * 1e6, + "num_tests": args.num_tests, + "flush_l2": bool(args.flush_l2), + "cache_mode": _cache_mode(args.flush_l2), + }, + ) + + if rank_idx == 0: + rank0_median = statistics.median(rank0_observations) + max_rank_median = statistics.median(max_rank_observations) + print( + f"[mxfp4/{model_name}] M={num_tokens:4d} " + f"obs={repeats:2d} " + f"max-rank median={max_rank_median * 1e6:8.1f} us " + f"range={min(max_rank_observations) * 1e6:.1f}-" + f"{max(max_rank_observations) * 1e6:.1f} us " + f"cache={_cache_mode(args.flush_l2)}", + flush=True, + ) + _emit_json( + "BENCH_SUMMARY_JSON", + { + **case_metadata, + "observations": repeats, + "rank0_median_us": rank0_median * 1e6, + "persistent_rank0_median_us": rank0_median * 1e6, + "max_rank_median_us": max_rank_median * 1e6, + "max_rank_min_us": min(max_rank_observations) * 1e6, + "max_rank_max_us": max(max_rank_observations) * 1e6, + "num_tests": args.num_tests, + "num_warmups": args.num_warmups, + "flush_l2": bool(args.flush_l2), + "cache_mode": _cache_mode(args.flush_l2), + "timing_scope": "persistent_kernel_only", + }, + ) + + _barrier(group) + finally: + buffer.destroy() + + +def _benchmark_worker( + local_rank: int, + num_local_ranks: int, + args: argparse.Namespace, +) -> None: + from deep_gemm.testing import get_arch_major + from deep_gemm.utils.dist import init_dist + + if args.no_dist: + if local_rank != 0 or num_local_ranks != 1: + raise ValueError("--no-dist requires exactly one worker") + torch.cuda.set_device(0) + torch.set_default_device("cuda") + rank_idx, num_ranks, group = 0, 1, _SingleProcessGroup() + else: + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + + try: + if get_arch_major() != 9: + raise RuntimeError( + f"SM90 Humming MegaMoE benchmark requires SM90, " + f"got SM{get_arch_major()}0" + ) + + plan = { + "implementation": args.implementation, + "kernel": KERNEL_NAME, + "models": args.model_config, + "batches": args.batches, + "num_ranks": num_ranks, + "profile_only": args.profile_only, + "flush_l2": bool(args.flush_l2), + "cache_mode": _cache_mode(args.flush_l2), + } + if rank_idx == 0: + _emit_json("BENCH_PLAN_JSON", plan) + + for model_name in args.model_config: + model = MODEL_CONFIGS[model_name] + num_experts = args.num_experts_override or model["num_experts"] + if rank_idx == 0: + print( + f"SM90 Humming MegaMoE: model={model_name} " + f"ranks={num_ranks} H={model['hidden']} " + f"I={model['intermediate_hidden']} E={num_experts} " + f"topk={model['num_topk']} " + f"shared={args.num_shared_experts}", + flush=True, + ) + for num_tokens in args.batches: + _benchmark_case( + args, + model_name, + num_tokens, + rank_idx, + num_ranks, + group, + ) + torch.cuda.empty_cache() + _barrier(group) + finally: + if not args.no_dist and dist.is_initialized(): + dist.destroy_process_group() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--num-processes", type=int, default=8) + parser.add_argument( + "--no-dist", + action="store_true", + help="run one GPU without NCCL; requires --num-processes 1", + ) + parser.add_argument("--local-rank-idx", type=int, default=None) + parser.add_argument("--implementation", choices=("mxfp4",), default="mxfp4") + parser.add_argument( + "--model-config", + nargs="+", + choices=sorted(MODEL_CONFIGS), + default=["flash", "pro"], + ) + parser.add_argument("--batches", type=int, nargs="+", default=DEFAULT_BATCHES) + parser.add_argument("--num-max-tokens-per-rank", type=int, default=8192) + parser.add_argument( + "--num-experts-override", + type=int, + default=None, + help="override model expert count, primarily for single-rank profiling", + ) + parser.add_argument("--num-shared-experts", type=int, default=0) + parser.add_argument("--num-warmups", type=int, default=1) + parser.add_argument("--small-repeats", type=int, default=50) + parser.add_argument("--large-repeats", type=int, default=3) + parser.add_argument( + "--repeats", + type=int, + default=None, + help="override repeat count for every token batch", + ) + parser.add_argument("--num-tests", type=int, default=20) + parser.add_argument( + "--flush-l2", + type=int, + choices=(0, 1), + default=1, + help="1 flushes L2 before each sample; 0 disables the explicit flush", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--masked-ratio", type=float, default=0.0) + parser.add_argument("--activation-clamp", type=float, default=10.0) + parser.add_argument("--fast-math", type=int, choices=(0, 1), default=1) + parser.add_argument( + "--profile-only", + "--ncu-profile-only", + dest="profile_only", + action="store_true", + help="launch exactly one selected case without Kineto timing", + ) + args = parser.parse_args() + + if args.num_processes <= 0: + parser.error("--num-processes must be positive") + if args.no_dist and args.num_processes != 1: + parser.error("--no-dist requires --num-processes 1") + if not args.batches or min(args.batches) <= 0: + parser.error("--batches must contain positive token counts") + if args.num_max_tokens_per_rank < max(args.batches): + parser.error("--num-max-tokens-per-rank must cover every batch") + if args.num_experts_override is not None and args.num_experts_override <= 0: + parser.error("--num-experts-override must be positive") + if args.num_shared_experts < 0: + parser.error("--num-shared-experts must be non-negative") + if args.num_warmups <= 0: + parser.error("--num-warmups must be positive") + if args.small_repeats <= 0 or args.large_repeats <= 0: + parser.error("repeat counts must be positive") + if args.repeats is not None and args.repeats <= 0: + parser.error("--repeats must be positive") + if args.num_tests <= 0: + parser.error("--num-tests must be positive") + if not 0 <= args.masked_ratio <= 1: + parser.error("--masked-ratio must be between 0 and 1") + if args.profile_only and (len(args.model_config) != 1 or len(args.batches) != 1): + parser.error("--profile-only requires exactly one --model-config and one batch") + if not args.profile_only and int(os.environ.get("DG_USE_NVIDIA_TOOLS", 0)): + parser.error( + "DG_USE_NVIDIA_TOOLS disables Kineto timing; use --profile-only " + "with NVIDIA tools or unset the variable" + ) + return args + + +if __name__ == "__main__": + benchmark_args = _parse_args() + if benchmark_args.local_rank_idx is not None: + _benchmark_worker( + benchmark_args.local_rank_idx, + benchmark_args.num_processes, + benchmark_args, + ) + elif benchmark_args.no_dist: + _benchmark_worker(0, 1, benchmark_args) + else: + torch.multiprocessing.spawn( + _benchmark_worker, + args=(benchmark_args.num_processes, benchmark_args), + nprocs=benchmark_args.num_processes, + ) diff --git a/tests/test_mega_moe_mxfp4_preprocess.py b/tests/test_mega_moe_mxfp4_preprocess.py new file mode 100644 index 0000000000..026bdab2b5 --- /dev/null +++ b/tests/test_mega_moe_mxfp4_preprocess.py @@ -0,0 +1,524 @@ +"""CPU contract tests for SM90 Humming-compatible MXFP4 weight transforms.""" + +import importlib.util +from pathlib import Path +import sys +import types + +import pytest +import torch + + +# Load the CUDA-independent transform module without importing deep_gemm._C. +_MODULE_PATH = Path(__file__).parents[1] / 'deep_gemm' / 'mega' / 'mxfp4.py' +_SPEC = importlib.util.spec_from_file_location('deep_gemm_mxfp4_contract', _MODULE_PATH) +mxfp4 = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(mxfp4) + + +def _valid_checkpoint_weights( + num_experts: int = 1, + hidden: int = 512, + intermediate: int = 256, +): + l1_w = torch.arange( + num_experts * 2 * intermediate * (hidden // 2), dtype=torch.int32 + ).to(torch.uint8).reshape(num_experts, 2 * intermediate, hidden // 2) + l1_sf = torch.full( + (num_experts, 2 * intermediate, hidden // 32), 109, dtype=torch.uint8) + l2_w = torch.arange( + num_experts * hidden * (intermediate // 2), dtype=torch.int32 + ).to(torch.uint8).reshape(num_experts, hidden, intermediate // 2) + l2_sf = torch.full( + (num_experts, hidden, intermediate // 32), 109, dtype=torch.uint8) + return (l1_w, l1_sf), (l2_w, l2_sf) + + +def _load_mega_api_for_cpu(): + package_name = '_deep_gemm_sm90_api_contract' + package = types.ModuleType(package_name) + package.__path__ = [str(_MODULE_PATH.parents[1])] + package._C = types.SimpleNamespace() + utils = types.ModuleType(f'{package_name}.utils') + utils.__path__ = [str(_MODULE_PATH.parents[1] / 'utils')] + math_module = types.ModuleType(f'{package_name}.utils.math') + math_module.align = lambda value, alignment: ( + (value + alignment - 1) // alignment * alignment) + module_name = f'{package_name}.mega' + spec = importlib.util.spec_from_file_location( + module_name, + _MODULE_PATH.parent / '__init__.py', + submodule_search_locations=[str(_MODULE_PATH.parent)], + ) + mega = importlib.util.module_from_spec(spec) + sys.modules.update({ + package_name: package, + f'{package_name}.utils': utils, + f'{package_name}.utils.math': math_module, + module_name: mega, + }) + spec.loader.exec_module(mega) + return mega, package_name + + +def _unload_fake_package(package_name: str) -> None: + for name in tuple(sys.modules): + if name == package_name or name.startswith(f'{package_name}.'): + sys.modules.pop(name, None) + + +def test_canonical_transform_returns_processed_triples_and_interleaves_l1(): + (l1_w, l1_sf), (l2_w, l2_sf) = _valid_checkpoint_weights() + # Avoid negative-zero E2M1 nibbles, which preprocessing intentionally + # canonicalizes to positive zero, while retaining varied sign bits. + l1_w.bitwise_or_(0x11) + l2_w.bitwise_or_(0x11) + l1_sf.copy_(( + 109 + torch.arange(l1_sf.size(1), dtype=torch.int32) % 4 + ).to(torch.uint8).reshape(1, -1, 1).expand_as(l1_sf)) + l2_sf.copy_(( + 109 + torch.arange(l2_sf.size(1), dtype=torch.int32) % 4 + ).to(torch.uint8).reshape(1, -1, 1).expand_as(l2_sf)) + e8m0_dtype = getattr(torch, 'float8_e8m0fnu', None) + l1_sf_input = l1_sf if e8m0_dtype is None else l1_sf.view(e8m0_dtype) + l2_sf_input = l2_sf if e8m0_dtype is None else l2_sf.view(e8m0_dtype) + + transformed_l1, transformed_l2 = ( + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1_w, l1_sf_input), (l2_w, l2_sf_input))) + half = l1_w.size(1) // 2 + row_order = torch.tensor([ + row + for start in range(0, half, 8) + for row in ( + *range(start, start + 8), + *range(half + start, half + start + 8), + ) + ]) + + assert len(transformed_l1) == len(transformed_l2) == 3 + assert transformed_l1[0].dtype == torch.int8 + assert transformed_l1[1].dtype == torch.uint8 + assert torch.equal( + mxfp4._restore_mxfp4_sign_bits_from_sm90(transformed_l1[0]) + .view(torch.uint8), + l1_w[:, row_order], + ) + assert torch.equal( + mxfp4._restore_mxfp4_scales_from_sm90(transformed_l1[1]), + l1_sf[:, row_order] - 108, + ) + assert transformed_l2[0].dtype == torch.int8 + assert torch.equal( + mxfp4._restore_mxfp4_sign_bits_from_sm90(transformed_l2[0]) + .view(torch.uint8), + l2_w, + ) + assert torch.equal( + mxfp4._restore_mxfp4_scales_from_sm90(transformed_l2[1]), + l2_sf - 108, + ) + assert transformed_l1[2].item() == pytest.approx(2.0 ** -19) + assert transformed_l2[2].item() == pytest.approx(2.0 ** -19) + assert all(tensor.is_contiguous() for tensor in transformed_l1 + transformed_l2) + + +def test_large_hidden_preserves_natural_scale_layout(): + (l1_w, l1_sf), (l2_w, l2_sf) = _valid_checkpoint_weights(hidden=9216) + l1_sf.copy_(( + 109 + torch.arange(l1_sf.numel(), dtype=torch.int32).reshape(l1_sf.shape) % 4 + ).to(torch.uint8)) + l2_sf.copy_(( + 109 + torch.arange(l2_sf.numel(), dtype=torch.int32).reshape(l2_sf.shape) % 4 + ).to(torch.uint8)) + + transformed_l1, transformed_l2 = ( + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1_w, l1_sf), (l2_w, l2_sf))) + + assert torch.equal( + transformed_l1[1], mxfp4._interleave_mxfp4_rows(l1_sf - 108)) + assert torch.equal(transformed_l2[1], l2_sf - 108) + + +def test_fp32_ue8m0_conversion_including_finite_endpoints(): + fp32_scales = torch.tensor( + [2.0 ** -127, 2.0 ** -126, 1.0, 2.0 ** 127], dtype=torch.float32) + assert torch.equal( + mxfp4._normalize_mxfp4_ue8m0(fp32_scales), + torch.tensor([0, 1, 127, 254], dtype=torch.uint8)) + + for invalid in (0.0, -1.0, 1.5, float('inf')): + with pytest.raises(ValueError, match='powers of two'): + mxfp4._normalize_mxfp4_ue8m0( + torch.tensor([invalid], dtype=torch.float32)) + + +def test_hidden_size_contract_matches_vectorized_combine_chunks(): + assert mxfp4._is_valid_sm90_mxfp4_hidden_size(512) + assert mxfp4._is_valid_sm90_mxfp4_hidden_size(8192) + assert not mxfp4._is_valid_sm90_mxfp4_hidden_size(8704) + assert mxfp4._is_valid_sm90_mxfp4_hidden_size(9216) + assert mxfp4._uses_coalesced_mxfp4_scales_sm90(4096) + assert mxfp4._uses_coalesced_mxfp4_scales_sm90(8192) + assert not mxfp4._uses_coalesced_mxfp4_scales_sm90(9216) + + +def test_processed_sign_layout_is_invertible_and_has_golden_word(): + source = torch.tensor([0x91, 0xa2, 0x3b, 0x4c], dtype=torch.uint8) + reordered = mxfp4._reorder_mxfp4_sign_bits_for_sm90(source) + assert torch.equal( + reordered, + torch.tensor([0x91, 0x2a, 0xb3, 0x4c], dtype=torch.uint8)) + assert torch.equal(mxfp4._restore_mxfp4_sign_bits_from_sm90(reordered), source) + + exhaustive = torch.arange(256, dtype=torch.uint8).repeat(4).reshape(1, 1, -1) + assert torch.equal( + mxfp4._restore_mxfp4_sign_bits_from_sm90( + mxfp4._reorder_mxfp4_sign_bits_for_sm90(exhaustive)), + exhaustive) + + +def test_processed_e8m0_requantization_matches_golden_contract(): + packed_codes = torch.tensor( + [0x10, 0x32, 0x54, 0x76, 0x98, 0xba, 0xdc, 0xfe] * 2, + dtype=torch.uint8) + weight = packed_codes.repeat(7, 1).reshape(1, 7, 16) + sf = torch.tensor( + [109, 108, 107, 106, 105, 104, 120], dtype=torch.uint8 + ).reshape(1, 7, 1) + processed_w, relative_sf, weight_scale_2 = ( + mxfp4._process_mxfp4_e8m0(weight, sf)) + expected = torch.tensor([ + [0x10, 0x32, 0x54, 0x76, 0x90, 0xba, 0xdc, 0xfe], + [0x10, 0x21, 0x32, 0x54, 0x90, 0xa9, 0xba, 0xdc], + [0x00, 0x11, 0x21, 0x32, 0x80, 0x99, 0xa9, 0xba], + [0x00, 0x00, 0x11, 0x21, 0x80, 0x88, 0x99, 0xa9], + [0x00, 0x00, 0x00, 0x11, 0x80, 0x88, 0x88, 0x99], + [0x00, 0x00, 0x00, 0x00, 0x80, 0x88, 0x88, 0x88], + [0x10, 0x32, 0x54, 0x76, 0x90, 0xba, 0xdc, 0xfe], + ], dtype=torch.uint8).repeat(1, 2) + + assert torch.equal( + mxfp4._restore_mxfp4_sign_bits_from_sm90(processed_w).view(torch.uint8)[0], + expected) + assert torch.equal( + relative_sf, + torch.tensor([1, 1, 1, 1, 1, 1, 12], dtype=torch.uint8).reshape(1, 7, 1)) + assert torch.equal(weight_scale_2, torch.tensor([2.0 ** -19], dtype=torch.float32)) + + +def test_canonical_transform_applies_checkpoint_weight_scale_2(): + (l1_w, l1_sf), (l2_w, l2_sf) = _valid_checkpoint_weights(num_experts=2) + l1_sf[:, -1] = 120 + l2_sf[:, -1] = 120 + transformed_l1, transformed_l2 = ( + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1_w, l1_sf), + (l2_w, l2_sf), + l1_weight_scale_2=torch.tensor([2.0, 4.0], dtype=torch.float32), + l2_weight_scale_2=torch.tensor([3.0, 5.0], dtype=torch.float32), + )) + + assert len(transformed_l1) == len(transformed_l2) == 3 + assert transformed_l1[0].dtype == transformed_l2[0].dtype == torch.int8 + assert transformed_l1[1].dtype == transformed_l2[1].dtype == torch.uint8 + assert torch.equal( + transformed_l1[2], + torch.tensor([2.0, 4.0], dtype=torch.float32) * (2.0 ** -19)) + assert torch.equal( + transformed_l2[2], + torch.tensor([3.0, 5.0], dtype=torch.float32) * (2.0 ** -19)) + mxfp4._validate_processed_mxfp4_kernel_weights(transformed_l1, transformed_l2) + + +def test_processed_payload_rejects_plain_triplets_and_invalid_semantics(): + checkpoint_l1, checkpoint_l2 = _valid_checkpoint_weights() + processed_l1, processed_l2 = ( + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + checkpoint_l1, checkpoint_l2)) + + with pytest.raises(ValueError, match='returned by the SM90 transform'): + mxfp4._validate_processed_mxfp4_kernel_weights( + tuple(processed_l1), processed_l2) + + raw_relative_sf = torch.full_like(processed_l1[1], 109) + with pytest.raises(ValueError, match=r'relative UE8M0 values in \[1, 12\]'): + mxfp4.MXFP4ProcessedWeights(( + processed_l1[0], raw_relative_sf, processed_l1[2])) + + for invalid_secondary in (float('nan'), float('inf'), 0.0, -1.0): + with pytest.raises(ValueError, match='finite positive'): + mxfp4.MXFP4ProcessedWeights(( + processed_l1[0], processed_l1[1], + torch.tensor([invalid_secondary], dtype=torch.float32), + )) + + +def test_checkpoint_weight_scale_2_overflow_is_rejected(): + (l1_w, l1_sf), (l2_w, l2_sf) = _valid_checkpoint_weights() + l1_sf.fill_(254) + l2_sf.fill_(254) + with pytest.raises(ValueError, match='finite positive'): + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1_w, l1_sf), + (l2_w, l2_sf), + l1_weight_scale_2=torch.tensor([8.0], dtype=torch.float32), + l2_weight_scale_2=torch.tensor([8.0], dtype=torch.float32), + ) + + +def test_invalid_mxfp4_contracts_fail_before_kernel_launch(): + l1, l2 = _valid_checkpoint_weights() + + with pytest.raises(TypeError, match='uint8 or int8'): + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1[0].float(), l1[1]), l2) + with pytest.raises(ValueError, match='packed K/2'): + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1[0][..., :-1], l1[1]), l2) + with pytest.raises(ValueError, match='incompatible L1/L2'): + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + l1, (l2[0][:, :-1], l2[1][:, :-1])) + with pytest.raises(ValueError, match='hidden size must be divisible by 512'): + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1[0][..., :128], l1[1][..., :8]), + (l2[0][:, :256], l2[1][:, :256]), + ) + with pytest.raises(ValueError, match='intermediate hidden size must be divisible by 256'): + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1[0][:, :256], l1[1][:, :256]), + (l2[0][..., :64], l2[1][..., :4]), + ) + with pytest.raises(ValueError, match='NaN code 255'): + bad_sf = torch.full_like(l1[1], 255) + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + (l1[0], bad_sf), l2) + with pytest.raises(ValueError, match='both be provided'): + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + l1, l2, l1_weight_scale_2=torch.ones(1)) + + +def test_kernel_payload_validation_accepts_only_processed_contiguous_triples(): + checkpoint_l1, checkpoint_l2 = _valid_checkpoint_weights(num_experts=2) + with pytest.raises(ValueError, match='processed triples'): + mxfp4._validate_processed_mxfp4_kernel_weights(checkpoint_l1, checkpoint_l2) + + processed_l1, processed_l2 = ( + mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + checkpoint_l1, checkpoint_l2)) + mxfp4._validate_processed_mxfp4_kernel_weights(processed_l1, processed_l2) + noncontiguous_scale = torch.ones(4, dtype=torch.float32)[::2] + assert not noncontiguous_scale.is_contiguous() + with pytest.raises(ValueError, match='contiguous before launch'): + mxfp4._validate_processed_mxfp4_kernel_weights( + (processed_l1[0], processed_l1[1], noncontiguous_scale), + processed_l2, + ) + + +def test_explicit_wrapper_rejects_the_sm100_buffer_abi_before_launch(): + mega, package_name = _load_mega_api_for_cpu() + try: + with pytest.raises(TypeError, match='requires an SM90SymmBuffer'): + mega.fp8_mega_moe(None, (), (), object()) + finally: + _unload_fake_package(package_name) + + +def test_explicit_wrapper_rejects_wrong_local_expert_shard_before_launch(): + mega, package_name = _load_mega_api_for_cpu() + + class FakeGroup: + @staticmethod + def size(): + return 2 + + buffer = mega.SM90SymmBuffer.__new__(mega.SM90SymmBuffer) + buffer.group = FakeGroup() + buffer.mma_type = 'fp8xmxfp4' + buffer.num_experts = 4 + checkpoint_l1, checkpoint_l2 = _valid_checkpoint_weights(num_experts=1) + transformed_l1, transformed_l2 = ( + mega.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + checkpoint_l1, checkpoint_l2)) + try: + with pytest.raises(ValueError, match='expected E_local=2, got 1'): + mega.fp8_mega_moe( + None, transformed_l1, transformed_l2, buffer) + finally: + _unload_fake_package(package_name) + + +def test_explicit_wrapper_forwards_only_processed_triples(): + mega, package_name = _load_mega_api_for_cpu() + calls = [] + + class FakeGroup: + @staticmethod + def size(): + return 1 + + @staticmethod + def rank(): + return 0 + + buffer = mega.SM90SymmBuffer.__new__(mega.SM90SymmBuffer) + buffer.group = FakeGroup() + buffer.mma_type = 'fp8xmxfp4' + buffer.num_experts = 1 + buffer.num_max_tokens_per_rank = 128 + buffer.num_topk = 1 + buffer.hidden = 512 + buffer.intermediate_hidden = 256 + buffer.num_shared_experts = 0 + buffer.buffer = object() + buffer.handle = types.SimpleNamespace(buffer_ptrs=[0x1234]) + checkpoint_l1, checkpoint_l2 = _valid_checkpoint_weights() + processed_l1, processed_l2 = ( + mega.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + checkpoint_l1, checkpoint_l2)) + mega._C = types.SimpleNamespace( + fp8_mxfp4_mega_moe=lambda *args: calls.append(args)) + y = object() + try: + mega.fp8_mega_moe( + y, processed_l1, processed_l2, buffer, + recipe=(1, 1, 32), activation='swiglu', + activation_clamp=10.0, fast_math=False) + finally: + _unload_fake_package(package_name) + + assert len(calls) == 1 + args = calls[0] + assert len(args) == 16 + assert args[0] is y + assert args[1] is processed_l1 and len(args[1]) == 3 + assert args[2] is processed_l2 and len(args[2]) == 3 + assert args[3] is None and args[4] is None + assert args[7] == [0x1234] + assert args[12] == (1, 1, 32) + assert args[13:] == ('swiglu', 10.0, False) + + +def test_sm90_buffer_uses_dedicated_alignment_and_twelve_view_abi(): + mega, package_name = _load_mega_api_for_cpu() + sizing_calls = [] + + class FakeBuffer: + def __init__(self): + self.zeroed = False + + def data_ptr(self): + return 0x1234 + + def zero_(self): + self.zeroed = True + return self + + class FakeGroup: + def __init__(self): + self.barriers = 0 + + def size(self): + return 1 + + def rank(self): + return 0 + + def barrier(self): + self.barriers += 1 + + views = tuple(object() for _ in range(12)) + + def get_size(*args): + sizing_calls.append(args) + return 4096, lambda _: views + + fake_torch = types.SimpleNamespace( + int8=torch.int8, + empty=lambda *args, **kwargs: FakeBuffer(), + cuda=types.SimpleNamespace(synchronize=lambda: None), + ) + mega.torch = fake_torch + mega._C = types.SimpleNamespace( + get_token_alignment_for_sm90_mega_moe=lambda: 128, + get_symm_buffer_size_for_sm90_mega_moe=get_size, + ) + group = FakeGroup() + try: + with pytest.raises(ValueError, match='Unsupported SM90 MegaMoE mma_type'): + mega.get_symm_buffer_for_sm90_mega_moe( + group, 16, 128, 2, 512, 256, mma_type='fp8xint4') + buffer = mega.get_symm_buffer_for_sm90_mega_moe( + group, + num_experts=16, + num_max_tokens_per_rank=17, + num_topk=2, + hidden=512, + intermediate_hidden=256, + ) + assert buffer.mma_type == 'fp8xmxfp4' + assert sizing_calls == [( + 1, 16, 128, 2, 512, 256, 'fp8xmxfp4', 'swiglu', 0)] + assert buffer.num_max_tokens_per_rank == 128 + assert group.barriers == 1 + assert ( + buffer.x, buffer.x_sf, + buffer.topk_idx, buffer.topk_weights, + buffer.shared_l1_acts, buffer.shared_l1_acts_sf, + buffer.shared_l2_acts, buffer.shared_l2_acts_sf, + buffer.l1_acts, buffer.l1_acts_sf, + buffer.l2_acts, buffer.l2_acts_sf, + ) == views + buffer.destroy() + assert all(getattr(buffer, name) is None for name in ( + 'x', 'x_sf', 'topk_idx', 'topk_weights', + 'shared_l1_acts', 'shared_l1_acts_sf', + 'shared_l2_acts', 'shared_l2_acts_sf', + 'l1_acts', 'l1_acts_sf', 'l2_acts', 'l2_acts_sf')) + shared_buffer = mega.get_symm_buffer_for_sm90_mega_moe( + group, 16, 128, 2, 512, 256, num_shared_experts=1) + assert shared_buffer.num_shared_experts == 1 + assert sizing_calls[-1] == ( + 1, 16, 128, 2, 512, 256, 'fp8xmxfp4', 'swiglu', 1) + shared_buffer.destroy() + finally: + _unload_fake_package(package_name) + + +def test_sm90_shared_weight_transform_interleaves_only_l1_fp8_rows(): + mega, package_name = _load_mega_api_for_cpu() + try: + row_values = ( + torch.arange(256, dtype=torch.float32) % 16 + ).to(torch.float8_e4m3fn) + l1_weight = row_values[:, None].expand(256, 128).contiguous() + l2_weight = torch.ones( + (128, 128), dtype=torch.float8_e4m3fn) + l1_scale = torch.tensor([[1.0], [2.0]], dtype=torch.float32) + l2_scale = torch.tensor([[3.0]], dtype=torch.float32) + + transformed_l1, transformed_l2 = ( + mega.transform_shared_weights_for_fp8_mega_moe_sm90( + (l1_weight, l1_scale), (l2_weight, l2_scale))) + + expected_rows = torch.cat((row_values[:8], row_values[128:136])) + assert torch.equal(transformed_l1[0][:16, 0], expected_rows) + assert transformed_l1[1].data_ptr() == l1_scale.data_ptr() + assert transformed_l2[0].data_ptr() == l2_weight.data_ptr() + assert transformed_l2[1].data_ptr() == l2_scale.data_ptr() + + for invalid_scale in (float('nan'), float('inf'), 0.0, -1.0): + bad_l1_scale = torch.full_like(l1_scale, invalid_scale) + with pytest.raises(ValueError, match='finite positive'): + mega.transform_shared_weights_for_fp8_mega_moe_sm90( + (l1_weight, bad_l1_scale), (l2_weight, l2_scale)) + + meta_l1_scale = torch.empty( + l1_scale.shape, dtype=torch.float32, device='meta') + with pytest.raises(ValueError, match='same device'): + mega.transform_shared_weights_for_fp8_mega_moe_sm90( + (l1_weight, meta_l1_scale), (l2_weight, l2_scale)) + finally: + _unload_fake_package(package_name) diff --git a/tests/test_mega_moe_sm90.py b/tests/test_mega_moe_sm90.py new file mode 100644 index 0000000000..52480ac59f --- /dev/null +++ b/tests/test_mega_moe_sm90.py @@ -0,0 +1,1062 @@ +"""Multi-rank runtime validation for SM90 FP8 x MXFP4 MegaMoE. + +The default smoke suite runs routed-only and shared-expert cases through the +single Humming processed-MXFP4 triplet contract. ``--suite standard`` also +covers heuristic token bands, fast-math modes, activation clamps, and 0/1/max +token boundaries. ``--suite full`` adds production-shaped and randomized cases. + +This file intentionally delays importing DeepGEMM until the spawned CUDA +worker starts, so ``--help`` and ``py_compile`` work on non-CUDA hosts. +""" + +import argparse +import math +import os +import random +import sys +from typing import Any, Dict, List, Optional, Tuple + +import pytest +import torch +import torch.distributed as dist + + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + + +Scenario = Tuple[str, Dict[str, Any]] + + +class _CorrectnessMismatch(AssertionError): + """A correctness failure already reduced to every participating rank.""" + + +class _SingleProcessGroup: + """Minimal rank-0 group for sanitizer runs that must avoid NCCL.""" + + @staticmethod + def size() -> int: + return 1 + + @staticmethod + def rank() -> int: + return 0 + + @staticmethod + def barrier() -> None: + torch.cuda.synchronize() + + +def _make_forced_ring_wrap_topk_idx( + num_tokens: int, + num_topk: int, + num_experts: int, + rank_idx: int, + num_ranks: int, + device: torch.device, +) -> torch.Tensor: + """Assign routes so every expert receives at least one global token.""" + if num_tokens * num_topk * num_ranks < num_experts: + raise ValueError( + 'forced ring-wrap routing needs at least one route per expert') + first_route = rank_idx * num_tokens * num_topk + return ( + torch.arange( + first_route, + first_route + num_tokens * num_topk, + dtype=torch.int64, + device=device, + ) % num_experts + ).reshape(num_tokens, num_topk) + + +def _quantize_grouped_mxfp4(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Create packed E2M1 ``[E,N,K/2]`` plus K32 UE8M0 FP32 powers.""" + from deep_gemm.utils import per_token_cast_to_fp4 + + num_experts, n, k = weight.shape + packed = torch.empty( + (num_experts, n, k // 2), dtype=torch.int8, device=weight.device) + sf = torch.empty( + (num_experts, n, k // 32), dtype=torch.float32, device=weight.device) + for expert_idx in range(num_experts): + packed[expert_idx], sf[expert_idx] = per_token_cast_to_fp4( + weight[expert_idx], use_ue8m0=True, gran_k=32) + return packed, sf + + +def _quantize_block_fp8(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D weight with natural FP32 block-(128, 128) scales.""" + from deep_gemm.utils import per_block_cast_to_fp8 + + fp8, sf = per_block_cast_to_fp8( + weight, use_ue8m0=False, gran_k=128) + assert sf.dtype == torch.float32 + assert sf.shape == (weight.size(0) // 128, weight.size(1) // 128) + return fp8, sf.contiguous() + + +def _dequant_block_fp8(weight: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """Restore a 2D FP8 tensor from natural block-(128, 128) scales.""" + n, k = weight.shape + assert n % 128 == 0 and k % 128 == 0 + assert sf.shape == (n // 128, k // 128) + blocks = weight.float().view(n // 128, 128, k // 128, 128) + return ( + blocks * sf[:, None, :, None] + ).view(n, k) + + +def _copy_shared_l1_sf( + dst: torch.Tensor, + src: torch.Tensor, +) -> None: + """Copy K128 FP32 scales into the shared L1 column-major view.""" + num_tokens, num_sf_columns = src.shape + assert dst.size(0) >= num_tokens and dst.size(1) == num_sf_columns + dst.zero_() + if num_tokens > 0: + dst[:num_tokens].copy_(src) + + +def _dequant_mxfp4(packed: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """Dequantize arbitrary-prefix packed E2M1 tensors with K32 scales.""" + *prefix, n, packed_k = packed.shape + k = packed_k * 2 + codes = torch.empty((*prefix, n, k), dtype=torch.int8, device=packed.device) + codes[..., 0::2] = packed & 0x0f + codes[..., 1::2] = (packed >> 4) & 0x0f + magnitudes = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=packed.device, + ) + values = magnitudes[(codes & 0x07).to(torch.long)] + values = torch.where((codes & 0x08) != 0, -values, values) + return ( + values.view(*prefix, n, k // 32, 32) * sf.unsqueeze(-1) + ).view(*prefix, n, k) + + +def _deinterleave_l1(tensor: torch.Tensor, granularity: int = 8) -> torch.Tensor: + """Restore natural ``[gate | up]`` rows from the SM90 SwiGLU layout.""" + num_experts, num_rows, *tail = tensor.shape + half = num_rows // 2 + interleaved = tensor.view( + num_experts, half // granularity, 2, granularity, *tail) + result = torch.empty_like(tensor) + result[:, :half].copy_( + interleaved[:, :, 0].reshape(num_experts, half, *tail)) + result[:, half:].copy_( + interleaved[:, :, 1].reshape(num_experts, half, *tail)) + return result + + +def _restore_sm90_mxfp4_scale_layout( + tensor: torch.Tensor, + hidden: int, +) -> torch.Tensor: + """Restore the processed ``[K128,N,4]`` payload to logical scale rows.""" + if hidden > 8192: + return tensor + num_experts, num_rows, num_k32_groups = tensor.shape + assert num_k32_groups % 4 == 0 + return tensor.view( + num_experts, num_k32_groups // 4, num_rows, 4 + ).permute(0, 2, 1, 3).contiguous().view(tensor.shape) + + +def _swiglu_fp32(gate_up: torch.Tensor, clamp: float) -> torch.Tensor: + half = gate_up.size(-1) // 2 + gate, up = gate_up[..., :half], gate_up[..., half:] + if math.isfinite(clamp): + gate = gate.clamp(max=clamp) + up = up.clamp(min=-clamp, max=clamp) + return torch.nn.functional.silu(gate) * up + + +def _gather_token_sizes( + num_local_tokens: int, + num_ranks: int, + group: dist.ProcessGroup, +) -> List[int]: + if num_ranks == 1: + return [num_local_tokens] + local_size = torch.tensor([num_local_tokens], dtype=torch.long, device='cuda') + all_sizes = torch.empty(num_ranks, dtype=torch.long, device='cuda') + dist.all_gather_into_tensor(all_sizes, local_size, group=group) + return all_sizes.cpu().tolist() + + +def _reference_mega_moe( + x_fp8_local: torch.Tensor, + x_sf_local: torch.Tensor, + topk_idx_local: torch.Tensor, + topk_weights_local: torch.Tensor, + l1_weight_local: torch.Tensor, + l1_sf_local: torch.Tensor, + l2_weight_local: torch.Tensor, + l2_sf_local: torch.Tensor, + shared_l1_weight: Optional[torch.Tensor], + shared_l1_sf: Optional[torch.Tensor], + shared_l2_weight: Optional[torch.Tensor], + shared_l2_sf: Optional[torch.Tensor], + rank_idx: int, + num_ranks: int, + group: dist.ProcessGroup, + num_experts: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, + activation_clamp: float, +) -> torch.Tensor: + """Cross-rank PyTorch oracle independent of the CUDA kernel.""" + from deep_gemm.utils.dist import uneven_all_gather + + sizes = _gather_token_sizes(x_fp8_local.size(0), num_ranks, group) + if sum(sizes) == 0: + return torch.empty((0, hidden), dtype=torch.bfloat16, device='cuda') + + if num_ranks == 1: + x_fp8 = x_fp8_local + x_sf = x_sf_local + topk_idx = topk_idx_local + topk_weights = topk_weights_local + else: + x_fp8 = uneven_all_gather(x_fp8_local, group=group) + x_sf = uneven_all_gather(x_sf_local, group=group) + topk_idx = uneven_all_gather(topk_idx_local, group=group) + topk_weights = uneven_all_gather(topk_weights_local, group=group) + num_global_tokens = x_fp8.size(0) + + x = ( + x_fp8.float().view(num_global_tokens, hidden // 128, 128) * + x_sf.unsqueeze(-1) + ).view(num_global_tokens, hidden) + combine = torch.zeros( + num_global_tokens, + num_topk, + hidden, + dtype=torch.float32, + device='cuda', + ) + num_local_experts = num_experts // num_ranks + local_expert_start = rank_idx * num_local_experts + local_expert_end = local_expert_start + num_local_experts + local_routes = ((topk_idx >= local_expert_start) & + (topk_idx < local_expert_end)) + valid_expert_ids = torch.unique(topk_idx[local_routes]).cpu().tolist() + for expert_id in valid_expert_ids: + routes = (topk_idx == expert_id).nonzero(as_tuple=False) + selected = routes[:, 0] + selected_slots = routes[:, 1] + local_expert = expert_id - local_expert_start + + # Dequantize each routed expert once. This keeps the oracle independent + # of the CUDA kernel while bounding production-shape memory to one + # expert instead of materializing a per-token weight batch. + l1_weight = _dequant_mxfp4( + l1_weight_local[local_expert], + l1_sf_local[local_expert], + ) + l1_output = torch.matmul(x[selected], l1_weight.transpose(0, 1)) + del l1_weight + l1_output = _swiglu_fp32(l1_output, activation_clamp) + l1_output *= topk_weights[selected, selected_slots].unsqueeze(-1) + + num_selected = selected.numel() + l1_view = l1_output.view(num_selected, intermediate_hidden // 64, 64) + l1_scale = l1_view.abs().amax(dim=-1).clamp(1e-4) / 448.0 + l1_quantized = ( + l1_view / l1_scale.unsqueeze(-1) + ).to(torch.float8_e4m3fn).float() + l2_input = ( + l1_quantized * l1_scale.unsqueeze(-1) + ).view(num_selected, intermediate_hidden) + + l2_weight = _dequant_mxfp4( + l2_weight_local[local_expert], + l2_sf_local[local_expert], + ) + l2_output = torch.matmul(l2_input, l2_weight.transpose(0, 1)) + del l2_weight + combine[selected, selected_slots] = l2_output.to(torch.bfloat16).float() + + # Each route is computed only by the rank that owns its expert. Summing + # outputs, instead of replicating all expert weights, keeps the oracle + # independent and memory-bounded at production shapes. + if num_ranks > 1: + dist.all_reduce(combine, op=dist.ReduceOp.SUM, group=group) + + combine = combine.to(torch.bfloat16).float() + if shared_l1_weight is not None: + assert shared_l1_sf is not None + assert shared_l2_weight is not None + assert shared_l2_sf is not None + shared_l1 = _dequant_block_fp8( + shared_l1_weight, shared_l1_sf) + shared_l1_output = torch.matmul(x, shared_l1.transpose(0, 1)) + del shared_l1 + shared_l1_output = _swiglu_fp32( + shared_l1_output, activation_clamp) + + shared_width = shared_l1_output.size(1) + shared_l1_view = shared_l1_output.view( + num_global_tokens, shared_width // 64, 64) + shared_l1_scale = ( + shared_l1_view.abs().amax(dim=-1).clamp(1e-4) / 448.0) + shared_l1_quantized = ( + shared_l1_view / shared_l1_scale.unsqueeze(-1) + ).to(torch.float8_e4m3fn).float() + shared_l2_input = ( + shared_l1_quantized * shared_l1_scale.unsqueeze(-1) + ).view(num_global_tokens, shared_width) + + shared_l2 = _dequant_block_fp8( + shared_l2_weight, shared_l2_sf) + shared_output = torch.matmul( + shared_l2_input, shared_l2.transpose(0, 1)) + del shared_l2 + combine = torch.cat( + [combine, shared_output.to(torch.bfloat16).float().unsqueeze(1)], + dim=1, + ) + + output = combine.sum(dim=1).to(torch.bfloat16) + rank_start = sum(sizes[:rank_idx]) + return output[rank_start:rank_start + sizes[rank_idx]].contiguous() + + +def _run_scenario( + name: str, + config: Dict[str, Any], + rank_idx: int, + num_ranks: int, + group: dist.ProcessGroup, + diff_tolerance: float, +) -> float: + import deep_gemm + from deep_gemm.mega import _restore_mxfp4_sign_bits_from_sm90 + from deep_gemm.testing import calc_diff + from deep_gemm.utils import per_token_cast_to_fp8 + + num_max_tokens = config['num_max_tokens_per_rank'] + num_tokens = config.get('num_tokens', num_max_tokens) + hidden = config['hidden'] + intermediate_hidden = config['intermediate_hidden'] + num_experts = config['num_experts'] + num_topk = config['num_topk'] + num_shared_experts = config.get('num_shared_experts', 0) + fast_math = config.get('fast_math', True) + activation_clamp = config.get('activation_clamp', 10.0) + masked_ratio = config.get('masked_ratio', 0.0) + repeat_count = config.get('repeat_count', 1) + + assert num_tokens <= num_max_tokens + assert num_experts % num_ranks == 0 + assert num_max_tokens % 128 == 0 + assert hidden % 512 == 0 and intermediate_hidden % 256 == 0 + assert num_shared_experts >= 0 + assert num_topk + int(num_shared_experts > 0) <= 32 + assert repeat_count > 0 + num_local_experts = num_experts // num_ranks + + seed = rank_idx * 100003 + sum( + (index + 1) * ord(character) for index, character in enumerate(name)) + torch.manual_seed(seed) + random.seed(seed) + + if num_tokens == 0: + x_fp8 = torch.empty( + (0, hidden), dtype=torch.float8_e4m3fn, device='cuda') + x_sf = torch.empty((0, hidden // 128), dtype=torch.float32, device='cuda') + else: + x_bf16 = torch.randn( + num_tokens, hidden, dtype=torch.bfloat16, device='cuda') + x_fp8, x_sf = per_token_cast_to_fp8( + x_bf16, + use_ue8m0=False, + gran_k=128, + use_packed_ue8m0=False, + ) + del x_bf16 + + l1_bf16 = torch.randn( + num_local_experts, + 2 * intermediate_hidden, + hidden, + dtype=torch.bfloat16, + device='cuda', + ) * 0.05 + l1_quantized = _quantize_grouped_mxfp4(l1_bf16) + del l1_bf16 + l2_bf16 = torch.randn( + num_local_experts, + hidden, + intermediate_hidden, + dtype=torch.bfloat16, + device='cuda', + ) * 0.05 + l2_quantized = _quantize_grouped_mxfp4(l2_bf16) + del l2_bf16 + + scores = torch.randn( + num_tokens, num_experts, dtype=torch.float32, device='cuda') + topk_weights, topk_idx = torch.topk( + scores, num_topk, dim=-1, largest=True, sorted=False) + if config.get('require_ring_wrap', False): + topk_idx = _make_forced_ring_wrap_topk_idx( + num_tokens, num_topk, num_experts, + rank_idx, num_ranks, scores.device) + if masked_ratio: + mask = torch.rand_like(topk_idx, dtype=torch.float32) < masked_ratio + topk_idx.masked_fill_(mask, -1) + topk_weights.masked_fill_(mask, 0.0) + + # Independent cumulative receive-count contract for local experts. + if num_ranks == 1: + global_topk_idx = topk_idx + else: + from deep_gemm.utils.dist import uneven_all_gather + global_topk_idx = uneven_all_gather(topk_idx, group=group) + valid_topk_idx = global_topk_idx[global_topk_idx >= 0].long() + expected_global_recv_stats = torch.bincount( + valid_topk_idx, minlength=num_experts) + local_expert_start = rank_idx * num_local_experts + local_recv_counts = expected_global_recv_stats[ + local_expert_start:local_expert_start + num_local_experts + ].to(torch.int32) + expected_recv_stats = local_recv_counts * repeat_count + + transformed_l1, transformed_l2 = ( + deep_gemm.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( + l1_quantized, l2_quantized)) + assert len(transformed_l1) == 3 and len(transformed_l2) == 3 + reference_l1_weight = _restore_mxfp4_sign_bits_from_sm90( + _deinterleave_l1(transformed_l1[0])) + reference_l1_sf = ( + transformed_l1[2][:, None, None] * + torch.exp2(_deinterleave_l1( + _restore_sm90_mxfp4_scale_layout( + transformed_l1[1], hidden)).float())) + reference_l2_weight = _restore_mxfp4_sign_bits_from_sm90( + transformed_l2[0]) + reference_l2_sf = ( + transformed_l2[2][:, None, None] * + torch.exp2( + _restore_sm90_mxfp4_scale_layout( + transformed_l2[1], hidden).float())) + + transformed_shared_l1 = transformed_shared_l2 = None + reference_shared_l1_weight = reference_shared_l1_sf = None + reference_shared_l2_weight = reference_shared_l2_sf = None + if num_shared_experts > 0: + # Shared weights are replicated, so use a rank-independent generator. + # Keeping this separate from the routed RNG also makes the routed-only + # scenarios byte-for-byte stable when shared coverage is added. + shared_generator = torch.Generator(device='cuda') + shared_generator.manual_seed( + 700001 + sum( + (index + 1) * ord(character) + for index, character in enumerate(name))) + shared_intermediate_hidden = ( + num_shared_experts * intermediate_hidden) + shared_l1_bf16 = torch.randn( + 2 * shared_intermediate_hidden, + hidden, + dtype=torch.bfloat16, + device='cuda', + generator=shared_generator, + ) * 0.05 + shared_l2_bf16 = torch.randn( + hidden, + shared_intermediate_hidden, + dtype=torch.bfloat16, + device='cuda', + generator=shared_generator, + ) * 0.05 + shared_l1_quantized = _quantize_block_fp8(shared_l1_bf16) + shared_l2_quantized = _quantize_block_fp8(shared_l2_bf16) + del shared_l1_bf16, shared_l2_bf16 + transformed_shared_l1, transformed_shared_l2 = ( + deep_gemm.transform_shared_weights_for_fp8_mega_moe_sm90( + shared_l1_quantized, shared_l2_quantized)) + reference_shared_l1_weight, reference_shared_l1_sf = ( + shared_l1_quantized) + reference_shared_l2_weight, reference_shared_l2_sf = ( + shared_l2_quantized) + + buffer = deep_gemm.get_symm_buffer_for_sm90_mega_moe( + group, + num_experts, + num_max_tokens, + num_topk, + hidden, + intermediate_hidden, + num_shared_experts=num_shared_experts, + ) + try: + # SM90 shares main's twelve-view host ABI even when shared experts are + # disabled. The routed L2 activation SF remains Hopper-specific FP32 + # K64: one logical scale column per 64 intermediate elements. + assert buffer.shared_l1_acts.data_ptr() == buffer.x.data_ptr() + if num_shared_experts == 0: + assert buffer.shared_l1_acts_sf is None + assert buffer.shared_l2_acts is None + assert buffer.shared_l2_acts_sf is None + else: + shared_intermediate_hidden = ( + num_shared_experts * intermediate_hidden) + assert buffer.shared_l1_acts_sf is not None + assert buffer.shared_l2_acts is not None + assert buffer.shared_l2_acts_sf is not None + assert buffer.shared_l1_acts_sf.shape[1] == hidden // 128 + assert buffer.shared_l2_acts.shape == ( + num_max_tokens, shared_intermediate_hidden) + assert buffer.shared_l2_acts_sf.shape[1] == ( + shared_intermediate_hidden // 64) + assert buffer.shared_l1_acts_sf.stride() == ( + 1, buffer.shared_l1_acts_sf.shape[0]) + assert buffer.shared_l2_acts_sf.stride() == ( + 1, buffer.shared_l2_acts_sf.shape[0]) + assert buffer.l1_acts.shape[1] == hidden + assert buffer.l1_acts_sf.shape[1] == hidden // 128 + assert buffer.l2_acts.shape[1] == intermediate_hidden + assert buffer.l2_acts_sf.shape[1] == intermediate_hidden // 64 + assert buffer.l1_acts_sf.stride() == (1, buffer.l1_acts_sf.shape[0]) + assert buffer.l2_acts_sf.stride() == (1, buffer.l2_acts_sf.shape[0]) + if config.get('require_ring_wrap', False): + num_logical_pool_blocks = int( + ((local_recv_counts + 63) // 64).sum().item()) + num_physical_ring_blocks = buffer.l1_acts.shape[0] // 64 + did_wrap = torch.tensor( + [int(num_logical_pool_blocks > num_physical_ring_blocks)], + dtype=torch.int32, + device='cuda', + ) + if num_ranks > 1: + dist.all_reduce(did_wrap, op=dist.ReduceOp.MIN, group=group) + if rank_idx == 0: + print( + f' [RING] {name:<40} ' + f'logical_blocks(rank0)={num_logical_pool_blocks} ' + f'physical_blocks={num_physical_ring_blocks} ' + f'all_ranks_wrap={did_wrap.item()}', + flush=True, + ) + if did_wrap.item() == 0: + raise AssertionError( + f'{name}: scenario did not force ring wrap; logical_blocks=' + f'{num_logical_pool_blocks}, physical_blocks=' + f'{num_physical_ring_blocks}') + + buffer.x[:num_tokens].copy_(x_fp8) + buffer.x_sf[:num_tokens].copy_(x_sf) + if num_shared_experts > 0: + # The SM90 descriptor indexes m_idx directly. The exposed view is + # already column-major, so tensor copy performs the physical layout. + _copy_shared_l1_sf(buffer.shared_l1_acts_sf, x_sf) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + output = torch.empty( + num_tokens, hidden, dtype=torch.bfloat16, device='cuda') + recv_stats = torch.zeros( + num_local_experts, dtype=torch.int32, device='cuda') + for _ in range(repeat_count): + deep_gemm.fp8_mega_moe( + output, + transformed_l1, + transformed_l2, + buffer, + shared_l1_weights=transformed_shared_l1, + shared_l2_weights=transformed_shared_l2, + cumulative_local_expert_recv_stats=recv_stats, + recipe=(1, 1, 32), + activation='swiglu', + activation_clamp=( + activation_clamp if math.isfinite(activation_clamp) else None), + fast_math=fast_math, + ) + torch.cuda.synchronize() + + stats_failed = torch.tensor( + [int(not torch.equal(recv_stats, expected_recv_stats))], + dtype=torch.int32, + device='cuda', + ) + if num_ranks > 1: + dist.all_reduce(stats_failed, op=dist.ReduceOp.MAX, group=group) + if stats_failed.item() != 0: + raise _CorrectnessMismatch( + f'{name}: cumulative expert receive stats mismatch; ' + f'actual={recv_stats.cpu().tolist()}, ' + f'expected={expected_recv_stats.cpu().tolist()}') + + reference = _reference_mega_moe( + x_fp8, + x_sf, + topk_idx, + topk_weights, + reference_l1_weight, + reference_l1_sf, + reference_l2_weight, + reference_l2_sf, + reference_shared_l1_weight, + reference_shared_l1_sf, + reference_shared_l2_weight, + reference_shared_l2_sf, + rank_idx, + num_ranks, + group, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + ) + diff = 0.0 if output.numel() == 0 else float(calc_diff(output, reference)) + local_failed = torch.tensor( + [int(not math.isfinite(diff) or diff >= diff_tolerance)], + dtype=torch.int32, + device='cuda', + ) + if num_ranks > 1: + dist.all_reduce(local_failed, op=dist.ReduceOp.MAX, group=group) + if local_failed.item() != 0: + raise _CorrectnessMismatch( + f'{name}: diff {diff:.6f} exceeded tolerance {diff_tolerance}') + return diff + finally: + buffer.destroy() + + +def _smoke_scenarios(num_ranks: int) -> List[Scenario]: + base = dict( + num_max_tokens_per_rank=128, + num_tokens=64, + hidden=512, + intermediate_hidden=512, + num_experts=8 * num_ranks, + num_topk=2, + fast_math=True, + activation_clamp=10.0, + ) + routed = [('smoke.routed', dict(base, repeat_count=2))] + shared = [ + (f'smoke.shared_s{num_shared_experts}', dict( + base, + num_shared_experts=num_shared_experts, + repeat_count=2, + )) + for num_shared_experts in (1, 2) + ] + return routed + shared + + +def _standard_scenarios(num_ranks: int) -> List[Scenario]: + scenarios: List[Scenario] = [] + base = dict( + hidden=512, + intermediate_hidden=512, + num_experts=8 * num_ranks, + num_topk=2, + activation_clamp=10.0, + ) + # Token-per-expert bands used by the SM90 heuristic. + for num_tokens in (64, 256, 512, 2048): + scenarios.append(( + f'heuristic.t{num_tokens}', + dict( + base, + num_max_tokens_per_rank=( + (num_tokens + 127) // 128 * 128), + num_tokens=num_tokens, + fast_math=True, + ), + )) + # Explicit fast-math and clamp behavior. + for fast_math in (False, True): + scenarios.append(( + f'fast_math.{int(fast_math)}', + dict( + base, + num_max_tokens_per_rank=128, + num_tokens=128, + fast_math=fast_math, + ), + )) + for clamp in (1.0, math.inf): + scenarios.append(( + f'clamp.{clamp}', + dict( + base, + num_max_tokens_per_rank=128, + num_tokens=128, + fast_math=True, + activation_clamp=clamp, + ), + )) + # 0, 1, and maximum-token boundaries plus masked routing. + for num_tokens in (0, 1, 128): + scenarios.append(( + f'token_boundary.{num_tokens}', + dict( + base, + num_max_tokens_per_rank=128, + num_tokens=num_tokens, + fast_math=True, + ), + )) + scenarios.append(( + 'masked_routes', + dict( + base, + num_max_tokens_per_rank=128, + num_tokens=128, + fast_math=True, + masked_ratio=0.5, + ), + )) + return scenarios + + +def _full_scenarios( + num_ranks: int, + stress_count: int, +) -> List[Scenario]: + scenarios: List[Scenario] = [ + ('ring_wrap.h2048', dict( + num_max_tokens_per_rank=128, + num_tokens=128, + hidden=2048, + intermediate_hidden=1024, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.flash_m128', dict( + num_max_tokens_per_rank=128, + num_tokens=128, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.flash_m8', dict( + num_max_tokens_per_rank=128, + num_tokens=8, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.flash_m16', dict( + num_max_tokens_per_rank=128, + num_tokens=16, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.flash_m32', dict( + num_max_tokens_per_rank=128, + num_tokens=32, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.flash_m64', dict( + num_max_tokens_per_rank=128, + num_tokens=64, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.flash_m1024', dict( + num_max_tokens_per_rank=1024, + num_tokens=1024, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + )), + ('production.pro_m256', dict( + num_max_tokens_per_rank=256, + num_tokens=256, + hidden=7168, + intermediate_hidden=3072, + num_experts=48 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + )), + ('production.pro_m512', dict( + num_max_tokens_per_rank=512, + num_tokens=512, + hidden=7168, + intermediate_hidden=3072, + num_experts=48 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + )), + ('production.pro_m8', dict( + num_max_tokens_per_rank=128, + num_tokens=8, + hidden=7168, + intermediate_hidden=3072, + num_experts=48 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + )), + ('production.pro_m16', dict( + num_max_tokens_per_rank=128, + num_tokens=16, + hidden=7168, + intermediate_hidden=3072, + num_experts=48 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.pro_m32', dict( + num_max_tokens_per_rank=128, + num_tokens=32, + hidden=7168, + intermediate_hidden=3072, + num_experts=48 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.pro_m64', dict( + num_max_tokens_per_rank=128, + num_tokens=64, + hidden=7168, + intermediate_hidden=3072, + num_experts=48 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ('production.pro_m128', dict( + num_max_tokens_per_rank=128, + num_tokens=128, + hidden=7168, + intermediate_hidden=3072, + num_experts=48 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + require_ring_wrap=True, + )), + ] + rng = random.Random(0xC0FFEE) + for index in range(stress_count): + num_tokens = rng.choice((32, 64, 128, 256, 512)) + scenarios.append((f'stress.{index:03d}', dict( + num_max_tokens_per_rank=(num_tokens + 127) // 128 * 128, + num_tokens=num_tokens, + hidden=rng.choice((512, 1024, 2048)), + intermediate_hidden=rng.choice((512, 1024, 2048)), + num_experts=8 * num_ranks, + num_topk=rng.choice((1, 2, 4)), + fast_math=rng.choice((False, True)), + activation_clamp=rng.choice((1.0, 10.0, math.inf)), + masked_ratio=rng.choice((0.0, 0.0, 0.3, 0.7)), + ))) + return scenarios + + +def _test_worker(local_rank: int, num_local_ranks: int, args: argparse.Namespace) -> None: + from deep_gemm.testing import get_arch_major + from deep_gemm.utils.dist import init_dist + + if args.no_dist: + assert local_rank == 0 and num_local_ranks == 1 + torch.cuda.set_device(0) + torch.set_default_device('cuda') + rank_idx, num_ranks, group = 0, 1, _SingleProcessGroup() + else: + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + if get_arch_major() != 9: + if rank_idx == 0: + print( + f'[SKIP] test_mega_moe_sm90.py requires SM90, got SM{get_arch_major()}0', + flush=True, + ) + if not args.no_dist: + dist.destroy_process_group() + return + + scenarios = _smoke_scenarios(num_ranks) + if args.suite in ('standard', 'full'): + scenarios.extend(_standard_scenarios(num_ranks)) + if args.suite == 'full': + scenarios.extend(_full_scenarios( + num_ranks, args.stress_count)) + if args.filter: + scenarios = [ + scenario for scenario in scenarios if args.filter in scenario[0] + ] + if args.fast_math is not None: + scenarios = [ + (name, dict(config, fast_math=bool(args.fast_math))) + for name, config in scenarios + ] + + if not scenarios: + if rank_idx == 0: + print('[FAIL] scenario filter selected zero tests', flush=True) + if not args.no_dist: + dist.destroy_process_group() + raise SystemExit(2) + + if rank_idx == 0: + print( + f'SM90 MXFP4 MegaMoE plan: {len(scenarios)} scenarios, ' + f'{num_ranks} ranks, suite={args.suite}', + flush=True, + ) + failures: List[str] = [] + successes = 0 + for name, config in scenarios: + try: + diff = _run_scenario( + name, + config, + rank_idx, + num_ranks, + group, + args.diff_tolerance, + ) + successes += 1 + if rank_idx == 0: + print( + f' [PASS] {name:<40} diff={diff:.6f} ' + f'fast_math={int(config.get("fast_math", True))}', + flush=True, + ) + except Exception as exception: + # A rank-local Python/JIT/launch exception can leave peers inside + # the NVLink kernel. Let `mp.spawn` terminate sibling processes + # immediately. Only globally reduced correctness mismatches are safe to collect, + # because `_run_scenario` first reduces them to every rank. + if num_ranks > 1 and not isinstance( + exception, _CorrectnessMismatch): + raise + failures.append(name) + if rank_idx == 0: + print(f' [FAIL] {name}: {exception}', flush=True) + if args.fail_fast or not isinstance( + exception, _CorrectnessMismatch): + break + + if rank_idx == 0: + num_executed = successes + len(failures) + print( + f'SUMMARY total={num_executed} success={successes} ' + f'failed={len(failures)} planned={len(scenarios)}', + flush=True, + ) + group.barrier() + if not args.no_dist: + dist.destroy_process_group() + if failures: + raise SystemExit(1) + + +def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description='Cross-rank SM90 FP8 x MXFP4 MegaMoE runtime validation') + parser.add_argument( + '--num-processes', type=int, default=2, + help='number of local GPU ranks to spawn (default: 2)') + parser.add_argument( + '--no-dist', action='store_true', + help='run one GPU without NCCL (intended for compute-sanitizer)') + parser.add_argument( + '--suite', choices=('smoke', 'standard', 'full'), default='smoke', + help='scenario breadth; all cases use Humming processed MXFP4 triplets') + parser.add_argument( + '--fast-math', type=int, choices=(0, 1), default=None, + help='override fast_math for every selected scenario') + parser.add_argument( + '--diff-tolerance', type=float, default=0.01, + help='maximum calc_diff value (default: 0.01)') + parser.add_argument( + '--stress-count', type=int, default=8, + help='number of randomized scenarios in the full suite') + parser.add_argument( + '--filter', default='', help='only run scenario names containing this text') + parser.add_argument( + '--fail-fast', action='store_true', help='stop after the first failed scenario') + args = parser.parse_args(argv) + if args.num_processes <= 0: + parser.error('--num-processes must be positive') + return args + + +def test_forced_ring_wrap_routes_cover_every_expert() -> None: + num_ranks = 8 + routes = torch.cat([ + _make_forced_ring_wrap_topk_idx( + num_tokens=8, + num_topk=6, + num_experts=32 * num_ranks, + rank_idx=rank_idx, + num_ranks=num_ranks, + device=torch.device('cpu'), + ).reshape(-1) + for rank_idx in range(num_ranks) + ]) + counts = torch.bincount(routes, minlength=32 * num_ranks) + assert torch.all(counts > 0) + + +@pytest.mark.parametrize('num_processes', [0, -1]) +def test_cli_rejects_nonpositive_process_counts(num_processes: int) -> None: + with pytest.raises(SystemExit) as exception: + _parse_args(['--num-processes', str(num_processes)]) + assert exception.value.code == 2 + + +if __name__ == '__main__': + cli_args = _parse_args() + if cli_args.no_dist: + if cli_args.num_processes != 1: + raise ValueError('--no-dist requires --num-processes 1') + _test_worker(0, 1, cli_args) + else: + torch.multiprocessing.spawn( + _test_worker, + args=(cli_args.num_processes, cli_args), + nprocs=cli_args.num_processes, + ) From f026a32667a88aacef2ce4982c15ec17fec34b19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E6=A5=BC?= Date: Fri, 21 Aug 2026 15:55:53 +0800 Subject: [PATCH 2/3] fix(sm90): harden cross-rank MegaMoE paths Squash the latest non-documentation changes from the auto branch: cross-rank-safe swap bounds, Flash M8/M16 specializations, distributed tail diagnostics, and focused coverage. --- csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp | 10 +- .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 470 ++++++++++-------- tests/bench_mega_moe_sm90.py | 75 +++ tests/test_mega_moe_sm90.py | 82 ++- 4 files changed, 406 insertions(+), 231 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index 507ed4e132..6e4115bea1 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -104,10 +104,16 @@ class SM90FP8MXFP4MegaMoERuntime final : public LaunchRuntime 1 ? 64 : local_swap_ab_tokens; const bool packed_bf16_swap_epilogue = args.num_shared_experts == 0 and ((args.hidden == 4096 and @@ -151,6 +157,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, + {}, {}, {}, {} >); }}; @@ -165,6 +172,7 @@ static void __instantiate_kernel() {{ args.fast_math ? "true" : "false", small_m_swap_ab ? "true" : "false", max_swap_ab_tokens, + local_swap_ab_tokens, packed_bf16_swap_epilogue ? "true" : "false", swizzle_l2_cd ? "true" : "false", overlap_mxfp4_scale_path ? "true" : "false", 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 index 3eb3804f42..a62fa1dae6 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -259,6 +259,7 @@ CUTLASS_DEVICE void sm90_nvlink_barrier( bool kFastMath, \ bool kSmallMSwapAB, \ uint32_t kMaxSwapABTokens, \ + uint32_t kLocalSwapABTokens, \ bool kPackedBF16SwapEpilogue, \ bool kSwizzleL2CD, \ bool kOverlapMXFP4ScalePath, \ @@ -336,7 +337,7 @@ CUTLASS_DEVICE void sm90_nvlink_barrier( kNumMaxTokensPerRank, kHidden, kIntermediateHidden, kNumExperts, kNumTopk, \ kNumSMs, kNumRanks, \ kActivationClamp, kFastMath, kSmallMSwapAB, \ - kMaxSwapABTokens, \ + kMaxSwapABTokens, kLocalSwapABTokens, \ kPackedBF16SwapEpilogue, kSwizzleL2CD, \ kOverlapMXFP4ScalePath, \ kUsePRMTMXFP4Exponent, \ @@ -383,6 +384,12 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { kMaxSwapABTokens == 8 or kMaxSwapABTokens == 16 or kMaxSwapABTokens == 32 or kMaxSwapABTokens == 64, "Swap-AB token bound must select a supported WGMMA bucket"); + DG_STATIC_ASSERT( + kLocalSwapABTokens == 8 or kLocalSwapABTokens == 16 or + kLocalSwapABTokens == 32 or kLocalSwapABTokens == 64, + "Local swap-AB token bucket must be supported"); + DG_STATIC_ASSERT(kMaxSwapABTokens >= kLocalSwapABTokens, + "Swap-AB storage must cover the local policy bucket"); // ===================================================================== // Template checks @@ -782,7 +789,7 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { kNumExpertsPerRank, kNumSMs, kNumRanks, kNumRingBlocks, kNumSharedExperts, 1, kSmallMSwapAB and - (kMaxSwapABTokens == 8 or kMaxSwapABTokens == 64)>; + (kLocalSwapABTokens == 8 or kLocalSwapABTokens == 64)>; auto scheduler = scheduler_t( workspace, task_info_full_barriers, task_info_empty_barriers, task_infos); @@ -811,7 +818,7 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { kGemmPhaseBoundaryBarrierIdx + 1; constexpr bool kCacheProExpertCounts = not kHasSharedExperts and kHidden == 7168 and - ((kSmallMSwapAB and kMaxSwapABTokens == 8) or + ((kSmallMSwapAB and kLocalSwapABTokens == 8) or (not kSmallMSwapAB and kBankPermuteMXFP4PairLoads and not kSwizzleL2CD)); @@ -1041,48 +1048,20 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { *sym_buffer.map( workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), dst_rank_idx) = expert_status & 0xffffffff; -#if not defined(DG_SM90_SPARSE_DISPATCH_COMPLETION) + // Keep the cross-rank completion protocol identical even when + // ranks select different local sparse-count specializations. ptx::atomic_add_sys( sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), expert_status); -#endif } } ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); -#if defined(DG_SM90_SPARSE_DISPATCH_COMPLETION) - sm90_nvlink_barrier( - workspace, sym_buffer, sm_idx, thread_idx, - [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, - false, false); - - if (sm_idx == 0) { - ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - for (uint32_t local_expert_idx = thread_idx; - local_expert_idx < kNumExpertsPerRank; - local_expert_idx += kNumDispatchThreads) { - uint32_t num_recv_tokens = 0; - #pragma unroll - for (uint32_t rank_idx = 0; rank_idx < kNumRanks; ++ rank_idx) - num_recv_tokens += static_cast( - *workspace.get_expert_recv_count_ptr( - rank_idx, local_expert_idx)); - *workspace.get_expert_recv_count_sum_ptr(local_expert_idx) = - (static_cast(kNumSMs * kNumRanks) << 32) | - num_recv_tokens; - } - } - sm90_grid_sync( - workspace, sm_idx, thread_idx, - [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }); -#else sm90_nvlink_barrier( workspace, sym_buffer, sm_idx, thread_idx, [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, false, true); -#endif // Selected Pro buckets make both dispatch warps and the B-loader // scheduler independently poll and load the same completed expert @@ -1152,8 +1131,9 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { // round-robin path for duplicate routes from any rank. bool used_single_slot_rank_selection = false; if constexpr (kSmallMSwapAB and kHidden == 4096 and - (kMaxSwapABTokens == 8 or kMaxSwapABTokens == 16 or - kMaxSwapABTokens == 32) and + (kLocalSwapABTokens == 8 or + kLocalSwapABTokens == 16 or + kLocalSwapABTokens == 32) and kNumRanks <= 32) { const uint32_t multi_slot_rank_mask = __ballot_sync( 0xffffffff, stored_rank_count[0] > 1); @@ -1834,14 +1814,14 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { constexpr bool kPrmtPairLoads = kSmallMSwapAB and ((kHidden == 4096 and - (kMaxSwapABTokens == 8 or - kMaxSwapABTokens == 16 or - kMaxSwapABTokens == 32 or - kMaxSwapABTokens == 64)) or + (kLocalSwapABTokens == 8 or + kLocalSwapABTokens == 16 or + kLocalSwapABTokens == 32 or + kLocalSwapABTokens == 64)) or (kHidden == 7168 and - (kMaxSwapABTokens == 8 or - kMaxSwapABTokens == 32 or - kMaxSwapABTokens == 64))); + (kLocalSwapABTokens == 8 or + kLocalSwapABTokens == 32 or + kLocalSwapABTokens == 64))); const uint32_t packed_k_pair_in_k32 = kBankPermutedPairLoads ? (kPrmtPairLoads ? @@ -2438,12 +2418,12 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { (kHidden == 4096 or kHidden == 7168) and kSmallMSwapAB and - (kMaxSwapABTokens == 8 or - kMaxSwapABTokens == 16 or + (kLocalSwapABTokens == 8 or + kLocalSwapABTokens == 16 or (kHidden == 4096 and - kMaxSwapABTokens == 32 and + kLocalSwapABTokens == 32 and kPackedBF16SwapEpilogue) or - kMaxSwapABTokens == 64)) { + kLocalSwapABTokens == 64)) { const nv_bfloat162 scale_pair = __floats2bfloat162_rn( combined_scale_0, @@ -2954,8 +2934,6 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { // across lanes/warps. Derive the dynamic per-token K64 // output scale with a two-level warpgroup reduction, then // write the same row-major FP8/SF contract consumed by L2. - float swap_swiglu[ - kSwapABWeightHalves][kSwapABTokenChunks][2]; auto silu = [](float x) { const float e = kFastMath ? __expf(-x) : expf(-x); const float sig = kFastMath ? @@ -2975,196 +2953,246 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { cute::max(x, -kActivationClamp), kActivationClamp); }; - const uint32_t num_swap_token_chunks = - math::ceil_div(valid_m, 8u); - // Pipeline SFA is reusable by the producer immediately - // after the math loop releases a stage. Use C/D storage, - // which is epilogue-exclusive, for the cross-warp amax. - auto* swap_scale_scratch = - reinterpret_cast(smem_cd_base); + auto run_swap_l1_epilogue = [&]< + uint32_t kEpilogueTokenChunks>() { + DG_STATIC_ASSERT( + kEpilogueTokenChunks == + kLocalSwapABTokens / 8 or + kEpilogueTokenChunks == + kSwapABTokenChunks, + "Unexpected swap-AB epilogue bucket"); + float swap_swiglu[ + kSwapABWeightHalves][kEpilogueTokenChunks][2]; + const uint32_t num_swap_token_chunks = + math::ceil_div(valid_m, 8u); + // Pipeline SFA is reusable by the producer immediately + // after the math loop releases a stage. Use C/D + // storage, which is epilogue-exclusive, for the + // cross-warp amax. + auto* swap_scale_scratch = + reinterpret_cast(smem_cd_base); - #pragma unroll - for (uint32_t chunk = 0; - chunk < kSwapABTokenChunks; ++ chunk) { - const uint32_t token_0 = - chunk * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - const bool active_chunk = - chunk < num_swap_token_chunks; - const float weight_0 = - active_chunk and token_0 < valid_m ? - *l1_topk_weights_buffer - .get_data_buffer(m_idx + token_0) - .template get_base_ptr() : - 0.0f; - const float weight_1 = - active_chunk and token_1 < valid_m ? - *l1_topk_weights_buffer - .get_data_buffer(m_idx + token_1) - .template get_base_ptr() : - 0.0f; #pragma unroll - for (uint32_t half = 0; - half < kSwapABWeightHalves; ++ half) { - const uint32_t accum_offset = - half * kSwapABHalfAccumPerThread + - chunk * 4; - float gate_0, gate_1, up_0, up_1; - if constexpr (kPackedBF16SwapEpilogue) { - const float2 gate_pair = - __bfloat1622float2( - mxfp4_final_bf16[ - accum_offset / 2]); - const float2 up_pair = - __bfloat1622float2( - mxfp4_final_bf16[ - accum_offset / 2 + 1]); - gate_0 = gate_pair.x; - gate_1 = gate_pair.y; - up_0 = up_pair.x; - up_1 = up_pair.y; - } else { - gate_0 = final_accum[accum_offset]; - gate_1 = final_accum[accum_offset + 1]; - up_0 = final_accum[accum_offset + 2]; - up_1 = final_accum[accum_offset + 3]; + for (uint32_t chunk = 0; + chunk < kEpilogueTokenChunks; ++ chunk) { + const uint32_t token_0 = + chunk * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + const bool active_chunk = + chunk < num_swap_token_chunks; + const float weight_0 = + active_chunk and token_0 < valid_m ? + *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_0) + .template get_base_ptr() : + 0.0f; + const float weight_1 = + active_chunk and token_1 < valid_m ? + *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_1) + .template get_base_ptr() : + 0.0f; + #pragma unroll + for (uint32_t half = 0; + half < kSwapABWeightHalves; ++ half) { + const uint32_t accum_offset = + half * kSwapABHalfAccumPerThread + + chunk * 4; + float gate_0, gate_1, up_0, up_1; + if constexpr (kPackedBF16SwapEpilogue) { + const float2 gate_pair = + __bfloat1622float2( + mxfp4_final_bf16[ + accum_offset / 2]); + const float2 up_pair = + __bfloat1622float2( + mxfp4_final_bf16[ + accum_offset / 2 + 1]); + gate_0 = gate_pair.x; + gate_1 = gate_pair.y; + up_0 = up_pair.x; + up_1 = up_pair.y; + } else { + gate_0 = final_accum[accum_offset]; + gate_1 = final_accum[accum_offset + 1]; + up_0 = final_accum[accum_offset + 2]; + up_1 = final_accum[accum_offset + 3]; + } + // Keep compile-time chunk indices so the M64 + // safety bound remains spill-free, but skip + // inactive M16 chunks' special-function work. + if constexpr ( + kHidden == 4096 and + kLocalSwapABTokens == 16) { + if (active_chunk) { + clamp_gate(gate_0); + clamp_gate(gate_1); + clamp_up(up_0); + clamp_up(up_1); + swap_swiglu[half][chunk][0] = + silu(gate_0) * up_0 * weight_0; + swap_swiglu[half][chunk][1] = + silu(gate_1) * up_1 * weight_1; + } else { + swap_swiglu[half][chunk][0] = 0.0f; + swap_swiglu[half][chunk][1] = 0.0f; + } + } else { + clamp_gate(gate_0); + clamp_gate(gate_1); + clamp_up(up_0); + clamp_up(up_1); + swap_swiglu[half][chunk][0] = + silu(gate_0) * up_0 * weight_0; + swap_swiglu[half][chunk][1] = + silu(gate_1) * up_1 * weight_1; + } } - clamp_gate(gate_0); - clamp_gate(gate_1); - clamp_up(up_0); - clamp_up(up_1); - swap_swiglu[half][chunk][0] = - silu(gate_0) * up_0 * weight_0; - swap_swiglu[half][chunk][1] = - silu(gate_1) * up_1 * weight_1; - } - float partial_0 = cute::max( - cute::abs(swap_swiglu[0][chunk][0]), - cute::abs(swap_swiglu[1][chunk][0])); - float partial_1 = cute::max( - cute::abs(swap_swiglu[0][chunk][1]), - cute::abs(swap_swiglu[1][chunk][1])); - #pragma unroll - for (uint32_t delta = 4; delta <= 16; delta *= 2) { - partial_0 = cute::max( - partial_0, - __shfl_xor_sync( - 0xffffffffu, partial_0, delta)); - partial_1 = cute::max( - partial_1, - __shfl_xor_sync( - 0xffffffffu, partial_1, delta)); - } - if (row_idx == 0 and active_chunk) { - swap_scale_scratch[ - warp_idx_in_wg * BLOCK_M + token_0] = - partial_0; - swap_scale_scratch[ - warp_idx_in_wg * BLOCK_M + token_1] = - partial_1; + float partial_0 = cute::max( + cute::abs(swap_swiglu[0][chunk][0]), + cute::abs(swap_swiglu[1][chunk][0])); + float partial_1 = cute::max( + cute::abs(swap_swiglu[0][chunk][1]), + cute::abs(swap_swiglu[1][chunk][1])); + #pragma unroll + for (uint32_t delta = 4; + delta <= 16; delta *= 2) { + partial_0 = cute::max( + partial_0, + __shfl_xor_sync( + 0xffffffffu, partial_0, delta)); + partial_1 = cute::max( + partial_1, + __shfl_xor_sync( + 0xffffffffu, partial_1, delta)); + } + if (row_idx == 0 and active_chunk) { + swap_scale_scratch[ + warp_idx_in_wg * BLOCK_M + token_0] = + partial_0; + swap_scale_scratch[ + warp_idx_in_wg * BLOCK_M + token_1] = + partial_1; + } } - } - ptx::sync_aligned( - kNumEpilogueThreads, - kEpilogueWGBarrierStartIdx); + ptx::sync_aligned( + kNumEpilogueThreads, + kEpilogueWGBarrierStartIdx); - if (warp_idx_in_wg == 0) { - #pragma unroll - for (uint32_t half = 0; half < 2; ++ half) { - const uint32_t token = lane_idx + half * 32; - if (token < valid_m) { - float amax = 0.0f; - #pragma unroll - for (uint32_t warp = 0; - warp < kNumEpilogueWarps; ++ warp) - amax = cute::max( - amax, - swap_scale_scratch[ - warp * BLOCK_M + token]); - float2 amax_pair = {amax, amax}; - float2 sf_pair, sf_inv_pair; - sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( - amax_pair, sf_pair, sf_inv_pair); - swap_scale_scratch[token] = sf_pair.x; - swap_scale_scratch[BLOCK_M + token] = - sf_inv_pair.x; + if (warp_idx_in_wg == 0) { + #pragma unroll + for (uint32_t half = 0; half < 2; ++ half) { + const uint32_t token = lane_idx + half * 32; + if (token < valid_m) { + float amax = 0.0f; + #pragma unroll + for (uint32_t warp = 0; + warp < kNumEpilogueWarps; ++ warp) + amax = cute::max( + amax, + swap_scale_scratch[ + warp * BLOCK_M + token]); + float2 amax_pair = {amax, amax}; + float2 sf_pair, sf_inv_pair; + sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( + amax_pair, sf_pair, sf_inv_pair); + swap_scale_scratch[token] = sf_pair.x; + swap_scale_scratch[BLOCK_M + token] = + sf_inv_pair.x; + } } } - } - ptx::sync_aligned( - kNumEpilogueThreads, - kEpilogueWGBarrierStartIdx); + ptx::sync_aligned( + kNumEpilogueThreads, + kEpilogueWGBarrierStartIdx); - float swap_sf_inv[kSwapABTokenChunks][2]; - #pragma unroll - for (uint32_t chunk = 0; - chunk < kSwapABTokenChunks; ++ chunk) { - const uint32_t token_0 = - chunk * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - swap_sf_inv[chunk][0] = token_0 < valid_m ? - swap_scale_scratch[BLOCK_M + token_0] : 0.0f; - swap_sf_inv[chunk][1] = token_1 < valid_m ? - swap_scale_scratch[BLOCK_M + token_1] : 0.0f; - } - if (warp_idx_in_wg == 0) { - auto sf_base_ptr = - l2_sf_buffer.get_base_ptr(); - const uint32_t base_k_sf_idx = n_block_idx; + float swap_sf_inv[kEpilogueTokenChunks][2]; #pragma unroll - for (uint32_t half = 0; half < 2; ++ half) { - const uint32_t token = lane_idx + half * 32; - if (token < valid_m) - sf_base_ptr[ - base_k_sf_idx * kNumSFRingTokens + - m_idx + token] = - swap_scale_scratch[token]; - } - } - // Every thread must retain its inverse scales before the - // row-major FP8 stores overwrite C/D scratch. - ptx::sync_aligned( - kNumEpilogueThreads, - kEpilogueWGBarrierStartIdx); - - #pragma unroll - for (uint32_t chunk = 0; - chunk < kSwapABTokenChunks; ++ chunk) { - if (chunk < num_swap_token_chunks) { + for (uint32_t chunk = 0; + chunk < kEpilogueTokenChunks; ++ chunk) { const uint32_t token_0 = chunk * 8 + col_idx * 2; const uint32_t token_1 = token_0 + 1; + swap_sf_inv[chunk][0] = token_0 < valid_m ? + swap_scale_scratch[BLOCK_M + token_0] : 0.0f; + swap_sf_inv[chunk][1] = token_1 < valid_m ? + swap_scale_scratch[BLOCK_M + token_1] : 0.0f; + } + if (warp_idx_in_wg == 0) { + auto sf_base_ptr = + l2_sf_buffer.get_base_ptr(); + const uint32_t base_k_sf_idx = n_block_idx; #pragma unroll - for (uint32_t half = 0; - half < kSwapABWeightHalves; ++ half) { - const uint32_t out_col = - half * 32u + - warp_idx_in_wg * 8 + row_idx; - if (token_0 < valid_m) { - const __nv_fp8_e4m3 q( - swap_swiglu[half][chunk][0] * - swap_sf_inv[chunk][0]); - reinterpret_cast( - smem_cd_l1_wg)[ - token_0 * - WG_SMEM_CD_L1_STRIDE_N + - out_col] = - *reinterpret_cast(&q); - } - if (token_1 < valid_m) { - const __nv_fp8_e4m3 q( - swap_swiglu[half][chunk][1] * - swap_sf_inv[chunk][1]); - reinterpret_cast( - smem_cd_l1_wg)[ - token_1 * - WG_SMEM_CD_L1_STRIDE_N + - out_col] = - *reinterpret_cast(&q); + for (uint32_t half = 0; half < 2; ++ half) { + const uint32_t token = lane_idx + half * 32; + if (token < valid_m) + sf_base_ptr[ + base_k_sf_idx * kNumSFRingTokens + + m_idx + token] = + swap_scale_scratch[token]; + } + } + // Every thread must retain its inverse scales before + // the row-major FP8 stores overwrite C/D scratch. + ptx::sync_aligned( + kNumEpilogueThreads, + kEpilogueWGBarrierStartIdx); + + #pragma unroll + for (uint32_t chunk = 0; + chunk < kEpilogueTokenChunks; ++ chunk) { + if (chunk < num_swap_token_chunks) { + const uint32_t token_0 = + chunk * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + #pragma unroll + for (uint32_t half = 0; + half < kSwapABWeightHalves; ++ half) { + const uint32_t out_col = + half * 32u + + warp_idx_in_wg * 8 + row_idx; + if (token_0 < valid_m) { + const __nv_fp8_e4m3 q( + swap_swiglu[half][chunk][0] * + swap_sf_inv[chunk][0]); + reinterpret_cast( + smem_cd_l1_wg)[ + token_0 * + WG_SMEM_CD_L1_STRIDE_N + + out_col] = + *reinterpret_cast( + &q); + } + if (token_1 < valid_m) { + const __nv_fp8_e4m3 q( + swap_swiglu[half][chunk][1] * + swap_sf_inv[chunk][1]); + reinterpret_cast( + smem_cd_l1_wg)[ + token_1 * + WG_SMEM_CD_L1_STRIDE_N + + out_col] = + *reinterpret_cast( + &q); + } } } } + }; + + if constexpr ( + kLocalSwapABTokens == kMaxSwapABTokens or + kHidden != 4096 or + kLocalSwapABTokens != 8) { + run_swap_l1_epilogue.template operator()< + kSwapABTokenChunks>(); + } else if (valid_m <= kLocalSwapABTokens) { + run_swap_l1_epilogue.template operator()< + kLocalSwapABTokens / 8>(); + } else { + run_swap_l1_epilogue.template operator()< + kSwapABTokenChunks>(); } } else { diff --git a/tests/bench_mega_moe_sm90.py b/tests/bench_mega_moe_sm90.py index f2614a85d4..81260f2b17 100644 --- a/tests/bench_mega_moe_sm90.py +++ b/tests/bench_mega_moe_sm90.py @@ -232,6 +232,45 @@ def _benchmark_case( topk_idx.masked_fill_(mask, -1) topk_weights.masked_fill_(mask, 0.0) + route_stats = None + if args.report_route_stats: + if num_ranks == 1: + global_topk_idx = topk_idx + else: + from deep_gemm.utils.dist import uneven_all_gather + + global_topk_idx = uneven_all_gather(topk_idx, group=group) + valid_topk_idx = global_topk_idx[global_topk_idx >= 0].long() + global_recv_counts = torch.bincount( + valid_topk_idx, minlength=num_experts + ).reshape(num_ranks, num_local_experts) + routed_m_blocks = torch.div( + global_recv_counts + 63, 64, rounding_mode="floor" + ) + if rank_idx == 0: + route_stats = [ + { + "rank": route_rank, + "recv_tokens": int( + global_recv_counts[route_rank].sum().item() + ), + "m_blocks": int( + routed_m_blocks[route_rank].sum().item() + ), + "experts_over_block_m": int( + (global_recv_counts[route_rank] > 64).sum().item() + ), + "max_expert_tokens": int( + global_recv_counts[route_rank].max().item() + ), + "expert_counts": global_recv_counts[ + route_rank + ].cpu().tolist(), + } + for route_rank in range(num_ranks) + ] + del global_topk_idx, valid_topk_idx + transformed_shared_l1: Optional[Tuple[torch.Tensor, torch.Tensor]] = None transformed_shared_l2: Optional[Tuple[torch.Tensor, torch.Tensor]] = None if num_shared_experts > 0: @@ -319,6 +358,11 @@ def run_kernel() -> torch.Tensor: "masked_ratio": args.masked_ratio, "seed": args.seed, } + if route_stats is not None: + _emit_json( + "BENCH_ROUTE_JSON", + {**case_metadata, "block_m": 64, "per_rank": route_stats}, + ) if args.profile_only: if rank_idx == 0: @@ -375,6 +419,23 @@ def run_kernel() -> torch.Tensor: if num_ranks > 1: dist.all_reduce(max_rank_time, op=dist.ReduceOp.MAX, group=group) + rank_times_us = None + if args.report_rank_times: + gathered_rank_times = [ + torch.zeros_like(max_rank_time) for _ in range(num_ranks) + ] + if num_ranks > 1: + dist.all_gather( + gathered_rank_times, max_rank_time.new_tensor(kernel_time), + group=group, + ) + else: + gathered_rank_times[0].copy_(max_rank_time) + rank_times_us = [ + rank_time.item() * 1e6 + for rank_time in gathered_rank_times + ] + rank0_observations.append(kernel_time) max_rank_observations.append(max_rank_time.item()) if rank_idx == 0: @@ -389,6 +450,10 @@ def run_kernel() -> torch.Tensor: "num_tests": args.num_tests, "flush_l2": bool(args.flush_l2), "cache_mode": _cache_mode(args.flush_l2), + **( + {"rank_times_us": rank_times_us} + if rank_times_us is not None else {} + ), }, ) @@ -535,6 +600,16 @@ def _parse_args() -> argparse.Namespace: help="1 flushes L2 before each sample; 0 disables the explicit flush", ) parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--report-rank-times", + action="store_true", + help="include every rank's local kernel time in BENCH_OBS_JSON", + ) + parser.add_argument( + "--report-route-stats", + action="store_true", + help="emit per-rank receive counts and routed M64 block counts", + ) parser.add_argument("--masked-ratio", type=float, default=0.0) parser.add_argument("--activation-clamp", type=float, default=10.0) parser.add_argument("--fast-math", type=int, choices=(0, 1), default=1) diff --git a/tests/test_mega_moe_sm90.py b/tests/test_mega_moe_sm90.py index 52480ac59f..8338e79f13 100644 --- a/tests/test_mega_moe_sm90.py +++ b/tests/test_mega_moe_sm90.py @@ -343,7 +343,12 @@ def _run_scenario( from deep_gemm.utils import per_token_cast_to_fp8 num_max_tokens = config['num_max_tokens_per_rank'] - num_tokens = config.get('num_tokens', num_max_tokens) + num_tokens_by_rank = config.get('num_tokens_by_rank') + if num_tokens_by_rank is None: + num_tokens = config.get('num_tokens', num_max_tokens) + else: + assert len(num_tokens_by_rank) == num_ranks + num_tokens = num_tokens_by_rank[rank_idx] hidden = config['hidden'] intermediate_hidden = config['intermediate_hidden'] num_experts = config['num_experts'] @@ -402,14 +407,26 @@ def _run_scenario( l2_quantized = _quantize_grouped_mxfp4(l2_bf16) del l2_bf16 - scores = torch.randn( - num_tokens, num_experts, dtype=torch.float32, device='cuda') - topk_weights, topk_idx = torch.topk( - scores, num_topk, dim=-1, largest=True, sorted=False) - if config.get('require_ring_wrap', False): - topk_idx = _make_forced_ring_wrap_topk_idx( - num_tokens, num_topk, num_experts, - rank_idx, num_ranks, scores.device) + hot_route_rank = config.get('hot_route_rank') + if hot_route_rank is None: + scores = torch.randn( + num_tokens, num_experts, dtype=torch.float32, device='cuda') + topk_weights, topk_idx = torch.topk( + scores, num_topk, dim=-1, largest=True, sorted=False) + if config.get('require_ring_wrap', False): + topk_idx = _make_forced_ring_wrap_topk_idx( + num_tokens, num_topk, num_experts, + rank_idx, num_ranks, scores.device) + else: + assert 0 <= hot_route_rank < num_ranks + assert num_topk <= num_local_experts + hot_expert_start = hot_route_rank * num_local_experts + topk_idx = ( + torch.arange(num_topk, dtype=torch.int64, device='cuda') + + hot_expert_start + ).expand(num_tokens, -1).clone() + topk_weights = torch.randn( + num_tokens, num_topk, dtype=torch.float32, device='cuda') if masked_ratio: mask = torch.rand_like(topk_idx, dtype=torch.float32) < masked_ratio topk_idx.masked_fill_(mask, -1) @@ -878,6 +895,53 @@ def _full_scenarios( require_ring_wrap=True, )), ] + if num_ranks > 1: + scenarios.extend([ + ('swap_ab_cross_rank_bound.flash_m8', dict( + num_max_tokens_per_rank=128, + num_tokens=8, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + hot_route_rank=0, + )), + ('swap_ab_cross_rank_bound.flash_m16', dict( + num_max_tokens_per_rank=128, + num_tokens=16, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + hot_route_rank=0, + )), + ('dispatch_mixed_protocol.flash_m32_m64', dict( + num_max_tokens_per_rank=128, + num_tokens_by_rank=(32,) + (64,) * (num_ranks - 1), + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + hot_route_rank=1, + )), + ('dispatch_mixed_protocol.flash_m1024_m64', dict( + num_max_tokens_per_rank=1024, + num_tokens_by_rank=(1024,) + (64,) * (num_ranks - 1), + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + hot_route_rank=1, + )), + ]) rng = random.Random(0xC0FFEE) for index in range(stress_count): num_tokens = rng.choice((32, 64, 128, 256, 512)) From 5d32216686b982a39eabccca9419af430a60cfc2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E6=A5=BC?= Date: Mon, 24 Aug 2026 11:37:24 +0800 Subject: [PATCH 3/3] perf(sm90): integrate latest MXFP4 MegaMoE optimizations Squash the current non-documentation delta from the auto-optimization branch on top of the latest development branch. This integrates the processed-MXFP4 runtime updates, accepted decode/scale/epilogue changes, workload-gated packed-FP16 WGMMA accumulation, and refreshed benchmark and correctness coverage. The packed path is selected from routed M64xN128xK128 macro-tile work rather than model names. Preliminary same-source H20 screening improves Flash M1024+ by 1.3-2.2% and Pro M512+ by 0.6-1.6%; the work gate excludes the regressing Flash M256/M512 and Pro M256 cases. Documentation and README changes are intentionally excluded. --- csrc/apis/sm90_mega.hpp | 20 +- deep_gemm/__init__.py | 2 +- .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 346 ++++++++++++++---- .../include/deep_gemm/layout/mega_moe.cuh | 18 +- deep_gemm/include/deep_gemm/mma/sm90.cuh | 47 +++ deep_gemm/include/deep_gemm/ptx/wgmma.cuh | 4 + deep_gemm/mega/__init__.py | 149 +++----- deep_gemm/mega/mxfp4.py | 56 +-- tests/bench_mega_moe_sm90.py | 2 +- tests/test_mega_moe_mxfp4_preprocess.py | 77 +--- tests/test_mega_moe_sm90.py | 106 +++--- 11 files changed, 474 insertions(+), 353 deletions(-) diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index b70bf06423..1299ef10b5 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -75,7 +75,7 @@ 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 std::string& mma_type, const std::string& activation, + const bool& use_fp8_dispatch, const std::string& activation, const int& num_shared_experts = 0) { DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); DG_HOST_ASSERT(num_ranks > 0 and num_ranks <= kSM90MegaMoEMaxRanks); @@ -91,12 +91,8 @@ get_symm_buffer_size_for_sm90_mega_moe( // TMA alignment. L2 K64 FP32 SF rows occupy I/16 bytes. DG_HOST_ASSERT(is_valid_hidden_for_sm90_mega_moe(hidden) and intermediate_hidden % 256 == 0); - // Keep the buffer API weight-format aware even though only MXFP4 is - // implemented today. Future FP8 and INT4 backends can select their layout - // here without introducing another public allocator. - DG_HOST_ASSERT(mma_type == "fp8xmxfp4"); + DG_HOST_ASSERT(use_fp8_dispatch); DG_HOST_ASSERT(activation == "swiglu"); - const auto num_ring_tokens = get_num_ring_tokens_for_sm90_mega_moe( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden); @@ -108,12 +104,10 @@ get_symm_buffer_size_for_sm90_mega_moe( nullptr, hidden, intermediate_hidden, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_ring_tokens, num_sf_ring_tokens, - /*with_sf=*/ true, num_shared_experts, scale_layout_spec, - kSM90MegaMoEBlockM); + /*with_sf=*/ true, num_shared_experts, scale_layout_spec); const auto shared_intermediate_hidden = intermediate_hidden * num_shared_experts; const auto num_max_shared_sf_tokens = - layout::get_num_shared_sf_tokens( - num_max_tokens_per_rank, kSM90MegaMoEBlockM); + layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank); // Slice function follows main's twelve-view ABI. Shared weights remain // FP8+FP32 on SM90; only routed weights use packed MXFP4+UE8M0. @@ -185,7 +179,7 @@ get_symm_buffer_size_for_sm90_mega_moe( // K32 relative UE8M0 scales plus one FP32 secondary scale per expert. // Top-level routing is the caller's responsibility (see // `deep_gemm/mega/__init__.py`). -static void run_sm90_fp8_mxfp4_mega_moe( +static void sm90_mega_moe( const torch::Tensor& y, const std::tuple& l1_weights_tuple, const std::tuple& l2_weights_tuple, @@ -364,7 +358,7 @@ static void run_sm90_fp8_mxfp4_mega_moe( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - "fp8xmxfp4", activation, num_shared_experts); + true, activation, num_shared_experts); // The live-ring offsets depend on H/I/S and the active worker-SM count. // Require the exact allocation contract so a buffer created for another // shape, shared-expert count, or `set_num_sms` value cannot be re-sliced @@ -415,7 +409,7 @@ static void fp8_mxfp4_mega_moe( const std::string& activation, const std::optional& activation_clamp_opt, const bool& fast_math) { - run_sm90_fp8_mxfp4_mega_moe( + sm90_mega_moe( y, l1_weights_tuple, l2_weights_tuple, shared_l1_weights_tuple_opt, shared_l2_weights_tuple_opt, cumulative_local_expert_recv_stats, diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index a75c7817fc..e7b0f2de59 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -92,7 +92,7 @@ transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90, transform_shared_weights_for_fp8_mega_moe_sm90, fp8_fp4_mega_moe, - fp8_mega_moe, + fp8_mxfp4_mega_moe, bf16_mega_moe, ) 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 index a62fa1dae6..6af299a622 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -441,8 +441,7 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { kNumRingTokens, kNumSFRingTokens, /*with_sf=*/ true, kNumSharedExperts, layout::ScaleLayoutSpec::sm90_fp32_k128_k64( - kHidden, kIntermediateHidden), - BLOCK_M); + kHidden, kIntermediateHidden)); const auto workspace = buffer.workspace; constexpr auto fp8_token_layout = layout::Data(kHidden); @@ -1641,6 +1640,35 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { // ---------------- GEMM ---------------- using WGMMA = L1WGMMA; constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; // 64 for M=64,N=128 + // Regular routed tasks can keep both the WGMMA fragment and the + // fast-math cross-K sum packed as FP16x2. Packed accumulation + // saves conversion/accumulator instructions but increases WGMMA + // barrier latency, so enable it only when the routed GEMMs expose + // enough aggregate work to amortize that latency. Count the + // M64xN128xK128 mainloop macro-tiles for the two L1 projections + // and L2. This depends on the routed workload and topology, not + // a model name or an exact hidden-size tuple. + constexpr uint64_t kRoutedMBlockEstimate = + (static_cast(kNumMaxTokensPerRank) * kNumTopk + + BLOCK_M - 1) / BLOCK_M; + constexpr uint64_t kWGMMAMacroTilesPerRoutedMBlock = + (2ull * kIntermediateHidden / BLOCK_N) * + (kHidden / BLOCK_K) + + (kHidden / BLOCK_N) * + (kIntermediateHidden / BLOCK_K); + constexpr uint64_t kRoutedWGMMAMacroTileEstimate = + kRoutedMBlockEstimate * kWGMMAMacroTilesPerRoutedMBlock; + constexpr uint64_t kMinPackedF16WGMMAMacroTiles = 1ull << 17; + constexpr bool kUsePackedF16Accum = + kFastMath and not kSmallMSwapAB and not kHasSharedExperts and + kRoutedWGMMAMacroTileEstimate >= + kMinPackedF16WGMMAMacroTiles; + using PackedF16WGMMA = typename + mma::sm90::FP8MMAF16Selector::type; + DG_STATIC_ASSERT( + PackedF16WGMMA::kNumAccum == kAccumPerThread and + PackedF16WGMMA::kNumPackedAccum * 2 == kAccumPerThread, + "Packed FP16 and FP32 WGMMA fragments must cover the same tile"); float final_accum[ kPackedBF16SwapEpilogue ? 1 : kAccumPerThread]; if constexpr (is_shared_phase) { @@ -1656,7 +1684,10 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { constexpr uint32_t kAccumStorage = kReuseSwapABFragment ? kSwapABHalfAccumPerThread : kAccumPerThread; float accum[kAccumStorage]; - nv_bfloat162 mxfp4_final_bf16[kAccumPerThread / 2]; + uint32_t packed_f16_accum[PackedF16WGMMA::kNumPackedAccum]; + using PersistentAccum2 = std::conditional_t< + kUsePackedF16Accum, __half2, nv_bfloat162>; + PersistentAccum2 mxfp4_final_bf16[kAccumPerThread / 2]; const auto run_mxfp4_gemm_loop = [&]() { { @@ -1664,9 +1695,14 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { constexpr uint32_t kWGThreads = 128; const uint32_t wg_thread_idx = warp_idx_in_wg * 32 + lane_idx; #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread / 2; ++ i) - mxfp4_final_bf16[i] = - __float2bfloat162_rn(0.0f); + for (uint32_t i = 0; i < kAccumPerThread / 2; ++ i) { + if constexpr (kUsePackedF16Accum) + mxfp4_final_bf16[i] = + __float2half2_rn(0.0f); + else + mxfp4_final_bf16[i] = + __float2bfloat162_rn(0.0f); + } const auto prepare_stage_weights = [&]( const uint32_t pipeline_stage, @@ -2077,9 +2113,17 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { uint32_t kNumWGMMAs>( const uint32_t pipeline_stage, const uint32_t expanded_slot) { - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ i) - ptx::warpgroup_fence_operand(accum[i]); + if constexpr (kUsePackedF16Accum) { + #pragma unroll + for (uint32_t i = 0; + i < PackedF16WGMMA::kNumPackedAccum; ++ i) + ptx::warpgroup_fence_operand( + packed_f16_accum[i]); + } else { + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + } ptx::warpgroup_arrive(); if constexpr (kUseIncrementalMXFP4Descriptor) { const auto desc_a_base = mma::sm90::make_smem_desc( @@ -2096,9 +2140,14 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { desc_a_base.desc_ + k * 2u); const cute::GmmaDescriptor desc_b( desc_b_base.desc_ + k * 2u); - WGMMA::wgmma( - desc_a, desc_b, accum, - k != 0); + if constexpr (kUsePackedF16Accum) + PackedF16WGMMA::wgmma( + desc_a, desc_b, packed_f16_accum, + k != 0); + else + WGMMA::wgmma( + desc_a, desc_b, accum, + k != 0); } } else { #pragma unroll @@ -2110,15 +2159,28 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { auto desc_b = mma::sm90::make_smem_desc( smem_b_expanded[expanded_slot] + k32_idx * kWeightGranK, 1); - WGMMA::wgmma( - desc_a, desc_b, accum, - k != 0); + if constexpr (kUsePackedF16Accum) + PackedF16WGMMA::wgmma( + desc_a, desc_b, packed_f16_accum, + k != 0); + else + WGMMA::wgmma( + desc_a, desc_b, accum, + k != 0); } } ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ i) - ptx::warpgroup_fence_operand(accum[i]); + if constexpr (kUsePackedF16Accum) { + #pragma unroll + for (uint32_t i = 0; + i < PackedF16WGMMA::kNumPackedAccum; ++ i) + ptx::warpgroup_fence_operand( + packed_f16_accum[i]); + } else { + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + } if constexpr (not kOverlapMXFP4ScalePath) ptx::warpgroup_wait<0>(); }; @@ -2163,10 +2225,24 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { // preparing the two promotion multipliers, then wait // immediately before the first accumulator access. if constexpr (kFastMath) { - const nv_bfloat162 combined_scale_bf16_0 = - __float2bfloat162_rn(combined_scale_0); - const nv_bfloat162 combined_scale_bf16_1 = - __float2bfloat162_rn(combined_scale_1); + const PersistentAccum2 combined_scale_packed_0 = + [&]() { + if constexpr (kUsePackedF16Accum) + return __float2half2_rn( + combined_scale_0); + else + return __float2bfloat162_rn( + combined_scale_0); + }(); + const PersistentAccum2 combined_scale_packed_1 = + [&]() { + if constexpr (kUsePackedF16Accum) + return __float2half2_rn( + combined_scale_1); + else + return __float2bfloat162_rn( + combined_scale_1); + }(); if constexpr (kOverlapMXFP4ScalePath) ptx::warpgroup_wait<0>(); // The final wait ends every shared-memory access @@ -2177,18 +2253,37 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { if constexpr (kOverlapMXFP4ScalePath) arrive_task_empty_barrier(pipeline_stage); } - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { - mxfp4_final_bf16[i * 2] = __hfma2( - combined_scale_bf16_0, - __floats2bfloat162_rn( - accum[i * 4], accum[i * 4 + 1]), - mxfp4_final_bf16[i * 2]); - mxfp4_final_bf16[i * 2 + 1] = __hfma2( - combined_scale_bf16_1, - __floats2bfloat162_rn( - accum[i * 4 + 2], accum[i * 4 + 3]), - mxfp4_final_bf16[i * 2 + 1]); + if constexpr (kUsePackedF16Accum) { + #pragma unroll + for (uint32_t i = 0; + i < PackedF16WGMMA::kNumPackedAccum; + ++ i) { + const auto fragment = + reinterpret_cast( + packed_f16_accum[i]); + mxfp4_final_bf16[i] = __hfma2( + (i & 1u) ? + combined_scale_packed_1 : + combined_scale_packed_0, + fragment, + mxfp4_final_bf16[i]); + } + } else { + #pragma unroll + for (uint32_t i = 0; + i < kAccumPerThread / 4; ++ i) { + mxfp4_final_bf16[i * 2] = __hfma2( + combined_scale_packed_0, + __floats2bfloat162_rn( + accum[i * 4], accum[i * 4 + 1]), + mxfp4_final_bf16[i * 2]); + mxfp4_final_bf16[i * 2 + 1] = __hfma2( + combined_scale_packed_1, + __floats2bfloat162_rn( + accum[i * 4 + 2], + accum[i * 4 + 3]), + mxfp4_final_bf16[i * 2 + 1]); + } } } else { if constexpr (kOverlapMXFP4ScalePath) @@ -2251,6 +2346,17 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { constexpr uint32_t kWeightHalfAccumStride = kPipelineWeightHalves ? kSwapAccum : kSwapABHalfAccumPerThread; + constexpr bool kVectorProM8ActivationScales = + kHidden == 7168 and + kLocalSwapABTokens == 8; + constexpr bool kPrecomputeFlashM16ScaleBeforeWait = + kHidden == 4096 and + kLocalSwapABTokens == 16; + constexpr bool kPrecomputeN8ScaleBeforeWait = + kFastMath and + (kVectorProM8ActivationScales or + kPrecomputeFlashM16ScaleBeforeWait) and + N_SWAP == 8; DG_STATIC_ASSERT( not kPipelineWeightHalves or 2 * kWeightHalfAccumStride <= kAccumStorage, @@ -2358,6 +2464,76 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { compensate_secondary ? mxfp4_secondary * 64.0f : mxfp4_secondary; + const nv_bfloat162 + precomputed_scale_pair = [&]() { + if constexpr ( + kPrecomputeN8ScaleBeforeWait) { + DG_STATIC_ASSERT( + kSwapAccum / 4 == 1, + "N8 precomputes one scale pair"); + const uint32_t token_0 = + swap_col_idx * 2; + const uint32_t token_1 = token_0 + 1; + // Every producer TMA fills the + // complete BLOCK_M SFA vector. In + // the exact Flash M16/N8 bucket, + // values for padded tokens are + // irrelevant because their + // accumulator columns are never + // written. Loading them directly + // removes two per-promotion bounds + // selects from the critical path. + const float2 scale_pair = [&]() { + if constexpr ( + kPrecomputeFlashM16ScaleBeforeWait) { + return ptx::ld_shared( + reinterpret_cast< + const float2*>( + smem_sfa[ + pipeline_stage] + + activation_sf_group * + kL2SFAHalfStride + + token_0)); + } else { + return token_0 < valid_m ? + ptx::ld_shared( + reinterpret_cast< + const float2*>( + smem_sfa[ + pipeline_stage] + + activation_sf_group * + kL2SFAHalfStride + + token_0)) : + make_float2(0.0f, 0.0f); + } + }(); + const float scale_a_0 = scale_pair.x; + const float scale_a_1 = [&]() { + if constexpr ( + kPrecomputeFlashM16ScaleBeforeWait) { + return scale_pair.y; + } else { + return token_1 < valid_m ? + scale_pair.y : 0.0f; + } + }(); + const float combined_scale_0 = + (compensate_secondary ? + scale_a_0 : + scale_a_0 * 64.0f) * + compensated_secondary; + const float combined_scale_1 = + (compensate_secondary ? + scale_a_1 : + scale_a_1 * 64.0f) * + compensated_secondary; + return __floats2bfloat162_rn( + combined_scale_0, + combined_scale_1); + } else { + return __float2bfloat162_rn(0.0f); + } + }(); ptx::warpgroup_wait(); #pragma unroll @@ -2366,30 +2542,56 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { const uint32_t token_0 = chunk * 8 + swap_col_idx * 2; const uint32_t token_1 = token_0 + 1; - const float scale_a_0 = - token_0 < valid_m ? - ptx::ld_shared( - smem_sfa[pipeline_stage] + - activation_sf_group * - kL2SFAHalfStride + - token_0) : - 0.0f; - const float scale_a_1 = - token_1 < valid_m ? - ptx::ld_shared( - smem_sfa[pipeline_stage] + - activation_sf_group * - kL2SFAHalfStride + - token_1) : - 0.0f; - const float combined_scale_0 = - (compensate_secondary ? - scale_a_0 : scale_a_0 * 64.0f) * - compensated_secondary; - const float combined_scale_1 = - (compensate_secondary ? - scale_a_1 : scale_a_1 * 64.0f) * - compensated_secondary; + float combined_scale_0, combined_scale_1; + if constexpr ( + not kPrecomputeN8ScaleBeforeWait) { + float scale_a_0, scale_a_1; + if constexpr ( + kVectorProM8ActivationScales) { + DG_STATIC_ASSERT( + kL2SFAHalfStride % 2 == 0, + "Vector SFA rows must stay aligned"); + const float2 scale_pair = + token_0 < valid_m ? + ptx::ld_shared( + reinterpret_cast< + const float2*>( + smem_sfa[ + pipeline_stage] + + activation_sf_group * + kL2SFAHalfStride + + token_0)) : + make_float2(0.0f, 0.0f); + scale_a_0 = scale_pair.x; + scale_a_1 = token_1 < valid_m ? + scale_pair.y : 0.0f; + } else { + scale_a_0 = token_0 < valid_m ? + ptx::ld_shared( + smem_sfa[pipeline_stage] + + activation_sf_group * + kL2SFAHalfStride + + token_0) : + 0.0f; + scale_a_1 = token_1 < valid_m ? + ptx::ld_shared( + smem_sfa[pipeline_stage] + + activation_sf_group * + kL2SFAHalfStride + + token_1) : + 0.0f; + } + combined_scale_0 = + (compensate_secondary ? + scale_a_0 : + scale_a_0 * 64.0f) * + compensated_secondary; + combined_scale_1 = + (compensate_secondary ? + scale_a_1 : + scale_a_1 * 64.0f) * + compensated_secondary; + } #pragma unroll for (uint32_t half_idx = 0; half_idx < kNumWeightHalves; @@ -2424,10 +2626,16 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { kLocalSwapABTokens == 32 and kPackedBF16SwapEpilogue) or kLocalSwapABTokens == 64)) { - const nv_bfloat162 scale_pair = - __floats2bfloat162_rn( - combined_scale_0, - combined_scale_1); + const nv_bfloat162 scale_pair = [&]() { + if constexpr ( + kPrecomputeN8ScaleBeforeWait) { + return precomputed_scale_pair; + } else { + return __floats2bfloat162_rn( + combined_scale_0, + combined_scale_1); + } + }(); mxfp4_final_bf16[pair_offset] = __hfma2( scale_pair, @@ -2707,8 +2915,14 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { } #pragma unroll for (uint32_t i = 0; i < kAccumPerThread / 2; ++ i) { - const float2 pair = - __bfloat1622float2(mxfp4_final_bf16[i]); + const float2 pair = [&]() { + if constexpr (kUsePackedF16Accum) + return __half22float2( + mxfp4_final_bf16[i]); + else + return __bfloat1622float2( + mxfp4_final_bf16[i]); + }(); final_accum[i * 2] = pair.x; final_accum[i * 2 + 1] = pair.y; } @@ -3354,8 +3568,8 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { shared_l2_sf_buffer.get_base_ptr() : l2_sf_buffer.get_base_ptr(); constexpr uint32_t kSharedSFStride = - layout::get_num_shared_sf_tokens( - kNumMaxTokensPerRank, BLOCK_M); + layout::get_num_max_shared_sf_tokens( + kNumMaxTokensPerRank); const uint32_t sf_stride = is_shared_phase ? kSharedSFStride : kNumSFRingTokens; const uint32_t token_r0 = m_idx + row_offset_r0; diff --git a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh index 0d205f8470..3534712f08 100644 --- a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh @@ -39,18 +39,10 @@ CUTLASS_HOST_DEVICE constexpr T get_num_padded_sf_pool_tokens(T num_max_pool_tok return get_num_sf_ring_tokens(num_max_pool_tokens, block_m); } -// Shared L2 input SF capacity for a fixed schedule block size. -template -CUTLASS_HOST_DEVICE constexpr T get_num_shared_sf_tokens( - const T& num_max_tokens_per_rank, const T& block_m) { - return math::constexpr_ceil_div(num_max_tokens_per_rank, block_m) * 128; -} - -// Backward-compatible worst case over all candidate BLOCK_M values. +// Shared L2 input SF capacity: worst-case aligned SF pages over all candidate BLOCK_M. template CUTLASS_HOST_DEVICE constexpr T get_num_max_shared_sf_tokens(const T& num_max_tokens_per_rank) { - return get_num_shared_sf_tokens( - num_max_tokens_per_rank, static_cast(kMinCandidateBlockM)); + return math::constexpr_ceil_div(num_max_tokens_per_rank, kMinCandidateBlockM) * 128; } // Per-token source metadata for combine write-back @@ -437,16 +429,14 @@ struct MegaMoEBuffer { const uint32_t& num_sf_ring_tokens, const bool& with_sf, const uint32_t& num_shared_experts = 0, - const ScaleLayoutSpec& scale_layout_spec = ScaleLayoutSpec(), - const uint32_t shared_sf_block_m = kMinCandidateBlockM) { + const ScaleLayoutSpec& scale_layout_spec = ScaleLayoutSpec()) { // Workspace workspace = Workspace(base, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_ring_tokens); // Shared const auto shared_intermediate_hidden = intermediate_hidden * num_shared_experts; - const auto num_max_shared_sf_tokens = with_sf ? get_num_shared_sf_tokens( - num_max_tokens_per_rank, shared_sf_block_m) : 0u; + const auto num_max_shared_sf_tokens = with_sf ? get_num_max_shared_sf_tokens(num_max_tokens_per_rank) : 0u; // A zero-initialized spec is the backward-compatible SM100 default. // Keeping the scale row sizes explicit prevents an SM90 K64 L2 scale diff --git a/deep_gemm/include/deep_gemm/mma/sm90.cuh b/deep_gemm/include/deep_gemm/mma/sm90.cuh index 2c061940de..4683ee78d6 100644 --- a/deep_gemm/include/deep_gemm/mma/sm90.cuh +++ b/deep_gemm/include/deep_gemm/mma/sm90.cuh @@ -74,6 +74,53 @@ struct FP8MMASelector { using type = decltype(select_type()); }; +template +struct FP8MMAF16 { + template + CUTLASS_DEVICE static void call_fma_impl( + uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t* d, + bool scale_d, + cute::index_sequence) { + using namespace cute::SM90::GMMA; + MMA::fma(desc_a, desc_b, d[Idx]..., + (scale_d ? ScaleOut::One : ScaleOut::Zero)); + } + + CUTLASS_DEVICE static void wgmma( + uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t* d, + bool scale_d) { + call_fma_impl(desc_a, desc_b, d, scale_d, + cute::make_index_sequence{}); + } + + static constexpr int M = 64; + static constexpr int N = N_; + static constexpr int K = 32; + static constexpr int kNumAccum = M * N / 128; + static constexpr int kNumPackedAccum = kNumAccum / 2; +}; + +template +struct FP8MMAF16Selector { + static constexpr auto select_mma() { + using namespace cute::SM90::GMMA; + DG_STATIC_ASSERT( + N == 128, + "Packed FP16 accumulation is only specialized for N128"); + return MMA_64x128x32_F16E4M3E4M3_SS_TN(); + } + + static constexpr auto select_type() { + return FP8MMAF16(); + } + + using type = decltype(select_type()); +}; + template struct BF16MMA { template diff --git a/deep_gemm/include/deep_gemm/ptx/wgmma.cuh b/deep_gemm/include/deep_gemm/ptx/wgmma.cuh index 8912a15766..6e3e5b34d5 100644 --- a/deep_gemm/include/deep_gemm/ptx/wgmma.cuh +++ b/deep_gemm/include/deep_gemm/ptx/wgmma.cuh @@ -16,6 +16,10 @@ CUTLASS_DEVICE void warpgroup_fence_operand(float& reg) { asm volatile("" : "+f"(reg) :: "memory"); } +CUTLASS_DEVICE void warpgroup_fence_operand(uint32_t& reg) { + asm volatile("" : "+r"(reg) :: "memory"); +} + template CUTLASS_DEVICE void warpgroup_wait() { DG_STATIC_ASSERT(N >= 0 and N <= 7, "WGMMA wait: N must be in range [0, 7]"); diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 5582b5c0eb..2a48357941 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -4,6 +4,7 @@ from typing import Tuple, Optional, Union from ..utils.math import align from .mxfp4 import ( + MXFP4ProcessedWeights as _MXFP4ProcessedWeights, _normalize_mxfp4_ue8m0, _process_mxfp4_e8m0, _reorder_mxfp4_sign_bits_for_sm90, @@ -23,11 +24,6 @@ from .. import _C -_SM90_MEGA_MOE_OP_NAMES = { - 'fp8xmxfp4': 'fp8_mxfp4_mega_moe', -} - - class SymmBuffer: def __init__(self, group: dist.ProcessGroup, num_experts: int, @@ -82,23 +78,21 @@ def destroy(self): class SM90SymmBuffer: """Symmetric buffer for the SM90 persistent MegaMoE backend. - The Hopper backend follows the common twelve-view live-ring/shared-expert - ABI while retaining its architecture-specific FP32 K128 input scales and - FP32 K64 intermediate scales. + The buffer describes the FP8 activation/communication workspace shared by + SM90 routed-weight backends. Routed-weight formats such as MXFP4 or WinT4 + are launch-API concerns and do not change these twelve views. """ def __init__(self, group: dist.ProcessGroup, num_experts: int, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, - mma_type: str = 'fp8xmxfp4', + use_fp8_dispatch: bool = True, activation: str = 'swiglu', num_shared_experts: int = 0): if num_shared_experts < 0: raise ValueError('num_shared_experts must be non-negative') - if mma_type not in _SM90_MEGA_MOE_OP_NAMES: - raise ValueError( - f'Unsupported SM90 MegaMoE mma_type `{mma_type}`; ' - f'supported values: {tuple(_SM90_MEGA_MOE_OP_NAMES)}') + if not use_fp8_dispatch: + raise ValueError('SM90 MegaMoE requires FP8 dispatch') if activation != 'swiglu': raise ValueError( f'Only `swiglu` activation is supported, got `{activation}`') @@ -110,13 +104,12 @@ def __init__(self, group: dist.ProcessGroup, self.hidden = hidden self.intermediate_hidden = intermediate_hidden self.num_shared_experts = num_shared_experts - self.mma_type = mma_type 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, - mma_type, activation, num_shared_experts, + use_fp8_dispatch, activation, num_shared_experts, ) allocator = torch if group.size() == 1 else symm_mem self.buffer = allocator.empty(num_bytes, dtype=torch.int8, device='cuda') @@ -181,7 +174,7 @@ 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, - mma_type: str = 'fp8xmxfp4', + use_fp8_dispatch: bool = True, activation: str = 'swiglu', num_shared_experts: int = 0) -> SM90SymmBuffer: """Allocate the twelve-view SM90 live-ring MegaMoE symmetric buffer. @@ -195,7 +188,7 @@ def get_symm_buffer_for_sm90_mega_moe(group: dist.ProcessGroup, group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - mma_type=mma_type, + use_fp8_dispatch=use_fp8_dispatch, activation=activation, num_shared_experts=num_shared_experts, ) @@ -266,7 +259,10 @@ def transform_shared_weights_for_fp8_mega_moe_sm90( Input scales are natural row-major FP32 block-(128, 128) tensors. Only the L1 FP8 rows are gate/up-interleaved at granularity 8; unlike the SM100 - helper, neither scale tensor is UTCCP-transposed or interleaved. + helper, neither scale tensor is UTCCP-transposed or interleaved. This + transform is intentionally named after the shared-weight format, not the + routed-weight backend, so MXFP4 and future WinT4 routed experts can share + the same replicated FP8 expert representation. """ if activation != 'swiglu': raise ValueError( @@ -276,17 +272,11 @@ def transform_shared_weights_for_fp8_mega_moe_sm90( raise TypeError('shared SM90 weights must be FP8/FP32-scale pairs') l1_weight, l1_scale = l1_weights l2_weight, l2_scale = l2_weights - if not all(isinstance(tensor, torch.Tensor) for tensor in ( - l1_weight, l1_scale, l2_weight, l2_scale)): - raise TypeError('shared SM90 weights and scales must be torch.Tensor objects') if l1_weight.dtype != torch.float8_e4m3fn or \ l2_weight.dtype != torch.float8_e4m3fn: raise TypeError('shared SM90 weights must use torch.float8_e4m3fn') if l1_scale.dtype != torch.float32 or l2_scale.dtype != torch.float32: raise TypeError('shared SM90 weight scales must use torch.float32') - if not (l1_weight.device == l1_scale.device == - l2_weight.device == l2_scale.device): - raise ValueError('shared SM90 weights and scales must be on the same device') if l1_weight.dim() != 2 or l2_weight.dim() != 2: raise ValueError('shared SM90 weights must be two-dimensional') shared_intermediate_hidden = l2_weight.size(1) @@ -305,11 +295,6 @@ def transform_shared_weights_for_fp8_mega_moe_sm90( if not all(tensor.is_contiguous() for tensor in ( l1_weight, l1_scale, l2_weight, l2_scale)): raise ValueError('shared SM90 weights and scales must be contiguous') - if not all( - bool(torch.isfinite(scale).all().item()) and - bool((scale > 0).all().item()) - for scale in (l1_scale, l2_scale)): - raise ValueError('shared SM90 weight scales must contain finite positive values') return (_interleave_weights(l1_weight), l1_scale), (l2_weight, l2_scale) @@ -340,11 +325,46 @@ def fp8_fp4_mega_moe(y: torch.Tensor, ) -def _validate_sm90_mxfp4_routed_weights( - l1_weights: Tuple[torch.Tensor, ...], - l2_weights: Tuple[torch.Tensor, ...], - sym_buffer: SM90SymmBuffer, -) -> None: +def fp8_mxfp4_mega_moe(y: torch.Tensor, + l1_weights: _MXFP4ProcessedWeights, + l2_weights: _MXFP4ProcessedWeights, + sym_buffer: SM90SymmBuffer, + shared_l1_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l2_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + recipe: Tuple[int, int, int] = (1, 1, 32), + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True): + """Run the SM90 Humming-compatible MXFP4 MegaMoE path. + + Routed weights must be processed triples + ``(processed_e2m1, relative_ue8m0, weight_scale_2)`` returned by + :func:`transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90`. Keeping this + explicit entry point avoids architecture-dependent Python dispatch while + the common ``fp8_fp4_mega_moe`` facade remains backward compatible with + the SM100 implementation. + + When shared experts are enabled, prepare their FP8/FP32 pairs with + :func:`transform_shared_weights_for_fp8_mega_moe_sm90` and copy the + input K128 FP32 scales into ``sym_buffer.shared_l1_acts_sf`` before launch. + ``shared_l1_acts`` itself aliases ``sym_buffer.x``. + """ + if not isinstance(sym_buffer, SM90SymmBuffer): + raise TypeError( + 'fp8_mxfp4_mega_moe requires an SM90SymmBuffer allocated by ' + 'get_symm_buffer_for_sm90_mega_moe') + if (shared_l1_weights is None) != (shared_l2_weights is None): + raise ValueError( + 'shared_l1_weights and shared_l2_weights must be provided together') + num_shared_experts = getattr(sym_buffer, 'num_shared_experts', 0) + if num_shared_experts == 0 and shared_l1_weights is not None: + raise ValueError( + 'shared weights require an SM90SymmBuffer allocated with ' + 'num_shared_experts > 0') + if num_shared_experts > 0 and shared_l1_weights is None: + raise ValueError( + 'an SM90SymmBuffer with shared experts requires both shared weight tuples') _validate_processed_mxfp4_kernel_weights(l1_weights, l2_weights) num_ranks = sym_buffer.group.size() if sym_buffer.num_experts % num_ranks != 0: @@ -371,61 +391,6 @@ def _validate_sm90_mxfp4_routed_weights( f'H/I contract: expected {expected_l1_shape} and ' f'{expected_l2_shape}, got {tuple(l1_weights[0].shape)} and ' f'{tuple(l2_weights[0].shape)}') - - -_SM90_MEGA_MOE_WEIGHT_VALIDATORS = { - 'fp8xmxfp4': _validate_sm90_mxfp4_routed_weights, -} - - -def fp8_mega_moe(y: torch.Tensor, - l1_weights: Tuple[torch.Tensor, ...], - l2_weights: Tuple[torch.Tensor, ...], - sym_buffer: SM90SymmBuffer, - shared_l1_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - shared_l2_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, - recipe: Tuple[int, int, int] = (1, 1, 32), - activation: str = 'swiglu', - activation_clamp: Optional[float] = None, - fast_math: bool = True): - """Run an SM90 FP8-activation MegaMoE kernel selected by ``mma_type``. - - ``sym_buffer.mma_type`` selects the routed-weight backend. The current - ``fp8xmxfp4`` backend requires processed triples - ``(processed_e2m1, relative_ue8m0, weight_scale_2)`` returned by - :func:`transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90`. Future - FP8/FP8 and FP8/INT4 backends can add dispatch branches without changing - the public launch API. - - When shared experts are enabled, prepare their FP8/FP32 pairs with - :func:`transform_shared_weights_for_fp8_mega_moe_sm90` and copy the - input K128 FP32 scales into ``sym_buffer.shared_l1_acts_sf`` before launch. - ``shared_l1_acts`` itself aliases ``sym_buffer.x``. - """ - if not isinstance(sym_buffer, SM90SymmBuffer): - raise TypeError( - 'fp8_mega_moe requires an SM90SymmBuffer allocated by ' - 'get_symm_buffer_for_sm90_mega_moe') - try: - op_name = _SM90_MEGA_MOE_OP_NAMES[sym_buffer.mma_type] - validate_routed_weights = _SM90_MEGA_MOE_WEIGHT_VALIDATORS[ - sym_buffer.mma_type] - except KeyError as exception: - raise ValueError( - f'Unsupported SM90 MegaMoE mma_type `{sym_buffer.mma_type}`') from exception - if (shared_l1_weights is None) != (shared_l2_weights is None): - raise ValueError( - 'shared_l1_weights and shared_l2_weights must be provided together') - num_shared_experts = getattr(sym_buffer, 'num_shared_experts', 0) - if num_shared_experts == 0 and shared_l1_weights is not None: - raise ValueError( - 'shared weights require an SM90SymmBuffer allocated with ' - 'num_shared_experts > 0') - if num_shared_experts > 0 and shared_l1_weights is None: - raise ValueError( - 'an SM90SymmBuffer with shared experts requires both shared weight tuples') - validate_routed_weights(l1_weights, l2_weights, sym_buffer) if num_shared_experts > 0: shared_intermediate_hidden = ( num_shared_experts * sym_buffer.intermediate_hidden) @@ -442,11 +407,11 @@ def fp8_mega_moe(y: torch.Tensor, f'{tuple(shared_l1_weights[0].shape)} and ' f'{tuple(shared_l2_weights[0].shape)}') try: - op = getattr(_C, op_name) + op = _C.fp8_mxfp4_mega_moe except AttributeError as exception: raise RuntimeError( - f'DeepGEMM was built without the SM90 MegaMoE binding `{op_name}` ' - f'required by `fp8_mega_moe(mma_type="{sym_buffer.mma_type}")`; ' + 'DeepGEMM was built without the SM90 MXFP4 MegaMoE binding ' + '`fp8_mxfp4_mega_moe`; ' 'rebuild the extension after enabling the SM90 backend') from exception op( y, diff --git a/deep_gemm/mega/mxfp4.py b/deep_gemm/mega/mxfp4.py index cbb0f385ed..0ea1dac148 100644 --- a/deep_gemm/mega/mxfp4.py +++ b/deep_gemm/mega/mxfp4.py @@ -11,7 +11,7 @@ MXFP4CheckpointWeights = Tuple[torch.Tensor, torch.Tensor] -_MXFP4Payload = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] +MXFP4ProcessedWeights = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] _SM90_MXFP4_COALESCED_SCALE_MAX_HIDDEN = 8192 @@ -25,41 +25,6 @@ def _require(condition: bool, message: str) -> None: raise ValueError(message) -class MXFP4ProcessedWeights(tuple): - """Validated tuple payload produced by the SM90 MXFP4 transform. - - A tuple subclass preserves the public/PyBind tuple ABI while distinguishing - transformed relative-UE8M0 payloads from raw checkpoint triplets. Semantic - validation happens once at model-load or deserialization time instead of - scanning the full scale tensor before every kernel launch. - """ - - __slots__ = () - - def __new__(cls, values: _MXFP4Payload) -> 'MXFP4ProcessedWeights': - if not isinstance(values, tuple) or len(values) != 3: - raise TypeError( - 'SM90 processed MXFP4 weights require a three-tensor tuple') - weight, relative_sf, secondary = values - if not all(isinstance(tensor, torch.Tensor) for tensor in values): - raise TypeError('SM90 processed MXFP4 entries must be torch.Tensor objects') - if relative_sf.dtype != torch.uint8: - raise TypeError('SM90 relative UE8M0 scales must have dtype uint8') - _require( - bool(((relative_sf >= 1) & (relative_sf <= 12)).all().item()), - 'SM90 processed MXFP4 requires relative UE8M0 values in [1, 12]') - if secondary.dtype != torch.float32: - raise TypeError('SM90 MXFP4 secondary scale must have dtype float32') - _require( - bool(torch.isfinite(secondary).all().item()) and - bool((secondary > 0).all().item()), - 'SM90 MXFP4 secondary scale must contain finite positive values') - return super().__new__(cls, (weight, relative_sf, secondary)) - - def __getnewargs__(self) -> Tuple[_MXFP4Payload]: - return (tuple(self),) - - def _normalize_mxfp4_ue8m0(sf: torch.Tensor) -> torch.Tensor: """Return natural-layout UE8M0 exponent bytes. @@ -254,7 +219,7 @@ def _process_mxfp4_e8m0( weight: torch.Tensor, sf: torch.Tensor, interleave_rows: bool = False, -) -> _MXFP4Payload: +) -> MXFP4ProcessedWeights: """Convert raw UE8M0 scales to bounded per-expert exponent offsets. The retained relative exponent range is 11, matching Humming's fused-E8M0 @@ -353,12 +318,7 @@ def _apply_weight_scale_2( f'{name}_weight_scale_2 must contain only finite values') _require(bool((scale > 0).all().item()), f'{name}_weight_scale_2 must contain only positive values') - result = (secondary * scale).contiguous() - _require( - bool(torch.isfinite(result).all().item()) and - bool((result > 0).all().item()), - f'{name}_weight_scale_2 result must contain finite positive values') - return result + return (secondary * scale).contiguous() def transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( @@ -399,15 +359,15 @@ def transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( l1_sf = _transpose_mxfp4_scales_for_sm90(l1_sf) l2_sf = _transpose_mxfp4_scales_for_sm90(l2_sf) - return MXFP4ProcessedWeights(( + return ( l1_w, l1_sf, l1_secondary, - )), MXFP4ProcessedWeights(( + ), ( l2_w.contiguous(), l2_sf.contiguous(), l2_secondary, - )) + ) def _validate_processed_mxfp4_kernel_weights( @@ -465,7 +425,3 @@ def validate_payload( f'{name}_weight_scale_2 must be on the weight device') _require(scale.is_contiguous(), f'{name}_weight_scale_2 must be contiguous before launch') - if not isinstance(l1_weights, MXFP4ProcessedWeights) or \ - not isinstance(l2_weights, MXFP4ProcessedWeights): - raise ValueError( - 'L1/L2 SM90 MXFP4 weights must be returned by the SM90 transform') diff --git a/tests/bench_mega_moe_sm90.py b/tests/bench_mega_moe_sm90.py index 81260f2b17..970a724673 100644 --- a/tests/bench_mega_moe_sm90.py +++ b/tests/bench_mega_moe_sm90.py @@ -322,7 +322,7 @@ def run_kernel() -> torch.Tensor: buffer.x_sf[:num_tokens].copy_(x_scale) buffer.topk_idx[:num_tokens].copy_(topk_idx) buffer.topk_weights[:num_tokens].copy_(topk_weights) - deep_gemm.fp8_mega_moe( + deep_gemm.fp8_mxfp4_mega_moe( output, transformed_l1, transformed_l2, diff --git a/tests/test_mega_moe_mxfp4_preprocess.py b/tests/test_mega_moe_mxfp4_preprocess.py index 026bdab2b5..c26b040b29 100644 --- a/tests/test_mega_moe_mxfp4_preprocess.py +++ b/tests/test_mega_moe_mxfp4_preprocess.py @@ -1,6 +1,7 @@ """CPU contract tests for SM90 Humming-compatible MXFP4 weight transforms.""" import importlib.util +import inspect from pathlib import Path import sys import types @@ -232,42 +233,6 @@ def test_canonical_transform_applies_checkpoint_weight_scale_2(): mxfp4._validate_processed_mxfp4_kernel_weights(transformed_l1, transformed_l2) -def test_processed_payload_rejects_plain_triplets_and_invalid_semantics(): - checkpoint_l1, checkpoint_l2 = _valid_checkpoint_weights() - processed_l1, processed_l2 = ( - mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( - checkpoint_l1, checkpoint_l2)) - - with pytest.raises(ValueError, match='returned by the SM90 transform'): - mxfp4._validate_processed_mxfp4_kernel_weights( - tuple(processed_l1), processed_l2) - - raw_relative_sf = torch.full_like(processed_l1[1], 109) - with pytest.raises(ValueError, match=r'relative UE8M0 values in \[1, 12\]'): - mxfp4.MXFP4ProcessedWeights(( - processed_l1[0], raw_relative_sf, processed_l1[2])) - - for invalid_secondary in (float('nan'), float('inf'), 0.0, -1.0): - with pytest.raises(ValueError, match='finite positive'): - mxfp4.MXFP4ProcessedWeights(( - processed_l1[0], processed_l1[1], - torch.tensor([invalid_secondary], dtype=torch.float32), - )) - - -def test_checkpoint_weight_scale_2_overflow_is_rejected(): - (l1_w, l1_sf), (l2_w, l2_sf) = _valid_checkpoint_weights() - l1_sf.fill_(254) - l2_sf.fill_(254) - with pytest.raises(ValueError, match='finite positive'): - mxfp4.transform_weights_for_fp8_mxfp4_fused_mega_moe_sm90( - (l1_w, l1_sf), - (l2_w, l2_sf), - l1_weight_scale_2=torch.tensor([8.0], dtype=torch.float32), - l2_weight_scale_2=torch.tensor([8.0], dtype=torch.float32), - ) - - def test_invalid_mxfp4_contracts_fail_before_kernel_launch(): l1, l2 = _valid_checkpoint_weights() @@ -321,7 +286,7 @@ def test_explicit_wrapper_rejects_the_sm100_buffer_abi_before_launch(): mega, package_name = _load_mega_api_for_cpu() try: with pytest.raises(TypeError, match='requires an SM90SymmBuffer'): - mega.fp8_mega_moe(None, (), (), object()) + mega.fp8_mxfp4_mega_moe(None, (), (), object()) finally: _unload_fake_package(package_name) @@ -336,7 +301,6 @@ def size(): buffer = mega.SM90SymmBuffer.__new__(mega.SM90SymmBuffer) buffer.group = FakeGroup() - buffer.mma_type = 'fp8xmxfp4' buffer.num_experts = 4 checkpoint_l1, checkpoint_l2 = _valid_checkpoint_weights(num_experts=1) transformed_l1, transformed_l2 = ( @@ -344,7 +308,7 @@ def size(): checkpoint_l1, checkpoint_l2)) try: with pytest.raises(ValueError, match='expected E_local=2, got 1'): - mega.fp8_mega_moe( + mega.fp8_mxfp4_mega_moe( None, transformed_l1, transformed_l2, buffer) finally: _unload_fake_package(package_name) @@ -365,7 +329,6 @@ def rank(): buffer = mega.SM90SymmBuffer.__new__(mega.SM90SymmBuffer) buffer.group = FakeGroup() - buffer.mma_type = 'fp8xmxfp4' buffer.num_experts = 1 buffer.num_max_tokens_per_rank = 128 buffer.num_topk = 1 @@ -382,7 +345,7 @@ def rank(): fp8_mxfp4_mega_moe=lambda *args: calls.append(args)) y = object() try: - mega.fp8_mega_moe( + mega.fp8_mxfp4_mega_moe( y, processed_l1, processed_l2, buffer, recipe=(1, 1, 32), activation='swiglu', activation_clamp=10.0, fast_math=False) @@ -405,6 +368,12 @@ def test_sm90_buffer_uses_dedicated_alignment_and_twelve_view_abi(): mega, package_name = _load_mega_api_for_cpu() sizing_calls = [] + assert tuple(inspect.signature( + mega.get_symm_buffer_for_sm90_mega_moe).parameters) == ( + 'group', 'num_experts', 'num_max_tokens_per_rank', 'num_topk', + 'hidden', 'intermediate_hidden', 'use_fp8_dispatch', 'activation', + 'num_shared_experts') + class FakeBuffer: def __init__(self): self.zeroed = False @@ -447,9 +416,14 @@ def get_size(*args): ) group = FakeGroup() try: - with pytest.raises(ValueError, match='Unsupported SM90 MegaMoE mma_type'): + with pytest.raises(ValueError, match='requires FP8 dispatch'): + mega.get_symm_buffer_for_sm90_mega_moe( + group, 16, 128, 2, 512, 256, + use_fp8_dispatch=False) + with pytest.raises(ValueError, match='Only `swiglu` activation'): mega.get_symm_buffer_for_sm90_mega_moe( - group, 16, 128, 2, 512, 256, mma_type='fp8xint4') + group, 16, 128, 2, 512, 256, + activation='gelu') buffer = mega.get_symm_buffer_for_sm90_mega_moe( group, num_experts=16, @@ -458,9 +432,8 @@ def get_size(*args): hidden=512, intermediate_hidden=256, ) - assert buffer.mma_type == 'fp8xmxfp4' assert sizing_calls == [( - 1, 16, 128, 2, 512, 256, 'fp8xmxfp4', 'swiglu', 0)] + 1, 16, 128, 2, 512, 256, True, 'swiglu', 0)] assert buffer.num_max_tokens_per_rank == 128 assert group.barriers == 1 assert ( @@ -481,7 +454,7 @@ def get_size(*args): group, 16, 128, 2, 512, 256, num_shared_experts=1) assert shared_buffer.num_shared_experts == 1 assert sizing_calls[-1] == ( - 1, 16, 128, 2, 512, 256, 'fp8xmxfp4', 'swiglu', 1) + 1, 16, 128, 2, 512, 256, True, 'swiglu', 1) shared_buffer.destroy() finally: _unload_fake_package(package_name) @@ -508,17 +481,5 @@ def test_sm90_shared_weight_transform_interleaves_only_l1_fp8_rows(): assert transformed_l1[1].data_ptr() == l1_scale.data_ptr() assert transformed_l2[0].data_ptr() == l2_weight.data_ptr() assert transformed_l2[1].data_ptr() == l2_scale.data_ptr() - - for invalid_scale in (float('nan'), float('inf'), 0.0, -1.0): - bad_l1_scale = torch.full_like(l1_scale, invalid_scale) - with pytest.raises(ValueError, match='finite positive'): - mega.transform_shared_weights_for_fp8_mega_moe_sm90( - (l1_weight, bad_l1_scale), (l2_weight, l2_scale)) - - meta_l1_scale = torch.empty( - l1_scale.shape, dtype=torch.float32, device='meta') - with pytest.raises(ValueError, match='same device'): - mega.transform_shared_weights_for_fp8_mega_moe_sm90( - (l1_weight, meta_l1_scale), (l2_weight, l2_scale)) finally: _unload_fake_package(package_name) diff --git a/tests/test_mega_moe_sm90.py b/tests/test_mega_moe_sm90.py index 8338e79f13..c2d7fd4eb5 100644 --- a/tests/test_mega_moe_sm90.py +++ b/tests/test_mega_moe_sm90.py @@ -16,7 +16,6 @@ import sys from typing import Any, Dict, List, Optional, Tuple -import pytest import torch import torch.distributed as dist @@ -49,29 +48,6 @@ def barrier() -> None: torch.cuda.synchronize() -def _make_forced_ring_wrap_topk_idx( - num_tokens: int, - num_topk: int, - num_experts: int, - rank_idx: int, - num_ranks: int, - device: torch.device, -) -> torch.Tensor: - """Assign routes so every expert receives at least one global token.""" - if num_tokens * num_topk * num_ranks < num_experts: - raise ValueError( - 'forced ring-wrap routing needs at least one route per expert') - first_route = rank_idx * num_tokens * num_topk - return ( - torch.arange( - first_route, - first_route + num_tokens * num_topk, - dtype=torch.int64, - device=device, - ) % num_experts - ).reshape(num_tokens, num_topk) - - def _quantize_grouped_mxfp4(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Create packed E2M1 ``[E,N,K/2]`` plus K32 UE8M0 FP32 powers.""" from deep_gemm.utils import per_token_cast_to_fp4 @@ -413,10 +389,6 @@ def _run_scenario( num_tokens, num_experts, dtype=torch.float32, device='cuda') topk_weights, topk_idx = torch.topk( scores, num_topk, dim=-1, largest=True, sorted=False) - if config.get('require_ring_wrap', False): - topk_idx = _make_forced_ring_wrap_topk_idx( - num_tokens, num_topk, num_experts, - rank_idx, num_ranks, scores.device) else: assert 0 <= hot_route_rank < num_ranks assert num_topk <= num_local_experts @@ -582,7 +554,7 @@ def _run_scenario( recv_stats = torch.zeros( num_local_experts, dtype=torch.int32, device='cuda') for _ in range(repeat_count): - deep_gemm.fp8_mega_moe( + deep_gemm.fp8_mxfp4_mega_moe( output, transformed_l1, transformed_l2, @@ -810,6 +782,26 @@ def _full_scenarios( activation_clamp=10.0, require_ring_wrap=True, )), + ('production.flash_m256', dict( + num_max_tokens_per_rank=256, + num_tokens=256, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + )), + ('production.flash_m512', dict( + num_max_tokens_per_rank=512, + num_tokens=512, + hidden=4096, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=6, + fast_math=True, + activation_clamp=10.0, + )), ('production.flash_m1024', dict( num_max_tokens_per_rank=1024, num_tokens=1024, @@ -840,6 +832,31 @@ def _full_scenarios( fast_math=True, activation_clamp=10.0, )), + # Compatibility probes intentionally exercise dimensions rather than + # model-name dispatch. They cover shared experts and wider top-k so + # optimization selectors cannot silently assume the DSV4 topology. + ('compat.h6144_i2048_e256_topk8_shared1_m512', dict( + num_max_tokens_per_rank=512, + num_tokens=512, + hidden=6144, + intermediate_hidden=2048, + num_experts=32 * num_ranks, + num_topk=8, + num_shared_experts=1, + fast_math=True, + activation_clamp=10.0, + )), + ('compat.h7168_i3072_e896_topk16_shared2_m512', dict( + num_max_tokens_per_rank=512, + num_tokens=512, + hidden=7168, + intermediate_hidden=3072, + num_experts=112 * num_ranks, + num_topk=16, + num_shared_experts=2, + fast_math=True, + activation_clamp=10.0, + )), ('production.pro_m8', dict( num_max_tokens_per_rank=128, num_tokens=8, @@ -1057,7 +1074,7 @@ def _test_worker(local_rank: int, num_local_ranks: int, args: argparse.Namespace raise SystemExit(1) -def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: +def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description='Cross-rank SM90 FP8 x MXFP4 MegaMoE runtime validation') parser.add_argument( @@ -1082,34 +1099,7 @@ def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: '--filter', default='', help='only run scenario names containing this text') parser.add_argument( '--fail-fast', action='store_true', help='stop after the first failed scenario') - args = parser.parse_args(argv) - if args.num_processes <= 0: - parser.error('--num-processes must be positive') - return args - - -def test_forced_ring_wrap_routes_cover_every_expert() -> None: - num_ranks = 8 - routes = torch.cat([ - _make_forced_ring_wrap_topk_idx( - num_tokens=8, - num_topk=6, - num_experts=32 * num_ranks, - rank_idx=rank_idx, - num_ranks=num_ranks, - device=torch.device('cpu'), - ).reshape(-1) - for rank_idx in range(num_ranks) - ]) - counts = torch.bincount(routes, minlength=32 * num_ranks) - assert torch.all(counts > 0) - - -@pytest.mark.parametrize('num_processes', [0, -1]) -def test_cli_rejects_nonpositive_process_counts(num_processes: int) -> None: - with pytest.raises(SystemExit) as exception: - _parse_args(['--num-processes', str(num_processes)]) - assert exception.value.code == 2 + return parser.parse_args() if __name__ == '__main__':