"""GLM-5.2-class fused MoE layer (CUDA, SM120 / Blackwell, bf16).

Fused vLLM-style weight layout with GLM routing (E=256 routed, top_k=8,
1 always-on shared expert). Reference math per expert:
    h = silu(x @ w1_gate.T) * (x @ w1_up.T);  y = h @ w2.T
out = sum_shared(y_s) + sum_k weight[t,k] * y_routed[e_k]

Two grouped-GEMM CUDA kernels (fused_moe.cu, WMMA m16n16k16 tensor cores):
  A) gate/up -> h buffer
  B) down -> weighted fp32 atomic accumulate into out

Routing is given; tokens are grouped (sorted) by expert on the host.
"""
from __future__ import annotations

import os
from pathlib import Path

_HERE = Path(__file__).resolve().parent


def _find_cuda_home():
    """Locate a real CUDA toolkit (nvcc), bypassing harness PATH shims."""
    for c in (os.environ.get("CUDA_HOME"), os.environ.get("TORCH_CUDA_HOME")):
        if c and os.path.exists(os.path.join(c, "bin", "nvcc")):
            return c
    for c in ("/usr/local/cuda-12.8", "/usr/local/cuda", "/usr/local/cuda-13",
              "/usr/local/cuda-13.0"):
        if os.path.exists(os.path.join(c, "bin", "nvcc")):
            return c
    return None


# Must run BEFORE importing torch.utils.cpp_extension (which caches CUDA_HOME
# at module-import time from the environment).
_ch = _find_cuda_home()
if _ch is not None:
    os.environ["CUDA_HOME"] = _ch

import torch  # noqa: E402
import torch.nn as nn  # noqa: E402
from torch.utils import cpp_extension  # noqa: E402

_EXT = None

# Kernel tile sizes (must match fused_moe.cu).
BM_A, BI_A, BK_A = 128, 32, 32
BM_B, BN_B, BK_B = 128, 64, 32
THREADS = 256


def _load_ext():
    global _EXT
    if _EXT is not None:
        return _EXT
    cuda_src = (_HERE / "fused_moe.cu").read_text()
    cpp_src = r"""
#include <torch/extension.h>
torch::Tensor run_moe(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w2_r,
                      torch::Tensor w1_s, torch::Tensor w2_s,
                      torch::Tensor toks_all, torch::Tensor wts_all,
                      torch::Tensor g_start,
                      torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
                      torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
                      int64_t E, int64_t H, int64_t I, int64_t T);
torch::Tensor gate_up_only(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w1_s,
                           torch::Tensor toks_all, torch::Tensor g_start,
                           torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
                           int64_t E, int64_t H, int64_t I);
torch::Tensor down_only(torch::Tensor h_all, torch::Tensor w2_r, torch::Tensor w2_s,
                        torch::Tensor wts, torch::Tensor toks_all, torch::Tensor g_start,
                        torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
                        int64_t E, int64_t H, int64_t I, int64_t T);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run_moe", &run_moe, "GLM-5.2 fused MoE (gate/up + down grouped GEMMs)");
    m.def("gate_up_only", &gate_up_only, "debug: gate/up -> h");
    m.def("down_only", &down_only, "debug: down -> out");
}
"""
    _EXT = cpp_extension.load_inline(
        name="fused_moe_ext",
        cpp_sources=cpp_src,
        cuda_sources=cuda_src,
        extra_cuda_cflags=["-arch=sm_120a", "-O3", "--expt-relaxed-constexpr"],
        verbose=False,
    )
    return _EXT


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))
        for p in self.parameters():
            nn.init.normal_(p, std=0.02)

    def forward(
        self,
        x: torch.Tensor,
        expert_ids: torch.Tensor,
        expert_weights: torch.Tensor,
    ) -> torch.Tensor:
        T, H = x.shape
        E, top_k, I = self.E, self.top_k, self.I
        dev = x.device
        ext = _load_ext()

        M_r = T * top_k

        # ---- group (sort) routed assignments by expert ---------------------
        ids32 = expert_ids.reshape(-1).to(torch.int32)
        toks_flat = (
            torch.arange(T, device=dev)
            .view(T, 1)
            .expand(T, top_k)
            .reshape(-1)
            .to(torch.int32)
        )

        order = torch.argsort(ids32, stable=True)
        toks_sorted = toks_flat[order]
        wts_sorted = expert_weights.reshape(-1)[order]

        cnts = torch.bincount(ids32[order], minlength=E)
        g_start = torch.cat(
            [
                torch.zeros(1, dtype=cnts.dtype, device=dev),
                torch.cumsum(cnts, 0),
            ]
        ).to(torch.int32)  # (E+1): routed groups only

        toks_all = toks_sorted
        wts_all = wts_sorted

        # gathered activations: contiguous rows so kernel A streams xg sequentially
        xg = x[toks_all.to(torch.int64)].contiguous()

        # ---- build tile maps for kernel A and B (vectorized) ---------------
        gcnts = (g_start[1:] - g_start[:-1]).to(torch.int64)

        def tiles(BM):
            nt = ((gcnts + BM - 1) // BM)
            total = int(nt.sum())
            tg = torch.repeat_interleave(
                torch.arange(E, device=dev, dtype=torch.int64), nt
            )
            gstarts = torch.cumsum(
                torch.cat([torch.zeros(1, dtype=torch.int64, device=dev), nt]), 0
            )
            pos = torch.arange(total, device=dev, dtype=torch.int64) - gstarts[tg]
            tr0 = g_start.to(torch.int64)[tg] + pos * BM
            return tg.to(torch.int32), tr0.to(torch.int32)

        tgA, tr0A = tiles(BM_A)
        tgB, tr0B = tiles(BM_B)

        out_f32 = ext.run_moe(xg,
            self.w1_routed.detach(),
            self.w2_routed.detach(),
            self.w1_shared.detach(),
            self.w2_shared.detach(),
            toks_all,
            wts_all,
            g_start,
            tgA, tr0A, tgB, tr0B,
            E, H, I, T,
        )

        # Shared expert: dense bf16 GEMMs (read once, vs the grouped kernel's
        # per-tile re-read), accumulated in fp32.
        w1s = self.w1_shared[0]
        w2s = self.w2_shared[0]
        gate = x @ w1s[:I].T
        up = x @ w1s[I:].T
        h_s = torch.nn.functional.silu(gate) * up
        y_s = h_s @ w2s.T
        return (out_f32 + y_s.float()).to(torch.bfloat16)


def get_init_inputs():
    return [4096, 256, 8, 1, 4096, 2048]


# ==================================================================
# ===== sidecar: fused_moe.cu (17626 bytes, loaded by solution.py) =====
# ==================================================================

// GLM-5.2 fused MoE — grouped-GEMM CUDA kernels (SM120 / Blackwell, bf16).
//
// Layout (vLLM-style fused weights):
//   w1_routed: (E, 2I, H) bf16 — [0:I)=gate, [I:2I)=up
//   w2_routed: (E, H, I) bf16
//   w1_shared / w2_shared: (n_shared, ...) — the always-on shared expert.
//
// Two kernels, unified "row" space:
//   rows [0, T*top_k)        -> routed assignments, sorted by expert.
//   rows [T*top_k, +T)       -> shared-expert tokens (all T), weight 1.0.
//   g_start[g]..g_start[g+1] -> rows belonging to group g (g in [0,E) routed, g=E shared).
//
//   Kernel A (gate_up): h[row, :] = silu(x@w1_gate.T) * (x@w1_up.T)   -> h_all (M_total, I)
//   Kernel B (down):    out[token] += wts[row] * (h[row] @ w2.T)      -> atomic fp32 accumulate
//
// WMMA (nvcuda::wmma) m16n16k16 bf16 tensor-core path with cp.async double
// buffering over the contraction dimension.

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <mma.h>

using namespace nvcuda;
using bf16 = __nv_bfloat16;

__device__ __forceinline__ float silu_f(float x) {
    return x / (1.0f + expf(-x));
}

// ---- cp.async helpers (16-byte copies, SM80+) ----------------------------
__device__ __forceinline__ void cp_async16(void* smem, const void* gmem) {
    unsigned s = (unsigned)__cvta_generic_to_shared(smem);
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(s), "l"(gmem));
}
__device__ __forceinline__ void cp_commit() {
    asm volatile("cp.async.commit_group;\n");
}
template <int N>
__device__ __forceinline__ void cp_wait() {
    asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
}

// wmma m16n16k16 f32 accumulator fragment -> (row, col) within 16x16 tile.
__device__ __forceinline__ void frag_row_col(int lane, int i, int& r, int& c) {
    r = (i & 2) ? (8 + (lane >> 2)) : (lane >> 2);
    c = ((lane & 3) << 1) | (i & 1) | ((i & 4) << 1);
}

// ---------------------------------------------------------------------------
// Kernel A: grouped gate/up GEMM -> h_all (M_total, I)
// ---------------------------------------------------------------------------
#define BM_A 128
#define BI_A 32
#define BK_A 32
#define THREADS_A 256
#define GS_A 4   // I-slices processed per block (xg rows reused from L2 across them)
// warps per block = THREADS_A / 32 = 16; rows per warp = BM_A / 16 = 16

__global__ void __launch_bounds__(THREADS_A)
gate_up_kernel(
    const bf16* __restrict__ x,          // (M_total, H) gathered xg
    const bf16* __restrict__ w1_r,       // (E, 2I, H)
    const bf16* __restrict__ w1_s,       // (n_shared, 2I, H)
    const int*  __restrict__ toks,       // (M_total)
    const int*  __restrict__ g_start,    // (G+1)
    const int*  __restrict__ tile_group, // (num_tiles)
    const int*  __restrict__ tile_row0,  // (num_tiles)
    bf16* __restrict__ h_all,            // (M_total, I)
    const int H, const int I, const int E, const int num_sgroups)
{
    // grid = (num_sgroups, num_tiles): a block handles GS_A consecutive I-slices
    // of one tile, re-reading its xg rows from L2 across the group.
    const int sgroup = blockIdx.x;
    const int tile   = blockIdx.y;
    const int g      = tile_group[tile];
    const int row0   = tile_row0[tile];
    const int g0     = g_start[g];
    const int g1     = g_start[g + 1];
    const int cnt    = g1 - g0;
    const int nrow   = min(BM_A, g1 - row0);

    const bf16* w1 = (g < E) ? (w1_r + (size_t)g * (size_t)2 * I * H) : w1_s;

    __shared__ __align__(16) bf16 xs[2][BM_A][BK_A];
    __shared__ __align__(16) bf16 wg[2][BI_A][BK_A];
    __shared__ __align__(16) bf16 wu[2][BI_A][BK_A];

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;

    wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc_g[BI_A / 16], acc_u[BI_A / 16];
    const int XBLOCKS = (BM_A * BK_A) / 8;
    const int WBLOCKS = (BI_A * BK_A) / 8;

    for (int gs_i = 0; gs_i < GS_A; gs_i++) {
        const int slice  = sgroup * GS_A + gs_i;
        const int ibase  = slice * BI_A;

#pragma unroll
        for (int i = 0; i < BI_A / 16; i++) {
            wmma::fill_fragment(acc_g[i], 0.f);
            wmma::fill_fragment(acc_u[i], 0.f);
        }

        // Prefetch chunk 0 into buffer 0.
        for (int b = tid; b < XBLOCKS; b += THREADS_A) {
            int r = b / (BK_A / 8);
            int k8 = b % (BK_A / 8);
            const bf16* src = (r < nrow) ? (x + (size_t)(row0 + r) * H + k8 * 8)
                                         : (x + k8 * 8);
            cp_async16(&xs[0][r][k8 * 8], src);
        }
        for (int b = tid; b < WBLOCKS; b += THREADS_A) {
            int n  = b / (BK_A / 8);
            int k8 = b % (BK_A / 8);
            cp_async16(&wg[0][n][k8 * 8], w1 + (size_t)(ibase + n) * H + k8 * 8);
            cp_async16(&wu[0][n][k8 * 8], w1 + (size_t)(I + ibase + n) * H + k8 * 8);
        }
        cp_commit();

        for (int k0 = 0; k0 < H; k0 += BK_A) {
            const int cur = (k0 / BK_A) & 1;
            const int nxt = cur ^ 1;
            if (k0 + BK_A < H) {
                for (int b = tid; b < XBLOCKS; b += THREADS_A) {
                    int r = b / (BK_A / 8);
                    int k8 = b % (BK_A / 8);
                    const bf16* src = (r < nrow) ? (x + (size_t)(row0 + r) * H + k0 + BK_A + k8 * 8)
                                                 : (x + k0 + BK_A + k8 * 8);
                    cp_async16(&xs[nxt][r][k8 * 8], src);
                }
                for (int b = tid; b < WBLOCKS; b += THREADS_A) {
                    int n  = b / (BK_A / 8);
                    int k8 = b % (BK_A / 8);
                    cp_async16(&wg[nxt][n][k8 * 8], w1 + (size_t)(ibase + n) * H + k0 + BK_A + k8 * 8);
                    cp_async16(&wu[nxt][n][k8 * 8], w1 + (size_t)(I + ibase + n) * H + k0 + BK_A + k8 * 8);
                }
                cp_commit();
            }
            if (k0 + BK_A < H) cp_wait<1>();
            else cp_wait<0>();
            __syncthreads();

            const int mbase = warp * 16;
#pragma unroll
            for (int ks = 0; ks < BK_A / 16; ks++) {
                wmma::fragment<wmma::matrix_a, 16, 16, 16, bf16, wmma::row_major> af;
                wmma::load_matrix_sync(af, &xs[cur][mbase][ks * 16], BK_A);
#pragma unroll
                for (int nt = 0; nt < BI_A / 16; nt++) {
                    wmma::fragment<wmma::matrix_b, 16, 16, 16, bf16, wmma::col_major> bg, bu;
                    wmma::load_matrix_sync(bg, &wg[cur][nt * 16][ks * 16], BK_A);
                    wmma::load_matrix_sync(bu, &wu[cur][nt * 16][ks * 16], BK_A);
                    wmma::mma_sync(acc_g[nt], af, bg, acc_g[nt]);
                    wmma::mma_sync(acc_u[nt], af, bu, acc_u[nt]);
                }
            }
            __syncthreads();
        }

        // Write h = silu(gate) * up directly from accumulator fragments.
#pragma unroll
        for (int nt = 0; nt < BI_A / 16; nt++) {
#pragma unroll
            for (int i = 0; i < 8; i++) {
                int rr, cc;
                frag_row_col(lane, i, rr, cc);
                const int grow = warp * 16 + rr;
                if (grow < nrow) {
                    float hv = silu_f(acc_g[nt].x[i]) * acc_u[nt].x[i];
                    h_all[(size_t)(row0 + grow) * I + ibase + nt * 16 + cc] = __float2bfloat16(hv);
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Kernel B: grouped down GEMM -> weighted atomic accumulate into out (T, H) fp32
// ---------------------------------------------------------------------------
#define BM_B 128
#define BN_B 64
#define BK_B 32
#define THREADS_B 256
#define HS_B 4   // H-slices processed per block (h rows reused from L2 across them)

__global__ void __launch_bounds__(THREADS_B)
down_kernel(
    const bf16* __restrict__ h_all,      // (M_total, I)
    const bf16* __restrict__ w2_r,       // (E, H, I)
    const bf16* __restrict__ w2_s,       // (n_shared, H, I)
    const bf16* __restrict__ wts,        // (M_total)
    const int*  __restrict__ toks,       // (M_total)
    const int*  __restrict__ g_start,
    const int*  __restrict__ tile_group,
    const int*  __restrict__ tile_row0,
    float* __restrict__ out,             // (T, H) fp32
    const int H, const int I, const int T, const int E, const int num_hgroups)
{
    // grid = (num_hgroups, num_tiles): a block processes HS_B consecutive H-slices
    // of one tile, re-reading its h rows from L2 across the group.
    const int hslice0 = blockIdx.x * HS_B;
    const int tile   = blockIdx.y;
    const int g      = tile_group[tile];
    const int row0   = tile_row0[tile];
    const int g0     = g_start[g];
    const int g1     = g_start[g + 1];
    const int cnt    = g1 - g0;
    const int nrow   = min(BM_B, g1 - row0);

    const bf16* w2 = (g < E) ? (w2_r + (size_t)g * H * I) : w2_s;

    __shared__ __align__(16) bf16 hs[2][BM_B][BK_B];
    __shared__ __align__(16) bf16 w2t[2][BN_B][BK_B];

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;

    // 2D warp tiling: warps_m x warps_n over the (BM_B, BN_B) output tile.
    const int warps_m = BM_B / 16;
    const int warps_n = (THREADS_B / 32) / warps_m;
    const int wm = warp % warps_m;
    const int wn = warp / warps_m;
    const int mbase = wm * 16;
    const int n_ntiles_per_warp = BN_B / (16 * warps_n);
    const int n_base = wn * (BN_B / warps_n);

    wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc[n_ntiles_per_warp];
    const int HBLOCKS = (BM_B * BK_B) / 8;
    const int WBLOCKS = (BN_B * BK_B) / 8;

    for (int hs_i = 0; hs_i < HS_B; hs_i++) {
        const int hb = (hslice0 + hs_i) * BN_B;
#pragma unroll
        for (int i = 0; i < n_ntiles_per_warp; i++) wmma::fill_fragment(acc[i], 0.f);

        // prefetch chunk 0
        for (int b = tid; b < HBLOCKS; b += THREADS_B) {
            int r = b / (BK_B / 8);
            int k8 = b % (BK_B / 8);
            const bf16* src = (r < nrow) ? (h_all + (size_t)(row0 + r) * I + k8 * 8)
                                         : (h_all + k8 * 8);
            cp_async16(&hs[0][r][k8 * 8], src);
        }
        for (int b = tid; b < WBLOCKS; b += THREADS_B) {
            int n  = b / (BK_B / 8);
            int k8 = b % (BK_B / 8);
            cp_async16(&w2t[0][n][k8 * 8], w2 + (size_t)(hb + n) * I + k8 * 8);
        }
        cp_commit();

        for (int k0 = 0; k0 < I; k0 += BK_B) {
            const int cur = (k0 / BK_B) & 1;
            const int nxt = cur ^ 1;
            if (k0 + BK_B < I) {
                for (int b = tid; b < HBLOCKS; b += THREADS_B) {
                    int r = b / (BK_B / 8);
                    int k8 = b % (BK_B / 8);
                    const bf16* src = (r < nrow) ? (h_all + (size_t)(row0 + r) * I + k0 + BK_B + k8 * 8)
                                                 : (h_all + k0 + BK_B + k8 * 8);
                    cp_async16(&hs[nxt][r][k8 * 8], src);
                }
                for (int b = tid; b < WBLOCKS; b += THREADS_B) {
                    int n  = b / (BK_B / 8);
                    int k8 = b % (BK_B / 8);
                    cp_async16(&w2t[nxt][n][k8 * 8], w2 + (size_t)(hb + n) * I + k0 + BK_B + k8 * 8);
                }
                cp_commit();
            }
            if (k0 + BK_B < I) cp_wait<1>();
            else cp_wait<0>();
            __syncthreads();

#pragma unroll
            for (int ks = 0; ks < BK_B / 16; ks++) {
                wmma::fragment<wmma::matrix_a, 16, 16, 16, bf16, wmma::row_major> af;
                wmma::load_matrix_sync(af, &hs[cur][mbase][ks * 16], BK_B);
#pragma unroll
                for (int nt = 0; nt < n_ntiles_per_warp; nt++) {
                    wmma::fragment<wmma::matrix_b, 16, 16, 16, bf16, wmma::col_major> bf;
                    wmma::load_matrix_sync(bf, &w2t[cur][n_base + nt * 16][ks * 16], BK_B);
                    wmma::mma_sync(acc[nt], af, bf, acc[nt]);
                }
            }
            __syncthreads();
        }

        // weighted atomic accumulate into out.
#pragma unroll
        for (int nt = 0; nt < n_ntiles_per_warp; nt++) {
#pragma unroll
            for (int i = 0; i < 8; i++) {
                int rr, cc;
                frag_row_col(lane, i, rr, cc);
                const int grow = mbase + rr;
                if (grow < nrow) {
                    int   tok = toks[row0 + grow];
                    float wt  = __bfloat162float(wts[row0 + grow]);
                    atomicAdd(&out[(size_t)tok * H + hb + n_base + nt * 16 + cc], wt * acc[nt].x[i]);
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Host entry points
// ---------------------------------------------------------------------------
torch::Tensor gate_up_only(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w1_s,
                           torch::Tensor toks_all, torch::Tensor g_start,
                           torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
                           int64_t E, int64_t H, int64_t I) {
    int64_t M_total = toks_all.size(0);
    auto h_all = torch::empty({M_total, I}, x.options().dtype(torch::kBFloat16));
    auto stream = at::cuda::getCurrentCUDAStream();
    const bf16* xp  = reinterpret_cast<const bf16*>(x.data_ptr<at::BFloat16>());
    const bf16* w1r = reinterpret_cast<const bf16*>(w1_r.data_ptr<at::BFloat16>());
    const bf16* w1s = reinterpret_cast<const bf16*>(w1_s.data_ptr<at::BFloat16>());
    const int* toksp = toks_all.data_ptr<int>();
    const int* gsp   = g_start.data_ptr<int>();
    const int* tgA   = tile_group_A.data_ptr<int>();
    const int* tr0A  = tile_row0_A.data_ptr<int>();
    bf16* hp = reinterpret_cast<bf16*>(h_all.data_ptr<at::BFloat16>());
    int num_tiles_A = (int)tile_group_A.numel();
    dim3 gridA(I / (BI_A * GS_A), num_tiles_A);
    gate_up_kernel<<<gridA, THREADS_A, 0, stream>>>(
        xp, w1r, w1s, toksp, gsp, tgA, tr0A, hp, (int)H, (int)I, (int)E, gridA.y);
    return h_all;
}

torch::Tensor down_only(torch::Tensor h_all, torch::Tensor w2_r, torch::Tensor w2_s,
                        torch::Tensor wts, torch::Tensor toks_all, torch::Tensor g_start,
                        torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
                        int64_t E, int64_t H, int64_t I, int64_t T) {
    int64_t M_total = toks_all.size(0);
    auto out = torch::zeros({T, H}, h_all.options().dtype(torch::kFloat32));
    auto stream = at::cuda::getCurrentCUDAStream();
    const bf16* hp  = reinterpret_cast<const bf16*>(h_all.data_ptr<at::BFloat16>());
    const bf16* w2r = reinterpret_cast<const bf16*>(w2_r.data_ptr<at::BFloat16>());
    const bf16* w2s = reinterpret_cast<const bf16*>(w2_s.data_ptr<at::BFloat16>());
    const bf16* wtsp = reinterpret_cast<const bf16*>(wts.data_ptr<at::BFloat16>());
    const int* toksp = toks_all.data_ptr<int>();
    const int* gsp   = g_start.data_ptr<int>();
    const int* tgB   = tile_group_B.data_ptr<int>();
    const int* tr0B  = tile_row0_B.data_ptr<int>();
    float* op = out.data_ptr<float>();
    int num_tiles_B = (int)tile_group_B.numel();
    dim3 gridB(H / (BN_B * HS_B), num_tiles_B);
    down_kernel<<<gridB, THREADS_B, 0, stream>>>(
        hp, w2r, w2s, wtsp, toksp, gsp, tgB, tr0B, op, (int)H, (int)I, (int)T, (int)E, gridB.y);
    return out;
}

torch::Tensor run_moe(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w2_r,
                      torch::Tensor w1_s, torch::Tensor w2_s,
                      torch::Tensor toks_all, torch::Tensor wts_all,
                      torch::Tensor g_start,
                      torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
                      torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
                      int64_t E, int64_t H, int64_t I, int64_t T) {
    int64_t M_total = toks_all.size(0);
    auto h_all = torch::empty({M_total, I}, x.options().dtype(torch::kBFloat16));
    auto out = torch::zeros({T, H}, x.options().dtype(torch::kFloat32));

    auto stream = at::cuda::getCurrentCUDAStream();
    const bf16* xp  = reinterpret_cast<const bf16*>(x.data_ptr<at::BFloat16>());
    const bf16* w1r = reinterpret_cast<const bf16*>(w1_r.data_ptr<at::BFloat16>());
    const bf16* w2r = reinterpret_cast<const bf16*>(w2_r.data_ptr<at::BFloat16>());
    const bf16* w1s = reinterpret_cast<const bf16*>(w1_s.data_ptr<at::BFloat16>());
    const bf16* w2s = reinterpret_cast<const bf16*>(w2_s.data_ptr<at::BFloat16>());
    const int* toksp = toks_all.data_ptr<int>();
    const bf16* wtsp = reinterpret_cast<const bf16*>(wts_all.data_ptr<at::BFloat16>());
    const int* gsp   = g_start.data_ptr<int>();
    const int* tgA   = tile_group_A.data_ptr<int>();
    const int* tr0A  = tile_row0_A.data_ptr<int>();
    const int* tgB   = tile_group_B.data_ptr<int>();
    const int* tr0B  = tile_row0_B.data_ptr<int>();

    bf16* hp = reinterpret_cast<bf16*>(h_all.data_ptr<at::BFloat16>());
    float* op = out.data_ptr<float>();

    int num_tiles_A = (int)tile_group_A.numel();
    int num_tiles_B = (int)tile_group_B.numel();
    dim3 gridA(I / (BI_A * GS_A), num_tiles_A);
    dim3 gridB(H / (BN_B * HS_B), num_tiles_B);

    gate_up_kernel<<<gridA, THREADS_A, 0, stream>>>(
        xp, w1r, w1s, toksp, gsp, tgA, tr0A, hp, (int)H, (int)I, (int)E, gridA.y);
    down_kernel<<<gridB, THREADS_B, 0, stream>>>(
        hp, w2r, w2s, wtsp, toksp, gsp, tgB, tr0B, op, (int)H, (int)I, (int)T, (int)E, gridB.y);

    return out;
}