"""CUDA implementation of the GLM-5.2 routed and shared MoE layer."""

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


_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#include <algorithm>
#include <vector>

using torch::Tensor;
using bf16 = __nv_bfloat16;

#define CUDA_OK(expr) TORCH_CHECK((expr) == cudaSuccess, "CUDA failure: ", cudaGetErrorString(cudaGetLastError()))
#define CUBLAS_OK(expr) TORCH_CHECK((expr) == CUBLAS_STATUS_SUCCESS, "cuBLAS failure at line ", __LINE__)

__global__ void histogram(const int64_t* ids, int* counts, int entries) {
  int p = blockIdx.x * blockDim.x + threadIdx.x;
  if (p < entries) atomicAdd(counts + ids[p], 1);
}

__global__ void pack_tokens(const bf16* x, const int64_t* ids,
                            int* counts, int* slots, bf16* packed,
                            int H, int K, int M, int entries) {
  int p = blockIdx.x;
  if (p >= entries) return;
  int e = static_cast<int>(ids[p]);
  __shared__ int shared_m;
  if (threadIdx.x == 0) {
    shared_m = atomicAdd(counts + e, 1);
    slots[p] = shared_m;
  }
  __syncthreads();
  const bf16* src = x + (p / K) * H;
  bf16* dst = packed + (static_cast<int64_t>(e) * M + shared_m) * H;
  for (int h = threadIdx.x; h < H; h += blockDim.x) dst[h] = src[h];
}

__global__ void silu_mul(const bf16* projected, bf16* hidden,
                         int rows, int I) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= rows * I) return;
  int row = idx / I, i = idx % I;
  float gate = __bfloat162float(projected[row * (2 * I) + i]);
  float up = __bfloat162float(projected[row * (2 * I) + I + i]);
  float silu = gate / (1.0f + __expf(-gate));
  hidden[idx] = __float2bfloat16_rn(silu * up);
}

__global__ void combine(const float* shared, const float* routed,
                        const int64_t* ids, const bf16* weights,
                        const int* slots, bf16* output,
                        int T, int H, int K, int M) {
  int t = blockIdx.x;
  if (t >= T) return;
  int expert[8], slot[8];
  float weight[8];
#pragma unroll
  for (int k = 0; k < 8; ++k) {
    expert[k] = static_cast<int>(ids[t * K + k]);
    slot[k] = slots[t * K + k];
    weight[k] = __bfloat162float(weights[t * K + k]);
  }
  for (int h = threadIdx.x; h < H; h += blockDim.x) {
    float acc = shared[static_cast<int64_t>(t) * H + h];
#pragma unroll
    for (int k = 0; k < 8; ++k)
      acc += weight[k] * routed[(static_cast<int64_t>(expert[k]) * M + slot[k]) * H + h];
    output[static_cast<int64_t>(t) * H + h] = __float2bfloat16_rn(acc);
  }
}

__global__ void setup_decode(const int64_t* ids,
                             const bf16* x, const bf16* w1r, const bf16* w2r,
                             const bf16* w1s, const bf16* w2s,
                             const bf16* h, bf16* p1, float* y,
                             const void** a1, const void** b1, void** c1,
                             const void** a2, const void** b2, void** c2,
                             int H, int I) {
  int j = threadIdx.x;
  if (j >= 9) return;
  int e = j ? static_cast<int>(ids[j - 1]) : 0;
  a1[j] = j ? static_cast<const void*>(w1r + static_cast<int64_t>(e) * 2 * I * H)
            : static_cast<const void*>(w1s);
  b1[j] = x;
  c1[j] = p1 + static_cast<int64_t>(j) * 2 * I;
  a2[j] = j ? static_cast<const void*>(w2r + static_cast<int64_t>(e) * H * I)
            : static_cast<const void*>(w2s);
  b2[j] = h + static_cast<int64_t>(j) * I;
  c2[j] = y + static_cast<int64_t>(j) * H;
}

__global__ void combine_decode(const float* y, const bf16* weights,
                               bf16* output, int H) {
  int h = blockIdx.x * blockDim.x + threadIdx.x;
  if (h >= H) return;
  float acc = y[h];
#pragma unroll
  for (int k = 0; k < 8; ++k)
    acc += __bfloat162float(weights[k]) * y[(k + 1) * H + h];
  output[h] = __float2bfloat16_rn(acc);
}

void lt_batched(cublasLtHandle_t lt, cudaStream_t stream,
                const void* A, const void* B, void* C,
                int m, int n, int k, int batch,
                int64_t stride_a, int64_t stride_b, int64_t stride_c,
                cudaDataType_t output_type, void* workspace, size_t workspace_size) {
  cublasLtMatmulDesc_t desc;
  cublasLtMatrixLayout_t la, lb, lc;
  cublasLtMatmulPreference_t pref;
  CUBLAS_OK(cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F, CUDA_R_32F));
  cublasOperation_t trans_a = CUBLAS_OP_T, trans_b = CUBLAS_OP_N;
  CUBLAS_OK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a)));
  CUBLAS_OK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b)));
  CUBLAS_OK(cublasLtMatrixLayoutCreate(&la, CUDA_R_16BF, k, m, k));
  CUBLAS_OK(cublasLtMatrixLayoutCreate(&lb, CUDA_R_16BF, k, n, k));
  CUBLAS_OK(cublasLtMatrixLayoutCreate(&lc, output_type, m, n, m));
  CUBLAS_OK(cublasLtMatrixLayoutSetAttribute(la, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch, sizeof(batch)));
  CUBLAS_OK(cublasLtMatrixLayoutSetAttribute(lb, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch, sizeof(batch)));
  CUBLAS_OK(cublasLtMatrixLayoutSetAttribute(lc, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch, sizeof(batch)));
  CUBLAS_OK(cublasLtMatrixLayoutSetAttribute(la, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_a, sizeof(stride_a)));
  CUBLAS_OK(cublasLtMatrixLayoutSetAttribute(lb, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_b, sizeof(stride_b)));
  CUBLAS_OK(cublasLtMatrixLayoutSetAttribute(lc, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_c, sizeof(stride_c)));
  CUBLAS_OK(cublasLtMatmulPreferenceCreate(&pref));
  CUBLAS_OK(cublasLtMatmulPreferenceSetAttribute(pref,
      CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspace_size, sizeof(workspace_size)));
  cublasLtMatmulHeuristicResult_t heuristic;
  int found = 0;
  CUBLAS_OK(cublasLtMatmulAlgoGetHeuristic(lt, desc, la, lb, lc, lc,
      pref, 1, &heuristic, &found));
  TORCH_CHECK(found > 0, "no cuBLASLt batched GEMM algorithm");
  const float alpha = 1.f, beta = 0.f;
  CUBLAS_OK(cublasLtMatmul(lt, desc, &alpha, A, la, B, lb, &beta,
      C, lc, C, lc, &heuristic.algo, workspace, workspace_size, stream));
  CUBLAS_OK(cublasLtMatmulPreferenceDestroy(pref));
  CUBLAS_OK(cublasLtMatrixLayoutDestroy(la));
  CUBLAS_OK(cublasLtMatrixLayoutDestroy(lb));
  CUBLAS_OK(cublasLtMatrixLayoutDestroy(lc));
  CUBLAS_OK(cublasLtMatmulDescDestroy(desc));
}

Tensor moe_forward(Tensor x, Tensor ids, Tensor weights,
                   Tensor w1r, Tensor w2r, Tensor w1s, Tensor w2s) {
  TORCH_CHECK(x.is_cuda() && ids.is_cuda() && weights.is_cuda(), "CUDA inputs required");
  TORCH_CHECK(x.is_contiguous() && ids.is_contiguous() && weights.is_contiguous(), "contiguous inputs required");
  const int T = x.size(0), H = x.size(1), K = ids.size(1);
  const int E = w1r.size(0), I = w2r.size(2);
  TORCH_CHECK(K == 8 && w1s.size(0) == 1, "expected top-8 and one shared expert");
  auto stream = at::cuda::getCurrentCUDAStream().stream();
  auto handle = at::cuda::getCurrentCUDABlasHandle();
  CUBLAS_OK(cublasSetStream(handle, stream));
  const float alpha = 1.f, beta = 0.f;
  auto output = torch::empty({T, H}, x.options());
  const bf16* xp = reinterpret_cast<const bf16*>(x.data_ptr());
  const int64_t* ip = ids.data_ptr<int64_t>();
  const bf16* wp = reinterpret_cast<const bf16*>(weights.data_ptr());
  const bf16* w1rp = reinterpret_cast<const bf16*>(w1r.data_ptr());
  const bf16* w2rp = reinterpret_cast<const bf16*>(w2r.data_ptr());
  const bf16* w1sp = reinterpret_cast<const bf16*>(w1s.data_ptr());
  const bf16* w2sp = reinterpret_cast<const bf16*>(w2s.data_ptr());
  bf16* op = reinterpret_cast<bf16*>(output.data_ptr());

  if (T == 1) {
    auto p1 = torch::empty({9, 2 * I}, x.options());
    auto h = torch::empty({9, I}, x.options());
    auto y = torch::empty({9, H}, x.options().dtype(torch::kFloat32));
    auto ptrs = torch::empty({54}, x.options().dtype(torch::kInt64));
    int64_t* ptr = ptrs.data_ptr<int64_t>();
    setup_decode<<<1, 32, 0, stream>>>(ip, xp, w1rp, w2rp, w1sp, w2sp,
        reinterpret_cast<const bf16*>(h.data_ptr()),
        reinterpret_cast<bf16*>(p1.data_ptr()), y.data_ptr<float>(),
        reinterpret_cast<const void**>(ptr),
        reinterpret_cast<const void**>(ptr + 9), reinterpret_cast<void**>(ptr + 18),
        reinterpret_cast<const void**>(ptr + 27),
        reinterpret_cast<const void**>(ptr + 36), reinterpret_cast<void**>(ptr + 45),
        H, I);
    CUBLAS_OK(cublasGemmBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
        2 * I, 1, H, &alpha, reinterpret_cast<const void* const*>(ptr), CUDA_R_16BF, H,
        reinterpret_cast<const void* const*>(ptr + 9), CUDA_R_16BF, H,
        &beta, reinterpret_cast<void**>(ptr + 18), CUDA_R_16BF, 2 * I, 9,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
    silu_mul<<<(9 * I + 255) / 256, 256, 0, stream>>>(
        reinterpret_cast<const bf16*>(p1.data_ptr()),
        reinterpret_cast<bf16*>(h.data_ptr()), 9, I);
    CUBLAS_OK(cublasGemmBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
        H, 1, I, &alpha, reinterpret_cast<const void* const*>(ptr + 27), CUDA_R_16BF, I,
        reinterpret_cast<const void* const*>(ptr + 36), CUDA_R_16BF, I,
        &beta, reinterpret_cast<void**>(ptr + 45), CUDA_R_32F, H, 9,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
    combine_decode<<<(H + 255) / 256, 256, 0, stream>>>(y.data_ptr<float>(), wp, op, H);
    return output;
  }

  auto counts = torch::zeros({E}, x.options().dtype(torch::kInt32));
  histogram<<<(T * K + 255) / 256, 256, 0, stream>>>(ip, counts.data_ptr<int>(), T * K);
  auto host_counts = counts.to(torch::kCPU);
  const int* cp = host_counts.data_ptr<int>();
  int maximum = 0;
  for (int e = 0; e < E; ++e) maximum = std::max(maximum, cp[e]);
  const int M = std::max(16, (maximum + 15) & ~15);
  auto packed = torch::empty({E, M, H}, x.options());
  auto slots = torch::empty({T, K}, x.options().dtype(torch::kInt32));
  CUDA_OK(cudaMemsetAsync(counts.data_ptr(), 0, E * sizeof(int), stream));
  pack_tokens<<<T * K, 256, 0, stream>>>(xp, ip, counts.data_ptr<int>(),
      slots.data_ptr<int>(), reinterpret_cast<bf16*>(packed.data_ptr()), H, K, M, T * K);

  auto p1r = torch::empty({E, M, 2 * I}, x.options());
  auto hr = torch::empty({E, M, I}, x.options());
  auto yr = torch::empty({E, M, H}, x.options().dtype(torch::kFloat32));
  auto p1s = torch::empty({T, 2 * I}, x.options());
  auto hs = torch::empty({T, I}, x.options());
  auto ys = torch::empty({T, H}, x.options().dtype(torch::kFloat32));
  const bool use_lt = T <= 1024;
  Tensor workspace;
  cublasLtHandle_t lt = nullptr;
  if (use_lt) {
    workspace = torch::empty({64 * 1024 * 1024}, x.options().dtype(torch::kUInt8));
    CUBLAS_OK(cublasLtCreate(&lt));
    lt_batched(lt, stream, w1rp, packed.data_ptr(), p1r.data_ptr(),
        2 * I, M, H, E, static_cast<int64_t>(2) * I * H,
        static_cast<int64_t>(M) * H, static_cast<int64_t>(M) * 2 * I,
        CUDA_R_16BF, workspace.data_ptr(), 64 * 1024 * 1024);
  } else {
    CUBLAS_OK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
        2 * I, M, H, &alpha, w1rp, CUDA_R_16BF, H,
        static_cast<long long>(2) * I * H, packed.data_ptr(), CUDA_R_16BF, H,
        static_cast<long long>(M) * H, &beta, p1r.data_ptr(), CUDA_R_16BF,
        2 * I, static_cast<long long>(M) * 2 * I, E,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
  }
  silu_mul<<<(E * M * I + 255) / 256, 256, 0, stream>>>(
      reinterpret_cast<const bf16*>(p1r.data_ptr()),
      reinterpret_cast<bf16*>(hr.data_ptr()), E * M, I);
  if (use_lt) {
    lt_batched(lt, stream, w2rp, hr.data_ptr(), yr.data_ptr(),
        H, M, I, E, static_cast<int64_t>(H) * I,
        static_cast<int64_t>(M) * I, static_cast<int64_t>(M) * H,
        CUDA_R_32F, workspace.data_ptr(), 64 * 1024 * 1024);
    CUBLAS_OK(cublasLtDestroy(lt));
  } else {
    CUBLAS_OK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
        H, M, I, &alpha, w2rp, CUDA_R_16BF, I,
        static_cast<long long>(H) * I, hr.data_ptr(), CUDA_R_16BF, I,
        static_cast<long long>(M) * I, &beta, yr.data_ptr(), CUDA_R_32F,
        H, static_cast<long long>(M) * H, E,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
  }

  CUBLAS_OK(cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
      2 * I, T, H, &alpha, w1sp, CUDA_R_16BF, H, xp, CUDA_R_16BF, H,
      &beta, p1s.data_ptr(), CUDA_R_16BF, 2 * I,
      CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
  silu_mul<<<(T * I + 255) / 256, 256, 0, stream>>>(
      reinterpret_cast<const bf16*>(p1s.data_ptr()),
      reinterpret_cast<bf16*>(hs.data_ptr()), T, I);
  CUBLAS_OK(cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
      H, T, I, &alpha, w2sp, CUDA_R_16BF, I, hs.data_ptr(), CUDA_R_16BF, I,
      &beta, ys.data_ptr(), CUDA_R_32F, H,
      CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));

  combine<<<T, 256, 0, stream>>>(ys.data_ptr<float>(), yr.data_ptr<float>(),
      ip, wp, slots.data_ptr<int>(), op, T, H, K, M);
  return output;
}
"""


_CPP = """
#include <torch/extension.h>
torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor weights,
                          torch::Tensor w1r, torch::Tensor w2r,
                          torch::Tensor w1s, torch::Tensor w2s);
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="glm52_moe_cuda_v1",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            functions=["moe_forward"],
            extra_cuda_cflags=["-O3", "-gencode=arch=compute_120,code=sm_120"],
            extra_ldflags=["-lcublas"],
            with_cuda=True,
            verbose=False,
        )
    return _extension


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:
        return _get_extension().moe_forward(
            x, expert_ids, expert_weights,
            self.w1_routed, self.w2_routed, self.w1_shared, self.w2_shared,
        )