from __future__ import annotations

import os
import ninja
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

os.environ["PATH"] = ninja.BIN_DIR + os.pathsep + os.environ.get("PATH", "")

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <stdint.h>

// ============================================================================
// PTX Helpers for SM120 (Blackwell)
// ============================================================================

__device__ __forceinline__ uint32_t smem_u32_ptr(const void* ptr) {
    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
}

__device__ __forceinline__ uint64_t make_policy_evict_first() {
    uint64_t policy;
    asm("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" : "=l"(policy));
    return policy;
}

__device__ __forceinline__ uint64_t make_policy_evict_last() {
    uint64_t policy;
    asm("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(policy));
    return policy;
}

__device__ __forceinline__ void cp_async_cg_16(uint32_t dst_smem, const void* src_gmem, uint64_t policy) {
    asm volatile("cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, %2;\n"
                 :: "r"(dst_smem), "l"(src_gmem), "l"(policy));
}

__device__ __forceinline__ void cp_async_ca_16(uint32_t dst_smem, const void* src_gmem, uint64_t policy) {
    asm volatile("cp.async.ca.shared.global.L2::cache_hint [%0], [%1], 16, %2;\n"
                 :: "r"(dst_smem), "l"(src_gmem), "l"(policy));
}

__device__ __forceinline__ void cp_async_commit() {
    asm volatile("cp.async.commit_group;\n" ::);
}

template <int N>
__device__ __forceinline__ void cp_async_wait_group() {
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}

__device__ __forceinline__ void ldmatrix_x4(
    uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3,
    uint32_t smem_ptr)
{
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
        : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
        : "r"(smem_ptr)
    );
}

__device__ __forceinline__ void mma_m16n8k16_bf16(
    float& d0, float& d1, float& d2, float& d3,
    uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
    uint32_t b0, uint32_t b1)
{
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)
    );
}

__device__ __forceinline__ float2 u32_to_f2(uint32_t u) {
    __nv_bfloat162 bf2;
    asm("mov.b32 %0, %1;" : "=r"(*reinterpret_cast<uint32_t*>(&bf2)) : "r"(u));
    return __bfloat1622float2(bf2);
}

__device__ __forceinline__ float dot8_bf16_u4(uint4 a, uint4 b) {
    float2 fa0 = u32_to_f2(a.x);
    float2 fb0 = u32_to_f2(b.x);
    float2 fa1 = u32_to_f2(a.y);
    float2 fb1 = u32_to_f2(b.y);
    float2 fa2 = u32_to_f2(a.z);
    float2 fb2 = u32_to_f2(b.z);
    float2 fa3 = u32_to_f2(a.w);
    float2 fb3 = u32_to_f2(b.w);
    return (fa0.x * fb0.x + fa0.y * fb0.y) +
           (fa1.x * fb1.x + fa1.y * fb1.y) +
           (fa2.x * fb2.x + fa2.y * fb2.y) +
           (fa3.x * fb3.x + fa3.y * fb3.y);
}

// ============================================================================
// T == 1 Ultra-Fast Decode Kernels (2 launches total for T == 1)
// ============================================================================

// Grid: (I / 8, top_k + n_shared) = (256, 9), Block: 128 threads (4 warps)
__global__ __launch_bounds__(128, 8)
void gemv_w1_silu_t1_kernel(
    const __nv_bfloat16* __restrict__ x,          // (1, H)
    const int64_t* __restrict__ expert_ids,       // (1, top_k)
    const __nv_bfloat16* __restrict__ w1_routed,  // (E, 2*I, H)
    const __nv_bfloat16* __restrict__ w1_shared,  // (n_shared, 2*I, H)
    __nv_bfloat16* __restrict__ h_out,            // (top_k + n_shared, I)
    int* __restrict__ tile_done,                  // (H / 8)
    int top_k,
    int H,
    int I)
{
    __shared__ __align__(16) uint4 s_x[512]; // 4096 bf16 = 512 uint4 (8 KB)

    int tid = threadIdx.x;
    int warp_id = tid >> 5;
    int lane = tid & 31;
    int slot = blockIdx.y;

    // Zero tile_done counters for kernel 2
    if (slot == 0 && blockIdx.x < 4) {
        int idx = blockIdx.x * 128 + tid;
        if (idx < (H >> 3)) {
            tile_done[idx] = 0;
        }
    }

    const uint4* x_u4 = reinterpret_cast<const uint4*>(x);
    int h_u4 = H >> 3;
    #pragma unroll 4
    for (int idx = tid; idx < h_u4; idx += 128) {
        s_x[idx] = x_u4[idx];
    }
    __syncthreads();

    const __nv_bfloat16* w1_e;
    if (slot < top_k) {
        int64_t e = expert_ids[slot];
        w1_e = w1_routed + (size_t)e * (2 * I) * H;
    } else {
        int s = slot - top_k;
        w1_e = w1_shared + (size_t)s * (2 * I) * H;
    }

    int i0 = blockIdx.x * 8 + warp_id * 2;
    int i1 = i0 + 1;

    const uint4* g0_ptr = reinterpret_cast<const uint4*>(w1_e + (size_t)i0 * H);
    const uint4* g1_ptr = reinterpret_cast<const uint4*>(w1_e + (size_t)i1 * H);
    const uint4* u0_ptr = reinterpret_cast<const uint4*>(w1_e + (size_t)(I + i0) * H);
    const uint4* u1_ptr = reinterpret_cast<const uint4*>(w1_e + (size_t)(I + i1) * H);

    float acc_g0 = 0.0f, acc_g1 = 0.0f, acc_u0 = 0.0f, acc_u1 = 0.0f;

    #pragma unroll 8
    for (int idx = lane; idx < h_u4; idx += 32) {
        uint4 xv = s_x[idx];
        uint4 gv0 = __ldg(&g0_ptr[idx]);
        uint4 gv1 = __ldg(&g1_ptr[idx]);
        uint4 uv0 = __ldg(&u0_ptr[idx]);
        uint4 uv1 = __ldg(&u1_ptr[idx]);
        acc_g0 += dot8_bf16_u4(xv, gv0);
        acc_g1 += dot8_bf16_u4(xv, gv1);
        acc_u0 += dot8_bf16_u4(xv, uv0);
        acc_u1 += dot8_bf16_u4(xv, uv1);
    }

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        acc_g0 += __shfl_down_sync(0xffffffff, acc_g0, offset);
        acc_g1 += __shfl_down_sync(0xffffffff, acc_g1, offset);
        acc_u0 += __shfl_down_sync(0xffffffff, acc_u0, offset);
        acc_u1 += __shfl_down_sync(0xffffffff, acc_u1, offset);
    }

    if (lane == 0) {
        float h0 = (acc_g0 / (1.0f + __expf(-acc_g0))) * acc_u0;
        float h1 = (acc_g1 / (1.0f + __expf(-acc_g1))) * acc_u1;
        __nv_bfloat162 res = __float22bfloat162_rn(make_float2(h0, h1));
        *reinterpret_cast<__nv_bfloat162*>(&h_out[slot * I + i0]) = res;
    }
}

// Grid: (H / 8, top_k + n_shared) = (512, 9), Block: 128 threads (4 warps)
__global__ __launch_bounds__(128, 8)
void gemv_w2_reduce_t1_kernel(
    const __nv_bfloat16* __restrict__ h_in,           // (top_k + n_shared, I)
    const int64_t* __restrict__ expert_ids,           // (1, top_k)
    const __nv_bfloat16* __restrict__ expert_weights, // (1, top_k)
    const __nv_bfloat16* __restrict__ w2_routed,      // (E, H, I)
    const __nv_bfloat16* __restrict__ w2_shared,      // (n_shared, H, I)
    float* __restrict__ partial_out,                  // (top_k + n_shared, H)
    int* __restrict__ tile_done,                      // (H / 8)
    __nv_bfloat16* __restrict__ out,                  // (1, H)
    int top_k,
    int total_slots,
    int H,
    int I)
{
    __shared__ __align__(16) uint4 s_h[256]; // 2048 bf16 = 256 uint4 (4 KB)
    __shared__ float s_warp_out[8];
    __shared__ int s_is_last;

    int tid = threadIdx.x;
    int warp_id = tid >> 5;
    int lane = tid & 31;
    int slot = blockIdx.y;

    const uint4* h_u4 = reinterpret_cast<const uint4*>(h_in + slot * I);
    int i_u4 = I >> 3;
    #pragma unroll 2
    for (int idx = tid; idx < i_u4; idx += 128) {
        s_h[idx] = h_u4[idx];
    }
    __syncthreads();

    const __nv_bfloat16* w2_e;
    float weight = 1.0f;
    if (slot < top_k) {
        int64_t e = expert_ids[slot];
        w2_e = w2_routed + (size_t)e * H * I;
        weight = __bfloat162float(expert_weights[slot]);
    } else {
        int s = slot - top_k;
        w2_e = w2_shared + (size_t)s * H * I;
    }

    int col0 = blockIdx.x * 8 + warp_id * 2;
    int col1 = col0 + 1;

    const uint4* w0_ptr = reinterpret_cast<const uint4*>(w2_e + (size_t)col0 * I);
    const uint4* w1_ptr = reinterpret_cast<const uint4*>(w2_e + (size_t)col1 * I);

    float acc0 = 0.0f, acc1 = 0.0f;
    #pragma unroll 8
    for (int idx = lane; idx < i_u4; idx += 32) {
        uint4 hv = s_h[idx];
        uint4 wv0 = __ldg(&w0_ptr[idx]);
        uint4 wv1 = __ldg(&w1_ptr[idx]);
        acc0 += dot8_bf16_u4(hv, wv0);
        acc1 += dot8_bf16_u4(hv, wv1);
    }

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        acc0 += __shfl_down_sync(0xffffffff, acc0, offset);
        acc1 += __shfl_down_sync(0xffffffff, acc1, offset);
    }

    if (lane == 0) {
        s_warp_out[warp_id * 2 + 0] = acc0 * weight;
        s_warp_out[warp_id * 2 + 1] = acc1 * weight;
    }
    __syncthreads();

    if (tid < 8) {
        partial_out[slot * H + blockIdx.x * 8 + tid] = s_warp_out[tid];
    }
    __threadfence();

    if (tid == 0) {
        int prev = atomicAdd(&tile_done[blockIdx.x], 1);
        s_is_last = (prev == (total_slots - 1)) ? 1 : 0;
    }
    __syncthreads();

    if (s_is_last && tid < 4) {
        int c0 = blockIdx.x * 8 + tid * 2;
        int c1 = c0 + 1;
        float sum0 = 0.0f, sum1 = 0.0f;
        #pragma unroll
        for (int s = 0; s < 9; ++s) {
            if (s < total_slots) {
                sum0 += partial_out[s * H + c0];
                sum1 += partial_out[s * H + c1];
            }
        }
        *reinterpret_cast<__nv_bfloat162*>(&out[c0]) = __float22bfloat162_rn(make_float2(sum0, sum1));
    }
}

// ============================================================================
// Single-CTA MoE Routing & Tile Table Builder (T > 1)
// ============================================================================

__global__ __launch_bounds__(512, 1)
void build_moe_metadata_kernel(
    const int64_t* __restrict__ expert_ids,           // (T, top_k)
    const __nv_bfloat16* __restrict__ expert_weights, // (T, top_k)
    int4* __restrict__ tile_table,                    // (max_tiles) -> {expert_id, start_row, m_valid, 0}
    int* __restrict__ sorted_token_ids,               // (T * top_k)
    float* __restrict__ sorted_weights,               // (T * top_k)
    int* __restrict__ d_num_tiles,                    // (1)
    int T,
    int E,
    int top_k,
    int n_shared,
    int block_m)
{
    __shared__ int s_counts[256];
    __shared__ int s_row_offsets[256];
    __shared__ int s_tile_offsets[256];
    __shared__ int s_write_cursor[256];
    __shared__ int s_warp_row_sums[8];
    __shared__ int s_warp_tile_sums[8];
    __shared__ int s_total_routed_tiles;

    int tid = threadIdx.x;
    if (tid < 256) {
        s_counts[tid] = 0;
    }
    __syncthreads();

    int total_routed = T * top_k;
    for (int idx = tid; idx < total_routed; idx += 512) {
        int e = static_cast<int>(expert_ids[idx]);
        atomicAdd(&s_counts[e], 1);
    }
    __syncthreads();

    // Prefix sum across E = 256 experts using first 8 warps (256 threads)
    if (tid < 256) {
        int warp_id = tid >> 5;
        int lane = tid & 31;
        int count = s_counts[tid];
        int tiles = (count + block_m - 1) / block_m;

        int r_scan = count;
        int t_scan = tiles;
        #pragma unroll
        for (int offset = 1; offset < 32; offset <<= 1) {
            int r_n = __shfl_up_sync(0xffffffff, r_scan, offset);
            int t_n = __shfl_up_sync(0xffffffff, t_scan, offset);
            if (lane >= offset) {
                r_scan += r_n;
                t_scan += t_n;
            }
        }
        if (lane == 31) {
            s_warp_row_sums[warp_id] = r_scan;
            s_warp_tile_sums[warp_id] = t_scan;
        }
    }
    __syncthreads();

    if (tid < 8) {
        int wr = s_warp_row_sums[tid];
        int wt = s_warp_tile_sums[tid];
        #pragma unroll
        for (int offset = 1; offset < 8; offset <<= 1) {
            int r_n = __shfl_up_sync(0xff, wr, offset);
            int t_n = __shfl_up_sync(0xff, wt, offset);
            if (tid >= offset) {
                wr += r_n;
                wt += t_n;
            }
        }
        s_warp_row_sums[tid] = wr;
        s_warp_tile_sums[tid] = wt;
        if (tid == 7) {
            s_total_routed_tiles = wt;
        }
    }
    __syncthreads();

    if (tid < 256) {
        int warp_id = tid >> 5;
        int lane = tid & 31;
        int count = s_counts[tid];
        int tiles = (count + block_m - 1) / block_m;

        int r_scan = count;
        int t_scan = tiles;
        #pragma unroll
        for (int offset = 1; offset < 32; offset <<= 1) {
            int r_n = __shfl_up_sync(0xffffffff, r_scan, offset);
            int t_n = __shfl_up_sync(0xffffffff, t_scan, offset);
            if (lane >= offset) {
                r_scan += r_n;
                t_scan += t_n;
            }
        }
        int warp_r_prefix = (warp_id > 0) ? s_warp_row_sums[warp_id - 1] : 0;
        int warp_t_prefix = (warp_id > 0) ? s_warp_tile_sums[warp_id - 1] : 0;

        int row_off = warp_r_prefix + r_scan - count;
        int tile_off = warp_t_prefix + t_scan - tiles;

        s_row_offsets[tid] = row_off;
        s_tile_offsets[tid] = tile_off;
        s_write_cursor[tid] = row_off;

        int step_m = (block_m == 64) ? 16 : 32;
        int base_m = (tiles > 0) ? ((count + tiles - 1) / tiles) : block_m;
        int tile_m = (tiles <= 4) ? min(block_m, (base_m + step_m - 1) & ~(step_m - 1)) : block_m;

        for (int m_tile = 0; m_tile < tiles; ++m_tile) {
            int start_r = row_off + m_tile * tile_m;
            int valid_m = min(tile_m, count - m_tile * tile_m);
            int b_l2_mode = (tiles == 1) ? 1 : 0; // 1 = evict_first, 0 = normal
            tile_table[tile_off + m_tile] = make_int4(tid, start_r, valid_m, b_l2_mode);
        }
    }
    __syncthreads();

    // Shared expert tiles
    int routed_tiles = s_total_routed_tiles;
    int shared_tiles_per_expert = (T + block_m - 1) / block_m;
    int total_shared_tiles = n_shared * shared_tiles_per_expert;
    for (int st = tid; st < total_shared_tiles; st += 512) {
        int s = st / shared_tiles_per_expert;
        int m_tile = st % shared_tiles_per_expert;
        int start_r = total_routed + s * T + m_tile * block_m;
        int valid_m = min(block_m, T - m_tile * block_m);
        int b_l2_mode = (m_tile == shared_tiles_per_expert - 1) ? 1 : 2; // 2 = evict_last, 1 = evict_first on last tile
        tile_table[routed_tiles + st] = make_int4(E + s, start_r, valid_m, b_l2_mode);
    }

    if (tid == 0) {
        d_num_tiles[0] = routed_tiles + total_shared_tiles;
    }

    // Scatter token_ids and weights
    for (int idx = tid; idx < total_routed; idx += 512) {
        int e = static_cast<int>(expert_ids[idx]);
        int dst = atomicAdd(&s_write_cursor[e], 1);
        sorted_token_ids[dst] = idx / top_k;
        sorted_weights[dst] = __bfloat162float(expert_weights[idx]);
    }
}

// ============================================================================
// Stage 1: Fused Gate + Up + SiLU Grouped GEMM (SM120 Tensor Core m16n8k16)
// ============================================================================

template <int BLOCK_M, int WARPS_M, int WARPS_N, int M_STEPS, int STAGES = 3>
__device__ __forceinline__ void moe_w1_silu_k_loop(
    const __nv_bfloat16* __restrict__ w1_e,
    __nv_bfloat16* __restrict__ h_sorted,
    uint4* __restrict__ smem_A,
    uint4* __restrict__ smem_B,
    const __nv_bfloat16* const* __restrict__ smem_a_ptr,
    int start_row,
    int m_valid,
    int i_base,
    int b_l2_mode,
    int H,
    int I)
{
    constexpr int M_ACTIVE = M_STEPS * WARPS_M * 16;
    constexpr int N_TILES_PER_WARP = 8 / WARPS_N; // 64 cols / (WARPS_N * 8)
    constexpr int A_U4_PER_STAGE = BLOCK_M * 8;
    constexpr int B_U4_PER_STAGE = 128 * 8;

    int tid = threadIdx.x;
    int warp_id = tid >> 5;
    int lane = tid & 31;
    int wm = warp_id / WARPS_N;
    int wn = warp_id % WARPS_N;

    float g_d0[M_STEPS][N_TILES_PER_WARP];
    float g_d1[M_STEPS][N_TILES_PER_WARP];
    float g_d2[M_STEPS][N_TILES_PER_WARP];
    float g_d3[M_STEPS][N_TILES_PER_WARP];
    float u_d0[M_STEPS][N_TILES_PER_WARP];
    float u_d1[M_STEPS][N_TILES_PER_WARP];
    float u_d2[M_STEPS][N_TILES_PER_WARP];
    float u_d3[M_STEPS][N_TILES_PER_WARP];

    #pragma unroll
    for (int ms = 0; ms < M_STEPS; ++ms) {
        #pragma unroll
        for (int nt = 0; nt < N_TILES_PER_WARP; ++nt) {
            g_d0[ms][nt] = 0.0f; g_d1[ms][nt] = 0.0f; g_d2[ms][nt] = 0.0f; g_d3[ms][nt] = 0.0f;
            u_d0[ms][nt] = 0.0f; u_d1[ms][nt] = 0.0f; u_d2[ms][nt] = 0.0f; u_d3[ms][nt] = 0.0f;
        }
    }

    int num_k_tiles = H >> 6; // H / 64 = 64
    uint64_t policy_A = make_policy_evict_last();
    uint64_t policy_B = (b_l2_mode == 1) ? make_policy_evict_first()
                      : ((b_l2_mode == 2) ? make_policy_evict_last() : 0ULL);

    auto issue_loads = [&](int kt, int stage) {
        int k_base = kt * 64;
        uint4* sA = smem_A + stage * A_U4_PER_STAGE;
        uint4* sB = smem_B + stage * B_U4_PER_STAGE;

        #pragma unroll
        for (int idx = tid; idx < M_ACTIVE * 8; idx += 256) {
            int r = idx >> 3;
            int c = idx & 7;
            int c_swz = c ^ (r & 7);
            const void* src = smem_a_ptr[r] + k_base + c * 8;
            cp_async_ca_16(smem_u32_ptr(&sA[r * 8 + c_swz]), src, policy_A);
        }

        #pragma unroll
        for (int idx = tid; idx < 128 * 8; idx += 256) {
            int r = idx >> 3;
            int c = idx & 7;
            int c_swz = c ^ (r & 7);
            int w_row = (r < 64) ? (i_base + r) : (I + i_base + (r - 64));
            const void* src = w1_e + (size_t)w_row * H + k_base + c * 8;
            cp_async_cg_16(smem_u32_ptr(&sB[r * 8 + c_swz]), src, policy_B);
        }
    };

    auto load_B_frag = [&](uint4* sB, int km, uint32_t* bg0, uint32_t* bg1, uint32_t* bu0, uint32_t* bu1) {
        if constexpr (N_TILES_PER_WARP == 1) {
            int rb = ((lane >> 4) * 64) + wn * 8 + (lane & 7);
            int cb = (km * 2) + ((lane >> 3) & 1);
            ldmatrix_x4(bg0[0], bg1[0], bu0[0], bu1[0],
                        smem_u32_ptr(&sB[rb * 8 + (cb ^ (rb & 7))]));
        } else {
            int rb_g = wn * 16 + (lane & 7) + ((lane >> 4) * 8);
            int rb_u = 64 + rb_g;
            int cb = (km * 2) + ((lane >> 3) & 1);
            ldmatrix_x4(bg0[0], bg1[0], bg0[1], bg1[1],
                        smem_u32_ptr(&sB[rb_g * 8 + (cb ^ (rb_g & 7))]));
            ldmatrix_x4(bu0[0], bu1[0], bu0[1], bu1[1],
                        smem_u32_ptr(&sB[rb_u * 8 + (cb ^ (rb_u & 7))]));
        }
    };

    auto load_A_frag = [&](uint4* sA, int km, int ms, uint32_t* a_frag) {
        int ra = (ms * WARPS_M + wm) * 16 + (lane & 15);
        int ca = (km * 2) + (lane >> 4);
        ldmatrix_x4(a_frag[0], a_frag[1], a_frag[2], a_frag[3],
                    smem_u32_ptr(&sA[ra * 8 + (ca ^ (ra & 7))]));
    };

    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        issue_loads(s, s);
        cp_async_commit();
    }
    cp_async_wait_group<STAGES - 2>();
    __syncthreads();

    uint32_t bg0[2][N_TILES_PER_WARP], bg1[2][N_TILES_PER_WARP];
    uint32_t bu0[2][N_TILES_PER_WARP], bu1[2][N_TILES_PER_WARP];
    uint32_t a_frag[2][4];

    load_B_frag(smem_B, 0, bg0[0], bg1[0], bu0[0], bu1[0]);
    load_A_frag(smem_A, 0, 0, a_frag[0]);

    for (int kt = 0; kt < num_k_tiles; ++kt) {
        int stage = kt % STAGES;
        uint4* sA = smem_A + stage * A_U4_PER_STAGE;
        uint4* sB = smem_B + stage * B_U4_PER_STAGE;

        #pragma unroll
        for (int km = 0; km < 4; ++km) {
            if (km == 0) {
                int kt_load = kt + STAGES - 1;
                if (kt_load < num_k_tiles) {
                    issue_loads(kt_load, kt_load % STAGES);
                }
                cp_async_commit();
            }

            #pragma unroll
            for (int ms = 0; ms < M_STEPS; ++ms) {
                constexpr int TOTAL_M = M_STEPS;
                int cur_a = (km * TOTAL_M + ms) & 1;
                int nxt_a = cur_a ^ 1;
                int cur_b = km & 1;
                int nxt_b = cur_b ^ 1;

                if (ms + 1 < TOTAL_M) {
                    load_A_frag(sA, km, ms + 1, a_frag[nxt_a]);
                } else if (km < 3) {
                    load_B_frag(sB, km + 1, bg0[nxt_b], bg1[nxt_b], bu0[nxt_b], bu1[nxt_b]);
                    load_A_frag(sA, km + 1, 0, a_frag[nxt_a]);
                } else if (kt + 1 < num_k_tiles) {
                    cp_async_wait_group<STAGES - 2>();
                    __syncthreads();
                    int next_stage = (kt + 1) % STAGES;
                    uint4* next_sA = smem_A + next_stage * A_U4_PER_STAGE;
                    uint4* next_sB = smem_B + next_stage * B_U4_PER_STAGE;
                    load_B_frag(next_sB, 0, bg0[nxt_b], bg1[nxt_b], bu0[nxt_b], bu1[nxt_b]);
                    load_A_frag(next_sA, 0, 0, a_frag[nxt_a]);
                }

                #pragma unroll
                for (int nt = 0; nt < N_TILES_PER_WARP; ++nt) {
                    mma_m16n8k16_bf16(g_d0[ms][nt], g_d1[ms][nt], g_d2[ms][nt], g_d3[ms][nt],
                                      a_frag[cur_a][0], a_frag[cur_a][1], a_frag[cur_a][2], a_frag[cur_a][3],
                                      bg0[cur_b][nt], bg1[cur_b][nt]);
                    mma_m16n8k16_bf16(u_d0[ms][nt], u_d1[ms][nt], u_d2[ms][nt], u_d3[ms][nt],
                                      a_frag[cur_a][0], a_frag[cur_a][1], a_frag[cur_a][2], a_frag[cur_a][3],
                                      bu0[cur_b][nt], bu1[cur_b][nt]);
                }
            }
        }
    }

    int gid = lane >> 2;
    int tid_in_g = lane & 3;

    #pragma unroll
    for (int ms = 0; ms < M_STEPS; ++ms) {
        int r0 = (ms * WARPS_M + wm) * 16 + gid;
        int r1 = r0 + 8;
        #pragma unroll
        for (int nt = 0; nt < N_TILES_PER_WARP; ++nt) {
            int col = i_base + (wn * N_TILES_PER_WARP + nt) * 8 + tid_in_g * 2;

            if (r0 < m_valid) {
                float h0 = (g_d0[ms][nt] / (1.0f + __expf(-g_d0[ms][nt]))) * u_d0[ms][nt];
                float h1 = (g_d1[ms][nt] / (1.0f + __expf(-g_d1[ms][nt]))) * u_d1[ms][nt];
                *reinterpret_cast<__nv_bfloat162*>(&h_sorted[(size_t)(start_row + r0) * I + col]) =
                    __float22bfloat162_rn(make_float2(h0, h1));
            }
            if (r1 < m_valid) {
                float h2 = (g_d2[ms][nt] / (1.0f + __expf(-g_d2[ms][nt]))) * u_d2[ms][nt];
                float h3 = (g_d3[ms][nt] / (1.0f + __expf(-g_d3[ms][nt]))) * u_d3[ms][nt];
                *reinterpret_cast<__nv_bfloat162*>(&h_sorted[(size_t)(start_row + r1) * I + col]) =
                    __float22bfloat162_rn(make_float2(h2, h3));
            }
        }
    }
}

template <int BLOCK_M, int WARPS_M, int WARPS_N, int STAGES = 2, int MIN_BLOCKS = 1>
__global__ __launch_bounds__(256, MIN_BLOCKS)
void moe_w1_silu_kernel(
    const __nv_bfloat16* __restrict__ x,          // (T, H)
    const __nv_bfloat16* __restrict__ w1_routed,  // (E, 2*I, H)
    const __nv_bfloat16* __restrict__ w1_shared,  // (n_shared, 2*I, H)
    const int4* __restrict__ tile_table,          // (num_tiles)
    const int* __restrict__ sorted_token_ids,     // (T * top_k)
    const int* __restrict__ d_num_tiles,          // (1)
    __nv_bfloat16* __restrict__ h_sorted,         // ((top_k + n_shared) * T, I)
    int T,
    int E,
    int top_k,
    int H,
    int I)
{
    int super_m = blockIdx.y >> 1;
    int linear_in_super = ((blockIdx.y & 1) << 5) + blockIdx.x;
    int n_tile = linear_in_super >> 1;
    int m_tile_idx = (super_m << 1) + (linear_in_super & 1);
    if (m_tile_idx >= d_num_tiles[0]) return;

    int4 tile_info = tile_table[m_tile_idx];
    int expert_id = tile_info.x;
    int start_row = tile_info.y;
    int m_valid   = tile_info.z;
    int is_last_m = tile_info.w;
    int i_base    = n_tile * 64;

    extern __shared__ __align__(128) uint8_t smem_raw[];
    uint4* smem_A = reinterpret_cast<uint4*>(smem_raw);
    uint4* smem_B = smem_A + STAGES * BLOCK_M * 8;
    const __nv_bfloat16** smem_a_ptr = reinterpret_cast<const __nv_bfloat16**>(smem_B + STAGES * 128 * 8);

    int tid = threadIdx.x;
    if (tid < BLOCK_M) {
        int r_clamped = (tid < m_valid) ? tid : 0;
        int token_id;
        if (expert_id < E) {
            token_id = sorted_token_ids[start_row + r_clamped];
        } else {
            int s = expert_id - E;
            token_id = (start_row - (top_k + s) * T) + r_clamped;
        }
        smem_a_ptr[tid] = x + (size_t)token_id * H;
    }
    __syncthreads();

    const __nv_bfloat16* w1_e = (expert_id < E)
        ? (w1_routed + (size_t)expert_id * (2 * I) * H)
        : (w1_shared + (size_t)(expert_id - E) * (2 * I) * H);

    constexpr int M_PER_STEP = WARPS_M * 16;
    constexpr int MAX_M_STEPS = BLOCK_M / M_PER_STEP;
    int num_m_steps = (m_valid + M_PER_STEP - 1) / M_PER_STEP;

    if (num_m_steps <= 1) {
        moe_w1_silu_k_loop<BLOCK_M, WARPS_M, WARPS_N, 1, STAGES>(
            w1_e, h_sorted, smem_A, smem_B, smem_a_ptr, start_row, m_valid, i_base, is_last_m, H, I);
    } else if (num_m_steps == 2) {
        moe_w1_silu_k_loop<BLOCK_M, WARPS_M, WARPS_N, 2, STAGES>(
            w1_e, h_sorted, smem_A, smem_B, smem_a_ptr, start_row, m_valid, i_base, is_last_m, H, I);
    } else if (num_m_steps == 3) {
        moe_w1_silu_k_loop<BLOCK_M, WARPS_M, WARPS_N, 3, STAGES>(
            w1_e, h_sorted, smem_A, smem_B, smem_a_ptr, start_row, m_valid, i_base, is_last_m, H, I);
    } else if (num_m_steps == 4 || MAX_M_STEPS <= 4) {
        moe_w1_silu_k_loop<BLOCK_M, WARPS_M, WARPS_N, 4, STAGES>(
            w1_e, h_sorted, smem_A, smem_B, smem_a_ptr, start_row, m_valid, i_base, is_last_m, H, I);
    } else {
        moe_w1_silu_k_loop<BLOCK_M, WARPS_M, WARPS_N, MAX_M_STEPS, STAGES>(
            w1_e, h_sorted, smem_A, smem_B, smem_a_ptr, start_row, m_valid, i_base, is_last_m, H, I);
    }
}

// ============================================================================
// Stage 2: Fused Down + Weight Mul + Coalesced float4 atomicAdd Grouped GEMM
// ============================================================================

template <int BLOCK_M, int WARPS_M, int WARPS_N, int M_STEPS, int STAGES = 2>
__device__ __forceinline__ void moe_w2_down_k_loop(
    const __nv_bfloat16* __restrict__ w2_e,
    float* __restrict__ out_fp32,
    uint4* __restrict__ smem_A,
    uint4* __restrict__ smem_B,
    const __nv_bfloat16* const* __restrict__ smem_a_ptr,
    const int* __restrict__ smem_token_id,
    const float* __restrict__ smem_weight,
    int m_valid,
    int n_base,
    int is_last_m,
    int H,
    int I)
{
    constexpr int M_ACTIVE = M_STEPS * WARPS_M * 16;
    constexpr int N_PAIRS_PER_WARP = 8 / WARPS_N;      // 16-col pairs per warp
    constexpr int N_TILES_PER_WARP = 2 * N_PAIRS_PER_WARP; // 8-col tiles per warp
    constexpr int A_U4_PER_STAGE = BLOCK_M * 8;
    constexpr int B_U4_PER_STAGE = 128 * 8;

    int tid = threadIdx.x;
    int warp_id = tid >> 5;
    int lane = tid & 31;
    int wm = warp_id / WARPS_N;
    int wn = warp_id % WARPS_N;

    float d0[M_STEPS][N_TILES_PER_WARP];
    float d1[M_STEPS][N_TILES_PER_WARP];
    float d2[M_STEPS][N_TILES_PER_WARP];
    float d3[M_STEPS][N_TILES_PER_WARP];

    #pragma unroll
    for (int ms = 0; ms < M_STEPS; ++ms) {
        #pragma unroll
        for (int nt = 0; nt < N_TILES_PER_WARP; ++nt) {
            d0[ms][nt] = 0.0f; d1[ms][nt] = 0.0f; d2[ms][nt] = 0.0f; d3[ms][nt] = 0.0f;
        }
    }

    int num_k_tiles = I >> 6; // I / 64 = 32
    uint64_t policy_A = make_policy_evict_last();
    uint64_t policy_B = (is_last_m == 1) ? make_policy_evict_first()
                      : ((is_last_m == 2) ? make_policy_evict_last() : 0ULL);

    auto issue_loads = [&](int kt, int stage) {
        int k_base = kt * 64;
        uint4* sA = smem_A + stage * A_U4_PER_STAGE;
        uint4* sB = smem_B + stage * B_U4_PER_STAGE;

        #pragma unroll
        for (int idx = tid; idx < M_ACTIVE * 8; idx += 256) {
            int r = idx >> 3;
            int c = idx & 7;
            int c_swz = c ^ (r & 7);
            const void* src = smem_a_ptr[r] + k_base + c * 8;
            cp_async_ca_16(smem_u32_ptr(&sA[r * 8 + c_swz]), src, policy_A);
        }

        #pragma unroll
        for (int idx = tid; idx < 128 * 8; idx += 256) {
            int r = idx >> 3;
            int c = idx & 7;
            int c_swz = c ^ (r & 7);
            const void* src = w2_e + (size_t)(n_base + r) * I + k_base + c * 8;
            cp_async_cg_16(smem_u32_ptr(&sB[r * 8 + c_swz]), src, policy_B);
        }
    };

    auto load_B_frag = [&](uint4* sB, int km, uint32_t* b0_dst, uint32_t* b1_dst) {
        #pragma unroll
        for (int np = 0; np < N_PAIRS_PER_WARP; ++np) {
            int rb = (wn * N_PAIRS_PER_WARP + np) * 16 + (lane & 7) + ((lane >> 4) * 8);
            int cb = (km * 2) + ((lane >> 3) & 1);
            ldmatrix_x4(b0_dst[2 * np + 0], b1_dst[2 * np + 0], b0_dst[2 * np + 1], b1_dst[2 * np + 1],
                        smem_u32_ptr(&sB[rb * 8 + (cb ^ (rb & 7))]));
        }
    };

    auto load_A_frag = [&](uint4* sA, int km, int ms, uint32_t* a_dst) {
        int ra = (ms * WARPS_M + wm) * 16 + (lane & 15);
        int ca = (km * 2) + (lane >> 4);
        ldmatrix_x4(a_dst[0], a_dst[1], a_dst[2], a_dst[3],
                    smem_u32_ptr(&sA[ra * 8 + (ca ^ (ra & 7))]));
    };

    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        issue_loads(s, s);
        cp_async_commit();
    }
    cp_async_wait_group<STAGES - 2>();
    __syncthreads();

    uint32_t b0[2][N_TILES_PER_WARP], b1[2][N_TILES_PER_WARP];
    uint32_t a_frag[2][4];

    load_B_frag(smem_B, 0, b0[0], b1[0]);
    load_A_frag(smem_A, 0, 0, a_frag[0]);

    for (int kt = 0; kt < num_k_tiles; ++kt) {
        int stage = kt % STAGES;
        uint4* sA = smem_A + stage * A_U4_PER_STAGE;
        uint4* sB = smem_B + stage * B_U4_PER_STAGE;

        #pragma unroll
        for (int km = 0; km < 4; ++km) {
            if (km == 0) {
                int kt_load = kt + STAGES - 1;
                if (kt_load < num_k_tiles) {
                    issue_loads(kt_load, kt_load % STAGES);
                }
                cp_async_commit();
            }

            #pragma unroll
            for (int ms = 0; ms < M_STEPS; ++ms) {
                constexpr int TOTAL_M = M_STEPS;
                int cur_a = (km * TOTAL_M + ms) & 1;
                int nxt_a = cur_a ^ 1;
                int cur_b = km & 1;
                int nxt_b = cur_b ^ 1;

                if (ms + 1 < TOTAL_M) {
                    load_A_frag(sA, km, ms + 1, a_frag[nxt_a]);
                } else if (km < 3) {
                    load_B_frag(sB, km + 1, b0[nxt_b], b1[nxt_b]);
                    load_A_frag(sA, km + 1, 0, a_frag[nxt_a]);
                } else if (kt + 1 < num_k_tiles) {
                    cp_async_wait_group<STAGES - 2>();
                    __syncthreads();
                    int next_stage = (kt + 1) % STAGES;
                    uint4* next_sA = smem_A + next_stage * A_U4_PER_STAGE;
                    uint4* next_sB = smem_B + next_stage * B_U4_PER_STAGE;
                    load_B_frag(next_sB, 0, b0[nxt_b], b1[nxt_b]);
                    load_A_frag(next_sA, 0, 0, a_frag[nxt_a]);
                }

                #pragma unroll
                for (int nt = 0; nt < N_TILES_PER_WARP; ++nt) {
                    mma_m16n8k16_bf16(d0[ms][nt], d1[ms][nt], d2[ms][nt], d3[ms][nt],
                                      a_frag[cur_a][0], a_frag[cur_a][1], a_frag[cur_a][2], a_frag[cur_a][3],
                                      b0[cur_b][nt], b1[cur_b][nt]);
                }
            }
        }
    }

    cp_async_wait_group<0>();
    __syncthreads();

    // Reuse smem_A + smem_B for coalesced float4 atomicAdd staging
    // STRIDE_C = 132 floats (16-byte aligned and bank-conflict free)
    constexpr int STRIDE_C = 132;
    float* smem_C = reinterpret_cast<float*>(smem_A);

    int gid = lane >> 2;
    int tid_in_g = lane & 3;

    #pragma unroll
    for (int ms = 0; ms < M_STEPS; ++ms) {
        int r0 = (ms * WARPS_M + wm) * 16 + gid;
        int r1 = r0 + 8;
        float w0 = (r0 < m_valid) ? smem_weight[r0] : 0.0f;
        float w1 = (r1 < m_valid) ? smem_weight[r1] : 0.0f;

        #pragma unroll
        for (int nt = 0; nt < N_TILES_PER_WARP; ++nt) {
            int col = (wn * N_TILES_PER_WARP + nt) * 8 + tid_in_g * 2;
            if (r0 < m_valid) {
                *reinterpret_cast<float2*>(&smem_C[r0 * STRIDE_C + col]) =
                    make_float2(d0[ms][nt] * w0, d1[ms][nt] * w0);
            }
            if (r1 < m_valid) {
                *reinterpret_cast<float2*>(&smem_C[r1 * STRIDE_C + col]) =
                    make_float2(d2[ms][nt] * w1, d3[ms][nt] * w1);
            }
        }
    }
    __syncthreads();

    // Each warp processes 1 row at a time: 32 lanes * float4 = 128 contiguous floats!
    for (int r = warp_id; r < m_valid; r += 8) {
        int token_id = smem_token_id[r];
        float4 v4 = *reinterpret_cast<const float4*>(&smem_C[r * STRIDE_C + lane * 4]);
        atomicAdd(reinterpret_cast<float4*>(&out_fp32[(size_t)token_id * H + n_base + lane * 4]), v4);
    }
}

template <int BLOCK_M, int WARPS_M, int WARPS_N, int STAGES = 2, int MIN_BLOCKS = 1>
__global__ __launch_bounds__(256, MIN_BLOCKS)
void moe_w2_down_kernel(
    const __nv_bfloat16* __restrict__ h_sorted,   // ((top_k + n_shared) * T, I)
    const __nv_bfloat16* __restrict__ w2_routed,  // (E, H, I)
    const __nv_bfloat16* __restrict__ w2_shared,  // (n_shared, H, I)
    const int4* __restrict__ tile_table,          // (num_tiles)
    const int* __restrict__ sorted_token_ids,     // (T * top_k)
    const float* __restrict__ sorted_weights,     // (T * top_k)
    const int* __restrict__ d_num_tiles,          // (1)
    float* __restrict__ out_fp32,                 // (T, H)
    int T,
    int E,
    int top_k,
    int H,
    int I)
{
    int super_m = blockIdx.y >> 1;
    int linear_in_super = ((blockIdx.y & 1) << 5) + blockIdx.x;
    int n_tile = linear_in_super >> 1;
    int m_tile_idx = (super_m << 1) + (linear_in_super & 1);
    if (m_tile_idx >= d_num_tiles[0]) return;

    int4 tile_info = tile_table[m_tile_idx];
    int expert_id = tile_info.x;
    int start_row = tile_info.y;
    int m_valid   = tile_info.z;
    int is_last_m = tile_info.w;
    int n_base    = n_tile * 128;

    constexpr int SMEM_STAGE_BYTES = STAGES * (BLOCK_M + 128) * 128;
    constexpr int SMEM_C_BYTES = BLOCK_M * 132 * 4;
    constexpr int SMEM_MAIN_BYTES = (SMEM_STAGE_BYTES > SMEM_C_BYTES) ? SMEM_STAGE_BYTES : SMEM_C_BYTES;

    extern __shared__ __align__(128) uint8_t smem_raw[];
    uint4* smem_A = reinterpret_cast<uint4*>(smem_raw);
    uint4* smem_B = smem_A + STAGES * BLOCK_M * 8;
    const __nv_bfloat16** smem_a_ptr = reinterpret_cast<const __nv_bfloat16**>(smem_raw + SMEM_MAIN_BYTES);
    int* smem_token_id = reinterpret_cast<int*>(smem_a_ptr + BLOCK_M);
    float* smem_weight = reinterpret_cast<float*>(smem_token_id + BLOCK_M);

    int tid = threadIdx.x;
    if (tid < BLOCK_M) {
        int r_clamped = (tid < m_valid) ? tid : 0;
        smem_a_ptr[tid] = h_sorted + (size_t)(start_row + r_clamped) * I;
        if (tid < m_valid) {
            if (expert_id < E) {
                smem_token_id[tid] = sorted_token_ids[start_row + tid];
                smem_weight[tid]   = sorted_weights[start_row + tid];
            } else {
                int s = expert_id - E;
                smem_token_id[tid] = (start_row - (top_k + s) * T) + tid;
                smem_weight[tid]   = 1.0f;
            }
        }
    }
    __syncthreads();

    const __nv_bfloat16* w2_e = (expert_id < E)
        ? (w2_routed + (size_t)expert_id * H * I)
        : (w2_shared + (size_t)(expert_id - E) * H * I);

    constexpr int M_PER_STEP = WARPS_M * 16;
    constexpr int MAX_M_STEPS = BLOCK_M / M_PER_STEP;
    int num_m_steps = (m_valid + M_PER_STEP - 1) / M_PER_STEP;

    if (num_m_steps <= 1) {
        moe_w2_down_k_loop<BLOCK_M, WARPS_M, WARPS_N, 1, STAGES>(
            w2_e, out_fp32, smem_A, smem_B, smem_a_ptr, smem_token_id, smem_weight, m_valid, n_base, is_last_m, H, I);
    } else if (num_m_steps == 2) {
        moe_w2_down_k_loop<BLOCK_M, WARPS_M, WARPS_N, 2, STAGES>(
            w2_e, out_fp32, smem_A, smem_B, smem_a_ptr, smem_token_id, smem_weight, m_valid, n_base, is_last_m, H, I);
    } else if (num_m_steps == 3) {
        moe_w2_down_k_loop<BLOCK_M, WARPS_M, WARPS_N, 3, STAGES>(
            w2_e, out_fp32, smem_A, smem_B, smem_a_ptr, smem_token_id, smem_weight, m_valid, n_base, is_last_m, H, I);
    } else if (num_m_steps == 4 || MAX_M_STEPS <= 4) {
        moe_w2_down_k_loop<BLOCK_M, WARPS_M, WARPS_N, 4, STAGES>(
            w2_e, out_fp32, smem_A, smem_B, smem_a_ptr, smem_token_id, smem_weight, m_valid, n_base, is_last_m, H, I);
    } else {
        moe_w2_down_k_loop<BLOCK_M, WARPS_M, WARPS_N, MAX_M_STEPS, STAGES>(
            w2_e, out_fp32, smem_A, smem_B, smem_a_ptr, smem_token_id, smem_weight, m_valid, n_base, is_last_m, H, I);
    }
}

// ============================================================================
// Stage 3: Convert FP32 -> BF16 and Reset out_fp32 to Zero in Place
// ============================================================================

__global__ __launch_bounds__(256, 4)
void convert_and_zero_out_kernel(
    float* __restrict__ out_fp32,
    __nv_bfloat16* __restrict__ out_bf16,
    int total_f4)
{
    int idx = blockIdx.x * 256 + threadIdx.x;
    if (idx < total_f4) {
        float4* f4_ptr = reinterpret_cast<float4*>(out_fp32);
        float4 v = f4_ptr[idx];
        f4_ptr[idx] = make_float4(0.0f, 0.0f, 0.0f, 0.0f);

        __nv_bfloat162 b0 = __float22bfloat162_rn(make_float2(v.x, v.y));
        __nv_bfloat162 b1 = __float22bfloat162_rn(make_float2(v.z, v.w));
        uint2 packed;
        packed.x = *reinterpret_cast<uint32_t*>(&b0);
        packed.y = *reinterpret_cast<uint32_t*>(&b1);
        reinterpret_cast<uint2*>(out_bf16)[idx] = packed;
    }
}

// ============================================================================
// Host Launcher
// ============================================================================

void glm52_moe_forward(
    torch::Tensor x,
    torch::Tensor expert_ids,
    torch::Tensor expert_weights,
    torch::Tensor w1_routed,
    torch::Tensor w2_routed,
    torch::Tensor w1_shared,
    torch::Tensor w2_shared,
    torch::Tensor out,
    torch::Tensor h_sorted,
    torch::Tensor out_fp32,
    torch::Tensor tile_table,
    torch::Tensor sorted_token_ids,
    torch::Tensor sorted_weights,
    torch::Tensor d_num_tiles,
    torch::Tensor t1_partial,
    torch::Tensor t1_tile_done)
{
    const int T = x.size(0);
    const int H = x.size(1);
    const int top_k = expert_ids.size(1);
    const int E = w1_routed.size(0);
    const int I = w2_routed.size(2);
    const int n_shared = w1_shared.size(0);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (T == 1) {
        int total_slots = top_k + n_shared;
        dim3 grid1(I / 8, total_slots);
        gemv_w1_silu_t1_kernel<<<grid1, 128, 0, stream>>>(
            reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
            expert_ids.data_ptr<int64_t>(),
            reinterpret_cast<const __nv_bfloat16*>(w1_routed.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w1_shared.data_ptr()),
            reinterpret_cast<__nv_bfloat16*>(h_sorted.data_ptr()),
            t1_tile_done.data_ptr<int>(),
            top_k, H, I);

        dim3 grid2(H / 8, total_slots);
        gemv_w2_reduce_t1_kernel<<<grid2, 128, 0, stream>>>(
            reinterpret_cast<const __nv_bfloat16*>(h_sorted.data_ptr()),
            expert_ids.data_ptr<int64_t>(),
            reinterpret_cast<const __nv_bfloat16*>(expert_weights.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w2_routed.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w2_shared.data_ptr()),
            t1_partial.data_ptr<float>(),
            t1_tile_done.data_ptr<int>(),
            reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
            top_k, total_slots, H, I);
        return;
    }

    static bool smem_configured = false;
    constexpr int SMEM_64     = 2 * (64 + 128) * 128 + 64 * 16;   // 50,176 B (2 CTAs/SM)
    constexpr int SMEM_160_W1 = 2 * (160 + 128) * 128 + 160 * 8;  // 75,008 B
    constexpr int SMEM_160_W2 = 160 * 132 * 4 + 160 * 16;         // 87,040 B
    if (!smem_configured) {
        cudaFuncSetAttribute(moe_w1_silu_kernel<64, 1, 8, 2, 2>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_64);
        cudaFuncSetAttribute(moe_w2_down_kernel<64, 1, 8, 2, 2>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_64);
        cudaFuncSetAttribute(moe_w1_silu_kernel<160, 2, 4, 2, 1>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_160_W1);
        cudaFuncSetAttribute(moe_w2_down_kernel<160, 2, 4, 2, 1>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_160_W2);
        smem_configured = true;
    }

    const int block_m = (T <= 1024) ? 64 : 160;
    const int max_routed_tiles = (T * top_k) / block_m + E;
    const int max_shared_tiles = n_shared * ((T + block_m - 1) / block_m);
    const int max_m_tiles = ((max_routed_tiles + max_shared_tiles) + 1) & ~1;

    build_moe_metadata_kernel<<<1, 512, 0, stream>>>(
        expert_ids.data_ptr<int64_t>(),
        reinterpret_cast<const __nv_bfloat16*>(expert_weights.data_ptr()),
        reinterpret_cast<int4*>(tile_table.data_ptr()),
        sorted_token_ids.data_ptr<int>(),
        sorted_weights.data_ptr<float>(),
        d_num_tiles.data_ptr<int>(),
        T, E, top_k, n_shared, block_m);

    dim3 grid_w1(I / 64, max_m_tiles);
    dim3 grid_w2(H / 128, max_m_tiles);

    if (block_m == 64) {
        moe_w1_silu_kernel<64, 1, 8, 2, 2><<<grid_w1, 256, SMEM_64, stream>>>(
            reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w1_routed.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w1_shared.data_ptr()),
            reinterpret_cast<const int4*>(tile_table.data_ptr()),
            sorted_token_ids.data_ptr<int>(),
            d_num_tiles.data_ptr<int>(),
            reinterpret_cast<__nv_bfloat16*>(h_sorted.data_ptr()),
            T, E, top_k, H, I);

        moe_w2_down_kernel<64, 1, 8, 2, 2><<<grid_w2, 256, SMEM_64, stream>>>(
            reinterpret_cast<const __nv_bfloat16*>(h_sorted.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w2_routed.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w2_shared.data_ptr()),
            reinterpret_cast<const int4*>(tile_table.data_ptr()),
            sorted_token_ids.data_ptr<int>(),
            sorted_weights.data_ptr<float>(),
            d_num_tiles.data_ptr<int>(),
            out_fp32.data_ptr<float>(),
            T, E, top_k, H, I);
    } else {
        moe_w1_silu_kernel<160, 2, 4, 2, 1><<<grid_w1, 256, SMEM_160_W1, stream>>>(
            reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w1_routed.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w1_shared.data_ptr()),
            reinterpret_cast<const int4*>(tile_table.data_ptr()),
            sorted_token_ids.data_ptr<int>(),
            d_num_tiles.data_ptr<int>(),
            reinterpret_cast<__nv_bfloat16*>(h_sorted.data_ptr()),
            T, E, top_k, H, I);

        moe_w2_down_kernel<160, 2, 4, 2, 1><<<grid_w2, 256, SMEM_160_W2, stream>>>(
            reinterpret_cast<const __nv_bfloat16*>(h_sorted.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w2_routed.data_ptr()),
            reinterpret_cast<const __nv_bfloat16*>(w2_shared.data_ptr()),
            reinterpret_cast<const int4*>(tile_table.data_ptr()),
            sorted_token_ids.data_ptr<int>(),
            sorted_weights.data_ptr<float>(),
            d_num_tiles.data_ptr<int>(),
            out_fp32.data_ptr<float>(),
            T, E, top_k, H, I);
    }

    int total_f4 = (T * H) >> 2;
    int blocks_conv = (total_f4 + 255) >> 8;
    convert_and_zero_out_kernel<<<blocks_conv, 256, 0, stream>>>(
        out_fp32.data_ptr<float>(),
        reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
        total_f4);
}
"""

_CPP_SRC = r"""
void glm52_moe_forward(
    torch::Tensor x,
    torch::Tensor expert_ids,
    torch::Tensor expert_weights,
    torch::Tensor w1_routed,
    torch::Tensor w2_routed,
    torch::Tensor w1_shared,
    torch::Tensor w2_shared,
    torch::Tensor out,
    torch::Tensor h_sorted,
    torch::Tensor out_fp32,
    torch::Tensor tile_table,
    torch::Tensor sorted_token_ids,
    torch::Tensor sorted_weights,
    torch::Tensor d_num_tiles,
    torch::Tensor t1_partial,
    torch::Tensor t1_tile_done);
"""

_ext = load_inline(
    name="glm52_fused_moe_sm120_ext",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=["glm52_moe_forward"],
    extra_cuda_cflags=["-O3", "-arch=sm_120", "--use_fast_math"],
)

_WORKSPACES: dict[tuple[torch.device, int], dict[str, torch.Tensor]] = {}


def _get_workspace(device: torch.device, max_t: int, E: int, top_k: int, n_shared: int, H: int, I: int) -> dict[str, torch.Tensor]:
    key = (device, max_t)
    ws = _WORKSPACES.get(key)
    if ws is None:
        max_total_tokens = max_t * (top_k + n_shared)
        max_tiles = max_total_tokens // 16 + E + n_shared * (max_t // 16 + 2)
        ws = {
            "h_sorted": torch.empty((max_total_tokens, I), dtype=torch.bfloat16, device=device),
            "out_fp32": torch.zeros((max_t, H), dtype=torch.float32, device=device),
            "tile_table": torch.empty((max_tiles, 4), dtype=torch.int32, device=device),
            "sorted_token_ids": torch.empty((max_t * top_k,), dtype=torch.int32, device=device),
            "sorted_weights": torch.empty((max_t * top_k,), dtype=torch.float32, device=device),
            "d_num_tiles": torch.empty((1,), dtype=torch.int32, device=device),
            "t1_partial": torch.empty((top_k + n_shared, H), dtype=torch.float32, device=device),
            "t1_tile_done": torch.zeros((H // 8,), dtype=torch.int32, device=device),
        }
        _WORKSPACES[key] = ws
    return ws


class Model(nn.Module):
    def __init__(
        self,
        T: int,
        E: int,
        top_k: int,
        n_shared: int,
        H: int,
        I: int,
    ):
        super().__init__()
        self.T, self.E, self.top_k = T, E, top_k
        self.n_shared, self.H, self.I = n_shared, H, I
        self.w1_routed = nn.Parameter(torch.empty(E, 2 * I, H, dtype=torch.bfloat16))
        self.w2_routed = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
        self.w1_shared = nn.Parameter(torch.empty(n_shared, 2 * I, H, dtype=torch.bfloat16))
        self.w2_shared = nn.Parameter(torch.empty(n_shared, H, I, dtype=torch.bfloat16))

    def forward(
        self,
        x: torch.Tensor,
        expert_ids: torch.Tensor,
        expert_weights: torch.Tensor,
    ) -> torch.Tensor:
        T, H = x.shape
        max_t = max(8192, T)
        ws = _get_workspace(x.device, max_t, self.E, self.top_k, self.n_shared, self.H, self.I)
        out = torch.empty((T, H), dtype=torch.bfloat16, device=x.device)
        _ext.glm52_moe_forward(
            x,
            expert_ids,
            expert_weights,
            self.w1_routed,
            self.w2_routed,
            self.w1_shared,
            self.w2_shared,
            out,
            ws["h_sorted"],
            ws["out_fp32"],
            ws["tile_table"],
            ws["sorted_token_ids"],
            ws["sorted_weights"],
            ws["d_num_tiles"],
            ws["t1_partial"],
            ws["t1_tile_done"],
        )
        return out