feat(mega): add fused shared side-LoRA kernels

csrc/apis/mega_backward.hpp

@@ -2,6 +2,7 @@
 
 #include "../jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp"
 #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe_backward.hpp"
+#include "../jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp"
 
 namespace deep_gemm::mega_backward {
 
@@ -476,6 +477,41 @@ static void bf16_mega_moe_backward_w13_combine(
 }
 
 static void register_apis(pybind11::module_& m) {
+    m.def(
+        "fp8_fp4_mega_moe_side_lora_backward",
+        &deep_gemm::sm100_fp8_fp4_mega_moe_side_lora_backward);
+    m.def(
+        "bf16_mega_moe_side_lora_backward",
+        &deep_gemm::sm100_bf16_mega_moe_side_lora_backward,
+        py::arg("gate_up_output"), py::arg("grad_h_output"),
+        py::arg("grad_gate_up_output"), py::arg("h_act_output"),
+        py::arg("h_weighted_output"), py::arg("x_pool_output"),
+        py::arg("grad_x_pool_output"), py::arg("grad_route_output"),
+        py::arg("grad_ye"), py::arg("grad_y_unweighted_output"),
+        py::arg("route_weights"), py::arg("w2_weights"),
+        py::arg("w13_weights"), py::arg("expert_counts"),
+        py::arg("grid_sync_counter"), py::arg("activation_limit"),
+        py::arg("activation"), py::arg("fast_math"),
+        py::arg("route_weight_mode"), py::arg("combine_order_mode"),
+        py::arg("down_unweighted_output"), py::arg("block_m"),
+        py::arg("direct_remote_grad_x"), py::arg("write_grad_x_pool"),
+        py::arg("clear_wgrad_padding"), py::arg("backward_grad_y"),
+        py::arg("backward_x"), py::arg("backward_topk_weights"),
+        py::arg("backward_grad_route"), py::arg("token_src_metadata"),
+        py::arg("backward_sym_buffer_ptrs"), py::arg("backward_rank"),
+        py::arg("num_max_tokens_per_rank"), py::arg("num_topk"),
+        py::arg("memory_mode"),
+        py::arg("side_lora_a1"), py::arg("side_lora_b1"),
+        py::arg("side_lora_a3"), py::arg("side_lora_b3"),
+        py::arg("side_lora_a2"), py::arg("side_lora_b2"),
+        py::arg("side_lora_q13"), py::arg("side_lora_q2"),
+        py::arg("side_lora_saved_h"), py::arg("side_lora_t13"),
+        py::arg("side_lora_t2"), py::arg("grad_side_lora_a1"),
+        py::arg("grad_side_lora_b1"), py::arg("grad_side_lora_a3"),
+        py::arg("grad_side_lora_b3"), py::arg("grad_side_lora_a2"),
+        py::arg("grad_side_lora_b2"), py::arg("expert_psum_rows"),
+        py::arg("padded_expert_counts"), py::arg("side_lora_scale"),
+        py::arg("kernel_trace") = py::none());
     m.def("fp8_fp4_mega_moe_backward_dgrad_swiglu_v2",
           &fp8_fp4_mega_moe_backward_dgrad_swiglu_v2,
           py::arg("gate_up_output"), py::arg("grad_h_output"),

csrc/jit_kernels/heuristics/mega_moe_side_lora.hpp

@@ -0,0 +1,100 @@
+#pragma once
+
+#include "mega_moe.hpp"
+#include "../impls/runtime_utils.hpp"
+
+namespace deep_gemm {
+
+// The side kernel adds shrink/expand phases around both expert GEMMs and
+// therefore retains a distinct scheduler. Keep its wave-capacity policy out
+// of the ordinary upstream scheduler/configuration.
+static int get_num_wave_pool_tokens(
+    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(
+        // All tokens come to all local experts in the wave
+        num_tokens_from_all_ranks * num_experts_per_wave,
+        // All routed tokens come to this local wave, and each expert needs a padding
+        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(
+    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) {
+
+    // Get max experts per wave limitation
+    int num_max_experts_per_wave = num_experts_per_rank;
+    while (num_max_experts_per_wave > 0 and
+           get_num_wave_pool_tokens(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");
+
+    // Reduce per-expert block count by this factor since uneven routing leaves some experts with fewer tokens
+    constexpr int kImbalanceFactor = 2;
+
+    // Count L1 blocks per expert assuming tokens are evenly spread across experts
+    const float num_expected_tokens_per_expert = static_cast<float>(num_tokens * num_topk) / num_experts_per_rank;
+    const int num_expected_m_blocks = std::max(ceil_div(static_cast<int>(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;
+
+    // Pick the smallest value whose total blocks (after imbalance reduction) can keep all SMs busy
+    int num_min_expected_experts_to_fill_sms = ceil_div(kImbalanceFactor * num_sms, num_expected_l1_blocks_per_expert);
+
+    // Most experts don't have tokens, calculate all experts at once
+    if (num_expected_tokens_per_expert < 1)
+        num_min_expected_experts_to_fill_sms = num_experts_per_rank;
+
+    // Ring capacity is the bottleneck
+    if (num_min_expected_experts_to_fill_sms >= num_max_experts_per_wave)
+        return num_max_experts_per_wave;
+
+    // When each expert nearly fills all SMs, use the smallest wave to maximize L2 cache reuse
+    if (num_expected_l1_blocks_per_expert >= num_sms) 
+        return num_min_expected_experts_to_fill_sms;
+
+    // Search to 2 * num_min_expected_experts_to_fill_sms for a value where the last partial
+    // wave has as many experts as possible relative to a full wave
+    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) {
+        int remainder = num_experts_per_rank % num_experts_per_wave;
+        float tail_ratio = (remainder == 0) ? 1.0f : static_cast<float>(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;
+}
+
+
+struct SideLoraMegaMoEConfig : MegaMoEConfig {
+    int num_experts_per_wave;
+};
+
+static SideLoraMegaMoEConfig get_side_lora_mega_moe_config(
+    const int num_ranks, const int num_experts, const int num_experts_per_rank,
+    const int num_max_tokens_per_rank, const int num_tokens, const int num_topk,
+    const int hidden, const int intermediate_hidden,
+    const int num_ring_tokens, const int num_sf_ring_tokens, const MmaKind mma_kind) {
+    const auto base = 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, mma_kind);
+    const int experts_per_wave = get_num_experts_per_wave_for_mega_moe(
+        num_experts_per_rank, num_tokens, num_topk, intermediate_hidden,
+        base.block_m, base.block_n, get_mega_moe_num_sms(),
+        num_ring_tokens, num_max_tokens_per_rank, num_ranks);
+    return {base, experts_per_wave};
+}
+
+}  // namespace deep_gemm

csrc/jit_kernels/impls/mega_moe_side_lora_launch.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "../../runtime/runtime.hpp"
+
+namespace deep_gemm {
+
+// Side kernels share DeepJIT's cache, current-stream launch and launch hooks.
+// This helper only constructs options; it owns no compiler or kernel cache.
+static deep_jit::cuda::LaunchOptions side_lora_launch_options(
+    const int num_ctas, const int num_threads,
+    const int smem_bytes, const int cluster_size) {
+    return {
+        .num_smem_bytes = smem_bytes,
+        .grid_dim = dim3(num_ctas, 1, 1),
+        .block_dim = dim3(num_threads, 1, 1),
+        .cluster_dim = dim3(cluster_size, 1, 1),
+    };
+}
+
+}  // namespace deep_gemm

csrc/jit_kernels/impls/sm100_bf16_gemm.hpp

@@ -52,6 +52,8 @@ class SM100BF16GemmRuntime final {
         bool combine_reduce = false;
         std::string combine_order_mode = "fixed_topk";
         uint32_t combine_num_extra_threads = 0;
+        bool mask_grouped_k_tail = false;
+        int k_alignment_override = 0;
     };
 
     static void compile_and_launch(const std::string& tag, const Args& args) {
@@ -76,7 +78,7 @@ static void __instantiate_kernel() {{
         {}, {}, {},
         {},
         {},
-        {}, {}, {}, {}
+        {}, {}, {}, {}, {}
     >);
 }};
 )",
@@ -91,15 +93,16 @@ static void __instantiate_kernel() {{
         args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads,
         args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
         args.gemm_config.launch_config.num_sms,
-        heuristics_runtime->get_mk_alignment_for_contiguous_layout(),
+        args.k_alignment_override != 0 ? args.k_alignment_override :
+            heuristics_runtime->get_mk_alignment_for_contiguous_layout(),
         args.gemm_config.layout.swap_ab, args.gemm_desc.ensure_zero_padding,
         to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation,
         to_string(args.gemm_desc.cd_dtype),
         args.epilogue.type,
         args.gemm_desc.tc_util,
         args.combine_num_ranks, args.fuse_combine,
         get_bf16_gemm_combine_order_mode_name(args.combine_order_mode),
-        args.combine_num_extra_threads));
+        args.combine_num_extra_threads, args.mask_grouped_k_tail));
 
         // Launch
         jit->launch(