"""GLM-5.2 fused MoE (256 routed top-8 + 1 shared), CUDA grouped GEMMs for SM120.

Pipeline (all CUDA, one extension):
  1. count/scatter: counting-sort the (token, k) pairs by expert, append the shared expert
     as expert E with all T tokens, build tile schedules for both GEMMs.
  2. gemm1: gathered x rows @ [gate|up]^T with fused silu(gate)*up -> h (bf16, sorted rows).
  3. gemm2: h @ down^T, scaled by routing weight, red.add into an fp32 accumulator.
  4. convert fp32 -> bf16.
"""
from __future__ import annotations

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

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

typedef __nv_bfloat16 bf16;

__device__ __forceinline__ uint32_t smem_u32(const void* p) {
  return (uint32_t)__cvta_generic_to_shared(p);
}
__device__ __forceinline__ void cp_async16(uint32_t dst, const void* src) {
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(dst), "l"(src));
}
__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); }
template <int N>
__device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); }

__device__ __forceinline__ void ldsm_x4(uint32_t* r, uint32_t addr) {
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
               : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(addr));
}
__device__ __forceinline__ void mma_bf16(float* c, const uint32_t* a, 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"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
}
__device__ __forceinline__ void red_add_v2(float* addr, float a, float b) {
  asm volatile("red.global.add.v2.f32 [%0], {%1,%2};\n" :: "l"(addr), "f"(a), "f"(b) : "memory");
}

// ---------------------------------------------------------------------------------
// prep: count + rank pairs per expert (multi-CTA), last CTA scans and builds schedules,
// then a scatter kernel writes the sorted token / weight lists.
// sch layout (S segments): [0,S) expert, [S,2S) m_begin, [2S,3S) m_count, [3S,4S) n_begin,
// [4S,5S) n_count, [5S, 6S+1) cumulative tile count.
// Routed experts are split into G groups; after each group comes one GMxGN block of
// shared-expert tiles, so compute-heavy shared tiles mix with weight-streaming routed ones.
// ---------------------------------------------------------------------------------
#define PREP_MAX_E 1024
#define PREP_MAX_S 2048
#define CNT_THREADS 512
#define CNT_PAIRS_PER_CTA 2048

struct SchedCfg { int BM, NT, GM, GN; };

inline int sched_groups(int T, int ns, SchedCfg c) {
  int mts = (T + c.BM - 1) / c.BM;
  return ((mts + c.GM - 1) / c.GM) * ((c.NT + c.GN - 1) / c.GN) * ns;
}

__device__ void build_schedule(int* sch, const int* s_cnt, int* s_tiles, int E, int ns, int T, SchedCfg c) {
  const int tid = threadIdx.x;
  const int mts = (T + c.BM - 1) / c.BM;
  const int nbm = (mts + c.GM - 1) / c.GM, nbn = (c.NT + c.GN - 1) / c.GN;
  const int NBs = nbm * nbn;
  const int G = NBs * ns;
  const int S = E + G;
  for (int gi = tid; gi < G; gi += blockDim.x) {
    int eb = (int)(((long)gi * E) / G), ee = (int)(((long)(gi + 1) * E) / G);
    int s = eb + gi;
    for (int e = eb; e < ee; e++, s++) {
      int mc = (s_cnt[e] + c.BM - 1) / c.BM;
      sch[s] = e; sch[S + s] = 0; sch[2 * S + s] = mc; sch[3 * S + s] = 0; sch[4 * S + s] = c.NT;
      s_tiles[s] = mc * c.NT;
    }
    int j = gi / NBs, b = gi % NBs, bm = b / nbn, bn = b % nbn;
    int mb = bm * c.GM, mc = min(c.GM, mts - mb), nb = bn * c.GN, nc = min(c.GN, c.NT - nb);
    sch[s] = E + j; sch[S + s] = mb; sch[2 * S + s] = mc; sch[3 * S + s] = nb; sch[4 * S + s] = nc;
    s_tiles[s] = mc * nc;
  }
  __syncthreads();
  if (tid < 32) {
    int per = (S + 31) / 32;
    int start = tid * per;
    int sum = 0;
    for (int j = 0; j < per; j++) { int idx = start + j; if (idx < S) sum += s_tiles[idx]; }
    int incl = sum;
    for (int o = 1; o < 32; o <<= 1) {
      int v = __shfl_up_sync(0xffffffffu, incl, o);
      if (tid >= o) incl += v;
    }
    int run = incl - sum;
    for (int j = 0; j < per; j++) {
      int idx = start + j;
      if (idx < S) { sch[5 * S + idx] = run; run += s_tiles[idx]; }
    }
    if (tid == 31) sch[5 * S + S] = incl;
  }
  __syncthreads();
}

// gcnt: [E] per-expert counters + [E+ns] done counter, all zero on entry.
__global__ void __launch_bounds__(CNT_THREADS) count_kernel(
    const int64_t* __restrict__ ids, int T, int K, int E, int ns, int* __restrict__ gcnt,
    int* __restrict__ rank, int* __restrict__ seg_off, int* __restrict__ sch1, SchedCfg c1,
    int* __restrict__ sch2, SchedCfg c2) {
  __shared__ int s_cnt[PREP_MAX_E + 8];
  __shared__ int s_off[PREP_MAX_E + 8];
  __shared__ int s_tiles[PREP_MAX_S];
  __shared__ int s_last;
  const int tid = threadIdx.x;
  const int P = T * K;
  const int p0 = blockIdx.x * CNT_PAIRS_PER_CTA;
  for (int i = tid; i < E; i += blockDim.x) s_cnt[i] = 0;
  __syncthreads();
  constexpr int PER = CNT_PAIRS_PER_CTA / CNT_THREADS;
  int ev[PER], lr[PER];
#pragma unroll
  for (int j = 0; j < PER; j++) {
    int i = p0 + tid + j * CNT_THREADS;
    ev[j] = (i < P) ? (int)ids[i] : -1;
  }
#pragma unroll
  for (int j = 0; j < PER; j++) lr[j] = (ev[j] >= 0) ? atomicAdd(&s_cnt[ev[j]], 1) : 0;
  __syncthreads();
  for (int i = tid; i < E; i += blockDim.x) {
    int n = s_cnt[i];
    s_off[i] = n ? atomicAdd(&gcnt[i], n) : 0;
  }
  __syncthreads();
#pragma unroll
  for (int j = 0; j < PER; j++) {
    int i = p0 + tid + j * CNT_THREADS;
    if (i < P) rank[i] = s_off[ev[j]] + lr[j];
  }
  // last CTA to finish: exclusive scan of counts + tile schedules
  __threadfence();
  __syncthreads();
  if (tid == 0) s_last = (atomicAdd(&gcnt[E + ns], 1) == (int)gridDim.x - 1);
  __syncthreads();
  if (!s_last) return;
  __threadfence();
  const int Et = E + ns;
  for (int i = tid; i < Et; i += blockDim.x) s_cnt[i] = (i < E) ? atomicAdd(&gcnt[i], 0) : T;
  __syncthreads();
  if (tid < 32) {
    int per = (Et + 31) / 32;
    int start = tid * per;
    int sum = 0;
    for (int j = 0; j < per; j++) { int idx = start + j; if (idx < Et) sum += s_cnt[idx]; }
    int incl = sum;
    for (int o = 1; o < 32; o <<= 1) {
      int v = __shfl_up_sync(0xffffffffu, incl, o);
      if (tid >= o) incl += v;
    }
    int run = incl - sum;
    for (int j = 0; j < per; j++) {
      int idx = start + j;
      if (idx < Et) { s_off[idx] = run; run += s_cnt[idx]; }
    }
    if (tid == 31) s_off[Et] = incl;
  }
  __syncthreads();
  for (int i = tid; i <= Et; i += blockDim.x) seg_off[i] = s_off[i];
  build_schedule(sch1, s_cnt, s_tiles, E, ns, T, c1);
  build_schedule(sch2, s_cnt, s_tiles, E, ns, T, c2);
}

__global__ void scatter_kernel(const int64_t* __restrict__ ids, const bf16* __restrict__ wts, int T, int K,
                               int E, int ns, const int* __restrict__ rank, const int* __restrict__ seg_off,
                               int* __restrict__ sorted_tok, float* __restrict__ sorted_w) {
  const int P = T * K;
  const int total = P + T * ns;
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += gridDim.x * blockDim.x) {
    if (i < P) {
      int e = (int)ids[i];
      int pos = seg_off[e] + rank[i];
      sorted_tok[pos] = i / K;
      sorted_w[pos] = __bfloat162float(wts[i]);
    } else {
      int r = i - P;
      int j = r / T, t = r % T;
      sorted_tok[seg_off[E + j] + t] = t;
      sorted_w[seg_off[E + j] + t] = 1.f;
    }
  }
}

// ---------------------------------------------------------------------------------
// Grouped GEMM. Rows (MMA M) = tokens of one expert, cols (MMA N) = weight rows.
// MODE 1: A = x gathered by token, B = w1 [gate|up] rows, epilogue silu*up -> h (bf16)
// MODE 2: A = h rows (contiguous), B = w2 rows, epilogue weight * acc red.add -> out f32
// Warps whose rows are all padding skip their MMAs (per-warp count of valid m16 blocks is a
// template parameter so the fragment loads stay unpredicated and double-buffered).
// ---------------------------------------------------------------------------------
template <int V> struct IC { static constexpr int value = V; };

template <int MV, class F>
__device__ __forceinline__ void dispatch_mv(int mv, F& f) {
  if constexpr (MV == 0) {
    f(IC<0>{});
  } else {
    if (mv == MV) f(IC<MV>{});
    else dispatch_mv<MV - 1>(mv, f);
  }
}

template <int MV, int MI, int NI, int BK, int ROWB>
__device__ __forceinline__ void warp_ktile(float (&acc)[MI][NI][4], uint32_t a_base, uint32_t b_base, int xr,
                                           int a_kc, int b_kc) {
  uint32_t af[2][MV][4];
  uint32_t bfr[2][NI][2];
  auto load = [&](int buf, int kk) {
#pragma unroll
    for (int j2 = 0; j2 < NI / 2; j2++) {
      uint32_t r[4];
      ldsm_x4(r, b_base + j2 * 16 * ROWB + (((kk * 2 + b_kc) ^ xr) << 4));
      bfr[buf][2 * j2][0] = r[0]; bfr[buf][2 * j2][1] = r[1];
      bfr[buf][2 * j2 + 1][0] = r[2]; bfr[buf][2 * j2 + 1][1] = r[3];
    }
#pragma unroll
    for (int i = 0; i < MV; i++) ldsm_x4(af[buf][i], a_base + i * 16 * ROWB + (((kk * 2 + a_kc) ^ xr) << 4));
  };
  load(0, 0);
#pragma unroll
  for (int kk = 0; kk < BK / 16; kk++) {
    if (kk + 1 < BK / 16) load((kk + 1) & 1, kk + 1);
#pragma unroll
    for (int i = 0; i < MV; i++)
#pragma unroll
      for (int j = 0; j < NI; j++) mma_bf16(acc[i][j], af[kk & 1][i], bfr[kk & 1][j][0], bfr[kk & 1][j][1]);
  }
}

template <int MODE, int BM, int BN, int BK, int STAGES, int WM, int WN, int MINB>
__global__ void __launch_bounds__(WM * WN * 32, MINB) moe_gemm_kernel(
    const bf16* __restrict__ Asrc, const bf16* __restrict__ Wr, const bf16* __restrict__ Ws,
    const int* __restrict__ seg_off, const int* __restrict__ sorted_tok,
    const float* __restrict__ sorted_w, const int* __restrict__ sch, int S, int E,
    int Kdim, int Nw, bf16* __restrict__ hout, float* __restrict__ out, int ldo) {
  constexpr int NTHREADS = WM * WN * 32;
  constexpr int WTM = BM / WM, WTN = BN / WN;
  constexpr int MI = WTM / 16, NI = WTN / 8;
  constexpr int CPR = BK / 8;
  constexpr int ROWB = BK * 2;
  constexpr int A_STAGE_B = BM * ROWB, B_STAGE_B = BN * ROWB;
  constexpr int A_CH = BM * CPR / NTHREADS, B_CH = BN * CPR / NTHREADS;
  static_assert(BM * CPR % NTHREADS == 0 && BN * CPR % NTHREADS == 0, "load split");
  static_assert(CPR == 8, "swizzle assumes 128B rows");
  static_assert(NI % 2 == 0, "NI even");

  extern __shared__ __align__(1024) uint8_t smem[];
  const uint32_t sA0 = smem_u32(smem);
  const uint32_t sB0 = sA0 + STAGES * A_STAGE_B;

  const int tid = threadIdx.x;
  const int lane = tid & 31, warp = tid >> 5;
  const int wm = warp / WN, wn = warp % WN;

  // ---- tile decode
  const int* cum = sch + 5 * S;
  const int g = blockIdx.x;
  if (g >= cum[S]) return;
  int lo = 0, hi = S - 1;
  while (lo < hi) {
    int mid = (lo + hi + 1) >> 1;
    if (cum[mid] <= g) lo = mid; else hi = mid - 1;
  }
  const int e = sch[lo];
  const int local = g - cum[lo];
  const int mcnt = sch[2 * S + lo];
  const int ntile = sch[3 * S + lo] + local / mcnt;
  const int m0 = (sch[S + lo] + local % mcnt) * BM;
  const int rowbeg = seg_off[e];
  const int cnt = seg_off[e + 1] - rowbeg;
  const int rowbase = rowbeg + m0;
  const int nvalid = min(BM, cnt - m0);
  const bf16* W = (e < E) ? (Wr + (size_t)e * Nw * Kdim) : (Ws + (size_t)(e - E) * Nw * Kdim);

  // ---- per-thread load assignments
  const bf16* a_ptr[A_CH];
  uint32_t a_so[A_CH];
  bool a_ok[A_CH];
#pragma unroll
  for (int i = 0; i < A_CH; i++) {
    int c = tid + i * NTHREADS;
    int r = c / CPR, ch = c % CPR;
    a_so[i] = r * ROWB + ((ch ^ (r & 7)) << 4);
    a_ok[i] = r < nvalid;
    int rr = a_ok[i] ? r : 0;
    size_t arow;
    if (MODE == 1) arow = (size_t)sorted_tok[rowbase + rr];
    else arow = (size_t)(rowbase + rr);
    a_ptr[i] = Asrc + arow * Kdim + ch * 8;
  }
  const bf16* b_ptr[B_CH];
  uint32_t b_so[B_CH];
#pragma unroll
  for (int i = 0; i < B_CH; i++) {
    int c = tid + i * NTHREADS;
    int j = c / CPR, ch = c % CPR;
    b_so[i] = j * ROWB + ((ch ^ (j & 7)) << 4);
    int wrow;
    if (MODE == 1) {
      constexpr int HALF = WTN / 2;
      int wj = j / WTN, s = j % WTN;
      int nh0 = ntile * (BN / 2);
      wrow = (s < HALF) ? (nh0 + wj * HALF + s) : (Nw / 2 + nh0 + wj * HALF + s - HALF);
    } else {
      wrow = ntile * BN + j;
    }
    b_ptr[i] = W + (size_t)wrow * Kdim + ch * 8;
  }

  auto load_stage = [&](uint32_t sa, uint32_t sb, int kt) {
    const int koff = kt * BK;
#pragma unroll
    for (int i = 0; i < A_CH; i++)
      if (a_ok[i]) cp_async16(sa + a_so[i], a_ptr[i] + koff);
#pragma unroll
    for (int i = 0; i < B_CH; i++) cp_async16(sb + b_so[i], b_ptr[i] + koff);
  };

  float acc[MI][NI][4];
#pragma unroll
  for (int i = 0; i < MI; i++)
#pragma unroll
    for (int j = 0; j < NI; j++)
#pragma unroll
      for (int q = 0; q < 4; q++) acc[i][j][q] = 0.f;

  int mi_valid = (nvalid - wm * WTM + 15) / 16;
  mi_valid = mi_valid < 0 ? 0 : (mi_valid > MI ? MI : mi_valid);

  const int KT = Kdim / BK;
#pragma unroll
  for (int s = 0; s < STAGES - 1; s++) {
    if (s < KT) load_stage(sA0 + s * A_STAGE_B, sB0 + s * B_STAGE_B, s);
    cp_async_commit();
  }

  // ldmatrix lane geometry (row & 7 == lane & 7 for both operands)
  const uint32_t a_off = (wm * WTM + (lane & 15)) * ROWB;
  const uint32_t b_off = (wn * WTN + ((lane >> 4) << 3) + (lane & 7)) * ROWB;
  const int a_kc = lane >> 4, b_kc = (lane >> 3) & 1, xr = lane & 7;

  auto run = [&](auto mv_tag) {
    constexpr int MV = decltype(mv_tag)::value;
    int st_r = 0, st_w = STAGES - 1;
    for (int kt = 0; kt < KT; kt++) {
      cp_async_wait<STAGES - 2>();
      __syncthreads();
      {
        int nk = kt + STAGES - 1;
        if (nk < KT) load_stage(sA0 + st_w * A_STAGE_B, sB0 + st_w * B_STAGE_B, nk);
        cp_async_commit();
        st_w = (st_w + 1 == STAGES) ? 0 : st_w + 1;
      }
      if constexpr (MV > 0)
        warp_ktile<MV, MI, NI, BK, ROWB>(acc, sA0 + st_r * A_STAGE_B + a_off, sB0 + st_r * B_STAGE_B + b_off, xr,
                                         a_kc, b_kc);
      st_r = (st_r + 1 == STAGES) ? 0 : st_r + 1;
    }
  };
  dispatch_mv<MI>(mi_valid, run);
  cp_async_wait<0>();

  // ---- epilogue
  const int gq = lane >> 2, tq = lane & 3;
  if (MODE == 1) {
    const int nh0 = ntile * (BN / 2);
#pragma unroll
    for (int i = 0; i < MI; i++) {
#pragma unroll
      for (int hh = 0; hh < 2; hh++) {
        int row = wm * WTM + i * 16 + gq + hh * 8;
        if (row < nvalid) {
          bf16* dst = hout + (size_t)(rowbase + row) * ldo;
#pragma unroll
          for (int j = 0; j < NI / 2; j++) {
            float g0 = acc[i][j][hh * 2], g1 = acc[i][j][hh * 2 + 1];
            float u0 = acc[i][j + NI / 2][hh * 2], u1 = acc[i][j + NI / 2][hh * 2 + 1];
            float h0 = g0 / (1.f + __expf(-g0)) * u0;
            float h1 = g1 / (1.f + __expf(-g1)) * u1;
            int col = nh0 + wn * (WTN / 2) + j * 8 + tq * 2;
            *reinterpret_cast<__nv_bfloat162*>(dst + col) = __floats2bfloat162_rn(h0, h1);
          }
        }
      }
    }
  } else {
#pragma unroll
    for (int i = 0; i < MI; i++) {
#pragma unroll
      for (int hh = 0; hh < 2; hh++) {
        int row = wm * WTM + i * 16 + gq + hh * 8;
        if (row < nvalid) {
          int tok = sorted_tok[rowbase + row];
          float wgt = sorted_w[rowbase + row];
          float* dst = out + (size_t)tok * ldo + ntile * BN + wn * WTN + tq * 2;
#pragma unroll
          for (int j = 0; j < NI; j++)
            red_add_v2(dst + j * 8, acc[i][j][hh * 2] * wgt, acc[i][j][hh * 2 + 1] * wgt);
        }
      }
    }
  }
}

__global__ void f32_to_bf16_kernel(const float4* __restrict__ in, uint2* __restrict__ out, size_t n4) {
  for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; i < n4; i += (size_t)gridDim.x * blockDim.x) {
    float4 v = in[i];
    __nv_bfloat162 a = __floats2bfloat162_rn(v.x, v.y);
    __nv_bfloat162 b = __floats2bfloat162_rn(v.z, v.w);
    uint2 o;
    o.x = *reinterpret_cast<uint32_t*>(&a);
    o.y = *reinterpret_cast<uint32_t*>(&b);
    out[i] = o;
  }
}

// ---------------------------------------------------------------------------------
constexpr int G_BM = 128, G_BN = 128, G_BK = 64, G_ST = 3, G_WM = 2, G_WN = 4;
constexpr int G_SMEM = G_ST * (G_BM + G_BN) * G_BK * 2;
constexpr int SH_GM = 8, SH_GN = 8;

torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor wts, torch::Tensor w1r,
                          torch::Tensor w2r, torch::Tensor w1s, torch::Tensor w2s) {
  const int T = x.size(0), H = x.size(1);
  const int K = ids.size(1);
  const int E = w1r.size(0);
  const int I = w2r.size(2);
  const int ns = w1s.size(0);
  TORCH_CHECK(H % G_BN == 0 && I % (G_BN / 2) == 0 && H % G_BK == 0 && I % G_BK == 0, "shape");
  TORCH_CHECK(E + ns <= PREP_MAX_E, "too many experts");
  auto dev = x.device();
  auto stream = at::cuda::getCurrentCUDAStream();
  auto opt_i = torch::TensorOptions().dtype(torch::kInt32).device(dev);
  auto opt_f = torch::TensorOptions().dtype(torch::kFloat32).device(dev);
  auto opt_b = torch::TensorOptions().dtype(torch::kBFloat16).device(dev);

  const int P = T * K;
  const int PT = T * (K + ns);
  SchedCfg c1{G_BM, I / (G_BN / 2), SH_GM, SH_GN};
  SchedCfg c2{G_BM, H / G_BN, SH_GM, SH_GN};
  const int S1 = E + sched_groups(T, ns, c1), S2 = E + sched_groups(T, ns, c2);
  TORCH_CHECK(S1 <= PREP_MAX_S && S2 <= PREP_MAX_S, "too many segments");

  const int n_gcnt = E + ns + 1;
  auto ws_i = torch::empty({(int64_t)n_gcnt + (E + ns + 2) + P + 2 * PT + (6 * S1 + 1) + (6 * S2 + 1)}, opt_i);
  int* gcnt = ws_i.data_ptr<int>();
  int* seg_off = gcnt + n_gcnt;
  int* rank = seg_off + (E + ns + 2);
  int* sorted_tok = rank + P;
  float* sorted_w = reinterpret_cast<float*>(sorted_tok + PT);
  int* sch1 = reinterpret_cast<int*>(sorted_w + PT);
  int* sch2 = sch1 + (6 * S1 + 1);

  auto hbuf = torch::empty({(int64_t)PT, I}, opt_b);
  auto out32 = torch::zeros({T, H}, opt_f);
  auto out = torch::empty({T, H}, opt_b);

  cudaMemsetAsync(gcnt, 0, n_gcnt * sizeof(int), stream);
  const int cnt_grid = std::max(1, (P + CNT_PAIRS_PER_CTA - 1) / CNT_PAIRS_PER_CTA);
  count_kernel<<<cnt_grid, CNT_THREADS, 0, stream>>>(ids.data_ptr<int64_t>(), T, K, E, ns, gcnt, rank, seg_off,
                                                      sch1, c1, sch2, c2);
  scatter_kernel<<<std::min(1024, (PT + 255) / 256), 256, 0, stream>>>(
      ids.data_ptr<int64_t>(), (const bf16*)wts.data_ptr(), T, K, E, ns, rank, seg_off, sorted_tok, sorted_w);

  const long routed_mt = std::min<long>((long)P, (long)P / G_BM + std::min<long>(E, (long)P));
  const long shared_mt = (long)ns * ((T + G_BM - 1) / G_BM);
  const long grid1 = (routed_mt + shared_mt) * c1.NT;
  const long grid2 = (routed_mt + shared_mt) * c2.NT;

  auto k1 = moe_gemm_kernel<1, G_BM, G_BN, G_BK, G_ST, G_WM, G_WN, 1>;
  auto k2 = moe_gemm_kernel<2, G_BM, G_BN, G_BK, G_ST, G_WM, G_WN, 1>;
  static bool attr_set = false;
  if (!attr_set) {
    cudaFuncSetAttribute(k1, cudaFuncAttributeMaxDynamicSharedMemorySize, G_SMEM);
    cudaFuncSetAttribute(k2, cudaFuncAttributeMaxDynamicSharedMemorySize, G_SMEM);
    attr_set = true;
  }
  k1<<<grid1, G_WM * G_WN * 32, G_SMEM, stream>>>(
      (const bf16*)x.data_ptr(), (const bf16*)w1r.data_ptr(), (const bf16*)w1s.data_ptr(), seg_off,
      sorted_tok, sorted_w, sch1, S1, E, H, 2 * I, (bf16*)hbuf.data_ptr(), nullptr, I);
  k2<<<grid2, G_WM * G_WN * 32, G_SMEM, stream>>>(
      (const bf16*)hbuf.data_ptr(), (const bf16*)w2r.data_ptr(), (const bf16*)w2s.data_ptr(), seg_off,
      sorted_tok, sorted_w, sch2, S2, E, I, H, nullptr, out32.data_ptr<float>(), H);
  const size_t n4 = (size_t)T * H / 4;
  f32_to_bf16_kernel<<<std::min<size_t>((n4 + 255) / 256, 188 * 8), 256, 0, stream>>>(
      (const float4*)out32.data_ptr<float>(), (uint2*)out.data_ptr(), n4);
  return out;
}
"""

CPP_SRC = """
torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor wts, torch::Tensor w1r,
                          torch::Tensor w2r, torch::Tensor w1s, torch::Tensor w2s);
"""

_ext = None


def _get_ext():
    global _ext
    if _ext is None:
        _ext = load_inline(
            name="glm52_fused_moe_v2",
            cpp_sources=CPP_SRC,
            cuda_sources=CUDA_SRC,
            functions=["moe_forward"],
            extra_cuda_cflags=[
                "-O3",
                "-gencode=arch=compute_120a,code=sm_120a",
                "-lineinfo",
            ],
            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:
        ext = _get_ext()
        return ext.moe_forward(
            x.contiguous(),
            expert_ids.contiguous().to(torch.int64),
            expert_weights.contiguous().to(torch.bfloat16),
            self.w1_routed.contiguous(),
            self.w2_routed.contiguous(),
            self.w1_shared.contiguous(),
            self.w2_shared.contiguous(),
        )