From ce30d5b56c6f3c5091e14ebef3412bf1b198102a Mon Sep 17 00:00:00 2001 From: morgendave Date: Thu, 13 Aug 2026 17:48:55 +0000 Subject: [PATCH 01/15] feat(mega): add fused shared side-LoRA kernels --- csrc/apis/mega.hpp | 384 +- csrc/apis/mega_backward.hpp | 36 + csrc/jit/compiler.hpp | 8 +- ...sm100_bf16_mega_moe_side_lora_backward.hpp | 2149 +++++++ .../sm100_bf16_mega_moe_side_lora_forward.hpp | 579 ++ .../impls/sm100_bf16_mega_moe_wgrad.hpp | 12 +- ...100_fp8_fp4_mega_moe_side_lora_forward.hpp | 517 ++ deep_gemm/__init__.py | 6 + .../impls/mega_moe_side_lora_params.cuh | 41 + ...sm100_bf16_mega_moe_side_lora_backward.cuh | 5656 +++++++++++++++++ .../sm100_bf16_mega_moe_side_lora_forward.cuh | 2585 ++++++++ ...100_fp8_fp4_mega_moe_side_lora_forward.cuh | 2230 +++++++ .../scheduler/mega_moe_side_lora.cuh | 262 + deep_gemm/mega/__init__.py | 278 +- deep_gemm/mega/backward.py | 304 +- tests/benchmark_mega_moe_native_side_lora.py | 306 + tests/run_mega_moe_side_lora_edge_matrix.py | 97 + tests/test_mega_moe_native_side_lora.py | 972 +++ 18 files changed, 16406 insertions(+), 16 deletions(-) create mode 100644 csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp create mode 100644 csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_forward.hpp create mode 100644 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/mega_moe_side_lora_params.cuh create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_forward.cuh create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.cuh create mode 100644 deep_gemm/include/deep_gemm/scheduler/mega_moe_side_lora.cuh create mode 100644 tests/benchmark_mega_moe_native_side_lora.py create mode 100644 tests/run_mega_moe_side_lora_edge_matrix.py create mode 100644 tests/test_mega_moe_native_side_lora.py diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index a876221710..e42fd76493 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -11,7 +11,9 @@ #endif #include "../jit/device_runtime.hpp" #include "../jit_kernels/impls/sm100_bf16_mega_moe.hpp" +#include "../jit_kernels/impls/sm100_bf16_mega_moe_side_lora_forward.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" +#include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.hpp" namespace deep_gemm::mega { @@ -22,7 +24,7 @@ using MegaMoELegacySlices = std::tuple< using MegaMoEExpandedSlices = std::tuple< torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, - torch::Tensor, torch::Tensor, torch::Tensor>; + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>; static int get_token_alignment_for_mega_moe() { return layout::kLCMCandidateBlockM; @@ -125,6 +127,12 @@ get_symm_buffer_size_for_mega_moe_v2( const auto backward_route_grad_buffer = layout::Buffer( input_topk_weights_layout, 1, num_max_tokens_per_rank, combine_token_buffer.get_end_ptr()); + // Unquantized source plane used only by the dedicated MXFP4 + BF16 + // side-LoRA specialization. It is appended so every existing offset and + // the legacy buffer ABI remain stable. + const auto side_lora_source_buffer = layout::Buffer( + bf16_token_layout, 1, with_sf ? num_max_tokens_per_rank : 0, + backward_route_grad_buffer.get_end_ptr()); // Check SF buffer requirements if (with_sf) { @@ -205,10 +213,25 @@ get_symm_buffer_size_for_mega_moe_v2( .dtype(torch::kFloat32) .device(buffer.device())) : torch::Tensor(); + auto side_lora_source = with_sf && + buffer.nbytes() >= static_cast( + reinterpret_cast( + side_lora_source_buffer.get_end_ptr())) + ? torch::from_blob( + math::advance_ptr( + buffer.data_ptr(), + reinterpret_cast( + side_lora_source_buffer.base)), + {num_max_tokens_per_rank, hidden}, + torch::TensorOptions() + .dtype(torch::kBFloat16) + .device(buffer.device())) + : torch::Tensor(); return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, - token_src_metadata, backward_grad_y, backward_grad_route); + token_src_metadata, backward_grad_y, + backward_grad_route, side_lora_source); }; - return {reinterpret_cast(backward_route_grad_buffer.get_end_ptr()), slice_input_buffers}; + return {reinterpret_cast(side_lora_source_buffer.get_end_ptr()), slice_input_buffers}; } // Keep the original raw _C slicer ABI and allocation size. Training callers @@ -235,7 +258,7 @@ get_symm_buffer_size_for_mega_moe( auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, token_src_metadata, backward_grad_y, - _backward_grad_route] = + _backward_grad_route, _side_lora_source] = expanded_slicer(buffer); return std::make_tuple( x, x_sf, topk_idx, topk_weights, l1_acts, @@ -243,7 +266,9 @@ get_symm_buffer_size_for_mega_moe( token_src_metadata, backward_grad_y); }; return { - expanded_num_bytes - route_plane_bytes, + expanded_num_bytes - route_plane_bytes - + static_cast(num_max_tokens_per_rank) * hidden * + (is_mma_with_sf(parse_mma_kind(mma_type)) ? 2 : 0), legacy_slicer}; } @@ -343,13 +368,16 @@ static void fp8_fp4_mega_moe( const auto num_required_bytes = expanded_num_required_bytes - static_cast(num_max_tokens_per_rank) * - num_topk * sizeof(float); + num_topk * sizeof(float) - + static_cast(num_max_tokens_per_rank) * hidden * + sizeof(at::BFloat16); DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); // Already registered tensors const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, - token_src_metadata, backward_grad_y, _backward_grad_route] = slice(sym_buffer); + token_src_metadata, backward_grad_y, _backward_grad_route, + _side_lora_source] = slice(sym_buffer); // Dispatch into different architectures if (arch_major == 10) { @@ -378,6 +406,158 @@ static void fp8_fp4_mega_moe( sym_buffer.zero_(); } +static void fp8_fp4_mega_moe_side_lora( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, const int& rank_idx, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + const std::tuple& recipe, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math, + const int& num_ring_tokens, + const std::optional& saved_l1_preact, + const std::string& route_weight_mode, + const std::optional& saved_down_unweighted, + const std::optional& num_config_tokens_opt, + const torch::Tensor& saved_x, + const torch::Tensor& saved_h_unweighted, + const torch::Tensor& side_lora_a1, + const torch::Tensor& side_lora_b1, + const torch::Tensor& side_lora_a3, + const torch::Tensor& side_lora_b3, + const torch::Tensor& side_lora_a2, + const torch::Tensor& side_lora_b2, + const torch::Tensor& side_lora_l1_scratch, + const torch::Tensor& side_lora_l2_scratch, + const torch::Tensor& side_lora_ready, + const float& side_lora_scale +) { + const auto [l1_weights, l1_weights_sf] = l1_weights_tuple; + const auto [l2_weights, l2_weights_sf] = l2_weights_tuple; + + // Config checks + const auto num_tokens = static_cast(y.size(0)); + const auto num_config_tokens = + num_config_tokens_opt.value_or(num_tokens); + const auto [rm, rn, rk] = recipe; + DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32); + DG_HOST_ASSERT(activation == "swiglu" or activation == "geglu"); + DG_HOST_ASSERT( + route_weight_mode == "pre_down" || + route_weight_mode == "post_down"); + + // Activation checks + const auto activation_clamp = + activation_clamp_opt.value_or(std::numeric_limits::infinity()); + DG_HOST_ASSERT(activation_clamp >= 0); + + // Tensor checks + DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K); + const auto arch_major = device_runtime->get_arch_major(); + const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = + check_grouped_ab_fp8_fp4(l1_weights, cute::UMMA::Major::K, arch_major); + const auto [num_experts_per_rank_, hidden_, intermediate_hidden] = + check_grouped_ab_fp8_fp4(l2_weights, cute::UMMA::Major::K, arch_major); + 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(num_config_tokens >= num_tokens); + DG_HOST_ASSERT(num_config_tokens <= num_max_tokens_per_rank); + + // Check weight SF layout for UE8M0 packing, MN-major, and TMA alignment + constexpr int kGranMN = 1, kGranK = 32; + 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, + num_experts_per_rank, true, false, torch::kInt); + + // Check stats counter + 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()); + } + + // Check buffer bytes + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts_ = num_experts_per_rank * num_ranks; + const auto num_max_pool_tokens = + layout::get_num_max_pool_tokens( + num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank); + if (saved_down_unweighted.has_value()) { + DG_HOST_ASSERT( + saved_down_unweighted->scalar_type() == + torch::kBFloat16); + DG_HOST_ASSERT(saved_down_unweighted->is_contiguous()); + DG_HOST_ASSERT(saved_down_unweighted->dim() == 2); + DG_HOST_ASSERT(saved_down_unweighted->size(1) == hidden); + DG_HOST_ASSERT(saved_down_unweighted->size(0) > 0); + DG_HOST_ASSERT( + saved_down_unweighted->size(0) <= + num_max_pool_tokens); + } + const auto [expanded_num_required_bytes, slice] = + get_symm_buffer_size_for_mega_moe_v2( + num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + "fp8xfp4", activation, num_ring_tokens); + const auto num_required_bytes = expanded_num_required_bytes; + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(num_experts == num_experts_); + + // Already registered tensors + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + token_src_metadata, backward_grad_y, _backward_grad_route, + side_lora_source] = slice(sym_buffer); + + // Dispatch into different architectures + if (arch_major == 10) { + sm100_fp8_fp4_mega_moe_side_lora_forward(y, + saved_l1_preact, + 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_config_tokens, num_topk, + hidden, intermediate_hidden, + activation, activation_clamp, fast_math, + route_weight_mode, + saved_down_unweighted, + side_lora_source, + saved_x, + saved_h_unweighted, + side_lora_a1, side_lora_b1, + side_lora_a3, side_lora_b3, + side_lora_a2, side_lora_b2, + side_lora_l1_scratch, + side_lora_l2_scratch, + side_lora_ready, + side_lora_scale); + } else { + DG_HOST_UNREACHABLE("Unsupported architecture"); + } + + // Zero the entire symmetric buffer for debug mode + // NOTES: caller must re-copy inputs into the buffer before each kernel call + if (get_env("DG_COMM_KERNEL_DEBUG")) + sym_buffer.zero_(); +} + + static void bf16_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_weights, @@ -482,13 +662,14 @@ static void bf16_mega_moe( // Already registered tensors const auto [x, _x_sf, topk_idx, topk_weights, l1_acts, _l1_acts_sf, l2_acts, _l2_acts_sf, - _token_src_metadata, _backward_grad_y, _backward_grad_route] = slice(sym_buffer); + _token_src_metadata, _backward_grad_y, + _backward_grad_route, _side_lora_source] = slice(sym_buffer); // Dispatch into different architectures if (arch_major == 10) { sm100_bf16_mega_moe(y, saved_l1_preact, - l1_acts, l2_acts, + l1_acts, l2_acts, l1_weights, l2_weights, cumulative_local_expert_recv_stats, sym_buffer_ptrs, @@ -517,6 +698,169 @@ static void bf16_mega_moe( sym_buffer.zero_(); } +static void bf16_mega_moe_side_lora( + const torch::Tensor& y, + const torch::Tensor& l1_weights, + const torch::Tensor& l2_weights, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, const int& rank_idx, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math, + const int& num_ring_tokens, + const std::optional& saved_l1_preact, + const std::string& route_weight_mode, + const std::optional& saved_h_unweighted, + const std::optional& saved_h_weighted, + const std::optional& saved_down_unweighted, + const int& num_config_tokens, + const std::string& combine_order_mode, + const std::optional& precomputed_route_counts, + const std::optional& active_pool_rows, + const std::optional& route_count_mismatch, + const std::optional& saved_x, + const std::optional& side_lora_a1, + const std::optional& side_lora_b1, + const std::optional& side_lora_a3, + const std::optional& side_lora_b3, + const std::optional& side_lora_a2, + const std::optional& side_lora_b2, + const std::optional& side_lora_l1_scratch, + const std::optional& side_lora_l2_scratch, + const std::optional& side_lora_ready, + const float& side_lora_scale +) { + // Config checks + const auto num_tokens = static_cast(y.size(0)); + DG_HOST_ASSERT(activation == "swiglu" or activation == "geglu"); + DG_HOST_ASSERT( + route_weight_mode == "pre_down" || + route_weight_mode == "post_down"); + DG_HOST_ASSERT( + combine_order_mode == "fixed_topk" || + combine_order_mode == "deepep" || + combine_order_mode == "deepep_v1"); + + // Activation checks + const auto activation_clamp = + activation_clamp_opt.value_or(std::numeric_limits::infinity()); + DG_HOST_ASSERT(activation_clamp >= 0); + + // Tensor checks + DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K); + const auto arch_major = device_runtime->get_arch_major(); + const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = get_shape<3>(l1_weights); + const auto [num_experts_per_rank_, hidden_, intermediate_hidden] = get_shape<3>(l2_weights); + DG_HOST_ASSERT(l1_weights.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(l2_weights.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(y.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(y.is_contiguous()); + DG_HOST_ASSERT(y.sizes() == torch::IntArrayRef({num_tokens, hidden})); + DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank); + DG_HOST_ASSERT(num_config_tokens >= num_tokens); + DG_HOST_ASSERT(num_config_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()); + + // Check stats counter + 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()); + } + DG_HOST_ASSERT( + precomputed_route_counts.has_value() == + active_pool_rows.has_value()); + DG_HOST_ASSERT( + precomputed_route_counts.has_value() == + route_count_mismatch.has_value()); + if (precomputed_route_counts.has_value()) { + DG_HOST_ASSERT( + precomputed_route_counts->scalar_type() == + torch::kInt); + DG_HOST_ASSERT( + precomputed_route_counts->numel() == num_experts); + DG_HOST_ASSERT(precomputed_route_counts->is_contiguous()); + DG_HOST_ASSERT( + route_count_mismatch->scalar_type() == torch::kInt); + DG_HOST_ASSERT(route_count_mismatch->numel() == 1); + DG_HOST_ASSERT(route_count_mismatch->is_contiguous()); + DG_HOST_ASSERT(*active_pool_rows > 0); + } + + // Check buffer bytes + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts_ = num_experts_per_rank * num_ranks; + const auto [expanded_num_required_bytes, slice] = + get_symm_buffer_size_for_mega_moe_v2( + num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + "bf16xbf16", activation, num_ring_tokens); + const auto num_required_bytes = + expanded_num_required_bytes - + static_cast(num_max_tokens_per_rank) * + num_topk * sizeof(float); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(num_experts == num_experts_); + + // Already registered tensors + const auto [x, _x_sf, topk_idx, topk_weights, l1_acts, _l1_acts_sf, + l2_acts, _l2_acts_sf, _token_src_metadata, + _backward_grad_y, _backward_grad_route, + _side_lora_source] = + slice(sym_buffer); + + // Dispatch into different architectures + if (arch_major == 10) { + sm100_bf16_mega_moe_side_lora_forward(y, + saved_l1_preact, + l1_acts, l2_acts, + l1_weights, l2_weights, + cumulative_local_expert_recv_stats, + sym_buffer_ptrs, + rank_idx, num_max_tokens_per_rank, + num_experts_per_rank, + num_tokens, num_config_tokens, num_topk, + hidden, intermediate_hidden, + activation, + activation_clamp, fast_math, + route_weight_mode, + saved_h_unweighted, + saved_h_weighted, + saved_down_unweighted, + combine_order_mode, + precomputed_route_counts, + active_pool_rows, + route_count_mismatch, + saved_x, + side_lora_a1, + side_lora_b1, + side_lora_a3, + side_lora_b3, + side_lora_a2, + side_lora_b2, + side_lora_l1_scratch, + side_lora_l2_scratch, + side_lora_ready, + side_lora_scale); + } else { + DG_HOST_UNREACHABLE("Unsupported architecture"); + } + + // Zero the entire symmetric buffer for debug mode + // NOTES: caller must re-copy inputs into the buffer before each kernel call + if (get_env("DG_COMM_KERNEL_DEBUG")) + sym_buffer.zero_(); +} + + static void register_apis(pybind11::module_& m) { #if DG_TENSORMAP_COMPATIBLE m.def("get_token_alignment_for_mega_moe", &get_token_alignment_for_mega_moe); @@ -545,7 +889,29 @@ static void register_apis(pybind11::module_& m) { py::arg("route_weight_mode") = "pre_down", py::arg("saved_down_unweighted") = py::none(), py::arg("num_config_tokens") = py::none()); + m.def( + "fp8_fp4_mega_moe_side_lora", + &fp8_fp4_mega_moe_side_lora, + py::arg("y"), py::arg("l1_weights_tuple"), + py::arg("l2_weights_tuple"), + py::arg("cumulative_local_expert_recv_stats"), + py::arg("sym_buffer"), py::arg("sym_buffer_ptrs"), + py::arg("rank_idx"), py::arg("num_max_tokens_per_rank"), + py::arg("num_experts"), py::arg("num_topk"), + py::arg("recipe"), py::arg("activation"), + py::arg("activation_clamp_opt"), py::arg("fast_math"), + py::arg("num_ring_tokens"), py::arg("saved_l1_preact"), + py::arg("route_weight_mode"), + py::arg("saved_down_unweighted"), + py::arg("num_config_tokens"), py::arg("saved_x"), + py::arg("saved_h_unweighted"), 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_l1_scratch"), + py::arg("side_lora_l2_scratch"), py::arg("side_lora_ready"), + py::arg("side_lora_scale")); m.def("bf16_mega_moe", &bf16_mega_moe); + m.def("bf16_mega_moe_side_lora", &bf16_mega_moe_side_lora); #endif } diff --git a/csrc/apis/mega_backward.hpp b/csrc/apis/mega_backward.hpp index e9feb0e931..0a3c34d80b 100644 --- a/csrc/apis/mega_backward.hpp +++ b/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 { @@ -477,6 +478,41 @@ static void bf16_mega_moe_backward_w13_combine( static void register_apis(pybind11::module_& m) { #if DG_TENSORMAP_COMPATIBLE + 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"), diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp index 7d85a5f556..16b5eb67d7 100644 --- a/csrc/jit/compiler.hpp +++ b/csrc/jit/compiler.hpp @@ -202,10 +202,14 @@ class NVCCCompiler final: public Compiler { // The override the compiler flags // Only NVCC >= 12.9 supports arch-specific family suffix const auto arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9); - flags = fmt::format("{} -I{} --gpu-architecture=sm_{} " + const auto source_cutlass_include = + library_root_path.parent_path() / + "third-party" / "cutlass" / "include"; + flags = fmt::format("{} -I{} -I{} --gpu-architecture=sm_{} " "--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi " "-O3 --expt-relaxed-constexpr --expt-extended-lambda", - flags, library_include_path.c_str(), arch); + flags, library_include_path.c_str(), + source_cutlass_include.c_str(), arch); } void compile(const std::string &code, const std::filesystem::path& dir_path, diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp new file mode 100644 index 0000000000..03b90c49db --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp @@ -0,0 +1,2149 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/math.hpp" +#include "runtime_utils.hpp" +#include "sm100_bf16_mega_moe_wgrad.hpp" + +#include +#include +#include + +namespace deep_gemm { + +// MegaMoE's physical expert boundaries are padded by the forward BLOCK_M. +// Select a grouped-GEMM tile that divides that exact boundary alignment; this +// keeps every tensor-core tile inside one expert without repacking the pool or +// launching a GEMM per expert. +static GemmConfig sm100_bf16_mega_moe_side_lora_rank_config( + const GemmDesc& desc, + const int pool_block_m) { + DG_HOST_ASSERT(pool_block_m >= 16 && pool_block_m <= 256 && + pool_block_m % 16 == 0); + // The swapped M-grouped kernel maps logical M onto UMMA N, which accepts + // every 16-row step through 256. Matching the forward tile exactly avoids + // both cross-expert tiles and any extra pool padding. + const Layout layout{ + .swap_ab = true, + .block_m = pool_block_m, + .block_n = 128, + .block_k = 64, + .cluster_m = 1, + .cluster_n = 1, + }; + const auto storage = SM100ArchSpec::get_storage_config(desc, layout); + return GemmConfig{ + .layout = layout, + .storage_config = storage, + .pipeline_config = + SM100ArchSpec::get_pipeline_config(desc, layout, storage), + .launch_config = SM100ArchSpec::get_launch_config(desc, layout), + }; +} + +static void sm100_bf16_mega_moe_side_lora_rank_gemm( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& d, + const torch::Tensor& expert_psum_rows, + const torch::Tensor&, + const int num_groups, + const int num_pool_rows, + const int n, + const int k, + const int pool_block_m, + const cute::UMMA::Major major_a, + const cute::UMMA::Major major_b) { + const GemmDesc desc{ + .gemm_type = GemmType::MGroupedContiguousWithPsumLayout, + .kernel_type = KernelType::KernelNoSF, + .m = num_pool_rows, + .n = n, + .k = k, + .num_groups = num_groups, + .a_dtype = a.scalar_type(), + .b_dtype = b.scalar_type(), + .cd_dtype = d.scalar_type(), + .major_a = major_a, + .major_b = major_b, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = "nk", + .ensure_zero_padding = false, + .expected_m = num_pool_rows / num_groups, + .expected_n = n, + .expected_k = k, + .expected_num_groups = num_groups, + }; + const auto config = + sm100_bf16_mega_moe_side_lora_rank_config(desc, pool_block_m); + const auto tensor_map_a = make_tma_a_desc( + major_a, a, num_pool_rows, k, + config.storage_config.load_block_m, config.layout.block_k, + static_cast(a.stride(get_non_contiguous_dim(major_a))), 1, + config.storage_config.swizzle_a_mode); + const auto tensor_map_b = make_tma_b_desc( + major_b, b, n, k, + config.storage_config.load_block_n, config.layout.block_k, + static_cast(b.stride(get_non_contiguous_dim(major_b))), + num_groups, config.storage_config.swizzle_b_mode); + const auto tensor_map_d = make_tma_cd_desc( + d, num_pool_rows, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(d.stride(-2)), 1, + config.storage_config.swizzle_cd_mode); + const SM100BF16GemmRuntime::Args args{ + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs( + config.launch_config.num_sms, + config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .grouped_layout = expert_psum_rows.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_cd = tensor_map_d, + }; + const auto code = SM100BF16GemmRuntime::generate(args); + const auto runtime = compiler->build( + "sm100_bf16_mega_moe_side_lora_rank_gemm", code); + SM100BF16GemmRuntime::launch(runtime, args); +} + +// Shared A1/A3/B2 factors are one matrix for the whole EP-local route pool. +// Keep their contractions in the native SM100 tensor-core path, but use one +// dense GEMM instead of repeating the same weight through every expert group. +static void sm100_bf16_mega_moe_side_lora_shared_gemm( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& d, + const int m, + const int n, + const int k, + const std::string& compiled_dims = "mnk") { + sm100_bf16_gemm( + a, b, std::nullopt, d, m, n, k, + get_major_type_ab(a), get_major_type_ab(b), compiled_dims); +} + +static std::string get_side_lora_backward_route_weight_mode_name( + const std::string& route_weight_mode) { + if (route_weight_mode == "pre_down") + return "RouteWeightMode::PreDown"; + if (route_weight_mode == "post_down") + return "RouteWeightMode::PostDown"; + DG_HOST_UNREACHABLE("Unsupported route weight mode"); +} + +static std::string get_side_lora_backward_combine_order_mode_name( + const std::string& combine_order_mode) { + if (combine_order_mode == "fixed_topk") + return "CombineOrderMode::FixedTopK"; + if (combine_order_mode == "deepep") + return "CombineOrderMode::DeepEP"; + if (combine_order_mode == "deepep_v1") + return "CombineOrderMode::DeepEPV1"; + DG_HOST_UNREACHABLE("Unsupported combine order mode"); +} + +class SM100BF16MegaMoESideLoraBackwardWaveRuntime final + : public LaunchRuntime { +public: + struct Args { + int hidden; + int intermediate_hidden; + int num_experts; + int num_pool_rows; + int num_acts_rows; + int num_sf_pool_rows; + int block_m; + int block_n; + int block_k; + int sf_block_m; + int sf_block_n; + int num_stages; + int num_sms; + int num_ranks; + bool bf16_mode = false; + std::string activation = "swiglu"; + bool fast_math = false; + std::string route_weight_mode = "pre_down"; + std::string combine_order_mode = "fixed_topk"; + + const int* expert_counts; + layout::SymBuffer<> backward_sym_buffer; + layout::Workspace backward_workspace; + const cutlass::bfloat16_t* backward_grad_y; + const cutlass::bfloat16_t* backward_x; + const float* backward_topk_weights; + float* backward_grad_route; + const layout::TokenSrcMetadata* token_src_metadata; + uint32_t num_topk; + uint32_t acts_sf_stride; + CUtensorMap tensor_map_acts; + CUtensorMap tensor_map_acts_sf; + CUtensorMap tensor_map_weights; + CUtensorMap tensor_map_weights_sf; + CUtensorMap tensor_map_output; + CUtensorMap tensor_map_grad_ye; + CUtensorMap tensor_map_w2_dequant; + CUtensorMap tensor_map_w2_weights; + CUtensorMap tensor_map_w2_scales; + CUtensorMap tensor_map_w13_dequant; + CUtensorMap tensor_map_w13_weights; + CUtensorMap tensor_map_w13_scales; + CUtensorMap tensor_map_grad_gate_up; + const cutlass::float_e4m3_t* acts_ptr; + const uint32_t* acts_sf_ptr; + const int8_t* w2_weights; + const float* w2_scales; + cutlass::bfloat16_t* w2_dequant_scratch; + const int8_t* w13_weights; + const float* w13_scales; + cutlass::bfloat16_t* w13_dequant_scratch; + const cutlass::bfloat16_t* gate_up_output; + cutlass::bfloat16_t* grad_ye_output; + cutlass::bfloat16_t* grad_y_unweighted_output; + cutlass::bfloat16_t* route_weights; + float* route_weights_fp32; + cutlass::bfloat16_t* grad_h_output; + cutlass::bfloat16_t* grad_gate_up_output; + cutlass::bfloat16_t* h_act_output; + cutlass::bfloat16_t* h_weighted_output; + cutlass::bfloat16_t* x_pool_output; + cutlass::bfloat16_t* grad_x_pool_output; + const cutlass::bfloat16_t* down_unweighted_output; + float* grad_route_output; + uint32_t* grid_sync_counter; + uint32_t launch_epoch; + float activation_limit; + MegaMoESideLoraBackwardParams side_lora{}; + bool compute_w13_dgrad; + bool direct_remote_grad_x; + bool write_grad_x_pool; + bool clear_wgrad_padding; + bool compute_route_grad = false; + bool trace_kernel = false; + bool vectorized_grad_x_store = false; + bool wide_grad_x_store = false; + bool gate_up_prepared = false; + uint64_t* kernel_trace = nullptr; + bool inputs_prepared = false; + bool dispatch_inputs_prepared = false; + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast( + &sm100_bf16_mega_moe_side_lora_backward_wave_impl< + {}, {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {} + >); +}}; +)", + args.hidden, args.intermediate_hidden, + args.num_experts, + args.block_m, args.block_n, args.block_k, + args.sf_block_m, args.sf_block_n, + args.num_stages, + args.num_sms, + args.num_ranks, + args.compute_w13_dgrad ? "true" : "false", + args.bf16_mode ? "true" : "false", + args.activation == "geglu" + ? "ActivationType::GeGLU" + : "ActivationType::SwiGLU", + args.fast_math ? "true" : "false", + get_side_lora_backward_route_weight_mode_name( + args.route_weight_mode), + get_side_lora_backward_combine_order_mode_name( + args.combine_order_mode), + args.inputs_prepared ? "true" : "false", + args.dispatch_inputs_prepared ? "true" : "false", + args.direct_remote_grad_x ? "true" : "false", + args.write_grad_x_pool ? "true" : "false", + args.clear_wgrad_padding ? "true" : "false", + args.compute_route_grad ? "true" : "false", + args.trace_kernel ? "true" : "false", + args.vectorized_grad_x_store ? "true" : "false", + args.wide_grad_x_store ? "true" : "false", + args.gate_up_prepared ? "true" : "false"); + } + + static void launch_impl( + const KernelHandle& kernel, + const LaunchConfigHandle& config, + Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel( + kernel, config, + args.expert_counts, + args.backward_sym_buffer, + args.backward_workspace, + args.backward_grad_y, + args.backward_x, + args.backward_topk_weights, + args.backward_grad_route, + args.token_src_metadata, + args.num_topk, + args.num_pool_rows, + args.num_acts_rows, + args.acts_sf_stride, + args.tensor_map_acts, + args.tensor_map_acts_sf, + args.tensor_map_weights, + args.tensor_map_weights_sf, + args.tensor_map_output, + args.tensor_map_grad_ye, + args.tensor_map_w2_dequant, + args.tensor_map_w2_weights, + args.tensor_map_w2_scales, + args.tensor_map_w13_dequant, + args.tensor_map_w13_weights, + args.tensor_map_w13_scales, + args.tensor_map_grad_gate_up, + args.acts_ptr, + args.acts_sf_ptr, + args.w2_weights, + args.w2_scales, + args.w2_dequant_scratch, + args.w13_weights, + args.w13_scales, + args.w13_dequant_scratch, + args.gate_up_output, + args.grad_ye_output, + args.grad_y_unweighted_output, + args.route_weights, + args.route_weights_fp32, + args.grad_h_output, + args.grad_gate_up_output, + args.h_act_output, + args.h_weighted_output, + args.x_pool_output, + args.grad_x_pool_output, + args.down_unweighted_output, + args.grad_route_output, + args.grid_sync_counter, + args.launch_epoch, + args.activation_limit, + args.side_lora, + args.kernel_trace)); + } +}; + +class SM100BF16MegaMoESideLoraGradXRuntime final + : public LaunchRuntime { +public: + struct Args { + int hidden; + int num_experts; + int block_m; + int num_ranks; + int num_sms; + bool write_grad_x_pool; + bool direct_remote_grad_x; + const int* expert_counts; + cutlass::bfloat16_t* grad_x_pool; + const layout::TokenSrcMetadata* token_src_metadata; + cutlass::bfloat16_t* combine_buffer; + layout::SymBuffer<> sym_buffer; + layout::Workspace workspace; + uint32_t num_pool_rows; + uint32_t num_topk; + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include +using namespace deep_gemm; +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast( + &sm100_bf16_mega_moe_side_lora_grad_x_impl< + {}, {}, {}, {}, {}, {}, {}>); +}} +)", args.hidden, args.num_experts, args.block_m, args.num_ranks, + args.num_sms, + args.write_grad_x_pool ? "true" : "false", + args.direct_remote_grad_x ? "true" : "false"); + } + + static void launch_impl( + const KernelHandle& kernel, + const LaunchConfigHandle& config, + Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel( + kernel, config, args.expert_counts, args.grad_x_pool, + args.token_src_metadata, args.combine_buffer, + args.sym_buffer, args.workspace, args.num_pool_rows, + args.num_topk)); + } +}; + +class SM100BF16MegaMoESideLoraScaleGradsRuntime final + : public LaunchRuntime { +public: + struct Args { + int hidden; + int intermediate_hidden; + int num_experts; + MegaMoESideLoraBackwardParams side_lora; + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include +using namespace deep_gemm; +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast( + &sm100_bf16_mega_moe_side_lora_scale_grads_impl<{}, {}, {}>); +}} +)", args.hidden, args.intermediate_hidden, args.num_experts); + } + + static void launch_impl( + const KernelHandle& kernel, + const LaunchConfigHandle& config, + Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel( + kernel, config, args.side_lora)); + } +}; + +class SM100BF16MegaMoESideLoraAxpy2Runtime final + : public LaunchRuntime { +public: + struct Args { + int num_sms; + cutlass::bfloat16_t* dst; + const cutlass::bfloat16_t* src1; + const cutlass::bfloat16_t* src3; + uint64_t num_elements; + float scale; + LaunchArgs launch_args; + }; + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include +using namespace deep_gemm; +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast( + &sm100_bf16_mega_moe_side_lora_axpy2_impl<{}>); +}} +)", args.num_sms); + } + static void launch_impl( + const KernelHandle& kernel, const LaunchConfigHandle& config, + Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel( + kernel, config, args.dst, args.src1, args.src3, + args.num_elements, args.scale)); + } +}; + +class SM100BF16MegaMoESideLoraClearPaddingRuntime final + : public LaunchRuntime { +public: + struct Args { + int intermediate_hidden; + int num_experts; + int block_m; + int num_sms; + const int* expert_counts; + cutlass::bfloat16_t* saved_h; + cutlass::bfloat16_t* q13; + cutlass::bfloat16_t* q2; + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include +using namespace deep_gemm; +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast( + &sm100_bf16_mega_moe_side_lora_clear_padding_impl<{}, {}, {}, {}>); +}} +)", args.intermediate_hidden, args.num_experts, args.block_m, + args.num_sms); + } + + static void launch_impl( + const KernelHandle& kernel, + const LaunchConfigHandle& config, + Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel( + kernel, config, args.expert_counts, args.saved_h, + args.q13, args.q2)); + } +}; + + +static void sm100_bf16_mega_moe_side_lora_backward( + const torch::Tensor& gate_up_output, + const torch::Tensor& grad_h_output, + const torch::Tensor& grad_gate_up_output, + const torch::Tensor& h_act_output, + const torch::Tensor& h_weighted_output, + const torch::Tensor& x_pool_output, + const torch::Tensor& grad_x_pool_output, + const torch::Tensor& grad_route_output, + const torch::Tensor& grad_ye, + const torch::Tensor& grad_y_unweighted_output, + const torch::Tensor& route_weights, + const torch::Tensor& w2_weights, + const torch::Tensor& w13_weights, + const torch::Tensor& expert_counts, + const torch::Tensor& grid_sync_counter, + const float& activation_limit, + const std::string& activation, + const bool& fast_math, + const std::string& route_weight_mode, + const std::string& combine_order_mode, + const torch::Tensor& down_unweighted_output, + const int& block_m, + const bool& direct_remote_grad_x, + const bool& write_grad_x_pool, + const bool& clear_wgrad_padding, + const torch::Tensor& backward_grad_y, + const torch::Tensor& backward_x, + const torch::Tensor& backward_topk_weights, + const std::optional& backward_grad_route, + const torch::Tensor& token_src_metadata, + const std::vector& backward_sym_buffer_ptrs, + const int& backward_rank, + const int& num_max_tokens_per_rank, + const int& num_topk, + const std::string& memory_mode, + const torch::Tensor& side_lora_a1, + const torch::Tensor& side_lora_b1, + const torch::Tensor& side_lora_a3, + const torch::Tensor& side_lora_b3, + const torch::Tensor& side_lora_a2, + const torch::Tensor& side_lora_b2, + const torch::Tensor& side_lora_q13, + const torch::Tensor& side_lora_q2, + const torch::Tensor& side_lora_saved_h, + const torch::Tensor& side_lora_t13, + const torch::Tensor& side_lora_t2, + const torch::Tensor& grad_side_lora_a1, + const torch::Tensor& grad_side_lora_b1, + const torch::Tensor& grad_side_lora_a3, + const torch::Tensor& grad_side_lora_b3, + const torch::Tensor& grad_side_lora_a2, + const torch::Tensor& grad_side_lora_b2, + const torch::Tensor& expert_psum_rows, + const torch::Tensor& padded_expert_counts, + const float& side_lora_scale, + const std::optional& kernel_trace = + std::nullopt) { + constexpr int block_n = 128; + constexpr int block_k = 128; + constexpr int dgrad_block_k = 64; + constexpr int store_block_m = 16; + constexpr int smem_capacity = 232448; + constexpr int num_epilogue_stages = 2; + constexpr int num_tma_store_stages = 2; + + const auto [num_experts, intermediate_hidden_2, hidden] = + get_shape<3>(w13_weights); + const auto [num_experts_w2, hidden_w2, intermediate_hidden] = + get_shape<3>(w2_weights); + const int num_pool_rows = static_cast(grad_ye.size(0)); + const int num_ranks = + static_cast(backward_sym_buffer_ptrs.size()); + const int sf_block_m = align(block_m, 128); + const int sf_block_n = block_n; + const int load_block_m = block_m / 2; + const int load_block_n = block_n; + const int num_dispatch_warps = num_ranks > 1 ? 4 : 0; + + DG_HOST_ASSERT(device_runtime->get_arch_major() == 10); + DG_HOST_ASSERT(activation == "swiglu" || activation == "geglu"); + DG_HOST_ASSERT( + route_weight_mode == "pre_down" || + route_weight_mode == "post_down"); + // The BF16 side-LoRA specialization reuses the now-dead grad-h and + // activation planes for canonical gate/up derivatives after its W13 + // dgrad consumes the forward-format interleave. Current training uses + // PRE_DOWN, whose phase plan keeps those planes distinct. + DG_HOST_ASSERT(route_weight_mode == "pre_down"); + DG_HOST_ASSERT( + memory_mode == "legacy" || + memory_mode == "phase_ordered" || + memory_mode == "dispatch_prepared"); + DG_HOST_ASSERT(num_ranks >= 1); + DG_HOST_ASSERT( + backward_rank >= 0 && backward_rank < num_ranks); + DG_HOST_ASSERT(num_experts == num_experts_w2); + DG_HOST_ASSERT(hidden == hidden_w2); + DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); + DG_HOST_ASSERT(block_m % 16 == 0); + DG_HOST_ASSERT(hidden % 256 == 0); + DG_HOST_ASSERT(intermediate_hidden % 256 == 0); + DG_HOST_ASSERT(num_pool_rows > 0); + DG_HOST_ASSERT(num_max_tokens_per_rank > 0); + DG_HOST_ASSERT(num_topk > 0); + constexpr int side_lora_rank = 128; + + const auto check_bf16_contiguous = [](const torch::Tensor& tensor) { + DG_HOST_ASSERT(tensor.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(tensor.is_contiguous()); + }; + check_bf16_contiguous(gate_up_output); + check_bf16_contiguous(grad_h_output); + check_bf16_contiguous(grad_gate_up_output); + check_bf16_contiguous(h_act_output); + check_bf16_contiguous(h_weighted_output); + check_bf16_contiguous(x_pool_output); + check_bf16_contiguous(grad_x_pool_output); + check_bf16_contiguous(grad_ye); + check_bf16_contiguous(grad_y_unweighted_output); + check_bf16_contiguous(down_unweighted_output); + check_bf16_contiguous(w2_weights); + check_bf16_contiguous(w13_weights); + check_bf16_contiguous(backward_grad_y); + check_bf16_contiguous(backward_x); + for (const auto* tensor : { + &side_lora_a1, &side_lora_b1, + &side_lora_a3, &side_lora_b3, + &side_lora_a2, &side_lora_b2, + &side_lora_q13, &side_lora_q2, + &side_lora_saved_h, + &side_lora_t13, &side_lora_t2, + &grad_side_lora_a1, &grad_side_lora_b1, + &grad_side_lora_a3, &grad_side_lora_b3, + &grad_side_lora_a2, &grad_side_lora_b2}) + check_bf16_contiguous(*tensor); + DG_HOST_ASSERT(backward_topk_weights.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(backward_topk_weights.is_contiguous()); + if (backward_grad_route.has_value()) { + DG_HOST_ASSERT( + backward_grad_route->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(backward_grad_route->is_contiguous()); + } + DG_HOST_ASSERT(route_weights.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(route_weights.is_contiguous()); + DG_HOST_ASSERT(grad_route_output.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(grad_route_output.is_contiguous()); + DG_HOST_ASSERT(expert_counts.scalar_type() == torch::kInt); + DG_HOST_ASSERT(expert_counts.is_contiguous()); + DG_HOST_ASSERT(expert_counts.numel() == num_experts); + DG_HOST_ASSERT(grid_sync_counter.scalar_type() == torch::kInt); + DG_HOST_ASSERT(grid_sync_counter.is_contiguous()); + DG_HOST_ASSERT(token_src_metadata.scalar_type() == torch::kInt); + DG_HOST_ASSERT(token_src_metadata.is_contiguous()); + DG_HOST_ASSERT(expert_psum_rows.scalar_type() == torch::kInt); + DG_HOST_ASSERT(expert_psum_rows.is_contiguous()); + DG_HOST_ASSERT(padded_expert_counts.scalar_type() == torch::kInt); + DG_HOST_ASSERT(padded_expert_counts.is_contiguous()); + + DG_HOST_ASSERT( + gate_up_output.sizes() == + torch::IntArrayRef( + {num_pool_rows, intermediate_hidden_2})); + DG_HOST_ASSERT(grad_gate_up_output.sizes() == + gate_up_output.sizes()); + DG_HOST_ASSERT( + grad_h_output.sizes() == + torch::IntArrayRef( + {num_pool_rows, intermediate_hidden})); + DG_HOST_ASSERT(h_act_output.sizes() == grad_h_output.sizes()); + DG_HOST_ASSERT( + h_weighted_output.sizes() == grad_h_output.sizes()); + DG_HOST_ASSERT( + x_pool_output.sizes() == + torch::IntArrayRef({num_pool_rows, hidden})); + DG_HOST_ASSERT( + write_grad_x_pool + ? grad_x_pool_output.sizes() == x_pool_output.sizes() + : grad_x_pool_output.sizes() == + torch::IntArrayRef({0, hidden})); + DG_HOST_ASSERT(grad_ye.sizes() == x_pool_output.sizes()); + DG_HOST_ASSERT( + grad_y_unweighted_output.sizes() == + x_pool_output.sizes()); + DG_HOST_ASSERT( + down_unweighted_output.sizes() == + x_pool_output.sizes()); + DG_HOST_ASSERT(route_weights.numel() == num_pool_rows); + DG_HOST_ASSERT(grad_route_output.numel() == num_pool_rows); + DG_HOST_ASSERT( + backward_grad_y.size(0) >= num_max_tokens_per_rank && + backward_grad_y.size(1) == hidden); + DG_HOST_ASSERT( + backward_x.size(0) >= num_max_tokens_per_rank && + backward_x.size(1) == hidden); + DG_HOST_ASSERT( + backward_topk_weights.size(0) >= + num_max_tokens_per_rank && + backward_topk_weights.size(1) == num_topk); + if (backward_grad_route.has_value()) { + DG_HOST_ASSERT( + backward_grad_route->size(0) >= + num_max_tokens_per_rank && + backward_grad_route->size(1) == num_topk); + } + DG_HOST_ASSERT( + token_src_metadata.size(0) >= num_pool_rows && + token_src_metadata.size(1) == 3); + DG_HOST_ASSERT(write_grad_x_pool || direct_remote_grad_x); + DG_HOST_ASSERT(side_lora_a1.sizes() == torch::IntArrayRef( + {side_lora_rank, hidden})); + DG_HOST_ASSERT(side_lora_a3.sizes() == side_lora_a1.sizes()); + DG_HOST_ASSERT(side_lora_b1.sizes() == torch::IntArrayRef( + {num_experts, intermediate_hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_b3.sizes() == side_lora_b1.sizes()); + DG_HOST_ASSERT(side_lora_a2.sizes() == torch::IntArrayRef( + {num_experts, side_lora_rank, intermediate_hidden})); + DG_HOST_ASSERT(side_lora_b2.sizes() == torch::IntArrayRef( + {hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_q13.sizes() == torch::IntArrayRef( + {num_pool_rows, 2, side_lora_rank})); + DG_HOST_ASSERT(side_lora_q2.sizes() == torch::IntArrayRef( + {num_pool_rows, side_lora_rank})); + DG_HOST_ASSERT(side_lora_saved_h.sizes() == torch::IntArrayRef( + {num_pool_rows, intermediate_hidden})); + DG_HOST_ASSERT(side_lora_t13.sizes() == side_lora_q13.sizes()); + DG_HOST_ASSERT(side_lora_t2.sizes() == side_lora_q2.sizes()); + DG_HOST_ASSERT(grad_side_lora_a1.sizes() == torch::IntArrayRef( + {hidden, side_lora_rank})); + DG_HOST_ASSERT(grad_side_lora_a3.sizes() == + grad_side_lora_a1.sizes()); + DG_HOST_ASSERT(grad_side_lora_b1.sizes() == torch::IntArrayRef( + {num_experts, side_lora_rank, intermediate_hidden})); + DG_HOST_ASSERT(grad_side_lora_b3.sizes() == + grad_side_lora_b1.sizes()); + DG_HOST_ASSERT(grad_side_lora_a2.sizes() == torch::IntArrayRef( + {num_experts, intermediate_hidden, side_lora_rank})); + DG_HOST_ASSERT(grad_side_lora_b2.sizes() == torch::IntArrayRef( + {side_lora_rank, hidden})); + DG_HOST_ASSERT(expert_psum_rows.numel() == num_experts); + DG_HOST_ASSERT(padded_expert_counts.numel() == num_experts); + if (kernel_trace.has_value()) { + DG_HOST_ASSERT(kernel_trace->is_cuda()); + DG_HOST_ASSERT( + kernel_trace->scalar_type() == torch::kInt64); + DG_HOST_ASSERT(kernel_trace->is_contiguous()); + DG_HOST_ASSERT(kernel_trace->device() == grad_ye.device()); + } + + const auto exact_alias = []( + const torch::Tensor& lhs, + const torch::Tensor& rhs) { + return lhs.data_ptr() == rhs.data_ptr() && + lhs.numel() == rhs.numel(); + }; + const auto overlaps = []( + const torch::Tensor& lhs, + const torch::Tensor& rhs) { + if (lhs.numel() == 0 || rhs.numel() == 0) + return false; + const auto lhs_begin = reinterpret_cast( + lhs.data_ptr()); + const auto rhs_begin = reinterpret_cast( + rhs.data_ptr()); + const auto lhs_end = + lhs_begin + lhs.numel() * lhs.element_size(); + const auto rhs_end = + rhs_begin + rhs.numel() * rhs.element_size(); + return std::max(lhs_begin, rhs_begin) < + std::min(lhs_end, rhs_end); + }; + const std::array, 10> + alias_tensors = { + gate_up_output, grad_gate_up_output, + grad_h_output, h_act_output, h_weighted_output, + x_pool_output, grad_x_pool_output, grad_ye, + grad_y_unweighted_output, down_unweighted_output}; + const auto allowed_overlap = [&](const int lhs, const int rhs) { + const auto pair_is = [=]( + const int first, const int second) { + return (lhs == first && rhs == second) || + (lhs == second && rhs == first); + }; + if (route_weight_mode == "pre_down" && + pair_is(7, 8)) + return exact_alias( + grad_ye, grad_y_unweighted_output); + if (route_weight_mode == "pre_down" && + pair_is(7, 9)) + return exact_alias( + grad_ye, down_unweighted_output); + if (route_weight_mode == "pre_down" && + pair_is(8, 9)) + return exact_alias( + grad_y_unweighted_output, + down_unweighted_output); + if (route_weight_mode == "post_down" && + pair_is(3, 4)) + return exact_alias( + h_act_output, h_weighted_output); + if (memory_mode != "phase_ordered") + return false; + if (pair_is(0, 1)) + return exact_alias( + gate_up_output, grad_gate_up_output); + if (route_weight_mode == "post_down") { + if (pair_is(7, 9)) + return exact_alias( + grad_ye, down_unweighted_output); + if (pair_is(2, 8)) + return + grad_h_output.data_ptr() == + grad_y_unweighted_output.data_ptr() && + grad_h_output.numel() <= + grad_y_unweighted_output.numel(); + if (pair_is(2, 3) || pair_is(2, 4)) + return exact_alias( + grad_h_output, + pair_is(2, 3) + ? h_act_output + : h_weighted_output); + if (pair_is(3, 8) || pair_is(4, 8)) + return + grad_y_unweighted_output.data_ptr() == + (pair_is(3, 8) + ? h_act_output.data_ptr() + : h_weighted_output.data_ptr()); + } else if (pair_is(3, 4)) { + return exact_alias( + h_act_output, h_weighted_output); + } + return false; + }; + for (int lhs = 0; lhs < alias_tensors.size(); ++lhs) { + for (int rhs = lhs + 1; + rhs < alias_tensors.size(); ++rhs) { + DG_HOST_ASSERT( + !overlaps( + alias_tensors[lhs].get(), + alias_tensors[rhs].get()) || + allowed_overlap(lhs, rhs)); + } + } + if (memory_mode == "phase_ordered") { + DG_HOST_ASSERT(exact_alias( + gate_up_output, grad_gate_up_output)); + if (route_weight_mode == "post_down") { + DG_HOST_ASSERT(exact_alias( + grad_ye, down_unweighted_output)); + DG_HOST_ASSERT( + grad_h_output.data_ptr() == + grad_y_unweighted_output.data_ptr() && + grad_h_output.numel() <= + grad_y_unweighted_output.numel()); + DG_HOST_ASSERT(exact_alias( + grad_h_output, h_act_output)); + DG_HOST_ASSERT(exact_alias( + grad_h_output, h_weighted_output)); + } else { + DG_HOST_ASSERT(exact_alias( + grad_ye, grad_y_unweighted_output)); + DG_HOST_ASSERT(exact_alias( + h_act_output, h_weighted_output)); + } + } + + const int num_w2_states = + num_experts * (hidden / dgrad_block_k) * + (intermediate_hidden / block_n); + const int num_w13_states = + num_experts * + (intermediate_hidden_2 / dgrad_block_k) * + (hidden / block_n); + DG_HOST_ASSERT( + grid_sync_counter.numel() >= + num_w2_states + num_w13_states + 2); + + const int smem_cd = + store_block_m * block_n * + static_cast(sizeof(cutlass::bfloat16_t)) * + num_tma_store_stages; + const int smem_per_stage = + load_block_m * block_k + + load_block_n * block_k + + sf_block_m * static_cast(sizeof(uint32_t)) + + sf_block_n * static_cast(sizeof(uint32_t)) + + 2 * static_cast(sizeof(uint64_t)); + const int smem_fixed = + num_dispatch_warps * hidden * + static_cast(sizeof(cutlass::bfloat16_t)) + + smem_cd + + 2 * num_epilogue_stages * + static_cast(sizeof(uint64_t)) + + num_dispatch_warps * + static_cast(sizeof(uint64_t)) + + static_cast(sizeof(uint32_t)); + const int num_stages = + std::min( + 32, + (smem_capacity - smem_fixed) / smem_per_stage); + const int smem_size = + align( + smem_fixed + num_stages * smem_per_stage, + 1024); + DG_HOST_ASSERT(num_stages >= 2); + + // Recompute descriptors are compile-time dead in BF16 mode. Keep valid + // descriptors in every slot so descriptor prefetch remains well-defined. + const auto tensor_map_acts = make_tma_2d_desc( + x_pool_output, hidden, num_pool_rows, + block_k, load_block_m, + static_cast(x_pool_output.stride(-2)), 128); + const auto tensor_map_w13 = make_tma_2d_desc( + w13_weights, hidden, + num_experts * intermediate_hidden_2, + block_k, load_block_n, + static_cast(w13_weights.stride(-2)), 128); + const auto tensor_map_w13_dgrad = make_tma_2d_desc( + w13_weights, hidden, + num_experts * intermediate_hidden_2, + load_block_n, dgrad_block_k, + static_cast(w13_weights.stride(-2)), 128); + const auto tensor_map_gate_up = make_tma_2d_desc( + gate_up_output, intermediate_hidden_2, num_pool_rows, + block_n, store_block_m, + static_cast(gate_up_output.stride(-2)), 128); + const auto tensor_map_grad_ye = make_tma_2d_desc( + grad_ye, hidden, num_pool_rows, + dgrad_block_k, load_block_m, + static_cast(grad_ye.stride(-2)), 128); + const auto tensor_map_w2 = make_tma_2d_desc( + w2_weights, intermediate_hidden, + num_experts * hidden, + load_block_n, dgrad_block_k, + intermediate_hidden, 128); + const auto tensor_map_grad_gate_up = make_tma_2d_desc( + grad_gate_up_output, intermediate_hidden_2, + num_pool_rows, + dgrad_block_k, load_block_m, + static_cast( + grad_gate_up_output.stride(-2)), 128); + + const int num_sms = device_runtime->get_num_sms(); + DG_HOST_ASSERT(num_sms % 2 == 0); + constexpr int num_trace_sites = 22; + constexpr int num_trace_values = 5; + if (kernel_trace.has_value()) { + DG_HOST_ASSERT( + kernel_trace->dim() == 3 && + kernel_trace->size(0) == num_trace_sites && + kernel_trace->size(1) == num_sms && + kernel_trace->size(2) == num_trace_values); + } + static std::atomic next_launch_epoch{1}; + uint32_t launch_epoch = + next_launch_epoch.fetch_add( + 1, std::memory_order_relaxed); + if (launch_epoch == 0) { + launch_epoch = + next_launch_epoch.fetch_add( + 1, std::memory_order_relaxed); + } + + const auto backward_sym_buffer = + layout::SymBuffer<>( + backward_sym_buffer_ptrs, backward_rank); + const auto backward_workspace = layout::Workspace( + reinterpret_cast( + backward_sym_buffer_ptrs[backward_rank]), + num_ranks, num_experts * num_ranks, + num_max_tokens_per_rank, num_topk, + layout::kMinCandidateBlockM); + + // The only pre-wave result is rank-width t2. Its expansion is consumed + // directly by the W2 dgrad epilogue below, so no [pool, I] side gradient + // is ever allocated or written. q2 was saved by the forward specialization. + const auto side_lora_b2_nt = side_lora_b2.transpose(0, 1); + sm100_bf16_mega_moe_side_lora_shared_gemm( + grad_ye, side_lora_b2_nt, side_lora_t2, + num_pool_rows, side_lora_rank, hidden, "nk"); + // Expand t2 through expert-local A2 on tensor cores before entering the + // persistent base dgrad wave. Reuse its required h_weighted output plane + // before the wave recomputes it; the wave consumes and overwrites each value + // with the final combined gradient. This replaces 128 scalar FMAs in the + // hot W2 epilogue without allocating a side-width gradient buffer. + const auto side_lora_a2_nt = side_lora_a2.transpose(1, 2); + sm100_bf16_mega_moe_side_lora_rank_gemm( + side_lora_t2, side_lora_a2_nt, h_weighted_output, + expert_psum_rows, padded_expert_counts, + num_experts, num_pool_rows, intermediate_hidden, side_lora_rank, + block_m, cute::UMMA::Major::K, + get_major_type_ab(side_lora_a2_nt)); + + MegaMoESideLoraBackwardParams side_lora_params{ + .saved_x = reinterpret_cast( + x_pool_output.data_ptr()), + .saved_h = reinterpret_cast( + side_lora_saved_h.data_ptr()), + .a1 = reinterpret_cast( + side_lora_a1.data_ptr()), + .b1 = reinterpret_cast( + side_lora_b1.data_ptr()), + .a3 = reinterpret_cast( + side_lora_a3.data_ptr()), + .b3 = reinterpret_cast( + side_lora_b3.data_ptr()), + .a2 = reinterpret_cast( + side_lora_a2.data_ptr()), + .b2 = reinterpret_cast( + side_lora_b2.data_ptr()), + .q1 = reinterpret_cast( + side_lora_q13.data_ptr()), + .q3 = reinterpret_cast( + side_lora_q13.data_ptr()) + side_lora_rank, + .q2 = reinterpret_cast( + side_lora_q2.data_ptr()), + .t1 = reinterpret_cast( + side_lora_t13.data_ptr()), + .t3 = reinterpret_cast( + side_lora_t13.data_ptr()) + side_lora_rank, + .t2 = reinterpret_cast( + side_lora_t2.data_ptr()), + .grad_a1 = reinterpret_cast( + grad_side_lora_a1.data_ptr()), + .grad_b1 = reinterpret_cast( + grad_side_lora_b1.data_ptr()), + .grad_a3 = reinterpret_cast( + grad_side_lora_a3.data_ptr()), + .grad_b3 = reinterpret_cast( + grad_side_lora_b3.data_ptr()), + .grad_a2 = reinterpret_cast( + grad_side_lora_a2.data_ptr()), + .grad_b2 = reinterpret_cast( + grad_side_lora_b2.data_ptr()), + .scale = side_lora_scale, + }; + + const SM100BF16MegaMoESideLoraBackwardWaveRuntime::Args args = { + .hidden = hidden, + .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, + .num_pool_rows = num_pool_rows, + .num_acts_rows = 0, + .num_sf_pool_rows = 1, + .block_m = block_m, + .block_n = block_n, + .block_k = block_k, + .sf_block_m = sf_block_m, + .sf_block_n = sf_block_n, + .num_stages = num_stages, + .num_sms = num_sms, + .num_ranks = num_ranks, + .bf16_mode = true, + .activation = activation, + .fast_math = fast_math, + .route_weight_mode = route_weight_mode, + .combine_order_mode = combine_order_mode, + .expert_counts = expert_counts.data_ptr(), + .backward_sym_buffer = backward_sym_buffer, + .backward_workspace = backward_workspace, + .backward_grad_y = + reinterpret_cast( + backward_grad_y.data_ptr()), + .backward_x = + reinterpret_cast( + backward_x.data_ptr()), + .backward_topk_weights = + backward_topk_weights.data_ptr(), + .backward_grad_route = + backward_grad_route.has_value() + ? backward_grad_route->data_ptr() + : nullptr, + .token_src_metadata = + reinterpret_cast< + const layout::TokenSrcMetadata*>( + token_src_metadata.data_ptr()), + .num_topk = static_cast(num_topk), + .acts_sf_stride = 0, + .tensor_map_acts = tensor_map_acts, + .tensor_map_acts_sf = tensor_map_acts, + .tensor_map_weights = tensor_map_w13, + .tensor_map_weights_sf = tensor_map_w13, + .tensor_map_output = tensor_map_gate_up, + .tensor_map_grad_ye = tensor_map_grad_ye, + .tensor_map_w2_dequant = tensor_map_w2, + .tensor_map_w2_weights = tensor_map_w2, + .tensor_map_w2_scales = tensor_map_w2, + .tensor_map_w13_dequant = tensor_map_w13_dgrad, + .tensor_map_w13_weights = tensor_map_w13, + .tensor_map_w13_scales = tensor_map_w13, + .tensor_map_grad_gate_up = + tensor_map_grad_gate_up, + .acts_ptr = nullptr, + .acts_sf_ptr = nullptr, + .w2_weights = nullptr, + .w2_scales = nullptr, + .w2_dequant_scratch = + reinterpret_cast( + w2_weights.data_ptr()), + .w13_weights = nullptr, + .w13_scales = nullptr, + .w13_dequant_scratch = + reinterpret_cast( + w13_weights.data_ptr()), + .gate_up_output = + reinterpret_cast( + gate_up_output.data_ptr()), + .grad_ye_output = + reinterpret_cast( + grad_ye.data_ptr()), + .grad_y_unweighted_output = + reinterpret_cast( + grad_y_unweighted_output + .data_ptr()), + .route_weights = + nullptr, + .route_weights_fp32 = route_weights.data_ptr(), + .grad_h_output = + reinterpret_cast( + grad_h_output.data_ptr()), + .grad_gate_up_output = + reinterpret_cast( + grad_gate_up_output.data_ptr()), + .h_act_output = + reinterpret_cast( + h_act_output.data_ptr()), + .h_weighted_output = + reinterpret_cast( + h_weighted_output.data_ptr()), + .x_pool_output = + reinterpret_cast( + x_pool_output.data_ptr()), + .grad_x_pool_output = + reinterpret_cast( + grad_x_pool_output.data_ptr()), + .down_unweighted_output = + reinterpret_cast( + down_unweighted_output + .data_ptr()), + .grad_route_output = + grad_route_output.data_ptr(), + .grid_sync_counter = + reinterpret_cast( + grid_sync_counter.data_ptr()), + .launch_epoch = launch_epoch, + .activation_limit = activation_limit, + .side_lora = side_lora_params, + .compute_w13_dgrad = true, + // The side publisher sends the final base+side value once. Emitting + // the base value remotely here would only be overwritten later. + .direct_remote_grad_x = false, + .write_grad_x_pool = true, + .clear_wgrad_padding = clear_wgrad_padding, + .compute_route_grad = true, + .trace_kernel = kernel_trace.has_value(), + .vectorized_grad_x_store = get_env( + "DG_BF16_MEGA_MOE_VECTORIZED_GRAD_X_STORE", + 1) == 1, + .wide_grad_x_store = get_env( + "DG_BF16_MEGA_MOE_WIDE_GRAD_X_STORE", + 0) == 1, + .kernel_trace = + kernel_trace.has_value() + ? reinterpret_cast( + kernel_trace->data_ptr()) + : nullptr, + .inputs_prepared = + memory_mode == "phase_ordered" && + route_weight_mode == "post_down", + .dispatch_inputs_prepared = + memory_mode == "phase_ordered" || + memory_mode == "dispatch_prepared", + .launch_args = + LaunchArgs(num_sms, 1024, smem_size, 2), + }; + const auto code = + SM100BF16MegaMoESideLoraBackwardWaveRuntime::generate(args); + const auto runtime = compiler->build( + fmt::format( + "sm100_bf16_mega_moe_side_lora_backward_trace{}_vec{}_wide{}", + kernel_trace.has_value(), + args.vectorized_grad_x_store, + args.wide_grad_x_store), + code); + SM100BF16MegaMoESideLoraBackwardWaveRuntime::launch( + runtime, args); + + // The activation gradients now exist. Contract each one to rank 128, + // then let the dedicated publisher consume both contractions directly + // into the base grad-x destinations. There is no full-width side grad-x. + const auto grad_gate = h_act_output; + const auto grad_up = grad_h_output; + const auto t1 = side_lora_t13.select(1, 0); + const auto t3 = side_lora_t13.select(1, 1); + const auto b1_nt = side_lora_b1.transpose(1, 2); + const auto b3_nt = side_lora_b3.transpose(1, 2); + sm100_bf16_mega_moe_side_lora_rank_gemm( + grad_gate, b1_nt, t1, expert_psum_rows, + padded_expert_counts, num_experts, num_pool_rows, side_lora_rank, + intermediate_hidden, block_m, cute::UMMA::Major::K, + get_major_type_ab(b1_nt)); + sm100_bf16_mega_moe_side_lora_rank_gemm( + grad_up, b3_nt, t3, expert_psum_rows, + padded_expert_counts, num_experts, num_pool_rows, side_lora_rank, + intermediate_hidden, block_m, cute::UMMA::Major::K, + get_major_type_ab(b3_nt)); + + // Clear forward padding before the adapter wgrads. B2 is formed now so + // the dead grad-ye plane can hold the second L1 dgrad expansion. + const SM100BF16MegaMoESideLoraClearPaddingRuntime::Args clear_args{ + .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, + .block_m = block_m, + .num_sms = num_sms, + .expert_counts = expert_counts.data_ptr(), + .saved_h = reinterpret_cast( + side_lora_saved_h.data_ptr()), + .q13 = reinterpret_cast( + side_lora_q13.data_ptr()), + .q2 = reinterpret_cast( + side_lora_q2.data_ptr()), + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto clear_code = + SM100BF16MegaMoESideLoraClearPaddingRuntime::generate(clear_args); + const auto clear_runtime = compiler->build( + "sm100_bf16_mega_moe_side_lora_clear_padding", clear_code); + SM100BF16MegaMoESideLoraClearPaddingRuntime::launch( + clear_runtime, clear_args); + sm100_bf16_mega_moe_side_lora_shared_gemm( + side_lora_q2.transpose(0, 1), grad_ye.transpose(0, 1), + grad_side_lora_b2, side_lora_rank, hidden, num_pool_rows, "mn"); + + // Reuse the now-dead saved-down and grad-ye planes for the two tensor-core + // L1 dgrad expansions. One vectorized native add consumes both, so no + // additional hidden-width side buffer or framework tensor op is needed. + const auto side_grad_x_scratch1 = down_unweighted_output; + const auto side_grad_x_scratch3 = grad_ye; + const auto a1_nt = side_lora_a1.transpose(0, 1); + const auto a3_nt = side_lora_a3.transpose(0, 1); + sm100_bf16_mega_moe_side_lora_shared_gemm( + t1, a1_nt, side_grad_x_scratch1, + num_pool_rows, hidden, side_lora_rank, "nk"); + sm100_bf16_mega_moe_side_lora_shared_gemm( + t3, a3_nt, side_grad_x_scratch3, + num_pool_rows, hidden, side_lora_rank, "nk"); + const SM100BF16MegaMoESideLoraAxpy2Runtime::Args axpy2_args{ + .num_sms = num_sms, + .dst = reinterpret_cast( + grad_x_pool_output.data_ptr()), + .src1 = reinterpret_cast( + side_grad_x_scratch1.data_ptr()), + .src3 = reinterpret_cast( + side_grad_x_scratch3.data_ptr()), + .num_elements = static_cast(num_pool_rows) * hidden, + .scale = side_lora_scale, + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto axpy2_code = + SM100BF16MegaMoESideLoraAxpy2Runtime::generate(axpy2_args); + const auto axpy2_runtime = compiler->build( + "sm100_bf16_mega_moe_side_lora_axpy2_grad_x", axpy2_code); + SM100BF16MegaMoESideLoraAxpy2Runtime::launch( + axpy2_runtime, axpy2_args); + + const SM100BF16MegaMoESideLoraGradXRuntime::Args grad_x_args{ + .hidden = hidden, + .num_experts = num_experts, + .block_m = block_m, + .num_ranks = num_ranks, + .num_sms = num_sms, + .write_grad_x_pool = write_grad_x_pool, + .direct_remote_grad_x = direct_remote_grad_x, + .expert_counts = expert_counts.data_ptr(), + .grad_x_pool = reinterpret_cast( + grad_x_pool_output.data_ptr()), + .token_src_metadata = reinterpret_cast< + const layout::TokenSrcMetadata*>( + token_src_metadata.data_ptr()), + .combine_buffer = reinterpret_cast( + backward_grad_y.data_ptr()), + .sym_buffer = backward_sym_buffer, + .workspace = backward_workspace, + .num_pool_rows = static_cast(num_pool_rows), + .num_topk = static_cast(num_topk), + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto grad_x_code = + SM100BF16MegaMoESideLoraGradXRuntime::generate(grad_x_args); + const auto grad_x_runtime = compiler->build( + "sm100_bf16_mega_moe_side_lora_grad_x", grad_x_code); + SM100BF16MegaMoESideLoraGradXRuntime::launch( + grad_x_runtime, grad_x_args); + + // Six rank-128 adapter wgrads replace the three frozen full-width base + // wgrads. A1/A3/B2 are shared, so each is one reduction across the route + // pool rather than one output matrix per local expert. + const auto q1 = side_lora_q13.select(1, 0); + const auto q3 = side_lora_q13.select(1, 1); + sm100_bf16_mega_moe_side_lora_shared_gemm( + x_pool_output.transpose(0, 1), t1.transpose(0, 1), + grad_side_lora_a1, hidden, side_lora_rank, num_pool_rows, "mn"); + sm100_bf16_mega_moe_wgrad_1sm( + q1, grad_gate, grad_side_lora_b1, + padded_expert_counts, block_m, {}, true, + "sm100_bf16_mega_moe_side_lora_wgrad_1sm"); + sm100_bf16_mega_moe_side_lora_shared_gemm( + x_pool_output.transpose(0, 1), t3.transpose(0, 1), + grad_side_lora_a3, hidden, side_lora_rank, num_pool_rows, "mn"); + sm100_bf16_mega_moe_wgrad_1sm( + q3, grad_up, grad_side_lora_b3, + padded_expert_counts, block_m, {}, true, + "sm100_bf16_mega_moe_side_lora_wgrad_1sm"); + sm100_bf16_mega_moe_wgrad_1sm( + side_lora_saved_h, side_lora_t2, grad_side_lora_a2, + padded_expert_counts, block_m, {}, true, + "sm100_bf16_mega_moe_side_lora_wgrad_1sm"); + + const SM100BF16MegaMoESideLoraScaleGradsRuntime::Args scale_args{ + .hidden = hidden, + .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, + .side_lora = side_lora_params, + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto scale_code = + SM100BF16MegaMoESideLoraScaleGradsRuntime::generate(scale_args); + const auto scale_runtime = compiler->build( + "sm100_bf16_mega_moe_side_lora_scale_grads", scale_code); + SM100BF16MegaMoESideLoraScaleGradsRuntime::launch( + scale_runtime, scale_args); +} + +static void sm100_fp8_fp4_mega_moe_side_lora_backward( + const torch::Tensor& gate_up_output, + const torch::Tensor& grad_h_output, + const torch::Tensor& grad_gate_up_output, + const torch::Tensor& h_act_output, + const torch::Tensor& h_weighted_output, + const torch::Tensor& x_pool_output, + const torch::Tensor& grad_x_pool_output, + const torch::Tensor& acts, + const torch::Tensor& acts_sf, + const torch::Tensor& l1_weights, + const torch::Tensor& l1_weights_sf, + const torch::Tensor& grad_ye, + const torch::Tensor& route_weights, + const torch::Tensor& w2_weights, + const torch::Tensor& w2_scales, + const torch::Tensor& w2_dequant_scratch, + const torch::Tensor& w13_weights, + const torch::Tensor& w13_scales, + const torch::Tensor& w13_dequant_scratch, + const torch::Tensor& expert_counts, + const torch::Tensor& grid_sync_counter, + const float& activation_limit, + const bool& compute_w13_dgrad, + const bool& direct_remote_grad_x, + const bool& write_grad_x_pool, + const bool& clear_wgrad_padding, + const int& block_m, + const std::vector& backward_sym_buffer_ptrs, + const int& backward_rank, + const int& num_max_tokens_per_rank, + const int& num_topk, + const std::optional& backward_grad_y, + const std::optional& backward_topk_weights, + const std::optional& backward_grad_route, + const std::optional& token_src_metadata, + const std::string& route_weight_mode, + const std::optional& grad_y_unweighted_output, + const std::optional& down_unweighted_output, + const std::optional& grad_route_output, + const torch::Tensor& side_lora_a1, + const torch::Tensor& side_lora_b1, + const torch::Tensor& side_lora_a3, + const torch::Tensor& side_lora_b3, + const torch::Tensor& side_lora_a2, + const torch::Tensor& side_lora_b2, + const torch::Tensor& side_lora_q13, + const torch::Tensor& side_lora_q2, + const torch::Tensor& side_lora_saved_h, + const torch::Tensor& side_lora_t13, + const torch::Tensor& side_lora_t2, + const torch::Tensor& grad_side_lora_a1, + const torch::Tensor& grad_side_lora_b1, + const torch::Tensor& grad_side_lora_a3, + const torch::Tensor& grad_side_lora_b3, + const torch::Tensor& grad_side_lora_a2, + const torch::Tensor& grad_side_lora_b2, + const torch::Tensor& expert_psum_rows, + const torch::Tensor& padded_expert_counts, + const float& side_lora_scale) { + constexpr int block_n = 128; + constexpr int block_k = 128; + constexpr int dgrad_block_k = 64; + constexpr int store_block_m = 16; + constexpr int gran_k = 32; + constexpr int smem_capacity = 232448; + constexpr int num_epilogue_stages = 2; + constexpr int num_tma_store_stages = 2; + + const auto [num_experts, intermediate_hidden_2, hidden] = + check_grouped_ab_fp8_fp4( + l1_weights, cute::UMMA::Major::K, + device_runtime->get_arch_major()); + const int intermediate_hidden = intermediate_hidden_2 / 2; + const int num_pool_rows = static_cast(grad_ye.size(0)); + const int num_acts_rows = static_cast(acts.size(0)); + const int num_sf_pool_rows = static_cast(acts_sf.size(0)); + const int sf_block_m = align(block_m, 128); + const int sf_block_n = block_n; + const int load_block_m = block_m / 2; + const int load_block_n = block_n; + const int num_ranks = backward_sym_buffer_ptrs.empty() + ? 1 + : static_cast(backward_sym_buffer_ptrs.size()); + const int num_dispatch_warps = num_ranks > 1 ? 4 : 0; + + DG_HOST_ASSERT(device_runtime->get_arch_major() == 10); + DG_HOST_ASSERT(num_ranks >= 1); + DG_HOST_ASSERT( + route_weight_mode == "pre_down" || + route_weight_mode == "post_down"); + DG_HOST_ASSERT(num_ranks == 1 || + (backward_rank >= 0 && backward_rank < num_ranks)); + DG_HOST_ASSERT(block_m % 16 == 0); + DG_HOST_ASSERT(hidden % block_k == 0); + DG_HOST_ASSERT(hidden % dgrad_block_k == 0); + DG_HOST_ASSERT(intermediate_hidden_2 % block_n == 0); + DG_HOST_ASSERT(acts.dim() == 2 and acts.size(1) == hidden); + DG_HOST_ASSERT(num_acts_rows > 0 && num_acts_rows <= num_pool_rows); + DG_HOST_ASSERT(num_sf_pool_rows > 0); + DG_HOST_ASSERT(grad_ye.dim() == 2); + DG_HOST_ASSERT(grad_ye.size(0) == num_pool_rows); + DG_HOST_ASSERT(grad_ye.size(1) == hidden); + DG_HOST_ASSERT(grad_ye.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(grad_ye.is_contiguous()); + DG_HOST_ASSERT( + gate_up_output.sizes() == + torch::IntArrayRef({num_pool_rows, intermediate_hidden_2})); + DG_HOST_ASSERT(grad_gate_up_output.sizes() == + gate_up_output.sizes()); + DG_HOST_ASSERT(grad_h_output.sizes() == + torch::IntArrayRef({num_pool_rows, intermediate_hidden})); + DG_HOST_ASSERT(h_act_output.sizes() == grad_h_output.sizes()); + DG_HOST_ASSERT(h_weighted_output.sizes() == grad_h_output.sizes()); + DG_HOST_ASSERT( + x_pool_output.sizes() == + torch::IntArrayRef({num_pool_rows, hidden})); + DG_HOST_ASSERT(grad_x_pool_output.sizes() == + torch::IntArrayRef({num_pool_rows, hidden})); + DG_HOST_ASSERT(gate_up_output.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(grad_h_output.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(grad_gate_up_output.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(h_act_output.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT( + h_weighted_output.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(x_pool_output.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(grad_x_pool_output.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(gate_up_output.is_contiguous()); + DG_HOST_ASSERT(grad_h_output.is_contiguous()); + DG_HOST_ASSERT(grad_gate_up_output.is_contiguous()); + DG_HOST_ASSERT(h_act_output.is_contiguous()); + DG_HOST_ASSERT(h_weighted_output.is_contiguous()); + DG_HOST_ASSERT(x_pool_output.is_contiguous()); + DG_HOST_ASSERT(grad_x_pool_output.is_contiguous()); + DG_HOST_ASSERT( + route_weights.scalar_type() == torch::kFloat || + route_weights.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(route_weights.numel() == num_pool_rows); + DG_HOST_ASSERT(route_weights.is_contiguous()); + const auto check_bf16_hidden_pool = + [&](const torch::Tensor& tensor) { + DG_HOST_ASSERT( + tensor.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(tensor.is_contiguous()); + DG_HOST_ASSERT( + tensor.sizes() == + torch::IntArrayRef( + {num_pool_rows, hidden})); + }; + if (grad_y_unweighted_output.has_value()) + check_bf16_hidden_pool( + *grad_y_unweighted_output); + if (down_unweighted_output.has_value()) { + DG_HOST_ASSERT( + down_unweighted_output->scalar_type() == + torch::kBFloat16); + DG_HOST_ASSERT(down_unweighted_output->is_contiguous()); + DG_HOST_ASSERT(down_unweighted_output->dim() == 2); + DG_HOST_ASSERT( + down_unweighted_output->size(1) == hidden); + DG_HOST_ASSERT( + down_unweighted_output->size(0) > 0 && + down_unweighted_output->size(0) <= num_pool_rows); + } + if (grad_route_output.has_value()) { + DG_HOST_ASSERT( + grad_route_output->scalar_type() == + torch::kFloat); + DG_HOST_ASSERT(grad_route_output->is_contiguous()); + DG_HOST_ASSERT( + grad_route_output->numel() == + num_pool_rows); + } + if (route_weight_mode == "post_down") { + DG_HOST_ASSERT( + route_weights.scalar_type() == torch::kFloat); + DG_HOST_ASSERT( + grad_y_unweighted_output.has_value()); + DG_HOST_ASSERT( + down_unweighted_output.has_value()); + DG_HOST_ASSERT(grad_route_output.has_value()); + } + DG_HOST_ASSERT( + w2_weights.scalar_type() == + torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(w2_weights.dim() == 3); + DG_HOST_ASSERT(w2_weights.size(0) == num_experts); + DG_HOST_ASSERT(w2_weights.size(1) == hidden); + DG_HOST_ASSERT(w2_weights.size(2) == intermediate_hidden / 2); + DG_HOST_ASSERT(w2_weights.is_contiguous()); + DG_HOST_ASSERT(w2_scales.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(w2_scales.dim() == 3); + DG_HOST_ASSERT(w2_scales.size(0) == num_experts); + DG_HOST_ASSERT(w2_scales.size(1) == hidden); + DG_HOST_ASSERT(w2_scales.size(2) == intermediate_hidden / gran_k); + DG_HOST_ASSERT(w2_scales.is_contiguous()); + DG_HOST_ASSERT( + w2_dequant_scratch.sizes() == + torch::IntArrayRef( + {num_experts, hidden, intermediate_hidden})); + DG_HOST_ASSERT( + w2_dequant_scratch.scalar_type() == + torch::kBFloat16); + DG_HOST_ASSERT(w2_dequant_scratch.is_contiguous()); + DG_HOST_ASSERT( + w13_weights.scalar_type() == + torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(w13_weights.dim() == 3); + DG_HOST_ASSERT(w13_weights.size(0) == 2 * num_experts); + DG_HOST_ASSERT(w13_weights.size(1) == intermediate_hidden); + DG_HOST_ASSERT(w13_weights.size(2) == hidden / 2); + DG_HOST_ASSERT(w13_weights.is_contiguous()); + DG_HOST_ASSERT(w13_scales.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(w13_scales.dim() == 3); + DG_HOST_ASSERT(w13_scales.size(0) == 2 * num_experts); + DG_HOST_ASSERT(w13_scales.size(1) == intermediate_hidden); + DG_HOST_ASSERT(w13_scales.size(2) == hidden / gran_k); + DG_HOST_ASSERT(w13_scales.is_contiguous()); + DG_HOST_ASSERT( + w13_dequant_scratch.sizes() == + torch::IntArrayRef( + {num_experts, intermediate_hidden_2, hidden})); + DG_HOST_ASSERT( + w13_dequant_scratch.scalar_type() == + torch::kBFloat16); + DG_HOST_ASSERT(w13_dequant_scratch.is_contiguous()); + DG_HOST_ASSERT(expert_counts.scalar_type() == torch::kInt); + DG_HOST_ASSERT(expert_counts.numel() == num_experts); + DG_HOST_ASSERT(expert_counts.is_contiguous()); + DG_HOST_ASSERT(grid_sync_counter.scalar_type() == torch::kInt); + DG_HOST_ASSERT( + grid_sync_counter.numel() >= + num_experts * + ((hidden / dgrad_block_k) * + (intermediate_hidden / block_n) + + (intermediate_hidden_2 / dgrad_block_k) * + (hidden / block_n)) + + 2); + DG_HOST_ASSERT(grid_sync_counter.is_contiguous()); + if (compute_w13_dgrad) + DG_HOST_ASSERT(write_grad_x_pool || direct_remote_grad_x); + else + DG_HOST_ASSERT(!direct_remote_grad_x); + if (direct_remote_grad_x) { + DG_HOST_ASSERT(compute_w13_dgrad); + DG_HOST_ASSERT(num_ranks > 1); + } + if (!backward_sym_buffer_ptrs.empty()) { + DG_HOST_ASSERT(backward_grad_y.has_value()); + DG_HOST_ASSERT(backward_topk_weights.has_value()); + DG_HOST_ASSERT(token_src_metadata.has_value()); + DG_HOST_ASSERT(num_max_tokens_per_rank > 0); + DG_HOST_ASSERT(num_topk > 0); + DG_HOST_ASSERT( + backward_grad_y->scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(backward_grad_y->dim() == 2); + DG_HOST_ASSERT(backward_grad_y->size(0) >= num_max_tokens_per_rank); + DG_HOST_ASSERT(backward_grad_y->size(1) == hidden); + DG_HOST_ASSERT(backward_grad_y->is_contiguous()); + DG_HOST_ASSERT( + backward_topk_weights->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(backward_topk_weights->dim() == 2); + DG_HOST_ASSERT(backward_topk_weights->size(0) >= + num_max_tokens_per_rank); + DG_HOST_ASSERT(backward_topk_weights->size(1) == num_topk); + DG_HOST_ASSERT(backward_topk_weights->is_contiguous()); + if (grad_route_output.has_value()) { + DG_HOST_ASSERT(backward_grad_route.has_value()); + DG_HOST_ASSERT( + backward_grad_route->scalar_type() == torch::kFloat); + DG_HOST_ASSERT(backward_grad_route->dim() == 2); + DG_HOST_ASSERT(backward_grad_route->size(0) >= + num_max_tokens_per_rank); + DG_HOST_ASSERT(backward_grad_route->size(1) == num_topk); + DG_HOST_ASSERT(backward_grad_route->is_contiguous()); + } + DG_HOST_ASSERT( + token_src_metadata->scalar_type() == torch::kInt); + DG_HOST_ASSERT(token_src_metadata->dim() == 2); + // Metadata is an immutable active-layout snapshot. Bucket-only output + // rows are never visited because expert_counts guards every access. + DG_HOST_ASSERT(token_src_metadata->size(0) >= num_acts_rows); + DG_HOST_ASSERT(token_src_metadata->size(1) == 3); + DG_HOST_ASSERT(token_src_metadata->is_contiguous()); + } + + check_sf_layout( + l1_weights_sf, intermediate_hidden_2, hidden, + 1, gran_k, num_experts, true, false, torch::kInt); + + const int smem_cd = + store_block_m * block_n * + static_cast(sizeof(cutlass::bfloat16_t)) * + num_tma_store_stages; + const int smem_per_stage = + load_block_m * block_k + + load_block_n * block_k + + sf_block_m * static_cast(sizeof(uint32_t)) + + sf_block_n * static_cast(sizeof(uint32_t)) + + 2 * static_cast(sizeof(uint64_t)); + const int smem_fixed = + num_dispatch_warps * hidden * + static_cast(sizeof(cutlass::bfloat16_t)) + + smem_cd + + 2 * num_epilogue_stages * + static_cast(sizeof(uint64_t)) + + num_dispatch_warps * + static_cast(sizeof(uint64_t)) + + static_cast(sizeof(uint32_t)); + const int num_stages = + std::min(32, (smem_capacity - smem_fixed) / smem_per_stage); + const int smem_size = + align(smem_fixed + num_stages * smem_per_stage, 1024); + DG_HOST_ASSERT(num_stages >= 2); + + const auto tensor_map_acts = make_tma_2d_desc( + acts, hidden, num_acts_rows, + block_k, load_block_m, + static_cast(acts.stride(-2)), 128); + const auto tensor_map_acts_sf = make_tma_sf_desc( + cute::UMMA::Major::MN, acts_sf, + num_sf_pool_rows, hidden, + sf_block_m, gran_k, 1, 0); + const auto tensor_map_l1_weights = make_tma_2d_desc( + l1_weights, hidden, num_experts * intermediate_hidden_2, + block_k, load_block_n, + static_cast(l1_weights.stride(-2)), 128); + const auto tensor_map_l1_weights_sf = make_tma_sf_desc( + cute::UMMA::Major::MN, l1_weights_sf, + intermediate_hidden_2, hidden, + block_n, gran_k, num_experts, 0); + const auto tensor_map_gate_up = make_tma_2d_desc( + gate_up_output, intermediate_hidden_2, num_pool_rows, + block_n, store_block_m, + static_cast(gate_up_output.stride(-2)), 128); + const auto tensor_map_grad_ye = make_tma_2d_desc( + grad_ye, hidden, num_pool_rows, + dgrad_block_k, load_block_m, + static_cast(grad_ye.stride(-2)), 128); + const auto tensor_map_w2_dequant = make_tma_2d_desc( + w2_dequant_scratch, intermediate_hidden, + num_experts * hidden, + load_block_n, dgrad_block_k, + intermediate_hidden, 128); + const auto tensor_map_w2_weights = make_tma_2d_desc( + w2_weights, intermediate_hidden / 2, + num_experts * hidden, + load_block_n / 2, 256, + intermediate_hidden / 2, 0); + const auto tensor_map_w2_scales = make_tma_2d_desc( + w2_scales, intermediate_hidden / gran_k, + num_experts * hidden, + load_block_n / gran_k, 256, + intermediate_hidden / gran_k, 0); + const auto tensor_map_w13_dequant = make_tma_2d_desc( + w13_dequant_scratch, hidden, + num_experts * intermediate_hidden_2, + load_block_n, dgrad_block_k, + hidden, 128); + const auto tensor_map_w13_weights = make_tma_2d_desc( + w13_weights, hidden / 2, + num_experts * intermediate_hidden_2, + load_block_n / 2, 256, + hidden / 2, 0); + const auto tensor_map_w13_scales = make_tma_2d_desc( + w13_scales, hidden / gran_k, + num_experts * intermediate_hidden_2, + load_block_n / gran_k, 256, + hidden / gran_k, 0); + const auto tensor_map_grad_gate_up = make_tma_2d_desc( + grad_gate_up_output, intermediate_hidden_2, + num_pool_rows, + dgrad_block_k, load_block_m, + static_cast(grad_gate_up_output.stride(-2)), 128); + // Each launch gets a unique readiness epoch; no host memset is required. + const int num_sms = device_runtime->get_num_sms(); + DG_HOST_ASSERT(num_sms % 2 == 0); + static std::atomic next_launch_epoch{1}; + uint32_t launch_epoch = + next_launch_epoch.fetch_add(1, std::memory_order_relaxed); + if (launch_epoch == 0) + launch_epoch = + next_launch_epoch.fetch_add(1, std::memory_order_relaxed); + layout::SymBuffer<> backward_sym_buffer{}; + void* backward_workspace_base = nullptr; + if (!backward_sym_buffer_ptrs.empty()) { + backward_sym_buffer = + layout::SymBuffer<>( + backward_sym_buffer_ptrs, backward_rank); + backward_workspace_base = + reinterpret_cast( + backward_sym_buffer_ptrs[backward_rank]); + } + const auto backward_workspace = layout::Workspace( + backward_workspace_base, num_ranks, + num_experts * num_ranks, + std::max(num_max_tokens_per_rank, 1), + std::max(num_topk, 1), + layout::kMinCandidateBlockM); + constexpr int side_lora_rank = 128; + DG_HOST_ASSERT(compute_w13_dgrad); + const auto check_side_bf16 = [](const torch::Tensor& tensor) { + DG_HOST_ASSERT(tensor.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(tensor.is_contiguous()); + }; + for (const auto* tensor : { + &side_lora_a1, &side_lora_b1, &side_lora_a3, + &side_lora_b3, &side_lora_a2, &side_lora_b2, + &side_lora_q13, &side_lora_q2, &side_lora_saved_h, + &side_lora_t13, &side_lora_t2, + &grad_side_lora_a1, &grad_side_lora_b1, + &grad_side_lora_a3, &grad_side_lora_b3, + &grad_side_lora_a2, &grad_side_lora_b2}) + check_side_bf16(*tensor); + DG_HOST_ASSERT(side_lora_a1.sizes() == torch::IntArrayRef( + {side_lora_rank, hidden})); + DG_HOST_ASSERT(side_lora_a3.sizes() == side_lora_a1.sizes()); + DG_HOST_ASSERT(side_lora_b1.sizes() == torch::IntArrayRef( + {num_experts, intermediate_hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_b3.sizes() == side_lora_b1.sizes()); + DG_HOST_ASSERT(side_lora_a2.sizes() == torch::IntArrayRef( + {num_experts, side_lora_rank, intermediate_hidden})); + DG_HOST_ASSERT(side_lora_b2.sizes() == torch::IntArrayRef( + {hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_q13.sizes() == torch::IntArrayRef( + {num_pool_rows, 2, side_lora_rank})); + DG_HOST_ASSERT(side_lora_q2.sizes() == torch::IntArrayRef( + {num_pool_rows, side_lora_rank})); + DG_HOST_ASSERT(side_lora_saved_h.sizes() == torch::IntArrayRef( + {num_pool_rows, intermediate_hidden})); + DG_HOST_ASSERT(side_lora_t13.sizes() == side_lora_q13.sizes()); + DG_HOST_ASSERT(side_lora_t2.sizes() == side_lora_q2.sizes()); + DG_HOST_ASSERT(expert_psum_rows.scalar_type() == torch::kInt && + expert_psum_rows.is_contiguous() && + expert_psum_rows.numel() == num_experts); + DG_HOST_ASSERT(padded_expert_counts.scalar_type() == torch::kInt && + padded_expert_counts.is_contiguous() && + padded_expert_counts.numel() == num_experts); + + DG_HOST_ASSERT(grad_side_lora_a1.sizes() == torch::IntArrayRef( + {hidden, side_lora_rank})); + DG_HOST_ASSERT(grad_side_lora_a3.sizes() == grad_side_lora_a1.sizes()); + DG_HOST_ASSERT(grad_side_lora_b2.sizes() == torch::IntArrayRef( + {side_lora_rank, hidden})); + + const auto side_lora_b2_nt = side_lora_b2.transpose(0, 1); + sm100_bf16_mega_moe_side_lora_shared_gemm( + grad_ye, side_lora_b2_nt, side_lora_t2, + num_pool_rows, side_lora_rank, hidden, "nk"); + const auto side_lora_a2_nt = side_lora_a2.transpose(1, 2); + sm100_bf16_mega_moe_side_lora_rank_gemm( + side_lora_t2, side_lora_a2_nt, h_weighted_output, + expert_psum_rows, padded_expert_counts, + num_experts, num_pool_rows, intermediate_hidden, side_lora_rank, + block_m, cute::UMMA::Major::K, + get_major_type_ab(side_lora_a2_nt)); + + MegaMoESideLoraBackwardParams side_lora_params{ + .saved_x = reinterpret_cast( + x_pool_output.data_ptr()), + .saved_h = reinterpret_cast( + side_lora_saved_h.data_ptr()), + .a1 = reinterpret_cast( + side_lora_a1.data_ptr()), + .b1 = reinterpret_cast( + side_lora_b1.data_ptr()), + .a3 = reinterpret_cast( + side_lora_a3.data_ptr()), + .b3 = reinterpret_cast( + side_lora_b3.data_ptr()), + .a2 = reinterpret_cast( + side_lora_a2.data_ptr()), + .b2 = reinterpret_cast( + side_lora_b2.data_ptr()), + .q1 = reinterpret_cast( + side_lora_q13.data_ptr()), + .q3 = reinterpret_cast( + side_lora_q13.data_ptr()) + side_lora_rank, + .q2 = reinterpret_cast( + side_lora_q2.data_ptr()), + .t1 = reinterpret_cast( + side_lora_t13.data_ptr()), + .t3 = reinterpret_cast( + side_lora_t13.data_ptr()) + side_lora_rank, + .t2 = reinterpret_cast( + side_lora_t2.data_ptr()), + .grad_a1 = reinterpret_cast( + grad_side_lora_a1.data_ptr()), + .grad_b1 = reinterpret_cast( + grad_side_lora_b1.data_ptr()), + .grad_a3 = reinterpret_cast( + grad_side_lora_a3.data_ptr()), + .grad_b3 = reinterpret_cast( + grad_side_lora_b3.data_ptr()), + .grad_a2 = reinterpret_cast( + grad_side_lora_a2.data_ptr()), + .grad_b2 = reinterpret_cast( + grad_side_lora_b2.data_ptr()), + .scale = side_lora_scale, + }; + const SM100BF16MegaMoESideLoraBackwardWaveRuntime::Args args = { + .hidden = hidden, + .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, + .num_pool_rows = num_pool_rows, + .num_acts_rows = num_acts_rows, + .num_sf_pool_rows = num_sf_pool_rows, + .block_m = block_m, + .block_n = block_n, + .block_k = block_k, + .sf_block_m = sf_block_m, + .sf_block_n = sf_block_n, + .num_stages = num_stages, + .num_sms = num_sms, + .num_ranks = num_ranks, + .route_weight_mode = route_weight_mode, + .expert_counts = expert_counts.data_ptr(), + .backward_sym_buffer = backward_sym_buffer, + .backward_workspace = backward_workspace, + .backward_grad_y = num_ranks > 1 + ? reinterpret_cast( + backward_grad_y->data_ptr()) + : nullptr, + .backward_x = nullptr, + .backward_topk_weights = num_ranks > 1 + ? backward_topk_weights->data_ptr() + : nullptr, + .backward_grad_route = + backward_grad_route.has_value() + ? backward_grad_route->data_ptr() + : nullptr, + .token_src_metadata = !backward_sym_buffer_ptrs.empty() + ? reinterpret_cast( + token_src_metadata->data_ptr()) + : nullptr, + .num_topk = static_cast(num_topk), + .acts_sf_stride = + static_cast(acts_sf.stride(1)), + .tensor_map_acts = tensor_map_acts, + .tensor_map_acts_sf = tensor_map_acts_sf, + .tensor_map_weights = tensor_map_l1_weights, + .tensor_map_weights_sf = tensor_map_l1_weights_sf, + .tensor_map_output = tensor_map_gate_up, + .tensor_map_grad_ye = tensor_map_grad_ye, + .tensor_map_w2_dequant = tensor_map_w2_dequant, + .tensor_map_w2_weights = tensor_map_w2_weights, + .tensor_map_w2_scales = tensor_map_w2_scales, + .tensor_map_w13_dequant = tensor_map_w13_dequant, + .tensor_map_w13_weights = tensor_map_w13_weights, + .tensor_map_w13_scales = tensor_map_w13_scales, + .tensor_map_grad_gate_up = tensor_map_grad_gate_up, + .acts_ptr = + reinterpret_cast( + acts.data_ptr()), + .acts_sf_ptr = + reinterpret_cast( + acts_sf.data_ptr()), + .w2_weights = + reinterpret_cast( + w2_weights.data_ptr()), + .w2_scales = w2_scales.data_ptr(), + .w2_dequant_scratch = + reinterpret_cast( + w2_dequant_scratch.data_ptr()), + .w13_weights = + reinterpret_cast( + w13_weights.data_ptr()), + .w13_scales = w13_scales.data_ptr(), + .w13_dequant_scratch = + reinterpret_cast( + w13_dequant_scratch.data_ptr()), + .gate_up_output = + reinterpret_cast( + gate_up_output.data_ptr()), + .grad_ye_output = + reinterpret_cast( + grad_ye.data_ptr()), + .grad_y_unweighted_output = + reinterpret_cast( + grad_y_unweighted_output + .value_or(grad_ye) + .data_ptr()), + .route_weights = + route_weights.scalar_type() == torch::kBFloat16 + ? reinterpret_cast( + route_weights.data_ptr()) + : nullptr, + .route_weights_fp32 = + route_weights.scalar_type() == torch::kFloat + ? route_weights.data_ptr() + : nullptr, + .grad_h_output = + reinterpret_cast( + grad_h_output.data_ptr()), + .grad_gate_up_output = + reinterpret_cast( + grad_gate_up_output.data_ptr()), + .h_act_output = + reinterpret_cast( + h_act_output.data_ptr()), + .h_weighted_output = + reinterpret_cast( + h_weighted_output.data_ptr()), + .x_pool_output = + reinterpret_cast( + x_pool_output.data_ptr()), + .grad_x_pool_output = + reinterpret_cast( + grad_x_pool_output.data_ptr()), + .down_unweighted_output = + down_unweighted_output.has_value() + ? reinterpret_cast< + const cutlass::bfloat16_t*>( + down_unweighted_output + ->data_ptr()) + : nullptr, + .grad_route_output = + grad_route_output.has_value() + ? grad_route_output->data_ptr() + : nullptr, + .grid_sync_counter = + reinterpret_cast( + grid_sync_counter.data_ptr()), + .launch_epoch = launch_epoch, + .activation_limit = activation_limit, + .side_lora = side_lora_params, + .compute_w13_dgrad = compute_w13_dgrad, + .direct_remote_grad_x = false, + .write_grad_x_pool = true, + .clear_wgrad_padding = clear_wgrad_padding, + .compute_route_grad = + grad_route_output.has_value(), + // The native side-LoRA forward already saved the full base+LoRA + // gate/up preactivation. Replaying the base-only MXFP4 W13 GEMM here + // would both waste work and erase the side delta before dSwiGLU. + .gate_up_prepared = true, + .inputs_prepared = route_weight_mode == "post_down", + .dispatch_inputs_prepared = true, + .launch_args = LaunchArgs( + num_sms, 1024, smem_size, 2), + }; + const auto code = + SM100BF16MegaMoESideLoraBackwardWaveRuntime::generate(args); + const auto runtime = compiler->build(fmt::format( + "sm100_fp8_fp4_mega_moe_backward_dgrad_swiglu_{}_r{}", + route_weight_mode, + grad_route_output.has_value()), code); + SM100BF16MegaMoESideLoraBackwardWaveRuntime::launch(runtime, args); + + const auto grad_gate = + grad_gate_up_output.slice(1, 0, intermediate_hidden); + const auto grad_up = grad_gate_up_output.slice( + 1, intermediate_hidden, intermediate_hidden_2); + const auto t1 = side_lora_t13.select(1, 0); + const auto t3 = side_lora_t13.select(1, 1); + const auto b1_nt = side_lora_b1.transpose(1, 2); + const auto b3_nt = side_lora_b3.transpose(1, 2); + sm100_bf16_mega_moe_side_lora_rank_gemm( + grad_gate, b1_nt, t1, expert_psum_rows, + padded_expert_counts, num_experts, num_pool_rows, + side_lora_rank, intermediate_hidden, block_m, + cute::UMMA::Major::K, + get_major_type_ab(b1_nt)); + sm100_bf16_mega_moe_side_lora_rank_gemm( + grad_up, b3_nt, t3, expert_psum_rows, + padded_expert_counts, num_experts, num_pool_rows, + side_lora_rank, intermediate_hidden, block_m, + cute::UMMA::Major::K, + get_major_type_ab(b3_nt)); + + const SM100BF16MegaMoESideLoraClearPaddingRuntime::Args clear_args{ + .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, + .block_m = block_m, + .num_sms = num_sms, + .expert_counts = expert_counts.data_ptr(), + .saved_h = reinterpret_cast( + side_lora_saved_h.data_ptr()), + .q13 = reinterpret_cast( + side_lora_q13.data_ptr()), + .q2 = reinterpret_cast( + side_lora_q2.data_ptr()), + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto clear_code = + SM100BF16MegaMoESideLoraClearPaddingRuntime::generate(clear_args); + const auto clear_runtime = compiler->build( + "sm100_fp8_fp4_mega_moe_side_lora_clear_padding", clear_code); + SM100BF16MegaMoESideLoraClearPaddingRuntime::launch( + clear_runtime, clear_args); + sm100_bf16_mega_moe_side_lora_shared_gemm( + side_lora_q2.transpose(0, 1), grad_ye.transpose(0, 1), + grad_side_lora_b2, side_lora_rank, hidden, num_pool_rows, "mn"); + + const auto side_grad_x_scratch1 = *down_unweighted_output; + const auto side_grad_x_scratch3 = grad_ye; + const auto a1_nt = side_lora_a1.transpose(0, 1); + const auto a3_nt = side_lora_a3.transpose(0, 1); + sm100_bf16_mega_moe_side_lora_shared_gemm( + t1, a1_nt, side_grad_x_scratch1, + num_pool_rows, hidden, side_lora_rank, "nk"); + sm100_bf16_mega_moe_side_lora_shared_gemm( + t3, a3_nt, side_grad_x_scratch3, + num_pool_rows, hidden, side_lora_rank, "nk"); + const SM100BF16MegaMoESideLoraAxpy2Runtime::Args axpy2_args{ + .num_sms = num_sms, + .dst = reinterpret_cast( + grad_x_pool_output.data_ptr()), + .src1 = reinterpret_cast( + side_grad_x_scratch1.data_ptr()), + .src3 = reinterpret_cast( + side_grad_x_scratch3.data_ptr()), + .num_elements = static_cast(num_pool_rows) * hidden, + .scale = side_lora_scale, + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto axpy2_code = + SM100BF16MegaMoESideLoraAxpy2Runtime::generate(axpy2_args); + const auto axpy2_runtime = compiler->build( + "sm100_fp8_fp4_mega_moe_side_lora_axpy2_grad_x", axpy2_code); + SM100BF16MegaMoESideLoraAxpy2Runtime::launch( + axpy2_runtime, axpy2_args); + + const SM100BF16MegaMoESideLoraGradXRuntime::Args grad_x_args{ + .hidden = hidden, + .num_experts = num_experts, + .block_m = block_m, + .num_ranks = num_ranks, + .num_sms = num_sms, + .write_grad_x_pool = write_grad_x_pool, + .direct_remote_grad_x = direct_remote_grad_x, + .expert_counts = expert_counts.data_ptr(), + .grad_x_pool = reinterpret_cast( + grad_x_pool_output.data_ptr()), + .token_src_metadata = !backward_sym_buffer_ptrs.empty() + ? reinterpret_cast( + token_src_metadata->data_ptr()) + : nullptr, + .combine_buffer = backward_grad_y.has_value() + ? reinterpret_cast( + backward_grad_y->data_ptr()) + : nullptr, + .sym_buffer = backward_sym_buffer, + .workspace = backward_workspace, + .num_pool_rows = static_cast(num_pool_rows), + .num_topk = static_cast(num_topk), + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto grad_x_code = + SM100BF16MegaMoESideLoraGradXRuntime::generate(grad_x_args); + const auto grad_x_runtime = compiler->build( + "sm100_fp8_fp4_mega_moe_side_lora_grad_x", grad_x_code); + SM100BF16MegaMoESideLoraGradXRuntime::launch( + grad_x_runtime, grad_x_args); + + const auto q1 = side_lora_q13.select(1, 0); + const auto q3 = side_lora_q13.select(1, 1); + sm100_bf16_mega_moe_side_lora_shared_gemm( + x_pool_output.transpose(0, 1), t1.transpose(0, 1), + grad_side_lora_a1, hidden, side_lora_rank, num_pool_rows, "mn"); + sm100_bf16_mega_moe_wgrad_1sm( + q1, grad_gate, grad_side_lora_b1, + padded_expert_counts, block_m, {}, true, + "sm100_fp8_fp4_mega_moe_side_lora_wgrad_1sm"); + sm100_bf16_mega_moe_side_lora_shared_gemm( + x_pool_output.transpose(0, 1), t3.transpose(0, 1), + grad_side_lora_a3, hidden, side_lora_rank, num_pool_rows, "mn"); + sm100_bf16_mega_moe_wgrad_1sm( + q3, grad_up, grad_side_lora_b3, + padded_expert_counts, block_m, {}, true, + "sm100_fp8_fp4_mega_moe_side_lora_wgrad_1sm"); + sm100_bf16_mega_moe_wgrad_1sm( + side_lora_saved_h, side_lora_t2, grad_side_lora_a2, + padded_expert_counts, block_m, {}, true, + "sm100_fp8_fp4_mega_moe_side_lora_wgrad_1sm"); + + const SM100BF16MegaMoESideLoraScaleGradsRuntime::Args scale_args{ + .hidden = hidden, + .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, + .side_lora = side_lora_params, + .launch_args = LaunchArgs(num_sms, 256, 0, 1), + }; + const auto scale_code = + SM100BF16MegaMoESideLoraScaleGradsRuntime::generate(scale_args); + const auto scale_runtime = compiler->build( + "sm100_fp8_fp4_mega_moe_side_lora_scale_grads", scale_code); + SM100BF16MegaMoESideLoraScaleGradsRuntime::launch( + scale_runtime, scale_args); +} + + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_forward.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_forward.hpp new file mode 100644 index 0000000000..c0ec212d60 --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_forward.hpp @@ -0,0 +1,579 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "runtime_utils.hpp" + +#include +#include + +#include "../heuristics/mega_moe.hpp" + +namespace deep_gemm { + +static std::string get_bf16_side_lora_activation_type_name( + const std::string& activation) { + if (activation == "swiglu") + return "ActivationType::SwiGLU"; + if (activation == "geglu") + return "ActivationType::GeGLU"; + DG_HOST_UNREACHABLE("Unsupported activation"); +} + +static std::string get_bf16_side_lora_route_weight_mode_name( + const std::string& route_weight_mode) { + if (route_weight_mode == "pre_down") + return "RouteWeightMode::PreDown"; + if (route_weight_mode == "post_down") + return "RouteWeightMode::PostDown"; + DG_HOST_UNREACHABLE("Unsupported route weight mode"); +} + +static std::string get_bf16_side_lora_combine_order_mode_name( + const std::string& combine_order_mode) { + if (combine_order_mode == "fixed_topk") + return "CombineOrderMode::FixedTopK"; + if (combine_order_mode == "deepep") + return "CombineOrderMode::DeepEP"; + if (combine_order_mode == "deepep_v1") + return "CombineOrderMode::DeepEPV1"; + DG_HOST_UNREACHABLE("Unsupported combine order mode"); +} + +class SM100BF16MegaMoESideLoraForwardRuntime final : public LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_max_tokens_per_rank; + int hidden, intermediate_hidden; + int num_experts, num_topk; + int num_ranks; + float activation_clamp; + bool fast_math; + std::string activation; + bool save_l1_preact; + bool save_stage_activations; + std::string route_weight_mode; + std::string combine_order_mode; + bool save_down_unweighted; + bool save_x; + int side_lora_rank; + MegaMoEConfig config; + + // Runtime arguments + void* y; + void* saved_l1_preact; + void* saved_h_unweighted; + void* saved_h_weighted; + void* saved_x; + void* saved_down_unweighted; + void* side_lora_a1; + void* side_lora_b1; + void* side_lora_a3; + void* side_lora_b3; + void* side_lora_a2; + void* side_lora_b2; + void* side_lora_l1_scratch; + void* side_lora_l2_scratch; + int* side_lora_ready; + float side_lora_scale; + int* cumulative_local_expert_recv_stats; + const int* precomputed_route_counts; + int* route_count_mismatch; + int num_tokens; + int num_saved_pool_tokens; + layout::SymBuffer<> sym_buffer_ptrs; + + // Tensormap + CUtensorMap tensor_map_l1_acts; + CUtensorMap tensor_map_l1_weights; + CUtensorMap tensor_map_l1_output; + CUtensorMap tensor_map_l2_acts; + CUtensorMap tensor_map_l2_weights; + CUtensorMap tensor_map_lora_a1; + CUtensorMap tensor_map_lora_a3; + CUtensorMap tensor_map_lora_b1; + CUtensorMap tensor_map_lora_b3; + CUtensorMap tensor_map_lora_a2; + CUtensorMap tensor_map_lora_b2; + CUtensorMap tensor_map_lora_l1_scratch; + CUtensorMap tensor_map_lora_l2_scratch; + CUtensorMap tensor_map_lora_l1_scratch_store; + CUtensorMap tensor_map_lora_l2_scratch_store; + CUtensorMap tensor_map_down_unweighted; + + // Launch configs + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_bf16_mega_moe_side_lora_forward_impl< + {}, + {}, {}, + {}, {}, + {}, + {}, {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {} + >); +}}; +)", args.num_max_tokens_per_rank, + args.hidden, args.intermediate_hidden, + args.num_experts, args.num_topk, + args.config.num_experts_per_wave, + args.config.block_m, args.config.block_n, args.config.block_k, + args.config.store_block_m, + args.config.num_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), + args.fast_math ? "true" : "false", + get_bf16_side_lora_activation_type_name(args.activation), + args.save_l1_preact ? "true" : "false", + args.save_stage_activations ? "true" : "false", + get_bf16_side_lora_route_weight_mode_name(args.route_weight_mode), + get_bf16_side_lora_combine_order_mode_name(args.combine_order_mode), + args.save_down_unweighted ? "true" : "false", + args.save_x ? "true" : "false", + args.side_lora_rank); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + // TODO: optimize `args` copy + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.y, + args.saved_l1_preact, + args.saved_h_unweighted, + args.saved_h_weighted, + args.saved_x, + args.saved_down_unweighted, + args.side_lora_a1, + args.side_lora_b1, + args.side_lora_a3, + args.side_lora_b3, + args.side_lora_a2, + args.side_lora_b2, + args.side_lora_l1_scratch, + args.side_lora_l2_scratch, + args.side_lora_ready, + args.side_lora_scale, + args.cumulative_local_expert_recv_stats, + args.precomputed_route_counts, + args.route_count_mismatch, + args.num_tokens, + args.num_saved_pool_tokens, + args.sym_buffer_ptrs, + args.tensor_map_l1_acts, + args.tensor_map_l1_weights, + args.tensor_map_l1_output, + args.tensor_map_l2_acts, + args.tensor_map_l2_weights, + args.tensor_map_lora_a1, + args.tensor_map_lora_a3, + args.tensor_map_lora_b1, + args.tensor_map_lora_b3, + args.tensor_map_lora_a2, + args.tensor_map_lora_b2, + args.tensor_map_lora_l1_scratch, + args.tensor_map_lora_l2_scratch, + args.tensor_map_lora_l1_scratch_store, + args.tensor_map_lora_l2_scratch_store, + args.tensor_map_down_unweighted + )); + } +}; + +static void sm100_bf16_mega_moe_side_lora_forward( + const torch::Tensor& y, + const std::optional& saved_l1_preact, + const torch::Tensor& l1_acts, const torch::Tensor& l2_acts, + const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, + 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_config_tokens, + const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const std::string& activation, + const float& activation_clamp, + const bool& fast_math, + const std::string& route_weight_mode, + const std::optional& saved_h_unweighted, + const std::optional& saved_h_weighted, + const std::optional& saved_down_unweighted, + const std::string& combine_order_mode, + const std::optional& precomputed_route_counts, + const std::optional& active_pool_rows, + const std::optional& route_count_mismatch, + const std::optional& saved_x, + const std::optional& side_lora_a1, + const std::optional& side_lora_b1, + const std::optional& side_lora_a3, + const std::optional& side_lora_b3, + const std::optional& side_lora_a2, + const std::optional& side_lora_b2, + const std::optional& side_lora_l1_scratch, + const std::optional& side_lora_l2_scratch, + const std::optional& side_lora_ready, + const float& side_lora_scale +) { + 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)); + + // Heuristics + const auto config = get_mega_moe_config( + num_ranks, num_experts, num_experts_per_rank, + num_max_tokens_per_rank, num_config_tokens, num_topk, + hidden, intermediate_hidden, + num_ring_tokens, 0, MmaKind::BF16); + const auto num_max_pool_tokens = + layout::get_num_max_pool_tokens( + num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank); + const auto num_saved_pool_tokens = + active_pool_rows.value_or(num_max_pool_tokens); + DG_HOST_ASSERT( + num_saved_pool_tokens > 0 && + num_saved_pool_tokens <= num_max_pool_tokens); + if (saved_l1_preact.has_value()) { + DG_HOST_ASSERT(saved_l1_preact->scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(saved_l1_preact->is_contiguous()); + DG_HOST_ASSERT( + saved_l1_preact->sizes() == + torch::IntArrayRef( + {num_saved_pool_tokens, 2 * intermediate_hidden})); + } + DG_HOST_ASSERT( + route_weight_mode == "pre_down" || + route_weight_mode == "post_down"); + DG_HOST_ASSERT( + combine_order_mode == "fixed_topk" || + combine_order_mode == "deepep" || + combine_order_mode == "deepep_v1"); + DG_HOST_ASSERT( + saved_h_unweighted.has_value() == + saved_h_weighted.has_value()); + if (saved_h_unweighted.has_value()) { + for (const auto* saved : + {&*saved_h_unweighted, &*saved_h_weighted}) { + DG_HOST_ASSERT( + saved->scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(saved->is_contiguous()); + DG_HOST_ASSERT( + saved->sizes() == torch::IntArrayRef( + {num_saved_pool_tokens, intermediate_hidden})); + } + } + if (saved_down_unweighted.has_value()) { + DG_HOST_ASSERT( + saved_down_unweighted->scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(saved_down_unweighted->is_contiguous()); + DG_HOST_ASSERT( + saved_down_unweighted->sizes() == + torch::IntArrayRef({num_saved_pool_tokens, hidden})); + } + if (saved_x.has_value()) { + DG_HOST_ASSERT(saved_x->scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(saved_x->is_contiguous()); + DG_HOST_ASSERT( + saved_x->sizes() == + torch::IntArrayRef({num_saved_pool_tokens, hidden})); + } + + const bool has_side_lora = side_lora_a1.has_value(); + DG_HOST_ASSERT(has_side_lora); + DG_HOST_ASSERT( + has_side_lora == side_lora_b1.has_value() && + has_side_lora == side_lora_a3.has_value() && + has_side_lora == side_lora_b3.has_value() && + has_side_lora == side_lora_a2.has_value() && + has_side_lora == side_lora_b2.has_value() && + has_side_lora == side_lora_l1_scratch.has_value() && + has_side_lora == side_lora_l2_scratch.has_value() && + has_side_lora == side_lora_ready.has_value()); + int side_lora_rank = 0; + if (has_side_lora) { + side_lora_rank = static_cast(side_lora_a1->size(0)); + DG_HOST_ASSERT(side_lora_rank == 128); + const auto check_bf16_contiguous = [&y]( + const torch::Tensor& tensor) { + DG_HOST_ASSERT(tensor.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(tensor.is_contiguous()); + DG_HOST_ASSERT(tensor.device() == y.device()); + }; + for (const auto* tensor : { + &*side_lora_a1, &*side_lora_b1, + &*side_lora_a3, &*side_lora_b3, + &*side_lora_a2, &*side_lora_b2, + &*side_lora_l1_scratch, + &*side_lora_l2_scratch}) + check_bf16_contiguous(*tensor); + DG_HOST_ASSERT(side_lora_a1->sizes() == torch::IntArrayRef( + {side_lora_rank, hidden})); + DG_HOST_ASSERT(side_lora_b1->sizes() == torch::IntArrayRef( + {num_experts_per_rank, intermediate_hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_a3->sizes() == side_lora_a1->sizes()); + DG_HOST_ASSERT(side_lora_b3->sizes() == side_lora_b1->sizes()); + DG_HOST_ASSERT(side_lora_a2->sizes() == torch::IntArrayRef( + {num_experts_per_rank, side_lora_rank, intermediate_hidden})); + DG_HOST_ASSERT(side_lora_b2->sizes() == torch::IntArrayRef( + {hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_l1_scratch->sizes() == torch::IntArrayRef( + {num_saved_pool_tokens, 2, side_lora_rank})); + DG_HOST_ASSERT(side_lora_l2_scratch->sizes() == torch::IntArrayRef( + {num_saved_pool_tokens, side_lora_rank})); + DG_HOST_ASSERT(side_lora_ready->scalar_type() == torch::kInt); + DG_HOST_ASSERT(side_lora_ready->is_contiguous()); + DG_HOST_ASSERT(side_lora_ready->device() == y.device()); + DG_HOST_ASSERT( + side_lora_ready->numel() >= + 4 * num_ring_tokens / config.block_m); + side_lora_ready->zero_(); + } + + // Make tensormap + const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, + hidden, config.num_ring_tokens, + config.block_k, config.load_block_m, + static_cast(l1_acts.stride(-2)), + config.swizzle_acts_mode); + 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, + static_cast(l1_weights.stride(-2)), + config.swizzle_weights_mode); + const auto tensor_map_l1_output = 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); + const auto tensor_map_l2_acts = 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 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, + static_cast(l2_weights.stride(-2)), + config.swizzle_weights_mode); + const auto tensor_map_lora_a1 = has_side_lora + ? make_tma_2d_desc(*side_lora_a1, hidden, + side_lora_rank, + config.block_k, config.load_block_n, + static_cast(side_lora_a1->stride(-2)), + config.swizzle_weights_mode) + : tensor_map_l1_weights; + const auto tensor_map_lora_a3 = has_side_lora + ? make_tma_2d_desc(*side_lora_a3, hidden, + side_lora_rank, + config.block_k, config.load_block_n, + static_cast(side_lora_a3->stride(-2)), + config.swizzle_weights_mode) + : tensor_map_l1_weights; + const auto tensor_map_lora_b1 = has_side_lora + ? make_tma_2d_desc(*side_lora_b1, side_lora_rank, + num_experts_per_rank * intermediate_hidden, + config.block_k, 8, + static_cast(side_lora_b1->stride(-2)), + config.swizzle_weights_mode) + : tensor_map_l1_weights; + const auto tensor_map_lora_b3 = has_side_lora + ? make_tma_2d_desc(*side_lora_b3, side_lora_rank, + num_experts_per_rank * intermediate_hidden, + config.block_k, 8, + static_cast(side_lora_b3->stride(-2)), + config.swizzle_weights_mode) + : tensor_map_l1_weights; + const auto tensor_map_lora_a2 = has_side_lora + ? make_tma_2d_desc(*side_lora_a2, intermediate_hidden, + num_experts_per_rank * side_lora_rank, + config.block_k, config.load_block_n, + static_cast(side_lora_a2->stride(-2)), + config.swizzle_weights_mode) + : tensor_map_l2_weights; + const auto tensor_map_lora_b2 = has_side_lora + ? make_tma_2d_desc(*side_lora_b2, side_lora_rank, + hidden, + config.block_k, config.load_block_n, + static_cast(side_lora_b2->stride(-2)), + config.swizzle_weights_mode) + : tensor_map_l2_weights; + const auto tensor_map_lora_l1_scratch = has_side_lora + ? make_tma_2d_desc(*side_lora_l1_scratch, + 2 * side_lora_rank, num_saved_pool_tokens, + config.block_k, config.load_block_m, + static_cast(side_lora_l1_scratch->stride(0)), + config.swizzle_acts_mode) + : tensor_map_l1_acts; + const auto tensor_map_lora_l2_scratch = has_side_lora + ? make_tma_2d_desc(*side_lora_l2_scratch, + side_lora_rank, num_saved_pool_tokens, + config.block_k, config.load_block_m, + static_cast(side_lora_l2_scratch->stride(0)), + config.swizzle_acts_mode) + : tensor_map_l2_acts; + const auto tensor_map_lora_l1_scratch_store = has_side_lora + ? make_tma_2d_desc(*side_lora_l1_scratch, + 2 * side_lora_rank, num_saved_pool_tokens, + config.block_k, config.store_block_m, + static_cast(side_lora_l1_scratch->stride(0)), + config.swizzle_acts_mode) + : tensor_map_l1_acts; + const auto tensor_map_lora_l2_scratch_store = has_side_lora + ? make_tma_2d_desc(*side_lora_l2_scratch, + side_lora_rank, num_saved_pool_tokens, + config.block_k, config.store_block_m, + static_cast(side_lora_l2_scratch->stride(0)), + config.swizzle_acts_mode) + : tensor_map_l2_acts; + const auto tensor_map_down_unweighted = + saved_down_unweighted.has_value() + ? make_tma_2d_desc( + *saved_down_unweighted, + hidden, saved_down_unweighted->size(0), + config.block_n, config.store_block_m, + static_cast(saved_down_unweighted->stride(-2)), + config.swizzle_acts_mode) + : tensor_map_l2_acts; + + // Stats can be optional + int* cumulative_local_expert_recv_stats_ptr = nullptr; + if (cumulative_local_expert_recv_stats.has_value()) + cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); + + // Launch + const auto physical_num_sms = device_runtime->get_num_sms(); + const auto num_sms = get_env( + "DG_BF16_MEGA_MOE_NUM_SMS", + physical_num_sms); + DG_HOST_ASSERT(num_sms > 0 && num_sms <= physical_num_sms); + const SM100BF16MegaMoESideLoraForwardRuntime::Args args = { + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .hidden = hidden, .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, .num_topk = num_topk, + .num_ranks = num_ranks, + .activation_clamp = activation_clamp, + .fast_math = fast_math, + .activation = activation, + .save_l1_preact = saved_l1_preact.has_value(), + .save_stage_activations = + saved_h_unweighted.has_value(), + .route_weight_mode = route_weight_mode, + .combine_order_mode = combine_order_mode, + .save_down_unweighted = + saved_down_unweighted.has_value(), + .save_x = saved_x.has_value(), + .side_lora_rank = side_lora_rank, + .config = config, + .y = y.data_ptr(), + .saved_l1_preact = saved_l1_preact.has_value() + ? saved_l1_preact->data_ptr() + : nullptr, + .saved_h_unweighted = + saved_h_unweighted.has_value() + ? saved_h_unweighted->data_ptr() + : nullptr, + .saved_h_weighted = + saved_h_weighted.has_value() + ? saved_h_weighted->data_ptr() + : nullptr, + .saved_x = saved_x.has_value() + ? saved_x->data_ptr() + : nullptr, + .saved_down_unweighted = + saved_down_unweighted.has_value() + ? saved_down_unweighted->data_ptr() + : nullptr, + .side_lora_a1 = has_side_lora + ? side_lora_a1->data_ptr() : nullptr, + .side_lora_b1 = has_side_lora + ? side_lora_b1->data_ptr() : nullptr, + .side_lora_a3 = has_side_lora + ? side_lora_a3->data_ptr() : nullptr, + .side_lora_b3 = has_side_lora + ? side_lora_b3->data_ptr() : nullptr, + .side_lora_a2 = has_side_lora + ? side_lora_a2->data_ptr() : nullptr, + .side_lora_b2 = has_side_lora + ? side_lora_b2->data_ptr() : nullptr, + .side_lora_l1_scratch = has_side_lora + ? side_lora_l1_scratch->data_ptr() : nullptr, + .side_lora_l2_scratch = has_side_lora + ? side_lora_l2_scratch->data_ptr() : nullptr, + .side_lora_ready = has_side_lora + ? side_lora_ready->data_ptr() : nullptr, + .side_lora_scale = side_lora_scale, + .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, + .precomputed_route_counts = + precomputed_route_counts.has_value() + ? precomputed_route_counts->data_ptr() + : nullptr, + .route_count_mismatch = + route_count_mismatch.has_value() + ? route_count_mismatch->data_ptr() + : nullptr, + .num_tokens = num_tokens, + .num_saved_pool_tokens = num_saved_pool_tokens, + .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + .tensor_map_l1_acts = tensor_map_l1_acts, + .tensor_map_l1_weights = tensor_map_l1_weights, + .tensor_map_l1_output = tensor_map_l1_output, + .tensor_map_l2_acts = tensor_map_l2_acts, + .tensor_map_l2_weights = tensor_map_l2_weights, + .tensor_map_lora_a1 = tensor_map_lora_a1, + .tensor_map_lora_a3 = tensor_map_lora_a3, + .tensor_map_lora_b1 = tensor_map_lora_b1, + .tensor_map_lora_b3 = tensor_map_lora_b3, + .tensor_map_lora_a2 = tensor_map_lora_a2, + .tensor_map_lora_b2 = tensor_map_lora_b2, + .tensor_map_lora_l1_scratch = + tensor_map_lora_l1_scratch, + .tensor_map_lora_l2_scratch = + tensor_map_lora_l2_scratch, + .tensor_map_lora_l1_scratch_store = + tensor_map_lora_l1_scratch_store, + .tensor_map_lora_l2_scratch_store = + tensor_map_lora_l2_scratch_store, + .tensor_map_down_unweighted = + tensor_map_down_unweighted, + .launch_args = LaunchArgs(num_sms, + config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, + config.smem_size, 2) + }; + + const auto code = SM100BF16MegaMoESideLoraForwardRuntime::generate(args); + const auto runtime = compiler->build( + "sm100_bf16_mega_moe_side_lora_forward", code); + SM100BF16MegaMoESideLoraForwardRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp index 29d4db6a04..868abf41eb 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp @@ -29,7 +29,9 @@ static void sm100_bf16_mega_moe_wgrad_1sm( const torch::Tensor& d, const torch::Tensor& padded_expert_counts, const int pool_block_m, - const MegaMoEBackwardCombineArgs& combine = {}) { + const MegaMoEBackwardCombineArgs& combine = {}, + const bool allow_row_strided_inputs = false, + const std::string& kernel_name = "sm100_bf16_mega_moe_wgrad_1sm") { const auto [num_groups, m, n] = get_shape<3>(d); const auto [pool_rows_a, m_] = get_shape<2>(a); const auto [pool_rows_b, n_] = get_shape<2>(b); @@ -42,7 +44,11 @@ static void sm100_bf16_mega_moe_wgrad_1sm( a.scalar_type() == torch::kBFloat16 and b.scalar_type() == torch::kBFloat16 and d.scalar_type() == torch::kBFloat16); - DG_HOST_ASSERT(a.is_contiguous() and b.is_contiguous() and d.is_contiguous()); + DG_HOST_ASSERT(d.is_contiguous()); + DG_HOST_ASSERT( + allow_row_strided_inputs + ? (a.stride(1) == 1 and b.stride(1) == 1) + : (a.is_contiguous() and b.is_contiguous())); DG_HOST_ASSERT( pool_block_m == 16 || pool_block_m == 32 || @@ -181,7 +187,7 @@ static void sm100_bf16_mega_moe_wgrad_1sm( }; const auto code = SM100BF16GemmRuntime::generate(args); const auto runtime = - compiler->build("sm100_bf16_mega_moe_wgrad_1sm", code); + compiler->build(kernel_name, code); SM100BF16GemmRuntime::launch(runtime, args); } diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.hpp new file mode 100644 index 0000000000..41c854758b --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.hpp @@ -0,0 +1,517 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "runtime_utils.hpp" + +#include +#include + +#include "../heuristics/mega_moe.hpp" + +namespace deep_gemm { + +// Map an activation name to its `deep_gemm::ActivationType` enumerator token +// (resolved inside the JIT-generated translation unit via `using namespace deep_gemm`). +static std::string get_mxfp4_side_lora_activation_type_name( + const std::string& activation) { + if (activation == "swiglu") + return "ActivationType::SwiGLU"; + if (activation == "geglu") + return "ActivationType::GeGLU"; + DG_HOST_UNREACHABLE("Unsupported activation"); +} + +static std::string get_mxfp4_side_lora_route_weight_mode_name( + const std::string& route_weight_mode) { + if (route_weight_mode == "pre_down") + return "RouteWeightMode::PreDown"; + if (route_weight_mode == "post_down") + return "RouteWeightMode::PostDown"; + DG_HOST_UNREACHABLE("Unsupported route weight mode"); +} + +class SM100FP8FP4MegaMoESideLoraForwardRuntime final : public LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_max_tokens_per_rank; + int hidden, intermediate_hidden; + int num_experts, num_topk; + int num_ranks; + float activation_clamp; + bool fast_math; + std::string activation; + bool save_l1_preact; + std::string route_weight_mode; + bool save_down_unweighted; + int side_lora_rank; + MegaMoEConfig config; + + // Runtime arguments + void* y; + void* saved_l1_preact; + void* saved_x; + void* saved_h; + void* saved_down_unweighted; + int* side_lora_ready; + float side_lora_scale; + int* cumulative_local_expert_recv_stats; + int num_tokens; + int num_saved_pool_tokens; + layout::SymBuffer<> sym_buffer_ptrs; + + // Tensormap + CUtensorMap tensor_map_l1_acts; + CUtensorMap tensor_map_l1_acts_sf; + CUtensorMap tensor_map_l1_weights; + CUtensorMap tensor_map_l1_weights_sf; + CUtensorMap tensor_map_l1_output; + CUtensorMap tensor_map_l2_acts; + CUtensorMap tensor_map_l2_acts_sf; + CUtensorMap tensor_map_l2_weights; + CUtensorMap tensor_map_l2_weights_sf; + CUtensorMap tensor_map_down_unweighted; + CUtensorMap tensor_map_saved_x; + CUtensorMap tensor_map_saved_h; + CUtensorMap tensor_map_lora_a1; + CUtensorMap tensor_map_lora_a3; + CUtensorMap tensor_map_lora_b1; + CUtensorMap tensor_map_lora_b3; + CUtensorMap tensor_map_lora_a2; + CUtensorMap tensor_map_lora_b2; + CUtensorMap tensor_map_lora_l1_scratch; + CUtensorMap tensor_map_lora_l2_scratch; + CUtensorMap tensor_map_lora_l1_scratch_store; + CUtensorMap tensor_map_lora_l2_scratch_store; + + // Launch configs + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_fp8_fp4_mega_moe_side_lora_forward_impl< + {}, + {}, {}, + {}, {}, + {}, + {}, {}, {}, + {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {} + >); +}}; +)", args.num_max_tokens_per_rank, + args.hidden, args.intermediate_hidden, + args.num_experts, args.num_topk, + args.config.num_experts_per_wave, + args.config.block_m, args.config.block_n, args.config.block_k, + args.config.store_block_m, + args.config.sf_block_m, args.config.sf_block_n, + 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), + args.fast_math ? "true" : "false", + get_mxfp4_side_lora_activation_type_name(args.activation), + args.save_l1_preact ? "true" : "false", + get_mxfp4_side_lora_route_weight_mode_name( + args.route_weight_mode), + args.save_down_unweighted ? "true" : "false", + args.side_lora_rank); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + // TODO: optimize `args` copy + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.y, + args.saved_l1_preact, + args.saved_x, + args.saved_h, + args.saved_down_unweighted, + args.side_lora_ready, + args.side_lora_scale, + args.cumulative_local_expert_recv_stats, + args.num_tokens, + args.num_saved_pool_tokens, + args.sym_buffer_ptrs, + args.tensor_map_l1_acts, + args.tensor_map_l1_acts_sf, + args.tensor_map_l1_weights, + args.tensor_map_l1_weights_sf, + args.tensor_map_l1_output, + args.tensor_map_l2_acts, + args.tensor_map_l2_acts_sf, + args.tensor_map_l2_weights, + args.tensor_map_l2_weights_sf, + args.tensor_map_down_unweighted, + args.tensor_map_saved_x, + args.tensor_map_saved_h, + args.tensor_map_lora_a1, + args.tensor_map_lora_a3, + args.tensor_map_lora_b1, + args.tensor_map_lora_b3, + args.tensor_map_lora_a2, + args.tensor_map_lora_b2, + args.tensor_map_lora_l1_scratch, + args.tensor_map_lora_l2_scratch, + args.tensor_map_lora_l1_scratch_store, + args.tensor_map_lora_l2_scratch_store + )); + } +}; + +static void sm100_fp8_fp4_mega_moe_side_lora_forward( + const torch::Tensor& y, + const std::optional& saved_l1_preact, + const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf, + const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf, + const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, + const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf, + const std::optional cumulative_local_expert_recv_stats, + const std::vector& sym_buffer_ptrs, + const int& rank_idx, const int& num_max_tokens_per_rank, + const int& num_experts_per_rank, + const int& num_tokens, const int& num_config_tokens, + const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const std::string& activation, + const float& activation_clamp, + const bool& fast_math, + const std::string& route_weight_mode, + const std::optional& saved_down_unweighted, + const torch::Tensor& side_lora_source, + const std::optional& saved_x, + const std::optional& saved_h_unweighted, + const std::optional& side_lora_a1, + const std::optional& side_lora_b1, + const std::optional& side_lora_a3, + const std::optional& side_lora_b3, + const std::optional& side_lora_a2, + const std::optional& side_lora_b2, + const std::optional& side_lora_l1_scratch, + const std::optional& side_lora_l2_scratch, + const std::optional& side_lora_ready, + const float& side_lora_scale +) { + 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 bool has_side_lora = side_lora_a1.has_value(); + const int side_lora_rank = has_side_lora ? 128 : 0; + DG_HOST_ASSERT(has_side_lora); + DG_HOST_ASSERT( + saved_x.has_value() && saved_h_unweighted.has_value() && + side_lora_b1.has_value() && side_lora_a3.has_value() && + side_lora_b3.has_value() && side_lora_a2.has_value() && + side_lora_b2.has_value() && side_lora_l1_scratch.has_value() && + side_lora_l2_scratch.has_value() && side_lora_ready.has_value()); + + // Heuristics + const auto config = get_mega_moe_config( + num_ranks, num_experts, num_experts_per_rank, + num_max_tokens_per_rank, num_config_tokens, num_topk, + hidden, intermediate_hidden, + num_ring_tokens, num_sf_ring_tokens, + MmaKind::MXFP8FP4); + const auto num_max_pool_tokens = + layout::get_num_max_pool_tokens( + num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank); + const int num_saved_pool_tokens = + static_cast(saved_x->size(0)); + const auto check_bf16_contiguous = [](const torch::Tensor& tensor) { + DG_HOST_ASSERT(tensor.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(tensor.is_contiguous()); + }; + for (const auto* tensor : { + &side_lora_source, &*saved_x, &*saved_h_unweighted, + &*side_lora_a1, &*side_lora_b1, &*side_lora_a3, + &*side_lora_b3, &*side_lora_a2, &*side_lora_b2, + &*side_lora_l1_scratch, &*side_lora_l2_scratch}) + check_bf16_contiguous(*tensor); + DG_HOST_ASSERT(saved_x->sizes() == torch::IntArrayRef( + {num_saved_pool_tokens, hidden})); + DG_HOST_ASSERT(saved_h_unweighted->sizes() == torch::IntArrayRef( + {num_saved_pool_tokens, intermediate_hidden})); + DG_HOST_ASSERT(side_lora_a1->sizes() == torch::IntArrayRef( + {side_lora_rank, hidden})); + DG_HOST_ASSERT(side_lora_a3->sizes() == side_lora_a1->sizes()); + DG_HOST_ASSERT(side_lora_b1->sizes() == torch::IntArrayRef( + {num_experts_per_rank, intermediate_hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_b3->sizes() == side_lora_b1->sizes()); + DG_HOST_ASSERT(side_lora_a2->sizes() == torch::IntArrayRef( + {num_experts_per_rank, side_lora_rank, intermediate_hidden})); + DG_HOST_ASSERT(side_lora_b2->sizes() == torch::IntArrayRef( + {hidden, side_lora_rank})); + DG_HOST_ASSERT(side_lora_l1_scratch->sizes() == torch::IntArrayRef( + {num_saved_pool_tokens, 2, side_lora_rank})); + DG_HOST_ASSERT(side_lora_l2_scratch->sizes() == torch::IntArrayRef( + {num_saved_pool_tokens, side_lora_rank})); + DG_HOST_ASSERT(side_lora_ready->scalar_type() == torch::kInt); + DG_HOST_ASSERT(side_lora_ready->is_contiguous()); + DG_HOST_ASSERT(side_lora_ready->numel() >= + 4 * config.num_ring_tokens / config.block_m); + side_lora_ready->zero_(); + if (saved_l1_preact.has_value()) { + DG_HOST_ASSERT(saved_l1_preact->scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(saved_l1_preact->is_contiguous()); + DG_HOST_ASSERT(saved_l1_preact->dim() == 2); + DG_HOST_ASSERT(saved_l1_preact->size(0) > 0 && + saved_l1_preact->size(0) <= num_max_pool_tokens); + DG_HOST_ASSERT(saved_l1_preact->size(0) % config.block_m == 0); + DG_HOST_ASSERT(saved_l1_preact->size(1) == + 2 * intermediate_hidden); + } + DG_HOST_ASSERT( + route_weight_mode == "pre_down" || + route_weight_mode == "post_down"); + if (saved_down_unweighted.has_value()) { + DG_HOST_ASSERT( + saved_down_unweighted->scalar_type() == + torch::kBFloat16); + DG_HOST_ASSERT(saved_down_unweighted->is_contiguous()); + DG_HOST_ASSERT(saved_down_unweighted->dim() == 2); + DG_HOST_ASSERT(saved_down_unweighted->size(1) == hidden); + DG_HOST_ASSERT(saved_down_unweighted->size(0) > 0); + DG_HOST_ASSERT( + saved_down_unweighted->size(0) % + config.block_m == 0); + DG_HOST_ASSERT( + saved_down_unweighted->size(0) <= + num_max_pool_tokens); + } + + // Make tensormap + constexpr int kGranK = 32; + const int sf_smem_outer_dim = config.block_k / (kGranK * 4); + const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, + hidden, config.num_ring_tokens, + config.block_k, config.load_block_m, + static_cast(l1_acts.stride(-2)), + config.swizzle_acts_mode); + const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, + config.num_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, + static_cast(l1_weights.stride(-2)), + config.swizzle_weights_mode); + 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); + // 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). + const auto tensor_map_l1_output = 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); + const auto tensor_map_l2_acts = 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 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, + static_cast(l2_weights.stride(-2)), + config.swizzle_weights_mode); + 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, + num_experts_per_rank, 0, 0, false, + sf_smem_outer_dim); + const auto tensor_map_down_unweighted = + saved_down_unweighted.has_value() + ? make_tma_2d_desc( + *saved_down_unweighted, + hidden, saved_down_unweighted->size(0), + config.block_n, config.store_block_m, + static_cast( + saved_down_unweighted->stride(-2)), + config.swizzle_acts_mode) + : tensor_map_l2_acts; + constexpr int kSideBlockK = 64; + const auto tensor_map_saved_x = has_side_lora + ? make_tma_2d_desc(*saved_x, hidden, saved_x->size(0), + kSideBlockK, config.load_block_m, + static_cast(saved_x->stride(0)), 128) + : tensor_map_l1_acts; + const auto tensor_map_saved_h = has_side_lora + ? make_tma_2d_desc(*saved_h_unweighted, intermediate_hidden, + saved_h_unweighted->size(0), kSideBlockK, + config.load_block_m, + static_cast(saved_h_unweighted->stride(0)), 128) + : tensor_map_l2_acts; + const auto make_side_weight_desc = [&](const torch::Tensor& tensor, + int k, int n) { + return make_tma_2d_desc(tensor, k, n, kSideBlockK, + config.load_block_n, + static_cast(tensor.stride(-2)), 128); + }; + const auto make_l1_expand_weight_desc = [&](const torch::Tensor& tensor) { + // B1/B3 are compact [E, I, R]. The kernel scatters eight-row TMA + // strips into the gate/up slots of the interleaved 128-column tile. + return make_tma_2d_desc(tensor, side_lora_rank, + num_experts_per_rank * intermediate_hidden, + kSideBlockK, 8, + static_cast(tensor.stride(-2)), 128); + }; + const auto tensor_map_lora_a1 = has_side_lora + ? make_side_weight_desc(*side_lora_a1, hidden, + side_lora_rank) + : tensor_map_l1_weights; + const auto tensor_map_lora_a3 = has_side_lora + ? make_side_weight_desc(*side_lora_a3, hidden, + side_lora_rank) + : tensor_map_l1_weights; + const auto tensor_map_lora_b1 = has_side_lora + ? make_l1_expand_weight_desc(*side_lora_b1) + : tensor_map_l1_weights; + const auto tensor_map_lora_b3 = has_side_lora + ? make_l1_expand_weight_desc(*side_lora_b3) + : tensor_map_l1_weights; + const auto tensor_map_lora_a2 = has_side_lora + ? make_side_weight_desc(*side_lora_a2, intermediate_hidden, + num_experts_per_rank * side_lora_rank) + : tensor_map_l2_weights; + const auto tensor_map_lora_b2 = has_side_lora + ? make_side_weight_desc(*side_lora_b2, side_lora_rank, + hidden) + : tensor_map_l2_weights; + const auto tensor_map_lora_l1_scratch = has_side_lora + ? make_tma_2d_desc(*side_lora_l1_scratch, 2 * side_lora_rank, + num_saved_pool_tokens, kSideBlockK, config.load_block_m, + static_cast(side_lora_l1_scratch->stride(0)), 128) + : tensor_map_l1_acts; + const auto tensor_map_lora_l2_scratch = has_side_lora + ? make_tma_2d_desc(*side_lora_l2_scratch, side_lora_rank, + num_saved_pool_tokens, kSideBlockK, config.load_block_m, + static_cast(side_lora_l2_scratch->stride(0)), 128) + : tensor_map_l2_acts; + // Shrink epilogues stage STORE_BLOCK_M rows, while the following expand + // GEMM loads LOAD_BLOCK_M rows. A TMA descriptor encodes that box height, + // so using the load descriptor for the store reads beyond the staged + // shared-memory tile whenever a full M block is present. + const auto tensor_map_lora_l1_scratch_store = has_side_lora + ? make_tma_2d_desc(*side_lora_l1_scratch, 2 * side_lora_rank, + num_saved_pool_tokens, kSideBlockK, config.store_block_m, + static_cast(side_lora_l1_scratch->stride(0)), 128) + : tensor_map_l1_acts; + const auto tensor_map_lora_l2_scratch_store = has_side_lora + ? make_tma_2d_desc(*side_lora_l2_scratch, side_lora_rank, + num_saved_pool_tokens, kSideBlockK, config.store_block_m, + static_cast(side_lora_l2_scratch->stride(0)), 128) + : tensor_map_l2_acts; + + // Stats can be optional + int* cumulative_local_expert_recv_stats_ptr = nullptr; + if (cumulative_local_expert_recv_stats.has_value()) + cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); + + // Launch + const auto num_sms = device_runtime->get_num_sms(); + const SM100FP8FP4MegaMoESideLoraForwardRuntime::Args args = { + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .hidden = hidden, .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, .num_topk = num_topk, + .num_ranks = num_ranks, + .activation_clamp = activation_clamp, + .fast_math = fast_math, + .activation = activation, + .save_l1_preact = saved_l1_preact.has_value(), + .route_weight_mode = route_weight_mode, + .save_down_unweighted = + saved_down_unweighted.has_value(), + .side_lora_rank = side_lora_rank, + .config = config, + .y = y.data_ptr(), + .saved_l1_preact = saved_l1_preact.has_value() + ? saved_l1_preact->data_ptr() + : nullptr, + .saved_x = has_side_lora ? saved_x->data_ptr() : nullptr, + .saved_h = has_side_lora + ? saved_h_unweighted->data_ptr() : nullptr, + .saved_down_unweighted = saved_down_unweighted.has_value() + ? saved_down_unweighted->data_ptr() : nullptr, + .side_lora_ready = has_side_lora + ? side_lora_ready->data_ptr() : nullptr, + .side_lora_scale = side_lora_scale, + .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, + .num_tokens = num_tokens, + .num_saved_pool_tokens = num_saved_pool_tokens, + .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + .tensor_map_l1_acts = tensor_map_l1_acts, + .tensor_map_l1_acts_sf = tensor_map_l1_acts_sf, + .tensor_map_l1_weights = tensor_map_l1_weights, + .tensor_map_l1_weights_sf = tensor_map_l1_weights_sf, + .tensor_map_l1_output = tensor_map_l1_output, + .tensor_map_l2_acts = tensor_map_l2_acts, + .tensor_map_l2_acts_sf = tensor_map_l2_acts_sf, + .tensor_map_l2_weights = tensor_map_l2_weights, + .tensor_map_l2_weights_sf = tensor_map_l2_weights_sf, + .tensor_map_down_unweighted = + tensor_map_down_unweighted, + .tensor_map_saved_x = tensor_map_saved_x, + .tensor_map_saved_h = tensor_map_saved_h, + .tensor_map_lora_a1 = tensor_map_lora_a1, + .tensor_map_lora_a3 = tensor_map_lora_a3, + .tensor_map_lora_b1 = tensor_map_lora_b1, + .tensor_map_lora_b3 = tensor_map_lora_b3, + .tensor_map_lora_a2 = tensor_map_lora_a2, + .tensor_map_lora_b2 = tensor_map_lora_b2, + .tensor_map_lora_l1_scratch = tensor_map_lora_l1_scratch, + .tensor_map_lora_l2_scratch = tensor_map_lora_l2_scratch, + .tensor_map_lora_l1_scratch_store = + tensor_map_lora_l1_scratch_store, + .tensor_map_lora_l2_scratch_store = + tensor_map_lora_l2_scratch_store, + .launch_args = LaunchArgs(num_sms, + config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, + config.smem_size, 2) + }; + + const auto code = SM100FP8FP4MegaMoESideLoraForwardRuntime::generate(args); + const auto runtime = compiler->build("sm100_fp8_fp4_mega_moe_side_lora_forward", code); + SM100FP8FP4MegaMoESideLoraForwardRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 11af41684f..22c2e5d06f 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -96,15 +96,21 @@ SymmBuffer, get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe, + transform_side_lora_for_mega_moe, fp8_fp4_mega_moe, + fp8_fp4_mega_moe_side_lora, bf16_mega_moe, + bf16_mega_moe_side_lora, ) from .mega.backward import ( + MegaMoESideLoraBackwardResult, bf16_mega_moe_backward_dgrad, + bf16_mega_moe_side_lora_backward, bf16_mega_moe_backward_w13, bf16_mega_moe_backward_w13_combine, bf16_mega_moe_backward_w2, bf16_mega_moe_backward_w2_combine, + fp8_fp4_mega_moe_side_lora_backward, fp8_fp4_mega_moe_backward_dgrad_swiglu, mega_moe_backward_combine_grad_x, ) diff --git a/deep_gemm/include/deep_gemm/impls/mega_moe_side_lora_params.cuh b/deep_gemm/include/deep_gemm/impls/mega_moe_side_lora_params.cuh new file mode 100644 index 0000000000..849833fc51 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/mega_moe_side_lora_params.cuh @@ -0,0 +1,41 @@ +#pragma once + +#include + +namespace deep_gemm { + +// Rank-128 native MegaMoE side-LoRA backward ABI. All pool tensors are +// expert-major with the same BLOCK_M padding as the base backward wave. +// Only rank-width intermediates cross phases; full-width side dgrad tensors +// are deliberately absent from this contract. +struct alignas(16) MegaMoESideLoraBackwardParams { + using bf16 = cutlass::bfloat16_t; + + const bf16* saved_x; + const bf16* saved_h; + + const bf16* a1; + const bf16* b1; + const bf16* a3; + const bf16* b3; + const bf16* a2; + const bf16* b2; + + bf16* q1; + bf16* q3; + bf16* q2; + bf16* t1; + bf16* t3; + bf16* t2; + + bf16* grad_a1; + bf16* grad_b1; + bf16* grad_a3; + bf16* grad_b3; + bf16* grad_a2; + bf16* grad_b2; + + float scale; +}; + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh new file mode 100644 index 0000000000..7ad2536df4 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh @@ -0,0 +1,5656 @@ +#pragma once + +#include + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +template < + uint32_t kHidden, uint32_t kNumExperts, uint32_t BLOCK_M, + uint32_t kNumSMs, uint32_t kNumThreads, + CombineOrderMode kCombineOrderMode> +__device__ __forceinline__ void +bf16_mega_moe_reduce_post_down_route( + const int* expert_counts, + const cutlass::bfloat16_t* grad_y_unweighted, + const cutlass::bfloat16_t* down_unweighted, + float* grad_route_output, + uint8_t* scratch) { + constexpr uint32_t kTritonRouteBlockH = [] { + uint32_t value = 1; + while (value < kHidden && value < 8192) + value <<= 1; + return value; + }(); + constexpr uint32_t kTritonRouteNumWarps = [] { + uint32_t value = kTritonRouteBlockH / 256; + value = value < 4 ? 4 : value; + return value > 32 ? 32 : value; + }(); + constexpr uint32_t kTritonRouteThreads = + kTritonRouteNumWarps * 32; + constexpr uint32_t kTritonRouteValuesPerThread = + kTritonRouteBlockH / kTritonRouteThreads; + DG_STATIC_ASSERT( + kTritonRouteValuesPerThread == 2 || + kTritonRouteValuesPerThread == 4 || + kTritonRouteValuesPerThread == 8, + "Unsupported Triton route reduction width"); + constexpr uint32_t kRouteInputPow2 = [] { + uint32_t value = 1; + constexpr uint32_t vectorized_columns = kHidden / 4; + while (value < 512 && + (value << 1) <= vectorized_columns) + value <<= 1; + return value; + }(); + + auto* route_lane_sums = reinterpret_cast(scratch); + auto* route_control = reinterpret_cast(scratch); + if constexpr ( + kCombineOrderMode != CombineOrderMode::FixedTopK) { + // A sub-CTA named barrier is not safe while the persistent kernel's + // earlier role-specific register/barrier phases are still live. + // Assign one route to the CTA instead. The first Triton-sized thread + // group keeps exactly the same lane-to-column map and butterfly tree; + // the remaining threads only participate in CTA phase barriers. + auto* route_warp_arrivals = + reinterpret_cast( + scratch + + kNumThreads * sizeof(float)); + if (threadIdx.x == 0) + *route_warp_arrivals = 0; + __syncthreads(); + uint32_t route_pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg(expert_counts + expert_idx)); + for (uint32_t token_idx = blockIdx.x; + token_idx < num_tokens; + token_idx += kNumSMs) { + const uint32_t pool_row = + route_pool_block_offset * BLOCK_M + + token_idx; + const uint32_t route_lane = threadIdx.x; + float grad_route = 0.0f; + if (route_lane < kTritonRouteThreads) { + float grad_y[ + kTritonRouteValuesPerThread]; + float down[ + kTritonRouteValuesPerThread]; + #pragma unroll + for (uint32_t i = 0; + i < + kTritonRouteValuesPerThread; + ++i) { + const uint32_t col = + route_lane + + i * kTritonRouteThreads; + grad_y[i] = + col < kHidden + ? static_cast( + grad_y_unweighted[ + static_cast< + uint64_t>( + pool_row) * + kHidden + + col]) + : 0.0f; + down[i] = + col < kHidden + ? static_cast( + down_unweighted[ + static_cast< + uint64_t>( + pool_row) * + kHidden + + col]) + : 0.0f; + } + if constexpr ( + kTritonRouteValuesPerThread == 2) { + grad_route = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[1], down[1])); + } else if constexpr ( + kTritonRouteValuesPerThread == 4) { + const float even = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], down[2])); + const float odd = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], down[3])); + grad_route = + __fadd_rn(even, odd); + } else { + const float pair_02 = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], down[2])); + const float pair_13 = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], down[3])); + const float pair_46 = __fmaf_rn( + grad_y[4], down[4], + __fmul_rn( + grad_y[6], down[6])); + const float pair_57 = __fmaf_rn( + grad_y[5], down[5], + __fmul_rn( + grad_y[7], down[7])); + grad_route = __fadd_rn( + __fadd_rn( + pair_02, pair_46), + __fadd_rn( + pair_13, pair_57)); + } + #pragma unroll + for (uint32_t offset = 16; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, offset)); + } + const uint32_t lane_in_warp = + route_lane & 31; + if (lane_in_warp == 0) { + route_lane_sums[ + route_lane / 32] = + grad_route; + __threadfence_block(); + atomicAdd( + route_warp_arrivals, 1u); + } + } + if (threadIdx.x < 32) { + if (threadIdx.x == 0) { + while (atomicAdd( + route_warp_arrivals, + 0u) != + kTritonRouteNumWarps) { + } + } + __syncwarp(); + grad_route = route_lane_sums[ + threadIdx.x & + (kTritonRouteNumWarps - 1)]; + #pragma unroll + for (uint32_t offset = + kTritonRouteNumWarps / 2; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, offset)); + } + } + if (threadIdx.x == 0) + grad_route_output[pool_row] = + grad_route; + __syncthreads(); + if (threadIdx.x == 0) + *route_warp_arrivals = 0; + __syncthreads(); + } + route_pool_block_offset += + math::ceil_div(num_tokens, BLOCK_M); + } + return; + } + + if (threadIdx.x == 0) { + uint32_t total_route_rows = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + total_route_rows += static_cast( + __ldg(expert_counts + expert_idx)); + } + route_control[0] = total_route_rows; + } + __syncthreads(); + const uint32_t total_route_rows = route_control[0]; + const uint32_t route_output_pow2 = + total_route_rows > 0 + ? 1u << (31 - __clz(total_route_rows)) + : 1u; + constexpr uint32_t kInitialRouteGroupThreads = + cute::min(kRouteInputPow2, 32u); + const uint32_t route_block_height = + cute::min( + route_output_pow2, + 512u / kInitialRouteGroupThreads); + const uint32_t route_group_threads = + kCombineOrderMode != CombineOrderMode::FixedTopK + ? kTritonRouteThreads + : cute::min( + kRouteInputPow2, + 512u / route_block_height); + const uint32_t num_route_groups_per_cta = + kNumThreads / route_group_threads; + const uint32_t route_group_idx = + threadIdx.x / route_group_threads; + const uint32_t route_group_lane_idx = + threadIdx.x & (route_group_threads - 1); + const uint32_t global_route_group = + blockIdx.x * num_route_groups_per_cta + + route_group_idx; + const uint32_t num_route_groups = + kNumSMs * num_route_groups_per_cta; + uint32_t route_pool_block_offset = 0; + + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = static_cast( + __ldg(expert_counts + expert_idx)); + for (uint32_t token_idx = global_route_group; + token_idx < num_tokens; + token_idx += num_route_groups) { + const uint32_t pool_row = + route_pool_block_offset * BLOCK_M + token_idx; + float grad_route = 0.0f; + if constexpr ( + kCombineOrderMode != CombineOrderMode::FixedTopK) { + float grad_y[kTritonRouteValuesPerThread]; + float down[kTritonRouteValuesPerThread]; + #pragma unroll + for (uint32_t i = 0; + i < kTritonRouteValuesPerThread; ++i) { + const uint32_t col = + route_group_lane_idx + + i * kTritonRouteThreads; + grad_y[i] = + col < kHidden + ? static_cast( + grad_y_unweighted[ + static_cast(pool_row) * + kHidden + + col]) + : 0.0f; + down[i] = + col < kHidden + ? static_cast( + down_unweighted[ + static_cast(pool_row) * + kHidden + + col]) + : 0.0f; + } + + if constexpr (kTritonRouteValuesPerThread == 2) { + grad_route = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn(grad_y[1], down[1])); + } else if constexpr ( + kTritonRouteValuesPerThread == 4) { + const float even = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn(grad_y[2], down[2])); + const float odd = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn(grad_y[3], down[3])); + grad_route = __fadd_rn(even, odd); + } else { + const float pair_02 = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn(grad_y[2], down[2])); + const float pair_13 = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn(grad_y[3], down[3])); + const float pair_46 = __fmaf_rn( + grad_y[4], down[4], + __fmul_rn(grad_y[6], down[6])); + const float pair_57 = __fmaf_rn( + grad_y[5], down[5], + __fmul_rn(grad_y[7], down[7])); + grad_route = __fadd_rn( + __fadd_rn(pair_02, pair_46), + __fadd_rn(pair_13, pair_57)); + } + + #pragma unroll + for (uint32_t offset = 16; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, grad_route, offset)); + } + const uint32_t warp_in_group = + route_group_lane_idx / 32; + const uint32_t lane_in_warp = + route_group_lane_idx & 31; + if (lane_in_warp == 0) { + route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + warp_in_group] = grad_route; + } + ptx::sync_aligned( + kTritonRouteThreads, route_group_idx); + if (warp_in_group == 0) { + grad_route = route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + (lane_in_warp & + (kTritonRouteNumWarps - 1))]; + #pragma unroll + for (uint32_t offset = + kTritonRouteNumWarps / 2; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, offset)); + } + } + } else { + float lane_sums[4] = { + 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t col_base = + route_group_lane_idx * 4; + col_base < kHidden; + col_base += route_group_threads * 4) { + #pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + const uint32_t col = col_base + i; + const float grad_y = static_cast( + grad_y_unweighted[ + static_cast(pool_row) * + kHidden + + col]); + const float down = static_cast( + down_unweighted[ + static_cast(pool_row) * + kHidden + + col]); + lane_sums[i] = __fadd_rn( + lane_sums[i], + __fmul_rn(grad_y, down)); + } + } + grad_route = __fadd_rn( + __fadd_rn(lane_sums[0], lane_sums[1]), + lane_sums[2]); + grad_route = + __fadd_rn(grad_route, lane_sums[3]); + route_lane_sums[threadIdx.x] = grad_route; + if (route_group_threads > 32) { + for (uint32_t offset = + route_group_threads / 2; + offset >= 32; offset >>= 1) { + ptx::sync_aligned( + route_group_threads, + route_group_idx); + if (route_group_lane_idx < offset) { + grad_route = __fadd_rn( + grad_route, + route_lane_sums[ + threadIdx.x + offset]); + route_lane_sums[threadIdx.x] = + grad_route; + } + } + } + if (route_group_lane_idx < 32) { + #pragma unroll + for (uint32_t offset = 16; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_down_sync( + 0xffffffff, + grad_route, offset)); + } + } + } + if (route_group_lane_idx == 0) + grad_route_output[pool_row] = grad_route; + if (route_group_threads > 32) { + ptx::sync_aligned( + route_group_threads, route_group_idx); + } else { + __syncwarp(); + } + } + route_pool_block_offset += + math::ceil_div(num_tokens, BLOCK_M); + } +} + +template < + uint32_t kHidden, uint32_t kNumExperts, uint32_t BLOCK_M, + uint32_t kNumSMs, uint32_t kNumRanks, + CombineOrderMode kCombineOrderMode, + bool kDoReverseDispatch = true, + bool kComputeRouteDot = true, + bool kWriteWeighted = true, + bool kWeightedSourceIsRhs = false, + bool kSynchronizeRanks = true, + bool kSynchronizeAfterDispatch = true, + bool kBarrierOnly = false, + bool kXPrepared = false, + uint32_t kRoutePreludeThreads = 256> +CUTLASS_GLOBAL __launch_bounds__(1024, 1) void +sm100_bf16_mega_moe_backward_post_down_prelude( + const int* expert_counts, + const __grid_constant__ layout::Workspace + backward_workspace, + const __grid_constant__ layout::SymBuffer + backward_sym_buffer, + const cutlass::bfloat16_t* backward_grad_y, + const cutlass::bfloat16_t* backward_x, + const float* backward_topk_weights, + float* backward_grad_route, + const layout::TokenSrcMetadata* token_src_metadata, + const uint32_t num_topk, + const uint32_t num_pool_rows, + cutlass::bfloat16_t* grad_y_unweighted_output, + cutlass::bfloat16_t* grad_y_weighted_output, + cutlass::bfloat16_t* x_pool_output, + float* route_weights_output, + const cutlass::bfloat16_t* down_unweighted, + float* grad_route_output) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) + constexpr uint32_t kNumThreads = 1024; + if constexpr (kSynchronizeRanks) { + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumThreads, 0, 71>( + backward_workspace, + backward_sym_buffer, + blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }); + } + if constexpr (kBarrierOnly) { + if constexpr (kSynchronizeAfterDispatch) { + // Profiling-only completion barrier. Production folds this into + // the dispatch launch below to avoid an extra kernel launch. + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumThreads, 1, 72>( + backward_workspace, + backward_sym_buffer, + blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }); + } + return; + } + constexpr uint32_t kTritonRouteBlockH = [] { + uint32_t value = 1; + while (value < kHidden && value < 8192) + value <<= 1; + return value; + }(); + constexpr uint32_t kTritonRouteNumWarps = [] { + uint32_t value = kTritonRouteBlockH / 256; + value = value < 4 ? 4 : value; + return value > 32 ? 32 : value; + }(); + constexpr uint32_t kTritonRouteThreads = + kTritonRouteNumWarps * 32; + constexpr uint32_t kTritonRouteValuesPerThread = + kTritonRouteBlockH / kTritonRouteThreads; + constexpr bool kVirtualizeRouteLanes = + kRoutePreludeThreads == 128; + DG_STATIC_ASSERT( + kRoutePreludeThreads == 128 || + kRoutePreludeThreads == 256, + "POST_DOWN route prelude requires 128 or 256 physical threads"); + DG_STATIC_ASSERT( + !kVirtualizeRouteLanes || + (kHidden == 2048 && + kCombineOrderMode != + CombineOrderMode::FixedTopK && + kComputeRouteDot), + "128-thread route prelude is only supported for the exact " + "non-fixed H=2048 route-dot path"); + constexpr uint32_t kExactRouteGroupThreads = + kVirtualizeRouteLanes + ? kRoutePreludeThreads + : kTritonRouteThreads; + constexpr uint32_t kRouteVirtualLanes = + kTritonRouteThreads / + kExactRouteGroupThreads; + DG_STATIC_ASSERT( + kRouteVirtualLanes == 1 || + kRouteVirtualLanes == 2, + "Unsupported POST_DOWN route lane virtualization"); + DG_STATIC_ASSERT( + kTritonRouteValuesPerThread == 2 || + kTritonRouteValuesPerThread == 4 || + kTritonRouteValuesPerThread == 8, + "Unsupported Triton route reduction width"); + constexpr uint32_t kRouteInputPow2 = [] { + uint32_t value = 1; + constexpr uint32_t vectorized_columns = + kHidden / 4; + while (value < 512 && + (value << 1) <= vectorized_columns) + value <<= 1; + return value; + }(); + constexpr uint32_t kInitialRouteGroupThreads = + cute::min(kRouteInputPow2, 32u); + + extern __shared__ __align__(1024) uint8_t scratch[]; + auto* route_lane_sums = + reinterpret_cast(scratch); + auto* route_control = + reinterpret_cast(scratch); + if (threadIdx.x == 0) { + uint32_t total_route_rows = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + total_route_rows += static_cast( + __ldg(expert_counts + expert_idx)); + } + route_control[0] = total_route_rows; + } + __syncthreads(); + const uint32_t total_route_rows = route_control[0]; + const uint32_t route_output_pow2 = + total_route_rows > 0 + ? 1u << (31 - __clz(total_route_rows)) + : 1u; + const uint32_t route_block_height = + cute::min( + route_output_pow2, + 512u / kInitialRouteGroupThreads); + // The dispatch-only launch copies vectorized BF16 payloads and does not + // need Triton's exact reduction lane map. Use more route groups per CTA + // so remote reads have enough independent rows to cover NVLink latency. + const uint32_t route_group_threads = + !kComputeRouteDot && !kWriteWeighted + ? cute::min(kRouteInputPow2, 128u) + : kCombineOrderMode != CombineOrderMode::FixedTopK + ? kExactRouteGroupThreads + : cute::min( + kRouteInputPow2, + 512u / route_block_height); + const uint32_t num_route_groups_per_cta = + kNumThreads / route_group_threads; + const uint32_t route_group_idx = + threadIdx.x / route_group_threads; + const uint32_t route_group_lane_idx = + threadIdx.x & (route_group_threads - 1); + const uint32_t global_route_group = + blockIdx.x * num_route_groups_per_cta + + route_group_idx; + const uint32_t num_route_groups = + kNumSMs * num_route_groups_per_cta; + uint32_t route_pool_block_offset = 0; + + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = static_cast( + __ldg(expert_counts + expert_idx)); + for (uint32_t token_idx = global_route_group; + token_idx < num_tokens; + token_idx += num_route_groups) { + const uint32_t pool_row = + route_pool_block_offset * BLOCK_M + + token_idx; + const cutlass::bfloat16_t* remote_grad_y; + const cutlass::bfloat16_t* remote_x; + if constexpr (kDoReverseDispatch) { + const auto metadata = + token_src_metadata[pool_row]; + remote_grad_y = + backward_sym_buffer.map( + backward_grad_y + + static_cast( + metadata.token_idx) * + kHidden, + metadata.rank_idx); + if constexpr (kXPrepared) { + remote_x = + x_pool_output + + static_cast(pool_row) * + kHidden; + } else { + remote_x = + backward_sym_buffer.map( + backward_x + + static_cast( + metadata.token_idx) * + kHidden, + metadata.rank_idx); + } + if (route_group_lane_idx == 0) { + const auto* remote_weight = + backward_sym_buffer.map( + backward_topk_weights + + static_cast( + metadata.token_idx) * + num_topk + + metadata.topk_idx, + metadata.rank_idx); + route_weights_output[pool_row] = + *remote_weight; + } + } else { + remote_grad_y = + grad_y_unweighted_output + + static_cast(pool_row) * + kHidden; + remote_x = + x_pool_output + + static_cast(pool_row) * + kHidden; + } + + float grad_route = 0.0f; + if constexpr (!kComputeRouteDot && !kWriteWeighted) { + constexpr uint32_t kBF16ValuesPerVector = + sizeof(uint4) / + sizeof(cutlass::bfloat16_t); + DG_STATIC_ASSERT( + kHidden % kBF16ValuesPerVector == 0, + "BF16 dispatch requires vector-aligned hidden"); + for (uint32_t col = + route_group_lane_idx * + kBF16ValuesPerVector; + col < kHidden; + col += route_group_threads * + kBF16ValuesPerVector) { + const uint64_t offset = + static_cast( + pool_row) * + kHidden + + col; + reinterpret_cast( + grad_y_unweighted_output)[ + offset / + kBF16ValuesPerVector] = + reinterpret_cast< + const uint4*>( + remote_grad_y)[ + col / + kBF16ValuesPerVector]; + if constexpr (!kXPrepared) { + reinterpret_cast( + x_pool_output)[ + offset / + kBF16ValuesPerVector] = + reinterpret_cast< + const uint4*>( + remote_x)[ + col / + kBF16ValuesPerVector]; + } + } + } else if constexpr ( + !kComputeRouteDot && kWriteWeighted) { + const float route_weight = + route_weights_output[pool_row]; + for (uint32_t col = + route_group_lane_idx; + col < kHidden; + col += route_group_threads) { + const uint64_t offset = + static_cast( + pool_row) * + kHidden + + col; + grad_y_weighted_output[offset] = + cutlass::bfloat16_t( + static_cast( + (kWeightedSourceIsRhs + ? down_unweighted[ + static_cast( + pool_row) * + kHidden + + col] + : remote_grad_y[col])) * + route_weight); + } + } else if constexpr ( + kCombineOrderMode != + CombineOrderMode::FixedTopK) { + if constexpr (kVirtualizeRouteLanes) { + // Each physical lane evaluates logical lanes p and + // p + 128. Their FMA and warp-XOR trees remain separate, + // then the four physical warps publish all eight logical + // Triton warp partials for the unchanged second level. + float weighted_values + [kRouteVirtualLanes] + [kTritonRouteValuesPerThread]; + #pragma unroll + for (uint32_t virtual_lane = 0; + virtual_lane < kRouteVirtualLanes; + ++virtual_lane) { + const uint32_t logical_route_lane = + route_group_lane_idx + + virtual_lane * + kExactRouteGroupThreads; + float grad_y[ + kTritonRouteValuesPerThread]; + float down[ + kTritonRouteValuesPerThread]; + #pragma unroll + for (uint32_t i = 0; + i < + kTritonRouteValuesPerThread; + ++i) { + const uint32_t col = + logical_route_lane + + i * kTritonRouteThreads; + grad_y[i] = + col < kHidden + ? static_cast( + remote_grad_y[col]) + : 0.0f; + down[i] = + col < kHidden + ? static_cast( + down_unweighted[ + static_cast( + pool_row) * + kHidden + + col]) + : 0.0f; + if constexpr ( + kDoReverseDispatch && + !kXPrepared) { + if (col < kHidden) { + x_pool_output[ + static_cast( + pool_row) * + kHidden + + col] = + remote_x[col]; + } + } + if constexpr (kWriteWeighted) { + weighted_values[virtual_lane][i] = + kWeightedSourceIsRhs + ? down[i] + : grad_y[i]; + } + } + float logical_grad_route; + if constexpr ( + kTritonRouteValuesPerThread == + 2) { + logical_grad_route = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[1], down[1])); + } else if constexpr ( + kTritonRouteValuesPerThread == + 4) { + const float even = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], down[2])); + const float odd = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], down[3])); + logical_grad_route = + __fadd_rn(even, odd); + } else { + const float pair_02 = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], down[2])); + const float pair_13 = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], down[3])); + const float pair_46 = __fmaf_rn( + grad_y[4], down[4], + __fmul_rn( + grad_y[6], down[6])); + const float pair_57 = __fmaf_rn( + grad_y[5], down[5], + __fmul_rn( + grad_y[7], down[7])); + logical_grad_route = __fadd_rn( + __fadd_rn( + pair_02, pair_46), + __fadd_rn( + pair_13, pair_57)); + } + #pragma unroll + for (uint32_t offset = 16; + offset > 0; offset >>= 1) { + logical_grad_route = __fadd_rn( + logical_grad_route, + __shfl_xor_sync( + 0xffffffff, + logical_grad_route, + offset)); + } + const uint32_t lane_in_warp = + route_group_lane_idx & 31; + const uint32_t logical_warp = + logical_route_lane / 32; + if (lane_in_warp == 0) { + route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + logical_warp] = + logical_grad_route; + } + } + ptx::sync_aligned( + kExactRouteGroupThreads, + route_group_idx); + const uint32_t physical_warp_in_group = + route_group_lane_idx / 32; + const uint32_t lane_in_warp = + route_group_lane_idx & 31; + if (physical_warp_in_group == 0) { + grad_route = route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + (lane_in_warp & + (kTritonRouteNumWarps - 1))]; + #pragma unroll + for (uint32_t offset = + kTritonRouteNumWarps / 2; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, offset)); + } + } + if constexpr (kWriteWeighted) { + const float route_weight = + route_weights_output[pool_row]; + #pragma unroll + for (uint32_t virtual_lane = 0; + virtual_lane < + kRouteVirtualLanes; + ++virtual_lane) { + #pragma unroll + for (uint32_t i = 0; + i < + kTritonRouteValuesPerThread; + ++i) { + const uint32_t col = + route_group_lane_idx + + virtual_lane * + kExactRouteGroupThreads + + i * + kTritonRouteThreads; + if (col < kHidden) { + grad_y_weighted_output[ + static_cast( + pool_row) * + kHidden + + col] = + cutlass::bfloat16_t( + weighted_values[ + virtual_lane] + [i] * + route_weight); + } + } + } + } + } else { + float grad_y[ + kTritonRouteValuesPerThread]; + float down[ + kTritonRouteValuesPerThread]; + #pragma unroll + for (uint32_t i = 0; + i < + kTritonRouteValuesPerThread; + ++i) { + const uint32_t col = + route_group_lane_idx + + i * kTritonRouteThreads; + grad_y[i] = + col < kHidden + ? static_cast( + remote_grad_y[col]) + : 0.0f; + down[i] = + col < kHidden + ? static_cast( + down_unweighted[ + static_cast( + pool_row) * + kHidden + + col]) + : 0.0f; + if constexpr ( + kDoReverseDispatch && + !kXPrepared) { + if (col < kHidden) { + x_pool_output[ + static_cast( + pool_row) * + kHidden + + col] = + remote_x[col]; + } + } + } + if constexpr ( + kTritonRouteValuesPerThread == + 2) { + grad_route = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[1], down[1])); + } else if constexpr ( + kTritonRouteValuesPerThread == + 4) { + const float even = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], down[2])); + const float odd = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], down[3])); + grad_route = + __fadd_rn(even, odd); + } else { + const float pair_02 = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], down[2])); + const float pair_13 = __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], down[3])); + const float pair_46 = __fmaf_rn( + grad_y[4], down[4], + __fmul_rn( + grad_y[6], down[6])); + const float pair_57 = __fmaf_rn( + grad_y[5], down[5], + __fmul_rn( + grad_y[7], down[7])); + grad_route = __fadd_rn( + __fadd_rn( + pair_02, pair_46), + __fadd_rn( + pair_13, pair_57)); + } + #pragma unroll + for (uint32_t offset = 16; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, offset)); + } + const uint32_t warp_in_group = + route_group_lane_idx / 32; + const uint32_t lane_in_warp = + route_group_lane_idx & 31; + if (lane_in_warp == 0) { + route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + warp_in_group] = + grad_route; + } + ptx::sync_aligned( + kTritonRouteThreads, + route_group_idx); + if (warp_in_group == 0) { + grad_route = route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + (lane_in_warp & + (kTritonRouteNumWarps - + 1))]; + #pragma unroll + for (uint32_t offset = + kTritonRouteNumWarps / + 2; + offset > 0; + offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, + offset)); + } + } + if constexpr (kWriteWeighted) { + const float route_weight = + route_weights_output[ + pool_row]; + #pragma unroll + for (uint32_t i = 0; + i < + kTritonRouteValuesPerThread; + ++i) { + const uint32_t col = + route_group_lane_idx + + i * + kTritonRouteThreads; + if (col < kHidden) { + grad_y_weighted_output[ + static_cast( + pool_row) * + kHidden + + col] = + cutlass::bfloat16_t( + (kWeightedSourceIsRhs + ? down[i] + : grad_y[i]) * + route_weight); + } + } + } + } + } else { + float lane_sums[4] = { + 0.0f, 0.0f, 0.0f, 0.0f}; + const float route_weight = + kWriteWeighted + ? route_weights_output[pool_row] + : 0.0f; + for (uint32_t col_base = + route_group_lane_idx * 4; + col_base < kHidden; + col_base += + route_group_threads * 4) { + #pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + const uint32_t col = + col_base + i; + const float grad_y = + static_cast( + remote_grad_y[col]); + const float down = + static_cast( + down_unweighted[ + static_cast( + pool_row) * + kHidden + + col]); + if constexpr (kWriteWeighted) { + grad_y_weighted_output[ + static_cast( + pool_row) * + kHidden + + col] = + cutlass::bfloat16_t( + (kWeightedSourceIsRhs + ? down + : grad_y) * + route_weight); + } + if constexpr (kDoReverseDispatch) { + grad_y_unweighted_output[ + static_cast( + pool_row) * + kHidden + + col] = + cutlass::bfloat16_t( + grad_y); + if constexpr (!kXPrepared) { + x_pool_output[ + static_cast( + pool_row) * + kHidden + + col] = + remote_x[col]; + } + } + lane_sums[i] = __fadd_rn( + lane_sums[i], + __fmul_rn( + grad_y, down)); + } + } + grad_route = __fadd_rn( + __fadd_rn( + lane_sums[0], + lane_sums[1]), + lane_sums[2]); + grad_route = __fadd_rn( + grad_route, lane_sums[3]); + route_lane_sums[threadIdx.x] = + grad_route; + if (route_group_threads > 32) { + for (uint32_t offset = + route_group_threads / 2; + offset >= 32; + offset >>= 1) { + ptx::sync_aligned( + route_group_threads, + route_group_idx); + if (route_group_lane_idx < + offset) { + grad_route = __fadd_rn( + grad_route, + route_lane_sums[ + threadIdx.x + + offset]); + route_lane_sums[ + threadIdx.x] = + grad_route; + } + } + } + if (route_group_lane_idx < 32) { + #pragma unroll + for (uint32_t offset = 16; + offset > 0; offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_down_sync( + 0xffffffff, + grad_route, offset)); + } + } + } + if constexpr (kComputeRouteDot) { + if (route_group_lane_idx == 0) { + grad_route_output[pool_row] = + grad_route; + if (backward_grad_route != nullptr) { + const auto metadata = + token_src_metadata[pool_row]; + auto* remote_grad_route = + backward_sym_buffer.map( + backward_grad_route + + static_cast( + metadata.token_idx) * + num_topk + + metadata.topk_idx, + metadata.rank_idx); + *remote_grad_route = grad_route; + } + } + if (route_group_threads > 32) { + ptx::sync_aligned( + route_group_threads, + route_group_idx); + } else { + __syncwarp(); + } + } + } + route_pool_block_offset += + math::ceil_div(num_tokens, BLOCK_M); + } + + uint32_t padded_pool_blocks = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = static_cast( + __ldg(expert_counts + expert_idx)); + const uint32_t num_blocks = + math::ceil_div(num_tokens, BLOCK_M); + const uint32_t num_padded_tokens = + num_blocks * BLOCK_M; + const uint32_t num_padding_rows = + num_padded_tokens - num_tokens; + for (uint64_t linear = + static_cast(blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast( + num_padding_rows) * + kHidden; + linear += + static_cast(kNumSMs) * + kNumThreads) { + const uint32_t padding_row = + linear / kHidden; + const uint32_t col = + linear - + static_cast(padding_row) * + kHidden; + const uint32_t pool_row = + padded_pool_blocks * BLOCK_M + + num_tokens + padding_row; + const uint64_t offset = + static_cast(pool_row) * + kHidden + + col; + if constexpr (kDoReverseDispatch) { + grad_y_unweighted_output[offset] = + cutlass::bfloat16_t(0.0f); + if constexpr (!kXPrepared) { + x_pool_output[offset] = + cutlass::bfloat16_t(0.0f); + } + } + if constexpr (kWriteWeighted) { + grad_y_weighted_output[offset] = + cutlass::bfloat16_t(0.0f); + } + } + for (uint32_t padding_row = + blockIdx.x * kNumThreads + + threadIdx.x; + padding_row < num_padding_rows; + padding_row += + kNumSMs * kNumThreads) { + const uint32_t pool_row = + padded_pool_blocks * BLOCK_M + + num_tokens + padding_row; + if constexpr (kDoReverseDispatch) + route_weights_output[pool_row] = 0.0f; + if constexpr (kComputeRouteDot) + grad_route_output[pool_row] = 0.0f; + } + padded_pool_blocks += num_blocks; + } + if constexpr (kSynchronizeAfterDispatch) { + // Every rank may reuse its local symmetric grad-y plane as the + // direct-write grad-x destination as soon as this producer returns. + // Publish completion only after all peers have finished their remote + // reads; an entry-only barrier permits checkpoint/replay rank skew to + // corrupt those in-flight pulls. + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumThreads, 1, 72>( + backward_workspace, + backward_sym_buffer, + blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }); + } + // Rows beyond the final padded expert block are capacity only. No + // downstream kernel addresses them, so clearing that high-water tail + // wastes bandwidth and grows with the cached pool margin. +#endif +} + +// Publish the two L1 side-LoRA dgrad branches without a [pool, hidden] +// temporary. The base wave has already populated the fixed local/remote +// combine slots; every valid expert-pool row owns exactly one such slot, so +// the split-additive BF16 update is race-free and needs no atomics. +template < + uint32_t kHidden, uint32_t kNumExperts, uint32_t BLOCK_M, + uint32_t kNumRanks, uint32_t kNumSMs, bool kWriteGradXPool, + bool kDirectRemoteGradX> +CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_grad_x_impl( + const int* expert_counts, + cutlass::bfloat16_t* grad_x_pool, + const layout::TokenSrcMetadata* token_src_metadata, + cutlass::bfloat16_t* combine_buffer, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ layout::Workspace workspace, + const uint32_t num_pool_rows, + const uint32_t num_topk) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) + using bf16 = cutlass::bfloat16_t; + constexpr uint32_t kOutputTileN = 128; + + if constexpr (kDirectRemoteGradX) { + // The direct-write planes previously held grad-y. Valid routes are + // overwritten below, but masked top-k slots have no expert owner and + // would otherwise leak that stale value into the final combine. + // Clear every local fixed-slot plane, then synchronize before any + // peer publishes its valid routes into this rank's symmetric plane. + constexpr uint32_t kVectorElems = sizeof(uint4) / sizeof(bf16); + const uint64_t num_vectors = + static_cast(workspace.num_max_tokens_per_rank) * + num_topk * kHidden / kVectorElems; + for (uint64_t vector_idx = + static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + vector_idx < num_vectors; + vector_idx += + static_cast(gridDim.x) * blockDim.x) { + reinterpret_cast(combine_buffer)[vector_idx] = + make_uint4(0, 0, 0, 0); + } + comm::nvlink_barrier( + workspace, sym_buffer, blockIdx.x, threadIdx.x, + []() { __syncthreads(); }); + } + + uint32_t pool_block_offset = 0; + uint32_t global_tile_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; expert_idx < kNumExperts; + ++expert_idx) { + const uint32_t num_tokens = static_cast( + __ldg(expert_counts + expert_idx)); + const uint32_t num_m_blocks = + math::ceil_div(num_tokens, BLOCK_M); + const uint32_t num_n_blocks = kHidden / kOutputTileN; + const uint32_t num_tiles = num_m_blocks * num_n_blocks; + uint32_t first_global_tile = blockIdx.x; + if (first_global_tile < global_tile_offset) { + first_global_tile += math::ceil_div( + global_tile_offset - first_global_tile, + static_cast(gridDim.x)) * gridDim.x; + } + for (uint32_t local_tile = + first_global_tile < global_tile_offset + num_tiles + ? first_global_tile - global_tile_offset + : num_tiles; + local_tile < num_tiles; + local_tile += gridDim.x) { + const uint32_t m_block = local_tile / num_n_blocks; + const uint32_t n_block = local_tile - m_block * num_n_blocks; + constexpr uint32_t kVectorElems = sizeof(uint4) / sizeof(bf16); + constexpr uint32_t kVectorsPerTile = + kOutputTileN / kVectorElems; + for (uint32_t linear = threadIdx.x; + linear < BLOCK_M * kVectorsPerTile; + linear += blockDim.x) { + const uint32_t local_m = linear / kVectorsPerTile; + const uint32_t local_vector = + linear - local_m * kVectorsPerTile; + const uint32_t expert_row = m_block * BLOCK_M + local_m; + if (expert_row >= num_tokens) + continue; + const uint32_t pool_row = + pool_block_offset + expert_row; + if (pool_row >= num_pool_rows) + continue; + const uint32_t hidden_col = n_block * kOutputTileN + + local_vector * kVectorElems; + const auto value = *reinterpret_cast( + grad_x_pool + + static_cast(pool_row) * kHidden + hidden_col); + if constexpr (kDirectRemoteGradX) { + const auto metadata = token_src_metadata[pool_row]; + auto* local_dst = combine_buffer + + ((static_cast(metadata.topk_idx) * + workspace.num_max_tokens_per_rank + + metadata.token_idx) * + kHidden + + hidden_col); + *reinterpret_cast( + sym_buffer.map(local_dst, metadata.rank_idx)) = value; + } + } + } + global_tile_offset += num_tiles; + pool_block_offset += num_m_blocks * BLOCK_M; + } + if constexpr (kNumRanks > 1 && kDirectRemoteGradX) { + comm::nvlink_barrier( + workspace, sym_buffer, blockIdx.x, threadIdx.x, + []() { __syncthreads(); }); + } +#endif +} + +template +CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_axpy2_impl( + cutlass::bfloat16_t* dst, + const cutlass::bfloat16_t* src1, + const cutlass::bfloat16_t* src3, + const uint64_t num_elements, + const float scale) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) + using bf16 = cutlass::bfloat16_t; + union alignas(16) BF16x8 { + uint4 packed; + bf16 element[8]; + }; + constexpr uint64_t kVectorElems = 8; + const uint64_t num_vectors = num_elements / kVectorElems; + for (uint64_t vector_idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + vector_idx < num_vectors; + vector_idx += static_cast(kNumSMs) * blockDim.x) { + BF16x8 d, s1, s3; + d.packed = reinterpret_cast(dst)[vector_idx]; + s1.packed = reinterpret_cast(src1)[vector_idx]; + s3.packed = reinterpret_cast(src3)[vector_idx]; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const bf16 side1(scale * static_cast(s1.element[i])); + const bf16 side3(scale * static_cast(s3.element[i])); + d.element[i] = bf16( + static_cast(d.element[i]) + + static_cast(side1) + static_cast(side3)); + } + reinterpret_cast(dst)[vector_idx] = d.packed; + } + for (uint64_t idx = num_vectors * kVectorElems + + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < num_elements; + idx += static_cast(kNumSMs) * blockDim.x) { + const bf16 side1(scale * static_cast(src1[idx])); + const bf16 side3(scale * static_cast(src3[idx])); + dst[idx] = bf16(static_cast(dst[idx]) + + static_cast(side1) + + static_cast(side3)); + } +#endif +} + +template +CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_scale_grads_impl( + const __grid_constant__ MegaMoESideLoraBackwardParams side_lora) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) + using bf16 = cutlass::bfloat16_t; + constexpr uint64_t kRank = 128; + // A1, A3, and B2 are shared across experts by the adapter layout. + constexpr uint64_t kHiddenElems = + static_cast(kHidden) * kRank; + constexpr uint64_t kIntermediateElems = + static_cast(kNumExperts) * kIntermediateHidden * kRank; + constexpr uint64_t kVectorElems = 8; + constexpr uint64_t kHiddenVectors = kHiddenElems / kVectorElems; + constexpr uint64_t kIntermediateVectors = + kIntermediateElems / kVectorElems; + constexpr uint64_t kScaledVectors = + 3 * kHiddenVectors + 3 * kIntermediateVectors; + union alignas(16) BF16x8 { + uint4 packed; + bf16 element[8]; + }; + for (uint64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < kScaledVectors; + linear += static_cast(gridDim.x) * blockDim.x) { + uint4* ptr; + uint64_t offset = linear; + if (offset < kHiddenVectors) { + ptr = reinterpret_cast(side_lora.grad_a1) + offset; + } else if ((offset -= kHiddenVectors) < kIntermediateVectors) { + ptr = reinterpret_cast(side_lora.grad_b1) + offset; + } else if ((offset -= kIntermediateVectors) < kHiddenVectors) { + ptr = reinterpret_cast(side_lora.grad_a3) + offset; + } else if ((offset -= kHiddenVectors) < kIntermediateVectors) { + ptr = reinterpret_cast(side_lora.grad_b3) + offset; + } else if ((offset -= kIntermediateVectors) < kIntermediateVectors) { + ptr = reinterpret_cast(side_lora.grad_a2) + offset; + } else { + offset -= kIntermediateVectors; + ptr = reinterpret_cast(side_lora.grad_b2) + offset; + } + BF16x8 values; + values.packed = *ptr; + #pragma unroll + for (int i = 0; i < 8; ++i) { + values.element[i] = bf16( + static_cast(values.element[i]) * side_lora.scale); + } + *ptr = values.packed; + } +#endif +} + +// Forward saves only logical expert rows. Clear the already-allocated padded +// tails before K-grouped adapter wgrads so they contribute exact zeros. +template +CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_clear_padding_impl( + const int* expert_counts, + cutlass::bfloat16_t* saved_h, + cutlass::bfloat16_t* q13, + cutlass::bfloat16_t* q2) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) + constexpr uint32_t kRankStorage = 3 * 128; + uint32_t pool_offset = 0; + for (uint32_t expert = 0; expert < kNumExperts; ++expert) { + const uint32_t count = static_cast( + __ldg(expert_counts + expert)); + const uint32_t capacity = math::ceil_div(count, BLOCK_M) * BLOCK_M; + const uint32_t padding = capacity - count; + const uint64_t elements = static_cast(padding) * + (kIntermediateHidden + kRankStorage); + for (uint64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < elements; + linear += static_cast(kNumSMs) * blockDim.x) { + const uint32_t padding_row = linear / + (kIntermediateHidden + kRankStorage); + const uint32_t column = linear - + static_cast(padding_row) * + (kIntermediateHidden + kRankStorage); + const uint32_t pool_row = pool_offset + count + padding_row; + if (column < kIntermediateHidden) { + saved_h[static_cast(pool_row) * + kIntermediateHidden + column] = + cutlass::bfloat16_t(0.0f); + } else { + const uint32_t rank_column = column - kIntermediateHidden; + if (rank_column < 256) { + q13[static_cast(pool_row) * 256 + + rank_column] = cutlass::bfloat16_t(0.0f); + } else { + q2[static_cast(pool_row) * 128 + + rank_column - 256] = cutlass::bfloat16_t(0.0f); + } + } + } + pool_offset += capacity; + } +#endif +} + +// Production MegaMoE backward wave. This persistent kernel consumes the +// forward kernel's block-padded expert pool directly and replays gate and up +// together as one W13 FP8xFP4 mainloop before computing the retained dgrads. +template < + uint32_t kHidden, uint32_t kIntermediateHidden, + uint32_t kNumExperts, + uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, + uint32_t SF_BLOCK_M, uint32_t SF_BLOCK_N, + uint32_t kNumStages, + uint32_t kNumSMs, + uint32_t kNumRanks = 1, + bool kCompileW13Dgrad = true, + bool kBF16Mode = false, + ActivationType kActivationType = ActivationType::SwiGLU, + bool kFastMath = false, + RouteWeightMode kRouteWeightMode = RouteWeightMode::PreDown, + CombineOrderMode kCombineOrderMode = CombineOrderMode::FixedTopK, + bool kInputsPrepared = false, + bool kDispatchInputsPrepared = false, + bool kDirectRemoteGradX = false, + bool kWriteGradXPool = true, + bool kClearWgradPadding = false, + bool kComputeRouteGrad = false, + bool kTraceKernel = false, + bool kVectorizedGradXStore = false, + bool kWideGradXStore = false, + bool kGateUpPrepared = false, + uint32_t kNumNonEpilogueThreads = 128, + uint32_t kNumEpilogueThreads = 128, + uint32_t kNumThreads = + kNumNonEpilogueThreads + kNumEpilogueThreads + + 768> +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void +sm100_bf16_mega_moe_side_lora_backward_wave_impl( + const int* expert_counts, + const __grid_constant__ layout::SymBuffer backward_sym_buffer, + const __grid_constant__ layout::Workspace backward_workspace, + const cutlass::bfloat16_t* backward_grad_y, + const cutlass::bfloat16_t* backward_x, + const float* backward_topk_weights, + float* backward_grad_route, + const layout::TokenSrcMetadata* token_src_metadata, + const uint32_t num_topk, + const uint32_t num_pool_rows, + const uint32_t num_acts_rows, + const uint32_t acts_sf_stride, + const __grid_constant__ cute::TmaDescriptor tensor_map_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_grad_ye, + const __grid_constant__ cute::TmaDescriptor tensor_map_w2_dequant, + const __grid_constant__ cute::TmaDescriptor tensor_map_w2_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_w2_scales, + const __grid_constant__ cute::TmaDescriptor tensor_map_w13_dequant, + const __grid_constant__ cute::TmaDescriptor tensor_map_w13_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_w13_scales, + const __grid_constant__ cute::TmaDescriptor tensor_map_grad_gate_up, + const cutlass::float_e4m3_t* acts_ptr, + const uint32_t* acts_sf_ptr, + const int8_t* w2_weights, + const float* w2_scales, + cutlass::bfloat16_t* w2_dequant_scratch, + const int8_t* w13_weights, + const float* w13_scales, + cutlass::bfloat16_t* w13_dequant_scratch, + const cutlass::bfloat16_t* gate_up_output, + cutlass::bfloat16_t* grad_ye_output, + cutlass::bfloat16_t* grad_y_unweighted_output, + cutlass::bfloat16_t* route_weights, + float* route_weights_fp32, + cutlass::bfloat16_t* grad_h_output, + cutlass::bfloat16_t* grad_gate_up_output, + cutlass::bfloat16_t* h_act_output, + cutlass::bfloat16_t* h_weighted_output, + cutlass::bfloat16_t* x_pool_output, + cutlass::bfloat16_t* grad_x_pool_output, + const cutlass::bfloat16_t* down_unweighted_output, + float* grad_route_output, + uint32_t* weight_tile_states, + const uint32_t launch_epoch, + const float activation_limit, + const __grid_constant__ MegaMoESideLoraBackwardParams side_lora, + uint64_t* kernel_trace) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Allocator = cute::TMEM::Allocator2Sm; + using a_dtype_t = cutlass::float_e4m3_t; + using b_dtype_t = cutlass::detail::float_e2m1_unpacksmem_t; + using cd_dtype_t = cutlass::bfloat16_t; + using dgrad_b_dtype_t = cutlass::bfloat16_t; + + constexpr uint32_t kNumEpilogueStages = 2; + constexpr uint32_t kNumTMAStoreStages = 2; + constexpr uint32_t kNumDispatchThreads = + kNumRanks > 1 ? 128 : 0; + constexpr uint32_t kNumDispatchWarps = + kNumDispatchThreads / 32; + constexpr uint32_t kDispatchWarpStart = + (kNumNonEpilogueThreads + kNumEpilogueThreads) / 32; + constexpr uint32_t kNumDgradEpilogueThreads = + kNumThreads - kNumNonEpilogueThreads; + constexpr uint32_t kGranK = 32; + constexpr uint32_t kNumUTCCPAlignedElems = 128; + constexpr uint32_t kNumBlockNs = (2 * kIntermediateHidden) / BLOCK_N; + constexpr uint32_t kNumDgradBlockNs = kIntermediateHidden / BLOCK_N; + constexpr uint32_t kNumW13DgradBlockNs = kHidden / BLOCK_N; + constexpr uint32_t kNumW13DgradSplits = + kBF16Mode ? 2 : 1; + constexpr uint32_t LAYOUT_AD_M = 128; + constexpr uint32_t UMMA_M = LAYOUT_AD_M * 2; + constexpr uint32_t UMMA_N = BLOCK_M; + constexpr uint32_t UMMA_K = 32; + constexpr uint32_t DGRAD_BLOCK_K = 64; + constexpr uint32_t DGRAD_UMMA_K = 16; + constexpr uint32_t kNumW2WeightTileStates = + kNumExperts * (kHidden / DGRAD_BLOCK_K) * + kNumDgradBlockNs; + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / 2; + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; + constexpr uint32_t STORE_BLOCK_M = 16; + constexpr uint32_t STORE_BLOCK_N = BLOCK_N; + constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); + constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t); + constexpr uint32_t kSwizzleCDMode = 128; + + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 128, "Invalid producer thread count"); + DG_STATIC_ASSERT(kNumEpilogueThreads == 128, "Invalid epilogue thread count"); + DG_STATIC_ASSERT(kNumRanks == 1 || kNumDispatchThreads == 128, + "Invalid backward dispatch thread count"); + DG_STATIC_ASSERT(BLOCK_M % 16 == 0 && BLOCK_N == 128 && BLOCK_K == 128, + "Invalid backward wave tile"); + DG_STATIC_ASSERT(kNumBlockNs % 2 == 0, "Cluster peers must receive adjacent N blocks"); + DG_STATIC_ASSERT(kNumDgradBlockNs % 2 == 0, + "Dgrad cluster peers must receive adjacent N blocks"); + DG_STATIC_ASSERT(kNumW13DgradBlockNs % 2 == 0, + "W13 dgrad cluster peers must receive adjacent N blocks"); + DG_STATIC_ASSERT(SF_BLOCK_M == math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems), + "Invalid SFA block"); + DG_STATIC_ASSERT(SF_BLOCK_N == BLOCK_N, "Invalid SFB block"); + DG_STATIC_ASSERT(kHidden % BLOCK_K == 0, "Invalid hidden size"); + DG_STATIC_ASSERT(kNumSMs % 2 == 0, "2-CTA clusters require an even SM count"); + + constexpr uint32_t kNumW13WeightTileStates = + kNumExperts * + ((2 * kIntermediateHidden) / DGRAD_BLOCK_K) * + kNumW13DgradBlockNs; + auto* phase_count = + weight_tile_states + kNumW2WeightTileStates + + kNumW13WeightTileStates; + auto* phase_sense = phase_count + 1; + constexpr uint32_t kTraceSiteCount = 23; + constexpr uint32_t kTraceValueCount = 5; + constexpr uint32_t kTraceBeginCycle = 0; + constexpr uint32_t kTraceEndCycle = 1; + constexpr uint32_t kTraceBeginGlobalNs = 2; + constexpr uint32_t kTraceEndGlobalNs = 3; + constexpr uint32_t kTraceSM = 4; + const auto globaltimer = [] { + uint64_t value; + asm volatile( + "mov.u64 %0, %%globaltimer;" : "=l"(value)); + return value; + }; + const auto trace_begin = [&](const uint32_t site) { + if constexpr (kTraceKernel) { + if (threadIdx.x == 0) { + auto* values = + kernel_trace + + (static_cast(site) * kNumSMs + + blockIdx.x) * + kTraceValueCount; + values[kTraceBeginCycle] = clock64(); + values[kTraceBeginGlobalNs] = globaltimer(); + values[kTraceSM] = ptx::get_sm_idx(); + } + } + }; + const auto trace_end = [&](const uint32_t site) { + if constexpr (kTraceKernel) { + if (threadIdx.x == 0) { + auto* values = + kernel_trace + + (static_cast(site) * kNumSMs + + blockIdx.x) * + kTraceValueCount; + values[kTraceEndCycle] = clock64(); + values[kTraceEndGlobalNs] = globaltimer(); + } + } + }; + if constexpr (kTraceKernel) { + DG_STATIC_ASSERT( + kTraceSiteCount == 23, + "Update the host trace-site schema with the kernel"); + trace_begin(0); + } + const auto full_grid_phase_barrier = + [&](const uint32_t trace_site) { + trace_begin(trace_site); + if (threadIdx.x == 0) { + const uint32_t old_sense = + atomicAdd(phase_sense, 0u); + __threadfence(); + const uint32_t ticket = + atomicAdd(phase_count, 1u); + if (ticket == kNumSMs - 1) { + atomicExch(phase_count, 0u); + __threadfence(); + atomicAdd(phase_sense, 1u); + } else { + while (ptx::ld_acq(phase_sense) == + old_sense) { + } + } + } + __syncthreads(); + trace_end(trace_site); + }; + + const bool is_leader_cta = cute::block_rank_in_cluster() == 0; + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + + if (warp_idx == 0) { + cute::prefetch_tma_descriptor(&tensor_map_grad_ye); + cute::prefetch_tma_descriptor(&tensor_map_w2_dequant); + cute::prefetch_tma_descriptor(&tensor_map_w13_dequant); + cute::prefetch_tma_descriptor(&tensor_map_grad_gate_up); + if constexpr (!kBF16Mode) { + cute::prefetch_tma_descriptor(&tensor_map_acts); + cute::prefetch_tma_descriptor(&tensor_map_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_weights); + cute::prefetch_tma_descriptor(&tensor_map_weights_sf); + cute::prefetch_tma_descriptor(&tensor_map_output); + cute::prefetch_tma_descriptor(&tensor_map_w2_weights); + cute::prefetch_tma_descriptor(&tensor_map_w2_scales); + cute::prefetch_tma_descriptor(&tensor_map_w13_weights); + cute::prefetch_tma_descriptor(&tensor_map_w13_scales); + } + } + + constexpr uint32_t SMEM_CD_SIZE_PER_STAGE = + STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd_dtype_t); + constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_SIZE_PER_STAGE * kNumTMAStoreStages; + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = + LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = + LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(uint32_t); + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * sizeof(uint32_t); + constexpr uint32_t SMEM_DISPATCH_SIZE = + kNumDispatchWarps * kHidden * sizeof(cd_dtype_t); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto* smem_gemm_base = smem_buffer + SMEM_DISPATCH_SIZE; + auto smem_cd = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast( + smem_gemm_base + i * SMEM_CD_SIZE_PER_STAGE); + }); + auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast( + smem_gemm_base + SMEM_CD_SIZE + + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast( + smem_gemm_base + SMEM_CD_SIZE + + kNumStages * SMEM_A_SIZE_PER_STAGE + + i * SMEM_B_SIZE_PER_STAGE); + }); + // The dgrad phase aliases the recompute mainloop storage exactly: + // FP8 A [BLOCK_M/2, 128] == BF16 A [BLOCK_M/2, 64] + // packed-FP4 B [128, 128] == BF16 B [128, 64]. + // W2 is dequantized and transposed directly into the latter by a producer + // warp, so no persistent weight copy or host-side packing is needed. + auto smem_dgrad_a = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast( + smem_gemm_base + SMEM_CD_SIZE + + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_dgrad_b = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast( + smem_gemm_base + SMEM_CD_SIZE + + kNumStages * SMEM_A_SIZE_PER_STAGE + + i * SMEM_B_SIZE_PER_STAGE); + }); + DG_STATIC_ASSERT( + LOAD_BLOCK_M * DGRAD_BLOCK_K * sizeof(cd_dtype_t) == + SMEM_A_SIZE_PER_STAGE, + "Dgrad A alias size mismatch"); + DG_STATIC_ASSERT( + LOAD_BLOCK_N * DGRAD_BLOCK_K * sizeof(dgrad_b_dtype_t) == + SMEM_B_SIZE_PER_STAGE, + "Dgrad B alias size mismatch"); + auto sf_start_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); + }); + + auto barrier_start_ptr = reinterpret_cast(smem_sfb[kNumStages]); + auto full_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + i; }); + auto empty_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); + auto tmem_full_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + 2 * kNumStages + i; + }); + auto tmem_empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + 2 * kNumStages + kNumEpilogueStages + i; + }); + auto dispatch_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + 2 * kNumStages + + 2 * kNumEpilogueStages + i; + }); + auto tmem_ptr_in_smem = reinterpret_cast( + barrier_start_ptr + 2 * kNumStages + + 2 * kNumEpilogueStages + kNumDispatchWarps); + + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; + constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; + constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; + constexpr uint32_t kNumTmemCols = + utils::get_num_aligned_tmem_cols< + kNumAccumTmemCols + kNumSFATmemCols + kNumSFBTmemCols>(); + constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFB = + kNumAccumTmemCols + kNumSFATmemCols; + DG_STATIC_ASSERT(kNumTmemCols <= 512, "Backward recompute exceeds TMEM"); + + if constexpr (!kBF16Mode) { + // Dequantize W2 exactly once per launch into an ephemeral + // [expert, dim, H] BF16 workspace. Keeping the source orientation + // makes both packed-FP4 reads and BF16 writes coalesced; dgrad consumes + // it as an MN-major transposed operand. + constexpr uint32_t kDequantTileK = 256; + constexpr uint32_t kDequantTileN = LOAD_BLOCK_N; + constexpr uint32_t kDequantPairsPerTile = + kDequantTileK * kDequantTileN / 2; + constexpr uint32_t kDequantSFsPerK = + kDequantTileN / 32; + constexpr uint32_t kDequantSFsPerTile = + kDequantTileK * kDequantSFsPerK; + constexpr uint32_t kDequantWeightBytes = + kDequantTileK * (kDequantTileN / 2); + constexpr uint32_t kDequantScaleBytes = + kDequantSFsPerTile * sizeof(float); + constexpr uint32_t kNumDequantKTiles = + kHidden / kDequantTileK; + constexpr uint32_t kNumDequantNTiles = + kIntermediateHidden / kDequantTileN; + constexpr uint32_t kNumDequantTiles = + kNumExperts * kNumDequantKTiles * + kNumDequantNTiles; + auto* dequant_weights = + reinterpret_cast(smem_buffer); + auto* dequant_scales = + reinterpret_cast( + smem_buffer + kDequantWeightBytes); + auto* dequant_scale_half2 = + reinterpret_cast( + smem_buffer + kDequantWeightBytes + + kDequantScaleBytes); + comm::cluster_sync_with_relaxed_arrive(); + if (warp_idx == 0 && cute::elect_one_sync()) { + full_barriers[0]->init(1); + if constexpr (kCompileW13Dgrad) + full_barriers[1]->init(1); + cutlass::arch::fence_barrier_init(); + } + comm::cluster_sync_with_relaxed_arrive(); + uint32_t dequant_phase = 0; + + for (uint32_t tile_idx = blockIdx.x; + tile_idx < kNumDequantTiles; + tile_idx += kNumSMs) { + const uint32_t n_tile_idx = + tile_idx % kNumDequantNTiles; + const uint32_t k_expert_tile_idx = + tile_idx / kNumDequantNTiles; + const uint32_t k_tile_idx = + k_expert_tile_idx % kNumDequantKTiles; + const uint32_t expert_idx = + k_expert_tile_idx / kNumDequantKTiles; + const uint32_t global_k_base = + k_tile_idx * kDequantTileK; + const uint32_t global_n_base = + n_tile_idx * kDequantTileN; + + if (warp_idx == 0 && cute::elect_one_sync()) { + tma::copy< + kDequantTileN / 2, kDequantTileK, 0, + int8_t>( + &tensor_map_w2_weights, + full_barriers[0], dequant_weights, + global_n_base / 2, + expert_idx * kHidden + global_k_base); + tma::copy< + kDequantSFsPerK, kDequantTileK, 0, + float>( + &tensor_map_w2_scales, + full_barriers[0], dequant_scales, + global_n_base / 32, + expert_idx * kHidden + global_k_base); + full_barriers[0]->arrive_and_expect_tx( + kDequantWeightBytes + + kDequantScaleBytes); + } + full_barriers[0]->wait(dequant_phase); + __syncthreads(); + + for (uint32_t scale_idx = threadIdx.x; + scale_idx < kDequantSFsPerTile; + scale_idx += kNumThreads) { + const auto scale_half2 = + __float2half2_rn( + dequant_scales[scale_idx]); + dequant_scale_half2[scale_idx] = + *reinterpret_cast( + &scale_half2); + } + __syncthreads(); + + for (uint32_t pair_idx = threadIdx.x; + pair_idx < kDequantPairsPerTile; + pair_idx += kNumThreads) { + const uint32_t local_k = + pair_idx / (kDequantTileN / 2); + const uint32_t local_n_pair = + pair_idx % (kDequantTileN / 2); + const uint32_t global_k = + global_k_base + local_k; + const uint32_t global_n_pair = + global_n_base / 2 + local_n_pair; + const uint8_t packed = + static_cast( + dequant_weights[ + local_k * + (kDequantTileN / 2) + + local_n_pair]); + uint32_t fp16x2; + asm volatile( + "{\n" + ".reg .b8 fp4;\n" + ".reg .b8 unused1, unused2, unused3;\n" + "mov.b32 {fp4, unused1, unused2, unused3}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4;\n" + "}\n" + : "=r"(fp16x2) + : "r"(static_cast(packed))); + auto value_pair = + *reinterpret_cast<__half2*>(&fp16x2); + const uint32_t scale_half2_bits = + dequant_scale_half2[ + local_k * kDequantSFsPerK + + (local_n_pair * 2) / 32]; + const auto scale_half2 = + *reinterpret_cast( + &scale_half2_bits); + value_pair = + __hmul2(value_pair, scale_half2); + const auto value_pair_bf16 = + __float22bfloat162_rn( + __half22float2(value_pair)); + const uint32_t scaled_pair = + *reinterpret_cast( + &value_pair_bf16); + *reinterpret_cast( + w2_dequant_scratch + + (static_cast(expert_idx) * + kHidden + + global_k) * + kIntermediateHidden + + global_n_pair * 2) = + scaled_pair; + } + __syncthreads(); + if (threadIdx.x < + kDequantTileK / DGRAD_BLOCK_K) { + const uint32_t dgrad_k_block_idx = + k_tile_idx * + (kDequantTileK / + DGRAD_BLOCK_K) + + threadIdx.x; + const uint32_t weight_tile_idx = + (expert_idx * + (kHidden / DGRAD_BLOCK_K) + + dgrad_k_block_idx) * + kNumDgradBlockNs + + n_tile_idx; + asm volatile( + "st.release.gpu.global.u32 [%0], %1;" + :: "l"(weight_tile_states + + weight_tile_idx), + "r"(launch_epoch) + : "memory"); + } + __syncthreads(); + dequant_phase ^= 1; + } + + if constexpr (kCompileW13Dgrad) { + constexpr uint32_t kW13DequantTileK = 256; + constexpr uint32_t kW13DequantTileN = LOAD_BLOCK_N; + constexpr uint32_t kW13DequantPairsPerTile = + kW13DequantTileK * kW13DequantTileN / 2; + constexpr uint32_t kW13DequantSFsPerK = + kW13DequantTileN / 32; + constexpr uint32_t kW13DequantSFsPerTile = + kW13DequantTileK * kW13DequantSFsPerK; + constexpr uint32_t kW13DequantWeightBytes = + kW13DequantTileK * (kW13DequantTileN / 2); + constexpr uint32_t kW13DequantScaleBytes = + kW13DequantSFsPerTile * sizeof(float); + constexpr uint32_t kNumW13DequantKTiles = + (2 * kIntermediateHidden) / kW13DequantTileK; + constexpr uint32_t kNumW13DequantNTiles = + kHidden / kW13DequantTileN; + constexpr uint32_t kNumW13DequantTiles = + kNumExperts * kNumW13DequantKTiles * + kNumW13DequantNTiles; + const uint32_t w13_launch_epoch = + launch_epoch ^ 0x80000000u; + uint32_t w13_dequant_phase = 0; + + for (uint32_t tile_idx = blockIdx.x; + tile_idx < kNumW13DequantTiles; + tile_idx += kNumSMs) { + const uint32_t n_tile_idx = + tile_idx % kNumW13DequantNTiles; + const uint32_t k_expert_tile_idx = + tile_idx / kNumW13DequantNTiles; + const uint32_t k_tile_idx = + k_expert_tile_idx % kNumW13DequantKTiles; + const uint32_t expert_idx = + k_expert_tile_idx / kNumW13DequantKTiles; + const uint32_t global_k_base = + k_tile_idx * kW13DequantTileK; + const uint32_t global_n_base = + n_tile_idx * kW13DequantTileN; + + if (warp_idx == 0 && cute::elect_one_sync()) { + tma::copy< + kW13DequantTileN / 2, + kW13DequantTileK, 0, int8_t>( + &tensor_map_w13_weights, + full_barriers[1], + dequant_weights, + global_n_base / 2, + expert_idx * + (2 * kIntermediateHidden) + + global_k_base); + tma::copy< + kW13DequantSFsPerK, + kW13DequantTileK, 0, float>( + &tensor_map_w13_scales, + full_barriers[1], + dequant_scales, + global_n_base / 32, + expert_idx * + (2 * kIntermediateHidden) + + global_k_base); + full_barriers[1]->arrive_and_expect_tx( + kW13DequantWeightBytes + + kW13DequantScaleBytes); + } + full_barriers[1]->wait( + w13_dequant_phase); + __syncthreads(); + + for (uint32_t scale_idx = threadIdx.x; + scale_idx < kW13DequantSFsPerTile; + scale_idx += kNumThreads) { + const auto scale_half2 = + __float2half2_rn( + dequant_scales[scale_idx]); + dequant_scale_half2[scale_idx] = + *reinterpret_cast< + const uint32_t*>( + &scale_half2); + } + __syncthreads(); + + for (uint32_t pair_idx = threadIdx.x; + pair_idx < + kW13DequantPairsPerTile; + pair_idx += kNumThreads) { + const uint32_t local_k = + pair_idx / + (kW13DequantTileN / 2); + const uint32_t local_n_pair = + pair_idx % + (kW13DequantTileN / 2); + const uint32_t global_k = + global_k_base + local_k; + const uint32_t global_n_pair = + global_n_base / 2 + + local_n_pair; + const uint8_t packed = + static_cast( + dequant_weights[ + local_k * + (kW13DequantTileN / 2) + + local_n_pair]); + uint32_t fp16x2; + asm volatile( + "{\n" + ".reg .b8 fp4;\n" + ".reg .b8 unused1, unused2, unused3;\n" + "mov.b32 {fp4, unused1, unused2, unused3}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4;\n" + "}\n" + : "=r"(fp16x2) + : "r"( + static_cast( + packed))); + auto value_pair = + *reinterpret_cast<__half2*>( + &fp16x2); + const uint32_t scale_half2_bits = + dequant_scale_half2[ + local_k * + kW13DequantSFsPerK + + (local_n_pair * 2) / 32]; + value_pair = __hmul2( + value_pair, + *reinterpret_cast< + const __half2*>( + &scale_half2_bits)); + const auto value_pair_bf16 = + __float22bfloat162_rn( + __half22float2( + value_pair)); + *reinterpret_cast( + w13_dequant_scratch + + (static_cast( + expert_idx) * + (2 * + kIntermediateHidden) + + global_k) * + kHidden + + global_n_pair * 2) = + *reinterpret_cast< + const uint32_t*>( + &value_pair_bf16); + } + __syncthreads(); + + if (threadIdx.x < + kW13DequantTileK / + DGRAD_BLOCK_K) { + const uint32_t + dgrad_k_block_idx = + k_tile_idx * + (kW13DequantTileK / + DGRAD_BLOCK_K) + + threadIdx.x; + const uint32_t weight_tile_idx = + (expert_idx * + ((2 * + kIntermediateHidden) / + DGRAD_BLOCK_K) + + dgrad_k_block_idx) * + kNumW13DgradBlockNs + + n_tile_idx; + asm volatile( + "st.release.gpu.global.u32 [%0], %1;" + :: "l"(weight_tile_states + + kNumW2WeightTileStates + + weight_tile_idx), + "r"(w13_launch_epoch) + : "memory"); + } + __syncthreads(); + w13_dequant_phase ^= 1; + } + } + } + trace_begin(1); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(1); + if (warp_idx == 0 && cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) { + full_barriers[i]->init(4); + empty_barriers[i]->init(1); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueStages; ++i) { + tmem_full_barriers[i]->init(1); + tmem_empty_barriers[i]->init(2 * kNumEpilogueThreads); + } + #pragma unroll + for (uint32_t i = 0; i < kNumDispatchWarps; ++i) + dispatch_barriers[i]->init(1); + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 1) { + Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + trace_begin(2); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(2); + + // Every role walks this deterministic schedule independently. Pool offsets + // are prefixes of ceil(count/BLOCK_M), matching the forward MegaMoE layout. + const auto for_each_block = [&](const auto& func) { + uint32_t next_assigned_block = blockIdx.x; + uint32_t global_block = 0; + uint32_t pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast(__ldg(expert_counts + expert_idx)); + const uint32_t num_m_blocks = math::ceil_div(num_tokens, BLOCK_M); + const uint32_t expert_blocks = num_m_blocks * kNumBlockNs; + const uint32_t expert_end = global_block + expert_blocks; + + while (next_assigned_block < global_block) + next_assigned_block += kNumSMs; + while (next_assigned_block < expert_end) { + const uint32_t local_block = + next_assigned_block - global_block; + const uint32_t m_block_idx = local_block / kNumBlockNs; + const uint32_t n_block_idx = + local_block - m_block_idx * kNumBlockNs; + const uint32_t valid_m = cute::min( + num_tokens - m_block_idx * BLOCK_M, BLOCK_M); + func(expert_idx, pool_block_offset, m_block_idx, + n_block_idx, valid_m); + next_assigned_block += kNumSMs; + } + global_block = expert_end; + pool_block_offset += num_m_blocks; + } + }; + + uint32_t stage_idx = 0; + uint32_t phase = 0; + const auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + constexpr uint32_t kNumProducerRegisters = 40; + constexpr uint32_t kNumEpilogueRegisters = 208; + + if constexpr (!kBF16Mode && !kGateUpPrepared) { + if (warp_idx == 0) { + cutlass::arch::warpgroup_reg_dealloc(); + for_each_block([&](const uint32_t&, const uint32_t& pool_block_offset, + const uint32_t& m_block_idx, const uint32_t&, + const uint32_t& valid_m) { + const uint32_t pool_block_idx = + pool_block_offset + m_block_idx; + #pragma unroll + for (uint32_t k_block_idx = 0; + k_block_idx < kHidden / BLOCK_K; + advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + uint32_t m_idx = pool_block_idx * BLOCK_M; + if (!is_leader_cta) + m_idx += math::align(valid_m, 16u) / 2; + if (cute::elect_one_sync()) { + tma::copy( + &tensor_map_acts, full_barriers[stage_idx], + smem_a[stage_idx], k_block_idx * BLOCK_K, m_idx, 2); + tma::copy( + &tensor_map_acts_sf, full_barriers[stage_idx], + smem_sfa[stage_idx], + pool_block_idx * SF_BLOCK_M, k_block_idx, 2); + if (is_leader_cta) { + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE * 2 + + SF_BLOCK_M * sizeof(uint32_t) * 2); + } else { + full_barriers[stage_idx]->arrive(0u); + } + } + __syncwarp(); + } + }); + } else if (warp_idx == 1) { + cutlass::arch::warpgroup_reg_dealloc(); + for_each_block([&](const uint32_t& expert_idx, const uint32_t&, + const uint32_t&, const uint32_t& n_block_idx, + const uint32_t&) { + #pragma unroll + for (uint32_t k_block_idx = 0; + k_block_idx < kHidden / BLOCK_K; + advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + if (cute::elect_one_sync()) { + tma::copy( + &tensor_map_weights, full_barriers[stage_idx], + smem_b[stage_idx], k_block_idx * BLOCK_K, + expert_idx * 2 * kIntermediateHidden + + n_block_idx * BLOCK_N, + 2); + tma::copy( + &tensor_map_weights_sf, full_barriers[stage_idx], + smem_sfb[stage_idx], n_block_idx * BLOCK_N, + expert_idx * (kHidden / (kGranK * 4)) + + k_block_idx, + 2); + if (is_leader_cta) { + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_B_SIZE_PER_STAGE + + BLOCK_N * sizeof(uint32_t) * 2); + } else { + full_barriers[stage_idx]->arrive(0u); + } + } + __syncwarp(); + } + }); + } else if (warp_idx == 2) { + cutlass::arch::warpgroup_reg_dealloc(); + if (is_leader_cta) { + auto instr_desc = + cute::UMMA::make_instr_desc_block_scaled< + b_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, + UMMA_M, UMMA_N, cute::UMMA::Major::K, + cute::UMMA::Major::K>(); + auto sf_desc = mma::sm100::make_sf_desc(nullptr); + auto a_desc = mma::sm100::make_umma_desc< + cute::UMMA::Major::K, LOAD_BLOCK_M, BLOCK_K, + kSwizzleAMode>(smem_a[0], 0, 0); + auto b_desc = mma::sm100::make_umma_desc< + cute::UMMA::Major::K, LOAD_BLOCK_N, BLOCK_K, + kSwizzleBMode>(smem_b[0], 0, 0); + const uint32_t a_desc_lo = lane_idx < kNumStages + ? a_desc.lo + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 + : 0; + const uint32_t b_desc_lo = lane_idx < kNumStages + ? b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 + : 0; + uint32_t current_iter = 0; + + for_each_block([&](const uint32_t&, const uint32_t&, + const uint32_t&, const uint32_t&, + const uint32_t& valid_m) { + mma::sm100::update_instr_desc_with_umma_n( + instr_desc, math::align(valid_m, 16u)); + const uint32_t accum_stage = + current_iter % kNumEpilogueStages; + const uint32_t accum_phase = + (current_iter++ / kNumEpilogueStages) & 1; + tmem_empty_barriers[accum_stage]->wait( + accum_phase ^ 1); + ptx::tcgen05_after_thread_sync(); + + #pragma unroll + for (uint32_t k_block_idx = 0; + k_block_idx < kHidden / BLOCK_K; + advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + ptx::tcgen05_after_thread_sync(); + const uint32_t a_desc_base = + ptx::exchange(a_desc_lo, stage_idx); + const uint32_t b_desc_base = + ptx::exchange(b_desc_lo, stage_idx); + if (cute::elect_one_sync()) { + using utccp_t = + cute::SM100_UTCCP_4x32dp128bit_2cta; + #pragma unroll + for (uint32_t i = 0; + i < SF_BLOCK_M / + kNumUTCCPAlignedElems; + ++i) { + mma::sm100::replace_smem_desc_addr( + sf_desc, + smem_sfa[stage_idx] + + i * kNumUTCCPAlignedElems); + utccp_t::copy( + sf_desc, + kTmemStartColOfSFA + i * 4); + } + mma::sm100::replace_smem_desc_addr( + sf_desc, smem_sfb[stage_idx]); + utccp_t::copy(sf_desc, kTmemStartColOfSFB); + + #pragma unroll + for (uint32_t k = 0; + k < BLOCK_K / UMMA_K; ++k) { + const auto runtime_instr_desc = + 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, 0, k * UMMA_K); + b_desc.lo = + mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, + LOAD_BLOCK_N, kSwizzleBMode, + b_dtype_t>( + b_desc_base, 0, k * UMMA_K); + ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma( + b_desc, a_desc, + accum_stage * UMMA_N, + k_block_idx > 0 || k > 0, + runtime_instr_desc, + kTmemStartColOfSFB, + kTmemStartColOfSFA); + } + } + __syncwarp(); + + constexpr uint16_t kCTAMask = 0x3; + cutlass::arch::umma_arrive_multicast_2x1SM( + reinterpret_cast( + empty_barriers[stage_idx]), + kCTAMask); + if (k_block_idx == + kHidden / BLOCK_K - 1) { + cutlass::arch:: + umma_arrive_multicast_2x1SM( + reinterpret_cast( + tmem_full_barriers[ + accum_stage]), + kCTAMask); + } + __syncwarp(); + } + }); + if (current_iter > 0) { + const uint32_t last = current_iter - 1; + tmem_empty_barriers[ + last % kNumEpilogueStages] + ->wait((last / kNumEpilogueStages) & 1); + } + } + } else if (warp_idx == 3) { + cutlass::arch::warpgroup_reg_dealloc(); + } else if ( + warp_idx < + (kNumNonEpilogueThreads + + kNumEpilogueThreads) / + 32) { + cutlass::arch::warpgroup_reg_alloc(); + DG_TRAP_ONLY_DEVICE_ASSERT( + ptx::ld_shared(tmem_ptr_in_smem) == 0); + const uint32_t epilogue_warp_idx = warp_idx - 4; + uint32_t current_iter = 0; + uint32_t tma_stage_idx = 0; + + for_each_block([&](const uint32_t&, const uint32_t& pool_block_offset, + const uint32_t& m_block_idx, + const uint32_t& n_block_idx, + const uint32_t& valid_m) { + const uint32_t accum_stage = + current_iter % kNumEpilogueStages; + const uint32_t accum_phase = + (current_iter++ / kNumEpilogueStages) & 1; + tmem_full_barriers[accum_stage]->wait(accum_phase); + ptx::tcgen05_after_thread_sync(); + + epilogue::sm100_store_cd_swap_ab< + BLOCK_M, BLOCK_N, STORE_BLOCK_M, + STORE_BLOCK_N, kSwizzleCDMode, + kNumTMAStoreStages, kNumEpilogueThreads, + GemmType::Normal, false, cd_dtype_t, + epilogue::transform::EpilogueIdentity>( + smem_cd, tma_stage_idx, + accum_stage * UMMA_N, + (pool_block_offset + m_block_idx) * BLOCK_M, + n_block_idx * BLOCK_N, 0, + math::align(valid_m, 16u), + epilogue_warp_idx, lane_idx, + tmem_empty_barriers[accum_stage], + tensor_map_output); + }); + + // The dgrad phase consumes gate/up from global memory. Drain the final + // two TMA-store stages before publishing phase completion. + if (epilogue_warp_idx == 0) + cute::tma_store_wait<0>(); + __syncwarp(); + } + } + if ( + warp_idx >= kDispatchWarpStart && + warp_idx < kDispatchWarpStart + kNumDispatchWarps) { + // The 1024-thread dgrad launch already reserves extra warps. Reuse one + // warpgroup as the third role instead of increasing the launch size. + constexpr uint32_t kNumDispatchRegisters = + kBF16Mode ? 56 : (kGateUpPrepared ? 40 : 48); + cutlass::arch::warpgroup_reg_dealloc(); + if constexpr ( + kNumRanks > 1 && + !(kBF16Mode && kDispatchInputsPrepared)) { + const uint32_t dispatch_warp_idx = + warp_idx - kDispatchWarpStart; + const uint32_t dispatch_thread_idx = + dispatch_warp_idx * 32 + lane_idx; + constexpr uint32_t kDispatchGridSyncIndex = 0; + constexpr uint32_t kDispatchDoneGridSyncIndex = 1; + constexpr uint32_t kBeforeBackwardPullBarrierTag = 4; + constexpr uint32_t kDispatchNamedBarrierIdx = 15; + + // All ranks stage their local BF16 grad-y before launch. This + // system-scope barrier publishes those stores before remote TMA. + trace_begin(3); + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumDispatchThreads, + kDispatchGridSyncIndex, + kBeforeBackwardPullBarrierTag>( + backward_workspace, backward_sym_buffer, + blockIdx.x, dispatch_thread_idx, + [=]() { + ptx::sync_aligned( + kNumDispatchThreads, + kDispatchNamedBarrierIdx); + }, + true, true); + trace_end(3); + + auto* pull_buffer = + reinterpret_cast(smem_buffer) + + dispatch_warp_idx * kHidden; + auto* pull_mbarrier = + dispatch_barriers[dispatch_warp_idx]; + uint32_t pull_mbarrier_phase = 0; + uint32_t pool_block_offset = 0; + + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg(expert_counts + expert_idx)); + for (uint32_t token_idx = + blockIdx.x * kNumDispatchWarps + + dispatch_warp_idx; + token_idx < num_tokens; + token_idx += + kNumSMs * kNumDispatchWarps) { + const uint32_t pool_row = + pool_block_offset * BLOCK_M + + token_idx; + const auto metadata = + token_src_metadata[pool_row]; + const auto* remote_grad_y = + backward_sym_buffer.map( + backward_grad_y + + static_cast( + metadata.token_idx) * + kHidden, + metadata.rank_idx); + + if (cute::elect_one_sync()) { + ptx::tma_load_1d( + pull_buffer, remote_grad_y, + pull_mbarrier, + kHidden * sizeof(cd_dtype_t)); + } + __syncwarp(); + + if (cute::elect_one_sync()) { + const auto* remote_weight = + backward_sym_buffer.map( + backward_topk_weights + + static_cast( + metadata.token_idx) * + num_topk + + metadata.topk_idx, + metadata.rank_idx); + if (route_weights_fp32 != nullptr) { + route_weights_fp32[pool_row] = + *remote_weight; + } else { + route_weights[pool_row] = + cd_dtype_t(*remote_weight); + } + + ptx::mbarrier_arrive_and_set_tx( + pull_mbarrier, + kHidden * sizeof(cd_dtype_t)); + ptx::mbarrier_wait_and_flip_phase( + pull_mbarrier, + pull_mbarrier_phase); + ptx::tma_store_1d( + grad_y_unweighted_output + + static_cast( + pool_row) * + kHidden, + pull_buffer, + kHidden * sizeof(cd_dtype_t)); + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + } + __syncwarp(); + } + pool_block_offset += + math::ceil_div(num_tokens, BLOCK_M); + } + + // Stronger than the eventual per-expert handshake: every L2 tile + // sees every dispatched row. This barrier runs concurrently with + // recompute and joins only at the phase boundary below. + trace_begin(4); + comm::grid_sync< + kNumSMs, kDispatchDoneGridSyncIndex>( + backward_workspace, blockIdx.x, + dispatch_thread_idx, + [=]() { + ptx::sync_aligned( + kNumDispatchThreads, + kDispatchNamedBarrierIdx); + }); + trace_end(4); + } + } else if (warp_idx >= 12) { + // W13 wgrad needs the exact BF16 value represented by the forward + // FP8+UE8M0 pool. Produce it while the recompute MMA is running, using + // otherwise-idle warps. Padding rows are explicitly zeroed so the + // k-grouped wgrad mainloop can round K up to 64 without reading the + // following expert. + constexpr uint32_t kNumXPoolRegisters = + kBF16Mode ? 56 : 40; + cutlass::arch::warpgroup_reg_dealloc< + kNumXPoolRegisters>(); + constexpr uint32_t kFirstXPoolWarp = 12; + constexpr uint32_t kNumXPoolThreads = + kNumThreads - kFirstXPoolWarp * 32; + const uint32_t x_thread_idx = + (warp_idx - kFirstXPoolWarp) * 32 + lane_idx; + if constexpr ( + !kGateUpPrepared && + !(kBF16Mode && kDispatchInputsPrepared)) { + uint32_t pool_block_offset = 0; + uint32_t global_pool_block = 0; + #pragma unroll + for (uint32_t expert_idx = 0; expert_idx < kNumExperts; + ++expert_idx) { + const uint32_t num_tokens = + static_cast(__ldg(expert_counts + expert_idx)); + const uint32_t num_blocks = + math::ceil_div(num_tokens, BLOCK_M); + for (uint32_t m_block_idx = 0; m_block_idx < num_blocks; + ++m_block_idx, ++global_pool_block) { + if (global_pool_block % kNumSMs != blockIdx.x) + continue; + const uint32_t valid_m = cute::min( + num_tokens - m_block_idx * BLOCK_M, BLOCK_M); + const uint32_t pool_block = + pool_block_offset + m_block_idx; + for (uint32_t linear = x_thread_idx; + linear < BLOCK_M * kHidden; + linear += kNumXPoolThreads) { + const uint32_t row = linear / kHidden; + const uint32_t col = linear - row * kHidden; + const uint32_t pool_row = + pool_block * BLOCK_M + row; + cd_dtype_t value = cd_dtype_t(0.0f); + if (row < valid_m) { + if constexpr (kBF16Mode) { + const auto metadata = + token_src_metadata[pool_row]; + value = *backward_sym_buffer.map( + backward_x + + static_cast( + metadata.token_idx) * + kHidden + + col, + metadata.rank_idx); + if constexpr (kNumRanks == 1) { + grad_y_unweighted_output[ + static_cast(pool_row) * + kHidden + + col] = + *backward_sym_buffer.map( + backward_grad_y + + static_cast( + metadata.token_idx) * + kHidden + + col, + metadata.rank_idx); + if (col == 0) { + const float weight = + *backward_sym_buffer.map( + backward_topk_weights + + static_cast( + metadata.token_idx) * + num_topk + + metadata.topk_idx, + metadata.rank_idx); + route_weights_fp32[pool_row] = + weight; + } + } + } else { + if (pool_row >= num_acts_rows) { + asm volatile("trap;"); + } + const uint32_t idx = row % BLOCK_M; + const uint32_t sf_token = + pool_block * SF_BLOCK_M + + (idx & ~127u) + + (idx & 31u) * 4 + + ((idx >> 5) & 3u); + const uint32_t sf_group = col / 128; + const uint32_t sf_byte = (col / 32) & 3u; + const uint32_t packed_sf = + acts_sf_ptr[ + sf_group * acts_sf_stride + + sf_token]; + const uint32_t exponent = + (packed_sf >> (sf_byte * 8)) & + 0xffu; + const uint32_t scale_bits = + exponent << 23; + const float scale = + *reinterpret_cast( + &scale_bits); + value = cd_dtype_t( + static_cast( + acts_ptr[ + static_cast( + pool_row) * + kHidden + + col]) * + scale); + } + } + x_pool_output[ + static_cast(pool_row) * kHidden + + col] = value; + } + } + pool_block_offset += num_blocks; + } + } + } else { + constexpr uint32_t kNumIdleRegisters = + kBF16Mode ? 56 : 24; + cutlass::arch::warpgroup_reg_dealloc< + kNumIdleRegisters>(); + } + + { + __syncthreads(); + if constexpr ( + kBF16Mode && kNumRanks == 1 && + !kDispatchInputsPrepared) { + // In single-rank BF16 mode the x-pool warps also stage grad-y. + // Their pool-block assignment is independent of the dgrad tile + // assignment, so a cluster barrier is insufficient before W2 + // dgrad starts consuming the completed expert pool. + constexpr uint32_t kLocalDispatchDoneGridSyncIndex = 1; + trace_begin(5); + comm::grid_sync< + kNumSMs, kLocalDispatchDoneGridSyncIndex>( + backward_workspace, blockIdx.x, threadIdx.x, + []() { __syncthreads(); }); + trace_end(5); + } + if constexpr ( + (kBF16Mode || + kRouteWeightMode == + RouteWeightMode::PostDown) && + !kDispatchInputsPrepared) { + uint32_t grad_pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg(expert_counts + expert_idx)); + const uint32_t num_padded_tokens = + math::ceil_div(num_tokens, BLOCK_M) * + BLOCK_M; + for (uint64_t linear = + static_cast(blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast( + num_padded_tokens) * + kHidden; + linear += + static_cast(kNumSMs) * + kNumThreads) { + const uint32_t token_idx = + linear / kHidden; + const uint32_t col = + linear - + static_cast(token_idx) * + kHidden; + const uint32_t pool_row = + grad_pool_block_offset * BLOCK_M + + token_idx; + grad_ye_output[ + static_cast(pool_row) * + kHidden + + col] = + token_idx >= num_tokens + ? cd_dtype_t(0.0f) + : kRouteWeightMode == + RouteWeightMode::PostDown + ? cd_dtype_t( + static_cast( + grad_y_unweighted_output[ + static_cast( + pool_row) * + kHidden + + col]) * + (route_weights_fp32 != nullptr + ? route_weights_fp32[ + pool_row] + : static_cast( + route_weights[ + pool_row]))) + : grad_y_unweighted_output[ + static_cast( + pool_row) * + kHidden + + col]; + } + grad_pool_block_offset += + math::ceil_div(num_tokens, BLOCK_M); + } + if constexpr (kBF16Mode) { + constexpr uint32_t + kW2GradInputGridSyncIndex = 0; + trace_begin(6); + comm::grid_sync< + kNumSMs, + kW2GradInputGridSyncIndex>( + backward_workspace, blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }); + trace_end(6); + } else { + // Standalone MXFP4 backward has no symmetric Workspace. + // Reuse its launch-epoch grid state for the same publication + // barrier before W2 dgrad consumes weighted grad-y. + full_grid_phase_barrier(6); + } + } + if constexpr (kDirectRemoteGradX) { + if constexpr (kNumRanks > 1) { + // backward_grad_y aliases combine plane zero. All ranks must + // finish remotely pulling it before any W13 dgrad epilogue + // reuses the combine planes for direct grad-x writes. + if constexpr ( + !(kBF16Mode && kDispatchInputsPrepared)) { + constexpr uint32_t + kBeforeDirectGradXGridSyncIndex = 2; + constexpr uint32_t + kBeforeDirectGradXBarrierTag = 7; + trace_begin(7); + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumThreads, + kBeforeDirectGradXGridSyncIndex, + kBeforeDirectGradXBarrierTag>( + backward_workspace, + backward_sym_buffer, + blockIdx.x, threadIdx.x, + []() { __syncthreads(); }); + trace_end(7); + } + + } + + if constexpr ( + kCombineOrderMode == + CombineOrderMode::FixedTopK) { + // FixedTopK consumes every physical slot, including invalid + // routes. Clear all slot planes only after all grad-y pulls + // have completed, then publish the clear before any direct + // remote stores. This also makes repeated and single-rank + // calls independent of stale valid routes. + auto* combine_buffer = + const_cast(backward_grad_y); + const uint64_t num_plane_values = + static_cast(num_topk) * + backward_workspace.num_max_tokens_per_rank * + kHidden; + for (uint64_t linear = + static_cast(blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < num_plane_values; + linear += + static_cast(kNumSMs) * + kNumThreads) { + combine_buffer[linear] = + cd_dtype_t(0.0f); + } + + if constexpr (kNumRanks > 1) { + // Do not let a rank remotely write direct grad-x until + // every destination has finished clearing its local + // slot planes. + constexpr uint32_t + kAfterGradYClearGridSyncIndex = 3; + constexpr uint32_t + kAfterGradYClearBarrierTag = 8; + trace_begin(8); + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumThreads, + kAfterGradYClearGridSyncIndex, + kAfterGradYClearBarrierTag>( + backward_workspace, + backward_sym_buffer, + blockIdx.x, threadIdx.x, + []() { __syncthreads(); }); + trace_end(8); + } + } + } + if constexpr (!kBF16Mode) { + if (warp_idx >= kDispatchWarpStart && + warp_idx < + kDispatchWarpStart + + kNumDispatchWarps) { + // Dispatch used 48 registers; transition down to the common + // dgrad epilogue budget with dealloc, not reg_alloc + // (allocating a lower count is illegal on SM100). + cutlass::arch::warpgroup_reg_dealloc<40>(); + } else if (warp_idx >= kDispatchWarpStart) { + cutlass::arch::warpgroup_reg_alloc<40>(); + } + } + const auto for_each_dgrad_block = [&](const auto& func) { + uint32_t next_assigned_block = blockIdx.x; + uint32_t global_block = 0; + uint32_t pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg(expert_counts + expert_idx)); + const uint32_t num_m_blocks = + math::ceil_div(num_tokens, BLOCK_M); + const uint32_t expert_blocks = + num_m_blocks * kNumDgradBlockNs; + const uint32_t expert_end = + global_block + expert_blocks; + + while (next_assigned_block < global_block) + next_assigned_block += kNumSMs; + while (next_assigned_block < expert_end) { + const uint32_t local_block = + next_assigned_block - global_block; + const uint32_t m_block_idx = + local_block / kNumDgradBlockNs; + const uint32_t n_block_idx = + local_block - + m_block_idx * kNumDgradBlockNs; + const uint32_t valid_m = cute::min( + num_tokens - m_block_idx * BLOCK_M, + BLOCK_M); + func( + expert_idx, pool_block_offset, + m_block_idx, n_block_idx, valid_m); + next_assigned_block += kNumSMs; + } + global_block = expert_end; + pool_block_offset += num_m_blocks; + } + }; + + trace_begin(9); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(9); + + trace_begin(10); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(10); + + // Reinitialize the drained pipelines in-place. The FP32 accumulator + // columns are phase-aliased; dgrad does not need the SFA/SFB columns. + if (warp_idx == 0 && cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) { + // A and transposed-W2 TMA warps in both CTAs. + full_barriers[i]->init(4); + empty_barriers[i]->init(1); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueStages; ++i) { + tmem_full_barriers[i]->init(1); + tmem_empty_barriers[i]->init( + 2 * kNumDgradEpilogueThreads); + } + cutlass::arch::fence_barrier_init(); + } + trace_begin(11); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(11); + + stage_idx = 0; + phase = 0; + if (warp_idx == 0) { + // BF16 grad_y producer. The 2-SM TMA instruction writes each + // CTA's half-M operand and completes transactions on CTA0's + // cluster barrier. + for_each_dgrad_block( + [&](const uint32_t& expert_idx, + const uint32_t& pool_block_offset, + const uint32_t& m_block_idx, const uint32_t&, + const uint32_t& valid_m) { + const uint32_t pool_block_idx = + pool_block_offset + m_block_idx; + #pragma unroll 1 + for (uint32_t k_block_idx = 0; + k_block_idx < kHidden / DGRAD_BLOCK_K; + advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + uint32_t m_idx = pool_block_idx * BLOCK_M; + if (!is_leader_cta) + m_idx += math::align(valid_m, 16u) / 2; + if (cute::elect_one_sync()) { + tma::copy< + DGRAD_BLOCK_K, LOAD_BLOCK_M, + DGRAD_BLOCK_K * sizeof(cd_dtype_t), + cd_dtype_t>( + &tensor_map_grad_ye, + full_barriers[stage_idx], + smem_dgrad_a[stage_idx], + k_block_idx * DGRAD_BLOCK_K, + m_idx, 2); + if (is_leader_cta) { + full_barriers[stage_idx] + ->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE * 2); + } else { + full_barriers[stage_idx]->arrive(0u); + } + } + __syncwarp(); + } + }); + } else if (warp_idx == 1) { + // Load the in-kernel dequantized transposed W2 workspace. It is + // shared by all M tiles for this expert instead of reconverting + // the same packed weights for every token block. + for_each_dgrad_block( + [&](const uint32_t& expert_idx, const uint32_t&, + const uint32_t&, const uint32_t& n_block_idx, + const uint32_t&) { + #pragma unroll 1 + for (uint32_t k_block_idx = 0; + k_block_idx < kHidden / DGRAD_BLOCK_K; + advance_pipeline(k_block_idx)) { + const uint32_t weight_tile_idx = + (expert_idx * + (kHidden / DGRAD_BLOCK_K) + + k_block_idx) * + kNumDgradBlockNs + + n_block_idx; + if constexpr (!kBF16Mode) { + while (ptx::ld_acq( + weight_tile_states + + weight_tile_idx) != + launch_epoch) { + } + } + constexpr bool weight_tile_ready = true; + empty_barriers[stage_idx]->wait(phase ^ 1); + if (weight_tile_ready) { + if (cute::elect_one_sync()) { + tma::copy< + LOAD_BLOCK_N, + DGRAD_BLOCK_K, + DGRAD_BLOCK_K * + sizeof( + dgrad_b_dtype_t), + dgrad_b_dtype_t>( + &tensor_map_w2_dequant, + full_barriers[stage_idx], + smem_dgrad_b[stage_idx], + n_block_idx * + BLOCK_N, + expert_idx * + kHidden + + k_block_idx * + DGRAD_BLOCK_K, + 2); + if (is_leader_cta) { + full_barriers[stage_idx] + ->arrive_and_expect_tx( + SMEM_B_SIZE_PER_STAGE * + 2); + } else { + full_barriers[stage_idx] + ->arrive(0u); + } + } + } else { + constexpr uint32_t + kPairsPerTile = + LOAD_BLOCK_N * + DGRAD_BLOCK_K / 2; + auto* smem_b_bytes = + reinterpret_cast( + smem_dgrad_b[stage_idx]); + for (uint32_t pair_idx = lane_idx; + pair_idx < kPairsPerTile; + pair_idx += 32) { + const uint32_t local_k = + pair_idx / + (LOAD_BLOCK_N / 2); + const uint32_t + local_n_pair = + pair_idx % + (LOAD_BLOCK_N / 2); + const uint32_t global_k = + k_block_idx * + DGRAD_BLOCK_K + + local_k; + const uint32_t + global_n_pair = + n_block_idx * + (LOAD_BLOCK_N / + 2) + + local_n_pair; + const uint8_t packed = + static_cast( + __ldg( + w2_weights + + (static_cast< + uint64_t>( + expert_idx) * + kHidden + + global_k) * + (kIntermediateHidden / + 2) + + global_n_pair)); + const float scale = __ldg( + w2_scales + + (static_cast( + expert_idx) * + kHidden + + global_k) * + (kIntermediateHidden / + 32) + + n_block_idx * + (LOAD_BLOCK_N / 32) + + (local_n_pair * 2) / + 32); + uint32_t fp16x2; + asm volatile( + "{\n" + ".reg .b8 fp4;\n" + ".reg .b8 unused1, unused2, unused3;\n" + "mov.b32 {fp4, unused1, unused2, unused3}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4;\n" + "}\n" + : "=r"(fp16x2) + : "r"( + static_cast< + uint32_t>( + packed))); + auto value_pair = + *reinterpret_cast< + __half2*>(&fp16x2); + value_pair = __hmul2( + value_pair, + __float2half2_rn( + scale)); + const auto + value_pair_bf16 = + __float22bfloat162_rn( + __half22float2( + value_pair)); + const uint32_t + scaled_pair = + *reinterpret_cast< + const uint32_t*>( + &value_pair_bf16); + #pragma unroll + for (uint32_t i = 0; + i < 2; ++i) { + const uint32_t local_n = + local_n_pair * 2 + i; + const uint32_t row = + local_n & 7; + const uint32_t col_byte = + local_k * + sizeof( + dgrad_b_dtype_t); + const uint32_t + byte_offset = + (local_n >> 3) * + 8 * 128 + + row * 128 + + ((col_byte >> 4) ^ + row) * + 16 + + (col_byte & 15); + *reinterpret_cast< + uint16_t*>( + smem_b_bytes + + byte_offset) = + static_cast< + uint16_t>( + scaled_pair >> + (i * 16)); + } + } + cutlass::arch:: + fence_view_async_shared(); + if (cute::elect_one_sync()) + full_barriers[stage_idx] + ->arrive(0u); + } + __syncwarp(); + } + }); + } else if (warp_idx == 2) { + if (is_leader_cta) { + auto instr_desc = + cute::UMMA::make_instr_desc< + dgrad_b_dtype_t, cd_dtype_t, float, + UMMA_M, UMMA_N, + cute::UMMA::Major::MN, + cute::UMMA::Major::K>(); + auto a_desc = mma::sm100::make_umma_desc< + cute::UMMA::Major::K, LOAD_BLOCK_M, + DGRAD_BLOCK_K, + DGRAD_BLOCK_K * sizeof(cd_dtype_t)>( + smem_dgrad_a[0], 0, 0); + auto b_desc = mma::sm100::make_umma_desc< + cute::UMMA::Major::MN, LOAD_BLOCK_N, + DGRAD_BLOCK_K, + DGRAD_BLOCK_K * sizeof(dgrad_b_dtype_t)>( + smem_dgrad_b[0], 0, 0); + const uint32_t a_desc_lo = lane_idx < kNumStages + ? a_desc.lo + + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 + : 0; + const uint32_t b_desc_lo = lane_idx < kNumStages + ? b_desc.lo + + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 + : 0; + uint32_t current_iter = 0; + + for_each_dgrad_block( + [&](const uint32_t&, const uint32_t&, + const uint32_t&, const uint32_t&, + const uint32_t& valid_m) { + mma::sm100::update_instr_desc_with_umma_n( + instr_desc, + math::align(valid_m, 16u)); + const auto runtime_instr_desc = + cute::UMMA::make_runtime_instr_desc( + instr_desc); + const uint32_t accum_stage = + current_iter % kNumEpilogueStages; + const uint32_t accum_phase = + (current_iter++ / + kNumEpilogueStages) & + 1; + tmem_empty_barriers[accum_stage]->wait( + accum_phase ^ 1); + ptx::tcgen05_after_thread_sync(); + + #pragma unroll 1 + for (uint32_t k_block_idx = 0; + k_block_idx < + kHidden / DGRAD_BLOCK_K; + advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + ptx::tcgen05_after_thread_sync(); + const uint32_t a_desc_base = + ptx::exchange( + a_desc_lo, stage_idx); + const uint32_t b_desc_base = + ptx::exchange( + b_desc_lo, stage_idx); + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t k = 0; + k < + DGRAD_BLOCK_K / + DGRAD_UMMA_K; + ++k) { + a_desc.lo = + mma::sm100:: + advance_umma_desc_lo< + cute::UMMA::Major::K, + LOAD_BLOCK_M, + DGRAD_BLOCK_K * + sizeof( + cd_dtype_t), + cd_dtype_t>( + a_desc_base, 0, + k * + DGRAD_UMMA_K); + b_desc.lo = + mma::sm100:: + advance_umma_desc_lo< + cute::UMMA::Major::MN, + LOAD_BLOCK_N, + DGRAD_BLOCK_K * + sizeof( + dgrad_b_dtype_t), + dgrad_b_dtype_t>( + b_desc_base, 0, + k * + DGRAD_UMMA_K); + ptx:: + SM100_MMA_F16BF16_2x1SM_SS:: + fma( + b_desc, a_desc, + accum_stage * + UMMA_N, + k_block_idx > 0 || + k > 0, + runtime_instr_desc); + } + } + __syncwarp(); + constexpr uint16_t kCTAMask = 0x3; + cutlass::arch:: + umma_arrive_multicast_2x1SM( + reinterpret_cast( + empty_barriers[ + stage_idx]), + kCTAMask); + if (k_block_idx == + kHidden / DGRAD_BLOCK_K - 1) { + cutlass::arch:: + umma_arrive_multicast_2x1SM( + reinterpret_cast( + tmem_full_barriers[ + accum_stage]), + kCTAMask); + } + __syncwarp(); + } + }); + if (current_iter > 0) { + const uint32_t last = current_iter - 1; + tmem_empty_barriers[ + last % kNumEpilogueStages] + ->wait( + (last / kNumEpilogueStages) & 1); + } + } + } else if (warp_idx >= 4) { + const uint32_t epilogue_warp_idx = warp_idx - 4; + const uint32_t epilogue_thread_idx = + epilogue_warp_idx * 32 + lane_idx; + uint32_t current_iter = 0; + + for_each_dgrad_block( + [&](const uint32_t& expert_idx, + const uint32_t& pool_block_offset, + const uint32_t& m_block_idx, + const uint32_t& n_block_idx, + const uint32_t& valid_m) { + const uint32_t accum_stage = + current_iter % kNumEpilogueStages; + const uint32_t accum_phase = + (current_iter++ / + kNumEpilogueStages) & + 1; + tmem_full_barriers[accum_stage]->wait( + accum_phase); + ptx::tcgen05_after_thread_sync(); + const uint32_t effective_m = + math::align(valid_m, 16u); + + for (uint32_t s = 0; + s < effective_m / STORE_BLOCK_M; ++s) { + cutlass::arch::NamedBarrier::sync( + kNumDgradEpilogueThreads, 0); + if (epilogue_warp_idx < + kNumEpilogueThreads / 32) { + #pragma unroll 8 + for (uint32_t i = 0; + i < STORE_BLOCK_M / 8; ++i) { + const uint32_t tmem_addr = + accum_stage * UMMA_N + + s * STORE_BLOCK_M + i * 8; + uint32_t values[8]; + cute::SM100_TMEM_LOAD_16dp256b1x:: + copy( + tmem_addr, values[0], + values[1], values[2], + values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x:: + copy( + tmem_addr | 0x00100000, + values[4], values[5], + values[6], values[7]); + cutlass::arch:: + fence_view_async_tmem_load(); + + constexpr uint32_t kBankBytes = 16; + const uint32_t outer_atom = + (epilogue_warp_idx / 2) * + STORE_BLOCK_M * 128; + const uint32_t inner_atom = + i * 8 * 128; + const uint32_t row = lane_idx % 8; + const uint32_t col = + (epilogue_warp_idx % 2) * 4 + + lane_idx / 8; + auto* smem_ptr = + reinterpret_cast( + smem_cd[0]) + + outer_atom + inner_atom + + row * (kBankBytes * 8) + + (col ^ row) * kBankBytes; + ptx::SM90_U32x4_STSM_T::copy( + math::cast_into_bf16_and_pack( + values[0], values[1]), + math::cast_into_bf16_and_pack( + values[2], values[3]), + math::cast_into_bf16_and_pack( + values[4], values[5]), + math::cast_into_bf16_and_pack( + values[6], values[7]), + smem_ptr); + } + } + cutlass::arch::NamedBarrier::sync( + kNumDgradEpilogueThreads, 0); + + #pragma unroll + for (uint32_t linear = + epilogue_thread_idx; + linear < + STORE_BLOCK_M * BLOCK_N; + linear += + kNumDgradEpilogueThreads) { + const uint32_t row = + linear / BLOCK_N; + const uint32_t n = + linear - row * BLOCK_N; + const uint32_t local_m = + s * STORE_BLOCK_M + row; + if (local_m >= valid_m) + continue; + + const uint32_t n_atom = n / 64; + const uint32_t n_in_atom = + n - n_atom * 64; + const uint32_t row_in_atom = + row & 7; + const uint32_t smem_byte_offset = + n_atom * + STORE_BLOCK_M * 128 + + (row >> 3) * 8 * 128 + + row_in_atom * 128 + + ((n_in_atom >> 3) ^ + row_in_atom) * + 16 + + (n_in_atom & 7) * + sizeof(cd_dtype_t); + const cd_dtype_t grad_h_w2 = + *reinterpret_cast< + cd_dtype_t*>( + reinterpret_cast< + uint8_t*>( + smem_cd[0]) + + smem_byte_offset); + const uint32_t pool_row = + (pool_block_offset + + m_block_idx) * + BLOCK_M + + local_m; + const uint32_t hidden_col = + n_block_idx * BLOCK_N + n; + // Consume the rank-width contraction directly at + // the base W2 dgrad epilogue. Both branches are + // rounded to BF16 before their sum, matching an + // explicit base + side-LoRA graph without ever + // materializing a [pool, intermediate] side dgrad. + // The dedicated host path tensor-core-expands + // t2 @ A2 into this otherwise-uninitialized output + // plane. Consume it once, then overwrite it below + // with the combined base+side grad_h. + const float side_grad_h_accum = + static_cast(h_weighted_output[ + static_cast(pool_row) * + kIntermediateHidden + hidden_col]); + const cd_dtype_t side_grad_h_bf16 = + cd_dtype_t(side_lora.scale * + side_grad_h_accum); + const cd_dtype_t combined_grad_h_w2 = + cd_dtype_t( + static_cast(grad_h_w2) + + static_cast(side_grad_h_bf16)); + const float route_weight = + route_weights_fp32 != nullptr + ? route_weights_fp32[pool_row] + : static_cast( + route_weights[pool_row]); + const cd_dtype_t grad_h_bf16 = + kRouteWeightMode == + RouteWeightMode::PostDown + ? combined_grad_h_w2 + : cd_dtype_t( + static_cast( + combined_grad_h_w2) * + route_weight); + const float grad_h = + static_cast(grad_h_bf16); + // This is the W2 dgrad output before any pre-down + // route multiplication. In post-down mode its GEMM + // input was already weighted BF16 grad-y. + grad_h_output[ + static_cast( + pool_row) * + kIntermediateHidden + + hidden_col] = + combined_grad_h_w2; + const uint32_t chunk = + hidden_col / 8; + const uint32_t in_chunk = + hidden_col & 7; + const uint32_t gate_col = + kBF16Mode + ? hidden_col + : chunk * 16 + in_chunk; + const uint32_t up_col = + kBF16Mode + ? kIntermediateHidden + + hidden_col + : gate_col + 8; + const float gate_unclamped = + static_cast( + gate_up_output[ + static_cast( + pool_row) * + (2 * + kIntermediateHidden) + + gate_col]); + const float up_unclamped = + static_cast( + gate_up_output[ + static_cast( + pool_row) * + (2 * + kIntermediateHidden) + + up_col]); + + const bool has_activation_clamp = + kBF16Mode + ? activation_limit != + cute::numeric_limits< + float>::infinity() + : activation_limit > 0.0f; + const bool gate_in_range = + !has_activation_clamp || + gate_unclamped <= + activation_limit; + const bool up_in_range = + !has_activation_clamp || + (up_unclamped >= + -activation_limit && + up_unclamped <= + activation_limit); + const float gate = + has_activation_clamp + ? cute::min( + gate_unclamped, + activation_limit) + : gate_unclamped; + const float up = + has_activation_clamp + ? cute::min( + cute::max( + up_unclamped, + -activation_limit), + activation_limit) + : up_unclamped; + float z; + float dz_dgate; + if constexpr ( + kActivationType == + ActivationType::GeGLU) { + constexpr float kAlpha = + 1.5957691216057308f; + constexpr float kBeta = 0.044715f; + // Python evaluates 3.0 * beta in FP64 before + // converting the scalar to FP32. Multiplying + // the already-rounded kBeta by 3.0f is one ULP + // lower and changes BF16 ties in GeGLU dgate. + constexpr float kThreeBeta = 0.134145f; + const float gate_sq = + __fmul_rn(gate, gate); + z = __fmul_rn( + __fmul_rn(kAlpha, gate), + __fadd_rn( + 1.0f, + __fmul_rn( + kBeta, gate_sq))); + dz_dgate = __fmul_rn( + kAlpha, + __fadd_rn( + 1.0f, + __fmul_rn( + kThreeBeta, + gate_sq))); + } else { + z = gate; + dz_dgate = 1.0f; + } + const float neg_exp = + (!kBF16Mode && !kGateUpPrepared) || kFastMath + ? __expf(-z) + : expf(-z); + const float denom = + __fadd_rn(1.0f, neg_exp); + const float sig = + 1.0f / denom; + cd_dtype_t h_act_bf16; + cd_dtype_t grad_gate_bf16; + cd_dtype_t grad_up_bf16; + if constexpr ( + (kBF16Mode || kGateUpPrepared) && + kActivationType == + ActivationType::SwiGLU) { + // Native grouped experts materialize BF16 + // clamp -> SiLU -> SiLU*up. Autograd likewise + // rounds grad_h*up before silu_backward. + const cd_dtype_t silu_bf16 = + cd_dtype_t(gate / denom); + h_act_bf16 = cd_dtype_t(__fmul_rn( + static_cast(silu_bf16), up)); + const cd_dtype_t grad_silu_bf16 = + cd_dtype_t(__fmul_rn(grad_h, up)); + const float one_minus_sig = + __fsub_rn(1.0f, sig); + const float silu_inner = __fadd_rn( + 1.0f, __fmul_rn(gate, one_minus_sig)); + const float grad_silu_sig = __fmul_rn( + static_cast(grad_silu_bf16), sig); + grad_gate_bf16 = cd_dtype_t( + gate_in_range + ? __fmul_rn(grad_silu_sig, silu_inner) + : 0.0f); + grad_up_bf16 = cd_dtype_t( + up_in_range + ? __fmul_rn( + grad_h, + static_cast(silu_bf16)) + : 0.0f); + } else { + const float activated_gate = + __fmul_rn(gate, sig); + h_act_bf16 = + cd_dtype_t( + __fmul_rn( + activated_gate, up)); + const float one_minus_sig = + __fsub_rn(1.0f, sig); + const float gate_sig = + __fmul_rn(gate, sig); + const float activation_grad = + __fadd_rn( + sig, + __fmul_rn( + __fmul_rn( + gate_sig, + one_minus_sig), + dz_dgate)); + grad_gate_bf16 = + cd_dtype_t( + gate_in_range + ? __fmul_rn( + __fmul_rn( + grad_h, + up), + activation_grad) + : 0.0f); + grad_up_bf16 = + cd_dtype_t( + up_in_range + ? __fmul_rn( + grad_h, + activated_gate) + : 0.0f); + } + h_act_output[ + static_cast( + pool_row) * + kIntermediateHidden + + hidden_col] = + h_act_bf16; + // PRE_DOWN may phase-alias h_act and h_weighted. + // Preserve unweighted h until its route-gradient + // reduction, then overwrite it in a later phase. + if (!( + kBF16Mode && + kRouteWeightMode == + RouteWeightMode::PreDown && + h_act_output == + h_weighted_output)) { + h_weighted_output[ + static_cast( + pool_row) * + kIntermediateHidden + + hidden_col] = + kRouteWeightMode == + RouteWeightMode::PostDown + ? h_act_bf16 + : cd_dtype_t( + static_cast( + h_act_bf16) * + route_weight); + } + if constexpr (kBF16Mode) { + // The BF16 forward consumes W13 in MegaMoE's + // 8-row [gate, up] interleave. Emit dSwiGLU in + // that same K order so W13 dgrad can consume + // the already-live forward weight directly; + // no full expert-weight deinterleave is needed + // between recomputation and backward. + const uint32_t interleaved_col = + (hidden_col / 8) * 16 + + (hidden_col % 8); + grad_gate_up_output[ + static_cast(pool_row) * + (2 * kIntermediateHidden) + + interleaved_col] = + grad_gate_bf16; + grad_gate_up_output[ + static_cast(pool_row) * + (2 * kIntermediateHidden) + + interleaved_col + 8] = + grad_up_bf16; + } else { + grad_gate_up_output[ + static_cast(pool_row) * + (2 * kIntermediateHidden) + + hidden_col] = + grad_gate_bf16; + grad_gate_up_output[ + static_cast(pool_row) * + (2 * kIntermediateHidden) + + kIntermediateHidden + hidden_col] = + grad_up_bf16; + } + } + } + ptx::tcgen05_before_thread_sync(); + tmem_empty_barriers[accum_stage]->arrive(0u); + }); + + } + + __syncthreads(); + if constexpr (kCompileW13Dgrad) { + // W13 dgrad consumes grad_gate_up rows produced by every CTA in + // the preceding L2-dgrad/SwiGLU phase. Cluster synchronization is + // insufficient here: an early cluster can otherwise read rows + // whose owning cluster has not stored them yet. + full_grid_phase_barrier(12); + + if constexpr (kBF16Mode) { + // In phase-ordered mode these outputs may still contain the + // forward gate values or reverse-dispatched grad-y in every + // row that the active activation tiles did not visit. Clear + // per-expert block padding and the unused capacity tail only + // after all active gate reads have completed. + uint32_t padding_pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg( + expert_counts + + expert_idx)); + const uint32_t num_blocks = + math::ceil_div( + num_tokens, BLOCK_M); + const uint32_t num_padded_tokens = + num_blocks * BLOCK_M; + const uint32_t padding_rows = + num_padded_tokens - num_tokens; + for (uint64_t linear = + static_cast( + blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast( + padding_rows) * + (2 * + kIntermediateHidden); + linear += + static_cast( + kNumSMs) * + kNumThreads) { + const uint32_t row = + linear / + (2 * kIntermediateHidden); + const uint32_t col = + linear - + static_cast(row) * + (2 * + kIntermediateHidden); + const uint32_t pool_row = + padding_pool_block_offset * + BLOCK_M + + num_tokens + row; + grad_gate_up_output[ + static_cast( + pool_row) * + (2 * + kIntermediateHidden) + + col] = + cd_dtype_t(0.0f); + } + for (uint64_t linear = + static_cast( + blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast( + padding_rows) * + kIntermediateHidden; + linear += + static_cast( + kNumSMs) * + kNumThreads) { + const uint32_t row = + linear / + kIntermediateHidden; + const uint32_t col = + linear - + static_cast(row) * + kIntermediateHidden; + const uint32_t pool_row = + padding_pool_block_offset * + BLOCK_M + + num_tokens + row; + h_weighted_output[ + static_cast( + pool_row) * + kIntermediateHidden + + col] = + cd_dtype_t(0.0f); + } + padding_pool_block_offset += + num_blocks; + } + const uint32_t capacity_tail_start = + padding_pool_block_offset * BLOCK_M; + const uint32_t capacity_tail_rows = + num_pool_rows - capacity_tail_start; + for (uint64_t linear = + static_cast( + blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast( + capacity_tail_rows) * + (2 * + kIntermediateHidden); + linear += + static_cast( + kNumSMs) * + kNumThreads) { + grad_gate_up_output[ + static_cast( + capacity_tail_start) * + (2 * + kIntermediateHidden) + + linear] = + cd_dtype_t(0.0f); + } + for (uint64_t linear = + static_cast( + blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast( + capacity_tail_rows) * + kIntermediateHidden; + linear += + static_cast( + kNumSMs) * + kNumThreads) { + h_weighted_output[ + static_cast( + capacity_tail_start) * + kIntermediateHidden + + linear] = + cd_dtype_t(0.0f); + } + full_grid_phase_barrier(13); + } + + if constexpr ( + kComputeRouteGrad && !kInputsPrepared) { + // The activation epilogue spans multiple N-tile CTAs. Reduce + // each route term only after all tiles are visible so the + // router gradient has a fixed FP32 summation order instead of + // depending on cross-CTA atomic arrival order. + constexpr uint32_t kRouteColumns = + kRouteWeightMode == + RouteWeightMode::PostDown + ? kHidden + : kIntermediateHidden; + // The reference POST_DOWN path uses a Triton tl.sum with a + // power-of-two BLOCK_H and BLOCK_H / 256 warps (clamped to + // [4, 32]). For the production hidden sizes this gives 2, 4, + // or 8 elements per thread. Preserve that exact logical + // layout; changing the columns assigned to a lane changes the + // FP32 reduction result. + constexpr uint32_t kTritonRouteBlockH = [] { + uint32_t value = 1; + while (value < kRouteColumns && value < 8192) + value <<= 1; + return value; + }(); + constexpr uint32_t kTritonRouteNumWarps = [] { + uint32_t value = + kTritonRouteBlockH / 256; + value = value < 4 ? 4 : value; + return value > 32 ? 32 : value; + }(); + constexpr uint32_t kTritonRouteThreads = + kTritonRouteNumWarps * 32; + constexpr uint32_t + kTritonRouteValuesPerThread = + kTritonRouteBlockH / + kTritonRouteThreads; + DG_STATIC_ASSERT( + kTritonRouteValuesPerThread == 2 || + kTritonRouteValuesPerThread == 4 || + kTritonRouteValuesPerThread == 8, + "Unsupported Triton route reduction width"); + constexpr uint32_t kRouteInputPow2 = [] { + uint32_t value = 1; + const uint32_t vectorized_columns = + kRouteColumns / 4; + while (value < 512 && + (value << 1) <= + vectorized_columns) + value <<= 1; + return value; + }(); + auto* route_lane_sums = + reinterpret_cast(smem_gemm_base); + auto* route_control = + reinterpret_cast(smem_gemm_base); + if (threadIdx.x == 0) { + uint32_t total_route_rows = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; + ++expert_idx) { + total_route_rows += + static_cast( + __ldg( + expert_counts + + expert_idx)); + } + route_control[0] = + total_route_rows; + } + __syncthreads(); + const uint32_t total_route_rows = + route_control[0]; + const uint32_t route_output_pow2 = + total_route_rows > 0 + ? 1u << (31 - __clz( + total_route_rows)) + : 1u; + constexpr uint32_t + kInitialRouteGroupThreads = + cute::min( + kRouteInputPow2, 32u); + const uint32_t route_block_height = + cute::min( + route_output_pow2, + 512u / + kInitialRouteGroupThreads); + const uint32_t route_group_threads = + cute::min( + kRouteInputPow2, + 512u / route_block_height); + const uint32_t num_route_groups_per_cta = + kNumThreads / route_group_threads; + const uint32_t route_group_idx = + threadIdx.x / route_group_threads; + const uint32_t route_group_lane_idx = + threadIdx.x & + (route_group_threads - 1); + const uint32_t global_route_group = + blockIdx.x * + num_route_groups_per_cta + + route_group_idx; + const uint32_t num_route_groups = + kNumSMs * + num_route_groups_per_cta; + uint32_t route_pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg(expert_counts + expert_idx)); + for (uint32_t token_idx = global_route_group; + token_idx < num_tokens; + token_idx += num_route_groups) { + const uint32_t pool_row = + route_pool_block_offset * BLOCK_M + + token_idx; + float grad_route = 0.0f; + if constexpr (false) { + float grad_y[ + kTritonRouteValuesPerThread]; + float down[ + kTritonRouteValuesPerThread]; + #pragma unroll + for (uint32_t i = 0; + i < + kTritonRouteValuesPerThread; + ++i) { + const uint32_t col = + route_group_lane_idx + + i * kTritonRouteThreads; + grad_y[i] = + col < kHidden + ? static_cast( + grad_y_unweighted_output[ + static_cast( + pool_row) * + kHidden + + col]) + : 0.0f; + down[i] = + col < kHidden + ? static_cast( + down_unweighted_output[ + static_cast( + pool_row) * + kHidden + + col]) + : 0.0f; + } + + if constexpr ( + kTritonRouteValuesPerThread == + 2) { + grad_route = __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[1], down[1])); + } else if constexpr ( + kTritonRouteValuesPerThread == + 4) { + const float even = + __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], + down[2])); + const float odd = + __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], + down[3])); + grad_route = + __fadd_rn(even, odd); + } else { + const float pair_02 = + __fmaf_rn( + grad_y[0], down[0], + __fmul_rn( + grad_y[2], + down[2])); + const float pair_13 = + __fmaf_rn( + grad_y[1], down[1], + __fmul_rn( + grad_y[3], + down[3])); + const float pair_46 = + __fmaf_rn( + grad_y[4], down[4], + __fmul_rn( + grad_y[6], + down[6])); + const float pair_57 = + __fmaf_rn( + grad_y[5], down[5], + __fmul_rn( + grad_y[7], + down[7])); + grad_route = __fadd_rn( + __fadd_rn( + pair_02, pair_46), + __fadd_rn( + pair_13, pair_57)); + } + + // Triton first performs a butterfly reduction + // within each physical warp. + #pragma unroll + for (uint32_t offset = 16; + offset > 0; + offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, + offset)); + } + + const uint32_t warp_in_group = + route_group_lane_idx / 32; + const uint32_t lane_in_warp = + route_group_lane_idx & 31; + if (lane_in_warp == 0) { + route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + warp_in_group] = + grad_route; + } + ptx::sync_aligned( + kTritonRouteThreads, + route_group_idx); + + // Triton loads the power-of-two set of warp + // partials into the first warp and reduces it with + // the same butterfly tree. + if (warp_in_group == 0) { + grad_route = + route_lane_sums[ + route_group_idx * + kTritonRouteNumWarps + + (lane_in_warp & + (kTritonRouteNumWarps - + 1))]; + #pragma unroll + for (uint32_t offset = + kTritonRouteNumWarps / + 2; + offset > 0; + offset >>= 1) { + grad_route = __fadd_rn( + grad_route, + __shfl_xor_sync( + 0xffffffff, + grad_route, + offset)); + } + } + } else { + float lane_sums[4] = { + 0.0f, 0.0f, 0.0f, 0.0f}; + if constexpr ( + kRouteWeightMode == + RouteWeightMode::PostDown) { + for (uint32_t col_base = + route_group_lane_idx * + 4; + col_base < kHidden; + col_base += + route_group_threads * + 4) { + #pragma unroll + for (uint32_t i = 0; i < 4; + ++i) { + const uint32_t col = + col_base + i; + const float grad_y = + static_cast( + grad_y_unweighted_output[ + static_cast< + uint64_t>( + pool_row) * + kHidden + + col]); + const float down = + static_cast( + down_unweighted_output[ + static_cast< + uint64_t>( + pool_row) * + kHidden + + col]); + lane_sums[i] = + __fadd_rn( + lane_sums[i], + __fmul_rn( + grad_y, + down)); + } + } + } else { + for (uint32_t col_base = + route_group_lane_idx * + 4; + col_base < + kIntermediateHidden; + col_base += + route_group_threads * + 4) { + #pragma unroll + for (uint32_t i = 0; i < 4; + ++i) { + const uint32_t col = + col_base + i; + const float grad_h = + static_cast( + grad_h_output[ + static_cast< + uint64_t>( + pool_row) * + kIntermediateHidden + + col]); + const float h_act = + static_cast( + h_act_output[ + static_cast< + uint64_t>( + pool_row) * + kIntermediateHidden + + col]); + lane_sums[i] = + __fadd_rn( + lane_sums[i], + __fmul_rn( + grad_h, + h_act)); + } + } + } + grad_route = __fadd_rn( + __fadd_rn( + lane_sums[0], + lane_sums[1]), + lane_sums[2]); + grad_route = __fadd_rn( + grad_route, lane_sums[3]); + + route_lane_sums[threadIdx.x] = + grad_route; + if (route_group_threads > 32) { + for (uint32_t offset = + route_group_threads / + 2; + offset >= 32; + offset >>= 1) { + ptx::sync_aligned( + route_group_threads, + route_group_idx); + if (route_group_lane_idx < + offset) { + grad_route = + __fadd_rn( + grad_route, + route_lane_sums[ + threadIdx.x + + offset]); + route_lane_sums[ + threadIdx.x] = + grad_route; + } + } + } + if (route_group_lane_idx < 32) { + #pragma unroll + for (uint32_t offset = 16; + offset > 0; + offset >>= 1) { + grad_route = + __fadd_rn( + grad_route, + __shfl_down_sync( + 0xffffffff, + grad_route, + offset)); + } + } + } + if (route_group_lane_idx == 0) { + grad_route_output[pool_row] = + grad_route; + if (backward_grad_route != nullptr) { + const auto metadata = + token_src_metadata[pool_row]; + auto* remote_grad_route = + backward_sym_buffer.map( + backward_grad_route + + static_cast( + metadata.token_idx) * + num_topk + + metadata.topk_idx, + metadata.rank_idx); + *remote_grad_route = grad_route; + } + } + if (route_group_threads > 32) { + ptx::sync_aligned( + route_group_threads, + route_group_idx); + } else { + __syncwarp(); + } + } + route_pool_block_offset += + math::ceil_div(num_tokens, BLOCK_M); + } + } + if constexpr ( + kBF16Mode && + kRouteWeightMode == RouteWeightMode::PreDown && + !kInputsPrepared) { + if (h_act_output == h_weighted_output) { + // Every route reduction must consume unweighted h before + // the shared storage becomes the W2-wgrad input. + full_grid_phase_barrier(14); + uint32_t pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg( + expert_counts + + expert_idx)); + for (uint64_t linear = + static_cast( + blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast( + num_tokens) * + kIntermediateHidden; + linear += + static_cast( + kNumSMs) * + kNumThreads) { + const uint32_t token_idx = + linear / + kIntermediateHidden; + const uint32_t col = + linear - + static_cast( + token_idx) * + kIntermediateHidden; + const uint32_t pool_row = + pool_block_offset * BLOCK_M + + token_idx; + const uint64_t offset = + static_cast( + pool_row) * + kIntermediateHidden + + col; + h_weighted_output[offset] = + cd_dtype_t( + static_cast( + h_act_output[offset]) * + route_weights_fp32[ + pool_row]); + } + pool_block_offset += + math::ceil_div( + num_tokens, BLOCK_M); + } + } + } + + // Phase 3: dequantize canonical [W1; W3] once per launch, then + // consume it as the transposed BF16 operand for W13 dgrad. This + // phase starts only after L2 dgrad/SwiGLU has drained both TMEM + // accumulator stages, so the same 512-column allocation is reused. + + const uint32_t w13_launch_epoch = + launch_epoch ^ 0x80000000u; + + const auto for_each_w13_dgrad_block = + [&](const auto& func) { + uint32_t next_assigned_block = + blockIdx.x; + uint32_t global_block = 0; + uint32_t pool_block_offset = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; + ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg( + expert_counts + + expert_idx)); + const uint32_t num_m_blocks = + math::ceil_div( + num_tokens, BLOCK_M); + const uint32_t expert_blocks = + num_m_blocks * + kNumW13DgradBlockNs; + const uint32_t expert_end = + global_block + + expert_blocks; + + while (next_assigned_block < + global_block) + next_assigned_block += + kNumSMs; + while (next_assigned_block < + expert_end) { + const uint32_t local_block = + next_assigned_block - + global_block; + const uint32_t + m_block_idx = + local_block / + kNumW13DgradBlockNs; + const uint32_t + n_block_idx = + local_block - + m_block_idx * + kNumW13DgradBlockNs; + const uint32_t valid_m = + cute::min( + num_tokens - + m_block_idx * + BLOCK_M, + BLOCK_M); + func( + expert_idx, + pool_block_offset, + m_block_idx, + n_block_idx, + valid_m); + next_assigned_block += + kNumSMs; + } + global_block = expert_end; + pool_block_offset += + num_m_blocks; + } + }; + + trace_begin(15); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(15); + if (warp_idx == 0 && + cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; + i < kNumStages; ++i) { + full_barriers[i]->init(4); + empty_barriers[i]->init(1); + } + #pragma unroll + for (uint32_t i = 0; + i < kNumEpilogueStages; ++i) { + tmem_full_barriers[i]->init(1); + tmem_empty_barriers[i]->init( + 2 * + kNumDgradEpilogueThreads); + } + cutlass::arch::fence_barrier_init(); + } + trace_begin(16); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(16); + trace_begin(21); + + stage_idx = 0; + phase = 0; + if (warp_idx == 0) { + for_each_w13_dgrad_block( + [&](const uint32_t&, + const uint32_t& + pool_block_offset, + const uint32_t& + m_block_idx, + const uint32_t&, + const uint32_t& valid_m) { + const uint32_t pool_block_idx = + pool_block_offset + + m_block_idx; + #pragma unroll + for (uint32_t split_idx = 0; + split_idx < + kNumW13DgradSplits; + ++split_idx) { + #pragma unroll 1 + for (uint32_t k_block_idx = 0; + k_block_idx < + (2 * + kIntermediateHidden) / + (DGRAD_BLOCK_K * + kNumW13DgradSplits); + advance_pipeline( + k_block_idx)) { + empty_barriers[stage_idx] + ->wait(phase ^ 1); + uint32_t m_idx = + pool_block_idx * + BLOCK_M; + if (!is_leader_cta) + m_idx += + math::align( + valid_m, 16u) / + 2; + if (cute::elect_one_sync()) { + tma::copy< + DGRAD_BLOCK_K, + LOAD_BLOCK_M, + DGRAD_BLOCK_K * + sizeof( + cd_dtype_t), + cd_dtype_t>( + &tensor_map_grad_gate_up, + full_barriers[ + stage_idx], + smem_dgrad_a[ + stage_idx], + split_idx * + ((2 * + kIntermediateHidden) / + kNumW13DgradSplits) + + k_block_idx * + DGRAD_BLOCK_K, + m_idx, 2); + if (is_leader_cta) { + full_barriers[ + stage_idx] + ->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE * + 2); + } else { + full_barriers[ + stage_idx] + ->arrive(0u); + } + } + __syncwarp(); + } + } + }); + } else if (warp_idx == 1) { + for_each_w13_dgrad_block( + [&](const uint32_t& expert_idx, + const uint32_t&, + const uint32_t&, + const uint32_t& + n_block_idx, + const uint32_t&) { + #pragma unroll + for (uint32_t split_idx = 0; + split_idx < + kNumW13DgradSplits; + ++split_idx) { + #pragma unroll 1 + for (uint32_t k_block_idx = 0; + k_block_idx < + (2 * + kIntermediateHidden) / + (DGRAD_BLOCK_K * + kNumW13DgradSplits); + advance_pipeline( + k_block_idx)) { + const uint32_t + global_k_block_idx = + split_idx * + ((2 * + kIntermediateHidden) / + (DGRAD_BLOCK_K * + kNumW13DgradSplits)) + + k_block_idx; + const uint32_t + weight_tile_idx = + (expert_idx * + ((2 * + kIntermediateHidden) / + DGRAD_BLOCK_K) + + global_k_block_idx) * + kNumW13DgradBlockNs + + n_block_idx; + if constexpr (!kBF16Mode) { + while (ptx::ld_acq( + weight_tile_states + + kNumW2WeightTileStates + + weight_tile_idx) != + w13_launch_epoch) { + } + } + empty_barriers[stage_idx] + ->wait(phase ^ 1); + if (cute::elect_one_sync()) { + tma::copy< + LOAD_BLOCK_N, + DGRAD_BLOCK_K, + DGRAD_BLOCK_K * + sizeof( + dgrad_b_dtype_t), + dgrad_b_dtype_t>( + &tensor_map_w13_dequant, + full_barriers[ + stage_idx], + smem_dgrad_b[ + stage_idx], + n_block_idx * + BLOCK_N, + expert_idx * + (2 * + kIntermediateHidden) + + global_k_block_idx * + DGRAD_BLOCK_K, + 2); + if (is_leader_cta) { + full_barriers[ + stage_idx] + ->arrive_and_expect_tx( + SMEM_B_SIZE_PER_STAGE * + 2); + } else { + full_barriers[ + stage_idx] + ->arrive(0u); + } + } + __syncwarp(); + } + } + }); + } else if (warp_idx == 2) { + if (is_leader_cta) { + auto instr_desc = + cute::UMMA::make_instr_desc< + dgrad_b_dtype_t, + cd_dtype_t, float, + UMMA_M, UMMA_N, + cute::UMMA::Major::MN, + cute::UMMA::Major::K>(); + auto a_desc = + mma::sm100::make_umma_desc< + cute::UMMA::Major::K, + LOAD_BLOCK_M, + DGRAD_BLOCK_K, + DGRAD_BLOCK_K * + sizeof(cd_dtype_t)>( + smem_dgrad_a[0], 0, 0); + auto b_desc = + mma::sm100::make_umma_desc< + cute::UMMA::Major::MN, + LOAD_BLOCK_N, + DGRAD_BLOCK_K, + DGRAD_BLOCK_K * + sizeof( + dgrad_b_dtype_t)>( + smem_dgrad_b[0], 0, 0); + const uint32_t a_desc_lo = + lane_idx < kNumStages + ? a_desc.lo + + lane_idx * + SMEM_A_SIZE_PER_STAGE / + 16 + : 0; + const uint32_t b_desc_lo = + lane_idx < kNumStages + ? b_desc.lo + + lane_idx * + SMEM_B_SIZE_PER_STAGE / + 16 + : 0; + uint32_t current_iter = 0; + + for_each_w13_dgrad_block( + [&](const uint32_t&, + const uint32_t&, + const uint32_t&, + const uint32_t&, + const uint32_t& + valid_m) { + mma::sm100:: + update_instr_desc_with_umma_n( + instr_desc, + math::align( + valid_m, 16u)); + const auto + runtime_instr_desc = + cute::UMMA:: + make_runtime_instr_desc( + instr_desc); + #pragma unroll + for (uint32_t split_idx = 0; + split_idx < + kNumW13DgradSplits; + ++split_idx) { + const uint32_t accum_stage = + current_iter % + kNumEpilogueStages; + const uint32_t accum_phase = + (current_iter++ / + kNumEpilogueStages) & + 1; + tmem_empty_barriers[ + accum_stage] + ->wait( + accum_phase ^ 1); + ptx::tcgen05_after_thread_sync(); + + #pragma unroll 1 + for (uint32_t + k_block_idx = 0; + k_block_idx < + (2 * + kIntermediateHidden) / + (DGRAD_BLOCK_K * + kNumW13DgradSplits); + advance_pipeline( + k_block_idx)) { + full_barriers[ + stage_idx] + ->wait(phase); + ptx::tcgen05_after_thread_sync(); + const uint32_t + a_desc_base = + ptx::exchange( + a_desc_lo, + stage_idx); + const uint32_t + b_desc_base = + ptx::exchange( + b_desc_lo, + stage_idx); + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t k = 0; + k < + DGRAD_BLOCK_K / + DGRAD_UMMA_K; + ++k) { + a_desc.lo = + mma::sm100:: + advance_umma_desc_lo< + cute::UMMA::Major::K, + LOAD_BLOCK_M, + DGRAD_BLOCK_K * + sizeof( + cd_dtype_t), + cd_dtype_t>( + a_desc_base, + 0, + k * + DGRAD_UMMA_K); + b_desc.lo = + mma::sm100:: + advance_umma_desc_lo< + cute::UMMA::Major::MN, + LOAD_BLOCK_N, + DGRAD_BLOCK_K * + sizeof( + dgrad_b_dtype_t), + dgrad_b_dtype_t>( + b_desc_base, + 0, + k * + DGRAD_UMMA_K); + ptx:: + SM100_MMA_F16BF16_2x1SM_SS:: + fma( + b_desc, + a_desc, + accum_stage * + UMMA_N, + k_block_idx > + 0 || + k > 0, + runtime_instr_desc); + } + } + __syncwarp(); + constexpr uint16_t + kCTAMask = 0x3; + cutlass::arch:: + umma_arrive_multicast_2x1SM( + reinterpret_cast< + uint64_t*>( + empty_barriers[ + stage_idx]), + kCTAMask); + if (k_block_idx == + (2 * + kIntermediateHidden) / + (DGRAD_BLOCK_K * + kNumW13DgradSplits) - + 1) { + cutlass::arch:: + umma_arrive_multicast_2x1SM( + reinterpret_cast< + uint64_t*>( + tmem_full_barriers[ + accum_stage]), + kCTAMask); + } + __syncwarp(); + } + } + }); + if (current_iter > 0) { + const uint32_t last = + current_iter - 1; + tmem_empty_barriers[ + last % + kNumEpilogueStages] + ->wait( + (last / + kNumEpilogueStages) & + 1); + } + } + } else if (warp_idx >= 4) { + const uint32_t epilogue_warp_idx = + warp_idx - 4; + const uint32_t epilogue_thread_idx = + epilogue_warp_idx * 32 + + lane_idx; + uint32_t current_iter = 0; + + for_each_w13_dgrad_block( + [&](const uint32_t&, + const uint32_t& + pool_block_offset, + const uint32_t& + m_block_idx, + const uint32_t& + n_block_idx, + const uint32_t& valid_m) { + const uint32_t accum_stage = + current_iter % + kNumEpilogueStages; + const uint32_t accum_phase = + (current_iter / + kNumEpilogueStages) & + 1; + current_iter += + kNumW13DgradSplits; + tmem_full_barriers[accum_stage]->wait( + accum_phase); + if constexpr (kBF16Mode) + tmem_full_barriers[ + accum_stage ^ 1] + ->wait(accum_phase); + ptx::tcgen05_after_thread_sync(); + const uint32_t effective_m = + math::align(valid_m, 16u); + + for (uint32_t s = 0; + s < + effective_m / + STORE_BLOCK_M; + ++s) { + cutlass::arch:: + NamedBarrier::sync( + kNumDgradEpilogueThreads, + 0); + // The four proven loader warps own the 4 KiB + // TMEM-to-shared mapping. Extra dgrad epilogue + // warps participate only in the global scatter. + if (epilogue_warp_idx < + kNumEpilogueThreads / + 32) { + #pragma unroll + for (uint32_t i = 0; + i < + STORE_BLOCK_M / + 8; + ++i) { + const uint32_t tmem_addr = + accum_stage * UMMA_N + + s * STORE_BLOCK_M + + i * 8; + uint32_t w1_values[8]; + uint32_t w3_values[8]; + cute:: + SM100_TMEM_LOAD_16dp256b1x:: + copy( + tmem_addr, + w1_values[0], + w1_values[1], + w1_values[2], + w1_values[3]); + cute:: + SM100_TMEM_LOAD_16dp256b1x:: + copy( + tmem_addr | + 0x00100000, + w1_values[4], + w1_values[5], + w1_values[6], + w1_values[7]); + if constexpr (kBF16Mode) { + const uint32_t + w3_tmem_addr = + (accum_stage ^ + 1) * + UMMA_N + + s * + STORE_BLOCK_M + + i * 8; + cute:: + SM100_TMEM_LOAD_16dp256b1x:: + copy( + w3_tmem_addr, + w3_values[0], + w3_values[1], + w3_values[2], + w3_values[3]); + cute:: + SM100_TMEM_LOAD_16dp256b1x:: + copy( + w3_tmem_addr | + 0x00100000, + w3_values[4], + w3_values[5], + w3_values[6], + w3_values[7]); + } + cutlass::arch:: + fence_view_async_tmem_load(); + + constexpr uint32_t + kBankBytes = 16; + const uint32_t + outer_atom = + (epilogue_warp_idx / + 2) * + STORE_BLOCK_M * + 128; + const uint32_t + inner_atom = + i * 8 * 128; + const uint32_t row = + lane_idx % 8; + const uint32_t col = + (epilogue_warp_idx % + 2) * + 4 + + lane_idx / 8; + auto* smem_ptr = + reinterpret_cast< + uint8_t*>( + smem_cd[0]) + + outer_atom + + inner_atom + + row * + (kBankBytes * + 8) + + (col ^ row) * + kBankBytes; + const auto add_bf16_pair = + [](uint32_t a, + uint32_t b, + uint32_t c, + uint32_t d) { + const uint32_t + w1_packed = + math:: + cast_into_bf16_and_pack( + a, + b); + const uint32_t + w3_packed = + math:: + cast_into_bf16_and_pack( + c, + d); + const auto w1 = + *reinterpret_cast< + const nv_bfloat162*>( + &w1_packed); + const auto w3 = + *reinterpret_cast< + const nv_bfloat162*>( + &w3_packed); + const auto sum = + __hadd2_rn( + w1, w3); + return *reinterpret_cast< + const uint32_t*>( + &sum); + }; + if constexpr (kBF16Mode) { + ptx:: + SM90_U32x4_STSM_T:: + copy( + add_bf16_pair( + w1_values[0], + w1_values[1], + w3_values[0], + w3_values[1]), + add_bf16_pair( + w1_values[2], + w1_values[3], + w3_values[2], + w3_values[3]), + add_bf16_pair( + w1_values[4], + w1_values[5], + w3_values[4], + w3_values[5]), + add_bf16_pair( + w1_values[6], + w1_values[7], + w3_values[6], + w3_values[7]), + smem_ptr); + } else { + ptx:: + SM90_U32x4_STSM_T:: + copy( + math:: + cast_into_bf16_and_pack( + w1_values[0], + w1_values[1]), + math:: + cast_into_bf16_and_pack( + w1_values[2], + w1_values[3]), + math:: + cast_into_bf16_and_pack( + w1_values[4], + w1_values[5]), + math:: + cast_into_bf16_and_pack( + w1_values[6], + w1_values[7]), + smem_ptr); + } + } + } + cutlass::arch:: + NamedBarrier::sync( + kNumDgradEpilogueThreads, + 0); + + if constexpr ( + kWideGradXStore) { + DG_STATIC_ASSERT( + BLOCK_N % 8 == 0, + "Wide grad-x stores require eight-column alignment"); + DG_STATIC_ASSERT( + kHidden % 8 == 0, + "Wide grad-x stores require aligned output rows"); + #pragma unroll + for (uint32_t linear = + epilogue_thread_idx; + linear < + STORE_BLOCK_M * + (BLOCK_N / 8); + linear += + kNumDgradEpilogueThreads) { + const uint32_t row = + linear / + (BLOCK_N / 8); + const uint32_t n = + (linear - + row * + (BLOCK_N / 8)) * + 8; + const uint32_t local_m = + s * STORE_BLOCK_M + + row; + if (local_m >= valid_m) + continue; + const uint32_t n_atom = + n / 64; + const uint32_t + n_in_atom = + n - + n_atom * 64; + const uint32_t + row_in_atom = + row & 7; + const uint32_t + smem_byte_offset = + n_atom * + STORE_BLOCK_M * + 128 + + (row >> 3) * + 8 * 128 + + row_in_atom * 128 + + ((n_in_atom >> 3) ^ + row_in_atom) * + 16; + const auto packed = + *reinterpret_cast< + const uint4*>( + reinterpret_cast< + const uint8_t*>( + smem_cd[0]) + + smem_byte_offset); + const uint32_t pool_row = + (pool_block_offset + + m_block_idx) * + BLOCK_M + + local_m; + const uint32_t out_col = + n_block_idx * + BLOCK_N + + n; + if constexpr ( + kWriteGradXPool) { + *reinterpret_cast< + uint4*>( + grad_x_pool_output + + static_cast< + uint64_t>( + pool_row) * + kHidden + + out_col) = packed; + } + if constexpr ( + kDirectRemoteGradX) { + const auto metadata = + token_src_metadata[ + pool_row]; + auto* combine_buffer = + const_cast< + cd_dtype_t*>( + backward_grad_y); + auto* dst = + combine_buffer + + ((static_cast< + uint64_t>( + metadata + .topk_idx) * + backward_workspace + .num_max_tokens_per_rank + + metadata + .token_idx) * + kHidden + + out_col); + *reinterpret_cast< + uint4*>( + backward_sym_buffer + .map( + dst, + metadata + .rank_idx)) = + packed; + } + } + } else if constexpr ( + kVectorizedGradXStore) { + #pragma unroll + for (uint32_t linear = + epilogue_thread_idx; + linear < + STORE_BLOCK_M * + (BLOCK_N / 2); + linear += + kNumDgradEpilogueThreads) { + const uint32_t row = + linear / + (BLOCK_N / 2); + const uint32_t n = + (linear - + row * + (BLOCK_N / 2)) * + 2; + const uint32_t local_m = + s * STORE_BLOCK_M + + row; + if (local_m >= valid_m) + continue; + const uint32_t + row_in_atom = + row & 7; + const auto load_bf16_bits = + [&](const uint32_t + element_n) { + const uint32_t + n_atom = + element_n / + 64; + const uint32_t + n_in_atom = + element_n - + n_atom * + 64; + const uint32_t + smem_byte_offset = + n_atom * + STORE_BLOCK_M * + 128 + + (row >> 3) * + 8 * + 128 + + row_in_atom * + 128 + + ((n_in_atom >> + 3) ^ + row_in_atom) * + 16 + + (n_in_atom & + 7) * + sizeof( + cd_dtype_t); + return *reinterpret_cast< + uint16_t*>( + reinterpret_cast< + uint8_t*>( + smem_cd[0]) + + smem_byte_offset); + }; + const uint32_t packed = + static_cast( + load_bf16_bits(n)) | + (static_cast( + load_bf16_bits( + n + 1)) + << 16); + const uint32_t pool_row = + (pool_block_offset + + m_block_idx) * + BLOCK_M + + local_m; + const uint32_t out_col = + n_block_idx * + BLOCK_N + + n; + if constexpr ( + kWriteGradXPool) { + *reinterpret_cast< + uint32_t*>( + grad_x_pool_output + + static_cast< + uint64_t>( + pool_row) * + kHidden + + out_col) = packed; + } + if constexpr ( + kDirectRemoteGradX) { + const auto metadata = + token_src_metadata[ + pool_row]; + auto* combine_buffer = + const_cast< + cd_dtype_t*>( + backward_grad_y); + auto* dst = + combine_buffer + + ((static_cast< + uint64_t>( + metadata + .topk_idx) * + backward_workspace + .num_max_tokens_per_rank + + metadata + .token_idx) * + kHidden + + out_col); + *reinterpret_cast< + uint32_t*>( + backward_sym_buffer + .map( + dst, + metadata + .rank_idx)) = + packed; + } + } + } else { + #pragma unroll + for (uint32_t linear = + epilogue_thread_idx; + linear < + STORE_BLOCK_M * + BLOCK_N; + linear += + kNumDgradEpilogueThreads) { + const uint32_t row = + linear / BLOCK_N; + const uint32_t n = + linear - + row * BLOCK_N; + const uint32_t local_m = + s * + STORE_BLOCK_M + + row; + if (local_m >= valid_m) + continue; + const uint32_t n_atom = + n / 64; + const uint32_t + n_in_atom = + n - + n_atom * 64; + const uint32_t + row_in_atom = + row & 7; + const uint32_t + smem_byte_offset = + n_atom * + STORE_BLOCK_M * + 128 + + (row >> 3) * + 8 * 128 + + row_in_atom * + 128 + + ((n_in_atom >> 3) ^ + row_in_atom) * + 16 + + (n_in_atom & 7) * + sizeof( + cd_dtype_t); + const uint32_t pool_row = + (pool_block_offset + + m_block_idx) * + BLOCK_M + + local_m; + const uint32_t out_col = + n_block_idx * + BLOCK_N + + n; + const auto value = + *reinterpret_cast< + cd_dtype_t*>( + reinterpret_cast< + uint8_t*>( + smem_cd[0]) + + smem_byte_offset); + if constexpr ( + kWriteGradXPool) { + grad_x_pool_output[ + static_cast< + uint64_t>( + pool_row) * + kHidden + + out_col] = value; + } + if constexpr ( + kDirectRemoteGradX) { + const auto metadata = + token_src_metadata[ + pool_row]; + auto* combine_buffer = + const_cast< + cd_dtype_t*>( + backward_grad_y); + auto* dst = + combine_buffer + + ((static_cast< + uint64_t>( + metadata + .topk_idx) * + backward_workspace + .num_max_tokens_per_rank + + metadata + .token_idx) * + kHidden + + out_col); + *backward_sym_buffer + .map( + dst, + metadata + .rank_idx) = + value; + } + } + } + } + ptx::tcgen05_before_thread_sync(); + tmem_empty_barriers[accum_stage] + ->arrive(0u); + if constexpr (kBF16Mode) + tmem_empty_barriers[ + accum_stage ^ 1] + ->arrive(0u); + }); + } + } + + trace_end(21); + if constexpr (kNumRanks > 1) { + if constexpr ( + kDirectRemoteGradX || kComputeRouteGrad) { + // Publish every direct grad-x and route-gradient NVLink store + // before any destination rank consumes its source planes. + constexpr uint32_t + kDirectGradXDoneGridSyncIndex = 1; + constexpr uint32_t + kDirectGradXDoneBarrierTag = 9; + if constexpr (kTraceKernel) { + // Decompose the otherwise identical NVLink barrier so the + // trace distinguishes local compute/grid skew from the + // cross-rank signal and its publication grid sync. + trace_begin(17); + comm::grid_sync< + kNumSMs, + kDirectGradXDoneGridSyncIndex>( + backward_workspace, + blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }); + trace_end(17); + + if (blockIdx.x == 0) + trace_begin(18); + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumThreads, + kDirectGradXDoneGridSyncIndex, + kDirectGradXDoneBarrierTag>( + backward_workspace, + backward_sym_buffer, + blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }, + false, false); + if (blockIdx.x == 0) + trace_end(18); + + trace_begin(19); + comm::grid_sync< + kNumSMs, + kDirectGradXDoneGridSyncIndex>( + backward_workspace, + blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }); + trace_end(19); + } else { + comm::nvlink_barrier< + kNumRanks, kNumSMs, kNumThreads, + kDirectGradXDoneGridSyncIndex, + kDirectGradXDoneBarrierTag>( + backward_workspace, + backward_sym_buffer, + blockIdx.x, + threadIdx.x, + []() { __syncthreads(); }); + } + } + } + + if constexpr ( + kBF16Mode && + kRouteWeightMode == RouteWeightMode::PreDown) { + // W13 dgrad consumes the forward-format 8-row [gate, up] + // interleave above. Once every CTA has drained those reads, + // publish the two canonical activation-gradient planes directly + // into backward buffers that are dead at this point. Keeping + // this conversion inside the persistent wave avoids the old + // expert-weight deinterleave/cat and gives the rank-128 adapter + // GEMMs ordinary contiguous operands. + full_grid_phase_barrier(22); + for (uint64_t linear = + static_cast(blockIdx.x) * + kNumThreads + + threadIdx.x; + linear < + static_cast(num_pool_rows) * + kIntermediateHidden; + linear += + static_cast(kNumSMs) * + kNumThreads) { + const uint32_t pool_row = + linear / kIntermediateHidden; + const uint32_t col = + linear - + static_cast(pool_row) * + kIntermediateHidden; + const uint32_t interleaved_col = + (col / 8) * 16 + (col % 8); + const uint64_t src_offset = + static_cast(pool_row) * + (2 * kIntermediateHidden) + + interleaved_col; + h_act_output[linear] = + grad_gate_up_output[src_offset]; + grad_h_output[linear] = + grad_gate_up_output[src_offset + 8]; + } + } + + const auto clear_wgrad_padding_rows = [&]() { + uint32_t pad_pool_block_offset = 0; + uint32_t pad_global_block = 0; + #pragma unroll + for (uint32_t expert_idx = 0; + expert_idx < kNumExperts; ++expert_idx) { + const uint32_t num_tokens = + static_cast( + __ldg(expert_counts + expert_idx)); + const uint32_t num_blocks = + math::ceil_div(num_tokens, BLOCK_M); + if (num_blocks != 0) { + const uint32_t last_valid = + num_tokens - (num_blocks - 1) * BLOCK_M; + const uint32_t pool_block = + pad_pool_block_offset + num_blocks - 1; + if (pad_global_block % kNumSMs == + blockIdx.x) { + for (uint32_t linear = threadIdx.x; + linear < + (BLOCK_M - last_valid) * + (kHidden + + 3 * kIntermediateHidden); + linear += kNumThreads) { + const uint32_t row_delta = + linear / + (kHidden + + 3 * kIntermediateHidden); + const uint32_t col = + linear - + row_delta * + (kHidden + + 3 * + kIntermediateHidden); + const uint32_t pool_row = + pool_block * BLOCK_M + + last_valid + row_delta; + if (col < kHidden) { + grad_ye_output[ + static_cast( + pool_row) * + kHidden + + col] = cd_dtype_t(0.0f); + } else if ( + col < + kHidden + + kIntermediateHidden) { + h_weighted_output[ + static_cast( + pool_row) * + kIntermediateHidden + + col - kHidden] = + cd_dtype_t(0.0f); + } else { + grad_gate_up_output[ + static_cast( + pool_row) * + (2 * + kIntermediateHidden) + + col - kHidden - + kIntermediateHidden] = + cd_dtype_t(0.0f); + } + } + } + ++pad_global_block; + } + pad_pool_block_offset += num_blocks; + } + }; + + // Standalone Kernel B consumes only these three padded operands. Valid + // rows are fully overwritten above; clear only the final partial block + // of each expert instead of memset'ing every active scratch prefix. + if constexpr (kClearWgradPadding) + clear_wgrad_padding_rows(); + + __syncthreads(); + trace_begin(20); + comm::cluster_sync_with_relaxed_arrive(); + trace_end(20); + trace_end(0); + if (warp_idx == 0) + Allocator().free(0, kNumTmemCols); + } +#endif +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_forward.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_forward.cuh new file mode 100644 index 0000000000..ca895ee173 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_forward.cuh @@ -0,0 +1,2585 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +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 STORE_BLOCK_M, + uint32_t kNumRingTokens, + uint32_t kNumStages, + uint32_t kNumBytesPerPull, + uint32_t kNumDispatchThreads, uint32_t kNumNonEpilogueThreads, + uint32_t kNumEpilogueThreads, + uint32_t kNumSMs, uint32_t kNumRanks, + float kActivationClamp, + bool kFastMath, + ActivationType kActivationType, + bool kSaveL1Preact, + bool kSaveStageActivations, + RouteWeightMode kRouteWeightMode, + CombineOrderMode kCombineOrderMode, + bool kSaveDownUnweighted, + bool kSaveX, + uint32_t kSideLoraRank, + uint32_t L1_SHAPE_N = kIntermediateHidden * 2, + uint32_t L1_SHAPE_K = kHidden, + uint32_t L2_SHAPE_N = kHidden, + uint32_t L2_SHAPE_K = kIntermediateHidden, + uint32_t kNumDispatchWarps = kNumDispatchThreads / 32, + uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32, + uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32, + uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4, + uint32_t kNumThreads = kNumDispatchThreads + kNumNonEpilogueThreads + kNumEpilogueThreads, + uint32_t kNumTokensPerWarp = 32 / kNumTopk, + uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks, + uint32_t kNumRingBlocks = kNumRingTokens / BLOCK_M +> +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void +sm100_bf16_mega_moe_side_lora_forward_impl(void* y, + nv_bfloat16* saved_l1_preact, + nv_bfloat16* saved_h_unweighted, + nv_bfloat16* saved_h_weighted, + nv_bfloat16* saved_x, + nv_bfloat16* saved_down_unweighted, + const nv_bfloat16* side_lora_a1, + const nv_bfloat16* side_lora_b1, + const nv_bfloat16* side_lora_a3, + const nv_bfloat16* side_lora_b3, + const nv_bfloat16* side_lora_a2, + const nv_bfloat16* side_lora_b2, + nv_bfloat16* side_lora_l1_scratch, + nv_bfloat16* side_lora_l2_scratch, + int* side_lora_ready, + const float side_lora_scale, + int* cumulative_local_expert_recv_stats, + const int* precomputed_route_counts, + int* route_count_mismatch, + const uint32_t num_tokens, + const uint32_t num_saved_pool_tokens, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_a1, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_a3, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_b1, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_b3, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_a2, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_b2, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l1_scratch, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l2_scratch, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l1_scratch_store, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l2_scratch_store, + const __grid_constant__ cute::TmaDescriptor tensor_map_down_unweighted) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Allocator = cute::TMEM::Allocator2Sm; + + // Template checks + DG_STATIC_ASSERT(kNumDispatchThreads % 128 == 0, "Invalid number of dispatch threads"); + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 128, "Invalid number of MMA non-epilogue threads"); + DG_STATIC_ASSERT(kNumEpilogueThreads % 128 == 0, "Invalid number of MMA epilogue and combine threads"); + DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + DG_STATIC_ASSERT(kSideLoraRank == 128, "Native side-LoRA specialization requires rank 128"); + constexpr bool kHasSideLora = true; + DG_STATIC_ASSERT( + !kHasSideLora || + kSideLoraRank == 128, + "The tensor-core side-LoRA path is specialized for rank 128"); + + // Thread indices + const bool is_leader_cta = cute::block_rank_in_cluster() == 0; + const uint32_t sm_idx = blockIdx.x; + const uint32_t thread_idx = threadIdx.x; + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == 0) { + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); + if constexpr (kHasSideLora) { + cute::prefetch_tma_descriptor(&tensor_map_lora_a1); + cute::prefetch_tma_descriptor(&tensor_map_lora_a3); + cute::prefetch_tma_descriptor(&tensor_map_lora_b1); + cute::prefetch_tma_descriptor(&tensor_map_lora_b3); + cute::prefetch_tma_descriptor(&tensor_map_lora_a2); + cute::prefetch_tma_descriptor(&tensor_map_lora_b2); + cute::prefetch_tma_descriptor(&tensor_map_lora_l1_scratch); + cute::prefetch_tma_descriptor(&tensor_map_lora_l2_scratch); + cute::prefetch_tma_descriptor(&tensor_map_lora_l1_scratch_store); + cute::prefetch_tma_descriptor(&tensor_map_lora_l2_scratch_store); + } + if constexpr (kSaveDownUnweighted) + cute::prefetch_tma_descriptor( + &tensor_map_down_unweighted); + } + + // Workspaces + const auto workspace = layout::Workspace( + sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk, kNumRingTokens); + + // Token and buffer layouts + constexpr auto bf16_token_layout = layout::Data(kHidden * sizeof(nv_bfloat16)); + constexpr auto bf16_intermediate_token_layout = layout::Data(kIntermediateHidden * sizeof(nv_bfloat16)); + constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * sizeof(int64_t), false); + constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * sizeof(float), false); + constexpr auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + + // Registered inputs + const auto input_token_buffer = layout::Buffer( + bf16_token_layout, 1, kNumMaxTokensPerRank, + workspace.get_end_ptr()); + const auto input_topk_idx_buffer = layout::Buffer( + input_topk_idx_layout, 1, kNumMaxTokensPerRank, + input_token_buffer.get_end_ptr()); + const auto input_topk_weights_buffer = layout::Buffer( + input_topk_weights_layout, 1, kNumMaxTokensPerRank, + input_topk_idx_buffer.get_end_ptr()); + + // L1 inputs + const auto l1_token_buffer = layout::Buffer( + bf16_token_layout, 1, kNumRingTokens, + input_topk_weights_buffer.get_end_ptr()); + const auto l1_topk_weights_buffer = layout::Buffer( + l1_topk_weights_layout, 1, kNumRingTokens, + l1_token_buffer.get_end_ptr()); + + // L2 inputs + const auto l2_token_buffer = layout::Buffer( + bf16_intermediate_token_layout, 1, kNumRingTokens, + l1_topk_weights_buffer.get_end_ptr() + ); + + // Combine inputs + const auto combine_token_buffer = layout::Buffer( + bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, + l2_token_buffer.get_end_ptr() + ); + + // Data types + using a_dtype_t = cutlass::bfloat16_t; + using b_dtype_t = cutlass::bfloat16_t; + using d_dtype_t = cutlass::bfloat16_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 + constexpr uint32_t UMMA_BLOCK_K = 64; + constexpr uint32_t UMMA_K = 16; + 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"); + + // Swizzle configs + constexpr uint32_t kSwizzleAMode = 128; + constexpr uint32_t kSwizzleBMode = 128; + constexpr uint32_t kSwizzleCDMode = 128; + DG_STATIC_ASSERT(BLOCK_N * sizeof(nv_bfloat16) % kSwizzleCDMode == 0, "Invalid block N"); + + // Epilogue configs + constexpr uint32_t kNumEpilogueStages = 2; + constexpr uint32_t kNumTMAStoreStages = 2; + + // Shared memory + constexpr uint32_t kSharedMemoryAlignment = 1024; + extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[]; + + // Shared memory sizes + // NOTES: BF16 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; + + // Tensor memory size + constexpr uint32_t kNumAccumTmemCols = + UMMA_N * kNumEpilogueStages; + constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); + DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); + + // Assign shared memory + struct SharedStorage { + alignas(kSharedMemoryAlignment) uint32_t expert_token_count[kNumExperts]; + alignas(kSharedMemoryAlignment) uint8_t dispatch_send_buffer[kNumDispatchWarps][kNumBytesPerPull]; + union { + alignas(kSharedMemoryAlignment) d_dtype_t l1[kNumEpilogueWarpgroups][kNumTMAStoreStages][STORE_BLOCK_M * L1_OUT_BLOCK_N]; + alignas(kSharedMemoryAlignment) d_dtype_t l2[kNumEpilogueWarpgroups][STORE_BLOCK_M * BLOCK_N]; + } smem_d; + alignas(kSharedMemoryAlignment) a_dtype_t smem_a[kNumStages][LOAD_BLOCK_M * BLOCK_K]; + alignas(kSharedMemoryAlignment) b_dtype_t smem_b[kNumStages][LOAD_BLOCK_N * BLOCK_K]; + 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]; + 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)); + + // 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()) { + // The bytes must be 8 bytes aligned + 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) + 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 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 + shared_storage.tmem_full_barriers[i].init(1); + // Arrive only at the leader CTA + shared_storage.tmem_empty_barriers[i].init(2 * kNumEpilogueThreads); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++ i) + shared_storage.combine_barriers[i].init(1); + } + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 3) { + // Allocate tensor memory + 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` + comm::cluster_sync_with_relaxed_arrive(); + + // Task scheduler + auto scheduler = sched::SideLoraMegaMoEScheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, + kNumExpertsPerWave, + kNumSMs, kNumRanks, kHasSideLora>(workspace); + + // MMA pipeline and TMA phases + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + // Intra-SM Barrier indices + constexpr uint32_t kDispatchBarrierIdx = 0; + constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; + constexpr uint32_t kEpilogueFullBarrierIdx = 2; + constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; + + // NVLink barrier tags + constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; + constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + + // Adjust registers + // NOTES: more experts per rank will cost more schedulers' registers + 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, + "Too many registers"); + + // Grid sync index assignments (dispatch and epilogue use separate counters to avoid conflicts) + constexpr uint32_t kDispatchGridSyncIndex = 0; + constexpr uint32_t kEpilogueGridSyncIndex = 1; + + // Different warp roles + if (warp_idx < kNumDispatchWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // Dispatch warps + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + // TODO: figure out better unrolling + // Now, `unroll` is better than `unroll 8` + #pragma unroll + for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; + i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) { + // Allocate slots for each token-topk + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) + process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + }; + + // Count experts' tokens + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + 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(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); + + if constexpr (kSaveL1Preact) { + // Training wgrads must traverse the same stable + // expert/source-rank/token/top-k order as native DeepEP. Partition + // the source routes into contiguous global-warp ranges, build an + // expert prefix for every range, then compact each range in route + // order. This preserves the exact order without making every + // expert warp rescan the complete route tensor. + const uint32_t global_warp_idx = + sm_idx * kNumDispatchWarps + warp_idx; + constexpr uint32_t kNumGlobalWarps = + kNumSMs * kNumDispatchWarps; + constexpr uint32_t kWarpSize = 32; + constexpr uint32_t kScratchValues = + kNumGlobalWarps * kNumExperts; + const bool use_parallel_stable_compaction = + 3 * workspace.num_max_pool_tokens >= kScratchValues; + if (use_parallel_stable_compaction) { + auto warp_expert_prefix = reinterpret_cast( + workspace.get_token_src_metadata_ptr()); + const uint32_t num_routes = num_tokens * kNumTopk; + const uint32_t routes_per_warp = math::align( + math::ceil_div(num_routes, kNumGlobalWarps), + kWarpSize); + const uint32_t route_begin = + global_warp_idx * routes_per_warp; + const uint32_t route_end = + cute::min(route_begin + routes_per_warp, num_routes); + auto warp_prefix_row = + warp_expert_prefix + + global_warp_idx * kNumExperts; + + DG_STATIC_ASSERT( + kNumBytesPerPull >= + kNumExperts * sizeof(uint32_t), + "Dispatch scratch is too small for stable counters"); + auto local_expert_count = + reinterpret_cast( + shared_storage + .dispatch_send_buffer[warp_idx]); + for (uint32_t expert_idx = lane_idx; + expert_idx < kNumExperts; + expert_idx += kWarpSize) { + local_expert_count[expert_idx] = 0; + } + __syncwarp(); + for (uint32_t route_base = route_begin; + route_base < route_end; + route_base += kWarpSize) { + const uint32_t token_topk_idx = + route_base + lane_idx; + int expert_idx = -1; + if (token_topk_idx < route_end) { + expert_idx = static_cast( + __ldg( + input_topk_idx_buffer + .get_base_ptr() + + token_topk_idx)); + } + if (expert_idx >= 0) + atomicAdd_block( + local_expert_count + expert_idx, 1u); + } + __syncwarp(); + for (uint32_t expert_idx = lane_idx; + expert_idx < kNumExperts; + expert_idx += kWarpSize) { + warp_prefix_row[expert_idx] = + local_expert_count[expert_idx]; + } + + comm::grid_sync< + kNumSMs, kDispatchGridSyncIndex>( + workspace, sm_idx, thread_idx, + [=]() { + ptx::sync_aligned( + kNumDispatchThreads, + kDispatchBarrierIdx); + }); + + if ( + global_warp_idx < kNumExperts && + lane_idx == 0) { + uint32_t prefix = 0; + for (uint32_t source_warp = 0; + source_warp < kNumGlobalWarps; + ++source_warp) { + auto count_ptr = + warp_expert_prefix + + source_warp * kNumExperts + + global_warp_idx; + const uint32_t count = *count_ptr; + *count_ptr = prefix; + prefix += count; + } + } + + comm::grid_sync< + kNumSMs, kDispatchGridSyncIndex>( + workspace, sm_idx, thread_idx, + [=]() { + ptx::sync_aligned( + kNumDispatchThreads, + kDispatchBarrierIdx); + }); + + for (uint32_t expert_idx = lane_idx; + expert_idx < kNumExperts; + expert_idx += kWarpSize) { + local_expert_count[expert_idx] = 0; + } + __syncwarp(); + + for (uint32_t route_base = route_begin; + route_base < route_end; + route_base += kWarpSize) { + const uint32_t token_topk_idx = + route_base + lane_idx; + int expert_idx = -1; + if (token_topk_idx < route_end) { + expert_idx = static_cast( + __ldg( + input_topk_idx_buffer + .get_base_ptr() + + token_topk_idx)); + } + const uint32_t active_mask = + __ballot_sync( + 0xffffffff, expert_idx >= 0); + if (expert_idx >= 0) { + const uint32_t matches = + __match_any_sync( + active_mask, expert_idx); + const uint32_t lanes_before = + (1u << lane_idx) - 1u; + const uint32_t dst_slot_idx = + warp_prefix_row[expert_idx] + + local_expert_count[expert_idx] + + __popc(matches & lanes_before); + const uint32_t dst_rank_idx = + expert_idx / kNumExpertsPerRank; + 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; + if ( + lane_idx == + static_cast( + __ffs(matches) - 1)) { + local_expert_count[expert_idx] += + __popc(matches); + } + } + __syncwarp(); + } + } else { + // Tiny workspaces cannot hold the global-warp prefix table. + // Retain the original one-warp-per-expert exact path. + for (uint32_t target_expert = global_warp_idx; + target_expert < kNumExperts; + target_expert += kNumGlobalWarps) { + uint32_t dst_slot_base = 0; + for (uint32_t route_base = 0; + route_base < num_tokens * kNumTopk; + route_base += kWarpSize) { + const uint32_t token_topk_idx = + route_base + lane_idx; + int expert_idx = -1; + if ( + token_topk_idx < + num_tokens * kNumTopk) { + expert_idx = static_cast( + __ldg( + input_topk_idx_buffer + .get_base_ptr() + + token_topk_idx)); + } + const uint32_t matches = + __ballot_sync( + 0xffffffff, + expert_idx == + static_cast( + target_expert)); + if ( + expert_idx == + static_cast(target_expert)) { + const uint32_t lanes_before = + (1u << lane_idx) - 1u; + const uint32_t dst_slot_idx = + dst_slot_base + + __popc( + matches & lanes_before); + const uint32_t dst_rank_idx = + target_expert / + kNumExpertsPerRank; + const auto dst_ptr = + workspace + .get_src_token_topk_idx_ptr( + target_expert % + kNumExpertsPerRank, + sym_buffer.rank_idx, + dst_slot_idx); + *sym_buffer.map( + dst_ptr, dst_rank_idx) = + token_topk_idx; + } + dst_slot_base += __popc(matches); + } + } + } + } else { + // Inference keeps round-robin pull order for NVLink balance. + 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( + 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; + }); + } + + // Grid sync + comm::grid_sync( + workspace, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + + // Write expert count + if (sm_idx == 0) { + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + const uint32_t actual_count = + static_cast( + expert_status & 0xffffffff); + uint32_t published_count = actual_count; + if (precomputed_route_counts != nullptr) { + const int expected_count = + precomputed_route_counts[i]; + if ( + expected_count < 0 || + static_cast( + expected_count) != actual_count + ) { + atomicExch( + route_count_mismatch, 1); + } + published_count = + expected_count >= 0 + ? static_cast( + expected_count) + : actual_count; + } + const uint64_t published_status = + (expert_status & + 0xffffffff00000000ull) | + published_count; + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = published_count; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), + published_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Barrier before pulling + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* After the grid sync above, there is no more writes by other SMs (except 0) */ false, + /* After the NVLink barrier, there is a grid sync */ true + ); + + // Ensure the epilogue barrier cannot run with the pull barrier + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // 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 = &shared_storage.dispatch_barriers[warp_idx]; + + // Cache expert token counts in registers (same pattern as scheduler) + scheduler.fetch_expert_recv_count(); + + // Per-rank counts for current expert (re-loaded when expert changes) + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalWarps) { + // Advance expert until within the range + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++ current_expert_idx >= kNumExpertsPerRank) + break; + + // Update pool block offset for the new expert + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + + // Move start and end to the next expert + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + + // Finish all tokens + if (current_expert_idx >= kNumExpertsPerRank) + break; + + // Load per-rank counts when expert changes + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t j = i * 32 + lane_idx; + // TODO: this is not coalesced + stored_rank_count[i] = j < kNumRanks ? + static_cast(*workspace.get_expert_recv_count_ptr(j, current_expert_idx)) : 0; + } + } + + uint32_t current_rank_in_expert_idx; + uint32_t token_idx_in_expert = token_idx - expert_start_idx; + uint32_t token_idx_in_rank; + if constexpr (kSaveL1Preact) { + // Match native DeepEP's rank-major stable route order. + uint32_t slot_idx = token_idx_in_expert; + current_rank_in_expert_idx = 0; + token_idx_in_rank = 0; + #pragma unroll + for (uint32_t rank_idx = 0; + rank_idx < kNumRanks; ++rank_idx) { + const uint32_t rank_count = + __shfl_sync( + 0xffffffff, + stored_rank_count[ + rank_idx / 32], + rank_idx % 32); + if (slot_idx < rank_count) { + current_rank_in_expert_idx = + rank_idx; + token_idx_in_rank = slot_idx; + break; + } + slot_idx -= rank_count; + } + } else { + // Round-robin rank selection via iterative min-peeling. + uint32_t remaining[kNumRanksPerLane]; + #pragma unroll + for (uint32_t i = 0; + i < kNumRanksPerLane; ++i) + remaining[i] = + stored_rank_count[i]; + uint32_t offset = 0; + uint32_t slot_idx = + token_idx_in_expert; + while (true) { + uint32_t num_actives_in_lane = 0; + uint32_t min_in_lane = 0xffffffff; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + num_actives_in_lane += + remaining[i] > 0; + if (remaining[i] > 0) + min_in_lane = cute::min( + min_in_lane, + remaining[i]); + } + const uint32_t num_active_ranks = + __reduce_add_sync( + 0xffffffff, + num_actives_in_lane); + const uint32_t length = + __reduce_min_sync( + 0xffffffff, min_in_lane); + const uint32_t num_round_tokens = + length * num_active_ranks; + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = + slot_idx % num_active_ranks; + uint32_t num_seen_ranks = 0; + current_rank_in_expert_idx = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t mask = + __ballot_sync( + 0xffffffff, + remaining[i] > 0); + const uint32_t num_active_lanes = + __popc(mask); + if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes) + current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); + num_seen_ranks += num_active_lanes; + } + token_idx_in_rank = + offset + + (slot_idx / + num_active_ranks); + break; + } + slot_idx -= num_round_tokens; + offset += length; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] -= cute::min(remaining[i], length); + } + } + + // Read source token-topk index (written by remote dispatch via NVLink) + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + // Hidden bytes are divided into chunks + constexpr uint32_t kHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumChunks = kHiddenBytes / kNumBytesPerPull; + DG_STATIC_ASSERT(kHiddenBytes % kNumBytesPerPull == 0, "Invalid hidden"); + + // TMA load token from remote rank and store into local + const auto pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + const uint32_t pool_block_idx = pool_token_idx / BLOCK_M; + + // Wait for ring buffer slot to be available (previous consumer must have finished all N blocks) + constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; + const auto l1_empty_count_target = (pool_block_idx / kNumRingBlocks) * kNumL1BlockNs; + if (l1_empty_count_target > 0) { + const auto empty_ptr = workspace.get_l1_empty_count_ptr(pool_block_idx % kNumRingBlocks); + while (ptx::ld_acq(empty_ptr) < l1_empty_count_target); + } + + const auto src_base_ptr = sym_buffer.map( + input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), current_rank_in_expert_idx); + const auto dst_base_ptr = l1_token_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).get_base_ptr(); + 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 * kNumBytesPerPull), + pull_buffer.get_base_ptr(), kNumBytesPerPull + ); + cute::tma_store_arrive(); + if constexpr (kSaveX) { + ptx::tma_store_1d( + saved_x + + static_cast(pool_token_idx) * + kHidden + + i * kNumBytesPerPull / + sizeof(nv_bfloat16), + pull_buffer.get_base_ptr(), kNumBytesPerPull + ); + cute::tma_store_arrive(); + } + ptx::tma_store_wait<0>(); + }; + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumChunks; ++ i) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + math::advance_ptr(src_base_ptr, i * kNumBytesPerPull), + pull_mbarrier, kNumBytesPerPull + ); + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kNumBytesPerPull); + i != (kNumChunks - 1) ? issue_and_wait_pull_store(i) : void(); + } + } + __syncwarp(); + + // Store weights and metadata, then finish the token copy. + if (cute::elect_one_sync()) { + // Load weights + const auto weight = *sym_buffer.map( + input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).template get_base_ptr() = weight; + + // 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 + issue_and_wait_pull_store(kNumChunks - 1); + } + __syncwarp(); + + if (cute::elect_one_sync()) { + const bool is_last_token = + token_idx == expert_end_idx - 1; + ptx::red_add_rel( + workspace.get_l1_full_count_ptr(pool_block_idx % kNumRingBlocks), + is_last_token ? BLOCK_M - (token_idx_in_expert % BLOCK_M) : 1u + ); + } + __syncwarp(); + } + + // Clean workspace for the next usage, and also do cumulative stats + // NOTES: it is overlapped with combine reduction epilogue + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + DG_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { + // SM 0: clear expert send count + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + } else { + // Other SMs: clean blocks + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + // Read expert token count before clearing + const auto num_recv_tokens = static_cast( + *workspace.get_expert_recv_count_sum_ptr(i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + + // Compute expert pool block offset + expert_pool_block_offset = scheduler.get_pool_block_offset(i); + + // Wait read count ready + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + if constexpr (kSaveX) { + // The dispatch path stores only real routes. W13 wgrad + // consumes BLOCK_M-padded expert ranges, so initialize + // just the short per-expert tails instead of clearing the + // entire multi-GiB saved pool before every forward. + const uint32_t num_padding_tokens = + num_recv_m_blocks * BLOCK_M - num_recv_tokens; + const uint64_t padding_start = + static_cast(expert_pool_block_offset) * + BLOCK_M + + num_recv_tokens; + for (uint64_t linear = thread_idx; + linear < + static_cast(num_padding_tokens) * + kHidden; + linear += kNumDispatchThreads) { + saved_x[padding_start * kHidden + linear] = + nv_bfloat16(0.0f); + } + } + + // Clean expert token count, and add cumulative results + DG_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); + if (warp_idx == 0) { + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + } else if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + + // Clean per-rank token count + for (uint32_t j = thread_idx; j < kNumRanks; j += kNumDispatchThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + __syncwarp(); + + // Clean L1 and L2 full stuffs and ring buffer counts + for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { + *workspace.get_l1_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l1_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + } + __syncwarp(); + } + } + + // Wait for all ranks to finish cleaning + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* Before the NVLink barrier, there is a grid sync */ true, + /* At the end of kernel does not need to sync */ false + ); + } else if (warp_idx == kNumDispatchWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM TMA load warp for tokens + 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 cute::TmaDescriptor* tensor_map_a_ptr = + &tensor_map_l1_acts; + if (block_phase == sched::BlockPhase::LoraL1Expand) + tensor_map_a_ptr = &tensor_map_lora_l1_scratch; + else if (block_phase == sched::BlockPhase::LoraL2Shrink || + block_phase == sched::BlockPhase::Linear2) + tensor_map_a_ptr = &tensor_map_l2_acts; + else if (block_phase == sched::BlockPhase::LoraL2Expand) + tensor_map_a_ptr = &tensor_map_lora_l2_scratch; + + // Compute pool block offset for this expert + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + + // Wait the token arrival + if (block_phase == sched::BlockPhase::LoraL1Shrink || + block_phase == sched::BlockPhase::Linear1) { + const auto ptr = workspace.get_l1_full_count_ptr(ring_block_idx); + const auto num_expected_tokens = BLOCK_M * (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != num_expected_tokens); + if constexpr (kHasSideLora) { + if (block_phase == sched::BlockPhase::Linear1) { + const auto generation = + pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq( + reinterpret_cast( + side_lora_ready) + ring_block_idx) < + 2 * generation); + } + } + } else if (block_phase == sched::BlockPhase::LoraL1Expand) { + const auto generation = pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq( + reinterpret_cast(side_lora_ready) + + ring_block_idx) < 2 * generation); + } else if (block_phase == sched::BlockPhase::LoraL2Shrink || + block_phase == sched::BlockPhase::Linear2) { + const auto ptr = workspace.get_l2_full_count_ptr(ring_block_idx); + const auto num_expected_blocks = L2_SHAPE_K / (BLOCK_N / 2) * (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != num_expected_blocks); + if constexpr (kHasSideLora) { + if (block_phase == sched::BlockPhase::Linear2) { + const auto generation = + pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq( + reinterpret_cast( + side_lora_ready) + + 2 * kNumRingBlocks + ring_block_idx) < + generation); + } + } + } else { + const auto generation = pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq( + reinterpret_cast(side_lora_ready) + + 2 * kNumRingBlocks + ring_block_idx) < generation); + } + + // Issue TMA + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); + + // Compute token offset from ring block index + const uint32_t base_k_blocks = + block_phase == sched::BlockPhase::Linear1 ? + L1_SHAPE_K / BLOCK_K : + block_phase == sched::BlockPhase::Linear2 ? + L2_SHAPE_K / BLOCK_K : 0; + const bool is_linear_side_expand = kHasSideLora && + (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) && + k_block_idx >= base_k_blocks; + const uint32_t logical_k_block_idx = + is_linear_side_expand ? + k_block_idx - base_k_blocks : k_block_idx; + const cute::TmaDescriptor* current_tensor_map_a_ptr = + tensor_map_a_ptr; + if (is_linear_side_expand) + current_tensor_map_a_ptr = + block_phase == sched::BlockPhase::Linear1 ? + &tensor_map_lora_l1_scratch : + &tensor_map_lora_l2_scratch; + + uint32_t load_m_idx = is_linear_side_expand + ? pool_block_idx * BLOCK_M + : ring_block_idx * BLOCK_M; + uint32_t k_idx = logical_k_block_idx * BLOCK_K; + if (is_linear_side_expand && + block_phase == sched::BlockPhase::Linear1) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / BLOCK_K; + const uint32_t projection_idx = + logical_k_block_idx / kRankBlocks; + const uint32_t rank_block_idx = + logical_k_block_idx % kRankBlocks; + k_idx = projection_idx * kSideLoraRank + + rank_block_idx * BLOCK_K; + } else if (block_phase == sched::BlockPhase::LoraL1Expand) { + k_idx += (n_block_idx & 1u) * kSideLoraRank; + } + + // Add 2 CTA offsets for non-leader CTA + if (not is_leader_cta) + load_m_idx += scheduler.template get_valid_m() / 2; + + // TMA copy tokens, then arrive at full barrier + if (cute::elect_one_sync()) { + tma::copy( + current_tensor_map_a_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_a[stage_idx], k_idx, load_m_idx, 2); + if (is_leader_cta) { + // Multicast + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(sizeof(shared_storage.smem_a[0]) * 2); + } else { + shared_storage.full_barriers[stage_idx].arrive(0u); + } + } + __syncwarp(); + } + }); + } else if (warp_idx == kNumDispatchWarps + 1) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM TMA load warp for weights + 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 cute::TmaDescriptor* tensor_map_b_ptr = + &tensor_map_l1_weights; + uint32_t shape_n = L1_SHAPE_N; + uint32_t expert_n_block_idx = n_block_idx; + if (block_phase == sched::BlockPhase::LoraL1Shrink) { + tensor_map_b_ptr = (n_block_idx & 1u) ? + &tensor_map_lora_a3 : &tensor_map_lora_a1; + shape_n = kSideLoraRank; + expert_n_block_idx = 0; + } else if (block_phase == sched::BlockPhase::LoraL1Expand) { + tensor_map_b_ptr = (n_block_idx & 1u) ? + &tensor_map_lora_b3 : &tensor_map_lora_b1; + shape_n = kIntermediateHidden; + expert_n_block_idx = n_block_idx / 2; + } else if (block_phase == sched::BlockPhase::LoraL2Shrink) { + tensor_map_b_ptr = &tensor_map_lora_a2; + shape_n = kSideLoraRank; + expert_n_block_idx = 0; + } else if (block_phase == sched::BlockPhase::LoraL2Expand) { + tensor_map_b_ptr = &tensor_map_lora_b2; + shape_n = kHidden; + } else if (block_phase == sched::BlockPhase::Linear2) { + tensor_map_b_ptr = &tensor_map_l2_weights; + shape_n = L2_SHAPE_N; + } + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); + + const uint32_t base_k_blocks = + block_phase == sched::BlockPhase::Linear1 ? + L1_SHAPE_K / BLOCK_K : + block_phase == sched::BlockPhase::Linear2 ? + L2_SHAPE_K / BLOCK_K : 0; + const bool is_linear_side_expand = kHasSideLora && + (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) && + k_block_idx >= base_k_blocks; + const uint32_t logical_k_block_idx = + is_linear_side_expand ? + k_block_idx - base_k_blocks : k_block_idx; + const cute::TmaDescriptor* current_tensor_map_b_ptr = + tensor_map_b_ptr; + uint32_t current_shape_n = shape_n; + uint32_t current_expert_n_block_idx = expert_n_block_idx; + if (is_linear_side_expand) { + if (block_phase == sched::BlockPhase::Linear1) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / BLOCK_K; + const uint32_t projection_idx = + logical_k_block_idx / kRankBlocks; + current_tensor_map_b_ptr = projection_idx ? + &tensor_map_lora_b3 : &tensor_map_lora_b1; + current_shape_n = kIntermediateHidden; + current_expert_n_block_idx = n_block_idx; + } else { + current_tensor_map_b_ptr = &tensor_map_lora_b2; + current_shape_n = kHidden; + } + } + + // Compute weight offset + const bool shared_side_weight = + block_phase == sched::BlockPhase::LoraL1Shrink || + (block_phase == sched::BlockPhase::Linear2 && + is_linear_side_expand); + uint32_t n_idx = + (shared_side_weight ? 0 : + local_expert_idx * current_shape_n) + + current_expert_n_block_idx * BLOCK_N; + uint32_t k_idx = logical_k_block_idx * BLOCK_K; + if (is_linear_side_expand && + block_phase == sched::BlockPhase::Linear1) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / BLOCK_K; + k_idx = (logical_k_block_idx % kRankBlocks) * BLOCK_K; + } + + // TMA copy weights + const bool is_l1_expand = is_linear_side_expand && + block_phase == sched::BlockPhase::Linear1; + if (is_l1_expand) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / BLOCK_K; + const uint32_t projection_idx = + logical_k_block_idx / kRankBlocks; + auto* side_tile = shared_storage.smem_b[stage_idx]; + for (uint32_t idx = lane_idx; + idx < LOAD_BLOCK_N * BLOCK_K; idx += 32) + side_tile[idx] = b_dtype_t(0.0f); + __syncwarp(); + cutlass::arch::fence_view_async_shared(); + if (cute::elect_one_sync()) { + constexpr uint32_t kGran = 8; + #pragma unroll + for (uint32_t chunk = 0; + chunk < BLOCK_N / (2 * kGran); ++chunk) { + tma::copy( + current_tensor_map_b_ptr, + &shared_storage.full_barriers[stage_idx], + side_tile + + (chunk * 2 * kGran + + projection_idx * kGran) * BLOCK_K, + k_idx, + local_expert_idx * kIntermediateHidden + + n_block_idx * (BLOCK_N / 2) + + chunk * kGran, + 2); + } + if (is_leader_cta) + shared_storage.full_barriers[stage_idx] + .arrive_and_expect_tx( + sizeof(shared_storage.smem_b[0])); + else + shared_storage.full_barriers[stage_idx].arrive(0u); + } + } else if (cute::elect_one_sync()) { + tma::copy( + current_tensor_map_b_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_b[stage_idx], k_idx, n_idx, 2); + if (is_leader_cta) { + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(sizeof(shared_storage.smem_b[0]) * 2); + } else { + shared_storage.full_barriers[stage_idx].arrive(0u); + } + } + __syncwarp(); + } + }); + } else if (warp_idx == kNumDispatchWarps + 2) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM MMA issue warp (only the leader CTA will run) + if (is_leader_cta) { + // Make instruction descriptor + // NOTES: always swap A/B + auto instr_desc = cute::UMMA::make_instr_desc(); + + DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); + uint32_t lane_stage = lane_idx < kNumStages ? lane_idx : 0u; + auto a_desc = mma::sm100::make_umma_desc(shared_storage.smem_a[lane_stage], 0, 0); + auto b_desc = mma::sm100::make_umma_desc(shared_storage.smem_b[lane_stage], 0, 0); + uint32_t a_desc_lo = a_desc.lo; + uint32_t b_desc_lo = b_desc.lo; + + // Checks for MMA instructions + DG_STATIC_ASSERT((UMMA_M == 64 and UMMA_N % 8 == 0 and 8 <= UMMA_N and UMMA_N <= 256) or + (UMMA_M == 128 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256) or + (UMMA_M == 256 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256), + "Invalid MMA instruction shape"); + + // Persistently schedule over blocks + uint32_t current_iter_idx = 0; + uint32_t side_empty_wait_phase[2] = {1, 1}; + 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) { + // Dynamic update of UMMA N based on effective M + mma::sm100::update_instr_desc_with_umma_n(instr_desc, scheduler.template get_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; + if constexpr (kHasSideLora) { + shared_storage.tmem_empty_barriers[accum_stage_idx] + .wait(side_empty_wait_phase[accum_stage_idx]); + side_empty_wait_phase[accum_stage_idx] ^= 1; + } else { + shared_storage.tmem_empty_barriers[accum_stage_idx] + .wait(accum_phase ^ 1); + } + ptx::tcgen05_after_thread_sync(); + + // Empty barrier arrival + auto empty_barrier_arrive = [&](const bool& do_tmem_full_arrive) { + auto umma_arrive = [](const uint64_t* barrier) { + constexpr uint16_t kCTAMask = (1 << 2) - 1; + cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); + }; + 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(&shared_storage.tmem_full_barriers[accum_stage_idx])); + __syncwarp(); + }; + + // Launch MMAs + #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 + 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(b_desc_lo, stage_idx); + const uint32_t base_k_blocks = + block_phase == sched::BlockPhase::Linear1 ? + L1_SHAPE_K / BLOCK_K : + block_phase == sched::BlockPhase::Linear2 ? + L2_SHAPE_K / BLOCK_K : 0; + const bool is_linear_side_expand = kHasSideLora && + (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) && + k_block_idx >= base_k_blocks; + const uint32_t logical_k_block_idx = + is_linear_side_expand ? + k_block_idx - base_k_blocks : k_block_idx; + if (is_linear_side_expand && + logical_k_block_idx == 0) { + const uint32_t side_stage_idx = + 1 - accum_stage_idx; + shared_storage.tmem_empty_barriers[side_stage_idx] + .wait(side_empty_wait_phase[side_stage_idx]); + side_empty_wait_phase[side_stage_idx] ^= 1; + ptx::tcgen05_after_thread_sync(); + } + const uint32_t accum_tmem_col = + is_linear_side_expand ? + (1 - accum_stage_idx) * UMMA_N : + accum_stage_idx * UMMA_N; + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t umma_k_block_idx = 0; umma_k_block_idx < BLOCK_K / UMMA_BLOCK_K; ++ umma_k_block_idx) { + // Issue UMMA + #pragma unroll + for (uint32_t k = 0; k < UMMA_BLOCK_K / UMMA_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, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_M, k * UMMA_K); + b_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, b_dtype_t>(b_desc_base_lo, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_N, k * UMMA_K); + ptx::SM100_MMA_F16BF16_2x1SM_SS::fma( + b_desc, a_desc, accum_tmem_col, + logical_k_block_idx > 0 or + umma_k_block_idx > 0 or k > 0, + static_cast(static_cast(instr_desc)) << 32); + } + } + } + __syncwarp(); + + // Commit to the mbarrier object + // No explicit `tcgen05.fence::before_thread_sync` is needed, as this is implicitly performed by `tcgen05.commit` + empty_barrier_arrive(k_block_idx == num_k_blocks - 1); + } + }); + + // To safely deconstruct barriers, we need another round of waits + if (current_iter_idx > 0) { + if constexpr (kHasSideLora) { + #pragma unroll + for (uint32_t i = 0; i < 2; ++i) + shared_storage.tmem_empty_barriers[i] + .wait(side_empty_wait_phase[i]); + } else { + const auto accum_phase_idx = + ((current_iter_idx - 1) / + kNumEpilogueStages) & 1; + shared_storage.tmem_empty_barriers[ + (current_iter_idx - 1) % + kNumEpilogueStages].wait(accum_phase_idx); + } + } + } + } else if (warp_idx == kNumDispatchWarps + 3) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + } else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_alloc(); + + // 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(&shared_storage.tmem_ptr_in_smem) == 0); + + // GEMM epilogue warps + const auto epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const auto epilogue_wg_idx = epilogue_warp_idx / 4; + const auto epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const auto warp_idx_in_wg = epilogue_warp_idx % 4; + DG_STATIC_ASSERT((kNumDispatchWarps + kNumMMANonEpilogueWarps) % 4 == 0 and + kNumEpilogueWarps % 4 == 0, "Invalid epilogue warps"); + + // TODO: support effective block M + // NOTES: + // - 2 warpgroups divide the whole BM into BM / 2 + // - 4 warps divide the whole BN into BN / 4 + // - BM / 2 is further divided into stored blocks, i.e. with `STORE_BLOCK_M` size + // - `STORE_BLOCK_M` in further divided into `ATOM_M` + constexpr uint32_t WG_BLOCK_M = BLOCK_M / kNumEpilogueWarpgroups; + constexpr uint32_t ATOM_M = 8; + constexpr uint32_t kNumBankGroupBytes = 16u; + constexpr uint32_t kNumAtomsPerStore = STORE_BLOCK_M / ATOM_M; + DG_STATIC_ASSERT(BLOCK_M % kNumEpilogueWarpgroups == 0, "Invalid block M"); + DG_STATIC_ASSERT(WG_BLOCK_M % STORE_BLOCK_M == 0, "Invalid warpgroup block M"); + DG_STATIC_ASSERT(STORE_BLOCK_M % ATOM_M == 0, "Invalid store block M"); + DG_STATIC_ASSERT(BLOCK_N == 128, "Invalid block N"); + + // Ensure the epilogue barrier cannot run with the pull barrier + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Persistently schedule over blocks + uint32_t current_iter_idx = 0; + 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) { + // Wait UMMA arrival + const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; + const auto accum_phase = (current_iter_idx ++ / kNumEpilogueStages) & 1; + shared_storage.tmem_full_barriers[accum_stage_idx].wait(accum_phase); + ptx::tcgen05_after_thread_sync(); + + // Compute offsets + // NOTES: use shuffle here to let NVCC know warp divergence won't happen + const uint32_t valid_m = ptx::exchange(scheduler.template get_valid_m(), 0); + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t ring_m_idx = ring_block_idx * BLOCK_M; // Ring-buffer offset for reusable data buffers + const uint32_t pool_m_idx = pool_block_idx * BLOCK_M; // Full-pool offset for non-ring metadata + uint32_t n_idx = n_block_idx * BLOCK_N; + auto release_tmem = [&]() { + ptx::tcgen05_before_thread_sync(); + shared_storage.tmem_empty_barriers[accum_stage_idx] + .arrive(0u); + if constexpr (kHasSideLora) { + if (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) + shared_storage.tmem_empty_barriers[ + 1 - accum_stage_idx].arrive(0u); + } + }; + + const bool is_lora_phase = + block_phase == sched::BlockPhase::LoraL1Shrink || + block_phase == sched::BlockPhase::LoraL2Shrink; + + if (is_lora_phase) { + // Materialize tensor-core side GEMMs at BF16 boundaries. The + // L2 shrink uses both CTAs for the 256-column UMMA shape, but + // the second half is a duplicate and is discarded. + const bool do_store = + block_phase != sched::BlockPhase::LoraL2Shrink || + n_block_idx == 0; + const cute::TmaDescriptor* tensor_map_d_ptr = + &tensor_map_lora_l1_scratch_store; + uint32_t out_n_idx = n_block_idx * BLOCK_N; + if (block_phase == sched::BlockPhase::LoraL2Shrink) { + tensor_map_d_ptr = &tensor_map_lora_l2_scratch_store; + out_n_idx = 0; + } + + if (!do_store) { + release_tmem(); + } else { + #pragma unroll + for (uint32_t s = 0; + s < WG_BLOCK_M / STORE_BLOCK_M; ++s) { + if (epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M >= valid_m) { + release_tmem(); + break; + } + + #pragma unroll + for (uint32_t i = 0; + i < STORE_BLOCK_M / ATOM_M; ++i) { + const uint32_t tmem_addr = + accum_stage_idx * UMMA_N + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + i * ATOM_M; + uint32_t values[ATOM_M]; + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + tmem_addr, values[0], values[1], + values[2], values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + tmem_addr | 0x00100000, values[4], + values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + + if (s == WG_BLOCK_M / STORE_BLOCK_M - 1 && + i == STORE_BLOCK_M / ATOM_M - 1) { + release_tmem(); + } + + const uint32_t row = lane_idx % 8; + const uint32_t col = + (warp_idx_in_wg % 2) * 4 + lane_idx / 8; + const auto smem_ptr = + shared_storage.smem_d.l2[epilogue_wg_idx] + + (warp_idx_in_wg / 2) * STORE_BLOCK_M * + (kSwizzleCDMode / sizeof(d_dtype_t)) + + i * ATOM_M * + (kSwizzleCDMode / sizeof(d_dtype_t)) + + row * ((kNumBankGroupBytes * 8) / + sizeof(d_dtype_t)) + + (col ^ row) * + (kNumBankGroupBytes / sizeof(d_dtype_t)); + ptx::SM90_U32x4_STSM_T::copy( + math::cast_into_bf16_and_pack( + values[0], values[1]), + math::cast_into_bf16_and_pack( + values[2], values[3]), + math::cast_into_bf16_and_pack( + values[4], values[5]), + math::cast_into_bf16_and_pack( + values[6], values[7]), + smem_ptr); + } + ptx::sync_aligned( + 128, kEpilogueWGBarrierStartIdx + + epilogue_wg_idx); + + if (warp_idx_in_wg == 0 && + cute::elect_one_sync()) { + cute::tma_store_fence(); + #pragma unroll + for (uint32_t atom = 0; + atom < BLOCK_N * sizeof(d_dtype_t) / + kSwizzleCDMode; + ++atom) { + cute::SM90_TMA_STORE_2D::copy( + tensor_map_d_ptr, + shared_storage.smem_d + .l2[epilogue_wg_idx] + + atom * STORE_BLOCK_M * + (kSwizzleCDMode / + sizeof(d_dtype_t)), + out_n_idx + + atom * (kSwizzleCDMode / + sizeof(d_dtype_t)), + pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M); + cute::tma_store_arrive(); + } + } + __syncwarp(); + // The next store slice reuses this warpgroup's same + // shared-memory tile. TMA stores are asynchronous, so + // wait for the issuing warp and then release all four + // producer warps before any of them overwrite it. + if (warp_idx_in_wg == 0) + cute::tma_store_wait<0>(); + ptx::sync_aligned( + 128, kEpilogueWGBarrierStartIdx + + epilogue_wg_idx); + } + } + + ptx::tma_store_wait<0>(); + ptx::sync_aligned( + kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (do_store && epilogue_warp_idx == 0 && + cute::elect_one_sync()) { + uint32_t ready_plane = 0; + if (block_phase == sched::BlockPhase::LoraL1Expand) + ready_plane = 1; + else if (block_phase == sched::BlockPhase::LoraL2Shrink) + ready_plane = 2; + else if (block_phase == sched::BlockPhase::LoraL2Expand) + ready_plane = 3; + ptx::red_add_rel( + reinterpret_cast(side_lora_ready) + + ready_plane * kNumRingBlocks + ring_block_idx, + 1u); + } + __syncwarp(); + } else if (block_phase == sched::BlockPhase::Linear1) { + // Wait L2 block empty + const auto l2_empty_ptr = workspace.get_l2_empty_count_ptr(ring_block_idx); + const auto num_expected_blocks = (L2_SHAPE_N / BLOCK_N) * (pool_block_idx / kNumRingBlocks); + while (ptx::ld_acq(l2_empty_ptr) != num_expected_blocks); + + // Unified L1 epilogue: gated activation 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]) + + // TopK weight for this lane + float stored_cached_weight = 0; + + #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) { + release_tmem(); + break; + } + + // Iterate all atoms in the store block + nv_bfloat162 bf16x2_output[kNumAtomsPerStore * 2]; + #pragma unroll + for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { + const uint32_t j = s * kNumAtomsPerStore + i; + + // Load weights from global into register cache per 32 tokens + DG_STATIC_ASSERT(32 % ATOM_M == 0, "Invalid block size"); + if ((j * ATOM_M) % 32 == 0 and (WG_BLOCK_M % 32 == 0 or j * ATOM_M + lane_idx < WG_BLOCK_M)) { + stored_cached_weight = *l1_topk_weights_buffer + .get_data_buffer(ring_m_idx + epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M + lane_idx) + .template get_base_ptr(); + } + + // Load weights from register cache + const 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) + }; + + // Load from TMEM + uint32_t tmem_addr = accum_stage_idx * UMMA_N + epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M; + uint32_t values[ATOM_M]; + uint32_t side_values[ATOM_M] = {}; + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, + values[0], values[1], values[2], values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, + values[4], values[5], values[6], values[7]); + if constexpr (kHasSideLora) { + const uint32_t side_tmem_addr = + (1 - accum_stage_idx) * UMMA_N + + epilogue_wg_idx * WG_BLOCK_M + + j * ATOM_M; + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr, side_values[0], + side_values[1], side_values[2], + side_values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr | 0x00100000, + side_values[4], side_values[5], + side_values[6], side_values[7]); + } + cutlass::arch::fence_view_async_tmem_load(); + + // Signal tensor memory consumed on the last atom + if (j == WG_BLOCK_M / ATOM_M - 1) { + release_tmem(); + } + + // Apply gated activation: act(gate) * up. + // Gate/up pairs: (0, 2), (1, 3), (4, 6), (5, 7) + auto fp32_values = reinterpret_cast(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 (kHasSideLora) { + const auto side_fp32_values = + reinterpret_cast(side_values); + const auto gate_side = + __bfloat1622float2( + __float22bfloat162_rn(make_float2( + side_fp32_values[k * 4], + side_fp32_values[k * 4 + 1]))); + const auto up_side = + __bfloat1622float2( + __float22bfloat162_rn(make_float2( + side_fp32_values[k * 4 + 2], + side_fp32_values[k * 4 + 3]))); + const auto gate_base = + __bfloat1622float2(bf16_gate); + const auto up_base = + __bfloat1622float2(bf16_up); + bf16_gate = __float22bfloat162_rn( + __fadd2_rn( + gate_base, + __fmul2_rn( + {side_lora_scale, + side_lora_scale}, + gate_side))); + bf16_up = __float22bfloat162_rn( + __fadd2_rn( + up_base, + __fmul2_rn( + {side_lora_scale, + side_lora_scale}, + up_side))); + } + + if constexpr (kSaveL1Preact) { + // Persist exact BF16 W13 output before clamp in + // the caller-owned full-pool tensor. Convert + // from the interleaved MMA column order back to + // [gate | up]. + const uint32_t hidden_col = + warp_idx_in_wg * 16 + + (lane_idx / 4) * 2 + k; + const uint32_t chunk = hidden_col / 8; + const uint32_t in_chunk = hidden_col & 7; + const uint32_t gate_col = + chunk * 16 + in_chunk; + const auto output_col = [=](uint32_t col) { + const uint32_t low = col & 31; + return n_idx + (col & ~31u) + + ((low & 1) << 4) + + ((low >> 1) & 3) + + (low & 8) + + ((low & 16) >> 2); + }; + const uint32_t physical_gate_col = + output_col(gate_col); + const uint32_t deinterleaved_col = + (physical_gate_col / 16) * 8 + + (physical_gate_col & 7); + const uint32_t row_base = + pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + + i * ATOM_M + + (lane_idx % 4) * 2; + const uint32_t gate_bits = + *reinterpret_cast(&bf16_gate); + const uint32_t up_bits = + *reinterpret_cast(&bf16_up); + #pragma unroll + for (uint32_t r = 0; r < 2; ++ r) { + const uint32_t m_out = row_base + r; + if ( + m_out < pool_m_idx + valid_m && + m_out < num_saved_pool_tokens + ) { + auto* dst = reinterpret_cast( + saved_l1_preact + + static_cast(m_out) * + (2 * kIntermediateHidden)); + dst[deinterleaved_col] = + static_cast( + gate_bits >> (r * 16)); + dst[kIntermediateHidden + + deinterleaved_col] = + static_cast( + up_bits >> (r * 16)); + } + } + } + + // Clamp + if constexpr (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}); + } + + const auto gate = __bfloat1622float2(bf16_gate); + const auto up = __bfloat1622float2(bf16_up); + float2 z; + if constexpr (kActivationType == ActivationType::GeGLU) { + constexpr float kAlpha = 1.5957691216057308f; + constexpr float kBeta = 0.044715f; + const auto gate_sq = __fmul2_rn(gate, gate); + z = __fmul2_rn( + __fmul2_rn( + {kAlpha, kAlpha}, gate), + __fadd2_rn( + {1.0f, 1.0f}, + __fmul2_rn( + {kBeta, kBeta}, gate_sq))); + } else { + z = gate; + } + + const auto neg_exp = make_float2( + kFastMath ? __expf(-z.x) : expf(-z.x), + kFastMath ? __expf(-z.y) : expf(-z.y)); + const auto denom = + __fadd2_rn({1.0f, 1.0f}, neg_exp); + float2 activated; + if constexpr (kFastMath) { + activated = __fmul2_rn( + gate, + {math::fast_rcp(denom.x), + math::fast_rcp(denom.y)}); + } else { + if constexpr ( + kActivationType == + ActivationType::SwiGLU) { + // CUDA aten::silu evaluates x / (1 + + // exp(-x)) directly. Multiplying x by a + // separately rounded reciprocal differs + // at BF16 product ties (for example + // x=0.78515625). + activated = { + gate.x / denom.x, + gate.y / denom.y}; + } else { + // The reference GeGLU expression materializes + // sigmoid(z) before multiplying by gate. + const float2 sig = { + 1.0f / denom.x, + 1.0f / denom.y}; + activated = + __fmul2_rn(gate, sig); + } + } + // The unclamped SwiGLU contract materializes + // F.silu(gate_bf16) before the separate BF16 + // multiply by up_bf16. Clamped SwiGLU and + // GeGLU instead evaluate activation*up in FP32 and + // round only the product. + const auto activated_for_mul = + kActivationType == + ActivationType::SwiGLU && + kActivationClamp == + cute::numeric_limits< + float>::infinity() + ? __bfloat1622float2( + __float22bfloat162_rn( + activated)) + : activated; + const auto h_fp32 = + __fmul2_rn( + activated_for_mul, up); + const auto h_bf16 = + __float22bfloat162_rn(h_fp32); + // The clamped path applies its pre-down score while + // clamped activation is still FP32, then rounds the + // weighted W2 input once. Other model families + // materialize the activation as BF16 first. + const auto h_for_weight = + kRouteWeightMode == + RouteWeightMode::PreDown && + kActivationClamp != + cute::numeric_limits< + float>::infinity() + ? h_fp32 + : __bfloat1622float2(h_bf16); + const auto h_weighted_bf16 = + __float22bfloat162_rn( + __fmul2_rn( + h_for_weight, + weights)); + if constexpr ( + kRouteWeightMode == + RouteWeightMode::PreDown) { + bf16x2_output[i * 2 + k] = + h_weighted_bf16; + } else { + bf16x2_output[i * 2 + k] = + h_bf16; + } + + if constexpr (kSaveStageActivations) { + const uint32_t hidden_col = + warp_idx_in_wg * 16 + + (lane_idx / 4) * 2 + k; + const uint32_t chunk = + hidden_col / 8; + const uint32_t in_chunk = + hidden_col & 7; + const uint32_t interleaved_col = + chunk * 16 + in_chunk; + const uint32_t low = + interleaved_col & 31; + const uint32_t physical_col = + n_idx + + (interleaved_col & ~31u) + + ((low & 1) << 4) + + ((low >> 1) & 3) + + (low & 8) + + ((low & 16) >> 2); + const uint32_t canonical_col = + (physical_col / 16) * 8 + + (physical_col & 7); + const uint32_t row_base = + pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + + i * ATOM_M + + (lane_idx % 4) * 2; + const uint32_t h_bits = + *reinterpret_cast< + const uint32_t*>( + &h_bf16); + const uint32_t weighted_bits = + *reinterpret_cast< + const uint32_t*>( + &h_weighted_bf16); + #pragma unroll + for (uint32_t r = 0; + r < 2; ++r) { + const uint32_t m_out = + row_base + r; + if ( + m_out < pool_m_idx + valid_m && + m_out < num_saved_pool_tokens + ) { + const uint64_t output_idx = + static_cast( + m_out) * + kIntermediateHidden + + canonical_col; + *reinterpret_cast< + uint16_t*>( + saved_h_unweighted + + output_idx) = + static_cast( + h_bits >> + (r * 16)); + *reinterpret_cast< + uint16_t*>( + saved_h_weighted + + output_idx) = + static_cast( + weighted_bits >> + (r * 16)); + } + } + } + } + } + + // Wait shared memory release from previous TMA store + const uint32_t tma_stage_idx = s % kNumTMAStoreStages; + ptx::tma_store_wait(); + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + #pragma unroll + for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { + // STSM + uint32_t row = lane_idx % 8; + uint32_t col = warp_idx_in_wg * 2 + lane_idx / 8; + const auto smem_ptr = shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx] + + (i * ATOM_M + row) * L1_OUT_BLOCK_N + + (col ^ row) * (kNumBankGroupBytes / sizeof(d_dtype_t)); + ptx::SM90_U32x2_STSM_T<__nv_bfloat162>::copy( + bf16x2_output[i * 2 + 0], + bf16x2_output[i * 2 + 1], + smem_ptr + ); + } + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Issue TMA store after all atoms in this store block + if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + 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, + shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx], + out_n_idx, + ring_m_idx + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M); + cute::tma_store_arrive(); + } + __syncwarp(); + } + + // Notify L2 and increment L1 empty count + // TODO: less epilogue sync scope + ptx::tma_store_wait<0>(); + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + workspace.get_l2_full_count_ptr(ring_block_idx), 1u); + + // Increment L1 empty count for this physical slot (one per N block) + ptx::red_add( + workspace.get_l1_empty_count_ptr(ring_block_idx), 1u); + } + __syncwarp(); + } else { + // Increment L2 empty count for this physical slot (one per N block) + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_add( + workspace.get_l2_empty_count_ptr(ring_block_idx), 1u); + } + __syncwarp(); + + DG_STATIC_ASSERT(STORE_BLOCK_M % 8 == 0, "Invalid store M"); + constexpr uint32_t kNumRowsPerWarp = STORE_BLOCK_M / 8; + + // 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) { + // Early break if the entire store block is beyond the valid token range + // TODO: check performance + if (epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M >= valid_m) { + release_tmem(); + break; + } + + #pragma unroll + for (uint32_t i = 0; i < STORE_BLOCK_M / ATOM_M; ++ i) { + // Load from TMEM using .16x256b shape to satisfy STSM layout requirements + // Start from lane index 0 and 16 + uint32_t tmem_addr = accum_stage_idx * UMMA_N + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M; + uint32_t values[ATOM_M]; + uint32_t side_values[ATOM_M] = {}; + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, + values[0], values[1], values[2], values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, + values[4], values[5], values[6], values[7]); + if constexpr (kHasSideLora) { + const uint32_t side_tmem_addr = + (1 - accum_stage_idx) * UMMA_N + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + i * ATOM_M; + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr, side_values[0], + side_values[1], side_values[2], + side_values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr | 0x00100000, + side_values[4], side_values[5], + side_values[6], side_values[7]); + } + cutlass::arch::fence_view_async_tmem_load(); + + if constexpr (kHasSideLora) { + auto* base_fp32 = + reinterpret_cast(values); + const auto* side_fp32 = + reinterpret_cast(side_values); + #pragma unroll + for (uint32_t value_idx = 0; + value_idx < ATOM_M; ++value_idx) { + base_fp32[value_idx] = + __bfloat162float(__float2bfloat16_rn( + base_fp32[value_idx])) + + side_lora_scale * + __bfloat162float(__float2bfloat16_rn( + side_fp32[value_idx])); + } + } + + // 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) + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Signal tensor memory consumed + if (s == WG_BLOCK_M / STORE_BLOCK_M - 1 and i == STORE_BLOCK_M / ATOM_M - 1) { + release_tmem(); + } + + // Store into shared memory + // 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 = (warp_idx_in_wg % 2) * 4 + lane_idx / 8; + const auto smem_ptr = shared_storage.smem_d.l2[epilogue_wg_idx] + + (warp_idx_in_wg / 2) * STORE_BLOCK_M * (kSwizzleCDMode / sizeof(d_dtype_t)) + + i * ATOM_M * (kSwizzleCDMode / sizeof(d_dtype_t)) + + row * ((kNumBankGroupBytes * 8) / sizeof(d_dtype_t)) + + (col ^ row) * (kNumBankGroupBytes / sizeof(d_dtype_t)); + ptx::SM90_U32x4_STSM_T::copy( + math::cast_into_bf16_and_pack(values[0], values[1]), + math::cast_into_bf16_and_pack(values[2], values[3]), + math::cast_into_bf16_and_pack(values[4], values[5]), + math::cast_into_bf16_and_pack(values[6], values[7]), + smem_ptr + ); + } + + // Wait shared memory ready + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + if constexpr (kSaveDownUnweighted) { + const uint32_t saved_store_row = + pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M; + // A cached active-pool plan may underestimate a later + // routing histogram. The caller publishes the required + // size through route_count_mismatch and replays the + // capture before backward; do not turn that recoverable + // cache miss into an out-of-bounds TMA store. + if (saved_store_row + STORE_BLOCK_M <= + num_saved_pool_tokens) { + if (warp_idx_in_wg == 0 && + cute::elect_one_sync()) { + cute::tma_store_fence(); + #pragma unroll + for (uint32_t atom = 0; + atom < + BLOCK_N * sizeof(d_dtype_t) / + kSwizzleCDMode; + ++atom) { + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_down_unweighted, + shared_storage.smem_d + .l2[epilogue_wg_idx] + + atom * STORE_BLOCK_M * + (kSwizzleCDMode / + sizeof(d_dtype_t)), + n_idx + + atom * + (kSwizzleCDMode / + sizeof(d_dtype_t)), + saved_store_row); + cute::tma_store_arrive(); + } + } + if (warp_idx_in_wg == 0) { + cute::tma_store_wait<0>(); + } + __syncwarp(); + } + } + + // Write into remote buffers + // 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; + + #pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++ j) { + const uint32_t row_in_store = j * 8 + warp_idx_in_wg * 2 + lane_idx / 16; + const uint32_t m_idx_in_block = epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + row_in_store; + + // Skip padding rows beyond the actual token count for this expert + if (m_idx_in_block >= valid_m) + break; + + const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_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; + + // Read from shared memory + const auto smem_ptr = shared_storage.smem_d.l2[epilogue_wg_idx] + + (lane_idx % 16 / 8) * STORE_BLOCK_M * (kSwizzleCDMode / sizeof(d_dtype_t)) + + row_in_store * (kSwizzleCDMode / sizeof(d_dtype_t)) + + (bank_group_idx ^ row_in_atom) * (kNumBankGroupBytes / sizeof(d_dtype_t)); + auto packed = ptx::ld_shared( + reinterpret_cast(smem_ptr)); + if constexpr ( + kHasSideLora && + kSaveDownUnweighted) { + // The regular TMA save above captured the base + // down projection before the side update. Replace + // this vector with the combined BF16 boundary used + // by post-down router-gradient backward. + if ( + pool_m_idx + m_idx_in_block < + num_saved_pool_tokens) { + auto* saved_ptr = + saved_down_unweighted + + static_cast( + pool_m_idx + + m_idx_in_block) * + kHidden + + n_idx + + (lane_idx % 16) * 8; + *reinterpret_cast(saved_ptr) = + packed; + } + } + if constexpr ( + kRouteWeightMode == + RouteWeightMode::PostDown && + kCombineOrderMode == + CombineOrderMode::FixedTopK) { + // Keep forward offsets unchanged: recover the + // immutable source token/slot score directly from + // the existing symmetric input plane. + const float route_weight = + *sym_buffer.map( + input_topk_weights_buffer + .get_base_ptr() + + static_cast( + src_metadata.token_idx) * + kNumTopk + + src_metadata.topk_idx, + src_metadata.rank_idx); + auto* values = + reinterpret_cast( + &packed); + #pragma unroll + for (uint32_t value_idx = 0; + value_idx < 8; ++value_idx) { + values[value_idx] = + __float2bfloat16_rn( + __bfloat162float( + values[value_idx]) * + route_weight); + } + } + + // Write into remote + const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_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))); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } + } + + // Ensure the next epilogue safe to use shared memory + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + }); + + // Deallocate tensor memory + // NOTES: must be called by the same logical warp ID on both CTAs + if (epilogue_warp_idx == 0) + Allocator().free(0, kNumTmemCols); + + // NVLink barrier (grid sync + cross-rank signal + grid sync): ~4 us + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + + // Barrier with dispatch warps, so that they can do clean workspace + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Combine: reduce top-k results and write back + // NOTES: reuse shared memory from start up to the barriers + // 1 token, 1 topk latency: ~3 us + constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); + + // 3 slots of chunk is needed: 2 load stages and 1 store + constexpr uint32_t kNumChunkSlots = 3; + constexpr uint32_t kNumMaxRegistersForBuffer = 128; + + // NOTES: either 1 or 2 chunks for simplicity + // NOTES: Restrict on both smem and register + constexpr uint32_t kNumChunks = + 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 <= 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 <= 32, "Top-k must fit in a single warp"); + + // Verify combined shared memory budget at runtime + DG_DEVICE_ASSERT(kNumChunkSlots * kNumEpilogueWarps * kNumChunkBytes <= + static_cast(offsetof(SharedStorage, dispatch_barriers))); + + // Per-warp buffer: 2 stage load buffers + 1 store buffer + const auto combine_load_buffer = utils::PatternVisitor([&](const uint32_t& i) { + return math::advance_ptr(smem_buffer, (epilogue_warp_idx + i * kNumEpilogueWarps) * kNumChunkBytes); + }); + const auto combine_store_buffer = math::advance_ptr(smem_buffer, (epilogue_warp_idx + kNumEpilogueWarps * 2) * kNumChunkBytes); + + // Per-warp barriers + auto combine_load_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return &shared_storage.combine_barriers[i + epilogue_warp_idx * 2]; + }); + + // Iterate over all tokens + uint32_t combine_phase = 0; + uint32_t load_stage_idx = 0; + for (uint32_t token_idx = sm_idx * kNumEpilogueWarps + epilogue_warp_idx; + token_idx < num_tokens; + token_idx += kNumSMs * kNumEpilogueWarps) { + // Read top-k slot indices: each lane reads one slot, then broadcast via exchange + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + const int stored_topk_slot_idx = lane_idx < kNumTopk ? + static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; + const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); + + // Iterate all chunks + for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { + const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; + + // Accumulate all top-k contributions for this chunk in float registers + float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {}; + if constexpr ( + kCombineOrderMode == + CombineOrderMode::FixedTopK) { + // Move mask and load + uint32_t mask = total_mask; + const auto move_mask_and_load = + [&](const uint32_t& i) { + if (mask) { + // Move + const uint32_t slot_idx = + __ffs(mask) - 1; + mask ^= 1 << slot_idx; + + // Load + if (cute::elect_one_sync()) { + const auto src_ptr = + math::advance_ptr( + combine_token_buffer + .get_rank_buffer( + slot_idx) + .get_data_buffer( + token_idx) + .get_base_ptr(), + chunk_byte_offset); + ptx::tma_load_1d( + combine_load_buffer[i], + src_ptr, + combine_load_barriers[i], + kNumChunkBytes); + ptx::mbarrier_arrive_and_set_tx( + combine_load_barriers[i], + kNumChunkBytes); + } + __syncwarp(); + return true; + } + return false; + }; + + // Load the first selection + bool do_reduce = + move_mask_and_load(load_stage_idx); + while (do_reduce) { + // Prefetch next top-k while accumulating current. + do_reduce = move_mask_and_load( + load_stage_idx ^ 1); + + combine_load_barriers[load_stage_idx]->wait( + combine_phase); + #pragma unroll + for (uint32_t j = 0; + j < kNumUint4PerLane; ++j) { + const auto uint4_values = + combine_load_buffer[ + load_stage_idx][ + j * 32 + lane_idx]; + const auto bf16_values = + reinterpret_cast< + const nv_bfloat162*>( + &uint4_values); + #pragma unroll + for (uint32_t l = 0; + l < kNumElemsPerUint4; ++l) { + ptx::accumulate( + reduced[ + j * + kNumElemsPerUint4 + + l], + bf16_values[l]); + } + } + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } + } else { + // Deterministic fixed-top-k combine paths both + // reduce each destination-rank partial in source top-k + // slot order. V1 then uses rank order while V2 uses + // last-slot rank order. + constexpr uint32_t kNumExpertsPerRank = + kNumExperts / kNumRanks; + const int dst_rank = + stored_topk_slot_idx >= 0 + ? stored_topk_slot_idx / + static_cast( + kNumExpertsPerRank) + : -1; + const uint32_t same_rank_mask = + __match_any_sync( + 0xffffffff, dst_rank); + uint32_t rank_master_mask = + __ballot_sync( + 0xffffffff, + stored_topk_slot_idx >= 0 && + lane_idx == + (kCombineOrderMode == + CombineOrderMode:: + DeepEPV1 + ? __ffs( + same_rank_mask) - + 1 + : 31 - + __clz( + same_rank_mask))); + while (rank_master_mask) { + uint32_t rank_master_slot = + __ffs(rank_master_mask) - 1; + if constexpr ( + kCombineOrderMode == + CombineOrderMode::DeepEPV1) { + int selected_rank = + ptx::exchange( + dst_rank, + rank_master_slot); + uint32_t candidates = + rank_master_mask & + ~(1u << + rank_master_slot); + while (candidates) { + const uint32_t + candidate_slot = + __ffs(candidates) - + 1; + candidates &= + candidates - 1; + const int candidate_rank = + ptx::exchange( + dst_rank, + candidate_slot); + if (candidate_rank < + selected_rank) { + selected_rank = + candidate_rank; + rank_master_slot = + candidate_slot; + } + } + } + rank_master_mask &= + ~(1u << rank_master_slot); + const int current_rank = + ptx::exchange( + dst_rank, + rank_master_slot); + uint32_t rank_mask = __ballot_sync( + 0xffffffff, + dst_rank == current_rank); + + float2 rank_reduced[ + kNumUint4PerLane * + kNumElemsPerUint4] = {}; + while (rank_mask) { + uint32_t slot_idx = + __ffs(rank_mask) - 1; + rank_mask &= ~(1u << slot_idx); + if (cute::elect_one_sync()) { + const auto src_ptr = + math::advance_ptr( + combine_token_buffer + .get_rank_buffer( + slot_idx) + .get_data_buffer( + token_idx) + .get_base_ptr(), + chunk_byte_offset); + ptx::tma_load_1d( + combine_load_buffer[0], + src_ptr, + combine_load_barriers[0], + kNumChunkBytes); + ptx::mbarrier_arrive_and_set_tx( + combine_load_barriers[0], + kNumChunkBytes); + } + __syncwarp(); + combine_load_barriers[0]->wait( + combine_phase); + #pragma unroll + for (uint32_t j = 0; + j < kNumUint4PerLane; ++j) { + const auto uint4_values = + combine_load_buffer[0][ + j * 32 + lane_idx]; + const auto bf16_values = + reinterpret_cast< + const nv_bfloat162*>( + &uint4_values); + const float route_weight = + kRouteWeightMode == + RouteWeightMode:: + PostDown + ? input_topk_weights_buffer + .get_data_buffer( + token_idx) + .template + get_base_ptr< + float>()[slot_idx] + : 1.0f; + #pragma unroll + for (uint32_t l = 0; + l < kNumElemsPerUint4; + ++l) { + if constexpr ( + kRouteWeightMode == + RouteWeightMode:: + PostDown) { + const float2 source = + __bfloat1622float2( + bf16_values[l]); + auto& value = + rank_reduced[ + j * + kNumElemsPerUint4 + + l]; + value.x = __fmaf_rn( + source.x, + route_weight, + value.x); + value.y = __fmaf_rn( + source.y, + route_weight, + value.y); + } else { + ptx::accumulate( + rank_reduced[ + j * + kNumElemsPerUint4 + + l], + bf16_values[l]); + } + } + } + combine_phase ^= 1; + } + + #pragma unroll + for (uint32_t value_idx = 0; + value_idx < + kNumUint4PerLane * + kNumElemsPerUint4; + ++value_idx) { + const nv_bfloat162 rank_partial = + __float22bfloat162_rn( + rank_reduced[value_idx]); + ptx::accumulate( + reduced[value_idx], + rank_partial); + } + } + } + + // Cast + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); + + // Wait share memory release and write + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); + } + __syncwarp(); + + // TMA store the token chunk + if (cute::elect_one_sync()) { + cute::tma_store_fence(); + ptx::tma_store_1d( + math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + chunk_byte_offset), + combine_store_buffer, kNumChunkBytes); + cute::tma_store_arrive(); + } + __syncwarp(); + } + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_100f"); +#endif +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.cuh new file mode 100644 index 0000000000..3c9073331d --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_side_lora_forward.cuh @@ -0,0 +1,2230 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +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 STORE_BLOCK_M, + uint32_t SF_BLOCK_M, uint32_t SF_BLOCK_N, + uint32_t kNumRingTokens, + uint32_t kNumSFRingTokens, + uint32_t kNumStages, + uint32_t kNumBytesPerPull, + uint32_t kNumDispatchThreads, uint32_t kNumNonEpilogueThreads, + uint32_t kNumEpilogueThreads, + uint32_t kNumSMs, uint32_t kNumRanks, + float kActivationClamp, + bool kFastMath, + ActivationType kActivationType, + bool kSaveL1Preact, + RouteWeightMode kRouteWeightMode = RouteWeightMode::PreDown, + bool kSaveDownUnweighted = false, + uint32_t kSideLoraRank = 0, + uint32_t L1_SHAPE_N = kIntermediateHidden * 2, + uint32_t L1_SHAPE_K = kHidden, + uint32_t L2_SHAPE_N = kHidden, + uint32_t L2_SHAPE_K = kIntermediateHidden, + uint32_t kNumDispatchWarps = kNumDispatchThreads / 32, + uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32, + uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32, + uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4, + uint32_t kNumThreads = kNumDispatchThreads + kNumNonEpilogueThreads + kNumEpilogueThreads, + uint32_t kNumTokensPerWarp = 32 / kNumTopk, + uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks, + uint32_t kNumRingBlocks = kNumRingTokens / BLOCK_M +> +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void +sm100_fp8_fp4_mega_moe_side_lora_forward_impl(void* y, + nv_bfloat16* saved_l1_preact, + nv_bfloat16* saved_x, + nv_bfloat16* saved_h, + nv_bfloat16* saved_down_unweighted, + int* side_lora_ready, + const float side_lora_scale, + int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, + const uint32_t num_saved_pool_tokens, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_down_unweighted, + const __grid_constant__ cute::TmaDescriptor tensor_map_saved_x, + const __grid_constant__ cute::TmaDescriptor tensor_map_saved_h, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_a1, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_a3, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_b1, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_b3, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_a2, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_b2, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l1_scratch, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l2_scratch, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l1_scratch_store, + const __grid_constant__ cute::TmaDescriptor tensor_map_lora_l2_scratch_store) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Allocator = cute::TMEM::Allocator2Sm; + + // Template checks + DG_STATIC_ASSERT(kNumDispatchThreads % 128 == 0, "Invalid number of dispatch threads"); + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 128, "Invalid number of MMA non-epilogue threads"); + DG_STATIC_ASSERT(kNumEpilogueThreads % 128 == 0, "Invalid number of MMA epilogue and combine threads"); + DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + constexpr bool kHasSideLora = kSideLoraRank > 0; + DG_STATIC_ASSERT(!kHasSideLora || kSideLoraRank == 128, + "MXFP4 side LoRA requires rank 128"); + + // Thread indices + const bool is_leader_cta = cute::block_rank_in_cluster() == 0; + const uint32_t sm_idx = blockIdx.x; + const uint32_t thread_idx = threadIdx.x; + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == 0) { + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights_sf); + if constexpr (kSaveDownUnweighted) + cute::prefetch_tma_descriptor( + &tensor_map_down_unweighted); + if constexpr (kHasSideLora) { + cute::prefetch_tma_descriptor(&tensor_map_saved_x); + cute::prefetch_tma_descriptor(&tensor_map_saved_h); + cute::prefetch_tma_descriptor(&tensor_map_lora_a1); + cute::prefetch_tma_descriptor(&tensor_map_lora_a3); + cute::prefetch_tma_descriptor(&tensor_map_lora_b1); + cute::prefetch_tma_descriptor(&tensor_map_lora_b3); + cute::prefetch_tma_descriptor(&tensor_map_lora_a2); + cute::prefetch_tma_descriptor(&tensor_map_lora_b2); + cute::prefetch_tma_descriptor(&tensor_map_lora_l1_scratch); + cute::prefetch_tma_descriptor(&tensor_map_lora_l2_scratch); + cute::prefetch_tma_descriptor( + &tensor_map_lora_l1_scratch_store); + cute::prefetch_tma_descriptor( + &tensor_map_lora_l2_scratch_store); + } + } + + // Workspaces + const auto workspace = layout::Workspace( + sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk, kNumRingTokens); + + // Token and buffer layouts + constexpr auto fp8_token_layout = layout::Data(kHidden); + constexpr auto bf16_token_layout = layout::Data(kHidden * sizeof(nv_bfloat16)); + constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden); + constexpr auto fp8_sf_layout = layout::Data(kHidden / 32); + constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 32); + constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * sizeof(int64_t), false); + constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * sizeof(float), false); + constexpr auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + + // Registered inputs + const auto input_token_buffer = layout::Buffer( + fp8_token_layout, 1, kNumMaxTokensPerRank, + workspace.get_end_ptr()); + const auto input_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, kNumMaxTokensPerRank, + input_token_buffer.get_end_ptr()); + const auto input_topk_idx_buffer = layout::Buffer( + input_topk_idx_layout, 1, kNumMaxTokensPerRank, + input_sf_buffer.get_end_ptr()); + const auto input_topk_weights_buffer = layout::Buffer( + input_topk_weights_layout, 1, kNumMaxTokensPerRank, + input_topk_idx_buffer.get_end_ptr()); + + // SF and its buffer configs + constexpr uint32_t kGranK = 32; + 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"); + + // UTCCP 4x32 transpose index mapping within each 128-element group + const auto transform_sf_token_idx = [](const uint32_t& token_idx_in_expert) { + const uint32_t idx = token_idx_in_expert % BLOCK_M; + return token_idx_in_expert / BLOCK_M * SF_BLOCK_M + + (idx & ~127u) + (idx & 31u) * 4 + ((idx >> 5) & 3u); + }; + + // L1 inputs + const auto l1_token_buffer = layout::Buffer( + fp8_token_layout, 1, kNumRingTokens, + input_topk_weights_buffer.get_end_ptr()); + const auto l1_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, kNumSFRingTokens, + l1_token_buffer.get_end_ptr()); + const auto l1_topk_weights_buffer = layout::Buffer( + l1_topk_weights_layout, 1, kNumRingTokens, + l1_sf_buffer.get_end_ptr()); + + // L2 inputs + const auto l2_token_buffer = layout::Buffer( + fp8_intermediate_token_layout, 1, kNumRingTokens, + l1_topk_weights_buffer.get_end_ptr() + ); + const auto l2_sf_buffer = layout::Buffer( + fp8_intermediate_sf_layout, 1, kNumSFRingTokens, + l2_token_buffer.get_end_ptr() + ); + + // Combine inputs + const auto combine_token_buffer = layout::Buffer( + bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, + l2_sf_buffer.get_end_ptr() + ); + const auto backward_route_grad_buffer = layout::Buffer( + input_topk_weights_layout, 1, kNumMaxTokensPerRank, + combine_token_buffer.get_end_ptr()); + const auto side_lora_source_buffer = layout::Buffer( + bf16_token_layout, 1, kNumMaxTokensPerRank, + backward_route_grad_buffer.get_end_ptr()); + + // 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; + + // 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 + constexpr uint32_t UMMA_BLOCK_K = 128; + constexpr uint32_t UMMA_K = 32; + constexpr uint32_t SIDE_BLOCK_K = 64; + constexpr uint32_t SIDE_UMMA_K = 16; + 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"); + + // Swizzle configs + 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"); + + // Epilogue configs + constexpr uint32_t kNumEpilogueStages = 2; + constexpr uint32_t kNumTMAStoreStages = 2; + + // Shared memory + constexpr uint32_t kSharedMemoryAlignment = 1024; + extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[]; + + // 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; + constexpr uint32_t AMAX_REDUCTION_WARP_BUFFER_SIZE = STORE_BLOCK_M / 2; // float2 + + struct SharedStorage { + union AStage { + alignas(kSharedMemoryAlignment) + a_dtype_t base[LOAD_BLOCK_M * BLOCK_K]; + alignas(kSharedMemoryAlignment) + cutlass::bfloat16_t side[LOAD_BLOCK_M * SIDE_BLOCK_K]; + }; + union BStage { + alignas(kSharedMemoryAlignment) + b_dtype_t base[LOAD_BLOCK_N * BLOCK_K]; + alignas(kSharedMemoryAlignment) + cutlass::bfloat16_t side[LOAD_BLOCK_N * SIDE_BLOCK_K]; + }; + alignas(kSharedMemoryAlignment) uint32_t expert_token_count[kNumExperts]; + alignas(kSharedMemoryAlignment) uint8_t dispatch_send_buffer[kNumDispatchWarps][kNumBytesPerPull]; + union { + alignas(kSharedMemoryAlignment) cutlass::float_e4m3_t l1[kNumEpilogueWarpgroups][kNumTMAStoreStages][STORE_BLOCK_M * L1_OUT_BLOCK_N]; + alignas(kSharedMemoryAlignment) nv_bfloat16 l2[kNumEpilogueWarpgroups][STORE_BLOCK_M * BLOCK_N]; + } smem_d; + alignas(kSharedMemoryAlignment) AStage smem_a[kNumStages]; + alignas(kSharedMemoryAlignment) BStage smem_b[kNumStages]; + uint32_t smem_sfa[kNumStages][SF_BLOCK_M * (BLOCK_K / 128)]; + uint32_t smem_sfb[kNumStages][SF_BLOCK_N * (BLOCK_K / 128)]; + float2 amax_reduction[kNumEpilogueWarps][AMAX_REDUCTION_WARP_BUFFER_SIZE]; + 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]; + 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; + constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; + constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; + constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; + DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); + + // 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()) { + // The bytes must be 8 bytes aligned + 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) + 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 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 + shared_storage.tmem_full_barriers[i].init(1); + // Arrive only at the leader CTA + shared_storage.tmem_empty_barriers[i].init(2 * kNumEpilogueThreads); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++ i) + shared_storage.combine_barriers[i].init(1); + } + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 3) { + // Allocate tensor memory + 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` + comm::cluster_sync_with_relaxed_arrive(); + + // Task scheduler + auto scheduler = sched::SideLoraMegaMoEScheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, + kNumExpertsPerWave, + kNumSMs, kNumRanks, kHasSideLora, SIDE_BLOCK_K>(workspace); + + // MMA pipeline and TMA phases + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + // Intra-SM Barrier indices + constexpr uint32_t kDispatchBarrierIdx = 0; + constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; + constexpr uint32_t kEpilogueFullBarrierIdx = 2; + constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; + + // NVLink barrier tags + constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; + constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + + // Adjust registers + // NOTES: more experts per rank will cost more schedulers' registers + 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, + "Too many registers"); + + // Grid sync index assignments (dispatch and epilogue use separate counters to avoid conflicts) + constexpr uint32_t kDispatchGridSyncIndex = 0; + constexpr uint32_t kEpilogueGridSyncIndex = 1; + + // Different warp roles + if (warp_idx < kNumDispatchWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // Dispatch warps + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + // TODO: figure out better unrolling + // Now, `unroll` is better than `unroll 8` + #pragma unroll + for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; + i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) { + // Allocate slots for each token-topk + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) + process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + }; + + // Count experts' tokens + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + 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(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); + + // 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(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; + }); + + // Grid sync + comm::grid_sync( + workspace, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + + // Write expert count + if (sm_idx == 0) { + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Barrier before pulling + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* After the grid sync above, there is no more writes by other SMs (except 0) */ false, + /* After the NVLink barrier, there is a grid sync */ true + ); + + // Ensure the epilogue barrier cannot run with the pull barrier + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // 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 = &shared_storage.dispatch_barriers[warp_idx]; + + // Cache expert token counts in registers (same pattern as scheduler) + scheduler.fetch_expert_recv_count(); + + // Per-rank counts for current expert (re-loaded when expert changes) + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalWarps) { + // Advance expert until within the range + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++ current_expert_idx >= kNumExpertsPerRank) + break; + + // Update pool block offset for the new expert + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + + // Move start and end to the next expert + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + + // Finish all tokens + if (current_expert_idx >= kNumExpertsPerRank) + break; + + // Load per-rank counts when expert changes + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t j = i * 32 + lane_idx; + // TODO: this is not coalesced + stored_rank_count[i] = j < kNumRanks ? + static_cast(*workspace.get_expert_recv_count_ptr(j, current_expert_idx)) : 0; + } + } + + // Round-robin rank selection via iterative min-peeling + uint32_t current_rank_in_expert_idx; + uint32_t remaining[kNumRanksPerLane]; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] = stored_rank_count[i]; + uint32_t offset = 0; + uint32_t token_idx_in_expert = token_idx - expert_start_idx; + uint32_t slot_idx = token_idx_in_expert; + uint32_t token_idx_in_rank; + while (true) { + // Compute active count and min across all ranks + // NOTES: reduce within each lane first, then warp-reduce once + uint32_t num_actives_in_lane = 0; + uint32_t min_in_lane = 0xffffffff; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + num_actives_in_lane += remaining[i] > 0; + if (remaining[i] > 0) + min_in_lane = cute::min(min_in_lane, remaining[i]); + } + const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane); + const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane); + + // Hit in the current round + const uint32_t num_round_tokens = length * num_active_ranks; + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = slot_idx % num_active_ranks; + uint32_t num_seen_ranks = 0; + current_rank_in_expert_idx = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0); + const uint32_t num_active_lanes = __popc(mask); + if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes) + current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); + num_seen_ranks += num_active_lanes; + } + token_idx_in_rank = offset + (slot_idx / num_active_ranks); + break; + } + + // Move into the next round + slot_idx -= num_round_tokens; + offset += length; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] -= cute::min(remaining[i], length); + } + + // Read source token-topk index (written by remote dispatch via NVLink) + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + // Hidden bytes are divided into chunks + constexpr uint32_t kNumChunks = kHidden / kNumBytesPerPull; + DG_STATIC_ASSERT(kNumChunks * kNumBytesPerPull == kHidden, "kNumBytesPerPull must divide hidden"); + + // TMA load token from remote rank and store into local + 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; + + // Wait for ring buffer slot to be available (previous consumer must have finished all N blocks) + constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; + const auto l1_empty_count_target = (pool_block_idx / kNumRingBlocks) * kNumL1BlockNs; + if (l1_empty_count_target > 0) { + const auto empty_ptr = workspace.get_l1_empty_count_ptr(pool_block_idx % kNumRingBlocks); + while (ptx::ld_acq(empty_ptr) < l1_empty_count_target); + } + + const auto src_base_ptr = sym_buffer.map( + input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), current_rank_in_expert_idx); + const auto dst_base_ptr = l1_token_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).get_base_ptr(); + 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 * kNumBytesPerPull), + pull_buffer.get_base_ptr(), kNumBytesPerPull + ); + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + }; + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumChunks; ++ i) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + math::advance_ptr(src_base_ptr, i * kNumBytesPerPull), + pull_mbarrier, kNumBytesPerPull + ); + 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 last chunk's TMA load from remote) + constexpr uint32_t kNumSFUint32 = kHidden / 128; + DG_STATIC_ASSERT(kNumSFUint32 > 0 and kHidden % 128 == 0, "Invalid SF"); + const auto remote_sf_ptr = sym_buffer.map( + input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t token_idx_in_block = token_idx_in_expert % BLOCK_M; + const auto sf_ring_token_idx = ring_block_idx * SF_BLOCK_M + + transform_sf_token_idx(token_idx_in_block); + #pragma unroll + for (uint32_t i = 0; i < math::constexpr_ceil_div(kNumSFUint32, 32u); ++ i) { + const uint32_t j = i * 32 + lane_idx; + if (j < kNumSFUint32) + local_sf_ptr[j * kNumSFRingTokens + sf_ring_token_idx] = remote_sf_ptr[j]; + } + __syncwarp(); + + // Finish the overlapped final MXFP8 chunk, then reuse the same + // dispatch TMA pipe to pull the exact BF16 QLoRA operand. The old + // peer uint4 loop serialized thousands of NVLink loads per route + // and dominated EP8 despite the compact FP4 base. + if (cute::elect_one_sync()) + issue_and_wait_pull_store(kNumChunks - 1); + __syncwarp(); + + if constexpr (kHasSideLora) { + constexpr uint32_t kSideBytes = + kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumSideChunks = + kSideBytes / kNumBytesPerPull; + DG_STATIC_ASSERT( + kNumSideChunks * kNumBytesPerPull == kSideBytes, + "kNumBytesPerPull must divide BF16 side input"); + const auto remote_side_x = sym_buffer.map( + side_lora_source_buffer.get_data_buffer(src_token_idx) + .get_base_ptr(), + current_rank_in_expert_idx); + auto* local_saved_x = saved_x + + static_cast(pool_token_idx) * kHidden; + const auto issue_and_wait_side_store = + [&](const uint32_t& i) { + ptx::mbarrier_wait_and_flip_phase( + pull_mbarrier, pull_mbarrier_phase); + ptx::tma_store_1d( + math::advance_ptr( + local_saved_x, + i * kNumBytesPerPull), + pull_buffer.get_base_ptr(), + kNumBytesPerPull); + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + }; + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumSideChunks; ++i) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + math::advance_ptr( + remote_side_x, + i * kNumBytesPerPull), + pull_mbarrier, kNumBytesPerPull); + ptx::mbarrier_arrive_and_set_tx( + pull_mbarrier, kNumBytesPerPull); + if (i != kNumSideChunks - 1) + issue_and_wait_side_store(i); + } + issue_and_wait_side_store(kNumSideChunks - 1); + } + __syncwarp(); + } + + // Store weights and metadata + if (cute::elect_one_sync()) { + // Load weights + const auto weight = *sym_buffer.map( + input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).template get_base_ptr() = weight; + + // 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}; + + const bool is_last_token = (token_idx == expert_end_idx - 1); + ptx::red_add_rel( + workspace.get_l1_full_count_ptr(pool_block_idx % kNumRingBlocks), + is_last_token ? BLOCK_M - (token_idx_in_expert % BLOCK_M) : 1u + ); + } + __syncwarp(); + } + + // Clean workspace for the next usage, and also do cumulative stats + // NOTES: it is overlapped with combine reduction epilogue + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + DG_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { + // SM 0: clear expert send count + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + } else { + // Other SMs: clean blocks + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + // Read expert token count before clearing + const auto num_recv_tokens = static_cast( + *workspace.get_expert_recv_count_sum_ptr(i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + + // Compute expert pool block offset + expert_pool_block_offset = scheduler.get_pool_block_offset(i); + + // Wait read count ready + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Clean expert token count, and add cumulative results + DG_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); + if (warp_idx == 0) { + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + } else if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + + // Clean per-rank token count + for (uint32_t j = thread_idx; j < kNumRanks; j += kNumDispatchThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + __syncwarp(); + + // Clean L1 and L2 full stuffs and ring buffer counts + for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { + *workspace.get_l1_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l1_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + } + __syncwarp(); + } + } + + // Wait for all ranks to finish cleaning + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* Before the NVLink barrier, there is a grid sync */ true, + /* At the end of kernel does not need to sync */ false + ); + } else if (warp_idx == kNumDispatchWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM TMA load warp for tokens with SFA + 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 bool is_side = + block_phase != sched::BlockPhase::Linear1 && + block_phase != sched::BlockPhase::Linear2; + const cute::TmaDescriptor* tensor_map_a_ptr = &tensor_map_l1_acts; + const cute::TmaDescriptor* tensor_map_sfa_ptr = &tensor_map_l1_acts_sf; + if (block_phase == sched::BlockPhase::LoraL1Shrink) + tensor_map_a_ptr = &tensor_map_saved_x; + else if (block_phase == sched::BlockPhase::LoraL1Expand) + tensor_map_a_ptr = &tensor_map_lora_l1_scratch; + else if (block_phase == sched::BlockPhase::LoraL2Shrink) + tensor_map_a_ptr = &tensor_map_saved_h; + else if (block_phase == sched::BlockPhase::LoraL2Expand) + tensor_map_a_ptr = &tensor_map_lora_l2_scratch; + else if (block_phase == sched::BlockPhase::Linear2) { + tensor_map_a_ptr = &tensor_map_l2_acts; + tensor_map_sfa_ptr = &tensor_map_l2_acts_sf; + } + + const auto shape_k = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_K : L1_SHAPE_K; + const auto shape_sfa_k = math::ceil_div(shape_k, kGranK * 4u); + + // Compute pool block offset for this expert + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + + // Wait the entire token arrival for linear 1 + if (block_phase == sched::BlockPhase::LoraL1Shrink || + block_phase == sched::BlockPhase::Linear1) { + const auto ptr = workspace.get_l1_full_count_ptr(ring_block_idx); + const auto num_expected_tokens = BLOCK_M * (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != num_expected_tokens); + if constexpr (kHasSideLora) { + if (block_phase == sched::BlockPhase::Linear1) { + const auto generation = + pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq( + reinterpret_cast( + side_lora_ready) + ring_block_idx) < + 2 * generation); + } + } + } else if (block_phase == sched::BlockPhase::LoraL1Expand) { + const auto generation = pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq(reinterpret_cast( + side_lora_ready) + ring_block_idx) < + 2 * generation); + } else if (block_phase == sched::BlockPhase::LoraL2Shrink || + block_phase == sched::BlockPhase::Linear2) { + const auto ptr = workspace.get_l2_full_count_ptr(ring_block_idx); + const auto num_expected_blocks = (L2_SHAPE_K / BLOCK_N) * 2 * (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != num_expected_blocks); + if constexpr (kHasSideLora) { + if (block_phase == sched::BlockPhase::Linear2) { + const auto generation = + pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq( + reinterpret_cast( + side_lora_ready) + + 2 * kNumRingBlocks + ring_block_idx) < + generation); + } + } + } else { + const auto generation = pool_block_idx / kNumRingBlocks + 1; + while (ptx::ld_acq(reinterpret_cast( + side_lora_ready) + 2 * kNumRingBlocks + + ring_block_idx) < generation); + } + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); + + const uint32_t base_k_blocks = + block_phase == sched::BlockPhase::Linear1 ? + L1_SHAPE_K / BLOCK_K : + block_phase == sched::BlockPhase::Linear2 ? + L2_SHAPE_K / BLOCK_K : 0; + const bool is_linear_side_expand = kHasSideLora && + (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) && + k_block_idx >= base_k_blocks; + const bool is_side_iter = is_side || + is_linear_side_expand; + const uint32_t logical_k_block_idx = + is_linear_side_expand ? + k_block_idx - base_k_blocks : k_block_idx; + const cute::TmaDescriptor* current_tensor_map_a_ptr = + tensor_map_a_ptr; + if (is_linear_side_expand) + current_tensor_map_a_ptr = + block_phase == sched::BlockPhase::Linear1 ? + &tensor_map_lora_l1_scratch : + &tensor_map_lora_l2_scratch; + + // Compute token offsets from ring block index + uint32_t ring_m_idx = ring_block_idx * BLOCK_M; + uint32_t load_m_idx = is_side_iter + ? pool_block_idx * BLOCK_M : ring_m_idx; + uint32_t k_idx = logical_k_block_idx * + (is_side_iter ? SIDE_BLOCK_K : BLOCK_K); + if (is_linear_side_expand && + block_phase == sched::BlockPhase::Linear1) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / SIDE_BLOCK_K; + const uint32_t projection_idx = + logical_k_block_idx / kRankBlocks; + const uint32_t rank_block_idx = + logical_k_block_idx % kRankBlocks; + k_idx = projection_idx * kSideLoraRank + + rank_block_idx * SIDE_BLOCK_K; + } else if (block_phase == sched::BlockPhase::LoraL1Expand) { + k_idx += (n_block_idx & 1u) * kSideLoraRank; + } + uint32_t sfa_ring_m_idx = ring_block_idx * SF_BLOCK_M; + uint32_t sfa_k_idx = k_block_idx * (BLOCK_K / 128); + + // Add 2 CTA offsets for non-leader CTA + if (not is_leader_cta) + load_m_idx += scheduler.template get_valid_m() / 2; + + // TMA copy tokens and SFA, then arrive at full barrier + if (cute::elect_one_sync()) { + if (is_side_iter) { + tma::copy( + current_tensor_map_a_ptr, &shared_storage.full_barriers[stage_idx], + shared_storage.smem_a[stage_idx].side, k_idx, load_m_idx, 2); + if (is_leader_cta) + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx( + sizeof(shared_storage.smem_a[0].side) * 2); + else + shared_storage.full_barriers[stage_idx].arrive(0u); + } else { + tma::copy( + tensor_map_a_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_a[stage_idx].base, k_idx, load_m_idx, 2); + tma::copy( + tensor_map_sfa_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_sfa[stage_idx], sfa_ring_m_idx, sfa_k_idx, 2); + if (is_leader_cta) + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(sizeof(shared_storage.smem_a[0].base) * 2 + sizeof(SharedStorage::smem_sfa[0]) * 2); + else + shared_storage.full_barriers[stage_idx].arrive(0u); + } + } + __syncwarp(); + } + }); + } else if (warp_idx == kNumDispatchWarps + 1) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM TMA load warp for weights with SF + 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 bool is_side = + block_phase != sched::BlockPhase::Linear1 && + block_phase != sched::BlockPhase::Linear2; + const cute::TmaDescriptor* tensor_map_b_ptr = &tensor_map_l1_weights; + const cute::TmaDescriptor* tensor_map_sfb_ptr = &tensor_map_l1_weights_sf; + uint32_t shape_n = L1_SHAPE_N; + uint32_t expert_n_block_idx = n_block_idx; + if (block_phase == sched::BlockPhase::LoraL1Shrink) { + tensor_map_b_ptr = (n_block_idx & 1u) ? + &tensor_map_lora_a3 : &tensor_map_lora_a1; + shape_n = kSideLoraRank; + expert_n_block_idx = 0; + } else if (block_phase == sched::BlockPhase::LoraL1Expand) { + tensor_map_b_ptr = (n_block_idx & 1u) ? + &tensor_map_lora_b3 : &tensor_map_lora_b1; + shape_n = kIntermediateHidden; + expert_n_block_idx = n_block_idx / 2; + } else if (block_phase == sched::BlockPhase::LoraL2Shrink) { + tensor_map_b_ptr = &tensor_map_lora_a2; + shape_n = kSideLoraRank; + expert_n_block_idx = 0; + } else if (block_phase == sched::BlockPhase::LoraL2Expand) { + tensor_map_b_ptr = &tensor_map_lora_b2; + shape_n = kHidden; + } else if (block_phase == sched::BlockPhase::Linear2) { + tensor_map_b_ptr = &tensor_map_l2_weights; + tensor_map_sfb_ptr = &tensor_map_l2_weights_sf; + shape_n = L2_SHAPE_N; + } + + const auto shape_k = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_K : L1_SHAPE_K; + const auto shape_sfb_k = math::ceil_div(shape_k, kGranK * 4u); + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); + + const uint32_t base_k_blocks = + block_phase == sched::BlockPhase::Linear1 ? + L1_SHAPE_K / BLOCK_K : + block_phase == sched::BlockPhase::Linear2 ? + L2_SHAPE_K / BLOCK_K : 0; + const bool is_linear_side_expand = kHasSideLora && + (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) && + k_block_idx >= base_k_blocks; + const bool is_side_iter = is_side || + is_linear_side_expand; + const uint32_t logical_k_block_idx = + is_linear_side_expand ? + k_block_idx - base_k_blocks : k_block_idx; + const cute::TmaDescriptor* current_tensor_map_b_ptr = + tensor_map_b_ptr; + uint32_t current_shape_n = shape_n; + uint32_t current_expert_n_block_idx = expert_n_block_idx; + if (is_linear_side_expand) { + if (block_phase == sched::BlockPhase::Linear1) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / SIDE_BLOCK_K; + const uint32_t projection_idx = + logical_k_block_idx / kRankBlocks; + current_tensor_map_b_ptr = projection_idx ? + &tensor_map_lora_b3 : &tensor_map_lora_b1; + current_shape_n = kIntermediateHidden; + current_expert_n_block_idx = n_block_idx; + } else { + current_tensor_map_b_ptr = &tensor_map_lora_b2; + current_shape_n = kHidden; + } + } + + // Compute weight offset + const bool shared_side_weight = + block_phase == sched::BlockPhase::LoraL1Shrink || + (block_phase == sched::BlockPhase::Linear2 && + is_linear_side_expand); + uint32_t n_idx = + (shared_side_weight ? 0 : + local_expert_idx * current_shape_n) + + current_expert_n_block_idx * BLOCK_N; + uint32_t k_idx = logical_k_block_idx * + (is_side_iter ? SIDE_BLOCK_K : BLOCK_K); + if (is_linear_side_expand && + block_phase == sched::BlockPhase::Linear1) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / SIDE_BLOCK_K; + k_idx = (logical_k_block_idx % kRankBlocks) * + SIDE_BLOCK_K; + } + uint32_t sfb_n_idx = n_block_idx * BLOCK_N; + uint32_t sfb_k_idx = local_expert_idx * shape_sfb_k + k_block_idx * (BLOCK_K / 128); + + // TMA copy weights with SF. L1 side expansion loads compact + // B1/B3 strips into their gate/up slots and zeros the other + // half, so two MMAs can accumulate into the base-interleaved + // side tile without expanding the persistent adapter layout. + const bool is_l1_expand = is_linear_side_expand && + block_phase == sched::BlockPhase::Linear1; + if (is_l1_expand) { + constexpr uint32_t kRankBlocks = + kSideLoraRank / SIDE_BLOCK_K; + const uint32_t projection_idx = + logical_k_block_idx / kRankBlocks; + auto* side_tile = shared_storage.smem_b[stage_idx].side; + for (uint32_t idx = lane_idx; + idx < LOAD_BLOCK_N * SIDE_BLOCK_K; idx += 32) + side_tile[idx] = cutlass::bfloat16_t(0.0f); + __syncwarp(); + cutlass::arch::fence_view_async_shared(); + if (cute::elect_one_sync()) { + constexpr uint32_t kGran = 8; + #pragma unroll + for (uint32_t chunk = 0; + chunk < BLOCK_N / (2 * kGran); ++chunk) { + tma::copy( + current_tensor_map_b_ptr, + &shared_storage.full_barriers[stage_idx], + side_tile + + (chunk * 2 * kGran + + projection_idx * kGran) * SIDE_BLOCK_K, + k_idx, + local_expert_idx * kIntermediateHidden + + n_block_idx * (BLOCK_N / 2) + + chunk * kGran, + 2); + } + if (is_leader_cta) + shared_storage.full_barriers[stage_idx] + .arrive_and_expect_tx( + sizeof(shared_storage.smem_b[0].side)); + else + shared_storage.full_barriers[stage_idx].arrive(0u); + } + } else if (cute::elect_one_sync()) { + if (is_side_iter) { + tma::copy( + current_tensor_map_b_ptr, &shared_storage.full_barriers[stage_idx], + shared_storage.smem_b[stage_idx].side, k_idx, n_idx, 2); + if (is_leader_cta) + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx( + sizeof(shared_storage.smem_b[0].side) * 2); + else + shared_storage.full_barriers[stage_idx].arrive(0u); + } else { + tma::copy( + tensor_map_b_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_b[stage_idx].base, 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(shared_storage.smem_b[0].base) + sizeof(SharedStorage::smem_sfb[0]) * 2); + else + shared_storage.full_barriers[stage_idx].arrive(0u); + } + } + __syncwarp(); + } + }); + } else if (warp_idx == kNumDispatchWarps + 2) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM MMA issue warp (only the leader CTA will run) + if (is_leader_cta) { + // Make instruction descriptor with block scaling + // NOTES: always swap A/B + auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< + b_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, + UMMA_M, UMMA_N, + cute::UMMA::Major::K, cute::UMMA::Major::K + >(); + auto side_instr_desc = cute::UMMA::make_instr_desc< + cutlass::bfloat16_t, cutlass::bfloat16_t, float, 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"); + auto a_desc = mma::sm100::make_umma_desc(shared_storage.smem_a[0].base, 0, 0); + auto b_desc = mma::sm100::make_umma_desc(shared_storage.smem_b[0].base, 0, 0); + auto side_a_desc = mma::sm100::make_umma_desc(shared_storage.smem_a[0].side, 0, 0); + auto side_b_desc = mma::sm100::make_umma_desc(shared_storage.smem_b[0].side, 0, 0); + uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * sizeof(SharedStorage::AStage) / 16 : 0u; + uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * sizeof(SharedStorage::BStage) / 16 : 0u; + uint32_t side_a_desc_lo = lane_idx < kNumStages ? side_a_desc.lo + lane_idx * sizeof(SharedStorage::AStage) / 16 : 0u; + uint32_t side_b_desc_lo = lane_idx < kNumStages ? side_b_desc.lo + lane_idx * sizeof(SharedStorage::BStage) / 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 + (UMMA_M == 128 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256) or + (UMMA_M == 256 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256), + "Invalid MMA instruction shape"); + + // Persistently schedule over blocks + uint32_t current_iter_idx = 0; + uint32_t side_empty_wait_phase[2] = {1, 1}; + 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) { + // Dynamic update of UMMA N based on effective M + mma::sm100::update_instr_desc_with_umma_n(instr_desc, scheduler.template get_valid_m()); + mma::sm100::update_instr_desc_with_umma_n(side_instr_desc, scheduler.template get_valid_m()); + const bool is_side = + block_phase != sched::BlockPhase::Linear1 && + block_phase != sched::BlockPhase::Linear2; + + // Wait tensor memory empty barrier arrival + const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; + const auto accum_phase = (current_iter_idx ++ / kNumEpilogueStages) & 1; + if constexpr (kHasSideLora) { + shared_storage.tmem_empty_barriers[accum_stage_idx] + .wait(side_empty_wait_phase[accum_stage_idx]); + side_empty_wait_phase[accum_stage_idx] ^= 1; + } else { + shared_storage.tmem_empty_barriers[accum_stage_idx] + .wait(accum_phase ^ 1); + } + ptx::tcgen05_after_thread_sync(); + + // Empty barrier arrival + auto empty_barrier_arrive = [&](const bool& do_tmem_full_arrive) { + auto umma_arrive = [](const uint64_t* barrier) { + constexpr uint16_t kCTAMask = (1 << 2) - 1; + cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); + }; + 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(&shared_storage.tmem_full_barriers[accum_stage_idx])); + __syncwarp(); + }; + + // Launch MMAs + #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 + shared_storage.full_barriers[stage_idx].wait(phase); + ptx::tcgen05_after_thread_sync(); + + const uint32_t base_k_blocks = + block_phase == sched::BlockPhase::Linear1 ? + L1_SHAPE_K / BLOCK_K : + block_phase == sched::BlockPhase::Linear2 ? + L2_SHAPE_K / BLOCK_K : 0; + const bool is_linear_side_expand = kHasSideLora && + (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) && + k_block_idx >= base_k_blocks; + const bool is_side_iter = is_side || + is_linear_side_expand; + const uint32_t logical_k_block_idx = + is_linear_side_expand ? + k_block_idx - base_k_blocks : k_block_idx; + if (is_linear_side_expand && + logical_k_block_idx == 0) { + const uint32_t side_stage_idx = + 1 - accum_stage_idx; + shared_storage.tmem_empty_barriers[side_stage_idx] + .wait(side_empty_wait_phase[side_stage_idx]); + side_empty_wait_phase[side_stage_idx] ^= 1; + ptx::tcgen05_after_thread_sync(); + } + const auto a_desc_base_lo = ptx::exchange( + is_side_iter ? side_a_desc_lo : a_desc_lo, + stage_idx); + const auto b_desc_base_lo = ptx::exchange( + is_side_iter ? side_b_desc_lo : b_desc_lo, + stage_idx); + if (cute::elect_one_sync()) { + if (is_side_iter) { + #pragma unroll + for (uint32_t k = 0; k < SIDE_BLOCK_K / SIDE_UMMA_K; ++k) { + side_a_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_M, 128, + cutlass::bfloat16_t>(a_desc_base_lo, 0, + k * SIDE_UMMA_K); + side_b_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_N, 128, + cutlass::bfloat16_t>(b_desc_base_lo, 0, + k * SIDE_UMMA_K); + ptx::SM100_MMA_F16BF16_2x1SM_SS::fma( + side_b_desc, side_a_desc, + is_linear_side_expand ? + (1 - accum_stage_idx) * UMMA_N : + accum_stage_idx * UMMA_N, + logical_k_block_idx > 0 || k > 0, + static_cast( + static_cast(side_instr_desc)) << 32); + } + } else { + #pragma unroll + for (uint32_t umma_k_block_idx = 0; umma_k_block_idx < BLOCK_K / UMMA_BLOCK_K; ++ umma_k_block_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] + umma_k_block_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] + umma_k_block_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 + #pragma unroll + for (uint32_t k = 0; k < UMMA_BLOCK_K / UMMA_K; ++ k) { + const auto runtime_instr_desc = + 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, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_M * sizeof(a_dtype_t), k * UMMA_K); + b_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, b_dtype_t>(b_desc_base_lo, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_N * sizeof(b_dtype_t), k * UMMA_K); + ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or umma_k_block_idx > 0 or k > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); + } + } + } + } + __syncwarp(); + + // Commit to the mbarrier object + // No explicit `tcgen05.fence::before_thread_sync` is needed, as this is implicitly performed by `tcgen05.commit` + empty_barrier_arrive(k_block_idx == num_k_blocks - 1); + } + }); + + // To safely deconstruct barriers, we need another round of waits + if (current_iter_idx > 0) { + if constexpr (kHasSideLora) { + #pragma unroll + for (uint32_t i = 0; i < 2; ++i) + shared_storage.tmem_empty_barriers[i] + .wait(side_empty_wait_phase[i]); + } else { + const auto accum_phase_idx = + ((current_iter_idx - 1) / + kNumEpilogueStages) & 1; + shared_storage.tmem_empty_barriers[ + (current_iter_idx - 1) % + kNumEpilogueStages].wait(accum_phase_idx); + } + } + } + } else if (warp_idx == kNumDispatchWarps + 3) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + } else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_alloc(); + + // 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(&shared_storage.tmem_ptr_in_smem) == 0); + + // GEMM epilogue warps + const auto epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const auto epilogue_wg_idx = epilogue_warp_idx / 4; + const auto epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const auto warp_idx_in_wg = epilogue_warp_idx % 4; + DG_STATIC_ASSERT((kNumDispatchWarps + kNumMMANonEpilogueWarps) % 4 == 0 and + kNumEpilogueWarps % 4 == 0, "Invalid epilogue warps"); + + // TODO: support effective block M + // NOTES: + // - 2 warpgroups divide the whole BM into BM / 2 + // - 4 warps divide the whole BN into BN / 4 + // - BM / 2 is further divided into stored blocks, i.e. with `STORE_BLOCK_M` size + // - `STORE_BLOCK_M` in further divided into `ATOM_M` + constexpr uint32_t WG_BLOCK_M = BLOCK_M / kNumEpilogueWarpgroups; + constexpr uint32_t ATOM_M = 8; + constexpr uint32_t kNumBankGroupBytes = 16u; + constexpr uint32_t kNumAtomsPerStore = STORE_BLOCK_M / ATOM_M; + DG_STATIC_ASSERT(BLOCK_M % kNumEpilogueWarpgroups == 0, "Invalid block M"); + DG_STATIC_ASSERT(WG_BLOCK_M % STORE_BLOCK_M == 0, "Invalid warpgroup block M"); + DG_STATIC_ASSERT(STORE_BLOCK_M % ATOM_M == 0, "Invalid store block M"); + DG_STATIC_ASSERT(BLOCK_N == 128, "Invalid block N"); + + // Ensure the epilogue barrier cannot run with the pull barrier + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Persistently schedule over blocks + uint32_t current_iter_idx = 0; + 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) { + // Wait UMMA arrival + const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; + const auto accum_phase = (current_iter_idx ++ / kNumEpilogueStages) & 1; + shared_storage.tmem_full_barriers[accum_stage_idx].wait(accum_phase); + ptx::tcgen05_after_thread_sync(); + + // Compute offsets + // NOTES: use shuffle here to let NVCC know warp divergence won't happen + const uint32_t valid_m = ptx::exchange(scheduler.template get_valid_m(), 0); + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t ring_m_idx = ring_block_idx * BLOCK_M; // Ring-buffer offset for reusable data buffers + const uint32_t pool_m_idx = pool_block_idx * BLOCK_M; // Full-pool offset for non-ring metadata + uint32_t n_idx = n_block_idx * BLOCK_N; + auto release_tmem = [&]() { + ptx::tcgen05_before_thread_sync(); + shared_storage.tmem_empty_barriers[accum_stage_idx] + .arrive(0u); + if constexpr (kHasSideLora) { + if (block_phase == sched::BlockPhase::Linear1 || + block_phase == sched::BlockPhase::Linear2) + shared_storage.tmem_empty_barriers[ + 1 - accum_stage_idx].arrive(0u); + } + }; + + const bool is_lora_phase = + block_phase == sched::BlockPhase::LoraL1Shrink || + block_phase == sched::BlockPhase::LoraL2Shrink; + + if (is_lora_phase) { + const bool do_store = + block_phase != sched::BlockPhase::LoraL2Shrink || + n_block_idx == 0; + const cute::TmaDescriptor* tensor_map_d_ptr = + &tensor_map_lora_l1_scratch_store; + uint32_t out_n_idx = n_block_idx * BLOCK_N; + if (block_phase == sched::BlockPhase::LoraL2Shrink) { + tensor_map_d_ptr = + &tensor_map_lora_l2_scratch_store; + out_n_idx = 0; + } + + if (!do_store) { + release_tmem(); + } else { + #pragma unroll + for (uint32_t s = 0; + s < WG_BLOCK_M / STORE_BLOCK_M; ++s) { + if (epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M >= valid_m) { + release_tmem(); + break; + } + #pragma unroll + for (uint32_t i = 0; + i < STORE_BLOCK_M / ATOM_M; ++i) { + const uint32_t tmem_addr = + accum_stage_idx * UMMA_N + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + i * ATOM_M; + uint32_t values[ATOM_M]; + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + tmem_addr, values[0], values[1], + values[2], values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + tmem_addr | 0x00100000, values[4], + values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + if (s == WG_BLOCK_M / STORE_BLOCK_M - 1 && + i == STORE_BLOCK_M / ATOM_M - 1) { + release_tmem(); + } + const uint32_t row = lane_idx % 8; + const uint32_t col = + (warp_idx_in_wg % 2) * 4 + lane_idx / 8; + const auto smem_ptr = + shared_storage.smem_d.l2[epilogue_wg_idx] + + (warp_idx_in_wg / 2) * STORE_BLOCK_M * + (kSwizzleCDMode / sizeof(nv_bfloat16)) + + i * ATOM_M * + (kSwizzleCDMode / sizeof(nv_bfloat16)) + + row * ((kNumBankGroupBytes * 8) / + sizeof(nv_bfloat16)) + + (col ^ row) * + (kNumBankGroupBytes / sizeof(nv_bfloat16)); + ptx::SM90_U32x4_STSM_T::copy( + math::cast_into_bf16_and_pack(values[0], values[1]), + math::cast_into_bf16_and_pack(values[2], values[3]), + math::cast_into_bf16_and_pack(values[4], values[5]), + math::cast_into_bf16_and_pack(values[6], values[7]), + smem_ptr); + } + ptx::sync_aligned( + 128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + if (warp_idx_in_wg == 0 && cute::elect_one_sync()) { + cute::tma_store_fence(); + #pragma unroll + for (uint32_t atom = 0; + atom < BLOCK_N * sizeof(nv_bfloat16) / + kSwizzleCDMode; ++atom) { + cute::SM90_TMA_STORE_2D::copy( + tensor_map_d_ptr, + shared_storage.smem_d.l2[epilogue_wg_idx] + + atom * STORE_BLOCK_M * + (kSwizzleCDMode / + sizeof(nv_bfloat16)), + out_n_idx + atom * + (kSwizzleCDMode / + sizeof(nv_bfloat16)), + pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M); + cute::tma_store_arrive(); + } + } + __syncwarp(); + // Scratch slices share one staging tile. Do not let + // the next 32-row slice overwrite it while the prior + // asynchronous TMA store is still reading it. + if (warp_idx_in_wg == 0) + cute::tma_store_wait<0>(); + ptx::sync_aligned( + 128, kEpilogueWGBarrierStartIdx + + epilogue_wg_idx); + } + } + ptx::tma_store_wait<0>(); + ptx::sync_aligned( + kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (do_store && epilogue_warp_idx == 0 && + cute::elect_one_sync()) { + uint32_t ready_plane = 0; + if (block_phase == sched::BlockPhase::LoraL1Expand) + ready_plane = 1; + else if (block_phase == sched::BlockPhase::LoraL2Shrink) + ready_plane = 2; + else if (block_phase == sched::BlockPhase::LoraL2Expand) + ready_plane = 3; + ptx::red_add_rel( + reinterpret_cast(side_lora_ready) + + ready_plane * kNumRingBlocks + ring_block_idx, + 1u); + } + __syncwarp(); + } else if (block_phase == sched::BlockPhase::Linear1) { + // Wait L2 block empty + const auto l2_empty_ptr = workspace.get_l2_empty_count_ptr(ring_block_idx); + const auto num_expected_blocks = (L2_SHAPE_N / BLOCK_N) * (pool_block_idx / kNumRingBlocks); + while (ptx::ld_acq(l2_empty_ptr) != num_expected_blocks); + + // Unified L1 epilogue: gated activation (SwiGLU/GeGLU) in-place using + // granularity 8 interleaved weights. + // With `SM100_TMEM_LOAD_16dp256b1x`, gate/up pairs are: + float stored_cached_weight = 0; + + #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) { + release_tmem(); + break; + } + + // Iterate all atoms in the store block + float2 activation_values[kNumAtomsPerStore][2]; + float2 amax_values[kNumAtomsPerStore]; + #pragma unroll + for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { + const uint32_t j = s * kNumAtomsPerStore + i; + + // Load weights from global into register cache per 32 tokens + DG_STATIC_ASSERT(32 % ATOM_M == 0, "Invalid block size"); + if ((j * ATOM_M) % 32 == 0 and (WG_BLOCK_M % 32 == 0 or j * ATOM_M + lane_idx < WG_BLOCK_M)) { + stored_cached_weight = *l1_topk_weights_buffer + .get_data_buffer(ring_m_idx + epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M + lane_idx) + .template get_base_ptr(); + } + + // Load weights from register cache + const 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) + }; + + // Load from TMEM + uint2 raw_values[4]; + uint2 side_raw_values[4] = {}; + uint32_t tmem_addr = accum_stage_idx * UMMA_N + epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M; + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, + 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, + raw_values[2].x, raw_values[2].y, raw_values[3].x, raw_values[3].y); + if constexpr (kHasSideLora) { + const uint32_t side_tmem_addr = + (1 - accum_stage_idx) * UMMA_N + + epilogue_wg_idx * WG_BLOCK_M + + j * ATOM_M; + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr, side_raw_values[0].x, + side_raw_values[0].y, + side_raw_values[1].x, + side_raw_values[1].y); + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr | 0x00100000, + side_raw_values[2].x, + side_raw_values[2].y, + side_raw_values[3].x, + side_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) { + release_tmem(); + } + + // Apply gated activation: act(gate) * up (SwiGLU or GeGLU) + auto fp32_values = reinterpret_cast(raw_values); + #pragma unroll + for (uint32_t k = 0; k < 2; ++ k) { + auto bf16_gate = __float22bfloat162_rn(fp32_values[k * 2 + 0]); + auto bf16_up = __float22bfloat162_rn(fp32_values[k * 2 + 1]); + + if constexpr (kHasSideLora) { + const auto side_fp32_values = + reinterpret_cast( + side_raw_values); + const auto gate_side = + __bfloat1622float2( + __float22bfloat162_rn( + side_fp32_values[k * 2 + 0])); + const auto up_side = + __bfloat1622float2( + __float22bfloat162_rn( + side_fp32_values[k * 2 + 1])); + bf16_gate = __float22bfloat162_rn(__fadd2_rn( + __bfloat1622float2(bf16_gate), + __fmul2_rn({side_lora_scale, + side_lora_scale}, gate_side))); + bf16_up = __float22bfloat162_rn(__fadd2_rn( + __bfloat1622float2(bf16_up), + __fmul2_rn({side_lora_scale, + side_lora_scale}, up_side))); + } + + if constexpr (kSaveL1Preact) { + // Each lane owns two rows and two adjacent + // hidden columns. Persist the exact BF16 bits + // before clamp without adding a staged path. + const uint32_t hidden_col = + warp_idx_in_wg * 16 + + (lane_idx / 4) * 2 + k; + const uint32_t chunk = hidden_col / 8; + const uint32_t in_chunk = + hidden_col & 7; + const uint32_t gate_col = + chunk * 16 + in_chunk; + const uint32_t up_col = + gate_col + 8; + const auto output_col = + [=](uint32_t col) { + const uint32_t low = + col & 31; + return n_idx + + (col & ~31u) + + ((low & 1) << 4) + + ((low >> 1) & 3) + + (low & 8) + + ((low & 16) >> 2); + }; + const uint32_t row_base = + pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + + i * ATOM_M + + (lane_idx % 4) * 2; + const uint32_t gate_bits = + *reinterpret_cast< + const uint32_t*>( + &bf16_gate); + const uint32_t up_bits = + *reinterpret_cast< + const uint32_t*>( + &bf16_up); + #pragma unroll + for (uint32_t r = 0; r < 2; + ++r) { + const uint32_t m_out = + row_base + r; + if (m_out < + pool_m_idx + valid_m) { + auto* dst = + reinterpret_cast< + uint16_t*>( + saved_l1_preact + + static_cast< + uint64_t>( + m_out) * + (2 * + kIntermediateHidden)); + dst[output_col( + gate_col)] = + static_cast< + uint16_t>( + gate_bits >> + (r * 16)); + dst[output_col(up_col)] = + static_cast< + uint16_t>( + up_bits >> + (r * 16)); + } + } + } + + // Clamp + if constexpr (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}); + } + + const auto gate = __bfloat1622float2(bf16_gate); + const auto up = __bfloat1622float2(bf16_up); + + // Gated activation, applied to the gate projection. Both variants + // reduce to `gate * sigmoid(z) * up` for an activation-specific `z`: + // - SwiGLU: SiLU(gate) => z = gate + // - GeGLU: GELU(gate) (tanh) => z = alpha * (gate + beta * gate^3) + // since 0.5 * (1 + tanh(t)) == sigmoid(2t) + float2 z; + if constexpr (kActivationType == ActivationType::GeGLU) { + constexpr float kAlpha = 1.5957691216057308f; // 2 * sqrt(2 / pi) + constexpr float kBeta = 0.044715f; + const auto gate_sq = __fmul2_rn(gate, gate); + z = __fmul2_rn( + {kAlpha, kAlpha}, + __fmul2_rn(gate, __fadd2_rn({1.0f, 1.0f}, __fmul2_rn({kBeta, kBeta}, gate_sq)))); + } else { + z = gate; + } + + const auto neg_exp = make_float2( + kFastMath ? __expf(-z.x) : expf(-z.x), + kFastMath ? __expf(-z.y) : expf(-z.y)); + const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_exp); + float2 activated; + if constexpr (kFastMath) { + activated = __fmul2_rn(gate, {math::fast_rcp(denom.x), math::fast_rcp(denom.y)}); + } else { + activated = {gate.x / denom.x, gate.y / denom.y}; + } + const float2 h_unweighted = + __fmul2_rn(activated, up); + const float2 h_for_w2 = + kRouteWeightMode == RouteWeightMode::PreDown + ? __fmul2_rn(h_unweighted, weights) + : h_unweighted; + if constexpr (kHasSideLora) { + const uint32_t hidden_col = + warp_idx_in_wg * 16 + + (lane_idx / 4) * 2 + k; + // `hidden_col` is the raw TMEM lane order. The + // normal FP8 L2-activation store canonicalizes + // it through STSM before TMA, but this BF16 tap + // writes global memory directly. Apply the same + // permutation explicitly so the saved tensor is + // [pool, intermediate] and can be consumed by + // the side shrink and its backward wgrad. + const uint32_t interleaved_gate_col = + (hidden_col / 8) * 16 + + (hidden_col & 7); + const uint32_t low = + interleaved_gate_col & 31; + const uint32_t canonical_interleaved_col = + (interleaved_gate_col & ~31u) + + ((low & 1) << 4) + + ((low >> 1) & 3) + + (low & 8) + + ((low & 16) >> 2); + const uint32_t canonical_hidden_col = + (canonical_interleaved_col / 16) * 8 + + (canonical_interleaved_col & 7); + const uint32_t output_col = + n_block_idx * (BLOCK_N / 2) + + canonical_hidden_col; + const uint32_t row0 = pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + i * ATOM_M + + (lane_idx % 4) * 2; + saved_h[static_cast(row0) * + kIntermediateHidden + output_col] = + __float2bfloat16_rn(h_for_w2.x); + if (row0 + 1 < pool_m_idx + valid_m) + saved_h[static_cast(row0 + 1) * + kIntermediateHidden + output_col] = + __float2bfloat16_rn(h_for_w2.y); + } + if constexpr ( + kRouteWeightMode == + RouteWeightMode::PreDown) { + // Keep the legacy expression intact: PRE_DOWN + // must remain bitwise identical. + activation_values[i][k] = h_for_w2; + } else { + activation_values[i][k] = h_unweighted; + } + } + + // Amax reduction (thread-level) + float2 thread_local_amax = {0.f, 0.f}; + #pragma unroll + for (uint32_t k = 0; k < 2; ++ k) { + thread_local_amax.x = cute::max(thread_local_amax.x, cute::abs(activation_values[i][k].x)); + thread_local_amax.y = cute::max(thread_local_amax.y, cute::abs(activation_values[i][k].y)); + } + + // Amax reduction (warp-level) + amax_values[i].x = math::warp_reduce<4, true>( + thread_local_amax.x, math::ReduceMax()); + amax_values[i].y = math::warp_reduce<4, true>( + thread_local_amax.y, math::ReduceMax()); + + // Reduce amax (warp-pair-level) + if (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 `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 and store into shared memory + #pragma unroll + for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { + // 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); + + // Calculate SF + float2 sf, sf_inv; + math::get_e4m3_sf_and_sf_inv(amax_values[i], sf, sf_inv); + + // Cast + const float2 upper = __fmul2_rn(activation_values[i][0], sf_inv); + const float2 lower = __fmul2_rn(activation_values[i][1], sf_inv); + const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(upper.x, upper.y, lower.x, lower.y)); + + // STSM + uint32_t row = lane_idx; + uint32_t col = warp_idx_in_wg; + const auto smem_ptr = reinterpret_cast(shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx]) + + i * ATOM_M * L1_OUT_BLOCK_N + + row * L1_OUT_BLOCK_N + // 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 `l2_sf_buffer` as UE8M0 (MN-major layout) + // Only one warp per pair writes (both hold the same SF after cross-warp reduce) + // 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; + const uint32_t k_uint_idx = k_idx / 4, byte_idx = k_idx % 4; + const uint32_t mn_stride = kNumSFRingTokens * sizeof(uint32_t); + const auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + // NOTES: consecutive tokens (t, t + 1) are in the same 32-group, so `sf_idx` differs by 4 + // NOTES: originally there was: + // - `const uint32_t token_idx_in_expert = m_block_idx * BLOCK_M + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M + lane_idx * 2 + // - `scheduler.get_current_pool_block_offset() * SF_BLOCK_M + transform_sf_token_idx(token_idx_in_expert)` + // We find out that + // 1. `m_block_idx * BLOCK_M` mod `BLOCK_M` is 0, and `epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M + lane_idx * 2` is always < `BLOCK_M`, so we can put `m_block_idx * BLOCK_M` outside + // 2. `lane_idx * 2` controls the lowest 3 bit of `token_idx_in_expert`, and `transform_sf_token_idx` is a bitwise-independent transformation if the input is less than `BLOCK_M`, so we can put `lane_idx * 2` outside + // This reduce the number of computation instructions. + const uint32_t token_base_idx = epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M; + __builtin_assume(token_base_idx < BLOCK_M); + const auto sf_ring_token_idx = ring_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_ring_token_idx * static_cast(sizeof(uint32_t)) + byte_idx; + sf_base_ptr[sf_addr] = + (*reinterpret_cast(&sf.x) >> 23); + sf_base_ptr[sf_addr + 4 * static_cast(sizeof(uint32_t))] = + (*reinterpret_cast(&sf.y) >> 23); + } + __syncwarp(); + } + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Issue TMA store after all atoms in this store block + if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + 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, + shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx], + out_n_idx, + ring_m_idx + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M); + cute::tma_store_arrive(); + } + __syncwarp(); + } + + // Notify L2 and increment L1 empty count + // TODO: less epilogue sync scope + ptx::tma_store_wait<0>(); + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + workspace.get_l2_full_count_ptr(ring_block_idx), 1u); + + // Increment L1 empty count for this physical slot (one per N block) + ptx::red_add( + workspace.get_l1_empty_count_ptr(ring_block_idx), 1u); + } + __syncwarp(); + } else { + // Increment L2 empty count for this physical slot (one per N block) + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_add( + workspace.get_l2_empty_count_ptr(ring_block_idx), 1u); + } + __syncwarp(); + + DG_STATIC_ASSERT(STORE_BLOCK_M % 8 == 0, "Invalid store M"); + constexpr uint32_t kNumRowsPerWarp = STORE_BLOCK_M / 8; + + // 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) { + // Early break if the entire store block is beyond the valid token range + // TODO: check performance + if (epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M >= valid_m) { + release_tmem(); + break; + } + + #pragma unroll + for (uint32_t i = 0; i < STORE_BLOCK_M / ATOM_M; ++ i) { + // Load from TMEM using .16x256b shape to satisfy STSM layout requirements + // Start from lane index 0 and 16 + uint32_t tmem_addr = accum_stage_idx * UMMA_N + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M; + uint32_t values[ATOM_M]; + uint32_t side_values[ATOM_M] = {}; + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, + values[0], values[1], values[2], values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, + values[4], values[5], values[6], values[7]); + if constexpr (kHasSideLora) { + const uint32_t side_tmem_addr = + (1 - accum_stage_idx) * UMMA_N + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M + i * ATOM_M; + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr, side_values[0], + side_values[1], side_values[2], + side_values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy( + side_tmem_addr | 0x00100000, + side_values[4], side_values[5], + side_values[6], side_values[7]); + } + cutlass::arch::fence_view_async_tmem_load(); + + if constexpr (kHasSideLora) { + auto* base_fp32 = + reinterpret_cast(values); + const auto* side_fp32 = + reinterpret_cast(side_values); + #pragma unroll + for (uint32_t value_idx = 0; + value_idx < ATOM_M; ++value_idx) { + base_fp32[value_idx] = + __bfloat162float(__float2bfloat16_rn( + base_fp32[value_idx])) + + side_lora_scale * + __bfloat162float(__float2bfloat16_rn( + side_fp32[value_idx])); + } + } + + // 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) + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Signal tensor memory consumed + if (s == WG_BLOCK_M / STORE_BLOCK_M - 1 and i == STORE_BLOCK_M / ATOM_M - 1) { + release_tmem(); + } + + // Store into shared memory + // 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 = 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) + + (col ^ row) * kNumBankGroupBytes; + ptx::SM90_U32x4_STSM_T::copy( + math::cast_into_bf16_and_pack(values[0], values[1]), + math::cast_into_bf16_and_pack(values[2], values[3]), + math::cast_into_bf16_and_pack(values[4], values[5]), + math::cast_into_bf16_and_pack(values[6], values[7]), + smem_ptr + ); + } + + // Wait shared memory ready + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + if constexpr (kSaveDownUnweighted) { + const uint32_t saved_store_row = + pool_m_idx + + epilogue_wg_idx * WG_BLOCK_M + + s * STORE_BLOCK_M; + DG_DEVICE_ASSERT( + saved_store_row + STORE_BLOCK_M <= + num_saved_pool_tokens); + if (warp_idx_in_wg == 0 && + cute::elect_one_sync()) { + cute::tma_store_fence(); + #pragma unroll + for (uint32_t atom = 0; + atom < + BLOCK_N * + sizeof(nv_bfloat16) / + kSwizzleCDMode; + ++atom) { + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_down_unweighted, + shared_storage.smem_d + .l2[epilogue_wg_idx] + + atom * STORE_BLOCK_M * + (kSwizzleCDMode / + sizeof( + nv_bfloat16)), + n_idx + + atom * + (kSwizzleCDMode / + sizeof( + nv_bfloat16)), + saved_store_row); + cute::tma_store_arrive(); + } + } + if (warp_idx_in_wg == 0) + cute::tma_store_wait<0>(); + __syncwarp(); + } + + // Write into remote buffers + // 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; + + #pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++ j) { + const uint32_t row_in_store = j * 8 + warp_idx_in_wg * 2 + lane_idx / 16; + const uint32_t m_idx_in_block = epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + row_in_store; + + // Skip padding rows beyond the actual token count for this expert + if (m_idx_in_block >= valid_m) + break; + + const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_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; + + // Read from shared memory + 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; + auto packed = ptx::ld_shared(reinterpret_cast(smem_ptr)); + if constexpr (kHasSideLora && + kSaveDownUnweighted) { + const uint32_t output_col_base = + n_idx + (lane_idx % 16) * 8; + auto* saved_ptr = saved_down_unweighted + + static_cast(pool_m_idx + + m_idx_in_block) * kHidden + + output_col_base; + *reinterpret_cast(saved_ptr) = packed; + } + if constexpr ( + kRouteWeightMode == + RouteWeightMode::PostDown) { + // The route score remains in its immutable source + // token/slot plane. Loading it from metadata avoids + // adding a full-pool field to the forward layout. + const float route_weight = + *sym_buffer.map( + input_topk_weights_buffer + .get_base_ptr() + + static_cast( + src_metadata.token_idx) * + kNumTopk + + src_metadata.topk_idx, + src_metadata.rank_idx); + auto* values = + reinterpret_cast< + nv_bfloat16*>(&packed); + #pragma unroll + for (uint32_t value_idx = 0; + value_idx < 8; + ++value_idx) { + values[value_idx] = + __float2bfloat16_rn( + __bfloat162float( + values[value_idx]) * + route_weight); + } + } + + // Write into remote + const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_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))); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } + } + + // Ensure the next epilogue safe to use shared memory + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + }); + + // Deallocate tensor memory + // NOTES: must be called by the same logical warp ID on both CTAs + if (epilogue_warp_idx == 0) + Allocator().free(0, kNumTmemCols); + + // NVLink barrier (grid sync + cross-rank signal + grid sync): ~4 us + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + + // Barrier with dispatch warps, so that they can do clean workspace + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Combine: reduce top-k results and write back + // NOTES: reuse shared memory from start up to the barriers + // 1 token, 1 topk latency: ~3 us + constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); + + // 3 slots of chunk is needed: 2 load stages and 1 store + constexpr uint32_t kNumChunkSlots = 3; + constexpr uint32_t kNumMaxRegistersForBuffer = 128; + + // NOTES: either 1 or 2 chunks for simplicity + // NOTES: Restrict on both smem and register + constexpr uint32_t kNumChunks = + 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 <= 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 <= 32, "Top-k must fit in a single warp"); + + // Verify combined shared memory budget at runtime + 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) { + return math::advance_ptr(smem_buffer, (epilogue_warp_idx + i * kNumEpilogueWarps) * kNumChunkBytes); + }); + const auto combine_store_buffer = math::advance_ptr(smem_buffer, (epilogue_warp_idx + kNumEpilogueWarps * 2) * kNumChunkBytes); + + // Per-warp barriers + auto combine_load_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return &shared_storage.combine_barriers[i + epilogue_warp_idx * 2]; + }); + + // Iterate over all tokens + uint32_t combine_phase = 0; + uint32_t load_stage_idx = 0; + for (uint32_t token_idx = sm_idx * kNumEpilogueWarps + epilogue_warp_idx; + token_idx < num_tokens; + token_idx += kNumSMs * kNumEpilogueWarps) { + // Read top-k slot indices: each lane reads one slot, then broadcast via exchange + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + const int stored_topk_slot_idx = lane_idx < kNumTopk ? + static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; + const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); + + // Iterate all chunks + for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { + const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; + + // Move mask and load + uint32_t mask = total_mask; + const auto move_mask_and_load = [&](const uint32_t& i) { + if (mask) { + // Move + const uint32_t slot_idx = __ffs(mask) - 1; + mask ^= 1 << slot_idx; + + // Load + if (cute::elect_one_sync()) { + const auto src_ptr = math::advance_ptr( + combine_token_buffer.get_rank_buffer(slot_idx) + .get_data_buffer(token_idx).get_base_ptr(), + chunk_byte_offset); + ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], kNumChunkBytes); + ptx::mbarrier_arrive_and_set_tx(combine_load_barriers[i], kNumChunkBytes); + } + __syncwarp(); + return true; + } + return false; + }; + + // Load the first selection + bool do_reduce = move_mask_and_load(load_stage_idx); + + // Accumulate all top-k contributions for this chunk in float registers + float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {}; + while (do_reduce) { + // Prefetch next top-k into the buffer while current is being accumulated + do_reduce = move_mask_and_load(load_stage_idx ^ 1); + + // Accumulate + combine_load_barriers[load_stage_idx]->wait(combine_phase); + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; + const auto bf16_values = reinterpret_cast(&uint4_values); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); + } + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } + + // Cast + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); + + // Wait share memory release and write + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); + } + __syncwarp(); + + // TMA store the token chunk + if (cute::elect_one_sync()) { + cute::tma_store_fence(); + ptx::tma_store_1d( + math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + chunk_byte_offset), + combine_store_buffer, kNumChunkBytes); + cute::tma_store_arrive(); + } + __syncwarp(); + } + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_100f"); +#endif +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_moe_side_lora.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_moe_side_lora.cuh new file mode 100644 index 0000000000..bab7d7e3da --- /dev/null +++ b/deep_gemm/include/deep_gemm/scheduler/mega_moe_side_lora.cuh @@ -0,0 +1,262 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace deep_gemm::sched { + +// Computation phase for the current block +enum class BlockPhase { + None = 0, + LoraL1Shrink = 1, + LoraL1Expand = 2, + Linear1 = 3, + LoraL2Shrink = 4, + LoraL2Expand = 5, + Linear2 = 6 +}; + +template +struct SideLoraMegaMoEScheduler { + 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(L1_SHAPE_K % SIDE_BLOCK_K == 0, "Invalid side shape"); + DG_STATIC_ASSERT(L2_SHAPE_K % SIDE_BLOCK_K == 0, "Invalid side shape"); + DG_STATIC_ASSERT(128 % SIDE_BLOCK_K == 0, "Invalid side rank tile"); + DG_STATIC_ASSERT(kNumExpertsPerWave > 0 and kNumExpertsPerWave <= kNumExpertsPerRank, "Invalid wave config"); + + // NOTES: N block counts must be even so that 2 adjacent CTAs in a cluster + // always land on the same m_block_idx with n_block_idx differing by 1 + DG_STATIC_ASSERT(kNumSMs % 2 == 0, "Number of SMs must be even for 2-CTA cluster"); + DG_STATIC_ASSERT(kNumL1BlockNs % 2 == 0, "L1 N block count must be even for 2-CTA cluster"); + DG_STATIC_ASSERT(kNumL2BlockNs % 2 == 0, "L2 N block count must be even for 2-CTA cluster"); + + // Rank-128 side LoRA maps naturally to the 2-CTA, 128-column Mega tile: + // A1/A3 occupy the two shrink CTAs, B1/B3 alternate within each expand + // pair, and A2 is duplicated across the pair (only CTA 0 is stored). + static constexpr uint32_t kNumLoraL1ShrinkBlockNs = 2; + static constexpr uint32_t kNumLoraL1ExpandBlockNs = kNumL1BlockNs; + static constexpr uint32_t kNumLoraL2ShrinkBlockNs = 2; + static constexpr uint32_t kNumLoraL2ExpandBlockNs = kNumL2BlockNs; + + // Arrival counts + const layout::Workspace& workspace; + + // Scheduler state + BlockPhase next_phase = kHasSideLora ? BlockPhase::LoraL1Shrink : BlockPhase::Linear1; + + // Current expert and block indices + uint32_t current_local_expert_idx = 0; + uint32_t current_num_tokens = 0; + uint32_t current_pool_block_offset = 0; + uint32_t block_idx = 0; + uint32_t m_block_idx = 0; + uint32_t n_block_idx = 0; + + // Pre-cached per-expert token counts (filled during `for_each_block` init) + // Layout: `stored_num_tokens_per_expert[i]` holds expert (i * 32 + lane_idx)'s count + uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {}; + + CUTLASS_DEVICE explicit SideLoraMegaMoEScheduler(const layout::Workspace& workspace): workspace(workspace) { + block_idx = blockIdx.x; + } + + CUTLASS_DEVICE uint32_t get_wave_expert_end_idx() const { + // Align up to wave boundary, clamped for the last partial wave + 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; + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ? + stored_num_tokens_per_expert[i] : valid_value; + } + return ptx::exchange(valid_value, expert_idx % 32); + } + + // Get pool block offset for a given expert index from a per-lane token count array + CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) { + uint32_t num_blocks = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + if (i * 32 + ptx::get_lane_idx() < expert_idx) + num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M); + } + return __reduce_add_sync(0xffffffff, num_blocks); + } + + CUTLASS_DEVICE void advance_expert_idx() { + current_pool_block_offset += get_current_num_m_blocks(); + current_local_expert_idx += 1; + current_num_tokens = get_num_tokens(current_local_expert_idx); + } + + CUTLASS_DEVICE void set_expert_idx(const uint32_t& expert_idx) { + current_local_expert_idx = expert_idx; + current_num_tokens = get_num_tokens(expert_idx); + current_pool_block_offset = get_pool_block_offset(expert_idx); + } + + CUTLASS_DEVICE uint32_t get_current_pool_block_offset() const { + return current_pool_block_offset; + } + + CUTLASS_DEVICE uint32_t get_current_num_m_blocks() const { + return math::ceil_div(current_num_tokens, BLOCK_M); + } + + template + CUTLASS_DEVICE uint32_t get_valid_m() const { + const auto m = cute::min(current_num_tokens - m_block_idx * BLOCK_M, BLOCK_M); + return kDoUMMAAligned ? math::align(m, 16u) : m; + } + + CUTLASS_DEVICE bool fetch_next_block_with_n_count(const uint32_t num_n_blocks) { + 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 / num_n_blocks; + if (m_block_idx < num_m_blocks) + return true; + + // Current expert is fully assigned, move to the next + block_idx -= num_m_blocks * num_n_blocks; + advance_expert_idx(); + } + return false; + } + + CUTLASS_DEVICE uint32_t get_phase_num_n_blocks(const BlockPhase phase) const { + if (phase == BlockPhase::LoraL1Shrink) + return kNumLoraL1ShrinkBlockNs; + if (phase == BlockPhase::LoraL1Expand) + return kNumLoraL1ExpandBlockNs; + if (phase == BlockPhase::Linear1) + return kNumL1BlockNs; + if (phase == BlockPhase::LoraL2Shrink) + return kNumLoraL2ShrinkBlockNs; + return kNumL2BlockNs; + } + + CUTLASS_DEVICE uint32_t get_phase_num_k_blocks(const BlockPhase phase) const { + if (phase == BlockPhase::LoraL1Shrink) + return kNumLoraL1ShrinkBlockKs; + if (phase == BlockPhase::LoraL1Expand || + phase == BlockPhase::LoraL2Expand) + return kNumLoraExpandBlockKs; + if (phase == BlockPhase::LoraL2Shrink) + return kNumLoraL2ShrinkBlockKs; + if (phase == BlockPhase::Linear2) + return kNumL2BlockKs + + (kHasSideLora ? kNumLoraExpandBlockKs : 0); + return kNumL1BlockKs + + // Gate and up have different shared-A projections. A 2-CTA UMMA + // splits M across the CTAs, so they must be issued as two K passes + // rather than assigning q1/B1 to CTA 0 and q3/B3 to CTA 1. + (kHasSideLora ? 2 * kNumLoraExpandBlockKs : 0); + } + + CUTLASS_DEVICE BlockPhase get_phase_after(const BlockPhase phase) const { + if constexpr (kHasSideLora) { + if (phase == BlockPhase::LoraL1Shrink) + return BlockPhase::Linear1; + if (phase == BlockPhase::Linear1) + return BlockPhase::LoraL2Shrink; + if (phase == BlockPhase::LoraL2Shrink) + return BlockPhase::Linear2; + return BlockPhase::LoraL1Shrink; + } else { + return phase == BlockPhase::Linear1 ? + BlockPhase::Linear2 : BlockPhase::Linear1; + } + } + + // Core state machine: assigns the next block + CUTLASS_DEVICE cute::tuple get_next_block() { + while (true) { + if (current_local_expert_idx >= kNumExpertsPerRank) + break; + + const auto phase = next_phase; + const auto num_n_blocks = get_phase_num_n_blocks(phase); + if (fetch_next_block_with_n_count(num_n_blocks)) { + n_block_idx = block_idx - m_block_idx * num_n_blocks; + block_idx += kNumSMs; + return {phase, current_local_expert_idx, m_block_idx, n_block_idx}; + } else { + next_phase = get_phase_after(phase); + // Every side/base phase in a wave revisits the same experts. + // Linear2 is the last phase and advances to the next wave. + if (phase != BlockPhase::Linear2) + set_expert_idx(math::align( + current_local_expert_idx - 1, kNumExpertsPerWave)); + } + } + + // All waves and experts are fully processed + return {BlockPhase::None, 0, 0, 0}; + } + + CUTLASS_DEVICE void fetch_expert_recv_count() { + // NOTES: each lane caches experts at indices (i * 32 + lane_idx) + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + const auto expert_idx = i * 32 + ptx::get_lane_idx(); + uint64_t value = 0; + if (expert_idx < kNumExpertsPerRank) { + do { + value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx)); + } while (static_cast(value >> 32) != kNumSMs * kNumRanks); + } + stored_num_tokens_per_expert[i] = static_cast(value); + } + __syncwarp(); + } + + template + CUTLASS_DEVICE void for_each_block(Func&& func) { + // Wait for all expert counters to be finalized + fetch_expert_recv_count(); + + // Initialize current expert with 0 + set_expert_idx(0); + + // Iterate over all blocks + // TODO: add swizzle within expert waves for better L2 cache utilization + while (true) { + CUTE_TIE_DECL(get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx); + if (block_phase == BlockPhase::None) + break; + + func(block_phase, current_local_expert_idx, + get_phase_num_k_blocks(block_phase), + m_block_idx, n_block_idx); + } + } +}; + +} // namespace deep_gemm::sched diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 0cf7847101..97ddda2599 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -82,6 +82,7 @@ def __init__(self, group: dist.ProcessGroup, self.token_src_metadata, self.backward_grad_y, self.backward_grad_route, + self.side_lora_source, ) = slice_input_buffers(self.buffer) def destroy(self): @@ -93,6 +94,7 @@ def destroy(self): self.token_src_metadata = None self.backward_grad_y = None self.backward_grad_route = None + self.side_lora_source = None def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, @@ -196,6 +198,56 @@ def transform_weights_for_mega_moe( return l1_transformed, l2_transformed +def transform_side_lora_for_mega_moe( + side_lora: Tuple[torch.Tensor, torch.Tensor, torch.Tensor, + torch.Tensor, torch.Tensor, torch.Tensor] +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, + torch.Tensor, torch.Tensor, torch.Tensor]: + """Create the K-major rank-128 views consumed by the fused TC path. + + Input tensors use the conventional ``(A1, B1, A3, B3, A2, B2)`` + layouts from the shared side-LoRA contract: A1/A3 are independent + ``[H, R]`` matrices shared across experts, B1/B3 and A2 are expert-local, + and B2 is one shared ``[R, H]`` matrix. Keep and refresh this transformed + tuple after optimizer updates; it must not be rebuilt for every forward. + """ + if len(side_lora) != 6: + raise ValueError( + "side_lora must contain (A1, B1, A3, B3, A2, B2)") + rank = side_lora[0].size(-1) + if rank != 128: + raise ValueError( + "the tensor-core MegaMoE side-LoRA path currently requires " + "rank 128") + a1, b1, a3, b3, a2, b2 = side_lora + if a1.dim() != 2 or a3.dim() != 2 or b2.dim() != 2: + raise ValueError( + "side-LoRA requires shared A1/A3 and shared B2") + if b1.dim() != 3 or b3.dim() != 3 or a2.dim() != 3: + raise ValueError( + "side-LoRA requires expert-local B1/B3/A2") + hidden, rank = a1.shape + experts, b1_rank, intermediate = b1.shape + expected_shapes = ( + ("A3", a3.shape, (hidden, rank)), + ("B1", b1.shape, (experts, rank, intermediate)), + ("B3", b3.shape, (experts, rank, intermediate)), + ("A2", a2.shape, (experts, intermediate, rank)), + ("B2", b2.shape, (rank, hidden)), + ) + for name, actual, expected in expected_shapes: + if tuple(actual) != expected: + raise ValueError( + f"invalid {name} shape {tuple(actual)}; expected {expected}") + if b1_rank != rank: + raise ValueError("B1 rank must match A1 rank") + if any(tensor.dtype != torch.bfloat16 for tensor in side_lora): + raise TypeError("side-LoRA tensors must all be BF16") + return tuple(tensor.transpose(-1, -2).contiguous() + for tensor in side_lora) + + + def fp8_fp4_mega_moe( y: torch.Tensor, l1_weights: Tuple[torch.Tensor, torch.Tensor], @@ -257,6 +309,90 @@ def fp8_fp4_mega_moe( num_config_tokens, ) + +def fp8_fp4_mega_moe_side_lora( + y: torch.Tensor, + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], + sym_buffer: SymmBuffer, + side_lora_input: torch.Tensor, + side_lora: Tuple[torch.Tensor, torch.Tensor, torch.Tensor, + torch.Tensor, torch.Tensor, torch.Tensor], + saved_x: torch.Tensor, + saved_h_unweighted: torch.Tensor, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + recipe: Tuple[int, int, int] = (1, 1, 32), + activation: str = "swiglu", + activation_clamp: Optional[float] = None, + fast_math: bool = True, + saved_l1_preact: Optional[torch.Tensor] = None, + route_weight_mode: RouteWeightMode = RouteWeightMode.PRE_DOWN, + saved_down_unweighted: Optional[torch.Tensor] = None, + num_config_tokens: Optional[int] = None, + side_lora_scale: float = 1.0, + side_lora_scratch: Optional[ + Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the dedicated MXFP4 base + BF16 rank-128 side-LoRA kernel. + + Only q1/q3/q2 are persisted. Rank expansion accumulators remain in TMEM + beside the MXFP4 base accumulators and are BF16-rounded before addition; + no hidden- or intermediate-width side delta is allocated. + """ + route_weight_mode = RouteWeightMode(route_weight_mode) + if len(side_lora) != 6: + raise ValueError( + "side_lora must contain (A1, B1, A3, B3, A2, B2)") + if side_lora[0].size(0) != 128: + raise ValueError("MXFP4 side-LoRA requires rank 128") + if side_lora_input.dtype != torch.bfloat16: + raise TypeError("side_lora_input must be BF16") + if side_lora_input.shape != (y.size(0), sym_buffer.hidden): + raise ValueError("side_lora_input must have shape [tokens, hidden]") + if sym_buffer.side_lora_source.numel() == 0: + raise ValueError("symmetric buffer has no MXFP4 side-LoRA source plane") + sym_buffer.side_lora_source[:y.size(0)].copy_( + side_lora_input.contiguous()) + pool_rows = saved_x.size(0) + if side_lora_scratch is None: + side_lora_scratch = ( + torch.empty( + (pool_rows, 2, 128), dtype=torch.bfloat16, + device=y.device), + torch.empty( + (pool_rows, 128), dtype=torch.bfloat16, + device=y.device), + torch.zeros( + 4 * sym_buffer.num_ring_tokens // + _C.get_token_alignment_for_mega_moe(), + dtype=torch.int32, device=y.device), + ) + if len(side_lora_scratch) != 3: + raise ValueError("side_lora_scratch must contain (q13, q2, ready)") + q13, q2, ready = side_lora_scratch + has_explicit_config_tokens = num_config_tokens is not None + if num_config_tokens is None: + num_config_tokens = y.size(0) + if not has_explicit_config_tokens and sym_buffer.group.size() > 1: + rank_uniform_num_tokens = torch.tensor( + num_config_tokens, dtype=torch.int32, device=y.device) + dist.all_reduce( + rank_uniform_num_tokens, op=dist.ReduceOp.MAX, + group=sym_buffer.group) + num_config_tokens = int(rank_uniform_num_tokens.item()) + _C.fp8_fp4_mega_moe_side_lora( + y, l1_weights, l2_weights, + cumulative_local_expert_recv_stats, + sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, + sym_buffer.group.rank(), sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_experts, sym_buffer.num_topk, recipe, + activation, activation_clamp, fast_math, + sym_buffer.num_ring_tokens, saved_l1_preact, + route_weight_mode.value, saved_down_unweighted, + num_config_tokens, saved_x, saved_h_unweighted, + *side_lora, q13, q2, ready, float(side_lora_scale)) + return q13, q2, ready + def bf16_mega_moe(y: torch.Tensor, l1_weights: torch.Tensor, l2_weights: torch.Tensor, @@ -279,6 +415,106 @@ def bf16_mega_moe(y: torch.Tensor, saved_x: Optional[torch.Tensor] = None): """Run BF16 MegaMoE with an explicit route-weight boundary. + The optional stage saves expose unweighted/weighted activation and W2 + output boundaries for strict parity checks. The fused backward currently + supports the production ``pre_down`` route-weight boundary only. + + Training callers may provide an exact local source-route histogram, + rank-uniform ``active_pool_rows``, and a scalar mismatch flag to size saved + pools from actual receive counts. The kernel publishes the precomputed + counts, verifies them against its internal dispatch count, and sets the + flag before any caller can accept a truncated result. + """ + route_weight_mode = RouteWeightMode(route_weight_mode) + combine_order_mode = CombineOrderMode(combine_order_mode) + if ( + (saved_h_unweighted is None) != + (saved_h_weighted is None) + ): + raise ValueError( + "both activation stage outputs must be provided together") + active_plan = ( + precomputed_route_counts is not None, + active_pool_rows is not None, + route_count_mismatch is not None, + ) + if any(active_plan) and not all(active_plan): + raise ValueError( + "precomputed_route_counts, active_pool_rows, and " + "route_count_mismatch must be provided together") + has_precomputed_config_tokens = num_config_tokens is not None + if num_config_tokens is None: + num_config_tokens = y.size(0) + if ( + not has_precomputed_config_tokens + and sym_buffer.group.size() > 1 + ): + # The config selects BLOCK_M, which defines the persistent launch and + # pool packing. Empty source ranks can still receive expert rows, so + # every rank must select the config from the same source-token extent. + rank_uniform_num_tokens = torch.tensor( + num_config_tokens, dtype=torch.int32, device=y.device) + dist.all_reduce( + rank_uniform_num_tokens, + op=dist.ReduceOp.MAX, + group=sym_buffer.group) + num_config_tokens = int(rank_uniform_num_tokens.item()) + _C.bf16_mega_moe( + y, + l1_weights, + l2_weights, + cumulative_local_expert_recv_stats, + sym_buffer.buffer, + sym_buffer.handle.buffer_ptrs, + sym_buffer.group.rank(), + sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_experts, + sym_buffer.num_topk, + activation, activation_clamp, + fast_math, + sym_buffer.num_ring_tokens, + saved_l1_preact, + route_weight_mode.value, + saved_h_unweighted, + saved_h_weighted, + saved_down_unweighted, + num_config_tokens, + combine_order_mode.value, + precomputed_route_counts, + active_pool_rows, + route_count_mismatch, + saved_x, + ) + +def bf16_mega_moe_side_lora(y: torch.Tensor, + l1_weights: torch.Tensor, + l2_weights: torch.Tensor, + sym_buffer: SymmBuffer, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True, + saved_l1_preact: Optional[torch.Tensor] = None, + route_weight_mode: RouteWeightMode = RouteWeightMode.PRE_DOWN, + saved_h_unweighted: Optional[torch.Tensor] = None, + saved_h_weighted: Optional[torch.Tensor] = None, + saved_down_unweighted: Optional[torch.Tensor] = None, + combine_order_mode: CombineOrderMode = + CombineOrderMode.FIXED_TOPK, + precomputed_route_counts: Optional[torch.Tensor] = None, + active_pool_rows: Optional[int] = None, + route_count_mismatch: Optional[torch.Tensor] = None, + num_config_tokens: Optional[int] = None, + saved_x: Optional[torch.Tensor] = None, + side_lora: Optional[Tuple[torch.Tensor, torch.Tensor, + torch.Tensor, torch.Tensor, + torch.Tensor, torch.Tensor]] = None, + side_lora_scale: float = 1.0, + side_lora_scratch: Optional[ + Tuple[torch.Tensor, torch.Tensor, + torch.Tensor]] = None): + """Run the dedicated native rank-128 BF16 side-LoRA MegaMoE kernel. + The optional stage saves expose unweighted/weighted activation and W2 output boundaries for strict parity checks. ``saved_down_unweighted`` is also used by post-down backward for the exact router gradient. @@ -288,6 +524,12 @@ def bf16_mega_moe(y: torch.Tensor, pools from actual receive counts. The kernel publishes the precomputed counts, verifies them against its internal dispatch count, and sets the flag before any caller can accept a truncated result. + + ``side_lora`` is the K-major tuple returned by + :func:`transform_side_lora_for_mega_moe`. The optional scratch tuple is + ``(l1_shrink, l2_shrink, ready)``. Callers on a hot + path should retain and reuse it; otherwise correctly-sized buffers are + allocated. """ route_weight_mode = RouteWeightMode(route_weight_mode) combine_order_mode = CombineOrderMode(combine_order_mode) @@ -323,7 +565,39 @@ def bf16_mega_moe(y: torch.Tensor, op=dist.ReduceOp.MAX, group=sym_buffer.group) num_config_tokens = int(rank_uniform_num_tokens.item()) - _C.bf16_mega_moe( + if side_lora is None: + raise ValueError("the dedicated side-LoRA kernel requires adapters") + else: + if len(side_lora) != 6: + raise ValueError( + "side_lora must contain (A1, B1, A3, B3, A2, B2)") + side_lora_rank = side_lora[0].size(0) + if side_lora_rank != 128: + raise ValueError( + "side_lora must be the rank-128 K-major tuple returned by " + "transform_side_lora_for_mega_moe") + if side_lora_scratch is None: + side_lora_scratch = ( + torch.empty( + (active_pool_rows or + sym_buffer.token_src_metadata.size(0), + 2, side_lora_rank), + dtype=torch.bfloat16, device=y.device), + torch.empty( + (active_pool_rows or + sym_buffer.token_src_metadata.size(0), + side_lora_rank), + dtype=torch.bfloat16, device=y.device), + torch.empty( + (4, sym_buffer.num_ring_tokens // 8), + dtype=torch.int32, device=y.device), + ) + if len(side_lora_scratch) != 3: + raise ValueError( + "side_lora_scratch must contain " + "(l1_shrink, l2_shrink, ready)") + side_lora_args = (*side_lora, *side_lora_scratch) + _C.bf16_mega_moe_side_lora( y, l1_weights, l2_weights, @@ -348,4 +622,6 @@ def bf16_mega_moe(y: torch.Tensor, active_pool_rows, route_count_mismatch, saved_x, + *side_lora_args, + float(side_lora_scale), ) diff --git a/deep_gemm/mega/backward.py b/deep_gemm/mega/backward.py index 915c57d1c6..86bbb7abb4 100644 --- a/deep_gemm/mega/backward.py +++ b/deep_gemm/mega/backward.py @@ -1,6 +1,6 @@ import os import time -from typing import Any, Callable, Optional, Tuple +from typing import Any, Callable, NamedTuple, Optional, Tuple import torch import torch.distributed as dist @@ -939,6 +939,308 @@ def fp8_fp4_mega_moe_backward_dgrad_swiglu( ) +class MegaMoESideLoraBackwardResult(NamedTuple): + """Outputs of the native side-LoRA training backward. + + ``grad_side_lora`` uses the conventional trainable layouts supplied to + :func:`transform_side_lora_for_mega_moe`, not the transposed inference + layouts. Shared A1/A3/B2 gradients are EP-local partials, matching the + caller contract that reduces replicated 2-D factors across EP ranks. + The frozen base W1/W2/W3 gradients are deliberately absent. + """ + + grad_x_pool: torch.Tensor + grad_x: Optional[torch.Tensor] + grad_route: torch.Tensor + grad_ye: torch.Tensor + grad_h: torch.Tensor + grad_gate_up: torch.Tensor + h_act: torch.Tensor + h_weighted: torch.Tensor + x_pool: torch.Tensor + route_weights: torch.Tensor + t13: torch.Tensor + t2: torch.Tensor + grad_side_lora: Tuple[torch.Tensor, ...] + + +def _allocate_side_lora_backward_outputs( + gate_up: torch.Tensor, + side_lora: Tuple[torch.Tensor, ...], + hidden: int, + intermediate_hidden: int, + write_grad_x_pool: bool = True, +) -> tuple[torch.Tensor, ...]: + pool_rows = gate_up.size(0) + options = dict(dtype=torch.bfloat16, device=gate_up.device) + grad_ye = torch.empty((pool_rows, hidden), **options) + grad_h = torch.empty((pool_rows, intermediate_hidden), **options) + grad_gate_up = torch.empty_like(gate_up) + h_act = torch.empty_like(grad_h) + h_weighted = torch.empty_like(grad_h) + x_pool = torch.empty_like(grad_ye) + grad_x_pool = ( + torch.empty_like(grad_ye) + if write_grad_x_pool + else torch.empty((0, hidden), **options) + ) + route_weights = torch.empty( + pool_rows, dtype=torch.float32, device=gate_up.device) + grad_route = torch.empty_like(route_weights) + t13 = torch.empty((pool_rows, 2, 128), **options) + t2 = torch.empty((pool_rows, 128), **options) + # Every inference tensor is K-major. Transposing only its shape yields the + # conventional optimizer layout without performing a runtime data copy. + # Some expert-local output tiles can be structurally empty for a routing + # realization. The native wgrad kernels overwrite every visited tile but + # intentionally skip those empty tiles, so these buffers must start at + # zero rather than exposing allocator contents to the optimizer. + grad_side_lora = tuple( + torch.zeros( + (*tensor.shape[:-2], tensor.size(-1), tensor.size(-2)), + dtype=torch.bfloat16, + device=gate_up.device, + ) + for tensor in side_lora + ) + return ( + grad_ye, grad_h, grad_gate_up, h_act, h_weighted, x_pool, + grad_x_pool, route_weights, grad_route, t13, t2, + grad_side_lora, + ) + + +def bf16_mega_moe_side_lora_backward( + gate_up_output: torch.Tensor, + saved_h: torch.Tensor, + saved_down_unweighted: torch.Tensor, + q13: torch.Tensor, + q2: torch.Tensor, + side_lora: Tuple[torch.Tensor, ...], + w2_weights: torch.Tensor, + w13_weights: torch.Tensor, + expert_counts: torch.Tensor, + padded_expert_counts: torch.Tensor, + grad_y: torch.Tensor, + sym_buffer: Any, + block_m: int, + activation_limit: float = float("inf"), + activation: str = "swiglu", + fast_math: bool = False, + route_weight_mode: RouteWeightMode = RouteWeightMode.PRE_DOWN, + combine_order_mode: CombineOrderMode = CombineOrderMode.FIXED_TOPK, + side_lora_scale: float = 1.0, + direct_remote_grad_x: Optional[bool] = None, + combine_grad_x: bool = True, + write_grad_x_pool: Optional[bool] = None, + out: Optional[MegaMoESideLoraBackwardResult] = None, + grid_sync_counter: Optional[torch.Tensor] = None, + expert_psum_rows: Optional[torch.Tensor] = None, +) -> MegaMoESideLoraBackwardResult: + """Run the dedicated BF16 base-dgrad + rank-128 LoRA backward. + + The native path produces six adapter gradients and intentionally performs + no frozen base-weight gradients. It also avoids hidden-width LoRA delta or + gradient scratch buffers. ``w13_weights`` uses the same 8-row gate/up + interleave returned by :func:`transform_weights_for_mega_moe`; the + internal ``grad_gate_up`` result is emitted in that matching interleave. + """ + route_weight_mode = RouteWeightMode(route_weight_mode) + combine_order_mode = CombineOrderMode(combine_order_mode) + if route_weight_mode is RouteWeightMode.POST_DOWN: + raise NotImplementedError( + "BF16 side-LoRA backward currently supports pre_down routing only") + if len(side_lora) != 6 or side_lora[0].size(0) != 128: + raise ValueError("side_lora must be a rank-128 transformed tuple") + if direct_remote_grad_x is None: + direct_remote_grad_x = sym_buffer.group.size() > 1 + if write_grad_x_pool is None: + write_grad_x_pool = not direct_remote_grad_x + if not write_grad_x_pool and not direct_remote_grad_x: + raise ValueError("grad-x requires a local or direct remote output") + # The native specialization materializes the base dgrad pool, folds both + # side branches into it, then publishes the final value once. This is the + # ordinary base-dgrad output, not an additional side-LoRA wide scratch. + write_grad_x_pool = True + outputs = ( + _allocate_side_lora_backward_outputs( + gate_up_output, side_lora, w13_weights.size(2), + w2_weights.size(2), write_grad_x_pool) + if out is None else ( + out.grad_ye, out.grad_h, out.grad_gate_up, out.h_act, + out.h_weighted, out.x_pool, out.grad_x_pool, + out.route_weights, out.grad_route, out.t13, out.t2, + out.grad_side_lora) + ) + (grad_ye, grad_h, grad_gate_up, h_act, h_weighted, x_pool, + grad_x_pool, route_weights, grad_route, t13, t2, + grad_side_lora) = outputs + _direct_grad_x_planes(sym_buffer).zero_() + sym_buffer.backward_grad_y[:grad_y.size(0)].copy_( + grad_y.to(torch.bfloat16).contiguous()) + sym_buffer.backward_grad_route.zero_() + num_grid_states = ( + expert_counts.numel() * + ((w13_weights.size(2) // 64) * (w2_weights.size(2) // 128) + + (w13_weights.size(1) // 64) * (w13_weights.size(2) // 128)) + 2) + if grid_sync_counter is None: + grid_sync_counter = torch.zeros( + num_grid_states, dtype=torch.int32, device=gate_up_output.device) + if expert_psum_rows is None: + expert_psum_rows = padded_expert_counts.cumsum(0).to(torch.int32) + # Publish and pull grad-y/x/route metadata before the rank shrink. The + # following persistent specialization is told that dispatch is complete, + # so it does not repeat communication. This is a native CUDA prelude; no + # framework GEMM participates in the production backward. + _C.bf16_mega_moe_backward_post_down_prelude_v2( + grad_ye, grad_ye, x_pool, route_weights, grad_route, + saved_down_unweighted, expert_counts, + sym_buffer.backward_grad_y, sym_buffer.x, + sym_buffer.topk_weights, sym_buffer.backward_grad_route, + sym_buffer.token_src_metadata, sym_buffer.handle.buffer_ptrs, + sym_buffer.group.rank(), sym_buffer.num_topk, block_m, + combine_order_mode.value, True, False, False, True, True, + False, False, 256) + _C.bf16_mega_moe_side_lora_backward( + gate_up_output, grad_h, grad_gate_up, h_act, h_weighted, + x_pool, grad_x_pool, grad_route, grad_ye, grad_ye, + route_weights, w2_weights, w13_weights, expert_counts, + grid_sync_counter, float(activation_limit), activation, + bool(fast_math), route_weight_mode.value, + combine_order_mode.value, saved_down_unweighted, block_m, + bool(direct_remote_grad_x), + bool(write_grad_x_pool), True, + sym_buffer.backward_grad_y, sym_buffer.x, + sym_buffer.topk_weights, sym_buffer.backward_grad_route, + sym_buffer.token_src_metadata, sym_buffer.handle.buffer_ptrs, + sym_buffer.group.rank(), sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_topk, "dispatch_prepared", *side_lora, q13, q2, saved_h, + t13, t2, *grad_side_lora, expert_psum_rows, + padded_expert_counts, float(side_lora_scale), None) + grad_x = None + if direct_remote_grad_x and combine_grad_x: + grad_x = ( + torch.empty_like(grad_y, dtype=torch.bfloat16) + if out is None or out.grad_x is None else out.grad_x) + mega_moe_backward_combine_grad_x( + grad_x, sym_buffer, combine_order_mode) + return MegaMoESideLoraBackwardResult( + grad_x_pool, grad_x, grad_route, grad_ye, grad_h, + grad_gate_up, h_act, h_weighted, x_pool, route_weights, + t13, t2, grad_side_lora) + + +def fp8_fp4_mega_moe_side_lora_backward( + gate_up_output: torch.Tensor, + saved_h: torch.Tensor, + saved_down_unweighted: torch.Tensor, + q13: torch.Tensor, + q2: torch.Tensor, + side_lora: Tuple[torch.Tensor, ...], + l1_acts: torch.Tensor, + l1_acts_sf: torch.Tensor, + l1_weights: Tuple[torch.Tensor, torch.Tensor], + w13_weights: Tuple[torch.Tensor, torch.Tensor], + w2_weights: Tuple[torch.Tensor, torch.Tensor], + w13_dequant_scratch: torch.Tensor, + w2_dequant_scratch: torch.Tensor, + expert_counts: torch.Tensor, + padded_expert_counts: torch.Tensor, + grad_y: torch.Tensor, + sym_buffer: Any, + block_m: int, + activation_limit: float = float("inf"), + route_weight_mode: RouteWeightMode = RouteWeightMode.PRE_DOWN, + side_lora_scale: float = 1.0, + direct_remote_grad_x: Optional[bool] = None, + combine_grad_x: bool = True, + write_grad_x_pool: Optional[bool] = None, + out: Optional[MegaMoESideLoraBackwardResult] = None, + grid_sync_counter: Optional[torch.Tensor] = None, + expert_psum_rows: Optional[torch.Tensor] = None, +) -> MegaMoESideLoraBackwardResult: + """Run the dedicated MXFP4 base-dgrad + BF16 side-LoRA backward.""" + route_weight_mode = RouteWeightMode(route_weight_mode) + if route_weight_mode is RouteWeightMode.POST_DOWN: + raise NotImplementedError( + "MXFP4 side-LoRA backward currently supports pre_down routing only") + if len(side_lora) != 6 or side_lora[0].size(0) != 128: + raise ValueError("side_lora must be a rank-128 transformed tuple") + if direct_remote_grad_x is None: + direct_remote_grad_x = sym_buffer.group.size() > 1 + if write_grad_x_pool is None: + write_grad_x_pool = not direct_remote_grad_x + if not write_grad_x_pool and not direct_remote_grad_x: + raise ValueError("grad-x requires a local or direct remote output") + write_grad_x_pool = True + hidden = w2_weights[0].size(1) + intermediate_hidden = w2_dequant_scratch.size(2) + outputs = ( + _allocate_side_lora_backward_outputs( + gate_up_output, side_lora, hidden, intermediate_hidden, + write_grad_x_pool) + if out is None else ( + out.grad_ye, out.grad_h, out.grad_gate_up, out.h_act, + out.h_weighted, out.x_pool, out.grad_x_pool, + out.route_weights, out.grad_route, out.t13, out.t2, + out.grad_side_lora) + ) + (grad_ye, grad_h, grad_gate_up, h_act, h_weighted, x_pool, + grad_x_pool, route_weights, grad_route, t13, t2, + grad_side_lora) = outputs + _direct_grad_x_planes(sym_buffer).zero_() + sym_buffer.backward_grad_y[:grad_y.size(0)].copy_( + grad_y.to(torch.bfloat16).contiguous()) + sym_buffer.backward_grad_route.zero_() + num_grid_states = ( + expert_counts.numel() * + ((hidden // 64) * (intermediate_hidden // 128) + + ((2 * intermediate_hidden) // 64) * (hidden // 128)) + 2) + if grid_sync_counter is None: + grid_sync_counter = torch.zeros( + num_grid_states, dtype=torch.int32, device=gate_up_output.device) + if expert_psum_rows is None: + expert_psum_rows = padded_expert_counts.cumsum(0).to(torch.int32) + _C.bf16_mega_moe_backward_post_down_prelude_v2( + grad_ye, grad_ye, x_pool, route_weights, grad_route, + saved_down_unweighted, expert_counts, + sym_buffer.backward_grad_y, sym_buffer.side_lora_source, + sym_buffer.topk_weights, sym_buffer.backward_grad_route, + sym_buffer.token_src_metadata, sym_buffer.handle.buffer_ptrs, + sym_buffer.group.rank(), sym_buffer.num_topk, block_m, + CombineOrderMode.FIXED_TOPK.value, True, False, + False, True, True, False, False, 256) + _C.fp8_fp4_mega_moe_side_lora_backward( + gate_up_output, grad_h, grad_gate_up, h_act, h_weighted, + x_pool, grad_x_pool, l1_acts, l1_acts_sf, + l1_weights[0], l1_weights[1], grad_ye, route_weights, + w2_weights[0], w2_weights[1], w2_dequant_scratch, + w13_weights[0], w13_weights[1], w13_dequant_scratch, + expert_counts, grid_sync_counter, float(activation_limit), True, + bool(direct_remote_grad_x), + bool(write_grad_x_pool), True, block_m, + sym_buffer.handle.buffer_ptrs, + sym_buffer.group.rank(), sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_topk, + sym_buffer.backward_grad_y, sym_buffer.topk_weights, + sym_buffer.backward_grad_route, sym_buffer.token_src_metadata, + route_weight_mode.value, grad_ye, saved_down_unweighted, + grad_route, *side_lora, q13, q2, saved_h, t13, t2, + *grad_side_lora, expert_psum_rows, padded_expert_counts, + float(side_lora_scale)) + grad_x = None + if direct_remote_grad_x and combine_grad_x: + grad_x = ( + torch.empty_like(grad_y, dtype=torch.bfloat16) + if out is None or out.grad_x is None else out.grad_x) + mega_moe_backward_combine_grad_x( + grad_x, sym_buffer, CombineOrderMode.FIXED_TOPK) + return MegaMoESideLoraBackwardResult( + grad_x_pool, grad_x, grad_route, grad_ye, grad_h, + grad_gate_up, h_act, h_weighted, x_pool, route_weights, + t13, t2, grad_side_lora) + + def bf16_mega_moe_backward_w2( grad_w2_output: torch.Tensor, grad_ye: torch.Tensor, diff --git a/tests/benchmark_mega_moe_native_side_lora.py b/tests/benchmark_mega_moe_native_side_lora.py new file mode 100644 index 0000000000..8632a1ba74 --- /dev/null +++ b/tests/benchmark_mega_moe_native_side_lora.py @@ -0,0 +1,306 @@ +"""B300 EP benchmark for native MegaMoE side-LoRA specializations.""" + +import argparse +import json + +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.utils.dist import init_dist + + +def _block_m(tokens: int, ranks: int, topk: int, experts: int) -> int: + expected = tokens * ranks * topk / experts + if expected <= 8.5: + return 16 + if expected <= 16.5: + return 32 + if expected <= 32.5: + return 64 + if expected <= 64.5: + return 96 + if expected <= 96.5: + return 128 + return 192 + + +def _adapters(experts: int, hidden: int, intermediate: int): + rank = 128 + return ( + torch.randn(hidden, rank, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(experts, rank, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(hidden, rank, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(experts, rank, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(experts, intermediate, rank, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(rank, hidden, device="cuda", dtype=torch.bfloat16) * 0.02, + ) + + +def _time(fn, group, warmup: int, iterations: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + dist.barrier(group=group) + elapsed = 0.0 + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + sample = torch.tensor( + start.elapsed_time(end), dtype=torch.float64, device="cuda") + dist.all_reduce(sample, op=dist.ReduceOp.MAX, group=group) + elapsed += float(sample.item()) + return elapsed / iterations + + +def run_bf16(local_rank: int, world: int, args) -> None: + rank, ranks, group = init_dist(local_rank, world) + if ranks != args.num_processes: + raise RuntimeError("unexpected EP world size") + torch.manual_seed(9000 + rank) + tokens = args.tokens + hidden = args.hidden + intermediate = args.intermediate + experts = args.experts + topk = args.topk + local_experts = experts // ranks + block_m = _block_m(tokens, ranks, topk, experts) + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, experts, tokens, topk, hidden, intermediate, + mma_type="bf16xbf16", activation="swiglu") + + x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) * 0.1 + w13 = torch.randn( + local_experts, 2 * intermediate, hidden, + device="cuda", dtype=torch.bfloat16) * 0.02 + w2 = torch.randn( + local_experts, hidden, intermediate, + device="cuda", dtype=torch.bfloat16) * 0.02 + transformed_w13, transformed_w2 = deep_gemm.transform_weights_for_mega_moe( + w13, w2) + adapters = _adapters(local_experts, hidden, intermediate) + side = deep_gemm.transform_side_lora_for_mega_moe(adapters) + + route_id = ( + rank * tokens * topk + + torch.arange(tokens, device="cuda")[:, None] * topk + + torch.arange(topk, device="cuda")[None, :]) + topk_idx = (route_id % experts).to(torch.int64) + topk_weights = torch.sigmoid( + torch.randn(tokens, topk, device="cuda", dtype=torch.float32)) + buffer.x[:tokens].copy_(x) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + + local_start = rank * local_experts + mask = (topk_idx >= local_start) & (topk_idx < local_start + local_experts) + counts = torch.bincount( + topk_idx[mask] - local_start, minlength=local_experts).to(torch.int32) + dist.all_reduce(counts, group=group) + padded = ((counts + block_m - 1) // block_m * block_m).to(torch.int32) + pool_rows = int(padded.sum().item()) + route_counts = torch.bincount( + topk_idx.flatten(), minlength=experts).to(torch.int32) + + options = dict(device="cuda", dtype=torch.bfloat16) + y = torch.empty((tokens, hidden), **options) + base_gate = torch.empty((pool_rows, 2 * intermediate), **options) + base_h = torch.empty((pool_rows, intermediate), **options) + base_hw = torch.empty_like(base_h) + base_down = torch.empty((pool_rows, hidden), **options) + mismatch = torch.zeros(1, device="cuda", dtype=torch.int32) + + def base_forward(): + deep_gemm.bf16_mega_moe( + y, transformed_w13, transformed_w2, buffer, + saved_l1_preact=base_gate, + saved_h_unweighted=base_h, + saved_h_weighted=base_hw, + saved_down_unweighted=base_down, + precomputed_route_counts=route_counts, + active_pool_rows=pool_rows, + route_count_mismatch=mismatch, + num_config_tokens=tokens, + fast_math=True) + + base_forward_ms = _time( + base_forward, group, args.warmup, args.iterations) + if mismatch.item() != 0: + raise RuntimeError("forward route histogram mismatch") + + grad_y = torch.randn_like(y) + grad_ye = torch.empty((pool_rows, hidden), **options) + grad_h = torch.empty((pool_rows, intermediate), **options) + grad_gate = torch.empty_like(base_gate) + h_act = torch.empty_like(base_h) + h_weighted = torch.empty_like(base_h) + x_pool = torch.empty((pool_rows, hidden), **options) + grad_x_pool = torch.empty((0, hidden), **options) + grad_route = torch.empty(pool_rows, device="cuda", dtype=torch.float32) + route_weights = torch.empty_like(grad_route) + grid_states = local_experts * ( + (hidden // 64) * (intermediate // 128) + + ((2 * intermediate) // 64) * (hidden // 128)) + 2 + grid = torch.zeros(grid_states, device="cuda", dtype=torch.int32) + grad_w13 = torch.empty_like(w13) + grad_w2 = torch.empty_like(w2) + grad_x = torch.empty_like(y) + + def original_full_backward(): + deep_gemm.bf16_mega_moe_backward_dgrad( + base_gate, grad_h, grad_gate, h_act, h_weighted, + x_pool, grad_x_pool, grad_route, grad_ye, route_weights, + w2, w13, counts, grid, grad_y, buffer, + activation_limit=float("inf"), block_m=block_m, + fast_math=True, direct_remote_grad_x=True, + write_grad_x_pool=False, clear_wgrad_padding=True, + route_weight_mode=deep_gemm.RouteWeightMode.PRE_DOWN, + grad_y_unweighted_output=grad_ye, + down_unweighted_output=base_down, + combine_order_mode=deep_gemm.CombineOrderMode.FIXED_TOPK, + memory_mode="legacy", rank_uniform_block_m=True) + deep_gemm.bf16_mega_moe_backward_w2_combine( + grad_w2, grad_ye, h_weighted, padded, block_m, + grad_x, buffer) + deep_gemm.bf16_mega_moe_backward_w13_combine( + grad_w13, grad_gate, x_pool, padded, block_m, + grad_x, buffer) + + original_backward_ms = _time( + original_full_backward, group, args.warmup, args.iterations) + + side_gate = torch.empty_like(base_gate) + side_h = torch.empty_like(base_h) + side_hw = torch.empty_like(base_h) + side_down = torch.empty_like(base_down) + saved_x = torch.empty_like(base_down) + q13 = torch.empty((pool_rows, 2, 128), **options) + q2 = torch.empty((pool_rows, 128), **options) + ready = torch.empty( + (4, buffer.num_ring_tokens // 8), device="cuda", dtype=torch.int32) + + def side_forward(): + deep_gemm.bf16_mega_moe_side_lora( + y, transformed_w13, transformed_w2, buffer, + saved_l1_preact=side_gate, + saved_h_unweighted=side_h, + saved_h_weighted=side_hw, + saved_down_unweighted=side_down, + precomputed_route_counts=route_counts, + active_pool_rows=pool_rows, + route_count_mismatch=mismatch, + num_config_tokens=tokens, + saved_x=saved_x, side_lora=side, + side_lora_scale=args.scale, + side_lora_scratch=(q13, q2, ready), fast_math=True) + + side_forward_ms = _time( + side_forward, group, args.warmup, args.iterations) + + # First call allocates reusable outputs; timed calls reuse every tensor and + # the immutable padded-prefix plan. + side_result = deep_gemm.bf16_mega_moe_side_lora_backward( + side_gate, side_hw, side_down, q13, q2, side, + w2, transformed_w13, counts, padded, grad_y, buffer, block_m, + fast_math=True, side_lora_scale=args.scale, + direct_remote_grad_x=True, write_grad_x_pool=True) + side_grid = torch.zeros_like(grid) + expert_psum = padded.cumsum(0).to(torch.int32) + + def side_backward(): + nonlocal side_result + side_result = deep_gemm.bf16_mega_moe_side_lora_backward( + side_gate, side_hw, side_down, q13, q2, side, + w2, transformed_w13, counts, padded, grad_y, buffer, block_m, + fast_math=True, side_lora_scale=args.scale, + direct_remote_grad_x=True, write_grad_x_pool=True, + out=side_result, grid_sync_counter=side_grid, + expert_psum_rows=expert_psum) + # The caller represents the replicated 2-D A1/A3/B2 + # factors as DTensors with Partial gradients. Include those reductions + # so this is an end-to-end EP training comparison, not only kernel time. + for shared_grad in ( + side_result.grad_side_lora[0], + side_result.grad_side_lora[2], + side_result.grad_side_lora[5]): + dist.all_reduce(shared_grad, group=group) + + side_backward_ms = _time( + side_backward, group, args.warmup, args.iterations) + profile_table = None + if args.profile_side: + dist.barrier(group=group) + if rank == 0: + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA], + record_shapes=True, + ) as profile: + side_backward() + profile_table = profile.key_averages().table( + sort_by="self_cuda_time_total", row_limit=40) + else: + side_backward() + dist.barrier(group=group) + result = { + "mode": "bf16", + "device": torch.cuda.get_device_name(), + "ep_ranks": ranks, + "tokens_per_rank": tokens, + "hidden": hidden, + "intermediate": intermediate, + "experts": experts, + "local_experts": local_experts, + "topk": topk, + "block_m": block_m, + "pool_rows_per_rank": pool_rows, + "routes_per_local_expert": counts.tolist(), + "forward_ms": { + "original": base_forward_ms, + "side_lora": side_forward_ms, + "ratio": side_forward_ms / base_forward_ms, + }, + "backward_ms": { + "original_full_wgrad": original_backward_ms, + "side_lora_no_base_wgrad": side_backward_ms, + "speedup": original_backward_ms / side_backward_ms, + }, + "forward_backward_ms": { + "original_full_wgrad": base_forward_ms + original_backward_ms, + "side_lora_no_base_wgrad": side_forward_ms + side_backward_ms, + "speedup": ((base_forward_ms + original_backward_ms) / + (side_forward_ms + side_backward_ms)), + }, + "frozen_base_wgrads_emitted": False, + "full_width_side_scratch_emitted": False, + "shared_factor_ep_reductions_included": True, + } + if rank == 0: + print(json.dumps(result, indent=2), flush=True) + if profile_table is not None: + print(profile_table, flush=True) + dist.barrier(group=group) + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-processes", type=int, default=8) + parser.add_argument("--tokens", type=int, default=8192) + parser.add_argument("--hidden", type=int, default=7168) + parser.add_argument("--intermediate", type=int, default=3072) + parser.add_argument("--experts", type=int, default=384) + parser.add_argument("--topk", type=int, default=6) + parser.add_argument("--scale", type=float, default=0.25) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--profile-side", action="store_true") + args = parser.parse_args() + torch.multiprocessing.spawn( + run_bf16, args=(args.num_processes, args), + nprocs=args.num_processes) diff --git a/tests/run_mega_moe_side_lora_edge_matrix.py b/tests/run_mega_moe_side_lora_edge_matrix.py new file mode 100644 index 0000000000..3b21308e30 --- /dev/null +++ b/tests/run_mega_moe_side_lora_edge_matrix.py @@ -0,0 +1,97 @@ +"""Run the native MegaMoE side-LoRA forward/backward edge matrix. + +The numerical worker lives in ``test_mega_moe_native_side_lora.py``. Keeping +the matrix here makes every production contract explicit while each case gets +a fresh distributed process group and symmetric buffer. +""" + +import argparse +import subprocess +import sys +from pathlib import Path + + +WORKER = Path(__file__).with_name("test_mega_moe_native_side_lora.py") +COMMON = ( + "--hidden", "1024", + "--intermediate", "512", + "--experts", "8", +) + + +def _case(name: str, *arguments: str) -> tuple[str, tuple[str, ...]]: + return name, (*COMMON, *arguments) + + +EP1_CASES = ( + _case( + "bf16_single_route_single_token", + "--mode", "bf16", "--tokens", "1", "--topk", "1"), + _case( + "bf16_empty_experts_masked_bm_minus_one", + "--mode", "bf16", "--tokens", "15", "--topk", "2", + "--routing", "empty_experts", "--masked-ratio", "0.2", + "--activation-limit", "2.0"), + _case( + "bf16_geglu_bm_boundary", + "--mode", "bf16", "--tokens", "16", "--topk", "1", + "--routing", "skewed", "--activation", "geglu", + "--activation-limit", "1.5"), + _case( + "bf16_swiglu_bm_plus_one", + "--mode", "bf16", "--tokens", "17", "--topk", "2"), + _case( + "bf16_zero_scale_base_preservation", + "--mode", "bf16", "--tokens", "17", "--topk", "2", + "--scale", "0"), + _case( + "mxfp4_balanced_bm_plus_one", + "--mode", "mxfp4", "--tokens", "17", "--topk", "2"), + _case( + "mxfp4_zero_scale_base_preservation", + "--mode", "mxfp4", "--tokens", "17", "--topk", "2", + "--scale", "0"), + _case( + "mxfp4_empty_experts_masked_clamped", + "--mode", "mxfp4", "--tokens", "15", "--topk", "2", + "--routing", "empty_experts", "--masked-ratio", "0.2", + "--activation-limit", "2.0"), +) + + +EP_CASES = ( + _case( + "bf16_remote_masked_ep", + "--mode", "bf16", "--tokens", "17", "--topk", "2", + "--routing", "remote", "--masked-ratio", "0.2"), + _case( + "mxfp4_remote_masked_ep", + "--mode", "mxfp4", "--tokens", "17", "--topk", "2", + "--routing", "remote", "--masked-ratio", "0.2"), +) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--ep-processes", type=int, default=2, + help="number of GPUs for remote-route cases; use 1 to skip them") + args = parser.parse_args() + if args.ep_processes < 1: + parser.error("ep-processes must be positive") + + cases = list(EP1_CASES) + if args.ep_processes > 1: + cases.extend(EP_CASES) + for name, arguments in cases: + processes = ( + args.ep_processes if name.endswith("_ep") else 1) + command = ( + sys.executable, str(WORKER), + "--num-processes", str(processes), *arguments) + print(f"\n=== {name} ===", flush=True) + subprocess.run(command, check=True) + + +if __name__ == "__main__": + main() diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py new file mode 100644 index 0000000000..1bcb0e5731 --- /dev/null +++ b/tests/test_mega_moe_native_side_lora.py @@ -0,0 +1,972 @@ +"""Native MegaMoE side-LoRA correctness and EP benchmark driver. + +This intentionally exercises the dedicated kernels. PyTorch matmuls are used +only to construct a numerical reference after the native call has completed. +""" + +import argparse +import json +import math + +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils import ( + cast_back_from_fp4, per_token_cast_to_fp4, per_token_cast_to_fp8, + unpack_ue8m0_from_int) +from deep_gemm.utils.dist import init_dist + + +def _block_m(tokens: int, ranks: int, topk: int, experts: int) -> int: + expected = tokens * ranks * topk / experts + if expected <= 8.5: + return 16 + if expected <= 16.5: + return 32 + if expected <= 32.5: + return 64 + if expected <= 64.5: + return 96 + if expected <= 96.5: + return 128 + return 192 + + +def _active_rows(counts: torch.Tensor, padded: torch.Tensor) -> torch.Tensor: + rows = [] + offset = 0 + for count, capacity in zip(counts.cpu().tolist(), padded.cpu().tolist()): + rows.extend(range(offset, offset + count)) + offset += capacity + return torch.tensor(rows, dtype=torch.long, device=counts.device) + + +def _make_routing( + tokens: int, + topk: int, + experts: int, + rank: int, + ranks: int, + scenario: str, + masked_ratio: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build deterministic, duplicate-free routes for edge-case coverage.""" + token_ids = torch.arange(tokens, device="cuda").unsqueeze(1) + slots = torch.arange(topk, device="cuda").unsqueeze(0) + global_token_ids = rank * tokens + token_ids + if scenario == "balanced": + topk_idx = (global_token_ids * topk + slots) % experts + elif scenario == "skewed": + # Keep most traffic on the first experts while retaining distinct + # top-k destinations for every token. + hot_experts = max(topk, min(experts, 2 * topk)) + topk_idx = (global_token_ids + slots) % hot_experts + elif scenario == "empty_experts": + # Only the first top-k experts receive traffic. This exercises zero + # count experts and padding without violating the no-duplicate rule. + topk_idx = slots.expand(tokens, -1).clone() + elif scenario == "remote": + # Rotate the balanced assignment so EP ranks send to remote owners. + topk_idx = ( + global_token_ids * topk + slots + experts // ranks) % experts + else: + raise ValueError(f"unsupported routing scenario: {scenario}") + + route_ids = global_token_ids * topk + slots + topk_weights = (0.125 + (route_ids % 7).float() / 10).contiguous() + if masked_ratio: + period = max(2, round(1.0 / masked_ratio)) + mask = route_ids.remainder(period) == 0 + topk_idx.masked_fill_(mask, -1) + topk_weights.masked_fill_(mask, 0.0) + return topk_idx.contiguous(), topk_weights + + +def _counts( + topk_idx: torch.Tensor, + experts: int, + local_start: int, + local_experts: int, + group, +) -> tuple[torch.Tensor, torch.Tensor]: + valid = topk_idx >= 0 + source_counts = torch.bincount( + topk_idx[valid], minlength=experts).to(torch.int32) + global_counts = source_counts.clone() + dist.all_reduce(global_counts, group=group) + return ( + source_counts, + global_counts[local_start:local_start + local_experts].contiguous(), + ) + + +def _all_gather_equal(tensor: torch.Tensor, group) -> torch.Tensor: + gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group))] + dist.all_gather(gathered, tensor.contiguous(), group=group) + return torch.stack(gathered) + + +def _rank_uniform_max(value: int, group) -> int: + value_tensor = torch.tensor(value, dtype=torch.int32, device="cuda") + dist.all_reduce(value_tensor, op=dist.ReduceOp.MAX, group=group) + return int(value_tensor.item()) + + +def _scatter_to_sources( + rows: torch.Tensor, + metadata: torch.Tensor, + ranks: int, + tokens: int, + group, +) -> torch.Tensor: + planes = torch.zeros( + (ranks * tokens, *rows.shape[1:]), + dtype=rows.dtype, device=rows.device) + source_rows = metadata[:, 0] * tokens + metadata[:, 1] + planes.index_add_(0, source_rows, rows) + dist.all_reduce(planes, group=group) + return planes.view(ranks, tokens, *rows.shape[1:]) + + +def _scatter_route_grads( + rows: torch.Tensor, + metadata: torch.Tensor, + ranks: int, + tokens: int, + topk: int, + group, +) -> torch.Tensor: + flat = torch.zeros( + ranks * tokens * topk, dtype=rows.dtype, device=rows.device) + indices = ( + (metadata[:, 0] * tokens + metadata[:, 1]) * topk + + metadata[:, 2]) + flat.index_add_(0, indices, rows) + dist.all_reduce(flat, group=group) + return flat.view(ranks, tokens, topk) + + +def _apply_activation( + gate: torch.Tensor, + up: torch.Tensor, + activation: str, + activation_limit: float, +) -> torch.Tensor: + gate_clamped = torch.clamp(gate, max=activation_limit) + up_clamped = torch.clamp( + up, min=-activation_limit, max=activation_limit) + if activation == "swiglu": + activated = torch.nn.functional.silu(gate_clamped) + elif activation == "geglu": + alpha = 1.5957691216057308 + beta = 0.044715 + gate_sq = gate_clamped * gate_clamped + activated = gate_clamped * torch.sigmoid( + (alpha * gate_clamped) * (1.0 + beta * gate_sq)) + else: + raise ValueError(f"unsupported activation: {activation}") + if activation == "swiglu" and math.isinf(activation_limit): + return activated * up_clamped + return (activated.float() * up_clamped.float()).to(torch.bfloat16) + + +def _cast_fp4(weights: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + groups, n, k = weights.shape + packed = torch.empty((groups, n, k // 2), dtype=torch.int8, device="cuda") + scales = torch.empty((groups, n, k // 32), dtype=torch.float32, device="cuda") + for group in range(groups): + packed[group], scales[group] = per_token_cast_to_fp4( + weights[group], use_ue8m0=True, gran_k=32) + scales = deep_gemm.transform_sf_into_required_layout( + scales, n, k, (1, 32), groups) + return packed, scales + + +def _cast_fp4_backward(weights: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + packed = torch.empty( + (*weights.shape[:-1], weights.size(-1) // 2), + dtype=torch.int8, device="cuda") + scales = torch.empty( + (*weights.shape[:-1], weights.size(-1) // 32), + dtype=torch.float32, device="cuda") + for expert in range(weights.size(0)): + packed[expert], scales[expert] = per_token_cast_to_fp4( + weights[expert], use_ue8m0=True, gran_k=32) + return packed.view(torch.float8_e4m3fn), scales + + +def _adapters(experts: int, hidden: int, intermediate: int): + rank = 128 + return ( + (torch.randn(hidden, rank, device="cuda", dtype=torch.bfloat16) * 0.02), + (torch.randn(experts, rank, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02), + (torch.randn(hidden, rank, device="cuda", dtype=torch.bfloat16) * 0.02), + (torch.randn(experts, rank, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02), + (torch.randn(experts, intermediate, rank, device="cuda", dtype=torch.bfloat16) * 0.02), + (torch.randn(rank, hidden, device="cuda", dtype=torch.bfloat16) * 0.02), + ) + + +def _dequant_fp4(weights: torch.Tensor) -> torch.Tensor: + output = torch.empty_like(weights, dtype=torch.float32) + for expert in range(weights.size(0)): + packed, scales = per_token_cast_to_fp4( + weights[expert], use_ue8m0=True, gran_k=32) + output[expert] = cast_back_from_fp4(packed, scales, gran_k=32) + return output + + +def _dequant_fp8(x: torch.Tensor, packed_scales: torch.Tensor) -> torch.Tensor: + rows, width = x.shape + scales = unpack_ue8m0_from_int(packed_scales)[:, :width // 32] + return ( + x.float().view(rows, width // 32, 32) * scales.unsqueeze(2) + ).reshape(rows, width) + + +def _relative(actual: torch.Tensor, expected: torch.Tensor) -> float: + if expected.numel() == 0: + return 0.0 + if float(expected.detach().float().norm()) == 0.0: + return float(actual.detach().float().norm()) + return float(calc_diff(actual.detach().float(), expected.detach().float())) + + +def _accuracy(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + actual_f = actual.detach().float() + expected_f = expected.detach().float() + delta = actual_f - expected_f + actual_sq = actual_f.square().sum().double() + expected_sq = expected_f.square().sum().double() + if float(expected_sq) == 0.0: + actual_norm = math.sqrt(float(actual_sq)) + return { + "relative_l2": actual_norm, + "cosine_similarity": 1.0 if actual_norm == 0.0 else 0.0, + "max_abs": float(delta.abs().max()) if delta.numel() else 0.0, + } + return { + "relative_l2": math.sqrt( + float(delta.square().sum().double() / expected_sq) + ), + "cosine_similarity": float( + (actual_f * expected_f).sum().double() + / torch.sqrt(actual_sq * expected_sq) + ), + "max_abs": float(delta.abs().max()), + } + + +def _mx_interleave_pair(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + width = left.size(1) + return torch.stack( + (left.view(-1, width // 8, 8), right.view(-1, width // 8, 8)), + dim=2).reshape(-1, 2 * width) + + +def _mx_deinterleave_pair(value: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + width = value.size(1) // 2 + chunks = value.view(-1, width // 8, 2, 8) + return chunks[:, :, 0].reshape(-1, width), chunks[:, :, 1].reshape(-1, width) + + +def test_side_lora_backward_rejects_unsupported_post_down() -> None: + """Do not silently run a numerically invalid backward boundary.""" + calls = ( + ( + deep_gemm.bf16_mega_moe_side_lora_backward, + dict( + gate_up_output=None, saved_h=None, + saved_down_unweighted=None, q13=None, q2=None, + side_lora=None, w2_weights=None, w13_weights=None, + expert_counts=None, padded_expert_counts=None, + grad_y=None, sym_buffer=None, block_m=16, + ), + ), + ( + deep_gemm.fp8_fp4_mega_moe_side_lora_backward, + dict( + gate_up_output=None, saved_h=None, + saved_down_unweighted=None, q13=None, q2=None, + side_lora=None, l1_acts=None, l1_acts_sf=None, + l1_weights=None, w13_weights=None, w2_weights=None, + w13_dequant_scratch=None, w2_dequant_scratch=None, + expert_counts=None, padded_expert_counts=None, + grad_y=None, sym_buffer=None, block_m=16, + ), + ), + ) + for function, keywords in calls: + try: + function( + **keywords, + route_weight_mode=deep_gemm.RouteWeightMode.POST_DOWN) + except NotImplementedError as error: + assert "pre_down" in str(error) + else: + raise AssertionError( + f"{function.__name__} accepted unsupported post_down") + + +def test_side_lora_transform_validates_shared_layout() -> None: + hidden, intermediate, experts, rank = 256, 128, 4, 128 + side_lora = ( + torch.empty(hidden, rank, dtype=torch.bfloat16), + torch.empty(experts, rank, intermediate, dtype=torch.bfloat16), + torch.empty(hidden, rank, dtype=torch.bfloat16), + torch.empty(experts, rank, intermediate, dtype=torch.bfloat16), + torch.empty(experts, intermediate, rank, dtype=torch.bfloat16), + torch.empty(rank, hidden, dtype=torch.bfloat16), + ) + transformed = deep_gemm.transform_side_lora_for_mega_moe(side_lora) + assert [tuple(tensor.shape) for tensor in transformed] == [ + (rank, hidden), + (experts, intermediate, rank), + (rank, hidden), + (experts, intermediate, rank), + (experts, rank, intermediate), + (hidden, rank), + ] + + invalid_b2 = (*side_lora[:-1], side_lora[-1][:, :-1]) + try: + deep_gemm.transform_side_lora_for_mega_moe(invalid_b2) + except ValueError as error: + assert "B2" in str(error) + else: + raise AssertionError("accepted an inconsistent shared B2 shape") + + +def run_bf16_correctness(local_rank: int, world: int, args) -> None: + rank, ranks, group = init_dist(local_rank, world) + torch.manual_seed(1234 + rank) + tokens, hidden, intermediate = args.tokens, args.hidden, args.intermediate + experts, topk = args.experts, args.topk + local_experts = experts // ranks + block_m = _block_m(tokens, ranks, topk, experts) + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, experts, tokens, topk, hidden, intermediate, + mma_type="bf16xbf16", activation=args.activation) + + x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) * 0.1 + w13 = torch.randn(local_experts, 2 * intermediate, hidden, device="cuda", dtype=torch.bfloat16) * 0.02 + w2 = torch.randn(local_experts, hidden, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02 + transformed_w13, transformed_w2 = deep_gemm.transform_weights_for_mega_moe(w13, w2) + adapters = _adapters(local_experts, hidden, intermediate) + side = deep_gemm.transform_side_lora_for_mega_moe(adapters) + + topk_idx, topk_weights = _make_routing( + tokens, topk, experts, rank, ranks, + args.routing, args.masked_ratio) + route_counts, counts = _counts( + topk_idx, experts, rank * local_experts, local_experts, group) + padded = ((counts + block_m - 1) // block_m * block_m).to(torch.int32) + local_pool_rows = int(padded.sum().item()) + pool_rows = _rank_uniform_max(local_pool_rows, group) + active = _active_rows(counts, padded) + + saved_gate_up = torch.full((pool_rows, 2 * intermediate), float("nan"), device="cuda", dtype=torch.bfloat16) + saved_h = torch.full((pool_rows, intermediate), float("nan"), device="cuda", dtype=torch.bfloat16) + saved_h_weighted = torch.full_like(saved_h, float("nan")) + saved_down = torch.full((pool_rows, hidden), float("nan"), device="cuda", dtype=torch.bfloat16) + saved_x = torch.full_like(saved_down, float("nan")) + q13 = torch.empty((pool_rows, 2, 128), device="cuda", dtype=torch.bfloat16) + q2 = torch.empty((pool_rows, 128), device="cuda", dtype=torch.bfloat16) + ready = torch.zeros((4, buffer.num_ring_tokens // 8), device="cuda", dtype=torch.int32) + mismatch = torch.zeros(1, device="cuda", dtype=torch.int32) + + buffer.x[:tokens].copy_(x) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + y = torch.empty_like(x) + deep_gemm.bf16_mega_moe_side_lora( + y, transformed_w13, transformed_w2, buffer, + saved_l1_preact=saved_gate_up, + saved_h_unweighted=saved_h, + saved_h_weighted=saved_h_weighted, + saved_down_unweighted=saved_down, + saved_x=saved_x, + precomputed_route_counts=route_counts, + active_pool_rows=pool_rows, + route_count_mismatch=mismatch, + num_config_tokens=tokens, + side_lora=side, + side_lora_scale=args.scale, + side_lora_scratch=(q13, q2, ready), + activation=args.activation, + activation_clamp=args.activation_limit, + route_weight_mode=args.route_weight_mode, + fast_math=False) + torch.cuda.synchronize() + assert mismatch.item() == 0 + base_preservation = None + + # Build an autograd reference on the exact dispatched rows and preserve + # the same BF16 materialization boundaries as the native epilogues. + metadata = buffer.token_src_metadata[active].long() + x_ref = saved_x[active].detach().clone().requires_grad_(True) + grad_y = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + all_grad_y = _all_gather_equal(grad_y, group) + all_route_weights = _all_gather_equal(topk_weights, group) + route_ref = all_route_weights[ + metadata[:, 0], metadata[:, 1], metadata[:, 2] + ].to(torch.bfloat16).detach().clone().requires_grad_(True) + adapter_ref = tuple(t.detach().clone().requires_grad_(True) for t in adapters) + output_rows, down_rows = [], [] + gate_rows, up_rows, h_rows, h_weighted_rows = [], [], [], [] + q1_rows, q3_rows, q2_rows = [], [], [] + cursor = 0 + for expert, count in enumerate(counts.cpu().tolist()): + xe = x_ref[cursor:cursor + count] + a1, b1, a3, b3, a2, b2 = ( + adapter_ref[0], adapter_ref[1][expert], + adapter_ref[2], adapter_ref[3][expert], + adapter_ref[4][expert], adapter_ref[5]) + gate_base = (xe @ w13[expert, :intermediate].t()).to(torch.bfloat16) + up_base = (xe @ w13[expert, intermediate:].t()).to(torch.bfloat16) + q1 = (xe @ a1).to(torch.bfloat16) + q3 = (xe @ a3).to(torch.bfloat16) + gate = torch.add(gate_base, (q1 @ b1).to(torch.bfloat16), alpha=args.scale) + up = torch.add(up_base, (q3 @ b3).to(torch.bfloat16), alpha=args.scale) + h = _apply_activation( + gate, up, args.activation, args.activation_limit) + routes = route_ref[cursor:cursor + count] + h_weighted = ( + h.float() * routes.float().unsqueeze(1)).to(torch.bfloat16) + h_for_w2 = ( + h_weighted + if args.route_weight_mode == "pre_down" + else h.to(torch.bfloat16)) + q_down = (h_for_w2 @ a2).to(torch.bfloat16) + down_base = (h_for_w2 @ w2[expert].t()).to(torch.bfloat16) + down = torch.add(down_base, (q_down @ b2).to(torch.bfloat16), alpha=args.scale) + down_rows.append(down) + output_rows.append( + down if args.route_weight_mode == "pre_down" else + (down.float() * routes.float().unsqueeze(1)).to(torch.bfloat16)) + gate_rows.append(gate) + up_rows.append(up) + h_rows.append(h) + h_weighted_rows.append(h_weighted) + q1_rows.append(q1) + q3_rows.append(q3) + q2_rows.append(q_down) + cursor += count + route_output = torch.cat(output_rows) + down_unweighted_ref = torch.cat(down_rows) + route_grad = all_grad_y[metadata[:, 0], metadata[:, 1]] + (route_output.float() * route_grad.float()).sum().backward() + + forward_diff = _relative( + saved_gate_up[active], + torch.cat((torch.cat(gate_rows), torch.cat(up_rows)), dim=1)) + h_diff = max( + _relative(saved_h[active], torch.cat(h_rows)), + _relative(saved_h_weighted[active], torch.cat(h_weighted_rows))) + down_diff = _relative(saved_down[active], down_unweighted_ref) + expected_y = _scatter_to_sources( + route_output, metadata, ranks, tokens, group)[rank] + output_diff = _relative(y, expected_y) + q_diff = max( + _relative(q13[active, 0], torch.cat(q1_rows)), + _relative(q13[active, 1], torch.cat(q3_rows)), + _relative(q2[active], torch.cat(q2_rows)), + ) + result = deep_gemm.bf16_mega_moe_side_lora_backward( + saved_gate_up, + saved_h_weighted if args.route_weight_mode == "pre_down" else saved_h, + saved_down, q13, q2, side, + w2, transformed_w13, counts, padded, grad_y, buffer, + block_m, activation_limit=args.activation_limit, + activation=args.activation, fast_math=False, + route_weight_mode=args.route_weight_mode, + side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) + torch.cuda.synchronize() + grad_diffs = [ + _relative(actual, expected.grad) + for actual, expected in zip(result.grad_side_lora, adapter_ref) + ] + grad_x_diff = _relative(result.grad_x_pool[active], x_ref.grad) + source_grad_x = _scatter_to_sources( + x_ref.grad.to(torch.bfloat16), metadata, ranks, tokens, group)[rank] + combined_grad_x_diff = ( + _relative(result.grad_x, source_grad_x) if ranks > 1 else 0.0) + route_grad_diff = _relative(result.grad_route[active], route_ref.grad) + route_grad_max_abs = float( + (result.grad_route[active] - route_ref.grad).abs().max()) + source_route_grad = _scatter_route_grads( + route_ref.grad, metadata, ranks, tokens, topk, group)[rank] + combined_route_grad_diff = _relative( + buffer.backward_grad_route[:tokens, :topk], source_route_grad) + combined_route_grad_max_abs = float(( + buffer.backward_grad_route[:tokens, :topk] - source_route_grad + ).abs().max()) + route_grad_close = ( + route_grad_diff < 0.03 or route_grad_max_abs < 0.01) + combined_route_grad_close = ( + combined_route_grad_diff < 0.03 or + combined_route_grad_max_abs < 0.01) + diagnostics = { + "actual_grad_norms": [ + float(t.float().norm()) for t in result.grad_side_lora], + "reference_grad_norms": [ + float(t.grad.float().norm()) for t in adapter_ref], + "t13_norm": float(result.t13[active].float().norm()), + "t2_norm": float(result.t2[active].float().norm()), + } + if rank == 0 and ( + max(grad_diffs) >= 0.03 or grad_x_diff >= 0.03 or + combined_grad_x_diff >= 0.03 or not route_grad_close or + not combined_route_grad_close + ): + print(json.dumps({ + "adapter_grad_diffs": grad_diffs, + "grad_x_diff": grad_x_diff, + "combined_grad_x_diff": combined_grad_x_diff, + "route_grad_diff": route_grad_diff, + "route_grad_max_abs": route_grad_max_abs, + "combined_route_grad_diff": combined_route_grad_diff, + "combined_route_grad_max_abs": combined_route_grad_max_abs, + "zero_scale_base_preservation": base_preservation, + **diagnostics, + }, indent=2), flush=True) + forward_tolerance = 0.04 if args.scale == 0.0 else 0.02 + assert forward_diff < forward_tolerance, forward_diff + assert h_diff < forward_tolerance, h_diff + assert down_diff < forward_tolerance, down_diff + if args.scale != 0.0: + assert output_diff < forward_tolerance, output_diff + assert q_diff < forward_tolerance, q_diff + assert max(grad_diffs) < 0.03, grad_diffs + assert grad_x_diff < 0.03, grad_x_diff + assert combined_grad_x_diff < 0.03, combined_grad_x_diff + assert route_grad_close, (route_grad_diff, route_grad_max_abs) + assert combined_route_grad_close, ( + combined_route_grad_diff, combined_route_grad_max_abs) + if args.scale == 0.0: + base_y = torch.empty_like(y) + buffer.x[:tokens].copy_(x) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + deep_gemm.bf16_mega_moe( + base_y, transformed_w13, transformed_w2, buffer, + activation=args.activation, + activation_clamp=args.activation_limit, + route_weight_mode=args.route_weight_mode, + num_config_tokens=tokens, fast_math=False) + torch.cuda.synchronize() + base_preservation = _accuracy(y, base_y) + assert base_preservation["relative_l2"] == 0.0, base_preservation + if rank == 0: + print(json.dumps({ + "mode": "bf16", "ranks": ranks, "tokens_per_rank": tokens, + "routing": args.routing, + "masked_ratio": args.masked_ratio, + "activation": args.activation, + "activation_limit": args.activation_limit, + "route_weight_mode": args.route_weight_mode, + "pool_rows": pool_rows, "block_m": block_m, + "forward_diff": forward_diff, "activation_diff": h_diff, + "down_diff": down_diff, "output_diff": output_diff, + "side_q_diff": q_diff, "adapter_grad_diffs": grad_diffs, + "grad_x_diff": grad_x_diff, + "combined_grad_x_diff": combined_grad_x_diff, + "route_grad_diff": route_grad_diff, + "route_grad_max_abs": route_grad_max_abs, + "combined_route_grad_diff": combined_route_grad_diff, + "combined_route_grad_max_abs": combined_route_grad_max_abs, + "zero_scale_base_preservation": base_preservation, + **diagnostics, + }, indent=2), flush=True) + buffer.destroy() + dist.destroy_process_group() + + +def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: + rank, ranks, group = init_dist(local_rank, world) + if args.activation != "swiglu": + raise ValueError("MXFP4 side-LoRA backward currently supports SwiGLU") + torch.manual_seed(4321 + rank) + tokens, hidden, intermediate = args.tokens, args.hidden, args.intermediate + experts, topk = args.experts, args.topk + local_experts = experts // ranks + block_m = _block_m(tokens, ranks, topk, experts) + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, experts, tokens, topk, hidden, intermediate, + mma_type="fp8xfp4", activation="swiglu") + + x_bf16 = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) * 0.1 + x_fp8 = per_token_cast_to_fp8( + x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + w13_bf16 = torch.randn(local_experts, 2 * intermediate, hidden, device="cuda", dtype=torch.bfloat16) * 0.02 + w2_bf16 = torch.randn(local_experts, hidden, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02 + w13 = _cast_fp4(w13_bf16) + w2 = _cast_fp4(w2_bf16) + transformed_w13, transformed_w2 = deep_gemm.transform_weights_for_mega_moe(w13, w2) + backward_w13_raw = _cast_fp4_backward(w13_bf16) + backward_w2 = _cast_fp4_backward(w2_bf16) + backward_w13 = ( + backward_w13_raw[0].view( + 2 * local_experts, intermediate, hidden // 2), + backward_w13_raw[1].view( + 2 * local_experts, intermediate, hidden // 32), + ) + adapters = _adapters(local_experts, hidden, intermediate) + side = deep_gemm.transform_side_lora_for_mega_moe(adapters) + topk_idx, topk_weights = _make_routing( + tokens, topk, experts, rank, ranks, + args.routing, args.masked_ratio) + _, counts = _counts( + topk_idx, experts, rank * local_experts, local_experts, group) + padded = ((counts + block_m - 1) // block_m * block_m).to(torch.int32) + local_pool_rows = int(padded.sum().item()) + pool_rows = _rank_uniform_max(local_pool_rows, group) + active = _active_rows(counts, padded) + + saved_gate_up = torch.full((pool_rows, 2 * intermediate), float("nan"), device="cuda", dtype=torch.bfloat16) + saved_h = torch.full((pool_rows, intermediate), float("nan"), device="cuda", dtype=torch.bfloat16) + saved_down = torch.full((pool_rows, hidden), float("nan"), device="cuda", dtype=torch.bfloat16) + saved_x = torch.full_like(saved_down, float("nan")) + q13 = torch.empty((pool_rows, 2, 128), device="cuda", dtype=torch.bfloat16) + q2 = torch.empty((pool_rows, 128), device="cuda", dtype=torch.bfloat16) + ready = torch.zeros(4 * buffer.num_ring_tokens // 8, device="cuda", dtype=torch.int32) + buffer.x[:tokens].copy_(x_fp8[0]) + buffer.x_sf[:tokens].copy_(x_fp8[1]) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + y = torch.empty_like(x_bf16) + deep_gemm.fp8_fp4_mega_moe_side_lora( + y, transformed_w13, transformed_w2, buffer, + side_lora_input=x_bf16, side_lora=side, + saved_x=saved_x, saved_h_unweighted=saved_h, + saved_l1_preact=saved_gate_up, + saved_down_unweighted=saved_down, + num_config_tokens=tokens, side_lora_scale=args.scale, + side_lora_scratch=(q13, q2, ready), + activation_clamp=args.activation_limit, + route_weight_mode=args.route_weight_mode, + fast_math=False) + torch.cuda.synchronize() + base_preservation = None + + metadata = buffer.token_src_metadata[active].long() + x_deq = _dequant_fp8(*x_fp8) + all_x_deq = _all_gather_equal(x_deq, group) + all_route_weights = _all_gather_equal(topk_weights, group) + x_base_ref = all_x_deq[ + metadata[:, 0], metadata[:, 1] + ].detach().clone().requires_grad_(True) + x_side_ref = saved_x[active].detach().clone().requires_grad_(True) + adapter_ref = tuple( + tensor.detach().clone().requires_grad_(True) for tensor in adapters) + grad_y = torch.randn_like(y) + all_grad_y = _all_gather_equal(grad_y, group) + route_ref = all_route_weights[ + metadata[:, 0], metadata[:, 1], metadata[:, 2] + ].to(torch.bfloat16).detach().clone().requires_grad_(True) + w13_deq = _dequant_fp4(w13_bf16) + w2_deq = _dequant_fp4(w2_bf16) + gate_rows, up_rows = [], [] + h_unweighted_rows, h_for_w2_rows = [], [] + down_rows, down_unweighted_rows = [], [] + q1_rows, q3_rows, q2_rows = [], [], [] + cursor = 0 + for expert, count in enumerate(counts.cpu().tolist()): + xe_side = x_side_ref[cursor:cursor + count] + xe_base = x_base_ref[cursor:cursor + count] + a1, b1, a3, b3, a2, b2 = ( + adapter_ref[0], adapter_ref[1][expert], + adapter_ref[2], adapter_ref[3][expert], + adapter_ref[4][expert], adapter_ref[5]) + q1 = (xe_side @ a1).to(torch.bfloat16) + q3 = (xe_side @ a3).to(torch.bfloat16) + gate = torch.add( + (xe_base @ w13_deq[expert, :intermediate].t()).to(torch.bfloat16), + (q1 @ b1).to(torch.bfloat16), alpha=args.scale) + up = torch.add( + (xe_base @ w13_deq[expert, intermediate:].t()).to(torch.bfloat16), + (q3 @ b3).to(torch.bfloat16), alpha=args.scale) + h_unweighted = _apply_activation( + gate, up, "swiglu", args.activation_limit) + route = route_ref[cursor:cursor + count] + h_weighted = ( + h_unweighted.float() * route.float().unsqueeze(1) + ).to(torch.bfloat16) + h = ( + h_weighted + if args.route_weight_mode == "pre_down" + else h_unweighted.to(torch.bfloat16)) + q_down = (h @ a2).to(torch.bfloat16) + down = torch.add( + (h.float() @ w2_deq[expert].t()).to(torch.bfloat16), + (q_down @ b2).to(torch.bfloat16), alpha=args.scale) + gate_rows.append(gate); up_rows.append(up) + h_unweighted_rows.append(h_unweighted) + h_for_w2_rows.append(h) + route_output = ( + down if args.route_weight_mode == "pre_down" else + (down.float() * route.float().unsqueeze(1)).to(torch.bfloat16)) + down_unweighted_rows.append(down) + down_rows.append(route_output); q1_rows.append(q1); q3_rows.append(q3) + q2_rows.append(q_down) + gate.retain_grad(); up.retain_grad(); h_unweighted.retain_grad(); h.retain_grad() + q1.retain_grad(); q3.retain_grad(); q_down.retain_grad() + cursor += count + gate_ref = _mx_interleave_pair(torch.cat(gate_rows), torch.cat(up_rows)) + down_ref = torch.cat(down_rows) + expected_y = _scatter_to_sources( + down_ref, metadata, ranks, tokens, group)[rank] + route_grad_y = all_grad_y[metadata[:, 0], metadata[:, 1]] + (down_ref.float() * route_grad_y.float()).sum().backward() + # The production backward is free to phase-reuse forward-only saved + # storage. Preserve the native forward boundary for post-run diagnostics. + saved_gate_up_before_backward = saved_gate_up.clone() + forward_diffs = { + "gate_up": _relative(saved_gate_up[active], gate_ref), + "h": _relative(saved_h[active], torch.cat(h_for_w2_rows)), + "down": _relative( + saved_down[active], torch.cat(down_unweighted_rows)), + "output": _relative(y, expected_y), + "q": max( + _relative(q13[active, 0], torch.cat(q1_rows)), + _relative(q13[active, 1], torch.cat(q3_rows)), + _relative(q2[active], torch.cat(q2_rows))), + } + result = deep_gemm.fp8_fp4_mega_moe_side_lora_backward( + saved_gate_up, saved_h, saved_down, q13, q2, side, + buffer.l1_acts[:pool_rows], buffer.l1_acts_sf, + transformed_w13, backward_w13, backward_w2, + torch.empty((local_experts, 2 * intermediate, hidden), device="cuda", dtype=torch.bfloat16), + torch.empty((local_experts, hidden, intermediate), device="cuda", dtype=torch.bfloat16), + counts, padded, grad_y, buffer, block_m, + activation_limit=args.activation_limit, + route_weight_mode=args.route_weight_mode, + side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) + torch.cuda.synchronize() + + expected_grads = [tensor.grad for tensor in adapter_ref] + expected_grad_x = ( + x_base_ref.grad.float() + x_side_ref.grad.float() + ).to(torch.bfloat16) + grad_diffs = [ + _relative(actual, expected) + for actual, expected in zip(result.grad_side_lora, expected_grads)] + grad_x_diff = _relative(result.grad_x_pool[active], expected_grad_x) + source_grad_x = _scatter_to_sources( + expected_grad_x, metadata, ranks, tokens, group)[rank] + combined_grad_x_diff = ( + _relative(result.grad_x, source_grad_x) if ranks > 1 else 0.0) + route_grad_diff = _relative(result.grad_route[active], route_ref.grad) + source_route_grad = _scatter_route_grads( + route_ref.grad, metadata, ranks, tokens, topk, group)[rank] + combined_route_grad_diff = _relative( + buffer.backward_grad_route[:tokens, :topk], source_route_grad) + grad_accuracy = [ + _accuracy(actual, expected) + for actual, expected in zip(result.grad_side_lora, expected_grads) + ] + grad_x_accuracy = _accuracy(result.grad_x_pool[active], expected_grad_x) + backward_boundary_accuracy = { + "grad_h": _accuracy( + result.grad_h[active], + torch.cat([tensor.grad for tensor in h_for_w2_rows]), + ), + "grad_gate_up": _accuracy( + result.grad_gate_up[active], + torch.cat( + ( + torch.cat([tensor.grad for tensor in gate_rows]), + torch.cat([tensor.grad for tensor in up_rows]), + ), + dim=1, + ), + ), + "t2": _accuracy( + result.t2[active].float() * args.scale, + torch.cat([tensor.grad for tensor in q2_rows]), + ), + "t1": _accuracy( + result.t13[active, 0].float() * args.scale, + torch.cat([tensor.grad for tensor in q1_rows]), + ), + "t3": _accuracy( + result.t13[active, 1].float() * args.scale, + torch.cat([tensor.grad for tensor in q3_rows]), + ), + } + # Diagnose the exact fused dSwiGLU boundary independently of either base + # dgrad kernel. This reference starts from the native W2+LoRA grad-h and + # the native forward's saved gate/up values, so any discrepancy here is + # owned by the activation-backward phase itself. + saved_gate, saved_up = _mx_deinterleave_pair( + saved_gate_up_before_backward[active]) + active_route = route_ref.detach() + grad_h_for_activation = result.grad_h[active] + if args.route_weight_mode == "pre_down": + grad_h_for_activation = ( + grad_h_for_activation.float() * + active_route.float().unsqueeze(1)).to(torch.bfloat16) + native_gate = saved_gate.detach().clone().requires_grad_(True) + native_up = saved_up.detach().clone().requires_grad_(True) + native_h = _apply_activation( + native_gate, native_up, "swiglu", args.activation_limit) + (native_h.float() * grad_h_for_activation.float()).sum().backward() + exact_grad_gate = native_gate.grad + exact_grad_up = native_up.grad + backward_boundary_accuracy["grad_gate_up_from_native_boundaries"] = _accuracy( + result.grad_gate_up[active], + torch.cat((exact_grad_gate, exact_grad_up), dim=1), + ) + autograd_grad_gate = torch.cat([tensor.grad for tensor in gate_rows]) + autograd_grad_up = torch.cat([tensor.grad for tensor in up_rows]) + backward_boundary_accuracy["native_boundary_vs_autograd_gate"] = _accuracy( + exact_grad_gate, autograd_grad_gate) + backward_boundary_accuracy["native_boundary_vs_autograd_up"] = _accuracy( + exact_grad_up, autograd_grad_up) + backward_boundary_accuracy["kernel_vs_native_boundary_gate"] = _accuracy( + result.grad_gate_up[active, :intermediate], exact_grad_gate) + backward_boundary_accuracy["kernel_vs_native_boundary_up"] = _accuracy( + result.grad_gate_up[active, intermediate:], exact_grad_up) + backward_boundary_accuracy["native_vs_autograd_grad_h_unweighted"] = _accuracy( + grad_h_for_activation, + torch.cat([tensor.grad for tensor in h_unweighted_rows]), + ) + backward_boundary_accuracy["saved_vs_autograd_gate_value"] = _accuracy( + saved_gate, torch.cat(gate_rows)) + backward_boundary_accuracy["saved_vs_autograd_up_value"] = _accuracy( + saved_up, torch.cat(up_rows)) + backward_boundary_accuracy["saved_vs_autograd_silu_value"] = _accuracy( + native_h, _apply_activation( + torch.cat(gate_rows), torch.cat(up_rows), + "swiglu", args.activation_limit)) + backward_boundary_accuracy["saved_gate_up_after_backward_reuse"] = _accuracy( + saved_gate_up[active], saved_gate_up_before_backward[active]) + # Validate the six native wgrad contractions on the kernel's own exact + # forward/backward boundaries. This removes unrelated MXFP4 base-GEMM + # ordering drift from the adapter-gradient check. + native_boundary_grad_refs = [ + (x_side_ref.detach().float().t() @ result.t13[active, 0].float() + * args.scale).bfloat16(), + [], + (x_side_ref.detach().float().t() @ result.t13[active, 1].float() + * args.scale).bfloat16(), + [], + [], + (q2[active].float().t() + @ (route_grad_y.float() * + (active_route.float().unsqueeze(1) + if args.route_weight_mode == "post_down" else 1.0)) + * args.scale).bfloat16(), + ] + cursor = 0 + for expert, count in enumerate(counts.cpu().tolist()): + rows = slice(cursor, cursor + count) + native_boundary_grad_refs[1].append( + (q13[active][rows, 0].float().t() + @ result.grad_gate_up[active][rows, :intermediate].float() + * args.scale).bfloat16()) + native_boundary_grad_refs[3].append( + (q13[active][rows, 1].float().t() + @ result.grad_gate_up[active][rows, intermediate:].float() + * args.scale).bfloat16()) + native_boundary_grad_refs[4].append( + (saved_h[active][rows].float().t() + @ result.t2[active][rows].float() + * args.scale).bfloat16()) + cursor += count + native_boundary_grad_refs[1] = torch.stack(native_boundary_grad_refs[1]) + native_boundary_grad_refs[3] = torch.stack(native_boundary_grad_refs[3]) + native_boundary_grad_refs[4] = torch.stack(native_boundary_grad_refs[4]) + native_boundary_adapter_accuracy = [ + _accuracy(actual, expected) + for actual, expected in zip( + result.grad_side_lora, native_boundary_grad_refs, strict=True) + ] + assert max(forward_diffs.values()) < 0.04, forward_diffs + assert max(grad_diffs) < 0.04, grad_diffs + assert grad_x_diff < 0.04, grad_x_diff + assert combined_grad_x_diff < 0.04, combined_grad_x_diff + assert route_grad_diff < 0.04, route_grad_diff + assert combined_route_grad_diff < 0.04, combined_route_grad_diff + assert backward_boundary_accuracy[ + "saved_gate_up_after_backward_reuse"]["relative_l2"] == 0.0 + assert backward_boundary_accuracy[ + "grad_gate_up_from_native_boundaries"]["cosine_similarity"] > 0.99999 + assert min( + metric["cosine_similarity"] + for metric in native_boundary_adapter_accuracy + ) > 0.9999, native_boundary_adapter_accuracy + if args.scale == 0.0: + base_y = torch.empty_like(y) + buffer.x[:tokens].copy_(x_fp8[0]) + buffer.x_sf[:tokens].copy_(x_fp8[1]) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + deep_gemm.fp8_fp4_mega_moe( + base_y, transformed_w13, transformed_w2, buffer, + activation_clamp=args.activation_limit, + route_weight_mode=args.route_weight_mode, + num_config_tokens=tokens, fast_math=False) + torch.cuda.synchronize() + base_preservation = _accuracy(y, base_y) + assert base_preservation["relative_l2"] == 0.0, base_preservation + if rank == 0: + print(json.dumps({ + "mode": "mxfp4", "ranks": ranks, + "tokens_per_rank": tokens, "pool_rows": pool_rows, + "routing": args.routing, + "masked_ratio": args.masked_ratio, + "activation_limit": args.activation_limit, + "route_weight_mode": args.route_weight_mode, + "block_m": block_m, "forward_diffs": forward_diffs, + "adapter_grad_diffs": grad_diffs, + "grad_x_diff": grad_x_diff, + "combined_grad_x_diff": combined_grad_x_diff, + "route_grad_diff": route_grad_diff, + "combined_route_grad_diff": combined_route_grad_diff, + "zero_scale_base_preservation": base_preservation, + "adapter_grad_accuracy": grad_accuracy, + "grad_x_accuracy": grad_x_accuracy, + "backward_boundary_accuracy": backward_boundary_accuracy, + "native_boundary_adapter_accuracy": native_boundary_adapter_accuracy, + }, indent=2), flush=True) + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-processes", type=int, default=1) + parser.add_argument("--tokens", type=int, default=128) + parser.add_argument("--hidden", type=int, default=2048) + parser.add_argument("--intermediate", type=int, default=1024) + parser.add_argument("--experts", type=int, default=8) + parser.add_argument("--topk", type=int, default=2) + parser.add_argument("--scale", type=float, default=0.25) + parser.add_argument( + "--routing", + choices=("balanced", "skewed", "empty_experts", "remote"), + default="balanced") + parser.add_argument("--masked-ratio", type=float, default=0.0) + parser.add_argument( + "--activation", choices=("swiglu", "geglu"), default="swiglu") + parser.add_argument("--activation-limit", type=float, default=float("inf")) + parser.add_argument( + "--route-weight-mode", choices=("pre_down", "post_down"), + default="pre_down") + parser.add_argument("--mode", choices=("bf16", "mxfp4"), default="bf16") + args = parser.parse_args() + if args.experts % args.num_processes: + parser.error("experts must be divisible by num-processes") + if not 1 <= args.topk <= args.experts: + parser.error("topk must be in [1, experts]") + if not 0.0 <= args.masked_ratio < 1.0: + parser.error("masked-ratio must be in [0, 1)") + if args.activation_limit < 0: + parser.error("activation-limit must be non-negative") + torch.multiprocessing.spawn( + run_bf16_correctness if args.mode == "bf16" else run_mxfp4_correctness, + args=(args.num_processes, args), + nprocs=args.num_processes) From 64b1dff56b82d6c5f2de09cdf1a0b564ad8a49da Mon Sep 17 00:00:00 2001 From: morgendave Date: Thu, 13 Aug 2026 21:15:13 +0000 Subject: [PATCH 02/15] fix(mega): zero sparse side-LoRA reduction tails --- ...sm100_bf16_mega_moe_side_lora_backward.hpp | 35 +++++- ...sm100_bf16_mega_moe_side_lora_backward.cuh | 111 ++++++++++++++---- tests/run_mega_moe_side_lora_edge_matrix.py | 47 +++++++- tests/test_mega_moe_native_side_lora.py | 59 +++++++++- 4 files changed, 218 insertions(+), 34 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp index 03b90c49db..30e5dc66a1 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp @@ -487,14 +487,20 @@ class SM100BF16MegaMoESideLoraClearPaddingRuntime final : public LaunchRuntime { public: struct Args { + int hidden; int intermediate_hidden; int num_experts; int block_m; int num_sms; const int* expert_counts; + uint32_t num_pool_rows; cutlass::bfloat16_t* saved_h; cutlass::bfloat16_t* q13; cutlass::bfloat16_t* q2; + cutlass::bfloat16_t* t13; + cutlass::bfloat16_t* t2; + cutlass::bfloat16_t* x_pool; + cutlass::bfloat16_t* grad_ye; LaunchArgs launch_args; }; @@ -504,9 +510,9 @@ class SM100BF16MegaMoESideLoraClearPaddingRuntime final using namespace deep_gemm; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast( - &sm100_bf16_mega_moe_side_lora_clear_padding_impl<{}, {}, {}, {}>); + &sm100_bf16_mega_moe_side_lora_clear_padding_impl<{}, {}, {}, {}, {}>); }} -)", args.intermediate_hidden, args.num_experts, args.block_m, +)", args.hidden, args.intermediate_hidden, args.num_experts, args.block_m, args.num_sms); } @@ -515,8 +521,9 @@ static void __instantiate_kernel() {{ const LaunchConfigHandle& config, Args args) { DG_CUDA_UNIFIED_CHECK(launch_kernel( - kernel, config, args.expert_counts, args.saved_h, - args.q13, args.q2)); + kernel, config, args.expert_counts, args.num_pool_rows, args.saved_h, + args.q13, args.q2, args.t13, args.t2, + args.x_pool, args.grad_ye)); } }; @@ -1231,17 +1238,27 @@ static void sm100_bf16_mega_moe_side_lora_backward( // Clear forward padding before the adapter wgrads. B2 is formed now so // the dead grad-ye plane can hold the second L1 dgrad expansion. const SM100BF16MegaMoESideLoraClearPaddingRuntime::Args clear_args{ + .hidden = hidden, .intermediate_hidden = intermediate_hidden, .num_experts = num_experts, .block_m = block_m, .num_sms = num_sms, .expert_counts = expert_counts.data_ptr(), + .num_pool_rows = static_cast(num_pool_rows), .saved_h = reinterpret_cast( side_lora_saved_h.data_ptr()), .q13 = reinterpret_cast( side_lora_q13.data_ptr()), .q2 = reinterpret_cast( side_lora_q2.data_ptr()), + .t13 = reinterpret_cast( + side_lora_t13.data_ptr()), + .t2 = reinterpret_cast( + side_lora_t2.data_ptr()), + .x_pool = reinterpret_cast( + x_pool_output.data_ptr()), + .grad_ye = reinterpret_cast( + grad_ye.data_ptr()), .launch_args = LaunchArgs(num_sms, 256, 0, 1), }; const auto clear_code = @@ -2025,17 +2042,27 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( get_major_type_ab(b3_nt)); const SM100BF16MegaMoESideLoraClearPaddingRuntime::Args clear_args{ + .hidden = hidden, .intermediate_hidden = intermediate_hidden, .num_experts = num_experts, .block_m = block_m, .num_sms = num_sms, .expert_counts = expert_counts.data_ptr(), + .num_pool_rows = static_cast(num_pool_rows), .saved_h = reinterpret_cast( side_lora_saved_h.data_ptr()), .q13 = reinterpret_cast( side_lora_q13.data_ptr()), .q2 = reinterpret_cast( side_lora_q2.data_ptr()), + .t13 = reinterpret_cast( + side_lora_t13.data_ptr()), + .t2 = reinterpret_cast( + side_lora_t2.data_ptr()), + .x_pool = reinterpret_cast( + x_pool_output.data_ptr()), + .grad_ye = reinterpret_cast( + grad_ye.data_ptr()), .launch_args = LaunchArgs(num_sms, 256, 0, 1), }; const auto clear_code = diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh index 7ad2536df4..7fdd1ea222 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh @@ -1514,52 +1514,117 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_scale_grads_impl( #endif } +template +CUTLASS_DEVICE void sm100_bf16_mega_moe_side_lora_clear_padding_element( + const uint32_t pool_row, + const uint32_t column, + cutlass::bfloat16_t* saved_h, + cutlass::bfloat16_t* q13, + cutlass::bfloat16_t* q2, + cutlass::bfloat16_t* t13, + cutlass::bfloat16_t* t2, + cutlass::bfloat16_t* x_pool, + cutlass::bfloat16_t* grad_ye) { + if (column < kIntermediateHidden) { + saved_h[static_cast(pool_row) * kIntermediateHidden + + column] = cutlass::bfloat16_t(0.0f); + return; + } + uint32_t remaining = column - kIntermediateHidden; + const uint64_t rank_row = static_cast(pool_row) * 128; + const uint64_t rank_pair_row = static_cast(pool_row) * 256; + if (remaining < 256) { + q13[rank_pair_row + remaining] = cutlass::bfloat16_t(0.0f); + return; + } + remaining -= 256; + if (remaining < 128) { + q2[rank_row + remaining] = cutlass::bfloat16_t(0.0f); + return; + } + remaining -= 128; + if (remaining < 256) { + t13[rank_pair_row + remaining] = cutlass::bfloat16_t(0.0f); + return; + } + remaining -= 256; + if (remaining < 128) { + t2[rank_row + remaining] = cutlass::bfloat16_t(0.0f); + return; + } + remaining -= 128; + const uint64_t hidden_row = static_cast(pool_row) * kHidden; + if (remaining < kHidden) { + x_pool[hidden_row + remaining] = cutlass::bfloat16_t(0.0f); + } else { + grad_ye[hidden_row + remaining - kHidden] = + cutlass::bfloat16_t(0.0f); + } +} + // Forward saves only logical expert rows. Clear the already-allocated padded // tails before K-grouped adapter wgrads so they contribute exact zeros. -template CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_clear_padding_impl( const int* expert_counts, + const uint32_t num_pool_rows, cutlass::bfloat16_t* saved_h, cutlass::bfloat16_t* q13, - cutlass::bfloat16_t* q2) { + cutlass::bfloat16_t* q2, + cutlass::bfloat16_t* t13, + cutlass::bfloat16_t* t2, + cutlass::bfloat16_t* x_pool, + cutlass::bfloat16_t* grad_ye) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) - constexpr uint32_t kRankStorage = 3 * 128; + // Every operand below participates in a GEMM whose K dimension spans the + // complete block-padded route pool. Clearing only one operand is not + // sufficient: IEEE 0 * NaN still produces NaN, so allocator contents in + // the other operand can poison a shared A1/A3/B2 reduction. + constexpr uint32_t kRankStorage = 6 * 128; + constexpr uint32_t kHiddenStorage = 2 * kHidden; + constexpr uint32_t kRowStorage = + kIntermediateHidden + kRankStorage + kHiddenStorage; uint32_t pool_offset = 0; for (uint32_t expert = 0; expert < kNumExperts; ++expert) { const uint32_t count = static_cast( __ldg(expert_counts + expert)); const uint32_t capacity = math::ceil_div(count, BLOCK_M) * BLOCK_M; const uint32_t padding = capacity - count; - const uint64_t elements = static_cast(padding) * - (kIntermediateHidden + kRankStorage); + const uint64_t elements = + static_cast(padding) * kRowStorage; for (uint64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; linear < elements; linear += static_cast(kNumSMs) * blockDim.x) { - const uint32_t padding_row = linear / - (kIntermediateHidden + kRankStorage); + const uint32_t padding_row = linear / kRowStorage; const uint32_t column = linear - - static_cast(padding_row) * - (kIntermediateHidden + kRankStorage); + static_cast(padding_row) * kRowStorage; const uint32_t pool_row = pool_offset + count + padding_row; - if (column < kIntermediateHidden) { - saved_h[static_cast(pool_row) * - kIntermediateHidden + column] = - cutlass::bfloat16_t(0.0f); - } else { - const uint32_t rank_column = column - kIntermediateHidden; - if (rank_column < 256) { - q13[static_cast(pool_row) * 256 + - rank_column] = cutlass::bfloat16_t(0.0f); - } else { - q2[static_cast(pool_row) * 128 + - rank_column - 256] = cutlass::bfloat16_t(0.0f); - } - } + sm100_bf16_mega_moe_side_lora_clear_padding_element< + kHidden, kIntermediateHidden>( + pool_row, column, saved_h, q13, q2, t13, t2, + x_pool, grad_ye); } pool_offset += capacity; } + // Distributed buffers are sized to the maximum route-pool length across + // ranks. A sparse rank can therefore have a whole unowned suffix beyond + // its final local expert; shared wgrads still reduce across that suffix. + const uint64_t suffix_elements = + static_cast(num_pool_rows - pool_offset) * kRowStorage; + for (uint64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < suffix_elements; + linear += static_cast(kNumSMs) * blockDim.x) { + const uint32_t suffix_row = linear / kRowStorage; + const uint32_t column = linear - + static_cast(suffix_row) * kRowStorage; + sm100_bf16_mega_moe_side_lora_clear_padding_element< + kHidden, kIntermediateHidden>( + pool_offset + suffix_row, column, saved_h, q13, q2, t13, + t2, x_pool, grad_ye); + } #endif } diff --git a/tests/run_mega_moe_side_lora_edge_matrix.py b/tests/run_mega_moe_side_lora_edge_matrix.py index 3b21308e30..d96a4190ab 100644 --- a/tests/run_mega_moe_side_lora_edge_matrix.py +++ b/tests/run_mega_moe_side_lora_edge_matrix.py @@ -23,6 +23,24 @@ def _case(name: str, *arguments: str) -> tuple[str, tuple[str, ...]]: return name, (*COMMON, *arguments) +def _production_case( + mode: str, tokens: int, +) -> tuple[str, tuple[str, ...]]: + return ( + f"{mode}_production_boundary_{tokens}_ep4", + ( + "--hidden", "4096", + "--intermediate", "2048", + "--experts", "256", + "--topk", "6", + "--mode", mode, + "--tokens", str(tokens), + "--routing", "remote", + "--masked-ratio", "0.2", + ), + ) + + EP1_CASES = ( _case( "bf16_single_route_single_token", @@ -71,11 +89,32 @@ def _case(name: str, *arguments: str) -> tuple[str, tuple[str, ...]]: ) +# Every point immediately below and above a production scheduler transition, +# plus a zero-active-rank case and a larger steady-state case. These exact +# DSV4 Flash widths reproduce padding and rank-uniform route-pool suffixes that +# the compact default cases cannot exercise. +PRODUCTION_BOUNDARY_TOKENS = ( + 1, 15, 16, 17, 90, 91, 176, 177, 346, 347, 688, 689, 1029, 1030, 2048, +) +PRODUCTION_EP4_CASES = tuple( + _production_case(mode, tokens) + for mode in ("bf16", "mxfp4") + for tokens in PRODUCTION_BOUNDARY_TOKENS +) + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument( "--ep-processes", type=int, default=2, help="number of GPUs for remote-route cases; use 1 to skip them") + parser.add_argument( + "--production-boundaries", action="store_true", + help=( + "also run the exact 4096x2048, 256-expert EP4 boundary sweep " + "on four GPUs" + ), + ) args = parser.parse_args() if args.ep_processes < 1: parser.error("ep-processes must be positive") @@ -83,9 +122,13 @@ def main() -> None: cases = list(EP1_CASES) if args.ep_processes > 1: cases.extend(EP_CASES) + if args.production_boundaries: + cases.extend(PRODUCTION_EP4_CASES) for name, arguments in cases: - processes = ( - args.ep_processes if name.endswith("_ep") else 1) + if name.endswith("_ep4"): + processes = 4 + else: + processes = args.ep_processes if name.endswith("_ep") else 1 command = ( sys.executable, str(WORKER), "--num-processes", str(processes), *arguments) diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index 1bcb0e5731..45e50ede0e 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -234,6 +234,11 @@ def _relative(actual: torch.Tensor, expected: torch.Tensor) -> float: return float(calc_diff(actual.detach().float(), expected.detach().float())) +def _max_abs_diff(actual: torch.Tensor, expected: torch.Tensor) -> float: + delta = actual.detach().float() - expected.detach().float() + return float(delta.abs().max()) if delta.numel() else 0.0 + + def _accuracy(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: actual_f = actual.detach().float() expected_f = expected.detach().float() @@ -488,21 +493,52 @@ def run_bf16_correctness(local_rank: int, world: int, args) -> None: _relative(actual, expected.grad) for actual, expected in zip(result.grad_side_lora, adapter_ref) ] + grad_accuracy = [ + _accuracy(actual, expected.grad) + for actual, expected in zip(result.grad_side_lora, adapter_ref) + ] grad_x_diff = _relative(result.grad_x_pool[active], x_ref.grad) + grad_x_accuracy = _accuracy(result.grad_x_pool[active], x_ref.grad) + # Validate the three shared native contractions on the kernel's own saved + # forward/backward boundaries. The separate autograd comparison above is + # still useful, but at tiny route counts its BF16 reduction order can + # dominate cosine even when the fused contraction itself is exact. + native_boundary_shared_grad_refs = [ + (saved_x[active].float().t() @ result.t13[active, 0].float() + * args.scale).bfloat16(), + (saved_x[active].float().t() @ result.t13[active, 1].float() + * args.scale).bfloat16(), + (q2[active].float().t() + @ (route_grad.float() * + (route_ref.detach().float().unsqueeze(1) + if args.route_weight_mode == "post_down" else 1.0)) + * args.scale).bfloat16(), + ] + native_boundary_shared_adapter_accuracy = [ + _accuracy(actual, expected) + for actual, expected in zip( + ( + result.grad_side_lora[0], + result.grad_side_lora[2], + result.grad_side_lora[5], + ), + native_boundary_shared_grad_refs, + strict=True, + ) + ] source_grad_x = _scatter_to_sources( x_ref.grad.to(torch.bfloat16), metadata, ranks, tokens, group)[rank] combined_grad_x_diff = ( _relative(result.grad_x, source_grad_x) if ranks > 1 else 0.0) route_grad_diff = _relative(result.grad_route[active], route_ref.grad) - route_grad_max_abs = float( - (result.grad_route[active] - route_ref.grad).abs().max()) + route_grad_max_abs = _max_abs_diff( + result.grad_route[active], route_ref.grad) source_route_grad = _scatter_route_grads( route_ref.grad, metadata, ranks, tokens, topk, group)[rank] combined_route_grad_diff = _relative( buffer.backward_grad_route[:tokens, :topk], source_route_grad) - combined_route_grad_max_abs = float(( - buffer.backward_grad_route[:tokens, :topk] - source_route_grad - ).abs().max()) + combined_route_grad_max_abs = _max_abs_diff( + buffer.backward_grad_route[:tokens, :topk], source_route_grad) route_grad_close = ( route_grad_diff < 0.03 or route_grad_max_abs < 0.01) combined_route_grad_close = ( @@ -540,6 +576,10 @@ def run_bf16_correctness(local_rank: int, world: int, args) -> None: assert output_diff < forward_tolerance, output_diff assert q_diff < forward_tolerance, q_diff assert max(grad_diffs) < 0.03, grad_diffs + assert min( + metric["cosine_similarity"] + for metric in native_boundary_shared_adapter_accuracy + ) > 0.9999, native_boundary_shared_adapter_accuracy assert grad_x_diff < 0.03, grad_x_diff assert combined_grad_x_diff < 0.03, combined_grad_x_diff assert route_grad_close, (route_grad_diff, route_grad_max_abs) @@ -571,7 +611,12 @@ def run_bf16_correctness(local_rank: int, world: int, args) -> None: "forward_diff": forward_diff, "activation_diff": h_diff, "down_diff": down_diff, "output_diff": output_diff, "side_q_diff": q_diff, "adapter_grad_diffs": grad_diffs, + "adapter_grad_accuracy": grad_accuracy, + "native_boundary_shared_adapter_accuracy": ( + native_boundary_shared_adapter_accuracy + ), "grad_x_diff": grad_x_diff, + "grad_x_accuracy": grad_x_accuracy, "combined_grad_x_diff": combined_grad_x_diff, "route_grad_diff": route_grad_diff, "route_grad_max_abs": route_grad_max_abs, @@ -886,7 +931,11 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: ] assert max(forward_diffs.values()) < 0.04, forward_diffs assert max(grad_diffs) < 0.04, grad_diffs + assert min( + metric["cosine_similarity"] for metric in grad_accuracy + ) > 0.9999, grad_accuracy assert grad_x_diff < 0.04, grad_x_diff + assert grad_x_accuracy["cosine_similarity"] > 0.9999, grad_x_accuracy assert combined_grad_x_diff < 0.04, combined_grad_x_diff assert route_grad_diff < 0.04, route_grad_diff assert combined_route_grad_diff < 0.04, combined_route_grad_diff From 49fb5b213ac567697c9ed0d889a86dd8ec17b5b7 Mon Sep 17 00:00:00 2001 From: morgendave Date: Fri, 14 Aug 2026 00:51:44 +0000 Subject: [PATCH 03/15] fix(mega): harden side-LoRA scratch contracts --- ...sm100_bf16_mega_moe_side_lora_backward.hpp | 17 +++------ deep_gemm/mega/__init__.py | 3 +- deep_gemm/mega/backward.py | 16 ++++++++ tests/run_mega_moe_side_lora_edge_matrix.py | 3 +- tests/test_mega_moe_native_side_lora.py | 37 +++++++++++++++++-- 5 files changed, 57 insertions(+), 19 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp index 30e5dc66a1..f620a55cb3 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp @@ -1523,18 +1523,11 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( if (grad_y_unweighted_output.has_value()) check_bf16_hidden_pool( *grad_y_unweighted_output); - if (down_unweighted_output.has_value()) { - DG_HOST_ASSERT( - down_unweighted_output->scalar_type() == - torch::kBFloat16); - DG_HOST_ASSERT(down_unweighted_output->is_contiguous()); - DG_HOST_ASSERT(down_unweighted_output->dim() == 2); - DG_HOST_ASSERT( - down_unweighted_output->size(1) == hidden); - DG_HOST_ASSERT( - down_unweighted_output->size(0) > 0 && - down_unweighted_output->size(0) <= num_pool_rows); - } + // The MXFP4 side path reuses this dead forward save as a full hidden-width + // dgrad scratch. It is therefore mandatory even for pre-down routing and + // must cover the entire rank-uniform route pool. + DG_HOST_ASSERT(down_unweighted_output.has_value()); + check_bf16_hidden_pool(*down_unweighted_output); if (grad_route_output.has_value()) { DG_HOST_ASSERT( grad_route_output->scalar_type() == diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 97ddda2599..2eef23dc35 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -363,8 +363,7 @@ def fp8_fp4_mega_moe_side_lora( (pool_rows, 128), dtype=torch.bfloat16, device=y.device), torch.zeros( - 4 * sym_buffer.num_ring_tokens // - _C.get_token_alignment_for_mega_moe(), + 4 * sym_buffer.num_ring_tokens // 8, dtype=torch.int32, device=y.device), ) if len(side_lora_scratch) != 3: diff --git a/deep_gemm/mega/backward.py b/deep_gemm/mega/backward.py index 86bbb7abb4..ec2d73a4dc 100644 --- a/deep_gemm/mega/backward.py +++ b/deep_gemm/mega/backward.py @@ -1175,6 +1175,22 @@ def fp8_fp4_mega_moe_side_lora_backward( write_grad_x_pool = True hidden = w2_weights[0].size(1) intermediate_hidden = w2_dequant_scratch.size(2) + expected_down_shape = (gate_up_output.size(0), hidden) + if saved_down_unweighted.dtype != torch.bfloat16: + raise TypeError("saved_down_unweighted must be BF16") + if saved_down_unweighted.device != gate_up_output.device: + raise ValueError( + "saved_down_unweighted must be on the same device as " + "gate_up_output" + ) + if not saved_down_unweighted.is_contiguous(): + raise ValueError("saved_down_unweighted must be contiguous") + if tuple(saved_down_unweighted.shape) != expected_down_shape: + raise ValueError( + "saved_down_unweighted must cover the full route pool with " + f"shape {expected_down_shape}; got " + f"{tuple(saved_down_unweighted.shape)}" + ) outputs = ( _allocate_side_lora_backward_outputs( gate_up_output, side_lora, hidden, intermediate_hidden, diff --git a/tests/run_mega_moe_side_lora_edge_matrix.py b/tests/run_mega_moe_side_lora_edge_matrix.py index d96a4190ab..cd2ff31444 100644 --- a/tests/run_mega_moe_side_lora_edge_matrix.py +++ b/tests/run_mega_moe_side_lora_edge_matrix.py @@ -64,7 +64,8 @@ def _production_case( "--scale", "0"), _case( "mxfp4_balanced_bm_plus_one", - "--mode", "mxfp4", "--tokens", "17", "--topk", "2"), + "--mode", "mxfp4", "--tokens", "17", "--topk", "2", + "--default-side-lora-scratch", "--check-short-saved-down"), _case( "mxfp4_zero_scale_base_preservation", "--mode", "mxfp4", "--tokens", "17", "--topk", "2", diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index 45e50ede0e..f20dd39b70 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -404,6 +404,8 @@ def run_bf16_correctness(local_rank: int, world: int, args) -> None: activation_clamp=args.activation_limit, route_weight_mode=args.route_weight_mode, fast_math=False) + if args.default_side_lora_scratch: + assert ready.numel() == 4 * buffer.num_ring_tokens // 8 torch.cuda.synchronize() assert mismatch.item() == 0 base_preservation = None @@ -682,14 +684,17 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: buffer.topk_idx[:tokens].copy_(topk_idx) buffer.topk_weights[:tokens].copy_(topk_weights) y = torch.empty_like(x_bf16) - deep_gemm.fp8_fp4_mega_moe_side_lora( + q13, q2, ready = deep_gemm.fp8_fp4_mega_moe_side_lora( y, transformed_w13, transformed_w2, buffer, side_lora_input=x_bf16, side_lora=side, saved_x=saved_x, saved_h_unweighted=saved_h, saved_l1_preact=saved_gate_up, saved_down_unweighted=saved_down, num_config_tokens=tokens, side_lora_scale=args.scale, - side_lora_scratch=(q13, q2, ready), + side_lora_scratch=( + None if args.default_side_lora_scratch + else (q13, q2, ready) + ), activation_clamp=args.activation_limit, route_weight_mode=args.route_weight_mode, fast_math=False) @@ -779,12 +784,34 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: _relative(q13[active, 1], torch.cat(q3_rows)), _relative(q2[active], torch.cat(q2_rows))), } + w13_dequant_scratch = torch.empty( + (local_experts, 2 * intermediate, hidden), + device="cuda", dtype=torch.bfloat16) + w2_dequant_scratch = torch.empty( + (local_experts, hidden, intermediate), + device="cuda", dtype=torch.bfloat16) + if args.check_short_saved_down: + try: + deep_gemm.fp8_fp4_mega_moe_side_lora_backward( + saved_gate_up, saved_h, saved_down[:-1], q13, q2, side, + buffer.l1_acts[:pool_rows], buffer.l1_acts_sf, + transformed_w13, backward_w13, backward_w2, + w13_dequant_scratch, w2_dequant_scratch, + counts, padded, grad_y, buffer, block_m, + activation_limit=args.activation_limit, + route_weight_mode=args.route_weight_mode, + side_lora_scale=args.scale, + direct_remote_grad_x=ranks > 1) + except ValueError as error: + assert "full route pool" in str(error), str(error) + else: + raise AssertionError( + "short saved_down_unweighted was not rejected") result = deep_gemm.fp8_fp4_mega_moe_side_lora_backward( saved_gate_up, saved_h, saved_down, q13, q2, side, buffer.l1_acts[:pool_rows], buffer.l1_acts_sf, transformed_w13, backward_w13, backward_w2, - torch.empty((local_experts, 2 * intermediate, hidden), device="cuda", dtype=torch.bfloat16), - torch.empty((local_experts, hidden, intermediate), device="cuda", dtype=torch.bfloat16), + w13_dequant_scratch, w2_dequant_scratch, counts, padded, grad_y, buffer, block_m, activation_limit=args.activation_limit, route_weight_mode=args.route_weight_mode, @@ -1006,6 +1033,8 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: "--route-weight-mode", choices=("pre_down", "post_down"), default="pre_down") parser.add_argument("--mode", choices=("bf16", "mxfp4"), default="bf16") + parser.add_argument("--default-side-lora-scratch", action="store_true") + parser.add_argument("--check-short-saved-down", action="store_true") args = parser.parse_args() if args.experts % args.num_processes: parser.error("experts must be divisible by num-processes") From 80ca366456e5f15ebf739d0ca0f0e435f67dc5a7 Mon Sep 17 00:00:00 2001 From: morgendave Date: Mon, 17 Aug 2026 23:53:36 +0000 Subject: [PATCH 04/15] fix(mega): honor MXFP4 backward activation contract --- ...sm100_bf16_mega_moe_side_lora_backward.hpp | 10 ++++- deep_gemm/mega/backward.py | 7 +++- tests/run_mega_moe_side_lora_edge_matrix.py | 5 +++ tests/test_mega_moe_native_side_lora.py | 41 ++++++++++++++----- 4 files changed, 50 insertions(+), 13 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp index f620a55cb3..db30018a63 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp @@ -1394,6 +1394,8 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( const torch::Tensor& expert_counts, const torch::Tensor& grid_sync_counter, const float& activation_limit, + const std::string& activation, + const bool& fast_math, const bool& compute_w13_dgrad, const bool& direct_remote_grad_x, const bool& write_grad_x_pool, @@ -1459,9 +1461,11 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( DG_HOST_ASSERT(device_runtime->get_arch_major() == 10); DG_HOST_ASSERT(num_ranks >= 1); + DG_HOST_ASSERT(activation == "swiglu" || activation == "geglu"); DG_HOST_ASSERT( route_weight_mode == "pre_down" || route_weight_mode == "post_down"); + DG_HOST_ASSERT(route_weight_mode == "pre_down"); DG_HOST_ASSERT(num_ranks == 1 || (backward_rank >= 0 && backward_rank < num_ranks)); DG_HOST_ASSERT(block_m % 16 == 0); @@ -1879,6 +1883,8 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( .num_stages = num_stages, .num_sms = num_sms, .num_ranks = num_ranks, + .activation = activation, + .fast_math = fast_math, .route_weight_mode = route_weight_mode, .expert_counts = expert_counts.data_ptr(), .backward_sym_buffer = backward_sym_buffer, @@ -2008,8 +2014,8 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( const auto code = SM100BF16MegaMoESideLoraBackwardWaveRuntime::generate(args); const auto runtime = compiler->build(fmt::format( - "sm100_fp8_fp4_mega_moe_backward_dgrad_swiglu_{}_r{}", - route_weight_mode, + "sm100_fp8_fp4_mega_moe_side_lora_backward_{}_fast{}_{}_r{}", + activation, fast_math, route_weight_mode, grad_route_output.has_value()), code); SM100BF16MegaMoESideLoraBackwardWaveRuntime::launch(runtime, args); diff --git a/deep_gemm/mega/backward.py b/deep_gemm/mega/backward.py index ec2d73a4dc..b1205b6b63 100644 --- a/deep_gemm/mega/backward.py +++ b/deep_gemm/mega/backward.py @@ -1150,6 +1150,8 @@ def fp8_fp4_mega_moe_side_lora_backward( sym_buffer: Any, block_m: int, activation_limit: float = float("inf"), + activation: str = "swiglu", + fast_math: bool = True, route_weight_mode: RouteWeightMode = RouteWeightMode.PRE_DOWN, side_lora_scale: float = 1.0, direct_remote_grad_x: Optional[bool] = None, @@ -1160,6 +1162,8 @@ def fp8_fp4_mega_moe_side_lora_backward( expert_psum_rows: Optional[torch.Tensor] = None, ) -> MegaMoESideLoraBackwardResult: """Run the dedicated MXFP4 base-dgrad + BF16 side-LoRA backward.""" + if activation not in ("swiglu", "geglu"): + raise ValueError(f"unsupported activation: {activation}") route_weight_mode = RouteWeightMode(route_weight_mode) if route_weight_mode is RouteWeightMode.POST_DOWN: raise NotImplementedError( @@ -1232,7 +1236,8 @@ def fp8_fp4_mega_moe_side_lora_backward( l1_weights[0], l1_weights[1], grad_ye, route_weights, w2_weights[0], w2_weights[1], w2_dequant_scratch, w13_weights[0], w13_weights[1], w13_dequant_scratch, - expert_counts, grid_sync_counter, float(activation_limit), True, + expert_counts, grid_sync_counter, float(activation_limit), + activation, bool(fast_math), True, bool(direct_remote_grad_x), bool(write_grad_x_pool), True, block_m, sym_buffer.handle.buffer_ptrs, diff --git a/tests/run_mega_moe_side_lora_edge_matrix.py b/tests/run_mega_moe_side_lora_edge_matrix.py index cd2ff31444..ffbbe6525a 100644 --- a/tests/run_mega_moe_side_lora_edge_matrix.py +++ b/tests/run_mega_moe_side_lora_edge_matrix.py @@ -66,6 +66,11 @@ def _production_case( "mxfp4_balanced_bm_plus_one", "--mode", "mxfp4", "--tokens", "17", "--topk", "2", "--default-side-lora-scratch", "--check-short-saved-down"), + _case( + "mxfp4_geglu_bm_boundary", + "--mode", "mxfp4", "--tokens", "16", "--topk", "1", + "--routing", "skewed", "--activation", "geglu", + "--activation-limit", "1.5"), _case( "mxfp4_zero_scale_base_preservation", "--mode", "mxfp4", "--tokens", "17", "--topk", "2", diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index f20dd39b70..0a51ee3be8 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -315,6 +315,25 @@ def test_side_lora_backward_rejects_unsupported_post_down() -> None: f"{function.__name__} accepted unsupported post_down") +def test_mxfp4_side_lora_backward_rejects_unknown_activation() -> None: + try: + deep_gemm.fp8_fp4_mega_moe_side_lora_backward( + gate_up_output=None, saved_h=None, + saved_down_unweighted=None, q13=None, q2=None, + side_lora=None, l1_acts=None, l1_acts_sf=None, + l1_weights=None, w13_weights=None, w2_weights=None, + w13_dequant_scratch=None, w2_dequant_scratch=None, + expert_counts=None, padded_expert_counts=None, + grad_y=None, sym_buffer=None, block_m=16, + activation="relu", + ) + except ValueError as error: + assert "unsupported activation" in str(error) + else: + raise AssertionError( + "MXFP4 side-LoRA backward accepted an unknown activation") + + def test_side_lora_transform_validates_shared_layout() -> None: hidden, intermediate, experts, rank = 256, 128, 4, 128 side_lora = ( @@ -633,8 +652,6 @@ def run_bf16_correctness(local_rank: int, world: int, args) -> None: def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: rank, ranks, group = init_dist(local_rank, world) - if args.activation != "swiglu": - raise ValueError("MXFP4 side-LoRA backward currently supports SwiGLU") torch.manual_seed(4321 + rank) tokens, hidden, intermediate = args.tokens, args.hidden, args.intermediate experts, topk = args.experts, args.topk @@ -642,7 +659,7 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: block_m = _block_m(tokens, ranks, topk, experts) buffer = deep_gemm.get_symm_buffer_for_mega_moe( group, experts, tokens, topk, hidden, intermediate, - mma_type="fp8xfp4", activation="swiglu") + mma_type="fp8xfp4", activation=args.activation) x_bf16 = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) * 0.1 x_fp8 = per_token_cast_to_fp8( @@ -739,7 +756,7 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: (xe_base @ w13_deq[expert, intermediate:].t()).to(torch.bfloat16), (q3 @ b3).to(torch.bfloat16), alpha=args.scale) h_unweighted = _apply_activation( - gate, up, "swiglu", args.activation_limit) + gate, up, args.activation, args.activation_limit) route = route_ref[cursor:cursor + count] h_weighted = ( h_unweighted.float() * route.float().unsqueeze(1) @@ -799,6 +816,8 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: w13_dequant_scratch, w2_dequant_scratch, counts, padded, grad_y, buffer, block_m, activation_limit=args.activation_limit, + activation=args.activation, + fast_math=False, route_weight_mode=args.route_weight_mode, side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) @@ -814,6 +833,8 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: w13_dequant_scratch, w2_dequant_scratch, counts, padded, grad_y, buffer, block_m, activation_limit=args.activation_limit, + activation=args.activation, + fast_math=False, route_weight_mode=args.route_weight_mode, side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) torch.cuda.synchronize() @@ -868,9 +889,9 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: torch.cat([tensor.grad for tensor in q3_rows]), ), } - # Diagnose the exact fused dSwiGLU boundary independently of either base - # dgrad kernel. This reference starts from the native W2+LoRA grad-h and - # the native forward's saved gate/up values, so any discrepancy here is + # Diagnose the exact fused gated-activation boundary independently of + # either base dgrad kernel. This reference starts from the native W2+LoRA + # grad-h and the native forward's saved gate/up values, so discrepancies are # owned by the activation-backward phase itself. saved_gate, saved_up = _mx_deinterleave_pair( saved_gate_up_before_backward[active]) @@ -883,7 +904,7 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: native_gate = saved_gate.detach().clone().requires_grad_(True) native_up = saved_up.detach().clone().requires_grad_(True) native_h = _apply_activation( - native_gate, native_up, "swiglu", args.activation_limit) + native_gate, native_up, args.activation, args.activation_limit) (native_h.float() * grad_h_for_activation.float()).sum().backward() exact_grad_gate = native_gate.grad exact_grad_up = native_up.grad @@ -909,10 +930,10 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: saved_gate, torch.cat(gate_rows)) backward_boundary_accuracy["saved_vs_autograd_up_value"] = _accuracy( saved_up, torch.cat(up_rows)) - backward_boundary_accuracy["saved_vs_autograd_silu_value"] = _accuracy( + backward_boundary_accuracy["saved_vs_autograd_activation_value"] = _accuracy( native_h, _apply_activation( torch.cat(gate_rows), torch.cat(up_rows), - "swiglu", args.activation_limit)) + args.activation, args.activation_limit)) backward_boundary_accuracy["saved_gate_up_after_backward_reuse"] = _accuracy( saved_gate_up[active], saved_gate_up_before_backward[active]) # Validate the six native wgrad contractions on the kernel's own exact From 929739f2df882292bf804c084020256b89602b25 Mon Sep 17 00:00:00 2001 From: morgendave Date: Mon, 17 Aug 2026 23:57:45 +0000 Subject: [PATCH 05/15] fix(mega): preserve accurate MXFP4 backward default --- deep_gemm/mega/backward.py | 2 +- tests/test_mega_moe_native_side_lora.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/deep_gemm/mega/backward.py b/deep_gemm/mega/backward.py index b1205b6b63..21a8358101 100644 --- a/deep_gemm/mega/backward.py +++ b/deep_gemm/mega/backward.py @@ -1151,7 +1151,7 @@ def fp8_fp4_mega_moe_side_lora_backward( block_m: int, activation_limit: float = float("inf"), activation: str = "swiglu", - fast_math: bool = True, + fast_math: bool = False, route_weight_mode: RouteWeightMode = RouteWeightMode.PRE_DOWN, side_lora_scale: float = 1.0, direct_remote_grad_x: Optional[bool] = None, diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index 0a51ee3be8..c4917bc192 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -5,6 +5,7 @@ """ import argparse +import inspect import json import math @@ -315,7 +316,11 @@ def test_side_lora_backward_rejects_unsupported_post_down() -> None: f"{function.__name__} accepted unsupported post_down") -def test_mxfp4_side_lora_backward_rejects_unknown_activation() -> None: +def test_mxfp4_side_lora_backward_signature_and_unknown_activation() -> None: + signature = inspect.signature( + deep_gemm.fp8_fp4_mega_moe_side_lora_backward) + assert signature.parameters["activation"].default == "swiglu" + assert signature.parameters["fast_math"].default is False try: deep_gemm.fp8_fp4_mega_moe_side_lora_backward( gate_up_output=None, saved_h=None, From ccaa442e9f71ba7a83e9acafd21189034536ebad Mon Sep 17 00:00:00 2001 From: morgendave Date: Mon, 24 Aug 2026 19:05:30 +0000 Subject: [PATCH 06/15] test(mega): distinguish MXFP4 native accuracy boundaries --- tests/test_mega_moe_native_side_lora.py | 29 +++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index c4917bc192..9d867374d9 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -806,6 +806,16 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: _relative(q13[active, 1], torch.cat(q3_rows)), _relative(q2[active], torch.cat(q2_rows))), } + forward_accuracy = { + "gate_up": _accuracy(saved_gate_up[active], gate_ref), + "h": _accuracy(saved_h[active], torch.cat(h_for_w2_rows)), + "down": _accuracy( + saved_down[active], torch.cat(down_unweighted_rows)), + "output": _accuracy(y, expected_y), + "q1": _accuracy(q13[active, 0], torch.cat(q1_rows)), + "q3": _accuracy(q13[active, 1], torch.cat(q3_rows)), + "q2": _accuracy(q2[active], torch.cat(q2_rows)), + } w13_dequant_scratch = torch.empty( (local_experts, 2 * intermediate, hidden), device="cuda", dtype=torch.bfloat16) @@ -983,10 +993,21 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: result.grad_side_lora, native_boundary_grad_refs, strict=True) ] assert max(forward_diffs.values()) < 0.04, forward_diffs + assert min( + metric["cosine_similarity"] for metric in forward_accuracy.values() + ) > 0.999, forward_accuracy assert max(grad_diffs) < 0.04, grad_diffs + # The full PyTorch reference materializes the gated activation as BF16 + # before PRE_DOWN routing. The native MXFP4 forward intentionally preserves + # its serving contract: activation * up * route is rounded only once. Tiny + # route-count cases can therefore amplify a sub-0.1% forward difference in + # the down-adapter gradients. The strict check below remains on the six + # contractions evaluated at the exact native forward/backward boundaries. + full_reference_min_cosine = ( + 0.999 if args.activation == "geglu" else 0.9999) assert min( metric["cosine_similarity"] for metric in grad_accuracy - ) > 0.9999, grad_accuracy + ) > full_reference_min_cosine, grad_accuracy assert grad_x_diff < 0.04, grad_x_diff assert grad_x_accuracy["cosine_similarity"] > 0.9999, grad_x_accuracy assert combined_grad_x_diff < 0.04, combined_grad_x_diff @@ -994,8 +1015,11 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: assert combined_route_grad_diff < 0.04, combined_route_grad_diff assert backward_boundary_accuracy[ "saved_gate_up_after_backward_reuse"]["relative_l2"] == 0.0 + activation_boundary_min_cosine = ( + 0.99998 if args.activation == "geglu" else 0.99999) assert backward_boundary_accuracy[ - "grad_gate_up_from_native_boundaries"]["cosine_similarity"] > 0.99999 + "grad_gate_up_from_native_boundaries" + ]["cosine_similarity"] > activation_boundary_min_cosine assert min( metric["cosine_similarity"] for metric in native_boundary_adapter_accuracy @@ -1023,6 +1047,7 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: "activation_limit": args.activation_limit, "route_weight_mode": args.route_weight_mode, "block_m": block_m, "forward_diffs": forward_diffs, + "forward_accuracy": forward_accuracy, "adapter_grad_diffs": grad_diffs, "grad_x_diff": grad_x_diff, "combined_grad_x_diff": combined_grad_x_diff, From 1b922096789a41cdf18fd242b5b88576edd62e6d Mon Sep 17 00:00:00 2001 From: morgendave Date: Tue, 1 Sep 2026 00:15:45 +0000 Subject: [PATCH 07/15] test(mega): isolate side LoRA repeatability boundaries --- tests/test_mega_moe_native_side_lora.py | 143 ++++++++++++++++++++++-- 1 file changed, 132 insertions(+), 11 deletions(-) diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index 9d867374d9..1120a7b517 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -841,17 +841,20 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: else: raise AssertionError( "short saved_down_unweighted was not rejected") - result = deep_gemm.fp8_fp4_mega_moe_side_lora_backward( - saved_gate_up, saved_h, saved_down, q13, q2, side, - buffer.l1_acts[:pool_rows], buffer.l1_acts_sf, - transformed_w13, backward_w13, backward_w2, - w13_dequant_scratch, w2_dequant_scratch, - counts, padded, grad_y, buffer, block_m, - activation_limit=args.activation_limit, - activation=args.activation, - fast_math=False, - route_weight_mode=args.route_weight_mode, - side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) + def run_side_backward(): + return deep_gemm.fp8_fp4_mega_moe_side_lora_backward( + saved_gate_up, saved_h, saved_down, q13, q2, side, + buffer.l1_acts[:pool_rows], buffer.l1_acts_sf, + transformed_w13, backward_w13, backward_w2, + w13_dequant_scratch, w2_dequant_scratch, + counts, padded, grad_y, buffer, block_m, + activation_limit=args.activation_limit, + activation=args.activation, + fast_math=False, + route_weight_mode=args.route_weight_mode, + side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) + + result = run_side_backward() torch.cuda.synchronize() expected_grads = [tensor.grad for tensor in adapter_ref] @@ -1024,6 +1027,122 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: metric["cosine_similarity"] for metric in native_boundary_adapter_accuracy ) > 0.9999, native_boundary_adapter_accuracy + repeatability = None + if args.check_repeatability: + first_adapter_grads = tuple( + grad.detach().clone() for grad in result.grad_side_lora) + first_grad_x = result.grad_x.detach().clone() + first_grad_route = buffer.backward_grad_route[:tokens, :topk].clone() + first_metadata = buffer.token_src_metadata[:pool_rows].clone() + + # Repeat backward without changing any saved forward boundary. This + # isolates the new side-LoRA backward from MegaMoE's dispatch order. + fixed_boundary_result = run_side_backward() + fixed_boundary = { + "adapter_grads": [ + _accuracy(actual, expected) + for actual, expected in zip( + fixed_boundary_result.grad_side_lora, + first_adapter_grads, + strict=True, + ) + ], + "grad_x": _accuracy(fixed_boundary_result.grad_x, first_grad_x), + "grad_route": _accuracy( + buffer.backward_grad_route[:tokens, :topk], + first_grad_route, + ), + } + + # Repeat the entire native forward and backward with identical source + # tensors and routes. Any additional drift includes MegaMoE dispatch + # row assignment and is not owned by the backward contractions alone. + buffer.x[:tokens].copy_(x_fp8[0]) + buffer.x_sf[:tokens].copy_(x_fp8[1]) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + deep_gemm.fp8_fp4_mega_moe_side_lora( + y, transformed_w13, transformed_w2, buffer, + side_lora_input=x_bf16, side_lora=side, + saved_x=saved_x, saved_h_unweighted=saved_h, + saved_l1_preact=saved_gate_up, + saved_down_unweighted=saved_down, + num_config_tokens=tokens, side_lora_scale=args.scale, + side_lora_scratch=(q13, q2, ready), + activation_clamp=args.activation_limit, + route_weight_mode=args.route_weight_mode, + fast_math=False) + repeated_result = run_side_backward() + side_forward_backward_repeatability = { + "metadata_exact_fraction": float( + (buffer.token_src_metadata[:pool_rows] == first_metadata) + .float() + .mean() + .item() + ), + "adapter_grads": [ + _accuracy(actual, expected) + for actual, expected in zip( + repeated_result.grad_side_lora, + first_adapter_grads, + strict=True, + ) + ], + "grad_x": _accuracy(repeated_result.grad_x, first_grad_x), + "grad_route": _accuracy( + buffer.backward_grad_route[:tokens, :topk], + first_grad_route, + ), + } + + # Run the ordinary MegaMoE forward twice as a direct control for the + # dispatch-row ordering used by the separate side-LoRA forward. + base_y_first = torch.empty_like(y) + buffer.x[:tokens].copy_(x_fp8[0]) + buffer.x_sf[:tokens].copy_(x_fp8[1]) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + deep_gemm.fp8_fp4_mega_moe( + base_y_first, transformed_w13, transformed_w2, buffer, + activation_clamp=args.activation_limit, + route_weight_mode=args.route_weight_mode, + num_config_tokens=tokens, fast_math=False) + base_metadata_first = buffer.token_src_metadata[:pool_rows].clone() + base_y_second = torch.empty_like(y) + buffer.x[:tokens].copy_(x_fp8[0]) + buffer.x_sf[:tokens].copy_(x_fp8[1]) + buffer.topk_idx[:tokens].copy_(topk_idx) + buffer.topk_weights[:tokens].copy_(topk_weights) + deep_gemm.fp8_fp4_mega_moe( + base_y_second, transformed_w13, transformed_w2, buffer, + activation_clamp=args.activation_limit, + route_weight_mode=args.route_weight_mode, + num_config_tokens=tokens, fast_math=False) + torch.cuda.synchronize() + repeatability = { + "fixed_saved_boundary": fixed_boundary, + "forward_backward": side_forward_backward_repeatability, + "ordinary_megamoe_forward": { + "metadata_exact_fraction": float( + (buffer.token_src_metadata[:pool_rows] == base_metadata_first) + .float() + .mean() + .item() + ), + "output": _accuracy(base_y_second, base_y_first), + }, + } + assert all( + metric["relative_l2"] == 0.0 + for metric in fixed_boundary["adapter_grads"] + ), fixed_boundary + assert fixed_boundary["grad_x"]["relative_l2"] == 0.0, fixed_boundary + assert fixed_boundary["grad_route"]["relative_l2"] == 0.0, fixed_boundary + assert repeatability["ordinary_megamoe_forward"]["output"]["relative_l2"] == 0.0 + assert max( + metric["relative_l2"] + for metric in side_forward_backward_repeatability["adapter_grads"] + ) < 1e-4, side_forward_backward_repeatability if args.scale == 0.0: base_y = torch.empty_like(y) buffer.x[:tokens].copy_(x_fp8[0]) @@ -1058,6 +1177,7 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: "grad_x_accuracy": grad_x_accuracy, "backward_boundary_accuracy": backward_boundary_accuracy, "native_boundary_adapter_accuracy": native_boundary_adapter_accuracy, + "self_repeatability": repeatability, }, indent=2), flush=True) buffer.destroy() dist.destroy_process_group() @@ -1086,6 +1206,7 @@ def run_mxfp4_correctness(local_rank: int, world: int, args) -> None: parser.add_argument("--mode", choices=("bf16", "mxfp4"), default="bf16") parser.add_argument("--default-side-lora-scratch", action="store_true") parser.add_argument("--check-short-saved-down", action="store_true") + parser.add_argument("--check-repeatability", action="store_true") args = parser.parse_args() if args.experts % args.num_processes: parser.error("experts must be divisible by num-processes") From 4461ad8763d761e2a9eac5fabb2a4663faafa170 Mon Sep 17 00:00:00 2001 From: morgendave Date: Tue, 1 Sep 2026 20:41:41 +0000 Subject: [PATCH 08/15] perf: reuse saved side-LoRA storage for grad-x --- deep_gemm/mega/backward.py | 30 ++++++++++++++++++++----- tests/test_mega_moe_native_side_lora.py | 28 +++++++++++++++++++++++ 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/deep_gemm/mega/backward.py b/deep_gemm/mega/backward.py index 21a8358101..cb6af1b521 100644 --- a/deep_gemm/mega/backward.py +++ b/deep_gemm/mega/backward.py @@ -970,6 +970,7 @@ def _allocate_side_lora_backward_outputs( hidden: int, intermediate_hidden: int, write_grad_x_pool: bool = True, + reuse_gate_for_grad_x: bool = False, ) -> tuple[torch.Tensor, ...]: pool_rows = gate_up.size(0) options = dict(dtype=torch.bfloat16, device=gate_up.device) @@ -979,11 +980,18 @@ def _allocate_side_lora_backward_outputs( h_act = torch.empty_like(grad_h) h_weighted = torch.empty_like(grad_h) x_pool = torch.empty_like(grad_ye) - grad_x_pool = ( - torch.empty_like(grad_ye) - if write_grad_x_pool - else torch.empty((0, hidden), **options) - ) + if write_grad_x_pool and reuse_gate_for_grad_x: + if gate_up.numel() != pool_rows * hidden: + raise ValueError( + "gate_up storage cannot cover the grad-x pool" + ) + grad_x_pool = gate_up.view(pool_rows, hidden) + else: + grad_x_pool = ( + torch.empty_like(grad_ye) + if write_grad_x_pool + else torch.empty((0, hidden), **options) + ) route_weights = torch.empty( pool_rows, dtype=torch.float32, device=gate_up.device) grad_route = torch.empty_like(route_weights) @@ -1160,6 +1168,7 @@ def fp8_fp4_mega_moe_side_lora_backward( out: Optional[MegaMoESideLoraBackwardResult] = None, grid_sync_counter: Optional[torch.Tensor] = None, expert_psum_rows: Optional[torch.Tensor] = None, + reuse_gate_for_grad_x: bool = False, ) -> MegaMoESideLoraBackwardResult: """Run the dedicated MXFP4 base-dgrad + BF16 side-LoRA backward.""" if activation not in ("swiglu", "geglu"): @@ -1198,7 +1207,8 @@ def fp8_fp4_mega_moe_side_lora_backward( outputs = ( _allocate_side_lora_backward_outputs( gate_up_output, side_lora, hidden, intermediate_hidden, - write_grad_x_pool) + write_grad_x_pool, + reuse_gate_for_grad_x=reuse_gate_for_grad_x) if out is None else ( out.grad_ye, out.grad_h, out.grad_gate_up, out.h_act, out.h_weighted, out.x_pool, out.grad_x_pool, @@ -1208,6 +1218,14 @@ def fp8_fp4_mega_moe_side_lora_backward( (grad_ye, grad_h, grad_gate_up, h_act, h_weighted, x_pool, grad_x_pool, route_weights, grad_route, t13, t2, grad_side_lora) = outputs + if reuse_gate_for_grad_x: + if ( + grad_x_pool.data_ptr() != gate_up_output.data_ptr() + or grad_x_pool.numel() != gate_up_output.numel() + ): + raise ValueError( + "reuse_gate_grad_x requires grad_x_pool to reuse gate_up" + ) _direct_grad_x_planes(sym_buffer).zero_() sym_buffer.backward_grad_y[:grad_y.size(0)].copy_( grad_y.to(torch.bfloat16).contiguous()) diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index 1120a7b517..e29948d188 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -13,6 +13,7 @@ import torch.distributed as dist import deep_gemm +from deep_gemm.mega import backward as mega_backward from deep_gemm.testing import calc_diff from deep_gemm.utils import ( cast_back_from_fp4, per_token_cast_to_fp4, per_token_cast_to_fp8, @@ -321,6 +322,7 @@ def test_mxfp4_side_lora_backward_signature_and_unknown_activation() -> None: deep_gemm.fp8_fp4_mega_moe_side_lora_backward) assert signature.parameters["activation"].default == "swiglu" assert signature.parameters["fast_math"].default is False + assert signature.parameters["reuse_gate_for_grad_x"].default is False try: deep_gemm.fp8_fp4_mega_moe_side_lora_backward( gate_up_output=None, saved_h=None, @@ -339,6 +341,32 @@ def test_mxfp4_side_lora_backward_signature_and_unknown_activation() -> None: "MXFP4 side-LoRA backward accepted an unknown activation") +def test_mxfp4_side_lora_backward_can_reuse_saved_gate_for_grad_x() -> None: + rows, hidden, intermediate, experts, rank = 5, 16, 8, 2, 128 + gate_up = torch.empty(rows, 2 * intermediate, dtype=torch.bfloat16) + transformed_side_lora = ( + torch.empty(rank, hidden, dtype=torch.bfloat16), + torch.empty(experts, intermediate, rank, dtype=torch.bfloat16), + torch.empty(rank, hidden, dtype=torch.bfloat16), + torch.empty(experts, intermediate, rank, dtype=torch.bfloat16), + torch.empty(experts, rank, intermediate, dtype=torch.bfloat16), + torch.empty(hidden, rank, dtype=torch.bfloat16), + ) + + outputs = mega_backward._allocate_side_lora_backward_outputs( + gate_up, + transformed_side_lora, + hidden, + intermediate, + reuse_gate_for_grad_x=True, + ) + + grad_x_pool = outputs[6] + assert grad_x_pool.data_ptr() == gate_up.data_ptr() + assert grad_x_pool.numel() == gate_up.numel() + assert outputs[2].data_ptr() != gate_up.data_ptr() + + def test_side_lora_transform_validates_shared_layout() -> None: hidden, intermediate, experts, rank = 256, 128, 4, 128 side_lora = ( From f8809a1dacacb4fb27b8a48dcca3a6e4cea57ce7 Mon Sep 17 00:00:00 2001 From: morgendave Date: Tue, 1 Sep 2026 21:41:19 +0000 Subject: [PATCH 09/15] perf(mega): bound MXFP4 scale ring storage --- csrc/apis/mega.hpp | 17 +++--- csrc/jit_kernels/heuristics/mega_moe.hpp | 68 ++++++++++++++++++++++++ deep_gemm/mega/__init__.py | 6 +++ tests/test_mega_moe.py | 62 +++++++++++++++++++++ tests/test_mega_moe_native_side_lora.py | 3 ++ 5 files changed, 148 insertions(+), 8 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index e42fd76493..c73a5bf28e 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -90,14 +90,13 @@ get_symm_buffer_size_for_mega_moe_v2( input_topk_weights_layout, 1, num_max_tokens_per_rank, input_topk_idx_buffer.get_end_ptr()); - // Padded SF pool tokens - int num_sf_ring_tokens = 0; - for (int block_m: layout::kCandidateBlockM) { - num_sf_ring_tokens = std::max( - num_sf_ring_tokens, - layout::get_num_sf_ring_tokens(num_ring_tokens, block_m) - ); - } + // Padded SF pool tokens. Small BLOCK_M configurations are reachable only + // for small live batches, so they cannot consume every block in a large + // token ring. Size across the live-token regimes selected by the kernel. + const int num_sf_ring_tokens = with_sf ? + get_num_max_required_sf_ring_tokens_for_mega_moe( + num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, num_ring_tokens) : 0; // L1 input buffer const auto l1_token_buffer = layout::Buffer( @@ -865,6 +864,8 @@ static void register_apis(pybind11::module_& m) { #if DG_TENSORMAP_COMPATIBLE m.def("get_token_alignment_for_mega_moe", &get_token_alignment_for_mega_moe); m.def("get_ring_limit_for_mega_moe", &get_ring_limit_for_mega_moe); + m.def("get_num_max_required_sf_ring_tokens_for_mega_moe", + &get_num_max_required_sf_ring_tokens_for_mega_moe); m.def("get_symm_buffer_size_for_mega_moe", &get_symm_buffer_size_for_mega_moe); m.def("get_symm_buffer_size_for_mega_moe_v2", &get_symm_buffer_size_for_mega_moe_v2); diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index c35aaa0bbf..4f7ed0d6d4 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -193,6 +193,67 @@ static int get_num_experts_per_wave_for_mega_moe( return best_num_experts_per_wave; } +// Scale factors use an SF_BLOCK_M-padded row for every reachable pool block. +// A small BLOCK_M expands each scale row substantially, but those block sizes +// are only selected for small live batches and therefore cannot reach every +// block in a large token ring. Bound the SF storage by the live pool span +// instead of charging every block configuration for the full ring. +static int get_num_required_sf_ring_tokens_for_mega_moe( + const int& num_ranks, const int& num_experts_per_rank, + const int& num_tokens, const int& num_topk, + const int& num_ring_tokens, const int& block_m) { + const auto num_live_pool_tokens = layout::get_num_max_pool_tokens( + num_ranks, num_tokens, num_topk, num_experts_per_rank); + const auto num_reachable_ring_tokens = + std::min(num_ring_tokens, num_live_pool_tokens); + return layout::get_num_sf_ring_tokens( + num_reachable_ring_tokens, block_m); +} + +static int get_num_max_required_sf_ring_tokens_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_ring_tokens) { + DG_HOST_ASSERT(num_experts % num_ranks == 0); + const auto num_experts_per_rank = num_experts / num_ranks; + int num_sf_ring_tokens = 0; + + // BLOCK_M is monotonic in the live token count. Find the last live batch + // selecting each candidate with a binary search, then size that regime by + // its largest reachable pool span. This also excludes candidate values + // that the MXFP4 heuristic never selects (currently BLOCK_M=8). + for (const int candidate_block_m: layout::kCandidateBlockM) { + int lower = 0; + int upper = num_max_tokens_per_rank; + int last_matching_tokens = -1; + while (lower <= upper) { + const int num_tokens = lower + (upper - lower) / 2; + 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, MmaKind::MXFP8FP4); + if (block_m <= candidate_block_m) { + if (block_m == candidate_block_m) + last_matching_tokens = num_tokens; + lower = num_tokens + 1; + } else { + upper = num_tokens - 1; + } + } + if (last_matching_tokens < 0) + continue; + num_sf_ring_tokens = std::max( + num_sf_ring_tokens, + get_num_required_sf_ring_tokens_for_mega_moe( + num_ranks, num_experts_per_rank, + last_matching_tokens, num_topk, + num_ring_tokens, candidate_block_m)); + } + return num_sf_ring_tokens; +} + static std::pair get_pipeline_config_for_mega_moe( const int& smem_capacity, const int& num_experts, const int& hidden, @@ -269,6 +330,13 @@ static MegaMoEConfig get_mega_moe_config( // 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); + if (is_mma_with_sf(mma_kind)) { + DG_HOST_ASSERT( + num_sf_ring_tokens >= + get_num_required_sf_ring_tokens_for_mega_moe( + num_ranks, num_experts_per_rank, num_tokens, num_topk, + num_ring_tokens, block_m)); + } const int block_n = 128; const int load_block_m = block_m / 2; const int load_block_n = block_n; diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 2eef23dc35..e0c06237a7 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -45,6 +45,12 @@ def __init__(self, group: dist.ProcessGroup, self.hidden = hidden self.intermediate_hidden = intermediate_hidden self.num_ring_tokens = num_ring_tokens + self.num_sf_ring_tokens = ( + _C.get_num_max_required_sf_ring_tokens_for_mega_moe( + group.size(), num_experts, num_max_tokens_per_rank, + num_topk, num_ring_tokens) + if mma_type == 'fp8xfp4' else 0 + ) # Allocate a symmetric buffer num_bytes, slice_input_buffers = \ diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index d6b2628014..49ba5f55d0 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -20,6 +20,68 @@ from deep_gemm.testing import bench_kineto, calc_diff +def test_mxfp4_sf_ring_sizing_tracks_reachable_block_regimes(): + def align(value, alignment): + return (value + alignment - 1) // alignment * alignment + + def block_m(tokens, ranks, topk, experts): + expected = tokens * ranks * topk / experts + if expected <= 8.5: + return 16 + if expected <= 16.5: + return 32 + if expected <= 32.5: + return 64 + if expected <= 64.5: + return 96 + if expected <= 96.5: + return 128 + return 192 + + def brute_force(ranks, experts, max_tokens, topk, ring_tokens): + experts_per_rank = experts // ranks + required = 0 + for tokens in range(max_tokens + 1): + selected_block_m = block_m(tokens, ranks, topk, experts) + live_pool_tokens = align( + ranks * tokens * min(topk, experts_per_rank) + + experts_per_rank * (192 - 1), + 384, + ) + reachable_ring_tokens = min(ring_tokens, live_pool_tokens) + sf_block_m = align(selected_block_m, 128) + required = max( + required, + reachable_ring_tokens // selected_block_m * sf_block_m, + ) + return required + + shapes = ( + # DSV4-Flash EP4, 262K global tokens: one expert per wave. + (4, 256, 65_664, 6, 262_656), + # The old generic prefill ring at the same production shape. + (4, 256, 65_664, 6, 786_432), + # A small-batch shape exercises every compact BLOCK_M regime. + (8, 256, 1_152, 6, 9_216), + ) + for shape in shapes: + actual = ( + deep_gemm._C.get_num_max_required_sf_ring_tokens_for_mega_moe( + *shape + ) + ) + assert actual == brute_force(*shape) + legacy_worst_case = shape[-1] // 8 * 128 + assert actual < legacy_worst_case + + assert ( + deep_gemm._C.get_num_max_required_sf_ring_tokens_for_mega_moe( + *shapes[0] + ) + == 350_208 + ) + + def test_fp8_backward_canonicalizes_block_m_and_clears_padding( monkeypatch, ): diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index e29948d188..ce13ff525c 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -1062,6 +1062,7 @@ def run_side_backward(): first_grad_x = result.grad_x.detach().clone() first_grad_route = buffer.backward_grad_route[:tokens, :topk].clone() first_metadata = buffer.token_src_metadata[:pool_rows].clone() + first_side_y = y.detach().clone() # Repeat backward without changing any saved forward boundary. This # isolates the new side-LoRA backward from MegaMoE's dispatch order. @@ -1108,6 +1109,7 @@ def run_side_backward(): .mean() .item() ), + "output": _accuracy(y, first_side_y), "adapter_grads": [ _accuracy(actual, expected) for actual, expected in zip( @@ -1167,6 +1169,7 @@ def run_side_backward(): assert fixed_boundary["grad_x"]["relative_l2"] == 0.0, fixed_boundary assert fixed_boundary["grad_route"]["relative_l2"] == 0.0, fixed_boundary assert repeatability["ordinary_megamoe_forward"]["output"]["relative_l2"] == 0.0 + assert side_forward_backward_repeatability["output"]["relative_l2"] == 0.0 assert max( metric["relative_l2"] for metric in side_forward_backward_repeatability["adapter_grads"] From 97b060ba1d813906ba301c68f69485e72b831e52 Mon Sep 17 00:00:00 2001 From: morgendave Date: Tue, 1 Sep 2026 22:15:25 +0000 Subject: [PATCH 10/15] test(mega): cover expanded legacy buffer ABI --- tests/test_mega_moe.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 49ba5f55d0..6fd383d673 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -397,7 +397,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): buffer.num_ring_tokens)) legacy_buffer = buffer.buffer.narrow(0, 0, legacy_num_bytes) legacy_v2_slices = expanded_slicer(legacy_buffer) - assert len(legacy_v2_slices) == 11 + assert len(legacy_v2_slices) == 12 assert legacy_v2_slices[-1] is None legacy_forward_buffer = copy.copy(buffer) legacy_forward_buffer.buffer = legacy_buffer @@ -414,11 +414,13 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): legacy_forward_buffer.backward_grad_y, ) = legacy_slicer(legacy_buffer) legacy_forward_buffer.backward_grad_route = None + appended_bytes = buffer.backward_grad_route.nbytes + if buffer.side_lora_source is not None: + appended_bytes += buffer.side_lora_source.nbytes assert ( - legacy_num_bytes + - buffer.backward_grad_route.nbytes == + legacy_num_bytes + appended_bytes == expanded_num_bytes - ), 'v2 must append the route plane without shifting legacy storage' + ), 'v2 must append backward-only planes without shifting legacy storage' # Cast weights into FP4 def _cast_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: From f19fe0c436a49f93afb9ad93f8559bbee851d196 Mon Sep 17 00:00:00 2001 From: Zhiwei Zhang Date: Wed, 2 Sep 2026 08:23:37 +0800 Subject: [PATCH 11/15] perf(mega): fuse side-LoRA backward tail safely --- ...sm100_bf16_mega_moe_side_lora_backward.hpp | 53 ++---- ...sm100_bf16_mega_moe_side_lora_backward.cuh | 177 ++++++++++++------ tests/test_mega_moe_native_side_lora.py | 15 +- 3 files changed, 146 insertions(+), 99 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp index db30018a63..b0afb8831a 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp @@ -385,6 +385,9 @@ class SM100BF16MegaMoESideLoraGradXRuntime final bool direct_remote_grad_x; const int* expert_counts; cutlass::bfloat16_t* grad_x_pool; + const cutlass::bfloat16_t* side_grad_x_scratch1; + const cutlass::bfloat16_t* side_grad_x_scratch3; + float side_lora_scale; const layout::TokenSrcMetadata* token_src_metadata; cutlass::bfloat16_t* combine_buffer; layout::SymBuffer<> sym_buffer; @@ -415,6 +418,8 @@ static void __instantiate_kernel() {{ Args args) { DG_CUDA_UNIFIED_CHECK(launch_kernel( kernel, config, args.expert_counts, args.grad_x_pool, + args.side_grad_x_scratch1, args.side_grad_x_scratch3, + args.side_lora_scale, args.token_src_metadata, args.combine_buffer, args.sym_buffer, args.workspace, args.num_pool_rows, args.num_topk)); @@ -1284,25 +1289,6 @@ static void sm100_bf16_mega_moe_side_lora_backward( sm100_bf16_mega_moe_side_lora_shared_gemm( t3, a3_nt, side_grad_x_scratch3, num_pool_rows, hidden, side_lora_rank, "nk"); - const SM100BF16MegaMoESideLoraAxpy2Runtime::Args axpy2_args{ - .num_sms = num_sms, - .dst = reinterpret_cast( - grad_x_pool_output.data_ptr()), - .src1 = reinterpret_cast( - side_grad_x_scratch1.data_ptr()), - .src3 = reinterpret_cast( - side_grad_x_scratch3.data_ptr()), - .num_elements = static_cast(num_pool_rows) * hidden, - .scale = side_lora_scale, - .launch_args = LaunchArgs(num_sms, 256, 0, 1), - }; - const auto axpy2_code = - SM100BF16MegaMoESideLoraAxpy2Runtime::generate(axpy2_args); - const auto axpy2_runtime = compiler->build( - "sm100_bf16_mega_moe_side_lora_axpy2_grad_x", axpy2_code); - SM100BF16MegaMoESideLoraAxpy2Runtime::launch( - axpy2_runtime, axpy2_args); - const SM100BF16MegaMoESideLoraGradXRuntime::Args grad_x_args{ .hidden = hidden, .num_experts = num_experts, @@ -1314,6 +1300,11 @@ static void sm100_bf16_mega_moe_side_lora_backward( .expert_counts = expert_counts.data_ptr(), .grad_x_pool = reinterpret_cast( grad_x_pool_output.data_ptr()), + .side_grad_x_scratch1 = reinterpret_cast( + side_grad_x_scratch1.data_ptr()), + .side_grad_x_scratch3 = reinterpret_cast( + side_grad_x_scratch3.data_ptr()), + .side_lora_scale = side_lora_scale, .token_src_metadata = reinterpret_cast< const layout::TokenSrcMetadata*>( token_src_metadata.data_ptr()), @@ -2084,25 +2075,6 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( sm100_bf16_mega_moe_side_lora_shared_gemm( t3, a3_nt, side_grad_x_scratch3, num_pool_rows, hidden, side_lora_rank, "nk"); - const SM100BF16MegaMoESideLoraAxpy2Runtime::Args axpy2_args{ - .num_sms = num_sms, - .dst = reinterpret_cast( - grad_x_pool_output.data_ptr()), - .src1 = reinterpret_cast( - side_grad_x_scratch1.data_ptr()), - .src3 = reinterpret_cast( - side_grad_x_scratch3.data_ptr()), - .num_elements = static_cast(num_pool_rows) * hidden, - .scale = side_lora_scale, - .launch_args = LaunchArgs(num_sms, 256, 0, 1), - }; - const auto axpy2_code = - SM100BF16MegaMoESideLoraAxpy2Runtime::generate(axpy2_args); - const auto axpy2_runtime = compiler->build( - "sm100_fp8_fp4_mega_moe_side_lora_axpy2_grad_x", axpy2_code); - SM100BF16MegaMoESideLoraAxpy2Runtime::launch( - axpy2_runtime, axpy2_args); - const SM100BF16MegaMoESideLoraGradXRuntime::Args grad_x_args{ .hidden = hidden, .num_experts = num_experts, @@ -2114,6 +2086,11 @@ static void sm100_fp8_fp4_mega_moe_side_lora_backward( .expert_counts = expert_counts.data_ptr(), .grad_x_pool = reinterpret_cast( grad_x_pool_output.data_ptr()), + .side_grad_x_scratch1 = reinterpret_cast( + side_grad_x_scratch1.data_ptr()), + .side_grad_x_scratch3 = reinterpret_cast( + side_grad_x_scratch3.data_ptr()), + .side_lora_scale = side_lora_scale, .token_src_metadata = !backward_sym_buffer_ptrs.empty() ? reinterpret_cast( token_src_metadata->data_ptr()) diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh index 7fdd1ea222..b568beb287 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh @@ -1310,6 +1310,9 @@ template < CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_grad_x_impl( const int* expert_counts, cutlass::bfloat16_t* grad_x_pool, + const cutlass::bfloat16_t* side_grad_x_scratch1, + const cutlass::bfloat16_t* side_grad_x_scratch3, + const float side_lora_scale, const layout::TokenSrcMetadata* token_src_metadata, cutlass::bfloat16_t* combine_buffer, const __grid_constant__ layout::SymBuffer sym_buffer, @@ -1319,6 +1322,12 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_grad_x_impl( #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) || defined(__CLION_IDE__) using bf16 = cutlass::bfloat16_t; constexpr uint32_t kOutputTileN = 128; + constexpr uint32_t kVectorElems = sizeof(uint4) / sizeof(bf16); + union alignas(16) BF16x8 { + uint4 packed; + bf16 element[kVectorElems]; + }; + const float2 scale2{side_lora_scale, side_lora_scale}; if constexpr (kDirectRemoteGradX) { // The direct-write planes previously held grad-y. Valid routes are @@ -1326,7 +1335,6 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_grad_x_impl( // would otherwise leak that stale value into the final combine. // Clear every local fixed-slot plane, then synchronize before any // peer publishes its valid routes into this rank's symmetric plane. - constexpr uint32_t kVectorElems = sizeof(uint4) / sizeof(bf16); const uint64_t num_vectors = static_cast(workspace.num_max_tokens_per_rank) * num_topk * kHidden / kVectorElems; @@ -1369,7 +1377,6 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_grad_x_impl( local_tile += gridDim.x) { const uint32_t m_block = local_tile / num_n_blocks; const uint32_t n_block = local_tile - m_block * num_n_blocks; - constexpr uint32_t kVectorElems = sizeof(uint4) / sizeof(bf16); constexpr uint32_t kVectorsPerTile = kOutputTileN / kVectorElems; for (uint32_t linear = threadIdx.x; @@ -1387,9 +1394,36 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_grad_x_impl( continue; const uint32_t hidden_col = n_block * kOutputTileN + local_vector * kVectorElems; - const auto value = *reinterpret_cast( - grad_x_pool + - static_cast(pool_row) * kHidden + hidden_col); + const uint64_t element_offset = + static_cast(pool_row) * kHidden + hidden_col; + BF16x8 value, side1, side3; + value.packed = *reinterpret_cast( + grad_x_pool + element_offset); + side1.packed = *reinterpret_cast( + side_grad_x_scratch1 + element_offset); + side3.packed = *reinterpret_cast( + side_grad_x_scratch3 + element_offset); + auto* value2 = reinterpret_cast(&value.packed); + const auto* side12 = + reinterpret_cast(&side1.packed); + const auto* side32 = + reinterpret_cast(&side3.packed); + #pragma unroll + for (int i = 0; i < kVectorElems / 2; ++i) { + const auto scaled1 = __float22bfloat162_rn( + __fmul2_rn(scale2, __bfloat1622float2(side12[i]))); + const auto scaled3 = __float22bfloat162_rn( + __fmul2_rn(scale2, __bfloat1622float2(side32[i]))); + value2[i] = __float22bfloat162_rn(__fadd2_rn( + __fadd2_rn( + __bfloat1622float2(value2[i]), + __bfloat1622float2(scaled1)), + __bfloat1622float2(scaled3))); + } + if constexpr (kWriteGradXPool) { + *reinterpret_cast(grad_x_pool + element_offset) = + value.packed; + } if constexpr (kDirectRemoteGradX) { const auto metadata = token_src_metadata[pool_row]; auto* local_dst = combine_buffer + @@ -1399,7 +1433,8 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_grad_x_impl( kHidden + hidden_col); *reinterpret_cast( - sym_buffer.map(local_dst, metadata.rank_idx)) = value; + sym_buffer.map(local_dst, metadata.rank_idx)) = + value.packed; } } } @@ -1427,6 +1462,7 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_axpy2_impl( uint4 packed; bf16 element[8]; }; + const float2 scale2{scale, scale}; constexpr uint64_t kVectorElems = 8; const uint64_t num_vectors = num_elements / kVectorElems; for (uint64_t vector_idx = @@ -1437,13 +1473,20 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_axpy2_impl( d.packed = reinterpret_cast(dst)[vector_idx]; s1.packed = reinterpret_cast(src1)[vector_idx]; s3.packed = reinterpret_cast(src3)[vector_idx]; + auto* d2 = reinterpret_cast(&d.packed); + const auto* s12 = reinterpret_cast(&s1.packed); + const auto* s32 = reinterpret_cast(&s3.packed); #pragma unroll - for (int i = 0; i < 8; ++i) { - const bf16 side1(scale * static_cast(s1.element[i])); - const bf16 side3(scale * static_cast(s3.element[i])); - d.element[i] = bf16( - static_cast(d.element[i]) + - static_cast(side1) + static_cast(side3)); + for (int i = 0; i < 4; ++i) { + const auto side1 = __float22bfloat162_rn( + __fmul2_rn(scale2, __bfloat1622float2(s12[i]))); + const auto side3 = __float22bfloat162_rn( + __fmul2_rn(scale2, __bfloat1622float2(s32[i]))); + d2[i] = __float22bfloat162_rn(__fadd2_rn( + __fadd2_rn( + __bfloat1622float2(d2[i]), + __bfloat1622float2(side1)), + __bfloat1622float2(side3))); } reinterpret_cast(dst)[vector_idx] = d.packed; } @@ -1515,9 +1558,9 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_scale_grads_impl( } template -CUTLASS_DEVICE void sm100_bf16_mega_moe_side_lora_clear_padding_element( +CUTLASS_DEVICE void sm100_bf16_mega_moe_side_lora_clear_padding_vector( const uint32_t pool_row, - const uint32_t column, + const uint32_t vector_column, cutlass::bfloat16_t* saved_h, cutlass::bfloat16_t* q13, cutlass::bfloat16_t* q2, @@ -1525,40 +1568,51 @@ CUTLASS_DEVICE void sm100_bf16_mega_moe_side_lora_clear_padding_element( cutlass::bfloat16_t* t2, cutlass::bfloat16_t* x_pool, cutlass::bfloat16_t* grad_ye) { - if (column < kIntermediateHidden) { - saved_h[static_cast(pool_row) * kIntermediateHidden + - column] = cutlass::bfloat16_t(0.0f); + constexpr uint32_t kVectorElems = sizeof(uint4) / sizeof(cutlass::bfloat16_t); + constexpr uint32_t kIntermediateVectors = + kIntermediateHidden / kVectorElems; + constexpr uint32_t kRankVectors = 128 / kVectorElems; + constexpr uint32_t kRankPairVectors = 256 / kVectorElems; + constexpr uint32_t kHiddenVectors = kHidden / kVectorElems; + const uint4 zero{}; + if (vector_column < kIntermediateVectors) { + reinterpret_cast(saved_h)[ + static_cast(pool_row) * kIntermediateVectors + + vector_column] = zero; return; } - uint32_t remaining = column - kIntermediateHidden; - const uint64_t rank_row = static_cast(pool_row) * 128; - const uint64_t rank_pair_row = static_cast(pool_row) * 256; - if (remaining < 256) { - q13[rank_pair_row + remaining] = cutlass::bfloat16_t(0.0f); + uint32_t remaining = vector_column - kIntermediateVectors; + const uint64_t rank_row = + static_cast(pool_row) * kRankVectors; + const uint64_t rank_pair_row = + static_cast(pool_row) * kRankPairVectors; + if (remaining < kRankPairVectors) { + reinterpret_cast(q13)[rank_pair_row + remaining] = zero; return; } - remaining -= 256; - if (remaining < 128) { - q2[rank_row + remaining] = cutlass::bfloat16_t(0.0f); + remaining -= kRankPairVectors; + if (remaining < kRankVectors) { + reinterpret_cast(q2)[rank_row + remaining] = zero; return; } - remaining -= 128; - if (remaining < 256) { - t13[rank_pair_row + remaining] = cutlass::bfloat16_t(0.0f); + remaining -= kRankVectors; + if (remaining < kRankPairVectors) { + reinterpret_cast(t13)[rank_pair_row + remaining] = zero; return; } - remaining -= 256; - if (remaining < 128) { - t2[rank_row + remaining] = cutlass::bfloat16_t(0.0f); + remaining -= kRankPairVectors; + if (remaining < kRankVectors) { + reinterpret_cast(t2)[rank_row + remaining] = zero; return; } - remaining -= 128; - const uint64_t hidden_row = static_cast(pool_row) * kHidden; - if (remaining < kHidden) { - x_pool[hidden_row + remaining] = cutlass::bfloat16_t(0.0f); + remaining -= kRankVectors; + const uint64_t hidden_row = + static_cast(pool_row) * kHiddenVectors; + if (remaining < kHiddenVectors) { + reinterpret_cast(x_pool)[hidden_row + remaining] = zero; } else { - grad_ye[hidden_row + remaining - kHidden] = - cutlass::bfloat16_t(0.0f); + reinterpret_cast(grad_ye)[ + hidden_row + remaining - kHiddenVectors] = zero; } } @@ -1581,29 +1635,35 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_clear_padding_impl( // complete block-padded route pool. Clearing only one operand is not // sufficient: IEEE 0 * NaN still produces NaN, so allocator contents in // the other operand can poison a shared A1/A3/B2 reduction. - constexpr uint32_t kRankStorage = 6 * 128; - constexpr uint32_t kHiddenStorage = 2 * kHidden; - constexpr uint32_t kRowStorage = - kIntermediateHidden + kRankStorage + kHiddenStorage; + constexpr uint32_t kVectorElems = + sizeof(uint4) / sizeof(cutlass::bfloat16_t); + static_assert(kHidden % kVectorElems == 0); + static_assert(kIntermediateHidden % kVectorElems == 0); + constexpr uint32_t kRankStorageVectors = 6 * 128 / kVectorElems; + constexpr uint32_t kHiddenStorageVectors = + 2 * kHidden / kVectorElems; + constexpr uint32_t kRowStorageVectors = + kIntermediateHidden / kVectorElems + kRankStorageVectors + + kHiddenStorageVectors; uint32_t pool_offset = 0; for (uint32_t expert = 0; expert < kNumExperts; ++expert) { const uint32_t count = static_cast( __ldg(expert_counts + expert)); const uint32_t capacity = math::ceil_div(count, BLOCK_M) * BLOCK_M; const uint32_t padding = capacity - count; - const uint64_t elements = - static_cast(padding) * kRowStorage; + const uint64_t vectors = + static_cast(padding) * kRowStorageVectors; for (uint64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - linear < elements; + linear < vectors; linear += static_cast(kNumSMs) * blockDim.x) { - const uint32_t padding_row = linear / kRowStorage; - const uint32_t column = linear - - static_cast(padding_row) * kRowStorage; + const uint32_t padding_row = linear / kRowStorageVectors; + const uint32_t vector_column = linear - + static_cast(padding_row) * kRowStorageVectors; const uint32_t pool_row = pool_offset + count + padding_row; - sm100_bf16_mega_moe_side_lora_clear_padding_element< + sm100_bf16_mega_moe_side_lora_clear_padding_vector< kHidden, kIntermediateHidden>( - pool_row, column, saved_h, q13, q2, t13, t2, + pool_row, vector_column, saved_h, q13, q2, t13, t2, x_pool, grad_ye); } pool_offset += capacity; @@ -1611,19 +1671,20 @@ CUTLASS_GLOBAL void sm100_bf16_mega_moe_side_lora_clear_padding_impl( // Distributed buffers are sized to the maximum route-pool length across // ranks. A sparse rank can therefore have a whole unowned suffix beyond // its final local expert; shared wgrads still reduce across that suffix. - const uint64_t suffix_elements = - static_cast(num_pool_rows - pool_offset) * kRowStorage; + const uint64_t suffix_vectors = + static_cast(num_pool_rows - pool_offset) * + kRowStorageVectors; for (uint64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - linear < suffix_elements; + linear < suffix_vectors; linear += static_cast(kNumSMs) * blockDim.x) { - const uint32_t suffix_row = linear / kRowStorage; - const uint32_t column = linear - - static_cast(suffix_row) * kRowStorage; - sm100_bf16_mega_moe_side_lora_clear_padding_element< + const uint32_t suffix_row = linear / kRowStorageVectors; + const uint32_t vector_column = linear - + static_cast(suffix_row) * kRowStorageVectors; + sm100_bf16_mega_moe_side_lora_clear_padding_vector< kHidden, kIntermediateHidden>( - pool_offset + suffix_row, column, saved_h, q13, q2, t13, - t2, x_pool, grad_ye); + pool_offset + suffix_row, vector_column, saved_h, q13, q2, + t13, t2, x_pool, grad_ye); } #endif } diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index ce13ff525c..fe16c96678 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -1057,9 +1057,16 @@ def run_side_backward(): ) > 0.9999, native_boundary_adapter_accuracy repeatability = None if args.check_repeatability: + def effective_grad_x(backward_result): + return ( + backward_result.grad_x + if backward_result.grad_x is not None + else backward_result.grad_x_pool + ) + first_adapter_grads = tuple( grad.detach().clone() for grad in result.grad_side_lora) - first_grad_x = result.grad_x.detach().clone() + first_grad_x = effective_grad_x(result).detach().clone() first_grad_route = buffer.backward_grad_route[:tokens, :topk].clone() first_metadata = buffer.token_src_metadata[:pool_rows].clone() first_side_y = y.detach().clone() @@ -1076,7 +1083,8 @@ def run_side_backward(): strict=True, ) ], - "grad_x": _accuracy(fixed_boundary_result.grad_x, first_grad_x), + "grad_x": _accuracy( + effective_grad_x(fixed_boundary_result), first_grad_x), "grad_route": _accuracy( buffer.backward_grad_route[:tokens, :topk], first_grad_route, @@ -1118,7 +1126,8 @@ def run_side_backward(): strict=True, ) ], - "grad_x": _accuracy(repeated_result.grad_x, first_grad_x), + "grad_x": _accuracy( + effective_grad_x(repeated_result), first_grad_x), "grad_route": _accuracy( buffer.backward_grad_route[:tokens, :topk], first_grad_route, From b56ba1f57c89e1b6ac228bb9db4d409960e04484 Mon Sep 17 00:00:00 2001 From: Zhiwei Zhang Date: Wed, 2 Sep 2026 15:23:01 +0800 Subject: [PATCH 12/15] Reuse saved expert input in side LoRA backward --- deep_gemm/mega/backward.py | 33 ++++++++++-- tests/test_mega_moe_native_side_lora.py | 69 ++++++++++++++++++++++++- 2 files changed, 95 insertions(+), 7 deletions(-) diff --git a/deep_gemm/mega/backward.py b/deep_gemm/mega/backward.py index cb6af1b521..fc959413e9 100644 --- a/deep_gemm/mega/backward.py +++ b/deep_gemm/mega/backward.py @@ -971,6 +971,7 @@ def _allocate_side_lora_backward_outputs( intermediate_hidden: int, write_grad_x_pool: bool = True, reuse_gate_for_grad_x: bool = False, + saved_x_pool: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, ...]: pool_rows = gate_up.size(0) options = dict(dtype=torch.bfloat16, device=gate_up.device) @@ -979,7 +980,21 @@ def _allocate_side_lora_backward_outputs( grad_gate_up = torch.empty_like(gate_up) h_act = torch.empty_like(grad_h) h_weighted = torch.empty_like(grad_h) - x_pool = torch.empty_like(grad_ye) + if saved_x_pool is not None: + if saved_x_pool.dtype != torch.bfloat16: + raise TypeError("saved_x_pool must be BF16") + if saved_x_pool.device != gate_up.device: + raise ValueError("saved_x_pool must be on the same device as gate_up") + if not saved_x_pool.is_contiguous(): + raise ValueError("saved_x_pool must be contiguous") + if tuple(saved_x_pool.shape) != (pool_rows, hidden): + raise ValueError( + "saved_x_pool must have shape " + f"{(pool_rows, hidden)}; got {tuple(saved_x_pool.shape)}" + ) + x_pool = saved_x_pool + else: + x_pool = torch.empty_like(grad_ye) if write_grad_x_pool and reuse_gate_for_grad_x: if gate_up.numel() != pool_rows * hidden: raise ValueError( @@ -1044,6 +1059,7 @@ def bf16_mega_moe_side_lora_backward( out: Optional[MegaMoESideLoraBackwardResult] = None, grid_sync_counter: Optional[torch.Tensor] = None, expert_psum_rows: Optional[torch.Tensor] = None, + saved_x_pool: Optional[torch.Tensor] = None, ) -> MegaMoESideLoraBackwardResult: """Run the dedicated BF16 base-dgrad + rank-128 LoRA backward. @@ -1073,7 +1089,8 @@ def bf16_mega_moe_side_lora_backward( outputs = ( _allocate_side_lora_backward_outputs( gate_up_output, side_lora, w13_weights.size(2), - w2_weights.size(2), write_grad_x_pool) + w2_weights.size(2), write_grad_x_pool, + saved_x_pool=saved_x_pool) if out is None else ( out.grad_ye, out.grad_h, out.grad_gate_up, out.h_act, out.h_weighted, out.x_pool, out.grad_x_pool, @@ -1083,6 +1100,8 @@ def bf16_mega_moe_side_lora_backward( (grad_ye, grad_h, grad_gate_up, h_act, h_weighted, x_pool, grad_x_pool, route_weights, grad_route, t13, t2, grad_side_lora) = outputs + if saved_x_pool is not None and x_pool.data_ptr() != saved_x_pool.data_ptr(): + raise ValueError("out.x_pool must alias saved_x_pool when it is provided") _direct_grad_x_planes(sym_buffer).zero_() sym_buffer.backward_grad_y[:grad_y.size(0)].copy_( grad_y.to(torch.bfloat16).contiguous()) @@ -1108,7 +1127,7 @@ def bf16_mega_moe_side_lora_backward( sym_buffer.token_src_metadata, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), sym_buffer.num_topk, block_m, combine_order_mode.value, True, False, False, True, True, - False, False, 256) + False, saved_x_pool is not None, 256) _C.bf16_mega_moe_side_lora_backward( gate_up_output, grad_h, grad_gate_up, h_act, h_weighted, x_pool, grad_x_pool, grad_route, grad_ye, grad_ye, @@ -1169,6 +1188,7 @@ def fp8_fp4_mega_moe_side_lora_backward( grid_sync_counter: Optional[torch.Tensor] = None, expert_psum_rows: Optional[torch.Tensor] = None, reuse_gate_for_grad_x: bool = False, + saved_x_pool: Optional[torch.Tensor] = None, ) -> MegaMoESideLoraBackwardResult: """Run the dedicated MXFP4 base-dgrad + BF16 side-LoRA backward.""" if activation not in ("swiglu", "geglu"): @@ -1208,7 +1228,8 @@ def fp8_fp4_mega_moe_side_lora_backward( _allocate_side_lora_backward_outputs( gate_up_output, side_lora, hidden, intermediate_hidden, write_grad_x_pool, - reuse_gate_for_grad_x=reuse_gate_for_grad_x) + reuse_gate_for_grad_x=reuse_gate_for_grad_x, + saved_x_pool=saved_x_pool) if out is None else ( out.grad_ye, out.grad_h, out.grad_gate_up, out.h_act, out.h_weighted, out.x_pool, out.grad_x_pool, @@ -1218,6 +1239,8 @@ def fp8_fp4_mega_moe_side_lora_backward( (grad_ye, grad_h, grad_gate_up, h_act, h_weighted, x_pool, grad_x_pool, route_weights, grad_route, t13, t2, grad_side_lora) = outputs + if saved_x_pool is not None and x_pool.data_ptr() != saved_x_pool.data_ptr(): + raise ValueError("out.x_pool must alias saved_x_pool when it is provided") if reuse_gate_for_grad_x: if ( grad_x_pool.data_ptr() != gate_up_output.data_ptr() @@ -1247,7 +1270,7 @@ def fp8_fp4_mega_moe_side_lora_backward( sym_buffer.token_src_metadata, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), sym_buffer.num_topk, block_m, CombineOrderMode.FIXED_TOPK.value, True, False, - False, True, True, False, False, 256) + False, True, True, False, saved_x_pool is not None, 256) _C.fp8_fp4_mega_moe_side_lora_backward( gate_up_output, grad_h, grad_gate_up, h_act, h_weighted, x_pool, grad_x_pool, l1_acts, l1_acts_sf, diff --git a/tests/test_mega_moe_native_side_lora.py b/tests/test_mega_moe_native_side_lora.py index fe16c96678..9d2d26f179 100644 --- a/tests/test_mega_moe_native_side_lora.py +++ b/tests/test_mega_moe_native_side_lora.py @@ -323,6 +323,10 @@ def test_mxfp4_side_lora_backward_signature_and_unknown_activation() -> None: assert signature.parameters["activation"].default == "swiglu" assert signature.parameters["fast_math"].default is False assert signature.parameters["reuse_gate_for_grad_x"].default is False + assert signature.parameters["saved_x_pool"].default is None + bf16_signature = inspect.signature( + deep_gemm.bf16_mega_moe_side_lora_backward) + assert bf16_signature.parameters["saved_x_pool"].default is None try: deep_gemm.fp8_fp4_mega_moe_side_lora_backward( gate_up_output=None, saved_h=None, @@ -367,6 +371,64 @@ def test_mxfp4_side_lora_backward_can_reuse_saved_gate_for_grad_x() -> None: assert outputs[2].data_ptr() != gate_up.data_ptr() +def test_side_lora_backward_can_reuse_forward_saved_x_pool() -> None: + rows, hidden, intermediate, experts, rank = 5, 16, 8, 2, 128 + gate_up = torch.empty(rows, 2 * intermediate, dtype=torch.bfloat16) + saved_x = torch.randn(rows, hidden, dtype=torch.bfloat16) + transformed_side_lora = ( + torch.empty(rank, hidden, dtype=torch.bfloat16), + torch.empty(experts, intermediate, rank, dtype=torch.bfloat16), + torch.empty(rank, hidden, dtype=torch.bfloat16), + torch.empty(experts, intermediate, rank, dtype=torch.bfloat16), + torch.empty(experts, rank, intermediate, dtype=torch.bfloat16), + torch.empty(hidden, rank, dtype=torch.bfloat16), + ) + + outputs = mega_backward._allocate_side_lora_backward_outputs( + gate_up, + transformed_side_lora, + hidden, + intermediate, + saved_x_pool=saved_x, + ) + + assert outputs[5].data_ptr() == saved_x.data_ptr() + assert torch.equal(outputs[5], saved_x) + + +def test_side_lora_backward_rejects_invalid_saved_x_pool() -> None: + rows, hidden, intermediate, experts, rank = 5, 16, 8, 2, 128 + gate_up = torch.empty(rows, 2 * intermediate, dtype=torch.bfloat16) + transformed_side_lora = ( + torch.empty(rank, hidden, dtype=torch.bfloat16), + torch.empty(experts, intermediate, rank, dtype=torch.bfloat16), + torch.empty(rank, hidden, dtype=torch.bfloat16), + torch.empty(experts, intermediate, rank, dtype=torch.bfloat16), + torch.empty(experts, rank, intermediate, dtype=torch.bfloat16), + torch.empty(hidden, rank, dtype=torch.bfloat16), + ) + + for invalid, error in ( + (torch.empty(rows, hidden), TypeError), + (torch.empty(rows, hidden + 1, dtype=torch.bfloat16), ValueError), + (torch.empty(hidden, rows, dtype=torch.bfloat16).t(), ValueError), + ): + try: + mega_backward._allocate_side_lora_backward_outputs( + gate_up, + transformed_side_lora, + hidden, + intermediate, + saved_x_pool=invalid, + ) + except error: + pass + else: + raise AssertionError( + f"invalid saved_x_pool {tuple(invalid.shape)} was accepted" + ) + + def test_side_lora_transform_validates_shared_layout() -> None: hidden, intermediate, experts, rank = 256, 128, 4, 128 side_lora = ( @@ -541,7 +603,8 @@ def run_bf16_correctness(local_rank: int, world: int, args) -> None: block_m, activation_limit=args.activation_limit, activation=args.activation, fast_math=False, route_weight_mode=args.route_weight_mode, - side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) + side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1, + saved_x_pool=saved_x if args.reuse_saved_x_pool else None) torch.cuda.synchronize() grad_diffs = [ _relative(actual, expected.grad) @@ -880,7 +943,8 @@ def run_side_backward(): activation=args.activation, fast_math=False, route_weight_mode=args.route_weight_mode, - side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1) + side_lora_scale=args.scale, direct_remote_grad_x=ranks > 1, + saved_x_pool=saved_x if args.reuse_saved_x_pool else None) result = run_side_backward() torch.cuda.synchronize() @@ -1247,6 +1311,7 @@ def effective_grad_x(backward_result): parser.add_argument("--default-side-lora-scratch", action="store_true") parser.add_argument("--check-short-saved-down", action="store_true") parser.add_argument("--check-repeatability", action="store_true") + parser.add_argument("--reuse-saved-x-pool", action="store_true") args = parser.parse_args() if args.experts % args.num_processes: parser.error("experts must be divisible by num-processes") From ef6154a6e5170ede44a57f8dbf1616b1c5462f35 Mon Sep 17 00:00:00 2001 From: morgendave Date: Sun, 20 Sep 2026 20:03:01 +0000 Subject: [PATCH 13/15] Align MegaMoE training with upstream K-grouped TMA and block layouts --- .../impls/sm100_bf16_mega_moe_wgrad.hpp | 9 ++-- tests/test_mega_moe_training.py | 43 +++---------------- tests/test_mega_moe_wgrad_layout.py | 31 +++++++++++++ 3 files changed, 42 insertions(+), 41 deletions(-) create mode 100644 tests/test_mega_moe_wgrad_layout.py diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp index a46ead27d9..3d5addd6ef 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp @@ -155,9 +155,12 @@ static void sm100_bf16_mega_moe_wgrad_1sm( const auto tensor_map_b = make_tma_b_desc( cute::UMMA::Major::MN, b, n, pool_rows_b, kBlockN, kBlockK, static_cast(b.stride(0)), 1, kSwizzle); - const auto tensor_map_d = make_tma_cd_desc( - d, m, n, kBlockM, kStoreBlockN, - static_cast(d.stride(1)), num_groups, kSwizzleCD); + // Upstream K-grouped epilogues address [N, M, group] with a 3D TMA + // store. A flattened 2D descriptor compiles but traps on the first store. + const auto tensor_map_d = make_tma_3d_desc( + d, n, m, num_groups, kStoreBlockN, kBlockM, 1, + static_cast(d.stride(1)), static_cast(d.stride(0)), + kSwizzleCD); const SM100BF16GemmRuntime::Args args = { .gemm_desc = desc, diff --git a/tests/test_mega_moe_training.py b/tests/test_mega_moe_training.py index 48e8a660c2..61c984bffd 100644 --- a/tests/test_mega_moe_training.py +++ b/tests/test_mega_moe_training.py @@ -434,21 +434,9 @@ def create_inputs(): op=dist.ReduceOp.MAX, group=group) num_config_tokens = int(num_config_tokens_tensor.item()) - expected_tokens_per_expert = ( - num_config_tokens * num_ranks * num_topk / - num_experts) - if expected_tokens_per_expert <= 8.5: - pool_block_m = 16 - elif expected_tokens_per_expert <= 16.5: - pool_block_m = 32 - elif expected_tokens_per_expert <= 32.5: - pool_block_m = 64 - elif expected_tokens_per_expert <= 64.5: - pool_block_m = 96 - elif expected_tokens_per_expert <= 96.5: - pool_block_m = 128 - else: - pool_block_m = 192 + pool_block_m = deep_gemm._C.get_block_m_for_mega_moe( + num_ranks, num_experts, buffer.num_max_tokens_per_rank, + num_config_tokens, num_topk, args.mma_type) if args.expect_block_m: assert pool_block_m == args.expect_block_m, ( f'expected BLOCK_M={args.expect_block_m}, ' @@ -1550,28 +1538,7 @@ def check_native_forward_repeatability( def run_bf16_backward_test(): backward_base_allocated = torch.cuda.memory_allocated() torch.cuda.reset_peak_memory_stats() - config_num_tokens = torch.tensor( - num_tokens, dtype=torch.int32, device='cuda') - if num_ranks > 1: - dist.all_reduce( - config_num_tokens, - op=dist.ReduceOp.MAX, - group=group) - expected_tokens_per_expert = ( - int(config_num_tokens.item()) * - num_ranks * num_topk / num_experts) - if expected_tokens_per_expert <= 8.5: - block_m = 16 - elif expected_tokens_per_expert <= 16.5: - block_m = 32 - elif expected_tokens_per_expert <= 32.5: - block_m = 64 - elif expected_tokens_per_expert <= 64.5: - block_m = 96 - elif expected_tokens_per_expert <= 96.5: - block_m = 128 - else: - block_m = 192 + block_m = pool_block_m all_topk_idx = gather_rank_padded(topk_idx, -1) all_topk_weights = gather_rank_padded(topk_weights, 0) @@ -3194,7 +3161,7 @@ def run_baseline(): '--expect-block-m', type=int, default=0, - choices=[0, 16, 32, 64, 96, 128, 192], + choices=[0, 16, 32, 64, 96, 128, 192, 240], help='Assert the rank-uniform forward/backward BLOCK_M selection') parser.add_argument('--masked-ratio', type=float, default=0.0, help='Mask some expert selections') parser.add_argument('--routing', choices=['random', 'balanced', 'skew', 'extreme'], default='random') diff --git a/tests/test_mega_moe_wgrad_layout.py b/tests/test_mega_moe_wgrad_layout.py new file mode 100644 index 0000000000..97faf40df9 --- /dev/null +++ b/tests/test_mega_moe_wgrad_layout.py @@ -0,0 +1,31 @@ +"""MegaMoE wgrad must follow the upstream K-grouped 3D output ABI.""" + +import pytest +import torch + +import deep_gemm + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a Blackwell GPU") +@pytest.mark.parametrize("pool_block_m", [16, 32, 64, 96, 128, 192, 240]) +@pytest.mark.parametrize("output_n", [128, 256]) +def test_wgrad_group_dimension(pool_block_m, output_n): + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("requires SM100/SM103") + torch.manual_seed(123) + # Include an empty expert and unequal group extents to catch a descriptor + # that silently flattens (or aliases) the output group coordinate. + sizes = [pool_block_m, 0, 2 * pool_block_m, 3 * pool_block_m] + counts = torch.tensor(sizes, dtype=torch.int32, device="cuda") + a = torch.randn(sum(sizes), 512, dtype=torch.bfloat16, device="cuda") * 0.1 + b = torch.randn(sum(sizes), output_n, dtype=torch.bfloat16, device="cuda") * 0.1 + output = torch.full((4, 512, output_n), float("nan"), dtype=torch.bfloat16, device="cuda") + deep_gemm.bf16_mega_moe_backward_w13(output, a, b, counts, pool_block_m) + torch.cuda.synchronize() + offset = 0 + for expert, rows in enumerate(sizes): + reference = a[offset:offset + rows].float().T @ b[offset:offset + rows].float() + assert torch.isfinite(output[expert]).all() + torch.testing.assert_close(output[expert].float(), reference, rtol=0.004, atol=0.002) + offset += rows + assert torch.count_nonzero(output[1]) == 0 From f08dcb0b226d7bb1ef06f4f9a5661de1601833cf Mon Sep 17 00:00:00 2001 From: morgendave Date: Mon, 21 Sep 2026 02:01:50 +0000 Subject: [PATCH 14/15] perf(mega): use wide wgrad tiles with expert-tail MMA masking --- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 5 +- .../impls/sm100_bf16_mega_moe_wgrad.hpp | 19 ++++---- .../deep_gemm/impls/sm100_bf16_gemm.cuh | 18 ++++++- tests/compile_mega_moe_compat.py | 6 +++ tests/test_mega_moe_wgrad_layout.py | 48 +++++++++++++++++++ 5 files changed, 83 insertions(+), 13 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index 47544e74c6..2440b31761 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -52,6 +52,7 @@ 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; }; static void compile_and_launch(const std::string& tag, const Args& args) { @@ -76,7 +77,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, - {}, {}, {}, {} + {}, {}, {}, {}, {} >); }}; )", @@ -99,7 +100,7 @@ static void __instantiate_kernel() {{ 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( diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp index 3d5addd6ef..99c5f12bcc 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp @@ -57,23 +57,21 @@ static void sm100_bf16_mega_moe_wgrad_1sm( const int kBlockN = deep_jit::get_env( "DG_BF16_MEGA_MOE_WGRAD_BLOCK_N", n % 256 == 0 ? 256 : 128); - // The K-grouped scheduler addresses each expert in the shared physical - // pool. Its K tile must divide the forward pool alignment; otherwise the - // final tile of one expert reads rows from the next expert. In particular, - // BLOCK_M=96 previously contaminated Qwen top-8 wgrads while BLOCK_M=128 - // happened to pass. + // Preserve the physical expert offsets while amortizing TMA/barrier cost + // over K=64. For pools such as 96/240, mask whole 16-wide MMA atoms in the + // final tile rather than accumulating rows from the following expert. const int kBlockK = deep_jit::get_env( - "DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", - pool_block_m % 64 == 0 ? 64 : - pool_block_m % 32 == 0 ? 32 : 16); + "DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", 64); + DG_HOST_ASSERT(kBlockK == 16 or kBlockK == 32 or kBlockK == 64); + DG_HOST_ASSERT(pool_block_m % 16 == 0); + const bool mask_grouped_k_tail = pool_block_m % kBlockK != 0; const int kNumStages = deep_jit::get_env( "DG_BF16_MEGA_MOE_WGRAD_NUM_STAGES", 4); const int kSwizzle = kBlockK * static_cast(sizeof(cutlass::bfloat16_t)); const int kStoreBlockN = deep_jit::get_env( "DG_BF16_MEGA_MOE_WGRAD_STORE_BLOCK_N", 64); - // The output store width is independent of the A/B K tile. In particular, - // a 240-row pool requires K=16 but still uses the default 64-column store. + // The output store width is independent of the A/B K tile. const int kSwizzleCD = kStoreBlockN * sizeof(cutlass::bfloat16_t); DG_HOST_ASSERT(kSwizzleCD == 32 or kSwizzleCD == 64 or kSwizzleCD == 128); constexpr int kNumNonEpilogueThreads = 128; @@ -187,6 +185,7 @@ static void sm100_bf16_mega_moe_wgrad_1sm( .combine_order_mode = combine.order_mode, .combine_num_extra_threads = static_cast(num_extra_combine_threads), + .mask_grouped_k_tail = mask_grouped_k_tail, }; SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_mega_moe_wgrad_1sm", args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh index 3d505a2620..098771a0b0 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh @@ -37,7 +37,8 @@ template + uint32_t kNumExtraCombineThreads, + bool kMaskGroupedKTail = false> CUTLASS_GLOBAL void __launch_bounds__( kNumNonEpilogueThreads + kNumEpilogueThreads + kNumExtraCombineThreads, @@ -92,6 +93,10 @@ sm100_bf16_gemm_impl(int* grouped_layout, "Invalid block K"); DG_STATIC_ASSERT(BLOCK_K % UMMA_K == 0, "Block K must be divisible by UMMA K"); DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or kKAlignment % BLOCK_K == 0, "K alignment must be divisible by block K"); + DG_STATIC_ASSERT(not kMaskGroupedKTail or + (kGemmType == GemmType::KGroupedContiguous and + kMajorA == cute::UMMA::Major::MN and kMajorB == cute::UMMA::Major::MN), + "Masked K tails require physical count-based MN-major groups"); DG_STATIC_ASSERT(kNumMulticast == 1 or kNumMulticast == 2, "Only support 1/2 multicast"); DG_STATIC_ASSERT((kSwapAB and BLOCK_N == LAYOUT_AD_M) or (not kSwapAB and (BLOCK_M == 32 or BLOCK_M == 64 or BLOCK_M == LAYOUT_AD_M)), "Invalid block size"); @@ -321,6 +326,17 @@ sm100_bf16_gemm_impl(int* grouped_layout, if (cute::elect_one_sync()) { #pragma unroll for (uint32_t umma_k_idx = 0; umma_k_idx < BLOCK_K / UMMA_K; ++ umma_k_idx) { + // MegaMoE groups are padded to at least UMMA_K, not + // necessarily BLOCK_K (e.g. 240-row pools). TMA may + // read the next group's rows in the final tile, but + // those entire 16-wide MMA atoms must not contribute. + // Keep the existing empty-group initialization path; + // the epilogue explicitly writes zeros for that group. + if constexpr (kMaskGroupedKTail) { + if (scheduler.current_shape_k != 0 and + k_block_idx * BLOCK_K + umma_k_idx * UMMA_K >= scheduler.current_shape_k) + continue; + } const uint32_t atom_k_idx = umma_k_idx * UMMA_K / BLOCK_ATOM_K; const uint32_t inner_k_idx = umma_k_idx * UMMA_K % BLOCK_ATOM_K; a_desc.lo = mma::sm100::advance_umma_desc_lo( diff --git a/tests/compile_mega_moe_compat.py b/tests/compile_mega_moe_compat.py index c1a5fc566c..7cf6963f96 100644 --- a/tests/compile_mega_moe_compat.py +++ b/tests/compile_mega_moe_compat.py @@ -55,6 +55,12 @@ def cases(): "cutlass::bfloat16_t", "epilogue::transform::EpilogueIdentity", 100, 2, combine, "CombineOrderMode::DeepEPV1", 64 if combine else 0] yield f"wgrad_combine{int(combine)}", "sm100_bf16_gemm", "sm100_bf16_gemm_impl", args + masked_args = args.copy() + masked_args[7] = 64 + masked_args[9:11] = [128, 128] + masked_args[18] = 128 + masked_args.append(True) + yield f"wgrad_masked_tail_combine{int(combine)}", "sm100_bf16_gemm", "sm100_bf16_gemm_impl", masked_args for layout in (1, 2): yield (f"psum_layout{layout}", "smxx_layout", "transpose_and_pack_strided_fp32_into_ue8m0", diff --git a/tests/test_mega_moe_wgrad_layout.py b/tests/test_mega_moe_wgrad_layout.py index 97faf40df9..1900d3de83 100644 --- a/tests/test_mega_moe_wgrad_layout.py +++ b/tests/test_mega_moe_wgrad_layout.py @@ -29,3 +29,51 @@ def test_wgrad_group_dimension(pool_block_m, output_n): torch.testing.assert_close(output[expert].float(), reference, rtol=0.004, atol=0.002) offset += rows assert torch.count_nonzero(output[1]) == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a Blackwell GPU") +@pytest.mark.parametrize("pool_block_m", [16, 32, 64, 96, 128, 192, 240]) +@pytest.mark.parametrize("output_n", [128, 256]) +def test_wgrad_wide_k_masks_expert_tail(monkeypatch, pool_block_m, output_n): + """Wide loads may read the next expert; its MMA atoms must never execute.""" + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("requires SM100/SM103") + torch.manual_seed(941) + # Prefix offsets and tails cover all 16-row residues modulo 64. Empty + # first/middle/last groups exercise accumulator initialization and OOB TMA. + sizes = [0, pool_block_m, 2 * pool_block_m, 0, 3 * pool_block_m, pool_block_m, 0] + counts = torch.tensor(sizes, dtype=torch.int32, device="cuda") + a = torch.randn(sum(sizes), 512, dtype=torch.bfloat16, device="cuda") * 0.1 + b = torch.randn(sum(sizes), output_n, dtype=torch.bfloat16, device="cuda") * 0.1 + # Deliberately make adjacent experts very different, so a cross-expert + # contribution cannot hide under the approximate reference tolerance. + a[pool_block_m:3 * pool_block_m].mul_(32) + b[pool_block_m:3 * pool_block_m].mul_(32) + shape = (len(sizes), 512, output_n) + baseline = torch.full(shape, float("nan"), dtype=torch.bfloat16, device="cuda") + monkeypatch.setenv("DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", "16") + deep_gemm.bf16_mega_moe_backward_w13(baseline, a, b, counts, pool_block_m) + torch.cuda.synchronize() + for block_k in (32, 64): + monkeypatch.setenv("DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", str(block_k)) + output = torch.full_like(baseline, float("nan")) + for _ in range(2): + deep_gemm.bf16_mega_moe_backward_w13(output, a, b, counts, pool_block_m) + torch.cuda.synchronize() + assert torch.isfinite(output).all() + assert torch.equal(output, baseline) + for expert, rows in enumerate(sizes): + if rows == 0: + assert torch.count_nonzero(output[expert]) == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a Blackwell GPU") +def test_wgrad_rejects_invalid_k_tile(monkeypatch): + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("requires SM100/SM103") + monkeypatch.setenv("DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", "48") + a = torch.ones(240, 128, dtype=torch.bfloat16, device="cuda") + counts = torch.tensor([240], dtype=torch.int32, device="cuda") + output = torch.empty(1, 128, 128, dtype=torch.bfloat16, device="cuda") + with pytest.raises(RuntimeError, match="kBlockK"): + deep_gemm.bf16_mega_moe_backward_w13(output, a, a, counts, 240) From 0be759300fd5d198ef211e452ccb857217c41a1a Mon Sep 17 00:00:00 2001 From: morgendave Date: Mon, 21 Sep 2026 02:07:12 +0000 Subject: [PATCH 15/15] fix(mega): isolate wgrad tiles from global grouped alignment --- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 4 +++- .../impls/sm100_bf16_mega_moe_wgrad.hpp | 4 ++++ tests/test_mega_moe_wgrad_layout.py | 15 +++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index 2440b31761..2bb2aa1625 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -53,6 +53,7 @@ class SM100BF16GemmRuntime final { 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) { @@ -92,7 +93,8 @@ 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), diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp index 99c5f12bcc..67b96d6296 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp @@ -186,6 +186,10 @@ static void sm100_bf16_mega_moe_wgrad_1sm( .combine_num_extra_threads = static_cast(num_extra_combine_threads), .mask_grouped_k_tail = mask_grouped_k_tail, + // Count-based BF16 groups use explicit physical offsets and no SF. + // Do not inherit unrelated global grouped-GEMM alignment (e.g. 224), + // which need not be divisible by this kernel's K tile. + .k_alignment_override = kBlockK, }; SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_mega_moe_wgrad_1sm", args); } diff --git a/tests/test_mega_moe_wgrad_layout.py b/tests/test_mega_moe_wgrad_layout.py index 1900d3de83..848b35f51a 100644 --- a/tests/test_mega_moe_wgrad_layout.py +++ b/tests/test_mega_moe_wgrad_layout.py @@ -77,3 +77,18 @@ def test_wgrad_rejects_invalid_k_tile(monkeypatch): output = torch.empty(1, 128, 128, dtype=torch.bfloat16, device="cuda") with pytest.raises(RuntimeError, match="kBlockK"): deep_gemm.bf16_mega_moe_backward_w13(output, a, a, counts, 240) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a Blackwell GPU") +def test_wgrad_independent_of_global_grouped_alignment(monkeypatch): + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("requires SM100/SM103") + previous = deep_gemm.get_mk_alignment_for_contiguous_layout() + try: + # Dense/grouped BF16 users may select 224; MegaMoE's physical pool + # must remain independent, and its wgrad must not mutate that setting. + deep_gemm.set_mk_alignment_for_contiguous_layout(224) + test_wgrad_wide_k_masks_expert_tail(monkeypatch, 240, 128) + assert deep_gemm.get_mk_alignment_for_contiguous_layout() == 224 + finally: + deep_gemm.set_mk_alignment_for_contiguous_layout(previous)