diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp index c449fc5b21..1fa246b2a8 100644 --- a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -44,10 +44,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}; @@ -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; diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index 06e33d0270..cc3f5e7445 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -242,11 +242,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, 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..d3e6db5970 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 @@ -273,6 +273,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 +295,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; @@ -971,7 +972,10 @@ sm90_fp8_mega_moe_impl(void* y, // ---------------- 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, @@ -1365,38 +1369,43 @@ sm90_fp8_mega_moe_impl(void* y, 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>(); + 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>(); - #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]; + const uint32_t final_base = slab * kSwapSlabAccumStride; + #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 +1475,66 @@ 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) { #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(); @@ -1640,78 +1653,76 @@ 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 = *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] = + 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 = *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] = + 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; + 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; 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; + sf_base_ptr[(n_block_idx * 2u + 0u) * kNumPaddedSFPoolTokens + token_idx] = sf0; + sf_base_ptr[(n_block_idx * 2u + 1u) * 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; } } @@ -2013,32 +2024,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 +2085,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,11 +2130,20 @@ 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);