From d4bf04139fa65cef33af8937c09538cf0922d0b9 Mon Sep 17 00:00:00 2001 From: Steve Markgraf Date: Thu, 30 Jul 2026 01:03:20 +0200 Subject: [PATCH 1/2] vulkan: fmt=7 MXFP4 decode (e2m1 nibbles + per-32 group scales) for Kimi K3 experts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New fmt=7 branch in qmatmul.comp: e2m1 LUT nibble decode (bit3 = sign, low nibble = even column — same packing as the K3 checkpoint), one f32 scale per 32-input group reusing the #298 grouped-scale path; the host pre-expands the ue8m0 exponents to f32 at upload (mx4_scale), so the shader stays float-only. Host gates in tensor upload/scale sizing accept fmt=7 with word-aligned groups. Validated against quant.h matmul_mxfp4 on K3 expert dims (I=3584, O=3072, S=2, random nibbles + exponents 2^-7..2^4): rel_l2 2.6e-07 (tests/test_vk_mxfp4.c, skips without a Vulkan device). (cherry picked from commit ca5a1b254d8714361910f7d5eabaa072a0bdf11f) --- c/backend_vulkan.c | 11 +++++----- c/shaders/qmatmul.comp | 27 +++++++++++++++++++------ c/shaders/qmatmul.spv | Bin 0 -> 19120 bytes c/tests/test_vk_mxfp4.c | 43 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 70 insertions(+), 11 deletions(-) create mode 100644 c/shaders/qmatmul.spv create mode 100644 c/tests/test_vk_mxfp4.c diff --git a/c/backend_vulkan.c b/c/backend_vulkan.c index 9d22d2df4..d779607d8 100644 --- a/c/backend_vulkan.c +++ b/c/backend_vulkan.c @@ -211,7 +211,8 @@ static int rowwords(int fmt, int I) { * (one f32 per 64-input group). upload_tensor and tensor_free must agree on this. */ static size_t scale_floats(int fmt, int I, int O, int gs) { if (fmt == 5) return (size_t)O * (((size_t)I + 63) / 64); - if (fmt == 4) return (size_t)O * (((size_t)I + gs - 1) / gs); // per-group [O,ng] + if (fmt == 4 || fmt == 7) + return (size_t)O * (((size_t)I + gs - 1) / gs); // per-group [O,ng] return (size_t)O; } @@ -506,11 +507,11 @@ static int arena_suballoc(size_t bytes, VkBuffer *buf, void **ptr) { static int upload_tensor(ColiVkTensor **out, const void *weights, const float *scales, int fmt, int I, int O, int gs) { if (*out) return (*out)->fmt == fmt && (*out)->I == I && (*out)->O == O; - if (fmt != 1 && fmt != 2 && fmt != 5 && - !(fmt == 4 && gs >= 8 && gs % 8 == 0)) return 0; /* fmt=4: word-aligned groups only */ + if (fmt != 1 && fmt != 2 && fmt != 5 && /* fmt=4/7: word-aligned groups only */ + !((fmt == 4 || fmt == 7) && gs >= 8 && gs % 8 == 0)) return 0; ColiVkTensor *t = calloc(1, sizeof(*t)); if (!t) return 0; - t->fmt = fmt; t->I = I; t->O = O; t->rowWords = rowwords(fmt, I); t->gs = fmt == 4 ? gs : 0; + t->fmt = fmt; t->I = I; t->O = O; t->rowWords = rowwords(fmt, I); t->gs = (fmt == 4 || fmt == 7) ? gs : 0; size_t stride = (size_t)t->rowWords * 4; // padded row bytes size_t cpu_rb = fmt == 1 ? (size_t)I : fmt == 5 ? ((size_t)I + 63) / 64 * 24 : (size_t)(I + 1) / 2; @@ -923,7 +924,7 @@ static int upload_tensor_d2(ColiVkTensor **out, const void *weights, const float !(fmt == 4 && gs >= 8 && gs % 8 == 0)) return 0; ColiVkTensor *t = calloc(1, sizeof(*t)); if (!t) return 0; - t->fmt = fmt; t->I = I; t->O = O; t->rowWords = rowwords(fmt, I); t->gs = fmt == 4 ? gs : 0; + t->fmt = fmt; t->I = I; t->O = O; t->rowWords = rowwords(fmt, I); t->gs = (fmt == 4 || fmt == 7) ? gs : 0; t->dev = 1; size_t stride = (size_t)t->rowWords * 4; size_t cpu_rb = fmt == 1 ? (size_t)I diff --git a/c/shaders/qmatmul.comp b/c/shaders/qmatmul.comp index a42205856..25de4f201 100644 --- a/c/shaders/qmatmul.comp +++ b/c/shaders/qmatmul.comp @@ -46,6 +46,15 @@ float i8(uint word, int lane) { float i4(uint word, int lane) { return float(int((word >> (uint(lane) * 4u)) & 0xfu) - 8); } +// fmt=7 MXFP4 (Kimi K3 experts): e2m1 nibble, bit3 = sign, LOW nibble = even +// column (same packing order as i4). Scales arrive as f32 (host pre-expands +// the ue8m0 exponents), one per 32-input group. +float mx4(uint word, int lane) { + uint n = (word >> (uint(lane) * 4u)) & 0xfu; + const float lut[8] = float[8](0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0); + float v = lut[n & 7u]; + return (n & 8u) != 0u ? -v : v; +} void main() { int s = int(gl_WorkGroupID.y); @@ -97,16 +106,22 @@ void main() { } sum += a * scale[sb + uint(g)]; } - } else if (p.fmt == 4) { // grouped int4 (#298 semantics): nibble decode, one scale - // per gs inputs — gs % 8 == 0 (host-gated), so a packed word never straddles - // a group and its scale multiplies the 8-wide partial. Per-row scale must NOT run. + } else if (p.fmt == 4 || p.fmt == 7) { // grouped int4 (#298 semantics) or MXFP4: + // nibble decode, one scale per gs inputs — gs % 8 == 0 (host-gated), so a + // packed word never straddles a group and its scale multiplies the 8-wide + // partial. Per-row scale must NOT run. int words = (p.I + 7) / 8, ng = (p.I + p.gs - 1) / p.gs; uint sb = uint(o) * uint(ng); for (int wi = lane; wi < words; wi += sgsize) { uint pk = w[rowBase + uint(wi)]; int i0 = wi * 8; float a = 0.0; - for (int k = 0; k < 8; k++) { int i = i0 + k; - if (i < p.I) a += (staged ? xsh[i] : x[xoff + i]) * i4(pk, k); } + if (p.fmt == 4) { + for (int k = 0; k < 8; k++) { int i = i0 + k; + if (i < p.I) a += (staged ? xsh[i] : x[xoff + i]) * i4(pk, k); } + } else { + for (int k = 0; k < 8; k++) { int i = i0 + k; + if (i < p.I) a += (staged ? xsh[i] : x[xoff + i]) * mx4(pk, k); } + } sum += a * scale[sb + uint(i0 / p.gs)]; } } else { // int4 @@ -120,6 +135,6 @@ void main() { } } float tot = subgroupAdd(sum); - if (lane == 0) y[s * p.O + o] = (p.fmt == 5 || p.fmt == 4) ? tot : tot * scale[o]; + if (lane == 0) y[s * p.O + o] = (p.fmt == 5 || p.fmt == 4 || p.fmt == 7) ? tot : tot * scale[o]; } } diff --git a/c/shaders/qmatmul.spv b/c/shaders/qmatmul.spv new file mode 100644 index 0000000000000000000000000000000000000000..f64886140786470192e2d66f2b24252b7fa0de4c GIT binary patch literal 19120 zcmZ{r37DSK`NqFwCJV8|7D13m5c^J0ghoT9w%Ci9$TlQ1X=Xww4JlQ0p{=DGQ6<*e zmnyY{s%o{`Ds8nLrLC4OsI^6ZzjxkyoP5{+|GLig-1mK+^PK0b-}@%f(tq@z#(;i} zA&tR}y;>U8vr40XqaR4STW$Bi4rvUm^of%uPujApd)k)U@36HER%-M%w9U%28h~x3 z%xP=yz%*u2Iw=b%7f>#u+(@~FatGx;$~%-739uHShSHDiSE1OUaV;(U*B|>3q_wep zW8$PK`|We^l&)EA)27dx(lvkT%z2&j=T2#x*WNvA&h+l~OI{T*$Yz z>M!{YKv{{l?c>JI-)7hLZFU9gvjXP1SCslYr;}C3n$F%7S+NRDnLF=_$ z6&_fv$;{bP4)2_I;zVa{-wC}w#_DdLhX3Gd&G4Djd3@XKeLGIVbhmeQ=-KOQtd8zk z__j9YgS)!hW=>~z?8`h3q_~n|pBZzy`?%{lnUZ?fb>D)!qYkL}z^Z*KcwXnpE!7k{WF{78aG;RQN z&9rV{wRf49$@PY-Cd4*3VT|ipDEuc|VP1EmchI?Q z?zK(#y!L5Tn?cptq0RpD=kzUQH9xjnim&I+c28lMzrEVJdSm|1IP>RJbuWP3qc*hN z%W@`9{04aLiM{wO@Zl8aPV93sML*B4kGMaJ+rD)#=4xrsywX=xuW)N)ZSdT-d2MrQ zK4KZ((ijU*4%=E+=dd$4u|EV(>@?QI85f<4ergS zrLho*aPqxXpIW~c?D?_n zpTU#4xCX++K3$wY&#v{}kv^_z3x}p$x54m?YgKT@HC(@XJvRVnE;j)u&&|PYRos@w zw&3Ksi#~OK2ZPIbJ!Bc)(m1AQH@#?g`Z9j4jfGA8Qg}Iz#qv7e>-4SH{uXfZy&aq! z?^UmJd`$mh?H^x;w=|xHC-!sJ)v;d$C-zI=#C}b^j@_S2sf^va3~yLML)K;FL~M5u$E=}a-XKJ?Af(w>$t+pzTC&w8(VB&m_AzMss5|c zi~WqkTYvEy#BkgrDLy0NgRmBL>)pR2Dg7$zn_X|c`Y<9bU+KmjQ1vCZy?&e3@kfC5 z4__bre$CZKQ_i5QPFC&{{nW=&>^SYVg?a|UcL1+kbN$Crw2or});!E-1pFg?xVOFf zM2e5`pv6AG*$Y2RLR{ZFW| zG2-vNXbkJSvF74ht;dXeFFNo4SJB`1759UF#*;6oHT$>SDHQjl{-@S;>Zv~+>>*3O zpRc)V`vtJ;r)K*vV%;a|_9yp!RBc^m#rstM^}ybja@&0c>)otoJGu9xTJpFU?4$oT zvDO(&z8GsC^5l6H*l~ql4IbRDH~tCW8!6DLKe=~o$-PU%UB~-so_wB!yUyXyfXyMi z2W)%CojcWWzlA0@@3*mzL(Or^y))IU_pFND4~cs`LH(TyKLPIdYWQq;{I%y6+;zPI z?r)y{*86?wILy=g*f}&$?_0Tf`~__}xw1@0v2dr3F`is^AmAe!rCdFEsI&3$EWA1$R7tqm=CjG;zOG zQm^%UCEWIYuataj!PWg{DeL`iDY@S+CHMO!TvvAL^?vJw+n?V%;nw@jQ*yt1!u8*~;Oc%0rQY^_3zgh&p_2PO6mGn^ z1y}bQDfPzl8>!@eBb9t{!PWgvD(n4LD!JcFCHI>tT^T%4>pzRNUYl6wf{l?eod<3n)A?xXx71_$9N2d1<|;p%Vt>XLC;o+C#OY{zpJp8t0{?p z4Y-`MYvIjvb{(2}=InZ~?bOX-30BRR1F_$s_}-U2e^agLdjnV=-&?`AP~v+VSWRDl z@8$8m6MP3HzITDu^z~gpo^gE_T+ZV?aDC3H*YtZ}ebjA#H&)GcI|2KBitnOFQquqZ zwVwVTfUBj?2f@b4JUs+{fTC_6zK^J-kB7mtC}kgyz}3^o55Q{a<56(g#}DE9sK@6q zu=AY#`y;S^>gKT&tL7eg9P1v`dai$heS+d+{ZFyhYZL!xU}NN*{T%F`d6E*JUx3wf zu73%(ox0~*{$q;$8DE^XzXBWIIdDHcO;JmKzXsdiGnBOX4Ol&GehXG}EsXIj#W?ED z!Be#+_PuW+CGmd`HvYww#Qy_W{W{`}!9GV(>!EnRK9BWrY}Wmea$LXOv56D=Phex8 zLP_jDgVnQMe*vpGpV_Z3g4G?fIlq8aGlu#6m6E&uC9sdVsQ-pQCFXu--h+-^bi_>-`u(^b< z3^rEw&rq;F>gMWqhg$qs0jp(Ct_p76ldGYrC%@Ifwo`Y%55xMHt95Ho)XY_!zSab% zuf_NehhKnI*JpX~T3~g#>$witwJ~?Uk>rWDF4(bJ=l7C)XIdF&J@86c+gNut`$K(w zuzgscK1YH*_h;5Vqu^@#8rwQG*DvGJ%3WK>sZCxRfXg))4Oiui1YWf-9IyK|v{Cc+J-Ur6_4oLh>z&E4Y*6&QYzMEq8)w-8A z$8JIKv3|?CUYl53fsGNqHMl(c+rafvcfZItqZr%1#I|+qw*!}Fe|xyGbM|)tt2qz) z?+CVydd~h%;GBKy6MG!kI>+p{v^+WQ0=A!={awLoIs1NVtJ&7??414Gz>Y!deC>f9 zPw{cS_N?o*>1!{rF>Hn&n?zAFS8@8wD>!Y6c>Vv=+ z(Z>0;-nO3MgTdCxUH3!4uAyU`jFqRK!@$mob%$c*xg!n-dq>#Dy4(>*fbGNj^m!z> zyd#c+tLbZO>(pH5j9)8v4>%v%p%H#GVCi-iPgI>KV_8V6_t{OF938i8}|ZKAYkgXJXZoUnjVHmz)b% zOZ_Lnw#of653H8=#xAgJ)Q#VPRZIK%;M}9#V7c>m64o(ly%!c>Poel&e`;N?O{`CX zjS+qt*xYg_oeoy>p2*q!6x?>!ImhyoDfVao;>0@xZ2ZjKnP4^FMIGBhuv+{-4c7nB zdalm`tNC84|01xO{+`LtfXlo49Jnzvx1R;8Wp2*{tDQ@6jg4_W*f`cDkI#YaD{(## zR!f^NfNiIq-+(WI^BZ7&a=ZX+o%x@Qm8b4Y;4=R&!^`}?0#{4^7lGBv{4a(Z$GYTq z30Qqw#+NZ&3U^Fdi?4#!dhp3OzXn%NoUen8rJiwq1DtVMpRs%sY@Or05GzmLmxJv) z>$(`MX5ZOs--4@OK{3W>R@wpmoKk-=tR&)P4CeQCRaNAmE zKUZSa#Mfdc;v>EeoZsB*!R~FXXX{4ncPKudt()q4ZF0F8Y>b?nTfl0L&vPTcfnqFU zi_`Wtuz7{w4mMWCa0gf)b#uKHs}}z|!D_kd?gBU8b$6qwC%^B4ZKv)%`aP_VxmtG* zMa^8r_T~Ax5A429-`@wDYu5FCu$pH+`|JU@`a|^@d=P9Lb@zyImx9$FqNJ~f!S1c} z^$6H~68-~lIZuzm^-)jXKLo3%@5jK#t@OI@AHmhz$-}*5K99rI9mgZJ_G7qtTIV|4 zi&Zn9G5GxiJRj?P{tWEne5(JH@=J>IDb5(41ZV9Qx7XTW(UvOfL(8f;zqm6!e6e|&!fHaEHR_gk>@ z<2rZ_p9QNeVlL&s1FLyfJL_@&9kGR$8!qS`Owl zdK+Hp^!XpKk7rB$9g3QLi__PC!OeZWi>98w-Uq9_M{x{^qkX+GsOlTJ>f3SXZy%oV z{^-d=pMI1}^o59%M=RJo%0H7EfTo^#9SBx4kNlb3Ah`NC+B(ZSk z_vO{$YN=lXY@3{yVPG|XU%Xe=1lva4`2L-$TH3D#&ik@|zbbcK*T%XoT7T0=VArMi zSifFfuT8A=!Nv$533d(g8#)TCkGgwNz7ECM_9eEhd5;E{?|x(8#?Cl51gp6w`fmib zje6e2HU{Ti%=*M03%1Vr@b8f2$$3+-{p9y~Gq76rlz+#pW?TP$+4rHW_ZD#1yF9B~ zqN(TnYz0<(20zclHees;!@8|0M^l^+vHke>)_K?3uCCL%E<0d%r1-clJJt2tjA-SO>>l{?q=VLam|-Xw7OKJhWQ@iSlhgY{9*@7Dp~yiZu4eh&m&mwx4C zzxE&BgTUq{_dHAnySHt7Fjk%%4+Xo1)*XVC=g(sf1FuSZ+gO+1&cng>VSV~M0_+($ zvwpWY60D}Lv8_`x=d71j?w)n-w8`rzaJjcW4p-|TALCi4X8W82_pID=WIWG~zQ>S#vXJI+P8z%a%N5dH=mgk(bSXQY_RRrJu@9xANRj?b0})&Dz-1rOeZ*JW**qbHtL_C zoJ6sWIQeyfliy<=*-XT&;(EjAxyi?Q;%}t9j0h zHu;?cHV3(L@>%ezHTQk@T(EmPd-FW7d$WAUIUh|u?>L_W+fF^d-=7DoyTZ|qM{5JL~ijVbI*Y(=OS^_pk z_%&d2%sq82SRZxsl3z(Nwtb0h>z=qC{1i5^Zh#v*d+0{6TK3R)z_wA(JK#;=yaQUF z*f)c%a}KY=%G1xSVEf5C;B8>F%=s-~HQU~f&F{<|V8@_!zV5=_P4RKQzFXI8)7L#< zV`P552UfE$=SO}g#aPA`r|o@Ua|!=G*jU*=_k;COH`jZyYVm&ntd>3bAh>x?K7^*8 z{FZ`kr|$2}BUm4EweDexnz@S8*AKwyYcc+h!u`!q*T>(CAA;58uK$m~u8p}rhLtDY zkHL=By2r8d+{Zrw`wn9p>vB�k#k8)8|jY2r zat(eCSL-2;@vKv`efEoME58bBJZ<9r0$kohzl1NQWR0H!tLbNa>(q>w^Xu7?dmk9z zJ0S6&0q5*L4VLTsYplLn_wsMC&r*D>|6N_LO|0L8jS>C_@THXOogT10>h2f$Zz#sL zFR^W1`#*xqv;REY*g5+zfYoyL{{*&;dd~iz!8!ZZC-z^!);V9#VdcsBuVDMh+5a0@ zEoc8lu$paO!shJ140a4!=j$KXe^Pv$uYc9`+Vu4b*ch3gSHWsI`|`h2jAd+b+Ws4C zF5$0(jg@`z23Q|;bA1h~7XLTFYB~FFft%0%+i2>^?;Wu1)IIzE#rn9G*8PW~X0GD& z^)9%4zk3gE-^ulTus-Vf?I(QS`dH@i}&}Uo94ixhcJMZ2L+k@Rx?k)d60=aWwdwp#a-yOmFX3cg2FC|9qo^kMe zcUYHpJAHv&&>Js-{klVq3lGvLD&^) zyc~8#ifvcM=KnVv4%W|Hyqnje*e1EH4Ypl=!`A_;c@Fhom!dX;lD~gk4{RLk@^_yj z!RqT%jOF**D9T8RV>AzW`rQDmb`$-K!EQ)V8%=SKZH)DCuUoefMa}Vv6K4~!aq?&W zn}XHGQqr%QbDQ?swBHPD`~3ag=3uq=@E4c!?{AY&IsZf9u7&fz3dQ_Z!y4cD+lr!3 z`r8_;U*>xouv+GOJFwcel+5?`VB=Vq`Q8z%z5^xmy%ShHbFnj6Z5$=z+Xbxd_{8P- zhQWQx@vRQe_|^bBzBRG#uR++|C|dah?B0ERK=uD9A5!xJ3ZIYHTIOhXxcy{~_5j;g z=4d=vEpxOdSgo9+z2L^NE_3u@u)1^PUNBd+^m|~1Kkz>ou67W`yVUU5)3yz4+xW^$U)#iYD%g1L5wUyZf0t0fUH||9 literal 0 HcmV?d00001 diff --git a/c/tests/test_vk_mxfp4.c b/c/tests/test_vk_mxfp4.c new file mode 100644 index 000000000..e5faf176e --- /dev/null +++ b/c/tests/test_vk_mxfp4.c @@ -0,0 +1,43 @@ +/* fmt=7 (MXFP4) Vulkan matmul vs the CPU reference kernel (quant.h). + * Random e2m1 nibbles + ue8m0 group exponents, K3 expert dims. + * Skips (exit 0) when no Vulkan device is available. */ +#include +#include +#include +#include +#include "../quant.h" +#include "../backend_vulkan.h" + +int main(void){ + if(!coli_vk_init("shaders/qmatmul.spv")||!coli_vk_available()){ + fprintf(stderr,"vk-mxfp4: no Vulkan device — skipped\n"); return 0; + } + srand(7); + int S=2, I=3584, O=3072; /* K3 expert w1/w3 shape */ + int rb=(I+1)/2, ng=(I+31)/32; + uint8_t *q4=malloc((size_t)O*rb), *e8=malloc((size_t)O*ng); + float *x=malloc((size_t)S*I*sizeof(float)); + float *yc=calloc((size_t)S*O,sizeof(float)), *yg=calloc((size_t)S*O,sizeof(float)); + float *sc=malloc((size_t)O*ng*sizeof(float)); + for(size_t i=0;i<(size_t)O*rb;i++) q4[i]=(uint8_t)rand(); + for(size_t i=0;i<(size_t)O*ng;i++){ e8[i]=(uint8_t)(120+rand()%12); sc[i]=mx4_scale(e8[i]); } + for(int i=0;imx)mx=r; + } + double rel=sqrt(num/(den+1e-30)); + printf("vk-mxfp4: rel_l2 %.3e max_rel %.3e (S=%d I=%d O=%d)\n",rel,mx,S,I,O); + if(rel>1e-5){ fprintf(stderr,"vk-mxfp4: FAIL rel_l2 %.3e\n",rel); return 1; } + printf("vk-mxfp4: OK\n"); + coli_vk_shutdown(); + return 0; +} From b61089a812f9bcaf3b86eb65a22a244f654e60cc Mon Sep 17 00:00:00 2001 From: Steve Markgraf Date: Thu, 30 Jul 2026 01:12:43 +0200 Subject: [PATCH 2/2] =?UTF-8?q?kimi=5Fk3:=20Vulkan=20tier=20(K3=5FVK)=20?= =?UTF-8?q?=E2=80=94=20shared=20experts=20resident=20+=20fill-once=20MXFP4?= =?UTF-8?q?=20routed-expert=20tier?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Build with `make VK=1 kimi_k3` (backend_vulkan.o + shaders; default CPU build untouched). Two residency classes on the card, both with transparent CPU fallback and identical output: - dense W tensors: shared experts uploaded at init (they run every token and are the largest always-on RAM-bandwidth slice that fits VRAM); any uploaded W computes on the existing fmt=1/4 shaders at S==1 via the w_matmul hook. - routed experts: a fill-once tier in fmt=7 (native MXFP4 bytes straight from the RAM slot, ue8m0 exponents expanded to f32 group scales at upload; QAT weights never re-encoded). Experts enter from freshly-read slots on the main thread, K3_VK_UP per step (default 8), until the VRAM budget (K3_VK_GB, default driver budget) is full. At decode (C==1) tier-resident experts drop out of the disk union — skipping BOTH the 17.5 MB read and the CPU matmuls — and run as one paired w1/w3 submit + SiTU-GLU on CPU + w2 down. Chunked prefill stays on the CPU-batched path (and still warms the tier through the upload hook). K3_VK=0 disables at runtime; [K3-VK] stderr lines report residency and GPU hits (also after every serve-mode request). (cherry picked from commit 3bcdfd0974eef2e1953a6a849b6c785c31e93d79) --- c/Makefile | 4 +- c/kimi_k3.c | 178 +++++++++++++++++++++++++++++++++++++++++++++++- docs/kimi_k3.md | 30 +++++++- 3 files changed, 207 insertions(+), 5 deletions(-) diff --git a/c/Makefile b/c/Makefile index 0c4caac29..2f4819434 100644 --- a/c/Makefile +++ b/c/Makefile @@ -439,8 +439,8 @@ olmoe$(EXE): olmoe.c st.h json.h compat.h inkling$(EXE): inkling.c st.h json.h compat.h $(INK_CUDA_OBJ) $(CC) $(CFLAGS) inkling.c $(INK_CUDA_OBJ) -o inkling$(EXE) $(LDFLAGS) -kimi_k3$(EXE): kimi_k3.c st.h json.h tok.h tok_unicode.h tok_unicode_o200k.h compat.h quant.h - $(CC) $(CFLAGS) kimi_k3.c -o kimi_k3$(EXE) $(LDFLAGS) +kimi_k3$(EXE): kimi_k3.c st.h json.h tok.h tok_unicode.h tok_unicode_o200k.h compat.h quant.h $(VK_OBJ) $(VK_SPV) + $(CC) $(CFLAGS) kimi_k3.c $(VK_OBJ) -o kimi_k3$(EXE) $(LDFLAGS) # Use a baseline that matches the compiler target. macOS already targets a # portable baseline when ARCH is empty; forcing the x86 value there breaks diff --git a/c/kimi_k3.c b/c/kimi_k3.c index 4ea1e2882..24df36ef1 100644 --- a/c/kimi_k3.c +++ b/c/kimi_k3.c @@ -40,6 +40,13 @@ * K3_MLA_BITS=8|4|32 MLA projections (default 8) * K3_HEAD_BITS=8|4|32 lm_head (default 8) * K3_EXPERT_GB=N routed-expert LRU cache budget (default 8) + * K3_VK=0|1 Vulkan tier (build with `make VK=1 kimi_k3`; default + * 1 when built): shared experts VRAM-resident + a + * fill-once native-MXFP4 routed-expert tier; resident + * experts skip disk AND CPU at decode. CPU fallback + * everywhere; output identical. + * K3_VK_GB=N VRAM cap for the tier (default: driver budget) + * K3_VK_UP=N routed-expert uploads per step (default 8) * K3_DIRECT=0|1 O_DIRECT expert reads (default 1; buffered fallback) * K3_IDOT=0|1 int8-activation expert matmuls (default 1; 0 = float) * K3_PIPE=0|1 overlap expert loads with compute (default 1) @@ -94,7 +101,8 @@ typedef struct { } Cfg; /* ---------- RAM-resident weight, quantized at load ---------- */ -typedef struct { int fmt; float *f; int8_t *q8; uint8_t *q4; float *s; int O, I, gs; } W; +typedef struct { int fmt; float *f; int8_t *q8; uint8_t *q4; float *s; int O, I, gs; + void *vk; /* ColiVkTensor* once uploaded (K3_VK); NULL = CPU only */ } W; typedef struct { /* KDA layer */ W q, k, v, o, g; @@ -174,7 +182,62 @@ static void rmsnorm_(float *out, const float *x, const float *w, int D, float ep } /* ---------- W: load-time quantization + matvec ---------- */ +/* ---------- Vulkan tier (K3_VK, build with `make VK=1 kimi_k3`) ---------- + * Two residency classes on the card, both with transparent CPU fallback: + * - dense W tensors uploaded once at init (shared experts first): computed + * by the existing fmt=1/4 shaders whenever S==1 (w_matmul hook below); + * - a fill-once routed-expert tier in fmt=7 (native MXFP4, never re-encoded): + * experts enter from freshly-read RAM slots (K3_VK_UP per step) until the + * VRAM budget (K3_VK_GB) is reached; resident experts then skip BOTH the + * disk read and the CPU matmuls at decode (C==1). */ +static int g_k3_vk=0; /* backend live (K3_VK=0 disables) */ +#ifdef COLI_VULKAN +#include "backend_vulkan.h" +typedef struct { void *w1, *w2, *w3; } VkExp; /* ColiVkTensor* triple */ +static VkExp *g_vkexp; static int64_t g_vkexp_n; +static int g_vk_upcap=8, g_vk_up_left=0, g_vk_full=0; +static long g_vk_hit=0, g_vk_res=0; +static double g_vk_gb=0; /* K3_VK_GB cap (0 = driver budget) */ +static const char *k3_vk_spv(char *buf, size_t n){ + const char *env=getenv("COLI_VK_SHADERS"); + struct stat st; + if(env&&*env){ + if(!stat(env,&st)&&S_ISDIR(st.st_mode)){ snprintf(buf,n,"%s/qmatmul.spv",env); return buf; } + return env; + } +#ifdef __linux__ + ssize_t k=readlink("/proc/self/exe",buf,n-1); + if(k>0){ + buf[k]=0; + char *sl=strrchr(buf,'/'); + if(sl&&(size_t)(sl+1-buf)+sizeof("shaders/qmatmul.spv")<=n){ + strcpy(sl+1,"shaders/qmatmul.spv"); + if(!stat(buf,&st)) return buf; + } + } +#endif + return "shaders/qmatmul.spv"; +} +static int vk_budget_full(void){ + double used=0,budget=0; + if(!coli_vk_mem_budget(&used,&budget)) return 0; + double cap=budget-0.5; if(g_vk_gb>0&&g_vk_gb=cap; +} +static int w_vk_upload(W *w){ + if(w->vk) return 1; + int fmt = w->fmt==1?1 : w->fmt==4?4 : -1; + if(fmt<0) return 0; + return coli_vk_tensor_ensure((ColiVkTensor**)&w->vk, + fmt==1?(const void*)w->q8:(const void*)w->q4,w->s,fmt,w->I,w->O,w->gs); +} +#endif static void w_matmul(float *y, const float *x, const W *w, int S){ +#ifdef COLI_VULKAN + if(g_k3_vk&&S==1&&w->vk&& + coli_vk_matmul((ColiVkTensor**)&((W*)w)->vk,y,x,NULL,NULL,w->fmt,1,w->I,w->O,w->gs)) + return; +#endif if(w->fmt==0) matmul(y,x,w->f,S,w->I,w->O); else if(w->fmt==1) matmul_q(y,x,w->q8,w->s,S,w->I,w->O); else if(w->fmt==4) matmul_i4_grouped(y,x,w->q4,w->s,S,w->I,w->O,w->gs); @@ -519,6 +582,35 @@ static void model_init(Model *m, const char *snap, int n_layers_env){ free(rn); free(rp); } w_load(m,&m->lm_head,"lm_head.weight",c->vocab,c->hidden,hbits); } else fprintf(stderr,"[K3] final norm/lm_head not present — trace-only mode\n"); +#ifdef COLI_VULKAN + { const char *ev=getenv("K3_VK"); + if(!ev||atoi(ev)){ + char sbuf[512]; const char *spv=k3_vk_spv(sbuf,sizeof(sbuf)); + g_k3_vk=coli_vk_init(spv); + if(!g_k3_vk) fprintf(stderr,"[K3-VK] Vulkan unavailable (tried %s) — CPU only\n",spv); + } + if(g_k3_vk){ + g_vk_gb=getenv("K3_VK_GB")?atof(getenv("K3_VK_GB")):0; + g_vk_upcap=getenv("K3_VK_UP")?atoi(getenv("K3_VK_UP")):8; + g_vkexp_n=(int64_t)c->n_layers*c->n_experts; + g_vkexp=calloc((size_t)g_vkexp_n,sizeof(VkExp)); + if(!g_vkexp) g_k3_vk=0; + } + if(g_k3_vk){ + /* dense residency: shared experts first — they run every token and are + * the biggest always-on RAM-bandwidth slice that fits VRAM */ + int nsh=0; + for(int i=0;in_layers&&!vk_budget_full();i++){ + if(!m->L[i].sparse) continue; + Moe *sm2=&m->L[i].moe; + nsh+=w_vk_upload(&sm2->sh_gate)+w_vk_upload(&sm2->sh_up)+w_vk_upload(&sm2->sh_down); + } + double used=0,budget=0; coli_vk_mem_budget(&used,&budget); + fprintf(stderr,"[K3-VK] resident: %d shared-expert mats (%.1f/%.1f GB); routed MXFP4 tier fills at decode (K3_VK_UP=%d/step, cap %s)\n", + nsh,used,budget,g_vk_upcap,g_vk_gb>0?"K3_VK_GB":"driver budget"); + } + } +#endif expert_table_init(m); /* expert LRU cache, per-layer slots from the global budget */ double egb = getenv("K3_EXPERT_GB")?atof(getenv("K3_EXPERT_GB")):8.0; @@ -530,7 +622,8 @@ static void model_init(Model *m, const char *snap, int n_layers_env){ * regardless of K3_EXPERT_GB. */ if(cap<1) cap=1; if(cap>c->n_experts) cap=c->n_experts; - m->ecache=calloc(c->n_layers,sizeof(LCache)); + { int ncl=c->n_layers>0?c->n_layers:1; + m->ecache=calloc((size_t)ncl,sizeof(LCache)); } for(int i=0;in_layers;i++) if(m->L[i].sparse){ m->ecache[i].cap=cap; m->ecache[i].s=calloc(cap,sizeof(Slot)); for(int j2=0;j2ecache[i].s[j2].eid=-1; @@ -775,6 +868,59 @@ static void expert_apply(Model *m, Slot *s, const float *z, float wk, for(int i=0;ilatent;i++) u[i]+=wk*hz[i]; } +#ifdef COLI_VULKAN +/* GPU apply for a tier-resident expert (decode, S==1): w1/w3 in one paired + * submit, SiTU-GLU on CPU, w2 down. Returns 0 untouched on any failure so + * the caller can run the normal disk+CPU path. */ +static int vk_expert_apply(Model *m, int li, int eid, const float *z, float wk, + float *u, float *gate, float *up, float *hz){ + Cfg *c=&m->c; + VkExp *v=&g_vkexp[(int64_t)li*c->n_experts+eid]; + if(!v->w1||!v->w2||!v->w3) return 0; + if(!coli_vk_matmul_pair((ColiVkTensor**)&v->w1,gate,NULL,NULL,c->moe_inter, + (ColiVkTensor**)&v->w3,up,NULL,NULL,c->moe_inter, + 7,z,1,c->latent,32)) return 0; + for(int i=0;imoe_inter;i++) gate[i]=situf_(gate[i],up[i],c->situ_b1,c->situ_b2); + if(!coli_vk_matmul((ColiVkTensor**)&v->w2,hz,gate,NULL,NULL,7,1,c->moe_inter,c->latent,32)) + return 0; + for(int i=0;ilatent;i++) u[i]+=wk*hz[i]; + g_vk_hit++; + return 1; +} +/* Fill the tier from a freshly-read RAM slot (main thread only). The ue8m0 + * exponents expand to f32 group scales at upload (the shader is float-only). */ +static void vk_expert_try_upload(Model *m, int li, int eid, Slot *s){ + if(!g_k3_vk||g_vk_full||g_vk_up_left<=0) return; + VkExp *v=&g_vkexp[(int64_t)li*m->c.n_experts+eid]; + if(v->w1) return; + if(vk_budget_full()){ g_vk_full=1; + fprintf(stderr,"[K3-VK] expert tier full: %ld experts resident\n",g_vk_res); + return; } + uint8_t *w1p=s->buf, *w1s=w1p+m->e_w1p, *w2p=w1s+m->e_w1s, *w2s=w2p+m->e_w2p, + *w3p=w2s+m->e_w2s, *w3s=w3p+m->e_w1p; + int LT=m->c.latent, MI=m->c.moe_inter; + int64_t n1=m->e_w1s, n2=m->e_w2s; /* scale counts = scale bytes (u8) */ + float *sc=falloc(n1>n2?n1:n2); + int ok=1; + for(int64_t i=0;iw1,w1p,sc,7,LT,MI,32); + if(ok){ for(int64_t i=0;iw2,w2p,sc,7,MI,LT,32); } + if(ok){ for(int64_t i=0;iw3,w3p,sc,7,LT,MI,32); } + free(sc); + if(!ok){ + if(v->w1){ coli_vk_tensor_free(v->w1); v->w1=NULL; } + if(v->w2){ coli_vk_tensor_free(v->w2); v->w2=NULL; } + if(v->w3){ coli_vk_tensor_free(v->w3); v->w3=NULL; } + g_vk_full=1; + fprintf(stderr,"[K3-VK] expert tier full: %ld experts resident\n",g_vk_res); + return; + } + g_vk_res++; g_vk_up_left--; +} +#endif + /* ---------- async loader pool (K3_PIPE): expert preads overlap compute ---- * A batch of jobs is submitted per token+layer; the compute loop below waits * per-expert on its ready flag, so expert j's math runs while j+1.. load. @@ -869,6 +1015,9 @@ static void experts_apply_union(Model *m, int li, int nu, const int *uids, usleep(50); m->t_eload+=now_s()-t0; } +#ifdef COLI_VULKAN + if(g_k3_vk&&qof[j]>=0) vk_expert_try_upload(m,li,uids[base+j],use[j]); +#endif int f=pfirst[base+j]; for(int p2=0;p2contig){ if(er->fd[0]>=0) posix_fadvise(er->fd[0],er->off[0],m->e_slot,POSIX_FADV_WILLNEED); } else for(int k2=0;k2<6;k2++) if(er->fd[k2]>=0) posix_fadvise(er->fd[k2],er->off[k2],sizes[k2],POSIX_FADV_WILLNEED); } +#ifdef COLI_VULKAN + /* decode: tier-resident experts run on the GPU and drop out of the + * disk union — that skip is the I/O saving. C>1 prefill stays on the + * CPU-batched path (and still feeds the tier via the upload hook). */ + if(g_k3_vk&&C==1){ + int keep=0; + for(int j=0;jc; int D=c->hidden; +#ifdef COLI_VULKAN + g_vk_up_left=g_vk_upcap; /* routed-tier upload budget per step */ +#endif int nbmax=(c->n_layers+c->res_bs-1)/c->res_bs; float *hidden=falloc((int64_t)C*D), *bres=falloc((int64_t)C*nbmax*D); float *prefix=falloc((int64_t)C*D), *nrm=falloc((int64_t)C*D); @@ -1372,6 +1542,10 @@ static void serve_one(Model *m, Tok *T, ServeReq *q){ printf("PROF %.3f %d %d %.3f %.3f %.3f %.3f %.3f %d\n", dt,np,gen,disk,0.0,moe>disk?moe-disk:moe,m->t_attn-a0,m->t_head-h0,gen+1); fflush(stdout); +#ifdef COLI_VULKAN + if(g_k3_vk) fprintf(stderr,"[K3-VK] routed tier: %ld resident, %ld GPU hits so far\n", + g_vk_res,g_vk_hit); +#endif } static void serve_loop(Model *m, Tok *T){ diff --git a/docs/kimi_k3.md b/docs/kimi_k3.md index b9b9e8319..7e1272075 100644 --- a/docs/kimi_k3.md +++ b/docs/kimi_k3.md @@ -154,6 +154,9 @@ Judge quantization choices on real-text logits, not synthetic-vector norms. | `K3_MLA_BITS` | 8 | load-time bits for MLA projections | | `K3_HEAD_BITS` | 8 | load-time bits for lm_head | | `K3_EXPERT_GB` | 8 | routed-expert LRU budget | +| `K3_VK` | 1 | Vulkan tier when built with `make VK=1 kimi_k3` (0 = pure CPU) | +| `K3_VK_GB` | driver budget | VRAM cap for the Vulkan tier | +| `K3_VK_UP` | 8 | routed-expert uploads per step (fill-once tier) | | `K3_DIRECT` | 1 | O_DIRECT expert reads (0 = buffered + WILLNEED) | | `K3_IDOT` | 1 | int8-activation expert matmuls (0 = exact-float kernel) | | `K3_PIPE` | 1 | overlap expert loads with compute (loader threads) | @@ -220,12 +223,37 @@ is returned as `reasoning_content`, response text as `content`, and `<|end_of_msg|>` remains the model-owned stop token. `STOP` and `CANCEL` are honoured between generated tokens. +## Vulkan tier (`make VK=1 kimi_k3`) + +The shared Vulkan backend (`backend_vulkan.c`) gained an **fmt=7 MXFP4** +decode path for K3's expert format — e2m1 nibbles with the ue8m0 exponents +expanded to f32 per-32-group scales at upload, so the QAT bytes are uploaded +exactly as stored and never re-encoded (kernel vs `matmul_mxfp4`: rel_l2 +2.2e-07 on an RX 9070/RADV, 2.6e-07 on llvmpipe; +`tests/test_vk_mxfp4.c`). The engine keeps two residency classes on the +card, both with transparent CPU fallback and identical output: + +- **shared experts**, uploaded once at init (int4/int8, the existing + fmt-1/4 shaders): they run every token and are the largest always-on + dense slice that fits VRAM (7.5 GB for all 92 MoE layers at int4); +- a **fill-once routed-expert tier** in fmt=7: experts enter from + freshly-read RAM slots (`K3_VK_UP` per step) until the VRAM budget + (`K3_VK_GB`) is reached. At decode, tier-resident experts skip **both** + the 17.5 MB disk read and the CPU matmuls (one paired w1/w3 submit, + SiTU-GLU on CPU, w2 down). Chunked prefill stays on the CPU-batched path + and still warms the tier. + +K3's Quantile-Balancing-flat routing caps what any cache tier can do — the +tier's value scales with how long the server lives (fill-once) and with the +measured short-term reuse (temporal locality), not with marginal expert +heat. `K3_VK=0` disables the tier at runtime. + ## Current limitations - Decode is single-token (no speculative decoding — K3 has no MTP head). - Tool declarations/calls and image content are not exposed through the shared gateway yet; unsupported requests fail explicitly. -- CPU only (no CUDA/Metal/Vulkan tier). +- CPU + optional Vulkan tier (no CUDA/Metal). - The protocol, tokenizer, gateway, TUI, and Web client paths are locally testable without the 1.5 TB checkpoint. A release claim still requires one full-model multi-turn TUI/Web run on a host that owns the complete snapshot.