"""CUDA tensor-core MoE with GPU routing and a bandwidth-oriented decode path."""
import hashlib
import os

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

_CPP = r"""
#include <torch/extension.h>
torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor weights,
                         torch::Tensor w1, torch::Tensor w2,
                         torch::Tensor s1, torch::Tensor s2);
"""
_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
using bf = __nv_bfloat16;

__device__ __forceinline__ float warp_sum(float v) {
  #pragma unroll
  for (int d=16; d; d>>=1) v+=__shfl_down_sync(0xffffffff,v,d);
  return v;
}
__global__ void histogram(const int64_t* ids, int* counts, int R) {
  int j=blockIdx.x*blockDim.x+threadIdx.x;
  if(j<R) atomicAdd(counts+ids[j],1);
}
// One serial prefix has only E=256 entries. All subsequent work stays on GPU.
__global__ void layout(int* counts,int* offsets,int* cursor,int4* tiles,
                       int* ntile,int E,int T,int S,int BM) {
  if(threadIdx.x) return;
  int offset=0, b=0;
  for(int e=0;e<E+S;e++) {
    int n=e<E?counts[e]:T;
    offsets[e]=offset; cursor[e]=0;
    for(int m=0;m<n;m+=BM) tiles[b++]=make_int4(e,m,min(BM,n-m),offset+m);
    offset+=(n+BM-1)/BM*BM;
  }
  *ntile=b;
}
__global__ void scatter(const int64_t* ids,const int* offsets,int* cursor,int* order,int R) {
  int j=blockIdx.x*blockDim.x+threadIdx.x;
  if(j<R) {int e=ids[j]; int p=atomicAdd(cursor+e,1);order[offsets[e]+p]=j;}
}
__device__ __forceinline__ uint32_t smaddr(const void* p) {
  return (uint32_t)__cvta_generic_to_shared(p);
}
__device__ __forceinline__ void copy16(bf* d,const bf* s,bool valid) {
  uint32_t a=smaddr(d);
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"::"r"(a),"l"(s),"r"(valid?16:0));
}
__device__ __forceinline__ void ld_a(uint32_t* r,const bf* p) {
  uint32_t a=smaddr(p);
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
    :"=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]):"r"(a));
}
__device__ __forceinline__ void ld_b(uint32_t* r,const bf* p) {
  uint32_t a=smaddr(p);
  asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
    :"=r"(r[0]),"=r"(r[1]):"r"(a));
}
__device__ __forceinline__ void mma(float* c,const uint32_t* a,const uint32_t* b) {
  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};"
    :"+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"(b[0]),"r"(b[1]));
}
__device__ __forceinline__ int sw(int row,int col) {return row*64+(col^((row&7)*8));}

template<bool FIRST,int BM>
__device__ __forceinline__ void stage(bf* a,bf* b,const bf* input,const bf* wr,const bf* ws,
                                     const int* order,int4 tile,int K,int N,int H,int I,int E,int top,int ko,int T) {
  int tid=threadIdx.x;
  #pragma unroll
  for(int j=tid;j<BM*8;j+=128) {
    int r=j/8,c=(j%8)*8;
    int pos=tile.w+r;
    int [REDACTED credential assignment];
    if(r<tile.z) [REDACTED credential assignment]]/top:tile.y+r;
    const bf* src=FIRST?input+(int64_t)token*H+ko+c:input+(int64_t)pos*I+ko+c;
    copy16(a+sw(r,c),src,r<tile.z);
  }
  constexpr int NB=FIRST?128:64;
  #pragma unroll
  for(int j=tid;j<NB*8;j+=128) {
    int r=j/8,c=(j%8)*8;
    int nr=(r%64)+blockIdx.x*64;
    if(FIRST && r>=64) nr+=I;
    int stride=FIRST?2*I:H;
    const bf* w=tile.x<E?wr+(int64_t)tile.x*stride*K:ws+(int64_t)(tile.x-E)*stride*K;
    copy16(b+sw(r,c),w+(int64_t)nr*K+ko+c,nr<stride);
  }
  asm volatile("cp.async.commit_group;");
}

template<bool FIRST,int BM>
__global__ __launch_bounds__(128) void grouped(const bf* input,const bf* wr,const bf* ws,
             const int* order,const int4* tiles,const int* ntile,bf* hidden,float* result,
             int H,int I,int E,int top,int T) {
  if(blockIdx.y>=*ntile) return;
  int4 tile=tiles[blockIdx.y];
  constexpr int NB=FIRST?128:64;
  __shared__ bf storage[2*(BM+NB)*64];
  bf* a0=storage; bf* b0=a0+BM*64;
  bf* a1=b0+NB*64; bf* b1=a1+BM*64;
  int K=FIRST?H:I, N=FIRST?I:H;
  int lane=threadIdx.x%32,warp=threadIdx.x/32;
  int wm=(warp/2)*16, wn=(warp%2)*32;
  float acc[BM/32][4][4]={}; float up[BM/32][4][4]={};
  stage<FIRST,BM>(a0,b0,input,wr,ws,order,tile,K,N,H,I,E,top,0,T);
  stage<FIRST,BM>(a1,b1,input,wr,ws,order,tile,K,N,H,I,E,top,64,T);
  for(int ko=0;ko<K;ko+=64) {
    if(ko+64<K) asm volatile("cp.async.wait_group 1;");
    else asm volatile("cp.async.wait_group 0;");
    __syncthreads();
    int buf=(ko/64)&1;
    bf* a=buf?a1:a0;bf* b=buf?b1:b0;
    #pragma unroll
    for(int kk=0;kk<64;kk+=16) {
      uint32_t ar[BM/32][4];
      #pragma unroll
      for(int mm=0;mm<BM/32;mm++) {
        int ra=wm+mm*32+lane%16, ca=kk+(lane/16)*8;
        ld_a(ar[mm],a+sw(ra,ca));
      }
      #pragma unroll
      for(int nn=0;nn<4;nn++) {
        uint32_t br[2];
        int rb=wn+nn*8+lane%8,cb=kk+((lane/8)%2)*8;
        ld_b(br,b+sw(rb,cb));
        #pragma unroll
        for(int mm=0;mm<BM/32;mm++) mma(acc[mm][nn],ar[mm],br);
        if(FIRST) {
          ld_b(br,b+sw(rb+64,cb));
          #pragma unroll
          for(int mm=0;mm<BM/32;mm++) mma(up[mm][nn],ar[mm],br);
        }
      }
    }
    __syncthreads();
    if(ko+128<K) stage<FIRST,BM>(a,b,input,wr,ws,order,tile,K,N,H,I,E,top,ko+128,T);
  }
  #pragma unroll
  for(int mm=0;mm<BM/32;mm++) {
    #pragma unroll
    for(int nn=0;nn<4;nn++) {
      #pragma unroll
      for(int j=0;j<4;j++) {
        int r=wm+mm*32+lane/4+(j/2)*8;
        int c=blockIdx.x*64+wn+nn*8+(lane%4)*2+(j%2);
        if(r<tile.z && c<N) {
          if(FIRST) {
            float g=acc[mm][nn][j];float h=g/(1.0f+expf(-g))*up[mm][nn][j];
            hidden[(int64_t)(tile.w+r)*I+c]=__float2bfloat16_rn(h);
          } else {
            int dst=tile.x<E?order[tile.w+r]:T*top+(tile.x-E)*T+tile.y+r;
            result[(int64_t)dst*H+c]=acc[mm][nn][j];
          }
        }
      }
    }
  }
}

__global__ void combine(const float* result,const bf* weights,bf* out,int T,int H,int top,int S) {
  int j=blockIdx.x*blockDim.x+threadIdx.x;
  if(j<T*H) {
    int t=j/H,c=j%H;float v=0;
    for(int s=0;s<S;s++) v+=result[((int64_t)T*top+s*T+t)*H+c];
    for(int k=0;k<top;k++) v+=__bfloat162float(weights[t*top+k])*result[((int64_t)t*top+k)*H+c];
    out[j]=__float2bfloat16_rn(v);
  }
}

template<bool FIRST>
__global__ void decode(const bf* x,const bf* wr,const bf* ws,const int64_t* ids,bf* h,float* y,
                       int H,int I,int E,int top) {
  int expert=blockIdx.y;
  int e=expert<top?ids[expert]:expert-top;
  int row=blockIdx.x*8+threadIdx.x/32;
  int lane=threadIdx.x%32;
  int K=FIRST?H:I, N=FIRST?I:H;
  if(row>=N) return;
  const bf* w=expert<top?wr:ws;
  w+=(int64_t)e*(FIRST?2*I:H)*K+(int64_t)row*K;
  const bf* a=FIRST?x:h+expert*I;
  float g=0,u=0;
  for(int k=lane;k<K;k+=32) {
    float v=__bfloat162float(a[k]);g=fmaf(v,__bfloat162float(w[k]),g);
    if(FIRST) u=fmaf(v,__bfloat162float(w[(int64_t)I*K+k]),u);
  }
  g=warp_sum(g);if(FIRST) u=warp_sum(u);
  if(lane==0) {
    if(FIRST) h[expert*I+row]=__float2bfloat16_rn(g/(1.f+expf(-g))*u);
    else y[expert*H+row]=g;
  }
}

torch::Tensor moe_forward(torch::Tensor x,torch::Tensor ids,torch::Tensor weights,
                         torch::Tensor w1,torch::Tensor w2,torch::Tensor s1,torch::Tensor s2) {
  c10::cuda::CUDAGuard guard(x.device());
  int T=x.size(0), H=x.size(1),I=w2.size(2),E=w1.size(0),S=s1.size(0),top=ids.size(1);
  TORCH_CHECK(H%64==0 && I%64==0,"H and I must be multiples of 64");
  auto stream=at::cuda::getCurrentCUDAStream();
  auto out=torch::empty({T,H},x.options());
  auto result=torch::empty({(int64_t)T*(top+S),H},x.options().dtype(torch::kFloat32));
  bf* xp=(bf*)x.data_ptr();bf* w1p=(bf*)w1.data_ptr();bf* w2p=(bf*)w2.data_ptr();
  bf* s1p=(bf*)s1.data_ptr();bf* s2p=(bf*)s2.data_ptr();
  if(T==1) {
    auto hidden=torch::empty({top+S,I},x.options());
    decode<true><<<dim3((I+7)/8,top+S),256,0,stream>>>(xp,w1p,s1p,ids.data_ptr<int64_t>(),(bf*)hidden.data_ptr(),result.data_ptr<float>(),H,I,E,top);
    decode<false><<<dim3((H+7)/8,top+S),256,0,stream>>>(xp,w2p,s2p,ids.data_ptr<int64_t>(),(bf*)hidden.data_ptr(),result.data_ptr<float>(),H,I,E,top);
  } else {
    int BM=T>=2048?64:32;
    int R=T*top, cap=R+T*S+(E+S)*BM, maxtiles=(R+BM-1)/BM+E+S*((T+BM-1)/BM);
    auto io=x.options().dtype(torch::kInt32);
    auto counts=torch::empty({E},io);auto offsets=torch::empty({E+S},io);
    auto cursor=torch::empty({E+S},io);auto order=torch::empty({cap},io);
    auto tiles=torch::empty({maxtiles,4},io);auto ntile=torch::empty({1},io);
    auto hidden=torch::empty({cap,I},x.options());
    cudaMemsetAsync(counts.data_ptr(),0,E*sizeof(int),stream);
    histogram<<<(R+255)/256,256,0,stream>>>(ids.data_ptr<int64_t>(),counts.data_ptr<int>(),R);
    layout<<<1,32,0,stream>>>(counts.data_ptr<int>(),offsets.data_ptr<int>(),cursor.data_ptr<int>(),(int4*)tiles.data_ptr(),ntile.data_ptr<int>(),E,T,S,BM);
    scatter<<<(R+255)/256,256,0,stream>>>(ids.data_ptr<int64_t>(),offsets.data_ptr<int>(),cursor.data_ptr<int>(),order.data_ptr<int>(),R);
    if(BM==64) {
    grouped<true,64><<<dim3(I/64,maxtiles),128,0,stream>>>(xp,w1p,s1p,order.data_ptr<int>(),(int4*)tiles.data_ptr(),ntile.data_ptr<int>(),(bf*)hidden.data_ptr(),result.data_ptr<float>(),H,I,E,top,T);
    grouped<false,64><<<dim3(H/64,maxtiles),128,0,stream>>>((bf*)hidden.data_ptr(),w2p,s2p,order.data_ptr<int>(),(int4*)tiles.data_ptr(),ntile.data_ptr<int>(),(bf*)hidden.data_ptr(),result.data_ptr<float>(),H,I,E,top,T);
    } else {
    grouped<true,32><<<dim3(I/64,maxtiles),128,0,stream>>>(xp,w1p,s1p,order.data_ptr<int>(),(int4*)tiles.data_ptr(),ntile.data_ptr<int>(),(bf*)hidden.data_ptr(),result.data_ptr<float>(),H,I,E,top,T);
    grouped<false,32><<<dim3(H/64,maxtiles),128,0,stream>>>((bf*)hidden.data_ptr(),w2p,s2p,order.data_ptr<int>(),(int4*)tiles.data_ptr(),ntile.data_ptr<int>(),(bf*)hidden.data_ptr(),result.data_ptr<float>(),H,I,E,top,T);
    }
  }
  combine<<<(T*H+255)/256,256,0,stream>>>(result.data_ptr<float>(),(bf*)weights.data_ptr(),(bf*)out.data_ptr(),T,H,top,S);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return out;
}
"""
_ext = None

def _extension():
    global _ext
    if _ext is None:
        os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
        os.environ.setdefault("MAX_JOBS", "6")
        name = "glm_moe_cuda_" + hashlib.sha256((_CPP + _CUDA).encode()).hexdigest()[:12]
        _ext = load_inline(name=name, cpp_sources=_CPP, cuda_sources=_CUDA,
                           functions=["moe_forward"], extra_cflags=["-O3"],
                           extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"])
    return _ext


class Model(nn.Module):
    def __init__(self, T, E, top_k, n_shared, H, I):
        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, expert_ids, expert_weights):
        return _extension().moe_forward(x, expert_ids, expert_weights,
                                        self.w1_routed, self.w2_routed,
                                        self.w1_shared, self.w2_shared)