From 72767815bbf8dfe83ce507835bc5bf17db9958f6 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 11 Jun 2026 10:52:04 +0800 Subject: [PATCH 1/5] [Feature] Support MegaMoE (#7943) * [Feature] Support MegaMoE * update code * fix code style * fix code style * fix test * fix test * fix test * fix test * fix typo * fix code style * fix test * fix xpu test * fix test * fix test * fix code style * fix typo * fix typo --- custom_ops/.ninja_deps | Bin 0 -> 89432 bytes custom_ops/.ninja_log | 464 ++++++++++++++++++ custom_ops/gpu_ops/mega_moe_pre_dispatch.cu | 336 +++++++++++++ custom_ops/setup_ops.py | 1 + fastdeploy/config.py | 1 + fastdeploy/engine/args_utils.py | 11 + fastdeploy/engine/engine.py | 1 + .../layers/attention/flash_attn_backend.py | 2 +- .../batch_invariant_ops.py | 2 +- fastdeploy/model_executor/layers/moe/ep.py | 24 +- .../layers/moe/fused_moe_deepgemm_backend.py | 459 ++++++++++++++++- .../layers/quantization/__init__.py | 32 ++ .../layers/quantization/fp8_utils.py | 30 +- .../layers/quantization/mxfp4.py | 2 +- .../layers/quantization/nvfp4.py | 2 +- .../layers/quantization/wfp4afp8.py | 63 +++ .../model_executor/load_weight_utils.py | 2 +- fastdeploy/worker/worker_process.py | 8 +- test.py | 74 +++ tests/engine/test_engine.py | 1 + tests/model_executor/test_ep.py | 1 + tests/operators/test_mega_moe_pre_dispatch.py | 171 +++++++ tests/quantization/test_modelopt_nvfp4.py | 2 +- 23 files changed, 1676 insertions(+), 13 deletions(-) create mode 100644 custom_ops/.ninja_deps create mode 100644 custom_ops/.ninja_log create mode 100644 custom_ops/gpu_ops/mega_moe_pre_dispatch.cu create mode 100644 fastdeploy/model_executor/layers/quantization/wfp4afp8.py create mode 100644 test.py create mode 100644 tests/operators/test_mega_moe_pre_dispatch.py diff --git a/custom_ops/.ninja_deps b/custom_ops/.ninja_deps new file mode 100644 index 0000000000000000000000000000000000000000..88c3ef741f77ffd79f52cacb90d9350654f30196 GIT binary patch literal 89432 zcmeI*2b>etzBuqGD0Z?=cswu0A zM@vUX62q#JiD6?apO3P~ERpq4> z{++G*ce(%3UrUwvWzCqjW4d+;cj?qNnySrKCKGKtwr@8-KN*&)$^GGIDiw~6$;J|y zcruYGN!q{P{of%Aen0<*2%^#7Z2$cK?e?qViFh=bD35e3jU`59N=H|X4Ohp8MZ&4_ z(yDk_X}(50HYjA=j(Tx&Av(yC-MQsvYwk2F-5;J<{^ zHHmDzIu`8QkG|jT@4lbTmR2S+*|54=5g(Q-k01VF8$kn+jVHp{+Egs#7vKNG_se6W z;?Y>c*F)d^yD^o=DzsI>CPLr-(>BuaQL*&AZG7_&-yao=YA@$~|LcGFex@c}5oxp& zU;V@PvrgEW^t_$^@*lphNk~W1wG9D%;roAk+$tl9M64=Y9?L}2@sxcjsr>w(e$%LK zgzGB*XMWT5@9IbJo4Ny7SFP%Q8v67fwpX4%CKWY_Xg1flJC2|H!*A-2%O8I8@qhfA zRO%zYnH=ob|0Y-de{RF0$?EE4qO>Yfo2<#I)E|oRvcilpHYOTN1qUrS_dY1xvaF6o z)5)NLcwdZD3al)HoLG0v%OlxH?%WQ_^S#23wXC8l8OgTq5`-Fhw*X_UE7LMPrdXmp zTp80n*8GL$inVo62mU5E20QhR?-l5I7o_5RYb(&V3oy<`b=qM_RmH}XRK6w7Hh)|0 zJJo+RG{<*4{oQw?(W=sHW!mLPaPH(c;QMd--QVV^M!zeMWg~I@qWT`q8-CmU-`wUO zemn0^gkJx5;?DKN!r`b`>iScm*Z!St=YHqFm*#(48c9bhT@yru9)x=Es=v|gpW^?Q zzmdtr(m|F9gj79C-{s0M*sh` z9O9X9BAEy$BI$H;bUZOEI52vO*?3u1P_{43Ph75x)3InI8!OM1@ALEXHySc*u6rAL z&foZtk8vtHCLGU1f*_vt`+qz?{#-}W*=%qH^o-y3`cKO>6B}6*OXv|tqQOB9Vt9If z%AqEajzyxCk^JEZ&7Plfa6E(2gCK^d{JzlI_C{@HLu!3eoYw|rUO@fRKgZcM9P3ow zu_9NlPl$I);m6y#xvEK){AV_(nW4}_V!gG%j`<&AHLzFf*jGogrSm?PRL&^$GR`K$`O+?_e6RrP zY`%;~Mbhy|BIuPrAjSgjr(_H)M}lkW`wO-#lezJ!h}#X*#aL+fYx2k2m4`O;KJndH zVCQ_0kAnL(_4h(t6$L}q;Z!o2yMc4BmOEpD#J+Un z@~x=J#LDYN&ZiYYyqQ!aopEs&+#6sk~F1g?66R zw+-5Z((<^Pkd0ObWjeJ$V~r-$vC>g-x6vDh*h6=SwNTrHpj?wFvjN-_i0fzvWqZ4L z3pGas@g^c!H(2vA*WZi2P0W)EFt6q6LwR|;TEX60fHBrxW7S0%3f&^cLhHZO#~2ji z&4pWr4dW9xiLubm_qsAP!kMh9$`7RkQ=Wp;J#?dZ3-Nr?h_Vf1O>Y|E29UGl?L$kgP zIRaM|VEjR?ufHCyFY&s_O~cIB{#;EG{}k~TTE1xfB7y+ zN#&*DEx5kZpX+ke{P}rB9J)lzg_bY>rB2s9LSSGmVJdn$oT*8rlKDv_7Z*d_c9$iV zpQhUo`Q+k?dsL*VCUi4D z^G)5&b>$NIIf!Q#LwyMHh4MY480teVx8=t=P8V;%UyuB$@1~eTrximvXHuRd%U}1M zS`6`~W5epEnubmhZ^5?{^0NEn3OWC5-^s;L#<_cW`MF&u6+;=9<;DUUn&T6Tp^O{G z0ZJ+-6+^ta_Bas+lIzEoE7Hlav4qLXe;6bWXJY#K z6Ya~4BHEYwaq{||$v*>JQaPaj&x3jQy7G4shs6@E0jk1{A3i=t+=ZOu=bZ=h&fTMn zCgz|s-eLVcF7%rLABL= zf?Gu6il^*@*n{d{>gm3k>KapS^}7x)!uWS?w2;AJ#T9@4IkD1y^-COt4i$gF^=tmR z?i)ypR+i?U%(voHZN4%!Iv9u8mtt*8a}!|XfrITs!S&7l`gr^~AN0sQw3{34428zp zkJ-hM*Z!qmHOg^+*r!_iR3v$>p{?x^JZw2I_XuNXjQEQ!&-Kc!^Hfqfy12?dSB3Ji z56%rxj}rf^Vm==l?~MJZDbTu6iSM@P0;@RkLxG}Uk#74)Iuon?|nD-99gyR7Gh3oIG5|o!apO`@u=HUbt_AS zLRDff;OEKer)h;oh_j&E%lb057EQv+BEt)`PW)|ds%tyrVlSd|yKgY_y;98od~WBy z8X8ta$0i&bWg2eX!wK513b7aF`D#qAVRemAC??LDfb}5?uA?*@_m~>tc$;!cnh)rn5)0}cym`2^JXoDLVXIq zjB_#9Kj*Zocnfh1Hdn@V*AsP3uxa+e>x8`HCjVGYXcw_hF0gB%{H1WIMd^m+2G4to zu@LV&{}1IG&%F(-Je+^+iXO$z#Z^BVJv|Z%^%8p_wjT}sw)%UUjgsI_#S(L6tjfv( zxjtu4F&ARq{zEzE*1FGKGBrfrL)?XU{gFF28?A59U7Q6w=94SZc^4o%7IK-+3p2EX zI190D$Q}DX&4H+2#c+Eu7xeWz)Z|cpjl|-3tgiiG6y3T~GZF&$Bed-Y<8{sdw%L z1wm)DV-c4Bpt{BNJBYmyuND8U?&M=1oSP@xUi`C)qyGH$yz$06V?Wvz*YU_bpVPDV zyq7hWiho9N9FK-M@44_hg<~D6+gBvQV}dD#p%VKsy%_hS(F?QM7VEL7--nv9xoiEc ziuhRMUIwNfZCS*}!ah_a>W;-0MSLu}R=2YsC6$}oheGV@{SP@W*T(J|smi^)sZ9~a zUKY>h#-%qCdlB_Da^)WMTXQd%bq;S@gk?X{dXybyv`vaM_O3EO>msZ-BWo0VhcTUk z1h#Rp#@{L57uu*8>rPPqD$j<+Rra}|*g>uotbLp7`)nZof*pr1x2{7hoPVW&XxA5O z!H#d&uTAHL(C#5+a1p)+w3y6#804xDWaDV)o3v8!~8R z)-1w(=oPDoSmd`?FmWe1pVzPtGmGbZ$k(OZ7-uj}I5-*nXTQ3(X?6QjY}bFmaR^p9 zUrp>q)UM{Pp@Oo{t)8SJty(gI0k z|8Y-3p_Pm0oX8)ShBw%ULMs*5Igx9ddYXmlJEj%qRNG zSj+ViEIJppyoAo%vf?e6`FejDZ!l@Eq;eT?7Rauz&7jN##0UYQ?>bJN}4x8`Nzft+s=)`V zw4!Y?jk*ln+E$KA;bx~!pJ zY=8VhZu{!%K;t-mZh~?M#t|}A$y_V+lkG3;{F|z+jAXO{KQ`epNY$3d%4+iOBKe^S z;wY=l#&RzWuRkHw#qY&&LKD@=zl!HO@r+x@@dR;X8kNbnVwkv)W2iq*{~XIVO;8>+ ziFh)w4bTW$;-O#85XUry+_@n;-@Yat)hzW`;taG(imcN0l3jA@p$*#E`8q zJ!)PEp^wDSRBiU3WBE`l3%!jpE?Qq4AJ~3V^)ss)#qoXEOVNoJd))#jT`Sbo1jhv@42c!B5yQ78RUlf>jO|9UvG*sh|fC;ui1W6<&26( zalG0D<&dh4Mp9la9ZWcR#r7BW*hFx#6?E)hw(W)8Hq~Eu*2nNtlazl&I`?u*=fR6j zu>EX05{(59o;wIH*nU%;2aV%+z6t6=rYat7+^2r7364cHW-?^mnDMhsu>Igw==Ixw zrU~mpJQ2^<-?DnT3Ch7zYW2ksn%xB3Pt|6!xfic#6P{|qS&!1p!5_N$ZSnVM97&dIg(h0E%o_(S*DepA)M z#&Jw*f^x`Dfl1c{Z|u0c3Cp2z9CwLhVfUX>wN>$oShO};S10Zi!$fssVUB~&%ME{w zWvW;fdX9*ViUhNA?r4HzQC|+xx)^R3!@|x}-f5(X%8z;7X8TRmJ~WQw)+Q*2cx5AN$nD`UB_r<-kmVV?&krsoI28ZAGp&PHBP|a^vvzFZR+$y0Qsk$lVcX7-0#8t`I|0_1FF! z%jIIZY(e%Ny+q5T^@e9>G6v<_8fzJ0agt1FDpFov74v*-w5p~&R{EFE{6r(cQ-YU? zvnZbV$_;LoR9-6HsRdoGbTm3V3sZ{WMcX>@bcfaByaTmq2fjZ%_bRrVGv;4whi8-ju!NLsA zxeLTx6zzsf?wVNGyL!Ww$*S_)8#T@sdttX5$%ge*&J*jjf^OsL&coQ4Xe<|XSXU;I zsfuKC3rCzQ{-Vf-!Od?q`yBBWc0XfSJQzbdTbza6K8rJ)iBtqh+iM__iD>RZ?<{c_ z_PM8#bE7k8R$g7R)(T*Hx~ls)?7EHZ0~JoLymS%?HIAI=$fQfXTzD z+FY!s6@0AoR&701tkahWjZ?&-cOIS8alqW$)?Ccz#c4)!mS9O*uoO$P49l_{%d-M2 z(vp=}nN?Vo)mWW1Sd+C_n{{Z#x~#|gY`}(W#KyE{6EajM2R+$|UhGV7c41fgup7J6mwxQQp7dt`0~y3%_F@Qo zvk&_+l>OMB12~W{2N9u+DCNYcU>L!<5N9|esG^z#Nm7g?O@=HrjAArnsAVh%a|nlW z7>6^CBRGa4PIg7J7hjTfP^SOWvxroVJ z%q3jPWn9h`T*(x!;%ctpTCU@IZs104;%08)R&L{V?qDi+au;_qjeEG4`^95h>6<_lW-|`*b^8-Kf6F>6{zw#Tu6I#swvpCIY&Jrw13zlMO zmSI_zV|i9!MOv~FE3*o#vKp(i20`Pp7HhK(tyq`!Sf35pkd4@w)@;J2Y(^V4XA8Du zE80>*DeY)a2RhP;&U9gGy0Q)1vK`&no*mec?)0E1JJE}s>CG(GjIS&#MEfDPG*jcLs$Y|3V|VRN=%OSYmdC6v;R_H>{lo#;##wx%oFur1rsjqTZi z9qCRFda@I}*qPq!!mjjTH+H8l{n&#&>CXTLGKj(K#Sr#pANFM^`>{UVC*&D_GR+{W$P!Bp<#F79R;_i!)wF`fH)fCrhuLp;nQ%;ZrX<8fy31W)o5 zvw50lc$VjQo)>tLmw1_1c$L?9oi})sIlRT&yu)1Hvk#7dGUE&M$B z#|Wa)(1~?ROC27G@EX^`3%W^Ew3am&=R$^sVVO3UR zb=F`_)?#hep%v@09_zCK8?q4_)0$1#l+9?v=4`>1Y(-m2D5V|k=|D$1(U~r6O;@&I zTehPc+p_~Z(w!dkWG8yDGrie`UFpMa>`q_$u?Kt7p8*VH5QEu^A?(dQ?8{L0V}B0d zK*AhEgfgO(6QhD*R1#-6BdDU91W8hiBu$1aHH>03W2j{;2XhFAau|m*jw3jdqZrT8 z9K!@AaxBMjJSQ-T6FG^KIfYX>jng@UGdYX1IfrvOkMp^J3%Q8NT+Ah0%4J;66f%-|s&<`HJ{D39?t zvv`6hd5YOQ%`-g9b3D%ryvR$u%qzUgYrM`IyvZEi;%(kxF7NUl@ACm4@(~~N37_&A zpYsJ@@)ck64d3z|-}3`M@)JMv3%~Lkzq2S$7T2}q9Ea?`E>9+FvZGlc4K$?(vLmZ zll}}~AcGjpUJPMx_F-R!vLE|%00$E0AR?3zrJNWQ45N}b!x=#p)g(xgVkBuYWT{~k zqY3JHEn_*DLpYSfIGk}D!I2!rc#h^6CNPm>IgaBwfk~XmNu10noXTmO&KaD^S)9!| zoXdHf&jnn_MNH;mF5yxx<8rRxN~UlXS91;5avj%m12=LLH**WOavQgE2UEF|ySSTa z+{3-x$8_%J0Ul%q5AiUMFq21ljK`V96FkXN%;ssH;aQ&Jd0yZ}UgBk5;Z84j z=I|D8^A2-)m-l#|5BQLe_?S=ll+XB_FZhzL_?mC{mhbqUANY}<_?ch$mEZWCMY*}S za&vI}|4VaoxH>Z|JUX7O3}=(65i<30@>S6OEa3h9*F62FNTN|~U|sI4Sxj!2eNej~ zcaHLTQSMxnI~V27MY*$?wxBsnup}*5bkDS~?wQV6@VR6@ZhX$iPyQlf#Zo27iE`8- z3w|H}dh3$p1lu38@U|ax7wR8L3LRtH(_Pz+cHYI;daiwtpU3zd)@41`X9G55BQ~Zro3JUH(T2_0 zf-TvKwvn=-2X>@8J?P0!^kQdvvkSY@huzqnzVu@c z_M|@p7|0+7vlm0yn|;`qq3p;09KeBuIfw{lL@6gm1;eN$&TvLhMKuYMq!>w>3|VRz z#b|>1R?AopCaC|1au|m*jw3jdqZrT89K!@AaxBMjJSQ-T6FG^KIfYX>jng@UGdYX1 zIfrvOkMp^J3%Q8NT+Ah0%4J;66f%-|s&<`HJ{D39?tvv`6hd5YOQ%`-g9b3D%ryvR$u%qzUgYrM`I zyvZEi;%(kxF7NUl@ACm4@(~~N37_&ApYsJ@@)ck64d3z|-}3`M@)JMv3%~Lkzq9E2 zwz#fugZi+5ulMF7;m*fT>a%2SHghIE%;P)Zs)*HhM#qEuv`5QVL0hzd_x10ikh8~& zW8A`vBOR+wW`k+|))WYj7!{7Eq9v6_iD8D^J$5neVmA7Ekh>4|d2w+yqd7~kBrRBq zrCEk$S&rpdffZ@VO03K(tjcPv&Kj)ATCB}Fv|?S>V|_MYLpEY#TC)k8vKejIoGsXr zt!PULrL?0x9q33WI@5)%>B=^2%XV~Qdv;()y3>Q6>_jhirZ>B=D}C6F-RVm|_Fzx? zGk}2%VlaC#guU5^eHqGr?9TxlNSK3&P)3w;VpK4UO5zM>1XWa%AW4dmq{)z_hEa@W z47H5qU=HC>4&!jfaRf(l6yrIXW0=51j^#Lx=L9BkA}4V&r*JB#aXM#kCTDRr=Ws6P zaXuGtAr~>3i@AhLxs1!Xf-9NARb0(AT+4M_&kfwjP29{a+{$g-&K*qUPVVAvrg0DV zav#&Vp9gr589cy{F zEH>{tD>(KGIH%1=GMLcZ@p#53E`}}W3 zBb|qisLQ2?Iwr#-+l@a@Eu`;W$Dyy+N zYp^D3u{P__igj6!_1S<8*@%s4%_eNhX0%~*wqQ%PqAewq(vJ3Ypd+2=Oc%DME8DOw z+tH2f*?}GDP7ivr6TR4(-t5Az^kFx4r!W23gFWfb00uIM!R*Bl_GTaUWhncxKL>Ci zVGbff8BxlKQNb`Oi8GuLR8dWWBq>IcCPS7QMlqT()H0TXIfO$wjKdkn5gf@;jOS>M zVFD95mg6{{6PUz_oW#kT!l|6b>72otoWV zaW&U)E!S~9H*h02aWl7YE4OhwcQBPZxr@7*#y#B2eN5+m9^gS{@DLC42s3$<$9SAs zJi(JZ#cZDD8J^`ip63N#!V%Px*|` z`GPO`im&;GZ~2bz`GFt#iJ$p}U-^ySS(HnQHZs!iBawm6jH`BO> zd%2J4+|L6%$P6ChVIE;7kMbCgGm9s9lBbx>(>%koJje6Az>B=Z%e=y?yvFOi!JEwC zE#BrG=JGD@@jf5$As_KEpYSQ4@i|}cC13G1-|#Kp@jXBABR}yozwj%+@jHw1XmRDy z;P@}#d@>)2a6W!gpS3cPQL%7-8u6c|CnXP)wajM2R+$|UhGV7c41fgup7J6mwxQQ zp7dt`0~y3%_F@Qovk&_+l>OMB12~W{2N9u+DCNYcU>L!1i!+=NR8dWWBq>IcCPS7Q zMlqT()H0TXIfO$wjKdkn5gf@;jOS>MVFD95mg6{{6PUz_oW#kT!l|6b>72otoWVaW&U)E!S~9H*h02aWl7YE4OhwcQBPZxr@7* z#y#B2eN5+m9^gS{@DLC42s3$<$9SAsJi(JZ#cZDD8J^`ip63N#!V%Px*|``GPO`im&;GZ~2bz`GFt#iJ$p}U-^ySS#;f7 zT-Uw9byhI%{J$%Ir()3>tHj2_1XF+?hfXkxcx_CaK4AOsSj5vdKvgzW?G=y$ab#5)o(RF=+W7MU?&Kwk< ztS>OZ=ZP#X)@C$k36`V_~Td(373$ z#m@9*7j~r&yRkcc>Bk=INq+_~kU-!V*^m7>fCCA05E06VQcjEthEYkJ z;f$b)Y7#gd7WD;`!A;!EE!@g&+|C_LGlc4K$?(vLmZll}}~AcGjp zUJPMx_F-R!vLE|%00$E0AR?3zrJNWQ45N}b!x=#p)g(xgVkBuYWT{~kqZvajV>y^Z zIF!RUoN*k%ksQT%j^-F9Fp*<9j^jCjNu0<@oXjbl%4wX=8Jx*koXt6$%Xys71zgBQ zOy*)P;ZiQ+a<1S?rf?Nka}C#W9oKUMH*ymdpRbJzD-r!B<@D^|L4s&^z z_jsQV_>hnIm{0hW&-k1#_>!;qns4})@A#e{_>rIZnP2#o-}s$HxwZImYj7>pod1jY zwedn7S-EVoD!hQp;Xll?jUtqBGhz(xOCRreRW@6FV32Q5@Og1@H={XAup}*5ilteG zWm%5pS%DR4$x5uuDy+(Ctj-#&$y%(KlWfx`ZIum z3}P^QF@(L@hkY5!e(cWy97vdhh)_n9a$;04j7s7RX9QJLlORcok)+9xrG`8A zaUv&iGN*7Vr*S%Ga3*JQHs^3I=W#w4 za3L2lnTxrEOSz28xq>U1!c|<&HC)SeT+a>M$W7eLE!@g&+|C_L!cmvVtta}txZ~A~<}AUI zv|uTgW*L@cIhJPyR-`2>u`;W$Dyy+NYp^D3u{P__igj6!_1S<8*@%s4%_eNhX0%~* zwqQ%PqAewq(vJ3Ypd+2=Oc%DME8DOw+tH2f*?}GDP7ivr6TR4(-t5Az^kFx4r!W23 zgFWfb00uIM!R*Bl_GTaUWhncxKL>CiVGbff8BxlKQNb`Oi8GuLR8dWWBq>IcCPS7Q zMlqT()H0TXIfO$wjKdkn5gf@;jOS>MVFD95mg6{{6PUz_oW#kT!l|6b>72otoWVaW&U)E!S~9H*h02aWl7YE4OhwcQBPZxr@7* z#y#B2eN5+m9^gS{@DLC42s3$<$9SAsJi(JZ#cZDD8J^`ip63N#!V%Px*|``GPO`im&;GZ~2bz`GFt#iJ$p}U-^ySS(ICg zFSiEQvi}$E8_ydnZj{9uT`DG3-&6rz&;E6k(P*;ndiEs8t>2KZt_^bT$v!VG{$@01 z36`V_~Td(373$#m@9*7j~r&yRkcc z>Bk=INq+_~kU-!V*^m7>fCCA05E06VQcjEthEYkJ;f$b)Y7*q?oPUla zO@=HrjAArnsAVh%a|nlW7>6^CBRGa4P zIg7J7hjTfP^SOWvxroVJ%q3jPWn9h`T*(x!;%ctpTCU@IZs104;%08)R&L{V?qDi+ zau;_qjeEG4`^95h>6<_lW-|`*b^8-Kf6F>6{zw#Tu zvnb~lU(Wr%HFjK+Di3CG#S>n`5RR8;!qHSsNitDDx%Mw3j7CG%VhrBAeO}yY*No;Y!IHFKDVAm#mSs7XX9ZTIB`dKqtFS7ou{vw8CTp=a>(GjIS&#ME zfDPG*jcLs$Y|3V|VRN=%OSYmdC6v;R_H>{lo#;##wx%oFur1rsjqTZi9qCRFda@I} z*qPq!!mjjTH+H8l{n&#&>CXTLGKj(K#Sr#pANFM^`>{UkMyP3v4+{=AT=YAgGL1yp}5Az5!d6dU^oLM}(_ zANh%&`GsHkjo(?6Yl|z_2KD2=YTP)RjwCV_v2>%kPE~T$xCOtT^&{7RjzE^w)-L22 zf_dKI`bbJDQz_~Td(373$#m@9*7j~r&yRkcc>Bk=INq+_~kU-!V*^m7>fCCA05E06VQcjEthEYkJ;f$b)Y7!&~j_F9!WXMv(C`L1eTE=oPhj1u| zaX8~Rf+IPK@f^)DOkg6%avaBV0+TqAlQ@}EIF-{loijL-vpAb`IG6J{p9{E_fHCwP*ln9b8X!?Qfc^Sr=|yu{1A!mGT->%766%;7EG<{jqpF7NR^ zAMha`@iCw9DWCB +#include +#include + +#include +#include +#include + +#ifndef PD_BUILD_STATIC_OP +#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name) +#endif + +namespace { + +constexpr float kFP8E4M3Max = 448.0f; +constexpr uint32_t kVecElems = 8; + +template +__device__ __forceinline__ float WarpReduceMax(float value) { + static_assert(kNumThreads >= 1 && kNumThreads <= WARP_SIZE, + "kNumThreads must be in [1, 32]"); + static_assert((kNumThreads & (kNumThreads - 1)) == 0, + "kNumThreads must be a power of 2"); +#pragma unroll + for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) { + value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask, WARP_SIZE)); + } + return value; +} + +__device__ __forceinline__ uint32_t CastToUE8M0(float value) { + value = fabsf(value); + uint32_t bits = __float_as_uint(value); + uint32_t exp = (bits >> 23) & 0xffu; + const uint32_t mantissa = bits & 0x7fffffu; + exp += mantissa != 0; + exp = min(max(exp, 1u), 254u); + return exp; +} + +struct MegaMoEPreDispatchParams { + const __nv_bfloat16* __restrict__ x; + const int64_t* __restrict__ topk_idx; + const float* __restrict__ topk_weights; + + phi::dtype::float8_e4m3fn* __restrict__ buf_x; + int32_t* __restrict__ buf_x_sf; + int64_t* __restrict__ buf_topk_idx; + float* __restrict__ buf_topk_weights; + + uint32_t num_tokens; + uint32_t padded_max; + uint32_t hidden; + uint32_t num_groups; + uint32_t top_k; +}; + +template +__global__ __launch_bounds__(1024, 2) void MegaMoEPreDispatchKernel( + const MegaMoEPreDispatchParams params) { + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, + "unsupported group_size"); + static_assert(kGroupSize % kVecElems == 0, + "group_size must be a multiple of 8"); + constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems; + + const uint32_t bid = blockIdx.x; + const uint32_t tid = threadIdx.x; + + if (bid < params.num_tokens) { + const uint32_t token_id = bid; + const __nv_bfloat16* token_in = + params.x + static_cast(token_id) * params.hidden; + phi::dtype::float8_e4m3fn* token_out = + params.buf_x + static_cast(token_id) * params.hidden; + + const uint32_t base = tid * kVecElems; + float vals[kVecElems]; + float local_max = 0.0f; + +#pragma unroll + for (uint32_t i = 0; i < kVecElems; ++i) { + const float v = __bfloat162float(token_in[base + i]); + vals[i] = v; + local_max = fmaxf(local_max, fabsf(v)); + } + + local_max = WarpReduceMax(local_max); + + const float absmax = fmaxf(local_max, 1e-10f); + const float raw_scale = absmax / kFP8E4M3Max; + const uint32_t ue8m0_exp = CastToUE8M0(raw_scale); + const float inv_scale = __uint_as_float((127u + 127u - ue8m0_exp) << 23); + +#pragma unroll + for (uint32_t i = 0; i < kVecElems; ++i) { + token_out[base + i] = phi::dtype::float8_e4m3fn(vals[i] * inv_scale); + } + + const uint32_t group_id = tid / kThreadsPerGroup; + const uint32_t within_group_id = tid % kThreadsPerGroup; + if (within_group_id == 0 && group_id < params.num_groups) { + const uint32_t byte_off = token_id * params.num_groups + group_id; + reinterpret_cast(params.buf_x_sf)[byte_off] = + static_cast(ue8m0_exp); + } + + if (tid < params.top_k) { + const uint32_t off = token_id * params.top_k + tid; + params.buf_topk_idx[off] = static_cast(params.topk_idx[off]); + params.buf_topk_weights[off] = params.topk_weights[off]; + } + } +} + +void CheckShape2D(const paddle::Tensor& tensor, const char* name) { + PD_CHECK(tensor.shape().size() == 2, name, " must be a 2D tensor"); +} + +void CheckSameShape(const paddle::Tensor& lhs, + const paddle::Tensor& rhs, + const char* lhs_name, + const char* rhs_name) { + PD_CHECK(lhs.shape() == rhs.shape(), + lhs_name, + " shape must equal ", + rhs_name, + " shape"); +} + +template +void LaunchMegaMoEPreDispatch(const MegaMoEPreDispatchParams& params, + uint32_t num_total_blocks, + uint32_t num_threads, + cudaStream_t stream) { + MegaMoEPreDispatchKernel + <<>>(params); +} + +} // namespace + +void MegaMoePreDispatch(const paddle::Tensor& x, + const paddle::Tensor& topk_idx, + const paddle::Tensor& topk_weights, + const paddle::Tensor& buf_x, + const paddle::Tensor& buf_x_sf, + const paddle::Tensor& buf_topk_idx, + const paddle::Tensor& buf_topk_weights, + int64_t num_max_tokens_per_rank, + int64_t group_size) { + CheckShape2D(x, "x"); + CheckShape2D(topk_idx, "topk_idx"); + CheckShape2D(topk_weights, "topk_weights"); + CheckShape2D(buf_x, "buf_x"); + CheckShape2D(buf_x_sf, "buf_x_sf"); + CheckShape2D(buf_topk_idx, "buf_topk_idx"); + CheckShape2D(buf_topk_weights, "buf_topk_weights"); + CheckSameShape(topk_idx, topk_weights, "topk_idx", "topk_weights"); + CheckSameShape( + buf_topk_idx, buf_topk_weights, "buf_topk_idx", "buf_topk_weights"); + + PD_CHECK(x.dtype() == paddle::DataType::BFLOAT16, + "x must be bfloat16, but got ", + x.dtype()); + PD_CHECK(topk_idx.dtype() == paddle::DataType::INT64, + "topk_idx must be int64, but got ", + topk_idx.dtype()); + PD_CHECK(topk_weights.dtype() == paddle::DataType::FLOAT32, + "topk_weights must be float32, but got ", + topk_weights.dtype()); + PD_CHECK(buf_x.dtype() == paddle::DataType::FLOAT8_E4M3FN, + "buf_x must be float8_e4m3fn, but got ", + buf_x.dtype()); + PD_CHECK(buf_x_sf.dtype() == paddle::DataType::INT32, + "buf_x_sf must be int32, but got ", + buf_x_sf.dtype()); + PD_CHECK(buf_topk_idx.dtype() == paddle::DataType::INT64, + "buf_topk_idx must be int64, but got ", + buf_topk_idx.dtype()); + PD_CHECK(buf_topk_weights.dtype() == paddle::DataType::FLOAT32, + "buf_topk_weights must be float32, but got ", + buf_topk_weights.dtype()); + + const int64_t num_tokens_i64 = x.shape()[0]; + const int64_t hidden_i64 = x.shape()[1]; + const int64_t top_k_i64 = topk_idx.shape()[1]; + const int64_t padded_max_i64 = buf_x.shape()[0]; + + PD_CHECK(num_max_tokens_per_rank <= padded_max_i64, + "num_max_tokens_per_rank must not exceed buf_x.shape[0], but got ", + num_max_tokens_per_rank, + " vs ", + padded_max_i64); + PD_CHECK(num_tokens_i64 == topk_idx.shape()[0], + "x.shape[0] must equal topk_idx.shape[0]"); + PD_CHECK(buf_x.shape()[1] == hidden_i64, + "buf_x.shape[1] must equal hidden, but got ", + buf_x.shape()[1], + " vs ", + hidden_i64); + PD_CHECK(buf_topk_idx.shape()[0] == padded_max_i64, + "buf_topk_idx.shape[0] must equal padded_max"); + PD_CHECK(buf_topk_idx.shape()[1] == top_k_i64, + "buf_topk_idx.shape[1] must equal top_k"); + + PD_CHECK(group_size == 32 || group_size == 64 || group_size == 128, + "unsupported group_size: ", + group_size); + PD_CHECK(num_tokens_i64 <= num_max_tokens_per_rank, + "num_tokens must not exceed padded_max"); + PD_CHECK(hidden_i64 % group_size == 0, + "hidden must be a multiple of group_size"); + const int64_t num_groups_i64 = hidden_i64 / group_size; + PD_CHECK(num_groups_i64 % 4 == 0, "num_groups must be a multiple of 4"); + PD_CHECK(buf_x_sf.shape()[0] == padded_max_i64, + "buf_x_sf.shape[0] must equal padded_max"); + PD_CHECK(buf_x_sf.shape()[1] == num_groups_i64 / 4, + "buf_x_sf.shape[1] must equal hidden/group_size/4, but got ", + buf_x_sf.shape()[1], + " vs ", + num_groups_i64 / 4); + PD_CHECK(hidden_i64 % static_cast(kVecElems) == 0, + "hidden must be a multiple of 8 (16B bf16 loads)"); + const int64_t num_threads_i64 = hidden_i64 / static_cast(kVecElems); + PD_CHECK(num_threads_i64 <= 1024, + "hidden too large for single-block-per-row quant"); + PD_CHECK(num_threads_i64 >= top_k_i64, "top_k must fit into one quant CTA"); + + const uint32_t num_tokens = static_cast(num_tokens_i64); + const uint32_t padded_max = static_cast(padded_max_i64); + const uint32_t hidden = static_cast(hidden_i64); + const uint32_t num_groups = static_cast(num_groups_i64); + const uint32_t top_k = static_cast(top_k_i64); + const uint32_t num_threads = static_cast(num_threads_i64); + const uint32_t num_total_blocks = num_tokens; + + const MegaMoEPreDispatchParams params{ + reinterpret_cast(x.data()), + topk_idx.data(), + topk_weights.data(), + const_cast( + buf_x.data()), + const_cast(buf_x_sf.data()), + const_cast(buf_topk_idx.data()), + const_cast(buf_topk_weights.data()), + num_tokens, + padded_max, + hidden, + num_groups, + top_k, + }; + + if (num_total_blocks > 0) { + auto stream = x.stream(); + switch (group_size) { + case 32: + LaunchMegaMoEPreDispatch<32>( + params, num_total_blocks, num_threads, stream); + break; + case 64: + LaunchMegaMoEPreDispatch<64>( + params, num_total_blocks, num_threads, stream); + break; + case 128: + LaunchMegaMoEPreDispatch<128>( + params, num_total_blocks, num_threads, stream); + break; + default: + PD_THROW("unsupported group_size: ", group_size); + } + } + + // return {buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights}; +} + +std::vector MegaMoePreDispatchInferDtype( + const paddle::DataType& x_dtype, + const paddle::DataType& topk_idx_dtype, + const paddle::DataType& topk_weights_dtype, + const paddle::DataType& buf_x_dtype, + const paddle::DataType& buf_x_sf_dtype, + const paddle::DataType& buf_topk_idx_dtype, + const paddle::DataType& buf_topk_weights_dtype) { + return { + buf_x_dtype, buf_x_sf_dtype, buf_topk_idx_dtype, buf_topk_weights_dtype}; +} + +std::vector> MegaMoePreDispatchInferShape( + const std::vector& x_shape, + const std::vector& topk_idx_shape, + const std::vector& topk_weights_shape, + const std::vector& buf_x_shape, + const std::vector& buf_x_sf_shape, + const std::vector& buf_topk_idx_shape, + const std::vector& buf_topk_weights_shape) { + return { + buf_x_shape, buf_x_sf_shape, buf_topk_idx_shape, buf_topk_weights_shape}; +} + +PD_BUILD_STATIC_OP(mega_moe_pre_dispatch) + .Inputs({"x", + "topk_idx", + "topk_weights", + "buf_x", + "buf_x_sf", + "buf_topk_idx", + "buf_topk_weights"}) + .Outputs({"buf_x_out", + "buf_x_sf_out", + "buf_topk_idx_out", + "buf_topk_weights_out"}) + .Attrs({"num_max_tokens_per_rank: int64_t", "group_size: int64_t"}) + .SetInplaceMap({{"buf_x", "buf_x_out"}, + {"buf_x_sf", "buf_x_sf_out"}, + {"buf_topk_idx", "buf_topk_idx_out"}, + {"buf_topk_weights", "buf_topk_weights_out"}}) + .SetKernelFn(PD_KERNEL(MegaMoePreDispatch)) + .SetInferShapeFn(PD_INFER_SHAPE(MegaMoePreDispatchInferShape)) + .SetInferDtypeFn(PD_INFER_DTYPE(MegaMoePreDispatchInferDtype)); diff --git a/custom_ops/setup_ops.py b/custom_ops/setup_ops.py index 1e700f87634..866d7bc6351 100644 --- a/custom_ops/setup_ops.py +++ b/custom_ops/setup_ops.py @@ -339,6 +339,7 @@ def find_end_files(directory, end_str): "gpu_ops/gelu_tanh.cu", "gpu_ops/reasoning_phase_token_constraint.cu", "gpu_ops/get_attn_mask_q.cu", + "gpu_ops/mega_moe_pre_dispatch.cu", ] sm_versions = get_sm_version(archs) # Some kernels in this file require SM75+ instructions. Exclude them when building SM70 (V100). diff --git a/fastdeploy/config.py b/fastdeploy/config.py index a684442d09e..20fa4b48ad3 100644 --- a/fastdeploy/config.py +++ b/fastdeploy/config.py @@ -646,6 +646,7 @@ def __init__( self.enable_expert_parallel = False self.enable_chunked_moe = False self.chunked_moe_size = 256 + self.enable_mega_moe = False self.local_data_parallel_id = 0 # Engine worker queue port diff --git a/fastdeploy/engine/args_utils.py b/fastdeploy/engine/args_utils.py index d8541f96d80..449f4bb2eed 100644 --- a/fastdeploy/engine/args_utils.py +++ b/fastdeploy/engine/args_utils.py @@ -327,6 +327,11 @@ class EngineArgs: Whether use chunked moe. """ + enable_mega_moe: bool = False + """ + Whether use MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear. + """ + chunked_moe_size: int = 256 """ Chunk size of moe input. @@ -1075,6 +1080,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: default=EngineArgs.enable_chunked_moe, help="Use chunked moe.", ) + parallel_group.add_argument( + "--enable-mega-moe", + action="store_true", + default=EngineArgs.enable_mega_moe, + help="Use MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear.", + ) parallel_group.add_argument( "--chunked-moe-size", type=int, diff --git a/fastdeploy/engine/engine.py b/fastdeploy/engine/engine.py index 68bfac79dfd..85d7459815f 100644 --- a/fastdeploy/engine/engine.py +++ b/fastdeploy/engine/engine.py @@ -641,6 +641,7 @@ def _start_worker_service(self): worker_store_true_flag = { "enable_expert_parallel": self.cfg.parallel_config.enable_expert_parallel, "enable_chunked_moe": self.cfg.parallel_config.enable_chunked_moe, + "enable_mega_moe": self.cfg.parallel_config.enable_mega_moe, "enable_prefix_caching": self.cfg.cache_config.enable_prefix_caching, "enable_chunked_prefill": self.cfg.cache_config.enable_chunked_prefill, "do_profile": self.do_profile, diff --git a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py index d4ab8f5f59b..3ad7e2738e4 100644 --- a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py +++ b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py @@ -87,7 +87,7 @@ def init_flash_attn_version(): sm_version = get_sm_version() if sm_version >= 100: try: - paddle.compat.enable_torch_proxy(scope={"cutlass"}) + paddle.enable_compat(scope={"cutlass"}) from flash_mask.cute.interface import flashmask_attention as fa4 global flashmask_attention_v4 diff --git a/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py b/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py index c0df764c07c..5a2675ec07c 100644 --- a/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py +++ b/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py @@ -806,7 +806,7 @@ def enable_batch_invariant_mode(): return if hasattr(paddle, "compat") and hasattr(paddle.compat, "enable_torch_proxy"): - paddle.compat.enable_torch_proxy() + paddle.enable_compat() # TODO(liujundong): Enabling torch proxy here has a global effect. # Do NOT call this function from module import time, # otherwise it may affect other test cases during pytest collection. diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index f0b563bf95a..589ee0d415e 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/fastdeploy/model_executor/layers/moe/ep.py @@ -41,7 +41,7 @@ def load_deep_ep() -> ModuleType: try: if envs.FD_USE_PFCC_DEEP_EP: # Enable torch proxy before importing deep_ep (required by PFCC/PaddleFleet variants) - paddle.compat.enable_torch_proxy(scope={"deep_ep"}) + paddle.enable_compat(scope={"deep_ep"}) try: import paddlefleet.ops.deep_ep as deep_ep # type: ignore @@ -521,7 +521,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): expert_in_rank_num_list=expert_in_rank_num_list, tokens_per_expert_stats_list=tokens_per_expert_stats_list, bias=layer.gate_correction_bias, - moe_topk=self.top_k, + moe_topk=layer.top_k, apply_norm_weight=True, enable_softmax_top_k_fused=False, redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1, @@ -548,7 +548,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select( gate_out, layer.gate_correction_bias, - self.top_k, + layer.top_k, True, False, ) @@ -777,3 +777,21 @@ def combine(self, ffn_out, topk_idx, topk_weights, handle, **kwargs): combine_hook() return combined_hidden_states + + +class FakeEPRunner: + """ """ + + def __init__(self, *args, **kwargs): + pass + + def dispatch(self, *args, **kwargs): + """ """ + pass + + def combine(self, *args, **kwargs): + """ """ + pass + + def clean_low_latency_buffer(self): + pass diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index 1c1bb85f71c..105a1f0167e 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -23,8 +23,10 @@ from paddleformers.utils.log import logger import fastdeploy -from fastdeploy.model_executor.layers.moe.ep import deep_ep +from fastdeploy.model_executor.layers.moe.ep import EPRunner, FakeEPRunner, deep_ep from fastdeploy.model_executor.layers.quantization.fp8_utils import ( + _interleave_weights, + _transpose_sf_for_utccp, deep_gemm, paddlefleet_ops, ) @@ -32,10 +34,18 @@ from fastdeploy.model_executor.ops.gpu import ( count_tokens_per_expert_func, depermute_prefill_combine, + mega_moe_pre_dispatch, prefill_permute_to_masked_gemm, ) +from fastdeploy.model_executor.utils import ( + TensorTracker, + free_tensor, + get_sm_version, + set_weight_attrs, + weight_fully_copied, +) from fastdeploy.platforms import current_platform -from fastdeploy.utils import register_custom_python_op +from fastdeploy.utils import ceil_div, register_custom_python_op, singleton from fastdeploy.worker.tbo import let_another_thread_run from .fused_moe_backend_base import MoEMethodBase @@ -868,3 +878,448 @@ def apply_tp( 1.0, ) return tmp_ffn_out + + +@singleton +class MegaMoEBuffer: + """ + A wrapper class for DeepEP engine. + Manages buffer lifecycle based on role and phase. + """ + + def __init__( + self, + ep_group, + num_experts: int, + num_max_tokens_per_rank: int, + top_k: int, + hidden_size: int, + moe_intermediate_size: int, + ): + self.buffer = deep_gemm.get_symm_buffer_for_mega_moe( + ep_group, + num_experts, + num_max_tokens_per_rank, + top_k, + hidden_size, + moe_intermediate_size, + ) + + +class DeepGemmMegaMoEMethod(DeepGemmFusedMoeMethod): + def __init__(self, quant_config): + if not get_sm_version() >= 100: + raise ValueError("MegaMoE now only support sm100+ devices.") + super().__init__(quant_config) + self.added_scale_attrs = ["up_gate_proj_weight_scale", "down_proj_weight_scale"] + self.quant_config.deepgemm_scale_ue8m0 = True + self.gran_k = 32 + + def create_weights(self, layer: nn.Layer, **extra_weight_attrs): + """ + Triton MoE create weight process. + """ + logger.info("MegaMoE create_weights...") + self.model_format = extra_weight_attrs.get("model_format") + self.up_gate_proj_quant_weight_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size * 2, + layer.hidden_size, + ] + self.down_proj_quant_weight_shape = [ + layer.num_local_experts, + layer.hidden_size, + layer.moe_intermediate_size, + ] + self.up_gate_proj_pretranspose_weight_shape = [ + layer.num_local_experts, + layer.hidden_size, + layer.moe_intermediate_size * 2, + ] + self.down_proj_pretranspose_weight_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size, + layer.hidden_size, + ] + if self.model_format != "torch": + self.up_gate_proj_bf16_weight_shape = self.up_gate_proj_pretranspose_weight_shape + self.down_proj_bf16_weight_shape = self.down_proj_pretranspose_weight_shape + else: + self.up_gate_proj_bf16_weight_shape = self.up_gate_proj_quant_weight_shape + self.down_proj_bf16_weight_shape = self.down_proj_quant_weight_shape + self.up_gate_proj_packed_weight_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size * 2, + layer.hidden_size // 2, # 4-bit packing + ] + self.down_proj_packed_weight_shape = [ + layer.num_local_experts, + layer.hidden_size, + layer.moe_intermediate_size // 2, # 4-bit packing + ] + up_num_scales = ceil_div(layer.hidden_size, self.gran_k) + down_num_scales = ceil_div(layer.moe_intermediate_size, self.gran_k) + self.up_gate_proj_scale_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size * 2, + (up_num_scales + 3) // 4, + ] + self.down_proj_scale_shape = [ + layer.num_local_experts, + layer.hidden_size, + (down_num_scales + 3) // 4, + ] + self.up_gate_proj_weight_shape = self.up_gate_proj_quant_weight_shape + self.down_proj_weight_shape = self.down_proj_quant_weight_shape + + if self.quant_config.is_checkpoint_bf16 and layer.fd_config.load_config.load_choices == "default_v1": + if self.model_format != "torch": + up_gate_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.up_gate_proj_bf16_weight_shape, output_dim=True), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 1, "down": 0, "up": 1}, + } + down_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.down_proj_bf16_weight_shape, output_dim=False), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 1, "down": 0, "up": 1}, + } + else: + up_gate_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.up_gate_proj_bf16_weight_shape, output_dim=False), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 0, "down": 1, "up": 0}, + } + down_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.down_proj_bf16_weight_shape, output_dim=True), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 0, "down": 1, "up": 0}, + } + layer.up_gate_proj_weight = layer.create_parameter( + shape=self.up_gate_proj_bf16_weight_shape, + dtype=layer.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ) + + layer.down_proj_weight = layer.create_parameter( + shape=self.down_proj_bf16_weight_shape, + dtype=layer.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ) + + set_weight_attrs( + layer.up_gate_proj_weight, + up_gate_proj_attrs, + ) + set_weight_attrs( + layer.down_proj_weight, + down_proj_attrs, + ) + else: + # offline quant + self.up_gate_proj_weight_shape = self.up_gate_proj_packed_weight_shape + self.down_proj_weight_shape = self.down_proj_packed_weight_shape + up_gate_proj_attrs = {} + down_proj_attrs = {} + + self.weight_dtype = paddle.int8 + up_gate_proj_weight_name = self.added_weight_attrs[0] + down_proj_weight_name = self.added_weight_attrs[1] + up_gate_proj_scale_name = self.added_scale_attrs[0] + down_proj_scale_name = self.added_scale_attrs[1] + + setattr( + layer, + up_gate_proj_weight_name, + layer.create_parameter( + shape=self.up_gate_proj_packed_weight_shape, + dtype=self.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + setattr( + layer, + down_proj_weight_name, + layer.create_parameter( + shape=self.down_proj_packed_weight_shape, + dtype=self.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + # weight_scale + setattr( + layer, + up_gate_proj_scale_name, + layer.create_parameter( + shape=self.up_gate_proj_scale_shape, + dtype="int32", + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + setattr( + layer, + down_proj_scale_name, + layer.create_parameter( + shape=self.down_proj_scale_shape, + dtype="int32", + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + + set_weight_attrs( + getattr(layer, up_gate_proj_weight_name), + up_gate_proj_attrs, + ) + set_weight_attrs( + getattr(layer, up_gate_proj_scale_name), + up_gate_proj_attrs, + ) + + set_weight_attrs( + getattr(layer, down_proj_weight_name), + down_proj_attrs, + ) + set_weight_attrs( + getattr(layer, down_proj_scale_name), + down_proj_attrs, + ) + + def init_ep(self, layer: nn.Layer) -> None: + logger.info("Use MegaMoE backend") + if layer.ep_size <= 1: + raise ValueError( + "Ep size must be greater than 1 when use MegaMoE backend. Please set --enable-expert-parallel" + ) + + config = layer.fd_config + splitwise_role = config.scheduler_config.splitwise_role + + if splitwise_role == "mixed" or splitwise_role == "prefill": + self.num_max_tokens_per_rank = config.scheduler_config.max_num_batched_tokens + elif splitwise_role == "decode": + num_spec_tokens = config.speculative_config.num_speculative_tokens + self.num_max_tokens_per_rank = config.scheduler_config.max_num_seqs * (num_spec_tokens + 1) + else: + raise ValueError(f"Unsupported splitwise role: {splitwise_role}") + + self.mega_moe_buffer = MegaMoEBuffer( + layer.fd_config.parallel_config.ep_group, + layer.num_experts, + self.num_max_tokens_per_rank, + layer.top_k, + layer.hidden_size, + layer.moe_intermediate_size, + ).buffer + self.num_max_tokens_per_rank = self.mega_moe_buffer.num_max_tokens_per_rank + self.cumulative_local_expert_recv_stats = paddle.zeros((layer.num_local_experts,), dtype=paddle.int32) + + self.ep_prefill_runner = FakeEPRunner() + self.ep_decoder_runner = FakeEPRunner() + + def process_weights_after_loading(self, layer): + def cast_grouped_weights_to_fp4(bf16_weights: paddle.Tensor): + num_groups, n, k = bf16_weights.shape + w = paddle.empty((num_groups, n, k // 2), dtype=paddle.int8) + w_sf = paddle.empty((num_groups, n, k // self.gran_k), dtype=paddle.float32) + for i in range(num_groups): + w[i], w_sf[i] = deep_gemm.per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=self.gran_k) + w = w.contiguous() + w_sf = w_sf.contiguous() + + # pack four scales into one and transform to specific stride. + # for example: + # shape: [48, 6144, 224] -> [48, 6144, 56] + # stride: [1376256, 224, 1] -> [344064, 1, 6144] + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, self.gran_k), num_groups) + return w, w_sf + + def _process_quantize_mega_moe(weight_type): + weight_idx = 0 if weight_type == "gate_up" else 1 + weight_name = self.added_weight_attrs[weight_idx] + scale_name = self.added_scale_attrs[weight_idx] + weight = getattr(layer, weight_name) + if not hasattr(weight, "tensor_track") or weight.tensor_track is None: + return + + if self.model_format != "torch": + weight = weight.transpose([0, 2, 1]).contiguous() + + expected_weight_shape = ( + self.up_gate_proj_quant_weight_shape if weight_type == "gate_up" else self.down_proj_quant_weight_shape + ) + expected_packed_shape = ( + self.up_gate_proj_packed_weight_shape + if weight_type == "gate_up" + else self.down_proj_packed_weight_shape + ) + expected_scale_shape = ( + self.up_gate_proj_scale_shape if weight_type == "gate_up" else self.down_proj_scale_shape + ) + + if list(weight.shape) != list(expected_weight_shape): + raise ValueError( + f"MegaMoE {weight_type} BF16 weight shape mismatch for {weight_name}: " + f"got {list(weight.shape)}, expected {list(expected_weight_shape)}" + ) + if weight.dtype != paddle.bfloat16: + weight = weight.astype(paddle.bfloat16) + weight = weight.contiguous() + + weight_quantized, scale = cast_grouped_weights_to_fp4(weight) + if weight_type == "gate_up": + weight_quantized, scale = _interleave_weights((weight_quantized, scale)) + scale = _transpose_sf_for_utccp(scale) + + if list(weight_quantized.shape) != list(expected_packed_shape): + raise ValueError( + f"MegaMoE {weight_type} packed weight shape mismatch: " + f"got {list(weight_quantized.shape)}, expected {list(expected_packed_shape)}" + ) + if list(scale.shape) != list(expected_scale_shape): + raise ValueError( + f"MegaMoE {weight_type} scale shape mismatch: " + f"got {list(scale.shape)}, expected {list(expected_scale_shape)}" + ) + if weight_quantized.dtype != paddle.int8: + raise ValueError(f"MegaMoE {weight_type} packed weight dtype mismatch: got {weight_quantized.dtype}") + if scale.dtype != paddle.int32: + raise ValueError(f"MegaMoE {weight_type} scale dtype mismatch: got {scale.dtype}") + + free_tensor(getattr(layer, weight_name)) + setattr( + layer, + weight_name, + layer.create_parameter( + shape=weight_quantized.shape, + dtype=paddle.int8, + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + setattr( + layer, + scale_name, + layer.create_parameter( + shape=scale.shape, + dtype=paddle.int32, + default_initializer=paddle.nn.initializer.Constant(0), + ).as_strided(scale.shape, scale.stride()), + ) + getattr(layer, weight_name).copy_(weight_quantized, False) + getattr(layer, scale_name).copy_(scale, False) + + if not self.quant_config.is_checkpoint_bf16: + return + if hasattr(layer, "up_gate_proj_weight") and weight_fully_copied(layer.up_gate_proj_weight): + _process_quantize_mega_moe("gate_up") + if hasattr(layer, "down_proj_weight") and weight_fully_copied(layer.down_proj_weight): + _process_quantize_mega_moe("down") + + def process_prequanted_weights(self, layer: nn.Layer, state_dict, is_rearrange: bool = False): + """ + Paddle cutlass process prequanted weights. + """ + up_gate_proj_expert_weight_key = layer.weight_key_map.get("up_gate_proj_expert_weight_key", None) + down_proj_expert_weight_key = layer.weight_key_map.get("down_proj_expert_weight_key", None) + up_gate_proj_expert_weight_scale_key = layer.weight_key_map.get("up_gate_proj_expert_weight_scale_key", None) + down_proj_expert_weight_scale_key = layer.weight_key_map.get("down_proj_expert_weight_scale_key", None) + + up_gate_proj_weights, down_proj_weights, logical_expert_ids, _ = layer.load_experts_weight( + state_dict, up_gate_proj_expert_weight_key, down_proj_expert_weight_key, is_rearrange + ) + # self.check(layer, up_gate_proj_weights, down_proj_weights) + up_gate_proj_weight_scale = [] + down_proj_weight_scale = [] + + if isinstance(state_dict, list): + state_dict = dict(state_dict) + + for expert_idx in logical_expert_ids: + up_gate_proj_expert_weight_scale_key_name = up_gate_proj_expert_weight_scale_key.format(expert_idx) + down_proj_expert_weight_scale_key_name = down_proj_expert_weight_scale_key.format(expert_idx) + + up_gate_weight_scale = get_tensor( + ( + state_dict.pop(up_gate_proj_expert_weight_scale_key_name) + if up_gate_proj_expert_weight_scale_key_name in state_dict + else up_gate_proj_expert_weight_scale_key_name + ), + layer.fd_config.model_config.model, + ) + down_weight_scale = get_tensor( + ( + state_dict.pop(down_proj_expert_weight_scale_key_name) + if down_proj_expert_weight_scale_key_name in state_dict + else down_proj_expert_weight_scale_key_name + ), + layer.fd_config.model_config.model, + ) + + up_gate_proj_weight_scale.append(up_gate_weight_scale) + down_proj_weight_scale.append(down_weight_scale) + + up_gate_proj_weight = paddle.stack(up_gate_proj_weights, axis=0) + down_proj_weight = paddle.stack(down_proj_weights, axis=0) + up_gate_proj_weight_scale = paddle.stack(up_gate_proj_weight_scale, axis=0).transpose([0, 2, 1]) + down_proj_weight_scale = paddle.stack(down_proj_weight_scale, axis=0).transpose([0, 2, 1]) + + name_tensor_map = { + self.added_weight_attrs[0]: up_gate_proj_weight, + self.added_weight_attrs[1]: down_proj_weight, + self.added_scale_attrs[0]: up_gate_proj_weight_scale, + self.added_scale_attrs[1]: down_proj_weight_scale, + } + for name, tensor in name_tensor_map.items(): + getattr(layer, name).data = tensor + + def apply_ep_prefill(self, layer, x, gate, topk_ids_hookfunc, shared_experts): + return self.apply_mega_moe(layer, x, gate, topk_ids_hookfunc, shared_experts) + + def apply_ep_decode(self, layer, x, gate, topk_ids_hookfunc, shared_experts): + return self.apply_mega_moe(layer, x, gate, topk_ids_hookfunc, shared_experts) + + def apply_mega_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts): + hidden_size = layer.hidden_size + num_tokens = x.shape[0] + + gate_out = gate(x).cast("float32") + + # 1. Select topk experts and weights. + topk_idx, topk_weights = EPRunner.moe_select(None, layer, gate_out) + + buffer_capacity = self.mega_moe_buffer.x.shape[0] + if num_tokens > buffer_capacity: + raise ValueError(f"MegaMoE buffer capacity exceeded: num_tokens={num_tokens}, capacity={buffer_capacity}") + + # copy x, topk_idx, topk_weights to mega_moe_buffer and quantization. + mega_moe_pre_dispatch( + x, + topk_idx, + topk_weights, + self.mega_moe_buffer.x, + self.mega_moe_buffer.x_sf, + self.mega_moe_buffer.topk_idx, + self.mega_moe_buffer.topk_weights, + self.num_max_tokens_per_rank, + self.gran_k, # group_size + ) + + l1_weight = getattr(layer, self.added_weight_attrs[0]) + l1_scale = getattr(layer, self.added_scale_attrs[0]) + l2_weight = getattr(layer, self.added_weight_attrs[1]) + l2_scale = getattr(layer, self.added_scale_attrs[1]) + y = paddle.empty((num_tokens, hidden_size), dtype=paddle.bfloat16) + + swiglu_limit = getattr(layer.fd_config.model_config, "swiglu_limit", 10) + deep_gemm.fp8_fp4_mega_moe( + y, + (l1_weight, l1_scale), + (l2_weight, l2_scale), + self.mega_moe_buffer, + cumulative_local_expert_recv_stats=self.cumulative_local_expert_recv_stats, + recipe=(1, 1, self.gran_k), + activation="swiglu", + activation_clamp=swiglu_limit, + fast_math=True, + ) + + return y diff --git a/fastdeploy/model_executor/layers/quantization/__init__.py b/fastdeploy/model_executor/layers/quantization/__init__.py index 678873f76dc..2660470eb9c 100644 --- a/fastdeploy/model_executor/layers/quantization/__init__.py +++ b/fastdeploy/model_executor/layers/quantization/__init__.py @@ -30,6 +30,7 @@ "weight_only", "block_wise_fp8", "w4afp8", + "wfp4afp8", "w8a8", "w4a8", "wfp8afp8", @@ -68,10 +69,39 @@ def _is_full_quantization_config(quantization_dict): return False +def _is_mega_moe_quantization_config(quantization_config): + return isinstance(quantization_config, dict) and quantization_config.get("moe_quant_type") == "wfp4afp8" + + +def _get_mega_moe_quantization_config(): + return { + "quantization": "mix_quant", + "kv_cache_quant_type": "block_wise_fp8", + "dense_quant_type": "block_wise_fp8", + "moe_quant_type": "wfp4afp8", + "is_quantized": False, + } + + def parse_quant_config(args, model_config, is_ernie, is_v1_loader): if args.quantization is not None and isinstance(args.quantization, str): args.quantization = parse_quantization(args.quantization) + enable_mega_moe = getattr(args, "enable_mega_moe", False) + if enable_mega_moe: + mega_moe_quantization_config = _get_mega_moe_quantization_config() + + if args.quantization is None and model_config.quantization_config is None: + args.quantization = mega_moe_quantization_config + if args.quantization is not None and not _is_mega_moe_quantization_config(args.quantization): + raise ValueError("--enable-mega-moe requires moe_quant_type=wfp4afp8.") + if model_config.quantization_config is not None and not _is_mega_moe_quantization_config( + model_config.quantization_config + ): + raise ValueError( + "--enable-mega-moe conflicts with model quantization_config. It requires moe_quant_type=wfp4afp8." + ) + # Determine whether CLI --quantization is a simple method name or a full JSON quantization_config cli_quantization = args.quantization cli_is_full_config = ( @@ -213,6 +243,7 @@ def get_quantization_config(quantization: str) -> Type[QuantConfigBase]: from .w4afp8 import W4AFP8Config from .w8a8 import W8A8Config from .weight_only import WeightOnlyConfig, WINT4Config, WINT8Config + from .wfp4afp8 import WFP4AFP8Config from .wfp8afp8 import WFP8AFP8Config from .wint2 import WINT2Config @@ -229,6 +260,7 @@ def get_quantization_config(quantization: str) -> Type[QuantConfigBase]: "w8a8": W8A8Config, "w4a8": W4A8Config, "wfp8afp8": WFP8AFP8Config, + "wfp4afp8": WFP4AFP8Config, "tensor_wise_fp8": TensorWiseFP8Config, "kvcache": KvCacheQuantConfig, "mix_quant": MixQuantConfig, diff --git a/fastdeploy/model_executor/layers/quantization/fp8_utils.py b/fastdeploy/model_executor/layers/quantization/fp8_utils.py index 99c2dc8ebf5..ccc8726efdc 100644 --- a/fastdeploy/model_executor/layers/quantization/fp8_utils.py +++ b/fastdeploy/model_executor/layers/quantization/fp8_utils.py @@ -68,7 +68,7 @@ def load_deep_gemm(): if current_platform.is_cuda(): if get_sm_version() >= 100: # SM100 should use PFCC DeepGemm - paddle.compat.enable_torch_proxy(scope={"deep_gemm"}) + paddle.enable_compat(scope={"deep_gemm"}) try: import logging @@ -260,3 +260,31 @@ def fused_stack_transpose_quant(expert_weight_list, use_ue8m0=False): raise RuntimeError("'fuse_stack_transpose_fp8_quant' is not available in the current paddlefleet_ops.") return w, scale + + +def _interleave_weights(l1_weights): + # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up] + def interleave(t, gran: int = 8) -> paddle.Tensor: + g, n, *rest = t.shape + half = n // 2 + gate = t[:, :half].reshape(g, half // gran, gran, *rest) + up = t[:, half:].reshape(g, half // gran, gran, *rest) + return paddle.stack([gate, up], dim=2).reshape(g, n, *rest).contiguous() + + return interleave(l1_weights[0]), interleave(l1_weights[1]) + + +def _transpose_sf_for_utccp(sf: paddle.Tensor) -> paddle.Tensor: + num_groups, mn, packed_sf_k = sf.shape + assert sf.dtype == paddle.int and mn % 128 == 0 + # sf is MN-major: strides [mn*packed_sf_k, 1, mn] + # We need to do the 4x32 transpose in data while preserving MN-major strides + sf_c = sf.contiguous() # make C-contiguous for reshape/transpose + result_c = ( + sf_c.reshape(num_groups, -1, 4, 32, packed_sf_k) + .transpose(2, 3) + .reshape(num_groups, mn, packed_sf_k) + .contiguous() + ) + # Convert back to MN-major layout: transpose last two dims, make contiguous, transpose back + return result_c.transpose(1, 2).contiguous().transpose(1, 2) diff --git a/fastdeploy/model_executor/layers/quantization/mxfp4.py b/fastdeploy/model_executor/layers/quantization/mxfp4.py index 9fa02866210..8732b1cc4d1 100644 --- a/fastdeploy/model_executor/layers/quantization/mxfp4.py +++ b/fastdeploy/model_executor/layers/quantization/mxfp4.py @@ -35,7 +35,7 @@ from ..moe import FusedMoE from .quant_base import QuantConfigBase, QuantMethodBase -paddle.compat.enable_torch_proxy(scope={"flashinfer"}) +paddle.enable_compat(scope={"flashinfer"}) logger = get_logger("config", "config.log") diff --git a/fastdeploy/model_executor/layers/quantization/nvfp4.py b/fastdeploy/model_executor/layers/quantization/nvfp4.py index f8293bf69ed..7a22b3e1260 100644 --- a/fastdeploy/model_executor/layers/quantization/nvfp4.py +++ b/fastdeploy/model_executor/layers/quantization/nvfp4.py @@ -38,7 +38,7 @@ # Only import flashinfer on supported GPUs (B卡) if is_nvfp4_supported(): - paddle.compat.enable_torch_proxy(scope={"flashinfer"}) + paddle.enable_compat(scope={"flashinfer"}) from flashinfer import fp4_quantize, mm_fp4 from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe diff --git a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py new file mode 100644 index 00000000000..4c8a462d476 --- /dev/null +++ b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py @@ -0,0 +1,63 @@ +""" +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" + +from typing import Optional + +from paddleformers.utils.log import logger + +from ..moe import FusedMoE +from .quant_base import QuantConfigBase, QuantMethodBase + +QUANT_SCALING_FACTOR = 6 + + +class WFP4AFP8Config(QuantConfigBase): + """ + quantization config for weight 4bits and activation fp8 + """ + + def __init__(self, weight_scale_dict, act_scale_dict, is_permuted, is_quantized) -> None: + super().__init__() + self.weight_scale_dict = weight_scale_dict + self.act_scale_dict = act_scale_dict + self.quant_max_bound = 6 + self.quant_min_bound = -6 + self.quant_round_type = 1 + self.is_permuted = is_permuted + self.is_quantized = is_quantized + self.is_checkpoint_bf16 = not is_quantized + + def name(self) -> str: + return "wfp4afp8" + + @classmethod + def from_config(cls, config: dict) -> "WFP4AFP8Config": + weight_scale_dict = config.get("weight_scale_dict", None) + act_scale_dict = config.get("act_scale_dict", None) + is_permuted = config.get("is_permuted", True) + is_quantized = config.get("is_quantized", False) + return cls(weight_scale_dict, act_scale_dict, is_permuted, is_quantized) + + def get_quant_method(self, layer) -> Optional[QuantMethodBase]: + logger.debug("Currently only support DeepGEMMMegaMoE for wfp4afp8") + if isinstance(layer, FusedMoE): + from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import ( + DeepGemmMegaMoEMethod, + ) + + return DeepGemmMegaMoEMethod(self) + else: + raise NotImplementedError(f"wfp4afp8 quant method not supported for {type(layer)}") diff --git a/fastdeploy/model_executor/load_weight_utils.py b/fastdeploy/model_executor/load_weight_utils.py index 2f181d7740a..5f19ec5a3f3 100644 --- a/fastdeploy/model_executor/load_weight_utils.py +++ b/fastdeploy/model_executor/load_weight_utils.py @@ -135,7 +135,7 @@ def get_weight_iterator(model_path: str, fd_config: Optional[FDConfig] = None): disable_mmap=extra_config.get("disable_mmap", False), ) else: - if is_layers_are_grouped or (parallel_config is not None and parallel_config.tensor_parallel_size == 1): + if is_layers_are_grouped: weights_iterator = safetensors_weights_iterator(files_list) else: weights_iterator = safetensors_weights_iterator_ordered(ordered_weight_map) diff --git a/fastdeploy/worker/worker_process.py b/fastdeploy/worker/worker_process.py index e744c4c7e1f..810115492eb 100644 --- a/fastdeploy/worker/worker_process.py +++ b/fastdeploy/worker/worker_process.py @@ -861,6 +861,12 @@ def parse_args(): action="store_true", help="enable chunked moe", ) + parser.add_argument( + "--enable_mega_moe", + action="store_true", + dest="enable_mega_moe", + help="enable MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear", + ) parser.add_argument( "--chunked_moe_size", type=int, @@ -1281,7 +1287,7 @@ def run_worker_proc() -> None: # Enable batch-invariant mode for deterministic inference. # This must happen AFTER worker creation but BEFORE model loading, - # because enable_batch_invariant_mode() calls paddle.compat.enable_torch_proxy() + # because enable_batch_invariant_mode() calls paddle.enable_compat() # which makes torch appear available via proxy. If called before worker creation, # the gpu_model_runner import chain (ernie4_5_vl_processor → paddleformers → # transformers) will fail when transformers tries to query torch metadata. diff --git a/test.py b/test.py new file mode 100644 index 00000000000..753a76d1589 --- /dev/null +++ b/test.py @@ -0,0 +1,74 @@ +from pathlib import Path +from safetensors import safe_open +import json +import re + +def natural_key(s: str): + return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", s)] + +def layers_are_grouped(keys): + + seen = set() + current_layer = None + + for k in keys: + m = re.search(r"layers\.(\d+)", k) + if not m: + continue + + layer = int(m.group(1)) + + if layer != current_layer: + if layer in seen: + return False + seen.add(layer) + current_layer = layer + + return True + +def values_are_naturally_ordered(values): + """Check if values are sorted in natural order.""" + return list(values) == sorted(values, key=natural_key) + + +def get_all_weights_file(model_path: str): + """ + get_all_safetensors + """ + model_path = Path(model_path) + use_safetensors = True + files_list = [str(file) for file in model_path.glob("*.pdparams") if file.name != "scheduler.pdparams"] + if len(files_list) > 0: + ordered_weight_map = {} + use_safetensors = False + # dont care about the order of the files + return files_list, {}, use_safetensors, False + else: + safe_model_path = model_path / "model.safetensors" + if safe_model_path.exists(): + with safe_open(safe_model_path, framework="np", device="cpu") as f: + key_name_list = sorted(f.keys(), key=natural_key) + ordered_weight_map = {key: "model.safetensors" for key in key_name_list} + is_layers_are_grouped = True + files_list = [str(safe_model_path)] + return files_list, ordered_weight_map, use_safetensors, is_layers_are_grouped + else: + index_file = model_path / "model.safetensors.index.json" + with index_file.open("r") as f: + weight_map = json.load(f)["weight_map"] + keys = list(weight_map.keys()) + values = list(weight_map.values()) + is_keys_orders = layers_are_grouped(keys) + is_values_naturally_ordered = values_are_naturally_ordered(values) + is_layers_are_grouped = is_keys_orders and is_values_naturally_ordered + ordered_weight_map = { + key: str(model_path / weight_map[key]) for key in sorted(weight_map.keys(), key=natural_key) + } + weight_files_in_index = {str(model_path / weight_map[name]) for name in weight_map} + files_list = sorted(weight_files_in_index) + return files_list, ordered_weight_map, use_safetensors, is_layers_are_grouped + + +files_list, ordered_weight_map, use_safetensors, is_layers_are_grouped = get_all_weights_file("/root/paddlejob/share-storage/gpfs/system-public/dangweichong/models/eb5_v4p1_v4p2_merge_0413v1_step120") + +print("is_layers_are_grouped: ", is_layers_are_grouped) \ No newline at end of file diff --git a/tests/engine/test_engine.py b/tests/engine/test_engine.py index 001140f1993..b0e1ec6bf46 100644 --- a/tests/engine/test_engine.py +++ b/tests/engine/test_engine.py @@ -36,6 +36,7 @@ def _make_cfg(**ov): pc = ns(tensor_parallel_size=1, tensor_parallel_rank=0, device_ids="0", data_parallel_size=1) pc.expert_parallel_size, pc.chunked_moe_size, pc.engine_worker_queue_port = 1, 0, [6778] pc.enable_expert_parallel = pc.enable_chunked_moe = pc.disable_custom_all_reduce = False + pc.enable_mega_moe = False pc.use_internode_ll_two_stage = pc.disable_sequence_parallel_moe = False pc.shutdown_comm_group_if_worker_idle = False pc.ep_prefill_use_worst_num_tokens = False diff --git a/tests/model_executor/test_ep.py b/tests/model_executor/test_ep.py index 950314a5e57..8c7d2f01909 100644 --- a/tests/model_executor/test_ep.py +++ b/tests/model_executor/test_ep.py @@ -496,6 +496,7 @@ def fake_topk_select(*_args, **_kwargs): redundant_table_manger=None, topk_method="aux", gate_correction_bias=None, + top_k=2, ) gate_out = paddle.randn([1, 4], dtype="float32") diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py new file mode 100644 index 00000000000..9d80c7ddee2 --- /dev/null +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -0,0 +1,171 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from dataclasses import dataclass + +import numpy as np +import paddle + +from fastdeploy.model_executor.ops.gpu import mega_moe_pre_dispatch + + +@dataclass +class FakeBuffer: + x: paddle.Tensor + x_sf: paddle.Tensor + topk_idx: paddle.Tensor + topk_weights: paddle.Tensor + + +def ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def align(x: int, y: int) -> int: + return ceil_div(x, y) * y + + +def ceil_to_ue8m0(x: paddle.Tensor): + bits = x.abs().astype("float32").view(paddle.int32) + mask_ff = paddle.to_tensor(0xFF, dtype=paddle.int32) + mask_mantissa = paddle.to_tensor(0x7FFFFF, dtype=paddle.int32) + exp = ((bits >> 23) & mask_ff) + ((bits & mask_mantissa) != 0).astype("int32") + return (exp.clip(1, 254) << 23).view(paddle.float32) + + +def pack_ue8m0_to_int(x: paddle.Tensor): + assert x.dtype == paddle.float32 and x.shape[-1] % 4 == 0 + x_bits = x.view(paddle.int32) + mantissa_mask = paddle.to_tensor((1 << 23) - 1, dtype=paddle.int32) + assert bool(((x_bits & mantissa_mask) == 0).all()) + return (x_bits >> 23).astype(paddle.uint8).view(paddle.int32) + + +def per_token_cast_to_fp8( + x: paddle.Tensor, + use_ue8m0: bool, + gran_k: int = 128, + use_packed_ue8m0: bool = False, +): + assert len(x.shape) == 2 + m, n = x.shape + padded_n = align(n, gran_k) + x_padded = paddle.zeros((m, padded_n), dtype=x.dtype) + x_padded[:, :n] = x + x_view = x_padded.reshape([m, padded_n // gran_k, gran_k]) + x_amax = x_view.abs().astype("float32").amax(axis=2).reshape([m, padded_n // gran_k]).clip(min=1e-4) + sf = x_amax / 448.0 + sf = ceil_to_ue8m0(sf) if use_ue8m0 else sf + x_fp8 = (x_view * (1.0 / sf.unsqueeze(2))).astype(paddle.float8_e4m3fn).reshape([m, padded_n])[:, :n] + return x_fp8.contiguous(), pack_ue8m0_to_int(sf) if use_packed_ue8m0 else sf + + +class TestMegaMoEPreDispatch(unittest.TestCase): + @classmethod + def setUpClass(cls): + paddle.seed(2025) + + def setUp(self): + self.num_experts = 160 + self.num_max_tokens_per_rank = 8192 + self.top_k = 6 + self.hidden_size = 7168 + self.moe_intermediate_size = 3584 + self.group_size = 32 + self.num_tokens = 128 + + self.x = paddle.randn([self.num_tokens, self.hidden_size], dtype=paddle.bfloat16) + scores = paddle.randn((self.num_tokens, self.num_experts), dtype=paddle.float32) + self.topk_weights, self.topk_idx = paddle.topk(scores, self.top_k, axis=-1, largest=True, sorted=False) + self.topk_idx = self.topk_idx.astype("int64") + self.topk_weights = self.topk_weights.astype("float32") + + def _new_buffer(self): + x = paddle.zeros([self.num_max_tokens_per_rank, self.hidden_size], dtype=paddle.bfloat16).astype( + "float8_e4m3fn" + ) + x_sf = paddle.zeros([self.num_max_tokens_per_rank, self.hidden_size // self.group_size // 4], paddle.int32) + topk_idx = paddle.zeros([self.num_max_tokens_per_rank, self.top_k], dtype=paddle.int64) + topk_weights = paddle.zeros([self.num_max_tokens_per_rank, self.top_k], dtype=paddle.float32) + fake_buffer = FakeBuffer(x=x, x_sf=x_sf, topk_idx=topk_idx, topk_weights=topk_weights) + + return fake_buffer + + def mega_moe_pre_dispatch_ref(self, x: paddle.Tensor, topk_idx: paddle.Tensor, topk_weights: paddle.Tensor): + x_fp8, x_scale_tensor = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True) + return ( + x_fp8, + x_scale_tensor, + topk_idx.astype("int64"), + topk_weights.astype("float32"), + ) + + def test_mega_moe_pre_dispatch(self): + buffer = self._new_buffer() + buffer.topk_idx[self.num_tokens :] = -2 + buffer.topk_weights[self.num_tokens :] = 3.0 + + mega_moe_pre_dispatch( + self.x, + self.topk_idx, + self.topk_weights, + buffer.x, + buffer.x_sf, + buffer.topk_idx, + buffer.topk_weights, + self.num_max_tokens_per_rank, + self.group_size, + ) + paddle.device.synchronize() + + x_ref, x_sf_ref, topk_idx_ref, topk_weights_ref = self.mega_moe_pre_dispatch_ref( + self.x, self.topk_idx, self.topk_weights + ) + + np.testing.assert_allclose( + buffer.x[: self.num_tokens].astype("float32").numpy(), + x_ref.astype("float32").numpy(), + rtol=0, + atol=0, + ) + np.testing.assert_array_equal( + buffer.x_sf[: self.num_tokens].numpy(), + x_sf_ref.numpy(), + ) + np.testing.assert_array_equal( + buffer.topk_idx[: self.num_tokens].numpy(), + topk_idx_ref.numpy(), + ) + np.testing.assert_allclose( + buffer.topk_weights[: self.num_tokens].numpy(), + topk_weights_ref.numpy(), + rtol=0, + atol=0, + ) + padded_max = buffer.x.shape[0] + np.testing.assert_array_equal( + buffer.topk_idx[self.num_tokens :].numpy(), + np.full((padded_max - self.num_tokens, self.top_k), -2, dtype=np.int64), + ) + np.testing.assert_allclose( + buffer.topk_weights[self.num_tokens :].numpy(), + np.full((padded_max - self.num_tokens, self.top_k), 3.0, dtype=np.float32), + rtol=0, + atol=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/quantization/test_modelopt_nvfp4.py b/tests/quantization/test_modelopt_nvfp4.py index ebf3b1de064..2090280ade7 100644 --- a/tests/quantization/test_modelopt_nvfp4.py +++ b/tests/quantization/test_modelopt_nvfp4.py @@ -163,7 +163,7 @@ def test_module_import_with_flashinfer(self): return_value=True, ), mock.patch.dict(sys.modules, {"flashinfer": mock_flashinfer, "flashinfer.fused_moe": mock_fused_moe}), - mock.patch("paddle.compat.enable_torch_proxy"), + mock.patch("paddle.enable_compat"), ): importlib.reload(nvfp4_module) From 8b58eafc6c350034fa64ffa6bf602e963211189e Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Thu, 11 Jun 2026 10:55:59 +0800 Subject: [PATCH 2/5] delete log --- custom_ops/.ninja_deps | Bin 89432 -> 0 bytes custom_ops/.ninja_log | 464 ----------------------------------------- 2 files changed, 464 deletions(-) delete mode 100644 custom_ops/.ninja_deps delete mode 100644 custom_ops/.ninja_log diff --git a/custom_ops/.ninja_deps b/custom_ops/.ninja_deps deleted file mode 100644 index 88c3ef741f77ffd79f52cacb90d9350654f30196..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 89432 zcmeI*2b>etzBuqGD0Z?=cswu0A zM@vUX62q#JiD6?apO3P~ERpq4> z{++G*ce(%3UrUwvWzCqjW4d+;cj?qNnySrKCKGKtwr@8-KN*&)$^GGIDiw~6$;J|y zcruYGN!q{P{of%Aen0<*2%^#7Z2$cK?e?qViFh=bD35e3jU`59N=H|X4Ohp8MZ&4_ z(yDk_X}(50HYjA=j(Tx&Av(yC-MQsvYwk2F-5;J<{^ zHHmDzIu`8QkG|jT@4lbTmR2S+*|54=5g(Q-k01VF8$kn+jVHp{+Egs#7vKNG_se6W z;?Y>c*F)d^yD^o=DzsI>CPLr-(>BuaQL*&AZG7_&-yao=YA@$~|LcGFex@c}5oxp& zU;V@PvrgEW^t_$^@*lphNk~W1wG9D%;roAk+$tl9M64=Y9?L}2@sxcjsr>w(e$%LK zgzGB*XMWT5@9IbJo4Ny7SFP%Q8v67fwpX4%CKWY_Xg1flJC2|H!*A-2%O8I8@qhfA zRO%zYnH=ob|0Y-de{RF0$?EE4qO>Yfo2<#I)E|oRvcilpHYOTN1qUrS_dY1xvaF6o z)5)NLcwdZD3al)HoLG0v%OlxH?%WQ_^S#23wXC8l8OgTq5`-Fhw*X_UE7LMPrdXmp zTp80n*8GL$inVo62mU5E20QhR?-l5I7o_5RYb(&V3oy<`b=qM_RmH}XRK6w7Hh)|0 zJJo+RG{<*4{oQw?(W=sHW!mLPaPH(c;QMd--QVV^M!zeMWg~I@qWT`q8-CmU-`wUO zemn0^gkJx5;?DKN!r`b`>iScm*Z!St=YHqFm*#(48c9bhT@yru9)x=Es=v|gpW^?Q zzmdtr(m|F9gj79C-{s0M*sh` z9O9X9BAEy$BI$H;bUZOEI52vO*?3u1P_{43Ph75x)3InI8!OM1@ALEXHySc*u6rAL z&foZtk8vtHCLGU1f*_vt`+qz?{#-}W*=%qH^o-y3`cKO>6B}6*OXv|tqQOB9Vt9If z%AqEajzyxCk^JEZ&7Plfa6E(2gCK^d{JzlI_C{@HLu!3eoYw|rUO@fRKgZcM9P3ow zu_9NlPl$I);m6y#xvEK){AV_(nW4}_V!gG%j`<&AHLzFf*jGogrSm?PRL&^$GR`K$`O+?_e6RrP zY`%;~Mbhy|BIuPrAjSgjr(_H)M}lkW`wO-#lezJ!h}#X*#aL+fYx2k2m4`O;KJndH zVCQ_0kAnL(_4h(t6$L}q;Z!o2yMc4BmOEpD#J+Un z@~x=J#LDYN&ZiYYyqQ!aopEs&+#6sk~F1g?66R zw+-5Z((<^Pkd0ObWjeJ$V~r-$vC>g-x6vDh*h6=SwNTrHpj?wFvjN-_i0fzvWqZ4L z3pGas@g^c!H(2vA*WZi2P0W)EFt6q6LwR|;TEX60fHBrxW7S0%3f&^cLhHZO#~2ji z&4pWr4dW9xiLubm_qsAP!kMh9$`7RkQ=Wp;J#?dZ3-Nr?h_Vf1O>Y|E29UGl?L$kgP zIRaM|VEjR?ufHCyFY&s_O~cIB{#;EG{}k~TTE1xfB7y+ zN#&*DEx5kZpX+ke{P}rB9J)lzg_bY>rB2s9LSSGmVJdn$oT*8rlKDv_7Z*d_c9$iV zpQhUo`Q+k?dsL*VCUi4D z^G)5&b>$NIIf!Q#LwyMHh4MY480teVx8=t=P8V;%UyuB$@1~eTrximvXHuRd%U}1M zS`6`~W5epEnubmhZ^5?{^0NEn3OWC5-^s;L#<_cW`MF&u6+;=9<;DUUn&T6Tp^O{G z0ZJ+-6+^ta_Bas+lIzEoE7Hlav4qLXe;6bWXJY#K z6Ya~4BHEYwaq{||$v*>JQaPaj&x3jQy7G4shs6@E0jk1{A3i=t+=ZOu=bZ=h&fTMn zCgz|s-eLVcF7%rLABL= zf?Gu6il^*@*n{d{>gm3k>KapS^}7x)!uWS?w2;AJ#T9@4IkD1y^-COt4i$gF^=tmR z?i)ypR+i?U%(voHZN4%!Iv9u8mtt*8a}!|XfrITs!S&7l`gr^~AN0sQw3{34428zp zkJ-hM*Z!qmHOg^+*r!_iR3v$>p{?x^JZw2I_XuNXjQEQ!&-Kc!^Hfqfy12?dSB3Ji z56%rxj}rf^Vm==l?~MJZDbTu6iSM@P0;@RkLxG}Uk#74)Iuon?|nD-99gyR7Gh3oIG5|o!apO`@u=HUbt_AS zLRDff;OEKer)h;oh_j&E%lb057EQv+BEt)`PW)|ds%tyrVlSd|yKgY_y;98od~WBy z8X8ta$0i&bWg2eX!wK513b7aF`D#qAVRemAC??LDfb}5?uA?*@_m~>tc$;!cnh)rn5)0}cym`2^JXoDLVXIq zjB_#9Kj*Zocnfh1Hdn@V*AsP3uxa+e>x8`HCjVGYXcw_hF0gB%{H1WIMd^m+2G4to zu@LV&{}1IG&%F(-Je+^+iXO$z#Z^BVJv|Z%^%8p_wjT}sw)%UUjgsI_#S(L6tjfv( zxjtu4F&ARq{zEzE*1FGKGBrfrL)?XU{gFF28?A59U7Q6w=94SZc^4o%7IK-+3p2EX zI190D$Q}DX&4H+2#c+Eu7xeWz)Z|cpjl|-3tgiiG6y3T~GZF&$Bed-Y<8{sdw%L z1wm)DV-c4Bpt{BNJBYmyuND8U?&M=1oSP@xUi`C)qyGH$yz$06V?Wvz*YU_bpVPDV zyq7hWiho9N9FK-M@44_hg<~D6+gBvQV}dD#p%VKsy%_hS(F?QM7VEL7--nv9xoiEc ziuhRMUIwNfZCS*}!ah_a>W;-0MSLu}R=2YsC6$}oheGV@{SP@W*T(J|smi^)sZ9~a zUKY>h#-%qCdlB_Da^)WMTXQd%bq;S@gk?X{dXybyv`vaM_O3EO>msZ-BWo0VhcTUk z1h#Rp#@{L57uu*8>rPPqD$j<+Rra}|*g>uotbLp7`)nZof*pr1x2{7hoPVW&XxA5O z!H#d&uTAHL(C#5+a1p)+w3y6#804xDWaDV)o3v8!~8R z)-1w(=oPDoSmd`?FmWe1pVzPtGmGbZ$k(OZ7-uj}I5-*nXTQ3(X?6QjY}bFmaR^p9 zUrp>q)UM{Pp@Oo{t)8SJty(gI0k z|8Y-3p_Pm0oX8)ShBw%ULMs*5Igx9ddYXmlJEj%qRNG zSj+ViEIJppyoAo%vf?e6`FejDZ!l@Eq;eT?7Rauz&7jN##0UYQ?>bJN}4x8`Nzft+s=)`V zw4!Y?jk*ln+E$KA;bx~!pJ zY=8VhZu{!%K;t-mZh~?M#t|}A$y_V+lkG3;{F|z+jAXO{KQ`epNY$3d%4+iOBKe^S z;wY=l#&RzWuRkHw#qY&&LKD@=zl!HO@r+x@@dR;X8kNbnVwkv)W2iq*{~XIVO;8>+ ziFh)w4bTW$;-O#85XUry+_@n;-@Yat)hzW`;taG(imcN0l3jA@p$*#E`8q zJ!)PEp^wDSRBiU3WBE`l3%!jpE?Qq4AJ~3V^)ss)#qoXEOVNoJd))#jT`Sbo1jhv@42c!B5yQ78RUlf>jO|9UvG*sh|fC;ui1W6<&26( zalG0D<&dh4Mp9la9ZWcR#r7BW*hFx#6?E)hw(W)8Hq~Eu*2nNtlazl&I`?u*=fR6j zu>EX05{(59o;wIH*nU%;2aV%+z6t6=rYat7+^2r7364cHW-?^mnDMhsu>Igw==Ixw zrU~mpJQ2^<-?DnT3Ch7zYW2ksn%xB3Pt|6!xfic#6P{|qS&!1p!5_N$ZSnVM97&dIg(h0E%o_(S*DepA)M z#&Jw*f^x`Dfl1c{Z|u0c3Cp2z9CwLhVfUX>wN>$oShO};S10Zi!$fssVUB~&%ME{w zWvW;fdX9*ViUhNA?r4HzQC|+xx)^R3!@|x}-f5(X%8z;7X8TRmJ~WQw)+Q*2cx5AN$nD`UB_r<-kmVV?&krsoI28ZAGp&PHBP|a^vvzFZR+$y0Qsk$lVcX7-0#8t`I|0_1FF! z%jIIZY(e%Ny+q5T^@e9>G6v<_8fzJ0agt1FDpFov74v*-w5p~&R{EFE{6r(cQ-YU? zvnZbV$_;LoR9-6HsRdoGbTm3V3sZ{WMcX>@bcfaByaTmq2fjZ%_bRrVGv;4whi8-ju!NLsA zxeLTx6zzsf?wVNGyL!Ww$*S_)8#T@sdttX5$%ge*&J*jjf^OsL&coQ4Xe<|XSXU;I zsfuKC3rCzQ{-Vf-!Od?q`yBBWc0XfSJQzbdTbza6K8rJ)iBtqh+iM__iD>RZ?<{c_ z_PM8#bE7k8R$g7R)(T*Hx~ls)?7EHZ0~JoLymS%?HIAI=$fQfXTzD z+FY!s6@0AoR&701tkahWjZ?&-cOIS8alqW$)?Ccz#c4)!mS9O*uoO$P49l_{%d-M2 z(vp=}nN?Vo)mWW1Sd+C_n{{Z#x~#|gY`}(W#KyE{6EajM2R+$|UhGV7c41fgup7J6mwxQQp7dt`0~y3%_F@Qo zvk&_+l>OMB12~W{2N9u+DCNYcU>L!<5N9|esG^z#Nm7g?O@=HrjAArnsAVh%a|nlW z7>6^CBRGa4PIg7J7hjTfP^SOWvxroVJ z%q3jPWn9h`T*(x!;%ctpTCU@IZs104;%08)R&L{V?qDi+au;_qjeEG4`^95h>6<_lW-|`*b^8-Kf6F>6{zw#Tu6I#swvpCIY&Jrw13zlMO zmSI_zV|i9!MOv~FE3*o#vKp(i20`Pp7HhK(tyq`!Sf35pkd4@w)@;J2Y(^V4XA8Du zE80>*DeY)a2RhP;&U9gGy0Q)1vK`&no*mec?)0E1JJE}s>CG(GjIS&#MEfDPG*jcLs$Y|3V|VRN=%OSYmdC6v;R_H>{lo#;##wx%oFur1rsjqTZi z9qCRFda@I}*qPq!!mjjTH+H8l{n&#&>CXTLGKj(K#Sr#pANFM^`>{UVC*&D_GR+{W$P!Bp<#F79R;_i!)wF`fH)fCrhuLp;nQ%;ZrX<8fy31W)o5 zvw50lc$VjQo)>tLmw1_1c$L?9oi})sIlRT&yu)1Hvk#7dGUE&M$B z#|Wa)(1~?ROC27G@EX^`3%W^Ew3am&=R$^sVVO3UR zb=F`_)?#hep%v@09_zCK8?q4_)0$1#l+9?v=4`>1Y(-m2D5V|k=|D$1(U~r6O;@&I zTehPc+p_~Z(w!dkWG8yDGrie`UFpMa>`q_$u?Kt7p8*VH5QEu^A?(dQ?8{L0V}B0d zK*AhEgfgO(6QhD*R1#-6BdDU91W8hiBu$1aHH>03W2j{;2XhFAau|m*jw3jdqZrT8 z9K!@AaxBMjJSQ-T6FG^KIfYX>jng@UGdYX1IfrvOkMp^J3%Q8NT+Ah0%4J;66f%-|s&<`HJ{D39?t zvv`6hd5YOQ%`-g9b3D%ryvR$u%qzUgYrM`IyvZEi;%(kxF7NUl@ACm4@(~~N37_&A zpYsJ@@)ck64d3z|-}3`M@)JMv3%~Lkzq2S$7T2}q9Ea?`E>9+FvZGlc4K$?(vLmZ zll}}~AcGjpUJPMx_F-R!vLE|%00$E0AR?3zrJNWQ45N}b!x=#p)g(xgVkBuYWT{~k zqY3JHEn_*DLpYSfIGk}D!I2!rc#h^6CNPm>IgaBwfk~XmNu10noXTmO&KaD^S)9!| zoXdHf&jnn_MNH;mF5yxx<8rRxN~UlXS91;5avj%m12=LLH**WOavQgE2UEF|ySSTa z+{3-x$8_%J0Ul%q5AiUMFq21ljK`V96FkXN%;ssH;aQ&Jd0yZ}UgBk5;Z84j z=I|D8^A2-)m-l#|5BQLe_?S=ll+XB_FZhzL_?mC{mhbqUANY}<_?ch$mEZWCMY*}S za&vI}|4VaoxH>Z|JUX7O3}=(65i<30@>S6OEa3h9*F62FNTN|~U|sI4Sxj!2eNej~ zcaHLTQSMxnI~V27MY*$?wxBsnup}*5bkDS~?wQV6@VR6@ZhX$iPyQlf#Zo27iE`8- z3w|H}dh3$p1lu38@U|ax7wR8L3LRtH(_Pz+cHYI;daiwtpU3zd)@41`X9G55BQ~Zro3JUH(T2_0 zf-TvKwvn=-2X>@8J?P0!^kQdvvkSY@huzqnzVu@c z_M|@p7|0+7vlm0yn|;`qq3p;09KeBuIfw{lL@6gm1;eN$&TvLhMKuYMq!>w>3|VRz z#b|>1R?AopCaC|1au|m*jw3jdqZrT89K!@AaxBMjJSQ-T6FG^KIfYX>jng@UGdYX1 zIfrvOkMp^J3%Q8NT+Ah0%4J;66f%-|s&<`HJ{D39?tvv`6hd5YOQ%`-g9b3D%ryvR$u%qzUgYrM`I zyvZEi;%(kxF7NUl@ACm4@(~~N37_&ApYsJ@@)ck64d3z|-}3`M@)JMv3%~Lkzq9E2 zwz#fugZi+5ulMF7;m*fT>a%2SHghIE%;P)Zs)*HhM#qEuv`5QVL0hzd_x10ikh8~& zW8A`vBOR+wW`k+|))WYj7!{7Eq9v6_iD8D^J$5neVmA7Ekh>4|d2w+yqd7~kBrRBq zrCEk$S&rpdffZ@VO03K(tjcPv&Kj)ATCB}Fv|?S>V|_MYLpEY#TC)k8vKejIoGsXr zt!PULrL?0x9q33WI@5)%>B=^2%XV~Qdv;()y3>Q6>_jhirZ>B=D}C6F-RVm|_Fzx? zGk}2%VlaC#guU5^eHqGr?9TxlNSK3&P)3w;VpK4UO5zM>1XWa%AW4dmq{)z_hEa@W z47H5qU=HC>4&!jfaRf(l6yrIXW0=51j^#Lx=L9BkA}4V&r*JB#aXM#kCTDRr=Ws6P zaXuGtAr~>3i@AhLxs1!Xf-9NARb0(AT+4M_&kfwjP29{a+{$g-&K*qUPVVAvrg0DV zav#&Vp9gr589cy{F zEH>{tD>(KGIH%1=GMLcZ@p#53E`}}W3 zBb|qisLQ2?Iwr#-+l@a@Eu`;W$Dyy+N zYp^D3u{P__igj6!_1S<8*@%s4%_eNhX0%~*wqQ%PqAewq(vJ3Ypd+2=Oc%DME8DOw z+tH2f*?}GDP7ivr6TR4(-t5Az^kFx4r!W23gFWfb00uIM!R*Bl_GTaUWhncxKL>Ci zVGbff8BxlKQNb`Oi8GuLR8dWWBq>IcCPS7QMlqT()H0TXIfO$wjKdkn5gf@;jOS>M zVFD95mg6{{6PUz_oW#kT!l|6b>72otoWV zaW&U)E!S~9H*h02aWl7YE4OhwcQBPZxr@7*#y#B2eN5+m9^gS{@DLC42s3$<$9SAs zJi(JZ#cZDD8J^`ip63N#!V%Px*|` z`GPO`im&;GZ~2bz`GFt#iJ$p}U-^ySS(HnQHZs!iBawm6jH`BO> zd%2J4+|L6%$P6ChVIE;7kMbCgGm9s9lBbx>(>%koJje6Az>B=Z%e=y?yvFOi!JEwC zE#BrG=JGD@@jf5$As_KEpYSQ4@i|}cC13G1-|#Kp@jXBABR}yozwj%+@jHw1XmRDy z;P@}#d@>)2a6W!gpS3cPQL%7-8u6c|CnXP)wajM2R+$|UhGV7c41fgup7J6mwxQQ zp7dt`0~y3%_F@Qovk&_+l>OMB12~W{2N9u+DCNYcU>L!1i!+=NR8dWWBq>IcCPS7Q zMlqT()H0TXIfO$wjKdkn5gf@;jOS>MVFD95mg6{{6PUz_oW#kT!l|6b>72otoWVaW&U)E!S~9H*h02aWl7YE4OhwcQBPZxr@7* z#y#B2eN5+m9^gS{@DLC42s3$<$9SAsJi(JZ#cZDD8J^`ip63N#!V%Px*|``GPO`im&;GZ~2bz`GFt#iJ$p}U-^ySS#;f7 zT-Uw9byhI%{J$%Ir()3>tHj2_1XF+?hfXkxcx_CaK4AOsSj5vdKvgzW?G=y$ab#5)o(RF=+W7MU?&Kwk< ztS>OZ=ZP#X)@C$k36`V_~Td(373$ z#m@9*7j~r&yRkcc>Bk=INq+_~kU-!V*^m7>fCCA05E06VQcjEthEYkJ z;f$b)Y7#gd7WD;`!A;!EE!@g&+|C_LGlc4K$?(vLmZll}}~AcGjp zUJPMx_F-R!vLE|%00$E0AR?3zrJNWQ45N}b!x=#p)g(xgVkBuYWT{~kqZvajV>y^Z zIF!RUoN*k%ksQT%j^-F9Fp*<9j^jCjNu0<@oXjbl%4wX=8Jx*koXt6$%Xys71zgBQ zOy*)P;ZiQ+a<1S?rf?Nka}C#W9oKUMH*ymdpRbJzD-r!B<@D^|L4s&^z z_jsQV_>hnIm{0hW&-k1#_>!;qns4})@A#e{_>rIZnP2#o-}s$HxwZImYj7>pod1jY zwedn7S-EVoD!hQp;Xll?jUtqBGhz(xOCRreRW@6FV32Q5@Og1@H={XAup}*5ilteG zWm%5pS%DR4$x5uuDy+(Ctj-#&$y%(KlWfx`ZIum z3}P^QF@(L@hkY5!e(cWy97vdhh)_n9a$;04j7s7RX9QJLlORcok)+9xrG`8A zaUv&iGN*7Vr*S%Ga3*JQHs^3I=W#w4 za3L2lnTxrEOSz28xq>U1!c|<&HC)SeT+a>M$W7eLE!@g&+|C_L!cmvVtta}txZ~A~<}AUI zv|uTgW*L@cIhJPyR-`2>u`;W$Dyy+NYp^D3u{P__igj6!_1S<8*@%s4%_eNhX0%~* zwqQ%PqAewq(vJ3Ypd+2=Oc%DME8DOw+tH2f*?}GDP7ivr6TR4(-t5Az^kFx4r!W23 zgFWfb00uIM!R*Bl_GTaUWhncxKL>CiVGbff8BxlKQNb`Oi8GuLR8dWWBq>IcCPS7Q zMlqT()H0TXIfO$wjKdkn5gf@;jOS>MVFD95mg6{{6PUz_oW#kT!l|6b>72otoWVaW&U)E!S~9H*h02aWl7YE4OhwcQBPZxr@7* z#y#B2eN5+m9^gS{@DLC42s3$<$9SAsJi(JZ#cZDD8J^`ip63N#!V%Px*|``GPO`im&;GZ~2bz`GFt#iJ$p}U-^ySS(ICg zFSiEQvi}$E8_ydnZj{9uT`DG3-&6rz&;E6k(P*;ndiEs8t>2KZt_^bT$v!VG{$@01 z36`V_~Td(373$#m@9*7j~r&yRkcc z>Bk=INq+_~kU-!V*^m7>fCCA05E06VQcjEthEYkJ;f$b)Y7*q?oPUla zO@=HrjAArnsAVh%a|nlW7>6^CBRGa4P zIg7J7hjTfP^SOWvxroVJ%q3jPWn9h`T*(x!;%ctpTCU@IZs104;%08)R&L{V?qDi+ zau;_qjeEG4`^95h>6<_lW-|`*b^8-Kf6F>6{zw#Tu zvnb~lU(Wr%HFjK+Di3CG#S>n`5RR8;!qHSsNitDDx%Mw3j7CG%VhrBAeO}yY*No;Y!IHFKDVAm#mSs7XX9ZTIB`dKqtFS7ou{vw8CTp=a>(GjIS&#ME zfDPG*jcLs$Y|3V|VRN=%OSYmdC6v;R_H>{lo#;##wx%oFur1rsjqTZi9qCRFda@I} z*qPq!!mjjTH+H8l{n&#&>CXTLGKj(K#Sr#pANFM^`>{UkMyP3v4+{=AT=YAgGL1yp}5Az5!d6dU^oLM}(_ zANh%&`GsHkjo(?6Yl|z_2KD2=YTP)RjwCV_v2>%kPE~T$xCOtT^&{7RjzE^w)-L22 zf_dKI`bbJDQz_~Td(373$#m@9*7j~r&yRkcc>Bk=INq+_~kU-!V*^m7>fCCA05E06VQcjEthEYkJ;f$b)Y7!&~j_F9!WXMv(C`L1eTE=oPhj1u| zaX8~Rf+IPK@f^)DOkg6%avaBV0+TqAlQ@}EIF-{loijL-vpAb`IG6J{p9{E_fHCwP*ln9b8X!?Qfc^Sr=|yu{1A!mGT->%766%;7EG<{jqpF7NR^ zAMha`@iCw9DWCB Date: Thu, 11 Jun 2026 11:01:18 +0800 Subject: [PATCH 3/5] fix code style --- .../model_executor/load_weight_utils.py | 1 - test.py | 74 ------------------- 2 files changed, 75 deletions(-) delete mode 100644 test.py diff --git a/fastdeploy/model_executor/load_weight_utils.py b/fastdeploy/model_executor/load_weight_utils.py index 5f19ec5a3f3..1ea906a2f3a 100644 --- a/fastdeploy/model_executor/load_weight_utils.py +++ b/fastdeploy/model_executor/load_weight_utils.py @@ -127,7 +127,6 @@ def get_weight_iterator(model_path: str, fd_config: Optional[FDConfig] = None): if use_safetensors: load_config = fd_config.load_config if fd_config else None extra_config = load_config.model_loader_extra_config if load_config else None - parallel_config = fd_config.parallel_config if fd_config else None if extra_config is not None and extra_config.get("enable_multithread_load", False): weights_iterator = multi_thread_safetensors_weights_iterator( files_list, diff --git a/test.py b/test.py deleted file mode 100644 index 753a76d1589..00000000000 --- a/test.py +++ /dev/null @@ -1,74 +0,0 @@ -from pathlib import Path -from safetensors import safe_open -import json -import re - -def natural_key(s: str): - return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", s)] - -def layers_are_grouped(keys): - - seen = set() - current_layer = None - - for k in keys: - m = re.search(r"layers\.(\d+)", k) - if not m: - continue - - layer = int(m.group(1)) - - if layer != current_layer: - if layer in seen: - return False - seen.add(layer) - current_layer = layer - - return True - -def values_are_naturally_ordered(values): - """Check if values are sorted in natural order.""" - return list(values) == sorted(values, key=natural_key) - - -def get_all_weights_file(model_path: str): - """ - get_all_safetensors - """ - model_path = Path(model_path) - use_safetensors = True - files_list = [str(file) for file in model_path.glob("*.pdparams") if file.name != "scheduler.pdparams"] - if len(files_list) > 0: - ordered_weight_map = {} - use_safetensors = False - # dont care about the order of the files - return files_list, {}, use_safetensors, False - else: - safe_model_path = model_path / "model.safetensors" - if safe_model_path.exists(): - with safe_open(safe_model_path, framework="np", device="cpu") as f: - key_name_list = sorted(f.keys(), key=natural_key) - ordered_weight_map = {key: "model.safetensors" for key in key_name_list} - is_layers_are_grouped = True - files_list = [str(safe_model_path)] - return files_list, ordered_weight_map, use_safetensors, is_layers_are_grouped - else: - index_file = model_path / "model.safetensors.index.json" - with index_file.open("r") as f: - weight_map = json.load(f)["weight_map"] - keys = list(weight_map.keys()) - values = list(weight_map.values()) - is_keys_orders = layers_are_grouped(keys) - is_values_naturally_ordered = values_are_naturally_ordered(values) - is_layers_are_grouped = is_keys_orders and is_values_naturally_ordered - ordered_weight_map = { - key: str(model_path / weight_map[key]) for key in sorted(weight_map.keys(), key=natural_key) - } - weight_files_in_index = {str(model_path / weight_map[name]) for name in weight_map} - files_list = sorted(weight_files_in_index) - return files_list, ordered_weight_map, use_safetensors, is_layers_are_grouped - - -files_list, ordered_weight_map, use_safetensors, is_layers_are_grouped = get_all_weights_file("/root/paddlejob/share-storage/gpfs/system-public/dangweichong/models/eb5_v4p1_v4p2_merge_0413v1_step120") - -print("is_layers_are_grouped: ", is_layers_are_grouped) \ No newline at end of file From 66fac9a2d0950694df9c43ba273b1f961b74a5b0 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Thu, 11 Jun 2026 11:35:27 +0800 Subject: [PATCH 4/5] fix paddle compat --- fastdeploy/__init__.py | 5 ++++- .../layers/attention/dsa_attention_backend.py | 5 ++++- .../model_executor/layers/attention/flash_attn_backend.py | 5 ++++- .../layers/attention/mla_attention_backend.py | 5 ++++- .../layers/batch_invariant_ops/batch_invariant_ops.py | 4 +++- fastdeploy/model_executor/layers/moe/ep.py | 5 ++++- fastdeploy/model_executor/layers/quantization/fp8_utils.py | 5 ++++- fastdeploy/model_executor/layers/quantization/mxfp4.py | 5 ++++- fastdeploy/model_executor/layers/quantization/nvfp4.py | 6 ++++-- tests/quantization/test_modelopt_nvfp4.py | 1 + 10 files changed, 36 insertions(+), 10 deletions(-) diff --git a/fastdeploy/__init__.py b/fastdeploy/__init__.py index 733c2c6ccde..5972411319b 100644 --- a/fastdeploy/__init__.py +++ b/fastdeploy/__init__.py @@ -106,7 +106,10 @@ def _configure_logger(name=None): # cause some unexpected issues in triton kernels. We use enable_compat_on_triton_kernel # for these cases. if not _is_package_installed("torch"): - paddle.enable_compat(scope={"triton"}) + try: + paddle.enable_compat(scope={"triton"}) + except Exception as e: + paddle.compat.enable_torch_proxy(scope={"triton"}) if envs.FD_DEBUG != 1: # Log level has been configured above diff --git a/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py b/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py index aaaf7018dbe..0865109b205 100644 --- a/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py @@ -26,7 +26,10 @@ from fastdeploy.platforms import current_platform if current_platform.is_cuda(): - paddle.enable_compat(scope={"flash_mla"}) + try: + paddle.enable_compat(scope={"flash_mla"}) + except Exception as e: + paddle.compat.enable_torch_proxy(scope={"flash_mla"}) from fastdeploy.model_executor.layers.attention.ops import ( get_block_shape_and_split_kv_block, diff --git a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py index 3ad7e2738e4..3226934c3a6 100644 --- a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py +++ b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py @@ -87,7 +87,10 @@ def init_flash_attn_version(): sm_version = get_sm_version() if sm_version >= 100: try: - paddle.enable_compat(scope={"cutlass"}) + try: + paddle.enable_compat(scope={"cutlass"}) + except Exception as e: + paddle.compat.enable_torch_proxy(scope={"cutlass"}) from flash_mask.cute.interface import flashmask_attention as fa4 global flashmask_attention_v4 diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 6ada2c9468f..8f0f22d8ccf 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -18,7 +18,10 @@ import paddle -paddle.enable_compat(scope={"flash_mla"}) # Enable torch proxy before importing flash_mla +try: + paddle.enable_compat(scope={"flash_mla"}) +except Exception as e: + paddle.compat.enable_torch_proxy(scope={"flash_mla"}) import math import os from dataclasses import dataclass, field diff --git a/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py b/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py index 5a2675ec07c..6ce3d691708 100644 --- a/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py +++ b/fastdeploy/model_executor/layers/batch_invariant_ops/batch_invariant_ops.py @@ -806,12 +806,14 @@ def enable_batch_invariant_mode(): return if hasattr(paddle, "compat") and hasattr(paddle.compat, "enable_torch_proxy"): - paddle.enable_compat() + paddle.compat.enable_torch_proxy() # TODO(liujundong): Enabling torch proxy here has a global effect. # Do NOT call this function from module import time, # otherwise it may affect other test cases during pytest collection. # (ex: Could not import module 'PretrainedTokenizer' or No module named 'paddle.distributed.tensor') # Other side effects have not been observed yet, but they should be watched out for in the future. + elif hasattr(paddle, "enable_compat"): + paddle.enable_compat() else: raise RuntimeError( "Unable to enable batch-invariant mode: Paddle version is too old. " "Please upgrade PaddlePaddle." diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index 589ee0d415e..b7d68e1ec94 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/fastdeploy/model_executor/layers/moe/ep.py @@ -41,7 +41,10 @@ def load_deep_ep() -> ModuleType: try: if envs.FD_USE_PFCC_DEEP_EP: # Enable torch proxy before importing deep_ep (required by PFCC/PaddleFleet variants) - paddle.enable_compat(scope={"deep_ep"}) + try: + paddle.enable_compat(scope={"deep_ep"}) + except Exception as e: + paddle.compat.enable_torch_proxy(scope={"deep_ep"}) try: import paddlefleet.ops.deep_ep as deep_ep # type: ignore diff --git a/fastdeploy/model_executor/layers/quantization/fp8_utils.py b/fastdeploy/model_executor/layers/quantization/fp8_utils.py index ccc8726efdc..de9fc557e36 100644 --- a/fastdeploy/model_executor/layers/quantization/fp8_utils.py +++ b/fastdeploy/model_executor/layers/quantization/fp8_utils.py @@ -68,7 +68,10 @@ def load_deep_gemm(): if current_platform.is_cuda(): if get_sm_version() >= 100: # SM100 should use PFCC DeepGemm - paddle.enable_compat(scope={"deep_gemm"}) + try: + paddle.enable_compat(scope={"deep_gemm"}) + except Exception as e: + paddle.compat.enable_torch_proxy(scope={"deep_gemm"}) try: import logging diff --git a/fastdeploy/model_executor/layers/quantization/mxfp4.py b/fastdeploy/model_executor/layers/quantization/mxfp4.py index 8732b1cc4d1..ccf08548fb1 100644 --- a/fastdeploy/model_executor/layers/quantization/mxfp4.py +++ b/fastdeploy/model_executor/layers/quantization/mxfp4.py @@ -35,7 +35,10 @@ from ..moe import FusedMoE from .quant_base import QuantConfigBase, QuantMethodBase -paddle.enable_compat(scope={"flashinfer"}) +try: + paddle.enable_compat(scope={"flashinfer"}) +except Exception as e: + paddle.compat.enable_torch_proxy(scope={"flashinfer"}) logger = get_logger("config", "config.log") diff --git a/fastdeploy/model_executor/layers/quantization/nvfp4.py b/fastdeploy/model_executor/layers/quantization/nvfp4.py index 7a22b3e1260..c48f8e71665 100644 --- a/fastdeploy/model_executor/layers/quantization/nvfp4.py +++ b/fastdeploy/model_executor/layers/quantization/nvfp4.py @@ -38,8 +38,10 @@ # Only import flashinfer on supported GPUs (B卡) if is_nvfp4_supported(): - paddle.enable_compat(scope={"flashinfer"}) - + try: + paddle.enable_compat(scope={"flashinfer"}) + except Exception as e: + paddle.compat.enable_torch_proxy(scope={"flashinfer"}) from flashinfer import fp4_quantize, mm_fp4 from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe diff --git a/tests/quantization/test_modelopt_nvfp4.py b/tests/quantization/test_modelopt_nvfp4.py index 2090280ade7..30afb9b90dd 100644 --- a/tests/quantization/test_modelopt_nvfp4.py +++ b/tests/quantization/test_modelopt_nvfp4.py @@ -164,6 +164,7 @@ def test_module_import_with_flashinfer(self): ), mock.patch.dict(sys.modules, {"flashinfer": mock_flashinfer, "flashinfer.fused_moe": mock_fused_moe}), mock.patch("paddle.enable_compat"), + mock.patch("paddle.compat.enable_torch_proxy"), ): importlib.reload(nvfp4_module) From e562c606f0d3befd238b9d2f540e4e7654239d33 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Thu, 11 Jun 2026 14:30:52 +0800 Subject: [PATCH 5/5] fix code style --- fastdeploy/__init__.py | 2 +- .../model_executor/layers/attention/dsa_attention_backend.py | 2 +- .../model_executor/layers/attention/flash_attn_backend.py | 2 +- .../model_executor/layers/attention/mla_attention_backend.py | 2 +- fastdeploy/model_executor/layers/moe/ep.py | 2 +- fastdeploy/model_executor/layers/quantization/fp8_utils.py | 2 +- fastdeploy/model_executor/layers/quantization/mxfp4.py | 2 +- fastdeploy/model_executor/layers/quantization/nvfp4.py | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) diff --git a/fastdeploy/__init__.py b/fastdeploy/__init__.py index 5972411319b..62e85063f9b 100644 --- a/fastdeploy/__init__.py +++ b/fastdeploy/__init__.py @@ -108,7 +108,7 @@ def _configure_logger(name=None): if not _is_package_installed("torch"): try: paddle.enable_compat(scope={"triton"}) - except Exception as e: + except Exception: paddle.compat.enable_torch_proxy(scope={"triton"}) if envs.FD_DEBUG != 1: diff --git a/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py b/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py index 0865109b205..547898bd983 100644 --- a/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/dsa_attention_backend.py @@ -28,7 +28,7 @@ if current_platform.is_cuda(): try: paddle.enable_compat(scope={"flash_mla"}) - except Exception as e: + except Exception: paddle.compat.enable_torch_proxy(scope={"flash_mla"}) from fastdeploy.model_executor.layers.attention.ops import ( diff --git a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py index 3226934c3a6..61fb540ed81 100644 --- a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py +++ b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py @@ -89,7 +89,7 @@ def init_flash_attn_version(): try: try: paddle.enable_compat(scope={"cutlass"}) - except Exception as e: + except Exception: paddle.compat.enable_torch_proxy(scope={"cutlass"}) from flash_mask.cute.interface import flashmask_attention as fa4 diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 8f0f22d8ccf..b6353539c97 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -20,7 +20,7 @@ try: paddle.enable_compat(scope={"flash_mla"}) -except Exception as e: +except Exception: paddle.compat.enable_torch_proxy(scope={"flash_mla"}) import math import os diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index b7d68e1ec94..159c758006d 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/fastdeploy/model_executor/layers/moe/ep.py @@ -43,7 +43,7 @@ def load_deep_ep() -> ModuleType: # Enable torch proxy before importing deep_ep (required by PFCC/PaddleFleet variants) try: paddle.enable_compat(scope={"deep_ep"}) - except Exception as e: + except Exception: paddle.compat.enable_torch_proxy(scope={"deep_ep"}) try: import paddlefleet.ops.deep_ep as deep_ep # type: ignore diff --git a/fastdeploy/model_executor/layers/quantization/fp8_utils.py b/fastdeploy/model_executor/layers/quantization/fp8_utils.py index de9fc557e36..45dfb5c0983 100644 --- a/fastdeploy/model_executor/layers/quantization/fp8_utils.py +++ b/fastdeploy/model_executor/layers/quantization/fp8_utils.py @@ -70,7 +70,7 @@ def load_deep_gemm(): # SM100 should use PFCC DeepGemm try: paddle.enable_compat(scope={"deep_gemm"}) - except Exception as e: + except Exception: paddle.compat.enable_torch_proxy(scope={"deep_gemm"}) try: import logging diff --git a/fastdeploy/model_executor/layers/quantization/mxfp4.py b/fastdeploy/model_executor/layers/quantization/mxfp4.py index ccf08548fb1..e2930f99e12 100644 --- a/fastdeploy/model_executor/layers/quantization/mxfp4.py +++ b/fastdeploy/model_executor/layers/quantization/mxfp4.py @@ -37,7 +37,7 @@ try: paddle.enable_compat(scope={"flashinfer"}) -except Exception as e: +except Exception: paddle.compat.enable_torch_proxy(scope={"flashinfer"}) logger = get_logger("config", "config.log") diff --git a/fastdeploy/model_executor/layers/quantization/nvfp4.py b/fastdeploy/model_executor/layers/quantization/nvfp4.py index c48f8e71665..8edb73a9bea 100644 --- a/fastdeploy/model_executor/layers/quantization/nvfp4.py +++ b/fastdeploy/model_executor/layers/quantization/nvfp4.py @@ -40,7 +40,7 @@ if is_nvfp4_supported(): try: paddle.enable_compat(scope={"flashinfer"}) - except Exception as e: + except Exception: paddle.compat.enable_torch_proxy(scope={"flashinfer"}) from flashinfer import fp4_quantize, mm_fp4 from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe