#pragma once

#include <cutlass/arch/barrier.h>
#include <cutlass/cutlass.h>

#include <deep_gemm/common/math.cuh>
#include <deep_gemm/common/exception.cuh>
#include <deep_gemm/common/tma_copy.cuh>
#include <deep_gemm/common/tma_utils.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/common/utils.cuh>
#include <deep_gemm/comm/barrier.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include <deep_gemm/mma/sm90.cuh>
#include <deep_gemm/ptx/ld_st.cuh>
#include <deep_gemm/ptx/tma.cuh>
#include <deep_gemm/ptx/utils.cuh>
#include <deep_gemm/ptx/wgmma.cuh>
#include <deep_gemm/scheduler/sm90_mega_moe.cuh>

namespace deep_gemm {

template <
    uint32_t kNumMaxTokensPerRank,
    uint32_t kHidden, uint32_t kIntermediateHidden,
    uint32_t kNumExperts, uint32_t kNumTopk,
    uint32_t kNumExpertsPerWave,
    uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
    uint32_t kStoreBlockM,
    uint32_t kNumMaxPoolTokens,
    uint32_t kNumPaddedSFPoolTokens,
    uint32_t kNumStages,
    uint32_t kNumDispatchThreads, uint32_t kNumTMAThreads,
    uint32_t kNumMathThreads,
    bool kCooperativeMode, bool kUseNMajorL2,
    uint32_t kNumSMs, uint32_t kNumRanks,
    float kActivationClamp, bool kFastMath,
    uint32_t L1_SHAPE_N = kIntermediateHidden * 2,
    uint32_t L1_SHAPE_K = kHidden,
    uint32_t L2_SHAPE_N = kHidden,
    uint32_t L2_SHAPE_K = kIntermediateHidden,
    uint32_t kNumThreads = kNumDispatchThreads + kNumTMAThreads + kNumMathThreads,
    uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks>
CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void
sm90_fp8_mega_moe_impl(void* y,
                       int* cumulative_local_expert_recv_stats,
                       const uint32_t num_tokens,
                       const __grid_constant__ layout::SymBuffer<kNumRanks> sym_buffer,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf,
                        const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights,
                        const void* l1_weights_sf,
                        const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output,
                        const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts,
                        const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf,
                        const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights,
                        const void* l2_weights_sf,
                        void* l1_accum_debug) {
    DG_STATIC_ASSERT(kNumThreads == 384, "SM90 MegaMoE expects 384 threads");
    DG_STATIC_ASSERT(BLOCK_N == 128, "SM90 MegaMoE expects BLOCK_N=128");
    DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128");
    DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks");

#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
    DG_STATIC_ASSERT(kNumDispatchThreads == 64, "SM90 MegaMoE expects 64 dispatch threads");
    DG_STATIC_ASSERT(kNumTMAThreads == 64, "SM90 MegaMoE expects 64 TMA threads");
    DG_STATIC_ASSERT(kNumMathThreads == 256, "SM90 MegaMoE expects 256 math threads");
    DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of top-k experts");
    DG_STATIC_ASSERT(kHidden % 128 == 0, "SM90 activation SF expects per-128K groups");
    DG_STATIC_ASSERT(kHidden % sizeof(uint4) == 0, "Token copy expects 16-byte alignment");

    using Barrier = cutlass::arch::ClusterTransactionBarrier;
    using a_dtype_t = __nv_fp8_e4m3;
    using b_dtype_t = __nv_fp8_e4m3;
    using WGMMA = typename mma::sm90::FP8MMASelector<BLOCK_N>::type;

    constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32;
    constexpr uint32_t kNumTMAWarps = kNumTMAThreads / 32;
    constexpr uint32_t kNumMathWarps = kNumMathThreads / 32;
    constexpr uint32_t kMathWarpStart = kNumDispatchWarps + kNumTMAWarps;
    constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk;
    constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk;
    constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, 128u);
    constexpr uint32_t kNumMaxPoolBlocks = kNumMaxPoolTokens / layout::kMinCandidateBlockM;
    constexpr uint32_t kNumTokenUint4 = kHidden / sizeof(uint4);
    constexpr uint32_t kNumSFValues = kHidden / 128;
    constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t);
    constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t);
    constexpr uint32_t kNumL1WeightSFGroupsN = L1_SHAPE_N / 128;
    constexpr uint32_t kNumL1WeightSFGroupsK = L1_SHAPE_K / 128;
    constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2;
    constexpr uint32_t kL2ActsGranK = 64;
    constexpr uint32_t kMathBarrierIdx = 2;
    DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts");
    DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size");
    DG_STATIC_ASSERT(BLOCK_N == WGMMA::N and BLOCK_K % WGMMA::K == 0, "Invalid WGMMA tile shape");
    DG_STATIC_ASSERT(L1_OUT_BLOCK_N == kL2ActsGranK, "SM90 Phase 4 expects per-64 L2 activation SF");

    const uint32_t thread_idx = threadIdx.x;
    const uint32_t sm_idx = blockIdx.x;
    const uint32_t warp_idx = thread_idx / 32;
    const uint32_t lane_idx = ptx::get_lane_idx();

    if (warp_idx == kNumDispatchWarps and cute::elect_one_sync()) {
        cute::prefetch_tma_descriptor(&tensor_map_l1_acts);
        cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf);
        cute::prefetch_tma_descriptor(&tensor_map_l1_weights);
    }
    __syncwarp();

    constexpr uint32_t kDispatchBarrierIdx = 1;
    constexpr uint32_t kDispatchGridSyncIndex = 0;
    constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1;
    constexpr uint32_t kBeforeDispatchPullBarrierTag = 2;
    const auto dispatch_sync = []() {
        ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
    };

    const auto workspace = layout::Workspace(
        sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk);

    constexpr auto fp8_token_layout = layout::Data(kHidden);
    constexpr auto fp8_sf_layout = layout::Data(kHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
    constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden);
    constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / kL2ActsGranK * static_cast<uint32_t>(sizeof(float)), false);
    constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(int64_t)), false);
    constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(float)), false);
    constexpr auto l1_topk_weights_layout = layout::Data(static_cast<uint32_t>(sizeof(float)), false);

    const auto input_token_buffer = layout::Buffer(
        fp8_token_layout, 1, kNumMaxTokensPerRank,
        workspace.get_end_ptr());
    const auto input_sf_buffer = layout::Buffer(
        fp8_sf_layout, 1, kNumMaxTokensPerRank,
        input_token_buffer.get_end_ptr());
    const auto input_topk_idx_buffer = layout::Buffer(
        input_topk_idx_layout, 1, kNumMaxTokensPerRank,
        input_sf_buffer.get_end_ptr());
    const auto input_topk_weights_buffer = layout::Buffer(
        input_topk_weights_layout, 1, kNumMaxTokensPerRank,
        input_topk_idx_buffer.get_end_ptr());

    const auto l1_token_buffer = layout::Buffer(
        fp8_token_layout, 1, kNumMaxPoolTokens,
        input_topk_weights_buffer.get_end_ptr());
    const auto l1_sf_buffer = layout::Buffer(
        fp8_sf_layout, 1, kNumPaddedSFPoolTokens,
        l1_token_buffer.get_end_ptr());
    const auto l1_topk_weights_buffer = layout::Buffer(
        l1_topk_weights_layout, 1, kNumMaxPoolTokens,
        l1_sf_buffer.get_end_ptr());
    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());
    (void)l2_sf_buffer;

    constexpr uint32_t kSharedMemoryAlignment = 1024;
    extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[];

    constexpr uint32_t SMEM_EXPERT_COUNT_SIZE =
        math::constexpr_align<uint32_t>(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment);
    constexpr uint32_t SMEM_SEND_BUFFER_SIZE =
        math::constexpr_align<uint32_t>(fp8_token_layout.get_num_bytes() * kNumDispatchWarps, kSharedMemoryAlignment);
    constexpr uint32_t SMEM_CD_L1_SIZE =
        (kCooperativeMode ? 2u : 1u) * kStoreBlockM * (BLOCK_N / 2) * 2;
    constexpr uint32_t SMEM_CD_L2_SIZE = BLOCK_M * BLOCK_N * static_cast<uint32_t>(sizeof(nv_bfloat16));
    constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE;
    constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(a_dtype_t);
    constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(b_dtype_t);
    constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(float);
    DG_STATIC_ASSERT(SMEM_A_SIZE_PER_STAGE % kSharedMemoryAlignment == 0, "Invalid A smem alignment");
    DG_STATIC_ASSERT(SMEM_B_SIZE_PER_STAGE % kSharedMemoryAlignment == 0, "Invalid B smem alignment");

    const auto smem_expert_count = reinterpret_cast<uint32_t*>(smem_buffer);
    auto smem_gemm_base = math::advance_ptr<uint8_t>(
        smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE);
    auto smem_a = utils::PatternVisitor([=](const uint32_t& i) {
        return math::advance_ptr<a_dtype_t>(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE);
    });
    auto smem_b = utils::PatternVisitor([=](const uint32_t& i) {
        return math::advance_ptr<b_dtype_t>(smem_gemm_base, SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
    });
    auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) {
        return math::advance_ptr<float>(smem_gemm_base,
            SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE) + i * SMEM_SFA_SIZE_PER_STAGE);
    });
    auto barrier_start_ptr = reinterpret_cast<Barrier*>(math::advance_ptr<uint8_t>(
        smem_gemm_base,
        SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE)));
    auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) {
        return barrier_start_ptr + i;
    });
    auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) {
        return barrier_start_ptr + kNumStages + i;
    });

    if constexpr (BLOCK_M == 128) {
        if (sm_idx == 0) {
            auto debug_ptr = reinterpret_cast<float*>(l1_accum_debug);
            for (uint32_t i = thread_idx; i < 128 * 128; i += kNumThreads)
                debug_ptr[i] = 0.0f;
        }

        if (warp_idx == kNumDispatchWarps + 1 and cute::elect_one_sync()) {
            #pragma unroll
            for (uint32_t i = 0; i < kNumStages; ++ i) {
                full_barriers[i]->init(1);
                empty_barriers[i]->init(kNumMathWarps);
            }
            cutlass::arch::fence_barrier_init();
        }
        __syncthreads();
    }

    if (thread_idx < kNumDispatchThreads) {
        // Clean local dispatch workspace from any previous call. Barrier state is intentionally
        // persistent and must not be reset because grid/NVLink barriers use phase counters.
        for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
            *workspace.get_expert_send_count_ptr(i) = 0;
        for (uint32_t i = thread_idx; i < kNumRanks * kNumExpertsPerRank; i += kNumDispatchThreads)
            *workspace.get_expert_recv_count_ptr(i / kNumExpertsPerRank, i % kNumExpertsPerRank) = 0;
        for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads)
            *workspace.get_expert_recv_count_sum_ptr(i) = 0;
        for (uint32_t i = thread_idx; i < kNumMaxPoolBlocks; i += kNumDispatchThreads) {
            *workspace.get_l1_arrival_count_ptr(i) = 0;
            *workspace.get_l2_arrival_mask_ptr(i) = 0;
        }
        comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
                             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);

    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<int>(
                    __ldg(input_topk_idx_buffer.get_base_ptr<int64_t>() + i * kNumTopk + lane_idx));
                if (expert_idx >= 0 and expert_idx < static_cast<int>(kNumExperts))
                    process(i * kNumTopk + lane_idx, static_cast<uint32_t>(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<uint64_t>(smem_expert_count[i]);
        smem_expert_count[i] = static_cast<uint32_t>(
            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<kNumSMs, kDispatchGridSyncIndex>(
        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<kNumRanks, kNumSMs, kNumDispatchThreads,
                         kDispatchGridSyncIndex, kBeforeDispatchPullBarrierTag>(
        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<uint32_t>(*workspace.get_expert_recv_count_sum_ptr(i));
            ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast<int>(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<int>(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<uint32_t>(current_expert_idx));
        }

        if (current_expert_idx >= static_cast<int>(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<uint32_t>(*workspace.get_expert_recv_count_ptr(j, static_cast<uint32_t>(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<uint32_t>(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<uint4>(),
            current_rank_in_expert_idx);
        const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr<uint4>();
        #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<float>(),
            current_rank_in_expert_idx);
        const auto local_sf_ptr = l1_sf_buffer.get_base_ptr<float>();
        #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<float>() + src_token_topk_idx,
                current_rank_in_expert_idx);
            *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr<float>() = 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) {
                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);

                uint32_t stage_idx = 0, phase = 0;
                auto advance_pipeline = [&](uint32_t& k_block_idx) {
                    ++ k_block_idx;
                    stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
                    phase ^= stage_idx == 0;
                };

                scheduler.for_each_block([&](const sched::SM90BlockPhase& 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) {
                    if (block_phase != sched::SM90BlockPhase::Linear1)
                        return;

                    const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx;
                    const uint32_t valid_m = scheduler.get_valid_m();
                    const auto arrival_ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx);
                    while (ptx::ld_acq(arrival_ptr) != valid_m);

                    #pragma unroll 1
                    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()) {
                            auto& full_barrier = *full_barriers[stage_idx];
                            const uint32_t k_idx = k_block_idx * BLOCK_K;
                            const uint32_t m_idx = pool_block_idx * BLOCK_M;
                            const uint32_t n_idx = local_expert_idx * L1_SHAPE_N + n_block_idx * BLOCK_N;
                            const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M;
                            tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(
                                &tensor_map_l1_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx);
                            tma::copy<SF_BLOCK_M, 1, 0>(
                                &tensor_map_l1_acts_sf, &full_barrier, smem_sfa[stage_idx], sfa_m_idx, k_block_idx);
                            tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode>(
                                &tensor_map_l1_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx);
                            full_barrier.arrive_and_expect_tx(
                                SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
                        }
                        __syncwarp();
                    }
                });
            }
        }
    } else {
        if constexpr (BLOCK_M == 128) {
            const uint32_t math_warp_idx = warp_idx - kMathWarpStart;
            const uint32_t math_wg_idx = math_warp_idx / 4;
            const uint32_t warp_idx_in_wg = math_warp_idx % 4;
            const uint32_t row_idx = lane_idx / 4;
            const uint32_t col_idx = lane_idx % 4;
            const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx;
            const uint32_t r_1 = r_0 + 8;
            const auto l1_weights_sf_ptr = reinterpret_cast<const float*>(l1_weights_sf);
            auto l1_accum_debug_ptr = reinterpret_cast<float*>(l1_accum_debug);
            auto l2_acts_ptr = l2_token_buffer.get_base_ptr<__nv_fp8_e4m3>();
            auto l2_acts_sf_ptr = l2_sf_buffer.get_base_ptr<float>();

            const auto get_l1_weight_sf_group = [](const uint32_t& interleaved_n) {
                constexpr uint32_t kInterleaveGran = 8;
                const uint32_t pair_group = interleaved_n / (2 * kInterleaveGran);
                const uint32_t offset = interleaved_n - pair_group * (2 * kInterleaveGran);
                const uint32_t natural_n = offset < kInterleaveGran ?
                    pair_group * kInterleaveGran + offset :
                    (L1_SHAPE_N / 2) + pair_group * kInterleaveGran + offset - kInterleaveGran;
                return natural_n / 128;
            };

            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);

            uint32_t stage_idx = 0, phase = 0;
            auto advance_pipeline = [&](uint32_t& k_block_idx) {
                ++ k_block_idx;
                stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
                phase ^= stage_idx == 0;
            };

            scheduler.for_each_block([&](const sched::SM90BlockPhase& 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) {
                if (block_phase != sched::SM90BlockPhase::Linear1)
                    return;

                const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx;
                const uint32_t valid_m = scheduler.get_valid_m();
                float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};

                const auto empty_barrier_arrive = [&](const uint32_t& s) {
                    if (lane_idx == 0)
                        empty_barriers[s]->arrive();
                };

                #pragma unroll 1
                for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) {
                    full_barriers[stage_idx]->wait(phase);

                    const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_0);
                    const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_1);

                    #pragma unroll
                    for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
                        ptx::warpgroup_fence_operand(accum[i]);
                    ptx::warpgroup_arrive();
                    #pragma unroll
                    for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
                        auto desc_a = mma::sm90::make_smem_desc(
                            smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1);
                        auto desc_b = mma::sm90::make_smem_desc(
                            smem_b[stage_idx] + k * WGMMA::K, 1);
                        WGMMA::wgmma(desc_a, desc_b, accum, k);
                    }
                    ptx::warpgroup_commit_batch();
                    #pragma unroll
                    for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
                        ptx::warpgroup_fence_operand(accum[i]);
                    ptx::warpgroup_wait<0>();

                    empty_barrier_arrive(stage_idx);

                    #pragma unroll
                    for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
                        const uint32_t col_0 = i * 8 + col_idx * 2;
                        const uint32_t col_1 = col_0 + 1;
                        const uint32_t n_0 = n_block_idx * BLOCK_N + col_0;
                        const uint32_t n_1 = n_block_idx * BLOCK_N + col_1;
                        const auto sf_base = l1_weights_sf_ptr +
                            local_expert_idx * kNumL1WeightSFGroupsN * kNumL1WeightSFGroupsK;
                        const float scale_b_0 = __ldg(
                            sf_base + get_l1_weight_sf_group(n_0) * kNumL1WeightSFGroupsK + k_block_idx);
                        const float scale_b_1 = __ldg(
                            sf_base + get_l1_weight_sf_group(n_1) * kNumL1WeightSFGroupsK + k_block_idx);
                        final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
                        final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
                        final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
                        final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
                    }
                }

                const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0;
                const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1;
                const uint32_t pool_token_idx_0 = pool_block_idx * BLOCK_M + row_0;
                const uint32_t pool_token_idx_1 = pool_block_idx * BLOCK_M + row_1;
                const float topk_weight_0 = row_0 < valid_m ?
                    *l1_topk_weights_buffer.get_data_buffer(pool_token_idx_0).get_base_ptr<float>() : 0.0f;
                const float topk_weight_1 = row_1 < valid_m ?
                    *l1_topk_weights_buffer.get_data_buffer(pool_token_idx_1).get_base_ptr<float>() : 0.0f;

                const auto apply_swiglu = [&](float gate, float up, const float& topk_weight) {
                    if constexpr (kActivationClamp != cute::numeric_limits<float>::infinity()) {
                        gate = fminf(fmaxf(gate, -kActivationClamp), kActivationClamp);
                        up = fminf(fmaxf(up, -kActivationClamp), kActivationClamp);
                    }

                    const float denom = 1.0f + (kFastMath ? __expf(-gate) : expf(-gate));
                    const float silu = gate * (kFastMath ? math::fast_rcp(denom) : 1.0f / denom);
                    return silu * up * topk_weight;
                };

                float local_amax_0 = 0.0f, local_amax_1 = 0.0f;
                #pragma unroll
                for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; i += 2) {
                    const float v_00 = apply_swiglu(final_accum[i * 4 + 0], final_accum[(i + 1) * 4 + 0], topk_weight_0);
                    const float v_01 = apply_swiglu(final_accum[i * 4 + 1], final_accum[(i + 1) * 4 + 1], topk_weight_0);
                    const float v_10 = apply_swiglu(final_accum[i * 4 + 2], final_accum[(i + 1) * 4 + 2], topk_weight_1);
                    const float v_11 = apply_swiglu(final_accum[i * 4 + 3], final_accum[(i + 1) * 4 + 3], topk_weight_1);
                    local_amax_0 = fmaxf(local_amax_0, fmaxf(fabsf(v_00), fabsf(v_01)));
                    local_amax_1 = fmaxf(local_amax_1, fmaxf(fabsf(v_10), fabsf(v_11)));
                }

                const float row_amax_0 = math::warp_reduce<4, false>(local_amax_0, math::ReduceMax<float>());
                const float row_amax_1 = math::warp_reduce<4, false>(local_amax_1, math::ReduceMax<float>());
                const float sf_0 = fmaxf(row_amax_0 / 448.0f, 1.0e-12f);
                const float sf_1 = fmaxf(row_amax_1 / 448.0f, 1.0e-12f);
                const float sf_inv_0 = kFastMath ? math::fast_rcp(sf_0) : 1.0f / sf_0;
                const float sf_inv_1 = kFastMath ? math::fast_rcp(sf_1) : 1.0f / sf_1;

                if (col_idx == 0) {
                    const uint32_t sf_pool_token_idx_0 = pool_block_idx * SF_BLOCK_M + row_0;
                    const uint32_t sf_pool_token_idx_1 = pool_block_idx * SF_BLOCK_M + row_1;
                    if (row_0 < valid_m)
                        l2_acts_sf_ptr[n_block_idx * kNumPaddedSFPoolTokens + sf_pool_token_idx_0] = sf_0;
                    if (row_1 < valid_m)
                        l2_acts_sf_ptr[n_block_idx * kNumPaddedSFPoolTokens + sf_pool_token_idx_1] = sf_1;
                }

                const uint32_t out_n_base = n_block_idx * L1_OUT_BLOCK_N;
                #pragma unroll
                for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; i += 2) {
                    const uint32_t out_col = out_n_base + (i / 2) * 8 + col_idx * 2;
                    if (row_0 < valid_m) {
                        const float v_0 = apply_swiglu(final_accum[i * 4 + 0], final_accum[(i + 1) * 4 + 0], topk_weight_0);
                        const float v_1 = apply_swiglu(final_accum[i * 4 + 1], final_accum[(i + 1) * 4 + 1], topk_weight_0);
                        l2_acts_ptr[pool_token_idx_0 * L2_SHAPE_K + out_col + 0] = __nv_fp8_e4m3(v_0 * sf_inv_0);
                        l2_acts_ptr[pool_token_idx_0 * L2_SHAPE_K + out_col + 1] = __nv_fp8_e4m3(v_1 * sf_inv_0);
                    }
                    if (row_1 < valid_m) {
                        const float v_0 = apply_swiglu(final_accum[i * 4 + 2], final_accum[(i + 1) * 4 + 2], topk_weight_1);
                        const float v_1 = apply_swiglu(final_accum[i * 4 + 3], final_accum[(i + 1) * 4 + 3], topk_weight_1);
                        l2_acts_ptr[pool_token_idx_1 * L2_SHAPE_K + out_col + 0] = __nv_fp8_e4m3(v_0 * sf_inv_1);
                        l2_acts_ptr[pool_token_idx_1 * L2_SHAPE_K + out_col + 1] = __nv_fp8_e4m3(v_1 * sf_inv_1);
                    }
                }

                __threadfence();
                ptx::sync_aligned(kNumMathThreads, kMathBarrierIdx);
                if (math_warp_idx == 0 and cute::elect_one_sync()) {
                    ptx::red_or_rel_gpu(
                        workspace.get_l2_arrival_mask_ptr(pool_block_idx),
                        1ull << n_block_idx);
                }
                __syncwarp();

                if (local_expert_idx == 0 and m_block_idx == 0 and n_block_idx == 0) {
                    #pragma unroll
                    for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
                        const uint32_t col = i * 8 + col_idx * 2;
                        if (row_0 < valid_m) {
                            l1_accum_debug_ptr[row_0 * BLOCK_N + col + 0] = final_accum[i * 4 + 0];
                            l1_accum_debug_ptr[row_0 * BLOCK_N + col + 1] = final_accum[i * 4 + 1];
                        }
                        if (row_1 < valid_m) {
                            l1_accum_debug_ptr[row_1 * BLOCK_N + col + 0] = final_accum[i * 4 + 2];
                            l1_accum_debug_ptr[row_1 * BLOCK_N + col + 1] = final_accum[i * 4 + 3];
                        }
                    }
                }
            });
        }
    }
#endif
    return;
}

} // namespace deep_gemm
