SM90 (Hopper) FP4 MegaMoE fused kernel with swapAB small-batch path

csrc/apis/sm90_mega.hpp

@@ -1,10 +1,20 @@
 #pragma once
 
+#include <algorithm>
 #include <functional>
+#include <limits>
+#include <optional>
+#include <string>
+#include <tuple>
+#include <vector>
 
 #include "mega.hpp"
+#include "../jit/device_runtime.hpp"
+#include "../jit_kernels/impls/sm90_fp8_fp4_mega_moe.hpp"
 #include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp"
 #include "../jit_kernels/impls/sm90_mega_moe_pre_dispatch.hpp"
+#include "../utils/layout.hpp"
+#include "../utils/system.hpp"
 
 namespace deep_gemm::mega {
 
@@ -30,6 +40,92 @@ static void mega_moe_pre_dispatch_sm90(
         num_tokens, group_size, routed_scaling_factor);
 }
 
+static bool is_packed_fp4_storage_sm90(const torch::Tensor& t) {
+    return t.scalar_type() == kPackedFP4 or t.scalar_type() == torch::kByte;
+}
+
+static std::tuple<int, int, int> check_grouped_ab_sm90_fp4_mega_moe(const torch::Tensor& ab) {
+    const auto [num_groups, mn, packed_k] = get_shape<3>(ab);
+    DG_HOST_ASSERT(is_packed_fp4_storage_sm90(ab));
+    DG_HOST_ASSERT(get_major_type_ab(ab) == cute::UMMA::Major::K);
+    DG_HOST_ASSERT(packed_k > 0 and packed_k % 64 == 0);
+    return {num_groups, mn, packed_k * 2};
+}
+
+static void check_sm90_fp4_sfb_layout(const torch::Tensor& sf,
+                                      const int& mn, const int& k,
+                                      const int& num_groups) {
+    DG_HOST_ASSERT(sf.scalar_type() == torch::kInt);
+    DG_HOST_ASSERT(sf.dim() == 3);
+    DG_HOST_ASSERT(sf.size(0) == num_groups);
+    DG_HOST_ASSERT(sf.size(1) == mn);
+    DG_HOST_ASSERT(sf.size(2) == ceil_div(k, 128));
+    DG_HOST_ASSERT(sf.is_contiguous());
+}
+
+struct FP4SM90APIDefaults {
+    bool math_wg_participates_in_decode;
+    int num_math_wg_decode_warps;
+    int first_decode_assist_warp;
+    bool wide_load_decode;
+    bool early_b_decode;
+    bool decode_done_mbarrier;
+    bool l2_arrival_counter;
+    bool ss_nsplit;
+    bool swap_ab;
+    bool swap_ab_fast_amax;
+};
+
+static FP4SM90APIDefaults get_fp4_sm90_api_defaults(
+    const int& num_experts_per_rank, const int& num_tokens, const int& num_topk,
+    const int& intermediate_hidden, const bool& use_situ) {
+    // Simplified, shape-agnostic defaults, mirroring the FP8 path's style
+    // (`should_use_swap_ab_for_mega_moe_sm90`): one decode/prefill split plus a
+    // single swapAB threshold. The historical per-(shape x e-band) table was
+    // tuned point-by-point on benchmark batches; on the shapes that matter it
+    // collapsed to constants plus a few sliver bands, so it is retired.
+    (void)intermediate_hidden;
+    const float expected_tokens_per_expert =
+        static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
+    // Decode -> prefill boundary is activation-specific (Kimi/SiTU keeps its
+    // own tuning so DSV4/SwiGLU stays on e=80), and the SiTU 128/64 recipe
+    // additionally has a non-monotonic prefill mid band; share the predicate
+    // with the block-config heuristics so the feature bundle never mixes.
+    const bool prefill_band = is_fp4_sm90_prefill_band(
+        expected_tokens_per_expert, use_situ);
+    const bool decode_band =
+        expected_tokens_per_expert > 0.0f and !prefill_band;
+    // swapAB on/off kill-switch (default ON). Set DG_SM90_FP4_SWAP_AB=0 to force
+    // the non-swap path for A/B accuracy comparison.
+    const bool swap_ab_env_enabled = get_env<int>("DG_SM90_FP4_SWAP_AB", 1) != 0;
+    // Measured crossover on H20: swapAB wins at e=18.3 (+3.1% on Kimi
+    // 128/64) and loses from e~23 on (-6..-18%), so the fallback bound is 20
+    // (FP8 uses 30; the FP4 kernel pays extra decode work on the swapped
+    // path, and larger e also lands wider WGMMA-N buckets). The original
+    // swiglu tuning had 16; doc 15.15. DG_SM90_FP4_SWAP_AB_MAX_E overrides.
+    // With the SiTU 128/64 cost model active, the swapAB decision comes from
+    // the same argmin as the prefill band so all three kinds stay coherent.
+    const float swap_ab_max_e = static_cast<float>(
+        get_env<int>("DG_SM90_FP4_SWAP_AB_MAX_E", 20));
+    const bool swap_ab = use_fp4_sm90_situ_cost_model(use_situ)
+        ? (get_fp4_sm90_situ_config_kind(expected_tokens_per_expert) ==
+           FP4SM90ConfigKind::kSwapAB)
+        : (swap_ab_env_enabled and decode_band and
+           expected_tokens_per_expert < swap_ab_max_e);
+    return {
+        /*math_wg_participates_in_decode=*/ false,
+        /*num_math_wg_decode_warps=*/ 0,
+        /*first_decode_assist_warp=*/ 2,
+        /*wide_load_decode=*/ decode_band,
+        /*early_b_decode=*/ prefill_band,
+        /*decode_done_mbarrier=*/ expected_tokens_per_expert > 0.0f,
+        /*l2_arrival_counter=*/ false,
+        /*ss_nsplit=*/ prefill_band,
+        /*swap_ab=*/ swap_ab,
+        /*swap_ab_fast_amax=*/ false
+    };
+}
+
 static std::tuple<int64_t, std::function<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(const torch::Tensor&)>>
 get_symm_buffer_size_for_sm90_mega_moe(
     const int& num_ranks, const int& num_experts,
@@ -38,16 +134,21 @@ get_symm_buffer_size_for_sm90_mega_moe(
     const bool& use_fp8_dispatch, const std::string& activation) {
     DG_HOST_ASSERT(num_experts % num_ranks == 0);
     DG_HOST_ASSERT(use_fp8_dispatch);
-    DG_HOST_ASSERT(activation == "swiglu");
+    DG_HOST_ASSERT(activation == "swiglu" or activation == "situ");
+    const bool use_situ = activation == "situ";
+    const auto [sm90_l1_act_sf_gran_k, sm90_l2_act_sf_gran_k] =
+        get_act_sf_grans_for_mega_moe_sm90_fp4(use_situ);
 
     const auto workspace = layout::SM90Workspace(
         nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
 
     const auto fp8_token_layout = layout::Data(hidden);
     const auto bf16_token_layout = layout::Data(hidden * 2);
     const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden);
-    const auto fp8_sf_layout = layout::Data(hidden / 32);
-    const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 16);
+    const auto fp8_sf_layout =
+        layout::Data(hidden * static_cast<int>(sizeof(float)) / sm90_l1_act_sf_gran_k);
+    const auto fp8_intermediate_sf_layout =
+        layout::Data(intermediate_hidden * static_cast<int>(sizeof(float)) / sm90_l2_act_sf_gran_k);
     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);
@@ -96,6 +197,8 @@ get_symm_buffer_size_for_sm90_mega_moe(
         l2_sf_buffer.get_end_ptr());
 
     DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
+    DG_HOST_ASSERT(hidden % sm90_l1_act_sf_gran_k == 0);
+    DG_HOST_ASSERT(intermediate_hidden % sm90_l2_act_sf_gran_k == 0);
 
     auto slice_input_buffers = [=](const torch::Tensor& buffer) {
         auto x = torch::from_blob(
@@ -104,7 +207,7 @@ get_symm_buffer_size_for_sm90_mega_moe(
             torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
         auto x_sf = torch::from_blob(
             math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_sf_buffer.base)),
-            {num_max_tokens_per_rank, hidden / 128},
+            {num_max_tokens_per_rank, hidden / sm90_l1_act_sf_gran_k},
             torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
         auto topk_idx = torch::from_blob(
             math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_topk_idx_buffer.base)),
@@ -120,7 +223,7 @@ get_symm_buffer_size_for_sm90_mega_moe(
             torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
         auto l1_acts_sf = torch::from_blob(
             math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_sf_buffer.base)),
-            {num_max_padded_sf_pool_tokens, hidden / 128},
+            {num_max_padded_sf_pool_tokens, hidden / sm90_l1_act_sf_gran_k},
             {1, num_max_padded_sf_pool_tokens},
             torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
         auto l2_acts = torch::from_blob(
@@ -129,7 +232,7 @@ get_symm_buffer_size_for_sm90_mega_moe(
             torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
         auto l2_acts_sf = torch::from_blob(
             math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_sf_buffer.base)),
-            {num_max_padded_sf_pool_tokens, intermediate_hidden / 64},
+            {num_max_padded_sf_pool_tokens, intermediate_hidden / sm90_l2_act_sf_gran_k},
             {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);
@@ -221,4 +324,124 @@ static void fp8_mega_moe(
         sym_buffer.zero_();
 }
 
+static void fp8_fp4_mega_moe_sm90(
+    const torch::Tensor& y,
+    const std::tuple<torch::Tensor, torch::Tensor>& l1_weights_tuple,
+    const std::tuple<torch::Tensor, torch::Tensor>& l2_weights_tuple,
+    const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
+    const torch::Tensor& sym_buffer,
+    const std::vector<int64_t>& sym_buffer_ptrs, const int& rank_idx,
+    const int& num_max_tokens_per_rank,
+    const int& num_experts, const int& num_topk,
+    const std::tuple<int, int, int>& recipe,
+    const std::string& activation,
+    const std::optional<float>& activation_clamp_opt,
+    const std::optional<float>& activation_alpha_opt,
+    const std::optional<float>& activation_linear_beta_opt,
+    const bool& fast_math
+) {
+    const auto [l1_weights, l1_weights_sf] = l1_weights_tuple;
+    const auto [l2_weights, l2_weights_sf] = l2_weights_tuple;
+
+    const auto arch_major = device_runtime->get_arch_major();
+    DG_HOST_ASSERT(arch_major == 9);
+
+    const auto num_tokens = static_cast<int>(y.size(0));
+    const auto [rm, rn, rk] = recipe;
+    DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32);
+    DG_HOST_ASSERT(activation == "swiglu" or activation == "situ");
+    const bool use_situ = activation == "situ";
+
+    const auto activation_clamp =
+        activation_clamp_opt.value_or(std::numeric_limits<float>::infinity());
+    DG_HOST_ASSERT(activation_clamp >= 0);
+    const auto activation_alpha = activation_alpha_opt.value_or(4.0f);
+    const auto activation_linear_beta = activation_linear_beta_opt.value_or(25.0f);
+    DG_HOST_ASSERT(not use_situ or
+                   (activation_alpha > 0.0f and activation_linear_beta > 0.0f));
+
+    const auto [num_experts_per_rank, intermediate_hidden_2, hidden] =
+        check_grouped_ab_sm90_fp4_mega_moe(l1_weights);
+    const auto [num_experts_per_rank_, hidden_, intermediate_hidden] =
+        check_grouped_ab_sm90_fp4_mega_moe(l2_weights);
+    DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank);
+    DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_);
+    DG_HOST_ASSERT(hidden == hidden_);
+    DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden);
+    DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous());
+    DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
+    DG_HOST_ASSERT(intermediate_hidden / 64 <= 64);
+
+    check_sm90_fp4_sfb_layout(l1_weights_sf, intermediate_hidden * 2, hidden,
+                              num_experts_per_rank);
+    check_sm90_fp4_sfb_layout(l2_weights_sf, hidden, intermediate_hidden,
+                              num_experts_per_rank);
+
+    if (cumulative_local_expert_recv_stats.has_value()) {
+        DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt);
+        DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank);
+        DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous());
+    }
+
+    const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
+    const auto num_experts_ = num_experts_per_rank * num_ranks;
+    const auto [num_required_bytes, slice] = get_symm_buffer_size_for_sm90_mega_moe(
+        num_ranks, num_experts,
+        num_max_tokens_per_rank, num_topk,
+        hidden, intermediate_hidden,
+        true, activation);
+    DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast<size_t>(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);
+    (void)x;
+    (void)x_sf;
+    (void)topk_idx;
+    (void)topk_weights;
+
+    DG_HOST_ASSERT(get_env<int>("DG_USE_FP4_ACTS") == 0);
+    DG_HOST_ASSERT(get_env<int>("DG_USE_FP8_COMBINE") == 0);
+
+    const auto fp4_defaults = get_fp4_sm90_api_defaults(
+        num_experts_per_rank, num_tokens, num_topk, intermediate_hidden,
+        use_situ);
+    sm90_fp8_fp4_mega_moe(y,
+                          l1_acts, l1_acts_sf,
+                          l2_acts, l2_acts_sf,
+                          l1_weights, l2_weights,
+                          l1_weights_sf, l2_weights_sf,
+                          cumulative_local_expert_recv_stats,
+                          sym_buffer_ptrs,
+                          rank_idx, num_max_tokens_per_rank,
+                          num_experts_per_rank,
+                          num_tokens, num_topk,
+                          hidden, intermediate_hidden,
+                          activation,
+                          activation_alpha,
+                          activation_linear_beta,
+                          activation_clamp, fast_math,
+                          fp4_defaults.math_wg_participates_in_decode,
+                          fp4_defaults.num_math_wg_decode_warps,
+                          fp4_defaults.first_decode_assist_warp,
+                          fp4_defaults.wide_load_decode,
+                          fp4_defaults.early_b_decode,
+                          fp4_defaults.decode_done_mbarrier,
+                          fp4_defaults.l2_arrival_counter,
+                          fp4_defaults.ss_nsplit,
+                          fp4_defaults.swap_ab,
+                          fp4_defaults.swap_ab_fast_amax);
+
+    if (get_env<int>("DG_COMM_KERNEL_DEBUG"))
+        sym_buffer.zero_();
+}
+
+static void register_sm90_apis(pybind11::module_& m) {
+#if DG_TENSORMAP_COMPATIBLE
+    m.def("get_token_alignment_for_sm90_mega_moe", &get_token_alignment_for_sm90_mega_moe);
+    m.def("get_symm_buffer_size_for_sm90_mega_moe", &get_symm_buffer_size_for_sm90_mega_moe);
+    m.def("fp8_fp4_mega_moe_sm90", &fp8_fp4_mega_moe_sm90);
+    m.def("fp8_mega_moe", &fp8_mega_moe);
+#endif
+}
+
 } // namespace deep_gemm::mega