submission 114034

#!/usr/bin/env python3
from torch.utils import cpp_extension
import os

sources = [
    "module.cpp",
    "test_kernel.cu",
    "minifloats.cu",
    "torch.cpp",
    "nvfp4_gemv/dispatch.cu",
]

def compile_module(sources=None, cuda_sources=None, verbose=False):
    if verbose:
        print("compile_module()")

    current_dir = os.path.dirname(os.path.abspath(__file__))

    extra_cflags = [
        "-fdiagnostics-color=always",
        "-O2",
        "-march=native",
        "-Wall",
    ]
    extra_cuda_cflags = [
        "-arch=sm_100a",
        "--generate-line-info",
        "--resource-usage",
        "--restrict",
    ]
    for flag in extra_cflags:
        extra_cuda_cflags.append('-Xcompiler')
        extra_cuda_cflags.append(flag)

    kwargs = dict(
        name="nvfp4_module",
        # extra_include_paths=include_dirs,
        extra_cflags=extra_cflags,
        extra_cuda_cflags=extra_cuda_cflags,
        # extra_ldflags=extra_ldflags,
        verbose=verbose,
    )
    
    if sources is not None:
        return cpp_extension.load(**kwargs,
            sources=[f"{current_dir}/{x}" for x in sources],
            build_directory=os.environ.get("BUILD_DIR", None),
        )
    elif cuda_sources is not None:
        return cpp_extension.load_inline(**kwargs,
            cpp_sources='',
            cuda_sources=cuda_sources,
        )
    else:
        raise ValueError("Either sources or cuda_sources must be provided.")


import base64, gzip
cuda_sources = """ABzY8#+xc>0{`tiYjfK;lHc_!Fq*Q5bQoEr6phxes%&g;TyY+k*jsz$@v-=TES4owOHz)NtoL8tAGlw3-FT27Me1S8Gk3djC6WkqH@X{*?gj{;vAauNKJl*sG`wg$j>6Wku?c#OpM%gFO?^rZys7Wl@ia_>F>MWx)S4vqgXjo=eL4t2O5T3_a$z5z{Ppb339;>Al%%!|KYV%@c$EDdM6pMMa0Jk{o#ACDwUE&2AOvXg=JNQp{pRh-n|B#@UZm0cNlcR@h(euI{{7a#&tRZA(=;MU>ZH?z91<0AUl#v#NKp1KQ_q7cUEqEvkND%L_&1ouL6{CSP#uk9V;xla*F>%eU`O#D`DH8lMceXq;-q9N+1v7e3!6qgE~5U0rYDH}B?VS)G%l@fENzzqK8j`AK0P}>v2EiqsOM)N02TUWRaAN(q}s{f&MxfJ*JtOKpHG(Ia@2>$&ouOdL8Fn<@nqzr^dOy0Xy}Y7xj1U5#_dI6lA_I!3lc}u(6^@(?SkAJ@I0Q5(qJ;8I{6M{W13FmkZ9UHd689-m&DY`t^&G~{oZdptjT~M+&!^*)+V?O_?Jwm7Yt@38AV@dk`iy|gg0O&(8`#GI6%lCh?5k-`RK%&4m5N}<AmDypn-*xdtg8|&kFCM7ao?Q83ey(W$c+j)+4~bzLn{GGe|}XX=R#rqi94x*Y1e=M_7}j@zhI6;zbkc+vxpbI|Zde6(Pf5a&kw*6iQ%N#Hj&a(PW=jg3e;zLpuaJPZT3Ru@m0zH~4drvQ{PugQZXuO`xPq+BPz>w!X*d5{|ENDmrIPV+7y<u{4GX4~VC6WRCk?k_ZNhjcGzJdR|Y++p|x?>q<61Q!k2rO(>+T>-!S=f&X|AlK-<Ly1t=yqe-#ZD&Q}(&>O~47<{7)qX0e_Mc_xr!DtjDKs59*IJ6hiI~h5Xg!&gCings{8meU!S`i_jx#jPH&n9UU-IA?d%N&sZ`R|1JZ|oA@RYrlE9g}tj6>5dAhy9{GtDd9Inyr=0`C=QNVrj+z3z365V+M>rstwv&&7viBqk=Fhe>JOSBsBKPo)BB(UuDZ)X14-MYPrT{E!Ak3m}H6KnpU~QCO2S_Ol>X4_wB?1>!F&8l<OfkvQ4C`;oV@;wQuP7Zr2ZziIaN6P35i)w0j`wQKc2=tp<_9AP6VZlzDTM|1JuA=~SKlxgvMIC>c0ag$m>j@%RI}=W+fs%zxhIKhN<eDpydH8qp-9tch*TH}pv30b_NQn9tEaVrJn@pJ#3X?Ud#fywvPfOgF;xXr=!3>3t{bV5aysI8+n<TU_nzLu6A}4}&{)U?6BZOCy3dkF0?yzAc@1f*{F>LUVynO&E~o{#<bziX{L<>k0@fgJ>f}TSAD0C9t4WLcRc2vPTmi+u5T)e8{aR3j<Pg_AO)yt)qEWSwgUXS(QvfOkTOcxTDCs&CJ8L!#lLXd&{QQ*d*>OrC0rHPGJhXbu565>CKgOozvUhfsgRc^T<ml_J)Va*f!g3hrIMXs4#ov{P-`elh$RQKb=~qpIWD2(<jz@vvv95)9W|y>_1}@ee(A7{c8(<n&j>4<74$_>$GDL&2vD@Fkxk}K!@s`-A2|`JV2;|4aP8Ir~cG~Os^m>@2+jfAeH5=o@r$0j0P6SoWL9bx&-Fq*h*rtayt&ScB?&rYzsbhuBe|Uv2!zaK+kbB8nGgn!xL^gEM&pg0qMinix+xkGqNeiV$DELcz3DjjWlAFvGdi^oEGbr;JRhokmK9-hrimWN`1dX$X!CnT}DU|dI*|?rY!@g8lpnLJv0=v6m2sdY53s#8i#J(OT5-hXQ&UE@mSnM*srtWPp$52uohu3<c!%%$7)@6ELL#({_>n?*&NQDf%2_agXPu`_>-quLpB6Fo_Vl^1g0N<(siI`aA>rzK}p#!Pp<(pGcZ#D6Kdr8^Kvx_nTqV$_Q(&tDp_{OB#kB4#rivE#G7)GF~nH#jR7P&&5_@_8AYx$YK30w&T36zw0++KN4I4@802aT27FVP#U5!M#!U?f2`M4bER>kndJZ~jYRF4bKEVW<&LZKwWYfZV$sZ&fgzUXNCyt%jL4jq%IuiY8EtPc<jg`8!t4a?4VJL5Pm<GC48UX3%l^_evLZMj&X<k34#=PX2HHNnGuc5K7)L8fN8tbm8v3`}t;C~nXzk>fg`2Q=@VE=hFs1#aNkX99BrwX$3q+St8Ro|xFaYwqfV&Qx@@ElF11{inD9`|o%FOTyojjVGXmiE!KPq2Hgw#;G=44JrQZ)GtT2?>Gej}sD{-7^Y_>2WTmwN%bB-YUL&3Yqc({xN0SkOb&QH!cBM+4D-!RVC=IE<tw<3BZw6N)SInf`{i=khq~CjRS9UD-x9}3VVB^ox;+>EsaAuveUiNB9bgGF4>T6zGRc>OVnzGVgnbHY&H}B&MI1CgG07%$u7(u<<t3tg5ef0!lf8qxdpQX^qpB8$yZM@u9W67I{8f4wz=`xHebCGs<_ppPQ<#EvKljjS;;-j4Oo~x^Y2?TWwoqacc0g#!<afgRPrCC+L05|td4Avj-@QT*}OQBl#hck;xYg*B371=Xn>P-I-+A5rU_t@+kC+e{B#n0!zNT!^EUXrOU#gpvtPvHB5@Y>vyy|1W?uH<Fy~b-F0Ej9Fgm>Xta~)3GDDIuxO(-%>I(K>0ax^=)~>~4@5mh+1~Z_=6v)tc0wdhU<%ol-l<J7vb<5GO1hdx#V(w#zxk${NCx$I>avEP3nWI|e9VJ!3u3-`zdkWn;ejV3Nrua^yW~qirV{jd1ec+)z2+SR#Ntt&fm)UBBVuYM|ThO)H!u|{rllJ|o_+txrng?Xs_vYk<*&gV+cz6j9`Ka*<M&lLDtz9&j3N_)XNTwN2zhdOB^YtXQoTN!H9<ufCcpBQ=c~q}}i($du0CPR!7}7eHV1-rO#DUMtOFs2nKuJG!aBz^p8VbFmqa&^8)@7Jeww#C^73Ux@+f|{-Ri0;weBLxkqw%u`KHdzEdxfRs_k8L11OIx*zs`-Khdo*p6J&oDI<gRUtOID}orq%RM|l+5FM)k1fw|RS?^&b+h%9jZ0bxjFW698gQ3*<6K6y;XR~<`uiC{o9fpP#Sx$9?OyyiMzfP0V>VPp%vYx>IJEPuyd3T*F`kcJz60?&4mgvP0s0d1+V^boy3Gpiee;SiVRTTOt047#b;g~}p3d<EGaI9>`N^iXbak&PHie5_h}F5*9TD1KZJUOW7X3G|v#n7ojj>XmU=%vje9$Be2wPjxOJ*D<{;F4wU=7n$q$UlyS2B&ZD0%KG^RLNts`>%U%_WDtXcb%+!C#5J(!PrL}wSs@_vi5fWKz<cfc+BS0p=v=O=;`e^FL&)80oqyrOk;}a}X;3{;>Ag#qbSQD1D}*N%)<8O^Vq&p4rC2Yi$d^fEGV!?B^^_E(?3zlM2=FkMeLxvTG6Q+Z`B;>c$ZY_tr6SDcN?^xqqpiB@<abD0tluebv3?{jKN6S9#3dh|;Fy5ZwbB<v*?bDa7<p<MgG5$OWw`85OJ`WPyke0KG|0Zq>nu;PHkGI-DTR2)uGhA(R#s6?E7nZQu3RCxTaX-=#W$8rxmA)^wy;*ZGoSCYmn1w=%yw0>WG+EArH;zt@eLL*(kz)W8KsP=kXqr^p29YZWYD^=0yHZCXFqbVA34~+c@8$AaS|ly^Cw}lp~q%n@{wl(HjuN=e`IDgnc4GZOoyecp(<CRf1iTe0heBxDAi^wzi(o>cAQ@n<=4({Dx&<N82^aurYgpN9R`rWa<2{wd|!?>YWDgg4gB{@hJ=q?KOI^-A*xA<G6re~IyXyZkjjwX&rj2?oHpg!sdFDf6J^l^oIzfyB-T&3)=Rk_nRH1fSmy?b2tjN%sK(d-b2(cvY89EqvjyEguC-ci_}*0)>37L_WztuquxbWql+%#^{~6&p@~0!(@+Ol_ALg0NTrg_W+?7A@_)ep6PH8fOzLV8plbOY&G)Z%_fm9CU=9TmYa!EXK7scLClel|(Tn?LTkI}e4{pX+0J{+5-{r=<e<@t$r@$vJUxAw{3E>1ptIs5p*{_y(!i9yCX&q+}YEuRjwW=s>BHVv{pi6cDaZ{Nhx6n!=n8yu3I*vy)B;v2wn!rP`^U;j0VN8Co5Iv;wnIUy>xSP~mjD|`|L+2m*?>Breg?M+yM(m2csHe(GyBr}NSFqpK4&8iC{%BlQ(_cJ#k56Y_5TfWYY$Oc>&3oM9{D7gXC>a!75c0aOQp|WtOGhh(5aQd=t_~PXp)uVfp-8l}zfFIJ?SZ}U%P;ahN`OzDUU9w*qg|k+8?9{oT`F`vmG#c4Co*fcR*$`x%88X%DjqkFcG>?K{kZ<aO46iha!UTh7!IF{)L!c&RWPm_PpSVHFCQ?Cog`rh~<mhHw&#V%Fws-<@zty(xxm98`-<gZud@tZMOoE$``s@@x1oCasjIL+1A%293e=a=5Bv9wiFM-ysL|VYeWbuO^RmI-G{1j+m$7%TiR<h$`jlnyQJ@9RW5Pqj)d*TFfqKv7E!BZM*T80K4l7mD4i2EnhGDcv|Vit|T>ODtXlqxF&g);$}->%KCECFWr1Y9mVa=0w6Z_qoV$N$~JwCcBWh$J1&RWiPI>3zy?11UNLUOAbi!|EEmlprZq1Q9<T4-SOE?Co6u<tQd&=N5X6{nl`k{3ONuH+&nvjYj=tF5CjxfsQ+huTC(D?SaJ~ZIfi^Ou&#9U<)3;c+SX)FMwU*Bh$Ec!E0_$(Z}|>i&(%hBd#(Ee0GTe8#nx1xbLJ63Q@HUl0l#WqL9qsc=&h&8cHeQ+91&+1<fbybr1xg-OnH2j759oYmg4bB`$#{3(t-lVx}okGjX`GKk{OYZ%A%2ku~-yOBf+dARQN;ZgEf;xxhMJw*yRn5m(6wCfJ0hEKwMRFQcy^UN9qsV*-f%XFeTK&<{9gnC1BqO#nU}0<aWE0vwX7Yv#&!@V=Mb9daB=P;fS1p%bI<my|$DKr2WHOzfg~Ms{}jiHiK993c`;dol5MHQ8H+v>nl;Ayvr-{bs!Det}HCmQ25nj9mTWmo{)y<qIK$uspt0J}YY2Pn#Pq_ScYlw|ojmK!H)nTp`r)Oq;tz1+f(*HhM@6PZXUoG||Cw-B6I=3x#!5xFxa|wN5GrUzZ)z9)$M2l@XUz(NpvZVu=)*-HZnF4h-;u0|1GLg9=T9n`oMda%og~`(RfpGj{>~ni>7&88oD~IFTt((UkEDae><lSRXbFlb>7a%l3y(!eTS4l3m6T9p4Ae^1GHG7qReaF9sR<xXc<;ysC_{m&}6h>8CEP0ZeHkWxybr6xeUEB290BQ{=3&g4Sf33~d+aXjY~$F*qTW$op+A50d$ig44`TA4z|ruqR4*#$2{Xvu6utvS#UclFl?HgE|&j56q22wH+~_(J&RlrSc3MMinHkoGW=sImcpSU*1Fa`Q$$?&pw}Mcx?*5EzXa?8~v88W&gyh^sJCzo*I4S!Ae}QXS0jlwB~2}_h}F_+aLoc7*T(ZG|A75!oIw*jD0`ImEN`v*#{mhclJd3z!iAKrUhL-Xs*$wxNTmI4RULc(KW_0cSRgi!AS;{ElH0v{Z?)ajSKx1pp$6=mg~+ys{<O-5F-R$K7#@~*qe;LyE=XIG4vf=*BMV!)b6Scb-?R%*{S~29_K$5t+-F58dgyCsz{l=l7>wLZ|v-rd_7K(^H+4*WVW}LPH}$=U!p_X=X0wfnj?4qGYCA6%CeLK(2-U#V%;L5xiTUyH5X8=HBd{fb{3$@1Y2!5>$D=7?s{b08Zxg|7Q`ze_G%Hu#MWs+ynq<_(;?}QfBb_uS=9WgCTb$oPfq4?^H?pX#r1Bf5Xzk_>~@Zc>M&$*7r9?km*<0#U$?}syBfdl)A@BD#V?pfiDOn$#j_mRj9~`#RF-i-l$q{}Miw*e&N1z-#I(B!rrouf7RxN^ILrm;39TzttEgWAuMvr@->%*j`Pi1?^;QrMFHTb>bn7K_*OSoQw1n=HBt-GNlG%BxAf<GWN#Q-+-n=<$1KQPwCe{z}g`vYzVvWDZW&0HJLR06P6iAlU-I6WbUx7PDGlcUTf$!mzz);9;q0;ys)2t;RnpL&+YLvE$Vir)$8WfYEm?9+LTo4yWRq-T>n~?<sJicmGRdcS-%T=mKt|F3KBa&Mb$yG#hYeaG@MUu7XQDw7GJ1(GRmAP;*;i!bDl8;rY5~_{u5T&{#Wuc~kXqFJoTErS@T~*q!gzPSr)?Fm6Tb9->NbA;0>z1W;3(~r^(k|{>ufAtJRZ_F6s!Bw!D4?f^=hX=36-DzD!Mqx=yh@?e?Qn(NX~M6y9OgNcIVUf6*|?o6=2<bzC+4eMzMpBn3iC*Bz8GEKSo}sm5sKgF{cgo?bZqC@`F>TJ;?0!uEhU>5d&&#B72eu1RfzYZ%33eJ8dV-SL#6ktY<;w;SJtt)P)(iB*M;g9bIZ_jL6c#$3omk&#RX(oc{E7#O%&eIIAzN+c@CRKPypaEU}lv1MN0ijN_`^2uR<ngC+{KoW;rjZmJO_;iK{UieP>a1r|ah5MY3y}if<*FV$Z3Gn_HWeRb5_HcXudTH(KHDRe6U`BJM=DQ+8oxc6$^=UzoffWqVNYl|_I<u6(gaeqTp^3y0iB3TYHWI^-0-|7!O7Ft@kxqu1$o;p-vaBqiTrf*o8kw#4eWA|HR7#$Fg8;YxzPh0SdH5qo2eL|EP>U&d8nvBzxd3*w2xD}A{aji1O`PRDxHtD<8mLW(|Kfv2|BE$+ur39P`(E7PjS@2)3ILoC(NQ9|dtw%h(wPoQ&3{ipS81wFN6D|&Sv!yvtGRbTwNq%W!{b-Vpld!h>6n(xSk>Qs6I_l80H8yhH+-df{CAk}HqD$}tZNyuD%m}Hcckt?v54Uj4(Rv#W2cE?)Or?R)i<x48H0+!gUSbr1t{{l?q=^<HX000"""
cuda_sources = gzip.decompress(base64.b85decode(cuda_sources)).decode("utf-8")
module = compile_module(cuda_sources=cuda_sources)

from task import input_t, output_t
import torch


'''
  - {"m": 7168, "k": 16384, "l":1, "seed": 1111}
a   ([7168, 8192, 1]) (8192, 1, 58720256)
b   ([128, 8192, 1]) (8192, 1, 1048576)
sfa ([7168, 1024, 1]) (1024, 1, 7340032)
sfb ([128, 1024, 1]) (1024, 1, 131072)
c   ([7168, 1, 1]) (1, 1, 7168)

mma_permute_order = (3, 4, 1, 5, 2, 0)
logical: mmc32, mmb4, mma, kkb4, kka, l
memory: l, mma, kka, mmc32, mmb4, kkb4
sfa_permuted    ([32, 4, 56, 4, 256, 1]) (16, 4, 131072, 1, 512, 7340032)
sfb_permuted    ([32, 4, 1, 4, 256, 1]) (16, 4, 131072, 1, 512, 131072)
'''

def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp8 gemv
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float4e2m1fn] of shape [m, k, l],
            b: torch.Tensor[float4e2m1fn] of shape [1, k, l],
            sfa: torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l], used by reference implementation
            sfb: torch.Tensor[float8_e4m3fnuz] of shape [1, k // 16, l], used by reference implementation
            sfa_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_m, 4, rest_k, l],
            sfb_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_n, 4, rest_k, l],
            c: torch.Tensor[float16] of shape [m, 1, l]
    Returns:
        Tensor containing output in float16
        c: torch.Tensor[float16] of shape [m, 1, l]
    """
    # c: [l, m, 1] is pre-allocated memory to avoid timing allocation overhead.
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data
    m, k_div_16, l = sfa.shape
    k = k_div_16 * 16
    
    a = module.fp4_to_fp16(a)
    b = module.fp4_to_fp16(b)

    # sfa2 = sfa_permuted.permute(2, 1, 0, 4, 3, 5).reshape(m, k_div_16, l)
    # sfb2 = sfb_permuted[0, 0, 0, :, :, :].permute(1, 0, 2).reshape(k_div_16, l)
    # torch.testing.assert_close(sfa, sfa2, rtol=0, atol=0)
    # torch.testing.assert_close(sfb[0,:,:], sfb2, rtol=0, atol=0)

    sfa = sfa.to(torch.float16)
    sfb = sfb.to(torch.float16)
    
    for l_idx in range(l):
        a_mat = a[:, :, l_idx, :].reshape(m, k_div_16, 16)
        b_col = b[0, :, l_idx, :].reshape(k_div_16, 16)
        sfa_mat = sfa[:, :, l_idx].reshape(m, k_div_16, 1)
        sfb_col = sfb[0, :, l_idx].reshape(k_div_16, 1)
        
        a_mat *= sfa_mat
        b_col *= sfb_col
        c[:, 0, l_idx] = torch.matmul(a_mat.reshape(m, k), b_col.reshape(k))

    return c

custom_kernel = module.nvfp4_gemv