From 6b0b0d59f68cd10828d3d64874a7a64e1c123319 Mon Sep 17 00:00:00 2001 From: linuofan2025 <2531980518@qq.com> Date: Thu, 19 Mar 2026 20:23:36 +0800 Subject: [PATCH] NPU KV Cache --- engine/50e8005addb244c88d30f9daccaabf8a.png | Bin 0 -> 18791 bytes engine/BlockManager.py | 56 ++++++ engine/Design and test.md | 87 +++++++++ engine/KVCache.py | 56 ++++++ engine/TransformerDecoderOnly.py | 185 ++++++++++++++++++++ engine/test_kv_cache.py | 51 ++++++ 6 files changed, 435 insertions(+) create mode 100644 engine/50e8005addb244c88d30f9daccaabf8a.png create mode 100644 engine/BlockManager.py create mode 100644 engine/Design and test.md create mode 100644 engine/KVCache.py create mode 100644 engine/TransformerDecoderOnly.py create mode 100644 engine/test_kv_cache.py diff --git a/engine/50e8005addb244c88d30f9daccaabf8a.png b/engine/50e8005addb244c88d30f9daccaabf8a.png new file mode 100644 index 0000000000000000000000000000000000000000..3faf8a260d837173f74f22e12bb1b0ee26fdedd5 GIT binary patch literal 18791 zcmX_o1yogC*ENU_-72MYN_R;}cSuXONOzZ{bayvMcXvs5cb9bczwi5v@vA7~a=d4s zbJkvKt~uwrLGrRc(oEMMN+??MjKiu&*vdIXhu+y~)oPI^3 z9rBpSCLuBVC4wQMdDTqwBCMwOK}6uIlS7`k7K5|E=aXQPp+Dvr!ZpF(RBzES)UZta z>uUV-$Js3=DSjsB-alio)W~OkrKcO&bGeW<-W~pODt$4^b(wzKBDQVMo%R+UJyT}i zo?G9*;JH*Z1>;fU7Zfx%;{E--bPD^(nyDsM_E!`FqOVI0_q!e^>`^M1D~GU zYt-quxVSh!pFAgkAz4{eC?g}I?d@%Nw8)kgo~o)UDk`du4ePUmf(*pHqq%ZI*H(g@3GsH% z=lK>F7O(rEYBF;2ES2Jcu`x;-nh>ydbM}v5&82)5Z#3~MD=Rgfo*n!G0-~d$Hmn>L zr3+ZC(`3b7Uts*Mhliy-9}XzltyYmc96}g zVNnE&lyc{u&~hIyhsW0tF<*_XJ@C9s%nlB^+Kx|$2*GU0lFMZfg+_U z(-~^8&?k?SatUyT=h7}C;_Lc1FJCEQm>_XAoa8H-1FgH+6J$PLA>h1+a)8;@7xwW0 zJh%{S1`0a5u%A=;=Tmb0mG%b;ZTGt$Db~IuRaGfSN=k}LOTX1Un4AN5?&DOx3@HZ( zc3^j>meWcv%NP$#2WbRPR|rF!I0zB`*C&>-m!0n_rRx6Lwf=#9ocN{Ju!Y%8Hh%j z;=+X$@?Vs=-_(IH8Rorm!Qr9ll`AhfIXQ@hJ7w(KXKU`X!kI%P#}x-thNkX~4%6Cf zC2D^QP8q@^7~vN5-Ota<7bDkPs;}VzD#dwnYkY49e4rNnZ8!>A zTYZQw6#~W*$~oNVA00gK#D-dVeD8lO%0T$$f>A6Z`B^l5(?Z0PZh+o zu&8KVGhAe+Zf0fWH^vL^UvMVaV9RHsDk8Btxd~9jR2mIoJ=~tIW#6Y{WT1UDH8pK! zxu~jQ-rU;yMje+n6+)2dv0=?wp+HQZG9<&XV!zsK0iqV{_LzmY-kS8c^-XU+ zfBqaBAO9w;CWZFvg{Ookkea=8YM+Z^UW{M5Ls?q=6AoXS*p@3hL^?3Pr;7 z32xUQiWv@tYZuI^-Xof`YwS5z-3L1y%|e%zH9+kdRH`k8)wgg_R7>oOj({w}!I7ZV z?1+Jh*%kX)5sK4h^KEw^8us8_SM%z)t{#Y{B-Wpk|D7grOUrYr@iQ=k9a81Ft_aLrST#mSO{xh6h2imm_!5|=ew||kwl??_|W@&yT)Lz8Xp@w zUaG-CLQ47_Bn*|}55+Tw?^c$-MPN_0W@Qndp`KVUGc)7xI5Yafp`qe9(JVWy#X5d; z_))!p!>H@c*db051-~L!8=x+m#x0_$`H4McSvRHJ=S*@*^GA8WbNR0$FTIH|a}xXEoF zN?HS~6Xz$jYkt9)TR4=fFeew4k%Ici-No2v1O$Z3fyT1;)eF+!zvKS=`SV33Pv);7 zbGpaPcfO}Hq1~|zxw$f}=Te1~{Tf)y(#2m2S#Ok7>5derwpzWycR9@itvQ+1dn)TW zAcup*y=TR?7%E#}aCf#r6(@OfywrbpzV$^IW%YB){%jNTc9q%fB)P(WVF5nH*xMq4 zf`Z(!Q%uav(F3=#6nuQ?LGLIX;kN>0P1m8HBEHDy=IB;0R3y^jgo*0?wiYAOU59ob z1r_RxG|K!00?u6=CRw++gKix+k+@^$bLq)Z&orLtJt$Ay_JYCT!UYahS|0?M+1)FB zGf-bb^EwC5^qw0G{-MwXyEnRU1YS!0-v6U_2;=lsh=yq}tgadj>h^Yb$P7JBLB!Fq zu6<^o^)2!f;zEn^&^wk>$~-5u;BW%C^>T)`K=5AOQxq2`gAU8CLk_fFKj=JMbE{7~ zv}I${@|{izKJmJyMsXk;_Bg=&_1qDnMNvt2v!%A$E7=<0m)qUp;NKlMnA`1mU#Yg7 zmH%R6XF30e16_iQllRG!knxNrYobqsVQlBvtHhu``@_>(Ytkof^&Z005i)YF@J26> z!1MYD6lD+7R20g03UzX<$C)aw{5Fk9AuS#Fk{WE?%fsuY% z3A@7Q@eg}-Sy)ru{{X;5AS%B8!_|Ry6hpK)BMx2S*tE+rUbMc#XYPNkp|wl{7zAur z-oJDZ1ekYUuJN@8)4NHleH!-CDREgG)rv(%^0_mkPF8Y$iL!1;W@6+~`?b(DOUQ=i zXlTaDYr|(Yz@{8DHa7A+Cj4k6(IAV;{Nf24(+2Gz{CloUZ`q>WO@{sJMw_C&%xN&U z(gp@88hD)JCI)>lPK;)=c2LsK?Haf&`iyqfFVG#Ka3|~I-&AEU`3`vRY`0XRH*>N& zr6(DDKY;JRISp(|kR8TS|-`m*Hi3Kn^}cu z1qn{|!#P!3Z13$&RsR@MS<|_^yey~A2@xmRxI@d+E^YJ(Tlq;^C#FkKxMj^{l& z%a73(U8fw7q6-VD#$_CPSr^qkldbH$`0JPb>vRf-*?I_H@5IWv3UW6)*6sw~o=8KL zE1DJV-0x-F*5l!$%C-19Ny1>m_1|gySH4-Hfj+b-XU8UuFKTR@45TR(th#$_4rtfyYyE#!-NF$>6|?S?fpBb46m0)aGdAlX>$X>N*cp) zA3EyoO<`{TQ!iCX#A;wMR=%#cQB9rADtE}^oAvwoJIGQav1H|^5Os7KWsYjOc9)~v z!JJB3nV;-|;U#%@0LoVFoHpdkV7@|6m7xO0W{fxHiJO)+R|?+Ct@i_0!JI165`Ax! zPxrsoRw*j}+x!0iaFu`jStQG$Q%aq!1ly~8)%PpPi4hrQ5aW)yn{3whPk?{V&k;Nc zvHl<*`cMqBeUXs1rzPBZLGqsZA$z(0Dr0c|acqL$0{w8D(KYvKJEWYyV1Js4S?ozR z5+g4N!|+P_xcBI^Mbci0wO}mO_@9+m|C+W-{K-l)%SN5V`wS{FuFoG@+wj6Dd{+2P zq9n<4$*wc^du(J?;+eq#|6T;7v<&EYo7^}kwK_U9vi{V)X`Z)ZldBz3c2mOP{UT0G zQRCq}s6ookK5Ld$ln~OU!Z5u~Xrr2RsZlhz}?_|`}62X)>p(G~g<7^S1)iY0uWn*z%5qhcl;v-^8z zXoVs_Z0wdUa5aOzuN)}NQ+=vBUI#N~XgJO^pNe?{r2jQ0|5%S~gGw#o-NA5?{LTn# zrq#@T=Huhz7{lZ+zkXC$>=ecrGTCE&s&$)%iXxN~1cg3}(zkdS`cbpgdzabqc&)S5 zX{CRR{Ql`Cp6R_Je8<3`f+e0S5htIG=3jF&b4+^6!kvuJ=}(tD+FpCg zWsfqj-LN`=+BzeUOPXJM2~XTtQ+a;mC1H(1CSbhsVK$cWNS#}AiTkcdNJ+sRzn!ae zczEP&U1f%29ub=k<-_BBVw*FJW`n( zjF4~oQHPa4Db~UN~hm?82t>kcrhaF)+^e1X9H!YU^BZ zBj%YICTI0^?1rWIU)v37n?I}T4aR=9U_XIcCD&aiSeKS4I@7{0%fBsmQ?xGh#@BGH zFc?C8FZth0+>CqG;k!RHTBGUU)w6*=a!^!sv_}2$_&BD>FKw~K#w$$@d`YBVzAnMz zGz2nD@7lkxsLE+dO10NGF!m=LsKjvE}8k-@nZfoCVIiX z6BQZj&8>NCbzP1KfcrQiwV@U6`KAcFT92fRG7C1${+fx~I<-L8B~#hn%o5O#z=686K>@=ZR4mrk21 zJem9c*9J2axtZzXQo^U!+tsa;=F^*{d*x(ZueCQ3dWB^|Q!6Z5p3pF-ne)kus;hO7 z|4&!Zdti;eM2aX>dwY@fh$!TMPc6fEp*v=u1)IEInu&}EU9gAzSN`abT_MnAw7Dd{ zJ|hzqG^U1zxjJz1ocO4&PZXU;?+BlsF2}AUjzU&^Un6#%nJKy7BDW|%^juvRwv^mf zK|Nh~9wOz-r1Vi{xNqJ@8OR3u$}6L1wGX3iO8Dd|-27YU)=8k%1#nZXM5SEIm$HBX z5vL&mblVrhEMriNcC*r;TYOO&$eA~1?~5Q@J9%j1GRV-YuMH4n|->%Bwb+?m9LRZI9jaA3fssiw4hH2NYF?ZR&%idIr!Eyr__N zh`{4Gd(-r-A>cm=n3zyu2KAK7_O7qjsq#U?SdpsO(b*c!E90ZVGnc8Pw+{>sZg#>E zlCjUaefC-+9pkw~hC?GbD@Oe>s>BTTA$uGOj`K%&74h$^w+L9Rjlr-XTJy7uw@9V0 z4Q_3yz~|!hImfa+BelK#H^AFH z^FdDy&oqE`!R53js$cW`@=^i2Fyunk*%qb`(y#3IAY(vVdS+Ndf%I_ND@3su{*8u z`{Jj}j@Mi1oFuV-^zH1JUO@U*q_oUp!A>3bX9b7jD`MI73uW0yHjG#dvRnKvx}5&W*9Q$oyf7b|K_fQ=F8G#hl7X#!D;_DG4)PoTRL8?`Rd8=Ezc#9rpRd+HO>IbQj94Y`<=d zetox0l9PUYvpNC_))!x^r_y$-UmdS6xBSLq>BpeS@SB(xH%g<-h-B7xL*5aW2ot-O z4m(d5ygt1aQ>zn{&1`T#vnOO+eze2{@j#XtpcHW(NICyR&K)y3iVR=I8Zzo_mrSS}&koDP?4o2w(Bt=v!)c z!0p|i$oX00F#MzkqH<>_!O7&*r9yD<$SbA$s-S9Y$&=WZmt(b4UaL1PNAxUjmaTq% z(}hL-xXYYQwu=xg;D6NYZz zl$E2Fo#3H_%0BwfXiB- zUexLV^QYOV zRVX6LQ?&Jq_{mMb=;>|jD(zgF?iMq5AXlceyY8or+rSh8d2=&@4p0LtIM3!;Pqet8eY`b zvouCkXLom=ZxXE%;s;?MeLP|xL6%R(EfI-_E}k*&06QimxJ)H-APCKWOh!;sEd|$ zGFyaO1Nt5+dXwf4^-RKLUqHM-_&~=OAD=fF87$ls4}UH?TF>gSov$Xu%9HFB_wckY z9cr6sSoY{mF}B84TvRb(Hjuc(ejV9!hl-M(pP%#b*MB(*DioirLX zNBarAhPTkx#0(4!0`D0df-S280$_a%6K;i?eY8eB%dos$vn<$`jCRi7sJ?pq{CQx6 zZ%3!7s?&c|-Y#n@`3ABO##Y3%Y}j_po0F9#eO@lEdmn>32LDdobr7=n#_Y)6zkkCM zBtCez&zPu|=(MLpkVfTJpJU=$rQ-B@LI`N9Ko;BVK>)SX-_yKNGhCX&tdhs6`O*2d zNug*Pv}lx{KNpFFba;d|QjyjJFiYwH2y||UxLcYA4GGa{{1tzwY9BK z)3gf)AS1=)ayM)0;P2>Y+a^|To?5jr3Sg$2DD=m81#(KGVBk=p%7U;~i2H>CQ7vxS zTry$Xq4r*^jjNNIPntr&b79u+frn7)xeA|ujHYG*Dhi;fi$MOTUx3!D_%TKRSag@O zPB?Jo24FRYhKAN>>(KEGI->;{p8sQiyZn*9w|nwph4{zDzOSmTUTCx{h%Bwi&GmbJ zyk#*S#*wCY2gDK_WLUr+@33vbs#WTXbToKnP0ap1(P|#}`xinYf)Wa0&=_^0Kyv*F zR@%|Q023zayI^j7FqxOi?I0=cKHl;;V8Y7mc4H2RaDXWU3#Qn?gg7*Cts?vWzUVvWt+>$A75WEdcPUM@Z?dv^dEClD7Z9ujNhn%LU+7z zlH_2S1%-v;&K7@4O3a;L2~Dzf-h}k}bZ=y4b;dzuPt5+-){&8sK~X=J?F+>PE3&j1 zb~|kqv$bUgs6!-63EmfbdFQ0<(K6lZ#RV4ilQ_r>0E=x7XC@^ z`Qe&Yr|pw(=laZ1DcdycYOB}AvWljb2R@5<}60kd;9|{m!ymUTQ zstwotoZ__IhIess!8g_v=;+|}9x&A3ElKh#>IU+I`ASvX`1p7T!PNZx?_T|lhoka- z6IOtETkz5xnE*EFrYUdPOt5o$Kbp!V1a1(Z0m&#SA=>0SMr*W8g(Bw-D+Z!C$%zAd z5cR2PU!0Oz5d2sBDbPfxzfy)m&oh4Z_u2o|7fz1)v^J+!ER zWn3N346W6VoBi%d526hT4SlN_6$~hBH9T{Im}B#cRVxB``*SGfB;?BZo$^?5GEg12 zZMgtL->4}IHWsm#px^bbl+$>K+lq!J)p)OQH8EDmGq*ud{TnS=oDe3CS*_MgW2Lwm zbX>v{CYL&jkL|(xfx@=;_lWUy8hwdO25&&H=hv#wS(KdUChi`MB%Bzu{#;I&Y$lZx zv;Q0`kS?Q@8hQBLWw(#JBpfz2;pe(Vo}zt>e?S0Yc=_mWkp1)zQmV}tsxC1#-Zs?j zzlSCG_)*HgEE{mO=a6jF)zzhvV#A$Azs_qFkd?plHT`f9kbU+u_J}ZWzk|ZEDk{d; zOz%>`1DrN>DDGrkSXfXcI$zvrrd#`_?#y0QaV3n4nD@qC!AAT0Z5nbOJ`W%f z#S_qF(|Hc8l-+^aCs|CFAd|w*WVy`gRXleh@81M#4LFXQlZRjMVF4{NF06@o4X0U; zRJ`B%2BMkmcen}M9+)5$NTp3#3l!-37p^9}aGXw8kUw#EtR=j$CxL-WU@595J8#DU zoqPoU3#aaSLT2am*II44-3IAtJhnhj?w%j7tACH%3CA`F!>y~f$%M*ZQ78;}*F2-d zZhM?x&i;1iUQ2cB*WGD&WWHmYfJx1(tc=3ecJ5p=t&x%zRV-dQcAR^Qm}QEa(;u;0 zCsD7*oP76^a*D(2X&+K`+}l75NSI!8cDu_F4x}Tg1bPc5N@uVwzJLEg>7Bu0MF2#X zjH9H)L>!ZW8o}3J&jU-b{lOe6@_?%mXvY=y6ZKWk5v7icj*RWdNcT*gOK1rnP?J+{ zXN`8FxCER{gNr!@242{Bj}?RcABv`WAfVA|wc*&j42> zj_h7G6@XHBzwV2c$YGVP1DjCde7oYv^s-FJX_06FdiMU1`4*f&R6q%Fe*gYGl#oCD z{<5jQ92oz?=GcD^@)J-E%31)}E*TEkyMf8fcJ^jz->HLDM~?kz*qCBY7mi zK->yqgdWpt^?wU~^Kky|OC!)=*~BuvIU(QeOQF{Yn1WkS83Wu~n%~vcl`I>6lsHM& z%eWVYva+kRv7iE)C;b%0J4~OSOYuY8v$zX!^v|BhcBH5Fu-ayY$MV|a#R6GM==rEQ zv+^`#XQ$?c+1X5KT%g@7EH4#t&&CmHvC{qFU&rmgg(>2E zw`($iL!#{LY;g&RgVM^2t?G+t9PQII9jkn;vPy%1YJk%1$l8@55Ivlw1 zA#Lbj!-`s|3Irhn_kp;S)MY)lXSY%gWv?_n)j z@xBO0yISd0FTkUs-qhnA`!roBs%nL}y_u?`%!^Ss-Ti`(7kY_>e$As*!M5O6fx&lQ*wZo0Nis>gxvas7-62AY z^@2K?l65<5Lw!)M{;K<>3>4_u>4xYC%$FPerdEo;ECAggTDEW93&mQd9TpYJE_$TL4IcoUqmH7GliM_PAEAvseN#k`Hdo)xU+9Z#kk- zF)`haM+c6se)3V?U;CVyx zN`j6jm+mq_hf?z-y>_1b-Ap)n*TAK%qmUj9!p4SWlAkb{~KJeS^B0t`q{@Cd92lW;2 zt|TJ~SQx?>goWZvd$Zo-4pqKJHH6Tean+%|p#5@A;n* ztWnr%myS7~(5F6S8j!f0H14He)?=>bH#I#3_iOj}^fY9zy+@ojW&8H+8?ZQB0PP); zs)%9+z`h}-pn#1M7Y0H#F0&y#XaJ4IGX;=8aiRm{2n6-0sHmQvp7H7*ascu0Kb^tI zjd&misriWJ2$ZT!)taepZf*j$4lt%cAf-k=X1p?*j1xhi8F-PN{)FepS{`G{h@Zdz z{bPYi+@C+x5fNe7LA|W$Tr!d3!UkISB$bIo{%Do|#-pi0pp#wr zk@26V4Eb$;nA?gjV}*V%;Ic3xufSNMZ)& zo1GXKTCXrcI|kB(BIrv2p?C*Sp#Rp<;n949jX}p>80Xag!=b%Xm+8i!IdTU@xTmja z_Ye3*0&E?iM@;s3ys>$@vw=i9WIZ#cy}FEmOL4Y24_RhDjPGc6go=P5M(~}m-`(fU@w4JbJDPd3gV>M@wiKu5q+|^`rQzm z$rr@D=n=PDg{(QGCC8Ur^kBP5sne%ZVkqh>*A2UqwxJD8+*Lr=5$NxQrvZ8kIJe|X zO{qZ}&*$-2;rZ+eBBxlTbQ4@-0U&h&HpQY<`)XsuRKMy<0~BH7foO7QxHQ56S1a;0 zd~PzF$6z@=zVcy?FKq>a<&=9l5$fiuRB`UrTW*Ib$%R%}8S?V#6z=8YN1v(agtC;f z<*AFtkM2*#iIJ}*>Vw^x>dH%cj&7 z{rMxA_>H{dUTxSIi-!l_Ut!=yjvqAo7X=bA;wP(V1VIwQ0sa#VY|2#pKd(#_nw@f_ z)PTH&>s?U5=iO_*Ter~-%!AIZhcNEv&$r9`6frnq0U$Of8Q9`Hi;BpcPDg(%)n7v1 zPU`H$t4a~4L7IY;lM~0c^DCfO3QQ3>7V(&Hkbe>9ClJvQ6ND_duiD;~Kc=bH)mkv& zjkpZCiH4%zL=bQ*-B~iFXq33c$m*@HCtqn7Onp;R`z=m(anQK3{sU8^pmr>ED^ z@#QmsBhbLa>QD4uXEWW~8i<&ljVEVk*8jhkyL1xfn{SB1jiEGPN&zuz@YNNs7w`*g zSN>E~bfk2|>1nUL)21?~X59WUBN@P?3Y0lM$X7eLDPU*!?cbqkqY#=e|DEj)tV|r4 z=3uLI7D0u(KvJu^@aQn=Ir)m?^jg}|Eu{6salpHyeJ9@pGe6#{$_w|tKO5*fGlD-HUeW|bVt1U| z8P~_%x8_V*+&;YZ!jvSf727-b1M6bgOlrX82DejSPbL>N5GL;!9rM1HUne2Qk^|Kt zAktt28^C-!(5G)S^M9lIuza($1R7PoRyh11nEU_U$0bxhtr4EzhvN??@+4V{Ff!qx zNSvR4V&D?Ut3(Z%5C>aiBUl$Uw;+nikP+&$jZKdUE0v%S*Zn0EY;Lo03?xIG;JdI4 ztTaDfx%Pwwydj%3TWf$OgE*?p+VgCQfr97`(m|Kz9JD&fDlXu{_|$5C|0v+?A-$E| zO<8fHqGl*)D+(%gyRinr4%O4!Q^Pr{&-zAQNn2LnCoWb_)UQ_83$W|mukJ-H59cBz zb52j2|7v>77e%_Z7maII%*Ro&dbqoe0_gP#6%Ke~FK>5{08b52#JvVerbnLnpYbCN z&PhnBD=#Gn%Ha(lcyJD7kj3S4s@l@&^_i|^WV$A`w6m0KiQ2s_SDMpQ^Ftc)9)~nO zr@WD6XHYOJ6e*s;#KV(xcIE_j6KAv=j@vaSeqaH8QOQ}cDNASxQCl`8tZ&X(wkQOB zx9!Qx)$HJ1aP^9zz#eA?9s?#aCYjwazte&WeZ9NJwwB}%!knoCDtdYfQqr8>8(^mn z70p@u@u2x!KJ#z6vU-I=|LXl}vV3OFx@)Z8A6)iLy_bhgw_(={auv^-1tZoIv>LUw zvtHC;6*wCH#n6Sr^`Vhw2KATaYuw9K%lofCw`@w|QIuJZ1`?>K=?nIicJd-_F@Ymh zgk!$b7w*5Ux&`w@uk*zGf&xV?t>fF}UA;p_&?z~Mh+cEyqbhkkWaDxD(Hf>WWbmVr zJI3h09OH-X=m|n&aaj)IxxJ8l-)Wnk@k;cuZN5_u&TF@c*aJRVqNs54@?uSuSzk>D zDmZxP`^yU~a74LmX9Yc7_ufJ^Sc(&hYI1o0Q|0TK)ZJ;MEWpHDX3ieca$GI>?A~&s z@rUMmdM3s1@?_cWH?XUbm5RL?T3h?7lp8>q0T`KE4tcEw+m5HDSGE3wLG|qBfb+j4 zk1qG8G|=+YhCCPJ!WF)GodMMY)n5NZHArul?a-=?N3y+MUm)D{e7)1DHjouSUGg^D zgz<7rV*7EtSo;TqE7%E;#+xP?*N|`xjW+YX_eDX?x=l^{<`h3X`s6)?}}o zR;{&(jhEWn;FyDy^uLMc^$_P8dWoUKDB=zNbqj}LsY;C36}mxRsNO`j2+zyi7L@s7 zO%NKNGhy|D58$tXH5-UVSK6c76Ds-ANnI^2mknFtU&;o`SPHrYhQm0%70(I%+`Ht0 zhb5S!OHHWS`|6A96Sr?R>N-mYmc(@FNnMQ3Q*H7eppx4Q#|KUO8AO+VD}3nMcqM=# zq9E}TO=%VRF6Dtg5M+~kZSC@{{&;@v&#=!{-X8uOo}@DTxbqG1WGXDn?dzISzDs9_ zC3^~>CO5JI3A#*cfo&J!RtGzrl!xcjSCgtaJky}xzt0*_RV5+*NaZtX*+3Y7g zp-oNPN-rvZ0~=O5KMr60i&|$s5&ejRLhHJ#XFTE>zH@4|9>pI@X*FcoHs6vo;ItA< zab@0qf3pI3EO_7tI-cP@4hD&WLZf8cOEpF;ld=qN{`q{x6b@;ceHM$+3|5@z>ywMLqJ>H zuwrWwjiiM)Qef8V>f(;uIyUEbHcLO=T|Jm>`L}0@q>}C-&U#zHu{QT~hm^LrywljP zEm@!wzuN>`wWcnlV9S#%|DRK6D|E>B-55{nl#xC_9X1@~r^vrV&-QOvLu8bPtK%ul zhs=~1FmnM$DMl}!+&)@e`RX8EI<37)`+1l>+mx7W)jaL#E!utm=$H2kX>!penSsIn zU$4DsO;)T_B8M5`4lE3Blh|w37zckozzd*L1}QF!9StZ&5W6t@>pgG@KbQK*0h4K3 z-?`%mTocs}rmOJJr&+Od&%JWc3%r8_t-MzrI>A8ufe;8FL%fTX7_gwmCX2f`?e@pRK6SsiV~UM zpu&}%)}<@Mvmeb!J!0@5?(93GNT-WgaOl^RyCI$Q9V;)lECbv|@c?c$yKCE)-;EQ^ z-?h1Ge`#EaS#!5!h!&`3(=aqO+D6EgNtNMlp8p~}ZGVM^L;Z9X^a9NzR={R`d2mei zfjZLF>#ee6eHmXXd!PVP@<(23#Udpzb@NNMpxL;(Djq97d$#afU@^cbCiSNz20hnm z!}uJV$H@_fW?@ZD>cFP`I9cjoq{@){0gxD?va;dr3}ZB(mc7t@S3h7`FlqcAH+j2N zR#qaKrIaa8O+{zURaJT4l-1{*)i-_<NPUDNu!ii z)|rc#=X0G=0?1#O5@&<|7LV0ymYmKXqUWHUw|Y57HBQ#tywT#lhF{%-h>Sd3 zin77Z0TUT}Hstq{s`L?+o>)rOi)SrYdDGJ~gqu z8(5XVAQEUNt7qIOqr?k~iV#>Q5Tdg8Sgz1{^xD~DuSb32%l1fEy01$Ywd_f;DgQ2xj9e2dK9+u`RK@gw`Um9*s6MAeq~Tga}&IWFjatKJ*cr zueQV-O7gz*aw7daVAJv(;Q$2z768rMug9TrIkv=|%)NDYcF`n#f^ zQBStRcN%`n_Mo!SYq@jM-l$6fZ=5QN^}ct=30#BTG>Z2SMnYurcvR zk@H1Sk@@PMojY7!Kz;4q@BLu9z31SxulXHJ&HF*w6)`-)_I9N{ovl!p91iVi7yuIH zt;&SY01?ROc=W@>^jEnQ#}8~-8t$bFL78;L0aT!*tyPy?|5_Q|2@?>Viqk(B!G3vI zHN*4zMBC&Pk-x7*-j67UP*jA9PnW%{yY4%aasG5ES8-5$3bC^5lL-3NsJ`d}ibPO- zAL=;dE6QqU1E;DY3rSb)zm;aVb=hra47}lG=OYw4Q|mnHw??XA6-sZ1K6Bvg$9Ssj z231YjhtLp4-*oPDF(x=w1JAd|=dPl6bu=&L6O0gwtykJs2Rqlj17`Z+hBmM9j z5dD|H%UD0O9TpI96V-a7|CijtYEp+D8U~IyD0eARu&JR*GQPb!9z2-t!0Q~H)A!oy zZRz8>7^GG>;`VclZObhwOr)?v717VlXlHZnZNUI^{rw2rgk081ODoh=_*=2~ew%-Z zNyFS{7!+$L-98uuWWw2qgx9v{F^-@2?57UKy!V)ams&-hzJ_HPgBE{-b-6&L`~{hA zC=`DUFrc8IsFT`~MU)YJKRZNPm5WJ5O}z;StBn}BS^b}9%f*oec2^w#8^qL(S(UAv zjU%p2en)^@}>yUW5cztvVfE*lyL<%yX3exh)WNmhHyKX3L(@K*(G-wW#1vnQfHH)3ojMi>JHMDlyLqi{0kbyG)$0@ytlE0DgBc zcNRR7!r`<}oL~A)L7}Vl_9gx9?yYPprvRA23dDH_1ZJenSSN3&aM2W+(UqwfE^T7W zpzUNvO`V}%KHu+T59QNWFnz!Lo6 z18icz`moj<@AE-HrR9dC6Y*z=g^Ha;hz0iIo_4#9JfYrW@SJw8)bFe8J@ez{a-cv# zyF)xR5Z^6m!)}WSB=e$SLypoo)QFAFTX2`{r)0pmW+b={(G^erogK=eGVYPK)AoE- z2Gv3{;2Lu|dJ);}c~q-0!SM1rL5OduM$A)?xIUncZ@oeW6RaA2ujz8C3_#bOC{|?% z4h{zO2pNwRzSNO&jNS#Vk$L#X+{eOB4`Cr8UjVad!0aroT7@?t9=(c5rgMQMlEY|K zLQ4z3w^t;z>DjX;#qQqbkpdI*9Tkn;j`|X~MTx~Ocjxz4kJq|N)E|IHG{%af1Q6br z3+;B!n~n8}jYKJ>A78`NON~d@w4v8u4&ixM_vn<3Wso4m{3m^2k#yW>~^Ked7Jzzyk0%QnM@s2H@Dv5c)i*O zpb?@D!f!ao4D?iNbkEmY5dxDea1&DTz3@whgkygE7_wqn&}-hYYuc`5RxrH3X$H(t zh2`Z1PN%D25KYS(=r$~l%W4~0d!?cDY}f3$Rn=%J&8-OED^g#N7bI;uy$)7hU8zyn zgHb)@`bxKgLgjMW9^S(QM36LTkeJG_61G0#aV$H+9Wt_kb!`ta#%AdiHVHEu-kBJ7 z+Cr?a{FX~!56<2DR=*&L{DzegMi%#lM?mN$jt9a3ZD%KQlVg}gQA|z{J@c1~%#c7) zi21`8wT3rGs^F8%t{HHMeE9HTP7^kkMxw!4pc_$9*^x))FHH1W5II<^b>fkMpLX((Z?R90r1}v!OLeke*zcJiDkUh@v1S*%eMdC|2GdC+IH% zrwuDt5ffL1#>PWuyYb%+^MC9g1v=%;stlOEL`jGU`~Utt$1&~)2G690v6%|YO2l&& zvQ+*<56TSDd@HLO&Ql40NQ?SUmMI^WGB5529-@I^06RN6{bljg@wbnjX+XO30v^J0_z5IXg=)2l~lflbnaOum?18sExW!3bcCXYhLpHC zJh%#C5)vOU@_-C0l;VyH4THAMK_Q#^m-kkB(;pklMYB$vOz_a0D7VOmcrFjvc8bc% z`ZfxMbCkw&XnK0{fTaZx#Mx5VNLor-`|rwgolfn^FcI`Cah56#a)qx1$uJdll*iJ` z=PrSI3KNNASzwi#m`RjSu1<>uEK`xmQaD_kxIm**qwNFg1&J9MRbhHX1!Ips6Btf? zcyQgEag2VqS+QU)o>B|D+z<3sFc?i87u?e$4E7$FBqyPu$d}tocRalaT)FCSpXLH6_WL;v(i8Vr>KB@L*n2J9)oM2CdvFE~$theypCWyFjN#?9aU z_n%(U$%I9P4{0!&uGcS21~4eBxpD{Cy6(gwN~ z7&8F&V?E|1C{?5}{YoI+n3$SkuJ(S!!z(N<7Rr)ex90pE91M|e!HouGO}zy#HDtsX zP~5;H(niF@6xAcF6H0cJ$eAV^fp%nWYT9*RkNN(+pBk+w876VU5U{p`XGR()2@!q# z{r!P`uya0{G)f#iQDC?HpR?;z%z!X(H59|HPEY&In2Kb{%WG@PYHFsQK|$|>>s2U< z1p|8W3JN~JU}pc*#H7Qr-heSR?S_h-9V<^E2xz2v3J7T4)6+^NDm3gV$3CYUVRfQ& zwU3v9DCkgMv*dM|Q_8gES-=$o2tI%8m^EwvF;j_}oR&6pwbkt%Az_3x1;m-f%9fVK zBS?$2WzG4vyiPD{{TCcTW?5PIM2;A^_ZT6)e|Ky*NU=hKf`U5rE5Tf`zb3sHbzHVW z5hD&V*#1Cy0vQGX@pt}$zrc@*lpu*yoNN=kH!uKzs}&Fs0J;|tF(W1OQu#c9tsa=j zw&p7h1n0h+rIvBZC@Cqe3=9M}tYYaHD7=%?)55~SZ%Lwl zM@FIpy=i7Ao5@8biW9xzj8$y0$z9v@0V-DFH-U|5;$5)VAZUP03_el;IDmkB z@J(5{H^r)Lpt!tEa$$g3Qc4Ok)gd78!vVN;KrDe94lJEuA_UlHf`WrFk>%i3Fr@Bm z)*`NsmsBH(1>cfDym5dfo1U9{mt`X|!KU_6o*E3<$Bye-5S6Mfx~!Cg`0DBsR8Ulm z8Zi8Y$Mf|O#5G2#%8)lS3=f4ObsQ!(wwQg5Sz=jpJXVk}_|yyp8CziFpYCeEAF+B1 zW3J6dMh0;Ms?p|yRrfeWP_k2CB}vatE|9~xL~1fhTC?41Tsi{-6zc}wwt#IX+=e|j z2xB51+D|T#sfPik;JQF?`RH_l&GQc5cgFMs0U5A>pqPR*gUxsr0zMk!9_TZFhlkBC znZP^(@JU1`CjJKFK?_wzU2WjwC$KKf$~Efg(N|_e0fYjT&g0?cTutW05;Y@4Y4Jhl zRS;U@q>Rq`xlg~cL1e@us|p+F!-#Z<93Yg|*L6Y42YHNwm^e$M^z`Kkl*!9j``0wJH;)0Uvmf)xR+$m-3je9FD)k)6l`;Smna5bvoF_E8X%IxKahU zniJS$l`u?Vd3|jya19M;76@22bBpOD01slZ`F3ORpTF<*Hy1ti0?knGe{S)a^_^h4 zy8pZ*NgHFXM8CfcrFjfFrnn zfAfP1R^TdZ&_e6U67RkNJNTL98=2EKOKy$2`|iswiAJ?5=?zh9b+ZH(UpxWqMVtT* zkG{hcmRh%(G!n^W+a zPu{$7YiamaJCZb>b_{IcZQm1|of&YwSDe(Pwb?PUjZH|q*~${QIe8al4_M)lnSI6 zj0}v-bPWu3jf_GJjI9jKtPD)F4GgRd4C>k5p=!v@PsvQH#I1oZIxGWdg9hA&lFZ!H b;*!MN0^E8c{mQok^)Pt4`njxgN@xNAzHJ + +其中 Input Embedding 层是用于将 token 编码转化为向量,该层训练定义域为 $\{0,1,\dots,v\}$,值域为 $d_m$ 维向量的函数 $Emb(x)$,其时间复杂度是线性的。Input Embedding 层输出 $b$ 个 $l$ 行 $d_m$ 列矩阵组成的矩阵组,矩阵组的第 $i(0\leq i -1: + x = x + self.pe[pos][:x.size(0), :] # x: [batch_size, d_model] + else: + x = x + self.pe[:x.size(0), :] # x: [tgt_len + 1, batch_size, d_model] + return self.dropout(x) + +class ScaledDotProductAttention(nn.Module): + def __init__(self, d_k): + super(ScaledDotProductAttention, self).__init__() + self.d_k=d_k + + def forward(self, Q, K, V, attn_mask): + scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(self.d_k) # scores : [batch_size x n_heads x len_q(=len_k) x len_k(=len_q)] + scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is one. + attn = nn.Softmax(dim=-1)(scores) + context = torch.matmul(attn, V) + return context + +class MultiHeadAttention(nn.Module): + def __init__(self, d_model, d_k, d_v, n_heads): + super(MultiHeadAttention, self).__init__() + self.W_Q = nn.Linear(d_model, d_k * n_heads) + self.W_K = nn.Linear(d_model, d_k * n_heads) + self.W_V = nn.Linear(d_model, d_v * n_heads) + self.linear = nn.Linear(n_heads * d_v, d_model) + self.layer_norm = nn.LayerNorm(d_model) + self.d_model = d_model + self.d_k = d_k + self.d_v = d_v + self.n_heads = n_heads + + def forward(self, Q, batch_size, attn_mask, k_cache = None, v_cache = None): + residual = Q + if attn_mask is None: + q_s = self.W_Q(Q).view(batch_size, self.n_heads, self.d_k).flatten(0,1) # q_s: [batch_size x n_heads, d_k] + k_s = self.W_K(Q).view(batch_size, self.n_heads, self.d_k).flatten(0,1) # k_s: [batch_size x n_heads, d_k] + v_s = self.W_V(Q).view(batch_size, self.n_heads, self.d_v).flatten(0,1) # v_s: [batch_size x n_heads, d_v] + k_cache.write(k_s) + v_cache.write(v_s) + context = k_cache.transmatmul(q_s) + context = v_cache.matmul(nn.Softmax(dim=-1)(context)) # context: [batch_size x n_heads, d_v] + context = context.contiguous().view(batch_size, -1) # context: [batch_size, n_heads x d_v] + else: + q_s = self.W_Q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1,2) # q_s: [batch_size, n_heads, len_q, d_k] + k_s = self.W_K(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1,2) # k_s: [batch_size, n_heads, len_q, d_k] + v_s = self.W_V(Q).view(batch_size, -1, self.n_heads, self.d_v).transpose(1,2) # v_s: [batch_size, n_heads, len_q, d_v] + + attn_mask = attn_mask.unsqueeze(1).repeat(1, self.n_heads, 1, 1) # attn_mask : [batch_size, n_heads, len_q, len_k] + + context = ScaledDotProductAttention(self.d_k)(q_s, k_s, v_s, attn_mask) # context: [batch_size x n_heads x len_q x d_v] + context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_v) # context: [batch_size, len_q, n_heads x d_v] + output = self.linear(context) + return self.layer_norm(output + residual) # output: [batch_size, len_q, d_model] + +class PoswiseFeedForwardNet(nn.Module): + def __init__(self, d_model, d_ff): + super(PoswiseFeedForwardNet, self).__init__() + self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1) + self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1) + self.layer_norm = nn.LayerNorm(d_model) + + def forward(self, inputs): + residual = inputs # inputs : [batch_size, len_q, d_model] + output = nn.ReLU()(self.conv1(inputs.transpose(len(inputs.size()) - 2, len(inputs.size()) - 1))) + output = self.conv2(output).transpose(len(inputs.size()) - 2, len(inputs.size()) - 1) + return self.layer_norm(output + residual) + +class DecoderLayer(nn.Module): + def __init__(self, tgt_len, d_model, d_ff, d_k, d_v, n_heads): + super(DecoderLayer, self).__init__() + self.dec_enc_attn = MultiHeadAttention(d_model, d_k, d_v, n_heads) + self.pos_ffn = PoswiseFeedForwardNet(d_model, d_ff) + self.n_heads = n_heads + self.tgt_len = tgt_len + self.d_k = d_k + self.d_v = d_v + + def forward(self, dec_inputs, batch_size, dec_enc_attn_mask, block_size, k_cache, v_cache): + if block_size > 0: + dec_outputs = self.dec_enc_attn(dec_inputs, batch_size, None, k_cache, v_cache) + else: + dec_outputs = self.dec_enc_attn(dec_inputs, batch_size, dec_enc_attn_mask) + dec_outputs = self.pos_ffn(dec_outputs) + return dec_outputs + +class Decoder(nn.Module): + def __init__(self, tgt_vocab_size, tgt_len, d_model, d_ff, d_k, d_v, n_layers, n_heads): + super(Decoder, self).__init__() + self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model) + self.pos_emb = PositionalEncoding(tgt_len + 1, d_model) + self.layers = nn.ModuleList([DecoderLayer(tgt_len, d_model, d_ff, d_k, d_v, n_heads) for _ in range(n_layers)]) + + def forward(self, dec_inputs, batch_size, k_caches, v_caches, block_size = 0, pos = -1): # dec_inputs : [batch_size] or [batch_size, tgt_len] + dec_outputs = self.tgt_emb(dec_inputs) + if block_size > 0: + dec_outputs = self.pos_emb(dec_outputs, pos) # [batch_size, d_model] + dec_enc_attn_mask = None + else: + dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1), pos).transpose(0, 1) # [batch_size, tgt_len, d_model] + dec_enc_attn_mask = get_attn_pad_mask(dec_inputs) + + for layer, k_cache, v_cache in list(zip(self.layers, k_caches, v_caches)): + dec_outputs = layer(dec_outputs, batch_size, dec_enc_attn_mask, block_size, k_cache, v_cache) + return dec_outputs + +class Transformer(nn.Module): + def __init__(self, tgt_vocab_size, tgt_len = 10, d_model = 512, d_ff = 2048, d_k = 64, d_v = 64, n_layers = 6, n_heads = 8): + super(Transformer, self).__init__() + self.decoder = Decoder(tgt_vocab_size, tgt_len, d_model, d_ff, d_k, d_v, n_layers, n_heads) + self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False) + self.tgt_vocab_size = tgt_vocab_size + self.d_model = d_model + self.tgt_len = tgt_len + self.n_layers = n_layers + self.n_heads = n_heads + self.d_k = d_k + self.d_v = d_v + def forward(self, dec_inputs, block_size = 0): + batch_size = dec_inputs.size()[0] + if block_size > 0: + k_caches = [] + v_caches = [] + dec_inputs = dec_inputs.transpose(0, 1) + for i in range(self.n_layers): + k_caches.append(Cache(math.ceil(batch_size * len(dec_inputs) * self.n_heads / block_size), batch_size * self.n_heads, block_size, self.d_k)) + v_caches.append(Cache(math.ceil(batch_size * len(dec_inputs) * self.n_heads / block_size), batch_size * self.n_heads, block_size, self.d_v)) + dec_output = torch.empty(self.tgt_vocab_size, self.d_model) + dec_outputs = torch.empty(batch_size, 0, self.d_model) + pos = 0 + dec_logits = torch.empty(batch_size, 0, self.tgt_vocab_size) + for dec_input in dec_inputs: + for i in range(0, dec_input.shape[0]): + if dec_input[i] == 0: + dec_input[i] = dec_logit[:, :self.tgt_vocab_size - 1].data.max(1, keepdim=True).indices[i] + dec_output = self.decoder(dec_input, batch_size, k_caches, v_caches, block_size, pos) + dec_logit = self.projection(dec_output) # dec_logits : [batch_size, tgt_vocab_size] + pos = pos + 1 + dec_logits = torch.cat([dec_logits, dec_logit.unsqueeze(1)], 1) + for i in range(self.n_layers): + k_caches[i].delete() + v_caches[i].delete() + else: + k_caches = [None] * self.n_layers + v_caches = [None] * self.n_layers + dec_outputs = self.decoder(dec_inputs, batch_size, k_caches, v_caches) + dec_logits = self.projection(dec_outputs) # dec_logits : [batch_size, tgt_len, tgt_vocab_size] + return dec_logits \ No newline at end of file diff --git a/engine/test_kv_cache.py b/engine/test_kv_cache.py new file mode 100644 index 000000000..a90c2315a --- /dev/null +++ b/engine/test_kv_cache.py @@ -0,0 +1,51 @@ +from torch import nn, optim, LongTensor +from TransformerDecoderOnly import Transformer + +def make_training_batch(sentences): + output_batch = [] + target_batch = [] + for sentence in sentences: + output_batch.append([tgt_vocab[n] for n in sentence[0].split()]) + target_batch.append([tgt_vocab[n] for n in sentence[1].split()]) + return LongTensor(output_batch), LongTensor(target_batch) + +def make_learning_batch(sentences): + output_batch = [] + for sentence in sentences: + output_batch.append([tgt_vocab[n] for n in sentence.split()]) + return LongTensor(output_batch) + +if __name__ == '__main__': + + sentences = [['S i want a beer', 'i want a beer E'], ['S i want a beer', 'i want a beer E']] + + # Transformer Parameters + # Padding Should be 0 + # Unknown Should be tgt_vocab_size - 1 + tgt_vocab = {'P': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, 'S': 5, 'E': 6, 'U': 7} + + number_dict = {i: w for i, w in enumerate(tgt_vocab)} + + model = Transformer(len(tgt_vocab)) + + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=0.001) + + dec_inputs, target_batch = make_training_batch(sentences) + + for epoch in range(20): + optimizer.zero_grad() + outputs = model(dec_inputs) + loss = criterion(outputs.view(-1, outputs.size(-1)), target_batch.contiguous().view(-1)) + print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss)) + loss.backward() + optimizer.step() + + # Test + sentences = ['S P P P P'] * 10 + dec_inputs = make_learning_batch(sentences) + for epoch in range(1): + predict = model(dec_inputs, 16) + predict = predict[:, :, :len(tgt_vocab) - 1].data.max(2, keepdim=True)[1] + for sentence, output in list(zip(sentences, predict)): + print(sentence, '->', [number_dict[n.item()] for n in output.squeeze()]) \ No newline at end of file