diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index 90a165e0bd..2654ac7796 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -54,7 +54,7 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, } // (INT, 1, gran_k) on SM100/SM120: transform to TMA-aligned and MN-major - if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and (arch_major == 10 or arch_major == 12)) + if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 16 or gran_k == 32 or gran_k == 128) and (arch_major == 10 or arch_major == 12)) return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, true, false, torch::kInt); DG_HOST_UNREACHABLE("Unknown SF transformation"); diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index e7e899ec67..c46ea429e8 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -2,7 +2,7 @@ #include #include -// #include +#include #include #include @@ -14,8 +14,6 @@ #include "../jit_kernels/impls/sm100_bf16_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" #include "../jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp" -#include "../utils/math.hpp" -#include "../utils/system.hpp" namespace deep_gemm::mega { @@ -36,7 +34,8 @@ static int get_block_m_for_mega_moe( static std::tuple(const torch::Tensor&)>> + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor>(const torch::Tensor&)>> get_symm_buffer_size_for_mega_moe( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, @@ -48,8 +47,9 @@ get_symm_buffer_size_for_mega_moe( // SiTU is implemented only by the SM100 FP8xFP4 MegaMoE kernel. const auto mma_kind = parse_mma_kind(mma_type); - DG_HOST_ASSERT(activation == "swiglu" or + DG_HOST_ASSERT(activation == "swiglu" or activation == "swigluoai" or (mma_kind == MmaKind::MXFP8FP4 and activation == "situ")); + DG_HOST_ASSERT(num_shared_experts >= 0); // Ring capacity: worst-case live pool blocks over all candidate BLOCK_M; mirrors the kernel assert. // TODO: we temporarily assume the SM count is consistent with the runtime value @@ -80,18 +80,10 @@ get_symm_buffer_size_for_mega_moe( // Parse MMA type const auto with_sf = is_mma_with_sf(mma_kind); - // FP4 activations are currently a routed-expert-only optimization. Keep - // the upstream shared-expert path on its original FP8 representation. - const bool host_use_fp4_acts = with_sf and num_shared_experts == 0 and - get_env("DG_USE_FP4_ACTS") != 0; - - // Stream B (combine path): when `DG_USE_FP8_COMBINE=1`, the combine slot - // holds FP8 E4M3 (kHidden bytes/token) + a separate combine_sf slot - // holding UE8M0 SF bytes (kHidden/128 bytes/token, gran_k=128). When off, - // the combine slot holds BF16 (kHidden*2 bytes/token) and combine_sf is - // unused (zero-sized). + const bool host_use_fp8_combine = with_sf and get_env("DG_USE_FP8_COMBINE") != 0; - // Padded SF pool tokens + + // Compute num_sf_ring_tokens (max across all candidate block sizes) int num_sf_ring_tokens = 0; if (with_sf) { for (auto block_m: layout::kCandidateBlockM) { @@ -105,10 +97,21 @@ get_symm_buffer_size_for_mega_moe( 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, - num_shared_experts, host_use_fp4_acts, host_use_fp8_combine + num_topk, num_ring_tokens, num_sf_ring_tokens, mma_kind, + num_shared_experts, host_use_fp8_combine ); + // Activation layout: FP8 for MXFP8FP4, raw bytes holding 2 packed FP4 values for NVFP4/MXFP4 + const auto acts_dtype = with_sf ? + (get_element_bits(mma_kind) == 4 ? torch::kUInt8 : torch::kFloat8_e4m3fn) : torch::kBFloat16; + const int acts_storage_bits = static_cast(c10::elementSize(acts_dtype)) * 8; + const auto num_acts_cols = [=](const int& num_elems) { + return num_elems * get_element_bits(mma_kind) / acts_storage_bits; + }; + const auto num_sf_cols = [=](const int& num_elems) { + return num_elems / (get_sf_gran_k(mma_kind) * 4); + }; + // Check SF buffer requirements if (with_sf) { DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); @@ -118,20 +121,14 @@ get_symm_buffer_size_for_mega_moe( // Slice function: creates tensor views from the raw buffer. // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major - // Stream A0.0b: under `host_use_fp4_acts`, the `x` and `l1_acts` views - // expose packed E2M1 (`kPackedFP4` = `torch::kInt8`, 2 elements/byte) of - // shape `[..., hidden / 2]`. Underlying buffer bytes are the same as the - // sized `fp8_token_layout` slot, just half the row width. - const auto x_dtype = with_sf ? (host_use_fp4_acts ? kPackedFP4 : torch::kFloat8_e4m3fn) : torch::kBFloat16; - const int x_inner_cols = host_use_fp4_acts ? (hidden / 2) : hidden; 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, x_inner_cols}, - torch::TensorOptions().dtype(x_dtype).device(buffer.device())); + {num_max_tokens_per_rank, num_acts_cols(hidden)}, + torch::TensorOptions().dtype(acts_dtype).device(buffer.device())); auto x_sf = with_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}, + {num_max_tokens_per_rank, num_sf_cols(hidden)}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto topk_idx = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_topk_idx_buffer.base)), @@ -145,40 +142,44 @@ get_symm_buffer_size_for_mega_moe( auto shared_l1_acts = x; auto shared_l1_acts_sf = (with_sf and num_shared_experts > 0) ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l1_sf_buffer.base)), - {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), hidden / 128}, + {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), num_sf_cols(hidden)}, {1, layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank)}, torch::TensorOptions().dtype(torch::kInt).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(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : torch::Tensor(); + {num_max_tokens_per_rank, num_acts_cols(shared_intermediate_hidden)}, + torch::TensorOptions().dtype(acts_dtype).device(buffer.device())) : torch::Tensor(); auto shared_l2_acts_sf = (with_sf and num_shared_experts > 0) ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l2_sf_buffer.base)), - {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), shared_intermediate_hidden / 128}, + {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), num_sf_cols(shared_intermediate_hidden)}, {1, layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank)}, torch::TensorOptions().dtype(torch::kInt).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, x_inner_cols}, - torch::TensorOptions().dtype(x_dtype).device(buffer.device())); + {num_ring_tokens, num_acts_cols(hidden)}, + torch::TensorOptions().dtype(acts_dtype).device(buffer.device())); auto l1_acts_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l1_sf_buffer.base)), - {num_sf_ring_tokens, hidden / 128}, + {num_sf_ring_tokens, num_sf_cols(hidden)}, {1, num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); 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(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); + {num_ring_tokens, num_acts_cols(intermediate_hidden)}, + torch::TensorOptions().dtype(acts_dtype).device(buffer.device())); auto l2_acts_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l2_sf_buffer.base)), - {num_sf_ring_tokens, intermediate_hidden / 128}, + {num_sf_ring_tokens, num_sf_cols(intermediate_hidden)}, {1, num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); + auto x_scales = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_x_scales_buffer.base)), + {num_max_tokens_per_rank}, + 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); + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, x_scales); }; return {mega_buffer.get_num_bytes(), slice_input_buffers}; } @@ -195,9 +196,14 @@ static void fp8_fp4_mega_moe( const int& num_max_tokens_per_rank, const int& num_experts, const int& num_topk, const std::tuple& recipe, + const std::string& mma_type, const std::string& activation, const std::optional& activation_clamp_opt, - const bool& fast_math + const bool& fast_math, + const bool& use_x_scales, + const std::optional& l1_alphas, + const std::optional& l2_alphas, + const std::optional& l2_act_scales ) { const auto [l1_weights, l1_weights_sf] = l1_weights_tuple; const auto [l2_weights, l2_weights_sf] = l2_weights_tuple; @@ -205,11 +211,23 @@ static void fp8_fp4_mega_moe( // Config checks 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(shared_l1_weights_tuple_opt.has_value() == shared_l2_weights_tuple_opt.has_value()); - DG_HOST_ASSERT(activation == "swiglu" or activation == "situ"); + DG_HOST_ASSERT(rm == 1 and rn == 1 and (rk == 32 or rk == 16)); + const auto mma_kind = parse_mma_kind(mma_type); + if (mma_kind != MmaKind::MXFP8FP4 and mma_kind != MmaKind::MXFP4 and mma_kind != MmaKind::NVFP4) + DG_HOST_UNREACHABLE("`" + mma_type + "` activations are not implemented by the SM100 " + "mega-MoE kernel (only `fp8xfp4`, `mxf4xmxf4` and `nvfp4xnvfp4` are); " + "allocate the symmetric buffer with one of those"); + if (rk != get_sf_gran_k(mma_kind)) + DG_HOST_UNREACHABLE("recipe K granularity " + std::to_string(rk) + " does not match `" + + mma_type + "` (expected " + + std::to_string(get_sf_gran_k(mma_kind)) + ")"); + DG_HOST_ASSERT(not use_x_scales or mma_kind == MmaKind::NVFP4); + DG_HOST_ASSERT(activation == "swiglu" or activation == "swigluoai" or + (mma_kind == MmaKind::MXFP8FP4 and activation == "situ")); DG_HOST_ASSERT(activation != "situ" or not activation_clamp_opt.has_value()); const bool use_situ = activation == "situ"; + const float swiglu_alpha = activation == "swigluoai" ? 1.702f : 0.0f; + DG_HOST_ASSERT(shared_l1_weights_tuple_opt.has_value() == shared_l2_weights_tuple_opt.has_value()); // Activation checks const auto activation_clamp = @@ -232,8 +250,9 @@ static void fp8_fp4_mega_moe( DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); - // Check weight SF layout for UE8M0 packing, MN-major, and TMA alignment - constexpr int kGranMN = 1, kGranK = 32; + // Check weight SF layout for byte-packing, MN-major, and TMA alignment + constexpr int kGranMN = 1; + const int kGranK = rk; check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, num_experts_per_rank, true, false, torch::kInt); check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, @@ -244,16 +263,19 @@ static void fp8_fp4_mega_moe( 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(); - shared_intermediate_hidden = static_cast(shared_l2_weights.size(1)); + const bool is_packed_fp4 = mma_kind == MmaKind::MXFP4 or mma_kind == MmaKind::NVFP4; + shared_intermediate_hidden = static_cast(shared_l2_weights.size(1)) * + (is_packed_fp4 ? 2 : 1); num_shared_experts = shared_intermediate_hidden / intermediate_hidden; DG_HOST_ASSERT(shared_intermediate_hidden % intermediate_hidden == 0); DG_HOST_ASSERT(shared_l1_weights.dim() == 2 and shared_l2_weights.dim() == 2); DG_HOST_ASSERT(shared_l1_weights.size(0) == shared_intermediate_hidden * 2); - DG_HOST_ASSERT(shared_l1_weights.size(1) == hidden); + DG_HOST_ASSERT(shared_l1_weights.size(1) == (is_packed_fp4 ? hidden / 2 : hidden)); DG_HOST_ASSERT(shared_l2_weights.size(0) == hidden); - DG_HOST_ASSERT(shared_l1_weights.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(shared_l2_weights.scalar_type() == torch::kFloat8_e4m3fn); + const auto shared_weight_dtype = is_packed_fp4 ? kPackedFP4 : torch::kFloat8_e4m3fn; + DG_HOST_ASSERT(shared_l1_weights.scalar_type() == shared_weight_dtype); + DG_HOST_ASSERT(shared_l2_weights.scalar_type() == shared_weight_dtype); 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); DG_HOST_ASSERT(get_major_type_ab(shared_l2_weights) == cute::UMMA::Major::K); @@ -277,35 +299,45 @@ static void fp8_fp4_mega_moe( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - "fp8xfp4", activation, num_shared_experts + mma_type, activation, num_shared_experts ); - DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + if (sym_buffer.nbytes() < static_cast(num_required_bytes)) + DG_HOST_UNREACHABLE("symmetric buffer is " + std::to_string(sym_buffer.nbytes()) + + " bytes but `" + mma_type + "`/`" + activation + "` with " + + std::to_string(num_shared_experts) + " shared expert(s) needs " + + std::to_string(num_required_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); + if (l1_alphas.has_value()) { + DG_HOST_ASSERT(mma_kind == MmaKind::NVFP4); + DG_HOST_ASSERT(l1_alphas->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(l1_alphas->is_contiguous()); + DG_HOST_ASSERT(l1_alphas->dim() == 2); + DG_HOST_ASSERT(l1_alphas->size(0) == num_experts_per_rank); + DG_HOST_ASSERT(l1_alphas->size(1) == 2); + } + + if (l2_alphas.has_value()) { + DG_HOST_ASSERT(mma_kind == MmaKind::NVFP4); + DG_HOST_ASSERT(l2_alphas->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(l2_alphas->is_contiguous()); + DG_HOST_ASSERT(l2_alphas->dim() == 1); + DG_HOST_ASSERT(l2_alphas->size(0) == num_experts_per_rank); + } + + // Per-expert fc2 input global scales, `[num_experts_per_rank]` FP32 + if (l2_act_scales.has_value()) { + DG_HOST_ASSERT(mma_kind == MmaKind::NVFP4); + DG_HOST_ASSERT(l2_act_scales->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(l2_act_scales->is_contiguous()); + DG_HOST_ASSERT(l2_act_scales->dim() == 1); + DG_HOST_ASSERT(l2_act_scales->size(0) == num_experts_per_rank); + } + // 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); - - // Stream A0.1: pick up FP4-acts flag from `DG_USE_FP4_ACTS` env var. - // Default off — preserves byte-identical FP8-acts behavior. Setting - // `DG_USE_FP4_ACTS=1` flips L1's epilogue quant to E2M1 + UE8M0 SF. - const bool use_fp4_acts = num_shared_experts == 0 and - get_env("DG_USE_FP4_ACTS") != 0; - // Stream A0.5: when also `DG_USE_MXF4_KIND=1`, the L1 and L2 mainloops - // run `tcgen05.mma.kind::mxf4.block_scale.block32` instead of - // `kind::mxf8f6f4` — K=64 dense per call (vs K=32 with-padding), dense - // FP4 smem (`_ALIGN8B`, half the byte footprint), scale_vec::2X SF - // protocol with HALF-WORD address bits. Only honored when - // `DG_USE_FP4_ACTS=1` (kind::mxf4 is FP4-only). See A6 capstone / - // B2 standalone GEMM for the +20-22% headline. - const bool use_mxf4_kind = use_fp4_acts and get_env("DG_USE_MXF4_KIND") != 0; - // Stream B (combine path): when `DG_USE_FP8_COMBINE=1`, the L2 epilogue - // ships FP8 E4M3 + per-(token, N=128) UE8M0 SF over NVLink instead of - // BF16. The combine reduce dequantizes on the fly. NVLink bytes/token - // halve (from kHidden*2 → kHidden + kHidden/128). Independent of the - // FP4-acts / MXF4-kind flags above (those control the dispatch a2a + - // mainloops; this controls the combine a2a only). + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, x_scales] = slice(sym_buffer); const bool use_fp8_combine = get_env("DG_USE_FP8_COMBINE") != 0; // Dispatch into different architectures @@ -326,8 +358,16 @@ static void fp8_fp4_mega_moe( num_shared_experts, num_tokens, num_topk, hidden, intermediate_hidden, - activation_clamp, use_situ, fast_math, - use_fp4_acts, use_mxf4_kind, use_fp8_combine); + activation_clamp, swiglu_alpha, use_situ, fast_math, + use_x_scales, + l1_alphas.has_value() + ? l1_alphas->const_data_ptr() : nullptr, + l2_alphas.has_value() + ? l2_alphas->const_data_ptr() : nullptr, + l2_act_scales.has_value() + ? l2_act_scales->const_data_ptr() : nullptr, + mma_kind, + use_fp8_combine); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); } @@ -419,7 +459,7 @@ static void bf16_mega_moe( // 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); + l1_acts, _l1_acts_sf, l2_acts, _l2_acts_sf, _x_scales] = slice(sym_buffer); // Dispatch into different architectures if (arch_major == 10) { @@ -464,7 +504,9 @@ static void register_apis(pybind11::module_& m) { pybind11::arg("buf_topk_weights"), pybind11::arg("num_tokens"), pybind11::arg("group_size") = 32, - pybind11::arg("use_fp4_acts") = false); + pybind11::arg("mma_type") = "fp8xfp4", + pybind11::arg("buf_x_scales") = std::nullopt, + pybind11::arg("expert_scales") = std::nullopt); #endif } diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index 2e43626bb6..2bcf6faf77 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -30,24 +30,29 @@ static void mega_moe_pre_dispatch_sm90( num_tokens, group_size, routed_scaling_factor); } -static std::tuple(const torch::Tensor&)>> +static std::tuple(const torch::Tensor&)>> get_symm_buffer_size_for_sm90_mega_moe( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, const int& hidden, const int& intermediate_hidden, - const bool& use_fp8_dispatch, const std::string& activation) { + const bool& use_fp8_dispatch, const std::string& activation, + const int& num_shared_experts = 0) { DG_HOST_ASSERT(num_experts % num_ranks == 0); DG_HOST_ASSERT(use_fp8_dispatch); DG_HOST_ASSERT(activation == "swiglu"); + DG_HOST_ASSERT(num_shared_experts >= 0); const auto workspace = layout::SM90Workspace( nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk); + const auto shared_intermediate_hidden = intermediate_hidden * num_shared_experts; const auto fp8_token_layout = layout::Data(hidden); const auto bf16_token_layout = layout::Data(hidden * 2); const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden); const auto fp8_sf_layout = layout::Data(hidden / 32); const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 16); + const auto fp8_shared_intermediate_token_layout = layout::Data(shared_intermediate_hidden); + const auto fp8_shared_intermediate_sf_layout = layout::Data(shared_intermediate_hidden / 16); const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false); const auto l1_topk_weights_layout = layout::Data(sizeof(float), false); @@ -91,10 +96,22 @@ get_symm_buffer_size_for_sm90_mega_moe( fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens, l2_token_buffer.get_end_ptr()); + // The fused shared expert reduces through one extra combine slot on the local rank const auto combine_token_buffer = layout::Buffer( - bf16_token_layout, num_topk, num_max_tokens_per_rank, + bf16_token_layout, num_topk + (num_shared_experts > 0 ? 1 : 0), num_max_tokens_per_rank, l2_sf_buffer.get_end_ptr()); + // Fused shared-expert area, appended after the combine buffer so the routed + // regions keep their relative order and are zero-sized when the shared expert is + // disabled. Both are indexed by the local token and the SF buffer is K-major + // (per-64 K groups), so no SF-pool padding is needed. + const auto shared_l2_token_buffer = layout::Buffer( + fp8_shared_intermediate_token_layout, 1, num_shared_experts > 0 ? num_max_tokens_per_rank : 0, + combine_token_buffer.get_end_ptr()); + const auto shared_l2_sf_buffer = layout::Buffer( + fp8_shared_intermediate_sf_layout, 1, num_shared_experts > 0 ? num_max_tokens_per_rank : 0, + shared_l2_token_buffer.get_end_ptr()); + DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); auto slice_input_buffers = [=](const torch::Tensor& buffer) { @@ -132,15 +149,34 @@ get_symm_buffer_size_for_sm90_mega_moe( {num_max_padded_sf_pool_tokens, intermediate_hidden / 64}, {1, num_max_padded_sf_pool_tokens}, torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); + // Fused shared expert: post-SwiGLU FP8 pool plus its M-major per-64 float SF + // (token-contiguous inner stride so a (BLOCK_M, 1) TMA box is legal; same + // layout-class as the routed L2 acts SF pool). Zero-sized when the shared + // expert is off (kept defined so the returned tuple type never changes). + auto shared_l2_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(shared_l2_token_buffer.base)), + {num_shared_experts > 0 ? num_max_tokens_per_rank : 0, shared_intermediate_hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto shared_l2_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(shared_l2_sf_buffer.base)), + {num_shared_experts > 0 ? num_max_tokens_per_rank : 0, shared_intermediate_hidden / 64}, + {1, num_shared_experts > 0 ? num_max_tokens_per_rank : 0}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf); }; - return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; + return {reinterpret_cast( + num_shared_experts > 0 ? shared_l2_sf_buffer.get_end_ptr() + : combine_token_buffer.get_end_ptr()), + slice_input_buffers}; } static void fp8_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, @@ -186,6 +222,37 @@ static void fp8_mega_moe( check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, num_experts_per_rank, false, true, torch::kFloat); + // Fused shared expert: a single dense MLP (no expert dimension) whose intermediate + // size is `num_shared_experts * intermediate_hidden`. Both weight tuples must be + // given together; the SF layout matches the routed weights minus the group axis. + DG_HOST_ASSERT(shared_l1_weights_tuple_opt.has_value() == shared_l2_weights_tuple_opt.has_value()); + int num_shared_experts = 0; + torch::Tensor shared_l1_weights, shared_l1_weights_sf, shared_l2_weights, shared_l2_weights_sf; + 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(); + const auto shared_intermediate_hidden = static_cast(shared_l2_weights.size(1)); + DG_HOST_ASSERT(shared_intermediate_hidden % intermediate_hidden == 0); + num_shared_experts = shared_intermediate_hidden / intermediate_hidden; + // The shared L2 activation SF is K-major, so its per-token row (SIH / 64 + // floats) must stay 16-byte aligned for the TMA loads of the pool it feeds + DG_HOST_ASSERT(shared_intermediate_hidden % 256 == 0); + + DG_HOST_ASSERT(shared_l1_weights.dim() == 2 and shared_l2_weights.dim() == 2); + DG_HOST_ASSERT(shared_l1_weights.size(0) == shared_intermediate_hidden * 2); + DG_HOST_ASSERT(shared_l1_weights.size(1) == hidden); + DG_HOST_ASSERT(shared_l2_weights.size(0) == hidden); + DG_HOST_ASSERT(shared_l1_weights.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(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); + DG_HOST_ASSERT(get_major_type_ab(shared_l2_weights) == cute::UMMA::Major::K); + check_sf_layout(shared_l1_weights_sf, shared_intermediate_hidden * 2, hidden, kGranMN, kGranK, + std::nullopt, false, true, torch::kFloat); + check_sf_layout(shared_l2_weights_sf, hidden, shared_intermediate_hidden, kGranMN, kGranK, + std::nullopt, false, true, torch::kFloat); + } + if (cumulative_local_expert_recv_stats.has_value()) { DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank); @@ -198,23 +265,30 @@ static void fp8_mega_moe( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - true, activation); + true, activation, num_shared_experts); DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); - const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf] = slice(sym_buffer); sm90_fp8_mega_moe(y, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, l1_weights, l2_weights, l1_weights_sf, l2_weights_sf, + // The shared L1 activations are the local `x` region and its SF + x, x_sf, + shared_l2_acts, shared_l2_acts_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_tokens, num_topk, hidden, intermediate_hidden, + num_shared_experts, activation_clamp, fast_math); if (get_env("DG_COMM_KERNEL_DEBUG")) diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index b66e7b842e..edf70ba836 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -23,6 +23,8 @@ struct MegaMoEConfig { // SF block sizes (UTCCP 128-aligned) int sf_block_m, sf_block_n; + // SF granularity along K + int gran_k; // Ring capacity and SF ring token count int num_ring_tokens; @@ -46,6 +48,7 @@ struct MegaMoEConfig { << ", load_block_m=" << config.load_block_m << ", load_block_n=" << config.load_block_n << ", store_block_m=" << config.store_block_m << ", sf_block_m=" << config.sf_block_m << ", sf_block_n=" << config.sf_block_n + << ", gran_k=" << config.gran_k << ", num_ring_tokens=" << config.num_ring_tokens << ", num_sf_ring_tokens=" << config.num_sf_ring_tokens << ", swizzle_acts_mode=" << config.swizzle_acts_mode << ", swizzle_weights_mode=" << config.swizzle_weights_mode @@ -61,31 +64,56 @@ struct MegaMoEConfig { static MmaKind parse_mma_kind(const std::string& mma_type_str) { if (mma_type_str == "bf16xbf16") return MmaKind::BF16; - DG_HOST_ASSERT(mma_type_str == "fp8xfp4"); - return MmaKind::MXFP8FP4; + if (mma_type_str == "fp8xfp4") + return MmaKind::MXFP8FP4; + if (mma_type_str == "mxf4xmxf4") + return MmaKind::MXFP4; + DG_HOST_ASSERT(mma_type_str == "nvfp4xnvfp4"); + return MmaKind::NVFP4; } -static int get_num_mma_elem_bytes(const MmaKind& mma_kind) { - return mma_kind == MmaKind::BF16 ? 2 : 1; +static std::string to_mma_type_string(const MmaKind& mma_kind) { + switch (mma_kind) { + case MmaKind::BF16: return "bf16xbf16"; + case MmaKind::MXFP8FP4: return "fp8xfp4"; + case MmaKind::MXFP4: return "mxf4xmxf4"; + case MmaKind::NVFP4: return "nvfp4xnvfp4"; + } + DG_HOST_UNREACHABLE("Unknown MMA kind"); +} + +static std::string to_mma_kind_name(const MmaKind& mma_kind) { + switch (mma_kind) { + case MmaKind::BF16: return "MmaKind::BF16"; + case MmaKind::MXFP8FP4: return "MmaKind::MXFP8FP4"; + case MmaKind::MXFP4: return "MmaKind::MXFP4"; + case MmaKind::NVFP4: return "MmaKind::NVFP4"; + } + DG_HOST_UNREACHABLE("Unknown MMA kind"); } static bool is_mma_with_sf(const MmaKind& mma_kind) { - return mma_kind == MmaKind::MXFP8FP4; + return mma_kind == MmaKind::MXFP8FP4 or mma_kind == MmaKind::MXFP4 or mma_kind == MmaKind::NVFP4; +} + +static bool is_mxf4_mma_kind(const MmaKind& mma_kind) { + return mma_kind == MmaKind::MXFP4; +} + +static bool is_nvfp4_mma_kind(const MmaKind& mma_kind) { + return mma_kind == MmaKind::NVFP4; } static std::tuple get_block_config_for_mega_moe( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, const int& num_tokens, - const MmaKind& mma_kind, - const bool& use_mxf4_kind = false) { + const MmaKind& mma_kind) { auto [cluster_size, block_m, store_block_m, block_k, num_epilogue_warpgroups] = [&]() -> std::tuple { float num_expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; if (num_expected_tokens_per_expert <= 8.5) { - // Really small token-per-expert (e.g. RL long-tail rollout), use larger BLOCK_K for less synchronization. - // Under kind::mxf4, bump block_m so the dense FP4 A/B smem tiles remain comfortably aligned. - return use_mxf4_kind ? std::tuple{2, 32, 16, 128, 2} - : std::tuple{2, 16, 8, 256, 2}; + // Really small token-per-expert (e.g. RL long-tail rollout), use the smallest block_m and larger BLOCK_K for less synchronization + return {2, 16, 8, 256, 2}; } else if (num_expected_tokens_per_expert <= 16.5) { // Small batch size, small EP, decoding, e.g. 6/384 experts, EP8, bsz 128 return {2, 32, 16, 128, 2}; @@ -103,14 +131,7 @@ static std::tuple get_block_config_for_mega_moe( return {2, 192, 32, 128, 2}; } }(); - block_k /= get_num_mma_elem_bytes(mma_kind); - if (mma_kind == MmaKind::MXFP8FP4 and not use_mxf4_kind) { - // K-major SM100 descriptors only support swizzles up to 128 bytes. - // The non-MXF4 FP8/FP4 path stores both operands as 1 byte per K - // element in smem, so a 256-element tile would require an illegal - // 256-byte swizzle. - block_k = std::min(block_k, 128); - } + block_k = block_k * 8 / get_element_bits(mma_kind); // Check whether our `block_m` lies in `kCandidateBlockM` DG_HOST_ASSERT(std::any_of( @@ -125,16 +146,15 @@ static std::tuple get_block_config_for_mega_moe( static std::pair get_pipeline_config_for_mega_moe( const int& smem_capacity, const int& num_experts, const int& hidden, - const int& block_m, const int& block_n, const int& block_k, + const int& block_m, const int& block_n, const int& block_k, const int& num_bytes_per_pull, const int& store_block_m, const int& sf_block_m, const int& sf_block_n, const int& gran_k, const int& num_dispatch_warps, const int& num_epilogue_warps, - const MmaKind& mma_kind, - const bool& use_mxf4_kind = false) { + const MmaKind& mma_kind) { constexpr int kSmemAlignment = 1024; constexpr int kNumEpilogueStages = 2; constexpr int kNumTMAStoreStages = 2; - const int num_mma_elem_bytes = get_num_mma_elem_bytes(mma_kind); + const int elem_bits = get_element_bits(mma_kind); // Always multicast on A const int load_block_m = block_m / 2; @@ -149,7 +169,7 @@ static std::pair get_pipeline_config_for_mega_moe( // C/D output region: max of L1 output staging and L2 BF16 staging. const auto num_epilogue_warpgroups = num_epilogue_warps / 4; - const int smem_cd_l1 = num_epilogue_warpgroups * store_block_m * (block_n / 2) * kNumTMAStoreStages * get_num_mma_elem_bytes(mma_kind); + const int smem_cd_l1 = num_epilogue_warpgroups * store_block_m * (block_n / 2) * kNumTMAStoreStages * elem_bits / 8; const int smem_cd_l2 = num_epilogue_warpgroups * store_block_m * block_n * static_cast(sizeof(nv_bfloat16)); const int smem_cd = align(std::max(smem_cd_l1, smem_cd_l2), kSmemAlignment); @@ -173,19 +193,14 @@ static std::pair get_pipeline_config_for_mega_moe( const int smem_sfb_per_stage = is_mma_with_sf(mma_kind) ? sf_block_n * (block_k / gran_k) : 0; // Per-stage: A tile + B tile + optional SF tiles + full/empty barriers. - // Dense FP4 kind halves both A and B byte footprints in shared memory. - const int smem_a_size_per_stage = use_mxf4_kind - ? (load_block_m * block_k / 2) - : (load_block_m * block_k * num_mma_elem_bytes); - const int smem_b_size_per_stage = use_mxf4_kind - ? (block_n * block_k / 2) - : (block_n * block_k * num_mma_elem_bytes); + // NOTES: for `MXFP8FP4` the FP4 weights are unpacked into 8-bit containers in smem, + // so A and B always share the same per-element footprint + const int smem_a_size_per_stage = load_block_m * block_k * elem_bits / 8; + const int smem_b_size_per_stage = block_n * block_k * elem_bits / 8; DG_HOST_ASSERT(smem_a_size_per_stage % kSmemAlignment == 0); DG_HOST_ASSERT(smem_b_size_per_stage % kSmemAlignment == 0); const int smem_stage_barriers = 2 * 8; - const int smem_size_per_stage = smem_a_size_per_stage + smem_b_size_per_stage + - smem_sfa_per_stage + smem_sfb_per_stage + - smem_stage_barriers; + const int smem_size_per_stage = smem_a_size_per_stage + smem_b_size_per_stage + smem_sfa_per_stage + smem_sfb_per_stage + smem_stage_barriers; // Fixed total const int smem_fixed = smem_dispatch_size + smem_cd + smem_amax_reduction + smem_barriers + @@ -204,26 +219,21 @@ static MegaMoEConfig get_mega_moe_config( const int& hidden, const int& intermediate_hidden, const int& num_ring_tokens, const int& num_sf_ring_tokens, - const MmaKind& mma_kind, - const bool& use_fp4_acts = false, - const bool& use_mxf4_kind = false) { + const MmaKind& mma_kind) { // Block config const auto [cluster_size, block_m, store_block_m, block_k, num_epilogue_threads] = - get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens, mma_kind, use_mxf4_kind); + get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens, mma_kind); const int block_n = 128; const int load_block_m = block_m / 2; const int load_block_n = block_n; const auto [sf_block_m, sf_block_n] = is_mma_with_sf(mma_kind) ? - SM100ArchSpec::get_sf_uttcp_aligned_block_sizes(block_m, block_n, MmaKind::MXFP8FP4) : std::pair(0, 0); - // NOTES: FP8 activations and FP4 weights (unpacked to 8-bit in smem) both use 128B swizzle + SM100ArchSpec::get_sf_uttcp_aligned_block_sizes(block_m, block_n, mma_kind) : std::pair(0, 0); + // NOTES: FP8 activations and FP4 weights (unpacked to 8-bit in smem) both use 128B swizzle; + // NVFP4 keeps FP4 packed, so a 128B swizzle atom covers twice as many elements const int swizzle_acts_mode = 128; const int swizzle_weights_mode = 128; - const int gran_k = 32; - const int num_max_pool_tokens = layout::get_num_max_pool_tokens( - num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); - const bool use_full_pool_fp8_fp4_path = - mma_kind == MmaKind::MXFP8FP4 and num_ring_tokens >= num_max_pool_tokens; + const int gran_k = is_mma_with_sf(mma_kind) ? get_sf_gran_k(mma_kind) : 32; // Thread layout const int num_dispatch_threads = 128; @@ -231,10 +241,8 @@ static MegaMoEConfig get_mega_moe_config( // Pull: divide token bytes by 2 until <= kPullThreshold constexpr int kPullThreshold = 4096; - int num_bytes_per_pull = use_full_pool_fp8_fp4_path ? - hidden * get_num_mma_elem_bytes(mma_kind) : - (use_fp4_acts ? (hidden / 2) : hidden * get_num_mma_elem_bytes(mma_kind)); - while (not use_full_pool_fp8_fp4_path and num_bytes_per_pull > kPullThreshold) { + int num_bytes_per_pull = hidden * get_element_bits(mma_kind) / 8; + while (num_bytes_per_pull > kPullThreshold) { DG_HOST_ASSERT(num_bytes_per_pull % 2 == 0); num_bytes_per_pull /= 2; } @@ -246,12 +254,12 @@ static MegaMoEConfig get_mega_moe_config( block_m, block_n, block_k, num_bytes_per_pull, store_block_m, sf_block_m, sf_block_n, gran_k, num_dispatch_threads / 32, num_epilogue_threads / 32, - mma_kind, use_mxf4_kind); + mma_kind); const auto config = MegaMoEConfig { block_m, block_n, block_k, load_block_m, load_block_n, store_block_m, - sf_block_m, sf_block_n, + sf_block_m, sf_block_n, gran_k, num_ring_tokens, is_mma_with_sf(mma_kind) ? num_sf_ring_tokens : 0, swizzle_acts_mode, swizzle_weights_mode, num_stages, smem_size, diff --git a/csrc/jit_kernels/heuristics/sm100.hpp b/csrc/jit_kernels/heuristics/sm100.hpp index c8e9e2e07f..35b0e8fc6f 100644 --- a/csrc/jit_kernels/heuristics/sm100.hpp +++ b/csrc/jit_kernels/heuristics/sm100.hpp @@ -20,6 +20,8 @@ struct SM100ArchSpec { switch (mma_kind) { case MmaKind::BF16: return {0, 0}; case MmaKind::MXFP8FP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; + case MmaKind::MXFP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; + case MmaKind::NVFP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; default: DG_HOST_UNREACHABLE("Unknown dtype"); } } @@ -47,7 +49,7 @@ struct SM100ArchSpec { for (int swap_ab = 0; swap_ab < 2; ++ swap_ab) { // Block M/N candidates std::vector block_m_candidates; - std::vector block_n_candidates; + std::vector block_n_candidates; if (swap_ab) { int step = std::lcm(16, heuristics_runtime->get_block_m_multiple_of()); int end = 256; @@ -122,7 +124,9 @@ struct SM100ArchSpec { // Check tensor memory capacity const auto [sf_block_m, sf_block_n] = get_sf_uttcp_aligned_block_sizes(block_m, block_n, desc.get_mma_kind()); - const auto tmem_sf_cols = desc.get_mma_kind() == MmaKind::MXFP8FP4 ? sf_block_m / 32 + sf_block_n / 32 : 0; + const auto sf_granularity = get_sf_gran_k(desc.get_mma_kind()); + const auto tmem_sf_cols = sf_granularity ? + sf_block_m / sf_granularity + sf_block_n / sf_granularity : sf_granularity; const auto umma_n = swap_ab ? block_m : block_n; if (2 * umma_n + tmem_sf_cols > 512) continue; diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp index b149644791..d16862a906 100644 --- a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -22,7 +22,6 @@ struct MegaMoESM90Config { int num_max_pool_tokens; int num_padded_sf_pool_tokens; int swizzle_acts_mode, swizzle_weights_mode; - int num_experts_per_wave; int num_stages, smem_size; int num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads; @@ -33,7 +32,6 @@ struct MegaMoESM90Config { << ", num_max_pool_tokens=" << config.num_max_pool_tokens << ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens << ", swizzle_acts_mode=" << config.swizzle_acts_mode << ", swizzle_weights_mode=" << config.swizzle_weights_mode - << ", num_experts_per_wave=" << config.num_experts_per_wave << ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size << ", num_dispatch_threads=" << config.num_dispatch_threads << ", num_non_epilogue_threads=" << config.num_non_epilogue_threads @@ -44,10 +42,32 @@ struct MegaMoESM90Config { static std::tuple get_block_config_for_mega_moe_sm90( const int& num_ranks, const int& num_experts, - const int& num_topk, const int& num_tokens) { + const int& num_topk, const int& num_tokens, const int &intermediate_hidden) { const float expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; - const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; + // The relaxed 2-WG threshold enables the block_m=128 / 4-WG path only + // above a higher tokens/expert bar (instead of the original >= 64), + // trading two extra warpgroups for fewer register spills. On H20 the + // smaller SM count (78 vs 132 on H100/H200) makes the extra warpgroups + // costly, so the relaxation applies in two intermediate_hidden regimes: + // * pro (>= 3072): 4-WG only when expected_tokens_per_expert > 512 + // * flash (<= 2048): 4-WG only when expected_tokens_per_expert > 576, + // because 2-WG + BLOCK_N=256 outperforms 4-WG in part of the flash + // batch range -- 4-WG is reserved for the heaviest flash batches. + // On H200/H100 the larger SM count makes the extra warpgroups always win, + // so the original 4-WG-first (>= 64) threshold is kept for every shape, + // as well as for the H20 mid-range (2048 < intermediate_hidden < 3072). + const int num_sms = device_runtime->get_num_sms(); + const bool is_h20 = num_sms <= 84; + const bool apply_h20_pro_relaxation = is_h20 and intermediate_hidden >= 3072; + const bool apply_h20_flash_relaxation = is_h20 and intermediate_hidden <= 2048; + bool auto_split_mn; + if (apply_h20_pro_relaxation) + auto_split_mn = expected_tokens_per_expert > 512.0f; + else if (apply_h20_flash_relaxation) + auto_split_mn = expected_tokens_per_expert > 576.0f; + else + auto_split_mn = expected_tokens_per_expert >= 64.0f; if (auto_split_mn) return {128, 512}; @@ -61,99 +81,14 @@ static std::tuple get_block_config_for_mega_moe_sm90( return {block_m, num_epilogue_warpgroups * 128}; } -// SM90 retains the original wave scheduler and its ring-capacity heuristic. -// Keep these helpers local to the Hopper path: upstream's SM100 scheduler now -// sizes live task pools directly and no longer exposes the legacy helpers. -static int get_num_wave_pool_tokens_for_mega_moe_sm90( - const int& num_ranks, const int& num_topk, const int& num_max_tokens_per_rank, - const int& num_experts_per_wave, const int& block_m) { - DG_HOST_ASSERT(num_max_tokens_per_rank % block_m == 0); - const auto num_tokens_from_all_ranks = num_max_tokens_per_rank * num_ranks; - if (num_experts_per_wave == 1) - return num_tokens_from_all_ranks; - - return std::min( - num_tokens_from_all_ranks * num_experts_per_wave, - math::align( - num_tokens_from_all_ranks * num_topk + num_experts_per_wave * (block_m - 1), - block_m)); -} - -static int get_num_experts_per_wave_for_mega_moe_sm90_legacy( - const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, - const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms, - const int& num_ring_tokens, const int& num_max_tokens_per_rank, const int& num_ranks) { - int num_max_experts_per_wave = num_experts_per_rank; - while (num_max_experts_per_wave > 0 and - get_num_wave_pool_tokens_for_mega_moe_sm90( - num_ranks, num_topk, num_max_tokens_per_rank, - num_max_experts_per_wave, block_m) > num_ring_tokens) - --num_max_experts_per_wave; - DG_HOST_ASSERT(num_max_experts_per_wave > 0 and "Buffer size is too small"); - - constexpr int kImbalanceFactor = 2; - const float num_expected_tokens_per_expert = - static_cast(num_tokens * num_topk) / num_experts_per_rank; - const int num_expected_m_blocks = std::max( - ceil_div(static_cast(std::ceil(num_expected_tokens_per_expert)), block_m), 1); - const int num_l1_n_blocks = (2 * intermediate_hidden) / block_n; - const int num_expected_l1_blocks_per_expert = num_expected_m_blocks * num_l1_n_blocks; - int num_min_expected_experts_to_fill_sms = - ceil_div(kImbalanceFactor * num_sms, num_expected_l1_blocks_per_expert); - - if (num_expected_tokens_per_expert < 1) - num_min_expected_experts_to_fill_sms = num_experts_per_rank; - if (num_min_expected_experts_to_fill_sms >= num_max_experts_per_wave) - return num_max_experts_per_wave; - if (num_expected_l1_blocks_per_expert >= num_sms) - return num_min_expected_experts_to_fill_sms; - - const int num_sweep_max_experts_per_wave = std::min( - num_max_experts_per_wave, num_min_expected_experts_to_fill_sms * 2); - int best_num_experts_per_wave = num_min_expected_experts_to_fill_sms; - float best_tail_ratio = -1.0f; - for (int num_experts_per_wave = num_min_expected_experts_to_fill_sms; - num_experts_per_wave <= num_sweep_max_experts_per_wave; - ++num_experts_per_wave) { - const int remainder = num_experts_per_rank % num_experts_per_wave; - const float tail_ratio = remainder == 0 ? - 1.0f : static_cast(remainder) / num_experts_per_wave; - if (tail_ratio > best_tail_ratio) { - best_tail_ratio = tail_ratio; - best_num_experts_per_wave = num_experts_per_wave; - } - } - return best_num_experts_per_wave; -} - -static int get_num_experts_per_wave_for_mega_moe_sm90( - const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, - const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms, - const int& num_ring_tokens, const int& num_max_tokens_per_rank, const int& num_ranks) { - const float expected_tokens_per_expert = - static_cast(num_tokens) * num_topk / num_experts_per_rank; - if (expected_tokens_per_expert < 1.0f or expected_tokens_per_expert > 4.0f) - return num_experts_per_rank; - - if (block_m == 64 and intermediate_hidden >= 3072) { - const int num_n_blocks_per_expert = (2 * intermediate_hidden) / block_n; - const int single_wave_blocks = - num_experts_per_rank * num_n_blocks_per_expert; - if (single_wave_blocks >= 4 * num_sms) - return num_experts_per_rank; - } - return get_num_experts_per_wave_for_mega_moe_sm90_legacy( - num_experts_per_rank, num_tokens, num_topk, - intermediate_hidden, block_m, block_n, num_sms, - num_ring_tokens, num_max_tokens_per_rank, num_ranks); -} - static bool should_use_swap_ab_for_mega_moe_sm90( const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, const int& block_m, const int& num_epilogue_threads) { // swapAB is ENABLED by default (the L1 SF-pool stride bug that corrupted // pool blocks >= 1 was fixed: BLOCK_M -> SF_BLOCK_M in the swapAB L1 epilogue). // Kill-switch retained: set DG_SM90_FP8_SWAP_AB=0 to force the non-swap path. + // swapAB composes with the fused shared expert (shared's token-axis output + // matches swapAB's reduced M-axis), so no special-casing is needed here. if (get_env("DG_SM90_FP8_SWAP_AB", 1) == 0) return false; const float expected_tokens_per_expert = @@ -193,15 +128,37 @@ static std::pair get_pipeline_config_for_mega_moe_sm90( const int smem_per_stage = block_m * block_k + block_n * block_k + smem_sfa_per_stage + smem_sfb_per_stage; - const int smem_barriers_fixed = (num_dispatch_warps + 2 * num_epilogue_warps) * 8; + // The scheduler adds 2 task-info full/empty barrier pairs and two 32-byte + // task-info slots (see `sm90_fp8_mega_moe.cuh`). The scheduler `TaskInfo` is + // alignas(16)/32 B while a barrier slot is only 8 B, so the kernel pads one + // extra barrier when the preceding barrier count is odd + // (`kTaskInfoBarrierPad = kTaskInfoBaseBarriers & 1u`). Only + // `num_dispatch_warps` affects that parity (2*num_stages, 2*num_epilogue_warps + // and the 4 task-info barriers are all even), so mirror the same pad here. + const int smem_task_info_barriers = 4; // 2 full + 2 empty + const int smem_task_info_pad = (num_dispatch_warps & 1) * 8; + const int smem_barriers_fixed = + (num_dispatch_warps + 2 * num_epilogue_warps + smem_task_info_barriers) * 8 + + smem_task_info_pad; + const int smem_task_infos = 2 * 32; const int smem_barriers_per_stage = 2 * 8; - const int smem_fixed = smem_dispatch_size + smem_cd + smem_barriers_fixed; + const int smem_fixed = smem_dispatch_size + smem_cd + smem_barriers_fixed + smem_task_infos; const int num_stages = (smem_capacity - smem_fixed) / (smem_per_stage + smem_barriers_per_stage); DG_HOST_ASSERT(num_stages >= 2); const int smem_size = smem_fixed + num_stages * (smem_per_stage + smem_barriers_per_stage); DG_HOST_ASSERT(smem_size <= smem_capacity); + + // Cross-check against the kernel's exact barrier/task-info layout: the + // task-info ring (including the alignment pad) must end inside the + // allocated dynamic shared memory. + const int smem_task_info_end = + smem_dispatch_size + smem_cd + num_stages * smem_per_stage + + (num_dispatch_warps + 2 * num_stages + 2 * num_epilogue_warps + + smem_task_info_barriers + smem_task_info_pad / 8) * 8 + + smem_task_infos; + DG_HOST_ASSERT(smem_task_info_end <= smem_size); return {num_stages, smem_size}; } @@ -211,7 +168,7 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const int& hidden, const int& intermediate_hidden, const int& num_padded_sf_pool_tokens) { const auto [block_m, num_epilogue_threads] = get_block_config_for_mega_moe_sm90( - num_ranks, num_experts, num_topk, num_tokens); + num_ranks, num_experts, num_topk, num_tokens, intermediate_hidden); const float expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; const bool auto_split_mn = @@ -219,13 +176,13 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const bool decode_split_n_path = block_m == 64 and num_epilogue_threads == 256; const bool decode_use_block_n_256 = - decode_split_n_path and intermediate_hidden >= 3072 and + decode_split_n_path and expected_tokens_per_expert >= 0.25f and (2 * intermediate_hidden) % 256 == 0 and hidden % 256 == 0; const bool use_swap_ab = should_use_swap_ab_for_mega_moe_sm90( num_experts_per_rank, num_tokens, num_topk, block_m, num_epilogue_threads); - int block_n = use_swap_ab ? 128 + int block_n = use_swap_ab ? 256 : (auto_split_mn ? 256 : (decode_use_block_n_256 ? 256 : 128)); const int block_k = 128; @@ -235,22 +192,18 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const int swizzle_acts_mode = 128; const int swizzle_weights_mode = 128; - const int num_sms = device_runtime->get_num_sms(); - const int num_experts_per_wave = get_num_experts_per_wave_for_mega_moe_sm90( - num_experts_per_rank, num_tokens, num_topk, - intermediate_hidden, block_m, block_n, num_sms, - num_max_pool_tokens, num_max_tokens_per_rank, num_ranks); - - const bool reduce_decode_threads = num_epilogue_threads == 128; - const bool decode_split_n = - block_m == 64 and num_epilogue_threads == 256; - const bool shrink_non_epilogue = reduce_decode_threads or decode_split_n; - const int num_dispatch_threads = - (num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128; - const bool split_sfa_loader_warp = false; - const int num_non_epilogue_threads = - split_sfa_loader_warp ? 128 : - ((num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128); + // The scheduler needs a dedicated producer warp, so the non-epilogue section + // is exactly 3 warps (TMA-A, TMA-B, producer) and dispatch is a single warp: + // dispatch + non-epilogue = 32 + 96 = 128, a whole warpgroup that keeps the + // math warpgroups 128-thread aligned. This is the minimal aligned topology for + // every epilogue width: + // * 2-WG (epilogue=256): 32 + 96 + 256 = 384 threads, ceiling + // 65536/384 = 170 >= 168, so the epilogue accumulators do not spill. + // * 4-WG (epilogue=512): 32 + 96 + 512 = 640 threads. Halving dispatch to one + // warp is the cost of fitting the producer warp under 128-thread alignment; + // a 2-dispatch-warp variant would pad to 768 threads and spill worse. + const int num_dispatch_threads = 32; + const int num_non_epilogue_threads = 96; DG_HOST_ASSERT((num_dispatch_threads + num_non_epilogue_threads) % 128 == 0); const auto [num_stages, smem_size] = get_pipeline_config_for_mega_moe_sm90( @@ -265,7 +218,6 @@ static MegaMoESM90Config get_mega_moe_config_sm90( cluster_size, num_max_pool_tokens, num_padded_sf_pool_tokens, swizzle_acts_mode, swizzle_weights_mode, - num_experts_per_wave, num_stages, smem_size, num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads }; diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index e6d4dfb8fe..34962feb57 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -83,6 +83,7 @@ static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType& case torch::kFloat: return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; case torch::kBFloat16: return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; case torch::kFloat8_e4m3fn: return CU_TENSOR_MAP_DATA_TYPE_UINT8; + case torch::kUInt8: return CU_TENSOR_MAP_DATA_TYPE_UINT8; #if CUDA_VERSION >= 12080 case kPackedFP4: return fp4_unpacked_smem ? CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B : CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index 3e031f2197..1b14558121 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" @@ -23,26 +23,24 @@ class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime sym_buffer_ptrs; @@ -89,13 +87,13 @@ static void __instantiate_kernel() {{ {}, {}, {}, + {}, {}, {}, {}, {}, {}, + {}, {}, {}, {}, - {}, - {}, - {} + {}, {}, {}, {}, {} >); }}; )", args.num_max_tokens_per_rank, @@ -105,17 +103,20 @@ static void __instantiate_kernel() {{ args.config.block_m, args.config.block_n, args.config.block_k, args.config.store_block_m, args.config.sf_block_m, args.config.sf_block_n, + to_mma_kind_name(args.mma_kind), args.config.num_ring_tokens, args.config.num_sf_ring_tokens, args.config.num_stages, args.config.num_bytes_per_pull, args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads, args.launch_args.grid_dim.first, args.num_ranks, - to_string(args.activation_clamp), + to_string(args.activation_clamp), to_string(args.swiglu_alpha), args.use_situ ? "true" : "false", args.fast_math ? "true" : "false", - args.use_fp4_acts ? "true" : "false", - args.use_mxf4_kind ? "true" : "false", + args.use_x_scales ? "true" : "false", + args.with_l1_alphas ? "true" : "false", + args.with_l2_alphas ? "true" : "false", + args.with_l2_act_scales ? "true" : "false", args.use_fp8_combine ? "true" : "false"); } @@ -124,6 +125,9 @@ static void __instantiate_kernel() {{ DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, args.y, args.cumulative_local_expert_recv_stats, + args.l1_alphas, + args.l2_alphas, + args.l2_act_scales, args.num_tokens, args.sym_buffer_ptrs, args.tensor_map_l1_acts, @@ -166,151 +170,82 @@ static void sm100_fp8_fp4_mega_moe( const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, const float& activation_clamp, + const float& swiglu_alpha, const bool& use_situ, const bool& fast_math, - const bool& use_fp4_acts = false, - const bool& use_mxf4_kind = false, - const bool& use_fp8_combine = false + const bool& use_x_scales, + const float* l1_alphas, + const float* l2_alphas, + const float* l2_act_scales, + const MmaKind& mma_kind, + const bool& use_fp8_combine ) { 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; - // Stream A0.5 sanity: kind::mxf4 only accepts FP4 inputs. - DG_HOST_ASSERT(not use_mxf4_kind or use_fp4_acts); - DG_HOST_ASSERT(not use_fp4_acts or num_shared_experts == 0); // Heuristics const auto config = get_mega_moe_config( num_ranks, num_experts, num_experts_per_rank, num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_ring_tokens, num_sf_ring_tokens, - MmaKind::MXFP8FP4, use_fp4_acts, use_mxf4_kind); + mma_kind); // Make tensormap - constexpr int kGranK = 32; + const bool is_packed_fp4 = mma_kind == MmaKind::NVFP4 or mma_kind == MmaKind::MXFP4; + const int kGranK = config.gran_k; + const int elem_bits = get_element_bits(mma_kind); + const auto to_inner = [=](const int& num_elems) { return num_elems * elem_bits / 8; }; + const auto weight_tensor = [=](const torch::Tensor& t) { return is_packed_fp4 ? t.view(torch::kUInt8) : t; }; + const auto weight_inner = [=](const int& num_elems) { return is_packed_fp4 ? num_elems / 2 : num_elems; }; + const int block_k_inner = to_inner(config.block_k); const int sf_smem_outer_dim = config.block_k / (kGranK * 4); - // Stream A0.5: when `use_mxf4_kind` is on, BOTH L1 and L2 acts AND - // weights TMA descriptors switch from `_ALIGN16B` (FP4 with-padding, - // 8 data + 8 pad bytes per 16-byte atom) to `_ALIGN8B` (dense FP4, - // 2 nibbles/byte). The smem byte stride per K-row halves accordingly, - // and swizzle mode halves to match (128B → 64B). The gmem layout is - // unchanged — the underlying `l1_acts` / `l1_weights` storage is still - // packed FP4 nibbles; only how TMA expands them into smem changes. - const bool fp4_unpacked = not use_mxf4_kind; - const int swizzle_acts = use_mxf4_kind ? config.swizzle_acts_mode / 2 - : config.swizzle_acts_mode; - const int swizzle_weights = use_mxf4_kind ? config.swizzle_weights_mode / 2 - : config.swizzle_weights_mode; - // Stream A0.0b: when `use_fp4_acts` is on, the L1 token pool buffer - // (`l1_acts`) is already viewed as `kPackedFP4` (int8) by the symm-buffer - // slice (see `csrc/apis/mega.hpp`), with shape `[num_pool_tokens, hidden/2]` - // of packed E2M1 (low nibble = even col, high nibble = odd col). - // `make_tma_2d_desc` then auto-selects `CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B` - // via `aten_dtype_to_tensor_map_dtype` (runtime_utils.hpp:84-87) — or - // `_ALIGN8B` under `use_mxf4_kind` (Stream A0.5). - // - // TMA descriptor: `gmem_inner_dim = hidden` U4 elements (the descriptor - // reads `hidden/2` storage bytes per row); smem inner box has `config.block_k` - // elements and swizzle follows the selected dense/with-padding FP4 mode. const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, - hidden, config.num_ring_tokens, - config.block_k, config.load_block_m, + to_inner(hidden), config.num_ring_tokens, + block_k_inner, config.load_block_m, static_cast(l1_acts.stride(-2)), - swizzle_acts, /*swizzle_base=*/0, - /*allow_tf32=*/false, - /*fp4_unpacked_smem=*/fp4_unpacked); + config.swizzle_acts_mode); const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, config.num_sf_ring_tokens, hidden, config.sf_block_m, kGranK, 1, 0, 0, false, sf_smem_outer_dim); - const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights, - hidden, num_experts_per_rank * intermediate_hidden * 2, - config.block_k, config.load_block_n, + const auto tensor_map_l1_weights = make_tma_2d_desc(weight_tensor(l1_weights), + weight_inner(hidden), num_experts_per_rank * intermediate_hidden * 2, + block_k_inner, config.load_block_n, static_cast(l1_weights.stride(-2)), - swizzle_weights, /*swizzle_base=*/0, - /*allow_tf32=*/false, - /*fp4_unpacked_smem=*/fp4_unpacked); + config.swizzle_weights_mode, 0, false, not is_packed_fp4); const auto tensor_map_l1_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_weights_sf, intermediate_hidden * 2, hidden, config.block_n, kGranK, num_experts_per_rank, 0, 0, false, - sf_smem_outer_dim); + sf_smem_outer_dim); // NOTES: L1 output and L2 activations are essentially the same tensor. // Post-SwiGLU output has half the N width (`BLOCK_N / 2` per input tile), // so the swizzle mode is also halved (128 -> 64). - // - // Stream A0.2: when `use_fp4_acts` is on, the L1 epilogue emits packed - // E2M1 (FP4) where each byte holds 2 elements. The kernel writes a - // **dense canonical** smem layout (no swizzle XOR) — see the FP4 store - // branch in `sm100_fp8_fp4_mega_moe.cuh`. To match, we build the L1 - // output TMA descriptor with `swizzle = 0`. The gmem result is the - // canonical `[M, intermediate_hidden / 2]` packed FP4 layout, byte- - // identical to what `kernels/fused_gemm_swiglu_fp4_quant_1cta` produces - // (Stream A2). The L2 reader (built below) consumes this same canonical - // layout via `_ALIGN16B`. The per-row gmem byte footprint halves - // (`intermediate_hidden / 2` bytes vs `intermediate_hidden` for FP8); - // outer stride in the underlying buffer is unchanged. - const auto tensor_map_l1_output = use_fp4_acts - ? make_tma_2d_desc(l2_acts, - intermediate_hidden / 2, config.num_ring_tokens, - config.block_n / 4, config.store_block_m, - static_cast(l2_acts.stride(-2)), - /*swizzle_mode=*/0) - : make_tma_2d_desc(l2_acts, - intermediate_hidden, config.num_ring_tokens, - config.block_n / 2, config.store_block_m, - static_cast(l2_acts.stride(-2)), - config.swizzle_acts_mode / 2); - // Stream A0.2: when FP4 acts on, L2 reads packed E2M1 via `_ALIGN16B`. - // `make_tma_2d_desc` selects the descriptor dtype from the source - // tensor's `scalar_type`; `l2_acts` is allocated as FP8 (1 byte/elem). - // For the FP4 path we re-view the same byte buffer as `kPackedFP4` so - // the descriptor dtype is `CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B`. - // - // gmem layout (FP4 path, set up by L1 epilogue): - // - per row: first `intermediate_hidden / 2` bytes are packed E2M1 - // (low nibble = even col, high nibble = odd col — canonical MXFP4), - // remaining bytes in the row are stale FP8 from prior runs. - // - row stride: `l2_acts.stride(-2)` source bytes (= same as FP8 - // because the buffer view's underlying allocation hasn't changed). - // - // TMA descriptor tells the hardware: - // - `gmem_inner_dim = intermediate_hidden` U4 elements (= - // `intermediate_hidden / 2` source bytes are read per row). - // - `gmem_outer_stride = stride(-2)` source bytes (the actual storage - // row pitch — leaves the unused tail of each FP8-sized row alone). - // - smem inner box = `BLOCK_K = 128` elements (= 64 source bytes per - // row, expands to 128 smem bytes after `_ALIGN16B` doubling); 128B - // swizzle aligns with the per-stage atom (same as B-side, which has - // used this layout for FP4 weights from day one). - const auto tensor_map_l2_acts = use_fp4_acts - ? make_tma_2d_desc(l2_acts.view(kPackedFP4), - intermediate_hidden, config.num_ring_tokens, - config.block_k, config.load_block_m, - static_cast(l2_acts.stride(-2)), - swizzle_acts, /*swizzle_base=*/0, - /*allow_tf32=*/false, - /*fp4_unpacked_smem=*/fp4_unpacked) - : make_tma_2d_desc(l2_acts, - intermediate_hidden, config.num_ring_tokens, - config.block_k, config.load_block_m, - static_cast(l2_acts.stride(-2)), - config.swizzle_acts_mode); + const int l1_out_block_n_bytes = to_inner(config.block_n / 2); + const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, + to_inner(intermediate_hidden), config.num_ring_tokens, + l1_out_block_n_bytes, config.store_block_m, + static_cast(l2_acts.stride(-2)), + l1_out_block_n_bytes); + const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts, + to_inner(intermediate_hidden), config.num_ring_tokens, + block_k_inner, config.load_block_m, + static_cast(l2_acts.stride(-2)), + config.swizzle_acts_mode); const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, config.num_sf_ring_tokens, intermediate_hidden, config.sf_block_m, kGranK, 1, 0, 0, false, sf_smem_outer_dim); - const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights, - intermediate_hidden, num_experts_per_rank * hidden, - config.block_k, config.load_block_n, + const auto tensor_map_l2_weights = make_tma_2d_desc(weight_tensor(l2_weights), + weight_inner(intermediate_hidden), num_experts_per_rank * hidden, + block_k_inner, config.load_block_n, static_cast(l2_weights.stride(-2)), - swizzle_weights, /*swizzle_base=*/0, - /*allow_tf32=*/false, - /*fp4_unpacked_smem=*/fp4_unpacked); + config.swizzle_weights_mode, 0, false, not is_packed_fp4); const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf, hidden, intermediate_hidden, config.block_n, kGranK, @@ -319,8 +254,8 @@ static void sm100_fp8_fp4_mega_moe( const auto tensor_map_shared_l1_acts = num_shared_experts > 0 ? make_tma_2d_desc( shared_l1_acts, - hidden, num_max_tokens_per_rank, - config.block_k, config.load_block_m, + to_inner(hidden), num_max_tokens_per_rank, + block_k_inner, config.load_block_m, static_cast(shared_l1_acts.stride(-2)), config.swizzle_acts_mode) : tensor_map_l1_acts; const auto tensor_map_shared_l1_acts_sf = num_shared_experts > 0 ? make_tma_sf_desc( @@ -330,9 +265,9 @@ static void sm100_fp8_fp4_mega_moe( 1, 0, 0, false, sf_smem_outer_dim) : 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, - config.block_k, config.load_block_n, + weight_tensor(shared_l1_weights), + to_inner(hidden), shared_intermediate_hidden * 2, + block_k_inner, config.load_block_n, static_cast(shared_l1_weights.stride(-2)), config.swizzle_weights_mode) : tensor_map_l1_weights; const auto tensor_map_shared_l1_weights_sf = num_shared_experts > 0 ? make_tma_sf_desc( @@ -343,14 +278,14 @@ static void sm100_fp8_fp4_mega_moe( sf_smem_outer_dim) : tensor_map_l1_weights_sf; 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, - config.block_n / 2, config.store_block_m, + to_inner(shared_intermediate_hidden), num_max_tokens_per_rank, + l1_out_block_n_bytes, config.store_block_m, static_cast(shared_l2_acts.stride(-2)), - config.swizzle_acts_mode / 2) : tensor_map_l1_output; + l1_out_block_n_bytes) : 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, - config.block_k, config.load_block_m, + to_inner(shared_intermediate_hidden), num_max_tokens_per_rank, + block_k_inner, config.load_block_m, static_cast(shared_l2_acts.stride(-2)), config.swizzle_acts_mode) : tensor_map_l2_acts; const auto tensor_map_shared_l2_acts_sf = num_shared_experts > 0 ? make_tma_sf_desc( @@ -360,9 +295,9 @@ static void sm100_fp8_fp4_mega_moe( 1, 0, 0, false, sf_smem_outer_dim) : 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, - config.block_k, config.load_block_n, + weight_tensor(shared_l2_weights), + to_inner(shared_intermediate_hidden), hidden, + block_k_inner, config.load_block_n, static_cast(shared_l2_weights.stride(-2)), config.swizzle_weights_mode) : tensor_map_l2_weights; const auto tensor_map_shared_l2_weights_sf = num_shared_experts > 0 ? make_tma_sf_desc( @@ -385,15 +320,22 @@ static void sm100_fp8_fp4_mega_moe( .num_experts = num_experts, .num_shared_experts = num_shared_experts, .num_topk = num_topk, .num_ranks = num_ranks, + .mma_kind = mma_kind, .activation_clamp = activation_clamp, + .swiglu_alpha = swiglu_alpha, .use_situ = use_situ, .fast_math = fast_math, - .use_fp4_acts = use_fp4_acts, - .use_mxf4_kind = use_mxf4_kind, + .use_x_scales = use_x_scales, + .with_l1_alphas = l1_alphas != nullptr, + .with_l2_alphas = l2_alphas != nullptr, + .with_l2_act_scales = l2_act_scales != nullptr, .use_fp8_combine = use_fp8_combine, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, + .l1_alphas = l1_alphas, + .l2_alphas = l2_alphas, + .l2_act_scales = l2_act_scales, .num_tokens = num_tokens, .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), .tensor_map_l1_acts = tensor_map_l1_acts, diff --git a/csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp b/csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp index 9d6c347401..73fd889c08 100644 --- a/csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp +++ b/csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp @@ -2,9 +2,12 @@ #include +#include + #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" #include "../../jit/kernel_runtime.hpp" +#include "../heuristics/mega_moe.hpp" #include "../../utils/exception.hpp" #include "../../utils/format.hpp" #include "../../utils/math.hpp" @@ -13,21 +16,23 @@ namespace deep_gemm { // JIT runtime for `sm100_mega_moe_pre_dispatch` (see // `deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh`). -// Templated on (kGroupSize, kUseFp4Acts, kUsePDL); host fn picks the +// Templated on (kGroupSize, kMmaKind, kUsePDL); host fn picks the // instantiation from explicit args. class SM100MegaMoEPreDispatchRuntime final : public LaunchRuntime { public: struct Args { int group_size; - bool use_fp4_acts; + MmaKind mma_kind; bool use_pdl; // Runtime args (passed to the kernel via the params struct). const void* x; const void* topk_idx; const void* topk_weights; + const void* expert_scales; void* buf_x; void* buf_x_sf; + void* buf_x_scales; void* buf_topk_idx; void* buf_topk_weights; uint32_t num_tokens; @@ -51,14 +56,15 @@ static void __instantiate_kernel() {{ >); }}; )", args.group_size, - args.use_fp4_acts ? "true" : "false", + to_mma_kind_name(args.mma_kind), args.use_pdl ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.x, args.topk_idx, args.topk_weights, - args.buf_x, args.buf_x_sf, args.buf_topk_idx, args.buf_topk_weights, + args.x, args.topk_idx, args.topk_weights, args.expert_scales, + args.buf_x, args.buf_x_sf, args.buf_x_scales, + args.buf_topk_idx, args.buf_topk_weights, args.num_tokens, args.padded_max, args.hidden, args.num_groups, args.top_k)); } }; @@ -67,7 +73,7 @@ static void __instantiate_kernel() {{ // - x: (M, H) bf16, contiguous. // - topk_idx: (M, K) int32, contiguous. // - topk_weights: (M, K) float, contiguous. -// - buf_x: (P, H) fp8_e4m3 if !use_fp4_acts, else (P, H/2) int8 (packed FP4). +// - buf_x: (P, H) fp8_e4m3 for `fp8xfp4`, else (P, H/2) int8 (packed FP4). // - buf_x_sf: (P, G/4) int32, contiguous; G = H / group_size; each int32 // stores 4 UE8M0 bytes row-major. // - buf_topk_idx: (P, K) int64. @@ -87,8 +93,29 @@ static void mega_moe_pre_dispatch( const torch::Tensor& buf_topk_weights, const int& num_tokens, const int& group_size, - const bool& use_fp4_acts) { - DG_HOST_ASSERT(group_size == 32 || group_size == 64 || group_size == 128); + const std::string& mma_type, + const std::optional& buf_x_scales, + const std::optional& expert_scales) { + const auto mma_kind = parse_mma_kind(mma_type); + DG_HOST_ASSERT(mma_kind != MmaKind::BF16); + const bool is_packed_fp4 = mma_kind == MmaKind::MXFP4 or mma_kind == MmaKind::NVFP4; + if (mma_kind == MmaKind::NVFP4) { + // NVFP4 per-token acts: 16-elem UE4M3 block SFs plus one FP32 outer + // scale per token (written to `buf_x_scales`). + DG_HOST_ASSERT(group_size == 16); + DG_HOST_ASSERT(buf_x_scales.has_value()); + DG_HOST_ASSERT(buf_x_scales->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(buf_x_scales->is_contiguous()); + DG_HOST_ASSERT(buf_x_scales->numel() >= num_tokens); + } else { + DG_HOST_ASSERT(group_size == 32 || group_size == 64 || group_size == 128); + DG_HOST_ASSERT(!buf_x_scales.has_value()); + } + if (expert_scales.has_value()) { + DG_HOST_ASSERT(expert_scales->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(expert_scales->is_contiguous()); + DG_HOST_ASSERT(expert_scales->dim() == 1); + } DG_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(x.is_contiguous()); DG_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt32); @@ -123,7 +150,7 @@ static void mega_moe_pre_dispatch( DG_HOST_ASSERT(static_cast(buf_x_sf.size(0)) == padded_max); DG_HOST_ASSERT(static_cast(buf_x_sf.size(1)) == num_groups / 4); - if (use_fp4_acts) { + if (is_packed_fp4) { // Packed FP4: (P, hidden/2) bytes. The symm-buffer slice views this // as kPackedFP4 (int8); accept either int8 / uint8 / float8_e4m3fn // re-views since callers may bind the slot differently. @@ -149,13 +176,17 @@ static void mega_moe_pre_dispatch( SM100MegaMoEPreDispatchRuntime::Args args = { .group_size = group_size, - .use_fp4_acts = use_fp4_acts, + .mma_kind = mma_kind, .use_pdl = use_pdl, .x = x.const_data_ptr(), .topk_idx = topk_idx.const_data_ptr(), .topk_weights = topk_weights.const_data_ptr(), + .expert_scales = expert_scales.has_value() + ? expert_scales->const_data_ptr() : nullptr, .buf_x = buf_x.data_ptr(), .buf_x_sf = buf_x_sf.data_ptr(), + .buf_x_scales = buf_x_scales.has_value() + ? buf_x_scales->data_ptr() : nullptr, .buf_topk_idx = buf_topk_idx.data_ptr(), .buf_topk_weights = buf_topk_weights.data_ptr(), .num_tokens = static_cast(num_tokens), diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index 06e33d0270..1410eb0a8e 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -43,8 +43,8 @@ class SM90FP8MegaMoERuntime final : public LaunchRuntime bool reuse_accum_as_final; bool l2_arrival_counter; bool l2_epilogue_requires_full_sync; - bool split_phase_hot_path; bool use_swap_ab; + int num_shared_experts; MegaMoESM90Config config; // Runtime arguments @@ -66,6 +66,18 @@ class SM90FP8MegaMoERuntime final : public LaunchRuntime CUtensorMap tensor_map_l2_weights; const float* l2_weights_sf; + // Fused shared expert. When `num_shared_experts == 0` these mirror the + // routed descriptors: the kernel never reads them. Shared L1 acts SF needs + // no descriptor (the loader warp gathers the K-major `x_sf` column itself). + CUtensorMap tensor_map_shared_l1_acts; + 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_weights; + const float* shared_l2_weights_sf; + CUtensorMap tensor_map_shared_l2_acts_sf; + // Launch configs LaunchArgs launch_args; }; @@ -81,7 +93,6 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, - {}, {}, {}, {}, {}, {}, @@ -102,7 +113,6 @@ static void __instantiate_kernel() {{ args.num_max_tokens_per_rank, args.hidden, args.intermediate_hidden, args.num_experts, args.num_topk, - args.config.num_experts_per_wave, args.config.block_m, args.config.block_n, args.config.block_k, args.config.num_max_pool_tokens, args.config.num_padded_sf_pool_tokens, @@ -115,8 +125,8 @@ static void __instantiate_kernel() {{ args.reuse_accum_as_final ? "true" : "false", args.l2_arrival_counter ? "true" : "false", args.l2_epilogue_requires_full_sync ? "true" : "false", - args.split_phase_hot_path ? "true" : "false", - args.use_swap_ab ? "true" : "false"); + args.use_swap_ab ? "true" : "false", + args.num_shared_experts); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -133,7 +143,15 @@ static void __instantiate_kernel() {{ args.tensor_map_l2_acts, args.tensor_map_l2_acts_sf, args.tensor_map_l2_weights, - args.l2_weights_sf + args.l2_weights_sf, + args.tensor_map_shared_l1_acts, + 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_weights, + args.shared_l2_weights_sf, + args.tensor_map_shared_l2_acts_sf )); } }; @@ -144,20 +162,33 @@ static void sm90_fp8_mega_moe( const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf, const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf, + // Fused shared expert. `shared_l1_acts` is the local `x` region (and + // `shared_l1_acts_sf` its K-major per-128 SF); `shared_l2_acts` is the + // post-SwiGLU pool written by the fused L1 epilogue (zero-sized when the shared + // expert is off). The weight tensors are undefined when `num_shared_experts == 0`. + 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& 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_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, + const int& num_shared_experts, const float& activation_clamp, const bool& fast_math ) { const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts = num_experts_per_rank * num_ranks; const auto num_padded_sf_pool_tokens = static_cast(l1_acts_sf.size(0)); + const bool fuse_shared_experts = num_shared_experts > 0; + const int shared_intermediate_hidden = intermediate_hidden * num_shared_experts; - // Heuristics + // Heuristics. swapAB composes with the fused shared expert (shared's token + // M-axis matches swapAB's reduced output M-axis), so `use_swap_ab` is computed + // naturally for both paths. const auto config = get_mega_moe_config_sm90( num_ranks, num_experts, num_experts_per_rank, num_max_tokens_per_rank, num_tokens, num_topk, @@ -178,8 +209,6 @@ static void sm90_fp8_mega_moe( const bool default_split_mn_barrier_opt = config.block_m == 128 and config.block_n == 256 and config.num_epilogue_threads == 512; - const bool split_phase_hot_path = - config.block_m == 128 and config.block_n == 256 and hidden >= 7168; const bool decode_split_n_path = config.block_m == 64 and config.num_epilogue_threads == 256; const bool decode_split_n_bn256 = @@ -242,11 +271,12 @@ static void sm90_fp8_mega_moe( const int wg_l1_out_block_n = wg_block_n / 2; const bool split_n_shares_sf = split_n_warpgroups and wg_l1_out_block_n < kL2ActsSFGranK; + const bool l1_output_full_tile = split_n_shares_sf or use_swap_ab; const int l1_output_swizzle_mode = 0; const int l1_output_box_n = - split_n_shares_sf ? config.block_n / 2 : wg_l1_out_block_n; + l1_output_full_tile ? config.block_n / 2 : wg_l1_out_block_n; const int l1_output_box_m = - split_n_shares_sf ? config.block_m : wg_block_m; + l1_output_full_tile ? config.block_m : wg_block_m; const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, intermediate_hidden, config.num_max_pool_tokens, l1_output_box_n, l1_output_box_m, @@ -267,6 +297,80 @@ static void sm90_fp8_mega_moe( static_cast(l2_weights.stride(-2)), config.swizzle_weights_mode); + // ---- Fused shared expert descriptors ---- + // A: the local `x` region (K-major FP8) for L1 and the post-SwiGLU shared pool for + // L2, both with M = num_max_tokens_per_rank. Only the shared L2 activation SF gets a + // descriptor: the fused L1 epilogue writes it M-major, so a (BLOCK_M, 1) box is legal. + // The shared L1 activation SF (`x_sf`) is K-major, where that box would be 4 bytes and + // break TMA's 16-byte inner-box rule, so the loader warp gathers it into `smem_sfa`. + auto tensor_map_shared_l1_acts = tensor_map_l1_acts; + auto tensor_map_shared_l1_weights = tensor_map_l1_weights; + auto tensor_map_shared_l1_output = tensor_map_l1_output; + auto tensor_map_shared_l2_acts = tensor_map_l2_acts; + auto tensor_map_shared_l2_weights = tensor_map_l2_weights; + auto tensor_map_shared_l2_acts_sf = tensor_map_l2_acts_sf; + const float* shared_l1_weights_sf_ptr = l1_weights_sf.data_ptr(); + const float* shared_l2_weights_sf_ptr = l2_weights_sf.data_ptr(); + if (fuse_shared_experts) { + DG_HOST_ASSERT(shared_l1_acts.defined() and shared_l1_acts_sf.defined()); + DG_HOST_ASSERT(shared_l2_acts.defined() and shared_l2_acts_sf.defined()); + DG_HOST_ASSERT(static_cast(shared_l1_acts.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l1_acts.size(1)) == hidden); + DG_HOST_ASSERT(static_cast(shared_l1_acts_sf.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l1_acts_sf.size(1)) == hidden / kGranK); + DG_HOST_ASSERT(static_cast(shared_l2_acts.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l2_acts.size(1)) == shared_intermediate_hidden); + DG_HOST_ASSERT(static_cast(shared_l2_acts_sf.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l2_acts_sf.size(1)) == + shared_intermediate_hidden / kL2ActsSFGranK); + + tensor_map_shared_l1_acts = make_tma_2d_desc(shared_l1_acts, + hidden, num_max_tokens_per_rank, + config.block_k, config.block_m, + static_cast(shared_l1_acts.stride(-2)), + config.swizzle_acts_mode); + tensor_map_shared_l1_weights = make_tma_2d_desc(shared_l1_weights, + hidden, shared_intermediate_hidden * 2, + config.block_k, weight_tma_block_n, + static_cast(shared_l1_weights.stride(-2)), + config.swizzle_weights_mode); + tensor_map_shared_l1_output = 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)), + l1_output_swizzle_mode); + tensor_map_shared_l2_acts = make_tma_2d_desc(shared_l2_acts, + shared_intermediate_hidden, num_max_tokens_per_rank, + config.block_k, config.block_m, + static_cast(shared_l2_acts.stride(-2)), + config.swizzle_acts_mode); + tensor_map_shared_l2_weights = make_tma_2d_desc(shared_l2_weights, + shared_intermediate_hidden, hidden, + config.block_k, weight_tma_block_n, + static_cast(shared_l2_weights.stride(-2)), + config.swizzle_weights_mode); + shared_l1_weights_sf_ptr = shared_l1_weights_sf.data_ptr(); + shared_l2_weights_sf_ptr = shared_l2_weights_sf.data_ptr(); + // Shared L1 acts SF (`x_sf`) stays K-major: staging an M-major copy would cost a + // full transpose kernel on every call and the staged tensor would be freed while + // the launch is still in flight, so the loader warp gathers the column into + // `smem_sfa` itself (see `process_a_sfa_block`) and needs no descriptor. + // + // Shared L2 acts SF lives in the workspace `shared_l2_sf_buffer` and is written + // M-major during the launch by the fused L1 epilogue, so it can TMA-load with a + // (BLOCK_M, 1) box. Re-view the same memory with {1, nmt} strides (the from_blob + // view already attached those in `sm90_mega.hpp`) and build the descriptor. + auto shared_l2_acts_sf_mm = torch::from_blob( + shared_l2_acts_sf.data_ptr(), + {num_max_tokens_per_rank, shared_intermediate_hidden / 64}, + {1, num_max_tokens_per_rank}, + shared_l2_acts_sf.options()); + tensor_map_shared_l2_acts_sf = make_tma_sf_desc( + cute::UMMA::Major::MN, shared_l2_acts_sf_mm, + num_max_tokens_per_rank, shared_intermediate_hidden, + config.block_m, kL2ActsSFGranK, 1, 0); + } + // Stats can be optional int* cumulative_local_expert_recv_stats_ptr = nullptr; if (cumulative_local_expert_recv_stats.has_value()) @@ -285,8 +389,8 @@ static void sm90_fp8_mega_moe( .reuse_accum_as_final = reuse_accum_as_final, .l2_arrival_counter = l2_arrival_counter, .l2_epilogue_requires_full_sync = l2_epilogue_requires_full_sync, - .split_phase_hot_path = split_phase_hot_path, .use_swap_ab = use_swap_ab, + .num_shared_experts = num_shared_experts, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, @@ -301,6 +405,14 @@ static void sm90_fp8_mega_moe( .tensor_map_l2_acts_sf = tensor_map_l2_acts_sf, .tensor_map_l2_weights = tensor_map_l2_weights, .l2_weights_sf = l2_weights_sf.data_ptr(), + .tensor_map_shared_l1_acts = tensor_map_shared_l1_acts, + .tensor_map_shared_l1_weights = tensor_map_shared_l1_weights, + .shared_l1_weights_sf = shared_l1_weights_sf_ptr, + .tensor_map_shared_l1_output = tensor_map_shared_l1_output, + .tensor_map_shared_l2_acts = tensor_map_shared_l2_acts, + .tensor_map_shared_l2_weights = tensor_map_shared_l2_weights, + .shared_l2_weights_sf = shared_l2_weights_sf_ptr, + .tensor_map_shared_l2_acts_sf = tensor_map_shared_l2_acts_sf, .launch_args = LaunchArgs(num_sms, config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, config.smem_size, config.cluster_size) }; diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index 9cff5cd350..8141879cb7 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -639,21 +639,22 @@ int64_t dg_get_token_alignment_for_mega_moe() { return (int64_t)mega::get_token_alignment_for_mega_moe(); } -static Tensor mega_tensor_to_ffi(const torch::Tensor& tensor) { - if (not tensor.defined()) - return Tensor(); - // DLPack/TVM-FFI does not expose PyTorch's float8 dtype consistently. - // Preserve its storage as int8 and restore the logical dtype in Python. - const auto ffi_tensor = tensor.scalar_type() == torch::kFloat8_e4m3fn ? - tensor.view(torch::kInt8) : tensor; - return Tensor::FromDLPack(at::toDLPack(ffi_tensor)); +int64_t dg_get_block_m_for_mega_moe(int64_t num_ranks, int64_t num_experts, + int64_t num_max_tokens_per_rank, int64_t num_tokens, + int64_t num_topk, std::string mma_type) { + return static_cast(mega::get_block_m_for_mega_moe( + static_cast(num_ranks), + static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_tokens), + static_cast(num_topk), + mma_type)); } -using MegaBufferTensors = Tuple; +using MegaSliceResult = Tuple; -Tuple> +Tuple> dg_get_symm_buffer_size_for_mega_moe(int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, int64_t num_topk, int64_t hidden, int64_t intermediate_hidden, std::string mma_type, std::string activation, int64_t num_shared_experts) { @@ -670,25 +671,39 @@ dg_get_symm_buffer_size_for_mega_moe(int64_t num_ranks, int64_t num_experts, int ); auto slice_input_buffers = [=](TensorView buffer) { + const auto buffer_torch = convert_to_torch_tensor(buffer); 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] = fn(convert_to_torch_tensor(buffer)); - return MegaBufferTensors( - mega_tensor_to_ffi(x), mega_tensor_to_ffi(x_sf), - mega_tensor_to_ffi(topk_idx), mega_tensor_to_ffi(topk_weights), - mega_tensor_to_ffi(shared_l1_acts), mega_tensor_to_ffi(shared_l1_acts_sf), - mega_tensor_to_ffi(shared_l2_acts), mega_tensor_to_ffi(shared_l2_acts_sf), - mega_tensor_to_ffi(l1_acts), mega_tensor_to_ffi(l1_acts_sf), - mega_tensor_to_ffi(l2_acts), mega_tensor_to_ffi(l2_acts_sf) + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, x_scales] = fn(buffer_torch); + // DLPack cannot carry FP8/FP4 dtypes, so activation views cross the + // bridge as raw bytes; undefined views (BF16 SFs, absent shared + // experts) cross as empty byte tensors. + const auto as_bytes = [&](const torch::Tensor& t) { + const auto t_val = t.defined() + ? t + : torch::empty({0}, torch::TensorOptions().dtype(torch::kChar).device(buffer_torch.device())); + return Tensor::FromDLPack(at::toDLPack( + t_val.scalar_type() == torch::kFloat8_e4m3fn or t_val.scalar_type() == torch::kUInt8 + ? t_val.view(at::kChar) : t_val)); + }; + return MegaSliceResult( + as_bytes(x), as_bytes(x_sf), + as_bytes(topk_idx), as_bytes(topk_weights), + as_bytes(shared_l1_acts), as_bytes(shared_l1_acts_sf), + as_bytes(shared_l2_acts), as_bytes(shared_l2_acts_sf), + as_bytes(l1_acts), as_bytes(l1_acts_sf), + as_bytes(l2_acts), as_bytes(l2_acts_sf), + as_bytes(x_scales) ); }; - return Tuple>( + return Tuple>( num_bytes, slice_input_buffers); } -Tuple(TensorView)>> +Tuple(TensorView)>> dg_get_symm_buffer_size_for_sm90_mega_moe(int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, int64_t num_topk, int64_t hidden, - int64_t intermediate_hidden, bool use_fp8_dispatch, std::string activation) { + int64_t intermediate_hidden, bool use_fp8_dispatch, std::string activation, + int64_t num_shared_experts) { auto [num_bytes, fn] = mega::get_symm_buffer_size_for_sm90_mega_moe( static_cast(num_ranks), static_cast(num_experts), @@ -697,12 +712,16 @@ dg_get_symm_buffer_size_for_sm90_mega_moe(int64_t num_ranks, int64_t num_experts static_cast(hidden), static_cast(intermediate_hidden), use_fp8_dispatch, - activation + activation, + static_cast(num_shared_experts) ); auto slice_input_buffers = [=](TensorView buffer) { - auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = fn(convert_to_torch_tensor(buffer)); - return Tuple( + // The last two views are the fused shared-expert pool and its SF; they are + // zero-sized when the shared expert is disabled. + auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf] = fn(convert_to_torch_tensor(buffer)); + return Tuple( Tensor::FromDLPack(at::toDLPack(x.view(at::kChar))), Tensor::FromDLPack(at::toDLPack(x_sf)), Tensor::FromDLPack(at::toDLPack(topk_idx)), @@ -710,22 +729,23 @@ dg_get_symm_buffer_size_for_sm90_mega_moe(int64_t num_ranks, int64_t num_experts Tensor::FromDLPack(at::toDLPack(l1_acts.view(at::kChar))), Tensor::FromDLPack(at::toDLPack(l1_acts_sf)), Tensor::FromDLPack(at::toDLPack(l2_acts.view(at::kChar))), - Tensor::FromDLPack(at::toDLPack(l2_acts_sf)) + Tensor::FromDLPack(at::toDLPack(l2_acts_sf)), + Tensor::FromDLPack(at::toDLPack(shared_l2_acts.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(shared_l2_acts_sf)) ); }; - return Tuple(TensorView)>>( + return Tuple(TensorView)>>( num_bytes, slice_input_buffers); } -void dg_fp8_fp4_mega_moe(TensorView y, - TensorView l1_weights, TensorView l1_weights_sf, - TensorView l2_weights, TensorView l2_weights_sf, +void dg_fp8_fp4_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_sf, TensorView l2_weights, TensorView l2_weights_sf, Optional shared_l1_weights, Optional shared_l1_weights_sf, Optional shared_l2_weights, Optional shared_l2_weights_sf, Optional cumulative_local_expert_recv_stats, TensorView sym_buffer, Array sym_buffer_ptrs, int64_t rank_idx, int64_t num_max_tokens_per_rank, int64_t num_experts, int64_t num_topk, - Tuple recipe, std::string activation, Optional activation_clamp_opt, - bool fast_math) { + Tuple recipe, std::string mma_type, std::string activation, Optional activation_clamp_opt, + bool fast_math, bool use_x_scales, Optional l1_alphas, + Optional l2_alphas, Optional l2_act_scales) { auto c_val = cumulative_local_expert_recv_stats.has_value()? std::optional(convert_to_torch_tensor(cumulative_local_expert_recv_stats.value())) : std::nullopt; auto act_clamp_opt_val = activation_clamp_opt.has_value()? std::optional(static_cast(activation_clamp_opt.value())) : std::nullopt; std::vector sym_buffer_ptrs_val; @@ -736,27 +756,29 @@ void dg_fp8_fp4_mega_moe(TensorView y, } auto [recipe_a, recipe_b, recipe_c] = recipe; auto recipe_val = std::make_tuple(static_cast(recipe_a), static_cast(recipe_b), static_cast(recipe_c)); - std::optional> shared_l1_weights_val = std::nullopt; - std::optional> shared_l2_weights_val = std::nullopt; - if (shared_l1_weights.has_value() and shared_l1_weights_sf.has_value()) { - shared_l1_weights_val = std::make_tuple( - convert_to_torch_tensor(shared_l1_weights.value()), - convert_to_torch_tensor(shared_l1_weights_sf.value())); - } - if (shared_l2_weights.has_value() and shared_l2_weights_sf.has_value()) { - shared_l2_weights_val = std::make_tuple( - convert_to_torch_tensor(shared_l2_weights.value()), - convert_to_torch_tensor(shared_l2_weights_sf.value())); - } + DG_HOST_ASSERT(shared_l1_weights.has_value() == shared_l1_weights_sf.has_value()); + DG_HOST_ASSERT(shared_l2_weights.has_value() == shared_l2_weights_sf.has_value()); + auto shared_l1_val = shared_l1_weights.has_value() + ? std::optional>(std::make_tuple( + convert_to_torch_tensor(shared_l1_weights.value()), convert_to_torch_tensor(shared_l1_weights_sf.value()))) + : std::nullopt; + auto shared_l2_val = shared_l2_weights.has_value() + ? std::optional>(std::make_tuple( + convert_to_torch_tensor(shared_l2_weights.value()), convert_to_torch_tensor(shared_l2_weights_sf.value()))) + : std::nullopt; mega::fp8_fp4_mega_moe( convert_to_torch_tensor(y), - std::make_pair(convert_to_torch_tensor(l1_weights), convert_to_torch_tensor(l1_weights_sf)), - std::make_pair(convert_to_torch_tensor(l2_weights), convert_to_torch_tensor(l2_weights_sf)), - shared_l1_weights_val, shared_l2_weights_val, + std::make_tuple(convert_to_torch_tensor(l1_weights), convert_to_torch_tensor(l1_weights_sf)), + std::make_tuple(convert_to_torch_tensor(l2_weights), convert_to_torch_tensor(l2_weights_sf)), + shared_l1_val, shared_l2_val, c_val, convert_to_torch_tensor(sym_buffer), sym_buffer_ptrs_val, static_cast(rank_idx), static_cast(num_max_tokens_per_rank), static_cast(num_experts), - static_cast(num_topk), recipe_val, activation, act_clamp_opt_val, fast_math + static_cast(num_topk), recipe_val, mma_type, activation, act_clamp_opt_val, fast_math, + use_x_scales, + l1_alphas.has_value() ? std::optional(convert_to_torch_tensor(l1_alphas.value())) : std::nullopt, + l2_alphas.has_value() ? std::optional(convert_to_torch_tensor(l2_alphas.value())) : std::nullopt, + l2_act_scales.has_value() ? std::optional(convert_to_torch_tensor(l2_act_scales.value())) : std::nullopt ); } @@ -774,16 +796,14 @@ void dg_bf16_mega_moe(TensorView y, TensorView l1_weights, TensorView l2_weights for (Array::iterator it = sym_buffer_ptrs.begin(); it != sym_buffer_ptrs.end(); ++it) { sym_buffer_ptrs_val.push_back(*it); } - auto shared_l1_weights_val = shared_l1_weights.has_value() ? - std::make_optional(convert_to_torch_tensor(shared_l1_weights.value())) : std::nullopt; - auto shared_l2_weights_val = shared_l2_weights.has_value() ? - std::make_optional(convert_to_torch_tensor(shared_l2_weights.value())) : std::nullopt; + auto shared_l1_val = shared_l1_weights.has_value()? std::optional(convert_to_torch_tensor(shared_l1_weights.value())) : std::nullopt; + auto shared_l2_val = shared_l2_weights.has_value()? std::optional(convert_to_torch_tensor(shared_l2_weights.value())) : std::nullopt; mega::bf16_mega_moe( convert_to_torch_tensor(y), convert_to_torch_tensor(l1_weights), convert_to_torch_tensor(l2_weights), - shared_l1_weights_val, shared_l2_weights_val, + shared_l1_val, shared_l2_val, c_val, convert_to_torch_tensor(sym_buffer), sym_buffer_ptrs_val, static_cast(rank_idx), static_cast(num_max_tokens_per_rank), static_cast(num_experts), static_cast(num_topk), activation, act_clamp_opt_val, fast_math @@ -792,6 +812,8 @@ void dg_bf16_mega_moe(TensorView y, TensorView l1_weights, TensorView l2_weights void dg_fp8_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_sf, TensorView l2_weights, TensorView l2_weights_sf, + Optional shared_l1_weights, Optional shared_l1_weights_sf, + Optional shared_l2_weights, Optional shared_l2_weights_sf, Optional cumulative_local_expert_recv_stats, TensorView sym_buffer, Array sym_buffer_ptrs, int64_t rank_idx, int64_t num_max_tokens_per_rank, int64_t num_experts, int64_t num_topk, Tuple recipe, std::string activation, Optional activation_clamp_opt, bool fast_math) { @@ -806,10 +828,25 @@ void dg_fp8_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_ auto [recipe_a, recipe_b, recipe_c] = recipe; auto recipe_val = std::make_tuple(static_cast(recipe_a), static_cast(recipe_b), static_cast(recipe_c)); + // Fused shared expert: both weight tuples are optional and must come together + std::optional> shared_l1_weights_val = std::nullopt; + std::optional> shared_l2_weights_val = std::nullopt; + if (shared_l1_weights.has_value()) { + DG_HOST_ASSERT(shared_l1_weights_sf.has_value() and shared_l2_weights.has_value() and + shared_l2_weights_sf.has_value()); + shared_l1_weights_val = std::make_tuple( + convert_to_torch_tensor(shared_l1_weights.value()), + convert_to_torch_tensor(shared_l1_weights_sf.value())); + shared_l2_weights_val = std::make_tuple( + convert_to_torch_tensor(shared_l2_weights.value()), + convert_to_torch_tensor(shared_l2_weights_sf.value())); + } + mega::fp8_mega_moe( convert_to_torch_tensor(y), std::make_pair(convert_to_torch_tensor(l1_weights), convert_to_torch_tensor(l1_weights_sf)), std::make_pair(convert_to_torch_tensor(l2_weights), convert_to_torch_tensor(l2_weights_sf)), + shared_l1_weights_val, shared_l2_weights_val, c_val, convert_to_torch_tensor(sym_buffer), sym_buffer_ptrs_val, static_cast(rank_idx), static_cast(num_max_tokens_per_rank), static_cast(num_experts), static_cast(num_topk), recipe_val, activation, act_clamp_opt_val, fast_math @@ -820,7 +857,8 @@ void dg_mega_moe_pre_dispatch( TensorView x, TensorView topk_idx, TensorView topk_weights, TensorView buf_x, TensorView buf_x_sf, TensorView buf_topk_idx, TensorView buf_topk_weights, - int64_t num_tokens, int64_t group_size, bool use_fp4_acts) { + int64_t num_tokens, int64_t group_size, std::string mma_type, + Optional buf_x_scales, Optional expert_scales) { mega_moe_pre_dispatch( convert_to_torch_tensor(x), convert_to_torch_tensor(topk_idx), @@ -831,7 +869,13 @@ void dg_mega_moe_pre_dispatch( convert_to_torch_tensor(buf_topk_weights), static_cast(num_tokens), static_cast(group_size), - use_fp4_acts + mma_type, + buf_x_scales.has_value() + ? std::optional(convert_to_torch_tensor(buf_x_scales.value())) + : std::nullopt, + expert_scales.has_value() + ? std::optional(convert_to_torch_tensor(expert_scales.value())) + : std::nullopt ); } @@ -855,6 +899,7 @@ void dg_mega_moe_pre_dispatch_sm90( } TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_token_alignment_for_mega_moe, dg_get_token_alignment_for_mega_moe); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_block_m_for_mega_moe, dg_get_block_m_for_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_symm_buffer_size_for_mega_moe, dg_get_symm_buffer_size_for_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_symm_buffer_size_for_sm90_mega_moe, dg_get_symm_buffer_size_for_sm90_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_mega_moe, dg_fp8_fp4_mega_moe); diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index cbab419e13..518ec53cd4 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -1,6 +1,8 @@ import os import subprocess import torch +import tvm_ffi +from glob import glob # Set some default environment provided at setup try: @@ -12,73 +14,88 @@ except ImportError: pass -# Configs -from . import _C -from ._C import ( - set_num_sms, - get_num_sms, - set_tc_util, - get_tc_util, - set_ignore_compile_dims, - set_block_size_multiple_of, - set_pdl, - get_pdl, -) +_extension_paths = glob(os.path.join(os.path.dirname(__file__), '_C*.so')) +if not _extension_paths: + raise ImportError('DeepGEMM extension is missing; build the TVM-FFI _C module first.') +_C = tvm_ffi.load_module(max(_extension_paths, key=os.path.getmtime)) + + +def _bind_exports(*names: str) -> None: + for name in names: + globals()[name] = getattr(_C, name) + -# cuBLASLt Kernels -from ._C import ( - cublaslt_gemm_nt, cublaslt_gemm_nn, - cublaslt_gemm_tn, cublaslt_gemm_tt, +# Configs and cuBLASLt kernels +_bind_exports( + 'set_num_sms', 'get_num_sms', 'set_tc_util', 'get_tc_util', + 'set_pdl', 'get_pdl', + 'cublaslt_gemm_nt', 'cublaslt_gemm_nn', 'cublaslt_gemm_tn', 'cublaslt_gemm_tt', ) try: # DeepGEMM Kernels - from ._C import ( + _kernel_exports = ( # FP8 FP4 GEMMs - fp8_fp4_gemm_nt, fp8_fp4_gemm_nn, - fp8_fp4_gemm_tn, fp8_fp4_gemm_tt, - m_grouped_fp8_fp4_gemm_nt_contiguous, - m_grouped_fp8_fp4_gemm_nn_contiguous, - m_grouped_fp8_fp4_gemm_nt_masked, + 'fp8_fp4_gemm_nt', 'fp8_fp4_gemm_nn', + 'fp8_fp4_gemm_tn', 'fp8_fp4_gemm_tt', + 'm_grouped_fp8_fp4_gemm_nt_contiguous', + 'm_grouped_fp8_fp4_gemm_nn_contiguous', + 'm_grouped_fp8_fp4_gemm_nt_masked', # FP8 GEMMs - fp8_gemm_nt, fp8_gemm_nn, - fp8_gemm_tn, fp8_gemm_tt, - fp8_gemm_nt_skip_head_mid, - m_grouped_fp8_gemm_nt_contiguous, - m_grouped_fp8_gemm_nn_contiguous, - m_grouped_fp8_gemm_nt_masked, - k_grouped_fp8_gemm_nt_contiguous, - k_grouped_fp8_gemm_tn_contiguous, + 'fp8_gemm_nt', 'fp8_gemm_nn', + 'fp8_gemm_tn', 'fp8_gemm_tt', + 'fp8_gemm_nt_skip_head_mid', + 'm_grouped_fp8_gemm_nt_contiguous', + 'm_grouped_fp8_gemm_nn_contiguous', + 'm_grouped_fp8_gemm_nt_masked', + 'k_grouped_fp8_gemm_nt_contiguous', + 'k_grouped_fp8_gemm_tn_contiguous', # BF16 GEMMs - bf16_gemm_nt, bf16_gemm_nn, - bf16_gemm_tn, bf16_gemm_tt, - m_grouped_bf16_gemm_nt_contiguous, - m_grouped_bf16_gemm_nn_contiguous, - m_grouped_bf16_gemm_nt_masked, - k_grouped_bf16_gemm_tn_contiguous, + 'bf16_gemm_nt', 'bf16_gemm_nn', + 'bf16_gemm_tn', 'bf16_gemm_tt', + 'm_grouped_bf16_gemm_nt_contiguous', + 'm_grouped_bf16_gemm_nn_contiguous', + 'm_grouped_bf16_gemm_nt_masked', + 'k_grouped_bf16_gemm_tn_contiguous', # Einsum kernels - einsum, - fp8_einsum, + 'einsum', + 'fp8_einsum', # Attention kernels - fp8_fp4_mqa_logits, - get_paged_mqa_logits_metadata, - fp8_fp4_paged_mqa_logits, + 'fp8_fp4_mqa_logits', + 'get_paged_mqa_logits_metadata', + 'fp8_fp4_paged_mqa_logits', # Attention kernels (legacy) - fp8_mqa_logits, - fp8_paged_mqa_logits, + 'fp8_mqa_logits', + 'fp8_paged_mqa_logits', # Hyperconnection kernels - tf32_hc_prenorm_gemm, + 'tf32_hc_prenorm_gemm', # Layout kernels - transform_sf_into_required_layout, - # MegaMoE - get_block_m_for_mega_moe, + 'transform_sf_into_required_layout', ) + # Bind per-name: one kernel absent from this build (e.g. plain FP8 GEMMs + # not exported by the TVM-FFI bridge, or CUDA < 12.1) must not drop the rest. + for _name in _kernel_exports: + try: + _bind_exports(_name) + except AttributeError: + pass + + # Sugared recipe-tuple form matching the sgl_deep_gemm wheel API + # (the raw _C entry takes the recipe flattened into 3 ints). + if hasattr(_C, 'transform_sf_into_required_layout'): + def transform_sf_into_required_layout(sf, mn, k, recipe, num_groups=None, + is_sfa=None, disable_ue8m0_cast=False): + (recipe_a, recipe_b, recipe_c) = recipe if len(recipe) == 3 else (recipe[0], recipe[1], None) + return _C.transform_sf_into_required_layout( + sf, mn, k, recipe_a, recipe_b, recipe_c, num_groups, is_sfa, disable_ue8m0_cast) # Some alias for legacy supports # TODO: remove these later - fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked - bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked -except ImportError: + if 'm_grouped_fp8_gemm_nt_masked' in globals(): + fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked + if 'm_grouped_bf16_gemm_nt_masked' in globals(): + bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked +except AttributeError: # Expected behavior for CUDA runtime version before 12.1 pass @@ -88,6 +105,7 @@ get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe, fp8_fp4_mega_moe, + nvfp4_mega_moe, bf16_mega_moe, mega_moe_pre_dispatch, ) @@ -101,7 +119,8 @@ try: from . import legacy except Exception as e: - print(f'Failed to load legacy DeepGEMM A100 Triton kernels: {e}') + if not (isinstance(e, ImportError) and 'PyInit__C' in str(e)): + print(f'Failed to load legacy DeepGEMM A100 Triton kernels: {e}') # Initialize CPP modules def _find_cuda_home() -> str: diff --git a/deep_gemm/include/deep_gemm/common/math.cuh b/deep_gemm/include/deep_gemm/common/math.cuh index 6d5ece847e..f97421ff9d 100644 --- a/deep_gemm/include/deep_gemm/common/math.cuh +++ b/deep_gemm/include/deep_gemm/common/math.cuh @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -112,6 +113,31 @@ CUTLASS_DEVICE void get_e2m1_sf_and_sf_inv(const float2& amax, float2& sf, float sf.y = fast_pow2(exp_y), sf_inv.y = fast_pow2(-exp_y); } +CUTLASS_DEVICE void get_nvfp4_sf_and_sf_inv(const float2& amax, float2& sf, float2& sf_inv, uint2& sf_bits) { + constexpr float kInvMax = 1.0f / 6.0f; + const __nv_fp8_e4m3 qx(amax.x * kInvMax), qy(amax.y * kInvMax); + sf_bits = {qx.__x, qy.__x}; + sf = {static_cast(qx), static_cast(qy)}; + sf_inv.x = sf.x > 0.0f ? __frcp_rn(sf.x) : 0.0f; + sf_inv.y = sf.y > 0.0f ? __frcp_rn(sf.y) : 0.0f; +} + + +CUTLASS_DEVICE uint32_t cast_into_e2m1x2_pairs(const float2& lower, const float2& upper) { + uint32_t packed; + asm volatile( + "{\n\t" + ".reg .b8 byte0;\n\t" + ".reg .b8 byte1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %3, %1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %2;\n\t" + "mov.b32 %0, {byte0, byte1, byte0, byte1};\n\t" + "}\n" + : "=r"(packed) + : "f"(lower.x), "f"(lower.y), "f"(upper.x), "f"(upper.y)); + return packed; +} + // Pack two FP32 values into one FP4 (E2M1) byte: lower nibble = a, upper = b. // Matches PTX `cvt.rn.satfinite.e2m1x2.f32 d, b, a` (b → upper, a → lower). CUTLASS_DEVICE uint32_t cvt_pack_f32_to_e2m1x2(const float& a, const float& b) { diff --git a/deep_gemm/include/deep_gemm/common/tma_copy.cuh b/deep_gemm/include/deep_gemm/common/tma_copy.cuh index 2c5bf708d4..662e900a89 100644 --- a/deep_gemm/include/deep_gemm/common/tma_copy.cuh +++ b/deep_gemm/include/deep_gemm/common/tma_copy.cuh @@ -19,7 +19,8 @@ template (cute::TMA::CacheHintSm100::EVICT_NORMAL)) { DG_STATIC_ASSERT(static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL) == static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), "Invalid cache hint"); constexpr uint32_t BLOCK_INNER_ATOM = get_inner_block_atom_size(); @@ -29,7 +30,7 @@ copy(void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr #pragma unroll for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { cute::SM90_TMA_LOAD_2D::copy(desc_ptr, reinterpret_cast(barrier_ptr), - static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + cache_hint, smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, outer_idx); } @@ -39,7 +40,7 @@ copy(void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr #pragma unroll for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { cute::SM100_TMA_2SM_LOAD_2D::copy(desc_ptr, reinterpret_cast(barrier_ptr), - static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + cache_hint, smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, outer_idx); } diff --git a/deep_gemm/include/deep_gemm/common/types.cuh b/deep_gemm/include/deep_gemm/common/types.cuh index ee4e14c031..27baa225b8 100644 --- a/deep_gemm/include/deep_gemm/common/types.cuh +++ b/deep_gemm/include/deep_gemm/common/types.cuh @@ -7,12 +7,35 @@ namespace deep_gemm { enum class MmaKind { BF16 = 0, MXFP8FP4 = 1, + MXFP4 = 2, + NVFP4 = 3, }; constexpr CUTLASS_HOST_DEVICE int get_element_size(const MmaKind& mma_kind) { switch (mma_kind) { case MmaKind::BF16: return 2; case MmaKind::MXFP8FP4: return 1; + case MmaKind::MXFP4: return 1; + case MmaKind::NVFP4: return 1; + default: return 0; + } +} + +constexpr CUTLASS_HOST_DEVICE int get_element_bits(const MmaKind& mma_kind) { + switch (mma_kind) { + case MmaKind::BF16: return 16; + case MmaKind::MXFP8FP4: return 8; + case MmaKind::MXFP4: return 4; + case MmaKind::NVFP4: return 4; + default: return 0; + } +} + +constexpr CUTLASS_HOST_DEVICE int get_sf_gran_k(const MmaKind& mma_kind) { + switch (mma_kind) { + case MmaKind::MXFP8FP4: return 32; + case MmaKind::MXFP4: return 32; + case MmaKind::NVFP4: return 16; default: return 0; } } diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh index 73a8536f8d..ef45833562 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh @@ -102,7 +102,7 @@ sm100_bf16_mega_moe_impl(void* y, kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk, kNumRingTokens, 0, - /*with_sf=*/ false, + MmaKind::BF16, kNumSharedExperts ); const auto workspace = buffer.workspace; diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh index fed21f6d56..8c5d5976b1 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh @@ -26,6 +26,7 @@ template < uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, uint32_t STORE_BLOCK_M, uint32_t SF_BLOCK_M, uint32_t SF_BLOCK_N, + MmaKind kMmaKind, uint32_t kNumRingTokens, uint32_t kNumSFRingTokens, uint32_t kNumStages, @@ -34,38 +35,17 @@ template < uint32_t kNumEpilogueThreads, uint32_t kNumSMs, uint32_t kNumRanks, float kActivationClamp, + float kSwiGLUAlpha, bool kUseSitu, bool kFastMath, - // ====== Stream A0.1 — DG_USE_FP4_ACTS ====== - // When true, the L1 epilogue quantizes its SwiGLU outputs to E2M1 (FP4) + - // UE8M0 SF instead of E4M3 (FP8) + UE8M0 SF. The per-row gmem footprint - // halves (intermediate_hidden / 2 packed bytes vs intermediate_hidden FP8 - // bytes) and the smem CD staging is sized accordingly. The L2 phase still - // reads its activations as FP8 in this step (separate flag for A0.2), so - // end-to-end output is intentionally not bit-equivalent to the FP8 path — - // the accuracy harness compares L1's quantized output decoded back to BF16. - bool kUseFp4Acts = false, - // ====== Stream A0.5 — DG_USE_MXF4_KIND ====== - // When true (and `kUseFp4Acts` also true), L1 + L2 mainloops swap from - // `kind::mxf8f6f4.block_scale.block32` (K=32 with-padding FP4 smem) to - // `kind::mxf4.block_scale.block32` (K=64 dense FP4 smem). Per the - // `recipes/mxf4_vs_mxf8f6f4` microbench, `kind::mxf4` delivers 2× FLOPS/ - // cycle in isolation; the standalone GEMM (`kernels/fused_gemm_mxf4_native_1cta`) - // realizes +22%, the fused capstone (`kernels/fused_swiglu_mxf4_native_two_gemm`) - // realizes +20.6%. This kernel ports the same swap into the production - // mega_moe path. `kind::mxf4` is K-major-only (PTX ISA Table 53) and - // accepts only E2M1 inputs — see the host-side `DG_HOST_ASSERT(not - // use_mxf4_kind or use_fp4_acts)` in `mega.hpp`. - bool kUseMxf4Kind = false, - // ====== Stream B (combine path) — DG_USE_FP8_COMBINE ====== - // When true, the L2 epilogue ships FP8 E4M3 + per-(token, N=128) UE8M0 - // SF over NVLink instead of BF16. Byte footprint per token per slot: - // off: kHidden * 2 (BF16) - // on: kHidden + kHidden / kCombineGranK (FP8 + SF, kCombineGranK=128) - // Halves NVLink bytes/token on the second a2a. Independent of - // `kUseFp4Acts` / `kUseMxf4Kind` (which control the dispatch a2a + - // mainloops); this flag only changes the combine slot's layout + - // L2 epilogue write-back + combine-reduce read. + bool kUseXScales, + bool kWithL1Alphas, + bool kWithL2Alphas, + bool kWithL2ActScales, + // When true, the L2 epilogue ships FP8 E4M3 + a per-(token, N block) UE8M0 SF over + // NVLink instead of BF16, and the combine reduce dequantizes on the fly. Halves the + // second all-to-all's bytes per token. Independent of the MMA kind, which only + // controls the dispatch a2a and the mainloops. bool kUseFp8Combine = false, bool kHasShared = (kNumSharedExperts > 0), uint32_t L1_SHAPE_N = kIntermediateHidden * 2, @@ -87,6 +67,9 @@ template < CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void sm100_fp8_fp4_mega_moe_impl(void* y, int* cumulative_local_expert_recv_stats, + const float* __restrict__ l1_alphas, + const float* __restrict__ l2_alphas, + const float* __restrict__ l2_act_scales, const uint32_t num_tokens, const __grid_constant__ layout::SymBuffer sym_buffer, const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, @@ -146,27 +129,21 @@ sm100_fp8_fp4_mega_moe_impl(void* y, cute::prefetch_tma_descriptor(&tensor_map_shared_l2_weights_sf); } - // Stream A0.0b — DG_USE_FP4_ACTS L1 input path. The registered input - // token and routed L1 token buffers use packed E2M1 when enabled. - constexpr uint32_t kInputTokenBytes = kUseFp4Acts ? (kHidden / 2) : kHidden; - constexpr uint32_t kNumBytesPerPullForActs = - kNumBytesPerPull < kInputTokenBytes ? kNumBytesPerPull : kInputTokenBytes; - constexpr uint32_t kNumTokenPullChunks = kInputTokenBytes / kNumBytesPerPullForActs; - DG_STATIC_ASSERT(kNumTokenPullChunks * kNumBytesPerPullForActs == kInputTokenBytes, - "kNumBytesPerPullForActs must divide input token bytes"); - constexpr auto pull_layout = layout::Data(kNumBytesPerPullForActs); - - // Workspace and unified buffer layout. Shared experts and FP4 acts are - // deliberately mutually exclusive at the host boundary. + DG_STATIC_ASSERT(kMmaKind == MmaKind::NVFP4 or kMmaKind == MmaKind::MXFP4 or kMmaKind == MmaKind::MXFP8FP4, + "Invalid MMA kind"); + constexpr bool kIsNVFP4 = kMmaKind == MmaKind::NVFP4; + constexpr bool kIsMXFP4 = kMmaKind == MmaKind::MXFP4; + constexpr bool kIsFP4Acts = kIsNVFP4 or kIsMXFP4; + + // Workspaces and Buffer const auto buffer = layout::MegaMoEBuffer( sym_buffer.get_base_ptr(), kHidden, kIntermediateHidden, kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk, kNumRingTokens, kNumSFRingTokens, - /*with_sf=*/ true, + kMmaKind, kNumSharedExperts, - kUseFp4Acts, kUseFp8Combine ); const auto workspace = buffer.workspace; @@ -196,7 +173,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, }; // SF and its buffer configs - constexpr uint32_t kGranK = 32; + constexpr uint32_t kGranK = get_sf_gran_k(kMmaKind); + constexpr uint32_t kSFChunkK = kGranK * 4; constexpr uint32_t kNumUTCCPAlignedElems = 128; DG_STATIC_ASSERT(SF_BLOCK_M == math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems), "Invalid SF_BLOCK_M"); DG_STATIC_ASSERT(SF_BLOCK_N == BLOCK_N, "No padding is needed for SFB"); @@ -211,58 +189,47 @@ sm100_fp8_fp4_mega_moe_impl(void* y, constexpr uint32_t kCombineGranK = 128; DG_STATIC_ASSERT(kHidden % kCombineGranK == 0, "kHidden must be a multiple of 128 for FP8 combine SF"); // Data types - // NOTES: activations are FP8 (e4m3), weights are FP4 (e2m1) - using a_dtype_t = cutlass::float_e4m3_t; - using b_dtype_t = cutlass::detail::float_e2m1_unpacksmem_t; - using shared_b_dtype_t = cutlass::float_e4m3_t; - // Stream A0.2: when `kUseFp4Acts` is on, the L2 phase reads acts as - // E2M1 instead of E4M3. Both share the same byte footprint in smem - // (FP8 = 1 B, FP4 unpacksmem = 1 B with `_ALIGN16B` padding), so the - // smem A allocation, swizzle mode (128 B), and umma_desc stride math - // are identical. Only the *MMA instruction descriptor*'s A-dtype field - // and the source-side TMA `expect_tx` differ between phases. - using l2_a_dtype_t = cute::conditional_t; - // Stream A0.0b: same deal for L1 — when `kUseFp4Acts` is on, the L1 - // phase reads its A operand from the L1 token pool as packed E2M1. - // Same `_ALIGN16B` padded smem layout as L2; same MMA instruction - // descriptor flip from E4M3 to E2M1. - using l1_a_dtype_t = cute::conditional_t; + // NOTES: MXFP8FP4 uses FP8 (e4m3) activations with FP4 (e2m1) weights unpacked into + // 8-bit smem containers; NVFP4/MXFP4 use packed FP4 for both. A and B tiles are therefore + // always addressed in bytes, and only the MMA instruction descriptor sees the real formats. + // Dense `float_e2m1_t` resolves to `MXF4Format::E2M1`, `float_e2m1_unpacksmem_t` to + // `MXF8F6F4Format::E2M1` — mixing them up launches but hits `cudaErrorIllegalInstruction` + constexpr uint32_t kNumElemBits = kIsFP4Acts ? 4 : 8; + using a_dtype_t = uint8_t; + using b_dtype_t = uint8_t; + using shared_b_dtype_t = uint8_t; + using mma_a_fmt_t = cute::conditional_t; + using mma_b_fmt_t = cute::conditional_t; + using mma_shared_b_fmt_t = cute::conditional_t; + using sf_fmt_t = cute::conditional_t; // MMA configs // NOTES: always swap A/B, 2-CTA MMA, and matrices are K-major constexpr uint32_t LAYOUT_AD_M = 128; constexpr uint32_t UMMA_M = LAYOUT_AD_M * 2; constexpr uint32_t UMMA_N = BLOCK_M; // Swap AB - // Stream A0.5: kind::mxf4 runs K=64 dense per call (vs K=32 for - // kind::mxf8f6f4). The number of MMA calls per K-tile is BLOCK_K / UMMA_K. - constexpr uint32_t UMMA_K = kUseMxf4Kind ? 64 : 32; + constexpr uint32_t BLOCK_K_BYTES = BLOCK_K * kNumElemBits / 8; + constexpr uint32_t UMMA_BLOCK_K_BYTES = 128; + constexpr uint32_t UMMA_K_BYTES = 32; + constexpr uint32_t UMMA_K = UMMA_K_BYTES * 8 / kNumElemBits; + constexpr uint32_t kNumMMAsPerSFChunk = kSFChunkK / UMMA_K; + constexpr uint32_t kNumRoutedBTxTiles = kIsFP4Acts ? 2 : 1; constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / 2; // Multicast on A constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; DG_STATIC_ASSERT(BLOCK_M % 16 == 0, "Invalid block M"); DG_STATIC_ASSERT(BLOCK_N == LAYOUT_AD_M, "Invalid block N"); - DG_STATIC_ASSERT(BLOCK_K % 128 == 0, "Invalid block K"); + DG_STATIC_ASSERT(BLOCK_K_BYTES % UMMA_BLOCK_K_BYTES == 0, "Invalid block K"); + DG_STATIC_ASSERT(kNumMMAsPerSFChunk == (kIsNVFP4 ? 1u : kIsMXFP4 ? 2u : 4u), "Invalid SF chunking"); // Swizzle configs - // Stream A0.5: under `kUseMxf4Kind`, A and B smem use the dense FP4 - // layout (`_ALIGN8B`, 2 nibbles/byte) instead of the with-padding - // layout (`_ALIGN16B`, 1 byte per element). Per-K-row byte stride - // halves: BLOCK_K elements × 0.5 B/elem = BLOCK_K / 2 bytes. Swizzle - // mode tracks the row-byte width. - constexpr uint32_t kSwizzleAMode = kUseMxf4Kind - ? (BLOCK_K / 2) - : (BLOCK_K * static_cast(sizeof(a_dtype_t))); - constexpr uint32_t kSwizzleBMode = kUseMxf4Kind - ? (BLOCK_K / 2) - : (BLOCK_K * static_cast(sizeof(b_dtype_t))); - // Stream A0.2: l2_a_dtype must keep the same smem footprint as - // a_dtype so SMEM_A_SIZE_PER_STAGE / kSwizzleAMode are unchanged. - DG_STATIC_ASSERT(sizeof(l2_a_dtype_t) == sizeof(a_dtype_t), - "L2 A dtype must match A in smem footprint"); - DG_STATIC_ASSERT(sizeof(l1_a_dtype_t) == sizeof(a_dtype_t), - "L1 A dtype must match A in smem footprint"); + constexpr uint32_t kSwizzleAMode = 128; + constexpr uint32_t kSwizzleBMode = 128; constexpr uint32_t kSwizzleCDMode = 128; DG_STATIC_ASSERT(BLOCK_N % kSwizzleCDMode == 0, "Invalid block N"); + constexpr uint64_t kActsCacheHint = static_cast(BLOCK_M <= 128 ? + cute::TMA::CacheHintSm100::EVICT_LAST : cute::TMA::CacheHintSm100::EVICT_NORMAL); + // Epilogue configs constexpr uint32_t kNumEpilogueStages = 2; constexpr uint32_t kNumTMAStoreStages = 2; @@ -278,40 +245,40 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Shared memory sizes // NOTES: FP8 CD output for L1 (2 TMA stages, BLOCK_N/2 post-SwiGLU), BF16 output for L2 (no TMA, a single stage) constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; - // ====== Stream A0.1 ====== - // FP4 path packs 2 elements per byte → row footprint halves. We keep - // `L1_OUT_BLOCK_N` in *elements* and introduce a row-byte-stride that - // depends on the flag, so the existing offset arithmetic (`row * - // L1_OUT_BLOCK_N_BYTES`) still works for both paths. - constexpr uint32_t L1_OUT_ROW_BYTES = kUseFp4Acts ? (L1_OUT_BLOCK_N / 2) : L1_OUT_BLOCK_N; - 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(pull_layout.get_num_bytes() * kNumDispatchWarps, kSharedMemoryAlignment); - // Stream A0.5: under `kUseMxf4Kind`, dense FP4 smem (2 nibbles/byte) - // halves the per-stage byte footprint vs the with-padding layout. - constexpr uint32_t SMEM_A_SIZE_PER_STAGE = kUseMxf4Kind - ? (LOAD_BLOCK_M * BLOCK_K / 2) - : (LOAD_BLOCK_M * BLOCK_K * static_cast(sizeof(a_dtype_t))); - constexpr uint32_t SMEM_B_SIZE_PER_STAGE = kUseMxf4Kind - ? (LOAD_BLOCK_N * BLOCK_K / 2) - : (LOAD_BLOCK_N * BLOCK_K * static_cast(sizeof(b_dtype_t))); - constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * (BLOCK_K / 128) * sizeof(uint32_t); - constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * (BLOCK_K / 128) * sizeof(uint32_t); - // L1 CD smem: FP8 path = STORE_BLOCK_M * L1_OUT_BLOCK_N bytes/stage, - // FP4 path = STORE_BLOCK_M * L1_OUT_BLOCK_N / 2 bytes/stage. - constexpr uint32_t SMEM_CD_L1_SIZE = - kNumEpilogueWarpgroups * STORE_BLOCK_M * L1_OUT_ROW_BYTES * kNumTMAStoreStages; - constexpr uint32_t SMEM_CD_L2_SIZE = - kNumEpilogueWarpgroups * STORE_BLOCK_M * BLOCK_N * sizeof(nv_bfloat16); - constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE; - constexpr uint32_t SMEM_CD_L1_SIZE_PER_STAGE = SMEM_CD_L1_SIZE / kNumTMAStoreStages; - constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = - SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); - DG_STATIC_ASSERT(SMEM_CD_SIZE % kSharedMemoryAlignment == 0 and - SMEM_A_SIZE_PER_STAGE % kSharedMemoryAlignment == 0 and - SMEM_B_SIZE_PER_STAGE % kSharedMemoryAlignment == 0, - "Shared memory of CD/A/B must be aligned to 1024 bytes"); + constexpr uint32_t L1_OUT_BLOCK_N_BYTES = L1_OUT_BLOCK_N * kNumElemBits / 8; + constexpr uint32_t AMAX_REDUCTION_WARP_BUFFER_SIZE = STORE_BLOCK_M / 2; // float2 + + struct SharedStorage { + alignas(kSharedMemoryAlignment) uint32_t expert_token_count[kNumExperts]; + alignas(kSharedMemoryAlignment) uint8_t dispatch_send_buffer[kNumDispatchWarps][kNumBytesPerPull]; + union { + alignas(kSharedMemoryAlignment) uint8_t l1[kNumEpilogueWarpgroups][kNumTMAStoreStages][STORE_BLOCK_M * L1_OUT_BLOCK_N_BYTES]; + alignas(kSharedMemoryAlignment) nv_bfloat16 l2[kNumEpilogueWarpgroups][STORE_BLOCK_M * BLOCK_N]; + } smem_d; + alignas(kSharedMemoryAlignment) a_dtype_t smem_a[kNumStages][LOAD_BLOCK_M * BLOCK_K_BYTES]; + alignas(kSharedMemoryAlignment) b_dtype_t smem_b[kNumStages][LOAD_BLOCK_N * BLOCK_K_BYTES]; + uint32_t smem_sfa[kNumStages][SF_BLOCK_M * (BLOCK_K / kSFChunkK)]; + uint32_t smem_sfb[kNumStages][SF_BLOCK_N * (BLOCK_K / kSFChunkK)]; + float2 amax_reduction[kNumEpilogueWarps][AMAX_REDUCTION_WARP_BUFFER_SIZE]; + task_info_t task_infos[kNumScheduleStages]; + Barrier dispatch_barriers[kNumDispatchWarps]; + Barrier full_barriers[kNumStages]; + Barrier empty_barriers[kNumStages]; + Barrier tmem_full_barriers[kNumEpilogueStages]; + Barrier tmem_empty_barriers[kNumEpilogueStages]; + Barrier combine_barriers[kNumEpilogueWarps * 2]; + Barrier task_info_full_barriers[kNumScheduleStages]; + Barrier task_info_empty_barriers[kNumScheduleStages]; + uint32_t tmem_ptr_in_smem; + }; + constexpr uint32_t kNumReusableSmemBytes = offsetof(SharedStorage, dispatch_barriers); + SharedStorage &shared_storage = *reinterpret_cast(smem_buffer); + + // Send buffers + constexpr auto pull_layout = layout::Data(kNumBytesPerPull); + const auto smem_send_buffers = layout::Buffer( + pull_layout, kNumDispatchWarps, 1, + static_cast(shared_storage.dispatch_send_buffer)); // Tensor memory size constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; @@ -322,103 +289,53 @@ sm100_fp8_fp4_mega_moe_impl(void* y, constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); - // Assign shared memory for dispatch warps - const auto smem_expert_count = reinterpret_cast(smem_buffer); - const auto smem_send_buffers = layout::Buffer( - pull_layout, kNumDispatchWarps, 1, - math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE)); - - // GEMM shared memory: C/D, A, B - // NOTES: GEMM shared memory starts after the dispatch region, aligned to 1024 bytes - auto smem_gemm_base = math::advance_ptr( - smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE - ); - - // D/A/B shared memory - auto smem_cd = utils::PatternVisitor([=](const uint32_t& i) { - return math::advance_ptr(smem_gemm_base, i * SMEM_CD_L1_SIZE_PER_STAGE); - }); - auto smem_cd_l2 = smem_cd[0]; - auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { - return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); - }); - auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { - return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); - }); - - // SF shared memory: SFA and SFB per pipeline stage - auto sf_start_ptr = math::advance_ptr(smem_gemm_base, - SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); - auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { - return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); - }); - auto smem_sfb = utils::PatternVisitor([=](const uint32_t& i) { - return reinterpret_cast(sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * SMEM_SFB_SIZE_PER_STAGE); - }); - - // Epilogue amax reduction shared memory - auto smem_amax_reduction = reinterpret_cast(smem_sfb[kNumStages]); - - // Barriers and tensor memory pointer - auto barrier_start_ptr = reinterpret_cast(smem_amax_reduction + STORE_BLOCK_M * kNumEpilogueWarps / 2); - 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 tmem_full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumDispatchWarps + kNumStages * 2 + i); }); - auto tmem_empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumDispatchWarps + kNumStages * 2 + kNumEpilogueStages + i); }); - auto combine_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumDispatchWarps + kNumStages * 2 + kNumEpilogueStages * 2 + i); }); - auto task_info_full_barriers = barrier_start_ptr + kNumDispatchWarps + - kNumStages * 2 + kNumEpilogueStages * 2 + kNumEpilogueWarps * 2; - auto task_info_empty_barriers = task_info_full_barriers + kNumScheduleStages; - auto task_infos = reinterpret_cast( - barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + - kNumEpilogueStages * 2 + kNumEpilogueWarps * 2 + - kNumScheduleStages * 2); - auto tmem_ptr_in_smem = reinterpret_cast(task_infos + kNumScheduleStages); - // A cluster sync is essential for 2CTA tensor memory allocation comm::cluster_sync_with_relaxed_arrive(); // Initialization if (warp_idx == 0) { // Clean shared memory - if (cute::elect_one_sync()) - ptx::st_shared_bulk(smem_expert_count, kNumExperts * sizeof(uint32_t)); + if (cute::elect_one_sync()) { + ptx::st_shared_bulk( + shared_storage.expert_token_count, + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment) + ); + } } else if (warp_idx == 1) { // Init m-barriers for dispatch #pragma unroll for (uint32_t i = lane_idx; i < kNumDispatchWarps; i += 32) - dispatch_barriers[i]->init(1); + shared_storage.dispatch_barriers[i].init(1); cutlass::arch::fence_barrier_init(); } else if (warp_idx == 2) { // Init GEMM barriers if (cute::elect_one_sync()) { #pragma unroll for (uint32_t i = 0; i < kNumStages; ++ i) { - // Arrive at all CTAs - full_barriers[i]->init(2 * 2); - empty_barriers[i]->init(1); + // Arrive at 2 CTAs, A + B + shared_storage.full_barriers[i].init(2 * 2); + shared_storage.empty_barriers[i].init(1); } #pragma unroll for (uint32_t i = 0; i < kNumEpilogueStages; ++ i) { // Arrive at all CTAs - tmem_full_barriers[i]->init(1); + shared_storage.tmem_full_barriers[i].init(1); // Arrive only at the leader CTA - tmem_empty_barriers[i]->init(2 * kNumEpilogueThreads); + shared_storage.tmem_empty_barriers[i].init(2 * kNumEpilogueThreads); } #pragma unroll for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++ i) - combine_barriers[i]->init(1); + shared_storage.combine_barriers[i].init(1); #pragma unroll for (uint32_t i = 0; i < kNumScheduleStages; ++ i) { - task_info_full_barriers[i].init(1); - task_info_empty_barriers[i].init(kNumScheduleConsumerThreads); + shared_storage.task_info_full_barriers[i].init(1); + shared_storage.task_info_empty_barriers[i].init(kNumScheduleConsumerThreads); } } cutlass::arch::fence_barrier_init(); } else if (warp_idx == 3) { // Allocate tensor memory - Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + Allocator().allocate(kNumTmemCols, &shared_storage.tmem_ptr_in_smem); } // NOTES: Using `.relaxed` is allowed here since `fence_barrier_init` is `.release.cluster`, // and `barrier.cluster.wait.aligned` is by default `.acquire` @@ -434,9 +351,9 @@ sm100_fp8_fp4_mega_moe_impl(void* y, kNumRingBlocks, kNumSharedExperts>( workspace, - task_info_full_barriers, - task_info_empty_barriers, - task_infos + shared_storage.task_info_full_barriers, + shared_storage.task_info_empty_barriers, + shared_storage.task_infos ); // MMA pipeline and TMA phases @@ -461,9 +378,10 @@ sm100_fp8_fp4_mega_moe_impl(void* y, constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; // Adjust registers - constexpr uint32_t kNumDispatchRegisters = 48; - constexpr uint32_t kNumNonEpilogueRegisters = 40; - constexpr uint32_t kNumEpilogueRegisters = 208; + constexpr bool kUseMoreEpilogueRegisters = kNumExpertsPerRank <= 64; + constexpr uint32_t kNumDispatchRegisters = kUseMoreEpilogueRegisters ? 48 : 96; + constexpr uint32_t kNumNonEpilogueRegisters = kUseMoreEpilogueRegisters ? 40 : 88; + constexpr uint32_t kNumEpilogueRegisters = kUseMoreEpilogueRegisters ? 208 : 160; DG_STATIC_ASSERT(kNumDispatchRegisters * kNumDispatchThreads + kNumNonEpilogueRegisters * kNumNonEpilogueThreads + kNumEpilogueRegisters * kNumEpilogueThreads <= 64512, @@ -502,15 +420,15 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Count experts' tokens read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { - atomicAdd_block(smem_expert_count + expert_idx, 1); + atomicAdd_block(shared_storage.expert_token_count + expert_idx, 1); }); ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); // Get SM offset (~6.5 us) #pragma unroll for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { - const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); - smem_expert_count[i] = static_cast( + const uint64_t send_value = (1ull << 32) | static_cast(shared_storage.expert_token_count[i]); + shared_storage.expert_token_count[i] = static_cast( ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); } ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); @@ -518,7 +436,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Write source indices (~2 us with 512 tokens) 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_slot_idx = atomicAdd_block(shared_storage.expert_token_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; @@ -562,7 +480,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Pull token data and SF from remote ranks into local L1 buffer 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]; + const auto pull_mbarrier = &shared_storage.dispatch_barriers[warp_idx]; // Per-rank counts for current expert (re-loaded when expert changes) constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); @@ -662,9 +580,10 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; - // Hidden bytes are divided into chunks on the reusable-ring path. - // The full-pool path keeps the tuned one-shot pull/store sequence. - DG_STATIC_ASSERT(kNumTokenPullChunks >= 1, "Invalid token pull chunk count"); + // Hidden bytes are divided into chunks + constexpr uint32_t kNumHiddenTokenBytes = kHidden * kNumElemBits / 8; + constexpr uint32_t kNumChunks = kNumHiddenTokenBytes / kNumBytesPerPull; + DG_STATIC_ASSERT(kNumChunks * kNumBytesPerPull == kNumHiddenTokenBytes, "kNumBytesPerPull must divide the token bytes"); const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; const uint32_t pool_block_idx = pool_token_idx / BLOCK_M; @@ -684,8 +603,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const auto issue_and_wait_pull_store = [&](const uint32_t& i) { ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); ptx::tma_store_1d( - math::advance_ptr(dst_base_ptr, i * kNumBytesPerPullForActs), - pull_buffer.get_base_ptr(), kNumBytesPerPullForActs + math::advance_ptr(dst_base_ptr, i * kNumBytesPerPull), + pull_buffer.get_base_ptr(), kNumBytesPerPull ); cute::tma_store_arrive(); ptx::tma_store_wait<0>(); @@ -693,27 +612,27 @@ sm100_fp8_fp4_mega_moe_impl(void* y, if constexpr (kUseFullPoolFP8FP4Path) { if (cute::elect_one_sync()) { ptx::tma_load_1d( - pull_buffer.get_base_ptr(), src_base_ptr, pull_mbarrier, kInputTokenBytes); + pull_buffer.get_base_ptr(), src_base_ptr, pull_mbarrier, kNumHiddenTokenBytes); } } else { if (cute::elect_one_sync()) { #pragma unroll - for (uint32_t i = 0; i < kNumTokenPullChunks; ++ i) { + for (uint32_t i = 0; i < kNumChunks; ++ i) { ptx::tma_load_1d( pull_buffer.get_base_ptr(), - math::advance_ptr(src_base_ptr, i * kNumBytesPerPullForActs), - pull_mbarrier, kNumBytesPerPullForActs + math::advance_ptr(src_base_ptr, i * kNumBytesPerPull), + pull_mbarrier, kNumBytesPerPull ); - ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kNumBytesPerPullForActs); - i != (kNumTokenPullChunks - 1) ? issue_and_wait_pull_store(i) : void(); + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kNumBytesPerPull); + i != (kNumChunks - 1) ? issue_and_wait_pull_store(i) : void(); } } } __syncwarp(); - // Load and store SF (overlaps with TMA token load) - constexpr uint32_t kNumSFUint32 = kHidden / 128; - DG_STATIC_ASSERT(kNumSFUint32 > 0 and kHidden % 128 == 0, "Invalid SF"); + // Load and store SF (overlaps with last chunk's TMA load from remote) + constexpr uint32_t kNumSFUint32 = kHidden / kSFChunkK; + DG_STATIC_ASSERT(kNumSFUint32 > 0 and kHidden % kSFChunkK == 0, "Invalid SF"); const auto remote_sf_ptr = sym_buffer.map( buffer.input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), current_rank_in_expert_idx); @@ -730,7 +649,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, } __syncwarp(); - // Store weights and token data + // Store weights and metadata if (cute::elect_one_sync()) { // Load weights const auto weight = *sym_buffer.map( @@ -738,20 +657,26 @@ sm100_fp8_fp4_mega_moe_impl(void* y, current_rank_in_expert_idx); *buffer.l1_topk_weights_buffer.get_data_buffer(get_ring_token_idx(pool_token_idx)).template get_base_ptr() = weight; + if constexpr (kUseXScales) { + const auto x_scale = *sym_buffer.map( + buffer.input_x_scales_buffer.get_base_ptr() + src_token_idx, + current_rank_in_expert_idx); + *buffer.l1_x_scales_buffer.get_data_buffer(get_ring_token_idx(pool_token_idx)).template get_base_ptr() = x_scale; + } + if constexpr (kUseFullPoolFP8FP4Path) { - // Wait for TMA token load to complete - ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kInputTokenBytes); + // Wait for the one-shot TMA token load to complete + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kNumHiddenTokenBytes); ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); // Store token to local L1 buffer via TMA ptx::tma_store_1d( - dst_base_ptr, pull_buffer.get_base_ptr(), kInputTokenBytes); + dst_base_ptr, pull_buffer.get_base_ptr(), kNumHiddenTokenBytes); - // Write source metadata for combine write-back + // Write source metadata for combine write-back (logical pool token) *workspace.get_token_src_metadata_ptr(pool_token_idx) = {current_rank_in_expert_idx, src_token_idx, src_topk_idx}; - // Wait for token TMA store to complete cute::tma_store_arrive(); ptx::tma_store_wait<0>(); ptx::red_add_rel( @@ -762,10 +687,10 @@ sm100_fp8_fp4_mega_moe_impl(void* y, {current_rank_in_expert_idx, src_token_idx, src_topk_idx}; // Complete last chunk's store - issue_and_wait_pull_store(kNumTokenPullChunks - 1); + issue_and_wait_pull_store(kNumChunks - 1); const bool is_last_token = (token_idx == expert_end_idx - 1); ptx::red_add_rel( - workspace.get_l1_full_count_ptr(ring_block_idx), + workspace.get_l1_full_count_ptr(get_ring_block_idx(pool_block_idx)), is_last_token ? BLOCK_M - (token_idx_in_expert % BLOCK_M) : 1u ); } @@ -822,7 +747,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, *workspace.get_expert_recv_count_ptr(j, i) = 0; __syncwarp(); - // Clean L1 and L2 arrivals / ring-buffer counters + // Clean L1 and L2 full stuffs and ring buffer counts for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { const auto pool_block_idx = expert_pool_block_offset + j; if constexpr (kUseFullPoolFP8FP4Path) { @@ -899,64 +824,28 @@ sm100_fp8_fp4_mega_moe_impl(void* y, for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { // Wait consumer release - empty_barriers[stage_idx]->wait(phase ^ 1); + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); // Compute token offsets from block index uint32_t m_idx = block_idx * BLOCK_M; - uint32_t k_idx = k_block_idx * BLOCK_K; + uint32_t k_idx = k_block_idx * BLOCK_K_BYTES; const uint32_t sfa_m_idx = block_idx * SF_BLOCK_M; - uint32_t sfa_k_idx = k_block_idx * (BLOCK_K / 128); + uint32_t sfa_k_idx = k_block_idx * (BLOCK_K / kSFChunkK); // Add 2 CTA offsets for non-leader CTA if (not is_leader_cta) m_idx += task_info.get_umma_aligned_valid_m() / 2; - // TMA copy tokens and SFA, then arrive at full barrier. - // Stream A0.2 + A0.0b: under FP4 acts, BOTH L1 and L2 phases - // load A as packed E2M1 (`l1_a_dtype_t == l2_a_dtype_t == b_dtype_t`). - // Same per-byte smem layout as FP8 A (1 B/elem under `_ALIGN16B`), - // but source-side packed bytes are halved → expect_tx halved. + // TMA copy tokens and SFA, then arrive at full barrier if (cute::elect_one_sync()) { - if constexpr (kUseMxf4Kind) { - // Stream A0.5: dense FP4 smem (`_ALIGN8B`). The TMA - // descriptor's inner box covers BLOCK_K elements in - // BLOCK_K/2 bytes per row; one cluster-multicast TMA - // call fills the full A stage. Bypass `tma::copy` - // because its `BLOCK_INNER_ATOM = kSwizzleMode / - // sizeof(dtype_t)` math assumes ≥1-byte elements - // and would mis-stride sub-byte FP4 destinations. - cute::SM100_TMA_2SM_LOAD_2D::copy( - tensor_map_a_ptr, - reinterpret_cast(full_barriers[stage_idx]), - static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), - reinterpret_cast(smem_a[stage_idx]), - k_idx, m_idx); - } else if constexpr (kUseFp4Acts) { - // Both Linear1 (L1) and Linear2 (L2) take the FP4 path. - tma::copy( - tensor_map_a_ptr, full_barriers[stage_idx], - reinterpret_cast(smem_a[stage_idx]), - k_idx, m_idx, 2); - } else { - tma::copy( - tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], - k_idx, m_idx, 2); - } - tma::copy( - tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], sfa_m_idx, sfa_k_idx, 2); + tma::copy( + tensor_map_a_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_a[stage_idx], k_idx, m_idx, 2, 0, kActsCacheHint); + tma::copy( + tensor_map_sfa_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_sfa[stage_idx], sfa_m_idx, sfa_k_idx, 2); if (is_leader_cta) { - // Stream A0.5: under `kUseMxf4Kind`, smem A is dense - // FP4 (LOAD_BLOCK_M * BLOCK_K / 2 bytes per CTA, equal - // to source-side packed bytes — no `_ALIGN16B` doubling). - // For 2 CTAs (cluster multicast), tx-count is - // `2 * SMEM_A_SIZE_PER_STAGE` — same multiplier as the - // FP8 dense path. - const uint32_t expect_a_bytes = (kUseFp4Acts and not kUseMxf4Kind) - ? SMEM_A_SIZE_PER_STAGE // FP4 _ALIGN16B: source = LOAD_BLOCK_M * BLOCK_K / 2 per CTA × 2 CTAs (smem 2× larger) - : SMEM_A_SIZE_PER_STAGE * 2; // FP8 dense or FP4 dense (mxf4): source = smem footprint × 2 CTAs - full_barriers[stage_idx]->arrive_and_expect_tx(expect_a_bytes + SMEM_SFA_SIZE_PER_STAGE * 2); + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(sizeof(SharedStorage::smem_a[0]) * 2 + sizeof(SharedStorage::smem_sfa[0]) * 2); } else { - full_barriers[stage_idx]->arrive(0u); + shared_storage.full_barriers[stage_idx].arrive(0u); } } __syncwarp(); @@ -980,59 +869,42 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const auto shape_k = task_info.shape_k; const auto shape_n = task_info.shape_n; - const auto shape_sfb_k = math::ceil_div(shape_k, kGranK * 4u); + const auto shape_sfb_k = math::ceil_div(shape_k, kSFChunkK); const auto n_block_idx = task_info.n_cluster_idx * 2 + (is_leader_cta ? 0u : 1u); const auto num_k_blocks = math::ceil_div(shape_k, BLOCK_K); for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { // Wait consumer release - empty_barriers[stage_idx]->wait(phase ^ 1); + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); // Compute weight offset uint32_t n_idx = task_info.is_shared() ? n_block_idx * BLOCK_N : task_info.local_expert_idx * shape_n + n_block_idx * BLOCK_N; - uint32_t k_idx = k_block_idx * BLOCK_K; + uint32_t k_idx = k_block_idx * BLOCK_K_BYTES; uint32_t sfb_n_idx = n_block_idx * BLOCK_N; - uint32_t sfb_k_idx = task_info.is_shared() ? k_block_idx * (BLOCK_K / 128) : task_info.local_expert_idx * shape_sfb_k + k_block_idx * (BLOCK_K / 128); + uint32_t sfb_k_idx = task_info.is_shared() ? k_block_idx * (BLOCK_K / kSFChunkK) : task_info.local_expert_idx * shape_sfb_k + k_block_idx * (BLOCK_K / kSFChunkK); // TMA copy weights with SF if (cute::elect_one_sync()) { - if constexpr (kUseMxf4Kind) { - // Stream A0.5: dense FP4 smem; one cluster-multicast - // TMA call covers the full B stage. See A-side comment. - cute::SM100_TMA_2SM_LOAD_2D::copy( - tensor_map_b_ptr, - reinterpret_cast(full_barriers[stage_idx]), - static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), - reinterpret_cast(smem_b[stage_idx]), - k_idx, n_idx); - } else if (task_info.is_shared()) { - tma::copy( - tensor_map_b_ptr, full_barriers[stage_idx], - reinterpret_cast(smem_b[stage_idx]), - k_idx, n_idx, 2); - } else { - tma::copy( - tensor_map_b_ptr, full_barriers[stage_idx], smem_b[stage_idx], k_idx, n_idx, 2); - } - tma::copy( - tensor_map_sfb_ptr, full_barriers[stage_idx], smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx, 2); - if (is_leader_cta) { - // Stream A0.5: B-side tx-count for cluster-multicast - // counts SOURCE BYTES PER PEER × 2 PEERS (broadcast: both - // peers receive a copy of the same source bytes). For the - // existing FP4 unpacksmem path, that happens to equal - // `LOAD_BLOCK_N * BLOCK_K * 1B = SMEM_B_SIZE_PER_STAGE` - // (sizeof(b_dtype_t)=1 makes "smem footprint" a coincidental - // alias for source-bytes-summed). Under mxf4 dense FP4, - // SMEM_B_SIZE_PER_STAGE halves to `LOAD_BLOCK_N * BLOCK_K / 2`, - // so we need `* 2` to get the same source-bytes-summed value. - const uint32_t expect_b_bytes = - (kUseMxf4Kind or task_info.is_shared()) - ? SMEM_B_SIZE_PER_STAGE * 2 - : SMEM_B_SIZE_PER_STAGE; - full_barriers[stage_idx]->arrive_and_expect_tx(expect_b_bytes + SMEM_SFB_SIZE_PER_STAGE * 2); + if (task_info.is_shared()) { + tma::copy( + tensor_map_b_ptr, &shared_storage.full_barriers[stage_idx], reinterpret_cast(shared_storage.smem_b[stage_idx]), k_idx, n_idx, 2); + tma::copy( + tensor_map_sfb_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx, 2); + if (is_leader_cta) { + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(sizeof(SharedStorage::smem_b[0]) * 2 + sizeof(SharedStorage::smem_sfb[0]) * 2); + } else { + shared_storage.full_barriers[stage_idx].arrive(0u); + } } else { - full_barriers[stage_idx]->arrive(0u); + tma::copy( + tensor_map_b_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_b[stage_idx], k_idx, n_idx, 2); + tma::copy( + tensor_map_sfb_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx, 2); + if (is_leader_cta) { + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(sizeof(SharedStorage::smem_b[0]) * kNumRoutedBTxTiles + sizeof(SharedStorage::smem_sfb[0]) * 2); + } else { + shared_storage.full_barriers[stage_idx].arrive(0u); + } } } __syncwarp(); @@ -1047,73 +919,24 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Make instruction descriptor with block scaling // NOTES: always swap A/B auto routed_instr_desc = cute::UMMA::make_instr_desc_block_scaled< - b_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, + mma_b_fmt_t, mma_a_fmt_t, float, sf_fmt_t, UMMA_M, UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K >(); auto shared_instr_desc = cute::UMMA::make_instr_desc_block_scaled< - shared_b_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, - UMMA_M, UMMA_N, - cute::UMMA::Major::K, cute::UMMA::Major::K - >(); - // Stream A0.2 + A0.0b: when both L1 and L2 read FP4 acts under - // `kUseFp4Acts`, we need a separate instruction descriptor whose - // A-dtype field is E2M1 (not E4M3). All other fields (block-scale - // shape, UMMA M/N/K, K-major) are unchanged. The smem layout - // descriptors don't change because both dtypes have `sizeof = 1` - // (FP4 has the `_ALIGN16B` 1-byte-per-element padded smem layout). - // Single shared idesc — both `l1_a_dtype_t` and `l2_a_dtype_t` - // resolve to `b_dtype_t` (E2M1 unpacksmem) under the flag. - // - // Stream A0.5: under `kUseMxf4Kind`, the descriptor's a/b_format - // fields encode E2M1 as `MXF4Format::E2M1 = 1`, NOT - // `MXF8F6F4Format::E2M1 = 5`. CUTLASS picks the right enum via - // `to_UMMAFormat()`: passing `cute::float_e2m1_t` (dense) yields - // `MXF4Format::E2M1=1`; passing `cutlass::detail::float_e2m1_unpacksmem_t` - // yields `MXF8F6F4Format::E2M1=5`. Wrong encoding → the kernel - // launches but throws `cudaErrorIllegalInstruction` on first MMA. - using mxf4_e2m1_t = cute::float_e2m1_t; - using fp4_a_dtype_for_idesc = cute::conditional_t< - kUseMxf4Kind, mxf4_e2m1_t, b_dtype_t>; - using fp4_b_dtype_for_idesc = cute::conditional_t< - kUseMxf4Kind, mxf4_e2m1_t, l1_a_dtype_t>; - auto instr_desc_fp4 = cute::UMMA::make_instr_desc_block_scaled< - fp4_a_dtype_for_idesc, fp4_b_dtype_for_idesc, - float, cutlass::float_ue8m0_t, + mma_shared_b_fmt_t, mma_a_fmt_t, float, sf_fmt_t, UMMA_M, UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K >(); auto sf_desc = mma::sm100::make_sf_desc(nullptr); DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); - // Stream A0.5: under `kUseMxf4Kind`, smem A and B carry dense - // FP4 (2 nibbles/byte). The `make_umma_desc` helper asserts - // `kSwizzleMode == BLOCK_K * sizeof(dtype_t)`, so we pass a - // BLOCK_K of `BLOCK_K / 2` (the byte count) and `dtype_t = - // uint8_t` to get the right byte-stride math. The smem ptrs - // are reinterpreted to `uint8_t*` since the underlying buffer - // is just bytes. - cute::UMMA::SmemDescriptor a_desc, b_desc; - if constexpr (kUseMxf4Kind) { - a_desc = mma::sm100::make_umma_desc( - reinterpret_cast(smem_a[0]), 0, 0); - b_desc = mma::sm100::make_umma_desc( - reinterpret_cast(smem_b[0]), 0, 0); - } else { - a_desc = mma::sm100::make_umma_desc(smem_a[0], 0, 0); - b_desc = mma::sm100::make_umma_desc(smem_b[0], 0, 0); - } - uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 : 0u; - uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 : 0u; - auto shared_b_desc = b_desc; - if constexpr (kHasShared) { - shared_b_desc = mma::sm100::make_umma_desc< - cute::UMMA::Major::K, LOAD_BLOCK_N, BLOCK_K, kSwizzleBMode>( - reinterpret_cast(smem_b[0]), 0, 0); - } - uint32_t shared_b_desc_lo = lane_idx < kNumStages - ? shared_b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 - : 0u; + auto a_desc = mma::sm100::make_umma_desc(shared_storage.smem_a[0], 0, 0); + auto b_desc = mma::sm100::make_umma_desc(shared_storage.smem_b[0], 0, 0); + auto shared_b_desc = mma::sm100::make_umma_desc(reinterpret_cast(shared_storage.smem_b[0]), 0, 0); + uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * sizeof(SharedStorage::smem_a[0]) / 16 : 0u; + uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * sizeof(SharedStorage::smem_b[0]) / 16 : 0u; + uint32_t shared_b_desc_lo = lane_idx < kNumStages ? shared_b_desc.lo + lane_idx * sizeof(SharedStorage::smem_b[0]) / 16 : 0u; // Checks for MMA instructions DG_STATIC_ASSERT((UMMA_M == 64 and UMMA_N % 8 == 0 and 8 <= UMMA_N and UMMA_N <= 256) or @@ -1130,14 +953,11 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Dynamic update of UMMA N based on effective M auto& instr_desc = task_info.is_shared() ? shared_instr_desc : routed_instr_desc; mma::sm100::update_instr_desc_with_umma_n(instr_desc, task_info.get_umma_aligned_valid_m()); - if constexpr (kUseFp4Acts) - mma::sm100::update_instr_desc_with_umma_n( - instr_desc_fp4, task_info.get_umma_aligned_valid_m()); // Wait tensor memory empty barrier arrival const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; const auto accum_phase = (current_iter_idx ++ / kNumEpilogueStages) & 1; - tmem_empty_barriers[accum_stage_idx]->wait(accum_phase ^ 1); + shared_storage.tmem_empty_barriers[accum_stage_idx].wait(accum_phase ^ 1); ptx::tcgen05_after_thread_sync(); // Empty barrier arrival @@ -1146,11 +966,11 @@ sm100_fp8_fp4_mega_moe_impl(void* y, constexpr uint16_t kCTAMask = (1 << 2) - 1; cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); }; - umma_arrive(reinterpret_cast(empty_barriers[stage_idx])); + umma_arrive(reinterpret_cast(&shared_storage.empty_barriers[stage_idx])); // NOTES: the tensor memory accumulator pipeline has nothing to do with multicasting if (do_tmem_full_arrive) - umma_arrive(reinterpret_cast(tmem_full_barriers[accum_stage_idx])); + umma_arrive(reinterpret_cast(&shared_storage.tmem_full_barriers[accum_stage_idx])); __syncwarp(); }; @@ -1158,79 +978,59 @@ sm100_fp8_fp4_mega_moe_impl(void* y, #pragma unroll 2 for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { // Wait TMA load completion - full_barriers[stage_idx]->wait(phase); + shared_storage.full_barriers[stage_idx].wait(phase); ptx::tcgen05_after_thread_sync(); const auto a_desc_base_lo = ptx::exchange(a_desc_lo, stage_idx); const auto b_desc_base_lo = ptx::exchange(task_info.is_shared() ? shared_b_desc_lo : b_desc_lo, stage_idx); if (cute::elect_one_sync()) { - // UTCCP copy SFA and SFB to TMEM - using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta; - #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { - auto smem_ptr = smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); - } #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { - auto smem_ptr = smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); - } + for (uint32_t sf_chunk_idx = 0; sf_chunk_idx < BLOCK_K / kSFChunkK; ++ sf_chunk_idx) { + // UTCCP copy SFA and SFB to TMEM + using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta; + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = shared_storage.smem_sfa[stage_idx] + sf_chunk_idx * SF_BLOCK_M + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); + } + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = shared_storage.smem_sfb[stage_idx] + sf_chunk_idx * SF_BLOCK_N + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); + } - // Issue UMMA. Stream A0.2: L2 phase under FP4 acts - // uses `instr_desc_l2` (A=E2M1) instead of `instr_desc` - // (A=E4M3). The smem K-stride for A is the same - // (sizeof(l2_a_dtype_t) == sizeof(a_dtype_t) == 1) so - // `advance_umma_desc_lo` on `a_dtype_t` is correct - // for both phases. - // Stream A0.5: under `kUseMxf4Kind`, swap the MMA to - // `kind::mxf4` (cta_group::2). UMMA_K=64 (vs 32), - // so K_PER_TILE=2 (vs 4). The SF address top-2 bits - // are HALF-WORD offsets {0, 2} for scale_vec::2X - // (NOT byte offsets {0..3}); encode as `k * 2`, not `k`. - // Smem K-stride for the dense FP4 layout is `BLOCK_K/2` - // bytes/row, so `advance_umma_desc_lo` is templated on - // `uint8_t` and `BLOCK_K / 2` to match. - #pragma unroll - for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++ k) { - if constexpr (kUseMxf4Kind) { - const auto sf_id = k * 2u; // half-word offset for scale_vec::2X + // Issue UMMA + #pragma unroll + for (uint32_t k = 0; k < kNumMMAsPerSFChunk; ++ k) { + const uint32_t mma_idx = sf_chunk_idx * kNumMMAsPerSFChunk + k; + const uint32_t atom_byte_offset = (mma_idx * UMMA_K_BYTES / UMMA_BLOCK_K_BYTES) * UMMA_BLOCK_K_BYTES; + const uint32_t in_atom_byte_idx = mma_idx * UMMA_K_BYTES - atom_byte_offset; + // `kind::mxf4` uses `scale_vec::2X` + const uint32_t sf_id = kIsMXFP4 ? k * 2 : k; const auto runtime_instr_desc = - mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc_fp4, sf_id, sf_id); + mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sf_id, sf_id); a_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_M, kSwizzleAMode, uint8_t>( - a_desc_base_lo, 0, k * UMMA_K / 2); + cute::UMMA::Major::K, LOAD_BLOCK_M, kSwizzleAMode, a_dtype_t>(a_desc_base_lo, atom_byte_offset * LOAD_BLOCK_M, in_atom_byte_idx); b_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, uint8_t>( - b_desc_base_lo, 0, k * UMMA_K / 2); - ptx::SM100_MMA_MXF4_2x1SM_SS::fma( - b_desc, a_desc, accum_stage_idx * UMMA_N, - k_block_idx > 0 or k > 0, runtime_instr_desc, - kTmemStartColOfSFB, kTmemStartColOfSFA); - } else { - // Stream A0.0b: under `kUseFp4Acts`, both L1 and L2 read - // A as E2M1. Shared tasks retain their upstream FP8×FP8 - // descriptor; routed tasks select FP4 or the baseline descriptor. - const auto runtime_instr_desc = kUseFp4Acts - ? mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc_fp4, k, k) - : mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, k, k); - a_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_M, kSwizzleAMode, a_dtype_t>(a_desc_base_lo, 0, k * UMMA_K); - if (task_info.is_shared()) { - b_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, shared_b_dtype_t>( - b_desc_base_lo, 0, k * UMMA_K); + cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, b_dtype_t>(b_desc_base_lo, atom_byte_offset * LOAD_BLOCK_N, in_atom_byte_idx); + if constexpr (kIsNVFP4) { + ptx::SM100_MMA_MXF4NVF4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or mma_idx > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); + } else if constexpr (kIsMXFP4) { + ptx::SM100_MMA_MXF4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or mma_idx > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); } else { - b_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, b_dtype_t>( - b_desc_base_lo, 0, k * UMMA_K); + ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or mma_idx > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); } - ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma( - b_desc, a_desc, accum_stage_idx * UMMA_N, - k_block_idx > 0 or k > 0, runtime_instr_desc, - kTmemStartColOfSFB, kTmemStartColOfSFA); } } } @@ -1245,7 +1045,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // To safely deconstruct barriers, we need another round of waits if (current_iter_idx > 0) { const auto accum_phase_idx = ((current_iter_idx - 1) / kNumEpilogueStages) & 1; - tmem_empty_barriers[(current_iter_idx - 1) % kNumEpilogueStages]->wait(accum_phase_idx); + shared_storage.tmem_empty_barriers[(current_iter_idx - 1) % kNumEpilogueStages].wait(accum_phase_idx); } } } else if (warp_idx == kNumDispatchWarps + 3) { @@ -1262,7 +1062,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // NOTES: tensor memory addresses are simplified, as the hardware will ignore the warp index bits, // i.e., no need for `tmem_ptr |= (epilogue_warp_idx * 32) << 16`. // NOTES: we also forbid two CTAs to share the same SM and its tensor memory - DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(tmem_ptr_in_smem) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&shared_storage.tmem_ptr_in_smem) == 0); // GEMM epilogue warps const auto epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); @@ -1297,7 +1097,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Wait UMMA arrival const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; const auto accum_phase = (current_iter_idx ++ / kNumEpilogueStages) & 1; - tmem_full_barriers[accum_stage_idx]->wait(accum_phase); + shared_storage.tmem_full_barriers[accum_stage_idx].wait(accum_phase); ptx::tcgen05_after_thread_sync(); // Now we can release the task @@ -1327,17 +1127,29 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Unified L1 epilogue: SwiGLU in-place using granularity 8 interleaved weights // With `SM100_TMEM_LOAD_16dp256b1x`, gate/up pairs are: - // (values[0], values[2]), (values[1], values[3]), - // (values[4], values[6]), (values[5], values[7]) - // Shared experts have no routed top-k weight and therefore use 1. float stored_cached_weight = 1.0f; + float stored_cached_x_scale = 1.0f; + + float2 l1_alpha = {1.0f, 1.0f}; + if constexpr (kWithL1Alphas) { + if (not task_info.is_shared()) + l1_alpha = __ldg(reinterpret_cast(l1_alphas) + + task_info.local_expert_idx); + } + + // fc2 input global scale; the caller folds the inverse into `l2_alphas` + float l2_act_gs = 1.0f; + if constexpr (kWithL2ActScales) { + if (not task_info.is_shared()) + l2_act_gs = __ldg(l2_act_scales + task_info.local_expert_idx); + } #pragma unroll for (uint32_t s = 0; s < WG_BLOCK_M / STORE_BLOCK_M; ++ s) { // Early break if the entire store block is beyond the valid token range if (epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M >= valid_m) { ptx::tcgen05_before_thread_sync(); - tmem_empty_barriers[accum_stage_idx]->arrive(0u); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); break; } @@ -1357,53 +1169,81 @@ sm100_fp8_fp4_mega_moe_impl(void* y, .template get_base_ptr(); } + // Per-token L1 input outer scales, cached like the topk weights. + // Routed rows read the ring copy made at dispatch; shared rows + // read the local input scales directly. + if constexpr (kUseXScales) { + if ((j * ATOM_M) % 32 == 0 and + (WG_BLOCK_M % 32 == 0 or j * ATOM_M + lane_idx < WG_BLOCK_M)) { + const uint32_t row = epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M + lane_idx; + stored_cached_x_scale = task_info.is_shared() ? + *buffer.input_x_scales_buffer.get_data_buffer(m_idx + row).template get_base_ptr() : + *buffer.l1_x_scales_buffer.get_data_buffer(ring_m_idx + row).template get_base_ptr(); + } + } + const float2 x_scales = kUseXScales ? float2{ + ptx::exchange(stored_cached_x_scale, (j * ATOM_M) % 32 + (lane_idx % 4) * 2 + 0), + ptx::exchange(stored_cached_x_scale, (j * ATOM_M) % 32 + (lane_idx % 4) * 2 + 1) + } : float2{1.0f, 1.0f}; + const float2 gate_scales = kWithL1Alphas ? + __fmul2_rn(x_scales, {l1_alpha.x, l1_alpha.x}) : x_scales; + const float2 up_scales = kWithL1Alphas ? + __fmul2_rn(x_scales, {l1_alpha.y, l1_alpha.y}) : x_scales; + // Load weights from register cache - const float2 weights = { + float2 weights = { ptx::exchange(stored_cached_weight, (j * ATOM_M) % 32 + (lane_idx % 4) * 2 + 0), ptx::exchange(stored_cached_weight, (j * ATOM_M) % 32 + (lane_idx % 4) * 2 + 1) }; + if constexpr (kWithL2ActScales) + weights = __fmul2_rn(weights, {l2_act_gs, l2_act_gs}); // Load from TMEM + uint2 raw_values[4]; uint32_t tmem_addr = accum_stage_idx * UMMA_N + epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M; - uint32_t values[ATOM_M]; cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, - values[0], values[1], values[2], values[3]); + raw_values[0].x, raw_values[0].y, raw_values[1].x, raw_values[1].y); cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, - values[4], values[5], values[6], values[7]); + raw_values[2].x, raw_values[2].y, raw_values[3].x, raw_values[3].y); cutlass::arch::fence_view_async_tmem_load(); // Signal tensor memory consumed on the last atom if (j == WG_BLOCK_M / ATOM_M - 1) { ptx::tcgen05_before_thread_sync(); - tmem_empty_barriers[accum_stage_idx]->arrive(0u); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); } - // Apply the explicitly selected SwiGLU or Kimi-K3 SiTU activation. - // Gate/up pairs: (0, 2), (1, 3), (4, 6), (5, 7) + // Apply the selected activation: SwiGLU, OAI SwiGLU, or Kimi-K3 SiTU // SiTU: // act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate) // up' = kSituLinearBeta * tanh(up/kSituLinearBeta) // K3 config constants baked in (activation_situ_{beta,linear_beta}). constexpr float kSituBeta = 4.0f; constexpr float kSituLinearBeta = 25.0f; - auto fp32_values = reinterpret_cast(values); + auto fp32_values = reinterpret_cast(raw_values); #pragma unroll for (uint32_t k = 0; k < 2; ++ k) { - auto bf16_gate = __float22bfloat162_rn(make_float2(fp32_values[k * 4], fp32_values[k * 4 + 1])); - auto bf16_up = __float22bfloat162_rn(make_float2(fp32_values[k * 4 + 2], fp32_values[k * 4 + 3])); + if constexpr (kUseXScales or kWithL1Alphas) { + fp32_values[k * 2 + 0] = __fmul2_rn(fp32_values[k * 2 + 0], gate_scales); + fp32_values[k * 2 + 1] = __fmul2_rn(fp32_values[k * 2 + 1], up_scales); + } + auto bf16_gate = __float22bfloat162_rn(fp32_values[k * 2 + 0]); + auto bf16_up = __float22bfloat162_rn(fp32_values[k * 2 + 1]); // Clamp (SwiGLU-with-limit only; SiTU soft-clips below) - if constexpr (!kUseSitu && kActivationClamp != cute::numeric_limits::infinity()) { + if constexpr (not kUseSitu and kActivationClamp != cute::numeric_limits::infinity()) { bf16_gate = __hmin2(bf16_gate, {kActivationClamp, kActivationClamp}); bf16_up = __hmax2(bf16_up, {-kActivationClamp, -kActivationClamp}); bf16_up = __hmin2(bf16_up, {kActivationClamp, kActivationClamp}); } - // sigmoid(gate) + constexpr bool kIsOAISwiGLU = kSwiGLUAlpha != 0.0f; auto gate = __bfloat1622float2(bf16_gate); + const auto sigmoid_in = kIsOAISwiGLU ? + __fmul2_rn(gate, {kSwiGLUAlpha, kSwiGLUAlpha}) : gate; auto neg_gate_exp = make_float2( - kFastMath ? __expf(-gate.x) : expf(-gate.x), - kFastMath ? __expf(-gate.y) : expf(-gate.y)); + kFastMath ? __expf(-sigmoid_in.x) : expf(-sigmoid_in.x), + kFastMath ? __expf(-sigmoid_in.y) : expf(-sigmoid_in.y)); const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp); float2 sig; if constexpr (kFastMath) { @@ -1419,8 +1259,9 @@ sm100_fp8_fp4_mega_moe_impl(void* y, up = {kSituLinearBeta * tanhf(up.x / kSituLinearBeta), kSituLinearBeta * tanhf(up.y / kSituLinearBeta)}; } else { - // SwiGLU: silu(gate) * up gate = __fmul2_rn(gate, sig); + if constexpr (kIsOAISwiGLU) + up = __fadd2_rn(up, {1.0f, 1.0f}); } activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights); } @@ -1440,128 +1281,78 @@ sm100_fp8_fp4_mega_moe_impl(void* y, thread_local_amax.y, math::ReduceMax()); // Reduce amax (warp-pair-level) - if (lane_idx < 4) - smem_amax_reduction[epilogue_warp_idx * (STORE_BLOCK_M / 2) + i * (ATOM_M / 2) + lane_idx] = amax_values[i]; + if (not kIsNVFP4 and lane_idx < 4) + shared_storage.amax_reduction[epilogue_warp_idx][i * (ATOM_M / 2) + lane_idx] = amax_values[i]; __syncwarp(); } // Wait shared memory release from previous TMA store - // And fence `smem_amax_reduction` + // And fence `shared_storage.amax_reduction` const uint32_t tma_stage_idx = s % kNumTMAStoreStages; ptx::tma_store_wait(); ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); - // Cast to FP8 E4M3 (or FP4 E2M1 under `kUseFp4Acts`) and - // store into shared memory. + // Cast into the L1 output dtype and store into shared memory #pragma unroll for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { - // Reduce amax (warp-pair-level) - const float2 wp_amax = - smem_amax_reduction[(epilogue_warp_idx ^ 1) * (STORE_BLOCK_M / 2) + i * (ATOM_M / 2) + lane_idx % 4]; - amax_values[i].x = cute::max(amax_values[i].x, wp_amax.x); - amax_values[i].y = cute::max(amax_values[i].y, wp_amax.y); - - // Calculate SF (UE8M0 byte; only the finfo divisor differs: - // 1/448 for FP8 E4M3, 1/6 for FP4 E2M1). + // Calculate SF float2 sf, sf_inv; - if constexpr (kUseFp4Acts) { - math::get_e2m1_sf_and_sf_inv(amax_values[i], sf, sf_inv); + uint2 sf_bits; + if constexpr (kIsNVFP4) { + math::get_nvfp4_sf_and_sf_inv(amax_values[i], sf, sf_inv, sf_bits); } else { - math::get_e4m3_sf_and_sf_inv(amax_values[i], sf, sf_inv); + // Reduce amax (warp-pair-level) + const float2 wp_amax = + shared_storage.amax_reduction[epilogue_warp_idx ^ 1][i * (ATOM_M / 2) + lane_idx % 4]; + amax_values[i].x = cute::max(amax_values[i].x, wp_amax.x); + amax_values[i].y = cute::max(amax_values[i].y, wp_amax.y); + if constexpr (kIsMXFP4) + math::get_e2m1_sf_and_sf_inv(amax_values[i], sf, sf_inv); + else + math::get_e4m3_sf_and_sf_inv(amax_values[i], sf, sf_inv); + sf_bits = {*reinterpret_cast(&sf.x) >> 23, + *reinterpret_cast(&sf.y) >> 23}; } - // Apply scale, cast, store into shared memory. - const float2 upper = __fmul2_rn(activation_values[i][0], sf_inv); - const float2 lower = __fmul2_rn(activation_values[i][1], sf_inv); - if constexpr (kUseFp4Acts) { - // FP4 epilogue: write packed E2M1 nibbles to canonical - // dense smem (TMA descriptor built with swizzle=0 → - // byte-exact smem→gmem copy → canonical packed FP4 - // layout `[M, intermediate_hidden/2]` in gmem). - // - // Layout under SwapAB: `tcgen05.ld.16x256b.x1` puts - // lane T's accumulator values (upper.x, upper.y, - // lower.x, lower.y) at smem positions: - // upper.x → row 2*(T%4), col_in_stripe T/4 - // upper.y → row 2*(T%4)+1, col_in_stripe T/4 - // lower.x → row 2*(T%4), col_in_stripe T/4 + 8 - // lower.y → row 2*(T%4)+1, col_in_stripe T/4 + 8 - // (16-byte stripe per warp_idx_in_wg ∈ 0..3, 64 B row.) - // Adjacent N-cols therefore sit on lanes T and T XOR 4, - // so packing two values into one FP4 byte requires a - // `__shfl_xor 4` to pull the buddy. Half-warp gate - // (group = lane/4, group%2==0) means each "active" - // lane writes 4 bytes (upper.x, upper.y, lower.x, - // lower.y) and the inactive half is a donor. - // - // The cross-quad shuffle and half-warp gate are - // structural: they're a consequence of SwapAB's - // datapoint=N orientation. Replacing with - // `tcgen05.ld.32x32b.x8` would require dropping - // SwapAB at the mainloop level. See - // DeepGEMM/FP4_EPILOGUE_STORE_MICROBENCH.md for the - // full microbench analysis (P-A through P-D) and - // the negative results from bank-conflict - // elimination + atom-interleaving. - const float buddy_ux = __shfl_xor_sync(0xffffffffu, upper.x, 4); - const float buddy_uy = __shfl_xor_sync(0xffffffffu, upper.y, 4); - const float buddy_lx = __shfl_xor_sync(0xffffffffu, lower.x, 4); - const float buddy_ly = __shfl_xor_sync(0xffffffffu, lower.y, 4); - - const uint32_t frag = lane_idx % 4; // row-pair index 0..3 - const uint32_t group = lane_idx / 4; // col-group index 0..7 - const bool is_active = (group % 2u) == 0u; - - // Active lanes pack (own_val, buddy_val) into a byte - // (own=low nibble, buddy=high) and write 4 bytes per - // atom. `cvt_pack_f32_to_e2m1x2(a, b)` → {low=a, high=b}. - if (is_active) { - const uint8_t byte_ux = static_cast( - math::cvt_pack_f32_to_e2m1x2(upper.x, buddy_ux)); - const uint8_t byte_uy = static_cast( - math::cvt_pack_f32_to_e2m1x2(upper.y, buddy_uy)); - const uint8_t byte_lx = static_cast( - math::cvt_pack_f32_to_e2m1x2(lower.x, buddy_lx)); - const uint8_t byte_ly = static_cast( - math::cvt_pack_f32_to_e2m1x2(lower.y, buddy_ly)); - - constexpr uint32_t kFp4WarpStripeBytes = 8; // 16 elements / 2 - const uint32_t byte_pos_upper = group / 2u; // 0..3 - const uint32_t byte_pos_lower = 4u + group / 2u; // 4..7 - const uint32_t row_even = i * ATOM_M + 2u * frag; - const uint32_t row_odd = row_even + 1u; - const auto base = smem_cd[tma_stage_idx] - + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_ROW_BYTES - + warp_idx_in_wg * kFp4WarpStripeBytes; - auto write_byte = [&](uint32_t row, uint32_t bp, uint8_t v) { - auto p = base + row * L1_OUT_ROW_BYTES + bp; - asm volatile("st.shared.u8 [%0], %1;\n" - :: "l"(__cvta_generic_to_shared(p)), - "r"(static_cast(v))); - }; - write_byte(row_even, byte_pos_upper, byte_ux); - write_byte(row_odd, byte_pos_upper, byte_uy); - write_byte(row_even, byte_pos_lower, byte_lx); - write_byte(row_odd, byte_pos_lower, byte_ly); + // Cast + const float2 first = __fmul2_rn(activation_values[i][0], sf_inv); + const float2 second = __fmul2_rn(activation_values[i][1], sf_inv); + + const auto smem_base = reinterpret_cast(shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx]) + + i * ATOM_M * L1_OUT_BLOCK_N_BYTES; + if constexpr (kIsFP4Acts) { + const uint32_t packed = math::cast_into_e2m1x2_pairs(first, second); + const uint32_t byte_in_row = warp_idx_in_wg * (L1_OUT_BLOCK_N_BYTES / 4) + lane_idx / 4; + #pragma unroll + for (uint32_t t = 0; t < 2; ++ t) { + const uint32_t row = (lane_idx % 4) * 2 + t; + // 32B swizzle: 2 bank groups, XOR-ed by `row / 4` + const uint32_t swizzled_byte = + ((byte_in_row / kNumBankGroupBytes) ^ ((row / 4) & 1)) * kNumBankGroupBytes + + byte_in_row % kNumBankGroupBytes; + smem_base[row * L1_OUT_BLOCK_N_BYTES + swizzled_byte] = + static_cast(packed >> (t * 8)); } } else { - const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(upper.x, upper.y, lower.x, lower.y)); + const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(first.x, first.y, second.x, second.y)); // STSM uint32_t row = lane_idx; uint32_t col = warp_idx_in_wg; - const auto smem_ptr = smem_cd[tma_stage_idx] + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_BLOCK_N - + i * ATOM_M * L1_OUT_BLOCK_N - + row * L1_OUT_BLOCK_N - + (col ^ (row / 2)) * kNumBankGroupBytes; + const auto smem_ptr = smem_base + + row * L1_OUT_BLOCK_N_BYTES + // Use 64B swizzle for SwiGLU, so divided by 2 + + (col ^ (row / 2)) * kNumBankGroupBytes; ptx::SM100_U8x4_STSM_T<__nv_fp8x4_e4m3>::copy(fp8x4_values, smem_ptr); } - // Store SF to `buffer.l2_sf_buffer` as UE8M0 (MN-major layout) - // Only one warp per pair writes (both hold the same SF after cross-warp reduce) + // Store SF to `buffer.l2_sf_buffer` (MN-major layout) + // For MXFP8FP4 only one warp per pair writes (both hold the same SF after the + // cross-warp reduce); for NVFP4 every warp owns its own 16-element group // Each lane < 4 holds SF for 2 rows (sf.x and sf.y) - if (warp_idx_in_wg % 2 == 0 and lane_idx < 4) { - const uint32_t k_idx = n_block_idx * 2 + warp_idx_in_wg / 2; + if ((kIsNVFP4 or warp_idx_in_wg % 2 == 0) and lane_idx < 4) { + const uint32_t k_idx = kIsNVFP4 ? n_block_idx * 4 + warp_idx_in_wg + : n_block_idx * 2 + warp_idx_in_wg / 2; const uint32_t k_uint_idx = k_idx / 4, byte_idx = k_idx % 4; const uint32_t mn_stride = (task_info.is_shared() ? kNumSharedSFTokens : kNumSFRingTokens) * sizeof(uint32_t); const auto sf_base_ptr = task_info.is_shared() ? @@ -1579,31 +1370,22 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const auto sf_token_idx = block_idx * SF_BLOCK_M + transform_sf_token_idx(token_base_idx) + (lane_idx * 2) * 4; const auto sf_addr = k_uint_idx * mn_stride + sf_token_idx * static_cast(sizeof(uint32_t)) + byte_idx; - sf_base_ptr[sf_addr] = - (*reinterpret_cast(&sf.x) >> 23); + sf_base_ptr[sf_addr] = static_cast(sf_bits.x); sf_base_ptr[sf_addr + 4 * static_cast(sizeof(uint32_t))] = - (*reinterpret_cast(&sf.y) >> 23); + static_cast(sf_bits.y); } __syncwarp(); } ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); - // Issue TMA store after all atoms in this store block. - // FP8 path: out_n in elements-of-FP8 (= bytes), smem - // base offset by FP8 row width (L1_OUT_BLOCK_N). - // FP4 path: TMA descriptor's element type is uint8 with - // half the inner dim → out_n in packed bytes (= - // L1_OUT_BLOCK_N / 2), smem base offset by - // L1_OUT_ROW_BYTES = L1_OUT_BLOCK_N / 2 bytes. + // Issue TMA store after all atoms in this store block if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { - const uint32_t out_n_idx = kUseFp4Acts - ? (n_block_idx * (L1_OUT_BLOCK_N / 2)) - : (n_block_idx * L1_OUT_BLOCK_N); + uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N_BYTES; const auto tensor_map_l1_output_ptr = task_info.is_shared() ? &tensor_map_shared_l1_output : &tensor_map_l1_output; cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( tensor_map_l1_output_ptr, - smem_cd[tma_stage_idx] + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_ROW_BYTES, + shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx], out_n_idx, m_idx + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M); cute::tma_store_arrive(); @@ -1652,6 +1434,12 @@ sm100_fp8_fp4_mega_moe_impl(void* y, DG_STATIC_ASSERT(STORE_BLOCK_M % 8 == 0, "Invalid store M"); constexpr uint32_t kNumRowsPerWarp = STORE_BLOCK_M / 8; + float l2_alpha = 1.0f; + if constexpr (kWithL2Alphas) { + if (not task_info.is_shared()) + l2_alpha = __ldg(l2_alphas + task_info.local_expert_idx); + } + // L2 BF16 epilogue: write GEMM output to remote combine buffer via NVLink #pragma unroll for (uint32_t s = 0; s < WG_BLOCK_M / STORE_BLOCK_M; ++ s) { @@ -1659,7 +1447,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // TODO: check performance if (epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M >= valid_m) { ptx::tcgen05_before_thread_sync(); - tmem_empty_barriers[accum_stage_idx]->arrive(0u); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); break; } @@ -1675,6 +1463,13 @@ sm100_fp8_fp4_mega_moe_impl(void* y, values[4], values[5], values[6], values[7]); cutlass::arch::fence_view_async_tmem_load(); + if constexpr (kWithL2Alphas) { + auto fp32_values = reinterpret_cast(values); + #pragma unroll + for (uint32_t v = 0; v < ATOM_M; ++ v) + fp32_values[v] *= l2_alpha; + } + // Wait shared memory release from previous NVLink store // NOTES: skip for the first store block since the prior full barrier already ensures completion if (i == 0 and s > 0) @@ -1683,16 +1478,14 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Signal tensor memory consumed if (s == WG_BLOCK_M / STORE_BLOCK_M - 1 and i == STORE_BLOCK_M / ATOM_M - 1) { ptx::tcgen05_before_thread_sync(); - tmem_empty_barriers[accum_stage_idx]->arrive(0u); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); } // Store into shared memory - // NOTES: only use first 16 lanes for address - // NOTES: 2 warps share a BF16 swizzle atom + // NOTES: each lane provides its own address for stmatrix; 2 warps share a BF16 swizzle atom uint32_t row = lane_idx % 8; uint32_t col = (epilogue_warp_idx % 2) * 4 + lane_idx / 8; - const auto smem_ptr = smem_cd_l2 + - epilogue_wg_idx * STORE_BLOCK_M * BLOCK_N * static_cast(sizeof(nv_bfloat16)) + + const auto smem_ptr = reinterpret_cast(shared_storage.smem_d.l2[epilogue_wg_idx]) + (warp_idx_in_wg / 2) * STORE_BLOCK_M * kSwizzleCDMode + i * ATOM_M * kSwizzleCDMode + row * (kNumBankGroupBytes * 8) + @@ -1710,7 +1503,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); // Write into remote buffers - // One warp per row, now the layout is different from shared memory storing + // Each warp writes 2 rows (lane_idx/16 splits the warp into two halves, one per row) const uint32_t row_in_atom = (warp_idx_in_wg * 2 + lane_idx / 16) % ATOM_M; const uint32_t bank_group_idx = lane_idx % 8; @@ -1736,46 +1529,39 @@ sm100_fp8_fp4_mega_moe_impl(void* y, } // Read from shared memory - const auto smem_ptr = smem_cd_l2 + - epilogue_wg_idx * STORE_BLOCK_M * BLOCK_N * static_cast(sizeof(nv_bfloat16)) + + const auto smem_ptr = reinterpret_cast(shared_storage.smem_d.l2[epilogue_wg_idx]) + (lane_idx % 16 / 8) * STORE_BLOCK_M * kSwizzleCDMode + row_in_store * kSwizzleCDMode + (bank_group_idx ^ row_in_atom) * kNumBankGroupBytes; const auto packed = ptx::ld_shared(reinterpret_cast(smem_ptr)); + // Write into remote if constexpr (kUseFp8Combine) { - // Stream B: BF16 (in `packed`) → FP8 E4M3 + per-row UE8M0 SF. - // - // 16 lanes (lane_idx & ~15u) cover one row's - // BLOCK_N=128 elements (= 8 BF16 each). Compute - // per-row amax via warp_reduce over those 16 - // lanes, then quantize. + // BF16 (in `packed`) -> FP8 E4M3 + one UE8M0 SF per row tile. + // 16 lanes (lane_idx & ~15u) cover one row's BLOCK_N=128 + // elements (8 BF16 each), so the per-row amax is a reduction + // over that 16-lane group. const auto bf_pairs = reinterpret_cast(&packed); float local_amax = 0.0f; #pragma unroll - for (int q = 0; q < 4; ++q) { + for (uint32_t q = 0; q < 4; ++q) { const float2 vf = __bfloat1622float2(bf_pairs[q]); local_amax = cute::max(local_amax, cute::abs(vf.x)); local_amax = cute::max(local_amax, cute::abs(vf.y)); } - // Reduce within the 16-lane group sharing this row. - // Use a 16-lane mask (NOT 0xffffffff) because the - // outer `if (m_idx_in_block >= valid_m) break` may - // cause the OTHER half-warp's 16 lanes to exit - // early on padding rows. A full-warp shfl would - // deadlock waiting on those exited lanes. - const uint32_t row_mask = 0x0000FFFFu << (16u * (lane_idx / 16)); + // A 16-lane mask, not `0xffffffff`: the `m_idx_in_block >= valid_m` + // break above can retire the other half-warp on padding rows, and a + // full-warp shuffle would then wait on lanes that have already exited. + const uint32_t row_mask = 0x0000ffffu << (16u * (lane_idx / 16)); local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 1)); local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 2)); local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 4)); local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 8)); - // UE8M0 SF (E4M3, finfo_max = 448). const int log2_ceil = math::fast_log2_ceil(local_amax * (1.0f / 448.0f)); const float sf_inv = math::fast_pow2(-log2_ceil); const uint8_t sf_byte = static_cast(log2_ceil + 127); - // Scale, cast 4 BF16 pairs → 8 FP8 (= 2 fp8x4 = uint64). float4 lo, hi; const auto lo_pair = __bfloat1622float2(bf_pairs[0]); const auto lo_pair_b = __bfloat1622float2(bf_pairs[1]); @@ -1798,7 +1584,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Write 8 FP8 bytes (uint64) to remote, replacing // the BF16 16-byte write. const auto dst_token = buffer.combine_token_buffer.get_rank_buffer(dst_topk_idx) - .get_data_buffer(dst_token_idx); + .get_data_buffer(dst_token_idx); const auto dst_ptr = math::advance_ptr( dst_token.get_base_ptr(), n_idx * static_cast(sizeof(uint8_t)) + @@ -1808,7 +1594,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // 1 SF byte per row tile, written by lane 0 of the 16-lane group. if ((lane_idx & 15u) == 0) { const auto sf_token = buffer.combine_sf_buffer.get_rank_buffer(dst_topk_idx) - .get_data_buffer(dst_token_idx); + .get_data_buffer(dst_token_idx); const auto sf_ptr = math::advance_ptr( sf_token.get_base_ptr(), n_block_idx * static_cast(sizeof(uint8_t))); @@ -1817,7 +1603,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, } else { // Default BF16 path (16 bytes/lane = 8 BF16). const auto dst_token = buffer.combine_token_buffer.get_rank_buffer(dst_topk_idx) - .get_data_buffer(dst_token_idx); + .get_data_buffer(dst_token_idx); const auto dst_ptr = math::advance_ptr( dst_token.get_base_ptr(), n_idx * static_cast(sizeof(nv_bfloat16)) + (lane_idx % 16) * static_cast(sizeof(float4))); @@ -1859,20 +1645,19 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // NOTES: either 1 or 2 chunks for simplicity // NOTES: Restrict on both smem and register constexpr uint32_t kNumChunks = - kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes <= SMEM_BEFORE_BARRIER_SIZE and kHidden <= 32 * kNumMaxRegistersForBuffer ? 1 : 2; + kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes <= kNumReusableSmemBytes and kHidden <= 32 * kNumMaxRegistersForBuffer ? 1 : 2; constexpr uint32_t kNumChunkBytes = kNumHiddenBytes / kNumChunks; constexpr uint32_t kNumChunkUint4 = kNumChunkBytes / sizeof(uint4); constexpr uint32_t kNumUint4PerLane = kNumChunkUint4 / 32; DG_STATIC_ASSERT(kHidden % kNumChunks == 0, "Hidden must be divisible by number of chunks"); - DG_STATIC_ASSERT(kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes / kNumChunks <= SMEM_BEFORE_BARRIER_SIZE, "Hidden is too large"); + DG_STATIC_ASSERT(kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes / kNumChunks <= kNumReusableSmemBytes, "Hidden is too large"); DG_STATIC_ASSERT(kNumChunkBytes % 16 == 0, "Combine chunk must be TMA-aligned (16 bytes)"); DG_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, "Combine chunk must be divisible by 16 bytes"); DG_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, "Combine chunk must be a multiple of 32 16-byte elements (one per lane)"); DG_STATIC_ASSERT(kNumTopk + (kNumSharedExperts > 0 ? 1u : 0u) <= 32u, "Top-k + shared must fit in a single warp"); // Verify combined shared memory budget at runtime - DG_DEVICE_ASSERT(kNumChunkSlots * kNumEpilogueWarps * kNumChunkBytes <= static_cast( - reinterpret_cast(barrier_start_ptr) - smem_buffer)); + DG_DEVICE_ASSERT(kNumChunkSlots * kNumEpilogueWarps * kNumChunkBytes <= kNumReusableSmemBytes); // Per-warp buffer: 2 stage load buffers + 1 store buffer const auto combine_load_buffer = utils::PatternVisitor([&](const uint32_t& i) { @@ -1882,7 +1667,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Per-warp barriers auto combine_load_barriers = utils::PatternVisitor([&](const uint32_t& i) { - return combine_barriers[i + epilogue_warp_idx * 2]; + return &shared_storage.combine_barriers[i + epilogue_warp_idx * 2]; }); // Iterate over all tokens @@ -1908,7 +1693,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Per-uint4 load: BF16 → 8 BF16 = 4 float2 pairs. // FP8 → 16 FP8 = 8 float2 pairs (dequant'd). constexpr uint32_t kNumF32PairsPerLoadUint4 = - kUseFp8Combine ? 8u : 4u; + kUseFp8Combine ? 8u : 4u; // Per-element offset in the chunk for SF lookup: // sf_idx = (chunk * kNumLoadElemsPerChunk + elem_in_chunk) / 128 constexpr uint32_t kNumLoadElemsPerChunk = kHidden / kNumChunks; diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh index 303a4db74b..4c87ede3a6 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh @@ -5,23 +5,27 @@ #include #include +#include #include namespace deep_gemm { // Fused BF16 → quant + topk copy + pad-fill kernel that produces the exact // byte layout DeepGEMM's mega-MoE symmetric buffer expects in its `x`, -// `x_sf`, `topk_idx`, and `topk_weights` slots. Two variants: +// `x_sf`, `topk_idx`, and `topk_weights` slots. The acts format follows the +// consuming GEMM's `MmaKind`: // -// - `kUseFp4Acts == false` → FP8 (E4M3) acts; per-row stride = `hidden`. -// - `kUseFp4Acts == true` → packed FP4 (E2M1) acts; per-row stride -// = `hidden / 2`. Layout: byte holds 2 nibbles, -// low nibble = even col, high nibble = odd col, -// matching `deep_gemm.utils.per_token_cast_to_fp4`. +// - `MXFP8FP4` → FP8 (E4M3) acts + UE8M0 group scales; row stride = `hidden`. +// - `MXFP4` → packed FP4 (E2M1) acts + UE8M0 group scales; row stride +// = `hidden / 2`. Layout: byte holds 2 nibbles, low nibble = +// even col, high nibble = odd col, matching +// `deep_gemm.utils.per_token_cast_to_fp4`. +// - `NVFP4` → same packing as `MXFP4`, but UE4M3 SFs over 16-element +// groups plus one FP32 per-token outer scale in `buf_x_scales`. // -// Both paths share the UE8M0 SF byte layout: `byte_off = token*num_groups + -// group`, with the contiguous `(P, num_groups/4)` int32 slot storing 4 bytes -// per int32 in row-major order. +// All paths share the SF byte layout: `byte_off = token*num_groups + group`, +// with the contiguous `(P, num_groups/4)` int32 slot storing 4 bytes per int32 +// in row-major order. // // The FP4 quant matches `per_token_cast_to_fp4` (host helper) bytewise via // explicit bucketize boundaries — PTX `cvt.rn.satfinite.e2m1x2.f32` rounds @@ -55,14 +59,29 @@ __forceinline__ __device__ uint32_t pre_dispatch_e2m1_encode(float v) { return code; } -template +__forceinline__ __device__ float pre_dispatch_block_reduce_max(float val) { + __shared__ float smem[32]; + const uint32_t lane = threadIdx.x % 32u; + const uint32_t warp = threadIdx.x / 32u; + const uint32_t num_warps = (blockDim.x + 31u) / 32u; + val = math::warp_reduce<32, /*kIntergroupReduce=*/false>(val, math::ReduceMax{}); + if (lane == 0) + smem[warp] = val; + __syncthreads(); + val = (lane < num_warps) ? smem[lane] : 0.0f; + return math::warp_reduce<32, /*kIntergroupReduce=*/false>(val, math::ReduceMax{}); +} + +template __launch_bounds__(1024, 2) __global__ void mega_moe_pre_dispatch_kernel( const __nv_bfloat16* __restrict__ x, const int32_t* __restrict__ topk_idx, const float* __restrict__ topk_weights, + const float* __restrict__ expert_scales, void* __restrict__ buf_x, int32_t* __restrict__ buf_x_sf, + float* __restrict__ buf_x_scales, int64_t* __restrict__ buf_topk_idx, float* __restrict__ buf_topk_weights, const uint32_t num_tokens, @@ -70,8 +89,16 @@ __global__ void mega_moe_pre_dispatch_kernel( const uint32_t hidden, const uint32_t num_groups, const uint32_t top_k) { - static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, - "kGroupSize must be 32, 64, or 128"); + static_assert(kGroupSize == 16 || kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, + "kGroupSize must be 16, 32, 64, or 128"); + constexpr bool kNvfp4 = kMmaKind == MmaKind::NVFP4; + constexpr bool kPackedFp4 = get_element_bits(kMmaKind) == 4; + static_assert(kMmaKind == MmaKind::MXFP8FP4 || kMmaKind == MmaKind::MXFP4 || kNvfp4, + "Acts must be FP8 (MXFP8FP4) or packed FP4 (MXFP4 / NVFP4)"); + static_assert(!kNvfp4 || kGroupSize == 16, + "NVFP4 acts use 16-element UE4M3 blocks"); + static_assert(kNvfp4 || kGroupSize >= 32, + "UE8M0 modes keep the original 32/64/128 group sizes"); constexpr uint32_t kVecElems = 8; // 16-byte BF16 load per thread static_assert(kGroupSize % kVecElems == 0, "kGroupSize must be a multiple of 8"); constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems; @@ -110,17 +137,39 @@ __global__ void mega_moe_pre_dispatch_kernel( local_max = math::warp_reduce( local_max, math::ReduceMax{}); - // Match host `per_token_cast_to_fp4/fp8`: clamp absmax to 1e-4 - // before dividing by the dtype's max representable value. - const float absmax = fmaxf(local_max, 1e-4f); - constexpr float kFinfoMax = kUseFp4Acts ? 6.0f : 448.0f; - const float raw_scale = absmax / kFinfoMax; - const uint32_t ue8m0_exp = pre_dispatch_cast_to_ue8m0(raw_scale); - // 1 / 2^(ue8m0_exp - 127) = 2^(127 - ue8m0_exp); fp32 bits = - // (127 - ue8m0_exp + 127) << 23 = (254 - ue8m0_exp) << 23. - const float inv_scale = __uint_as_float((254u - ue8m0_exp) << 23u); - - if constexpr (kUseFp4Acts) { + float inv_scale; + uint32_t sf_byte; + if constexpr (kNvfp4) { + // NVFP4 per-token recipe (matches the flashinfer per-token + // quantizer's dequant contract): x ≈ e2m1 * ue4m3(sf) * pts with + // pts = row_absmax / (448 * 6) + // sf = e4m3_rn(group_absmax / (6 * pts)) (≤ 448 by construction) + // Encode divides by the ROUNDED sf so dequant is exact w.r.t. it. + const float row_max = fmaxf( + pre_dispatch_block_reduce_max(local_max), 1e-4f); + const float pts = row_max / (448.0f * 6.0f); + const __nv_fp8_storage_t sf_fp8 = __nv_cvt_float_to_fp8( + local_max / (6.0f * pts), __NV_SATFINITE, __NV_E4M3); + const __half_raw sf_hr = __nv_cvt_fp8_to_halfraw(sf_fp8, __NV_E4M3); + const float sf_dec = __half2float(*reinterpret_cast(&sf_hr)); + inv_scale = sf_dec > 0.0f ? 1.0f / (sf_dec * pts) : 0.0f; + sf_byte = static_cast(sf_fp8); + if (tid == 0u) + buf_x_scales[token_id] = pts; + } else { + // Match host `per_token_cast_to_fp4/fp8`: clamp absmax to 1e-4 + // before dividing by the dtype's max representable value. + const float absmax = fmaxf(local_max, 1e-4f); + constexpr float kFinfoMax = kPackedFp4 ? 6.0f : 448.0f; + const float raw_scale = absmax / kFinfoMax; + const uint32_t ue8m0_exp = pre_dispatch_cast_to_ue8m0(raw_scale); + // 1 / 2^(ue8m0_exp - 127) = 2^(127 - ue8m0_exp); fp32 bits = + // (127 - ue8m0_exp + 127) << 23 = (254 - ue8m0_exp) << 23. + sf_byte = ue8m0_exp; + inv_scale = __uint_as_float((254u - sf_byte) << 23u); + } + + if constexpr (kPackedFp4) { // 8 BF16 → 4 packed nibbles → 4 bytes (uint32_t). Output stride // per token is hidden/2; thread tid writes 4 bytes at offset // [tid*4, tid*4+4) in the output row. Pairing matches host @@ -153,23 +202,28 @@ __global__ void mega_moe_pre_dispatch_kernel( row_out[tid] = packed; } - // One thread per group writes its UE8M0 exponent byte. Row-major - // contiguous layout into `buf_x_sf` viewed as bytes: - // byte_off = token_id * num_groups + group_id. + // One thread per group writes its scale byte (UE8M0 exponent, or the + // UE4M3 block SF for NVFP4 acts). Row-major contiguous layout into + // `buf_x_sf` viewed as bytes: byte_off = token_id * num_groups + group_id. const uint32_t group_id = tid / kThreadsPerGroup; const uint32_t within_group_id = tid % kThreadsPerGroup; if (within_group_id == 0u && group_id < num_groups) { const uint32_t byte_off = token_id * num_groups + group_id; reinterpret_cast(buf_x_sf)[byte_off] = - static_cast(ue8m0_exp); + static_cast(sf_byte); } - // Copy this token's topk row. top_k is small (≤ num_threads enforced - // at host); each tid(topk_idx[off]); - buf_topk_weights[off] = topk_weights[off]; + const int32_t expert = topk_idx[off]; + buf_topk_idx[off] = static_cast(expert); + float w = topk_weights[off]; + if (expert_scales != nullptr && expert >= 0) + w *= expert_scales[expert]; + buf_topk_weights[off] = w; } } else { // ---- Pad path: trailing CTAs fill [num_tokens, padded_max) topk 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 02c623a2fd..871824c3ab 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 @@ -27,34 +27,6 @@ namespace deep_gemm { -template -__forceinline__ __device__ float sm90_fp8_mega_moe_clamp_gate(float x) { - if constexpr (kActivationClamp != cute::numeric_limits::infinity()) - x = cute::min(x, kActivationClamp); - return x; -} - -template -__forceinline__ __device__ float sm90_fp8_mega_moe_clamp_up(float x) { - if constexpr (kActivationClamp != cute::numeric_limits::infinity()) - x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); - return x; -} - -template -__forceinline__ __device__ float sm90_fp8_mega_moe_silu(float x) { - const float e = kFastMath ? __expf(-x) : expf(-x); - const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); - return x * sig; -} - -template -__forceinline__ __device__ float sm90_fp8_mega_moe_swiglu(float g, float u) { - g = sm90_fp8_mega_moe_clamp_gate(g); - u = sm90_fp8_mega_moe_clamp_up(u); - return sm90_fp8_mega_moe_silu(g) * u; -} - // Continuous FP32 activation scale. SM90 WGMMA has no hardware block-scale operand (the SF // is a plain FFMA in the epilogue), so the previous UE8M0 (power-of-two) scale bought nothing // on SM90 and only cost precision; the SF pool is already fp32, so this is byte/layout neutral. @@ -68,50 +40,15 @@ __forceinline__ __device__ void sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( sf.y = __fmul_rn(ay, kScale), sf_inv.y = 1.0f / sf.y; } -template -CUTLASS_DEVICE void sm90_fp8_mega_moe_for_each_block_split( - sched::LegacyMegaMoEScheduler& scheduler, - L1Func&& l1_func, L2Func&& l2_func) { - scheduler.fetch_expert_recv_count(); - scheduler.set_expert_idx(0); - - while (true) { - CUTE_TIE_DECL(scheduler.get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx); - if (block_phase == sched::BlockPhase::None) - break; - - if (block_phase == sched::BlockPhase::Linear1) { - l1_func(current_local_expert_idx, kNumL1BlockKs, m_block_idx, n_block_idx); - } else { - l2_func(current_local_expert_idx, kNumL2BlockKs, m_block_idx, n_block_idx); - } - } -} - // ============================================================================ // SM90 (Hopper) FP8 MegaMoE — full implementation // ---------------------------------------------------------------------------- // Pipeline (cluster=1, no TMA multicast): // * Dispatch warps: pull tokens (FP8) and SF (per-128 channel float) from // remote ranks via NVLink into the local L1 pool. +// * Producer warp: claims L1/L2 tasks (routed, plus SharedLinear1/2 when the +// shared expert is fused in) from global atomic counters and publishes them +// into a 2-stage SMEM task ring that every consumer warp drains in order. // * GEMM TMA-load warps (1 for A+SFA, 1 for B+SFB) feed the pipeline stages. // * Math warpgroups (totalling kNumEpilogueThreads) consume each // stage with WGMMA, accumulate into registers, then run the epilogue: @@ -131,7 +68,6 @@ template < uint32_t kNumMaxTokensPerRank, uint32_t kHidden, uint32_t kIntermediateHidden, uint32_t kNumExperts, uint32_t kNumTopk, - uint32_t kNumExpertsPerWave, uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, uint32_t kNumMaxPoolTokens, uint32_t kNumPaddedSFPoolTokens, @@ -145,8 +81,8 @@ template < bool kReuseAccumAsFinal, bool kL2ArrivalCounter, bool kL2EpilogueRequiresFullSync, - bool kSplitPhaseHotPath, bool kFP8SwapAB = false, + uint32_t kNumSharedExperts = 0, uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, uint32_t L2_SHAPE_N = kHidden, @@ -172,17 +108,34 @@ sm90_fp8_mega_moe_impl(void* y, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, - const float* __restrict__ l2_weights_sf) { + const float* __restrict__ l2_weights_sf, + // Fused shared expert (only read when `kNumSharedExperts > 0`; + // otherwise the host passes the routed descriptors as placeholders). + // Shared L1 acts SF has no descriptor: `x_sf` is K-major, so a + // (BLOCK_M, 1) box is illegal and the loader warp gathers the column + // into `smem_sfa` itself. Shared L2 acts SF is written M-major by the + // fused L1 epilogue, so it TMA-loads like the routed L2 SFA. + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_acts, + 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_weights, + const float* __restrict__ shared_l2_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l2_acts_sf) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; // ===================================================================== // Template checks // ===================================================================== - DG_STATIC_ASSERT(kNumDispatchThreads >= 64 and kNumDispatchThreads % 64 == 0, + DG_STATIC_ASSERT(kNumDispatchThreads >= 32 and kNumDispatchThreads % 32 == 0, "Invalid number of dispatch threads"); - DG_STATIC_ASSERT(kNumNonEpilogueThreads == 64 or kNumNonEpilogueThreads == 128, + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 64 or kNumNonEpilogueThreads == 96 or + kNumNonEpilogueThreads == 128 or kNumNonEpilogueThreads == 192, "Invalid number of GEMM TMA warps"); + DG_STATIC_ASSERT(kNumMMANonEpilogueWarps >= 3, + "The scheduler needs a dedicated producer warp"); DG_STATIC_ASSERT((kNumDispatchThreads + kNumNonEpilogueThreads) % 128 == 0, "Math warpgroup start must be 128-thread aligned"); DG_STATIC_ASSERT(kNumEpilogueThreads % 128 == 0, "Invalid number of math/epilogue threads"); @@ -209,6 +162,14 @@ sm90_fp8_mega_moe_impl(void* y, 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 (kNumSharedExperts > 0) { + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_acts); + 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_weights); + cute::prefetch_tma_descriptor(&tensor_map_shared_l2_acts_sf); + } } // ===================================================================== @@ -221,6 +182,26 @@ sm90_fp8_mega_moe_impl(void* y, DG_STATIC_ASSERT(kNumPaddedSFPoolTokens >= kNumPoolBlocks * SF_BLOCK_M, "Invalid SM90 MegaMoE SF pool capacity"); + // Fused shared expert. `M` is this rank's own token count, the intermediate size + // is scaled by the number of shared experts, and there is no expert dimension. + constexpr bool kHasSharedExperts = kNumSharedExperts > 0; + constexpr uint32_t kSharedIntermediateHidden = kIntermediateHidden * kNumSharedExperts; + constexpr uint32_t SHARED_L1_SHAPE_N = kSharedIntermediateHidden * 2; + constexpr uint32_t SHARED_L1_SHAPE_K = kHidden; + constexpr uint32_t SHARED_L2_SHAPE_N = kHidden; + constexpr uint32_t SHARED_L2_SHAPE_K = kSharedIntermediateHidden; + // The shared activation SF buffers are plain K-major (see the loader): their row + // strides must stay 16-byte aligned for TMA, which `layout::Data` already enforces. + DG_STATIC_ASSERT(not kHasSharedExperts or kSharedIntermediateHidden % 256 == 0, + "Shared intermediate hidden must be a multiple of 256"); + // swapAB's reduced-output tile shares the token M-axis the shared expert runs on, + // so the two compose: swapAB owns the SwiGLU/quantize/store epilogue, and the + // shared path merely selects its own weight/SF/output descriptors via `is_shared` + // ternaries (see `run_swap_ab_l1`/`run_swap_ab_l2` and the swapAB L1 epilogue). + DG_STATIC_ASSERT(not (kHasSharedExperts and kFP8SwapAB) or + (kSharedIntermediateHidden * 2) % BLOCK_N == 0, + "swapAB + shared expert requires the shared L1 N to tile evenly"); + const auto workspace = layout::SM90Workspace( sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk); @@ -250,8 +231,26 @@ sm90_fp8_mega_moe_impl(void* y, const auto l2_token_buffer = layout::Buffer(fp8_intermediate_token_layout, 1, kNumMaxPoolTokens, l1_topk_weights_buffer.get_end_ptr()); const auto l2_sf_buffer = layout::Buffer(fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens, l2_token_buffer.get_end_ptr()); - // Combine input area - const auto combine_token_buffer = layout::Buffer(bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, l2_sf_buffer.get_end_ptr()); + // Combine input area. The fused shared expert reduces through one extra slot + // (`topk_idx == kNumTopk`) written by the local rank only. + constexpr uint32_t kNumCombineSlots = kNumTopk + (kHasSharedExperts ? 1u : 0u); + const auto combine_token_buffer = layout::Buffer(bf16_token_layout, kNumCombineSlots, kNumMaxTokensPerRank, l2_sf_buffer.get_end_ptr()); + + // Fused shared-expert area, appended after the combine buffer so the routed + // regions keep their relative order and are zero-sized when the shared expert is + // disabled (the workspace itself always reserves the shared arrival counters, so + // absolute offsets shift by a few KB either way -- host and device agree because + // both derive them from the same `SM90Workspace`). + // The post-SwiGLU FP8 output and its per-64-K float SF are indexed by the local + // token index; the SF buffer is K-major (no SF-pool padding needed). + constexpr auto fp8_shared_intermediate_token_layout = layout::Data(kSharedIntermediateHidden); + constexpr auto fp8_shared_intermediate_sf_layout = layout::Data(kSharedIntermediateHidden / 16); + const auto shared_l2_token_buffer = layout::Buffer( + fp8_shared_intermediate_token_layout, 1, kHasSharedExperts ? kNumMaxTokensPerRank : 0, + combine_token_buffer.get_end_ptr()); + const auto shared_l2_sf_buffer = layout::Buffer( + fp8_shared_intermediate_sf_layout, 1, kHasSharedExperts ? kNumMaxTokensPerRank : 0, + shared_l2_token_buffer.get_end_ptr()); // ===================================================================== // GEMM data types and shape constants @@ -273,6 +272,7 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t kNumCombineWarps = kNumEpilogueWarps; using L1WGMMA = typename mma::sm90::FP8MMASelector::type; // M=64, N=WG_BLOCK_N, K=32 using L2WGMMA = typename mma::sm90::FP8MMASelector::type; + using SwapWGMMA64 = typename mma::sm90::FP8MMASelector<64>::type; constexpr uint32_t kL1OutputArrivalParts = 1; static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, "Unexpected WGMMA shape"); @@ -294,7 +294,7 @@ sm90_fp8_mega_moe_impl(void* y, // feeds that shared SF must be reduced across both warpgroups. constexpr bool kSplitNSharesSF = kSplitNWarpgroups and (WG_L1_OUT_BLOCK_N < 64); constexpr bool kSwapABEligible = - kFP8SwapAB and kSplitNWarpgroups and (BLOCK_M == 64) and (BLOCK_N == 128) and + kFP8SwapAB and kSplitNWarpgroups and (BLOCK_M == 64) and (BLOCK_N == 256) and (kWarpgroupSplitN == 2); constexpr bool kSwapABActive = kSwapABEligible; constexpr uint32_t kSwapABTokenChunks = BLOCK_M / 8; @@ -383,6 +383,36 @@ sm90_fp8_mega_moe_impl(void* y, 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; }); + // Interleaved-scheduler task ring: 2-stage task-info slots with dedicated + // full/empty barriers, placed right after the combine barriers (the host + // SMEM accounting reserves the same bytes, see `sm90_mega_moe.hpp`). + // SM90 reuses the unified interleaved `MegaMoEScheduler` in single-CTA mode + // (`kClusterSize = 1`): one task covers one CTA tile and the producer warp + // publishes tasks to local SMEM (no 2-CTA cluster, no ring bookkeeping -- + // `kNumRingBlocks` is assert-only on this path, so it is fed the pool block + // count). `TaskInfo` is alignas(16) / 32 B, but barrier slots are 8 B, so + // the barrier count preceding the ring must be even for the 32 B slots to + // land 16-byte aligned. Pad one unused barrier slot when it is odd -- this + // fires on the 1-dispatch-warp topology (parity is set by kNumDispatchWarps + // alone; 2*kNumStages and 2*kNumCombineWarps are even). + constexpr uint32_t kTaskInfoBaseBarriers = + kNumDispatchWarps + kNumStages * 2 + kNumCombineWarps * 2; + constexpr uint32_t kTaskInfoBarrierPad = kTaskInfoBaseBarriers & 1u; + using SchedulerT = sched::MegaMoEScheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, + kNumSMs, kNumRanks, + kNumPoolBlocks, + kNumSharedExperts, + /*kClusterSize=*/1, + layout::SM90Workspace>; + auto task_info_full_barriers = barrier_start_ptr + kTaskInfoBaseBarriers + kTaskInfoBarrierPad; + auto task_info_empty_barriers = task_info_full_barriers + SchedulerT::kNumScheduleStages; + auto task_infos = reinterpret_cast( + task_info_empty_barriers + SchedulerT::kNumScheduleStages); + // ===================================================================== // Initialization // ===================================================================== @@ -411,6 +441,13 @@ sm90_fp8_mega_moe_impl(void* y, #pragma unroll for (uint32_t i = 0; i < kNumCombineWarps * 2; ++ i) combine_barriers[i]->init(1); + #pragma unroll + for (uint32_t i = 0; i < SchedulerT::kNumScheduleStages; ++ i) { + // The producer warp publishes one task per slot + task_info_full_barriers[i].init(1); + // TMA-A + TMA-B warps and every math warp release each slot once + task_info_empty_barriers[i].init(2 + kNumEpilogueWarps); + } } cutlass::arch::fence_barrier_init(); } @@ -419,20 +456,11 @@ sm90_fp8_mega_moe_impl(void* y, // ===================================================================== // Scheduler (cluster=1) // ===================================================================== - constexpr uint32_t kNumExpertsPerLane = math::constexpr_ceil_div(kNumExpertsPerRank, 32u); - constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; - constexpr uint32_t kNumL2BlockNs = L2_SHAPE_N / BLOCK_N; - constexpr uint32_t kNumL1BlockKs = L1_SHAPE_K / BLOCK_K; - constexpr uint32_t kNumL2BlockKs = L2_SHAPE_K / BLOCK_K; - auto scheduler = sched::LegacyMegaMoEScheduler< - BLOCK_M, BLOCK_N, BLOCK_K, - L1_SHAPE_N, L1_SHAPE_K, - L2_SHAPE_N, L2_SHAPE_K, - kNumExpertsPerRank, kNumExpertsPerWave, - kNumSMs, kNumRanks, - kNumExpertsPerLane, kNumL1BlockNs, kNumL2BlockNs, - kNumL1BlockKs, kNumL2BlockKs, - layout::SM90Workspace>(workspace); + constexpr uint32_t kNumSharedL1BlockNs = SHARED_L1_SHAPE_N / BLOCK_N; + // The shared L2 N shape equals the routed one, which the scheduler already checks + DG_STATIC_ASSERT(not kHasSharedExperts or kNumSharedL1BlockNs > 0, + "BLOCK_N is too large for the shared-expert L1 shape"); + SchedulerT scheduler(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; @@ -454,10 +482,14 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; // Register reconfiguration counts (chosen to fit in 64512 reg budget). - // For the 256-epilogue-thread split-N decode path: - // 64*48 + 64*40 + 256*168 = 48640 <= 64512. - // For the 512-epilogue-thread split-MN path, trim dispatch and loader roles - // so launch bounds still leave enough WGMMA registers. + // The CTA topology is 1 dispatch warp + TMA-A/TMA-B/producer warps + the + // epilogue warpgroups: + // * 2-WG (epilogue=256): 32*48 + 96*40 + 256*168 = 48384 <= 64512, and + // at 384 threads the launch_bounds ceiling is 65536/384 = 170 >= 168, + // so the WGMMA accumulators do not spill (a 512-thread CTA would cap + // them at 128 < 168 and force local memory). + // * 4-WG (epilogue=512): 32*32 + 96*24 + 512*112 = 60672 <= 64512, with a + // 640-thread CTA ceiling of 65536/640 = 102. // Reduced-thread decode (kNumThreads<=256) raises the launch-bounds // register ceiling to 65536/256=256; grant the epilogue warpgroup the full // 256 so the accumulator double-buffer fits without spilling. @@ -720,6 +752,22 @@ sm90_fp8_mega_moe_impl(void* y, #pragma unroll for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) *workspace.get_expert_send_count_ptr(i) = 0; + // Reset the scheduler's global task counters for the next launch + if (warp_idx == 0 and cute::elect_one_sync()) { + *workspace.get_l1_task_count_ptr() = 0; + *workspace.get_l2_task_count_ptr() = 0; + if constexpr (kHasSharedExperts) { + *workspace.get_shared_l1_task_count_ptr() = 0; + *workspace.get_shared_l2_task_count_ptr() = 0; + } + } + + if constexpr (kHasSharedExperts) { + // Reset the per-M-block shared L1 arrival counters + const uint32_t num_shared_blocks = math::ceil_div(num_tokens, BLOCK_M); + for (uint32_t i = thread_idx; i < num_shared_blocks; i += kNumDispatchThreads) + *workspace.get_shared_l2_full_count_ptr(i) = 0; + } } else { for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { const auto num_recv_tokens = static_cast( @@ -730,11 +778,21 @@ sm90_fp8_mega_moe_impl(void* y, ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - DG_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); - if (warp_idx == 0) { + // Zero the per-expert recv-count sum on dispatch warp 0. The + // cumulative-stats red_add normally runs on dispatch warp 1; + // with the 384-thread topology the CTA collapses to + // a single dispatch warp, so fold it onto warp 0's elect-one lane. + if (warp_idx == 0) *workspace.get_expert_recv_count_sum_ptr(i) = 0; - } else if (warp_idx == 1) { - if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + if constexpr (kNumDispatchWarps >= 2) { + 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(); + } + } else { + if (warp_idx == 0 and 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(); } @@ -765,29 +823,31 @@ sm90_fp8_mega_moe_impl(void* y, } else if (warp_idx == kNumDispatchWarps) { cutlass::arch::warpgroup_reg_dealloc(); - auto process_a_sfa_block = [&](const auto& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - const auto tensor_map_a_ptr = block_phase == sched::BlockPhase::Linear2 - ? &tensor_map_l2_acts : &tensor_map_l1_acts; - const auto tensor_map_sfa_ptr = block_phase == sched::BlockPhase::Linear2 - ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; - - const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + auto process_a_sfa_block = [&](const auto& task_info) { + const auto block_phase = task_info.block_phase; + const auto num_k_blocks = task_info.shape_k / BLOCK_K; + const auto valid_m = task_info.valid_m; + const auto pool_block_idx = task_info.pool_block_idx; + const bool is_shared = task_info.is_shared(); + const bool is_l1 = task_info.is_l1(); + // Shared L1 reads the local `x` directly; shared L2 reads the post-SwiGLU + // shared pool written by the L1 epilogue. Shared activation SF is read from + // global memory by the math warps, so no SFA descriptor is needed for it. + const auto tensor_map_a_ptr = is_shared + ? (is_l1 ? &tensor_map_shared_l1_acts : &tensor_map_shared_l2_acts) + : (is_l1 ? &tensor_map_l1_acts : &tensor_map_l2_acts); + const auto tensor_map_sfa_ptr = is_l1 ? &tensor_map_l1_acts_sf : &tensor_map_l2_acts_sf; // Wait for the pool to be ready if (block_phase == sched::BlockPhase::Linear1) { const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); - const auto expected = scheduler.template get_valid_m(); - while (ptx::ld_acq(ptr) != expected); - } else { + while (ptx::ld_acq(ptr) != valid_m); + } else if (block_phase == sched::BlockPhase::Linear2) { constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; if constexpr (kL2ArrivalCounter) { const auto ptr = reinterpret_cast( workspace.get_l2_arrival_mask_ptr(pool_block_idx)); - const uint32_t active_m_wgs = math::ceil_div( - scheduler.template get_valid_m(), WG_BLOCK_M); + const uint32_t active_m_wgs = math::ceil_div(valid_m, WG_BLOCK_M); const uint32_t expected = kNumL1BlockNs * active_m_wgs * kWarpgroupSplitN * kL1OutputArrivalParts; while (ptx::ld_acq(ptr) != expected); @@ -797,22 +857,74 @@ sm90_fp8_mega_moe_impl(void* y, ? ~0ull : ((1ull << kNumL1BlockNs) - 1ull); while (ptx::ld_acq_gpu(ptr) != expected); } + } else if constexpr (kHasSharedExperts) { + // `SharedLinear1` has no dependency at all: `x` is resident before the + // launch. `SharedLinear2` waits for every shared L1 N tile of this M + // block, in the same counter mode the routed L2 uses. + if (block_phase == sched::BlockPhase::SharedLinear2) { + const auto ptr = workspace.get_shared_l2_full_count_ptr(pool_block_idx); + const uint32_t active_m_wgs = math::ceil_div(valid_m, WG_BLOCK_M); + const uint32_t expected = + kNumSharedL1BlockNs * active_m_wgs * kWarpgroupSplitN * kL1OutputArrivalParts; + while (ptx::ld_acq(ptr) != expected); + } } for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { empty_barriers[stage_idx]->wait(phase ^ 1); - if (cute::elect_one_sync()) { - const uint32_t m_idx = pool_block_idx * BLOCK_M; - const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M; - const uint32_t k_idx = k_block_idx * BLOCK_K; + // Shared tiles index the local token pool directly, so the padded + // SF-pool row mapping does not apply to them. + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // Shared L1 activation SF: `x_sf` is K-major, so a (BLOCK_M, 1) TMA box + // would be 4 bytes and break TMA's 16-byte inner-box rule. Gather the + // column into `smem_sfa` with the whole loader warp instead (BLOCK_M / 32 + // loads per lane, once per stage, versus one per math thread every stage), + // so the math warps keep the uniform `ld_shared` path. `x_sf` is written + // before the launch and never mutated, hence `__ldg`. The stores are + // published by the mbarrier arrive below (which only counts TMA bytes). + if constexpr (kHasSharedExperts) { + if (is_shared and is_l1) { + constexpr uint32_t kNumInputSFGroups = kHidden / kGranK; + const float* in_sf = + input_sf_buffer.template get_base_ptr() + k_block_idx; + #pragma unroll + for (uint32_t i = 0; i < BLOCK_M / 32; ++ i) { + const uint32_t row = i * 32 + lane_idx; + smem_sfa[stage_idx][row] = __ldg(in_sf + (m_idx + row) * kNumInputSFGroups); + } + } + __syncwarp(); + } + if (cute::elect_one_sync()) { // TMA load A tma::copy( tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx, 1); // TMA load SFA - if (block_phase == sched::BlockPhase::Linear1) { + if (is_shared) { + if (is_l1) { + // Gathered above by the whole warp; no TMA bytes to expect + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE); + } else { + // Shared L2 acts SF: the fused L1 epilogue writes it M-major so a + // (BLOCK_M, 1) TMA box is legal -- mirrors routed L2 SFA. Box is + // (BLOCK_M, 1), so issue two single-group TMAs at smem offsets 0 and + // BLOCK_M to match math's `+ 0 * BLOCK_M` / `+ 1 * BLOCK_M` reads. + tma::copy( + &tensor_map_shared_l2_acts_sf, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx * 2, 1); + tma::copy( + &tensor_map_shared_l2_acts_sf, full_barriers[stage_idx], + smem_sfa[stage_idx] + BLOCK_M, m_idx, k_block_idx * 2 + 1, 1); + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + 2 * BLOCK_M * sizeof(float)); + } + } else if (is_l1) { // L1 SFA per-128: load (BLOCK_M, 1) at K=k_block_idx tma::copy( tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], @@ -838,49 +950,33 @@ sm90_fp8_mega_moe_impl(void* y, } }; - if constexpr (kSplitPhaseHotPath) { - sm90_fp8_mega_moe_for_each_block_split( - scheduler, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_a_sfa_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_a_sfa_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); - } else { - scheduler.for_each_block([&](const sched::BlockPhase& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_a_sfa_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); + typename SchedulerT::task_info_t task_info; + while (scheduler.get_next_task(task_info)) { + process_a_sfa_block(task_info); } + } else if (warp_idx == kNumDispatchWarps + 1) { cutlass::arch::warpgroup_reg_dealloc(); - scheduler.for_each_block([&](const sched::BlockPhase& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - const auto tensor_map_b_ptr = - block_phase == sched::BlockPhase::Linear2 ? &tensor_map_l2_weights : &tensor_map_l1_weights; - - const uint32_t shape_n = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_N : L1_SHAPE_N; + auto process_b_block = [&](const auto& task_info) { + const auto local_expert_idx = task_info.local_expert_idx; + const auto num_k_blocks = task_info.shape_k / BLOCK_K; + const auto n_block_idx = task_info.n_cluster_idx; + const auto shape_n = task_info.shape_n; + const bool is_shared = task_info.is_shared(); + const bool is_l1 = task_info.is_l1(); + const auto tensor_map_b_ptr = is_shared + ? (is_l1 ? &tensor_map_shared_l1_weights : &tensor_map_shared_l2_weights) + : (is_l1 ? &tensor_map_l1_weights : &tensor_map_l2_weights); for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { empty_barriers[stage_idx]->wait(phase ^ 1); if (cute::elect_one_sync()) { - const uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + // The fused shared expert is a single dense MLP: no expert stride + const uint32_t n_idx = (is_shared ? 0u : local_expert_idx * shape_n) + + n_block_idx * BLOCK_N; const uint32_t k_idx = k_block_idx * BLOCK_K; // TMA load B (weight SF is now loaded directly by math warps from global) @@ -904,12 +1000,27 @@ sm90_fp8_mega_moe_impl(void* y, } __syncwarp(); } - }); + }; + + typename SchedulerT::task_info_t task_info; + while (scheduler.get_next_task(task_info)) { + process_b_block(task_info); + } + + + } else if (warp_idx == kNumDispatchWarps + 2) { + // Producer warp: claims routed/shared L1/L2 tasks from the global atomic + // counters and publishes them into the SMEM task ring + cutlass::arch::warpgroup_reg_dealloc(); + + scheduler.mainloop(num_tokens); } else if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) { - // Idle non-epilogue warps (kNumDispatchWarps+2, +3). They must still - // participate in the warpgroup-collective `setmaxnreg.dec.sync.aligned` - // so that the math warpgroup's `warpgroup_reg_alloc` can succeed. + // Idle/padding non-epilogue warps: none exist in the 32 + 96 topology the + // host selects (exactly TMA-A + TMA-B + producer). They must still take + // part in the warpgroup-collective `setmaxnreg.dec.sync.aligned` so the + // math warpgroup's `warpgroup_reg_alloc` succeeds if a wider + // non-epilogue section is ever configured. cutlass::arch::warpgroup_reg_dealloc(); } else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { @@ -940,12 +1051,14 @@ sm90_fp8_mega_moe_impl(void* y, // Sync with dispatch in the full communication path. ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); - auto process_math_block = [&](const auto& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - const uint32_t valid_m = scheduler.template get_valid_m(); - const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + auto process_math_block = [&](const auto& task_info) { + const auto local_expert_idx = task_info.local_expert_idx; + const auto num_k_blocks = task_info.shape_k / BLOCK_K; + const auto n_block_idx = task_info.n_cluster_idx; + const auto valid_m = task_info.valid_m; + const auto pool_block_idx = task_info.pool_block_idx; + const bool is_shared = task_info.is_shared(); + const bool is_l1 = task_info.is_l1(); const uint32_t m_idx = pool_block_idx * BLOCK_M; const uint32_t n_idx = n_block_idx * BLOCK_N; const uint32_t epilogue_wg_m_idx = epilogue_wg_idx / kWarpgroupSplitN; @@ -968,10 +1081,68 @@ sm90_fp8_mega_moe_impl(void* y, const bool valid_r0 = row_offset_r0 < valid_m; const bool valid_r1 = row_offset_r1 < valid_m; + // Fused shared expert: a single dense MLP over this rank's own tokens. + // `is_shared`/`is_l1` are derived from `task_info` at the top; `is_shared` + // folds to a compile-time false when the feature is off. Shared tiles + // always publish their L1 arrivals through a counter and never need the + // CTA-wide L2 epilogue sync; the routed modes are compile-time. + const bool use_arrival_counter = kL2ArrivalCounter or is_shared; + const bool needs_l2_full_sync = kL2EpilogueRequiresFullSync and not is_shared; + + // Activation SF for the current K block, read from the TMA/gather-staged SMEM + // tile. All four phases share one layout: L1-like phases put one per-128-K + // float per row at offset 0; L2-like phases put the two per-64-K groups at + // offsets 0 and BLOCK_M (`_hi` is only written for those). Shared tiles are + // staged by the same loader warp (`process_a_sfa_block`), so no phase test is + // needed here. + auto load_act_sf = [&](float& lo_0, float& lo_1, float& hi_0, float& hi_1) { + if (is_l1) { + lo_0 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + lo_1 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + lo_0 = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + lo_1 = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + hi_0 = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + hi_1 = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + }; + + // Block (128, 128) weight SF lookups. The shared expert has no per-expert + // stride and its gate/up N halves are `kSharedIntermediateHidden` apart. + auto load_l1_weight_sf = [&](const uint32_t& k_block_idx, float& gate_sf, float& up_sf) { + constexpr uint32_t kSFKBlocks = kHidden / 128; + const uint32_t gate_n = sf_n_block_idx / 2u; + if (is_shared) { + constexpr uint32_t kSFGateBlks = kSharedIntermediateHidden / 128; + const float* base = shared_l1_weights_sf + k_block_idx; + gate_sf = __ldg(base + gate_n * kSFKBlocks); + up_sf = __ldg(base + (kSFGateBlks + gate_n) * kSFKBlocks); + } else { + constexpr uint32_t kSFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kSFPerExpert = (kIntermediateHidden * 2 / 128) * kSFKBlocks; + const float* base = l1_weights_sf + local_expert_idx * kSFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kSFKBlocks); + up_sf = __ldg(base + (kSFGateBlks + gate_n) * kSFKBlocks); + } + }; + auto load_l2_weight_sf = [&](const uint32_t& k_block_idx) -> float { + if (is_shared) { + constexpr uint32_t kSFKBlocks = kSharedIntermediateHidden / 128; + return __ldg(shared_l2_weights_sf + sf_n_block_idx * kSFKBlocks + k_block_idx); + } + constexpr uint32_t kSFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kSFPerExpert = (kHidden / 128) * kSFKBlocks; + return __ldg(l2_weights_sf + local_expert_idx * kSFPerExpert + + sf_n_block_idx * kSFKBlocks + k_block_idx); + }; + // ---------------- GEMM ---------------- using WGMMA = L1WGMMA; constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; + constexpr uint32_t kSwapSlabAccumStride = kSwapABActive ? SwapWGMMA64::kNumAccum : 0; + constexpr uint32_t kScratchAccumPerThread = kSwapABActive ? kSwapSlabAccumStride : kAccumPerThread; float final_accum[kAccumPerThread] = {}; + float accum[kScratchAccumPerThread]; if constexpr (kReuseAccumAsFinal) { auto prescale_l1_final = [&](const float& scale_a_0, const float& scale_a_1, @@ -1089,34 +1260,15 @@ sm90_fp8_mega_moe_impl(void* y, float scale_a_0_lo, scale_a_1_lo; float scale_a_0_hi, scale_a_1_hi; - if (block_phase == sched::BlockPhase::Linear1) { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); - } else { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); - scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); - scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); - } + load_act_sf(scale_a_0_lo, scale_a_1_lo, scale_a_0_hi, scale_a_1_hi); - constexpr uint32_t kL1SFKBlocks = kHidden / 128; - constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; - constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; - if (block_phase == sched::BlockPhase::Linear1) { - const uint32_t gate_n = sf_n_block_idx / 2u; - const uint32_t up_n = kL1SFGateBlks + gate_n; - const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; - gate_sf = __ldg(base + gate_n * kL1SFKBlocks); - up_sf = __ldg(base + up_n * kL1SFKBlocks); - } else { - l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert - + sf_n_block_idx * kL2SFKBlocks + k_block_idx); - } + if (is_l1) + load_l1_weight_sf(k_block_idx, gate_sf, up_sf); + else + l2_sf = load_l2_weight_sf(k_block_idx); - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if (k_block_idx != 0) rescale_l1_final(prev_scale_a_0, prev_scale_a_1, prev_gate_sf, prev_up_sf, @@ -1197,7 +1349,7 @@ sm90_fp8_mega_moe_impl(void* y, } if (num_k_blocks != 0) { - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { postscale_l1_final(prev_scale_a_0, prev_scale_a_1, prev_gate_sf, prev_up_sf); } else { @@ -1210,34 +1362,15 @@ sm90_fp8_mega_moe_impl(void* y, float scale_a_0_lo, scale_a_1_lo; float scale_a_0_hi, scale_a_1_hi; - if (block_phase == sched::BlockPhase::Linear1) { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); - } else { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); - scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); - scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); - } + load_act_sf(scale_a_0_lo, scale_a_1_lo, scale_a_0_hi, scale_a_1_hi); - constexpr uint32_t kL1SFKBlocks = kHidden / 128; - constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; - constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; - if (block_phase == sched::BlockPhase::Linear1) { - const uint32_t gate_n = sf_n_block_idx / 2u; - const uint32_t up_n = kL1SFGateBlks + gate_n; - const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; - gate_sf = __ldg(base + gate_n * kL1SFKBlocks); - up_sf = __ldg(base + up_n * kL1SFKBlocks); - } else { - l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert - + sf_n_block_idx * kL2SFKBlocks + k_block_idx); - } + if (is_l1) + load_l1_weight_sf(k_block_idx, gate_sf, up_sf); + else + l2_sf = load_l2_weight_sf(k_block_idx); - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if (k_block_idx != 0) prescale_l1_final(scale_a_0_lo, scale_a_1_lo, gate_sf, up_sf); @@ -1317,16 +1450,7 @@ sm90_fp8_mega_moe_impl(void* y, // Read SF (must precede warpgroup_arrive) float scale_a_0_lo, scale_a_1_lo; float scale_a_0_hi, scale_a_1_hi; // Only used in L2 (per-64 K) - if (block_phase == sched::BlockPhase::Linear1) { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); - } else { - // L2: SFA layout is (K=2, M=BLOCK_M) MN-major; first half SF at offset 0, second at BLOCK_M - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); - scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); - scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); - } + load_act_sf(scale_a_0_lo, scale_a_1_lo, scale_a_0_hi, scale_a_1_hi); // ----- 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 @@ -1340,63 +1464,63 @@ sm90_fp8_mega_moe_impl(void* y, // logical 128x128 weight-SF tile, broadcast across the matching // WGMMA accumulators. // + // The fused shared expert uses the same layouts without the expert + // dimension (see `load_l1_weight_sf` / `load_l2_weight_sf`). + // // Load the weight scale after the barrier from all WG threads. // This keeps scale loads close to their WGMMA use and lets the // read-only cache coalesce the same-address accesses. - constexpr uint32_t kL1SFKBlocks = kHidden / 128; - constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; - constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; - if (block_phase == sched::BlockPhase::Linear1) { - const uint32_t gate_n = sf_n_block_idx / 2u; - const uint32_t up_n = kL1SFGateBlks + gate_n; - const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; - gate_sf = __ldg(base + gate_n * kL1SFKBlocks); - up_sf = __ldg(base + up_n * kL1SFKBlocks); - } else { - l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert - + sf_n_block_idx * kL2SFKBlocks + k_block_idx); - } + if (is_l1) + load_l1_weight_sf(k_block_idx, gate_sf, up_sf); + else + l2_sf = load_l2_weight_sf(k_block_idx); - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if constexpr (kSwapABActive) { auto run_swap_ab_l1 = [&]() { using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; - float swap_accum[kSwapAccum]; - - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++ k) { - auto desc_a = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k * SwapWGMMA::K, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_a[stage_idx] + k * SwapWGMMA::K, 1); - SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_wait<0>(); #pragma unroll - for (uint32_t i = 0; i < kSwapAccum / 4; ++ i) { - const uint32_t token_0 = i * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - const float scale_0 = token_0 < valid_m ? - ptx::ld_shared(smem_sfa[stage_idx] + token_0) : 0.0f; - const float scale_1 = token_1 < valid_m ? - ptx::ld_shared(smem_sfa[stage_idx] + token_1) : 0.0f; - final_accum[i * 4 + 0] += scale_0 * gate_sf * swap_accum[i * 4 + 0]; - final_accum[i * 4 + 2] += scale_0 * up_sf * swap_accum[i * 4 + 2]; - final_accum[i * 4 + 1] += scale_1 * gate_sf * swap_accum[i * 4 + 1]; - final_accum[i * 4 + 3] += scale_1 * up_sf * swap_accum[i * 4 + 3]; + for (uint32_t slab = 0; slab < 2; ++slab) { + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++k) { + const uint32_t slab_b_off = + smem_b_wg_offset + slab * SwapWGMMA::M * BLOCK_K; + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + slab_b_off + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + const uint32_t final_base = slab * kSwapSlabAccumStride; + // Routed and shared tiles both read the activation SF from + // `smem_sfa` (the producer stages the shared column there, + // see `process_a_sfa_block`), so this loop is phase-agnostic. + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + const float scale_0 = token_0 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + token_0) : 0.0f; + const float scale_1 = token_1 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + token_1) : 0.0f; + final_accum[final_base + i * 4 + 0] += scale_0 * gate_sf * accum[i * 4 + 0]; + final_accum[final_base + i * 4 + 1] += scale_1 * gate_sf * accum[i * 4 + 1]; + final_accum[final_base + i * 4 + 2] += scale_0 * up_sf * accum[i * 4 + 2]; + final_accum[final_base + i * 4 + 3] += scale_1 * up_sf * accum[i * 4 + 3]; + } } if (lane_idx == 0) @@ -1466,62 +1590,69 @@ sm90_fp8_mega_moe_impl(void* y, auto run_swap_ab_l2 = [&]() { using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; - float swap_accum[kSwapAccum]; - auto promote_swap_accum = [&](const uint32_t& sf_group) { + auto promote_swap_accum = [&](const uint32_t& sf_group, const uint32_t& final_base) { + // SFA layout is (K=2, M=BLOCK_M) M-major in `smem_sfa`, `lo`/`hi` + // slabs at `sf_group * BLOCK_M + ...`. Routed and shared tiles are + // both TMA-staged there, so this loop is phase-agnostic. #pragma unroll - for (uint32_t i = 0; i < kSwapAccum / 4; ++ i) { + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { const uint32_t token_0 = i * 8 + col_idx * 2; const uint32_t token_1 = token_0 + 1; const float scale_0 = token_0 < valid_m ? ptx::ld_shared(smem_sfa[stage_idx] + sf_group * BLOCK_M + token_0) : 0.0f; const float scale_1 = token_1 < valid_m ? ptx::ld_shared(smem_sfa[stage_idx] + sf_group * BLOCK_M + token_1) : 0.0f; - final_accum[i * 4 + 0] += scale_0 * l2_sf * swap_accum[i * 4 + 0]; - final_accum[i * 4 + 2] += scale_0 * l2_sf * swap_accum[i * 4 + 2]; - final_accum[i * 4 + 1] += scale_1 * l2_sf * swap_accum[i * 4 + 1]; - final_accum[i * 4 + 3] += scale_1 * l2_sf * swap_accum[i * 4 + 3]; + final_accum[final_base + i * 4 + 0] += scale_0 * l2_sf * accum[i * 4 + 0]; + final_accum[final_base + i * 4 + 1] += scale_1 * l2_sf * accum[i * 4 + 1]; + final_accum[final_base + i * 4 + 2] += scale_0 * l2_sf * accum[i * 4 + 2]; + final_accum[final_base + i * 4 + 3] += scale_1 * l2_sf * accum[i * 4 + 3]; } }; #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++ k) { - auto desc_a = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k * SwapWGMMA::K, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_a[stage_idx] + k * SwapWGMMA::K, 1); - SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_wait<0>(); - promote_swap_accum(0); + for (uint32_t slab = 0; slab < 2; ++slab) { + const uint32_t slab_b_off = smem_b_wg_offset + slab * SwapWGMMA::M * BLOCK_K; + const uint32_t final_base = slab * kSwapSlabAccumStride; - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++ k) { - const uint32_t k_off = (BLOCK_K / 2) + k * SwapWGMMA::K; - auto desc_a = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_a[stage_idx] + k_off, 1); - SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + slab_b_off + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + promote_swap_accum(0, final_base); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++k) { + const uint32_t k_off = (BLOCK_K / 2) + k * SwapWGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + slab_b_off + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k_off, 1); + SwapWGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + promote_swap_accum(1, final_base); } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_wait<0>(); - promote_swap_accum(1); if (lane_idx == 0) empty_barriers[stage_idx]->arrive(); @@ -1612,19 +1743,21 @@ sm90_fp8_mega_moe_impl(void* y, } } - // Skip epilogue when block is past valid M (still must release via empty) + // Skip epilogue when block is past valid M (still must release via empty). + // The taken paths below must issue the same aligned syncs, so both + // conditions are per-task (shared tiles always use counter arrivals). if (row_base >= valid_m) { - if (block_phase == sched::BlockPhase::Linear1) { - if constexpr (not kL2ArrivalCounter) + if (is_l1) { + if (not use_arrival_counter) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } else { - if constexpr (kL2EpilogueRequiresFullSync) + if (needs_l2_full_sync) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } return; } - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if constexpr (kSwapABActive) { auto silu = [](float x) -> float { const float e = kFastMath ? __expf(-x) : expf(-x); @@ -1640,78 +1773,88 @@ sm90_fp8_mega_moe_impl(void* y, x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); }; - const uint32_t out_col_base = - wg_l1_out_n_offset + warp_idx_in_wg * 8 + row_idx; - auto store_l1_swap_chunk = [&](const uint32_t& i) { - const uint32_t token_0 = i * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - if (token_0 < valid_m) { - float g0 = final_accum[i * 4 + 0]; - float u0 = final_accum[i * 4 + 2]; - clamp_gate(g0); - clamp_up(u0); - const float weight_0 = *l1_topk_weights_buffer - .get_data_buffer(m_idx + token_0) - .get_base_ptr(); - smem_cd_swap_l1_fp32[token_0 * L1_OUT_BLOCK_N + out_col_base] = - silu(g0) * u0 * weight_0; - } - if (token_1 < valid_m) { - float g1 = final_accum[i * 4 + 1]; - float u1 = final_accum[i * 4 + 3]; - clamp_gate(g1); - clamp_up(u1); - const float weight_1 = *l1_topk_weights_buffer - .get_data_buffer(m_idx + token_1) - .get_base_ptr(); - smem_cd_swap_l1_fp32[token_1 * L1_OUT_BLOCK_N + out_col_base] = - silu(g1) * u1 * weight_1; - } - }; + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + const uint32_t num_swap_chunks = n_swap / 8u; + const uint32_t wg_out_col_base = epilogue_wg_n_idx * (WG_BLOCK_N / 2); - const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; - store_l1_swap_chunk(0); - if (valid_m > 8) { + #pragma unroll + for (uint32_t slab = 0; slab < 2; ++slab) { + const uint32_t final_base = slab * kSwapSlabAccumStride; + const uint32_t out_col = wg_out_col_base + slab * 32u + + warp_idx_in_wg * 8u + row_idx; #pragma unroll - for (uint32_t i = 1; i < kSwapABTokenChunks; ++ i) { - if (i < num_swap_token_chunks) - store_l1_swap_chunk(i); + for (uint32_t i = 0; i < SwapWGMMA64::kNumAccum / 4; ++i) { + if (i >= num_swap_chunks) break; + const uint32_t token_0 = i * 8u + col_idx * 2u; + const uint32_t token_1 = token_0 + 1u; + if (token_0 < valid_m) { + float g0 = final_accum[final_base + i * 4 + 0]; + float u0 = final_accum[final_base + i * 4 + 2]; + clamp_gate(g0); + clamp_up(u0); + const float weight_0 = is_shared ? 1.0f : *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_0) + .template get_base_ptr(); + smem_cd_swap_l1_fp32[token_0 * L1_OUT_BLOCK_N + out_col] = + silu(g0) * u0 * weight_0; + } + if (token_1 < valid_m) { + float g1 = final_accum[final_base + i * 4 + 1]; + float u1 = final_accum[final_base + i * 4 + 3]; + clamp_gate(g1); + clamp_up(u1); + const float weight_1 = is_shared ? 1.0f : *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_1) + .template get_base_ptr(); + smem_cd_swap_l1_fp32[token_1 * L1_OUT_BLOCK_N + out_col] = + silu(g1) * u1 * weight_1; + } } } ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); for (uint32_t token = epilogue_thread_idx; token < valid_m; token += kNumEpilogueThreads) { - float amax = 0.0f; + constexpr uint32_t kHalfN = L1_OUT_BLOCK_N / 2; + float amax0 = 0.0f, amax1 = 0.0f; #pragma unroll - for (uint32_t col = 0; col < L1_OUT_BLOCK_N; ++ col) { + for (uint32_t col = 0; col < L1_OUT_BLOCK_N; ++col) { const float v = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col]; - amax = cute::max(amax, cute::abs(v)); + const float a = cute::abs(v); + if (col < kHalfN) amax0 = cute::max(amax0, a); + else amax1 = cute::max(amax1, a); } - 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); - const float sf = sf_pair.x; - const float sf_inv = sf_inv_pair.x; - auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); - // ROOT-CAUSE FIX: the L2-activation SF pool is strided by SF_BLOCK_M - // (=align(BLOCK_M,128)=128), which is how the L2 producer reads it - // (sfa_m_idx = pool_block_idx * SF_BLOCK_M) and how the non-swap L1 - // writes it. This swapAB path used BLOCK_M (64), so for pool_block_idx>=1 - // the SF landed in the wrong rows -> L2 read stale SF -> every pool block - // after the first was corrupted (block 0 was correct because 0*64==0*128). - const uint32_t token_idx = pool_block_idx * SF_BLOCK_M + token; - sf_base_ptr[n_block_idx * kNumPaddedSFPoolTokens + token_idx] = sf; + float2 sf_pair, sf_inv_pair; + sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( + make_float2(amax0, amax1), sf_pair, sf_inv_pair); + const float sf0 = sf_pair.x, sf1 = sf_pair.y; + const float sf_inv0 = sf_inv_pair.x, sf_inv1 = sf_inv_pair.y; + // Shared L2 activation SF (written here by the fused L1) is now M-major + // (token-contiguous inside each `k_sf_*` slab, mirroring the routed L2 SF + // pool) so producer can TMA-load (BLOCK_M, 1) tiles directly into smem_sfa. + const uint32_t k_sf_lo = n_block_idx * 2u + 0u; + const uint32_t k_sf_hi = n_block_idx * 2u + 1u; + if (is_shared) { + auto sf_base_ptr = shared_l2_sf_buffer.get_base_ptr(); + const uint32_t row = m_idx + token; + sf_base_ptr[k_sf_lo * kNumMaxTokensPerRank + row] = sf0; + sf_base_ptr[k_sf_hi * kNumMaxTokensPerRank + row] = sf1; + } else { + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + const uint32_t token_idx = pool_block_idx * SF_BLOCK_M + token; + sf_base_ptr[k_sf_lo * kNumPaddedSFPoolTokens + token_idx] = sf0; + sf_base_ptr[k_sf_hi * kNumPaddedSFPoolTokens + token_idx] = sf1; + } #pragma unroll for (uint32_t col = 0; col < L1_OUT_BLOCK_N; col += 2) { - const float v0 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col + 0] * sf_inv; + const float sf_inv = (col < kHalfN) ? sf_inv0 : sf_inv1; + const float v0 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col] * sf_inv; const float v1 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col + 1] * sf_inv; const __nv_fp8x2_e4m3 pair(make_float2(v0, v1)); - auto* ptr = reinterpret_cast( - smem_cd_swap_l1_fp8 + token * L1_OUT_BLOCK_N + col); - *ptr = pair.__x; + *reinterpret_cast( + smem_cd_swap_l1_fp8 + token * L1_OUT_BLOCK_N + col) = pair.__x; } } @@ -1720,7 +1863,7 @@ sm90_fp8_mega_moe_impl(void* y, if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( - &tensor_map_l1_output, + is_shared ? &tensor_map_shared_l1_output : &tensor_map_l1_output, smem_cd_swap_l1_fp8, n_block_idx * L1_OUT_BLOCK_N, m_idx); @@ -1729,12 +1872,16 @@ sm90_fp8_mega_moe_impl(void* y, __syncwarp(); ptx::tma_store_wait<0>(); - if constexpr (kL2ArrivalCounter) { - if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { - ptx::red_add_rel( - reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), - kWarpgroupSplitN); - } + // `use_arrival_counter` is true for shared tiles and for the routed + // counter mode; select the per-M-block arrival slot and publish the + // whole N-split warpgroup group from the storing WG (same shape as the + // non-swap L1 epilogue). The bitmask path otherwise stays untouched. + if (use_arrival_counter) { + auto arrival_ptr = is_shared + ? workspace.get_shared_l2_full_count_ptr(pool_block_idx) + : reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)); + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) + ptx::red_add_rel(arrival_ptr, kWarpgroupSplitN); } else { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { @@ -1744,7 +1891,9 @@ sm90_fp8_mega_moe_impl(void* y, } } __syncwarp(); - if constexpr (kL2ArrivalCounter) + // counter mode (incl. shared) needs the CTA-wide sync to protect + // the swapAB FP32/FP8 staging tiles from the next task's overwrite. + if (use_arrival_counter) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } else { @@ -1818,13 +1967,14 @@ sm90_fp8_mega_moe_impl(void* y, } } - // Apply token weight: SwiGLU * topk_weight (single load per row) - const float weight_r0 = valid_r0 ? *l1_topk_weights_buffer + // Apply token weight: SwiGLU * topk_weight (single load per row). + // The shared expert sees every token with weight 1.0. + const float weight_r0 = valid_r0 ? (is_shared ? 1.0f : *l1_topk_weights_buffer .get_data_buffer(m_idx + row_offset_r0) - .get_base_ptr() : 0.0f; - const float weight_r1 = valid_r1 ? *l1_topk_weights_buffer + .template get_base_ptr()) : 0.0f; + const float weight_r1 = valid_r1 ? (is_shared ? 1.0f : *l1_topk_weights_buffer .get_data_buffer(m_idx + row_offset_r1) - .get_base_ptr() : 0.0f; + .template get_base_ptr()) : 0.0f; #pragma unroll for (uint32_t p = 0; p < kNumPairs; ++ p) { swiglu_r0[p][0] *= weight_r0; @@ -1915,16 +2065,27 @@ sm90_fp8_mega_moe_impl(void* y, // In the shared-SF split both warpgroups own the same per-64 group and rows, so // only the first N-split warpgroup publishes the SF slot to avoid a write race. if (col_idx == 0 and (not kSplitNSharesSF or epilogue_wg_n_idx == 0)) { - auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); - // SF buffer is (kNumPaddedSFPoolTokens x kIntermediateHidden/64), MN-major: - // addr[k_idx * num_padded_sf_pool_tokens + token_idx] - const uint32_t token_r0 = pool_block_idx * SF_BLOCK_M + row_offset_r0; - const uint32_t token_r1 = pool_block_idx * SF_BLOCK_M + row_offset_r1; const uint32_t k_sf_idx = sf_n_block_idx; // one per-64 post-SwiGLU group - if (valid_r0) - sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0; - if (valid_r1) - sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1; + if (is_shared) { + // Shared SF buffer is (kNumMaxTokensPerRank x SIH/64) M-major and + // indexed by the local token (token-contiguous inner, stride 1), so no + // SF-pool padding applies: addr[k_idx * kNumMaxTokensPerRank + token_idx] + auto sf_base_ptr = shared_l2_sf_buffer.get_base_ptr(); + if (valid_r0) + sf_base_ptr[k_sf_idx * kNumMaxTokensPerRank + (m_idx + row_offset_r0)] = sf_r0; + if (valid_r1) + sf_base_ptr[k_sf_idx * kNumMaxTokensPerRank + (m_idx + row_offset_r1)] = sf_r1; + } else { + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + // SF buffer is (kNumPaddedSFPoolTokens x kIntermediateHidden/64), MN-major: + // addr[k_idx * num_padded_sf_pool_tokens + token_idx] + const uint32_t token_r0 = pool_block_idx * SF_BLOCK_M + row_offset_r0; + const uint32_t token_r1 = pool_block_idx * SF_BLOCK_M + row_offset_r1; + if (valid_r0) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0; + if (valid_r1) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1; + } } // Sync the warpgroup before TMA store. In the shared-tile split @@ -1951,7 +2112,7 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( - &tensor_map_l1_output, + is_shared ? &tensor_map_shared_l1_output : &tensor_map_l1_output, smem_cd_l1, out_n_idx, m_idx + row_base); @@ -1962,7 +2123,7 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N + wg_l1_out_n_offset; cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( - &tensor_map_l1_output, + is_shared ? &tensor_map_shared_l1_output : &tensor_map_l1_output, smem_cd_l1 + smem_cd_l1_wg_offset, out_n_idx, m_idx + row_base); @@ -1974,20 +2135,20 @@ sm90_fp8_mega_moe_impl(void* y, // Notify L2 that this L1 output (and SF) is ready. Counter mode lets // independent WG tiles publish arrivals without the CTA-wide barrier - // needed before the single bit-mask update. - if constexpr (kL2ArrivalCounter) { + // needed before the single bit-mask update. Shared tiles always use a + // counter, on their own per-M-block slot. + if (use_arrival_counter) { + auto arrival_ptr = is_shared + ? workspace.get_shared_l2_full_count_ptr(pool_block_idx) + : reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)); if constexpr (kSplitNSharesSF) { // The combined tile counts for both N-split warpgroups; the // storing warpgroup publishes all kWarpgroupSplitN arrivals // after its TMA store has drained. - if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { - ptx::red_add_rel( - reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), - kWarpgroupSplitN); - } + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) + ptx::red_add_rel(arrival_ptr, kWarpgroupSplitN); } else if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { - ptx::red_add_rel( - reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), 1); + ptx::red_add_rel(arrival_ptr, 1); } } else { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); @@ -2013,32 +2174,51 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t lane_in_row = lane_idx % 16; const uint32_t cols_per_lane = WG_BLOCK_N / 16; - if constexpr (kSwapABActive) { - auto store_bf16 = [&](const uint32_t& token, const uint32_t& col, float value) { - smem_cd_l2[smem_cd_l2_wg_offset + token * WG_BLOCK_N + col] = - __float2bfloat16_rn(value); - }; - - auto store_l2_swap_chunk = [&](const uint32_t& i) { - const uint32_t token_0 = i * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - if (token_0 < valid_m) { - store_bf16(token_0, r_0, final_accum[i * 4 + 0]); - store_bf16(token_0, r_1, final_accum[i * 4 + 2]); - } - if (token_1 < valid_m) { - store_bf16(token_1, r_0, final_accum[i * 4 + 1]); - store_bf16(token_1, r_1, final_accum[i * 4 + 3]); - } - }; + // XOR column swizzle (8-col granularity) for the row-major BF16 + // staging tile. The row stride is WG_BLOCK_N/2 banks, which is a + // multiple of 32 for every supported WG_BLOCK_N (64/128), so 8 + // lanes that share a col_idx (8 distinct row_idx) all hit the same + // bank -> 8-way conflict on each STS. XORing bits [3:5] of the + // column with (row & 7) spreads those 8 rows across 8 distinct + // banks. The swizzle MUST be applied on both the STS write and the + // LDS scatter read so the permutation cancels out; doing it on the + // write alone (as a port of SM100's layout) silently permutes the + // output columns -- SM100's swizzle is enforced by its TMA + // descriptor, this manual STS/LDS path has no such contract. The + // 8-col granularity is safe: the 2-BF16 STS pair and the + // cols_per_lane-BF16 LDS vector (4 or 8 BF16, the only sizes this + // path supports) never straddle an 8-col block, so the key is + // constant across each access. + auto swiz_col = [](uint32_t row, uint32_t col) -> uint32_t { + return col ^ ((row & 7) << 3); + }; - const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; - store_l2_swap_chunk(0); - if (valid_m > 8) { + if constexpr (kSwapABActive) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + const uint32_t num_swap_chunks = n_swap / 8u; + #pragma unroll + for (uint32_t slab = 0; slab < 2; ++slab) { + const uint32_t final_base = slab * kSwapSlabAccumStride; + const uint32_t slab_col_base = wg_n_offset + slab * SwapWGMMA64::M; #pragma unroll - for (uint32_t i = 1; i < kSwapABTokenChunks; ++ i) { - if (i < num_swap_token_chunks) - store_l2_swap_chunk(i); + for (uint32_t i = 0; i < SwapWGMMA64::kNumAccum / 4; ++i) { + if (i >= num_swap_chunks) break; + const uint32_t token_0 = i * 8u + col_idx * 2u; + const uint32_t token_1 = token_0 + 1u; + const uint32_t col_0 = slab_col_base + r_0; + const uint32_t col_1 = slab_col_base + r_1; + if (token_0 < valid_m) { + smem_cd_l2[token_0 * BLOCK_N + swiz_col(token_0, col_0)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 0]); + smem_cd_l2[token_0 * BLOCK_N + swiz_col(token_0, col_1)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 2]); + } + if (token_1 < valid_m) { + smem_cd_l2[token_1 * BLOCK_N + swiz_col(token_1, col_0)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 1]); + smem_cd_l2[token_1 * BLOCK_N + swiz_col(token_1, col_1)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 3]); + } } } } else { @@ -2055,7 +2235,7 @@ sm90_fp8_mega_moe_impl(void* y, auto smem_ptr = smem_cd_l2 + smem_cd_l2_wg_offset + row * WG_BLOCK_N - + col; + + swiz_col(row, col); // BF16 STS: 2 bf16 elements *reinterpret_cast(smem_ptr) = packed; }; @@ -2100,17 +2280,35 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t m_idx_in_block = row_base + row_in_wg; if (m_idx_in_block >= valid_m) break; - // Read cols_per_lane BF16 (= one ScatterVec) from smem - auto smem_ptr = smem_cd_l2 - + smem_cd_l2_wg_offset - + row_in_wg * WG_BLOCK_N - + lane_in_row * cols_per_lane; + // Read cols_per_lane BF16 (= one ScatterVec) from smem. + // swapAB uses a shared BLOCK_N-wide tile with column swizzle; + // non-swap uses per-WG slab with same swizzle. + nv_bfloat16* smem_ptr; + if constexpr (kSwapABActive) { + const uint32_t base_col = wg_n_offset + lane_in_row * cols_per_lane; + const uint32_t read_col = swiz_col(row_in_wg, base_col); + smem_ptr = smem_cd_l2 + row_in_wg * BLOCK_N + read_col; + } else { + uint32_t read_col = lane_in_row * cols_per_lane; + read_col = swiz_col(row_in_wg, read_col); + smem_ptr = smem_cd_l2 + smem_cd_l2_wg_offset + + row_in_wg * WG_BLOCK_N + read_col; + } const auto packed = *reinterpret_cast(smem_ptr); - const auto src_metadata = *workspace.get_token_src_metadata_ptr(m_idx + m_idx_in_block); - const uint32_t dst_rank_idx = src_metadata.rank_idx; - const uint32_t dst_token_idx = src_metadata.token_idx; - const uint32_t dst_topk_idx = src_metadata.topk_idx; + // The fused shared expert stays on this rank and reduces through the + // extra combine slot `kNumTopk`, keyed by the local token index. + uint32_t dst_rank_idx, dst_token_idx, dst_topk_idx; + if (is_shared) { + dst_rank_idx = sym_buffer.rank_idx; + dst_token_idx = m_idx + m_idx_in_block; + dst_topk_idx = kNumTopk; + } else { + const auto src_metadata = *workspace.get_token_src_metadata_ptr(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 auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) .get_data_buffer(dst_token_idx); auto dst_ptr = math::advance_ptr( @@ -2119,36 +2317,16 @@ sm90_fp8_mega_moe_impl(void* y, *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; } - if constexpr (kL2EpilogueRequiresFullSync) + if (needs_l2_full_sync) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } }; - if constexpr (kSplitPhaseHotPath) { - sm90_fp8_mega_moe_for_each_block_split( - scheduler, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_math_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_math_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); - } else { - scheduler.for_each_block([&](const sched::BlockPhase& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_math_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); + typename SchedulerT::task_info_t task_info; + while (scheduler.get_next_task(task_info)) { + process_math_block(task_info); } + // ---------------- COMBINE ---------------- // NVLink barrier first: signals remote ranks that this rank's GEMM @@ -2187,7 +2365,7 @@ sm90_fp8_mega_moe_impl(void* y, DG_STATIC_ASSERT(kNumChunkBytes % 16 == 0, "Combine chunk must be TMA-aligned (16 bytes)"); DG_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, "Combine chunk must be divisible by 16 bytes"); DG_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, "Combine chunk must be a multiple of 32 16-byte elements"); - DG_STATIC_ASSERT(kNumTopk <= 32, "Top-k must fit in a single warp"); + DG_STATIC_ASSERT(kNumCombineSlots <= 32, "Top-k (plus the shared slot) must fit in a single warp"); DG_TRAP_ONLY_DEVICE_ASSERT(kNumChunkSlots * kNumCombineWarps * kNumChunkBytes <= static_cast( reinterpret_cast(barrier_start_ptr) - smem_buffer)); @@ -2207,8 +2385,12 @@ sm90_fp8_mega_moe_impl(void* y, for (uint32_t token_idx = sm_idx * kNumCombineWarps + epilogue_warp_idx; token_idx < num_tokens; token_idx += kNumSMs * kNumCombineWarps) { + // Slots `[0, kNumTopk)` are the routed experts (negative expert id means the + // slot was masked out). With the fused shared expert, slot `kNumTopk` is + // always present and holds this rank's shared-expert contribution. const int stored_topk_slot_idx = lane_idx < kNumTopk ? - static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; + 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) { diff --git a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh index 4846c984c2..8dc7026501 100644 --- a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh @@ -4,6 +4,7 @@ #include #include +#include namespace deep_gemm::layout { @@ -261,8 +262,19 @@ struct SM90Workspace { uint32_t num_max_pool_tokens; uint32_t num_max_pool_blocks; + // Fused shared expert: one counter per M block of this rank's own tokens. + // Always reserved (a few KB) so the host and device layouts agree whether or + // not the shared expert is enabled. + uint32_t num_max_shared_pool_blocks; - static constexpr uint64_t kNumBarrierSignalBytes = 32; + // [ 0..15]: 4 x `uint32_t` grid sync counters + // [16..20]: `uint32_t` NVLink barrier counter + // [20..28]: 2 x `int` NVLink barrier signals (phase 0 and 1) + // [28..32]: `uint32_t` L1 task counter (interleaved scheduler) + // [32..36]: `uint32_t` L2 task counter (interleaved scheduler) + // [36..40]: `uint32_t` shared L1 task counter (fused shared expert) + // [40..44]: `uint32_t` shared L2 task counter (fused shared expert) + static constexpr uint64_t kNumBarrierSignalBytes = 48; CUTLASS_HOST_DEVICE SM90Workspace(void* base, @@ -278,6 +290,7 @@ struct SM90Workspace { num_max_pool_tokens = get_num_max_pool_tokens( num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); num_max_pool_blocks = num_max_pool_tokens / kMinCandidateBlockM; + num_max_shared_pool_blocks = math::ceil_div(num_max_tokens_per_rank, static_cast(kMinCandidateBlockM)); } CUTLASS_HOST_DEVICE @@ -290,6 +303,7 @@ struct SM90Workspace { num_bytes += num_max_pool_blocks * sizeof(uint64_t); num_bytes += num_experts_per_rank * num_ranks * num_max_recv_tokens_per_expert * sizeof(int); num_bytes += num_max_pool_tokens * sizeof(TokenSrcMetadata); + num_bytes += math::align(num_max_shared_pool_blocks, 4u) * sizeof(uint32_t); return math::align(num_bytes, 16); } @@ -318,6 +332,30 @@ struct SM90Workspace { base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int)); } + // Interleaved-scheduler global task counters. They are zeroed by the + // workspace cleanup path at the end of every kernel launch (and the + // workspace allocation is zero-initialized for the first launch). + CUTLASS_DEVICE + uint32_t* get_l1_task_count_ptr() const { + return math::advance_ptr(base, 28u); + } + + CUTLASS_DEVICE + uint32_t* get_l2_task_count_ptr() const { + return math::advance_ptr(base, 32u); + } + + // Fused shared-expert task counters, zeroed by the same cleanup path + CUTLASS_DEVICE + uint32_t* get_shared_l1_task_count_ptr() const { + return math::advance_ptr(base, 36u); + } + + CUTLASS_DEVICE + uint32_t* get_shared_l2_task_count_ptr() const { + return math::advance_ptr(base, 40u); + } + CUTLASS_DEVICE uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const { return math::advance_ptr(base, kNumBarrierSignalBytes) + expert_idx; @@ -360,6 +398,15 @@ struct SM90Workspace { const auto base = reinterpret_cast(get_src_token_topk_idx_ptr(num_experts_per_rank)); return base + pool_token_idx; } + + // Fused shared expert: per-M-block count of finished shared L1 tiles. The + // shared L2 A/SFA loader spins on it, mirroring `get_l2_arrival_mask_ptr` in + // counter mode for the routed path. + CUTLASS_DEVICE + uint32_t* get_shared_l2_full_count_ptr(const uint32_t& shared_block_idx = 0) const { + const auto base = get_token_src_metadata_ptr(num_max_pool_tokens); + return reinterpret_cast(base) + shared_block_idx; + } }; struct Data { @@ -474,6 +521,11 @@ struct MegaMoEBuffer { combine_token_buffer, combine_sf_buffer; + // Optional per-token FP32 outer scales for the L1 (fc13) input, applied on + // the L1 accumulator before activation. + Buffer input_x_scales_buffer, + l1_x_scales_buffer; + CUTLASS_HOST_DEVICE MegaMoEBuffer(void* base, const uint32_t& hidden, @@ -484,10 +536,10 @@ struct MegaMoEBuffer { const uint32_t& num_topk, const uint32_t& num_ring_tokens, const uint32_t& num_sf_ring_tokens, - const bool& with_sf, + const MmaKind& mma_kind, const uint32_t& num_shared_experts = 0, - const bool& use_fp4_acts = false, const bool& use_fp8_combine = false) { + const bool with_sf = get_sf_gran_k(mma_kind) != 0; // Workspace workspace = Workspace(base, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_ring_tokens); @@ -497,21 +549,25 @@ struct MegaMoEBuffer { const auto num_max_shared_sf_tokens = with_sf ? get_num_max_shared_sf_tokens(num_max_tokens_per_rank) : 0u; // Layouts - const uint32_t num_mma_elem_bytes = with_sf ? 1 : 2; - const auto input_token_layout = layout::Data( - use_fp4_acts ? hidden / 2 : hidden * num_mma_elem_bytes); + const uint32_t elem_bits = get_element_bits(mma_kind); + // NOTE: multiply before dividing — `elem_bits / 8` truncates to 0 for the + // 4-bit kinds (NVFP4/MXFP4), which would zero-size every activation slot. + const auto num_token_bytes = [=](const uint32_t& num_elems) { return num_elems * elem_bits / 8; }; + const uint32_t gran_k = with_sf ? get_sf_gran_k(mma_kind) : 1; + const auto input_token_layout = layout::Data(num_token_bytes(hidden)); const auto combine_token_layout = layout::Data( use_fp8_combine ? hidden : hidden * 2); const auto combine_sf_layout = layout::Data( use_fp8_combine ? hidden / 128 : 0, false); - 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 intermediate_token_layout = layout::Data(num_token_bytes(intermediate_hidden)); + const auto shared_intermediate_token_layout = layout::Data(num_token_bytes(shared_intermediate_hidden)); + const auto input_sf_layout = layout::Data(with_sf ? hidden / gran_k : 0); + const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / gran_k : 0); + const auto shared_intermediate_sf_layout = layout::Data(with_sf ? shared_intermediate_hidden / gran_k : 0); const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false); const auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + const auto x_scale_layout = layout::Data(sizeof(float), false); // Input buffers input_token_buffer = Buffer( @@ -566,11 +622,18 @@ struct MegaMoEBuffer { combine_sf_buffer = Buffer( combine_sf_layout, num_topk + (num_shared_experts > 0 ? 1u : 0u), num_max_tokens_per_rank, combine_token_buffer.get_end_ptr()); + + input_x_scales_buffer = Buffer( + x_scale_layout, 1, num_max_tokens_per_rank, + combine_sf_buffer.get_end_ptr()); + l1_x_scales_buffer = Buffer( + x_scale_layout, 1, num_ring_tokens, + input_x_scales_buffer.get_end_ptr()); } CUTLASS_HOST_DEVICE int64_t get_num_bytes() const { - return static_cast(combine_sf_buffer.get_end_ptr()) + return static_cast(l1_x_scales_buffer.get_end_ptr()) - static_cast(workspace.base); } }; diff --git a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh index fb7d3e6e12..166c034c09 100644 --- a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh +++ b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh @@ -171,6 +171,30 @@ struct SM100_MMA_MXF4_2x1SM_SS { } }; +struct SM100_MMA_MXF4NVF4_2x1SM_SS { + CUTLASS_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& scale_c, + uint64_t const& desc, + uint32_t const& tmem_sfa, + uint32_t const& tmem_sfb) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" +#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#else + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#endif + "}\n" + :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast(desc >> 32)), "r"(scale_c), + "r"(tmem_sfa), "r"(tmem_sfb)); + } +}; + struct SM100_MMA_F16BF16_WS_SS { CUTLASS_DEVICE static void fma(uint64_t const& desc_a, diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh index a4e87bbd1b..ac08ea04dd 100644 --- a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh @@ -88,6 +88,11 @@ enum class BlockPhase : uint32_t { SharedLinear2 = 4 }; +// The L1-like phases (Linear1 / SharedLinear1) read per-128-K activation SF and +// run the SwiGLU + FP8-quantize epilogue; the L2-like phases (Linear2 / +// SharedLinear2) read per-64-K SF and scatter BF16 into the combine buffer. The +// Shared* variants are the fused shared-expert counterparts of Linear1/Linear2. +// Phase membership is queried through `TaskInfo::is_l1()` / `TaskInfo::is_shared()`. template struct alignas(16) TaskInfo { BlockPhase block_phase; @@ -131,6 +136,11 @@ struct alignas(16) TaskInfo { CUTLASS_DEVICE uint32_t is_shared() const { return kHasSharedExperts ? (block_phase > BlockPhase::Linear2) : false; } + + // True for both routed and shared L1 phases (SwiGLU + FP8-quantize epilogue). + CUTLASS_DEVICE uint32_t is_l1() const { + return (block_phase == BlockPhase::Linear1) or (block_phase == BlockPhase::SharedLinear1); + } }; DG_STATIC_ASSERT(sizeof(sched::TaskInfo) == sizeof(sched::TaskInfo), "Invalid layout"); @@ -144,12 +154,19 @@ template + uint32_t kNumL1Clusters = kNumL1BlockNs / kClusterSize, + uint32_t kNumL2Clusters = kNumL2BlockNs / kClusterSize> struct MegaMoEScheduler { + // `kClusterSize` selects the task granularity: 2 for the SM100 2-CTA cluster + // (one task = a CTA pair, the original behaviour), 1 for the SM90 single-CTA + // path (one task = one CTA tile). With the default `kClusterSize = 2` every + // formula below reduces bit-for-bit to the pre-unification SM100 scheduler. + DG_STATIC_ASSERT(kClusterSize == 1 or kClusterSize == 2, "Invalid cluster size"); static constexpr bool kHasShared = kNumSharedExperts > 0; static constexpr uint32_t SHARED_L1_SHAPE_N = L1_SHAPE_N * kNumSharedExperts; static constexpr uint32_t SHARED_L1_SHAPE_K = L1_SHAPE_K; @@ -172,7 +189,7 @@ struct MegaMoEScheduler { DG_STATIC_ASSERT(kNumRingBlocks > 0, "Invalid ring buffer config"); // Workspace - const layout::Workspace& workspace; + const WorkspaceT& workspace; // Scheduler configs static constexpr uint32_t kNumScheduleStages = 2; @@ -192,10 +209,10 @@ struct MegaMoEScheduler { static constexpr uint32_t kNumSchedL1WavesDone = 0xffffffffu; uint32_t num_sched_l1_waves = 0; - CUTLASS_DEVICE explicit MegaMoEScheduler(const layout::Workspace& workspace): + CUTLASS_DEVICE explicit MegaMoEScheduler(const WorkspaceT& workspace): workspace(workspace) {} - CUTLASS_DEVICE MegaMoEScheduler(const layout::Workspace& workspace, + CUTLASS_DEVICE MegaMoEScheduler(const WorkspaceT& workspace, Barrier* task_info_full_barriers, Barrier* task_info_empty_barriers, task_info_t* task_infos): @@ -213,6 +230,15 @@ struct MegaMoEScheduler { CUTLASS_DEVICE bool get_next_task(task_info_t& task_info) { task_info_full_barriers[sched_stage_idx].wait(sched_phase); task_info = task_infos[sched_stage_idx]; + if constexpr (kClusterSize == 1) { + // Single-CTA (SM90): every consumer warp releases the slot immediately + // once it has copied the task into registers (`task_info_empty_barriers` + // is initialised to `2 + kNumEpilogueWarps` arrivals). The 2-CTA path + // keeps its deferred release via `release_task_info()`. + __syncwarp(); + if (cute::elect_one_sync()) + task_info_empty_barriers[sched_stage_idx].arrive(); + } advance_sched_pipeline(); return task_info.is_valid(); } @@ -263,9 +289,9 @@ struct MegaMoEScheduler { 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 num_total_l1_waves = math::ceil_div(num_total_l1_tasks, kNumSMs / kClusterSize); const uint32_t min_l1_warmup_waves = get_num_l1_warmup_waves( - num_total_m_blocks, kNumSMs / 2, kNumL1Clusters, kNumL2Clusters); + num_total_m_blocks, kNumSMs / kClusterSize, kNumL1Clusters, kNumL2Clusters); num_sched_l1_waves = cute::min(min_l1_warmup_waves, num_total_l1_waves); } @@ -350,12 +376,25 @@ 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 (kClusterSize == 1) { + // Single-CTA (SM90): the producer and the consumers live in the same + // CTA, so write the slot into local SMEM and arrive the full barrier. + if (lane_idx == 0) { + task_infos[sched_stage_idx] = task_info; + // The mbarrier arrive has release semantics, so the plain SMEM + // stores above are visible to consumers after their wait. + task_info_full_barriers[sched_stage_idx].arrive(); + } + } else { + // 2-CTA cluster (SM100): the producer (leader CTA) publishes the task + // to the SMEM of BOTH CTAs in the cluster via async cluster store. + 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] + ); + } } __syncwarp(); advance_sched_pipeline(); @@ -363,7 +402,7 @@ struct MegaMoEScheduler { 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 / kClusterSize; const uint32_t num_m_blocks = math::ceil_div(num_tokens, BLOCK_M); const uint32_t num_tasks = num_m_blocks * kNumNClusters; while (true) { @@ -414,175 +453,6 @@ struct MegaMoEScheduler { } }; -// Hopper MegaMoE retains the original wave scheduler. SM100 uses the task -// producer/consumer scheduler above so shared experts can participate in the -// same pipeline, while SM90 keeps its tuned single-CTA scheduling contract. -template -struct LegacyMegaMoEScheduler { - DG_STATIC_ASSERT(L1_SHAPE_N % BLOCK_N == 0, "Invalid shape"); - DG_STATIC_ASSERT(L2_SHAPE_N % BLOCK_N == 0, "Invalid shape"); - DG_STATIC_ASSERT(L1_SHAPE_K % BLOCK_K == 0, "Invalid shape"); - DG_STATIC_ASSERT(L2_SHAPE_K % BLOCK_K == 0, "Invalid shape"); - DG_STATIC_ASSERT(kNumExpertsPerWave > 0 and kNumExpertsPerWave <= kNumExpertsPerRank, "Invalid wave config"); - DG_STATIC_ASSERT(kNumSMs % 2 == 0, "Number of SMs must be even"); - DG_STATIC_ASSERT(kNumL1BlockNs % 2 == 0, "L1 N block count must be even"); - DG_STATIC_ASSERT(kNumL2BlockNs % 2 == 0, "L2 N block count must be even"); - - const WorkspaceT& workspace; - BlockPhase next_phase = BlockPhase::Linear1; - uint32_t current_local_expert_idx = 0; - uint32_t current_num_tokens = 0; - uint32_t current_pool_block_offset = 0; - uint32_t block_idx = 0; - uint32_t m_block_idx = 0; - uint32_t n_block_idx = 0; - uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {}; - - CUTLASS_DEVICE explicit LegacyMegaMoEScheduler(const WorkspaceT& workspace): workspace(workspace) { - block_idx = blockIdx.x; - } - - CUTLASS_DEVICE uint32_t get_wave_expert_end_idx() const { - const auto aligned = math::align(current_local_expert_idx + 1, kNumExpertsPerWave); - return cute::min(aligned, kNumExpertsPerRank); - } - - CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const { - uint32_t valid_value = 0; - #pragma unroll - for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { - valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ? - stored_num_tokens_per_expert[i] : valid_value; - } - return ptx::exchange(valid_value, expert_idx % 32); - } - - CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) { - uint32_t num_blocks = 0; - #pragma unroll - for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { - if (i * 32 + ptx::get_lane_idx() < expert_idx) - num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M); - } - return __reduce_add_sync(0xffffffff, num_blocks); - } - - CUTLASS_DEVICE void advance_expert_idx() { - current_pool_block_offset += get_current_num_m_blocks(); - current_local_expert_idx += 1; - current_num_tokens = get_num_tokens(current_local_expert_idx); - } - - CUTLASS_DEVICE void set_expert_idx(const uint32_t& expert_idx) { - current_local_expert_idx = expert_idx; - current_num_tokens = get_num_tokens(expert_idx); - current_pool_block_offset = get_pool_block_offset(expert_idx); - } - - CUTLASS_DEVICE uint32_t get_current_pool_block_offset() const { - return current_pool_block_offset; - } - - CUTLASS_DEVICE uint32_t get_current_num_m_blocks() const { - return math::ceil_div(current_num_tokens, BLOCK_M); - } - - template - CUTLASS_DEVICE uint32_t get_valid_m() const { - const auto m = cute::min(current_num_tokens - m_block_idx * BLOCK_M, BLOCK_M); - return kDoUMMAAligned ? math::align(m, 16u) : m; - } - - CUTLASS_DEVICE bool fetch_next_l1_block() { - const auto wave_end_expert_idx = get_wave_expert_end_idx(); - while (current_local_expert_idx < wave_end_expert_idx) { - const auto num_m_blocks = get_current_num_m_blocks(); - m_block_idx = block_idx / kNumL1BlockNs; - if (m_block_idx < num_m_blocks) - return true; - block_idx -= num_m_blocks * kNumL1BlockNs; - advance_expert_idx(); - } - return false; - } - - CUTLASS_DEVICE bool fetch_next_l2_block() { - const auto wave_end_expert_idx = get_wave_expert_end_idx(); - while (current_local_expert_idx < wave_end_expert_idx) { - const auto num_m_blocks = get_current_num_m_blocks(); - if (block_idx < num_m_blocks * kNumL2BlockNs) { - m_block_idx = block_idx / kNumL2BlockNs; - return true; - } - block_idx -= num_m_blocks * kNumL2BlockNs; - advance_expert_idx(); - } - return false; - } - - CUTLASS_DEVICE cute::tuple get_next_block() { - while (current_local_expert_idx < kNumExpertsPerRank) { - if (next_phase == BlockPhase::Linear1) { - if (fetch_next_l1_block()) { - n_block_idx = block_idx - m_block_idx * kNumL1BlockNs; - block_idx += kNumSMs; - return {BlockPhase::Linear1, current_local_expert_idx, m_block_idx, n_block_idx}; - } - next_phase = BlockPhase::Linear2; - set_expert_idx(math::align(current_local_expert_idx - 1, kNumExpertsPerWave)); - } else { - if (fetch_next_l2_block()) { - n_block_idx = block_idx - m_block_idx * kNumL2BlockNs; - block_idx += kNumSMs; - return {BlockPhase::Linear2, current_local_expert_idx, m_block_idx, n_block_idx}; - } - next_phase = BlockPhase::Linear1; - } - } - return {BlockPhase::None, 0, 0, 0}; - } - - CUTLASS_DEVICE void fetch_expert_recv_count() { - #pragma unroll - for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { - const auto expert_idx = i * 32 + ptx::get_lane_idx(); - uint64_t value = 0; - if (expert_idx < kNumExpertsPerRank) { - do { - value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx)); - } while (static_cast(value >> 32) != kNumSMs * kNumRanks); - } - stored_num_tokens_per_expert[i] = static_cast(value); - } - __syncwarp(); - } - - template - CUTLASS_DEVICE void for_each_block(Func&& func) { - fetch_expert_recv_count(); - set_expert_idx(0); - while (true) { - CUTE_TIE_DECL(get_next_block(), block_phase, local_expert_idx, block_m_idx, block_n_idx); - if (block_phase == BlockPhase::None) - break; - func(block_phase, local_expert_idx, - block_phase == BlockPhase::Linear2 ? kNumL2BlockKs : kNumL1BlockKs, - block_m_idx, block_n_idx); - } - } -}; - #endif } // namespace deep_gemm::sched diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 37ea27c786..48fe8334da 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -15,6 +15,31 @@ from .. import _C + +def _check_legacy_fp4_acts_env() -> None: + """The fork's `DG_USE_FP4_ACTS` / `DG_USE_MXF4_KIND` env vars are superseded by + `mma_type`, and are refused rather than ignored. + + They cannot be silently remapped. `DG_USE_FP4_ACTS` selected FP4 activations under + `kind::mxf8f6f4` while leaving weights in the gran-8 *unpacked* gate/up interleave; + the equivalent typed kinds (`mxf4xmxf4` / `nvfp4xnvfp4`) are dense `kind::mxf4` and + need the gran-16 *packed* interleave from `_interleave_weights_packed_fp4`. Honouring + the env var here would leave a caller who does not also pass `mma_type` to + `transform_weights_for_mega_moe` handing packed-layout kernels unpacked weights — + silently wrong results. Fail loudly instead. + """ + for name in ('DG_USE_FP4_ACTS', 'DG_USE_MXF4_KIND'): + if int(os.environ.get(name, '0')) != 0: + raise RuntimeError( + f'`{name}` is no longer supported; it is replaced by the symmetric ' + f"buffer's `mma_type`. Pass mma_type='mxf4xmxf4' (per-32 UE8M0 scales) " + f"or 'nvfp4xnvfp4' (per-16 UE4M3) to get_symm_buffer_for_mega_moe(), and " + f'the SAME mma_type to transform_weights_for_mega_moe() — the FP4 kinds ' + f'need the packed gate/up interleave, so weights transformed without it ' + f'would be laid out wrongly.' + ) + + class SymmBuffer: def __init__(self, group: dist.ProcessGroup, num_experts: int, @@ -23,15 +48,19 @@ def __init__(self, group: dist.ProcessGroup, num_shared_experts: int = 0, mma_type: str = 'fp8xfp4', activation: str = 'swiglu'): - assert activation in ('swiglu', 'situ'), f'Unsupported activation `{activation}`' + assert activation in ('swiglu', 'swigluoai', 'situ'), f'Unsupported activation `{activation}`' assert activation != 'situ' or mma_type == 'fp8xfp4', \ '`situ` activation is supported only for `fp8xfp4` MegaMoE' + _check_legacy_fp4_acts_env() 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 + self.activation = activation # Allocate a symmetric buffer num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe( @@ -52,36 +81,15 @@ def __init__(self, group: dist.ProcessGroup, self.group.barrier() torch.cuda.synchronize() - # Create input buffer views. TVM-FFI transports float8 storage as - # int8, so restore the logical dtype without changing the bytes. - raw_buffers = slice_input_buffers(self.buffer) - use_fp4_acts = (mma_type != 'bf16xbf16' and num_shared_experts == 0 and - int(os.environ.get('DG_USE_FP4_ACTS', '0')) != 0) - acts_dtype = (torch.bfloat16 if mma_type == 'bf16xbf16' else - (torch.int8 if use_fp4_acts else torch.float8_e4m3fn)) - l2_dtype = torch.bfloat16 if mma_type == 'bf16xbf16' else torch.float8_e4m3fn - - def as_torch(tensor, dtype=None): - if tensor is None: - return None - tensor = torch.from_dlpack(tensor) - return tensor.view(dtype) if dtype is not None and tensor.dtype != dtype else tensor - - (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) = raw_buffers - self.x = as_torch(x, acts_dtype) - self.x_sf = as_torch(x_sf) - self.topk_idx = as_torch(topk_idx) - self.topk_weights = as_torch(topk_weights) - self.shared_l1_acts = as_torch(shared_l1_acts, acts_dtype) - self.shared_l1_acts_sf = as_torch(shared_l1_acts_sf) - self.shared_l2_acts = as_torch(shared_l2_acts, l2_dtype) - self.shared_l2_acts_sf = as_torch(shared_l2_acts_sf) - self.l1_acts = as_torch(l1_acts, acts_dtype) - self.l1_acts_sf = as_torch(l1_acts_sf) - self.l2_acts = as_torch(l2_acts, l2_dtype) - self.l2_acts_sf = as_torch(l2_acts_sf) + # Create input buffer views (as torch tensors, not tvm-ffi tensors). + (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, + self.x_scales) = map( + torch.from_dlpack, slice_input_buffers(self.buffer)) def destroy(self): self.handle = None @@ -136,6 +144,22 @@ def _interleave_weights(t: torch.Tensor, gran: int = 8) -> torch.Tensor: return result.squeeze(0) if squeeze_group_dim else result +def _interleave_weights_packed_fp4(t: torch.Tensor) -> torch.Tensor: + assert t.dim() in (2, 3) + squeeze_group_dim = t.dim() == 2 + if squeeze_group_dim: + t = t.unsqueeze(0) + + g, n, *rest = t.shape + half = n // 2 + assert half % 16 == 0 + gate = t[:, :half].reshape(g, half // 16, 16, *rest) + up = t[:, half:].reshape(g, half // 16, 16, *rest) + result = torch.cat([gate[:, :, 0::2], up[:, :, 0::2], gate[:, :, 1::2], up[:, :, 1::2]], dim=2) + result = torch.empty_like(t).copy_(result.reshape(g, n, *rest)) + return result.squeeze(0) if squeeze_group_dim else result + + def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor: # Unsqueeze for 2D assert sf.dtype == torch.int and sf.dim() in (2, 3) @@ -156,17 +180,20 @@ def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor: def transform_weights_for_mega_moe( l1_weights: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], l2_weights: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], - activation: str = 'swiglu' + activation: str = 'swiglu', + mma_type: str = 'fp8xfp4' ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]: - assert activation in ('swiglu', 'situ'), f'Unsupported activation `{activation}`' + assert activation in ('swiglu', 'swigluoai', 'situ'), f'Unsupported activation `{activation}`' assert activation != 'situ' or ( isinstance(l1_weights, tuple) and isinstance(l2_weights, tuple) ), '`situ` activation is supported only for FP8xFP4 MegaMoE weights' if isinstance(l1_weights, tuple): - # FP8: interleave gate/up for weight and SF, then transpose L1 SF for UTCCP - l1_w = _interleave_weights(l1_weights[0]) - l1_sf = _transpose_sf_for_utccp(_interleave_weights(l1_weights[1])) + # FP8/MXFP4/NVFP4: interleave gate/up for weight and SF, then transpose L1 SF for UTCCP. + interleave = _interleave_weights_packed_fp4 if mma_type in ('mxf4xmxf4', 'nvfp4xnvfp4') \ + else _interleave_weights + l1_w = interleave(l1_weights[0]) + l1_sf = _transpose_sf_for_utccp(interleave(l1_weights[1])) l1_transformed = (l1_w, l1_sf) # L2: only transpose SF for UTCCP l2_transformed = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1])) @@ -188,32 +215,62 @@ def fp8_fp4_mega_moe(y: torch.Tensor, recipe: Tuple[int, int, int] = (1, 1, 32), activation: str = 'swiglu', activation_clamp: Optional[float] = None, - fast_math: bool = True): + fast_math: bool = True, + use_x_scales: bool = False, + l1_alphas: Optional[torch.Tensor] = None, + l2_alphas: Optional[torch.Tensor] = None, + l2_act_scales: Optional[torch.Tensor] = None): + if use_x_scales and sym_buffer.mma_type != 'nvfp4xnvfp4': + raise ValueError( + '`use_x_scales` is only supported for `nvfp4xnvfp4`; ' + f'got mma_type={sym_buffer.mma_type!r}' + ) + (l1_weights_data, l1_weights_sf) = l1_weights (l2_weights_data, l2_weights_sf) = l2_weights - assert (shared_l1_weights is None) == (shared_l2_weights is None) - if shared_l1_weights is None: - shared_l1_weights_data = shared_l1_weights_sf = None - shared_l2_weights_data = shared_l2_weights_sf = None - else: - (shared_l1_weights_data, shared_l1_weights_sf) = shared_l1_weights - (shared_l2_weights_data, shared_l2_weights_sf) = shared_l2_weights + (shared_l1_data, shared_l1_sf) = shared_l1_weights if shared_l1_weights is not None else (None, None) + (shared_l2_data, shared_l2_sf) = shared_l2_weights if shared_l2_weights is not None else (None, None) _C.fp8_fp4_mega_moe( y, l1_weights_data, l1_weights_sf, l2_weights_data, l2_weights_sf, - shared_l1_weights_data, shared_l1_weights_sf, - shared_l2_weights_data, shared_l2_weights_sf, + shared_l1_data, shared_l1_sf, + shared_l2_data, shared_l2_sf, 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, + recipe, sym_buffer.mma_type, activation, activation_clamp, - fast_math + fast_math, use_x_scales, l1_alphas, l2_alphas, l2_act_scales ) +def nvfp4_mega_moe(y: torch.Tensor, + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], + sym_buffer: SymmBuffer, + 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, + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True, + use_x_scales: bool = False, + l1_alphas: Optional[torch.Tensor] = None, + l2_alphas: Optional[torch.Tensor] = None, + l2_act_scales: Optional[torch.Tensor] = None): + fp8_fp4_mega_moe( + y, l1_weights, l2_weights, sym_buffer, + shared_l1_weights, shared_l2_weights, + cumulative_local_expert_recv_stats, + recipe=(1, 1, 16), + activation=activation, activation_clamp=activation_clamp, + fast_math=fast_math, use_x_scales=use_x_scales, + l1_alphas=l1_alphas, l2_alphas=l2_alphas, l2_act_scales=l2_act_scales + ) + + def bf16_mega_moe(y: torch.Tensor, l1_weights: torch.Tensor, l2_weights: torch.Tensor, @@ -242,6 +299,14 @@ def bf16_mega_moe(y: torch.Tensor, ) +def get_block_m_for_mega_moe(num_ranks: int, num_experts: int, + num_max_tokens_per_rank: int, num_tokens: int, + num_topk: int, mma_type: str = 'fp8xfp4') -> int: + """`BLOCK_M` the kernel will pick — callers need it to lay out shared-expert SFs.""" + return int(_C.get_block_m_for_mega_moe( + num_ranks, num_experts, num_max_tokens_per_rank, num_tokens, num_topk, mma_type)) + + def mega_moe_pre_dispatch(x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, @@ -251,11 +316,14 @@ def mega_moe_pre_dispatch(x: torch.Tensor, buf_topk_weights: torch.Tensor, num_tokens: int, group_size: int = 32, - use_fp4_acts: bool = False) -> None: + mma_type: str = 'fp8xfp4', + buf_x_scales: Optional[torch.Tensor] = None, + expert_scales: Optional[torch.Tensor] = None) -> None: _C.mega_moe_pre_dispatch( x, topk_idx, topk_weights, buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, - num_tokens, group_size, use_fp4_acts, + num_tokens, group_size, mma_type, + buf_x_scales, expert_scales, ) diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index a3caf5b051..55fa1412be 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -122,6 +122,45 @@ def per_token_cast_to_fp4(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128, return packed[:, :n // 2].contiguous(), sf +def cast_to_ue4m3(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float8_e4m3fn).float() + + +def pack_ue4m3_to_int(x: torch.Tensor) -> torch.Tensor: + """Pack 4 UE4M3 scaling factor bytes into one int32, matching the kernel's SF word.""" + assert x.size(-1) % 4 == 0 + return x.to(torch.float8_e4m3fn).contiguous().view(torch.uint8).view(torch.int) + + +def unpack_ue4m3_from_int(packed_sf: torch.Tensor) -> torch.Tensor: + return packed_sf.view(torch.uint8).view(torch.float8_e4m3fn).float() + + +def per_token_cast_to_nvfp4(x: torch.Tensor, gran_k: int = 16, + use_packed_ue4m3: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: + """NVFP4: E2M1 values with one UE4M3 (not pow2 UE8M0) scaling factor per `gran_k` elements.""" + m, n = x.shape + assert n % gran_k == 0 + x_view = x.view(m, -1, gran_k) + sf = cast_to_ue4m3(x_view.abs().float().amax(dim=2) / 6.0) + sf_inv = torch.where(sf > 0, 1.0 / sf, torch.zeros_like(sf)) + codes = _quantize_to_fp4_e2m1(x_view.float() * sf_inv.unsqueeze(2)).view(m, n) + codes2 = codes.view(m, n // 2, 2) + packed = (codes2[:, :, 0] & 0x0F) | ((codes2[:, :, 1] & 0x0F) << 4) + return packed.contiguous(), pack_ue4m3_to_int(sf) if use_packed_ue4m3 else sf + + +def transform_ue4m3_sf_into_required_layout(sf: torch.Tensor, mn: int) -> torch.Tensor: + """MN-major, TMA-aligned, int32-packed UE4M3 SFs — the weight-side layout the kernel reads. + + `transform_sf_into_required_layout` only knows how to pack UE8M0, so NVFP4 weight SFs are + packed here and handed to its already-packed `(INT, 1, gran_k)` branch. + """ + assert sf.dim() in (2, 3) and sf.size(-2) == mn + assert mn % 4 == 0, f'MN must be TMA-aligned for int32 SFs, got {mn}' + return pack_ue4m3_to_int(sf).transpose(-1, -2).contiguous().transpose(-1, -2) + + def transpose_packed_fp4(a: torch.Tensor) -> torch.Tensor: assert a.dtype == torch.int8 assert a.dim() == 2 diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index cccd2d33a2..b836aa1592 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -269,9 +269,11 @@ def k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, grouped_layout, c=None, compi SymmBuffer, transform_weights_for_mega_moe, fp8_fp4_mega_moe, + nvfp4_mega_moe, bf16_mega_moe, mega_moe_pre_dispatch, mega_moe_pre_dispatch_sm90, + get_block_m_for_mega_moe, ) @@ -289,7 +291,8 @@ def __init__(self, group, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, use_fp8_dispatch: bool = True, - activation: str = 'swiglu'): + activation: str = 'swiglu', + num_shared_experts: int = 0): import torch.distributed._symmetric_memory as symm_mem self.group = group @@ -298,12 +301,17 @@ def __init__(self, group, self.num_topk = num_topk self.hidden = hidden self.intermediate_hidden = intermediate_hidden + # Fused shared expert (SM90 only). 0 disables it and keeps + # the symmetric-buffer layout byte-identical to the routed-only path. + self.num_shared_experts = num_shared_experts + self.shared_intermediate_hidden = intermediate_hidden * num_shared_experts 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, - use_fp8_dispatch, activation + use_fp8_dispatch, activation, + num_shared_experts ) self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda') self.handle = symm_mem.rendezvous(self.buffer, group=group) @@ -312,7 +320,8 @@ def __init__(self, group, torch.cuda.synchronize() (x, x_sf, topk_idx, topk_weights, - l1_acts, l1_acts_sf, l2_acts, l2_acts_sf) = slice_input_buffers(self.buffer) + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf) = slice_input_buffers(self.buffer) self.x = _from_dlpack_if_needed(x, torch.float8_e4m3fn) self.x_sf = _from_dlpack_if_needed(x_sf) self.topk_idx = _from_dlpack_if_needed(topk_idx) @@ -321,6 +330,15 @@ def __init__(self, group, self.l1_acts_sf = _from_dlpack_if_needed(l1_acts_sf) self.l2_acts = _from_dlpack_if_needed(l2_acts, torch.float8_e4m3fn) self.l2_acts_sf = _from_dlpack_if_needed(l2_acts_sf) + # The fused shared expert reads its L1 activations from the same `x` region; + # only the post-SwiGLU pool and its SF are extra buffers (zero-sized when the + # shared expert is disabled). + self.shared_l1_acts = self.x + self.shared_l1_acts_sf = self.x_sf + shared_l2_acts = _from_dlpack_if_needed(shared_l2_acts, torch.float8_e4m3fn) + shared_l2_acts_sf = _from_dlpack_if_needed(shared_l2_acts_sf) + self.shared_l2_acts = shared_l2_acts if num_shared_experts > 0 else None + self.shared_l2_acts_sf = shared_l2_acts_sf if num_shared_experts > 0 else None def destroy(self): self.handle = None @@ -328,6 +346,10 @@ def destroy(self): self.group = None self.x = None self.x_sf = None + self.shared_l1_acts = None + self.shared_l1_acts_sf = None + self.shared_l2_acts = None + self.shared_l2_acts_sf = None def get_symm_buffer_for_sm90_mega_moe(group, @@ -335,7 +357,8 @@ def get_symm_buffer_for_sm90_mega_moe(group, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, use_fp8_dispatch: bool = True, - activation: str = 'swiglu') -> SM90SymmBuffer: + activation: str = 'swiglu', + num_shared_experts: int = 0) -> SM90SymmBuffer: from .utils.math import align num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) @@ -343,7 +366,8 @@ def get_symm_buffer_for_sm90_mega_moe(group, group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - use_fp8_dispatch, activation + use_fp8_dispatch, activation, + num_shared_experts ) @@ -351,10 +375,10 @@ def get_symm_buffer_for_mega_moe(group, num_experts: int, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, - num_shared_experts: int = 0, use_fp8_dispatch: Union[bool, None] = None, mma_type: str = 'fp8xfp4', - activation: str = 'swiglu'): + activation: str = 'swiglu', + num_shared_experts: int = 0): if use_fp8_dispatch is not None: assert use_fp8_dispatch == (mma_type.split('x')[0] == 'fp8') if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9: @@ -363,8 +387,10 @@ def get_symm_buffer_for_mega_moe(group, group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - True, activation + True, activation, + num_shared_experts ) + assert num_shared_experts == 0, 'Shared experts are only wired through the SM90 buffer' return mega.get_symm_buffer_for_mega_moe( group, num_experts, num_max_tokens_per_rank, num_topk, @@ -376,20 +402,39 @@ def get_symm_buffer_for_mega_moe(group, ) +def _interleave_gate_up_sm90(t: torch.Tensor, gran: int = 8) -> torch.Tensor: + """Interleave the gate/up halves of an L1 weight (or its SF) along N. + + Handles the routed `(G, 2*N, ...)` layout and the shared expert's dense + `(2*N, ...)` one, which the kernel's weight-SF indexing assumes for both. + """ + if t.dim() == 2: + return _interleave_gate_up_sm90(t.unsqueeze(0), gran).squeeze(0) + g, n, *rest = t.shape + half = n // 2 + gate = t[:, :half].reshape(g, half // gran, gran, *rest) + up = t[:, half:].reshape(g, half // gran, gran, *rest) + return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest)) + + def transform_weights_for_mega_moe_sm90( l1_weights: Tuple[torch.Tensor, torch.Tensor], l2_weights: Tuple[torch.Tensor, torch.Tensor] ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: l1_fp8, l1_sf = l1_weights + return (_interleave_gate_up_sm90(l1_fp8), l1_sf), l2_weights + - def _interleave_one(t, gran: int = 8) -> torch.Tensor: - g, n, *rest = t.shape - half = n // 2 - gate = t[:, :half].reshape(g, half // gran, gran, *rest) - up = t[:, half:].reshape(g, half // gran, gran, *rest) - return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest)) +def transform_shared_weights_for_mega_moe_sm90( + shared_l1_weights: Tuple[torch.Tensor, torch.Tensor], + shared_l2_weights: Tuple[torch.Tensor, torch.Tensor] +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + """Same gate/up interleave as the routed L1, on the dense shared-expert weights. - return (_interleave_one(l1_fp8), l1_sf), l2_weights + Shapes: L1 `(2 * shared_intermediate_hidden, hidden)`, L2 `(hidden, shared_intermediate_hidden)`. + """ + shared_l1_fp8, shared_l1_sf = shared_l1_weights + return (_interleave_gate_up_sm90(shared_l1_fp8), shared_l1_sf), shared_l2_weights def fp8_mega_moe(y: torch.Tensor, @@ -400,13 +445,23 @@ def fp8_mega_moe(y: torch.Tensor, recipe: Tuple[int, int, int] = (128, 128, 128), activation: str = 'swiglu', activation_clamp: Optional[float] = None, - fast_math: bool = True): + fast_math: bool = True, + shared_l1_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l2_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None): (l1_weights_data, l1_weights_sf) = l1_weights (l2_weights_data, l2_weights_sf) = l2_weights + # Fused shared expert: needs a symmetric buffer allocated with + # `num_shared_experts > 0` (see `get_symm_buffer_for_mega_moe`). + assert (shared_l1_weights is None) == (shared_l2_weights is None), \ + 'Shared-expert L1 and L2 weights must be passed together' + shared_l1_weights_data, shared_l1_weights_sf = shared_l1_weights or (None, None) + shared_l2_weights_data, shared_l2_weights_sf = shared_l2_weights or (None, None) _C.fp8_mega_moe( y, l1_weights_data, l1_weights_sf, l2_weights_data, l2_weights_sf, + shared_l1_weights_data, shared_l1_weights_sf, + shared_l2_weights_data, shared_l2_weights_sf, cumulative_local_expert_recv_stats, sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index ab983d0dd1..b89ea2506e 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/sgl_deep_gemm/run_tests.sh @@ -156,6 +156,7 @@ MEGA_MOE_BLACKWELL=( test_mega_moe_situ.py test_mega_moe_l1_fp4_accuracy.py test_mega_moe_l1_sentinel.py + test_mega_moe_nvfp4_alphas.py test_mega_moe_pre_dispatch.py ) MEGA_MOE_HOPPER=( diff --git a/sgl_deep_gemm/tests/test_mega_moe.py b/sgl_deep_gemm/tests/test_mega_moe.py index 1d8512761c..d29fd0f574 100644 --- a/sgl_deep_gemm/tests/test_mega_moe.py +++ b/sgl_deep_gemm/tests/test_mega_moe.py @@ -7,7 +7,8 @@ from typing import Tuple import deep_gemm -from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8 +from deep_gemm.utils import (per_token_cast_to_fp4, per_token_cast_to_fp8, per_token_cast_to_nvfp4, + transform_ue4m3_sf_into_required_layout) from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather from deep_gemm.testing import bench_kineto @@ -18,6 +19,7 @@ def import_baseline(): # noinspection PyBroadException try: import deep_ep + assert hasattr(deep_ep, 'ElasticBuffer'), 'deep_ep has no ElasticBuffer (DeepEP v2 required)' import importlib.util from tilelang.profiler.bench import do_bench spec = importlib.util.spec_from_file_location( @@ -42,6 +44,8 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): # Settings is_bf16xbf16 = args.mma_type == 'bf16xbf16' + use_nvfp4 = args.mma_type == 'nvfp4xnvfp4' + gran_k = 16 if use_nvfp4 else 32 num_max_tokens_per_rank = args.num_max_tokens_per_rank num_tokens = max(0, args.num_max_tokens_per_rank - random.randint(0, args.num_max_removed_tokens)) \ if args.num_tokens == 0 else args.num_tokens @@ -55,17 +59,22 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - mma_type=args.mma_type + mma_type=args.mma_type, + activation=args.activation ) # Cast weights into FP4 def _cast_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: num_groups, n, k = bf16_weights.shape w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) - w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + w_sf = torch.empty((num_groups, n, k // gran_k), device='cuda', dtype=torch.float) for i in range(num_groups): - w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) - w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + if use_nvfp4: + w[i], w_sf[i] = per_token_cast_to_nvfp4(bf16_weights[i], gran_k=gran_k) + else: + w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=gran_k) + w_sf = transform_ue4m3_sf_into_required_layout(w_sf, n) if use_nvfp4 else \ + deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, gran_k), num_groups) return w, w_sf # Create inputs @@ -96,7 +105,9 @@ def create_inputs(): # Stream A0.0b: when the flag is on, the symm buffer's `x` slot is sized # for packed E2M1 (`hidden/2` bytes/token), so we must quantize at the # source to match. - if os.environ.get('DG_USE_FP4_ACTS', '0') != '0': + if use_nvfp4: + x = per_token_cast_to_nvfp4(x, gran_k=gran_k, use_packed_ue4m3=True) + elif args.mma_type == 'mxf4xmxf4': x = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) else: x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) @@ -105,7 +116,7 @@ def create_inputs(): l2_weights = _cast_weights_to_fp4(l2_weights) transformed_l1_weights, transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)) + deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights, args.activation, args.mma_type)) # Run fused mega MoE # NOTES: copy x into buffer before each call because debug mode zeros the entire buffer @@ -123,8 +134,11 @@ def run_fused(): y=y, l1_weights=transformed_l1_weights, l2_weights=transformed_l2_weights, sym_buffer=buffer, cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, - activation_clamp=args.activation_clamp, + activation=args.activation, + activation_clamp=args.activation_clamp if args.activation != 'situ' else None, fast_math=bool(args.fast_math)) + if not is_bf16xbf16: + kernel_kwargs['recipe'] = (1, 1, gran_k) (deep_gemm.bf16_mega_moe if is_bf16xbf16 else deep_gemm.fp8_fp4_mega_moe)(**kernel_kwargs) return y, cumulative_local_expert_recv_stats_fused @@ -209,10 +223,14 @@ def run_baseline(): # Combine return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline - # Check correctness (must be bitwise identical) + # Check correctness (must be bitwise identical). num_correctness_tests = 1 if args.num_correctness_tests is None else args.num_correctness_tests + can_compare_bitwise = is_bf16xbf16 or args.mma_type == 'fp8xfp4' + if is_legacy_loaded and num_correctness_tests > 0 and not can_compare_bitwise: + dist_print(f'Skipping bitwise correctness vs legacy baseline (FP8-acts only, ' + f'undefined for `{args.mma_type}`)', once_in_node=True) # noinspection PyBroadException - if is_legacy_loaded and num_correctness_tests > 0: + if is_legacy_loaded and num_correctness_tests > 0 and can_compare_bitwise: dist_print('Running correctness tests:', once_in_node=True) for i in range(num_correctness_tests): create_inputs() @@ -299,7 +317,10 @@ def run_baseline(): parser.add_argument('--num-topk', type=int, default=6, help='Number of expert selections') parser.add_argument('--masked-ratio', type=float, default=0.0, help='Mask some expert selections') parser.add_argument('--fast-math', type=int, default=1, help='Enable fast math (0 or 1, default: 1)') - parser.add_argument('--mma-type', type=str, default='fp8xfp4', help='MMA type: fp8xfp4 or bf16xbf16') + parser.add_argument('--mma-type', type=str, default='fp8xfp4', + choices=('fp8xfp4', 'mxf4xmxf4', 'nvfp4xnvfp4', 'bf16xbf16'), help='MMA type') + parser.add_argument('--activation', type=str, default='swiglu', choices=('swiglu', 'situ', 'swigluoai'), + help='Activation: swiglu, situ, or swigluoai') # Test settings parser.add_argument('--num-correctness-tests', type=int, default=None, help='Pressure test') diff --git a/sgl_deep_gemm/tests/test_mega_moe_hopper.py b/sgl_deep_gemm/tests/test_mega_moe_hopper.py index f3579a7e91..5818a1ccda 100644 --- a/sgl_deep_gemm/tests/test_mega_moe_hopper.py +++ b/sgl_deep_gemm/tests/test_mega_moe_hopper.py @@ -11,6 +11,12 @@ per-128-K L2 activation SF, while the fused SM90 MegaMoE L1 epilogue writes per-64-K L2 activation SF to avoid cross-CTA synchronization. This is a same-pipeline performance reference, not a bitwise correctness oracle. +* shared expert (optional, ``--num-shared-experts``): DeepSeek-style shared + expert, fused into the kernel as the SM100-style ``SharedLinear1`` / + ``SharedLinear2`` scheduler phases and reduced through an extra combine slot, so + one launch produces routed + shared. Both baselines run the same shared expert + serially, matching ``tests/test_mega_moe.py``, which folds it in as a DeepEP + combine bias. * low-latency baseline (optional, ``--run-low-latency-baseline``): mirrors the sglang low-latency MoE pipeline (see ``sglang/srt/layers/moe/token_dispatcher/deepep.py::_DeepEPDispatcherImplLowLatency``): @@ -25,6 +31,8 @@ * accuracy mode (optional, ``--accuracy``): runs the former layered SM90 correctness suite with a PyTorch BF16/FP32 reference. It covers smoke, heuristic branches, shape sweeps, edge cases, and optional random stress. + Layer 6 adds the fused shared expert (``num_shared_experts >= 1``); the + reference is routed + dense shared MLP, compared with the fused kernel output. * output: TFLOPS, overlap-adjusted TFLOPS, HBM GB/s, NVLink GB/s, fused time, reduction estimate, and ``t_baseline / t_fused``. """ @@ -253,6 +261,18 @@ def _quantize_grouped_fp8_block_128_128( return w_fp8.view(g, n, k).contiguous(), sf.contiguous() +def _quantize_dense_fp8_block_128_128( + w: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(N, K) bf16 -> (N, K) fp8_e4m3fn plus (N/128, K/128) FP32 block SF. + + Used for the shared expert, which is a single dense MLP rather than a group + of per-rank experts. + """ + w_fp8, sf = _quantize_grouped_fp8_block_128_128(w.unsqueeze(0)) + return w_fp8.squeeze(0), sf.squeeze(0) + + # ============================================================================ # Section 3: layered accuracy reference and scenarios. # ============================================================================ @@ -380,6 +400,54 @@ def _reference_fused( return y_full_bf16[start:end].contiguous() +def _reference_shared( + x_fp8_local: torch.Tensor, + x_sf_local: torch.Tensor, + shared_l1: Tuple[torch.Tensor, torch.Tensor], + shared_l2: Tuple[torch.Tensor, torch.Tensor], + hidden: int, + shared_intermediate_hidden: int, + activation_clamp: float, +) -> torch.Tensor: + """PyTorch FP32 reference for the fused shared expert, on this rank's local + tokens (the shared expert is node-local — no all_gather, unlike the routed + path). Mirrors the kernel's fused-shared arithmetic: + + L1: (M, H) @ (2*SIH, H)^T -> (M, 2*SIH) + SwiGLU (gate clamp one-sided, up clamp two-sided), no topk weighting + L1 output FP8 round-trip with per-64-K scales (matches FUSED_L2_ACT_SF_GRAN) + L2: (M, SIH) @ (H, SIH)^T -> (M, H) -> bf16 + + The fused kernel adds the shared output into the routed output; this returns + the shared contribution so the caller can do ``y_ref routed + y_ref shared``. + """ + assert shared_intermediate_hidden % 64 == 0, ( + "shared_intermediate_hidden must be a multiple of 64 for the per-64-K " + "L2 activation SF granularity the fused epilogue uses" + ) + x = _dequant_per_token_per_128_k(x_fp8_local, x_sf_local) # (M, H) fp32 + + # L1 GEMM (dense, the shared expert is a single MLP). + l1_w = _dequant_block_128_128(shared_l1[0], shared_l1[1]) # (2*SIH, H) fp32 + l1_y = x @ l1_w.t() # (M, 2*SIH) + l1_y = _swiglu_fp32(l1_y, activation_clamp) # (M, SIH) + + # L1 output FP8 round-trip with per-64-K scales, matching the fused SM90 L2 + # activation SF granularity (FUSED_L2_ACT_SF_GRAN = 64). + s, ih = l1_y.shape + assert ih == shared_intermediate_hidden + v = l1_y.view(s, ih // 64, 64) + sf2 = v.abs().amax(dim=-1).clamp(1e-4) / FP8_E4M3_MAX + l2_in = ( + (v / sf2.unsqueeze(-1)).to(torch.float8_e4m3fn).float() + * sf2.unsqueeze(-1) + ).view(s, ih) + + # L2 GEMM -> bf16. + l2_w = _dequant_block_128_128(shared_l2[0], shared_l2[1]) # (H, SIH) fp32 + return (l2_in @ l2_w.t()).to(torch.bfloat16) + + def _run_accuracy_scenario( name: str, cfg: Dict[str, Any], @@ -397,6 +465,7 @@ def _run_accuracy_scenario( masked_ratio = cfg.get("masked_ratio", 0.0) activation_clamp = cfg.get("activation_clamp", 10.0) fast_math = cfg.get("fast_math", True) + num_shared_experts = cfg.get("num_shared_experts", 0) assert num_experts % num_ranks == 0, ( f"{name}: experts {num_experts} not divisible by ranks {num_ranks}" @@ -404,6 +473,14 @@ def _run_accuracy_scenario( num_experts_per_rank = num_experts // num_ranks assert num_tokens <= num_max assert hidden % 128 == 0 and intermediate_hidden % 128 == 0 + # Fused shared expert: each adds one routed intermediate size, fused into the + # mega kernel via SharedLinear1/2. + shared_intermediate_hidden = intermediate_hidden * num_shared_experts + if num_shared_experts > 0: + assert shared_intermediate_hidden % 64 == 0, ( + f"{name}: shared_intermediate_hidden={shared_intermediate_hidden} must be " + f"a multiple of 64 (fused L2 activation SF granularity)" + ) verbose = bool(int(os.environ.get("DG_TEST_VERBOSE", "0"))) @@ -445,6 +522,31 @@ def trace(stage: str): transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe_sm90( l1_weights, l2_weights ) + # Shared expert weights (dense MLP, block-(128,128) FP8 like the routed + # weights) and the matching interleave the kernel's weight-SF indexing assumes. + if num_shared_experts > 0: + shared_l1_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (shared_intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + shared_l2_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (hidden, shared_intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + transformed_shared_l1, transformed_shared_l2 = ( + deep_gemm.transform_shared_weights_for_mega_moe_sm90( + shared_l1_weights, shared_l2_weights + ) + ) + else: + shared_l1_weights = shared_l2_weights = None + transformed_shared_l1 = transformed_shared_l2 = None trace("alloc_symm_buffer") buffer = deep_gemm.get_symm_buffer_for_mega_moe( @@ -454,6 +556,7 @@ def trace(stage: str): num_topk, hidden, intermediate_hidden, + num_shared_experts=num_shared_experts, ) cum_stats = torch.zeros((num_experts_per_rank,), dtype=torch.int, device="cuda") @@ -475,6 +578,8 @@ def trace(stage: str): activation="swiglu", activation_clamp=activation_clamp if math.isfinite(activation_clamp) else None, fast_math=fast_math, + shared_l1_weights=transformed_shared_l1, + shared_l2_weights=transformed_shared_l2, ) torch.cuda.synchronize() @@ -497,6 +602,17 @@ def trace(stage: str): intermediate_hidden, activation_clamp, ) + # Add the shared-expert reference contribution: fused-out = routed + shared. + if num_shared_experts > 0: + y_ref = y_ref + _reference_shared( + x_fp8[0], + x_fp8[1], + shared_l1_weights, + shared_l2_weights, + hidden, + shared_intermediate_hidden, + activation_clamp, + ) diff = calc_diff(y_fused, y_ref) ok = diff < diff_tol @@ -614,6 +730,54 @@ def _accuracy_layer5_stress(num_ranks: int, num_tests: int) -> List[Tuple[str, D return out +def _accuracy_layer6_shared_expert(num_ranks: int) -> List[Tuple[str, Dict[str, Any]]]: + """Fused shared-expert correctness (fused-out = routed + shared). + + Each scenario sets num_shared_experts >= 1, which routes through the + fused-shared path and compares against routed reference + dense shared MLP. + shared_intermediate_hidden must be a multiple of 64, so intermediate_hidden is + picked from {512, 1024, 2048}. + """ + base = dict( + num_max_tokens_per_rank=128, + hidden=512, + intermediate_hidden=512, # * num_shared_experts stays a multiple of 64 + num_experts=8 * num_ranks, + num_topk=2, + num_shared_experts=1, + ) + out = [] + # Smoke: 1 shared expert, default clamp/fast_math/shape. + out.append(("L6.sh1.smoke", dict(base))) + # Two shared experts (shared_intermediate_hidden = 2*ih). + cfg = dict(base) + cfg.update(num_shared_experts=2) + out.append(("L6.sh2", cfg)) + # Larger hidden / intermediate with the shared expert on. + for hidden, ih in [(2048, 1024), (2048, 2048)]: + cfg = dict(base) + cfg.update(hidden=hidden, intermediate_hidden=ih) + out.append((f"L6.h{hidden}_ih{ih}", cfg)) + # Shared expert under masking (shared acts on local tokens regardless of routing). + for masked_ratio in (0.3, 0.7): + cfg = dict(base) + cfg.update(masked_ratio=masked_ratio) + out.append((f"L6.mask{masked_ratio:.1f}", cfg)) + # Shared expert with clamp variations. + for clamp in (1.0, math.inf): + cfg = dict(base) + cfg.update(activation_clamp=clamp) + out.append((f"L6.clamp{clamp}", cfg)) + # Shared expert with fewer/more routed experts selected. + for topk in (1, 4): + if topk > base["num_experts"]: + continue + cfg = dict(base) + cfg.update(num_topk=topk) + out.append((f"L6.topk{topk}", cfg)) + return out + + def _run_accuracy_tests(local_rank: int, num_local_ranks: int, args: argparse.Namespace): rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) @@ -636,6 +800,8 @@ def _run_accuracy_tests(local_rank: int, num_local_ranks: int, args: argparse.Na layers += _accuracy_layer4_edges(num_ranks) if 5 in args.layers: layers += _accuracy_layer5_stress(num_ranks, args.num_correctness_tests or 8) + if 6 in args.layers: + layers += _accuracy_layer6_shared_expert(num_ranks) if args.filter: layers = [(name, cfg) for name, cfg in layers if args.filter in name] @@ -644,6 +810,8 @@ def _run_accuracy_tests(local_rank: int, num_local_ranks: int, args: argparse.Na f"layers {sorted(args.layers)} on {num_ranks} ranks", once_in_node=True, ) + # --num-shared-experts is a config for the benchmark mode; in accuracy mode + # each scenario carries its own num_shared_experts (layer 6 turns it on). failures: List[str] = [] for name, cfg in layers: @@ -1115,6 +1283,12 @@ def _run_fused_only_sweep(local_rank: int, num_local_ranks: int, args: argparse. f"masked_ratio={args.masked_ratio} fast_math={bool(args.fast_math)}", once_in_node=True, ) + if args.num_shared_experts > 0: + dist_print( + " > note: --num-shared-experts is ignored here; this mode measures the " + "routed kernel alone", + once_in_node=True, + ) num_max_tokens_per_rank = max(batches) for num_tokens in batches: @@ -1192,6 +1366,14 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): f"SM90 fused kernel requires intermediate_hidden <= 4096, got {intermediate_hidden}" ) + # Shared-expert shape, following tests/test_mega_moe.py: one routed + # intermediate size per shared expert. When enabled it is fused into the mega + # kernel as the SharedLinear1/SharedLinear2 phases; 0 disables it. + num_shared_experts = args.num_shared_experts + shared_intermediate_hidden = intermediate_hidden * num_shared_experts + assert shared_intermediate_hidden % 128 == 0 + fused_shared = num_shared_experts > 0 + # ---- Create BF16 token and weight inputs ---- # x: local tokens for this rank. x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") @@ -1241,6 +1423,38 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): l1_weights, l2_weights ) + # Shared-expert weights: one dense MLP (not per-expert), quantized with the + # same block-(128, 128) FP8 recipe as the routed weights. The fused and + # baseline paths share them so the comparison stays apples-to-apples. + if num_shared_experts > 0: + shared_l1_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (shared_intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + shared_l2_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (hidden, shared_intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + else: + shared_l1_weights = shared_l2_weights = None + + # The fused kernel consumes the same shared weights with the gate/up gran-8 + # interleave applied to L1 (identical to the routed weight transform). + if fused_shared: + transformed_shared_l1, transformed_shared_l2 = ( + deep_gemm.transform_shared_weights_for_mega_moe_sm90( + shared_l1_weights, shared_l2_weights + ) + ) + else: + transformed_shared_l1 = transformed_shared_l2 = None + # SwiGLU clamp: finite values enable clamp; inf maps to None and disables it. clamp_arg = args.activation_clamp if math.isfinite(args.activation_clamp) else None run_baseline_enabled = args.run_baseline or bool(args.check_output_diff) @@ -1251,6 +1465,8 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) # ---- Allocate fused SymmBuffer and output buffer ---- + # The fused shared expert needs two extra symmetric-buffer regions (its + # post-SwiGLU pool and SF) plus one more combine slot. sym_buffer = deep_gemm.get_symm_buffer_for_mega_moe( group, num_experts, @@ -1258,10 +1474,15 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): num_topk, hidden, intermediate_hidden, + num_shared_experts=num_shared_experts if fused_shared else 0, ) y_fused = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") - def run_fused(): + # Output of the reference dense shared MLP (`run_shared`), used by the baselines + # and by `--check-output-diff`. Reused across calls: those paths never overlap. + y_shared = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + + def run_fused(with_shared: bool = True): # Match the SM100 test: DG_COMM_KERNEL_DEBUG=1 zeros the whole # sym_buffer at kernel exit, so inputs must be re-copied every call. sym_buffer.x[:num_tokens].copy_(x_fp8[0]) @@ -1269,6 +1490,7 @@ def run_fused(): sym_buffer.topk_idx[:num_tokens].copy_(topk_idx) sym_buffer.topk_weights[:num_tokens].copy_(topk_weights) + fuse_now = fused_shared and with_shared deep_gemm.fp8_mega_moe( y_fused, transformed_l1, @@ -1279,9 +1501,51 @@ def run_fused(): activation="swiglu", activation_clamp=clamp_arg, fast_math=bool(args.fast_math), + shared_l1_weights=transformed_shared_l1 if fuse_now else None, + shared_l2_weights=transformed_shared_l2 if fuse_now else None, ) return y_fused + def run_shared(): + """Dense FP8 shared-expert MLP writing into ``y_shared``. + + L1 GEMM -> SwiGLU + FP8 quantization -> L2 GEMM. There is no topk + weighting: every token passes through the shared expert with weight 1.0. + The activation SF stays row-major FP32; ``fp8_gemm_nt`` transposes it into + the MN-major TMA layout internally (see + ``layout::transform_sf_into_required_layout``). + """ + if num_tokens == 0: + return y_shared + + l1_out = torch.empty( + (num_tokens, shared_intermediate_hidden * 2), + dtype=torch.bfloat16, + device="cuda", + ) + deep_gemm.fp8_gemm_nt( + x_fp8, + shared_l1_weights, + l1_out, + recipe=(1, 128, 128), + disable_ue8m0_cast=True, + ) + l2_in = swiglu_apply_weight_to_fp8_triton( + x=l1_out, + topk_weights=None, + clamp_value=clamp_arg, + num_per_channels=BASELINE_L2_ACT_SF_GRAN, + use_ue8m0_scale=False, + ) + deep_gemm.fp8_gemm_nt( + l2_in, + shared_l2_weights, + y_shared, + recipe=(1, 128, 128), + disable_ue8m0_cast=True, + ) + return y_shared + # ---- Print config ---- dist_print("Config (SM90 fused MegaMoE):", once_in_node=True) dist_print(f" > Tokens: {num_tokens}/{num_max_tokens_per_rank}", once_in_node=True) @@ -1293,6 +1557,16 @@ def run_fused(): once_in_node=True, ) dist_print(f" > Masked ratio: {args.masked_ratio}", once_in_node=True) + dist_print( + f" > Shared experts: {num_shared_experts}" + + ( + f" (intermediate: {shared_intermediate_hidden}, fused into the mega " + f"kernel via SharedLinear1/2)" + if fused_shared + else " (disabled)" + ), + once_in_node=True, + ) dist_print( f" > Activation SF: fused L2 per-{FUSED_L2_ACT_SF_GRAN} FP32 pow2, " f"baseline L2 per-{BASELINE_L2_ACT_SF_GRAN} FP32 pow2 " @@ -1415,7 +1689,13 @@ def run_baseline(): ) # DeepEP combine: gather each token's topk expert outputs back to source rank. - return ep_buffer.combine(l2_y, handle=handle)[0] + combined = ep_buffer.combine(l2_y, handle=handle)[0] + # Non-overlapped baseline: the shared expert runs serially on the same + # stream. tests/test_mega_moe.py folds it into combine as a bias; the + # SM90 DeepEP shim here has no bias argument, so add it afterwards. + if num_shared_experts > 0: + combined.add_(run_shared()) + return combined # ---------------------------------------------------------------- # Low-latency baseline body. Mirrors the sglang @@ -1517,6 +1797,9 @@ def run_baseline_low_latency(): return_recv_hook=False, out=ll_combined, ) + # 6) Same serial shared expert as the normal-mode baseline. + if num_shared_experts > 0: + combined_x.add_(run_shared()) return combined_x # ---- Run once to check fused and optional baseline paths ---- @@ -1524,6 +1807,27 @@ def run_baseline_low_latency(): assert y.shape == (num_tokens, hidden) and y.dtype == torch.bfloat16, ( f"unexpected fused output shape/dtype: shape={y.shape}, dtype={y.dtype}" ) + if fused_shared and args.check_output_diff: + # Reference for the fused shared expert: the routed-only kernel output plus + # the Python dense shared MLP (the same weights, before the interleave). + y_fused_shared = y.clone() + y_ref = run_fused(with_shared=False).clone() + y_ref += run_shared() + diff = (y_fused_shared.float() - y_ref.float()).abs() + denom = y_ref.float().abs().mean().clamp_min(1e-12) + dist_print( + "Output diff (fused shared expert vs routed + two-stream shared):", + once_in_node=True, + ) + dist_print( + f" > max_abs={diff.max().item():.6e}, " + f"mean_abs={diff.mean().item():.6e}, " + f"mean_abs/mean_ref={diff.mean().div(denom).item():.6e}", + once_in_node=True, + ) + dist_print(once_in_node=True) + # Leave `y_fused` holding the fused output for the baseline diffs below + y = run_fused() if ep_buffer is not None: out_b = run_baseline() assert out_b.shape == (num_tokens, hidden) and out_b.dtype == torch.bfloat16, ( @@ -1574,7 +1878,9 @@ def run_baseline_low_latency(): num_touched_experts = int(torch.unique(local_expert_ids).numel()) # ---- benchmark ---- - # Fused: bench_kineto selects the sm90_fp8_mega_moe_impl GPU region only. + # Fused: bench_kineto selects the sm90_fp8_mega_moe_impl GPU region only, so + # this stays the pure routed-kernel time even with a shared expert running + # concurrently on the main stream. t_fused = bench_kineto( run_fused, SM90_KERNEL_NAME, @@ -1657,27 +1963,64 @@ def safe_div(a, b): num_nvlink_bytes = num_recv_tokens * (hidden + hidden // 32 + 4 + hidden * 2) nvlink_gbs = safe_div(num_nvlink_bytes / 1e9, t_fused) + # ---- Shared-expert FLOPs / HBM ---- + # Same three matmuls as a routed expert (L1 gate, L1 up, L2), but every local + # token goes through it, and the weights are streamed once (not per expert). + # The shared MLP is node-local, so it adds no NVLink traffic, and the fused + # epilogue keeps the SwiGLU input in registers (no BF16 staging round-trip). + num_shared_flops = 2 * num_tokens * hidden * shared_intermediate_hidden * 3 + num_shared_hbm_bytes = ( + 0 + if num_shared_experts == 0 + else ( + shared_intermediate_hidden * 2 * hidden # shared L1 weights (FP8) + + hidden * shared_intermediate_hidden # shared L2 weights (FP8) + + (shared_intermediate_hidden * 2 // WEIGHT_SF_GRAN_MN) + * (hidden // WEIGHT_SF_GRAN_K) + * 4 # shared L1 weight SF + + (hidden // WEIGHT_SF_GRAN_MN) + * (shared_intermediate_hidden // WEIGHT_SF_GRAN_K) + * 4 # shared L2 weight SF + + num_tokens * hidden + + num_tokens * (hidden // L1_ACT_SF_GRAN) * 4 # L1 input read (FP8 + SF) + + num_tokens * shared_intermediate_hidden + + num_tokens + * (shared_intermediate_hidden // BASELINE_L2_ACT_SF_GRAN) + * 4 # SwiGLU output write (FP8 + SF) + + num_tokens * shared_intermediate_hidden + + num_tokens + * (shared_intermediate_hidden // BASELINE_L2_ACT_SF_GRAN) + * 4 # L2 input read (FP8 + SF) + + num_tokens * hidden * 2 # L2 output write (BF16) + ) + ) + # Routed + shared: one launch produces both, so they share `t_fused`. + num_total_flops = ( + 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) + num_shared_flops + ) + num_total_hbm_bytes = num_hbm_bytes + num_shared_hbm_bytes + tflops_total = safe_div(num_total_flops / 1e12, t_fused) + hbm_gbs_total = safe_div(num_total_hbm_bytes / 1e9, t_fused) + # Serial lower bound for combine reduction, using 6.5e12 B/s as an estimate. t_reduction = num_tokens * hidden * 2 * (1 + num_topk) / 6.5e12 # Overlap adjustment: remove the non-overlapped serial reduction estimate. approx_factor = t_fused / max(t_fused - t_reduction, 1e-12) - # Baseline uses the same FLOPs and HBM byte estimate, with t_baseline. - tflops_baseline = safe_div( - 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_baseline - ) - hbm_gbs_baseline = safe_div(num_hbm_bytes / 1e9, t_baseline) + # Baselines run routed + shared serially, so they use the combined FLOPs and + # HBM byte estimate (identical to the routed-only one when shared is off). + tflops_baseline = safe_div(num_total_flops / 1e12, t_baseline) + hbm_gbs_baseline = safe_div(num_total_hbm_bytes / 1e9, t_baseline) nvlink_gbs_baseline = safe_div(num_nvlink_bytes / 1e9, t_baseline) # Low-latency baseline pads each expert's activation to ``M_max_ll``, so # the weights are streamed once per expert regardless of routing. NVLink # bytes match the normal-mode baseline (same per-routed-token volume). - tflops_baseline_ll = safe_div( - 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_baseline_ll - ) - hbm_gbs_baseline_ll = safe_div(num_hbm_bytes / 1e9, t_baseline_ll) + tflops_baseline_ll = safe_div(num_total_flops / 1e12, t_baseline_ll) + hbm_gbs_baseline_ll = safe_div(num_total_hbm_bytes / 1e9, t_baseline_ll) nvlink_gbs_baseline_ll = safe_div(num_nvlink_bytes / 1e9, t_baseline_ll) + def fmt_perf_line( name: str, t: float, @@ -1713,10 +2056,10 @@ def fmt_perf_line( ) dist_print( fmt_perf_line( - "[fused]", + "[fused+sh]" if fused_shared else "[fused]", t_fused, - tflops * approx_factor, - hbm_gbs * approx_factor, + (tflops_total if fused_shared else tflops) * approx_factor, + (hbm_gbs_total if fused_shared else hbm_gbs) * approx_factor, nvlink_gbs * approx_factor, reduction_us=t_reduction * 1e6, ) @@ -1835,6 +2178,17 @@ def fmt_perf_line( default=10.0, help="Clamp threshold for gate/up before SwiGLU; pass inf to disable", ) + parser.add_argument( + "--num-shared-experts", + type=int, + default=0, + help=( + "DeepSeek-style shared experts, each adding one routed intermediate " + "size, fused into the mega kernel as the SharedLinear1/2 phases (both " + "baselines run them serially as a dense FP8 MLP). 0 disables it; only " + "the default comparison mode uses this" + ), + ) parser.add_argument("--num-experts", type=int, default=384) parser.add_argument("--num-topk", type=int, default=6) parser.add_argument( @@ -1904,8 +2258,9 @@ def fmt_perf_line( "--layers", type=int, nargs="+", - default=[1, 2, 3, 4], - help="Accuracy layers to run with --accuracy (1..5); default: 1 2 3 4", + default=[1, 2, 3, 4, 6], + help="Accuracy layers to run with --accuracy (1..6); default: 1 2 3 4 6. " + "Layer 6 covers the fused shared expert (fused-out = routed + shared).", ) parser.add_argument( "--num-correctness-tests", diff --git a/sgl_deep_gemm/tests/test_mega_moe_l1_fp4_accuracy.py b/sgl_deep_gemm/tests/test_mega_moe_l1_fp4_accuracy.py index 1cd77fba6f..bd5019e1a8 100644 --- a/sgl_deep_gemm/tests/test_mega_moe_l1_fp4_accuracy.py +++ b/sgl_deep_gemm/tests/test_mega_moe_l1_fp4_accuracy.py @@ -29,7 +29,8 @@ from typing import Tuple import deep_gemm -from deep_gemm.utils import per_token_cast_to_fp8, per_token_cast_to_fp4 +from deep_gemm.utils import (per_token_cast_to_fp8, per_token_cast_to_fp4, + per_token_cast_to_nvfp4, transform_ue4m3_sf_into_required_layout) from deep_gemm.utils.dist import dist_print, init_dist @@ -71,6 +72,11 @@ def _decode_fp8_e4m3(fp8_bytes: torch.Tensor) -> torch.Tensor: return fp8_bytes.view(torch.float8_e4m3fn).to(torch.float) +def _decode_ue4m3(sf_bytes: torch.Tensor) -> torch.Tensor: + """Decode UE4M3 SF bytes (NVFP4) — same bit layout as unsigned `float8_e4m3fn`.""" + return sf_bytes.to(torch.uint8).view(torch.float8_e4m3fn).float() + + def _decode_ue8m0(sf_bytes: torch.Tensor) -> torch.Tensor: """Decode UE8M0 byte values to float32 multipliers (= 2^(byte - 127)).""" return ((sf_bytes.to(torch.int32) << 23).view(torch.float32)) @@ -156,7 +162,8 @@ def _dequant_l1_acts_fp4(l2_acts_bytes: torch.Tensor, intermediate_hidden: int, num_padded_sf_pool_tokens: int, valid_slots: int, - gran_k: int = 32) -> torch.Tensor: + gran_k: int = 16, + use_ue4m3: bool = True) -> torch.Tensor: """Decode the FP4 L1 output bytes from the same symm buffer slot. Per A0.1's TMA descriptor: only the first `intermediate_hidden / 2` bytes @@ -169,7 +176,7 @@ def _dequant_l1_acts_fp4(l2_acts_bytes: torch.Tensor, decoded = _decode_fp4_packed(raw_bytes) # (V, I) sf = _decode_sf_buffer_to_per_token( l2_acts_sf_bytes, num_padded_sf_pool_tokens, - intermediate_hidden, valid_slots, gran_k) + intermediate_hidden, valid_slots, gran_k, use_ue4m3=use_ue4m3) n_blocks = intermediate_hidden // gran_k decoded = decoded.view(valid_slots, n_blocks, gran_k) sf = sf.view(valid_slots, n_blocks, 1) @@ -180,7 +187,8 @@ def _decode_sf_buffer_to_per_token(sf_bytes_int32: torch.Tensor, num_padded_sf_pool_tokens: int, intermediate_hidden: int, valid_slots: int, - gran_k: int) -> torch.Tensor: + gran_k: int, + use_ue4m3: bool = False) -> torch.Tensor: """Read out per-token-K-block UE8M0 SF bytes from the M-major SF buffer. The SF buffer in the kernel uses an M-major / per-32-elements layout with a @@ -220,7 +228,7 @@ def _decode_sf_buffer_to_per_token(sf_bytes_int32: torch.Tensor, # word at that token's k_uint slot. word = sf_bytes_int32[sf_pool_token_idx, k_uint_idx] # int32 (V,) out[:, kb] = ((word >> (byte_idx * 8)) & 0xFF).to(torch.uint8) - return _decode_ue8m0(out) + return _decode_ue4m3(out) if use_ue4m3 else _decode_ue8m0(out) def _gather_l2_buffers(buffer): @@ -255,24 +263,39 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): (num_experts_per_rank,), dtype=torch.int, device='cuda') # FP8 / FP4 quantizations needed by the kernel + fp4_mma_type = args.fp4_mma_type + fp4_is_nvfp4 = fp4_mma_type == 'nvfp4xnvfp4' + fp4_gran_k = 16 if fp4_is_nvfp4 else 32 x_fp8 = per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - x_fp4 = per_token_cast_to_fp4(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + x_fp4 = per_token_cast_to_nvfp4(x_bf16, gran_k=16, use_packed_ue4m3=True) if fp4_is_nvfp4 \ + else per_token_cast_to_fp4(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - def cast_grouped_weights_to_fp4(bf16_weights): + def cast_grouped_weights_to_fp4(bf16_weights, use_nvfp4: bool): + gran_k = 16 if use_nvfp4 else 32 num_groups, n, k = bf16_weights.shape w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) - w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + w_sf = torch.empty((num_groups, n, k // gran_k), device='cuda', dtype=torch.float) for i in range(num_groups): - w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) - w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + if use_nvfp4: + w[i], w_sf[i] = per_token_cast_to_nvfp4(bf16_weights[i], gran_k=gran_k) + else: + w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=gran_k) + w_sf = transform_ue4m3_sf_into_required_layout(w_sf, n) if use_nvfp4 else \ + deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, gran_k), num_groups) return w, w_sf - l1_weights_fp4 = cast_grouped_weights_to_fp4(l1_weights_bf16) - l2_weights_fp4 = cast_grouped_weights_to_fp4(l2_weights_bf16) - transformed_l1_weights, transformed_l2_weights = \ - deep_gemm.transform_weights_for_mega_moe(l1_weights_fp4, l2_weights_fp4) + def transform_weights(mma_type: str): + use_nvfp4 = mma_type == 'nvfp4xnvfp4' + return deep_gemm.transform_weights_for_mega_moe( + cast_grouped_weights_to_fp4(l1_weights_bf16, use_nvfp4), + cast_grouped_weights_to_fp4(l2_weights_bf16, use_nvfp4), + 'swiglu', mma_type) + + weights_per_arm = {False: transform_weights('fp8xfp4'), + True: transform_weights(fp4_mma_type)} - def run_once(buffer, x_src): + def run_once(buffer, x_src, use_fp4_acts: bool = False): + transformed_l1_weights, transformed_l2_weights = weights_per_arm[use_fp4_acts] buffer.x[:num_tokens].copy_(x_src[0]) buffer.x_sf[:num_tokens].copy_(x_src[1]) buffer.topk_idx[:num_tokens].copy_(topk_idx) @@ -284,19 +307,19 @@ def run_once(buffer, x_src): transformed_l1_weights, transformed_l2_weights, buffer, cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats, + recipe=(1, 1, fp4_gran_k) if use_fp4_acts else (1, 1, 32), activation_clamp=activation_clamp, fast_math=bool(args.fast_math) ) return y, cumulative_local_expert_recv_stats.clone() - # Buffer layout depends on DG_USE_FP4_ACTS; set the env before allocating. def make_buffer(use_fp4_acts): - os.environ['DG_USE_FP4_ACTS'] = '1' if use_fp4_acts else '0' os.environ['DG_COMM_KERNEL_DEBUG'] = '0' # don't zero buffer between calls return deep_gemm.get_symm_buffer_for_mega_moe( group, num_experts, num_max_tokens_per_rank, num_topk, - hidden, intermediate_hidden + hidden, intermediate_hidden, + mma_type=fp4_mma_type if use_fp4_acts else 'fp8xfp4' ) # ---- BF16 reference for L1 SwiGLU output (per token×topk) ---- @@ -323,9 +346,9 @@ def make_buffer(use_fp4_acts): # ---- Run FP4 path (separate buffer, laid out for packed E2M1) ---- buffer = make_buffer(use_fp4_acts=True) - _ = run_once(buffer, x_fp4) + _ = run_once(buffer, x_fp4, use_fp4_acts=True) torch.cuda.synchronize() - y_fp4, recv_stats_fp4 = run_once(buffer, x_fp4) + y_fp4, recv_stats_fp4 = run_once(buffer, x_fp4, use_fp4_acts=True) torch.cuda.synchronize() l2_acts_fp4 = buffer.l2_acts.clone() l2_acts_sf_fp4 = buffer.l2_acts_sf.clone() @@ -417,7 +440,7 @@ def make_buffer(use_fp4_acts): fp4_dec = _dequant_l1_acts_fp4( l2_acts_fp4, l2_acts_sf_fp4, intermediate_hidden, num_padded_sf_pool_tokens, - total_local) + total_local, gran_k=fp4_gran_k, use_ue4m3=fp4_is_nvfp4) # Sanity: dump a few raw bytes from each path so we can compare visually # if the harness misaligns. @@ -488,6 +511,9 @@ def make_buffer(use_fp4_acts): parser.add_argument('--num-topk', type=int, default=2) parser.add_argument('--activation-clamp', type=float, default=10.0) parser.add_argument('--fast-math', type=int, default=1) + parser.add_argument('--fp4-mma-type', type=str, default='nvfp4xnvfp4', + choices=('nvfp4xnvfp4', 'mxf4xmxf4'), + help='MMA type for the FP4-acts arm') args = parser.parse_args() num_processes = args.num_processes diff --git a/sgl_deep_gemm/tests/test_mega_moe_l1_sentinel.py b/sgl_deep_gemm/tests/test_mega_moe_l1_sentinel.py index ca7a9cc7a2..bec18851c9 100644 --- a/sgl_deep_gemm/tests/test_mega_moe_l1_sentinel.py +++ b/sgl_deep_gemm/tests/test_mega_moe_l1_sentinel.py @@ -38,7 +38,8 @@ import torch.distributed as dist import deep_gemm -from deep_gemm.utils import per_token_cast_to_fp8, per_token_cast_to_fp4 +from deep_gemm.utils import (per_token_cast_to_fp8, per_token_cast_to_fp4, + per_token_cast_to_nvfp4, transform_ue4m3_sf_into_required_layout) from deep_gemm.utils.dist import dist_print, init_dist @@ -88,36 +89,49 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) cumulative = torch.zeros((num_experts_per_rank,), dtype=torch.int, device='cuda') - x_fp8 = per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - x_fp4 = per_token_cast_to_fp4(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - - def cast_grouped_weights_to_fp4(bf16_weights): + x_src_by_type = { + 'fp8xfp4': per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True), + 'nvfp4xnvfp4': per_token_cast_to_nvfp4(x_bf16, gran_k=16, use_packed_ue4m3=True), + 'mxf4xmxf4': per_token_cast_to_fp4(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True), + } + + def cast_grouped_weights_to_fp4(bf16_weights, mma_type: str): + use_nvfp4 = mma_type == 'nvfp4xnvfp4' + gran_k = 16 if use_nvfp4 else 32 num_groups, n, k = bf16_weights.shape w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) - w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + w_sf = torch.empty((num_groups, n, k // gran_k), device='cuda', dtype=torch.float) for i in range(num_groups): - w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) - w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + if use_nvfp4: + w[i], w_sf[i] = per_token_cast_to_nvfp4(bf16_weights[i], gran_k=gran_k) + else: + w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=gran_k) + w_sf = transform_ue4m3_sf_into_required_layout(w_sf, n) if use_nvfp4 else \ + deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, gran_k), num_groups) return w, w_sf - l1_weights_fp4 = cast_grouped_weights_to_fp4(l1_weights_bf16) - l2_weights_fp4 = cast_grouped_weights_to_fp4(l2_weights_bf16) - transformed_l1_weights, transformed_l2_weights = \ - deep_gemm.transform_weights_for_mega_moe(l1_weights_fp4, l2_weights_fp4) - - # Stream A0.0b: under `DG_USE_FP4_ACTS=1`, the symm buffer's `x` slot is - # sized for packed E2M1 (`hidden/2` bytes/token) — different from FP8. - # Allocate the buffer separately for each path and feed it the matching - # source tensor. - def make_buffer_and_run(use_fp4_acts: bool): - os.environ['DG_USE_FP4_ACTS'] = '1' if use_fp4_acts else '0' + # Each arm quantizes weights to its own SF granularity: per-32 UE8M0 for + # `fp8xfp4`/`mxf4xmxf4`, per-16 UE4M3 for `nvfp4xnvfp4`. + def transform_weights(mma_type: str): + return deep_gemm.transform_weights_for_mega_moe( + cast_grouped_weights_to_fp4(l1_weights_bf16, mma_type), + cast_grouped_weights_to_fp4(l2_weights_bf16, mma_type), + 'swiglu', mma_type) + + # FP4-acts kinds make the symm buffer's `x` slot packed E2M1 (`hidden/2` + # bytes/token), while `fp8xfp4` retains FP8 activations. + # Allocate the buffer separately for each path and feed it the matching source tensor. + def make_buffer_and_run(mma_type: str): os.environ['DG_COMM_KERNEL_DEBUG'] = '0' buf = deep_gemm.get_symm_buffer_for_mega_moe( group, num_experts, num_max_tokens_per_rank, num_topk, - hidden, intermediate_hidden + hidden, intermediate_hidden, + mma_type=mma_type ) - x_src = x_fp4 if use_fp4_acts else x_fp8 + x_src = x_src_by_type[mma_type] + transformed_l1_weights, transformed_l2_weights = transform_weights(mma_type) + recipe = (1, 1, 16) if mma_type == 'nvfp4xnvfp4' else (1, 1, 32) def run_once(): buf.x[:num_tokens].copy_(x_src[0]) @@ -129,6 +143,7 @@ def run_once(): deep_gemm.fp8_fp4_mega_moe( y, transformed_l1_weights, transformed_l2_weights, buf, cumulative_local_expert_recv_stats=cumulative, + recipe=recipe, activation_clamp=activation_clamp, fast_math=bool(args.fast_math) ) @@ -141,47 +156,48 @@ def run_once(): buf.destroy() return y_out - # Run FP8-acts first (warmup + measurement). - y_fp8 = make_buffer_and_run(use_fp4_acts=False) - # Run FP4-acts (separate buffer because the `x` slot footprint changes). - y_fp4 = make_buffer_and_run(use_fp4_acts=True) - - # End-to-end y comparison: this is the source of truth (no slot - # permutation ambiguity since y is indexed by global (token, hidden)). - y_diff = (y_fp4.float() - y_fp8.float()).abs() - y_rmse = y_diff.pow(2).mean().sqrt().item() + # Run FP8-acts first (warmup + measurement), then each FP4-acts kind + # (separate buffers because the `x` slot footprint changes). + y_fp8 = make_buffer_and_run('fp8xfp4') y_fp8_rms = y_fp8.float().pow(2).mean().sqrt().item() - rel_rmse = y_rmse / max(y_fp8_rms, 1e-12) - - dist_print(f'=== A0.2.1 sentinel — y rel-RMSE (FP4 vs FP8 acts) ===', - once_in_node=True) - dist_print(f' y_fp8 RMS: {y_fp8_rms:.4f}', once_in_node=True) - dist_print(f' y_rmse: {y_rmse:.4f}', once_in_node=True) - dist_print(f' rel-RMSE: {rel_rmse:.4f}', once_in_node=True) - max_rel_rmse = args.max_rel_rmse y_fp8_mag = y_fp8.float().abs().mean().item() - y_fp4_mag = y_fp4.float().abs().mean().item() - - dist_print(f' y_fp8 mean|.|: {y_fp8_mag:.4f}', once_in_node=True) - dist_print(f' y_fp4 mean|.|: {y_fp4_mag:.4f}', once_in_node=True) - dist_print(f' target: <= {max_rel_rmse:.2f} (layout sentinel)', - once_in_node=True) - dist_print(f' verdict: {"PASS" if rel_rmse <= max_rel_rmse else "FAIL"}', - once_in_node=True) - - # Spot-check first row to make the failure mode legible if it ever - # comes back: matched values at low N indices = layout correct; - # garbage = layout broken. - dist_print(f'\n y_fp8 [0, :8]: {y_fp8[0, :8].cpu().tolist()}', - once_in_node=True) - dist_print(f' y_fp4 [0, :8]: {y_fp4[0, :8].cpu().tolist()}', - once_in_node=True) - - assert torch.isfinite(y_fp4).all(), 'FP4 output contains NaN/Inf' - assert y_fp8_mag * 0.5 < y_fp4_mag < y_fp8_mag * 2.0, \ - f'FP4 magnitude miscalibrated: |y_fp4|={y_fp4_mag} vs |y_fp8|={y_fp8_mag}' - assert rel_rmse <= max_rel_rmse, \ - f'A0.2.1 layout regression: y rel-RMSE {rel_rmse:.4f} > {max_rel_rmse:.2f}' + max_rel_rmse = args.max_rel_rmse + + for mma_type in ('nvfp4xnvfp4', 'mxf4xmxf4'): + y_fp4 = make_buffer_and_run(mma_type) + + # End-to-end y comparison: this is the source of truth (no slot + # permutation ambiguity since y is indexed by global (token, hidden)). + y_diff = (y_fp4.float() - y_fp8.float()).abs() + y_rmse = y_diff.pow(2).mean().sqrt().item() + rel_rmse = y_rmse / max(y_fp8_rms, 1e-12) + y_fp4_mag = y_fp4.float().abs().mean().item() + + dist_print(f'=== A0.2.1 sentinel — y rel-RMSE ({mma_type} vs FP8 acts) ===', + once_in_node=True) + dist_print(f' y_fp8 RMS: {y_fp8_rms:.4f}', once_in_node=True) + dist_print(f' y_rmse: {y_rmse:.4f}', once_in_node=True) + dist_print(f' rel-RMSE: {rel_rmse:.4f}', once_in_node=True) + dist_print(f' y_fp8 mean|.|: {y_fp8_mag:.4f}', once_in_node=True) + dist_print(f' y_fp4 mean|.|: {y_fp4_mag:.4f}', once_in_node=True) + dist_print(f' target: <= {max_rel_rmse:.2f} (layout sentinel)', + once_in_node=True) + dist_print(f' verdict: {"PASS" if rel_rmse <= max_rel_rmse else "FAIL"}', + once_in_node=True) + + # Spot-check first row to make the failure mode legible if it ever + # comes back: matched values at low N indices = layout correct; + # garbage = layout broken. + dist_print(f'\n y_fp8 [0, :8]: {y_fp8[0, :8].cpu().tolist()}', + once_in_node=True) + dist_print(f' y_fp4 [0, :8]: {y_fp4[0, :8].cpu().tolist()}', + once_in_node=True) + + assert torch.isfinite(y_fp4).all(), f'{mma_type} output contains NaN/Inf' + assert y_fp8_mag * 0.5 < y_fp4_mag < y_fp8_mag * 2.0, \ + f'{mma_type} magnitude miscalibrated: |y_fp4|={y_fp4_mag} vs |y_fp8|={y_fp8_mag}' + assert rel_rmse <= max_rel_rmse, \ + f'A0.2.1 layout regression ({mma_type}): y rel-RMSE {rel_rmse:.4f} > {max_rel_rmse:.2f}' dist.barrier() dist.destroy_process_group() diff --git a/sgl_deep_gemm/tests/test_mega_moe_nvfp4_alphas.py b/sgl_deep_gemm/tests/test_mega_moe_nvfp4_alphas.py new file mode 100644 index 0000000000..8fd2409cc2 --- /dev/null +++ b/sgl_deep_gemm/tests/test_mega_moe_nvfp4_alphas.py @@ -0,0 +1,247 @@ +# NVFP4 outer-scale (global scale) plumbing for the mega-MoE kernel. +# +# Four tables carry the NVFP4 dequant alphas: +# +# x_scales per token, written by `mega_moe_pre_dispatch` (1 / gs_x1) +# l1_alphas [num_LOCAL_experts, 2] = 1 / (gs_w1, gs_w3), on the L1 accum +# l2_alphas [num_LOCAL_experts] = 1 / gs_w2, on the L2 accum +# expert_scales [num_GLOBAL_experts] folded into the topk weight, O(1) only +# +# The fc2 alpha is deliberately NOT routed through `expert_scales`: that lands on +# the L1 output, which is then re-quantized to NVFP4 with a bare per-block E4M3 +# SF, and a realistic `1 / gs_w2` (~1e-4) sinks that SF under E4M3's subnormal +# floor — `y` comes back all zeros (measured). `expert_scales` is exercised +# separately below with an O(1) ratio, which is all it can carry. +# +# Method: quantize the same BF16 weights twice — once with gs = 1 and no alphas, +# once with per-expert gs and the alphas that undo it — and compare `y`. +# - Power-of-two gs: every scaling is exact in FP32/E4M3, so the two runs must +# be BITWISE equal. +# - Realistic gs (= 448*6/amax, not powers of two): the two encodings round +# differently, so this arm is only a sanity band — measured 0.18 rel-L2 on +# `y`, which is what two independent NVFP4 encodings cost (7.7% apart on the +# weights alone, over two GEMMs). Its job is to show realistic gs magnitudes +# neither saturate the E4M3 SF nor zero it, not to catch subtle bugs; the +# pow2 arm and the roll controls are the sharp instruments here. +# - Control: rolling either alpha table by one expert must change `y`, +# otherwise the test would pass on a kernel that ignores the table. +# +# Run with >= 2 ranks: local- vs global-expert index confusion only shows up on +# a rank whose `expert_lo` is nonzero. +# +# Usage: +# PYTHONPATH=/workspace/chunan/DeepGEMM CUDA_VISIBLE_DEVICES=4,5 MASTER_PORT=29511 \ +# python3 sgl_deep_gemm/tests/test_mega_moe_nvfp4_alphas.py --num-processes 2 + +import argparse +import sys +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.utils import cast_to_ue4m3, transform_ue4m3_sf_into_required_layout +from deep_gemm.utils.math import _quantize_to_fp4_e2m1 +from deep_gemm.utils.dist import dist_print, init_dist + +GRAN_K = 16 +E4M3_MAX = 448.0 +FP4_MAX = 6.0 + + +def cast_to_nvfp4_with_gs(w: torch.Tensor, gs) -> tuple: + """NVFP4 with an outer scale. Dequant contract: `w ~= fp4 * sf / gs`. + + `gs` is a scalar or a per-row (m,) tensor — the L1 weight carries gs_w1 on + its gate rows and gs_w3 on its up rows. `gs = 1` reproduces + `per_token_cast_to_nvfp4` exactly. + """ + m, n = w.shape + v = w.view(m, -1, GRAN_K).float() + gs = torch.as_tensor(gs, dtype=torch.float, device=w.device).reshape(-1, 1).expand(m, 1) + sf = cast_to_ue4m3(v.abs().amax(dim=2) * (gs / FP4_MAX)) + scale = torch.where(sf > 0, gs / sf, torch.zeros_like(sf)) + codes = _quantize_to_fp4_e2m1(v * scale.unsqueeze(2)).view(m, n) + codes2 = codes.view(m, n // 2, 2) + packed = (codes2[:, :, 0] & 0x0F) | ((codes2[:, :, 1] & 0x0F) << 4) + return packed.contiguous(), sf + + +def cast_grouped(w: torch.Tensor, gs_rows) -> tuple: + """`w` is (E, N, K) BF16; `gs_rows` is one per-row gs tensor per expert.""" + num_experts, n, k = w.shape + packed = torch.empty((num_experts, n, k // 2), dtype=torch.int8, device=w.device) + sf = torch.empty((num_experts, n, k // GRAN_K), dtype=torch.float, device=w.device) + for e in range(num_experts): + packed[e], sf[e] = cast_to_nvfp4_with_gs(w[e], gs_rows[e]) + return packed, transform_ue4m3_sf_into_required_layout(sf, n) + + +def make_gs(amax_per_expert: torch.Tensor, pow2: bool) -> torch.Tensor: + """Realistic NVFP4 global scale, `448 * 6 / amax`, optionally snapped to a + power of two so the two runs stay bitwise comparable. + + Same-shaped random weights have near-identical amax, so the pow2 arm would + otherwise hand every expert the same scale and the roll-by-one control would + be vacuous. Spread it downwards per expert — never up, `gs > 448*6/amax` + saturates the E4M3 SF. + """ + gs = (E4M3_MAX * FP4_MAX) / amax_per_expert.clamp_min(1e-4) + if not pow2: + return gs + spread = torch.pow(2.0, -(torch.arange(gs.numel(), device=gs.device) % 4).float()) + return torch.pow(2.0, torch.floor(torch.log2(gs))) * spread + + +# noinspection PyUnboundLocalVariable +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(1000 + rank_idx) + + num_tokens = args.num_tokens + hidden, inter = args.hidden, args.intermediate_hidden + num_experts, num_topk = args.num_experts, args.num_topk + num_local_experts = num_experts // num_ranks + expert_lo = rank_idx * num_local_experts + + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, args.num_max_tokens_per_rank, num_topk, + hidden, inter, mma_type='nvfp4xnvfp4') + + # `x` is damped so the L1 output amax stays under the L2 activation SF's + # 448*6 ceiling; weights stay unit-variance so their own block SFs stay in + # E4M3's normal range in BOTH runs (a gs=1 encoding of small weights lands on + # E4M3 subnormals, where power-of-two rescaling is no longer exact). + x = (torch.randn((num_tokens, hidden), device='cuda') * args.x_scale).bfloat16() + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + topk_idx = topk_idx.int() + + # Weights are local, but a GLOBAL table (`expert_scales`) has to know every + # rank's values. Draw all experts from a shared seed and keep the local + # slice, so the global tables agree across ranks. + torch.manual_seed(7) + l1_all = torch.randn((num_experts, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') + l2_all = torch.randn((num_experts, hidden, inter), dtype=torch.bfloat16, device='cuda') + gate_amax = l1_all[:, :inter].abs().float().amax(dim=(1, 2)) + up_amax = l1_all[:, inter:].abs().float().amax(dim=(1, 2)) + l2_amax = l2_all.abs().float().amax(dim=(1, 2)) + l1_local = l1_all[expert_lo:expert_lo + num_local_experts].contiguous() + l2_local = l2_all[expert_lo:expert_lo + num_local_experts].contiguous() + del l1_all, l2_all + + base_w = deep_gemm.transform_weights_for_mega_moe( + cast_grouped(l1_local, [torch.ones(inter * 2, device='cuda')] * num_local_experts), + cast_grouped(l2_local, [torch.ones(hidden, device='cuda')] * num_local_experts), + 'swiglu', 'nvfp4xnvfp4') + + def run(weights, l1_alphas=None, l2_alphas=None, expert_scales=None, tw=None, + l2_act_scales=None): + deep_gemm.mega_moe_pre_dispatch( + x, topk_idx, topk_weights if tw is None else tw, + buffer.x, buffer.x_sf, buffer.topk_idx, buffer.topk_weights, + num_tokens=num_tokens, group_size=GRAN_K, mma_type='nvfp4xnvfp4', + buf_x_scales=buffer.x_scales, expert_scales=expert_scales) + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_fp4_mega_moe( + y=y, l1_weights=weights[0], l2_weights=weights[1], sym_buffer=buffer, + recipe=(1, 1, GRAN_K), activation='swiglu', fast_math=bool(args.fast_math), + use_x_scales=True, l1_alphas=l1_alphas, l2_alphas=l2_alphas, + l2_act_scales=l2_act_scales) + dist.barrier() + torch.cuda.synchronize() + return y + + y_base = run(base_w) + dist_print(f'Config: {num_ranks} ranks x {num_local_experts} local experts, ' + f'{num_tokens} tokens, hidden={hidden}, inter={inter}, ' + f'|y_base| mean {y_base.float().abs().mean().item():.4g}', once_in_node=True) + + for pow2 in (True, False): + gs_w1, gs_w3, gs_w2 = (make_gs(a, pow2) for a in (gate_amax, up_amax, l2_amax)) + scaled = deep_gemm.transform_weights_for_mega_moe( + cast_grouped(l1_local, [torch.cat([gs_w1[expert_lo + e].expand(inter), + gs_w3[expert_lo + e].expand(inter)]) + for e in range(num_local_experts)]), + cast_grouped(l2_local, [gs_w2[expert_lo + e].expand(hidden) + for e in range(num_local_experts)]), + 'swiglu', 'nvfp4xnvfp4') + + # Both LOCAL index space; `l1_alphas` columns are (gate, up). + local = slice(expert_lo, expert_lo + num_local_experts) + l1_alphas = torch.stack([1.0 / gs_w1[local], 1.0 / gs_w3[local]], dim=1).contiguous() + l2_alphas = (1.0 / gs_w2[local]).contiguous() + + y_scaled = run(scaled, l1_alphas, l2_alphas) + tag = 'pow2' if pow2 else 'realistic' + if pow2: + assert torch.equal(y_base, y_scaled), \ + f'[{tag}] power-of-two global scales must cancel bitwise, max |delta| = ' \ + f'{(y_base.float() - y_scaled.float()).abs().max().item():.3e}' + dist_print(f' > [{tag}] gs_w2 = {gs_w2.tolist()} — bitwise equal', once_in_node=True) + else: + rel = ((y_scaled.float() - y_base.float()).norm() / y_base.float().norm()).item() + assert rel < args.rel_tol, f'[{tag}] rel-L2 {rel:.4f} >= {args.rel_tol}' + dist_print(f' > [{tag}] rel-L2 = {rel:.5f} (< {args.rel_tol})', once_in_node=True) + + for name, wrong in (('l1_alphas', (l1_alphas.roll(1, dims=0).contiguous(), l2_alphas)), + ('l2_alphas', (l1_alphas, l2_alphas.roll(1).contiguous()))): + assert not torch.equal(y_scaled, run(scaled, *wrong)), \ + f'[{tag}] rolling `{name}` changed nothing — the kernel ignores it' + dist_print(f' > [{tag}] roll-by-one control: both tables have teeth', once_in_node=True) + + # `l2_act_scales` (LOCAL index): pow2 fc2 input gs is cancelled by `1/gs` + # in `l2_alphas`. Damp x so the scaled block SFs stay under E4M3's 448 + # satfinite ceiling. Exact pow2 invariance still breaks for blocks whose SF + # lands in E4M3's SUBNORMAL range (absolute 2**-9 grid, not scale + # invariant): a re-rounded SF moves single FP4 codes, perturbing outputs by + # at most ~1 BF16 ulp of the tensor scale. Bound accordingly, not bitwise. + x = (x.float() * 0.25).bfloat16() + y_damped = run(base_w) + gs_in = torch.pow(2.0, (torch.arange(num_local_experts, device='cuda', + dtype=torch.float) % 3)) + y_gs = run(base_w, l2_alphas=(1.0 / gs_in).contiguous(), + l2_act_scales=gs_in.contiguous()) + d = (y_gs.float() - y_damped.float()).abs() + tol = torch.finfo(torch.bfloat16).eps * y_damped.float().abs().max() + rel = (d.norm() / y_damped.float().norm()).item() + assert d.max() <= tol and rel < 1e-3, \ + f'l2_act_scales cancellation off: max |d| {d.max().item():.3e} ' \ + f'(tol {tol.item():.3e}), rel-L2 {rel:.3e}' + assert not torch.equal(y_damped, run(base_w, l2_act_scales=gs_in.contiguous())), \ + 'l2_act_scales alone changed nothing -- the kernel ignores it' + dist_print(f' > l2_act_scales: pow2 gs cancels to 1 tensor-scale ulp ' + f'(max |d| {d.max().item():.2e}, rel {rel:.1e}), teeth ok', + once_in_node=True) + + # `expert_scales`, GLOBAL index space: folding an O(1) per-expert ratio in the + # kernel must equal folding it into `topk_weights` on the host. Bitwise — + # both are the same FP32 multiply, so this pins the index space exactly. + ratios = (1.0 + 0.25 * torch.arange(num_experts, device='cuda', dtype=torch.float) + / num_experts).contiguous() + host_tw = topk_weights * torch.where(topk_idx >= 0, ratios[topk_idx.long().clamp_min(0)], 1.0) + assert torch.equal(run(base_w, expert_scales=ratios), run(base_w, tw=host_tw)), \ + '`expert_scales` does not match a host-side fold — check the global expert index' + dist_print(' > expert_scales: matches a host-side topk_weights fold', once_in_node=True) + + dist_print('OK', once_in_node=True) + dist.barrier() + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--num-processes', type=int, default=2) + parser.add_argument('--num-max-tokens-per-rank', type=int, default=1024) + parser.add_argument('--num-tokens', type=int, default=1024) + parser.add_argument('--hidden', type=int, default=1024) + parser.add_argument('--intermediate-hidden', type=int, default=512) + parser.add_argument('--num-experts', type=int, default=8) + parser.add_argument('--num-topk', type=int, default=2) + parser.add_argument('--x-scale', type=float, default=0.1) + parser.add_argument('--fast-math', type=int, default=1) + parser.add_argument('--rel-tol', type=float, default=0.30) + args = parser.parse_args() + assert args.num_experts % args.num_processes == 0 + torch.multiprocessing.spawn(test, args=(args.num_processes, args), nprocs=args.num_processes) + sys.exit(0) diff --git a/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch.py b/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch.py index 679e1e4271..6ab6743cd7 100644 --- a/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch.py +++ b/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch.py @@ -57,7 +57,8 @@ def _run_one(use_fp4_acts: bool, args: argparse.Namespace) -> None: deep_gemm.mega_moe_pre_dispatch( x, topk_idx, topk_weights, buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, - num_tokens=M, group_size=G, use_fp4_acts=use_fp4_acts, + num_tokens=M, group_size=G, + mma_type='mxf4xmxf4' if use_fp4_acts else 'fp8xfp4', ) torch.cuda.synchronize() diff --git a/sgl_deep_gemm/tests/test_mega_moe_situ.py b/sgl_deep_gemm/tests/test_mega_moe_situ.py index a89c82b17f..6f1528b977 100644 --- a/sgl_deep_gemm/tests/test_mega_moe_situ.py +++ b/sgl_deep_gemm/tests/test_mega_moe_situ.py @@ -135,13 +135,13 @@ def run(activation: str, activation_clamp=None): assert len(kernel_sources) == 2 assert any( re.search( - r"cute::numeric_limits::infinity\(\),\s+true,\s+true,", + r"cute::numeric_limits::infinity\(\),\s+0x0p[+-]\d+f,\s+true,\s+true", source, ) for source in kernel_sources ), "explicit SiTU did not instantiate kUseSitu=true" assert any( - re.search(r"0x1p-5f,\s+false,\s+true,", source) for source in kernel_sources + re.search(r"0x1p-5f,\s+0x0p[+-]\d+f,\s+false,\s+true", source) for source in kernel_sources ), "activation_clamp=0.03125 still instantiated kUseSitu=true" try: diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 7e5e8a537c..eff1931a46 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -145,7 +145,10 @@ def create_inputs(): # Cast inputs to FP8/FP4 with per-32 UE8M0 SF. FP4 activations # remain routed-expert-only; shared experts retain upstream FP8. assert hidden % 128 == 0 and intermediate_hidden % 128 == 0 and shared_intermediate_hidden % 128 == 0 - use_fp4_acts = (os.environ.get('DG_USE_FP4_ACTS', '0') != '0' + # FP4 activations are selected by the buffer's `mma_type` + # (`mxf4xmxf4` / `nvfp4xnvfp4`), not by an env var, and remain a + # routed-expert-only path. + use_fp4_acts = (args.mma_type in ('mxf4xmxf4', 'nvfp4xnvfp4') and num_shared_experts == 0) if use_fp4_acts: x = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) @@ -165,8 +168,12 @@ def create_inputs(): shared_l1_weights = _cast_fp8_for_mega_moe(shared_l1_weights)[0::2] shared_l2_weights = _cast_fp8_for_mega_moe(shared_l2_weights)[0::2] + # NOTES: the routed weights must be interleaved for the buffer's own MMA kind + # (the FP4 kinds need the packed gate/up interleave); the shared-expert weights + # stay FP8 and therefore keep the default `fp8xfp4` interleave. transformed_l1_weights, transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)) + deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights, + mma_type=args.mma_type)) if num_shared_experts > 0: transformed_shared_l1_weights, transformed_shared_l2_weights = ( deep_gemm.transform_weights_for_mega_moe(shared_l1_weights, shared_l2_weights))