From 71f012a356b7d5d82d15d2354dbf02612a637eba Mon Sep 17 00:00:00 2001 From: Xinyi Liu <94362768+XinyiLiu577086410@users.noreply.github.com> Date: Thu, 18 Jun 2026 15:54:00 +0800 Subject: [PATCH] style: fix sm90 megamoe dispatch indentation --- .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 410 +++++++++--------- 1 file changed, 205 insertions(+), 205 deletions(-) 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 d487168..8f16fc8 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 @@ -245,215 +245,215 @@ sm90_fp8_mega_moe_impl(void* y, kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag>( workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync); - for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) - smem_expert_count[i] = 0; - ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + smem_expert_count[i] = 0; + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - const auto read_topk_idx = [&](const auto& process) { - #pragma unroll 1 - for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp; - i < num_tokens; - i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) { - const uint32_t token_offset = lane_idx / kNumTopk; - const uint32_t token_idx = i + token_offset; - int expert_idx = -1; - if (token_idx < num_tokens and lane_idx < kNumActivateLanes) { - expert_idx = static_cast( - __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); - if (expert_idx >= 0 and expert_idx < static_cast(kNumExperts)) - process(i * kNumTopk + lane_idx, static_cast(expert_idx)); + const auto read_topk_idx = [&](const auto& process) { + #pragma unroll 1 + for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; + i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) { + const uint32_t token_offset = lane_idx / kNumTopk; + const uint32_t token_idx = i + token_offset; + int expert_idx = -1; + if (token_idx < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0 and expert_idx < static_cast(kNumExperts)) + process(i * kNumTopk + lane_idx, static_cast(expert_idx)); + } + __syncwarp(); + } + }; + + // Count local outgoing token-topk entries per global expert. + read_topk_idx([&](const uint32_t&, const uint32_t& expert_idx) { + atomicAdd_block(smem_expert_count + expert_idx, 1); + }); + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Convert per-SM counts into global per-rank offsets. High 32 bits count arrived SMs. + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); + smem_expert_count[i] = static_cast( + ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Write source token-topk indices into destination ranks' local-expert tables. + read_topk_idx([&](const uint32_t& token_topk_idx, const uint32_t& expert_idx) { + const uint32_t dst_rank_idx = expert_idx / kNumExpertsPerRank; + const uint32_t dst_local_expert_idx = expert_idx - dst_rank_idx * kNumExpertsPerRank; + const uint32_t dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1); + const auto dst_ptr = workspace.get_src_token_topk_idx_ptr( + dst_local_expert_idx, sym_buffer.rank_idx, dst_slot_idx); + *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; + }); + + // Wait until all local SMs have finished filling local send-count/source-index data. + comm::grid_sync( + workspace, sm_idx, thread_idx, dispatch_sync); + + // Publish per-rank expert counts and summed ready/count words to destination ranks. + if (sm_idx == 0) { + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint32_t dst_rank_idx = i / kNumExpertsPerRank; + const uint32_t dst_local_expert_idx = i - dst_rank_idx * kNumExpertsPerRank; + const uint64_t expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffffu; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // All ranks must observe the complete count/source-index tables before pulling. + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync, + /* sync_prologue */ false, /* sync_epilogue */ true); + + auto scheduler = sched::SM90MegaMoEScheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, + kNumExpertsPerWave, + kNumSMs, kNumRanks, + kUseNMajorL2>(workspace); + scheduler.fetch_expert_recv_count(); + + if (sm_idx == 0 and cumulative_local_expert_recv_stats != nullptr) { + for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads) { + const auto num_recv_tokens = static_cast(*workspace.get_expert_recv_count_sum_ptr(i)); + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalDispatchWarps = kNumSMs * kNumDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalDispatchWarps) { + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++ current_expert_idx >= static_cast(kNumExpertsPerRank)) + break; + + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(static_cast(current_expert_idx)); + } + + if (current_expert_idx >= static_cast(kNumExpertsPerRank)) + break; + + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t j = i * 32 + lane_idx; + stored_rank_count[i] = j < kNumRanks ? + static_cast(*workspace.get_expert_recv_count_ptr(j, static_cast(current_expert_idx))) : 0; + } + } + + uint32_t current_rank_in_expert_idx = 0; + uint32_t remaining[kNumRanksPerLane]; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] = stored_rank_count[i]; + + uint32_t offset = 0; + const uint32_t token_idx_in_expert = token_idx - expert_start_idx; + uint32_t slot_idx = token_idx_in_expert; + uint32_t token_idx_in_rank = 0; + while (true) { + uint32_t num_actives_in_lane = 0; + uint32_t min_in_lane = 0xffffffffu; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + num_actives_in_lane += remaining[i] > 0; + if (remaining[i] > 0) + min_in_lane = cute::min(min_in_lane, remaining[i]); + } + const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane); + const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane); + const uint32_t num_round_tokens = length * num_active_ranks; + + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = slot_idx % num_active_ranks; + uint32_t num_seen_ranks = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0); + const uint32_t num_active_lanes = __popc(mask); + if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes) + current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); + num_seen_ranks += num_active_lanes; + } + token_idx_in_rank = offset + slot_idx / num_active_ranks; + break; + } + + slot_idx -= num_round_tokens; + offset += length; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] -= cute::min(remaining[i], length); + } + + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + static_cast(current_expert_idx), current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx - src_token_idx * kNumTopk; + const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + const uint32_t sf_pool_token_idx = + (expert_pool_block_offset + token_idx_in_expert / BLOCK_M) * SF_BLOCK_M + + (token_idx_in_expert % BLOCK_M); + + const auto remote_token_ptr = sym_buffer.map( + input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(); + #pragma unroll + for (uint32_t i = lane_idx; i < kNumTokenUint4; i += 32) + local_token_ptr[i] = remote_token_ptr[i]; + + const auto remote_sf_ptr = sym_buffer.map( + input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + #pragma unroll + for (uint32_t i = lane_idx; i < kNumSFValues; i += 32) + local_sf_ptr[i * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[i]; + + __syncwarp(); + if (lane_idx == 0) { + const float weight = *sym_buffer.map( + input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; + *workspace.get_token_src_metadata_ptr(pool_token_idx) = { + current_rank_in_expert_idx, src_token_idx, src_topk_idx}; + } + __syncwarp(); + __threadfence(); + __syncwarp(); + if (lane_idx == 0) { + ptx::red_add_rel( + workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1); } __syncwarp(); } - }; - - // Count local outgoing token-topk entries per global expert. - read_topk_idx([&](const uint32_t&, const uint32_t& expert_idx) { - atomicAdd_block(smem_expert_count + expert_idx, 1); - }); - ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - - // Convert per-SM counts into global per-rank offsets. High 32 bits count arrived SMs. - for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { - const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); - smem_expert_count[i] = static_cast( - ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); - } - ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - - // Write source token-topk indices into destination ranks' local-expert tables. - read_topk_idx([&](const uint32_t& token_topk_idx, const uint32_t& expert_idx) { - const uint32_t dst_rank_idx = expert_idx / kNumExpertsPerRank; - const uint32_t dst_local_expert_idx = expert_idx - dst_rank_idx * kNumExpertsPerRank; - const uint32_t dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1); - const auto dst_ptr = workspace.get_src_token_topk_idx_ptr( - dst_local_expert_idx, sym_buffer.rank_idx, dst_slot_idx); - *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; - }); - - // Wait until all local SMs have finished filling local send-count/source-index data. - comm::grid_sync( - workspace, sm_idx, thread_idx, dispatch_sync); - - // Publish per-rank expert counts and summed ready/count words to destination ranks. - if (sm_idx == 0) { - for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { - const uint32_t dst_rank_idx = i / kNumExpertsPerRank; - const uint32_t dst_local_expert_idx = i - dst_rank_idx * kNumExpertsPerRank; - const uint64_t expert_status = *workspace.get_expert_send_count_ptr(i); - *sym_buffer.map( - workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), - dst_rank_idx) = expert_status & 0xffffffffu; - ptx::atomic_add_sys( - sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), - expert_status); - } - } - ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - - // All ranks must observe the complete count/source-index tables before pulling. - comm::nvlink_barrier( - workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync, - /* sync_prologue */ false, /* sync_epilogue */ true); - - auto scheduler = sched::SM90MegaMoEScheduler< - BLOCK_M, BLOCK_N, BLOCK_K, - L1_SHAPE_N, L1_SHAPE_K, - L2_SHAPE_N, L2_SHAPE_K, - kNumExpertsPerRank, - kNumExpertsPerWave, - kNumSMs, kNumRanks, - kUseNMajorL2>(workspace); - scheduler.fetch_expert_recv_count(); - - if (sm_idx == 0 and cumulative_local_expert_recv_stats != nullptr) { - for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads) { - const auto num_recv_tokens = static_cast(*workspace.get_expert_recv_count_sum_ptr(i)); - ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); - } - } - ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); - - constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); - int current_expert_idx = -1; - uint32_t stored_rank_count[kNumRanksPerLane] = {}; - uint32_t expert_start_idx = 0, expert_end_idx = 0; - uint32_t expert_pool_block_offset = 0; - - constexpr uint32_t kNumGlobalDispatchWarps = kNumSMs * kNumDispatchWarps; - for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalDispatchWarps) { - int old_expert_idx = current_expert_idx; - while (token_idx >= expert_end_idx) { - if (++ current_expert_idx >= static_cast(kNumExpertsPerRank)) - break; - - expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); - expert_start_idx = expert_end_idx; - expert_end_idx += scheduler.get_num_tokens(static_cast(current_expert_idx)); - } - - if (current_expert_idx >= static_cast(kNumExpertsPerRank)) - break; - - if (old_expert_idx != current_expert_idx) { - old_expert_idx = current_expert_idx; - #pragma unroll - for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { - const uint32_t j = i * 32 + lane_idx; - stored_rank_count[i] = j < kNumRanks ? - static_cast(*workspace.get_expert_recv_count_ptr(j, static_cast(current_expert_idx))) : 0; - } - } - - uint32_t current_rank_in_expert_idx = 0; - uint32_t remaining[kNumRanksPerLane]; - #pragma unroll - for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) - remaining[i] = stored_rank_count[i]; - - uint32_t offset = 0; - const uint32_t token_idx_in_expert = token_idx - expert_start_idx; - uint32_t slot_idx = token_idx_in_expert; - uint32_t token_idx_in_rank = 0; - while (true) { - uint32_t num_actives_in_lane = 0; - uint32_t min_in_lane = 0xffffffffu; - #pragma unroll - for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { - num_actives_in_lane += remaining[i] > 0; - if (remaining[i] > 0) - min_in_lane = cute::min(min_in_lane, remaining[i]); - } - const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane); - const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane); - const uint32_t num_round_tokens = length * num_active_ranks; - - if (slot_idx < num_round_tokens) { - const uint32_t slot_idx_in_round = slot_idx % num_active_ranks; - uint32_t num_seen_ranks = 0; - #pragma unroll - for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { - const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0); - const uint32_t num_active_lanes = __popc(mask); - if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes) - current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); - num_seen_ranks += num_active_lanes; - } - token_idx_in_rank = offset + slot_idx / num_active_ranks; - break; - } - - slot_idx -= num_round_tokens; - offset += length; - #pragma unroll - for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) - remaining[i] -= cute::min(remaining[i], length); - } - - const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( - static_cast(current_expert_idx), current_rank_in_expert_idx, token_idx_in_rank); - const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; - const uint32_t src_topk_idx = src_token_topk_idx - src_token_idx * kNumTopk; - const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; - const uint32_t sf_pool_token_idx = - (expert_pool_block_offset + token_idx_in_expert / BLOCK_M) * SF_BLOCK_M + - (token_idx_in_expert % BLOCK_M); - - const auto remote_token_ptr = sym_buffer.map( - input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), - current_rank_in_expert_idx); - const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(); - #pragma unroll - for (uint32_t i = lane_idx; i < kNumTokenUint4; i += 32) - local_token_ptr[i] = remote_token_ptr[i]; - - const auto remote_sf_ptr = sym_buffer.map( - input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), - current_rank_in_expert_idx); - const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); - #pragma unroll - for (uint32_t i = lane_idx; i < kNumSFValues; i += 32) - local_sf_ptr[i * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[i]; - - __syncwarp(); - if (lane_idx == 0) { - const float weight = *sym_buffer.map( - input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, - current_rank_in_expert_idx); - *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; - *workspace.get_token_src_metadata_ptr(pool_token_idx) = { - current_rank_in_expert_idx, src_token_idx, src_topk_idx}; - } - __syncwarp(); - __threadfence(); - __syncwarp(); - if (lane_idx == 0) { - ptx::red_add_rel( - workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1); - } - __syncwarp(); - } } else if (thread_idx < kNumDispatchThreads + kNumTMAThreads) { if constexpr (BLOCK_M == 128) { if (warp_idx == kNumDispatchWarps) {