diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index 2e43626bb6..2bcf6faf77 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -30,24 +30,29 @@ static void mega_moe_pre_dispatch_sm90( num_tokens, group_size, routed_scaling_factor); } -static std::tuple(const torch::Tensor&)>> +static std::tuple(const torch::Tensor&)>> get_symm_buffer_size_for_sm90_mega_moe( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, const int& hidden, const int& intermediate_hidden, - const bool& use_fp8_dispatch, const std::string& activation) { + const bool& use_fp8_dispatch, const std::string& activation, + const int& num_shared_experts = 0) { DG_HOST_ASSERT(num_experts % num_ranks == 0); DG_HOST_ASSERT(use_fp8_dispatch); DG_HOST_ASSERT(activation == "swiglu"); + DG_HOST_ASSERT(num_shared_experts >= 0); const auto workspace = layout::SM90Workspace( nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk); + const auto shared_intermediate_hidden = intermediate_hidden * num_shared_experts; const auto fp8_token_layout = layout::Data(hidden); const auto bf16_token_layout = layout::Data(hidden * 2); const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden); const auto fp8_sf_layout = layout::Data(hidden / 32); const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 16); + const auto fp8_shared_intermediate_token_layout = layout::Data(shared_intermediate_hidden); + const auto fp8_shared_intermediate_sf_layout = layout::Data(shared_intermediate_hidden / 16); const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false); const auto l1_topk_weights_layout = layout::Data(sizeof(float), false); @@ -91,10 +96,22 @@ get_symm_buffer_size_for_sm90_mega_moe( fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens, l2_token_buffer.get_end_ptr()); + // The fused shared expert reduces through one extra combine slot on the local rank const auto combine_token_buffer = layout::Buffer( - bf16_token_layout, num_topk, num_max_tokens_per_rank, + bf16_token_layout, num_topk + (num_shared_experts > 0 ? 1 : 0), num_max_tokens_per_rank, l2_sf_buffer.get_end_ptr()); + // Fused shared-expert area, appended after the combine buffer so the routed + // regions keep their relative order and are zero-sized when the shared expert is + // disabled. Both are indexed by the local token and the SF buffer is K-major + // (per-64 K groups), so no SF-pool padding is needed. + const auto shared_l2_token_buffer = layout::Buffer( + fp8_shared_intermediate_token_layout, 1, num_shared_experts > 0 ? num_max_tokens_per_rank : 0, + combine_token_buffer.get_end_ptr()); + const auto shared_l2_sf_buffer = layout::Buffer( + fp8_shared_intermediate_sf_layout, 1, num_shared_experts > 0 ? num_max_tokens_per_rank : 0, + shared_l2_token_buffer.get_end_ptr()); + DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); auto slice_input_buffers = [=](const torch::Tensor& buffer) { @@ -132,15 +149,34 @@ get_symm_buffer_size_for_sm90_mega_moe( {num_max_padded_sf_pool_tokens, intermediate_hidden / 64}, {1, num_max_padded_sf_pool_tokens}, torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); + // Fused shared expert: post-SwiGLU FP8 pool plus its M-major per-64 float SF + // (token-contiguous inner stride so a (BLOCK_M, 1) TMA box is legal; same + // layout-class as the routed L2 acts SF pool). Zero-sized when the shared + // expert is off (kept defined so the returned tuple type never changes). + auto shared_l2_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(shared_l2_token_buffer.base)), + {num_shared_experts > 0 ? num_max_tokens_per_rank : 0, shared_intermediate_hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto shared_l2_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(shared_l2_sf_buffer.base)), + {num_shared_experts > 0 ? num_max_tokens_per_rank : 0, shared_intermediate_hidden / 64}, + {1, num_shared_experts > 0 ? num_max_tokens_per_rank : 0}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf); }; - return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; + return {reinterpret_cast( + num_shared_experts > 0 ? shared_l2_sf_buffer.get_end_ptr() + : combine_token_buffer.get_end_ptr()), + slice_input_buffers}; } static void fp8_mega_moe( const torch::Tensor& y, const std::tuple& l1_weights_tuple, const std::tuple& l2_weights_tuple, + const std::optional>& shared_l1_weights_tuple_opt, + const std::optional>& shared_l2_weights_tuple_opt, const std::optional& cumulative_local_expert_recv_stats, const torch::Tensor& sym_buffer, const std::vector& sym_buffer_ptrs, const int& rank_idx, @@ -186,6 +222,37 @@ static void fp8_mega_moe( check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, num_experts_per_rank, false, true, torch::kFloat); + // Fused shared expert: a single dense MLP (no expert dimension) whose intermediate + // size is `num_shared_experts * intermediate_hidden`. Both weight tuples must be + // given together; the SF layout matches the routed weights minus the group axis. + DG_HOST_ASSERT(shared_l1_weights_tuple_opt.has_value() == shared_l2_weights_tuple_opt.has_value()); + int num_shared_experts = 0; + torch::Tensor shared_l1_weights, shared_l1_weights_sf, shared_l2_weights, shared_l2_weights_sf; + if (shared_l1_weights_tuple_opt.has_value()) { + std::tie(shared_l1_weights, shared_l1_weights_sf) = shared_l1_weights_tuple_opt.value(); + std::tie(shared_l2_weights, shared_l2_weights_sf) = shared_l2_weights_tuple_opt.value(); + const auto shared_intermediate_hidden = static_cast(shared_l2_weights.size(1)); + DG_HOST_ASSERT(shared_intermediate_hidden % intermediate_hidden == 0); + num_shared_experts = shared_intermediate_hidden / intermediate_hidden; + // The shared L2 activation SF is K-major, so its per-token row (SIH / 64 + // floats) must stay 16-byte aligned for the TMA loads of the pool it feeds + DG_HOST_ASSERT(shared_intermediate_hidden % 256 == 0); + + DG_HOST_ASSERT(shared_l1_weights.dim() == 2 and shared_l2_weights.dim() == 2); + DG_HOST_ASSERT(shared_l1_weights.size(0) == shared_intermediate_hidden * 2); + DG_HOST_ASSERT(shared_l1_weights.size(1) == hidden); + DG_HOST_ASSERT(shared_l2_weights.size(0) == hidden); + DG_HOST_ASSERT(shared_l1_weights.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(shared_l2_weights.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(shared_l1_weights.is_contiguous() and shared_l2_weights.is_contiguous()); + DG_HOST_ASSERT(get_major_type_ab(shared_l1_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(get_major_type_ab(shared_l2_weights) == cute::UMMA::Major::K); + check_sf_layout(shared_l1_weights_sf, shared_intermediate_hidden * 2, hidden, kGranMN, kGranK, + std::nullopt, false, true, torch::kFloat); + check_sf_layout(shared_l2_weights_sf, hidden, shared_intermediate_hidden, kGranMN, kGranK, + std::nullopt, false, true, torch::kFloat); + } + if (cumulative_local_expert_recv_stats.has_value()) { DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank); @@ -198,23 +265,30 @@ static void fp8_mega_moe( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - true, activation); + true, activation, num_shared_experts); DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); - const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf] = slice(sym_buffer); sm90_fp8_mega_moe(y, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, l1_weights, l2_weights, l1_weights_sf, l2_weights_sf, + // The shared L1 activations are the local `x` region and its SF + x, x_sf, + shared_l2_acts, shared_l2_acts_sf, + shared_l1_weights, shared_l2_weights, + shared_l1_weights_sf, shared_l2_weights_sf, cumulative_local_expert_recv_stats, sym_buffer_ptrs, rank_idx, num_max_tokens_per_rank, num_experts_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, + num_shared_experts, activation_clamp, fast_math); if (get_env("DG_COMM_KERNEL_DEBUG")) diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp index c449fc5b21..0ff47e99b3 100644 --- a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -22,7 +22,6 @@ struct MegaMoESM90Config { int num_max_pool_tokens; int num_padded_sf_pool_tokens; int swizzle_acts_mode, swizzle_weights_mode; - int num_experts_per_wave; int num_stages, smem_size; int num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads; @@ -33,7 +32,6 @@ struct MegaMoESM90Config { << ", num_max_pool_tokens=" << config.num_max_pool_tokens << ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens << ", swizzle_acts_mode=" << config.swizzle_acts_mode << ", swizzle_weights_mode=" << config.swizzle_weights_mode - << ", num_experts_per_wave=" << config.num_experts_per_wave << ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size << ", num_dispatch_threads=" << config.num_dispatch_threads << ", num_non_epilogue_threads=" << config.num_non_epilogue_threads @@ -44,10 +42,32 @@ struct MegaMoESM90Config { static std::tuple get_block_config_for_mega_moe_sm90( const int& num_ranks, const int& num_experts, - const int& num_topk, const int& num_tokens) { + const int& num_topk, const int& num_tokens, const int &intermediate_hidden) { const float expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; - const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; + // The relaxed 2-WG threshold enables the block_m=128 / 4-WG path only + // above a higher tokens/expert bar (instead of the original >= 64), + // trading two extra warpgroups for fewer register spills. On H20 the + // smaller SM count (78 vs 132 on H100/H200) makes the extra warpgroups + // costly, so the relaxation applies in two intermediate_hidden regimes: + // * pro (>= 3072): 4-WG only when expected_tokens_per_expert > 512 + // * flash (<= 2048): 4-WG only when expected_tokens_per_expert > 576, + // because 2-WG + BLOCK_N=256 outperforms 4-WG in part of the flash + // batch range -- 4-WG is reserved for the heaviest flash batches. + // On H200/H100 the larger SM count makes the extra warpgroups always win, + // so the original 4-WG-first (>= 64) threshold is kept for every shape, + // as well as for the H20 mid-range (2048 < intermediate_hidden < 3072). + const int num_sms = device_runtime->get_num_sms(); + const bool is_h20 = num_sms <= 84; + const bool apply_h20_pro_relaxation = is_h20 and intermediate_hidden >= 3072; + const bool apply_h20_flash_relaxation = is_h20 and intermediate_hidden <= 2048; + bool auto_split_mn; + if (apply_h20_pro_relaxation) + auto_split_mn = expected_tokens_per_expert > 512.0f; + else if (apply_h20_flash_relaxation) + auto_split_mn = expected_tokens_per_expert > 576.0f; + else + auto_split_mn = expected_tokens_per_expert >= 64.0f; if (auto_split_mn) return {128, 512}; @@ -61,34 +81,14 @@ static std::tuple get_block_config_for_mega_moe_sm90( return {block_m, num_epilogue_warpgroups * 128}; } -static int get_num_experts_per_wave_for_mega_moe_sm90( - const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, - const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms, - const int& num_ring_tokens, const int& num_max_tokens_per_rank, const int& num_ranks) { - const float expected_tokens_per_expert = - static_cast(num_tokens) * num_topk / num_experts_per_rank; - if (expected_tokens_per_expert < 1.0f or expected_tokens_per_expert > 4.0f) - return num_experts_per_rank; - - if (block_m == 64 and intermediate_hidden >= 3072) { - const int num_n_blocks_per_expert = (2 * intermediate_hidden) / block_n; - const int single_wave_blocks = - num_experts_per_rank * num_n_blocks_per_expert; - if (single_wave_blocks >= 4 * num_sms) - return num_experts_per_rank; - } - return get_num_experts_per_wave_for_mega_moe( - num_experts_per_rank, num_tokens, num_topk, - intermediate_hidden, block_m, block_n, num_sms, - num_ring_tokens, num_max_tokens_per_rank, num_ranks); -} - static bool should_use_swap_ab_for_mega_moe_sm90( const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, const int& block_m, const int& num_epilogue_threads) { // swapAB is ENABLED by default (the L1 SF-pool stride bug that corrupted // pool blocks >= 1 was fixed: BLOCK_M -> SF_BLOCK_M in the swapAB L1 epilogue). // Kill-switch retained: set DG_SM90_FP8_SWAP_AB=0 to force the non-swap path. + // swapAB composes with the fused shared expert (shared's token-axis output + // matches swapAB's reduced M-axis), so no special-casing is needed here. if (get_env("DG_SM90_FP8_SWAP_AB", 1) == 0) return false; const float expected_tokens_per_expert = @@ -128,15 +128,37 @@ static std::pair get_pipeline_config_for_mega_moe_sm90( const int smem_per_stage = block_m * block_k + block_n * block_k + smem_sfa_per_stage + smem_sfb_per_stage; - const int smem_barriers_fixed = (num_dispatch_warps + 2 * num_epilogue_warps) * 8; + // The scheduler adds 2 task-info full/empty barrier pairs and two 32-byte + // task-info slots (see `sm90_fp8_mega_moe.cuh`). `SM90MegaMoETaskInfo` is + // alignas(16)/32 B while a barrier slot is only 8 B, so the kernel pads one + // extra barrier when the preceding barrier count is odd + // (`kTaskInfoBarrierPad = kTaskInfoBaseBarriers & 1u`). Only + // `num_dispatch_warps` affects that parity (2*num_stages, 2*num_epilogue_warps + // and the 4 task-info barriers are all even), so mirror the same pad here. + const int smem_task_info_barriers = 4; // 2 full + 2 empty + const int smem_task_info_pad = (num_dispatch_warps & 1) * 8; + const int smem_barriers_fixed = + (num_dispatch_warps + 2 * num_epilogue_warps + smem_task_info_barriers) * 8 + + smem_task_info_pad; + const int smem_task_infos = 2 * 32; const int smem_barriers_per_stage = 2 * 8; - const int smem_fixed = smem_dispatch_size + smem_cd + smem_barriers_fixed; + const int smem_fixed = smem_dispatch_size + smem_cd + smem_barriers_fixed + smem_task_infos; const int num_stages = (smem_capacity - smem_fixed) / (smem_per_stage + smem_barriers_per_stage); DG_HOST_ASSERT(num_stages >= 2); const int smem_size = smem_fixed + num_stages * (smem_per_stage + smem_barriers_per_stage); DG_HOST_ASSERT(smem_size <= smem_capacity); + + // Cross-check against the kernel's exact barrier/task-info layout: the + // task-info ring (including the alignment pad) must end inside the + // allocated dynamic shared memory. + const int smem_task_info_end = + smem_dispatch_size + smem_cd + num_stages * smem_per_stage + + (num_dispatch_warps + 2 * num_stages + 2 * num_epilogue_warps + + smem_task_info_barriers + smem_task_info_pad / 8) * 8 + + smem_task_infos; + DG_HOST_ASSERT(smem_task_info_end <= smem_size); return {num_stages, smem_size}; } @@ -146,7 +168,7 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const int& hidden, const int& intermediate_hidden, const int& num_padded_sf_pool_tokens) { const auto [block_m, num_epilogue_threads] = get_block_config_for_mega_moe_sm90( - num_ranks, num_experts, num_topk, num_tokens); + num_ranks, num_experts, num_topk, num_tokens, intermediate_hidden); const float expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; const bool auto_split_mn = @@ -154,13 +176,13 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const bool decode_split_n_path = block_m == 64 and num_epilogue_threads == 256; const bool decode_use_block_n_256 = - decode_split_n_path and intermediate_hidden >= 3072 and + decode_split_n_path and expected_tokens_per_expert >= 0.25f and (2 * intermediate_hidden) % 256 == 0 and hidden % 256 == 0; const bool use_swap_ab = should_use_swap_ab_for_mega_moe_sm90( num_experts_per_rank, num_tokens, num_topk, block_m, num_epilogue_threads); - int block_n = use_swap_ab ? 128 + int block_n = use_swap_ab ? 256 : (auto_split_mn ? 256 : (decode_use_block_n_256 ? 256 : 128)); const int block_k = 128; @@ -170,22 +192,18 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const int swizzle_acts_mode = 128; const int swizzle_weights_mode = 128; - const int num_sms = device_runtime->get_num_sms(); - const int num_experts_per_wave = get_num_experts_per_wave_for_mega_moe_sm90( - num_experts_per_rank, num_tokens, num_topk, - intermediate_hidden, block_m, block_n, num_sms, - num_max_pool_tokens, num_max_tokens_per_rank, num_ranks); - - const bool reduce_decode_threads = num_epilogue_threads == 128; - const bool decode_split_n = - block_m == 64 and num_epilogue_threads == 256; - const bool shrink_non_epilogue = reduce_decode_threads or decode_split_n; - const int num_dispatch_threads = - (num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128; - const bool split_sfa_loader_warp = false; - const int num_non_epilogue_threads = - split_sfa_loader_warp ? 128 : - ((num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128); + // The scheduler needs a dedicated producer warp, so the non-epilogue section + // is exactly 3 warps (TMA-A, TMA-B, producer) and dispatch is a single warp: + // dispatch + non-epilogue = 32 + 96 = 128, a whole warpgroup that keeps the + // math warpgroups 128-thread aligned. This is the minimal aligned topology for + // every epilogue width: + // * 2-WG (epilogue=256): 32 + 96 + 256 = 384 threads, ceiling + // 65536/384 = 170 >= 168, so the epilogue accumulators do not spill. + // * 4-WG (epilogue=512): 32 + 96 + 512 = 640 threads. Halving dispatch to one + // warp is the cost of fitting the producer warp under 128-thread alignment; + // a 2-dispatch-warp variant would pad to 768 threads and spill worse. + const int num_dispatch_threads = 32; + const int num_non_epilogue_threads = 96; DG_HOST_ASSERT((num_dispatch_threads + num_non_epilogue_threads) % 128 == 0); const auto [num_stages, smem_size] = get_pipeline_config_for_mega_moe_sm90( @@ -200,7 +218,6 @@ static MegaMoESM90Config get_mega_moe_config_sm90( cluster_size, num_max_pool_tokens, num_padded_sf_pool_tokens, swizzle_acts_mode, swizzle_weights_mode, - num_experts_per_wave, num_stages, smem_size, num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads }; diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index 06e33d0270..1410eb0a8e 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -43,8 +43,8 @@ class SM90FP8MegaMoERuntime final : public LaunchRuntime bool reuse_accum_as_final; bool l2_arrival_counter; bool l2_epilogue_requires_full_sync; - bool split_phase_hot_path; bool use_swap_ab; + int num_shared_experts; MegaMoESM90Config config; // Runtime arguments @@ -66,6 +66,18 @@ class SM90FP8MegaMoERuntime final : public LaunchRuntime CUtensorMap tensor_map_l2_weights; const float* l2_weights_sf; + // Fused shared expert. When `num_shared_experts == 0` these mirror the + // routed descriptors: the kernel never reads them. Shared L1 acts SF needs + // no descriptor (the loader warp gathers the K-major `x_sf` column itself). + CUtensorMap tensor_map_shared_l1_acts; + CUtensorMap tensor_map_shared_l1_weights; + const float* shared_l1_weights_sf; + CUtensorMap tensor_map_shared_l1_output; + CUtensorMap tensor_map_shared_l2_acts; + CUtensorMap tensor_map_shared_l2_weights; + const float* shared_l2_weights_sf; + CUtensorMap tensor_map_shared_l2_acts_sf; + // Launch configs LaunchArgs launch_args; }; @@ -81,7 +93,6 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, - {}, {}, {}, {}, {}, {}, @@ -102,7 +113,6 @@ static void __instantiate_kernel() {{ args.num_max_tokens_per_rank, args.hidden, args.intermediate_hidden, args.num_experts, args.num_topk, - args.config.num_experts_per_wave, args.config.block_m, args.config.block_n, args.config.block_k, args.config.num_max_pool_tokens, args.config.num_padded_sf_pool_tokens, @@ -115,8 +125,8 @@ static void __instantiate_kernel() {{ args.reuse_accum_as_final ? "true" : "false", args.l2_arrival_counter ? "true" : "false", args.l2_epilogue_requires_full_sync ? "true" : "false", - args.split_phase_hot_path ? "true" : "false", - args.use_swap_ab ? "true" : "false"); + args.use_swap_ab ? "true" : "false", + args.num_shared_experts); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -133,7 +143,15 @@ static void __instantiate_kernel() {{ args.tensor_map_l2_acts, args.tensor_map_l2_acts_sf, args.tensor_map_l2_weights, - args.l2_weights_sf + args.l2_weights_sf, + args.tensor_map_shared_l1_acts, + args.tensor_map_shared_l1_weights, + args.shared_l1_weights_sf, + args.tensor_map_shared_l1_output, + args.tensor_map_shared_l2_acts, + args.tensor_map_shared_l2_weights, + args.shared_l2_weights_sf, + args.tensor_map_shared_l2_acts_sf )); } }; @@ -144,20 +162,33 @@ static void sm90_fp8_mega_moe( const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf, const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf, + // Fused shared expert. `shared_l1_acts` is the local `x` region (and + // `shared_l1_acts_sf` its K-major per-128 SF); `shared_l2_acts` is the + // post-SwiGLU pool written by the fused L1 epilogue (zero-sized when the shared + // expert is off). The weight tensors are undefined when `num_shared_experts == 0`. + const torch::Tensor& shared_l1_acts, const torch::Tensor& shared_l1_acts_sf, + const torch::Tensor& shared_l2_acts, const torch::Tensor& shared_l2_acts_sf, + const torch::Tensor& shared_l1_weights, const torch::Tensor& shared_l2_weights, + const torch::Tensor& shared_l1_weights_sf, const torch::Tensor& shared_l2_weights_sf, const std::optional cumulative_local_expert_recv_stats, const std::vector& sym_buffer_ptrs, const int& rank_idx, const int& num_max_tokens_per_rank, const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, + const int& num_shared_experts, const float& activation_clamp, const bool& fast_math ) { const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts = num_experts_per_rank * num_ranks; const auto num_padded_sf_pool_tokens = static_cast(l1_acts_sf.size(0)); + const bool fuse_shared_experts = num_shared_experts > 0; + const int shared_intermediate_hidden = intermediate_hidden * num_shared_experts; - // Heuristics + // Heuristics. swapAB composes with the fused shared expert (shared's token + // M-axis matches swapAB's reduced output M-axis), so `use_swap_ab` is computed + // naturally for both paths. const auto config = get_mega_moe_config_sm90( num_ranks, num_experts, num_experts_per_rank, num_max_tokens_per_rank, num_tokens, num_topk, @@ -178,8 +209,6 @@ static void sm90_fp8_mega_moe( const bool default_split_mn_barrier_opt = config.block_m == 128 and config.block_n == 256 and config.num_epilogue_threads == 512; - const bool split_phase_hot_path = - config.block_m == 128 and config.block_n == 256 and hidden >= 7168; const bool decode_split_n_path = config.block_m == 64 and config.num_epilogue_threads == 256; const bool decode_split_n_bn256 = @@ -242,11 +271,12 @@ static void sm90_fp8_mega_moe( const int wg_l1_out_block_n = wg_block_n / 2; const bool split_n_shares_sf = split_n_warpgroups and wg_l1_out_block_n < kL2ActsSFGranK; + const bool l1_output_full_tile = split_n_shares_sf or use_swap_ab; const int l1_output_swizzle_mode = 0; const int l1_output_box_n = - split_n_shares_sf ? config.block_n / 2 : wg_l1_out_block_n; + l1_output_full_tile ? config.block_n / 2 : wg_l1_out_block_n; const int l1_output_box_m = - split_n_shares_sf ? config.block_m : wg_block_m; + l1_output_full_tile ? config.block_m : wg_block_m; const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, intermediate_hidden, config.num_max_pool_tokens, l1_output_box_n, l1_output_box_m, @@ -267,6 +297,80 @@ static void sm90_fp8_mega_moe( static_cast(l2_weights.stride(-2)), config.swizzle_weights_mode); + // ---- Fused shared expert descriptors ---- + // A: the local `x` region (K-major FP8) for L1 and the post-SwiGLU shared pool for + // L2, both with M = num_max_tokens_per_rank. Only the shared L2 activation SF gets a + // descriptor: the fused L1 epilogue writes it M-major, so a (BLOCK_M, 1) box is legal. + // The shared L1 activation SF (`x_sf`) is K-major, where that box would be 4 bytes and + // break TMA's 16-byte inner-box rule, so the loader warp gathers it into `smem_sfa`. + auto tensor_map_shared_l1_acts = tensor_map_l1_acts; + auto tensor_map_shared_l1_weights = tensor_map_l1_weights; + auto tensor_map_shared_l1_output = tensor_map_l1_output; + auto tensor_map_shared_l2_acts = tensor_map_l2_acts; + auto tensor_map_shared_l2_weights = tensor_map_l2_weights; + auto tensor_map_shared_l2_acts_sf = tensor_map_l2_acts_sf; + const float* shared_l1_weights_sf_ptr = l1_weights_sf.data_ptr(); + const float* shared_l2_weights_sf_ptr = l2_weights_sf.data_ptr(); + if (fuse_shared_experts) { + DG_HOST_ASSERT(shared_l1_acts.defined() and shared_l1_acts_sf.defined()); + DG_HOST_ASSERT(shared_l2_acts.defined() and shared_l2_acts_sf.defined()); + DG_HOST_ASSERT(static_cast(shared_l1_acts.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l1_acts.size(1)) == hidden); + DG_HOST_ASSERT(static_cast(shared_l1_acts_sf.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l1_acts_sf.size(1)) == hidden / kGranK); + DG_HOST_ASSERT(static_cast(shared_l2_acts.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l2_acts.size(1)) == shared_intermediate_hidden); + DG_HOST_ASSERT(static_cast(shared_l2_acts_sf.size(0)) == num_max_tokens_per_rank); + DG_HOST_ASSERT(static_cast(shared_l2_acts_sf.size(1)) == + shared_intermediate_hidden / kL2ActsSFGranK); + + tensor_map_shared_l1_acts = make_tma_2d_desc(shared_l1_acts, + hidden, num_max_tokens_per_rank, + config.block_k, config.block_m, + static_cast(shared_l1_acts.stride(-2)), + config.swizzle_acts_mode); + tensor_map_shared_l1_weights = make_tma_2d_desc(shared_l1_weights, + hidden, shared_intermediate_hidden * 2, + config.block_k, weight_tma_block_n, + static_cast(shared_l1_weights.stride(-2)), + config.swizzle_weights_mode); + tensor_map_shared_l1_output = make_tma_2d_desc(shared_l2_acts, + shared_intermediate_hidden, num_max_tokens_per_rank, + l1_output_box_n, l1_output_box_m, + static_cast(shared_l2_acts.stride(-2)), + l1_output_swizzle_mode); + tensor_map_shared_l2_acts = make_tma_2d_desc(shared_l2_acts, + shared_intermediate_hidden, num_max_tokens_per_rank, + config.block_k, config.block_m, + static_cast(shared_l2_acts.stride(-2)), + config.swizzle_acts_mode); + tensor_map_shared_l2_weights = make_tma_2d_desc(shared_l2_weights, + shared_intermediate_hidden, hidden, + config.block_k, weight_tma_block_n, + static_cast(shared_l2_weights.stride(-2)), + config.swizzle_weights_mode); + shared_l1_weights_sf_ptr = shared_l1_weights_sf.data_ptr(); + shared_l2_weights_sf_ptr = shared_l2_weights_sf.data_ptr(); + // Shared L1 acts SF (`x_sf`) stays K-major: staging an M-major copy would cost a + // full transpose kernel on every call and the staged tensor would be freed while + // the launch is still in flight, so the loader warp gathers the column into + // `smem_sfa` itself (see `process_a_sfa_block`) and needs no descriptor. + // + // Shared L2 acts SF lives in the workspace `shared_l2_sf_buffer` and is written + // M-major during the launch by the fused L1 epilogue, so it can TMA-load with a + // (BLOCK_M, 1) box. Re-view the same memory with {1, nmt} strides (the from_blob + // view already attached those in `sm90_mega.hpp`) and build the descriptor. + auto shared_l2_acts_sf_mm = torch::from_blob( + shared_l2_acts_sf.data_ptr(), + {num_max_tokens_per_rank, shared_intermediate_hidden / 64}, + {1, num_max_tokens_per_rank}, + shared_l2_acts_sf.options()); + tensor_map_shared_l2_acts_sf = make_tma_sf_desc( + cute::UMMA::Major::MN, shared_l2_acts_sf_mm, + num_max_tokens_per_rank, shared_intermediate_hidden, + config.block_m, kL2ActsSFGranK, 1, 0); + } + // Stats can be optional int* cumulative_local_expert_recv_stats_ptr = nullptr; if (cumulative_local_expert_recv_stats.has_value()) @@ -285,8 +389,8 @@ static void sm90_fp8_mega_moe( .reuse_accum_as_final = reuse_accum_as_final, .l2_arrival_counter = l2_arrival_counter, .l2_epilogue_requires_full_sync = l2_epilogue_requires_full_sync, - .split_phase_hot_path = split_phase_hot_path, .use_swap_ab = use_swap_ab, + .num_shared_experts = num_shared_experts, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, @@ -301,6 +405,14 @@ static void sm90_fp8_mega_moe( .tensor_map_l2_acts_sf = tensor_map_l2_acts_sf, .tensor_map_l2_weights = tensor_map_l2_weights, .l2_weights_sf = l2_weights_sf.data_ptr(), + .tensor_map_shared_l1_acts = tensor_map_shared_l1_acts, + .tensor_map_shared_l1_weights = tensor_map_shared_l1_weights, + .shared_l1_weights_sf = shared_l1_weights_sf_ptr, + .tensor_map_shared_l1_output = tensor_map_shared_l1_output, + .tensor_map_shared_l2_acts = tensor_map_shared_l2_acts, + .tensor_map_shared_l2_weights = tensor_map_shared_l2_weights, + .shared_l2_weights_sf = shared_l2_weights_sf_ptr, + .tensor_map_shared_l2_acts_sf = tensor_map_shared_l2_acts_sf, .launch_args = LaunchArgs(num_sms, config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, config.smem_size, config.cluster_size) }; diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index bee2310a7f..2ff6bcb8f6 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -684,9 +684,10 @@ dg_get_symm_buffer_size_for_mega_moe(int64_t num_ranks, int64_t num_experts, int num_bytes, slice_input_buffers); } -Tuple(TensorView)>> +Tuple(TensorView)>> dg_get_symm_buffer_size_for_sm90_mega_moe(int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, int64_t num_topk, int64_t hidden, - int64_t intermediate_hidden, bool use_fp8_dispatch, std::string activation) { + int64_t intermediate_hidden, bool use_fp8_dispatch, std::string activation, + int64_t num_shared_experts) { auto [num_bytes, fn] = mega::get_symm_buffer_size_for_sm90_mega_moe( static_cast(num_ranks), static_cast(num_experts), @@ -695,12 +696,16 @@ dg_get_symm_buffer_size_for_sm90_mega_moe(int64_t num_ranks, int64_t num_experts static_cast(hidden), static_cast(intermediate_hidden), use_fp8_dispatch, - activation + activation, + static_cast(num_shared_experts) ); auto slice_input_buffers = [=](TensorView buffer) { - auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = fn(convert_to_torch_tensor(buffer)); - return Tuple( + // The last two views are the fused shared-expert pool and its SF; they are + // zero-sized when the shared expert is disabled. + auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf] = fn(convert_to_torch_tensor(buffer)); + return Tuple( Tensor::FromDLPack(at::toDLPack(x.view(at::kChar))), Tensor::FromDLPack(at::toDLPack(x_sf)), Tensor::FromDLPack(at::toDLPack(topk_idx)), @@ -708,10 +713,12 @@ dg_get_symm_buffer_size_for_sm90_mega_moe(int64_t num_ranks, int64_t num_experts Tensor::FromDLPack(at::toDLPack(l1_acts.view(at::kChar))), Tensor::FromDLPack(at::toDLPack(l1_acts_sf)), Tensor::FromDLPack(at::toDLPack(l2_acts.view(at::kChar))), - Tensor::FromDLPack(at::toDLPack(l2_acts_sf)) + Tensor::FromDLPack(at::toDLPack(l2_acts_sf)), + Tensor::FromDLPack(at::toDLPack(shared_l2_acts.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(shared_l2_acts_sf)) ); }; - return Tuple(TensorView)>>( + return Tuple(TensorView)>>( num_bytes, slice_input_buffers); } @@ -770,6 +777,8 @@ void dg_bf16_mega_moe(TensorView y, TensorView l1_weights, TensorView l2_weights void dg_fp8_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_sf, TensorView l2_weights, TensorView l2_weights_sf, + Optional shared_l1_weights, Optional shared_l1_weights_sf, + Optional shared_l2_weights, Optional shared_l2_weights_sf, Optional cumulative_local_expert_recv_stats, TensorView sym_buffer, Array sym_buffer_ptrs, int64_t rank_idx, int64_t num_max_tokens_per_rank, int64_t num_experts, int64_t num_topk, Tuple recipe, std::string activation, Optional activation_clamp_opt, bool fast_math) { @@ -784,10 +793,25 @@ void dg_fp8_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_ auto [recipe_a, recipe_b, recipe_c] = recipe; auto recipe_val = std::make_tuple(static_cast(recipe_a), static_cast(recipe_b), static_cast(recipe_c)); + // Fused shared expert: both weight tuples are optional and must come together + std::optional> shared_l1_weights_val = std::nullopt; + std::optional> shared_l2_weights_val = std::nullopt; + if (shared_l1_weights.has_value()) { + DG_HOST_ASSERT(shared_l1_weights_sf.has_value() and shared_l2_weights.has_value() and + shared_l2_weights_sf.has_value()); + shared_l1_weights_val = std::make_tuple( + convert_to_torch_tensor(shared_l1_weights.value()), + convert_to_torch_tensor(shared_l1_weights_sf.value())); + shared_l2_weights_val = std::make_tuple( + convert_to_torch_tensor(shared_l2_weights.value()), + convert_to_torch_tensor(shared_l2_weights_sf.value())); + } + mega::fp8_mega_moe( convert_to_torch_tensor(y), std::make_pair(convert_to_torch_tensor(l1_weights), convert_to_torch_tensor(l1_weights_sf)), std::make_pair(convert_to_torch_tensor(l2_weights), convert_to_torch_tensor(l2_weights_sf)), + shared_l1_weights_val, shared_l2_weights_val, c_val, convert_to_torch_tensor(sym_buffer), sym_buffer_ptrs_val, static_cast(rank_idx), static_cast(num_max_tokens_per_rank), static_cast(num_experts), static_cast(num_topk), recipe_val, activation, act_clamp_opt_val, fast_math diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh index 222908412e..7b9b3f21bf 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -68,50 +68,15 @@ __forceinline__ __device__ void sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( sf.y = __fmul_rn(ay, kScale), sf_inv.y = 1.0f / sf.y; } -template -CUTLASS_DEVICE void sm90_fp8_mega_moe_for_each_block_split( - sched::MegaMoEScheduler& scheduler, - L1Func&& l1_func, L2Func&& l2_func) { - scheduler.fetch_expert_recv_count(); - scheduler.set_expert_idx(0); - - while (true) { - CUTE_TIE_DECL(scheduler.get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx); - if (block_phase == sched::BlockPhase::None) - break; - - if (block_phase == sched::BlockPhase::Linear1) { - l1_func(current_local_expert_idx, kNumL1BlockKs, m_block_idx, n_block_idx); - } else { - l2_func(current_local_expert_idx, kNumL2BlockKs, m_block_idx, n_block_idx); - } - } -} - // ============================================================================ // SM90 (Hopper) FP8 MegaMoE — full implementation // ---------------------------------------------------------------------------- // Pipeline (cluster=1, no TMA multicast): // * Dispatch warps: pull tokens (FP8) and SF (per-128 channel float) from // remote ranks via NVLink into the local L1 pool. +// * Producer warp: claims L1/L2 tasks (routed, plus SharedLinear1/2 when the +// shared expert is fused in) from global atomic counters and publishes them +// into a 2-stage SMEM task ring that every consumer warp drains in order. // * GEMM TMA-load warps (1 for A+SFA, 1 for B+SFB) feed the pipeline stages. // * Math warpgroups (totalling kNumEpilogueThreads) consume each // stage with WGMMA, accumulate into registers, then run the epilogue: @@ -131,7 +96,6 @@ template < uint32_t kNumMaxTokensPerRank, uint32_t kHidden, uint32_t kIntermediateHidden, uint32_t kNumExperts, uint32_t kNumTopk, - uint32_t kNumExpertsPerWave, uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, uint32_t kNumMaxPoolTokens, uint32_t kNumPaddedSFPoolTokens, @@ -145,8 +109,8 @@ template < bool kReuseAccumAsFinal, bool kL2ArrivalCounter, bool kL2EpilogueRequiresFullSync, - bool kSplitPhaseHotPath, bool kFP8SwapAB = false, + uint32_t kNumSharedExperts = 0, uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, uint32_t L2_SHAPE_N = kHidden, @@ -172,17 +136,34 @@ sm90_fp8_mega_moe_impl(void* y, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, - const float* __restrict__ l2_weights_sf) { + const float* __restrict__ l2_weights_sf, + // Fused shared expert (only read when `kNumSharedExperts > 0`; + // otherwise the host passes the routed descriptors as placeholders). + // Shared L1 acts SF has no descriptor: `x_sf` is K-major, so a + // (BLOCK_M, 1) box is illegal and the loader warp gathers the column + // into `smem_sfa` itself. Shared L2 acts SF is written M-major by the + // fused L1 epilogue, so it TMA-loads like the routed L2 SFA. + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_weights, + const float* __restrict__ shared_l1_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l2_weights, + const float* __restrict__ shared_l2_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_shared_l2_acts_sf) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; // ===================================================================== // Template checks // ===================================================================== - DG_STATIC_ASSERT(kNumDispatchThreads >= 64 and kNumDispatchThreads % 64 == 0, + DG_STATIC_ASSERT(kNumDispatchThreads >= 32 and kNumDispatchThreads % 32 == 0, "Invalid number of dispatch threads"); - DG_STATIC_ASSERT(kNumNonEpilogueThreads == 64 or kNumNonEpilogueThreads == 128, + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 64 or kNumNonEpilogueThreads == 96 or + kNumNonEpilogueThreads == 128 or kNumNonEpilogueThreads == 192, "Invalid number of GEMM TMA warps"); + DG_STATIC_ASSERT(kNumMMANonEpilogueWarps >= 3, + "The scheduler needs a dedicated producer warp"); DG_STATIC_ASSERT((kNumDispatchThreads + kNumNonEpilogueThreads) % 128 == 0, "Math warpgroup start must be 128-thread aligned"); DG_STATIC_ASSERT(kNumEpilogueThreads % 128 == 0, "Invalid number of math/epilogue threads"); @@ -209,6 +190,14 @@ sm90_fp8_mega_moe_impl(void* y, cute::prefetch_tma_descriptor(&tensor_map_l2_acts); cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); cute::prefetch_tma_descriptor(&tensor_map_l2_weights); + if constexpr (kNumSharedExperts > 0) { + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_shared_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_shared_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_shared_l2_weights); + cute::prefetch_tma_descriptor(&tensor_map_shared_l2_acts_sf); + } } // ===================================================================== @@ -221,6 +210,26 @@ sm90_fp8_mega_moe_impl(void* y, DG_STATIC_ASSERT(kNumPaddedSFPoolTokens >= kNumPoolBlocks * SF_BLOCK_M, "Invalid SM90 MegaMoE SF pool capacity"); + // Fused shared expert. `M` is this rank's own token count, the intermediate size + // is scaled by the number of shared experts, and there is no expert dimension. + constexpr bool kHasSharedExperts = kNumSharedExperts > 0; + constexpr uint32_t kSharedIntermediateHidden = kIntermediateHidden * kNumSharedExperts; + constexpr uint32_t SHARED_L1_SHAPE_N = kSharedIntermediateHidden * 2; + constexpr uint32_t SHARED_L1_SHAPE_K = kHidden; + constexpr uint32_t SHARED_L2_SHAPE_N = kHidden; + constexpr uint32_t SHARED_L2_SHAPE_K = kSharedIntermediateHidden; + // The shared activation SF buffers are plain K-major (see the loader): their row + // strides must stay 16-byte aligned for TMA, which `layout::Data` already enforces. + DG_STATIC_ASSERT(not kHasSharedExperts or kSharedIntermediateHidden % 256 == 0, + "Shared intermediate hidden must be a multiple of 256"); + // swapAB's reduced-output tile shares the token M-axis the shared expert runs on, + // so the two compose: swapAB owns the SwiGLU/quantize/store epilogue, and the + // shared path merely selects its own weight/SF/output descriptors via `is_shared` + // ternaries (see `run_swap_ab_l1`/`run_swap_ab_l2` and the swapAB L1 epilogue). + DG_STATIC_ASSERT(not (kHasSharedExperts and kFP8SwapAB) or + (kSharedIntermediateHidden * 2) % BLOCK_N == 0, + "swapAB + shared expert requires the shared L1 N to tile evenly"); + const auto workspace = layout::SM90Workspace( sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk); @@ -250,8 +259,26 @@ sm90_fp8_mega_moe_impl(void* y, const auto l2_token_buffer = layout::Buffer(fp8_intermediate_token_layout, 1, kNumMaxPoolTokens, l1_topk_weights_buffer.get_end_ptr()); const auto l2_sf_buffer = layout::Buffer(fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens, l2_token_buffer.get_end_ptr()); - // Combine input area - const auto combine_token_buffer = layout::Buffer(bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, l2_sf_buffer.get_end_ptr()); + // Combine input area. The fused shared expert reduces through one extra slot + // (`topk_idx == kNumTopk`) written by the local rank only. + constexpr uint32_t kNumCombineSlots = kNumTopk + (kHasSharedExperts ? 1u : 0u); + const auto combine_token_buffer = layout::Buffer(bf16_token_layout, kNumCombineSlots, kNumMaxTokensPerRank, l2_sf_buffer.get_end_ptr()); + + // Fused shared-expert area, appended after the combine buffer so the routed + // regions keep their relative order and are zero-sized when the shared expert is + // disabled (the workspace itself always reserves the shared arrival counters, so + // absolute offsets shift by a few KB either way -- host and device agree because + // both derive them from the same `SM90Workspace`). + // The post-SwiGLU FP8 output and its per-64-K float SF are indexed by the local + // token index; the SF buffer is K-major (no SF-pool padding needed). + constexpr auto fp8_shared_intermediate_token_layout = layout::Data(kSharedIntermediateHidden); + constexpr auto fp8_shared_intermediate_sf_layout = layout::Data(kSharedIntermediateHidden / 16); + const auto shared_l2_token_buffer = layout::Buffer( + fp8_shared_intermediate_token_layout, 1, kHasSharedExperts ? kNumMaxTokensPerRank : 0, + combine_token_buffer.get_end_ptr()); + const auto shared_l2_sf_buffer = layout::Buffer( + fp8_shared_intermediate_sf_layout, 1, kHasSharedExperts ? kNumMaxTokensPerRank : 0, + shared_l2_token_buffer.get_end_ptr()); // ===================================================================== // GEMM data types and shape constants @@ -273,6 +300,7 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t kNumCombineWarps = kNumEpilogueWarps; using L1WGMMA = typename mma::sm90::FP8MMASelector::type; // M=64, N=WG_BLOCK_N, K=32 using L2WGMMA = typename mma::sm90::FP8MMASelector::type; + using SwapWGMMA64 = typename mma::sm90::FP8MMASelector<64>::type; constexpr uint32_t kL1OutputArrivalParts = 1; static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, "Unexpected WGMMA shape"); @@ -294,7 +322,7 @@ sm90_fp8_mega_moe_impl(void* y, // feeds that shared SF must be reduced across both warpgroups. constexpr bool kSplitNSharesSF = kSplitNWarpgroups and (WG_L1_OUT_BLOCK_N < 64); constexpr bool kSwapABEligible = - kFP8SwapAB and kSplitNWarpgroups and (BLOCK_M == 64) and (BLOCK_N == 128) and + kFP8SwapAB and kSplitNWarpgroups and (BLOCK_M == 64) and (BLOCK_N == 256) and (kWarpgroupSplitN == 2); constexpr bool kSwapABActive = kSwapABEligible; constexpr uint32_t kSwapABTokenChunks = BLOCK_M / 8; @@ -383,6 +411,22 @@ sm90_fp8_mega_moe_impl(void* y, auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages + i; }); auto combine_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + i; }); + // Interleaved-scheduler task ring: 2-stage task-info slots with dedicated + // full/empty barriers, placed right after the combine barriers (the host + // SMEM accounting reserves the same bytes, see `sm90_mega_moe.hpp`). + // `SM90MegaMoETaskInfo` is alignas(16) / 32 B, but barrier slots are 8 B, + // so the barrier count preceding the ring must be even for the 32 B slots + // to land 16-byte aligned. Pad one unused barrier slot when it is odd -- + // this fires on the 1-dispatch-warp topology (parity is set by + // kNumDispatchWarps alone; 2*kNumStages and 2*kNumCombineWarps are even). + constexpr uint32_t kTaskInfoBaseBarriers = + kNumDispatchWarps + kNumStages * 2 + kNumCombineWarps * 2; + constexpr uint32_t kTaskInfoBarrierPad = kTaskInfoBaseBarriers & 1u; + auto task_info_full_barriers = barrier_start_ptr + kTaskInfoBaseBarriers + kTaskInfoBarrierPad; + auto task_info_empty_barriers = task_info_full_barriers + sched::kNumSM90TaskInfoStages; + auto task_infos = reinterpret_cast( + task_info_empty_barriers + sched::kNumSM90TaskInfoStages); + // ===================================================================== // Initialization // ===================================================================== @@ -411,6 +455,13 @@ sm90_fp8_mega_moe_impl(void* y, #pragma unroll for (uint32_t i = 0; i < kNumCombineWarps * 2; ++ i) combine_barriers[i]->init(1); + #pragma unroll + for (uint32_t i = 0; i < sched::kNumSM90TaskInfoStages; ++ i) { + // The producer warp publishes one task per slot + task_info_full_barriers[i].init(1); + // TMA-A + TMA-B warps and every math warp release each slot once + task_info_empty_barriers[i].init(2 + kNumEpilogueWarps); + } } cutlass::arch::fence_barrier_init(); } @@ -419,20 +470,32 @@ sm90_fp8_mega_moe_impl(void* y, // ===================================================================== // Scheduler (cluster=1) // ===================================================================== - constexpr uint32_t kNumExpertsPerLane = math::constexpr_ceil_div(kNumExpertsPerRank, 32u); - constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; - constexpr uint32_t kNumL2BlockNs = L2_SHAPE_N / BLOCK_N; constexpr uint32_t kNumL1BlockKs = L1_SHAPE_K / BLOCK_K; constexpr uint32_t kNumL2BlockKs = L2_SHAPE_K / BLOCK_K; - auto scheduler = sched::MegaMoEScheduler< + constexpr uint32_t kNumSharedL1BlockNs = SHARED_L1_SHAPE_N / BLOCK_N; + constexpr uint32_t kNumSharedL1BlockKs = SHARED_L1_SHAPE_K / BLOCK_K; + constexpr uint32_t kNumSharedL2BlockKs = SHARED_L2_SHAPE_K / BLOCK_K; + // The shared L2 N shape equals the routed one, which the scheduler already checks + DG_STATIC_ASSERT(not kHasSharedExperts or kNumSharedL1BlockNs > 0, + "BLOCK_N is too large for the shared-expert L1 shape"); + // Number of K blocks for a task of the given phase + const auto get_num_k_blocks = [](const sched::BlockPhase& block_phase) -> uint32_t { + switch (block_phase) { + case sched::BlockPhase::Linear1: return kNumL1BlockKs; + case sched::BlockPhase::Linear2: return kNumL2BlockKs; + case sched::BlockPhase::SharedLinear1: return kNumSharedL1BlockKs; + default: return kNumSharedL2BlockKs; + } + }; + using SchedulerT = sched::MegaMoEInterleavedScheduler< BLOCK_M, BLOCK_N, BLOCK_K, L1_SHAPE_N, L1_SHAPE_K, L2_SHAPE_N, L2_SHAPE_K, - kNumExpertsPerRank, kNumExpertsPerWave, + kNumExpertsPerRank, kNumSMs, kNumRanks, - kNumExpertsPerLane, kNumL1BlockNs, kNumL2BlockNs, - kNumL1BlockKs, kNumL2BlockKs, - layout::SM90Workspace>(workspace); + kNumSharedExperts, + layout::SM90Workspace>; + SchedulerT scheduler(workspace, task_info_full_barriers, task_info_empty_barriers, task_infos); // Pipeline state shared by TMA loaders and math warpgroups uint32_t stage_idx = 0, phase = 0; @@ -454,10 +517,14 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; // Register reconfiguration counts (chosen to fit in 64512 reg budget). - // For the 256-epilogue-thread split-N decode path: - // 64*48 + 64*40 + 256*168 = 48640 <= 64512. - // For the 512-epilogue-thread split-MN path, trim dispatch and loader roles - // so launch bounds still leave enough WGMMA registers. + // The CTA topology is 1 dispatch warp + TMA-A/TMA-B/producer warps + the + // epilogue warpgroups: + // * 2-WG (epilogue=256): 32*48 + 96*40 + 256*168 = 48384 <= 64512, and + // at 384 threads the launch_bounds ceiling is 65536/384 = 170 >= 168, + // so the WGMMA accumulators do not spill (a 512-thread CTA would cap + // them at 128 < 168 and force local memory). + // * 4-WG (epilogue=512): 32*32 + 96*24 + 512*112 = 60672 <= 64512, with a + // 640-thread CTA ceiling of 65536/640 = 102. // Reduced-thread decode (kNumThreads<=256) raises the launch-bounds // register ceiling to 65536/256=256; grant the epilogue warpgroup the full // 256 so the accumulator double-buffer fits without spilling. @@ -720,6 +787,22 @@ sm90_fp8_mega_moe_impl(void* y, #pragma unroll for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) *workspace.get_expert_send_count_ptr(i) = 0; + // Reset the scheduler's global task counters for the next launch + if (warp_idx == 0 and cute::elect_one_sync()) { + *workspace.get_l1_task_count_ptr() = 0; + *workspace.get_l2_task_count_ptr() = 0; + if constexpr (kHasSharedExperts) { + *workspace.get_shared_l1_task_count_ptr() = 0; + *workspace.get_shared_l2_task_count_ptr() = 0; + } + } + + if constexpr (kHasSharedExperts) { + // Reset the per-M-block shared L1 arrival counters + const uint32_t num_shared_blocks = math::ceil_div(num_tokens, BLOCK_M); + for (uint32_t i = thread_idx; i < num_shared_blocks; i += kNumDispatchThreads) + *workspace.get_shared_l2_full_count_ptr(i) = 0; + } } else { for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { const auto num_recv_tokens = static_cast( @@ -730,11 +813,21 @@ sm90_fp8_mega_moe_impl(void* y, ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - DG_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); - if (warp_idx == 0) { + // Zero the per-expert recv-count sum on dispatch warp 0. The + // cumulative-stats red_add normally runs on dispatch warp 1; + // with the 384-thread topology the CTA collapses to + // a single dispatch warp, so fold it onto warp 0's elect-one lane. + if (warp_idx == 0) *workspace.get_expert_recv_count_sum_ptr(i) = 0; - } else if (warp_idx == 1) { - if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + if constexpr (kNumDispatchWarps >= 2) { + if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + } else { + if (warp_idx == 0 and cute::elect_one_sync() and + cumulative_local_expert_recv_stats != nullptr) ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); __syncwarp(); } @@ -768,26 +861,28 @@ sm90_fp8_mega_moe_impl(void* y, auto process_a_sfa_block = [&](const auto& block_phase, const uint32_t& local_expert_idx, const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - const auto tensor_map_a_ptr = block_phase == sched::BlockPhase::Linear2 - ? &tensor_map_l2_acts : &tensor_map_l1_acts; - const auto tensor_map_sfa_ptr = block_phase == sched::BlockPhase::Linear2 - ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; - - const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t& m_block_idx, const uint32_t& n_block_idx, + const uint32_t& valid_m, const uint32_t& pool_block_idx) { + const bool is_shared = kHasSharedExperts and sched::is_shared_phase(block_phase); + const bool is_l1 = sched::is_l1_phase(block_phase); + // Shared L1 reads the local `x` directly; shared L2 reads the post-SwiGLU + // shared pool written by the L1 epilogue. Shared activation SF is read from + // global memory by the math warps, so no SFA descriptor is needed for it. + const auto tensor_map_a_ptr = is_shared + ? (is_l1 ? &tensor_map_shared_l1_acts : &tensor_map_shared_l2_acts) + : (is_l1 ? &tensor_map_l1_acts : &tensor_map_l2_acts); + const auto tensor_map_sfa_ptr = is_l1 ? &tensor_map_l1_acts_sf : &tensor_map_l2_acts_sf; // Wait for the pool to be ready if (block_phase == sched::BlockPhase::Linear1) { const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); - const auto expected = scheduler.template get_valid_m(); - while (ptx::ld_acq(ptr) != expected); - } else { + while (ptx::ld_acq(ptr) != valid_m); + } else if (block_phase == sched::BlockPhase::Linear2) { constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; if constexpr (kL2ArrivalCounter) { const auto ptr = reinterpret_cast( workspace.get_l2_arrival_mask_ptr(pool_block_idx)); - const uint32_t active_m_wgs = math::ceil_div( - scheduler.template get_valid_m(), WG_BLOCK_M); + const uint32_t active_m_wgs = math::ceil_div(valid_m, WG_BLOCK_M); const uint32_t expected = kNumL1BlockNs * active_m_wgs * kWarpgroupSplitN * kL1OutputArrivalParts; while (ptx::ld_acq(ptr) != expected); @@ -797,22 +892,74 @@ sm90_fp8_mega_moe_impl(void* y, ? ~0ull : ((1ull << kNumL1BlockNs) - 1ull); while (ptx::ld_acq_gpu(ptr) != expected); } + } else if constexpr (kHasSharedExperts) { + // `SharedLinear1` has no dependency at all: `x` is resident before the + // launch. `SharedLinear2` waits for every shared L1 N tile of this M + // block, in the same counter mode the routed L2 uses. + if (block_phase == sched::BlockPhase::SharedLinear2) { + const auto ptr = workspace.get_shared_l2_full_count_ptr(pool_block_idx); + const uint32_t active_m_wgs = math::ceil_div(valid_m, WG_BLOCK_M); + const uint32_t expected = + kNumSharedL1BlockNs * active_m_wgs * kWarpgroupSplitN * kL1OutputArrivalParts; + while (ptx::ld_acq(ptr) != expected); + } } for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { empty_barriers[stage_idx]->wait(phase ^ 1); - if (cute::elect_one_sync()) { - const uint32_t m_idx = pool_block_idx * BLOCK_M; - const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M; - const uint32_t k_idx = k_block_idx * BLOCK_K; + // Shared tiles index the local token pool directly, so the padded + // SF-pool row mapping does not apply to them. + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // Shared L1 activation SF: `x_sf` is K-major, so a (BLOCK_M, 1) TMA box + // would be 4 bytes and break TMA's 16-byte inner-box rule. Gather the + // column into `smem_sfa` with the whole loader warp instead (BLOCK_M / 32 + // loads per lane, once per stage, versus one per math thread every stage), + // so the math warps keep the uniform `ld_shared` path. `x_sf` is written + // before the launch and never mutated, hence `__ldg`. The stores are + // published by the mbarrier arrive below (which only counts TMA bytes). + if constexpr (kHasSharedExperts) { + if (is_shared and is_l1) { + constexpr uint32_t kNumInputSFGroups = kHidden / kGranK; + const float* in_sf = + input_sf_buffer.template get_base_ptr() + k_block_idx; + #pragma unroll + for (uint32_t i = 0; i < BLOCK_M / 32; ++ i) { + const uint32_t row = i * 32 + lane_idx; + smem_sfa[stage_idx][row] = __ldg(in_sf + (m_idx + row) * kNumInputSFGroups); + } + } + __syncwarp(); + } + if (cute::elect_one_sync()) { // TMA load A tma::copy( tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx, 1); // TMA load SFA - if (block_phase == sched::BlockPhase::Linear1) { + if (is_shared) { + if (is_l1) { + // Gathered above by the whole warp; no TMA bytes to expect + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE); + } else { + // Shared L2 acts SF: the fused L1 epilogue writes it M-major so a + // (BLOCK_M, 1) TMA box is legal -- mirrors routed L2 SFA. Box is + // (BLOCK_M, 1), so issue two single-group TMAs at smem offsets 0 and + // BLOCK_M to match math's `+ 0 * BLOCK_M` / `+ 1 * BLOCK_M` reads. + tma::copy( + &tensor_map_shared_l2_acts_sf, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx * 2, 1); + tma::copy( + &tensor_map_shared_l2_acts_sf, full_barriers[stage_idx], + smem_sfa[stage_idx] + BLOCK_M, m_idx, k_block_idx * 2 + 1, 1); + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + 2 * BLOCK_M * sizeof(float)); + } + } else if (is_l1) { // L1 SFA per-128: load (BLOCK_M, 1) at K=k_block_idx tma::copy( tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], @@ -838,49 +985,36 @@ sm90_fp8_mega_moe_impl(void* y, } }; - if constexpr (kSplitPhaseHotPath) { - sm90_fp8_mega_moe_for_each_block_split( - scheduler, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_a_sfa_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_a_sfa_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); - } else { - scheduler.for_each_block([&](const sched::BlockPhase& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_a_sfa_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); + typename SchedulerT::task_info_t task_info; + while (scheduler.get_next_task(task_info)) { + process_a_sfa_block(task_info.block_phase, task_info.local_expert_idx, + get_num_k_blocks(task_info.block_phase), + task_info.m_block_idx, task_info.n_block_idx, + task_info.valid_m, task_info.pool_block_idx); } + } else if (warp_idx == kNumDispatchWarps + 1) { cutlass::arch::warpgroup_reg_dealloc(); - scheduler.for_each_block([&](const sched::BlockPhase& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - const auto tensor_map_b_ptr = - block_phase == sched::BlockPhase::Linear2 ? &tensor_map_l2_weights : &tensor_map_l1_weights; - - const uint32_t shape_n = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_N : L1_SHAPE_N; + auto process_b_block = [&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& n_block_idx, + const uint32_t& shape_n) { + const bool is_shared = kHasSharedExperts and sched::is_shared_phase(block_phase); + const bool is_l1 = sched::is_l1_phase(block_phase); + const auto tensor_map_b_ptr = is_shared + ? (is_l1 ? &tensor_map_shared_l1_weights : &tensor_map_shared_l2_weights) + : (is_l1 ? &tensor_map_l1_weights : &tensor_map_l2_weights); for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { empty_barriers[stage_idx]->wait(phase ^ 1); if (cute::elect_one_sync()) { - const uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + // The fused shared expert is a single dense MLP: no expert stride + const uint32_t n_idx = (is_shared ? 0u : local_expert_idx * shape_n) + + n_block_idx * BLOCK_N; const uint32_t k_idx = k_block_idx * BLOCK_K; // TMA load B (weight SF is now loaded directly by math warps from global) @@ -904,12 +1038,29 @@ sm90_fp8_mega_moe_impl(void* y, } __syncwarp(); } - }); + }; + + typename SchedulerT::task_info_t task_info; + while (scheduler.get_next_task(task_info)) { + process_b_block(task_info.block_phase, task_info.local_expert_idx, + get_num_k_blocks(task_info.block_phase), + task_info.n_block_idx, task_info.shape_n); + } + + + } else if (warp_idx == kNumDispatchWarps + 2) { + // Producer warp: claims routed/shared L1/L2 tasks from the global atomic + // counters and publishes them into the SMEM task ring + cutlass::arch::warpgroup_reg_dealloc(); + + scheduler.mainloop(num_tokens); } else if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) { - // Idle non-epilogue warps (kNumDispatchWarps+2, +3). They must still - // participate in the warpgroup-collective `setmaxnreg.dec.sync.aligned` - // so that the math warpgroup's `warpgroup_reg_alloc` can succeed. + // Idle/padding non-epilogue warps: none exist in the 32 + 96 topology the + // host selects (exactly TMA-A + TMA-B + producer). They must still take + // part in the warpgroup-collective `setmaxnreg.dec.sync.aligned` so the + // math warpgroup's `warpgroup_reg_alloc` succeeds if a wider + // non-epilogue section is ever configured. cutlass::arch::warpgroup_reg_dealloc(); } else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { @@ -943,9 +1094,8 @@ sm90_fp8_mega_moe_impl(void* y, auto process_math_block = [&](const auto& block_phase, const uint32_t& local_expert_idx, const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - const uint32_t valid_m = scheduler.template get_valid_m(); - const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t& m_block_idx, const uint32_t& n_block_idx, + const uint32_t& valid_m, const uint32_t& pool_block_idx) { const uint32_t m_idx = pool_block_idx * BLOCK_M; const uint32_t n_idx = n_block_idx * BLOCK_N; const uint32_t epilogue_wg_m_idx = epilogue_wg_idx / kWarpgroupSplitN; @@ -968,10 +1118,69 @@ sm90_fp8_mega_moe_impl(void* y, const bool valid_r0 = row_offset_r0 < valid_m; const bool valid_r1 = row_offset_r1 < valid_m; + // Fused shared expert: a single dense MLP over this rank's own tokens. + // `is_shared` folds to a compile-time false when the feature is off. + const bool is_shared = kHasSharedExperts and sched::is_shared_phase(block_phase); + const bool is_l1 = sched::is_l1_phase(block_phase); + // Shared tiles always publish their L1 arrivals through a counter and never + // need the CTA-wide L2 epilogue sync; the routed modes are compile-time. + const bool use_arrival_counter = kL2ArrivalCounter or is_shared; + const bool needs_l2_full_sync = kL2EpilogueRequiresFullSync and not is_shared; + + // Activation SF for the current K block, read from the TMA/gather-staged SMEM + // tile. All four phases share one layout: L1-like phases put one per-128-K + // float per row at offset 0; L2-like phases put the two per-64-K groups at + // offsets 0 and BLOCK_M (`_hi` is only written for those). Shared tiles are + // staged by the same loader warp (`process_a_sfa_block`), so no phase test is + // needed here. + auto load_act_sf = [&](float& lo_0, float& lo_1, float& hi_0, float& hi_1) { + if (is_l1) { + lo_0 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + lo_1 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + lo_0 = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + lo_1 = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + hi_0 = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + hi_1 = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + }; + + // Block (128, 128) weight SF lookups. The shared expert has no per-expert + // stride and its gate/up N halves are `kSharedIntermediateHidden` apart. + auto load_l1_weight_sf = [&](const uint32_t& k_block_idx, float& gate_sf, float& up_sf) { + constexpr uint32_t kSFKBlocks = kHidden / 128; + const uint32_t gate_n = sf_n_block_idx / 2u; + if (is_shared) { + constexpr uint32_t kSFGateBlks = kSharedIntermediateHidden / 128; + const float* base = shared_l1_weights_sf + k_block_idx; + gate_sf = __ldg(base + gate_n * kSFKBlocks); + up_sf = __ldg(base + (kSFGateBlks + gate_n) * kSFKBlocks); + } else { + constexpr uint32_t kSFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kSFPerExpert = (kIntermediateHidden * 2 / 128) * kSFKBlocks; + const float* base = l1_weights_sf + local_expert_idx * kSFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kSFKBlocks); + up_sf = __ldg(base + (kSFGateBlks + gate_n) * kSFKBlocks); + } + }; + auto load_l2_weight_sf = [&](const uint32_t& k_block_idx) -> float { + if (is_shared) { + constexpr uint32_t kSFKBlocks = kSharedIntermediateHidden / 128; + return __ldg(shared_l2_weights_sf + sf_n_block_idx * kSFKBlocks + k_block_idx); + } + constexpr uint32_t kSFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kSFPerExpert = (kHidden / 128) * kSFKBlocks; + return __ldg(l2_weights_sf + local_expert_idx * kSFPerExpert + + sf_n_block_idx * kSFKBlocks + k_block_idx); + }; + // ---------------- GEMM ---------------- using WGMMA = L1WGMMA; constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; + constexpr uint32_t kSwapSlabAccumStride = kSwapABActive ? SwapWGMMA64::kNumAccum : 0; + constexpr uint32_t kScratchAccumPerThread = kSwapABActive ? kSwapSlabAccumStride : kAccumPerThread; float final_accum[kAccumPerThread] = {}; + float accum[kScratchAccumPerThread]; if constexpr (kReuseAccumAsFinal) { auto prescale_l1_final = [&](const float& scale_a_0, const float& scale_a_1, @@ -1089,34 +1298,15 @@ sm90_fp8_mega_moe_impl(void* y, float scale_a_0_lo, scale_a_1_lo; float scale_a_0_hi, scale_a_1_hi; - if (block_phase == sched::BlockPhase::Linear1) { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); - } else { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); - scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); - scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); - } + load_act_sf(scale_a_0_lo, scale_a_1_lo, scale_a_0_hi, scale_a_1_hi); - constexpr uint32_t kL1SFKBlocks = kHidden / 128; - constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; - constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; - if (block_phase == sched::BlockPhase::Linear1) { - const uint32_t gate_n = sf_n_block_idx / 2u; - const uint32_t up_n = kL1SFGateBlks + gate_n; - const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; - gate_sf = __ldg(base + gate_n * kL1SFKBlocks); - up_sf = __ldg(base + up_n * kL1SFKBlocks); - } else { - l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert - + sf_n_block_idx * kL2SFKBlocks + k_block_idx); - } + if (is_l1) + load_l1_weight_sf(k_block_idx, gate_sf, up_sf); + else + l2_sf = load_l2_weight_sf(k_block_idx); - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if (k_block_idx != 0) rescale_l1_final(prev_scale_a_0, prev_scale_a_1, prev_gate_sf, prev_up_sf, @@ -1197,7 +1387,7 @@ sm90_fp8_mega_moe_impl(void* y, } if (num_k_blocks != 0) { - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { postscale_l1_final(prev_scale_a_0, prev_scale_a_1, prev_gate_sf, prev_up_sf); } else { @@ -1210,34 +1400,15 @@ sm90_fp8_mega_moe_impl(void* y, float scale_a_0_lo, scale_a_1_lo; float scale_a_0_hi, scale_a_1_hi; - if (block_phase == sched::BlockPhase::Linear1) { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); - } else { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); - scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); - scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); - } + load_act_sf(scale_a_0_lo, scale_a_1_lo, scale_a_0_hi, scale_a_1_hi); - constexpr uint32_t kL1SFKBlocks = kHidden / 128; - constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; - constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; - if (block_phase == sched::BlockPhase::Linear1) { - const uint32_t gate_n = sf_n_block_idx / 2u; - const uint32_t up_n = kL1SFGateBlks + gate_n; - const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; - gate_sf = __ldg(base + gate_n * kL1SFKBlocks); - up_sf = __ldg(base + up_n * kL1SFKBlocks); - } else { - l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert - + sf_n_block_idx * kL2SFKBlocks + k_block_idx); - } + if (is_l1) + load_l1_weight_sf(k_block_idx, gate_sf, up_sf); + else + l2_sf = load_l2_weight_sf(k_block_idx); - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if (k_block_idx != 0) prescale_l1_final(scale_a_0_lo, scale_a_1_lo, gate_sf, up_sf); @@ -1317,16 +1488,7 @@ sm90_fp8_mega_moe_impl(void* y, // Read SF (must precede warpgroup_arrive) float scale_a_0_lo, scale_a_1_lo; float scale_a_0_hi, scale_a_1_hi; // Only used in L2 (per-64 K) - if (block_phase == sched::BlockPhase::Linear1) { - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); - } else { - // L2: SFA layout is (K=2, M=BLOCK_M) MN-major; first half SF at offset 0, second at BLOCK_M - scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); - scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); - scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); - scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); - } + load_act_sf(scale_a_0_lo, scale_a_1_lo, scale_a_0_hi, scale_a_1_hi); // ----- Block (128, 128) weight SF (loaded directly from global) ----- // L1 weight SF shape: (E, 2*IH/128, H/128) MN-major. The N axis is @@ -1340,63 +1502,63 @@ sm90_fp8_mega_moe_impl(void* y, // logical 128x128 weight-SF tile, broadcast across the matching // WGMMA accumulators. // + // The fused shared expert uses the same layouts without the expert + // dimension (see `load_l1_weight_sf` / `load_l2_weight_sf`). + // // Load the weight scale after the barrier from all WG threads. // This keeps scale loads close to their WGMMA use and lets the // read-only cache coalesce the same-address accesses. - constexpr uint32_t kL1SFKBlocks = kHidden / 128; - constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; - constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; - constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; - if (block_phase == sched::BlockPhase::Linear1) { - const uint32_t gate_n = sf_n_block_idx / 2u; - const uint32_t up_n = kL1SFGateBlks + gate_n; - const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; - gate_sf = __ldg(base + gate_n * kL1SFKBlocks); - up_sf = __ldg(base + up_n * kL1SFKBlocks); - } else { - l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert - + sf_n_block_idx * kL2SFKBlocks + k_block_idx); - } + if (is_l1) + load_l1_weight_sf(k_block_idx, gate_sf, up_sf); + else + l2_sf = load_l2_weight_sf(k_block_idx); - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if constexpr (kSwapABActive) { auto run_swap_ab_l1 = [&]() { using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; - float swap_accum[kSwapAccum]; - - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++ k) { - auto desc_a = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k * SwapWGMMA::K, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_a[stage_idx] + k * SwapWGMMA::K, 1); - SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_wait<0>(); #pragma unroll - for (uint32_t i = 0; i < kSwapAccum / 4; ++ i) { - const uint32_t token_0 = i * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - const float scale_0 = token_0 < valid_m ? - ptx::ld_shared(smem_sfa[stage_idx] + token_0) : 0.0f; - const float scale_1 = token_1 < valid_m ? - ptx::ld_shared(smem_sfa[stage_idx] + token_1) : 0.0f; - final_accum[i * 4 + 0] += scale_0 * gate_sf * swap_accum[i * 4 + 0]; - final_accum[i * 4 + 2] += scale_0 * up_sf * swap_accum[i * 4 + 2]; - final_accum[i * 4 + 1] += scale_1 * gate_sf * swap_accum[i * 4 + 1]; - final_accum[i * 4 + 3] += scale_1 * up_sf * swap_accum[i * 4 + 3]; + for (uint32_t slab = 0; slab < 2; ++slab) { + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++k) { + const uint32_t slab_b_off = + smem_b_wg_offset + slab * SwapWGMMA::M * BLOCK_K; + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + slab_b_off + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + const uint32_t final_base = slab * kSwapSlabAccumStride; + // Routed and shared tiles both read the activation SF from + // `smem_sfa` (the producer stages the shared column there, + // see `process_a_sfa_block`), so this loop is phase-agnostic. + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + const float scale_0 = token_0 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + token_0) : 0.0f; + const float scale_1 = token_1 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + token_1) : 0.0f; + final_accum[final_base + i * 4 + 0] += scale_0 * gate_sf * accum[i * 4 + 0]; + final_accum[final_base + i * 4 + 1] += scale_1 * gate_sf * accum[i * 4 + 1]; + final_accum[final_base + i * 4 + 2] += scale_0 * up_sf * accum[i * 4 + 2]; + final_accum[final_base + i * 4 + 3] += scale_1 * up_sf * accum[i * 4 + 3]; + } } if (lane_idx == 0) @@ -1466,62 +1628,69 @@ sm90_fp8_mega_moe_impl(void* y, auto run_swap_ab_l2 = [&]() { using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; - float swap_accum[kSwapAccum]; - auto promote_swap_accum = [&](const uint32_t& sf_group) { + auto promote_swap_accum = [&](const uint32_t& sf_group, const uint32_t& final_base) { + // SFA layout is (K=2, M=BLOCK_M) M-major in `smem_sfa`, `lo`/`hi` + // slabs at `sf_group * BLOCK_M + ...`. Routed and shared tiles are + // both TMA-staged there, so this loop is phase-agnostic. #pragma unroll - for (uint32_t i = 0; i < kSwapAccum / 4; ++ i) { + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { const uint32_t token_0 = i * 8 + col_idx * 2; const uint32_t token_1 = token_0 + 1; const float scale_0 = token_0 < valid_m ? ptx::ld_shared(smem_sfa[stage_idx] + sf_group * BLOCK_M + token_0) : 0.0f; const float scale_1 = token_1 < valid_m ? ptx::ld_shared(smem_sfa[stage_idx] + sf_group * BLOCK_M + token_1) : 0.0f; - final_accum[i * 4 + 0] += scale_0 * l2_sf * swap_accum[i * 4 + 0]; - final_accum[i * 4 + 2] += scale_0 * l2_sf * swap_accum[i * 4 + 2]; - final_accum[i * 4 + 1] += scale_1 * l2_sf * swap_accum[i * 4 + 1]; - final_accum[i * 4 + 3] += scale_1 * l2_sf * swap_accum[i * 4 + 3]; + final_accum[final_base + i * 4 + 0] += scale_0 * l2_sf * accum[i * 4 + 0]; + final_accum[final_base + i * 4 + 1] += scale_1 * l2_sf * accum[i * 4 + 1]; + final_accum[final_base + i * 4 + 2] += scale_0 * l2_sf * accum[i * 4 + 2]; + final_accum[final_base + i * 4 + 3] += scale_1 * l2_sf * accum[i * 4 + 3]; } }; #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++ k) { - auto desc_a = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k * SwapWGMMA::K, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_a[stage_idx] + k * SwapWGMMA::K, 1); - SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_wait<0>(); - promote_swap_accum(0); + for (uint32_t slab = 0; slab < 2; ++slab) { + const uint32_t slab_b_off = smem_b_wg_offset + slab * SwapWGMMA::M * BLOCK_K; + const uint32_t final_base = slab * kSwapSlabAccumStride; - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++ k) { - const uint32_t k_off = (BLOCK_K / 2) + k * SwapWGMMA::K; - auto desc_a = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_a[stage_idx] + k_off, 1); - SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + slab_b_off + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + promote_swap_accum(0, final_base); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++k) { + const uint32_t k_off = (BLOCK_K / 2) + k * SwapWGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + slab_b_off + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k_off, 1); + SwapWGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + promote_swap_accum(1, final_base); } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kSwapAccum; ++ i) - ptx::warpgroup_fence_operand(swap_accum[i]); - ptx::warpgroup_wait<0>(); - promote_swap_accum(1); if (lane_idx == 0) empty_barriers[stage_idx]->arrive(); @@ -1612,19 +1781,21 @@ sm90_fp8_mega_moe_impl(void* y, } } - // Skip epilogue when block is past valid M (still must release via empty) + // Skip epilogue when block is past valid M (still must release via empty). + // The taken paths below must issue the same aligned syncs, so both + // conditions are per-task (shared tiles always use counter arrivals). if (row_base >= valid_m) { - if (block_phase == sched::BlockPhase::Linear1) { - if constexpr (not kL2ArrivalCounter) + if (is_l1) { + if (not use_arrival_counter) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } else { - if constexpr (kL2EpilogueRequiresFullSync) + if (needs_l2_full_sync) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } return; } - if (block_phase == sched::BlockPhase::Linear1) { + if (is_l1) { if constexpr (kSwapABActive) { auto silu = [](float x) -> float { const float e = kFastMath ? __expf(-x) : expf(-x); @@ -1640,78 +1811,88 @@ sm90_fp8_mega_moe_impl(void* y, x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); }; - const uint32_t out_col_base = - wg_l1_out_n_offset + warp_idx_in_wg * 8 + row_idx; - auto store_l1_swap_chunk = [&](const uint32_t& i) { - const uint32_t token_0 = i * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - if (token_0 < valid_m) { - float g0 = final_accum[i * 4 + 0]; - float u0 = final_accum[i * 4 + 2]; - clamp_gate(g0); - clamp_up(u0); - const float weight_0 = *l1_topk_weights_buffer - .get_data_buffer(m_idx + token_0) - .get_base_ptr(); - smem_cd_swap_l1_fp32[token_0 * L1_OUT_BLOCK_N + out_col_base] = - silu(g0) * u0 * weight_0; - } - if (token_1 < valid_m) { - float g1 = final_accum[i * 4 + 1]; - float u1 = final_accum[i * 4 + 3]; - clamp_gate(g1); - clamp_up(u1); - const float weight_1 = *l1_topk_weights_buffer - .get_data_buffer(m_idx + token_1) - .get_base_ptr(); - smem_cd_swap_l1_fp32[token_1 * L1_OUT_BLOCK_N + out_col_base] = - silu(g1) * u1 * weight_1; - } - }; + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + const uint32_t num_swap_chunks = n_swap / 8u; + const uint32_t wg_out_col_base = epilogue_wg_n_idx * (WG_BLOCK_N / 2); - const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; - store_l1_swap_chunk(0); - if (valid_m > 8) { + #pragma unroll + for (uint32_t slab = 0; slab < 2; ++slab) { + const uint32_t final_base = slab * kSwapSlabAccumStride; + const uint32_t out_col = wg_out_col_base + slab * 32u + + warp_idx_in_wg * 8u + row_idx; #pragma unroll - for (uint32_t i = 1; i < kSwapABTokenChunks; ++ i) { - if (i < num_swap_token_chunks) - store_l1_swap_chunk(i); + for (uint32_t i = 0; i < SwapWGMMA64::kNumAccum / 4; ++i) { + if (i >= num_swap_chunks) break; + const uint32_t token_0 = i * 8u + col_idx * 2u; + const uint32_t token_1 = token_0 + 1u; + if (token_0 < valid_m) { + float g0 = final_accum[final_base + i * 4 + 0]; + float u0 = final_accum[final_base + i * 4 + 2]; + clamp_gate(g0); + clamp_up(u0); + const float weight_0 = is_shared ? 1.0f : *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_0) + .template get_base_ptr(); + smem_cd_swap_l1_fp32[token_0 * L1_OUT_BLOCK_N + out_col] = + silu(g0) * u0 * weight_0; + } + if (token_1 < valid_m) { + float g1 = final_accum[final_base + i * 4 + 1]; + float u1 = final_accum[final_base + i * 4 + 3]; + clamp_gate(g1); + clamp_up(u1); + const float weight_1 = is_shared ? 1.0f : *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_1) + .template get_base_ptr(); + smem_cd_swap_l1_fp32[token_1 * L1_OUT_BLOCK_N + out_col] = + silu(g1) * u1 * weight_1; + } } } ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); for (uint32_t token = epilogue_thread_idx; token < valid_m; token += kNumEpilogueThreads) { - float amax = 0.0f; + constexpr uint32_t kHalfN = L1_OUT_BLOCK_N / 2; + float amax0 = 0.0f, amax1 = 0.0f; #pragma unroll - for (uint32_t col = 0; col < L1_OUT_BLOCK_N; ++ col) { + for (uint32_t col = 0; col < L1_OUT_BLOCK_N; ++col) { const float v = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col]; - amax = cute::max(amax, cute::abs(v)); + const float a = cute::abs(v); + if (col < kHalfN) amax0 = cute::max(amax0, a); + else amax1 = cute::max(amax1, a); } - float2 amax_pair = {amax, amax}; - float2 sf_pair, sf_inv_pair; - sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv(amax_pair, sf_pair, sf_inv_pair); - const float sf = sf_pair.x; - const float sf_inv = sf_inv_pair.x; - auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); - // ROOT-CAUSE FIX: the L2-activation SF pool is strided by SF_BLOCK_M - // (=align(BLOCK_M,128)=128), which is how the L2 producer reads it - // (sfa_m_idx = pool_block_idx * SF_BLOCK_M) and how the non-swap L1 - // writes it. This swapAB path used BLOCK_M (64), so for pool_block_idx>=1 - // the SF landed in the wrong rows -> L2 read stale SF -> every pool block - // after the first was corrupted (block 0 was correct because 0*64==0*128). - const uint32_t token_idx = pool_block_idx * SF_BLOCK_M + token; - sf_base_ptr[n_block_idx * kNumPaddedSFPoolTokens + token_idx] = sf; + float2 sf_pair, sf_inv_pair; + sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( + make_float2(amax0, amax1), sf_pair, sf_inv_pair); + const float sf0 = sf_pair.x, sf1 = sf_pair.y; + const float sf_inv0 = sf_inv_pair.x, sf_inv1 = sf_inv_pair.y; + // Shared L2 activation SF (written here by the fused L1) is now M-major + // (token-contiguous inside each `k_sf_*` slab, mirroring the routed L2 SF + // pool) so producer can TMA-load (BLOCK_M, 1) tiles directly into smem_sfa. + const uint32_t k_sf_lo = n_block_idx * 2u + 0u; + const uint32_t k_sf_hi = n_block_idx * 2u + 1u; + if (is_shared) { + auto sf_base_ptr = shared_l2_sf_buffer.get_base_ptr(); + const uint32_t row = m_idx + token; + sf_base_ptr[k_sf_lo * kNumMaxTokensPerRank + row] = sf0; + sf_base_ptr[k_sf_hi * kNumMaxTokensPerRank + row] = sf1; + } else { + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + const uint32_t token_idx = pool_block_idx * SF_BLOCK_M + token; + sf_base_ptr[k_sf_lo * kNumPaddedSFPoolTokens + token_idx] = sf0; + sf_base_ptr[k_sf_hi * kNumPaddedSFPoolTokens + token_idx] = sf1; + } #pragma unroll for (uint32_t col = 0; col < L1_OUT_BLOCK_N; col += 2) { - const float v0 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col + 0] * sf_inv; + const float sf_inv = (col < kHalfN) ? sf_inv0 : sf_inv1; + const float v0 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col] * sf_inv; const float v1 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col + 1] * sf_inv; const __nv_fp8x2_e4m3 pair(make_float2(v0, v1)); - auto* ptr = reinterpret_cast( - smem_cd_swap_l1_fp8 + token * L1_OUT_BLOCK_N + col); - *ptr = pair.__x; + *reinterpret_cast( + smem_cd_swap_l1_fp8 + token * L1_OUT_BLOCK_N + col) = pair.__x; } } @@ -1720,7 +1901,7 @@ sm90_fp8_mega_moe_impl(void* y, if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( - &tensor_map_l1_output, + is_shared ? &tensor_map_shared_l1_output : &tensor_map_l1_output, smem_cd_swap_l1_fp8, n_block_idx * L1_OUT_BLOCK_N, m_idx); @@ -1729,12 +1910,16 @@ sm90_fp8_mega_moe_impl(void* y, __syncwarp(); ptx::tma_store_wait<0>(); - if constexpr (kL2ArrivalCounter) { - if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { - ptx::red_add_rel( - reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), - kWarpgroupSplitN); - } + // `use_arrival_counter` is true for shared tiles and for the routed + // counter mode; select the per-M-block arrival slot and publish the + // whole N-split warpgroup group from the storing WG (same shape as the + // non-swap L1 epilogue). The bitmask path otherwise stays untouched. + if (use_arrival_counter) { + auto arrival_ptr = is_shared + ? workspace.get_shared_l2_full_count_ptr(pool_block_idx) + : reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)); + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) + ptx::red_add_rel(arrival_ptr, kWarpgroupSplitN); } else { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { @@ -1744,7 +1929,9 @@ sm90_fp8_mega_moe_impl(void* y, } } __syncwarp(); - if constexpr (kL2ArrivalCounter) + // counter mode (incl. shared) needs the CTA-wide sync to protect + // the swapAB FP32/FP8 staging tiles from the next task's overwrite. + if (use_arrival_counter) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } else { @@ -1818,13 +2005,14 @@ sm90_fp8_mega_moe_impl(void* y, } } - // Apply token weight: SwiGLU * topk_weight (single load per row) - const float weight_r0 = valid_r0 ? *l1_topk_weights_buffer + // Apply token weight: SwiGLU * topk_weight (single load per row). + // The shared expert sees every token with weight 1.0. + const float weight_r0 = valid_r0 ? (is_shared ? 1.0f : *l1_topk_weights_buffer .get_data_buffer(m_idx + row_offset_r0) - .get_base_ptr() : 0.0f; - const float weight_r1 = valid_r1 ? *l1_topk_weights_buffer + .template get_base_ptr()) : 0.0f; + const float weight_r1 = valid_r1 ? (is_shared ? 1.0f : *l1_topk_weights_buffer .get_data_buffer(m_idx + row_offset_r1) - .get_base_ptr() : 0.0f; + .template get_base_ptr()) : 0.0f; #pragma unroll for (uint32_t p = 0; p < kNumPairs; ++ p) { swiglu_r0[p][0] *= weight_r0; @@ -1915,16 +2103,27 @@ sm90_fp8_mega_moe_impl(void* y, // In the shared-SF split both warpgroups own the same per-64 group and rows, so // only the first N-split warpgroup publishes the SF slot to avoid a write race. if (col_idx == 0 and (not kSplitNSharesSF or epilogue_wg_n_idx == 0)) { - auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); - // SF buffer is (kNumPaddedSFPoolTokens x kIntermediateHidden/64), MN-major: - // addr[k_idx * num_padded_sf_pool_tokens + token_idx] - const uint32_t token_r0 = pool_block_idx * SF_BLOCK_M + row_offset_r0; - const uint32_t token_r1 = pool_block_idx * SF_BLOCK_M + row_offset_r1; const uint32_t k_sf_idx = sf_n_block_idx; // one per-64 post-SwiGLU group - if (valid_r0) - sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0; - if (valid_r1) - sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1; + if (is_shared) { + // Shared SF buffer is (kNumMaxTokensPerRank x SIH/64) M-major and + // indexed by the local token (token-contiguous inner, stride 1), so no + // SF-pool padding applies: addr[k_idx * kNumMaxTokensPerRank + token_idx] + auto sf_base_ptr = shared_l2_sf_buffer.get_base_ptr(); + if (valid_r0) + sf_base_ptr[k_sf_idx * kNumMaxTokensPerRank + (m_idx + row_offset_r0)] = sf_r0; + if (valid_r1) + sf_base_ptr[k_sf_idx * kNumMaxTokensPerRank + (m_idx + row_offset_r1)] = sf_r1; + } else { + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + // SF buffer is (kNumPaddedSFPoolTokens x kIntermediateHidden/64), MN-major: + // addr[k_idx * num_padded_sf_pool_tokens + token_idx] + const uint32_t token_r0 = pool_block_idx * SF_BLOCK_M + row_offset_r0; + const uint32_t token_r1 = pool_block_idx * SF_BLOCK_M + row_offset_r1; + if (valid_r0) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0; + if (valid_r1) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1; + } } // Sync the warpgroup before TMA store. In the shared-tile split @@ -1951,7 +2150,7 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( - &tensor_map_l1_output, + is_shared ? &tensor_map_shared_l1_output : &tensor_map_l1_output, smem_cd_l1, out_n_idx, m_idx + row_base); @@ -1962,7 +2161,7 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N + wg_l1_out_n_offset; cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( - &tensor_map_l1_output, + is_shared ? &tensor_map_shared_l1_output : &tensor_map_l1_output, smem_cd_l1 + smem_cd_l1_wg_offset, out_n_idx, m_idx + row_base); @@ -1974,20 +2173,20 @@ sm90_fp8_mega_moe_impl(void* y, // Notify L2 that this L1 output (and SF) is ready. Counter mode lets // independent WG tiles publish arrivals without the CTA-wide barrier - // needed before the single bit-mask update. - if constexpr (kL2ArrivalCounter) { + // needed before the single bit-mask update. Shared tiles always use a + // counter, on their own per-M-block slot. + if (use_arrival_counter) { + auto arrival_ptr = is_shared + ? workspace.get_shared_l2_full_count_ptr(pool_block_idx) + : reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)); if constexpr (kSplitNSharesSF) { // The combined tile counts for both N-split warpgroups; the // storing warpgroup publishes all kWarpgroupSplitN arrivals // after its TMA store has drained. - if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { - ptx::red_add_rel( - reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), - kWarpgroupSplitN); - } + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) + ptx::red_add_rel(arrival_ptr, kWarpgroupSplitN); } else if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { - ptx::red_add_rel( - reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), 1); + ptx::red_add_rel(arrival_ptr, 1); } } else { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); @@ -2013,32 +2212,51 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t lane_in_row = lane_idx % 16; const uint32_t cols_per_lane = WG_BLOCK_N / 16; - if constexpr (kSwapABActive) { - auto store_bf16 = [&](const uint32_t& token, const uint32_t& col, float value) { - smem_cd_l2[smem_cd_l2_wg_offset + token * WG_BLOCK_N + col] = - __float2bfloat16_rn(value); - }; - - auto store_l2_swap_chunk = [&](const uint32_t& i) { - const uint32_t token_0 = i * 8 + col_idx * 2; - const uint32_t token_1 = token_0 + 1; - if (token_0 < valid_m) { - store_bf16(token_0, r_0, final_accum[i * 4 + 0]); - store_bf16(token_0, r_1, final_accum[i * 4 + 2]); - } - if (token_1 < valid_m) { - store_bf16(token_1, r_0, final_accum[i * 4 + 1]); - store_bf16(token_1, r_1, final_accum[i * 4 + 3]); - } - }; + // XOR column swizzle (8-col granularity) for the row-major BF16 + // staging tile. The row stride is WG_BLOCK_N/2 banks, which is a + // multiple of 32 for every supported WG_BLOCK_N (64/128), so 8 + // lanes that share a col_idx (8 distinct row_idx) all hit the same + // bank -> 8-way conflict on each STS. XORing bits [3:5] of the + // column with (row & 7) spreads those 8 rows across 8 distinct + // banks. The swizzle MUST be applied on both the STS write and the + // LDS scatter read so the permutation cancels out; doing it on the + // write alone (as a port of SM100's layout) silently permutes the + // output columns -- SM100's swizzle is enforced by its TMA + // descriptor, this manual STS/LDS path has no such contract. The + // 8-col granularity is safe: the 2-BF16 STS pair and the + // cols_per_lane-BF16 LDS vector (4 or 8 BF16, the only sizes this + // path supports) never straddle an 8-col block, so the key is + // constant across each access. + auto swiz_col = [](uint32_t row, uint32_t col) -> uint32_t { + return col ^ ((row & 7) << 3); + }; - const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; - store_l2_swap_chunk(0); - if (valid_m > 8) { + if constexpr (kSwapABActive) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + const uint32_t num_swap_chunks = n_swap / 8u; + #pragma unroll + for (uint32_t slab = 0; slab < 2; ++slab) { + const uint32_t final_base = slab * kSwapSlabAccumStride; + const uint32_t slab_col_base = wg_n_offset + slab * SwapWGMMA64::M; #pragma unroll - for (uint32_t i = 1; i < kSwapABTokenChunks; ++ i) { - if (i < num_swap_token_chunks) - store_l2_swap_chunk(i); + for (uint32_t i = 0; i < SwapWGMMA64::kNumAccum / 4; ++i) { + if (i >= num_swap_chunks) break; + const uint32_t token_0 = i * 8u + col_idx * 2u; + const uint32_t token_1 = token_0 + 1u; + const uint32_t col_0 = slab_col_base + r_0; + const uint32_t col_1 = slab_col_base + r_1; + if (token_0 < valid_m) { + smem_cd_l2[token_0 * BLOCK_N + swiz_col(token_0, col_0)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 0]); + smem_cd_l2[token_0 * BLOCK_N + swiz_col(token_0, col_1)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 2]); + } + if (token_1 < valid_m) { + smem_cd_l2[token_1 * BLOCK_N + swiz_col(token_1, col_0)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 1]); + smem_cd_l2[token_1 * BLOCK_N + swiz_col(token_1, col_1)] = + __float2bfloat16_rn(final_accum[final_base + i * 4 + 3]); + } } } } else { @@ -2055,7 +2273,7 @@ sm90_fp8_mega_moe_impl(void* y, auto smem_ptr = smem_cd_l2 + smem_cd_l2_wg_offset + row * WG_BLOCK_N - + col; + + swiz_col(row, col); // BF16 STS: 2 bf16 elements *reinterpret_cast(smem_ptr) = packed; }; @@ -2100,17 +2318,35 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t m_idx_in_block = row_base + row_in_wg; if (m_idx_in_block >= valid_m) break; - // Read cols_per_lane BF16 (= one ScatterVec) from smem - auto smem_ptr = smem_cd_l2 - + smem_cd_l2_wg_offset - + row_in_wg * WG_BLOCK_N - + lane_in_row * cols_per_lane; + // Read cols_per_lane BF16 (= one ScatterVec) from smem. + // swapAB uses a shared BLOCK_N-wide tile with column swizzle; + // non-swap uses per-WG slab with same swizzle. + nv_bfloat16* smem_ptr; + if constexpr (kSwapABActive) { + const uint32_t base_col = wg_n_offset + lane_in_row * cols_per_lane; + const uint32_t read_col = swiz_col(row_in_wg, base_col); + smem_ptr = smem_cd_l2 + row_in_wg * BLOCK_N + read_col; + } else { + uint32_t read_col = lane_in_row * cols_per_lane; + read_col = swiz_col(row_in_wg, read_col); + smem_ptr = smem_cd_l2 + smem_cd_l2_wg_offset + + row_in_wg * WG_BLOCK_N + read_col; + } const auto packed = *reinterpret_cast(smem_ptr); - const auto src_metadata = *workspace.get_token_src_metadata_ptr(m_idx + m_idx_in_block); - const uint32_t dst_rank_idx = src_metadata.rank_idx; - const uint32_t dst_token_idx = src_metadata.token_idx; - const uint32_t dst_topk_idx = src_metadata.topk_idx; + // The fused shared expert stays on this rank and reduces through the + // extra combine slot `kNumTopk`, keyed by the local token index. + uint32_t dst_rank_idx, dst_token_idx, dst_topk_idx; + if (is_shared) { + dst_rank_idx = sym_buffer.rank_idx; + dst_token_idx = m_idx + m_idx_in_block; + dst_topk_idx = kNumTopk; + } else { + const auto src_metadata = *workspace.get_token_src_metadata_ptr(m_idx + m_idx_in_block); + dst_rank_idx = src_metadata.rank_idx; + dst_token_idx = src_metadata.token_idx; + dst_topk_idx = src_metadata.topk_idx; + } const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) .get_data_buffer(dst_token_idx); auto dst_ptr = math::advance_ptr( @@ -2119,36 +2355,19 @@ sm90_fp8_mega_moe_impl(void* y, *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; } - if constexpr (kL2EpilogueRequiresFullSync) + if (needs_l2_full_sync) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } }; - if constexpr (kSplitPhaseHotPath) { - sm90_fp8_mega_moe_for_each_block_split( - scheduler, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_math_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }, - [&](const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_math_block( - std::integral_constant{}, - local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); - } else { - scheduler.for_each_block([&](const sched::BlockPhase& block_phase, - const uint32_t& local_expert_idx, - const uint32_t& num_k_blocks, - const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - process_math_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); - }); + typename SchedulerT::task_info_t task_info; + while (scheduler.get_next_task(task_info)) { + process_math_block(task_info.block_phase, task_info.local_expert_idx, + get_num_k_blocks(task_info.block_phase), + task_info.m_block_idx, task_info.n_block_idx, + task_info.valid_m, task_info.pool_block_idx); } + // ---------------- COMBINE ---------------- // NVLink barrier first: signals remote ranks that this rank's GEMM @@ -2187,7 +2406,7 @@ sm90_fp8_mega_moe_impl(void* y, DG_STATIC_ASSERT(kNumChunkBytes % 16 == 0, "Combine chunk must be TMA-aligned (16 bytes)"); DG_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, "Combine chunk must be divisible by 16 bytes"); DG_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, "Combine chunk must be a multiple of 32 16-byte elements"); - DG_STATIC_ASSERT(kNumTopk <= 32, "Top-k must fit in a single warp"); + DG_STATIC_ASSERT(kNumCombineSlots <= 32, "Top-k (plus the shared slot) must fit in a single warp"); DG_TRAP_ONLY_DEVICE_ASSERT(kNumChunkSlots * kNumCombineWarps * kNumChunkBytes <= static_cast( reinterpret_cast(barrier_start_ptr) - smem_buffer)); @@ -2207,8 +2426,12 @@ sm90_fp8_mega_moe_impl(void* y, for (uint32_t token_idx = sm_idx * kNumCombineWarps + epilogue_warp_idx; token_idx < num_tokens; token_idx += kNumSMs * kNumCombineWarps) { + // Slots `[0, kNumTopk)` are the routed experts (negative expert id means the + // slot was masked out). With the fused shared expert, slot `kNumTopk` is + // always present and holds this rank's shared-expert contribution. const int stored_topk_slot_idx = lane_idx < kNumTopk ? - static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; + static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : + ((kHasSharedExperts and lane_idx == kNumTopk) ? static_cast(kNumTopk) : -1); const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { diff --git a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh index a6b693b690..93122989fe 100644 --- a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh @@ -216,8 +216,19 @@ struct SM90Workspace { uint32_t num_max_pool_tokens; uint32_t num_max_pool_blocks; + // Fused shared expert: one counter per M block of this rank's own tokens. + // Always reserved (a few KB) so the host and device layouts agree whether or + // not the shared expert is enabled. + uint32_t num_max_shared_pool_blocks; - static constexpr uint64_t kNumBarrierSignalBytes = 32; + // [ 0..15]: 4 x `uint32_t` grid sync counters + // [16..20]: `uint32_t` NVLink barrier counter + // [20..28]: 2 x `int` NVLink barrier signals (phase 0 and 1) + // [28..32]: `uint32_t` L1 task counter (interleaved scheduler) + // [32..36]: `uint32_t` L2 task counter (interleaved scheduler) + // [36..40]: `uint32_t` shared L1 task counter (fused shared expert) + // [40..44]: `uint32_t` shared L2 task counter (fused shared expert) + static constexpr uint64_t kNumBarrierSignalBytes = 48; CUTLASS_HOST_DEVICE SM90Workspace(void* base, @@ -233,6 +244,7 @@ struct SM90Workspace { num_max_pool_tokens = get_num_max_pool_tokens( num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); num_max_pool_blocks = num_max_pool_tokens / kMinCandidateBlockM; + num_max_shared_pool_blocks = math::ceil_div(num_max_tokens_per_rank, static_cast(kMinCandidateBlockM)); } CUTLASS_HOST_DEVICE @@ -245,6 +257,7 @@ struct SM90Workspace { num_bytes += num_max_pool_blocks * sizeof(uint64_t); num_bytes += num_experts_per_rank * num_ranks * num_max_recv_tokens_per_expert * sizeof(int); num_bytes += num_max_pool_tokens * sizeof(TokenSrcMetadata); + num_bytes += math::align(num_max_shared_pool_blocks, 4u) * sizeof(uint32_t); return math::align(num_bytes, 16); } @@ -273,6 +286,30 @@ struct SM90Workspace { base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int)); } + // Interleaved-scheduler global task counters. They are zeroed by the + // workspace cleanup path at the end of every kernel launch (and the + // workspace allocation is zero-initialized for the first launch). + CUTLASS_DEVICE + uint32_t* get_l1_task_count_ptr() const { + return math::advance_ptr(base, 28u); + } + + CUTLASS_DEVICE + uint32_t* get_l2_task_count_ptr() const { + return math::advance_ptr(base, 32u); + } + + // Fused shared-expert task counters, zeroed by the same cleanup path + CUTLASS_DEVICE + uint32_t* get_shared_l1_task_count_ptr() const { + return math::advance_ptr(base, 36u); + } + + CUTLASS_DEVICE + uint32_t* get_shared_l2_task_count_ptr() const { + return math::advance_ptr(base, 40u); + } + CUTLASS_DEVICE uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const { return math::advance_ptr(base, kNumBarrierSignalBytes) + expert_idx; @@ -315,6 +352,15 @@ struct SM90Workspace { const auto base = reinterpret_cast(get_src_token_topk_idx_ptr(num_experts_per_rank)); return base + pool_token_idx; } + + // Fused shared expert: per-M-block count of finished shared L1 tiles. The + // shared L2 A/SFA loader spins on it, mirroring `get_l2_arrival_mask_ptr` in + // counter mode for the routed path. + CUTLASS_DEVICE + uint32_t* get_shared_l2_full_count_ptr(const uint32_t& shared_block_idx = 0) const { + const auto base = get_token_src_metadata_ptr(num_max_pool_tokens); + return reinterpret_cast(base) + shared_block_idx; + } }; struct Data { diff --git a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh index b9bca55de9..f48393abe8 100644 --- a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh +++ b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh @@ -160,6 +160,12 @@ CUTLASS_DEVICE void st_shared_bulk(void* smem_ptr, const uint32_t& num_bytes) { } /// Global memory +CUTLASS_DEVICE uint32_t ld_volatile(const uint32_t* ptr) { + uint32_t ret; + asm volatile("ld.volatile.global.b32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + CUTLASS_DEVICE uint64_t ld_volatile(const uint64_t* ptr) { uint64_t ret; asm volatile("ld.volatile.global.b64 %0, [%1];" : "=l"(ret) : "l"(ptr)); diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh index 19dbfb0da7..4927df4f45 100644 --- a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -13,9 +14,26 @@ namespace deep_gemm::sched { enum class BlockPhase { None = 0, Linear1 = 1, - Linear2 = 2 + Linear2 = 2, + // Fused shared expert (SM90 interleaved scheduler only), ported from SM100. + // The shared expert is a single dense MLP over this rank's own tokens, so it + // has no expert dimension and no dispatch dependency. + SharedLinear1 = 3, + SharedLinear2 = 4 }; +// The L1-like phases read per-128-K activation SF and run the SwiGLU + FP8-quantize +// epilogue; the L2-like phases read per-64-K SF and scatter BF16 into the combine +// buffer. Both helpers also accept the `std::integral_constant` +// form used by the compile-time split-phase hot path. +CUTLASS_HOST_DEVICE constexpr bool is_l1_phase(const BlockPhase& block_phase) { + return block_phase == BlockPhase::Linear1 or block_phase == BlockPhase::SharedLinear1; +} + +CUTLASS_HOST_DEVICE constexpr bool is_shared_phase(const BlockPhase& block_phase) { + return block_phase == BlockPhase::SharedLinear1 or block_phase == BlockPhase::SharedLinear2; +} + template 0` the fused shared expert adds two more phases +// (`SharedLinear1` / `SharedLinear2`) with their own pair of global counters, +// published around the routed tasks exactly like the SM100 scheduler does. +// ============================================================================ + +// Get minimal L1 warmup waves to ensure no L1 -> L2 deadlock +// constexpr func, no runtime overhead in kernel +CUTLASS_HOST_DEVICE constexpr +int get_num_l1_warmup_waves( + const int& num_total_m_blocks, + const int& num_clusters, + const int& num_l1_n_clusters, + const int& num_l2_n_clusters) { + // The first L2 wave may touch multiple M blocks; all their L1 N tasks must be issued first. + const int num_first_l2_wave_m_blocks = math::constexpr_ceil_div(num_clusters, num_l2_n_clusters); + const int num_l1_warmup_clusters_for_first_l2_wave = math::constexpr_ceil_div( + num_first_l2_wave_m_blocks * num_l1_n_clusters, num_clusters); + + // When each M block has more L1 tasks than L2 tasks, the interleaved schedule + // leaves `(num_l1_n_clusters - num_l2_n_clusters)` extra L1 tasks behind + // per M block. To avoid deadlock, no L2 task for an M block should be scheduled + // before that block's L1 tasks have been issued. The last M block is the + // bottleneck: it needs its own `num_l1_n_clusters` L1 tasks, and the preceding + // `num_total_m_blocks - 1` M blocks accumulate + // `(num_total_m_blocks - 1) * (num_l1_n_clusters - num_l2_n_clusters)` + // pending L1 tasks, so we issue them during the warmup phase. Add one extra + // wave to cover partial-wave rounding. + const int num_interleave_cluster_diff_per_m_block = + num_l1_n_clusters > num_l2_n_clusters ? num_l1_n_clusters - num_l2_n_clusters : 0; + const int num_warmup_waves_for_interleave_schedule = math::constexpr_ceil_div( + num_l1_n_clusters + (num_total_m_blocks - 1) * num_interleave_cluster_diff_per_m_block, + num_clusters) + 1; + + return cute::max(num_l1_warmup_clusters_for_first_l2_wave, num_warmup_waves_for_interleave_schedule); +} + +// Number of task-info ring stages shared by the producer warp and consumers +constexpr uint32_t kNumSM90TaskInfoStages = 2; + +struct alignas(16) SM90MegaMoETaskInfo { + BlockPhase block_phase; + uint32_t local_expert_idx; + uint32_t m_block_idx; + uint32_t n_block_idx; + uint32_t pool_block_idx; + uint32_t valid_m; + uint32_t shape_n; + uint32_t shape_k; + + CUTLASS_HOST_DEVICE + SM90MegaMoETaskInfo(): SM90MegaMoETaskInfo(BlockPhase::None, 0, 0, 0, 0, 0, 0, 0) {} + + CUTLASS_HOST_DEVICE + SM90MegaMoETaskInfo(const BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& m_block_idx, + const uint32_t& n_block_idx, + const uint32_t& pool_block_idx, + const uint32_t& valid_m, + const uint32_t& shape_n, + const uint32_t& shape_k): + block_phase(block_phase), + local_expert_idx(local_expert_idx), + m_block_idx(m_block_idx), + n_block_idx(n_block_idx), + pool_block_idx(pool_block_idx), + valid_m(valid_m), shape_n(shape_n), shape_k(shape_k) {} + + CUTLASS_DEVICE + uint32_t is_valid() const { + return block_phase != BlockPhase::None; + } +}; + +DG_STATIC_ASSERT(sizeof(SM90MegaMoETaskInfo) == 32, "Invalid task info size"); + +template +struct MegaMoEInterleavedScheduler { + 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"); + + // Fused shared expert: one dense MLP whose N is scaled by the number of shared + // experts. `M` is this rank's own token count, so there is no expert dimension. + static constexpr bool kHasShared = kNumSharedExperts > 0; + static constexpr uint32_t SHARED_L1_SHAPE_N = L1_SHAPE_N * kNumSharedExperts; + static constexpr uint32_t SHARED_L1_SHAPE_K = L1_SHAPE_K; + static constexpr uint32_t SHARED_L2_SHAPE_N = L2_SHAPE_N; + static constexpr uint32_t SHARED_L2_SHAPE_K = L2_SHAPE_K * kNumSharedExperts; + DG_STATIC_ASSERT(not kHasShared or SHARED_L1_SHAPE_N % BLOCK_N == 0, "Invalid shared shape"); + DG_STATIC_ASSERT(not kHasShared or SHARED_L2_SHAPE_K % BLOCK_K == 0, "Invalid shared shape"); + + // Arrival counts + const WorkspaceT& workspace; + + // Scheduler configs + static constexpr uint32_t kNumScheduleStages = kNumSM90TaskInfoStages; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using task_info_t = SM90MegaMoETaskInfo; + uint32_t sched_stage_idx = 0; + uint32_t sched_phase = 0; + Barrier* task_info_full_barriers = nullptr; + Barrier* task_info_empty_barriers = nullptr; + task_info_t* task_infos = nullptr; + + // Pre-cached per-expert token counts + // Layout: `stored_num_tokens_per_expert[i]` holds expert (i * 32 + lane_idx)'s count + uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {}; + uint32_t num_total_m_blocks = 0; + + // Per-producer warmup task count; all producer warps together form one global wave. + static constexpr uint32_t kNumSchedL1WavesDone = 0xffffffffu; + uint32_t num_sched_l1_waves = 0; + + CUTLASS_DEVICE explicit MegaMoEInterleavedScheduler(const WorkspaceT& workspace): workspace(workspace) {} + + CUTLASS_DEVICE MegaMoEInterleavedScheduler(const WorkspaceT& workspace, + Barrier* task_info_full_barriers, + Barrier* task_info_empty_barriers, + task_info_t* task_infos): + workspace(workspace), + task_info_full_barriers(task_info_full_barriers), + task_info_empty_barriers(task_info_empty_barriers), + task_infos(task_infos) {} + + CUTLASS_DEVICE void advance_sched_pipeline() { + DG_STATIC_ASSERT(kNumScheduleStages == 2, "Invalid stages"); + sched_stage_idx ^= 1; + sched_phase ^= sched_stage_idx == 0; + } + + CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const { + uint32_t valid_value = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ? + stored_num_tokens_per_expert[i] : valid_value; + } + return ptx::exchange(valid_value, expert_idx % 32); + } + + // 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) const { + 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 uint32_t get_num_total_pool_blocks() const { + return get_pool_block_offset(kNumExpertsPerRank); + } + + 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(); + + num_total_m_blocks = get_num_total_pool_blocks(); + const uint32_t num_total_l1_tasks = num_total_m_blocks * kNumL1BlockNs; + const int num_total_l1_waves = static_cast(math::ceil_div(num_total_l1_tasks, kNumSMs)); + const int min_l1_warmup_waves = get_num_l1_warmup_waves( + static_cast(num_total_m_blocks), static_cast(kNumSMs), + static_cast(kNumL1BlockNs), static_cast(kNumL2BlockNs)); + num_sched_l1_waves = static_cast(cute::min(min_l1_warmup_waves, num_total_l1_waves)); + } + + // Resolve the owner expert / intra-expert indices of an absolute pool m-block + CUTLASS_DEVICE task_info_t create_task(const BlockPhase& block_phase, + const uint32_t& task_idx, + const uint32_t& num_n_blocks, + const uint32_t& shape_n, + const uint32_t& shape_k) const { + const uint32_t lane_idx = ptx::get_lane_idx(); + const uint32_t m_block_idx = task_idx / num_n_blocks; + const uint32_t n_block_idx = task_idx % num_n_blocks; + + task_info_t result(block_phase, 0, 0, n_block_idx, m_block_idx, 0, shape_n, shape_k); + uint32_t block_offset = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) { + // Reduce whether task fall in the expert + const uint32_t expert_idx = i * 32 + lane_idx; + const uint32_t num_tokens = stored_num_tokens_per_expert[i]; + const uint32_t num_m_blocks = math::ceil_div(num_tokens, BLOCK_M); + const uint32_t inclusive_num_m_blocks = math::warp_inclusive_sum(num_m_blocks, lane_idx); + const uint32_t lane_pool_block_offset = block_offset + inclusive_num_m_blocks - num_m_blocks; + const bool is_owner = expert_idx < kNumExpertsPerRank and + m_block_idx >= lane_pool_block_offset and m_block_idx < lane_pool_block_offset + num_m_blocks; + const uint32_t owner_mask = __ballot_sync(0xffffffff, is_owner); + + // Exchange the expert info + if (owner_mask) { + const uint32_t owner_lane_idx = static_cast(__ffs(owner_mask) - 1); + const uint32_t owner_m_block_idx = m_block_idx - lane_pool_block_offset; + const uint32_t owner_valid_m = cute::min(num_tokens - owner_m_block_idx * BLOCK_M, BLOCK_M); + result.local_expert_idx = ptx::exchange(expert_idx, owner_lane_idx); + result.m_block_idx = ptx::exchange(owner_m_block_idx, owner_lane_idx); + result.valid_m = ptx::exchange(owner_valid_m, owner_lane_idx); + } + block_offset += ptx::exchange(inclusive_num_m_blocks, 31); + } + return result; + } + + static CUTLASS_DEVICE uint32_t get_next_task_idx(const uint32_t* global_task_count_ptr) { + uint32_t result = 0; + if (cute::elect_one_sync()) + result = ptx::atomic_add_rel(global_task_count_ptr, 1u); + return ptx::exchange(result, 0); + } + + // Producer side: claim the next task from the global counters + CUTLASS_DEVICE task_info_t get_next_task() { + while (true) { + if (num_sched_l1_waves != kNumSchedL1WavesDone and num_sched_l1_waves) { + // One local L1 task per producer; globally this is one wave. + -- num_sched_l1_waves; + + // No more L1 tasks + const uint32_t l1_task_idx = get_next_task_idx(workspace.get_l1_task_count_ptr()); + if (l1_task_idx >= num_total_m_blocks * kNumL1BlockNs) { + num_sched_l1_waves = kNumSchedL1WavesDone; + continue; + } + + // Create task + return create_task(BlockPhase::Linear1, l1_task_idx, kNumL1BlockNs, L1_SHAPE_N, L1_SHAPE_K); + } else { + const uint32_t l2_task_idx = get_next_task_idx(workspace.get_l2_task_count_ptr()); + if (l2_task_idx >= num_total_m_blocks * kNumL2BlockNs) + break; + + // The next task should be L1 + if (num_sched_l1_waves != kNumSchedL1WavesDone) + num_sched_l1_waves = 1; + + // Create task + auto task_info = create_task(BlockPhase::Linear2, l2_task_idx, kNumL2BlockNs, L2_SHAPE_N, L2_SHAPE_K); + + // Wait until all required L1 tasks are fetched + const auto num_required_l1_tasks = (task_info.pool_block_idx + 1) * kNumL1BlockNs; + while (ptx::ld_volatile(workspace.get_l1_task_count_ptr()) < num_required_l1_tasks) {} + return task_info; + } + } + return task_info_t(); + } + + // Producer side: publish a task into the SMEM ring + CUTLASS_DEVICE void publish_task(const task_info_t& task_info, const uint32_t& lane_idx) { + if (lane_idx == 0) { + task_infos[sched_stage_idx] = task_info; + // The mbarrier arrive has release semantics, so the plain SMEM + // stores above are visible to consumers after their wait + task_info_full_barriers[sched_stage_idx].arrive(); + } + __syncwarp(); + advance_sched_pipeline(); + } + + // Producer side: publish every shared-expert task of one phase. + // The shared expert has no expert dimension and no ring buffer, so + // `local_expert_idx == 0` and `pool_block_idx == m_block_idx` (the M block index + // over this rank's own tokens). + template + CUTLASS_DEVICE void shared_mainloop(const uint32_t& num_tokens, const uint32_t& lane_idx, + const uint32_t* task_count_ptr) { + constexpr uint32_t kNumSharedNBlocks = kShapeN / BLOCK_N; + const uint32_t num_m_blocks = math::ceil_div(num_tokens, BLOCK_M); + const uint32_t num_tasks = num_m_blocks * kNumSharedNBlocks; + while (true) { + task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1); + + // Dynamic scheduling, like the routed path: reduces tailing across the + // different shared L1/L2 tile shapes + const uint32_t task_idx = get_next_task_idx(task_count_ptr); + if (task_idx >= num_tasks) + break; + + const uint32_t m_block_idx = task_idx / kNumSharedNBlocks; + const uint32_t n_block_idx = task_idx % kNumSharedNBlocks; + const uint32_t valid_m = cute::min(num_tokens - m_block_idx * BLOCK_M, BLOCK_M); + publish_task(task_info_t(kBlockPhase, 0, m_block_idx, n_block_idx, + m_block_idx, valid_m, kShapeN, kShapeK), lane_idx); + } + } + + // Task order per CTA is `[shared L1][routed L1/L2][shared L2]`, and consumers + // drain their CTA's ring in FIFO order. Shared L1 waits on nothing (its A operand + // is the local `x`, resident before the launch), and shared L2 sits behind every + // routed task, so a shared-L2 dependency can never block the production of + // anything ahead of it in any CTA -- the schedule stays deadlock-free. + CUTLASS_DEVICE void mainloop(const uint32_t& num_tokens) { + const auto lane_idx = ptx::get_lane_idx(); + + if constexpr (kHasShared) { + // Shared L1 does not depend on dispatch + shared_mainloop( + num_tokens, lane_idx, workspace.get_shared_l1_task_count_ptr()); + } + + // Wait dispatch's results + fetch_expert_recv_count(); + + // Generate routed tasks. Keep the wait -> claim -> publish ordering: + // `get_next_task()` advances global task counters and must not run + // before the schedule slot is released by consumers. + task_info_t task_info; + do { + task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1); + task_info = get_next_task(); + if (task_info.is_valid()) + publish_task(task_info, lane_idx); + } while (task_info.is_valid()); + + if constexpr (kHasShared) { + // Shared L2 consumes the shared L1 output; the A/SFA loader gates each + // tile on the per-M-block shared L1 arrival counter + shared_mainloop( + num_tokens, lane_idx, workspace.get_shared_l2_task_count_ptr()); + } + + // Sentinel + task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1); + publish_task(task_info_t(), lane_idx); + } + + // Consumer side: fetch the next published task (and immediately release + // the slot; all fields are copied into registers, so the producer may + // overwrite the slot as soon as every consumer warp has read it). + CUTLASS_DEVICE bool get_next_task(task_info_t& task_info) { + task_info_full_barriers[sched_stage_idx].wait(sched_phase); + task_info = task_infos[sched_stage_idx]; + __syncwarp(); + if (cute::elect_one_sync()) + task_info_empty_barriers[sched_stage_idx].arrive(); + advance_sched_pipeline(); + return task_info.is_valid(); + } +}; + } // namespace deep_gemm::sched diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index c3e73931e4..2cff99bf72 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -289,7 +289,8 @@ def __init__(self, group, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, use_fp8_dispatch: bool = True, - activation: str = 'swiglu'): + activation: str = 'swiglu', + num_shared_experts: int = 0): import torch.distributed._symmetric_memory as symm_mem self.group = group @@ -298,12 +299,17 @@ def __init__(self, group, self.num_topk = num_topk self.hidden = hidden self.intermediate_hidden = intermediate_hidden + # Fused shared expert (SM90 only). 0 disables it and keeps + # the symmetric-buffer layout byte-identical to the routed-only path. + self.num_shared_experts = num_shared_experts + self.shared_intermediate_hidden = intermediate_hidden * num_shared_experts num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe( group.size(), num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - use_fp8_dispatch, activation + use_fp8_dispatch, activation, + num_shared_experts ) self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda') self.handle = symm_mem.rendezvous(self.buffer, group=group) @@ -312,7 +318,8 @@ def __init__(self, group, torch.cuda.synchronize() (x, x_sf, topk_idx, topk_weights, - l1_acts, l1_acts_sf, l2_acts, l2_acts_sf) = slice_input_buffers(self.buffer) + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf, + shared_l2_acts, shared_l2_acts_sf) = slice_input_buffers(self.buffer) self.x = _from_dlpack_if_needed(x, torch.float8_e4m3fn) self.x_sf = _from_dlpack_if_needed(x_sf) self.topk_idx = _from_dlpack_if_needed(topk_idx) @@ -321,6 +328,15 @@ def __init__(self, group, self.l1_acts_sf = _from_dlpack_if_needed(l1_acts_sf) self.l2_acts = _from_dlpack_if_needed(l2_acts, torch.float8_e4m3fn) self.l2_acts_sf = _from_dlpack_if_needed(l2_acts_sf) + # The fused shared expert reads its L1 activations from the same `x` region; + # only the post-SwiGLU pool and its SF are extra buffers (zero-sized when the + # shared expert is disabled). + self.shared_l1_acts = self.x + self.shared_l1_acts_sf = self.x_sf + shared_l2_acts = _from_dlpack_if_needed(shared_l2_acts, torch.float8_e4m3fn) + shared_l2_acts_sf = _from_dlpack_if_needed(shared_l2_acts_sf) + self.shared_l2_acts = shared_l2_acts if num_shared_experts > 0 else None + self.shared_l2_acts_sf = shared_l2_acts_sf if num_shared_experts > 0 else None def destroy(self): self.handle = None @@ -328,6 +344,10 @@ def destroy(self): self.group = None self.x = None self.x_sf = None + self.shared_l1_acts = None + self.shared_l1_acts_sf = None + self.shared_l2_acts = None + self.shared_l2_acts_sf = None def get_symm_buffer_for_sm90_mega_moe(group, @@ -335,7 +355,8 @@ def get_symm_buffer_for_sm90_mega_moe(group, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, use_fp8_dispatch: bool = True, - activation: str = 'swiglu') -> SM90SymmBuffer: + activation: str = 'swiglu', + num_shared_experts: int = 0) -> SM90SymmBuffer: from .utils.math import align num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) @@ -343,7 +364,8 @@ def get_symm_buffer_for_sm90_mega_moe(group, group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - use_fp8_dispatch, activation + use_fp8_dispatch, activation, + num_shared_experts ) @@ -353,7 +375,8 @@ def get_symm_buffer_for_mega_moe(group, hidden: int, intermediate_hidden: int, use_fp8_dispatch: Union[bool, None] = None, mma_type: str = 'fp8xfp4', - activation: str = 'swiglu'): + activation: str = 'swiglu', + num_shared_experts: int = 0): if use_fp8_dispatch is not None: assert use_fp8_dispatch == (mma_type.split('x')[0] == 'fp8') if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9: @@ -362,8 +385,10 @@ def get_symm_buffer_for_mega_moe(group, group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - True, activation + True, activation, + num_shared_experts ) + assert num_shared_experts == 0, 'Shared experts are only wired through the SM90 buffer' return mega.get_symm_buffer_for_mega_moe( group, num_experts, num_max_tokens_per_rank, num_topk, @@ -374,20 +399,39 @@ def get_symm_buffer_for_mega_moe(group, ) +def _interleave_gate_up_sm90(t: torch.Tensor, gran: int = 8) -> torch.Tensor: + """Interleave the gate/up halves of an L1 weight (or its SF) along N. + + Handles the routed `(G, 2*N, ...)` layout and the shared expert's dense + `(2*N, ...)` one, which the kernel's weight-SF indexing assumes for both. + """ + if t.dim() == 2: + return _interleave_gate_up_sm90(t.unsqueeze(0), gran).squeeze(0) + g, n, *rest = t.shape + half = n // 2 + gate = t[:, :half].reshape(g, half // gran, gran, *rest) + up = t[:, half:].reshape(g, half // gran, gran, *rest) + return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest)) + + def transform_weights_for_mega_moe_sm90( l1_weights: Tuple[torch.Tensor, torch.Tensor], l2_weights: Tuple[torch.Tensor, torch.Tensor] ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: l1_fp8, l1_sf = l1_weights + return (_interleave_gate_up_sm90(l1_fp8), l1_sf), l2_weights + - def _interleave_one(t, gran: int = 8) -> torch.Tensor: - g, n, *rest = t.shape - half = n // 2 - gate = t[:, :half].reshape(g, half // gran, gran, *rest) - up = t[:, half:].reshape(g, half // gran, gran, *rest) - return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest)) +def transform_shared_weights_for_mega_moe_sm90( + shared_l1_weights: Tuple[torch.Tensor, torch.Tensor], + shared_l2_weights: Tuple[torch.Tensor, torch.Tensor] +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + """Same gate/up interleave as the routed L1, on the dense shared-expert weights. - return (_interleave_one(l1_fp8), l1_sf), l2_weights + Shapes: L1 `(2 * shared_intermediate_hidden, hidden)`, L2 `(hidden, shared_intermediate_hidden)`. + """ + shared_l1_fp8, shared_l1_sf = shared_l1_weights + return (_interleave_gate_up_sm90(shared_l1_fp8), shared_l1_sf), shared_l2_weights def fp8_mega_moe(y: torch.Tensor, @@ -398,13 +442,23 @@ def fp8_mega_moe(y: torch.Tensor, recipe: Tuple[int, int, int] = (128, 128, 128), activation: str = 'swiglu', activation_clamp: Optional[float] = None, - fast_math: bool = True): + fast_math: bool = True, + shared_l1_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l2_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None): (l1_weights_data, l1_weights_sf) = l1_weights (l2_weights_data, l2_weights_sf) = l2_weights + # Fused shared expert: needs a symmetric buffer allocated with + # `num_shared_experts > 0` (see `get_symm_buffer_for_mega_moe`). + assert (shared_l1_weights is None) == (shared_l2_weights is None), \ + 'Shared-expert L1 and L2 weights must be passed together' + shared_l1_weights_data, shared_l1_weights_sf = shared_l1_weights or (None, None) + shared_l2_weights_data, shared_l2_weights_sf = shared_l2_weights or (None, None) _C.fp8_mega_moe( y, l1_weights_data, l1_weights_sf, l2_weights_data, l2_weights_sf, + shared_l1_weights_data, shared_l1_weights_sf, + shared_l2_weights_data, shared_l2_weights_sf, cumulative_local_expert_recv_stats, sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), diff --git a/sgl_deep_gemm/tests/test_mega_moe_hopper.py b/sgl_deep_gemm/tests/test_mega_moe_hopper.py index f3579a7e91..5818a1ccda 100644 --- a/sgl_deep_gemm/tests/test_mega_moe_hopper.py +++ b/sgl_deep_gemm/tests/test_mega_moe_hopper.py @@ -11,6 +11,12 @@ per-128-K L2 activation SF, while the fused SM90 MegaMoE L1 epilogue writes per-64-K L2 activation SF to avoid cross-CTA synchronization. This is a same-pipeline performance reference, not a bitwise correctness oracle. +* shared expert (optional, ``--num-shared-experts``): DeepSeek-style shared + expert, fused into the kernel as the SM100-style ``SharedLinear1`` / + ``SharedLinear2`` scheduler phases and reduced through an extra combine slot, so + one launch produces routed + shared. Both baselines run the same shared expert + serially, matching ``tests/test_mega_moe.py``, which folds it in as a DeepEP + combine bias. * low-latency baseline (optional, ``--run-low-latency-baseline``): mirrors the sglang low-latency MoE pipeline (see ``sglang/srt/layers/moe/token_dispatcher/deepep.py::_DeepEPDispatcherImplLowLatency``): @@ -25,6 +31,8 @@ * accuracy mode (optional, ``--accuracy``): runs the former layered SM90 correctness suite with a PyTorch BF16/FP32 reference. It covers smoke, heuristic branches, shape sweeps, edge cases, and optional random stress. + Layer 6 adds the fused shared expert (``num_shared_experts >= 1``); the + reference is routed + dense shared MLP, compared with the fused kernel output. * output: TFLOPS, overlap-adjusted TFLOPS, HBM GB/s, NVLink GB/s, fused time, reduction estimate, and ``t_baseline / t_fused``. """ @@ -253,6 +261,18 @@ def _quantize_grouped_fp8_block_128_128( return w_fp8.view(g, n, k).contiguous(), sf.contiguous() +def _quantize_dense_fp8_block_128_128( + w: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(N, K) bf16 -> (N, K) fp8_e4m3fn plus (N/128, K/128) FP32 block SF. + + Used for the shared expert, which is a single dense MLP rather than a group + of per-rank experts. + """ + w_fp8, sf = _quantize_grouped_fp8_block_128_128(w.unsqueeze(0)) + return w_fp8.squeeze(0), sf.squeeze(0) + + # ============================================================================ # Section 3: layered accuracy reference and scenarios. # ============================================================================ @@ -380,6 +400,54 @@ def _reference_fused( return y_full_bf16[start:end].contiguous() +def _reference_shared( + x_fp8_local: torch.Tensor, + x_sf_local: torch.Tensor, + shared_l1: Tuple[torch.Tensor, torch.Tensor], + shared_l2: Tuple[torch.Tensor, torch.Tensor], + hidden: int, + shared_intermediate_hidden: int, + activation_clamp: float, +) -> torch.Tensor: + """PyTorch FP32 reference for the fused shared expert, on this rank's local + tokens (the shared expert is node-local — no all_gather, unlike the routed + path). Mirrors the kernel's fused-shared arithmetic: + + L1: (M, H) @ (2*SIH, H)^T -> (M, 2*SIH) + SwiGLU (gate clamp one-sided, up clamp two-sided), no topk weighting + L1 output FP8 round-trip with per-64-K scales (matches FUSED_L2_ACT_SF_GRAN) + L2: (M, SIH) @ (H, SIH)^T -> (M, H) -> bf16 + + The fused kernel adds the shared output into the routed output; this returns + the shared contribution so the caller can do ``y_ref routed + y_ref shared``. + """ + assert shared_intermediate_hidden % 64 == 0, ( + "shared_intermediate_hidden must be a multiple of 64 for the per-64-K " + "L2 activation SF granularity the fused epilogue uses" + ) + x = _dequant_per_token_per_128_k(x_fp8_local, x_sf_local) # (M, H) fp32 + + # L1 GEMM (dense, the shared expert is a single MLP). + l1_w = _dequant_block_128_128(shared_l1[0], shared_l1[1]) # (2*SIH, H) fp32 + l1_y = x @ l1_w.t() # (M, 2*SIH) + l1_y = _swiglu_fp32(l1_y, activation_clamp) # (M, SIH) + + # L1 output FP8 round-trip with per-64-K scales, matching the fused SM90 L2 + # activation SF granularity (FUSED_L2_ACT_SF_GRAN = 64). + s, ih = l1_y.shape + assert ih == shared_intermediate_hidden + v = l1_y.view(s, ih // 64, 64) + sf2 = v.abs().amax(dim=-1).clamp(1e-4) / FP8_E4M3_MAX + l2_in = ( + (v / sf2.unsqueeze(-1)).to(torch.float8_e4m3fn).float() + * sf2.unsqueeze(-1) + ).view(s, ih) + + # L2 GEMM -> bf16. + l2_w = _dequant_block_128_128(shared_l2[0], shared_l2[1]) # (H, SIH) fp32 + return (l2_in @ l2_w.t()).to(torch.bfloat16) + + def _run_accuracy_scenario( name: str, cfg: Dict[str, Any], @@ -397,6 +465,7 @@ def _run_accuracy_scenario( masked_ratio = cfg.get("masked_ratio", 0.0) activation_clamp = cfg.get("activation_clamp", 10.0) fast_math = cfg.get("fast_math", True) + num_shared_experts = cfg.get("num_shared_experts", 0) assert num_experts % num_ranks == 0, ( f"{name}: experts {num_experts} not divisible by ranks {num_ranks}" @@ -404,6 +473,14 @@ def _run_accuracy_scenario( num_experts_per_rank = num_experts // num_ranks assert num_tokens <= num_max assert hidden % 128 == 0 and intermediate_hidden % 128 == 0 + # Fused shared expert: each adds one routed intermediate size, fused into the + # mega kernel via SharedLinear1/2. + shared_intermediate_hidden = intermediate_hidden * num_shared_experts + if num_shared_experts > 0: + assert shared_intermediate_hidden % 64 == 0, ( + f"{name}: shared_intermediate_hidden={shared_intermediate_hidden} must be " + f"a multiple of 64 (fused L2 activation SF granularity)" + ) verbose = bool(int(os.environ.get("DG_TEST_VERBOSE", "0"))) @@ -445,6 +522,31 @@ def trace(stage: str): transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe_sm90( l1_weights, l2_weights ) + # Shared expert weights (dense MLP, block-(128,128) FP8 like the routed + # weights) and the matching interleave the kernel's weight-SF indexing assumes. + if num_shared_experts > 0: + shared_l1_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (shared_intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + shared_l2_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (hidden, shared_intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + transformed_shared_l1, transformed_shared_l2 = ( + deep_gemm.transform_shared_weights_for_mega_moe_sm90( + shared_l1_weights, shared_l2_weights + ) + ) + else: + shared_l1_weights = shared_l2_weights = None + transformed_shared_l1 = transformed_shared_l2 = None trace("alloc_symm_buffer") buffer = deep_gemm.get_symm_buffer_for_mega_moe( @@ -454,6 +556,7 @@ def trace(stage: str): num_topk, hidden, intermediate_hidden, + num_shared_experts=num_shared_experts, ) cum_stats = torch.zeros((num_experts_per_rank,), dtype=torch.int, device="cuda") @@ -475,6 +578,8 @@ def trace(stage: str): activation="swiglu", activation_clamp=activation_clamp if math.isfinite(activation_clamp) else None, fast_math=fast_math, + shared_l1_weights=transformed_shared_l1, + shared_l2_weights=transformed_shared_l2, ) torch.cuda.synchronize() @@ -497,6 +602,17 @@ def trace(stage: str): intermediate_hidden, activation_clamp, ) + # Add the shared-expert reference contribution: fused-out = routed + shared. + if num_shared_experts > 0: + y_ref = y_ref + _reference_shared( + x_fp8[0], + x_fp8[1], + shared_l1_weights, + shared_l2_weights, + hidden, + shared_intermediate_hidden, + activation_clamp, + ) diff = calc_diff(y_fused, y_ref) ok = diff < diff_tol @@ -614,6 +730,54 @@ def _accuracy_layer5_stress(num_ranks: int, num_tests: int) -> List[Tuple[str, D return out +def _accuracy_layer6_shared_expert(num_ranks: int) -> List[Tuple[str, Dict[str, Any]]]: + """Fused shared-expert correctness (fused-out = routed + shared). + + Each scenario sets num_shared_experts >= 1, which routes through the + fused-shared path and compares against routed reference + dense shared MLP. + shared_intermediate_hidden must be a multiple of 64, so intermediate_hidden is + picked from {512, 1024, 2048}. + """ + base = dict( + num_max_tokens_per_rank=128, + hidden=512, + intermediate_hidden=512, # * num_shared_experts stays a multiple of 64 + num_experts=8 * num_ranks, + num_topk=2, + num_shared_experts=1, + ) + out = [] + # Smoke: 1 shared expert, default clamp/fast_math/shape. + out.append(("L6.sh1.smoke", dict(base))) + # Two shared experts (shared_intermediate_hidden = 2*ih). + cfg = dict(base) + cfg.update(num_shared_experts=2) + out.append(("L6.sh2", cfg)) + # Larger hidden / intermediate with the shared expert on. + for hidden, ih in [(2048, 1024), (2048, 2048)]: + cfg = dict(base) + cfg.update(hidden=hidden, intermediate_hidden=ih) + out.append((f"L6.h{hidden}_ih{ih}", cfg)) + # Shared expert under masking (shared acts on local tokens regardless of routing). + for masked_ratio in (0.3, 0.7): + cfg = dict(base) + cfg.update(masked_ratio=masked_ratio) + out.append((f"L6.mask{masked_ratio:.1f}", cfg)) + # Shared expert with clamp variations. + for clamp in (1.0, math.inf): + cfg = dict(base) + cfg.update(activation_clamp=clamp) + out.append((f"L6.clamp{clamp}", cfg)) + # Shared expert with fewer/more routed experts selected. + for topk in (1, 4): + if topk > base["num_experts"]: + continue + cfg = dict(base) + cfg.update(num_topk=topk) + out.append((f"L6.topk{topk}", cfg)) + return out + + def _run_accuracy_tests(local_rank: int, num_local_ranks: int, args: argparse.Namespace): rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) @@ -636,6 +800,8 @@ def _run_accuracy_tests(local_rank: int, num_local_ranks: int, args: argparse.Na layers += _accuracy_layer4_edges(num_ranks) if 5 in args.layers: layers += _accuracy_layer5_stress(num_ranks, args.num_correctness_tests or 8) + if 6 in args.layers: + layers += _accuracy_layer6_shared_expert(num_ranks) if args.filter: layers = [(name, cfg) for name, cfg in layers if args.filter in name] @@ -644,6 +810,8 @@ def _run_accuracy_tests(local_rank: int, num_local_ranks: int, args: argparse.Na f"layers {sorted(args.layers)} on {num_ranks} ranks", once_in_node=True, ) + # --num-shared-experts is a config for the benchmark mode; in accuracy mode + # each scenario carries its own num_shared_experts (layer 6 turns it on). failures: List[str] = [] for name, cfg in layers: @@ -1115,6 +1283,12 @@ def _run_fused_only_sweep(local_rank: int, num_local_ranks: int, args: argparse. f"masked_ratio={args.masked_ratio} fast_math={bool(args.fast_math)}", once_in_node=True, ) + if args.num_shared_experts > 0: + dist_print( + " > note: --num-shared-experts is ignored here; this mode measures the " + "routed kernel alone", + once_in_node=True, + ) num_max_tokens_per_rank = max(batches) for num_tokens in batches: @@ -1192,6 +1366,14 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): f"SM90 fused kernel requires intermediate_hidden <= 4096, got {intermediate_hidden}" ) + # Shared-expert shape, following tests/test_mega_moe.py: one routed + # intermediate size per shared expert. When enabled it is fused into the mega + # kernel as the SharedLinear1/SharedLinear2 phases; 0 disables it. + num_shared_experts = args.num_shared_experts + shared_intermediate_hidden = intermediate_hidden * num_shared_experts + assert shared_intermediate_hidden % 128 == 0 + fused_shared = num_shared_experts > 0 + # ---- Create BF16 token and weight inputs ---- # x: local tokens for this rank. x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") @@ -1241,6 +1423,38 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): l1_weights, l2_weights ) + # Shared-expert weights: one dense MLP (not per-expert), quantized with the + # same block-(128, 128) FP8 recipe as the routed weights. The fused and + # baseline paths share them so the comparison stays apples-to-apples. + if num_shared_experts > 0: + shared_l1_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (shared_intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + shared_l2_weights = _quantize_dense_fp8_block_128_128( + torch.randn( + (hidden, shared_intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ) + else: + shared_l1_weights = shared_l2_weights = None + + # The fused kernel consumes the same shared weights with the gate/up gran-8 + # interleave applied to L1 (identical to the routed weight transform). + if fused_shared: + transformed_shared_l1, transformed_shared_l2 = ( + deep_gemm.transform_shared_weights_for_mega_moe_sm90( + shared_l1_weights, shared_l2_weights + ) + ) + else: + transformed_shared_l1 = transformed_shared_l2 = None + # SwiGLU clamp: finite values enable clamp; inf maps to None and disables it. clamp_arg = args.activation_clamp if math.isfinite(args.activation_clamp) else None run_baseline_enabled = args.run_baseline or bool(args.check_output_diff) @@ -1251,6 +1465,8 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) # ---- Allocate fused SymmBuffer and output buffer ---- + # The fused shared expert needs two extra symmetric-buffer regions (its + # post-SwiGLU pool and SF) plus one more combine slot. sym_buffer = deep_gemm.get_symm_buffer_for_mega_moe( group, num_experts, @@ -1258,10 +1474,15 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): num_topk, hidden, intermediate_hidden, + num_shared_experts=num_shared_experts if fused_shared else 0, ) y_fused = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") - def run_fused(): + # Output of the reference dense shared MLP (`run_shared`), used by the baselines + # and by `--check-output-diff`. Reused across calls: those paths never overlap. + y_shared = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + + def run_fused(with_shared: bool = True): # Match the SM100 test: DG_COMM_KERNEL_DEBUG=1 zeros the whole # sym_buffer at kernel exit, so inputs must be re-copied every call. sym_buffer.x[:num_tokens].copy_(x_fp8[0]) @@ -1269,6 +1490,7 @@ def run_fused(): sym_buffer.topk_idx[:num_tokens].copy_(topk_idx) sym_buffer.topk_weights[:num_tokens].copy_(topk_weights) + fuse_now = fused_shared and with_shared deep_gemm.fp8_mega_moe( y_fused, transformed_l1, @@ -1279,9 +1501,51 @@ def run_fused(): activation="swiglu", activation_clamp=clamp_arg, fast_math=bool(args.fast_math), + shared_l1_weights=transformed_shared_l1 if fuse_now else None, + shared_l2_weights=transformed_shared_l2 if fuse_now else None, ) return y_fused + def run_shared(): + """Dense FP8 shared-expert MLP writing into ``y_shared``. + + L1 GEMM -> SwiGLU + FP8 quantization -> L2 GEMM. There is no topk + weighting: every token passes through the shared expert with weight 1.0. + The activation SF stays row-major FP32; ``fp8_gemm_nt`` transposes it into + the MN-major TMA layout internally (see + ``layout::transform_sf_into_required_layout``). + """ + if num_tokens == 0: + return y_shared + + l1_out = torch.empty( + (num_tokens, shared_intermediate_hidden * 2), + dtype=torch.bfloat16, + device="cuda", + ) + deep_gemm.fp8_gemm_nt( + x_fp8, + shared_l1_weights, + l1_out, + recipe=(1, 128, 128), + disable_ue8m0_cast=True, + ) + l2_in = swiglu_apply_weight_to_fp8_triton( + x=l1_out, + topk_weights=None, + clamp_value=clamp_arg, + num_per_channels=BASELINE_L2_ACT_SF_GRAN, + use_ue8m0_scale=False, + ) + deep_gemm.fp8_gemm_nt( + l2_in, + shared_l2_weights, + y_shared, + recipe=(1, 128, 128), + disable_ue8m0_cast=True, + ) + return y_shared + # ---- Print config ---- dist_print("Config (SM90 fused MegaMoE):", once_in_node=True) dist_print(f" > Tokens: {num_tokens}/{num_max_tokens_per_rank}", once_in_node=True) @@ -1293,6 +1557,16 @@ def run_fused(): once_in_node=True, ) dist_print(f" > Masked ratio: {args.masked_ratio}", once_in_node=True) + dist_print( + f" > Shared experts: {num_shared_experts}" + + ( + f" (intermediate: {shared_intermediate_hidden}, fused into the mega " + f"kernel via SharedLinear1/2)" + if fused_shared + else " (disabled)" + ), + once_in_node=True, + ) dist_print( f" > Activation SF: fused L2 per-{FUSED_L2_ACT_SF_GRAN} FP32 pow2, " f"baseline L2 per-{BASELINE_L2_ACT_SF_GRAN} FP32 pow2 " @@ -1415,7 +1689,13 @@ def run_baseline(): ) # DeepEP combine: gather each token's topk expert outputs back to source rank. - return ep_buffer.combine(l2_y, handle=handle)[0] + combined = ep_buffer.combine(l2_y, handle=handle)[0] + # Non-overlapped baseline: the shared expert runs serially on the same + # stream. tests/test_mega_moe.py folds it into combine as a bias; the + # SM90 DeepEP shim here has no bias argument, so add it afterwards. + if num_shared_experts > 0: + combined.add_(run_shared()) + return combined # ---------------------------------------------------------------- # Low-latency baseline body. Mirrors the sglang @@ -1517,6 +1797,9 @@ def run_baseline_low_latency(): return_recv_hook=False, out=ll_combined, ) + # 6) Same serial shared expert as the normal-mode baseline. + if num_shared_experts > 0: + combined_x.add_(run_shared()) return combined_x # ---- Run once to check fused and optional baseline paths ---- @@ -1524,6 +1807,27 @@ def run_baseline_low_latency(): assert y.shape == (num_tokens, hidden) and y.dtype == torch.bfloat16, ( f"unexpected fused output shape/dtype: shape={y.shape}, dtype={y.dtype}" ) + if fused_shared and args.check_output_diff: + # Reference for the fused shared expert: the routed-only kernel output plus + # the Python dense shared MLP (the same weights, before the interleave). + y_fused_shared = y.clone() + y_ref = run_fused(with_shared=False).clone() + y_ref += run_shared() + diff = (y_fused_shared.float() - y_ref.float()).abs() + denom = y_ref.float().abs().mean().clamp_min(1e-12) + dist_print( + "Output diff (fused shared expert vs routed + two-stream shared):", + once_in_node=True, + ) + dist_print( + f" > max_abs={diff.max().item():.6e}, " + f"mean_abs={diff.mean().item():.6e}, " + f"mean_abs/mean_ref={diff.mean().div(denom).item():.6e}", + once_in_node=True, + ) + dist_print(once_in_node=True) + # Leave `y_fused` holding the fused output for the baseline diffs below + y = run_fused() if ep_buffer is not None: out_b = run_baseline() assert out_b.shape == (num_tokens, hidden) and out_b.dtype == torch.bfloat16, ( @@ -1574,7 +1878,9 @@ def run_baseline_low_latency(): num_touched_experts = int(torch.unique(local_expert_ids).numel()) # ---- benchmark ---- - # Fused: bench_kineto selects the sm90_fp8_mega_moe_impl GPU region only. + # Fused: bench_kineto selects the sm90_fp8_mega_moe_impl GPU region only, so + # this stays the pure routed-kernel time even with a shared expert running + # concurrently on the main stream. t_fused = bench_kineto( run_fused, SM90_KERNEL_NAME, @@ -1657,27 +1963,64 @@ def safe_div(a, b): num_nvlink_bytes = num_recv_tokens * (hidden + hidden // 32 + 4 + hidden * 2) nvlink_gbs = safe_div(num_nvlink_bytes / 1e9, t_fused) + # ---- Shared-expert FLOPs / HBM ---- + # Same three matmuls as a routed expert (L1 gate, L1 up, L2), but every local + # token goes through it, and the weights are streamed once (not per expert). + # The shared MLP is node-local, so it adds no NVLink traffic, and the fused + # epilogue keeps the SwiGLU input in registers (no BF16 staging round-trip). + num_shared_flops = 2 * num_tokens * hidden * shared_intermediate_hidden * 3 + num_shared_hbm_bytes = ( + 0 + if num_shared_experts == 0 + else ( + shared_intermediate_hidden * 2 * hidden # shared L1 weights (FP8) + + hidden * shared_intermediate_hidden # shared L2 weights (FP8) + + (shared_intermediate_hidden * 2 // WEIGHT_SF_GRAN_MN) + * (hidden // WEIGHT_SF_GRAN_K) + * 4 # shared L1 weight SF + + (hidden // WEIGHT_SF_GRAN_MN) + * (shared_intermediate_hidden // WEIGHT_SF_GRAN_K) + * 4 # shared L2 weight SF + + num_tokens * hidden + + num_tokens * (hidden // L1_ACT_SF_GRAN) * 4 # L1 input read (FP8 + SF) + + num_tokens * shared_intermediate_hidden + + num_tokens + * (shared_intermediate_hidden // BASELINE_L2_ACT_SF_GRAN) + * 4 # SwiGLU output write (FP8 + SF) + + num_tokens * shared_intermediate_hidden + + num_tokens + * (shared_intermediate_hidden // BASELINE_L2_ACT_SF_GRAN) + * 4 # L2 input read (FP8 + SF) + + num_tokens * hidden * 2 # L2 output write (BF16) + ) + ) + # Routed + shared: one launch produces both, so they share `t_fused`. + num_total_flops = ( + 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) + num_shared_flops + ) + num_total_hbm_bytes = num_hbm_bytes + num_shared_hbm_bytes + tflops_total = safe_div(num_total_flops / 1e12, t_fused) + hbm_gbs_total = safe_div(num_total_hbm_bytes / 1e9, t_fused) + # Serial lower bound for combine reduction, using 6.5e12 B/s as an estimate. t_reduction = num_tokens * hidden * 2 * (1 + num_topk) / 6.5e12 # Overlap adjustment: remove the non-overlapped serial reduction estimate. approx_factor = t_fused / max(t_fused - t_reduction, 1e-12) - # Baseline uses the same FLOPs and HBM byte estimate, with t_baseline. - tflops_baseline = safe_div( - 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_baseline - ) - hbm_gbs_baseline = safe_div(num_hbm_bytes / 1e9, t_baseline) + # Baselines run routed + shared serially, so they use the combined FLOPs and + # HBM byte estimate (identical to the routed-only one when shared is off). + tflops_baseline = safe_div(num_total_flops / 1e12, t_baseline) + hbm_gbs_baseline = safe_div(num_total_hbm_bytes / 1e9, t_baseline) nvlink_gbs_baseline = safe_div(num_nvlink_bytes / 1e9, t_baseline) # Low-latency baseline pads each expert's activation to ``M_max_ll``, so # the weights are streamed once per expert regardless of routing. NVLink # bytes match the normal-mode baseline (same per-routed-token volume). - tflops_baseline_ll = safe_div( - 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_baseline_ll - ) - hbm_gbs_baseline_ll = safe_div(num_hbm_bytes / 1e9, t_baseline_ll) + tflops_baseline_ll = safe_div(num_total_flops / 1e12, t_baseline_ll) + hbm_gbs_baseline_ll = safe_div(num_total_hbm_bytes / 1e9, t_baseline_ll) nvlink_gbs_baseline_ll = safe_div(num_nvlink_bytes / 1e9, t_baseline_ll) + def fmt_perf_line( name: str, t: float, @@ -1713,10 +2056,10 @@ def fmt_perf_line( ) dist_print( fmt_perf_line( - "[fused]", + "[fused+sh]" if fused_shared else "[fused]", t_fused, - tflops * approx_factor, - hbm_gbs * approx_factor, + (tflops_total if fused_shared else tflops) * approx_factor, + (hbm_gbs_total if fused_shared else hbm_gbs) * approx_factor, nvlink_gbs * approx_factor, reduction_us=t_reduction * 1e6, ) @@ -1835,6 +2178,17 @@ def fmt_perf_line( default=10.0, help="Clamp threshold for gate/up before SwiGLU; pass inf to disable", ) + parser.add_argument( + "--num-shared-experts", + type=int, + default=0, + help=( + "DeepSeek-style shared experts, each adding one routed intermediate " + "size, fused into the mega kernel as the SharedLinear1/2 phases (both " + "baselines run them serially as a dense FP8 MLP). 0 disables it; only " + "the default comparison mode uses this" + ), + ) parser.add_argument("--num-experts", type=int, default=384) parser.add_argument("--num-topk", type=int, default=6) parser.add_argument( @@ -1904,8 +2258,9 @@ def fmt_perf_line( "--layers", type=int, nargs="+", - default=[1, 2, 3, 4], - help="Accuracy layers to run with --accuracy (1..5); default: 1 2 3 4", + default=[1, 2, 3, 4, 6], + help="Accuracy layers to run with --accuracy (1..6); default: 1 2 3 4 6. " + "Layer 6 covers the fused shared expert (fused-out = routed + shared).", ) parser.add_argument( "--num-correctness-tests",