import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Any, Dict, Tuple, Optional, List
import triton.language as tl
import triton
import functools
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline
import lzma
import base64

import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

cuda_code = """{Wp48S^xk9=GL@E0stWa761SMbT8$j;7RQev|RuN36*Hjowa0`JA-;OMmpdSM@fDN)ed=~IkDS*f^C7omKppK6)r9#`GEu!gAiBVobhI~C5Ia!Nwj4-6!l|)jEmMZnn`ga@V`M+G3_>B<f<86^Jp^wo~sK4D>5zH|E151osTNI@HR^vS@+^JwU|V=$@-v-wP$-nGdLg~mR)!!%7VYPA^fT*QOEc@PTlZdrLRr;!@P=@F8s8<+y!k|cYm2<lvJhwecg(PyS<D1YeA{ReO)tufSoBP6B*$sk{7T*OEBin8sCm0LrplF-#=NbIzF?Az<f2E-=-Wdr~_nP5PN$#d=O2aeB4tLSRgQ<o{e{5f&_VDz%YTsMuUkIaZlHhZuTiLFq2|_r1uAeyxk1X@m5~k>1XXW=8;rr7h@VqZ?pv6Gf6OK<z=rh*wM6HI6lC>HQ2yjt3$owFo7-vI19Mze2T;cu$ng->qbj>-Z!H&wZt-}TJVi3zx`gHyBy1-<1exrf}MTL?5W}Y4)WyH=t3-r&#ne16SQIvyebJ1YSV0!O+(Mj1<^2&z-V(ysvqdfIEnx<4HpB6JiHLH<R-3Oc=*6y2~KHn9o0hftDsACl%2n(C4A_HY>CLrvCXr9V5M~vO~7)o@T(q#_PHpBbfr8mk=2RrzWLj+@=whEdKSYAvNvJxXl(CF_-PWmy$`8ciRS&(CwPij#dC$T91(Egp^1OvS74i@UOXOxoX%&ECbQ1OX)LS0upaqMk79NORk`q69m6cD!wFVbUCgJ}Y`P>K;*r-cLX%I`^|g!OS6n$2|6(<b`uY+e{n>uyOd_Lek*w=)b7-)meY?5Kbd%?n*6<Z@(Re7_cILf(0>D1+jRJGDGf0i+BXC;~xay@lewy(P>MHkuAziiQB6l0`ooi@v!BuW7c>yKYBAt^uG!oEc;QytfZEZx86n57J3!UleQ8oP3bFP*rmZBYEod8CE$CU>}SLLHFzSEJ1eU}2Ab%q!m>pDGWlgK+^VoEH*NJJ<@3G!-<)oF)k&qF5SD0hg9BFYtij7Yl(4RYP*Z(I#K8Q))4?u2g~DaDTA8yN~c+_JG%1wp?O){O+TBvUecN}Ncx)$H{3`JuZ(H|$JG<A;(u<di}(UbyW@Jf!C#=@I3HTxYC&8EX%ZasJ%CJFwvtCY*kpj1n$922fA~m-XehjOF2Q*^o4B`WirZ0{)c&n#54=r6cTc;YpVH{h#tAz_gZ#T)5b)9<t`so(0cZ)!8yWZnU>ehZ(?avpoX&%hc$gd}^LNEgMu`Kza;g-gvquApPctE+UDc3;SOnT`ZC3JFUos^DzD-y6n9qlF8qY9DKeh-Z&7K1=mG>5995tmmE!(8B5gJp#NU7{$SwW{xK8FNK5ruQRD%;Ehr2x1q2-Dh|9+@b@(HYIK^bH&~OV_NGG&*iK^DJ>M|g)uDWj*e*$h@?Xs!rCI810;CvPuIW}^A#8Q+^7x#4tap+!@hGGTp;t)0IzH<-9PbLfdSwm((((J25*HN`3#W75J!E8B5<tm(1sJ%4>bLWCxb$j9vHH8m-V@=qTKWj+x<E&uPo7|V>9c~xZqx#w6AmHHF`^tciICMk<p`8pFQ<*6V+|ASl660J3GrS95b6VD<*Ru&Pq(pu1;%jLyT<S~svR1$*Cf^T*fMU|L`LCr57Os0>(^hk$RNfz9XE^7^;AUK3(RcWBboar1lwDr!apJKZP$CAK{j7|gsDHVISs#JdT6@+?yujOI`h+S2R>{`CWBwbw2$lR|sFyA0BJM3Zf`UFj;lt(hqiALj$5ol?rt*IVmi*yd6>lJmicwo;nX9*)RZAb)(U~Q|D61M`-q3+$KWuR|j4b7#^JltE%jyXK76q4JfPeXFoKLbE7;#;#)=xev|A_Z_b!g!8eked?{IG+eTTET?jsqap90y++IM3OYww?cKcG0bU^GBeZxPD6co}YE;ZY_<0A|q-9taK2_eZIER;Nt_`RAbOxtGp2;m+oj=``c4SuYQ&T$V8O3*by*uY(HJl7Rc&=Unx6fkf1h0PeGq47kXXe*HbXituJy`>`m1UDk6^C_kC&IzSxZ$tz)TRbI0vQZQo;*JBRmOr}k1;8GWV}Z&mf6iW>Z~xl6lSnn#g=zZ&bbP3V5ccOPHF+6F$cQK)L>qRFpN55dK{pk`cj*~{b5k^C4z9rqUuHQ+mF-g#h}Zyl;r<=<Y4zCPxQG3GeXz=SUipnKI)ij_lADBxh=NVdh1$IIU|QG^`01mv^b#m5)1z?!_9<_{b-+e-1$;mSJ#L8WJ|pX>rsEKQ#vLV11wEDIY|S&81sh`}lD4p_R!UZUDKQ-Mz<6Z38(m!1VgqeyI@|2eIrdp6IpBF*o)f;g-Mh*FwrviNXRCA_UoJoz~<+-{cMklFkFycXg9c5cHSlNVIIE|k9slW6|ulD_^gZdzamc&FagI=`3uhs_)m(Wy{88o<Ci!rOCmA?C6PyLJVvE7_v#P5W{xjaNBU+YLWQP2pelv~&|g?-M$?9jWPg7+Z4VB<v-tirra|xVz@#PPE66?4z^QaiO@s2!n%Ijr$30!vn(k(=9)@s1(x=N(<p3fRKTC3TF~J2qF!ITW&6xn@a^H)nQ0j(t2AbE;)({hV5LuO$bRBN(0<bzqhnWd;%;pA%jhOEQPi*se4jbN1rYSp!UCsFw$L_v3g29W99t@)&;DFgUjWvoL3k<mHYP8tUUvS_B@~&r*F^d&w0G;5du9aAVDud1FG;sEf3sPDYP0>^y9l(O&Y~!K<VvFWp*6CeIkB>ns?yG>Z#ls$L2>(I)6;%MOWkX^l8R4BbpfL(FrN|)OR4id521v0Hg@TKG^g`>{DLL>%fuy#Bq{=Ee|0>T3LYzc60PMlJZb+nz0=|pJ#p>IRR^Ru5S*;r>6yxTG4)~c&WALSU}gErE+aLY!`!&Iqw_9%th+Uj(!X27(eHuORnyW->?Oxpfh9DuW@_w;SK@3Uebup2LetqyZ|j=LZ)r5OXvC#Dg)&*o57xjOcZtGy!9uMj~##8%^l9L47!4#iBu%)a@;uA(+WFY9j_Ng&S|M8xnN@gV4-c@X+F|F*iLuyIZ{BVB$U?KrI`*JuYO!U(L8^s$F{WnPFl$T;Oc-BYY#=_A`+Mcad1GzL1@?lY~PfG$B{le&0><Un+|^6xB3Hs|4ewx1U%FO)tq4k6yS^4^t+(q-{lSA&*+uQ+~J*fqoz<nVA=NzFgFMr!yYSNq_ipL3I}ChD1yxlsH1xuX%?W^MCW4qn6Myb%i<`1q0m9ZT7!y*8WMp<q8|=NMh=HQEajjThCyy8(i>03BqhzH!0i;rj-hilx=ObA%WD(6n%ceCu&Z`UbX%iEpF6o;EY3LkGGIO7ajpo`DUXoeG=tkb_ZQjOa)%cq!Qje|>KUg%K$y)C8_^n1O$ip_MG|#%6N*HuHt@3g`dh&6m<y|k)%tj)T_U~j!m65y29ESj#xwtT9=c$DuH>!z>!A6+^Mq9CO_agYi$<|!$xbyr`rzfv;=-?NpEQp+ADgq>m$o6@cYFz8aE*H^fQcJvEVlZ*AOY&>9uMNOp>yaXZGN6W-9Ahuv>xGS{c^`b@h%myC^Z@lTAD^blEe2(q@3b+fn@(#c&8f~_c0j$H8`LBC#wVji_H3=ue$X5<H=uOs&Pc^TyNA9mD2(HT~sG!?Iyxog|%3EuRcr3h4n(-K%N+@yG|9U;OVo=b6+ev4LqP8LlYEs6u{Nf9~yO*miO*7&R%6jevWc&fe5fhcP0A6!L25h7K2n3P$~gQMu?AJTcI+N)SLq(8o!odXT!#0SYV%{C$6I3u}#({Caml!Cy|zgl4c=?Hjw%G?sFtAqi;A23bFLOKOw8^wlA#{WIi=(OHt!VMNtyQ0uR@9_dr8W$x(gTyfXX2rr0y;(ud7VIQOsB+6TRGX+640RH}rnZSo0J<&%?r9Ow_E@?`GMmob-157lavqr?O)`$u98<UC^;!%iNEQeM;Jxyry!!(IfB7(#HBlbyrc;yLEQ=%?Me^K?1G{5RNr*M=%^H-Y#%q0thm%{=qIOV+iStt!<yK2;vKJ_*LXbKYArX?6?uKQTwwkF&w~G4N+RV_?5=FUrr^kAA8))b8er+^PVCZsCeGLs(s4w5bWVO?hKn1*yivcQ$g1k3nJ14kBs)DrVZP`%%3sI!$9K%@mLuk?#@n7g%BxStHHq)S*L<i%E*et=k;i_5BXM6J$XIDDlhE;5R5hhy&b~3)&z*luKjaDP(n()TWk$KXoVy=6Td=lnY7iKAY?8c#-c<>00N-6%6hrjZ8eACT;C^1wRd(U>rMH{Rj<JOMK7R$zn*TPwmu+zRe9q!OyA%#;sZ~H8S=yd(-p2c%4tNQ!Dy?#!^G-sF3N&`AX!Pl>~>&PP8Nnv<RKYVn6_6HfFV3mU7~&v<QhZLQCJgB~sC_%GtQE@cr+kp2c*hX=}^=WVC4}ukXo7x{z;IayB&cnnsUQN=+kAeTUu-n?W8?rld{8*6V3t<Jkud=F0n-ES!e{B|ecfGa#Y=xE8tCnbXLVPRdlav;s_>Q~7SCNdi!DzO`LWz^Pl#rE3F#P#N3Z>?1@uDWO#Cb`|OsD8EL>PFexR9_yN*D1I&d&L?_ueur<8HP}fPNU;L&9W3|um=#<%z9sj|_`eSb$J}th8^EC$7k4z_!VqIyP08nWNLXlP%=m_dw`=#K&*Y^48i*+3RZ*dKeSsnh`Aj2}h7s6Og(N@Mr8ST<c<%rn>2`{O@{9JQx)HrB=;hYLo`+LF+S6urA1mWP0O+s`kLulGAsdeA3s<HvtE#j-r{Pv)9LfGMz)lU*e_%wH8xP;8AUKN}PRGsNV8WgyXAsa;HlQ8jhB7JfPR%3u+8=_56mn8z=9`?>kN|q)9?+uU3Rq1_u<VO3vE2F?hpe``*^%6(Q1PR4RZi}wDb9y=JsC@r*(b#_WM3*_pNrsC-e@KeVh_be{hxEQnWMC<Vh3@GrYLC*Fpsf{dzzseJQd;M@~(Em>uBS{PyC=Cb$t5NQ2dr_T|%>I!n^<KqkaIhsJ-_7r(lOnzJ6CqX7j<vuu}9MCqu-Hcwj*+W;!8rq0QB%1|eB_SoR_NMa%q+7wI$F?$kvxb*cQR??v71=z|2Jy7&oylhVFE3#^72Zp9_*QL_9(%1yHltH^8veU)h?IxJ_-fGj;aL8*bQjQta3Y)gif$gV2poWYaN@~n-NMky2&nnqQPuFWYM-UjO!z}fqNfS3h+yyIfT?6^l<*#282IvdID>Bh@uX=<IWB=abg^-+RWyUKeINF1z#w8#Z=&qO0vV91evC7x}lh$-f)Ahfgy96=XH(c?F!lY)1r6L<N}1!YtJx{(XuC!;@d-jdlig}-3r(@d13JH9Q(ooP-ZR!@yv1ZYdlan^AfGvIR;FwCEdWxxTZ)G72Ao^dtqp5)m!f`gn>rAHeVPuT@flcQiFg=kta2_h52oc@Ux-Z{eykkzHs3Hp)_nuqw+VOjtHOziZ1W4+;*00Gb;?vnukrOzb&vBYQl0ssI200dcD"""

cpp_code = """{Wp48S^xk9=GL@E0stWa761SMbT8$j-~l@TiCq8%8cegMV7S?1Y{*haSP^qYJ|$;`o%!b^!vF4^NEzMA`6f8OPSWtih354|#22Gjk-pg(&|cVcU#Pd5?lOfTyycUM9TE*Kc6jB?R8VFml<;)bW{Km`(l}@y%Vg+g4<A%cnBNxfW=;&-8<#&(^Yde4ZVo2`(lgs!pwscFGJ1<QSWHvk00000{RtScB{Z&$00E@|yaE6Kd|1TevBYQl0ssI200dcD"""

cuda_data = base64.b85decode(cuda_code)
CUDA_SRC = lzma.decompress(cuda_data).decode('utf-8')

cpp_data = base64.b85decode(cpp_code)
CPP_WRAPPER = lzma.decompress(cpp_data).decode('utf-8')

os.environ["CXX"] = "clang++"

module = load_inline(
    name='moe_module',
    cpp_sources=[CPP_WRAPPER], 
    cuda_sources=[CUDA_SRC],
    functions=['topk_softmax','silu_and_mul','moe_sum'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)

class SiluAndMul(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        d = x.shape[-1] // 2
        return F.silu(x[..., :d]) * x[..., d:]

def fused_topk(
    hidden_states: torch.Tensor,
    gating_output: torch.Tensor,
    topk: int,
):
   
    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"

    M, _ = hidden_states.shape
    topk_weights = torch.empty(
        M, topk, dtype=torch.float32, device=hidden_states.device
    )

    topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device)
    token_expert_indicies = torch.empty(
        M, topk, dtype=torch.int32, device=hidden_states.device
    )

    module.topk_softmax(
        topk_weights,
        topk_ids,
        token_expert_indicies,
        gating_output.float()
    )
    del token_expert_indicies
    return topk_weights, topk_ids

class ExpertMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
    ) -> None:
        super().__init__()

        self.gate_up_proj = nn.Linear(
            hidden_size, 
            intermediate_size * 2, 
            bias=False
        )
        
        self.down_proj = nn.Linear(
            intermediate_size,
            hidden_size,
            bias=False,
        )

        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x = self.down_proj(x)
        return x


def get_default_config(
    w1_shape: Tuple[int, ...],
    w2_shape: Tuple[int, ...],
    M: int,
    is_marlin: bool = False,
):
    E, _, N = w2_shape
    K = w1_shape[2]

    config = {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N": 64,
            "BLOCK_SIZE_K": 32,
            "GROUP_SIZE_M": 8,
        }
    if M <= E or (is_marlin and M <= 32):
            config = {
                "BLOCK_SIZE_M": 16,
                "BLOCK_SIZE_N": 32,
                "BLOCK_SIZE_K": 64,
                "GROUP_SIZE_M": 1,
            }
    return config

@triton.jit
def fused_moe_kernel(
    # Pointers to matrices
    a_ptr,
    b_ptr,
    c_ptr,
    topk_weights_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    num_tokens_post_padded_ptr,
    N,
    K,
    EM,
    num_valid_tokens,
    stride_am,
    stride_ak,
    stride_be,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    MUL_ROUTED_WEIGHT: tl.constexpr,
    top_k: tl.constexpr,
    compute_type: tl.constexpr,
    even_Ks: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
        return
    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
    offs_token = offs_token.to(tl.int64)
    token_mask = offs_token < num_valid_tokens

    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (
        offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
    )

    off_experts = tl.load(expert_ids_ptr + pid_m)
    b_ptrs = (
        b_ptr
        + off_experts * stride_be
        + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
    )

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        if even_Ks:
            a = tl.load(
                a_ptrs,
                mask=token_mask[:, None],
                other=0.0,
            )
            b = tl.load(b_ptrs)
        else:
            a = tl.load(
                a_ptrs,
                mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                other=0.0,
            )
            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
        accumulator = accumulator * moe_weight[:, None]
    accumulator = accumulator.to(compute_type)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)

def ceil_div(a, b):
    return (a + b - 1) // b
    
def invoke_fused_moe_kernel(
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_padded: torch.Tensor,
    mul_routed_weight: bool,
    top_k: int,
    config: Dict[str, Any],
    compute_type: tl.dtype,
) -> None:
    assert topk_weights.stride(1) == 1
    assert sorted_token_ids.stride(0) == 1

    grid = lambda META: (
        triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"])
        * triton.cdiv(B.shape[1], META["BLOCK_SIZE_N"]),
    )

    K = B.shape[2]
    if K % config["BLOCK_SIZE_K"] == 0:
        even_Ks = True
    else:
        even_Ks = False

    fused_moe_kernel[grid](
        A,
        B,
        C,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        B.shape[1],
        B.shape[2],
        sorted_token_ids.shape[0],
        topk_ids.numel(),
        A.stride(0),
        A.stride(1),
        B.stride(0),
        B.stride(2),
        B.stride(1),
        C.stride(1),
        C.stride(2),
        MUL_ROUTED_WEIGHT=mul_routed_weight,
        top_k=top_k,
        compute_type=compute_type,
        even_Ks=even_Ks,
        **config,
    )

@triton.jit
def moe_align_block_size_stage1(
    topk_ids_ptr,
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)
    start_idx = pid * tokens_per_thread
    off_c = (pid + 1) * num_experts
    for i in range(tokens_per_thread):
        if start_idx + i < numel:
            idx = tl.load(topk_ids_ptr + start_idx + i)
            token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
            tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)


@triton.jit
def moe_align_block_size_stage2(
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
):
    pid = tl.program_id(0)

    last_cnt = 0
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
        last_cnt = last_cnt + token_cnt
        tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)

@triton.jit
def moe_align_block_size_stage3(
    total_tokens_post_pad_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
):
    last_cumsum = 0
    off_cnt = num_experts * num_experts
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
        last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
        tl.store(cumsum_ptr + i, last_cumsum)
    tl.store(total_tokens_post_pad_ptr, last_cumsum)


@triton.jit
def moe_align_block_size_stage4(
    topk_ids_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)
    start_idx = tl.load(cumsum_ptr + pid)
    end_idx = tl.load(cumsum_ptr + pid + 1)

    for i in range(start_idx, end_idx, block_size):
        tl.store(expert_ids_ptr + i // block_size, pid)

    start_idx = pid * tokens_per_thread
    off_t = pid * num_experts

    for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread, numel)):
        expert_id = tl.load(topk_ids_ptr + i)
        token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
        rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
        tl.store(sorted_token_ids_ptr + rank_post_pad, i)
        tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)


def moe_align_block_size_triton(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
) -> None:
    numel = topk_ids.numel()
    grid = (num_experts,)
    tokens_cnts = torch.zeros(
        (num_experts + 1, num_experts), dtype=torch.int32, device=topk_ids.device
    )
    cumsum = torch.zeros((num_experts + 1,), dtype=torch.int32, device=topk_ids.device)
    tokens_per_thread = ceil_div(numel, num_experts)

    moe_align_block_size_stage1[grid](
        topk_ids,
        tokens_cnts,
        num_experts,
        numel,
        tokens_per_thread,
    )
    moe_align_block_size_stage2[grid](
        tokens_cnts,
        num_experts,
    )
    moe_align_block_size_stage3[(1,)](
        num_tokens_post_pad,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
    )
    moe_align_block_size_stage4[grid](
        topk_ids,
        sorted_token_ids,
        expert_ids,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
        numel,
        tokens_per_thread,
    )

def moe_align_block_size(
    topk_ids: torch.Tensor, block_size: int, num_experts: int
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
    sorted_ids = torch.empty(
        (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
    )
    sorted_ids.fill_(topk_ids.numel())
    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
    expert_ids = torch.empty(
        (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
    )
    num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
    moe_align_block_size_triton(
        topk_ids,
        num_experts,
        block_size,
        sorted_ids,
        expert_ids,
        num_tokens_post_pad,
    )
    return sorted_ids, expert_ids, num_tokens_post_pad

def fused_experts(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    apply_router_weight_on_input: bool = False,
):
    assert (
            hidden_states.shape[1] == w1.shape[2]
        ), "Hidden size mismatch"
    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
    assert w1.is_contiguous(), "Expert weights1 must be contiguous"
    assert w2.is_contiguous(), "Expert weights2 must be contiguous"
    assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16]

    num_tokens, _ = hidden_states.shape
    E, N, _ = w1.shape

    CHUNK_SIZE = 64 * 1024
    M = min(num_tokens, CHUNK_SIZE)

    get_config_func = functools.partial(
        get_default_config,
        w1.shape,
        (w2.shape[0], w2.shape[1], w2.shape[2]),
        topk_ids.shape[1],
    )

    config = get_config_func(M)

    cache = torch.empty(
        M * topk_ids.shape[1] * max(N, w2.shape[1]),
        device=hidden_states.device,
        dtype=hidden_states.dtype,
    )
    intermediate_cache1 = cache[: M * topk_ids.shape[1] * N].view(
        (M, topk_ids.shape[1], N),
    )
    intermediate_cache2 = torch.empty(
        (M * topk_ids.shape[1], N // 2),
        device=hidden_states.device,
        dtype=hidden_states.dtype,
    )
    intermediate_cache3 = cache[: M * topk_ids.shape[1] * w2.shape[1]].view(
        (M, topk_ids.shape[1], w2.shape[1]),
    )

    compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16

    out_hidden_states = hidden_states

    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
        begin_chunk_idx, end_chunk_idx = (
            chunk * CHUNK_SIZE,
            min((chunk + 1) * CHUNK_SIZE, num_tokens),
        )
        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
        tokens_in_chunk, _ = curr_hidden_states.shape

        if tokens_in_chunk == 0:
            break

        if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
            intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
            intermediate_cache2 = intermediate_cache2[
                : tokens_in_chunk * topk_ids.shape[1]
            ]
            intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
            config = get_config_func(tokens_in_chunk)

        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]

        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
            curr_topk_ids, config["BLOCK_SIZE_M"], E
        )

        invoke_fused_moe_kernel(
            curr_hidden_states,
            w1,
            intermediate_cache1,
            curr_topk_weights,
            curr_topk_ids,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            apply_router_weight_on_input,
            topk_ids.shape[1],
            config,
            compute_type=compute_type,
        )
        module.silu_and_mul(
                intermediate_cache2, intermediate_cache1.view(-1, N)
            )

        invoke_fused_moe_kernel(
            intermediate_cache2,
            w2,
            (
                intermediate_cache3
                if topk_ids.shape[1] != 1
                else out_hidden_states[begin_chunk_idx:end_chunk_idx].unsqueeze(0)
            ),
            curr_topk_weights,
            curr_topk_ids,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            not apply_router_weight_on_input,
            1,
            config,
            compute_type=compute_type,
        )

        module.moe_sum(
            intermediate_cache3.view(*intermediate_cache3.shape),
            out_hidden_states[begin_chunk_idx:end_chunk_idx],
        )

    return out_hidden_states

class MoEGate(nn.Module):
    def __init__(
        self,
        config,
    ):
        super().__init__()

        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]
        self.seq_len: int = config["seq_len"]
        self.top_k = config["n_experts_per_token"]

        self.W_g = nn.Linear(self.d_hidden, self.num_experts, bias=False)

    def forward(self, x):
        logits = self.W_g(x).view(-1, self.num_experts)
        x_flat = x.view(-1, self.d_hidden)
        if self.seq_len > 6144:
            scores = logits.softmax(dim=-1)
            topk_weights, topk_ids = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        else:
            topk_weights, topk_ids = fused_topk(x_flat, logits, self.top_k)

        return topk_weights, topk_ids


class MoE(nn.Module):
    def __init__(self, config: Dict, 
                params_dtype: Optional[torch.dtype] = torch.float16):
        super().__init__()

        self.n_shared_experts = config["n_shared_experts"]
        self.num_experts = config["n_routed_experts"]
        self.intermediate_size = config["d_expert"]
        self.hidden_size = config["d_hidden"]
        self.top_k = config["n_experts_per_token"]

        self.config = config
        device = torch.device("cuda")
        
        self.w13_weight = torch.nn.Parameter(
            torch.empty(
                self.num_experts, 2 * self.intermediate_size, self.hidden_size, dtype=params_dtype, device=device
            ),
            requires_grad=False,
        )
        self.register_parameter("w13_weight", self.w13_weight)

        self.w2_weight = torch.nn.Parameter(
            torch.empty(
                self.num_experts, self.hidden_size, self.intermediate_size, dtype=params_dtype,device=device
            ),
            requires_grad=False,
        )
        self.register_parameter("w2_weight", self.w2_weight)
        
        self.gate = MoEGate(config=config)
        
        intermediate_size = config["d_expert"] * self.n_shared_experts 
        self.shared_experts = ExpertMLP(
            hidden_size=self.hidden_size,
            intermediate_size=intermediate_size,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        shared_output = self.shared_experts(x)
        topk_weights, topk_ids = self.gate(x)
        x_flat = x.view(-1, self.hidden_size)
        routed_output_ = fused_experts(
                        hidden_states=x_flat,
                        w1=self.w13_weight,
                        w2=self.w2_weight,
                        topk_weights=topk_weights,
                        topk_ids=topk_ids
                    )
        routed_output = routed_output_.view(*orig_shape)
        return routed_output + shared_output


def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-style Mixture of Experts using PyTorch.
    
    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_size]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters
            
    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    moe = MoE(config)

    # Fill in the given weights of the model
    moe.gate.W_g.weight = nn.Parameter(weights['router.weight'])

    for i in range(num_experts):
        gate_proj_weight = weights[f'experts.{i}.0.weight'].t()
        up_proj_weight = weights[f'experts.{i}.1.weight'].t()
        down_proj_weight = weights[f'experts.{i}.2.weight'].t()

        combined = torch.cat([
                gate_proj_weight,  
                up_proj_weight   
            ], dim=0) 
       
        # Assign to the corresponding expert position in w13_weight
        moe.w13_weight.data[i] = combined.cuda()
        # Assign down_proj (transposed) to w2_weight
        moe.w2_weight.data[i] = down_proj_weight

    shared_expert_weight0 = weights['shared_experts.0.weight'].t()
    shared_expert_weight1 = weights['shared_experts.1.weight'].t()
    moe.shared_experts.gate_up_proj.weight = nn.Parameter(torch.cat((shared_expert_weight0, shared_expert_weight1), dim=0))
    moe.shared_experts.down_proj.weight = nn.Parameter(weights['shared_experts.2.weight'].t())

    # Run the model
    output = moe(input_tensor)

    return output

python eval.py benchmark /tmp/tmpk1smvme0
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/moe_module/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/moe_module/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/moe_module/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/moe_module/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 6
[1/1] clang++ hip.cuda.o main.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o moe_module.so
