From 8cfc1cae58cbf40314c1f63a3a6b1e3a7a4134c1 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 10:23:02 +0000 Subject: [PATCH] Canonical encoding: derive special token IDs from official encoding module + tokenizer - Remove hardcoded THINK_START/THINK_END/USER_TOKEN/ASSISTANT_TOKEN IDs - Import token strings from encoding.deepseek_v4_encoding (official source) - Resolve IDs via tokenizer.convert_tokens_to_ids() at runtime - Use parse_message_from_completion_text() for structured output parsing - No more hand-rolled prompt construction or hardcoded token IDs - Clean up TEMP: replace old deepseek_v4_ref with dsv4thing.zip reference --- TEMP/CROSS_REFERENCE.md | 100 -- TEMP/deepseek_v4.zip | Bin 99035 -> 0 bytes TEMP/deepseek_v4_ref/deepseek_v4/__init__.py | 30 - .../deepseek_v4/amd/__init__.py | 2 - TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py | 972 ----------- TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py | 509 ------ TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py | 856 ---------- TEMP/deepseek_v4_ref/deepseek_v4/attention.py | 806 --------- .../deepseek_v4/common/__init__.py | 2 - .../deepseek_v4/common/ops/__init__.py | 28 - .../deepseek_v4/common/ops/cache_utils.py | 594 ------- .../common/ops/fused_compress_quant_cache.py | 666 -------- .../deepseek_v4/common/ops/fused_indexer_q.py | 438 ----- .../common/ops/fused_inv_rope_fp8_quant.py | 318 ---- .../common/ops/fused_mtp_input_rmsnorm.py | 203 --- .../common/ops/fused_qk_rmsnorm.py | 96 -- .../common/ops/save_partial_states.py | 101 -- .../deepseek_v4/common/rope.py | 36 - .../deepseek_v4_ref/deepseek_v4/compressor.py | 380 ----- .../deepseek_v4/nvidia/__init__.py | 2 - .../deepseek_v4/nvidia/flashmla.py | 424 ----- .../deepseek_v4/nvidia/model.py | 1476 ----------------- .../deepseek_v4_ref/deepseek_v4/nvidia/mtp.py | 516 ------ .../deepseek_v4/nvidia/ops/__init__.py | 13 - .../nvidia/ops/dequant_gather_k_cutedsl.py | 331 ---- .../nvidia/ops/fused_indexer_q_cutedsl.py | 610 ------- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 173 -- .../ops/sparse_attn_compress_cutedsl.py | 1380 --------------- .../deepseek_v4/quant_config.py | 158 -- TEMP/dsv4thing.zip | Bin 0 -> 20388 bytes TEMP/dsv4thing/README.md | 26 + TEMP/dsv4thing/config.json | 35 + TEMP/dsv4thing/convert.py | 168 ++ TEMP/dsv4thing/generate.py | 155 ++ TEMP/dsv4thing/kernel.py | 536 ++++++ TEMP/dsv4thing/model.py | 827 +++++++++ TEMP/dsv4thing/requirements.txt | 5 + single_shot_inference.py | 56 +- 38 files changed, 1799 insertions(+), 11229 deletions(-) delete mode 100644 TEMP/CROSS_REFERENCE.md delete mode 100644 TEMP/deepseek_v4.zip delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/__init__.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/__init__.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/attention.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/compressor.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py delete mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py create mode 100644 TEMP/dsv4thing.zip create mode 100644 TEMP/dsv4thing/README.md create mode 100644 TEMP/dsv4thing/config.json create mode 100644 TEMP/dsv4thing/convert.py create mode 100644 TEMP/dsv4thing/generate.py create mode 100644 TEMP/dsv4thing/kernel.py create mode 100644 TEMP/dsv4thing/model.py create mode 100644 TEMP/dsv4thing/requirements.txt diff --git a/TEMP/CROSS_REFERENCE.md b/TEMP/CROSS_REFERENCE.md deleted file mode 100644 index 7affd75f..00000000 --- a/TEMP/CROSS_REFERENCE.md +++ /dev/null @@ -1,100 +0,0 @@ -# vLLM Independent-Reference Cross-Checks (DSV4) - -**Purpose:** catch the architectural bugs your cosine backlog is blind to. Every check in `/home/openclaw/dev/nvfp4-megamoe-kernel/DEGENERATION_TESTS.md` and `/home/openclaw/dev/nvfp4-megamoe-kernel/CORRECTNESS_BACKLOG.md` compares your kernel to your *own* PyTorch reference — if both share a wrong assumption (RoPE position, RoPE style, norm ordering, compression overlap), they agree at cos 0.9999 while both being wrong. vLLM's DSV4 is an **independent** implementation; it can't share your blind spots. Use it as the tiebreaker. - -**How to use this:** these are mostly *static code diffs*, not GPU runs — fast. For each check, find the line in `single_shot_inference.py` (or `dsv4/`), find the cited line in the vLLM tree, and record match / mismatch + the fix. - -**Ground rules** -- The attached vLLM tree "works but is unstable" and strayed from the paper in places. So trust it as an oracle for the **math** of the compressor / mHC / RoPE / norm-ordering — it produces coherent output against the real weights — but NOT as a stability or paper-fidelity reference everywhere. -- **DO NOT mirror vLLM's attention.** `attention.py` / `nvidia/flashmla.py` shoehorn CSA/HCA onto the V3.2 MLA-sparse backend + a bolted-on SWA cache. That's the unstable stopgap. Your from-scratch shared-KV MQA + grouped-output is more paper-faithful — keep it. -- Where vLLM and your code disagree on a detail your cosines can't see (below), **vLLM is the tiebreaker worth investigating** — match it and re-test. - -Reference files: `deepseek_v4/nvidia/model.py`, `deepseek_v4/compressor.py`, `deepseek_v4/common/rope.py`, `deepseek_v4/common/ops/*`, `deepseek_v4/nvidia/ops/*cutedsl.py`. - ---- - -## Cross-check 1 — Final tail ordering: collapse → final RMSNorm → lm_head -*(Resolves the Test-2 / mHC question. If this is wrong, it explains the |X|=860 confusion.)* - -**vLLM reference (`nvidia/model.py`):** the expanded mHC residual stays `n_hc × d` "until `hc_head()` collapses it" (comment ~line 1130). Tail sequence: -- `mhc_post_tilelang(...)` collapse (`hc_head`, `n_hc·d → d`) — ~line 1172 -- `hidden_states = self.norm(hidden_states)` where `self.norm = RMSNorm(hidden_size)` — ~line 1191 -- `compute_logits` → `lm_head` — ~lines 1444–1448 - -So the canonical tail is **collapse → final RMSNorm → lm_head, in that order.** - -**Check:** trace the last layer's residual in `single_shot`. Does it (a) collapse the expanded mHC residual via `hc_head`/`mhc_post`, then (b) apply a model-final `RMSNorm`, then (c) `lm_head` — in that order? - -**Decision** -- Missing collapse, missing final norm, or reordered → **real bug, fix the ordering.** This is the legitimate "the final norm is broken" branch from the decode runbook — a small fix, NOT residual clipping. -- All present and correct → RMSNorm is scale-invariant, so |X|=860 is normalized away → **mHC growth is not the cause**, corroborating the Test-2 falsification. - ---- - -## Cross-check 2 — Compressed-entry RoPE position (block-aligned) — HIGHEST PRIORITY -*(Off-by-(r−1) here is invisible to same-input cosines and wrecks long-range attention.)* - -**vLLM reference (`compressor.py`):** the position used to RoPE a compressed entry is -``` -(positions // compress_ratio) * compress_ratio # block-aligned = FIRST position of the block -``` - -**Your earlier code used:** `((bi+1) * r - 1)` — the **LAST** position of the block. Difference = `r − 1`: **3 for CSA, 127 for HCA.** A constant 127-position offset on every HCA entry degrades long-range attention while every per-layer cosine stays high (your reference makes the same choice). - -**Check:** what position does `single_shot`'s compressor assign when RoPE-ing compressed KV entries **and** the indexer keys? Compare against the block-aligned floor formula. - -**Decision:** anything other than block-aligned `(pos // ratio) * ratio` → change to it and re-test. This is the single most likely concrete bug in this document. - ---- - -## Cross-check 3 — RoPE style: GPT-J interleaved, not NeoX split-half - -**vLLM reference (`compressor.py` compress kernel comment + `common/rope.py`):** -- `is_neox_style=False` → **interleaved pairs**, NOT split-half -- applied to the **last `rope_head_dim`** elements of `head_dim` -- `cos_sin_cache` layout `[max_pos, rope_head_dim]`: first half cos, second half sin (per-pair, each length `rope_head_dim // 2`) - -**Check:** is `single_shot`'s RoPE interleaved (GPT-J) or split-half (NeoX)? Verify **both** the forward RoPE on Q/KV **and** the inverse RoPE on the attention output, and that the cos/sin cache layout matches. - -**Decision:** split-half where the model expects interleaved (or a mismatched cos/sin layout) scrambles every roped vector — invisible to same-input cosines, wrong against the real weights. Match vLLM's interleaved layout. - ---- - -## Cross-check 4 — CSA overlap (2m) vs HCA no-overlap - -**vLLM reference (`compressor.py`):** `overlap = (compress_ratio == 4)` and `coff = 1 + (compress_ratio == 4)`. So CSA (ratio 4) uses **overlapped** compression — each compressed entry derived from `2m` KV entries via the `C^a` + `C^b` overlap (paper eqs. 11–12) — while HCA (ratio 128) does **not** overlap. - -**Check:** does `single_shot`'s CSA compressor produce overlapped blocks (each entry from `2m` entries with the `C^b` overlap), and HCA non-overlapped? - -**Decision:** if CSA isn't overlapping, the compressed representation differs from what the model was trained on → degraded attention. Match the overlap behavior. - ---- - -## Reporting - -Same harness/workflow as the decode runbook: edit locally → commit → push → pull on B200 → test; `TEST_LAYERS` as an **env var**, never a CLI arg. Most checks here are static diffs (no GPU). For each, paste: `single_shot` line, vLLM reference line, MATCH/MISMATCH, and the fix if mismatched. - ---- - -## Appendix A — vLLM reference → your kernel map (for the kernel work, not the bug) - -Independent CuteDSL/op references for the kernels you hand-wrote. BF16/FP8, not NVFP4 — so a *correctness/layout* reference, not liftable code (and useful evidence for the NVIDIA gaps writeup): - -| Your kernel | vLLM independent reference | -|---|---| -| compress → RMSNorm → RoPE → FP8 quant → cache-write | `nvidia/ops/sparse_attn_compress_cutedsl.py` (`compress_norm_rope_store_cutedsl`) | -| dequant + gather K (your `gather_mixed_*`) | `nvidia/ops/dequant_gather_k_cutedsl.py` | -| indexer Q | `nvidia/ops/fused_indexer_q_cutedsl.py`, `common/ops/fused_indexer_q.py` | -| inverse RoPE + FP8 quant | `common/ops/fused_inv_rope_fp8_quant.py` | -| QK RMSNorm | `common/ops/fused_qk_rmsnorm.py` | -| uncompressed tail / partial-state buffer (your decode buffering) | `common/ops/save_partial_states.py` | -| MoE dispatch prep | `nvidia/ops/prepare_megamoe.py` | -| mHC pre/post | `vllm…kernels.mhc.tilelang` (`mhc_pre`, `mhc_post`, `mhc_fused_post_pre`) — used in `nvidia/model.py` | - -## Appendix B — determinism gotcha worth stealing - -`compressor.py` disables PDL (`launch_pdl=False`) because the compress kernels consume a preceding GEMM's output without emitting/waiting on PDL grid-dependency primitives, and `launch_pdl=True` caused a read-after-write race → non-deterministic output. You target bitwise determinism, so this is a hazard already mapped: any kernel that depends on a prior kernel's output without an explicit grid-dependency primitive must not use PDL. - -## Appendix C — what NOT to take - -`attention.py`, `nvidia/flashmla.py`, and the `DeepseekV4SparseMLA` / `DeepseekV4SWACache` machinery are the MLA shoehorn — the unstable stopgap that motivated this whole project. Use vLLM's compressor / mHC / RoPE / norm-ordering as oracles; leave its attention impl alone. \ No newline at end of file diff --git a/TEMP/deepseek_v4.zip b/TEMP/deepseek_v4.zip deleted file mode 100644 index 77ae49cdf8c4a9b3e420ce54661e58bdc79997c5..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 99035 zcmb4qV{k7{yJT$JIr$PL?gQY-u#~h4u}ZI)XdD$#mvmcz@3%B%@Y;~7z_Xk1oZDkO%)yp z5_g}_TIPt*TI8Rp_-`aAAgKQf3G+W785mgETe%t-&^vnlSD^no!@oe|l(+1c1(E#r z8nDP1h>$}R2^Facs!PKnTSZBu1=Y+PR*y_4ygfGbD&Q4-!3-uc7~50Y&%qBUZguCl zT+F)}^WX+v#!45!uHu|olVF7RB2Xz-V0;`>IEuH_^ODAb+;kMC)p8Hyu12m=Q`r`m zi=$T|BQXwPj93T?LLOguswuhi=8zWUkk=Rmf|pJ!59S;88k6}4(g}RfYSWQj43Yg| z=aQqMkXXc%snwN1)8tJit*|F?7FvgZ(k2&0(zD=CNDVBoB~u=C(zyiUwV=bY>z63; zs@yDF+p6n?C~F-^svq);hMJ(_)8OHzhjA_ZY4bE8>UEX3h(!gIC2m<) zW?_CO@($e!SJ5fqY&(nB);4FLJOTpTIlZ{7R6{m-jriBP0G|#Q!tPwbzQCAwOHS^f zunh6KKy40Ty&^1%abb0CXGL+!4trQxjQE^lh62cw{(1WapQ#P5=)tmn;{a*fnc_h~ z8`mawFrTUP_(ja@`AE`+maX5Pf-)w2`VSy)4eDNLW>JNQ3tanUIG^jKp~4q+)CRp* zu=@>dwS0M{dIu)cwS*JYl6?FD#T4D2_{-%-zU;`Apu1=LE?q-( z(sj#dyCY^HYg~Zq-fmd_&I}4=k7h0}vTqPY8AvGh)7x`e5Fj8qFd(4+O~`*S=|55% z=|2$C$jx!|lSO5oxQNZGzj*>j}Q!4yKX$;#8lUx-& zLH)p>%IVj2+PsONPuEZ9{A>I>eg%FmAEN~c`W7_$f4%@UshD!@bDnN9`M1d{;e>n$ zgu6Wq&V5tP-^JTIsvD!_UgHfpCaZnTYFeyn?Y241)mkde<|7*`munQZ8fgIsHK+RR zJ85bo**SXZy4yM$Cm2)5T%=}ZN)eOJqo(PYQu!3qxeV(X@C|nv6Zu0rD(PNz%;s(g zfqN=y8ehQ%Ev4avdNyWln^+7+N`~phofOo4Q;&8~gAFaPYS$fi=iq*e&81p&ALA-% z44vcqUHfvqTcen28#fp`84!0`>i6xkUIe$vFx_AR)uo-hsTKz6UN~zT3BFDa^}s~U z_cr~)coAOCBQ%1zg}2l~Xw;SclMgRmN%^hs!k>TjR&DiLp%wT|Ot1#Qi zRoDHDHF(lU^2EA(wF5c+EMda`7Wq7twGLp@IHA^hK}~kZ;FSxqoSk#^a$DgKGU^5na{ARPE?Ck3EsXU$kq1 z>4kf?!5&h|HuDhv7xS~5mNN!mP>@#1i$QX^F0vaEeH zO`e$u^NtS*XzrJO(4M9tCvgKqv0zDJwvhM{rz3pLvk5{8_F6x1^KX#pneW&3(a}qU zNkZ>wHtQ!e9OxsX;@gBFP;Q(9fMMRaDM0f_txbA?)l<~L>2 zA!Dl`9nz^(u4|6I6&#RMJp2VDwn0U!k(N78|- z7Zbb&VI)M#ACkvs%u&q#XhYU3}ZLHrGe4i>x zu=i1|3st$}R~7vq1tKry6AoDIKwz^UAcNaMo-YS8mFt1Z76F=_l(;bw`ywR?&aQfc z3qGk^f(;c8K)=_1XQNpNQvn&?23dn1{=M?vT_dVE4<1!qFIw5}c6iq|L$U*BE=U)u zDaDoP3y#_H4vT>ix$u!6c4@j!gC&4W=C~1`ZY!_`?s;6Ki!)J8*>1uEj(*bmW4YYP0Te#u9 zvEULCmg59P8!4=dvB~BxCFFaQKpg^V?8P+Y;A5s7+orw*^t}ZXa+=SQ$|?3B#i7XG z%&!_G=5K?59x>iF)k12sLggUkf>M6rr4s>Jya*+v0AN-7DYQ2*`8%!{!EJQ2CH?VP z>r*9cgv7u>u4mOt7;*hrmVjvsBuAny{fr-2^LD`ff}lXRYgh- zKi?huSp(RNYadL<&HPuImyJk9v}=7%$~s>5gb3KFC9zI4H-uTa6|YyeF^Dfye{49{ zMugO3{sSSOK1E8ME^rMOftmDJ1wDbn|oI0~CBTfz?KU zaGMaK$T20FI;O)=Q6-@oEvB@1n_7a9U3S?H9LW95r3n>TScV5N+Z}BiiY%eUt;e6o zClDBXxOnz5FYJlFwYKAo==coNaA5#0lJM15Mc+kK=1>_5S>dO($j%#O4J?{ZDFEj$ z;&}~dS8f`y<$x#!cq0ZAR69!TCwmwyy3gOPiJ9tl!Ndz%^mv6)f2?emc4Ht5 zr#_=0nXX?)v%fl?q+BUoBE~8|*=nyhrlJ;;Puiek!_NlDU}PC@gF&`zj?V=>)9(Ei zW(T^)8=iFQ@cdkFm4A}fI}bIc-TTInLXDztVT>CTPa7V8CWS$n{6tBE9OB5dg4ek)bCY<+`d z7lKtC#3mEEf4uO$EmBOQOj9%J|FJx!_v@N2+Fp0gHP?YJdx*If*}h+&{rYBRSFYXF z4x7P9;APudFJ`OPzKwW09JkKcqU8%+}TB|ATV^WI{oEnq@K3jB6)1A!(R^lCd z%fH8;^^wIf;GQrjN2t332pNjDMV~I%x~rrwco~&GVHxmCPB9Qy4oL`jI1k?jdLN4G zL;R2&%yMPNw}8yK$74}w^-HC*9tV-f2AGV2K}Uq3GNVt|oqN)RSo2*aV>6^$HYuad zlltVFP7WHww}|a?oxUI}^$<$XH7*o2i(*2zKITIpB@jmfT44;wJo{EKaGx0^p&vJD zr#W~$R=quX5&M=;=*T4v;*cBvKnpSMXP0&x`9&TN4ER2uUdyoL#)E&#wQnj`B5Cf; zZrc-H3Rc0%eU1QT7GL2*gq2u_v*6m8E^v^sUe-WW$9isPzHqO}-^X}dzfC&JL8Muxi1OAgCy_WG zBpG8@SSX~cMK3)I0d|Fi2$qFR^+}Ro>SOaliHO#vw`@9@?sPow#40!0+U1pe)3`^o ztWXeK5Uj2J#*yLKDimBUH)EA_4j0Dh%I?^~|C6k$aMzI)udV)vDpelxbNFGD8h%Rf zPZhZQX0uF-^^I)SM60cBe?6F_DKwhTdsrmt(t9?C1aI3Ik5b=BQ%tP~EKlDEaWfy1 zM+ST9&r9r$`&d%+vPJ=7edtqqLu&1dwPc2o?PTdI)#)UAQ~@>o*rmMNSK*SAZR2!d z$hdZ`*1FwW@1((ezS&z_s=;gno9K>dA--jl{9&A;Dp|i49Gg6LIS<7QQJsybIeoOUl z6|#W7Hdn)KU2*Rk0$JT@j@n8>o#kPU1Y_;GRrFr%gIdh-2re`(KoLGjMDH_79KeIM z!a@gA5M3Gm23+%|yCXVEf5~|Sh8sqG`rT|pZiMZDJO`u%!rG~QINd;fr0jw11n3CT zfD>K3z(Drn4MLe=Z0cQ2|GN5RoTV3biec81+CPLMqSdP5M-vz5C^^(apy~o#YD}6m zW<`oWIhIj-ZuduvC=p!5b{VnlF1`o;7wBjZo48wgyA*+5j!6eS3CQ*Xkfh_U=MgD! z%U|z%+3-Z>LY-pqGj10^M4!3Bq}_&;Ro;{(gYY+f<$HMK)9-gmk^ zAnV5NHM_$^P?zB-xwunUe8dNBbd)#^dZypAu#q{RueT$ARNy~DJ_pCFVY?JKBqT{? z=wQ1PSCo_`;Q= zi6=LGlx^kOP$LR17 zB}Oij9Piud8)}clM#{-6w2B#s!5OGLmqwgMUx}C`teSf?fn*eM*1en>%Ith|RAL%r zy%|)_&0S8-H$a00qxGi40e+G3vBnP?kevHZRkrCK~C$_h_wyBpUcGLM0cXW zLs&e0BXiZk|B3+k&N(5;I<3qnL&@?!Q*|xZ4M?dRK9Qr*Beg{)3i-l5{ za41PIPL62)aWtKR`c{R`0YVdM*tNe^=bnp4CwP8F3vQi-VfpwVePyrXO3d|OOU{i| z_Opq(NUa=pR@C*8kiH`qw~7}|P(u)Ta}H}~5Xq@F(>A*8u-Cip&ICBgwp*A^TXCCi zV>Gjgg495Rdtf=txEuPrzPf0EC8?QKQ+KLA0y)chS}z9S53c8lMK#IMIVHtLj;JDW)^u9%bNUBUXxp6Seo2b z6lW>7&4WrZCcG;} zr}ST(^({lUEnezHE+)yV!t(c0o*6E5OY>q+;Nsmp+_Hm@raiJL@<8yS)SMI8 zao4y|?}`ZybY79d+e;lVnN1id+r;RU_Y$!_XY7(cLsfQks6MAh>%#Ygzmw7kG8u_} zDC4h0(#hO}xUUMGNLKFNC#Dtuv?1!e;NEUG0RLRCE)x6&;j!yAOyvm_%Iov#L8nc7 z=I5z$pT2;q@=gpKQpCg@*+FsMt$(}*G%xlF#t-RmNREDm2>jRt9y z-d!IdNa6vcTz161Ves?$vv0EGtKIT9TmuAsa6PDMVuFdhLn(*CF4#W8BGZz>a5o!u zbj!S=^V>)bv;a|^<0*=L_!k6ST2M>;*b|a3!h-?>wJc&t;%pur<#ikI>)^yBOgB^Y z^rzrsnrmG0cKHw^;%A~nH$oTB2R!SlQFo{Scy*X;P$#{T$gk%ym%_Uy`KyXCEwY`W zmYULPd&|vVDzovVe{qG9wAl=qBSi>6a;2SNm`0PI1s}PR<~kIx2X@KCqoA%^Fpj_&l+K8y(}3E+8TrWHan~_D7wnpAqNGA! zEPWC!l-CsIa`Uk%v#XrmlR%g63sfD|}HZ2=q0oSW;M9fdr0fWVv$0;}p$9n;Ksq2+cFC~_OMUo+nn5qY?`0=P}Lg&=FV zts;8Q8%mj43+Dx;UvoYi%QA{TNm3Sw^YAy9i@|g8pVD(8)`OyC>v*j~jGnA(-TP+0 z`s@z&cu%u3(^t5YIh)6JHo^$|C-?|&iR;4Q@#0C(!mD0gmyi^E64%@)XO;W$Jvcsc zqr&icK__B2xhSAjA0d`oPp~ab-J)#pEG37x<=s6Kkbr#5pPu+(DJ66RySk7#z_12d z{+K=u&G5SB&%Pw=--^gw`84wlndIl$0`A)pMGOxwzxnJbUsd_k!~J_^pS~T`?jERj zV?+Y8@~NpEw{rD^zRE`yQnK;lS#^reaDOI6xPeZp+H1HoUo9(WbpJ>ac$mncHZXC_ z0hthGv~2D#lI^^I74hJiQXz0Cx-uO&{v$Y2b_@JQoX_GGon6o76ij}w79*sz6ur2Z|ue7ET_X_IRr45Qvzn!81!Fk)k&?7#r(uov#3587vH}{jh#)w+e*3% z0lS&LnFbWfUM&eyWty7`kM+fn!l%qziw~)^#Tev@f2v2=L7z=P4O}K#j=j1HBIMnO zs4RA(d|1H1gj{k)A)il1c&v|TfzD?LSj6qjdE6$jvApDX{2iGrH!#}^46 z;*NRTKw3ekpOmEtqXP-zoD>+Toyp4|jghO7u--c1(EZZvVyg?s;^K#}fKFkB2h&D) zopBJ_HYeYzB_(mCR8)LI=+I{^1tllNq_Zcz9$@y9lB{*($4 z8ufgrO}pL@)jOqsSy{%f{6$ zX8dP3X2NGf^gmSj@~2|I1ml}XH)bd;vf*BlT9V^Bd1d%sH#|HC(3F&w4XisvnI`8g zXLURUbezTn^jwx>Ri)^i01weE$^)dNZ!hedlX`W-sR$2KoXsCi zVI^GreqST@?!B_Ipx*4zq=mz6ZM;)g)Ab^f5=y-Or&-AW-k+5~-5_A7sCkD;oUD&$ z%%pd5qOd>=&wa7S7=TFMe9-XdVj>q?P39+|?FQSv%cr8O1n4_Vai?l}t5rjlUG&w3 zH*iRgSi`L%iBd&^aF@`3{pZQ@-t+$WOjF}I;4KItBCh2S zMsyhmu?1nHX4yhk{}nE}Up=bxVmR<*-m$yz8k%>bygLlo>JLTuG8s=t;QrjyDT_}=tZ@~@WjEZ#43 zh`6L%g}iVR__bJd*f!gTovE;TXSx*6wg<^QzXBP*#s|BcJazOwt2_hGPsz79R^87Q z=;0QnAVJAclb$ZlC3(O5!;<;BabGyJp&gdPQ+VEAXt2))5cSMl+eh7cAU!QT!Uw0? zmYl@OzN!{l*j4PMr|3Myjh`m3HMe7G&}9xSziXPklUm)92l&rBWNBNsCd$G|e}y<- zdeeMSi*kjp#tga*{dEHG+HrQuWj}2{5MIlVnP&eHA{^;3bV()w#MJv53bMfJ0 zDvKtlS?M4I98Me{1o9ATHf(c!q4?}v->9|NF@gv@;0u%)2r@8~X-cY>Jud-nuZ!PE zUxxq&ZeRD%V|MHTCD>2h0=~j&oLHgMeowZ>waJwS zNB52n99X;@yoDe_coQi}{6?=_RY$3W#f&X`4w3)?9t`(Q#}{@{i9J>v5stg|&1UX! zdj^LdnI7Z&0Hc!<9a`X+7f5>$+zP}a(4V=D8S%tq_rIAyDXvGh=3_~ygmfgOk1xCq zIm+4~4Pi%OtU>f{RaC2sSqruEki$r+9t|+?nmh^-NR3GbtUePTM_$IdrvSaYpMZb4 zv^-to)q?myK$|iDolA@HpSiTIj{gIXc1QQpd5b&w`?g|miib^^grzBukScFUwG72V zzWNYP9?rMaai!68vGm5REk*QqsWhMWmzQ)p(ZRK{ z!nq&Sp|eY2#zH;Q18yWTn$ZwSPv1ldkO6fX{U-DPB-&Z5h?|xkwp9y_iurnZ z1#`$)z94ok22Zj=gM*{Y!XAq|v$6r(V(ZcLcNq+cs(tX!=2N0+0Z()@s`%PkFhs3( zFzZ+>U62amj>I?SUP3W8J9#TG@y=-4oxRL zXhH}iwUg_*^0kS(IKfGb+jn)eY>UdNX81&H0fIC&V$Wfe#Nj+iL2S~f8cYc@dXRDZ`W<0EW#}`C|V$%8OXf>T$55X2JQzzO8 zWp7_(Mzs^R1e5;A(bPWBxIqQCOX2M(i7}ZsMbsI?;yMQXY~wUmzQi5g-%f`?(ex!e zJEQ2m-f|AH)JCe(bmB=<-X__NJWB8A>yqkkr9J%Cpp68d;0m`DJQudjBnW?*KSX2^ z$;CYwl{|_U_{6MK-yhPHB$81QNCdYYjB~q}ew_r_huVK(801=5@aQZJu+pIHJRH4O zAl|uoh*Yh0rvZ(^4iXK&KKXmNC`+rlHF6~;kw4|{8UWUK&Pi*uW&HjA(qv9Lyw9m!mOBegog`eI1B)i0_D^zDSYyqw#-8iamyr_omdcHHJE10-#2jy=^;Kxjcc* zGf1^A%)G2(bgCTFK9ykXrx}t_R+KkU`suBHAH(DbyD&dwY9NUuXFE31!#Ha9(=%SA zA5LqR3fVE-GX5bYi4_cEn6#8BVTWHyQM-oLyI+d9B{|(6HihM~)Em0CMMOxgo0nrD z{3P|-E5H*>1iv&pM8@+UoeCUB$};?-mJRHg#X`eh{A!xLKW~|ZPT4JgS||IBdOtxt zL%^*1C^`tx?9v~;3Pd<1fghTzux*JBR=cLbQq?gdj-`S(fe966Qk)B1Kn`ue`~b%- z5E$05bd!%G2Sg6(2y=9F$vYC}SoE1Brn4GMVq9Nid~ z`B}cyxSjt>uqeb06Lgsr8WDujfoU-=Ku#!?1{)19v+C< z6g{aHg#RAuh&#dOZfxU-mTOCR z=VP4+8WAOqCQ!3SGRiOFf4ozR7(P8B42t z0fYAqDm2%SzVoFU;tuuQ%VARs0 z2e*Hhp!%?otU-P~k7suyPP@r{_k-Wkoi#xV!DJD8t&?v9f;!y#R0BT@*$Iez?~ddO z6ptA(a<~~&%+yTO1{{>}h<^COTfDvbTq-Sd?L+j_GAI3?|6wlxekAQ*mX% zuNzvH_6bMuCr?Gi0G|WQRWd77$|>Oi#q5_4OAgS6(v&4etF$3W(d(+Zmv) zj90b>wK~}=F5Dy9QsGd>-qG|t1ASSbZpGnAs1r~7ysQ_*I*OjQ2s6|cc!L`6aYD|I zmP=Mf1v7U237VVAmd%p*(#OFVs20N5PNm#3dSh#IJeGy@xFChG8(_J071T6%$-477 zoeHw!w?jlGe(v*gUu~TQG%`3({2iWfl(#~zn!ML_zF&{#_UlKzVgoNmF4NVx{>5f% zW^=XixrfzR!a~0;*E4s*`c!ZvItzlY=b%C&$8pbtdRGFL)8TGsk-b<(i@ zq{l=YRf}@-@3I6C$mdZ7Hc>3^nr_Va8qrIXOd1YX$K-15^G}k|<*M~oW~cQ-z?ayj zu^vxV?qI|;_0hKE9Qc-Z?nWp(wIdjLgT`iFt8uCw{;Ibb1SiH|>vB^-Q}v^5w&K}` zXJrmDAjGnoVCj4_QRg)>s&2fAHrjB%^5rzac zl?i54reAs}loTZWz&XbFyR@)^w&;X%o88Kib?^ZXeZ+(?z$dQ}WSa=vF=~rBr?wLk zIrugHh|aFFY@@s>5d#J+Ycsgj7mI(tk9!$6%|8h|s;KV`O7Nt%CzB-~36q~b5`F*U z_nWKpQ981w`>;zJ^<;=NOh{7v1fO8rl=-D}U7t6FWY#UNqOOmt0S`TzDHB9Y*365R z%4O?V1?HpqJ>6H`Dn+wvwX)@_cHc>`O{`|)_Lf4#SljQz&Z zs#LW(1IMj0@K1^q{4oW!x;q$Gz0}`;*`wjxiNnFaZp@rM$lW6e5RQJrb#8K9YO`qV zELo-j&e$<7-=>Rf+%3S%HXfW^PJ_0`_Z4L1py%>AS@P?mZ7ojv1uxfS%mSQ00cWHu z4Xw&w09kA$6s6#1i83(jmyE z3sT259q691S0`&ue|UXlwlG8xKK)p16$jdYsnnt{N2jrI!_^4k-wL zX4%M*$5qKsAQT++?p^ngV9vK&s5zAM(do?{r@P9yfQk|@7vd=GpS-p5uYcMQ1Z6MC=p~eXmI=+x z6Yhc=q2dWR>xb{^NIzY;>Gc2tl%^H4LVLOrLk6K>UwiWX> z9(Q_qhEbf0N6n}Lu4EDLtKA6-ONSd=7MvizmW}Gb)!Yei);Xh9s*)ZZ_Z_U&=R(k1 zi6X_7zb<%~vCn)~e)Heo^#=|o85F5dtb5l_6CtSfG{E@;)7l1?e4 z%+7J*O)OTBm(_msvEl^(=-OWgLVP zgjTG5`KEwRd6zv2wHew%&Wu)7{o+&9>FT0jsZqy=X2J~>Px~>{;FF}7vY)WJ&8-8W zegT`1I0jC2T|&3FtFQ|C7KC95CN9f9=uPK?l$(;nB+6;&!H%|3yHuSvMZNYEh-==m zI({yW+ntNOT$X`prvK$|ejYe4bWOq*j)W`UybE5TnaR=HHCO`V_;B&_;l}&zod;pG zr}fJ-_4J>=xL$ATMB(Xy;NUw~MzF{hxZzojkrXLSp-*q~W%>S$NKno_O}|b zwJp8whM|%W?4Z&OtJnaC38UsNjB7$En(Z`4_MUs zB8`iNb9wAFVP*Xb7CZ*r{2PJcD80VwP}=CqefzstboxmjIDU`$+utNbgh0_4c{)*= z&+H-=4M?@gQ2W(L`dQX&dq?5G*Vqyd@f+y~mfr2BJa-}#`*ael7&7#GyyM{$&ZdZ! z*Wa5VxaF3zwEVM@X*9b`t~w{(%aIizILUKrX1*7ax5+w)7rC7mC_D>>DnJh`Ss62{ zdu$D$Gv6^;Fx(o9_lmcYvm=r6ywAgu6sbYqk`#9sp!xgYj1<>fK^i`IA^v#szqJs% zZ3*XJrGqeG@R{qPZXXan&e=QE1SSbWP3y6#?bF@f#VAaQ*-RdAJH_8tp})HycuWg! zO7O_3_>|@T5U|b~`z443Bx*4{4&< zh;Ar*`Bow+)Ug%bvO7rF*6^&Lu}M2-G7 zy$HDMC^4+uvN`>I%I@|D_~Y*k!@RbC4gF&!1Qzc-gMJK^aWLvD8}NoouPGmIQxO6~ z6k{T;0pGYXa#yRsC#LeaBJWB~-x17^QzudH&X|&YI9`{WFTtu|7mg8nifYtr5#Pq7 z+W#7=xZXVuvq#+Ilqbth=wJ1N?;1&^Yo_hbx;epnUKZfcwtYAsl@eplI6#?6L>Rh{AZ62S|UUB3B?q)Pac?MX1Kf8E5BMf)u zmvQ97gh8$)I9!@#&@Zhf_}s<+ zHtv$;)4jUCJI$sU!or!6Y!}8L?tvmgemzhi4%i>gK>MWdWw4A=*!sCJi#r%(t7D3_ zg|h7q$TJLhzaFabr@M zIAHl7^@Hsn+1{{$YwG$L<%0tWn4_nBk1W}hD5F{mP$cI=6%-$;BxN{9M(1EDB>c@7 zhFLZ}Gq3z)b!}jm?80Rd0F%}OB7qX)YKy&A+T!7W&Y`pO*G94RdZ)K+*0T$5r6s}o zCe%MwRjyP{K5#BoV@TS5-3;8?Yz2N1@b4T+%Y3LR|37)L5rzLwq5O}5uYcuhoE=Q;{s$7}HJ^>kmSp_iUz+aZ zZl1{`be^2e?xpam=k>{Wn)Z}vzLY+fF3nca4PB*DS!H)#pU*u+z-vk=o`uDY`x|lk zx+QBM!~&4^j+gu#KmNDd8D~zpSl$iZ>{9XMsg}0f#5Z-;ZYfX0SdY#;_McB$1;nWp zku<$M*0b|9bKzqRX1bX2srpqpt(8(LDVLPwo}(kw=IW_|vSo7{8~rpHW|?f$Yznu$ zRZ^Y$scNMLEYZitsc4%;&u@c*-BCz^ex5f+*N+V3EOIu(Lp#PlPvJ*>(~l>f`s%Mz zY%xg~L<;+z@;U3)OWbtRTG?wyTG{m%ta`sa{*h#yyJ%yax|CVVJ1)UQA7yYTMrr%y z1N*4;f2*V_@iZUO=a@};BH-=PK8hOR*&=%)aP9vSMl<=Gy#L;yuDW=NFzuF%*_Z$^ zjst4jW93vab86K+DXA+G2bmF8!MN|9be8f@Nr3Wc0N%ZhLk%+b9fe@6;wf}aG|4NK z3r;bMKNC?&2vV7J)>Ubi+a@1?aqGN}a7&+Gs=fu-E_C)IM&=R=d^S78~7(=vcj7Ku(7a=n?c6U`le%*#OW^iYigl*;^LE!dG+Z_TE2g&y{i^Se` z?W4K$TqlDyJ5SIfCRG>)3Mv!V6n z_45TNK|GvZJU^`CyW26k!uOx3J)*WM=Wt^GjP}k? zFjDHTPbTZdx?MI^sDk(o{upE(Ri0sX*&1ilqmvWvB;!E4K%zf83; z9wr6n*Y2n(+Azxc13NWv@g>e&2#`A6{yYcW9uX}AZK+|tONw_Db*~HR!d!lmZvXr$ zpm%D%b_oSCO5tKwB&kU<-L$1qIxo)Sdsa{K$vEDKEe1D+gt&M3wIB1<^ZWwOWoMEe zC1-@6#mF!L%}6O{E$aTlYWv<9e!|HXcB11=8w9>e1cRQKA+IJNWxW}RbEP3J;=^$^ z4Fni1Ssqcgf6+#^eXy2=7|g?IyaS(gme8ifs&qE~W8f)T@f`5W+q;Z%^iUL<`|{U` zlRlIy2_<|OJt##!V~C1_<^d=&Y-EF+4wdihE2Fa);aXyx{9p?QGf^;9Q$2A@AeN4e zpmAFR#papHcUIV+OSOD`tP$)G19`mNFvw$raYp{J>2R%Sq-^a{T>m3$k&53#iO=D&|N!_jsSQEGYDpbvo7&Il!XaX z=So*X)DSG)BW||brksXK^7IZY5Q|%JejFDjfk z6$zYt#rtXx^IWv7exV>M8|hN5h@hMZvZQfVj~9CMN&;FjUR2tp2CUZ|P>K^?|C2}9 zK$A)vy!oI!Efx!sM7#P`cn$Q6+h>y?vlG7N{I0|UX5WhBNgEh&gF;kgokKoYSSR+SCQqV3UNex=K zH%}gkN{~ceYoSDcx1G+!uMc8msOP9yIaWgaoVnpXHFwlFOL#UPX0rO>$a&&sz_oO0xlS5gXUi?^7YJ1$ zhnJ8qAT=%}SCp9qeB2CCv;x1LWx=$Rma`gWrw0VRD=E9cBdX;WXqhwY3S?U~R#44wX z`1$`$zQ;T;1uWJ>`?UrQbs~&}F6fFi5{K>}gt$)4aAY?x%DV%mS|}4z;U+0ntYj)@ zNU}X{=VWWy*VR7xb<77Gy#K=F=<981FrR$+=}`C5kwN|~hjIVO^HkWSlxR&6?<0St zi#l4i?_M^#Y0+|d?Bh7#aanYD9~=9btai+?9JZX`V4oD6GX00 z>Sdx-GOhsix!@(p0ANVKO>(E94AtmR8|lx6fG#EOBLa~wsG!N-TWx(;7&}#0ylJZX zbX3Jf-B!cBWtAp3dVF&j3vIEDjjsos5xRf*>_J)2o^DO#%I}3gzh=y!I`qjv!-V4l zx8G#CeEfuQ%mv7p<$iAg1`gvL7%#U?d)8=y>>8gx`d(B{ZgQJUTh~J1>}%=X;bDaH z4FoOAA}T16v{daB@h>v6b+P!z8zAhTCrq_glS%8R2XD)L(aCIf^P8*DRONkj!AHkG zKY@JZXBsN({&eQNy}dmV7UZ)BDxBTI+?{KUG*gUbL7im)^B(E(L7`huVLghQh=-7R604b9)mN~(dJQL2APHIz)%4H??q zGrddmk?uZ4zf8?%=%HRqBvrR!2A$%Jn89Yk7AXYFAES7Ui9%x;m9}l$m8#7ArES}`ZQHhO+eW3;b$VuQdd}%M6MZxPz`l6Di2cQ0 z?^@4d>Z-vxHw1}JMF_llVIx3bqfwO<�$Q2;Sx;j*%I;C3e^Bu6H zaz>x+E+eIV-$yI2Uo5I1xGN&k0MsO&wT_g*TGa;?Q|z0Tdhvvmz$g3=mje8gQ9l>M z`|M%Ve|!e)5i;NdB<9+@IN93NZN{QoO*wGCC{qX%V5b;tldb$RAwU`^Vj!I}W+jYq zy&%gDo}Uj01k9adJ>Vf&v0?y_xT>&c4{vRjJKgkaq;+~wBJ_NIh*n7jQx}>%ViG8J zrdCu&f5?%eQ>T={xFN8w7UkX85fNKw8ZYhnx}{7n@$e9Q{rApN%#_m};HNesnBQgS z&>TFZjRUCc39iLtP7$_&=e zja^P2Y3NkS4TLO9Aw_~)vw~(HB5WRuY&0pD{fC>Ugh*}^b+ROi>%|EX^Ta-?py+5< z55^~@j=Jm(?LwQ!UKXxOG$WTJ__mr>rIS-}Q!J@swk2l`9ry+>{>j(Qv9|*k)MJoy z+>7ivj+Wj}^6fTFdty<@jTr=2MF^0(;W0{nZQu|76tft7q%jB~JT2inbDK3=6!=|X`yalZZp|&{QDm(1rn+JKg zRIn1B;OP)C*VjB|czNbp7OUj>W!5;{U})eaW@{K~!8HU>kRd-weHZ3|2t9BRB8hHFj$ISuWDrq3g)=^%t#E{ zdoQ82XRyc%@CU$Oz_+`76&H}ryc%9-Y^KqC0SCVQJ9M8KsnldHcicBCyQhUS&i-f;rc@csU#EEEXLUEoZbT=nZ28`)B_#(~TntFBq^U;GBGBhm1d;)hT zSK2HruhEJg)Y?Yi`Bsj01*jc*qsH?U4wl>d`@6Yhb7;z(v}LpB(UF6)6HpNTJz@Wj zBxFDzNrpUde_MdvLea@Kh@u%~WYASya#|`3B z#X|7;JWsM_jU5}8#A4@hpV1v!=4D7Uxs4rcDT~L&$|;Y5R2MAIE*uKQby7D4r;Q$F z#ZL=|@P92-L9L?Acd0xovB;G6kIj4yhBmTg(_O%0~yh5;ytbJPjzrfaSPiya)N~%7}8r{5bM0`g}WASiUoKCcOg&lbv zSau8~PbxwHYZS%ZuBY%;Tj;Nf;2f9mU@?knC=KeHx*KOtYNT6wr z#H$`Z%(alJ*Z%dJwE*bV(uvpL4oEmtQkb>8mJb6kqMU8$R#Gv%i(|q`YV%O;R(HbI zk?rW=pG@b-H2L1s^@AG+f7MtbWWBrb)Lpuq5-Iv8P=U;8&C00z8z%_jo(w^;c4;Hn z1vk?G4T3vrIGdr!8F2&8e~O@l0v?u2is&|PnFQQ<@D=LJxVvwF-~*F?6=l}a1@I0? zQF@YX7T__JuD%rd-ZH5DfGT`jOp#&(0C}AxRTc|zuw*u5*Yg@Vct+|r$@%=};oQPH z@c(>8Cl+Gz*r`i8+$F-)lb?i~YC18t3a@Rk4n^bA2r+ z^Ng7AKEhY)ByM-JyJ3Zh2*Ic^NH)db5wJW?yZn`zoN?l9xV5^N-Sp~qL(e9D-gJ7X z^egVGTWZ+XlWTCRjGd}Ta1VZCF)CZ;q80t+>iX3QqNh- ziz@w1!SMDgz055nk0)uDDV1-(4bbj2G ziZo-63XT26#@CS8cKi)fy^!BOu>av8<&@eS9hYY&X1E5z{bhk*5Z{CYWpv%~s$P^1_h^h!ffuI3NrxyRFhDhN}|eX2L;zrU^kUY}o+hlD$yuN?|Gr-T zWV73)^H#}^t&4ELdgCw2thFX$?FNiu>xJp_`}Xm(3zt1U52C#0JHDxa$gpEq^nxK@ zaP@z7E7}xu@%!ZYB4ekR#;DTDAiY#4_rjf#Gd3C{1SHBz6ca*)^9H@-slH|^l$7Bs z(98ovQS5}4hz3^ECe^SGPcZ)5HD;O8CX;b!+w0pW_Ji`2rYcR@OYYc3UjGw_b1vHe zvFY6~2uv+>@&wNNZR8nu;6=KWWWI{^3d8UL{M`{iz485Z#EJ8+MK3B{rR;=6cl_5Lerb1rcKQSxQ>6m`rk(SgRz3^?AX<{kWNb zq>Me!r%tBgu>i05sNC6&yyYZ)HB`FWDw+JRO0pD^KHrm%D(Z_tE2>EOt)}-F4f)ZB zSb*OSw7z!)MANscRS?fKGFX~Xng``ptmf%O$#L!}RDahHO;peYiDv{D1(CF04|}K@ z$!Zp2DQW{)w-=d(uGpy1P$L(OohBkKEvKE$b5`ozn&RF6rE5Qhc=J8pZ~#Nj^RQ;? zv?{NN4QT`SH-A4RUOVbZPgnnvrn)Y_P#&f);dSPPFL5=o3%Yl0B7_(=%LjD&)UD3IV>6Os{^wR>u7P$KM-Hu8~~pp{&ez;!d3N@`OsBj>`z0o(pK|xL1oz>!Yy( z>--g@*KK}?1#Tjt0UNbJAQQb!q{eL^BrsIR*~mJ1FmAyPnlFUB^IYA7uviOLJ-i$j z4fXdKEq8$1FC%;2TqKBa$MNowJo9;m7mD4rmqz*)IfaB`Bt3sv%d<{jY2_D! zAKNk(;1s7C0&o;;PwivArDZ71=vP^Ho1a06sm$RB9AiyFS{Z3zS(oC%tNHCgcXOJydSwgwP6(Us@y0-X|0p~4 zn^AI~Mr-Jedb-Kk6DEKhYXJ57Wx%mS)lo$l!JTRd2OPBUK!r*;-|05C_-v{AU?HcG zHaT-E+vbHx2v&<%E~C(jGFTM#nuhxtOY#vh_pBT2jbRY22*U#z6r$i+DN7^?JHw5>EjA@oMas^@remtb{kCCGW%DwPN(t1zvFXd!BFy$@fx<3HiGCqljw58yV0jkl;M3w*!XL5LL)Tq=k|1dc{b*;j%tkVh|)4}Jcv*Bg- zh#)Nlri*^gei@xuiajR(AV@z0!TM&cCZu^{hfZ`DbVLfin_*dtm9t@xGO~)V^LwwZ zV6pto5lIl|fME{c38Apox&>sgzys_=O*&IsBhTPA|b0`Q*QLNnIzG9LwP$Kl@bsKfHKvvn& zTBprc5pIV|gb!o4*Cx}!fv!Vz&kjy3{UjPtg;}EW$WZckPMN>2iNJ1A zpiXUW+5o9b+=tkoPeO-T$r(*yy=*=%vC~HN29ODekB9^j=e4E9I8!MOQ3Ef;W)VkNDY7JYKxBtNWM!ZCzgw8

Srw zp-hlaK(_wAUzrN`M@7q~@O(Y?OHO%#;PGobhN62ArLKA_`#=ZXCW8I6ll;8a{!lDi zfygR6y2Z!t=LcK40j+T+lij2BK!eJSfsOEo7f|_4)&4-9XVwZA(c6HN(-@}O2_zq-*6$P@yswoD~}{38If=))sbr4 zIbqYC+3rbtu~7J8NN0mWBSGqR`-%PgQQ?uPE1Ov5y}HPTKzi1csMsL*T`;i6-J;$L z&U%}4q z>>OxoQWK{)Tu@JMQPyP{0qEX|ywqs#prQ@aOCwu%yRZ3+STbJ3BZ?a!&r~Q>810}5 z(W$7gW1q(ZW{7iYgEOD)&(3~Dkl7uq%fNoxkHzqj zgsz=mk}{{m_fp(ybFbS1PlDp-ZRYAhnUk{2(THhr>4rkFcmW5bV^u5=Ei;&>_~*gg zf~~*~tbtQ1=xy;_Oxh>-W6hv~ULW@_hjsb5dqhZsA3e%6&{lZ#<<{|N#sr4xAQYa3 zgUGbG2GK3Epw8aGWT>N};b^x@h9J9-s;9+*9E+LC9%NJtz9j6zAQ^!!2TivXTPxnh ztId!#@?sSFou>o$$#|_@+b;|e{iCd|v0V~8)jWIY7InX8Cnc=U&r2=y{mW73zz6rg zLDOWip;t!HDQvQrge3B~cpKm{_w&wa&KZG2+>~$A@~+-jJ#yTy?mqePuorQS*1wUb z0rmq2x=6*dR*I(4lKAdNUhqEhC4Eg%?Dugl4Tfx-8_|sPRl5<2ck^`?g*>B;IIil6 zr^igLg865tHIXMck=DGZ3(WvrE56)h&F=1ahIXA2iMCiAzE4U;YR14D@XhbuFf-dk z)25vcc$7`Z+Og9NLh_hDYjO*WF>;>fB4ATjZPV1Dx|@qlv_hbN0@$YeWs&8p8K}iT zY=TRNVLtEZYg0>`5{z()K8u{jnDmd~hp;rf zEr=Z6f~{9T2#62ywYZS!_UQr;1jGXXS10-)bHC46KJ#ofSt$LK3H6gL6P+LN@` zr~l;oB}rb#{4!}z-bncaA(Q=}49nwwTAF9iUM;94-v7gYki&7O#fWTnJn5!B1;#8= z+L!=O)Id|m-Osp>1Vm2^aimk*a09A9MhZF)Ugud)_5z6TA=g|A_|zUMn5>;;Op-kZNz(lK!@_{ z)grqNLi%7ue=@G&A*EUOHq46%T^BbHhIw6F)LQ4pgG)o*3z0wN zUiq)@F5^5a7UQ1JxDzuQd=rxbZ7Dgz4@I1;dz?d-)rzig!5j_vxhtm{T@D2)Y`6Kp zwckYbgCPlc6*jG6#wN1pU$m_@X}rRgnT=Ar_D~&jwS_L4qz>&{0m_2M8FS+7{VT1^ zPjUTq@UH|En7M0yE1l7{sfCm&!*#BBjEJtO4HY-}z`y}>YU%=qDMdZMut~8W=vM9x z+A$Rd)5KBU`zpgc?uF3s3$Ra0sOY5ZbB&;gT==0rfY8cF7K#jI2~7yi{c5x4z5Lu= zdBU)UG450XIRlfxOn^=M#nbbZ=k; zDsAbDK$i0tM%qajq^>Rf^N#Rw{q^_s%$e$I0KvuC>>fG#Dz>cH*jSM(U6ganU82rr zP&r}hfZCk%*rh2*If9H;2cz+3mNH-V9Qjd#BXkctiu@p3DgPP6iLQR#P}g-a&)p}S zQnj$~#d3tNjoQXd>=QZ@LG_@Mc@W%OI?gQXP>|PhpH4U=`Azr0QvSf79&Tf=7~uVx zedRnws*&=tx)^`fVNu|nYVz(6a^7DFD!pZbZmIc&G!05nZ=S!xjwm6a!LMv?eFZV{ zGB$!M?EdM+ax%(R@IpyhdVfAkiD`TA80pH8n>%9Wa~Rnw(DljG^ZLC#`FOjV*8u3v9*rx&f$W>15QcU^DKm+ zmEP|D+d_O*gfxc|Sk+&5dU&tU?}QK7zcB|Jl3r@2a;{QH)f8+91AwCa4k2yx)*}Wb zX3z6S0@1y^G16pT#>t=NrwBFB?I>f?}h(({Z|Nc2 zh90i)*9i&WV-d!o*ZQO2JPXCas%R_vLcrX%q8_lznKcoeI^3Ivje68le#t9ci)`O? zFr~DKA7u>&o`T6lJ818WOC!~1j@0`og?qC>CTAIvDVJ>vmaPT)u})G4jLI^a>(_6h zZ}4Ss-Ne}iv%izR)~AoZkTC?tn2BR12E{Bz4G=K-wA@9!g-t%Q&r(?jsp@Id>}71# zyf>)|1oQkIh(%9k4aA}^CdQUOVez}3@NMen8k#M>B-VS{wh7h9_c57(K>GbFYkzq; zBPza3r4e>;S;8Zmj&BoXi*}PvV;@`1y6wHB0DWXaX>)pxSE7`--N-=EwG6mpSVpL$ zNfuS$#};B^%5TZ6x)j}Y)iPIsy%kARAAt{Pn&&z|(S7Pap6EjHricYzf55IxWXodp zlhx|S-`27A9~u4k=?jFGf9Id9unF9mT18VHX@m{?kT)c$7M|IlMzwnuPe3EUsTUEC z+{MOM>T~R-ygs(j6)l}Mxx|Rm3CG}NCl@Wgqoh&WkzhU!>3CNMDvcZCcUx8I0jY&i zGgSMbu}zm&Sw`2*jxIT_zgE$jn2*ejS02(rwh1HM~IAuWi zI|3)dbX4QPAZ7xx$WL}+6bBOx;0@aZq8b9W5<9Ur1tk6TGLfv(knO}_lMikI|AIjX6akn|K$2q(L0 z?H|_XXWzNe71Rz#kyRO(j2&Zl*X;Ofgu3!%x(Z8+pNT6XYp{5*T5DtR_IL6z25NIR z;sPGTCudJNkA1rDe%a_ERL2fM*dVU&hrRsx2X6JAcDd|<)!0}Gv{g=~Y`?8jTcJ_) zyGHf8i*7t2eD5{(zOJWgkcP|DPFQBgna=<)QL7{^M7pnqENcs&%h0>*CpxMZ3nk!U z2OqqH(gy6{V@+FCrtk(?2ZxL|k1<^!-QB~kjP@BfujJp6tS*7gym)&p^Rfc_c?e=# zGFo`2t+K5b^sVyXC339yX1!@Q`FbxeP;oc-NV>kQ>U_(lExwl;9cKJ((3(!MZQ=@4 z{F*>3!ULmRvc2|68XO4TBjKC&hm><)XRCyBS@xAbjdSpU5P<@L3C86J+D?g&M7PS4 zt?fhd;t{t95`o5U2F@$~BHK^8qL$`6$A#l8ISb-0<4mnz>N4nmkf5tYz9`S_54lLe zhpzH^F2AqP!38;4!c#du*QU3{)NVT^Y3fD zi)ut~#-z9h3I_Y_NhD1tg9zdIYv#KbwkZ-By>(1N(+UO>jM#JKNX1JO(fD*Y_;4UK z^>Z6*e}rP-`o^UK?ueu_x5lgvBs>N~pG6(&KpC#wAiu6KMGtt1rJYg-1JAhFAhH$G z6?2{NXjd6qde<3#JSpp-9Ma6yswuYqpb-;3z+!@$w<}73-QGP=%(+Jf+8DZ?^PL59 zPhT6M&zG*-j&aM7tM)~DczHYC&DoBUbPjmF!Jj!4}&le?9CHZ-QuIJ zQhOZLz`~ZMv)7m3o&RkDbx0ehK-ie`|D5FHJ<+S@CcSrZ(m9h**(}?rm z&qlZR|5Y|K>n0n`GRDO?hN**QtN~{>O&-C0Gp~3%{0yqn+GDh89*kZQJ)Vr2iJ~^% z-x{fANJ8{5#x%wA;oPu z%hM=diQPKIB2&IKxZ@8{SE_i5N$b|+u)f9C&cEwLsq-kbYu{G3wVNX$-510`Ut;gM z+{++xm*@h;`E%rUV-scYo zIl4E(o(YSRnFQSMXu1Vz%ST}>PZ=1STrthk8`{2T`z}w=v}H3b>&@_H;AD8TlG%}Z z02m>wWA**8i}-D9y;JXOkvD{;;SConTpm4b+CqbDI;{nWzABl+k)Hqjyi&TzO<|WW z3jXtWtH@}hk2nQ*&TiGv`U(scX*Z?4NOEeSuFx%X}juo&=1Pi51r$%t};pt45ALUGO(rIH(y+$Tm>Yg6< zE*c7bPxBoCy$OrG?s@*f=a6Hm5o$Po8w8jQ3s8AkS6^wN< zZ3jB~Bh@`2P-IFFr2q5T9=Mm!sik#p3mfY!1O&EH{ZW|aRA{k}sCJ$D>{)*eF7nsH zI<(4_Q=77o#r{2dioQ(9hl3}o&y@~spX1XaIJ~jH8tNh(A>GMV=7*SD(FrE;c|!aK zE7+6rvE)s{A+VDZW;>h0p+VLZ)+y4l(ixHoyeuZy*k;(;FI4?@9bCh-w{!Bx9Jl4E z+Wi$4L^|B6ja4D&dm2AZ^9zKp`MJK!*gcS@x@{E0X=@qd7c)5bRqQzOXJt*ya0q3n z3DXRXPybPT_2c`eb}r6jU8LgvH;+CNOUo zkw)y^9oU?(o)LnuD1Eq2z-8I!jBB$^jvy~djCTgdysblVPC8Tj3w$cU$|Vci6v7fy&{|H^YCG^^S+FY8;%JxW%S-gS)$Px&3&vp&3ec z0#zo1RMk3f2Dm~&o*=U~Z~xO}BxB?LI;NLZLf{a?(ErJXtIe=|fR52{XMo{ZsA#A)Qox56 z-QX@jyZxjA&zYLG;C z+~Mma)|YksSSJ(AUqlM{yqR5DJCdW%t0ReeB|34eb|E!?}^?za97};7|+uHmel(Z21=Li1(zo+Fto6i55 zulc{HV{7O5e=s4E|9nFKap?btDFr9=%MS=5h1_~a!0WIA8%9Bj90)0)7>1L%5#g|C zZTQzT{i8=U6)4;aYZV0>+Frj*UtN7#H=ZBauRMwQmG|FJ4J8c;=iM}V4hKSYu6EPZ z>S~K19|~63(xo7GA7s2+{Y-$QCty5g6|qN=6+&t4e;AcCeGd8Y+U5fWJo-y$j>rrL zMGJ~3T=tnhz?jn`(fo_F{of9V7Up6A+%XkdNGbJMk8MdVAa|@mYcu*8#3o#=t*`xb znhreSA?Q|!^vMxKSHo;y(5an5nC<=W@b=c1EF4#E>{(lVfj`V{A?Ih;<%1R^X5~mN z5=-3ZSTYoYYC=NCjpxn;{i!EEYx9J1blxP^c z1Qm4>>pnP!pkIJ<1alz!4zc{U>tC)Hem_k&PcUlB^%LIt7P{W7-(@f1X>ONMDj!Sm z)E*4tb1%e&51RVOmS(LrixKWH{KvkZEDKuE5*P?*;$JP-f8Fat{$GDuqd!Jw0DWgC zb1TRH&hnpk-Chl1pL_Id~c6fL{2G{aiCpj#jT&j+#%oiCrcc_-U}Wxl2;Lcp&$hCMq1PqybpK8;CBvU zl`&#d=$;<>dbMc91kSWQFsSIYyV*{XNkWw50A`U+*?39uO4C%U!LUqX8o5wYfI)2i zrV^TW_**mLJ31t`7>w7R7{kv}11vF#9fInRKK{iz*i*p&TBUEr-~vD4+%@1+?9g}v za@|E}Q0y1sNf%X}C}ZN=4N6!lgDW)1n;u9(8f_op*Jtp6$mDK5t}T_cK1LH2Yd zsL)|OUTbt->AeVm`I64zNQD(z39MF({!h5vJZc!aZ0)wii#fLq>*tHY1|j}0%gPSk zYrp2sLM@4dYSKDZpQ%1=P?}3%$&_emx>e-d7Lm9Y2&k%U%v5p3{5U@;DFhcvZKQ^1 z(c7zyFHk7AmJtx3I{;>!ZXP*0n`oy6Zd^V(KMEVyUsbTO87=9lqjflZ6_@N$mCEfJ z%5E;V3E&95ez9@RHJiu3#=^0_c$~kZg8g%KQX;)plsz9PwGu*+CPXw-|R+ zg_a|M2|3FyA&6)9wNsmJ$g{J6uWMJ^Xvmugk3qo48cX7_Y z18?x$7AyApd6yh2j=j>?@8Y*`9hTA==n{s!;rb*EVV=I}WE(ATXPm+2WNLDmqy?Hr zhP{T_&k*QDZJK+J!9^zO068G?ppq<3xD;!lT96>?qB2^`Uk|FMwt_T){mf^CG6c&+ zYxuj|uHlrW{~OpX z6W(d~n%lV(K?Bbrt-8-DhC5JgvxY@We>8|?VADpC!E>CSbAbP~jmF6#{}oLCL$E#Z z*PZg&Te`A7DFW2+8U~omi zP!PO0>BMd#xDojAy{AGbNush+cNatt2$5~*=nLlvkbdgSejMrz|Crvj?3k7~m;RYa zW3=P@D{Tm$Im_ROFJiZJTWFw|lOdTBO;ccZDLYw&vdTt+cY_$}nn`fo;Yq`?s6%X~ z!7Uu#IGAS~Y$=3Z-Dk6o_HZ-VvxG-L-K8}@FWILh;&=Y*EKOBF3GK4 zE=c!venkOt7-;AU<894&umUxU_YSeRseTd4#*TZ z^cKpx!wKCt2ba~FK9(NfS=BjKj0-i$)HmUkvo0cA%UM%C>F-~*PpL$$kahlquvuR< zVRr03k12`~DUC`Y;nw=@28BF@-mK`aOD@y?1jP&rK@7iJD$we(HbHXscEzeS`PoJEO@OOd-GKLq!NHUT_pH`uc% z<&<Ze8f z7!8z`E}({BsSvZR?kR}cA#kjkH>jW-ZZL7Sd?k>{qC;}h7~vbdg~;w$ik)6Eq*Zr~ z1^C9LUZkP>8iV6r$LB(KA1e-vdo*LbkyH>S0B543IQW2@W>_CRU^k~+B(fLF#!_?& zqKL!w?2mXjwlV7=N34l8cfOkodtjhteVPh>3Ird+ig$_k;`#cqkiORtQBR)Z$hZ61 zdXu5&M}+*3pjbprRhK_u>+`ys|9Qkf`Oyx#=ehSM6mnafjN%A*w%}Wu6#BhQv3vzA z@p5@{Ms78nlRME`$x)^9gHaE|Yz%@tDQAS(M<^wuiiWZ~o=AX4UJFgEe%Ce+zak>m zw*uj_JVPPJsHtsiLF8aB-Wqn;hw<%f)bC-vk%nH^_ELXuZuZJcTR{XlmRiROZ}1rN z^01bQh;QgpT`ih!_-PG|5*JnIIRw?q<)b{OWi!9bJ{#Z|chFe!56#eg|JeeOc>YC1 z>VX!bnQ>divPA47;c01qKERh>lMQYcgQ~W&0&U|`}B@g2H zhMJU31SMokB~!r@@SDz>(-iM; z;%iRcbL-oTXe&xC88xP5^+a=G>y9s-&m!$b^NNHjlg|in6C_xJd0v_;PBLbUuH)11 ztY+N(5Cr41h^Fg=Cw6&>$fS-5gVFmkouP92&xLiFlCkR}%=a>BrW{}UO=$~!Tp>10k(f~!%=piEw}>CoCfrpH6~qzo(& zr8Yg+9K-D`gg(A?iFFaRIjd8pw;=_33qu2PYi|noOsDl=!f#%lRPGw}6oQAPywT}K zn6_#1w8>v%rD-a9VyYi4B-)yAe&6;qWtO=YN-spYVijtXCO8jnbWq6EHAnoN5N2U0 zi985{(&%qza2ex&MZhl3l%pXvF7Bu{;h+~~KAE)Ely<(HfuE~TdPHY4{$x!G!H;HL zb-}G$gOE{K*I*eJRgGAD)!I-pM~U3-@op7iw$rA>l_~n0H>+0#z7QH>tvTL=&YM~; zP*S9IudT--R9y=8=b~-;8hln}Uss{mfKT3mi+XXC;7XeMn9*@5x0lB`39y|26nTO zE?NNdb`P~80{8AXbFh3obWB^RiMGGjNEuO zE?8fj@5>Iwo7dCV>FyiEz=KD22bHZ{!9TexFz9S}<-?cZ`Q_fwF2?TW z6jD!$D%w|$sC3WrlX|*Gp)N8}t{A9Zh$wqVs$zrRE7LrSG7fZ5TS^KMv%N_T_@ko_ zO}W46EM{4Ta(;XA=87Jh(w-ypkP$bg8M;KFS9veDMngi3ZCSfG(W;7VQ6H{5{a$8; z(4Bt7Q95V5o*Ga2FUL#A-Mi9vZxT^A!>arCzmhjXHZ4-mJZ6zhJ@6Mg@oY5YZZ_0% z{yKr^sxF-=mACQR;m+O~S(dM(b63|DPsh;1=Yqa*j`k#yz-^h*o`8@0`aW$CJabMj zX7KJBw?JMONiabvm(H|Ezh4IxN4}q%y_VW%3mrZFJ&Wi9q|o?e`|;rehGK#pLTc+h zAhT}%snwCc!Ibf})k4Kx{tP&~leRP)$etK%)YZ&G1=o_BdgEgY#s+&s3h z6{;3;?Ac_h=pW^msy+n~yF+q95D4*(m=!w1t;-lWDJ)Cr;hcn260i@;4eA5RLTu?o zd#<`ai_ClZ;(xhgD|io=bPepfyL>);?H{j~C)UDwZ1ONTPqoU|D3z#7%6ZS%&KcKD zVOXCwoab*&vlXrVnhsMp)Vu-7J?rWSq+h6@_u^Z7W8+ih@t_`iY(l^0@&$fr(Xa;a zN?1B+uj*6tMFR@%P8W=ys{_eh-^_kAZaULZdBc1Dh5zfB#NvuP-fL$a5MbrORX+tM z*LVYXor86I5ftkQQ%$K~R^i6cB83s-={(=Hppyl8_&6rZZ#43z7v2tmvb+G9ze}z6 zTUzisEBgW`@iJztvQIrsWjBY!1EYDat0Q(F)-ZmIU(FAamDZ?Z0Y^Nkh@K>et(7o~ zV~f1joO>ax!L8=(l6wTIMAJJ^(~F8u9`LL9eBUXrl@aN*t~xSR!*O>>@LKO~m)0l} z!-f{G_19c$p3h>S&JJ&|SISeO&LI9D3$ygENP9Zl&R$iMCYqSVFs!oLR4w&7YaaAF z#rF1&4qY1V`C7VxG()Cdgx=atzluNPyTa*bXJYRhwx3Or`|?2*Ws!ziqg%zr_&2Ca zy05S4f=y()Vfot$ZTE%lYispFy0f=;TTbls61s29>=$n=ilpWXXo$DQ5)LTd15&BD zNR(Q_IK=LL@%!KUik=$`1Rvfh0CO2Kxuc1!jp_C|LrWfsq>obX(@wl1PflkO7{8|N z@BjV?dujAg$uinA9aBE_N{=CfOdD$t?{TCb+%I)W6DJ1r4l$rX+-5zbJDd0|a zWxj4HhTFz+QImTq6*7EZP}plfQo%wBI?>03X zX@CL;H-@Ms>8aWj*}I8i(Y&UPev{)eRkP9%A(h`;`hgGlbB9G5fLlww&8`l});=qD zY~kdIje%gV9+};2M6K(Mj6D0=Io@CI+nvO?I@~}1UYFpHv>vgx@*X^~x!ya$W{3jC2-A zmrLi9otIdK^oVB?;Q`qBT-4f+mL6X^l3u&0Sz{V1LI!OuKarzMpw_`mfR^4b?Z!u^ zS?(`7s4zsvz{OS9`}gGge5qbn@DUPzUr|HAw(9QYb-1XNR`z$(L?f}-0Ye66MX%<@ zG8nV&?a{?-4Y5dIqo^oMk4Qu##1N!Dyfk~N>nSAv)TLo4D|&eFW? zj8UqD%U1t}>Gnc_ZyqaXnU5(MAj`265=n-Cc-gU7^sO!CkFFm>v`HpZ7b_oPyIXW$cqA|soD5eyE@objrFEFhpEA&=B)(j{-MqcZA9Ngw0t zcBf$w{2Ii*z;q(A0wSV8{4SXm6=^=!sf!Wt)rML;ttMMN#CUXewJ#nHF~|``+qGb2A}G; z{^6~L6#sW}D(inOrG5_B07plCd*?qkPWt~-WBu=_t2&o1nLzF*rx6&BvnAMo9LC=Z`_)ZFqCDbF zuC$bP)Xs+ijw3dix&^%nQt=#W+y@V%wRW2KEi@X%r~|` z%|1IZ^&<=YhLe}`+QbV6i#mrviz_{}hrlU)OKM}eg2%QqI-jaAl+WlyIWK)|@H^~H z-F4=EVA0$-QjY1EX?O#=s{Qp~3@p5AA)F{C+B6vtYG1L`?_LT)(fF`9e@$f$^V;Oi zbK%6T!cLk~r~zt}0#yJkVC^h|n_|bajcy1(1Edf^eX#vzm@_BUn72-E(`XrTRgUDwa)x;l^`WWtPB2rzF9nDvU7c7E;{SFb zR*+tEg{nKA#pWc9+Y$LemC7-lP>HLp+xIRAnSh&N!z_220;Kblq+qP|U$F?WupUas!-#2&9<-6Y1tGZWJchzqKGL}s0 zU!#!|wk=}``adJPIza!X`-1_+K@I>XADoMiADn^QCM>gaxjVS|@e%f^4SAI$luFS# zZP5bUDQJMW%{dgP5B(R@cXW_OK|Z4aZ3rQYNr0Y#!870EOK*{joYt z#Sf3Ul5{R|Cr|L;aD^Xb4eCgtcEP+xLK%z^xlh%gONsOk1+rERsJmLdEnMvpOqV*w z!)(j~N0_5RLX1k52k7hzLO2zZcw-*|oE!u%?wg>^hN_x2@Fy1nK8*EQC zX0oC$*pk_l;66JPk->S;JeYm33BMBb0w6EK6<>W+z#w09mOz96=W(=*3|$uLHXABj z4^uz9XA>8Ldap+(G+=KU6N4Jn#RsgT7=bFmI7NmeHMjcCXwAkc6qIRkvvRJougG`t zc=$hqqwmx^ZkRPm?J*sZq7&SMF+H2?;@Lw|!m<}fpWjMV7aCw7+zwt@cy;Dv(VyGTR$ zJEnh=%x=6`4h}ZPJ)!^+u#d(4ymNmxGA4kmR7Z~Q$M$H-n+$U?{k>r+phCL~j{d005l#kWPZGVqE`=?;kO5ON8!U9_7B`FxG@&me})MQj=ZxZyR3GhAyToF zx;#5SBhvQ2(g zJ<{yMNy)Y9-PTptJ_a-xs= zDoAT_3TIYka9nW62VGj5d`I^QV)X@dbBb9R?rjX7R$+Cuf*{NHpOBbEoZsn|z4^u7 zjX0;mQ`&EKb7wC#{IbVZ^V5>inqhi_n1~N3c?lX-fcKU^-DP0=HZ6PuzOf+(x0oX2 zJftOhKv8A$SHh<8SdxEOFPv*ZlFrbwpxTo~{4=Rk%og3ix%4R7u1ubA^ET`Z<%-+FP0-Y>Nw)Y7L?ep4n*l-~cJByF-=d)v}M^YecDr zJVRY82If^%%-jgO)-rj)XCOL9=gC#q6>HhH1)vNF5ae%#pE~OwvMuLp3rthCyfdyHK`g8~r&FPKn+T%{Aq|QZG7TD>MMS z*8!0?Oz#8o9nXQ_Dm4`D7JIWg2yL9v`o5iJAul=j15L?N+4fcvlhyWUiK_gCk4?N% z(TRJA9M`w8^M2azcP-<>iTU$NOo?J!KqJNzB=TCbgX{K}d(psd>BVizxxxi$m9gPop*hoO z?D5HHnjSE;W{dVmOJ8VERIVCU*jF5i3b!qas74TADgMSKbIH`$`?9`G)&H|>L zVi6=_DG&*QQL32+U-6xr>LLSI}&+0bu7>vTDs9d!s2vn)3vaM7@&#bg8GRT z@KR)=|EyLYJ(IYTbnJ!^R4R4~?V_r1IL1%AOq9MprSr_hi|AJ|%O48H4pq6>8m|WH z2U{MjS><#Jj2VSv4{RMY<*}PuU#qkj%srSUEe`3lp8FzGWhkJIQw&3iEJYt;o5MS? z+Kgp1L7;e=*<0JCcWWdoI=RJVuHvS)2e!fMu>?YlCQz;6Ak6S{CV~^rTH#3iWSqYh zpT$icUm1r`5$ESa^n%*SEGTX;31xqMj@IAdwd?aVH_YjP*5}L3Y#yA&57Y29RF?z+ zh~zUn6NWQNdkC}J;u3rueO7*=Bxs5`TeTrecX7x@y}`-x!;6Ryl|CC&`jk5Y6H1nm zq9HLvh?SN4)FeIz%y|h;GbboN%TooE_6>s52?a=47vcQg7 zfV_-$wF>*9xaJ294ZNvl@f>#but&~^n!yGcwc%bxi3Uib1e~j(-<13Y;XHDges#<+_>%*UqB8%ELtP}{}6fj|xbS1?`5%7f)HuD2;+!FT`1+Yi?bt!Ff z@ndTC%EcoOiM)o|ax%=o0@={gMfe|%l=C=5BNxE88{|BlTBvRd_>V^e_YQg5LYmO- zU~duP4jp|z1zQ#(6%7MlRfogDm~e%a^8a)Rl2|;*of_a;xTb&#yzm0hvc`T6<;`ll zVNbvWMM-GQ0I$)GFqk zi>15Drzof!o%DRwsq*01pgVQ!8V=lH%Oy`PjPl&P=Bd9TDE1-)yUW;OaEge-Yt){| zl^r1mqa6dopm&h*O-}iwG?@9=DUDrbo_H5Dt^(X@2AYlje@xP#0iy0M_x+& zDBJ6pYolYSfsXlE4H@}43*7IZX{DQNM||v{dBOiV*+N+{ZFkkIQB6Tc?~N-PZFU?@ zMGK7@gsm=-^Ljne8VoW^4f*={{A?&`oMEVCo`O#@d^ zfojQl$*vU57IxbsRgY?U?lM?j2THzd;`wBqp=rgW<@9<6U?8e1nvKo@k#Rn_!ti== z(P>h&qp>?Af^{n1RU8Z1Cok@sQvYUOUt_4Jq^r)pvK-^;I!H@N^O$KNj)NUg^9 z+36Y59eiQuOm1-TO$1?(M@H)6><>DmL;HH@#n_`k>W@1pwE_*gqm1D=#^JZ~<%k@$ zMXyJTc{LaE7&i&Vj3_Q;NRf9nxu6x<}4BOnnAhtBRuwNfZN9l*>s>lrcH zMPn8wP+;VtmxtbuMidssV~eJo1y;1$ob76BuDhwZ>bCNqXIPar_m6#ZcAgnJYG+?5 zi@U0g@hIIrEKoyHrz`D0-ewGie;sHK`U+9&1iJ@uwT#~gsFpu6Ch_QYEg5r|J|c>J zR3kP*cW;zWugjrb{4`%phh)7B2!0X8n0!;9!W27V1PahD)1*hP@$ZKEf*R-%E20Iv z5`QJI5as99(|iF)lX&n6@A;0%Y^&Ju`>T4cS26Y`ckhE3l2p$`u;NLx-d_n_5s1q_ zAA2O}yfnvCb?=nxY-RT8`9Q?qF=cDK@3e><11#}vCFYPVl%}$h5@ZPX=2=)gHsWQ` zWWA^sG%!Jr~G4hK-TDClAMFm3+qgBRMh^5gI_+MkY-luJfk_3%@R z^(zk|CY;gd#3G-dl!oh5KBAzMQwV|;gilb0?r5G5OIX92Nr=3EE=$FWeWK!xt7ow! z4iVHZrD-fJMP=kaOcdSF9>z|d~S0)npSF!dw-hX3%DE}w+$I{Nk)Wg(C-|>Gz`7CRDJ8!Zf{pjZcB{n&m5GkfyO7c#WIk#_E zI?Eos&W~n|ni3*5tSOMlDl}~CnP1m_0I@cid@T92F9|rrT9KVSA@Y2(->2eY+ zt20e6G=!5+GEs9-GAqqsYIgg5g&v7M181gW%XN)476i8#aoxC>iM9|_l4yd=RZlHA zXpTUVXHc9Wz3mX7{H|Bljh`fiEt-xK<+uoJuD`(W@5qpIAd;8w6r|!R1`r{-r-d?$I`6nNM z@(-Dk8!;p4H8N$H@fsIF0f+vbdiV@e6B6>y!C6LSI=P5=TIR*X% zF+BwA!*Ha02(Lrv`adz;w$k>R!2EQI(Q35QD4A@CBCQro{Yp053^dFhTy5B1?uNhL zugejajEA4tN{&8Cd_NyBlr`0#l&tA8E-vhoY|?rj!$}s7R4SR(5PvUPR&~ktPlk@y z50BSNh4)8=cb}YMjTMqCtVri*wmRjk0w$7WuPJNTQl}vR7CorbkO5d3%qZS6z^R#= z_9d+UKHdjJ%z*K09i>#;7x7tT4ui+|(?Z9^d^>$!D$1I#M}#!LBH^;aE^^`{E5`6y z`m~e33fl3W7J@d(Ci%rjkG|zGhon&@z-W7Wx%vkY^SZ zmr(Mt3G;%JvW{m%!8j7mCzoD*l0skcrVIlhqQ+l_4j`PgMH5gKBa%AS<&czML-M-K zwrD}!IdOvw5X64a;s!4hMHTm3U?;NZkT@#eN0!&lZ4lo_?y$>XoaWJ8|6lfa~XMsDC@S$NI^*SE+h(B|X zsR;oT)2$RywK!*(9Iv;XI|i>XUt)dFT{O;o=PxId7Z)7eIrEgjG#BF0$656+yjU-B zdv7?oy?m$L)D3(Pyl3=%EFX?s#-P%rRX;hXl+uBp-**@NjbE%lhnaMqXqSI5c5cU! zQHm0(K`)6BTqh1dmi%~8d_U%vUx4QU^+Q2JT$hZK9t7eq=vCQqYypCvBTjERKZ4j- zrBLHoB~t^+as^K4{`!;NmzypE1HHq?$R$X~S<6I&1t^Xe@7pb0s^|ZsHbvXz+EDFh z7`2-U)>O+b=FYIE;Yf*Ap>Z|3PiUlwUJT!^C+dQXqA_215pJK$i4Ap~^O4ab(i+F1|_9W}24n7?kg|h$7I-R=R01Zbqt6!vp}{pEYmcYu!S#8WSPwwM9kN zTk>+76RU^LzJe)}U2|eu6HNW{pA*@p*{hg+am$a`K}C z0j1;M-~?I!dNtl@{CA3g=wH$x2uh_BT*5CAyLg7P*h*2CvTCUhaj@q`C)~t!p{U)R zIRc!Y$C$SV-0cCW!33*#Q1F-$Jzi?uP|5)Mh&A$ez-a8L7P=z*M6kjJ5y3NTtHm_HvIyK@sJ>69Wh?rbBStCQ{GzWF&S!-qGB}={z*?=)} z%ef{LL-QQ;2GuevHhA{Lv~}gL+WWEK)u^I%Qmn_ghF!*5>k~Z&0)8JE2?d7B7g5_G zV#0KV7EP;;bkWE(Bnha<+FD8rtfHBKC0^^e{8C;lBZ23mp5l)#W!T~_ewRf9MH`qP zY9rG3N_s(V^gwWSr0~wG#9PfB!5gZwgDpY^9JO#9h9@y+knWUWl=k|(WX_I2?wj9nd1ea!{Iqe`}ZbCPgwH%@&d#4U) zlWJ)vLMNu1@mhKRTsg$cl1bqDRKQ<#R^+P^*K4%l$~uW|v`oMY+yE^&I5^jVlT#gX zzYAxDN^GbC+6ocxMePcS8zq7q%=VDHVVpsLGPbgcq+v*Pv401lG{ixS3N9BJQ5Yc1 zWyO{eR@KI<1(m4u=w;T{t6Go3Zt*pr9liRL!vO6geR6r$Q8tX0{i>~tXzE;3H3#qb zEOY}!-Zl;s+hJ;M&-dH;QY_w$G!Z7a@p6+~Q%(2XR8SS}1W9b(0GiB;!8fi(N6Vaoj(xX;&56qr7@T4K(^apGb>Q^rb(EkIsLq_jDzc}{ zM7v!buLnn!&_bXPizqEmR{ClUD1dLN{$HyS@P(IXlT^`kLr?~=#*V^^!H7n@(d1b#1CURX?-X%(;GG-OzJzVWh>^p2+dq*X?YMVZFtli(MQWMd{`*Ofk!Ywx~|~Ni4R# zruC-dA*^sd^q;sV#_oFH)m(dGI^6>mPrg=;N-Sa+igRHN(lewRAf#h1b+Ru$%=cKk zne=qom^Vn{I`BlxJZA|-6c-XuYZNOM-aH8d6q)Ub53%&tm**h?eAq^y3q*a8+>^io zidw?p1@GZm;q1*o%b z=lc03%LiDKK%zklj18a(Vmhf5Eo!k)0_+op`A24)G)pYZTG=$*vudbJ)H)aaVO_&@ z6!GwEKZ-hamXZ#jcUhct4)iyMQoG2&?)7y<_`WeAW1@DP0J7vjSk#q?h3r0?weFG~ zAi&7hUPOLNe^Xk@vAYQ^E}>Ua^hn?x$C{zznxD`OcmCA_tlYI^JvX}tCg)f5BZ|O4 zKKAz%mKy5JsgE1bAEviQGac(GW2$vR3Vxu=fz1vHIpebB(^j`^19FZW)WJkA+6+V4 zVIbX�wiHm&Y~B_LG%e0lGURcdC22+wrnz%%m^6IIx(WSn!3aj1&y`^l2QH!>p&y zP=mepv)CR;eaPaoC3Kw&1#r&}nvv>t)9RaRtX&ZrYhPs^TO_cVCe*HuJjWPuPZlgw zsaH0xvu-9Ul842_p3%exidElRa4?=~iz?-{OcsH?8EZn&Hke}lQF1n{pgI0g(=s_Ny=9J*f^9O@(3~ zh9)|tW5sUSrU7%`Im4P0=Y!r_(ITOEwE!UwsZ7~P7kfXQ<4d^;%bOr8iBBwRE9zP- zL?Y5AIbe)-2u@epXkw#I%eewtHpC#qY$Q}>`#qqHL`_oH;OzLPyC$$X z>B$|V0xPl$lqKN>HwfC$XJy7$SlL21$r;g(_fUU!yZi?cQM%Q$)&}tpSX~Y$Q6i0W zUxnU6@L@^XBOmE8ll)ro2HLVxynh9nG)t7oHfy74Z$*t?H;Oxe2;OtkL_;6HN^zpY zt@lj0R+>P;20K`WV@sc$vc%-^w&H6^H*{^-KH*lDFcE(c0mxZ`K|ov3O++afZ%QI4 zefGSm5|(8$$8CIEAz90-l5ZFxx5nixfko5gnwpiOZ&F+(5IayYl|td3S-N&-EXB{t zd|&7}d|Aig35-2@@cNg=ebaQfPRQ1+rb|+YBSia$&@-4Ody<{d=lt6%Vp#H3;gg_cw-dloI|Dx(17 zKc?LZ%P(DooOnH$O?H?cVASP}u2mWCSiX%;6drfg(w}u9$q%`MB~ROTn)0?v&tbf- z{mZMv-A;0o?|peQTvJ%NP8Hrv#g-Y#28k)r__QAA)Zni*pn z6=q~BYFPl1>q}A?+#S3v2sTT5=&`37!15rxApDc*vJGMG2ELfPA+h-uY<__dFYeba_+(#^pQeE~HU*4+^gG*J_*OzK<}Bp)3= zdKlF1JBl2)J|$i-ewL2znDbE3)nr-q>d%Q9HtzVX91bz6{H)-IWOlp=ttw)sqQI;?aV4YCPp@(j#mFgbW*yy zf*%Nn*j3!TB)FS>KMpYufcW7h27hB1V<4xCt%f`BXd5^j1k5vIk&|dK$?WrgoS%Y5 zKmn`-{-!>N95xvIv6Wb+H$u}KLc%-!SlT})kWhC-StOYBoRd*XPJ0(;ZV>}Ld}4P# zPEmKPYcBr0$(Pv}RPS4~Wi2n(*fbd2Bi0Je1Y4@xEwLEM(LZtB4EfI9a<2C0lw^yDkOQ&QnHdFkobh*b za$zDupG2`1XSC<2RV4h9Ku4h<8p8Jao+hej<*_NL)jKu##_cQ!g6fqfFx% z9fJm2iwaZHjO;Y|giPhB-(!~^@iY0`FCC28OpE!7ku!vTN1~(~ha~2wX~7Z}t(bcasf~xjp$f zpxN>&hB$xu_|JL4Q3lN6%NU(s>XebNP=Nb^7H7ylR;O!om^N~87NS+yuN*q)7FxVb zFkY7?yG@&wAZ9fBbH!ihWpMce5K|L0PFwGH^0dY{!q3xVuL&|=XUh^~vJJAW=&G*0 zNz-R<6K{QZI8*Pq+Gg<;J*NW!KJpF!A8G{keF(7#0RW&~;J;BLO#d@A;->Fp?_jEL z=D_iPFUR~ZIFc(YYvseSyH5Y%q*l7W0%$Gy5aM#$jXAcfRdUhVnrNcja&LP z;}++!;Pxadnkb3xn%^oSm8_YG;%?FBxfbcgRpDr)lTeWwp=%j=VNp@iN$Bh54JJYR z4jdfRl#mKZeDQ%uE{VFizXw6#jmKspseUl#=BwhT(q};^yFO^9l)t$Gq>NCFK-mZh z%0g^E90GjCItw(e;6ZVulhUCez_F>XFrp~eq{M2kl*a^nHoAH-NI%TC>H0HbemK=i zNS~;5Te(DSC(rkz^~@Z-88`Ccd2e78eA#%qvB&EsJS~XR3TgO;NX1Pu2-_l@7Fyc@ zmf?j-rfmfJ6J-Q|v0tTCCN&B^2&A5B3=A^$I`CkIe@RaXW0Be8d`HG9l`tk7C)VG$0-~{SOff3aaO^f;_6B=mxHTf`LzRA|xa)0e& zG3A&;E#-c5s?V|Su0t);4N`0g(*Nih`nYK;8jG!x+0@3n;B^qNZN`0`9yhgvvoZpyc-DdBad8 z3{9jh*y0l3>#w?F%dIh@9BT!~YXiXSM|^-@)V|!NHn3@@H9tgwO9+%DxuRT5M{&h= z0gSL-1g*;J)AK@4S;Udh$68f zH+iSQWp>WxsbJOdC8u>}ZMqsPQ0E;q4}5bCLCQ+szJMKrm1av2&~p;9yx7?H(#jD@ zNs%S!e1ij8Xc7aJ@N}419ns2~rFRo0n21mk*qj}wCN5g1Y^*ko+z&q)yU4fib4f|R zwwjX+xPa?zCLIL?b(9L3c#=@(C^SA9g#0s_GGKhY=0!VZbk~+;itl)Wbb$N}eP%Jt zT5A!=xf^Iob`d`kRf|%bxF<{pXiEN|ov1EqnUptW#%a}e;Dp2(#99k6R(`}bF_}(b zl;3hB>>sR{^`=tWM^V!UxzT^Gm~pXgA4%T@Mm&2nN@jugiz?FAP?me6VCQ8ql(G@d zuiT`=Y>8T^U%mY}a8f=UcS$&K6M1Q9J8Q|?#s(r)OkV4EuEV1YGr&q5OAw4?HYtu_k=Pf0hoZ!$oLB|4F3dcp4JZTJe-e9ss7f9p9H9 zXdrYvI}a#>4h(pK{J2#}l7Fkx{xwA#>r0g_o#3bSiwOb{y0##L=+V2!@uI&q_sK+zKHt*%^YX0<6LEg1&bSzd`*6a^$?3@MDvR;Z4v z29ltdYf9E~#Zsnc$Rlg;X{mZV{VQCY4=oWM>yP*)XvZ1gRi_7uO{mG)vNS&8{tbE* zTo1~vtSD4fl@EGvy>z}J+~#1*6!x}apF4N!DxQ%3Iv-J+R8w~-Rnn9E`7$237!5)< zW5v{plGviQtv%WODi7KVAB|_Hmn!(Mk&I1X)Z<|!3@~;0BgMYmO1L;0XKv2g<3o7} zuGrXaN{}dn$Nv>vRf6g^_ji6Z7nA(sBtw0gW%FOUdR$yfy8U*&qBSjICh&nXuy*y5 znkd%(d&mBTrX?Xo@dqKFqb+;Qwc_TVIeb??Rp9QL9BeqaL|g7B%WXWc3} zWh!Mj5+Kk6E)EW-7Z(E@96X5*9S#UK!Xb9~47sxC7!rP^I&w-1hgcNNA)xoapu@fX zNHKdv>W8!~)*|~;Cwg2S9JpTph~o(-25PktV`v)Q$`%cZMD8kX<3>U~ivfjL7;S}b zJlcWr{K`dnZZkTd60o$Y-;hC_#AY2g-FVZKzAu7zTovQ`*D_c?di)fDhl^T)qIQd3 zz2tGAf59|WkMZDG&4)MeH~|0zim)ZknY9frB+<~l0+9JAEhGTh)_ROP1iJGulJM8lJp+-qbJ`&G` zuQjWhY=@ux#nj=cjho-z61>56af&lBA3h>jhc|>tWDO>Q?)dT%mjSI2q$xuGGp~_9 zQ4EayMj$eOq0l-3x89;%ITjF)Q!A065;cbU=WLh-Jo0u1s*X_qrS<{^()3(q5ZrxRWS-sSGGOdEwQZt&{+kSUD$BjP z)}b9s`-z9|J^TDHb!Oqjxf-u0h+rO?O6YVFs&&rRk`>J!WWTmfU#hw7u!+4!>RIT@ z#h%sI@Y6YQxi}wbA*B&&Cu?lR9;w)Gyq>D^hEEc{OT`b)?FY~-ld+sFqG8v$9K!$~44USykO0=No?xwaNFaUj1-Ce^?7A^3c=hPYqhtul#dN?ucr}Vwc{(;CF zyBWIme%fRy)8I2bi=$ml8Tevz^n+{ki^o`l$%DpkCGXP;d&jjy?mD9%yP6M`HH$j| zX?ZN@Jk?RAjq-WQ7 zJlh~%9_+2~x?7ggaAnbYrIUZNypWkPqFKwjtsh!m)>mexHaD7|TH)kZy?k|Q7rHx6 zF9xr*roC(aM%_ifQ3FLvLD6me2CwO`akacnBDlY z(YJo7QzeIS%6$cU45BN!pQ@Y^vMAoDDz``hjCtNYm=eW8-EqJEpYi(1Z0#t>uhK5^uGyQZr0GZ-(*AZUoCAbhELYgc23BLAK|I8 z<|@)yK(&Q}fd*2opFy*hRFRx;d#@Yy4e%ZEO~94>a!%JlrHnxiNulCj%Lj~Nz5e2^d9GX9ljsi}ZF<_qc_8;_>%Nw;o zOpo}4xllSxdyp0;k-D7zRjC*1p5jHU!L%C-K`eTpsAXaeeDFrW{P;KI(<-y)nVB@L zJoXA^JNNAY#uV2GqecNWMe<7j`xk)p{jmOUW*{n{nG|vv5vBh1!|*)2Wlz#1r9{&J zFiTTbvrQ9`V5A;b0YFj2&CvdHN7Xa!z*C#xL3zORG@(&UsI%a2pkyy!pkd5p>nu29 zhWz?5+bdI|zz5mG$S*siAO%>9lHE$~Ald)>&W6NZvXmNsz#*Yd!5mPk1Np3cmdI5K zK<;00@oLSgx`-L0q-X#_zKYsrrtcEusYTT{WDGo;9-FAi!LvSVO4&iqH(Qsgg<<95 zi904mpD-&pp#_;W%r^Cu=u_$nry8api9rDEkM+)ZDfETnXF4Wr3kc^~{YT7&y<0Ty zX41o^SRU+Kq`a`W-Df*;Lj6%Dzrtroc{TV}+1Qv`d&bcapGSvkC#{99L1Rm~+me64 zH%-Tc_dA1UW?yL3X$(4*pifY}wB811{qdMQiWc~?tzMIWcG{GJdXuE+)%{ZWFl4o> z0>G5?IpmSmnvGdyKIBsPRG1>$Cvos}E!`&NkM1U~V+`8h=kk2?246huS@iN^HY~W7 zg1I<~{k4j1#BUR7m^(s0u85X-Zz60^c<4Qfr(?4H_Yz7`o)K&XE*mJuw%9qp#A#*v z6zc{8F8mU_JJl6Vx2vzyUgFpB`ZP>G->Kf)yOqcuXEA!d?ru;gFbioX=tK~Uu}RqT zYt@2;q{pMyw@N_~kKKL#7<^i+mN0A-p`hYvqK0UDTwict3JsF3d(kk`3|eRSic-BI z3lcatqE8>A(a;PjW*YTMQ3}{AOLS->+es4vv*8{RNBK&dacM(yU=+cZ9CIxYh-OY= zY}#v|XP|62-zaE0ikG%sOG+f0q&z1Oz#ClDs*j{OcL*b-IS+WCq(0lTkN_! zeZ-!tE9-{6L}I|ty12b`iKv>G#&kHL)@_!JH_vE|Dkkn7`77b4c$LD_QLOQVW(yDL%LTD>n`+|(4P#X55APW2M^2-;pJZp*8+E&^Qwgp z?t9u)a!6f~>4RXm3WJ#wcvei|!tiZKU?eP=ZKI5P+EZ+=A4l57+>@0^o!^bGz_?QV z;t9|ib!eUvB&_|mGuh!V$szuFo9{n{&!=FYvn5DlX(TT$%CxY>AJcOETq+3(qI){l zb_PXKNJFz>U3Oh%%Z;9C1{fu>8*xf)!)DfFEa|EQE6+%Slybr4+n^sjE}@^W-cT?Z zfd2@3F&QCUaA76Z%r~)8;Ms6zj5?V9R;YS)g||3CV4vYPFG`tK)o8F42b3NU90MXgb z0_tF(fGNp_5p4vMp zm!Cc0ffqWWD0lqI80x%IMqF;;rh2bG`v`DYnn3w!M-#`YeeTx#^S{x_i|)?vHZ_jm zf(hc`6AhI#N2)*Q48kv8tR9@2gQusbE)c#L@A`j;oO-I8AZaI#f32xGd07MkR zB{dg7!ULlo9u^+}NWvPFB2k%FsZ;c+p^22KXtPR0#43p7nmI@2#e=?q6~`Kj>|ym~ ztoEQa7?TG8sZ&UdPvxg7aK=67hDM-V9ogKV!}C#Z1@<3KFR{AzPdST{NRYn2p+og! zeF#K>d?bfi4yntpqMq zQGh2GGNG^!`!)XnZ{nQbORLjFQl->O-ShOt3Y5*&g^Yorc!^)vl7(5vvN*8z_L8FZ z4Q^D9N8C?n%v8*K3jbS@>9sg1-^D3gn zdUAA5uB&CEx2xw1&>V~jpyc=r!jjU1Rho5)-~ThBYsn`GK^igF#3U%OIFo@Xy3!mr zch{4NOd?+5E$5l6kQ6Z*g>iUz!e5>FI?||7nakfhG#ljshmwR8-QX5TqZsBSA26$s zt_m1|VWjp&L1}O*EP6D>LL}ZutXptT23B}% zD}Mge57xGyCYxCfX%_8z&Eoe)ROAP=&oxXQ%y81xhB}bev~N>nNuy!n0%U)FEMt7U zFD!8;pRalz*@-Y|bp!8$)ACkLP63Yf+ur~i^tCuQCWE9Eft<6`7>YT?=!y>a% zUWhkSfiIA-^%4raIVeG0L67`T0ISR=Rfa_$}Ww=wD%QNTy91v>2o zP))|{-c;IR>SV8$E-Lc4pcBy`p;=uVb+GNZ1mTWkA|>s{-n!sbu}%T>u=SUNKUV{Z ziD%Y)L4~tnd|O-dhnjP{udn}cQP_O%tU14adKcE`#%B9_72X@2gjwKlLxXnHUApk!MRJzd_4(DT5zR?4XV zKE1s6zSe;QqMA%mmDv~g)WJ9hVvDZ`%C}jNP)l(xU-;J#*5lupptH~}hwpp3>xt_; z2*dqNX9hhDa%Up9b*8&|VfjJdZu4dE5P#`gi)Ko00Qk z<+PZ$qzgXCh|+NdqmCnX)QpvLr(zkmP`vfoJk~EI+Cd)fJ+#ZQ=KmK#?scMR{1-KW zQj}!ahDZEwk%@nSQxqf_&&n;k4i%UO+Y0s~2v6C=I{VKIv;c#q2>)kcXchVqmBi9WYzvJ#E(4Secd1C#8cATX0zlbMB+ z9BzWkS{hCoL7^M(qIfmP3)qhtD=6od%K<~OYEkqO!UXi$`O?L;^absgF?a7{XEkva z^A8vH?7zcbx$cQ2bu;Zg9r`o9@%s4xT6;C)v)d^~_sK#y7$t{7yDnJ0yK&Ga&+cP= zx^rpyWqH`R`>sA84i~qL1anFzKi-O_ZNWP`&N-V-W~@{B0KCt)jtAV?vY&|B=v=X{zW3vB3R{5=4zk=E&q%Cfu=2O4 zru8DBK1&jH^Q8GfCCV4dNYm#bRHcGx&DT)C&?^xz3ck#MThiC(pR06FnjsxxPqoXL zv;yDn0v}=nq|D|59`-7f&!x}gKxMs7o&v4Wi}6*|b@RZIp_21HxAO*YEh7bCu63FM zs0;n#7BS220Y#Yl1?!8Tn*CVmHTi+ka7pl>T6U2gMC1; zMmaQFD18=_E`ec2r?VtW;F;ig?A8lb_cF2LoLCPbXCvx=GqGkItt()D5D!gWY}l|q zn%!SjSL-~+)PQvN6swV^i)m!bir&yrS$()MQdBTiNtf*2_ROfuTn3FZQBiV@)1(S% zo_vLnos_;C$@knsdOK%(Kw1->zQ`8@E8% zQI)f#Y_aM%vQYRe)GcjZxm)55>p(oYKBZ10$EfAg4 z|DY`j{fzgCgvPKGwTJlm@%ku6Wxy$%X{2|S`mnD#kU#02byy?B#T8AHg0m5r{#J+x zI>w3W6ChRz8z&M%W|RKa!USjKvPv3JCzo(CyGuLPcJSkRFQA~R59~TFC_4TFxy>6a z<&OAzXs&R_1-EM-wu+W6?wsxwXvgs6;nh#Ko?~U|IqPpVK7$aFu#wli^eD{5ky7Wy zbHkbqWoecsaKBIcv{&z9`L%}sK${H>$?S$&c21uWLITeLP*~8GfA!|=pJlI+YwOOa=Tz- zg*4b`>FD{pt|0psVvrC6NbKb4Ine+P2fGVQV=zrS2|8fViz=Ls&>g6UK6$0nN!ViU zA`CQ&MBadVd#beuhVT9*n2IY2)UW8;~O#WF=pH0b(XUWlzL+dXE2} zO6r?qt2zJ#0ND7grvGoRUCjSyCH)^S2z#ghMIkk7+d6NyA^r~efByC=#EOjak*6w# zWjkl3N*a^rG|swLiwz18nbiN&ekIaOH7m(|2j=<;d6Djb4j=&tY;ayZ4-!WFy#=VJ zpQTIkb#i?>PrdVOGmz{bxdBQ^+GjF(3(higsW&<_emUFe$@4-MyQG2=Gp9y^8^yEEq_g4P5e=LD0YRdqz-Q*}ZTp<{z9S`QsaI~ckP+r@QgFLsoY z{9Q2Nyw(_3)I*^9k8aUN_yGW1YR!Hcgd9tFqKE@l@ z{*#kdqjB8p5+_V=U@YkJm0vPrlr;{YzJ{-FHHgvRD z@=$3PF5IIPEeyhnPP3;7`jBIHH*5b^bp$P8 z_06(T>f@OT3fSQ!cfW1{Ha`6C{>~1byC+JVwrIv=H`l*^VaPKS2GR#X42EYPa{UFx zKC^M96B^er+h#bghadR2Wdov96RKD&I&8;UDYHFTBV7TRA@QgG= z5doSzfAj-)rwd&iPtf;ks}P_?4He}C(~s~YH_Wq}ORw^rS+d4m+Z){II*nN3Cus|W zwC5Yf8Ve5*Ila7mn7l2FoegF97*Ia~UZ+YTLt!ehUDPj%eJ!p=z@kv6(-v=KD>P19 z(-78n%P=jH?O^s`_ClqYPFNzDccPMkw`$_V0#k%|nDP*0!2@MZ3j+s*SmVMkIwr0e z#|_z(Qpj?|V`ZK=uH3{X;KZs;;4_>nDIiJ!9~N<%HqWVL&=f&y9SQ$}eG)_y9tNe# zj`7G6udQKYpxs&-HH>Ch(K$_dkGOOIEC5L@XcPspX0Xge&m7n$J`?^@p;kAez1Dv++ zpMitBW_xX-j?5@3F%1(<3y>s(9`S6YEbJq&hyB9%jOLkdh#{UR9q2(g;fyn3AlhEm zF7c`2b=cFTFYMZ09Jeh2Kn|h(oVagtU?Dz9L;*QNBM?^*#b6tbO51TZPhUg|E1HgG zq2pQdFhZ=GxQUH1Yysos=}Gah`@cv#r{GTFa81XyZQHh;Ol;eBGO>++>`ZLiwr$&< zVCQU|y*agKYq$2IZ@a3ytGep_xboUr^eQJSg;AFaRgpt?T+Y7n$Qxcje&6R?+FszoF}ujgZ8Dge9lBtK&TJd zc6q9|U)Bt3{S+vVc8m_gu@F=47V0 zv9(;94Cyhr6^~EguT9dmof!P zFDJ-^VSRFyoOTu*Lb+4dM}WLW1doc5@D7N0q(EwNzIvO|YB!>@O5aCfY)1D-lR19n z99V`;(*;H{+lJ02Kcp`MnF%(nv;z)D8%J037oW5bL2V(ciAqup<;mG$>4DV9@PaRT zRvkEG$RH5(@%Qv33FLi)&fS0F?(%#&Foozr)B+A7=#vSZP0!9LvsB9Ce+&Q z0{`4Z{`veovJXOgGu^brgVuyd)IR@6Ao%Oz_djN1j_c zr3pDCV)||%Zv{5ORyV7K`qv3~#8s7W%_+R($~lFxR09YjvjQ5RT{pI+&W6Am>*B?N zOSph38=HnsA2O^weTJ)UGr63?rqdw7uDphg_XBY<*^&zSz-v)cX5NM=B6euq3Va9rX5SzyM&B zjW|Gp8j7$3rbNQiCFVN9#e2F~3i;8WCjieZ*=IFR146z26=X(kA9;H>INhmntMG~= z3O@M61T;~?G<4-G&Qa-DKcLCb)K(^gfi`b)oNz=1RReQgE;d3p>d&CwZSR)}Fz$rh zd=3ed1Zh{p0t-;4pRUk`9J%BhK5Rla8z95cjfkxl=5?kF&dk4CBAzShol*b`Y~J$W zYhwfTA&lD_nm7qmem75$mM4n!U4jl#Fx2Gpg*9Am0ulRw&Cy&$owMG3O?d?5tJ@Jq zWZ}M>HM12|Z9!mFG-yeGz0a&lJV{3~ORui24XkUj0qjATYxGq7Y*HY)6a!O6+#>Mq zJIO`2sS{H_i#ljRwU;?{Lof=4cK+za&RK%~u_-&STB(($^2jG_`Q>paS@NK+1Unu_faVWZY*O1GD!>g)USL~gyS2~2?Z|16{u_KCT3gNd_!6WK@ zv9mnI*|2Xk!#g5`U`m01f-y(DnTw|QnV0FbSzvvF(x|p@)M+0Gw1cpLopzWa)e^Y^ z``4$P^7zLR}k$iQPAB3t71VgV4yUdQ1?`Y5;TUMO7 zMV==8fG(fEo%e&v^3KO%O{xla`=hV7{HkwNPT7ro4p9Z%+nDBE={B!x?To(~G1D)O zd1_yLi8xjtss&8Wtm0BJbAE?p3gj42fV94wo#txVq{0;VXjxsb!D=9&35(l-xi?E zH+*G5kv&wt6cAkmUqy4%fzjSSfpc0G+;7xY*KPBVh&WiI;`W{-zn(DV&uU8oJ*+o! zFGM;M@Dg+O+y1#c8Mun;2ftA38Jy8+M^lpv=3Se2Qw# zl3J^d%q;xe`YRvNKlgW|UN|RINis`YLs&tzD_^ z3-LB?&g82qp!yt4nyI&K3KAJ!mO@Y(nlh<7D)K&RI@p-2()6pxCaRS@DtI>JgjI^7 z8wxdhE6}FtwM8qojNCG4gGa!$Dzwj^bu7+}TL=RKmV*-cku9aMw#~u!#1IUVEJtEc ztB7?^sV++XH28x~dES!tbq1NGW0GUhwXmVrl_;UXO)>A0jVxp_QE5!?3|EL$9~Ncg z$}|u?3?stvv>i`su)exoj3Ff$bVLA43!65_?pUN2iGsnM9!&~0wfSoLGa-LP*SZV# z9CTMJKmDB#7Z-BkAh!Jl_P98?8Ug_vA1=BC=Pc4U^GGwJoGRc6S^4+_Ag!ai&{D}<}xe14rq<0wVdeCo>!%{qO%#& zN?6jVj3EAp9cRJKB`ecSD*3zXWPER(W>F>AW6MME!n(6Iwe~Hil<$@ z``Sct{|Q3TsgN_u47%bqYFD=h=`8%WO$CMjTt(*{bVF!>ynmDFIXS1pHL<>@$mj>dtQm<*SR6^a;&Z$_TI;Fy2-dLbIu8>aTkFvAqc;*Utc7o_8pkc}d zVUiB;#jFJCmV`bC39(8vbFratiiiVObX8<_UlJg41&Eh8=NPG>B%Sn}EVChIqq;T* zz#N1_b6%dl9&WNGSt=w^WpI>}w_T4|cHUw-YfUNcXx&eR#%gS7fS?Ahbz;UM+LY6Z{)K;o{{D=skW<(M99ZC5e05}v<8X#^3m(Y^s8 zI3a4cQ1mPaB;F=}-p$XY;0t8)lg~AA-tLt8d-LWzC~O0tZ>P8oqa&{`xXh`v;j!@O zR|5lT^Y%t#vSdoa-=iFF)FUkM=epmL8)nNta?bs3lG5#p3XJ+%bkIAJxE=IwiYV?f zucKs}K47>yySIU0tA)0;exHNAqXV;p&~G}H#$bUjot*8x=J{2 zqA>ZF=xDnX;OXJz|4`7Rq7H@N7&-SmE~hX}tFZhmCp}F& zE;lJh-=OSBCnYx{wamW2x~Cjg<~%ti?KZ<$MmHfhO9d@mO;xWP;DVm?_$U*_k}M

+6dS6ZO!joGjb+KL?~%jQNU2|6YjAzw5t8CI8LCnA;k;SlZbd z{a^5sW1a2);3dCb|M3z+680pNc7@De6>@&urwqwMP9~JAD~AQd$Vr}sQZPj1v{v7* zHy}jN3Hju_tLLKc2~r?+npdL^fnjXPk(KKPV61XCyH2#cwhcPALzD*?*WZ$zr8>EKI+?zaW$D-aI&@5fMuG z^igDHWBR*zPV+O8B$dJ48)6Ug^iN;yo*HLzJQ=uPrM3vXEgP$L*SA8mq>X_H(R zqmg6MO^8r@cYnDe($hx92T*Y@z!MQc(6)Wq#+6T|2R!9O*`*clERTGMrL?D?)@h1} zZfoTbrl`hZd|(}df^G?1;cixf1>{v!5t!IYlEI^Wv&xfw-zeEDTHX>utr9}64psen zrKm8k2LC1-e{7ne-0a9cr4e5%3In^m?LkYAnD4ma(X9JN=FM%H_aZr(P!SHw-U363C@wXlf|p} zx?#yVpyGIa`>ZPU}26l0c9FgAOLVgUP{xQDcMlTZd8s+XUZJ6kuWOd4VHM z)>K$J`Bek(l3Zp=gs~0;0=9s(K{%jExUT6ReL(u64iaN$NP}utRW& z<`3nVddy{vhI3_zPV*FYyZn9#K3^Z#d2bJ{dllfF0%U5&Bk(F8f2NgsJ_=C&pxG-% z7Nm1HhbrSI8Db1)>mAedGGr5|&e8XbJ;F&(+O1VPA6SHMG=OH54$4tH3QMT!Tg ziYlQs2R2 zquZcD4-y=!6?H`%2m3U506O{f2tp-lFb85-Sa1c~5_&^3HA}~kXYbf^;nrAd*k_ZT zb9!botEr@v!7$C1FWuv>Yti&_adBI*Gv<9rkSptP>qz3BMFgyJa{5KI6_Sx!;zXQd zj0O6MP*K!M9D6ynxC<8CQG)1%T{D3Ofq|1lCsC#LFT_X+O%@wQSrIlbiv+UTw&wI& z_x<1osN%@s(O_r=6+jJv6LT0Ba6$+RM}g%NZzwk!5h=TJKRA%Gw-cz(db&MbwM!JA zY?BCih@+FKk$+x|R$*ksil^+;BhmTW9 z%V(n@=iVDb$*n&u*t!DYSQNr8av&t^1-j15OBy|TXMPXKlQ>%6i})vwNg^;*dVe_m zMrIM7mpY{OdPzH_fp@|c=`CT(?@=*|ZkDm*XD)@NWhH5{u%d`4Siu%JT-mw23sRobyNHC+poFl^p>z`I6`%|gB~OKf zGYw_ws3^e65;>?!M1(7=^rE_Lq&lMU+@rt~{(Zr20OpmOdts~`j|aDAhWSZ!CG;OL zE}F(`zQ-j3VD4qGJFF|L(}}iCk3W)zP`pS^;kyECZTLx?r;Gs;U@}mJLZ~bqj zP%N+NefVw6*>G^h4jU%IIb8$wQMsl_Zz6*+(pL^TMQ?KOj_S> z9)<2U%Tmo7&n}Ks4xkY}4^W0hySaoOiPN_^J@WY~;;3BJZXLZc!1>e%AH1sfKZ&UY ztHJt9nLS`CA%&?QYve5@rS!DZ#Ma-OF3%nfV2r>CGL$Z3B~Nhjw?5@n3ZSqk$mp;S zhYT1?pAIJ%FtFZXB^LFDZ!?T^pQACGi`yz@>J26~Uhp?<}&aR-*Y$*v2$aYD$`TrHTUR=m>j!#R3gm z6?4B*(|)-uPgKcghknaS7uyOZQi5%&q8lqptsBn^z;m;NZ({-%FZNXYh)tmI_83aa z8Lvr^GPn-`L57SX7td!wt^~2BYe|-tHX8_Ka$j$xgN3(jI7mD2uP68GsK5CF?aW5;a~Ou+OMst#avn*H$gtT(AhX=SMoXF8mT#X=RxxRc=u41d zSQIjpafD_1caL2sA#UTL)p!Tnl&PY1#grP!E%rYuyv7AvoqkSuUNv76hb};>F+qv# z^jhjj?2`L4Nz0;$W6s}E+MaDtx^U*?+Y9}iZ`g;5EOex~`)@i4VPwwa@V4OaT6(Vn z=xZr397$8DwdvfljFyb)mF6)y`RdYlH%k$a(R`dcE?^T|OuQjU{GO=DT`Gjg_sDj{ zjTNO9!v>{d}K#r>dcbtp>s+C1fk=BZB4OpS#ZoCgrG1KqDAt# z}wIW9fX#Y7;hQN^vr4XH`Afm4B{BHp(QQ)139WHq>9 zg&#gm-Wosj%5@lg9B$7)A$X;HIFsylLj{2Y{x`hzO9t@biChlkEkKfDrdi=|l8~s{ z_m&F0UD6lsAx0L=C<{A?@6?E)6lRJyx(^DQ2uXAz3Ov^(F`P092klI!`cxGt!&dL7 zIJ%?luMcK0Y8t&M5sF+-g1&r7g^ z_mI#fAb^dseoxU3@+q|9v^KJ^>Wbx%-qqw)2~#h#IF3}fvDJaO>SmM-bw9s){;t;f z=gF^T>#{9Mk(BXoAdLWk=14C6)~Thz0JaQz&^scYtM|nW1Ar# zz%cEmU#G#~sIo|LNMg1CMgPRC|4_!0r*Aq{ryh!s&I(xD#7o*GLJ#i{B(IFQ57VSNQCb)LNfCFRjFK4k3{9t<&UYXSdkOy zQP~ml>so#}b`i4?iP{m@tyw*3!2Y-*PaGuO#6Zi((9L#+dvMc!z4;e)*RrHUVtV&c zpgm=LD>W*9@Qei&Y!v?!ew}rEgzC0K(*I)qri)Xo~D5jW)?E-hOV6O>Jy4 zKgjRv{S=kmIw5n8Y}$~*BB4BT;j+)ybeiN&4VK9uLF|+qc^zwuS(LF$R5zh>{NtZb z_<~I6sCJ4_?;%Edq&}YFb>5=Y`){mSj~$Y-n1s$?oJ8m+l5JHFHtee+=*zZ=v)pTu zL1*Qf59E-QN&Lq45`hOnUeQ5Pb5dk{PdnE4_k!DQzddRIG7<2O@76Sxl{B zqT_@0V3(O&2p|xtg%_!gwDSwM^8lMbGkaPc3Gk| znxFVb)+e&($ecFjChnSucTmT~k0}x|6`DpM-uf1RA$N{u2AGlb!ZDgfx@Q}nY3vdq zIIj!K!{?V{stL_i_H8Y#-%Lr1P2b+;S-s%vMJPNtGH*H2HYR8so4AT}ny0ICj2=~;%iN|X7p<5u65f~lHSBSR`EA)rq;cbmD}Jd&;)YclL!6QnoG?BS?ywfC1< z)9GmRq8uaGx@2gv-0&@~`m+dVik_8=YlT-wu{fpl4JMaUnx*Kj3{gG)Bz>XDcXZ*x zN-t}dIQzo9rW^q6Z9>Cz?s4T})&hHl8$7bJ`lT%-xYSynayJc+DLLtn$^2Mpowep~ zgZjq!M0v9;>WbI597+{GH#ZRg?C^`*&uh{~K6DU05ZpQqyN85)w=o^U>RjFIE>H+P zytw)+i_c5Bz&uN~los)W7bQjRncd+d=gzB(SQf`F@)9k)<6FhS4hlnrhaOE>3x$}p zffXjZLhIrAeN8>c*KIGi^_$nH+D(+3%z}_naWjhxW4(~IX}L{UPNHX2F4RoUi8no;>r`dXuF2RIZuO$M?+B-zW|!rIPiD0ypTc>woTf>xhez7 z@(32@UC8i+;If&DLD*M6vTDY!O)`}nS?>D<7gEQ; zB*qSIP{@xALWPfrM)7Hxejrd>T@?ZpWtOLV2TKPT8D!8 z-CKn9#^)Fv^r8(WvNTTfL{_nzo!x>jjhqC+{Zth^j-NgyekB*Uz=^-NJ;2ZiFs%Ls zF5U2vr|-7TY=DI@Q`Y!-7yT;D{C5=5{cU{B)HQGIz|OEzG_0*h{tEE`GBput{rLE8 zvHXpr?fXs&=-kM8yIOx38od(K2Hg5OQ2(~FA}n`B_y+wCIP|9$f;su0Ex}v(|Aaqq z|4UNjKUt2QgQ=PA{{{QF_Oo%>l1SR~_>Rh4k(X=T3K<62EA(UAW~&3t@rgJ1CJ)!4 zG9im_m&w8`fd+yRZP4rUULb9^7`F_ZGoD5G=ZVPIb0e-1O*s>uCnhE|%*^;n929#$ zRtGuX!vOUBpFBloX$~$Mc*&S<9qq;8m1pDqqzKPEjAxS-JvI-~lQ%xzoA{{D$+i^@ zHV}i`-lAY* zlTB0$Zo+9tmvJ1Qo6@`cX{}Fsij(ks3A>*1+(7PXYE2DsUG#=>kPXXW`g1x-EqCH> z#nv2YBk^T)$i;Mso_|)w>!S^IH$O$x|AO*)kLOwvY3)1|EXoZ`BSQHySfl#C^D{A4g7Vtg}*H!-YSDK5BMNt$6A z^e1zkIrwD&9nk^%~ARElZVr(3<|d#M_edo-HI+I&>`)0qu zV=>+x_0`3DX%&Y*_W1XTT}85BGedWOYizH@zZMmiF+9|Oqf~P=g+5dBK8EoNrpCXn zJk}@%-hbOFUH{DIa`TNEOuIEpuY!8ycd}16Z0Xlpr&IobOakJ+A?c_Mx{{71e)avv zP|E~{-cm{dkw%17yWfkkU+g(UxsMaQDX)!vK#{V}0y!|{r+l|t<$T(_uH3{~^M#FW ztG$=^@$P8lDU$U4WZe7O)y&)AtWv}SFD<3?+8iju);O94Ta(saeLx=ThR~E7P|7_r zlW=&C?h;?=IdrR&UYl0{Z2Zwk!L&LK?qnC=X_?o{n%c!WzKqi~Lgp1ZiCSlxK;#v~7P zkAs#v-52*sHbCQNn$tt5ZK+e>d2l$&swgHi|A9~%p+u=k#`ED4E7KKi^al+JZ>b~N zG--g)gAV9%@P>zi(mYVUM*{!gq1T27pfcb>b!AV#Wr1^lQm#DQGqn+@&-fx!cQGI= zbxIyxR<|ybs}Qcsp(X@^4Z-;z^C-y%;7R#pa9~e2bu7rM^rM+btLA7sgMtV!t_Lfe z=^`sZ@{Q><01Hp)-yaXrEYpA%f`YS%ry}qwsmo|}mw=U}J{J_Dbs$*{_k`kPWpxc( zLq|yWA)i`{_w-Q56H? zzA@*k*|xoHxH;eT(&WLv>|0c ztGYd45vWOKN%w-C4ecp6X)EY$K!F)w1hY2fTvT`2a^6#DF;mfK{dVY5J%V^73$_k+ zwxfsP3_FU5B7Du})7^h2oH~8`EUG5pmw9U_T{a$Sjime5Io0*ip5Ax zPwWy3z{{GcO3I%}|8Ini?$h|tOx^e9D#$Yn|5Wp{Q*ry(rBXY32{l6eEMx~HnMcMGS$G?uf-mQO zI5Zf;!wjvVei}5V-1q2!LKGUIp0!nmEPASm@2uJ;X;7f+H0YBlA|T$`DU!}7j(&f) zWpt`|pxR&dO9Mgu%S&UBDBd~lZ!#%7cLiJd5M`@Ji2_)LDLbRVLx-OWoU5>vh?%^M zEqq=bLDgANZ37lz)B?OGso02u>F~tyr?-i96KATK%_pAiKTj4gIBk_8T(0%5LhAXk zH}k7^MT_HU1{)u={tY)ZCCs-9i0lHK!WbU+FcjQjw!U$N$$w{1`i{*~i4XNbqiW1w$g2FISr9hR-kjQ3YZV77 zt1}V^^V<0}3Z+?8p74@@8V*2*d1P1JWn%1LVemZAE|xBt4i9~0WwbhrLF$Ld<~1^z3+tk+a^xg?URbtbJssFU zF{S;SMco+e7ZKL~&uC)%g3k%HK0)q6fKfoMpddFCj^zZDi{tI#E(v{l2t|-X{Sl~m zy`$x~_S)czNwP0s#W1O5ckd)^HgHNRmj9b`a&awy0({%u^-kM&c!On&oOQ#Pq6 z>lsQmo&Q9;ZG009V(~#CZYZC>lDl~5o*?2>YvYgQ&wh&d+u;ImJ_;I zC`BXoI{C~U?={XAS8i6z!8eT;7= z7f~vX87Ta0NU*wC47|vH<>FQ6nsh2ceMEiDcninyjdq z!pk-=cHlA12W5tBJHHq$W!pcp(C4#QL4}+s7m^#MhHyG^O|Q4eMk!L(z=GS&OS&3s(h>Rwsz|={Mv}(_2#=o-1*&A8HC630%OeSsluZv3joVTBhfd z@)7H?uE<>fR!cS@8r&DDeBZ4XTgSUM9P{sfI|L! zy?|D8GpmSkx9WR{R&QL!g`hpiB!%81&C245Nl|Ab^UcP{d;zw@v^rF(27E*-VX|UO zySfLD{Y-|-T9f{)uHB}8tfz!*yHMF(RGQ*c4jZ6e6;&_3$CJr- zvX9>Ff(3;&P~~j{{(u7lH*@L#Y@6-{wDz)HYC4Lefq$Ze{XK)gc@gT;7hGO$MQ%jy zH6R&N0hw!)I71LuIs9)~^O9^SuE0Yusx*IHteIw?T@Wg_iw z60szDm}Atxe3GESngJCm%FzL_Myl(?dBHO^XDV)Xu;B z+RR#)5v{f5$!1@`r_cz~lbPwA=zMzGXf zdW3=~R31o04C0mSoFc394SQ>iI-`?^fSXU7X*y8S3V{bLB=r4gv^F=VJ;r#oDC9-q%36b&h>a%9F{Hk_LTF zumFhm`$Qd+h^YVb%LlfX+s8@XxX=C&P34195!OB~ zsy5zWj7`3SgD~W`_-wW;yjrb$1?efE333RqeO3$C z+|7cEuVKhLg439^s!Pm*@^-y|F{l0Pl-l}1Vu zsyT$q-nKw2QzjsnLmMO%2ET`AU}d`{bGs}vd=J72Gtvy}cw+zAD1{ft6;w7sVAR)i zsKw&cL#z9}*zte=sy%U!fAAgE%?3lK3=AGVo*0d7zGwY&W1M7W@;=PYYRfi6O64pI zc|yD+sJ7n$SM^~HLpAk4IwvE`7o9llb7zmM3YzzttMZsPyKl3T!*MSBq)=w2h8WG# z!Ou&ee)tZhyI@Zl6eUr*QuQjJG{p7#%pZNe@1)q-bKR!gUeY%Fi;ZpiW(;Ds&y?&E zm#>c}d&zhTh;TX0O7?kB<9AUN(l~FDezBLdrkG5)WNQtz$AQ_Gy@!QH6Y~W$ws3D7 zI@u)rOqAUh-}=fd2tpKZ{3iKP;SOXwj1Yo~eNUZR#G1ch&fkOzv8tZ;Phk|53rq`QjQh__f5EvusuV=>>k(Bz~ zwPP^F2Mqk|R_IUhb2GpR$~M1$Vd6<{GSvvPq?*!iYoknK1vrY(CjSEk4Ad4kmFj8; zAz7oco5@%QiKkaAj+LI^dG*a2hL080*S#vtR%G`w)!JhhdW7+n$Q_*ib&2mdAM+we zx$Z%yJJO=xN>W7rLJKn(wlQzliwLoHkLACFzaE5^^U}{NJjh=%%5TZz7l3}R`bNAV zT<#NFWQPY?*~BcWWQE5^$S8WVpF|=`W9QoBFngnx9sf%CLII3{(g@3_o`j{#5 z*1?OTY*%3f%Bdg}&USfM;#h*g;$eV*{{kkPV9~6<;@;VoRnMPnL?8APMI~}Toik<( zC3-5x$avJJAIo!7AD+0gCBa}E##Cb3odd+{wG$Y^{Jc%|;%T0hV1Sn}5>$tupm>Vi zaSL0|&};yHpJ@&R1DRJ9{FyU6nizn?gkVh0Q1S%yE5y?Zf_f6{h8aE^rb=7{|;U zqcm3WdQa8>AnY#?b0%!jmly(2R)Gro;)7WZ(I1qXJte@=vT|8Ef)4E0%Go+)A+FIf zfVP}6(N;F~%6G&L8<*0xP(Kb@S`!R#-7`gm!|NNcRj|WZFF(Qmtgt0lQsfU)Y%v}D zy&ySqDpP?(b<+-cMw12)Hg)JA|oD*eBF1LNDfIrI4~#P6@Om;qCFQk&n?7bjyR{6LA(+k zS{}%1wAbfLvpH@AO^g8Ue**vY<%26ubeHNE{O~KU;n^&NpJcZpt>JXO@m_0D4&OKF)uaN>y%>h%1?=;RN&r7Ylg$G%C)Zj-Q zv>&1T@UrpNir2Il?12$i&{7@u>}dWJ6>G&|HBfPO**-KDO?6 z{Wg#NmC-rutQi>!s>GOvR8TU1934fiiq+4rAj01N>2)kz^!mXmd^EYq@1Zg@6PHI> z(!(Fc$oElJ=NJ)_-~LX_EU6hXN3y@AcASqP!YAL;yH);B zex0-rJ7N1#-waxQpd3T-21nf%&q!&C&Fw^Dk*bDkj&iYN@#Y}jf~spGGAA-K-inGX z)CL40_p>2`6E-ub66HIaDNe=LjhZl8tzuf9i4}GyymI-zc(VNk3QQ2>J94nKvl)f{ z(CM=lUDOM*n2vi!wVdN+#h%O9ZAID731KmvBSdKEb>qZUB|gJ9>B2PETf*1 z)bX!}7w_bJ$= z7g!7ZyBvm(WY_oaSLOSxhnkm@B<7J>GboQe^Irlav=iM=(9L?!!qh%6Aau-JEYzkF ztPtVx@z+ukj{+sH2WOS16F)o>=%W!^@b1LAzlBa9r)00Cs8Y3x#FDja0r#MEkqFLg+!bix@y^EbEr2{0q}F|qM7I9RhYKznT8S#Yi= z3*Wj*?>{q-re?HI;6Rs?qq9k@1R(8(5ctzSPx~*x}}%jG0G`<8`F{0mCRzM zncObX;~SW{Bd%uhe8}l|mDt+hAezQ{;gxs`pmWmNybfnnAGBPi&53ulKn0-3m0Hf>qm1hiBTlhy-xrX<5aHGUX#0Bz zHZO7F8)V%EpHGJQ`pK;5DjvE+6vaf?TX@Jc9U|Crzo?~y-y{~5E^}(t7cKCAK2!Fx zJA;_sjb*>}_cQ83Wxi6=$h!tbL&UX{t!JhqArLll{W%Y_upVxA-u=L08$rY@ANStti4*kmy96E^y2rI_# zm{P})l_f&x`|M*pwuu_gvJ$G7)N;Q&c{Qc~$$;fW623xFnL{7(l1MD%EocJqH2}Ls z+hTQ?aulr|jKMn4{_M^mO-8Jz4W=vbOif|Hq#x8>YI%8W*uFHCaVTA+J^PJO);lr* z0b7BJgFf1iZk-NVZ}Yk~D2%Tb2iBW)n|M4dlj8xwg@z$mwU8e*EJN+GQ@zYvX>KMp zQle)it47^pBAd)Or!Tbsmgp(dkR?l6%O#3r`_DLQxuAE>dK{*9JW|4ybIpf|T@?K_hi}hk?ODl$ z<>1Zbho(o*^}vt3&+T$&c-w8M*DO7|{^N}N{W^D{QVU zs|47e>pmdZv!tdn#(DC_D2VuXhF*Dw`?qIb@ddOC4C_fOz~&51#&# z(CIsYS{gyn)3Z32oy|XIL8n`rJBN*k-O|zVK+*?<-axi6E1s!u!J|ZBcQ*Oogbjzf zmsy^Odij@%ztLTPPI&p%h&$oWHGSBd@41MJ3K`aqTNhI>@aFvj9;+0izvlQhk9`Djd;VE1pVL24Z{UVSg3u%fIml(g-~P$AdRl`l@iNUt~!aMC|}R_~>>A_RWp^S-ME`V-gyfy2sa%7BAMJIFe! zo_P?EaRT=6q4pBGt#F5Hy>R_kc1u*@ug(Czv*CKafZk#Aj9TQoxQZ>n38bYP&Y<~V zWCT~^8rzPhp|I6;(Z;Kq&`5x)JtM}4{>pGk`8r_Irsm)L`R|=8RiPRgEF+1FhmJ2RW30~XQL9a`^rpry<`9jT;rhUiOn{fh4gY4 zNbu>n%*FmdS(kh|@W{)%FC-*Wi-}a%Hz)%scE*#;h+muU&@?dy2=*=@uV?;-`=ior zTiT1P7yQ7v+9)qK_lTY%`0y`P%V1_6rNtA;;rYbW8@(lTk!J1CAO32_ve8%x4qp9# zC8YqGew?*arvz3UDdaV#WWrqp5lirxVOY5}`=%|)n-u4iw<7%RCcR>#g%%D-z9+60 zek3BiU*`ltz~_DL`+_6b)zKMs(5F|Aho`K(gzIwxGwZwpZ3cxuuaCI|enC`D@7KuO zc4q}=qk9mj_u83t#%b#QzhxRIw%w7vCgZLz0FHqAWF#!x5(V7Cwvd-gStig5W8zv< z(^ti!EmdMWZSlH~J7a0S?>K4AwTv?m7XuEp=u$kJI%c7=hCy4qr@bIPQ^lQ_j6|7L z-MZwqOd~5~ccs(b@2{VFTH2qvBwP%Y@G>9?I5{DG1G+falGr%(IaQ{DT=)LbZf?nj zOY8AJuhIUWzD;kd=Psu5kn|L@^><&~@4ka|6R;J~hVh`JJ)D7s|@VBMRrVz+xpsMk%_&;y)lq8`ozWIQJ zrlH=pjp&pf<7iK$W6Rjv7YpBt&(m6WPlg#tECawNgCdj|w*F6iy#i=N{7tgC;UwitZ1)kXn>(ZcyWRw>4*l^-guoW=pk5WTNC zeX|QEuK3&3dx9_J3d@P>vSU4G?#AMq@EkjJq|D#Sr~(8sg-`&~3J73Zx3LA_6U%xD zBS+g|-=z_{yReaLYq%!ts1J?;6QTUNUHQ=Y03TlDl)s|K&5?Lj+{)Jl%!wbUUi)Nk zKQUIom` z4}pbwg*N$VSpeU^&$=*Rn#-|Y_%gc%HLg8!F*ccqk8g(H{|T5O;nFwr4R?-&l3`Y( zBlld#clNG`CI7{#?+@P;3RuRDpWaIYzBIx-m0N3W8ua-18o}?JOv$6SxmRz`+(k6o z3H*qd!WIi!?PSmbbuzeWExIxD3xAT4rxi zzMoF;a|54L`PIYxJhJ8y$RmR^R>$Fczk|*kyuj_V7HN@Y<6=5hk2#s~wR(jA2Wjur zELspC>2BM$ZQHhO<80fuZQHhO+cwX(F{kg1ewmo==!scBU_Dh-L}q1vNze3yd^Wk1 zlTG7_VBd=|KYRO`$m~&lnu(~#&aMGj4|yt?y3qa4*HSj)Auy!kkE3223HHYOxWIh! zyra)$L9k|R>ONv>DOG%GxUMM!m@Bgkm^(S|B@hS0TYwC7+6>wuY0Wx@`gqAot84z` zT*+xS@@QbRN=uAl|6n z6GKV(gtbtg0MU1PKS9bpd=9qKQelr8uYfx~l_7o4Fce>(Inm=b4N@0R<9${wjFCRx zaF-CRcAx$!P|_~WPshZ-RgS;#VVibMM;e#)+PD-0PoUUfz;e=K@M znf^@+!A8*3e3J2+OaAnJJJEXmtSizKdN4}1EpgO_c~r_c#%Ynp+wFcG(&F>~HRJ#?l!+e))^{L;s-vd!H9;f>Aj^VSjQdx4XYl%-)_X z&_c`I*n8!CRvg0I_Xxnx%-Hw!VH+_|z(J>P{6OqLpZJ838UDPAOkn$!eofVyY2-2D z&J1(o>rBj?w2c&M>^Nb+B=|#4HO~C3>M*RqhEo09D{Ih{Qrv zS|_g9Y~nVSM^Db^_zt-gAfj!}_`UFBAHe`+Bo(Eokh3cWc!CvO!^){ARdeF5u~Gyo z@&g}$nY(uem@b-!RA=P>k7CblLNm??7}bBhj|0$ zJK1bmndR@z(UkmZEX@Augw`rtqkieh^b}WQl!KeL7-Oh%Bz~`AV}SziI|9LIwWYS9FE=N)qc8FWws+;LS;Q@+t%@P^3wzP46&DQDFkd58#kp z56uzKsR33Us@H!hN*~tD>EifxQI_>8KDF|25hl{s;rkImnkw_~faM%aOEwrYezkPg z56QrcrKO`q$|Kpakhb>hlq(YBb_#D$j<#|tMT@-3UIVou%@R}~ucclO&;MO48rT>H zAnzc3-qZ4+iD9ib@4=k)OJVbXUxeR8?;CA98>=QihtB`-W%_H}%0V`BJytZ%z;^Rb z{|fc&(=>jTryPGymr8gK&zl!W$f9L1Kiyd?$rNaFJeCpb%#!5wl})8^`T&`OSVV@b z4Z3mGs6WZoLB@c8037ckqX+t67V|MwR!-0-)=EivCWQltG+kio_bgmv2rn{gi97^w z>N2+qR|iDwp`mbJG6US_p+&EYVCEkKL))&cXsKSOvAvbJMZ@$pgj(@q=)1E~)guEP ztlJJCGfj}|dBezN*^v zmi%ujrD#;%-$58OqIFL(mlbb>Ba<(`j}Q3_q6LNOO}7)9)-jx)K;(s9vul$rMGLRz0Awb>FaiJuS~VZMUXk{%2%|TDm~hs)%LO zs}C4NNO3esh-vC|qR(Ks`H1&Z1pGn_oXnHCsa5qg7k*yM5Z{gI;H@Vy|JC-5aWTEy zFi!3yse@UjDw3D**PSm*Gk!WFLErzMC~}?7w#^LZh`9Py4cjjV|USor>v4!L2tH=8JegG7{a%sS*vF=GE^P)bk9uOyIxnc zB`7iqDC+gDxCEmZIJl-=cVM)RqeeNzb&p9-f)|0fTDn8NxN+-`g6wpVz%04d#J<0S zjNnjh>+*w~^;*7Kiv)XE5&#TSC;IGThg!*Vb@=$jIAx_s(Vw4 z@apZ!<6!SD0EN77^NcU;5SGPE!~p7w&|4q+;T6sflf$2hRSi=W76$vw9Cs!Q6r7S- zim_?KdJuEzy*DUAI1%~?A`_c&4jq=W=JH?9jC;+EE9?|p?pLICiD(W2#edYYvm<$o zb1QZg@o7UWo^V|Qn_z7_{v#uy>Y^#ctEy%j-A5hr?*n#k)nrz`JLZ9p^%3)J4KMsK z##GHG8Pi9l>t{O|&ZmoowujL={uZ-id;i&1ii<}DyVC3L{};$MDR9f6palS++WzkX z+5bDb?>__C&i4OO zq<$1;i?qe>n>7*+jyJ_(vC$Nhe$&;J27Aa=B}8P{l~pSO;+}+H5HO#fSXFW00DFmn zf7kSpED1@{pPhs_-TMsK!N%AUIljES8F(~rARj2b2dRM$ z;$gNyJNccdX|HwCpOMrNr&O|xebUdygZ`;K_mpGISsAs!F&cd!HZc&ofJKZ2Oq3*j z3zB}!5NZIh;;JyAWW`=ODMG{#y12K&dKAP*APyi{ftzPmAX_TLL6Af#J$-C%Sa^NN zW9#=+{_w?ln0%T&pvsx=glegQHEUR%c0?-#u#m%}CSuKS2oho=0+a=t1YqtBwf#l> z1EB_`>VQ;|F2&+v&N$%}xL}a>4{-2_&he>8Rd@(Dfhvp>%gRzB;NI|bFpr6)WO^k= zQ#%2oJLR1&6!ybh2T1%s!JIVBkEG*FAQO$itkb%gwZI4qrjEv;ZxD5Qi`dX4a>a>8 zCKZ9$OMyG#RMi6F4@@K*u!>mzPb@SZuG~9-O*2z0aHbDg>3A~L2L=VUB9>xOjm1za z)X-uD|1f4kJ5aWQf+AEw&s0!U3ggM<>~s}coj7y+hih3}Vyf?yWDo@^F zJ@Z<^kr5iC7mX`s!V0Hw_v}=j@4nQ6S4sTQ4aTUE3R+-VN<{v^HAALvZSuoK6d6L? zgu$>xSNzC`@qmVbJ}&M8SKH6#UxcrQ*u=@_S<9z02{9R@?VAnrg8%gzx3C6eHgRkM zsSUD5C5jq&F}LoLDpakRgXvfEfT3fIQhtaNLMAZxKZoevW_}EenfrS{{lrx!hul18 zAYzz)b#=3LbTRStl!b^)gujz0^w2H%(CMyMq@ZV+aJpb4A?NI?X)6+h;n~D@OaSe6 zZVR`e?|+%t(_8-CaEj}GNU;0~ z3bBR)remGhY!n4Ph~RoG;5Dou7z7SxZVBymkAM?w-lvz{(+9)f?S213_~Yve|gm8R8Kfz>Py`So$i1X zpqAJi+eE_=mQR72gb7Kvj*S*d5?O%c9h z@9*29btw-Q6oh45mEcV96l>1jG2Xu)VsS4NUxQWh_|ktvFwQ?0V}*yUm4rd$fZEA$ z&iZPQ*`BWU3sZ`dE6c!>TL~s-5{98tr+|T`R3*(m%pQ zkvIAVvGnC!8BN%eVg8lgaY&Zi*}*n4z1}WBb&|tAo`Bnr;MM>9vk9ynAIA z1@{+(Nr$WY9HHEz&^QmvLYZ~t_D}7P>}N@f-2+Cdk^!*r$C0?3`_yWe@(N1x+loV% z(;Ro3^3>X__8#}w5P)Gunzr7k&wT*YSdzvI*S8K?@;$r=sGDSOLV;oj)U*mcOANWH z+A7e4D1-<@DHDRU(JcaWIDdU6qq%$y#-P~JFcO!Zedtptc-VDfT0EEzHzIOfgLDd~*6{{7k)*=o!>5&nMi~Nb3H(6ZDj@4eb`ZCS z_0B_YXC70$DRu*Wa*4 zeMFCDHtXNrA$_fePU2Oc84DHVF|TFoXb-O__>{b&Q~Rvg`3;gIM7aThLt90p5Ikks z%Y(@gpOG-Q!osiM6Sdz_&afMO^}r|AO>j3NrGwpyjr@&A{mWDz-& zMP^i+P8q2C-BbxNaj(-VVM!_vyw?lO?Eefq8IAG<3#N3eOftLP7I`f3lCsD&Rb^_c zL%HLc)6h!VPM-$M&nmVy>{|0sROaKzi<{GJ&yq{q<82{p$FO@5(D{kjwvWPi0Ws@5>iK|oH)@z8NE^ivjS@=1vXM7uN7E(H!oJfGy7>EYB@l+g zuFaA~Q2eyt=ZDN5DB8U`15@2qKbKy_z)GAT30bqxWH+uJ?J3S@pr?s&d{54gN|x&^ z$}~UFkKUNCd9Sgo(+ua`{lt8wm8ivkZGAd8nT@iJIG@Ols7X7Ydw$L?G1+f!uw4-y z8$oYByjo_ou5CWHi>YxkbCwI)4X!PU#}~NCcHxo~x^l%oah0^L@~I36URLIdhR5w< z&%dnV5q-qMlBqk_!!=>9=ZuKT_{vtATmF4vIBN(Y-_}GMe|;tob@^AGAY(TCFI&K~ zbv%S5ito!ncYq7LvoP%bT!RVuEViYMf9-$#CPk_AF$vulQJZ>5K22-kw;-HqMDx5Ql>QnEJopCiSj%$fsse zS}bIJejl&L;r`0(M6h7AB&d{ozb~AO^pWE`ww3p?Q+^ci(zIOg?Iqy8lu$Ww&*H`g zS>!JaYwi14-U$tX+D!A)M=l_axPdTxjWW#KL-@=>8t{gZeGp~-_7IF~R}!_pf+DER z&)2fH{0V`R#LL>R^7#4$2ASC>46sgEJq%>dBrSXqjG)Ya^+l*OU1gF`ENO~bNqD#)a*d~A0?=JJR597eLT*%E zAb8i-#P0X;X!>?L!kOF0;rDOz;B+U?3vc&@{Q*WTuXX7B%DkgZj=!ts;wXfjup0iX|ZFq*MV?z9sSw|9sNQ@@eo=Q3*!2M;ww&T z!bR_?dNG6}P9tl=h3_DiL3wSSnB^-`)#a<8ai`PT%W+PRt56t%Hv`#z91Bs#_Z|V0HK=HM{@_Y=T8nj%iF-Py4Yxck%A`lIZAusuwTd(ZL zGM_UMT@sZ1yAU|Z_AAFady1-Z^Piy#{ywCkhw!-dWiq#=9=LUKDoIiU(h`6A#Nq=hQi@^ zp26043E)SGGy1u~-FMLjbxQVfC3?pg(#MQNy6qWEJ{f!nKzE^@?wZAVr?}F~wsugz zU3|$=9A;D;_#33@u`sz~ThRI-O(-8;Lq~Lfa2w&As198eQr8=KF~<*l~G zkZcJ|IZa>P<5g1p?f-qW`G+Tkuht7qGoT{$8?cZ_cDG}p#~Wk}CVwZ3*{~~|=AA7P zG}dp{bRy{CoEUft>EO$G#fJ43lQS^HScHz9JfzVRH~--u*8+SxP)I^SA3>ihSRzgG=cYLZ0F=>mMK;(=0H$cH8%v;qjS;z9 zF!iV=#fAuxdN$w!C5p!0IdF=J9oavQmt|ecd?G-t$8M^GX+;5uTnqYj?t}& zU+|W}G>S*QuG9-t{PC$S#&E0KR08}$QOy-$xTQZRKa~Oo7^^3}?a8EFy+;30Yzd6} z9GW?k*%`I5Pnh4FaZ{NhRCU3zHWT`jq4lJ}dCvFG-31rp<5MPZLZSGOlqLj~Q`bsYt%AJkwoxt=jQkaN$t7?%8`_q^PX0dB<6GEG-h3ik8-`$x+?0|W3eZwEj= z6zuaIAY>wIdDnWC2A8<6SzqrVnJ&94BpLQawm83s++m<91C1}pyW9-=0fsrk@9eN> zf-SfM`3m32z}hlh{U(+Gh}~sYggt;<*-*3OCXRukM2ug9gg{>@p@dgv%*$E3b$^X{ zUS?c@{j;|pMixFBapOz&Av{klaH~q<@aW$$gcouzH~SxC)&|`m_4l;O0b!1{<$c)f zERaIH&XvzRjgE8ZB(7&@r<#f?f!XSCa#?vdIa$yFKLfY+>-~bH3k)rf@eK4OnF1D=TdPWyaCP%7hp=`MA}D&Fbl8kcSfp~ zbLN*%jV;!Cj*9WoSTrqdhG$9IstvA=C28cu>-UsEE!Fbbr7rMesUo{q^Kfidk8e?7 zu7QQ_5+~A1-*qRF{@rE2y86uuL`>#Raz#<owoa0!fBRs-j=2))&DhJ-b4w2#AZ;w*_hpb8p;k-yed*Atr6jWku) z)>2(#m)hWL=R|s+CzGE_e|g(Jy!N!p__!=CkMezJ+r5_7N}))5Z!CCCvXF0&>pn8V zwtU=$x8mxTBB8C;&=e%Ut`->o*4X}>uZrWJYIRMfONgf-|M)lY3e)?t(+;k))pag+ z0OK%q3kLDWkWqULje8f%a`%mLxWDXlI=Gqq*|hI6t?yt|sg-SGk6;9Act6goST(Y$ zZFz95)p?$=9uslCDmT8bF!}$7fy2}?NR?0k05}!@)42RUVMfva&m;*ud#C>g)F{z^ z{nP&o+^DV+uid|CoSkQCxLG6M1SOP=ZcH{P4EsPJ$sk)IQa)!&ywHc z%Lk)@g$I)0`SyM`ZOWavBK%XKH~*Xwn6*RpFgq_oGG$V-a;%~GG{#0zi+S*}WKL<0 zIg(0IK`h3MeIVl+TT}=RA5S#8Q%GOgP)TF5+qVj#U`3pD4inL^Hg=YVFU|-LWUzBHILXU^UCL^n=vP&$L%(^g^SeEN4)H zljc0Rt6nBnMf~8|fWD`3aMYJPV@)3rZSAAN$iR8J-z{uV)C&o}?hz!R#nFFAQoC_C zd`dn|_halWyypK8WQpodUQ;n306+`Se;zal00@8-f+3y%IB@^@G5yzP!o<|U#n9GS z-`vpI!qidUO5fPU+0?}8e?&6Tx{BLkL;CIY16|1X&tAXWyt4HX!(oY&t~-ChB7oE~ zXh5*YXw6uuB0AaOeTwVDR6@c}yxwW;_u=%DP&>Bgcs%NUa~I?Db${JDjvS*wwthGY zHBGeRq)t0_5j^6}njU;?GS9r^m0#cLcBgo^rRQ((>L9rxk9d?x_12x}MiDs`xy9tv z(fyW*WcnR6@^EdXeN48lqDl40g&15SqBq2ub#z_zR%I#>a?d0Dnu9=no5|!$r`WYB zb!aI#$TY{!y%A}rd=E&0{cXz&p}~~JmKv$Q@P6c*sL0lforufmPLG)|&z%VQ`D=6v zqsSDzVW_TBa}J6MuUS3eTW8^oJ?fzaAf-s!p~zU$MzeXykEAd=g%v5BSR7d;ja$>U zH1(;OtTvH83=VK>J=AJ4WB{g-{aYVymSrIhxX;uXgW`+WxDigmL3}qwt$_(D|0XHR zp*F*6qn~EbCl2ZT`nvcffSv?a0VvdA59i$?IB2+-QexMvHRf%v1op;pgB`<1nQQ!DOrV~$-mWovb(2vco z5Tk7{$_8Vsz2#6d3b9J5*2+rJxq6NHj!w&-#IaWoy8#saZ)fkc;Bu3>GlGqWl=|Vu zxpaa4u~Gp`@aP05Og+L433t@u!2K**n(Vfn3=AK0)b{DcA`~-Kpd3GK79_| z*pCtAXrU1l20Zkcu>U!`+LitnOqPeD@G_c&zif-}Rzso7|d;?dz$$y60< zXQ2_HJZFn^ZN?(*fwLpsr2|6JO{PgPXOBAWcd#Uc79bCxb=JW6&*c3w#nub%-qj83 z?b_Cjd%oxahrt-6mnJJ+*-#FmK*+c$HjR}z0~;V_zC^qAlR(qN0dQK107oEyifrg{ z-n?urE&sAuk8ixpw#GFI=m&m`DQG?)LL?vvpO6A`(Y%RMvc;o@B?O ze=!+MmTI#a#c^Va%4YfjeF1!ZJN1wddppe?I*a#kF530Sm)SM;Li-T#s=wX( zWWCZ}s5-N=OzE!0m6Hiz{?UjsB0xIu9%3Fo>W`hMiKE{HGyG|bcL)6(Mx z4rjaJf!c_Tx~~iZo-IV!UMYhc_Lj+E^Jc`zMSW5N{347*c?SXG@Hb_375zM_Csq<8Yq2d}{;SCAi2*VYbmD zr=TCTbZo-a?SAP#UnLMSf>?S3S`|V^=H*m___e+sj~|oijyeMMAaQaz2nfK$D_F;( z%J+W&S9KyGDxpOiGiTA*szw;)bP2_6WOq(rRKqmUVZF01QZ#kes=}(odf!Lm5ZQAn zE&Bk_Bz{`teavV<4oh4qwDq7OM$i;7BSFxcJ}*wP5+w1xo#1h-q!~k%M(6)*4ZG9~ ztRsTK2$fiv7-s|y!ZF$lOLr;M_EBFZj;M^2*$;iqWO@kBH8faRy*OCyu=hAD2#G4? z8dz2^I25=yG+L4Os>aaIkWtPFNVw`J2WYni?Yg^T<10G%cQ`f!zFt)J4T zE>qN~`N5SI{iAiMR15z_$S)s|t;_XHU$s|Yr#3xD*)W6UH|FaV5QF(q!F$hBn_Jal z9k`^U2LydKUmli>#5UpuOcye!7N%m#Q@WL|&eT=f?m9=Xgi!rb#Z?_OWX=HVN2r0A zUYB=bK;Z5Z1GoJo!cA|P9!vwaf^44Zk3hzkQkzdiOlVvd0skTz9bU;0dqc$1C(ma` zOHQS05k6Rj$xx78)K6QjsJKFQjo5fDdSBmqiTSSC9e;F&aYK}Sqm$MoQ;(mJ&93Gx z)C<}S;1V4#de+QVGQU?Z(a%QdH5a|-09wQ!PEjCzU@V9jSAG2Fyh!(qYG6PyP_2`_ zbugE)QDi*&`RiVzGc%CX|sXIk+U+2VQq8WzNf8PrLk z*+IO8i6A%ub;+3FmI2}+5E%SIeiYqCGF*>0yH`sz@g>AJ#2?U2tg%4exVtA;aYZ}c zX{~!75iPb9RDd)*d}JR0mghsR>c`RV6dt+c>(=B{QphtYiZYb$|GWBK?K*mkOQP!!txt%MqY`XD@e@)} zNVk|E7g}xve-ACnV$euuT*JxFjUqv=qW=JWa;V2d@31R5PX()u%FRJ1)nxYgU^{$- zyMSC=958H4k&I=;WYcT(sC}fthECNYjn{Qiw)nXsT%xgO7Mknteu1%X*FgAsW07z( zMfDN4_`$+`_vGAEejn|ap_?VG3~;MDg9A$|(-FIjT8l=G3beFrjFpXk$xw_h0sUpb z+H~65_bWA6mdr)xZLPK5)*Um{W@%-2c0X*zXBMSq>{L_ZcDR~1STTnSkV=NaYEkOs zgA4moXyPD)WO?d*`Q16PbAw+im4BCV0~uAVI{~d$ahf(jyNlpE+gPhhmOdfkE^j8Q zQ|Tcs064-#|GY!_%f0EyXQ^i9&f9#Ac!tb@uAG91mvHU~RhMJpJJlH9jTTJRw>+kC zW|Bv)fYl04jn}r(_HP!2r+y4k*(%r|{x4f71lu|xG~^I8Xb=9w`b7b>@JsxI6{t82Pb*OHxdVD$6+a3CgG?&$2A8d zvpZGp<=c|mIt7Dg9M^PU_Y8A+O^Ftii@ZEsHWG!cq!hEBNtE#K?`Kf?mk^f)f1j6ZDf)r+aR>nJb!bt)Hv!1G!;qqZ0ghP({TLfD9gN3KbybA1Pv}Y->uNG+{`UJ2oAm|QLcs^qd2(xbq zZ}QF(#)}z-7U)Eq37YGXaLnB^y3b8_@x&P9#Jrk2hs)uNZHGVOL%nH!HDbR-=B%7{4sCwOyRuORVfcuw}p`y7rL`>f* zq%+@I#t*1m0fxL@!p5QAiB+@c%Ecq&4#%*rQm+dVh)0WKUatUzrZcU^9{77O!7TJ) zNP%RUE_)el8TSgWcUVB5+`x^VZ1$wv%jK);FkU)-gDv3$oqBpDsX~Q#T5$PiW3T*{ zvbQG)Tq>(7Zmau&)V{2)7r?6C^Gcd~{p%n3-b^5)^UJaP<<61zzp>z@*1{QOJ4+=C zQD5!}rShl~dm7-jL;bhH4R3nw)mP6@ zE)V#YV5c?~^WnTN#dvLpvCERJE?BjBk}H2wQLpJ1xms+)cK4pPp+SDUzu~LCsP{io zqIutt97*}@2iABF9t=yMjD`w6-NBeVQyOD>6*V?q>0huheYCc~^}q914a`;DUnZ`P zLXGNQXkZ^;OdOsh=ag0;M~dD&r!2kld&hXpLIiCTUkWwm%g0!U!16zVDMW}oB#LfQ z#c8;gXQ-X=hN?^&shY?--E*yipYb4F@+dFZYicXLWaSf4ng=}A5VppGYji+uX+|PC%=q6H>yeeD8S^Nxl2rlgLGd3jq>FWR2i&Pgqi{lv3v1 z{~rEm_@215<1_#RkeJvb=h|e}K>~GtadENrWI7sleN*-KciRs!goo1o_dT23v^m_} zcM{KBu>JUZVZYpsB0pD@(J#@A-NU(qX7B#$A^Cz|Go8#7Z>ExI23oDOTF0@~_x@jF z;p74wFSSQBw|2cyYLH3xC{VPDX2uXhZt<$3uaXHavhe;4z}E>4<&9nauR*RoWqjxg z55{z%2_9+Mh(WHnRW9|H{4TobY*0!dN$}-~+uR`wMHo3=5&al(9IyOjQp4D;xFjcH z-yJ1QErpM7&I!<*Drzu3xMAm{KdG;(*mI!TfcP!_oj4?o>|-uQL+lZ1%}}VTCk)?@ zH6u?4-upoHT{U>sz7`L}TFG1DBLjfas^Imr4lwJ_9eRCx34XFEkp>w#BD2&_#Jd;4 z68_KE^~zf^ZX8D3I~0h1Jq!kW)Y4tin7{}Adv=Tl(u^5n+Fx6#+lEMaHF2O=R6~SE z!8rT(y&m`Xfh`)^LMEX19I4^I z5;R6sLWlY{?k8cE2k452NrkrH+xE?YxJeQKeG=%T;Va~$IZRq+=-Y|lk2M17y$Ex3 zgD2}gPsOSmAdZ^9Zr5Wh2W236D%5>4{DabQy~L(>DG=G>)bW+sjtVqE(a@}-85;C9 zVo>ueHB>*4L3Khi4u0Rcrh@{8OWT#~y>mrsuI*jb&@*4kC4f%E;}JgQRs2 z+HhKV}VT#ASD%N09{c`vOhVJ#R+uZBryRI5ar3d-;$CtC5=3vGK-m0QIjqb&Hz8(R$oZ|d|h<8OWSo5sodpyYwCzF6H8 zv^zu0H;&02q4Lf^Hl4qwb@~atz!3>T2=OBnLLD&POkrU1dNS)ZTQmeD01YMatSK5P z(uEz;xZm115fxqrxDK(!&mBdNy(1qXC}&j`J%M@rV)TT^Wy)6u1_riF4q`~SqXHC)+E$`Z4^fmO1!8q|Ao8Q2h0p|8(uPiz zM+9f}n1#^3Ls~`hy&#PtQiK(=#TO5bRvN~B-d}&6z)8k60X8YnQf!e3M>zZ~d?)BU zIuA{KhXPGv=Th#}{CS7gy~G^Gk$PwG6C;}p_~twS4?CV9mfYq-$ptVdcD9@50e(I3 z{;ijs`IvyvfKc`W;E$zvAdL8fqE$QNrs9A@(t|~g1cvU3ylaBxJ+h0s7`v44;nu8F{4O2_xj1t{F>kph`w6I!KIv z3pn6kQOEATE4ZwJjb8@9()7cw)}jEo`pXvHq65(Mb8ERSJPj>857=6H$OJTE3d9&} zL$N`m5#UbzmC#6_Uy@&+Z5_}IvXJ^bUH%(?zI?n>wR4q>%Sx#dtNzX(RLKDlVvIe@ zh6s{mU>85;O`L1LLLzR}&kRmH+M?^`{l`0p@hU7>w-%3Fz$|-8S;-K$u_#jMf|qs@$?0soh87d^i?ZKu);Zf8jU< z`nFdzz+mg|!U#N{h^dX)IPG0Yy;@{iYep6#Jt7=)XaC$v>RAj&1*ma4MEWoCx~M%e zQk7UAR=}o*P;p&G5RgmA5@aJyFQYEp#a6yr96QbzGe@6)3aUl5Z!G?q6xq%G0@0mO zngbWKQC=_b+C2FHwn8axrPSxMZwr1af79>USwB2~szzGO5xVu$Cv*@WRz zd%FUNKX`P!m^jp&F>=}ob*~WzzBGDcHKDS1jCgctSv)IX(!f}uxu~&A^W=wdHfqiBav?QQ{2n&2LAhApTWwWJXD>4-$p`oVf1su%qJ(Cc{* ziEcA10H>u(2TnSbLLGuM+ORMFPXVyN4rDn7p}sCRtMQbaM1uk39~(d>?<^+Fpk1ho7quOGm$N4?9aso}5+tq4!tuMUm1|+7BQs zC@;Yvd569%c}c(S8UpF#8>XnoPEPy)^9&Nc%8edSqw!VwRi$)K`d>H2)otXjfL_0WdZMpyAnI z$|qP&=^V1mi_{F93M(Y0&=Q#jTmorq3OM~>i75FoUjigh!xW(vJN-aCMDPV1&`@%= zi!x=b$C^D}#+t%kJg-{*96Wh!!%>1_&`dH*-h3|A8UpY8VC`1lpi&yRakb>BVg&p? z$7RA?T5j|TapsDYGC&2TMRs!5!O4z=>WX4u*IgpXi?|w94)WU6ln)9EMBFG!te}2# zL?g5{g|vDWuu4&n2vz^)Mvm-d|H=Y}gp>9ACcTM0=Q@KFAm_UvsK=u>@*9cL8wJBuqx5Hw77>>Lev@MzElJR9_g>D5OAA^O~A#hyw0ZL$jRa zfJr-*)4ubG8TMsua%Ry&WW%CqT0bN2Yv}A;HU8Tv+BL0a#aIP8a^|UnJ?UO!kS4(E zhB#(V7GzvVl0FeCDsI7wgg^oRi1A*T^;{;h2*>%%PXKBYkN5wxvt#<3N;4+%kmgf)_Oqv!KzE8+xsS9Va*Ddoc3bq`R?a64q;Niua zxFjyJuSUXpubn`V!WGr895)*R&(#S^!BJ5Ff^SP)7k~-b3;OS+gSyG%r=yEE_P5^Z z%OtA1)_*}=d&Bzn3BUD@Q+|qgx<1YzBr^ej*3b!{`_&50n`4j7yF3)TzMl%q*5Dc^ zmfN#&ROyJu{~A;X77Bg;m5O4B^nEtNcjAT?G7fgV>UU)&ihG*S5=VuL$~ zjzd#^%8Yk{LS*AJjz(&uNHyGwcSo<9>5ppbCEuj7 z_4t!1rBo^$l&)cDO`@(&__UnL?sgQk!M$1j7BH@=41QGORu-s!rDp{G&f(1{4c5H& z-lWctz%nT3Rrb$X`r{R~uB8O%Bu3SpVVo5@xS&keQ31!fYq4L{opG%Ok`7<48JKI?2&3;S+EKelRagc}xc z55^Jv-10Zs8mcgzHu4ij)IJ^MC{XlfBKC}jD^+7*&EFti$|O#sDp+N0Ad5cQe)(EE z+*L^%-3c!C;_$P3TV=z9P5Ue?X>%XbLpInCwGvcg$~jT1#_h4iSeI;Njbi z{Q~3e`_@E2#GKbXh22=`=)Y9r4BQeM}eWl8w4 z7_1Udq1v6uWog9V-vJz?$40(PA23;I5@g_Cf^6=$Ne_s zPPK=<}&*F{sb6PR$Z>%y)~9Dj}q6OWcTu5_Dhyn5;7Otq8Hr*M*vKJh@;L^U;c>(;lD!vH4B1AQwr2GF+4cjVJ*`%b`M%X-*C+n3XgcTe| zJdhd9t1cwrDDsza_vr-WAZ&L2fjC`5`t^_jl_iZI6I!?^M$V-j_DPmLH&w8pz3tyjGe*2(Hj*}ABi3}VFkSdH)^pDY_MdMJj1>iofU1i2fk zWP9Dhm)9cs77=_{igPUFvgT<5Mpm(ImyxbNe zU~sop#yH)CcN0W9I$y#uV-(lj=0{+vZyxjB1E!`@@$y29c#X?Lp)CFmO+(adt$5bs z92MOg_)Ax0KQd*?wZ?h3%DmVYUzEFgL8NqZXm$eJ?4TvdhPnmd!WK%7hqbn3eXxvv zn9Y0kJODpaMVMt$vND()|NOG-ZF{Ot6RSezyWj zW7qS99xY*xkp!7lfTu$BD6t2qgOlw`uZJgVN54`Eit3w_RVjrUS1@(CyaQ9z;#rbe zYw#(5?@W>*$^$=r1>6kx`VY%_7K?uMjKBtf{=_NJN%CRf2od6ivwu&ybUiDC+XouZ z=3}6a8FUnjT9|78`EUXX4pD5!gu$Sus7GvCJw4qQ?%LWiDB`QSy<#B{8RxA4{Z-O* zYb%3hwemf|PS}EY&WPlh-35Z#fTYMFR%Ix|U3s{|XgU;ec!E zQI-U50xLF`7OfPOY5OpsyLT+$Vm1|$g2 zM|I?bJslB`c?>&U^LXHCa|Wm+A+Bf%tG`8n{uhte{VWz=&r^T)7jh=+V7av6_}^EV zJ5J~xTaaIsua0oJZEjLCf1NH{-uaKrw(gE~c5+7EgvC(b{9ML=mxplI9eYjL;RyIU z(AT)GzUpj8dkS`AdCIU)EeWBfqo`tI)T z*jikf9%n^?RTCzv{G}as8sLXO9oHWk0vv?XS9WO^d~GIuY?W(qr#1Zf(|orHtCVf4 ztey~nd9?B^h{lbtuD_{Wws{ZMOo7vE-BYrhz0zW4-hk*CQKu{vi9k-5%>G#}7mQ7R zQf#hz4^WA^6_Dy-{x$#u5^K991bB6tbn=dV8n^AxW5Ry$;dA0(xTVE#-b8WmC^VwaW*6^$ZCAYZCd>I=U({hs?BPLV?+g0OY0CR=0xWeGS zm)Sz=?8bvoKqvJ3i2o}+=0vcWxS%eEB#e0e%M~NVV$8;k*D?o;A8AiJ4JPQ26znrO ze5xUFQFIdSRT9oD&u-hb*HtIq{%F*{t6;EY`N_?qRgt#2W8D*Ie{lJ>s!`st86){Z zd(*UR9ANFpd#i@mfUr}ZUFgYqM7O z&GpivZ(gxX->lF6jZ$h&{~u1|Uo-d!=zocx zGxh(@iKP6mok)8}Q+q>4Q+*p#b3+?D)Biy=u|-|mZi@}Ycdc&!e|7eiVR3EA+Kqc~ z4+KbXch>;H-QC^YJ-9>A;O_1cG)Uv_?(QD)aqi5yUveh7GiTbK{?X6-)?T$%)$YA( zy`^Ch&7fX|SHT$GI5$JmJLWPlQ;LNI4yrHDVtyzWE4K9Q0WY>Wn!$6w`vR$-YxeYY zI-GV^l=JEH{_+^RDEt1mc7mczdSVMdui#4M)%6XiIld|^?H5> z87Ap@e&YoK2H|gG2&;(p?stp=?Gijm?AIMR(%glG5{OcEKJMB2CNgPeGwYJiXR_(< z;St_{5SZDUgdXlk8LciEyL_giW4oqZHxa^IBLEfUkB?z{lY>d2v}eY3m_rNw(r$vEL7aVj>(4p?Ybj0%)9&O-U&4>N^~FZ5^m=#L7l^hCSlz5RlE&Kdmrg zthg$6ND-fE6%ipDcg2^y2=^8Ixb|}4%6c-em)8rTr=7L7rb!@R>qF5Uv&z(eV;vEb+(F4X#Uc+m<3sQ z3;8W3&p>YBrKr%sH$N5;XC!tV-0Sfh-}vnKBuQ$8=4<~BO4Vw%9YLEly17w`+3_ej zB*xD)A*-iKeV?AeFY97UYc*1;81xG8d{}D66kFv7CS)LXOEn$%EF@}+4z0>$n;ECI zn1iq#7+8*jJU!6@h%P%5t9L)yiiEwB^Ne4lhx$??@;O8KMQ%3$W8w^Jo$@nk5gxCt zwmrD!)tJet5(mB`w5?x*8WxX-xH0BB_4cH%e+-Gx1~85WC%|(8_Y>I^Jb9f z;ATd}wzf}uYIZxUb6VI!wp4ElnG`-@ zJPi8bH6)#>URi#&rqZG(dYvPhJhyz>C&LxCoFJ9q?2yzOSu0}U=3JnWli`Gi`Z_ia zlp)?#iAC{a%^RjDXSzNXiyWY8Tn)cDhgoM{)wII8mvl-CHR)iw2op7TMWKZVI!?5(3&~*d3r7-_e3EZ zLKB*!S0!)}hd7BKTnJiY;B89?uOQpE*gV#E!3L*#d3VCY3hJH) zc8cY|-Ve+*rKslW2ANeV43{9AG1CeQc20y2Pu+6k!1tNO<1QuH`SGkv@|R-)lRPsx zB5zB*p0W?lf&Gbo)eoj6)SYoxA(@e7a5=ei&?((L{d)UibEcXW(^Mt+umcmpZy=-4 zaTqI8$kykMzC*E0mb)c3WKOr?`c#}T%1T>BqTZ2eR{Io7bGy7#3ijtvWvE+x+%Y6u zow|hZyooq-qC6*{ix+)ka8xTHR!=logV~komXlQ>P%@9mJNzd1ZIrvxW4=j5keed; zFzsXKN_j9I&*u8dP91rfolF^&mdFPzNPeJYVy;jnvC9#L$%y3h4#UZdY6+*^_4LZy z#ca-_Vio*K%kkTJ-ooxqMw?u|6?yrw*T~yqT+|h|@5z)e$Sv-ujhbvEQ^QuG4cl6E z5-#@(7f%1t8!U_$RW<|y06=}o0l#mJ&hq>CZ~qeg?G3ec9UQE*^{p*m(t*AG-?<&d zZEn9P9J_i=i9|Y2L>^5-t~PEaIqSCF$Wk}|4jtH3GTAJH1S6*63qbI3)3Eyf=r{^@ z3C9}lt^T1}! zFwJ;09j`3v?Thy2n>Ic4oiPoSggLDQf#Yq8s8Awc8^qIhl2E17?7gl>ONP05c1W3$ zYg3fb))Fp_G9nX@Os%h%!A$l{=NN}XDek%)w~@k|U_Kc=dCth&rtxTC7H-DIF~;Jw z&=`kyX&?C4dVz`wi87?agGe!CD#^VZX`fUR&|3T9Vfa@+v6irLZ_;WNuyS&?jx|ek zWskMLTC?8z2;miV$4M}L^tmgwYrLAmxo2bV4foT{TCPm*O0?(R2095LY{A(Q4G>P~1kulkJGMz3Dj`a- zTXL3O9nX-LMC%_TJrUa>=s;bg%S_Ir4sD$1J7gi!VP*_9WbLuF1Uo^$`*sF_tj@zk zdOu(9ot|0STs_~pTOWLD-|T$f|Du>R7KyK{nd6DHVlDUg)BPs^x6b(x!LhQE}odGFQ<|ELA z-|+#^8t4Z=cqZ-GP;6MN+o)DKj;QG4PlhN68-KoTL{-{zFPJc<^wRo{jw=yrT2JGv z2j0sjxxob6r8iGlamJ3{H8u3}wQ~k@Ojbm@9 zW~&0aUU>5G+LNJf6p}uzrfkR$XhdEA31fxd+=hrVxGbc3Wt*;<;h=S zFx~@*m6FsM4<`5;AyC#^pmS7-KlY`LG;YA<`zUQ--1^JitBl)Q63$$3>|wy;8*ucd zFfvt`~l<8W9OvOnXtU!Jq*9eaVcvl<^Ve z_fm|ISQAxn9XUwRB%)fd>u>PzOFn3R&Zxka`FoDGjzkul(TqbG1!$FE3|?T(wb*fO zVHv0*VYJ)Twqnv-@G-E~Y{(Vi&s-@L#>g_XY9Uf4(_>dJE!w#$M>#V=+GG(vTU3!{ zPk=z5>?eXrN)1BxP^1GWLzk4svn=Am zg?!3nF>ZgPT)Jd%W#n=;Jb)m%j)|#ihxRDL=Z6Wk6pg7bZSZPWXky3&{E|U?7bh}Q zSMYgWkQ4~mf>zctpMy$*Rb+I|Bbq7X$BC_bU*VwM4Zf87j z3cr~9AGpBi7q!a<%6dBLev`#x+RWooQuL)UsdD7tJYCa~IDeDjWV8@d2Zf7}fp8H2 zyg%qC-Y48&&Z8h5`)Dds{w~7-dgH$C6yog72oyB^yGM*%mexZBA~m8*EMnm0Yf_#% zvJpOznCD`Gd3HkD8PisPm?MBd%F=zam}sP_mZ<$nDZQ)zXyEv@*2$uqyzdOiQYLI) zkufW`6T}e7BN#%;8WR;#RIzWmDfr1C+imGg^k7f!<^@4UF#SU&m3N~E^dmC8i>hd- zOidsXKSCEmny%25OBR4zVJb{P)~6xMF{NzMHEA`IomxM*C8b#SG}yTXJ`ouUVYPY` zLp<+Ll3K+XJrs#q;HxT0up}_4w(~)wSpq9%DS3Lj7;W^qpP)iPBCk?S+k8a$z((n7 zV!!8t#DL14q7apUXEN8FW;mKx46e6fUiA{4foEGxX{V@p>bu)3v@?m|zX)Fb7e%7+48pKTVHN&Qp0%J*0^Z4@_m^cRUuH zDTTDzO|fSr>WO&RykEi|@mT>Mk65X4YClV%iBQl`?%BO2;jhxF{~&0w&nANm`8|Jb zi%lPdv3tQ7w?BNLw!)(~C-${4yHu~sj;;<|Sw=w!eec;Oj8pk%Ftr8ucEEY$&H%(Q*aL_hc`wo)=Cl+;d#WaK>_uy)tE z7t(PS^wpMr>Y9=(WXa>j4FN)V1dglHN27HF{>C`ypzKi`+@Ad|S&Tvf;040YdPi5P zSR(cIG`$u?SljuvAR)s1F^FXIi!cS6811OiAqak2yYzWS1$117a$31Vd=SMx6OHM2 z<+EYIO73K?`SJfljN=y$LhIjaLoIw3=Nq$dJn9dkl_` zuIE(TtSPXPUd#dou%zf2r<8U}89aG>1hfYgFgQg-bW6Il*me-p7bg3GcgJb+dIhg! zNnXL97x~74>)0oMHOBko+$L3Lv|xByY*MypiACqpPf}dXXYv(Ya?#`K$4Cm{o;Lh* zGXhCGB3@{6*69$T%l=U+;KhAy_kfD>(lULqEWY`w8F0)lJjiJq`O}Nilz&iQzO|6B z0Eb0(_bz|NdHCq`GYm2kNnt8KivN~4qYvaL`GYzg9k`WX5Bkfswxg+#Rz*iakJCH^ zcmWW*@H}~esro~H4;q^NYL{*^QwL(){GNEt4n(~9D?R_suUz6g8_|i9E8F(~GP-hO zhdlJ8Ei4wcR2%p3bj69u4ajG$0HF{B>8MLb2VAYb9qVZX4m-riFGir0+U^qs$4cCm zK6y#~U!^G|);hGzJG%0X-^u5BEW+h&=r4!rh&OL7@aOFg@-@y<`)r-W3+%`6Qb_iX zP-$L?Duzd{ky~qOYpKv-I;H0}2|uosD1)aASc87jRJBCgJ?td{v5cf#pA@p2T;KMX z@yiEx24Zl|%UhdNea$P&M6GB1=4&n29d(?rp%#}{)5Jj8ool)u%oK^6EOf-mzVmtk zEn0eC??ud@_PYbaQ|4Y?>s^xij3@ftZZQfF9pn>c zHGR`P4<|||2)^EBLYH^3t;)cLghjotNXDQu&D@PRR)T9c+$EYgXsB?zi8ee|DGjNL z#KXaVqaC}Gr@w{zX{X&}vI-9VS^|sN=?h+ChthT3+uRz^8hx-q=ByOBI@ELQE}ENp z+Oa`sZ(@-}=e5N3WKf?QYh^IL4h~@xC;ky0i3}&f5mM^8miRL9hUm)zIffCUL0A*q~ zSe}Qr(XgF#)2c@^Xx2cG0#N>zUVVJOkWEn_VRl7a&-6;OwgI*{FI_i`F?=)rbpX z`*RTmXa$0gts6N@gMkLe>#t!jZeDRtA!tgK=0?dM=U+Y?6X z8-%{E!D|&2fKJRP@@;%Vq_@FSYLSR}s?`#9<%-i#ho9YX$s%0$gn)L*9G@&j z8!mZt(}jP!!iC+xb$tblrC8Ax+*9`bPTtI2u5SWZW$^JgE}krahJFa8_V54fEAKhz3SL$9Iao}QmT)8Tmvtqn(|Ld)^ zK+(YdbM7m4A(`)FSMkEQND80`+G0nICFIGF>@sUTN%(IsJM1J-^oVp};p*ExAAy~h zU)u1Fb25CX;B@u{YZ&x)O{uc0HsPC(%+9S$AjyU(Gh1tYtPUg<7Z*O=uD*G)7fD@P zE<~GjBY#@>$f#^v15fuo+PKn9OE`}`W)w~4P&j7-u_E!|5o6L}6ZY+JhU_<3MzWy{`-5C=JW z^6{=i&P*#`)3_jZL-P1li^G*v#F8S3ydYHuM&dU~QvfIf#b6{=PUa|upN~^FvKm;c z#{@Zh%`KMbQawcb>Tw{pb#!nydJL9luUa(a4(>cAn7=tF3a8KA-EFXFW)%fVv=L0@ zpYCs#czGj9+St&B^1%t4p^)&%&uI2pWp9@6a>Bn50FH!8rWpeiY<{g5q?lnkeRD3{QYS`70W`*o`c+*Jf%H zuKlKTdrGJmFY;WTLfLV}Xbrn1Q|Jk$=!l<$cb{wKkZQo>pU~~F*W)M)%Vx}yA9u+6 z`oN4}qcuoqXf)ZD$SonGrVCy6E=6*)WKPIEd*2Miqa2=>B4!8+K!4t^bwDyYfFye- zON5+U{2At5r60@CR9KiiACugG4s*htBJFBo2i8UjL7!A81sS)c7=T&Baf3Vu+#6q0 z7$zo5ck(lQ#5{&o48WH|#{y%`THAd_cstg&@>y(*E_cY&$ovwjgw|p%H4;_padZUT z4M9NTL9M=(JpV3S&AIC$EHWaaH?N*fwZ5} z*frQqn$*G<%$nnd8PY^Vux%QW@_oSXW(Ns*xPth~Y~*d26#0(a1@&l;)_LVl!(7f`*d;ghB#XLqPS!|9gvbnuxU|MgMXaY&`dnJFaoA*!q;8zf1p=Lo<7?ur}6Ecau0N*)qj z9~>Se#}qezM!CLLC3mFEQ(SlxUbSi5m*)g6!Ih|B__}}&?Z~FNG&1>6OdJ!Xbc}E) zynW|cPI%F1f;NDUCC7{e-2Sa1w|SsZX;Kj`b;4(Lb1!h3%02dnM)mP&*g)Cn_ZE$E zV&ww@mafCrhogh{?}bEC;K<(FD|mn`BnHlD@TU>a7tQr6>Z&0C^Y^~BA#7ZcUW44j znpyf^l!!5#hb%(wG=OM1-mrAuTP7Sn*_KO6n9l1|ADDzsC%Uz*2TaqcCg+T(>3wJw z9AG(PJONWPhJ_U4fJJk=ttWUB_)Xp5&Kw0&;(3>9GFI7oprB53)8D@Dgq{v@Rw$L9 zNV;Ui^F!}y87z*ob=sMIve^gm3#&YtrI3ZY8n}`%js|KLWai%6oV;v=E&Gto=gK69 z?KbXoccz4FgR5Yx`w2y2s+J4q=b8Nr8srm7{iNaG5!fv|&G;459)0P|=Y52bgQbkC ziZM96u-G;RcCKsRE{->Jke6BSsWSX1UtR;3I-Q?F-B#7M==QDhHcoiQWIak@iWh~^ zmgT8ZeHzPDSU&W;7pgx7k0U4igwT)(_qbj!s~B|)doi3-+x}U+XZwtMmUt#XYjwZe zEH4eAqR5ElEI-D`(yyi-spL|18AJTxsc-30BL|UHsaJZI0l4Lb;!AtkWv~>0Ta9@T zdP2@{Bo(8Zv8CBrE>f%100O+5&z?|O+YnYab;*FRE7bpReRHo?E_Z4<;h9RNwMz25ZOgx`E4?~M>JEd+d(z|I^@bg z3M$(nHRmV7&J^apMYozrS?TLcCEAfH6GWYpw&Rec-k&(_R}v-+smxw z4mQJa?Tg7(ySv>D=iFHzC0up4ULAxT-1SsrC%$kt3JI}6;m54{9+BE zeg(s1j!Rw6*;JMsI+~wf?5C)n_|{4DLzyKIXak&q+})Ro=yNC#Zh@n&mAz+c?Z!iD zzz=eAR^BU�Q;T)RtXuO&18WR;;mAR!duT?^8r6oor|x z5~yxyk1wI(Sxlus8_hpn60Yh^h2O#jh2@&Q#vf&dEVEhMdIdJGDUq|T`K{}n%|b7p zG$yIV?gscP^Ojdm!?T1;TJ$ozUdDTY(n~_Y-flfwI?~A}ZSQIAk!J@i~m2MI$cez-AqWJH2 z#HZhOdsO6dIJM1aATl>#)Fim$$zwi@RkGOxa5=O+Q^pGUEy(LQDW;=ja`A1oED*(1 zx0#?cdTDL;n~MoM9;OO;_cxTQF1t@_Nok^tx^XJDdbYq$;FZZnTpns#QCFKKBA29= z4fhRZV>L1%VUwycONEiqE)$LSt2ygAb;}m#L!ZaSdKfc*C#%^J8m3%k98~LZ2tz?J zR@Y<=Q9E|xUsmEDX5l}^@HKbO6fQq~*BYoA+G5+s``zuz`Hq#%&HL%w8WjcSLMEcf zIF0ovIA}>3?Ar;0W4~4_daY9B5~`rQI4j%9gek|n4eWFuuD)ZzRZgJTsWauy_n3s! zY)p>~N<(*8_r;EM8dbRJg?9z|*4ZZ4G@&%Dt0a&o?$!a4Lvk_c{rP3DEw46oE_`p0 z#?zi@-df*`6lYT>EC;%&Eu-10cu0ul2Tqe(->Q`Hs>KzXOzt+~>$ zEO?kP9qZ)7yN}GdBiJOlTih$Vdl-i}3ru$%pmyw0nx0XI8TqSkYNiUz8a-t%wU38K zqde?B&q->f_^s8XG0jJ`W?W5?->xG|XV$R-tlS4KKMvQ;NFv}&R%_6nh$w1NVsuCEqa(ET!Xq^s@sjTY-`3(%G zpRoc5It5t z)}!Egx6fx{-E#gAjWf&CVR46my>1u?TE*3D0`L704&fl*zM zD84%jQM8BbEU-1}E$F=xi!M*q+c&98;Xbz=$#l?DSe8z1UH6K8A8T=8gbWm0yUH;s zHqG8Ho@FeO(akkY7N6QQtk}>Hs6a`stkmpZCOLnM~D*4uAN7KE^|ee{LCc z|IHmj#xV2ATg1xec=WB^;>Tvre-w$(zpkju!UX{I2L7*`gt&jXN%#)|4t;AYBU9u5 zpe&?J#nOJ03F$dqJ4-x>P!|R4L*V`v3A{}rXm=f2xkIEbmSWwu;8a!ky8?v1Z_ZA{ z_d-?6oJp5K`Sr?=7e^;G)l&T>K+K9=TA?BP;a%Q$T&zUs8q@=<5Md1-&rA;%4jZ{6 zf?qR$8<{;u=H++(s<96VJq)a3O9{dhlHu=+fCJClKzli|d(x<36E6Jk8VoYuSs4X( z8VS&GW5tBDY_wW8o+349iI8w3TVczWP&~9+rB&; z+ev14fAMQnp?_&D~f@>%#3yKP>A=zsK7EE{%dtuiP7t?Ov z-nMCtT^|CkX_c=a5_%}BiQ?$l7}JsO*zx=;MY_#Y)!M#OtDwbvgwgCie7*XlkQDY! zwaSKVl*P+&wHeu^A^k{vXd+hk->1o%f;+Lk#6P{BZ?QNuRxeOi(S%L@W8X9rIF* z14Ee7j}~gS0eNDeV@i7amDD1owBV>tl>F;>yk`OZOm)EbQJu;koHDyHmitmH3BU{B#tVroe#%7PjAn>xpjtiafwDcHd?`Y{ z7Bb0hOqj=iJuG+R-CEi|LkxFPD{E-M$3c#c^Cb|{{-YEVt0BGrwT#JY ziG44#_NH@7X%&l}^bcE@BPdgfn|y4GHW2N>udie%>~zT>oKFpti4}|91KD-xcwZG_ z2qcjqyr*H!gwj#0W*WnBRIDX@X;|Q>;>#e1o%T}u#CQQCkBh*Yv9!GD-ClG)K@8wX ze%b7(KQ4^IBvk8{TF+P&kV}4n*c8nU<=36@(pk-=a8oY9a;pV%#~skK*?gp5<2z02 zLo8E6Z4>IbU;J?$AArXTI2lry&VfM*m%+)XG?tXOEjPV$97qndzygnWa4 zt(l94I@UBR5^0{7o8s{9OswzA&R7Clzi;6J7-JHi^@6kX)S|HgSLQ1S>V4ge{`gkx zn|sSCMU<{H71wf5M1R9!Cz7~S$(qyJ(oOsbfyPu!R5AN{|KVFE6jzF(A^nw)R$KuJ zj6^z+p1A}q{#qUCK^WAHX;*t?7#)HBW&dU&oJC#nGX zig(1Y)$_Y8?89-?G`&K`daN%yS`xe1V@t|be?bCceFB~BanXv)7~`^gadBI{H-T5C zbio16Q$TU~eqj(H7^||USaX68M~c-A4wJH<;hkoCvm9}TBgXQvX?WjT>LKM|QDqbg zUP@F6V>!^eQm^khlr_E)M*Pfk>YmTrwn=TS)=&rLBVNiUOASt+oKu~2&P6|)rXHs= z`fXQzwB7h3x$nh)uD*{9rszA7Ntz}D#;lS-7&T?d)@?1dePWEzU3yTCYd2%fQ z1~&Xfm^BD?a8nVy(U=BZ{lO>~ILl?)O;1e9JBAz4O zx}jmWewtMc2OQK~UO7KvC7Jdef? z6`3q%S<7<}`y$d*S&whQDo8L(Hn-r}dAFQx7VaYzGNRMpy$^~-_x8+$`Tb8RQ4f7O120D$1Gy)pX}@lS;>mYvgmsu$fa zNpJuF^9RIQU$MU+{$Ijxe*yo)n7^8f^aJ=e{>2x{AI;UZH29BAHl@v!)_q|xe|h8n zfLiNI4)(98f5YrQ$Aq(~ub2-H09XhEyukneQLgp9`~`ufwSnO;mCpV$D1V~Yj0Pz1 zytwp5kpY156M8b?znX9P>josh!DbU(?B+xV0Qj8#>P=3Tf5o=5*0=oCx#DlYp{%tf zxLyMQ?v8&QmPPTufIB!CS~-|nTm1*0AH5ckE`{c9{A%R^mX-347D8{OfBsH184t9b+J5j;eiYQ#A*Qm zOg~ZK`2LpaA2J;OTlV{Ki{|#pJhaTiyBmkg^4*($li5V;I-!T8L0Lb5tL#Dp>6M=>tOhIrf$a=OXz{TyZmCx3?jX9p}t7dQY2_})wG z`JWF(h2KP}AG7k$VfNFnf&YsLl~(~tANR6&-S`23pE>lVerddsg|5AcrG@T)FhGBX zrXSmd{d{(zyZuX!znlMml=c1D!51hXKZV)e(fnTm2*^L?KH None: - super().__init__() - - # If is_sequence_parallel, the input and output tensors are sharded - # across the ranks within the tp_group. In this case the weights are - # replicated and no collective ops are needed. - # Otherwise we use standard TP with an allreduce at the end. - self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, - [intermediate_size] * 2, - bias=False, - quant_config=quant_config, - disable_tp=is_sequence_parallel, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=False, - quant_config=quant_config, - reduce_results=reduce_results, - disable_tp=is_sequence_parallel, - prefix=f"{prefix}.down_proj", - ) - if hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported for now." - ) - if swiglu_limit is not None: - self.act_fn = SiluAndMulWithClamp(swiglu_limit) - else: - self.act_fn = SiluAndMul() - - def forward(self, x): - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -class DeepseekV4MoE(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - prefix: str = "", - ): - super().__init__() - - self.tp_size = get_tensor_model_parallel_world_size() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.prefix = prefix - - self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) - self.hidden_size = config.hidden_size - - self.n_routed_experts = config.n_routed_experts - self.n_activated_experts = config.num_experts_per_tok - self.moe_intermediate_size = config.moe_intermediate_size - self.swiglu_limit = config.swiglu_limit - self.renormalize = config.norm_topk_prob - self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") - - self.gate = GateLinear( - input_size=config.hidden_size, - output_size=config.n_routed_experts, - bias=False, - out_dtype=torch.float32, - prefix=f"{prefix}.gate", - ) - - self.gate.e_score_correction_bias = None - self.gate.tid2eid = None - is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers - self.hash_indices_dtype = torch.int32 - if is_hash_moe: - # hash MoE doesn't use e_score_correction_bias - # Use randint instead of empty to avoid garbage values causing - # invalid memory access in dummy mode (--load-format="dummy") - self.gate.tid2eid = nn.Parameter( - torch.randint( - 0, - config.n_routed_experts, - (config.vocab_size, config.num_experts_per_tok), - dtype=self.hash_indices_dtype, - ), - requires_grad=False, - ) - elif getattr(config, "topk_method", None) == "noaux_tc": - self.gate.e_score_correction_bias = nn.Parameter( - torch.empty(config.n_routed_experts, dtype=torch.float32), - requires_grad=False, - ) - - if config.n_shared_experts is None: - self.shared_experts = None - else: - intermediate_size = config.moe_intermediate_size * config.n_shared_experts - - self.shared_experts = DeepseekV4MLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - swiglu_limit=self.swiglu_limit, - quant_config=quant_config, - reduce_results=False, - prefix=f"{prefix}.shared_experts", - ) - - self.tp_rank = get_tensor_model_parallel_rank() - assert config.n_routed_experts % self.tp_size == 0 - - self.n_local_experts = config.n_routed_experts // self.tp_size - self.experts_start_idx = self.tp_rank * self.n_local_experts - self.experts_end_idx = self.experts_start_idx + self.n_local_experts - - self.experts = FusedMoE( - shared_experts=self.shared_experts, - gate=self.gate, - num_experts=config.n_routed_experts, - top_k=config.num_experts_per_tok, - hidden_size=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - renormalize=config.norm_topk_prob, - quant_config=quant_config, - prefix=f"{prefix}.experts", - scoring_func=self.scoring_func, - routed_scaling_factor=self.routed_scaling_factor, - e_score_correction_bias=self.gate.e_score_correction_bias, - hash_indices_table=self.gate.tid2eid, - swiglu_limit=self.swiglu_limit, - router_logits_dtype=torch.float32, - ) - - def forward( - self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None - ) -> torch.Tensor: - if self.gate.tid2eid is not None and input_ids is None: - raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") - - org_shape = hidden_states.shape - if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - final_hidden_states = self.experts( - hidden_states=hidden_states, - router_logits=hidden_states, - input_ids=input_ids, - ) - else: - router_logits, _ = self.gate(hidden_states) - final_hidden_states = self.experts( - hidden_states=hidden_states, - router_logits=router_logits, - input_ids=input_ids, - ) - - return final_hidden_states.view(org_shape) - - -class DeepseekV4Attention(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - prefix: str, - topk_indices_buffer: torch.Tensor | None = None, - aux_stream_list: list[torch.cuda.Stream] | None = None, - ): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - layer_id = extract_layer_index(prefix) - - self.layer_id = layer_id - self.hidden_size = config.hidden_size - self.n_heads = config.num_attention_heads - tp_size = get_tensor_model_parallel_world_size() - assert self.n_heads % tp_size == 0 - - self.n_local_heads = self.n_heads // tp_size - self.q_lora_rank = config.q_lora_rank - self.o_lora_rank = config.o_lora_rank - self.head_dim = config.head_dim - self.rope_head_dim = config.qk_rope_head_dim - self.nope_head_dim = self.head_dim - self.rope_head_dim - self.n_groups = config.o_groups - self.n_local_groups = self.n_groups // tp_size - self.window_size = config.sliding_window - # NOTE(zyongye) Compress ratio can't be 0 - # we do this for because MTP layer is not included - # in the compress ratio list - if layer_id < config.num_hidden_layers: - self.compress_ratio = max(1, config.compress_ratios[layer_id]) - else: - self.compress_ratio = 1 - self.eps = config.rms_norm_eps - self.max_position_embeddings = config.max_position_embeddings - - # Padded to min 64 heads for FlashMLA, initialized to -inf - # (no sink effect). Weight loading fills the first n_local_heads slots. - padded_heads = max(self.n_local_heads, 64) - self.attn_sink = nn.Parameter( - torch.full((padded_heads,), -float("inf"), dtype=torch.float32), - requires_grad=False, - ) - - self.fused_wqa_wkv = MergedColumnParallelLinear( - self.hidden_size, - [self.q_lora_rank, self.head_dim], - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.fused_wqa_wkv", - disable_tp=True, # fused ReplicatedLinear - ) - self.q_norm = RMSNorm(self.q_lora_rank, self.eps) - self.wq_b = ColumnParallelLinear( - self.q_lora_rank, - self.n_heads * self.head_dim, - bias=False, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.wq_b", - ) - - self.kv_norm = RMSNorm(self.head_dim, self.eps) - self.wo_a = ColumnParallelLinear( - self.n_heads * self.head_dim // self.n_groups, - self.n_groups * self.o_lora_rank, - bias=False, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.wo_a", - ) - self.wo_a.is_bmm = True - self.wo_a.bmm_batch_size = self.n_local_groups - self.wo_b = RowParallelLinear( - self.n_groups * self.o_lora_rank, - self.hidden_size, - bias=False, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.wo_b", - ) - self.softmax_scale = self.head_dim**-0.5 - self.scale_fmt = config.quantization_config["scale_fmt"] - - self.rope_parameters = config.rope_scaling - - # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it) - self.rotary_emb = build_deepseek_v4_rope( - config, - head_dim=self.head_dim, - rope_head_dim=self.rope_head_dim, - max_position_embeddings=self.max_position_embeddings, - compress_ratio=self.compress_ratio, - ) - - self.indexer = None - if self.compress_ratio == 4: - # Only C4A uses sparse attention and hence has indexer. - self.indexer = DeepseekV4Indexer( - vllm_config, - config=config, - hidden_size=self.hidden_size, - q_lora_rank=self.q_lora_rank, - quant_config=quant_config, - cache_config=vllm_config.cache_config, - topk_indices_buffer=topk_indices_buffer, - compress_ratio=self.compress_ratio, - prefix=f"{prefix}.indexer", - ) - - self.mla_attn = DeepseekV4MLA( - hidden_size=self.hidden_size, - num_heads=self.n_local_heads, - head_dim=self.head_dim, - scale=self.softmax_scale, - qk_nope_head_dim=self.nope_head_dim, - qk_rope_head_dim=self.rope_head_dim, - v_head_dim=self.head_dim, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.head_dim, - o_lora_rank=self.o_lora_rank, - vllm_config=vllm_config, - fused_wqa_wkv=self.fused_wqa_wkv, - q_norm=self.q_norm, - wq_b=self.wq_b, - kv_norm=self.kv_norm, - wo_a=self.wo_a, - wo_b=self.wo_b, - attn_sink=self.attn_sink, - rotary_emb=self.rotary_emb, - indexer=self.indexer, - indexer_rotary_emb=self.rotary_emb, - topk_indices_buffer=topk_indices_buffer, - aux_stream_list=aux_stream_list, - window_size=self.window_size, - compress_ratio=self.compress_ratio, - cache_config=vllm_config.cache_config, - quant_config=quant_config, - prefix=prefix, - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - llama_4_scaling: torch.Tensor | None, - ): - return self.mla_attn(positions, hidden_states, llama_4_scaling) - - -class DeepseekV4DecoderLayer(nn.Module): - def __init__( - self, - vllm_config, - prefix, - topk_indices_buffer: torch.Tensor | None = None, - aux_stream_list: list[torch.cuda.Stream] | None = None, - ): - super().__init__() - - # Lazy import to avoid top-level tilelang dependency. - # Registers both torch.ops.vllm.mhc_pre and mhc_post - import vllm.model_executor.layers.mhc # noqa: F401 - - config = vllm_config.model_config.hf_config - self.hidden_size = config.hidden_size - - self.rms_norm_eps = config.rms_norm_eps - self.attn = DeepseekV4Attention( - vllm_config, - prefix=f"{prefix}.attn", - topk_indices_buffer=topk_indices_buffer, - aux_stream_list=aux_stream_list, - ) - self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") - - self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) - self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) - self.hc_mult = config.hc_mult - self.hc_sinkhorn_iters = config.hc_sinkhorn_iters - self.hc_eps = config.hc_eps - self.hc_post_alpha = 2.0 - mix_hc = (2 + self.hc_mult) * self.hc_mult - hc_dim = self.hc_mult * self.hidden_size - self.hc_attn_fn = nn.Parameter( - torch.empty( - (mix_hc, hc_dim), - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_ffn_fn = nn.Parameter( - torch.empty( - (mix_hc, hc_dim), - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_attn_base = nn.Parameter( - torch.empty( - mix_hc, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_ffn_base = nn.Parameter( - torch.empty( - mix_hc, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_attn_scale = nn.Parameter( - torch.empty( - 3, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_ffn_scale = nn.Parameter( - torch.empty( - 3, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.mhc_pre = MHCPreOp() - self.mhc_post = MHCPostOp() - self.mhc_fused_post_pre = MHCFusedPostPreOp() - self.has_tilelang = has_tilelang() - - def hc_pre( - self, - x: torch.Tensor, - hc_fn: torch.Tensor, - hc_scale: torch.Tensor, - hc_base: torch.Tensor, - ): - post_mix, res_mix, layer_input = self.mhc_pre( - residual=x, - fn=hc_fn, - hc_scale=hc_scale, - hc_base=hc_base, - rms_eps=self.rms_norm_eps, - hc_pre_eps=self.hc_eps, - hc_sinkhorn_eps=self.hc_eps, - hc_post_mult_value=self.hc_post_alpha, - sinkhorn_repeat=self.hc_sinkhorn_iters, - ) - return layer_input, post_mix, res_mix - - def hc_post( - self, - x: torch.Tensor, - residual: torch.Tensor, - post: torch.Tensor, - comb: torch.Tensor, - ): - return self.mhc_post(x, residual, post, comb) - - def _forward_fused_post_pre( - self, - x: torch.Tensor, - positions: torch.Tensor, - input_ids: torch.Tensor | None, - post_mix: torch.Tensor | None = None, - res_mix: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - if residual is None: - # Run standalone hc_pre on first layer - residual = x - x, post_mix, res_mix = self.hc_pre( - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base - ) - else: - residual, post_mix, res_mix, x = self.mhc_fused_post_pre( - x, - residual, - post_mix, - res_mix, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - ) - - x = self.attn_norm(x) - x = self.attn(positions, x, None) - - residual, post_mix, res_mix, x = self.mhc_fused_post_pre( - x, - residual, - post_mix, - res_mix, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - ) - x = self.ffn_norm(x) - x = self.ffn(x, input_ids) - return x, residual, post_mix, res_mix - - def _forward_unfused_post_pre( - self, - x: torch.Tensor, - positions: torch.Tensor, - input_ids: torch.Tensor | None, - post_mix: torch.Tensor | None = None, - res_mix: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - ) -> tuple[ - torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None - ]: - residual = x - x, post, comb = self.hc_pre( - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base - ) - x = self.attn_norm(x) - x = self.attn(positions, x, None) - x = self.hc_post(x, residual, post, comb) - - residual = x - x, post, comb = self.hc_pre( - x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base - ) - x = self.ffn_norm(x) - x = self.ffn(x, input_ids) - x = self.hc_post(x, residual, post, comb) - return x, None, None, None - - def forward( - self, - x: torch.Tensor, - positions: torch.Tensor, - input_ids: torch.Tensor | None, - post_mix: torch.Tensor | None = None, - res_mix: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - ) -> tuple[ - torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None - ]: - if not self.has_tilelang: - return self._forward_unfused_post_pre( - x, positions, input_ids, post_mix, res_mix, residual - ) - return self._forward_fused_post_pre( - x, positions, input_ids, post_mix, res_mix, residual - ) - - -class DeepseekV4Model(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.config = config - self.vocab_size = config.vocab_size - self.hc_eps = config.hc_eps - self.hc_mult = config.hc_mult - self.hc_dim = self.hc_mult * config.hidden_size - self.rms_norm_eps = config.rms_norm_eps - - # Three aux streams: one per non-default input GEMM in - # DeepseekV4MLA.attn_gemm_parallel_execute - # (compressor kv_score, indexer.weights_proj, indexer.compressor - # kv_score). fused_wqa_wkv stays on the default stream. - # Disable them on ROCm because of hang issues. - aux_stream_list = ( - None - if current_platform.is_rocm() - else [torch.cuda.Stream() for _ in range(3)] - ) - - self.device = current_platform.device_type - # Reserved topk indices buffer for all Indexer layers to reuse. - self.topk_indices_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - config.index_topk, - dtype=torch.int32, - device=self.device, - ) - - if get_pp_group().is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - else: - self.embed_tokens = PPMissingLayer() - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: DeepseekV4DecoderLayer( - vllm_config, - prefix=prefix, - topk_indices_buffer=self.topk_indices_buffer, - aux_stream_list=aux_stream_list, - ), - prefix=f"{prefix}.layers", - ) - - if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) - else: - self.norm = PPMissingLayer() - - self.hc_head_fn = nn.Parameter( - torch.empty( - self.hc_mult, - self.hc_dim, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_head_base = nn.Parameter( - torch.empty( - self.hc_mult, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_head_scale = nn.Parameter( - torch.empty(1, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_op = HCHeadOp() - self.has_tilelang = has_tilelang() - # Pre-hc_head residual stream buffer for the MTP draft. Stable - # address (outside the cudagraph pool) so the copy_ in forward() - # refreshes it correctly across captured shapes. - # refreshes it correctly across captured shapes. Only allocated on - # the last PP rank — that's where MTP target hidden states are - # produced. - if get_pp_group().is_last_rank: - self._mtp_hidden_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - self.hc_dim, - dtype=vllm_config.model_config.dtype, - device=self.device, - ) - else: - self._mtp_hidden_buffer = None - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def make_empty_intermediate_tensors( - self, - batch_size: int, - dtype: torch.dtype, - device: torch.device, - ) -> IntermediateTensors: - # PP intermediate tensors carry the multi-stream hidden_states - # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the - # token embedding to hc_mult streams before the first decoder - # layer and keeps that shape until hc_head() collapses it. - return IntermediateTensors( - { - "hidden_states": torch.zeros( - (batch_size, self.hc_mult, self.config.hidden_size), - dtype=dtype, - device=device, - ), - } - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - - residual, post_mix, res_mix = None, None, None - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual, post_mix, res_mix = layer( - hidden_states, - positions, - input_ids, - post_mix, - res_mix, - residual, - ) - if layer is not None and self.has_tilelang: - hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) - - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) - - # Stash pre-hc_head residual for the MTP draft (captured copy_). - num_tokens = hidden_states.shape[0] - self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) - - hidden_states = self.hc_head_op( - hidden_states, - self.hc_head_fn, - self.hc_head_scale, - self.hc_head_base, - self.rms_norm_eps, - self.hc_eps, - ) - hidden_states = self.norm(hidden_states) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("gate_up_proj", "w1", 0), - ("gate_up_proj", "w3", 1), - ("attn.fused_wqa_wkv", "attn.wq_a", 0), - ("attn.fused_wqa_wkv", "attn.wkv", 1), - ("compressor.fused_wkv_wgate", "compressor.wkv", 0), - ("compressor.fused_wkv_wgate", "compressor.wgate", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - - # Pre-compute expert mapping ONCE. - expert_mapping = self.get_expert_mapping() - - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - # Skip non-stacked layers and experts (experts handled below). - if ".experts." in name: - continue - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - - if is_pp_missing_parameter(name, self): - break - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(name) - break - else: - if ".experts." in name: - # E8M0 scales are stored as float8_e8m0fnu in - # checkpoints but the MoE param is uint8. copy_() - # would do a numeric conversion (e.g. 2^-7 → 0), - # destroying the raw exponent bytes. - if ( - "weight_scale" in name - and loaded_weight.dtype == torch.float8_e8m0fnu - ): - loaded_weight = loaded_weight.view(torch.uint8) - for mapping in expert_mapping: - param_name, weight_name, expert_id, expert_shard_id = mapping - if weight_name not in name: - continue - name_mapped = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name_mapped, self): - continue - param = params_dict[name_mapped] - # We should ask the weight loader to return success or not - # here since otherwise we may skip experts with other - # available replicas. - weight_loader = typing.cast( - Callable[..., bool], param.weight_loader - ) - success = weight_loader( - param, - loaded_weight, - name_mapped, - shard_id=expert_shard_id, - expert_id=expert_id, - return_success=True, - ) - if success: - name = name_mapped - break - loaded_params.add(name_mapped) - continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue - else: - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - continue - - return loaded_params - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - return FusedMoE.make_expert_params_mapping( - self, - ckpt_gate_proj_name="w1", - ckpt_down_proj_name="w2", - ckpt_up_proj_name="w3", - num_experts=self.config.n_routed_experts, - ) - - -def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: - if expert_dtype == "fp4": - # MXFP4 experts use Mxfp4MoEMethod, which registers scales as - # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and - # shared experts use Fp8LinearMethod's block scales, which - # register as ``weight_scale_inv``. - scale_regex = { - re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", - re.compile(r"\.scale$"): ".weight_scale_inv", - } - else: - # FP8 experts use Fp8MoEMethod (block_quant=True), which registers - # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys - # there. - scale_regex = { - re.compile(r"\.scale$"): ".weight_scale_inv", - } - return WeightsMapper( - orig_to_new_prefix={ - "layers.": "model.layers.", - "embed.": "model.embed.", - "norm.": "model.norm.", - "hc_head": "model.hc_head", - "mtp.": "model.mtp.", - }, - orig_to_new_regex=scale_regex, - orig_to_new_suffix={ - "head.weight": "lm_head.weight", - "embed.weight": "embed_tokens.weight", - ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", - }, - orig_to_new_substr={ - ".attn.compressor.": ".attn.mla_attn.compressor.", - ".shared_experts.w2": ".shared_experts.down_proj", - }, - ) - - -class DeepseekV4ForCausalLM(nn.Module, SupportsPP): - model_cls = DeepseekV4Model - - # Default mapper assumes the original FP4-expert checkpoint layout. - # Overridden per-instance in __init__ when expert_dtype != "fp4". - hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4") - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - self.config = config - expert_dtype = getattr(config, "expert_dtype", "fp4") - if expert_dtype != "fp4": - self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype) - - self.model = self.model_cls( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - if get_pp_group().is_last_rank: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) - else: - self.lm_head = PPMissingLayer() - self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def get_mtp_target_hidden_states(self) -> torch.Tensor | None: - """Pre-hc_head residual stream buffer (max_num_batched_tokens, - hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" - return getattr(self.model, "_mtp_hidden_buffer", None) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) - loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - return loaded_params - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - return self.model.get_expert_mapping() diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py b/TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py deleted file mode 100644 index 5938cde6..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py +++ /dev/null @@ -1,509 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4). - -Split from ``deepseek_mtp.py`` because the V4 architecture introduces several -pieces that have no analogue in V3/V32: - * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of - the fused ``eh_proj``); - * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``; - * ``DeepseekV4DecoderLayer`` with its own aux-stream management; - * V4-specific checkpoint weight-name remapping in ``load_weights``. -""" - -import typing -from collections.abc import Callable, Iterable - -import regex as re -import torch -import torch.nn as nn - -from vllm.config import VllmConfig -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import FusedMoE -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ReplicatedLinear -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mhc import HCHeadOp -from vllm.model_executor.layers.vocab_parallel_embedding import ( - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.deepseek_mtp import SharedHead -from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name -from vllm.model_executor.models.utils import maybe_prefix -from vllm.models.deepseek_v4.common.ops import ( - fused_mtp_input_rmsnorm, - mtp_shared_head_rmsnorm, -) -from vllm.platforms import current_platform -from vllm.sequence import IntermediateTensors -from vllm.utils.import_utils import has_tilelang - -from .model import DeepseekV4DecoderLayer - -logger = init_logger(__name__) - -# MoE expert scales are fused into per-layer w13/w2 tensors. The exact -# parameter suffix depends on which FusedMoE method handles the experts: -# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``; -# - fp8 experts (Fp8MoEMethod with block_quant=True) register -# ``w{1,2,3}_weight_scale_inv``. -# Other FP8 linear scales (including shared experts) always use -# ``.weight_scale_inv``. Mirrors the per-instance mapper built by -# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py. -_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") - - -class DeepSeekV4MultiTokenPredictorLayer(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - topk_indices_buffer: torch.Tensor, - prefix: str, - aux_stream_list: list[torch.cuda.Stream] | None = None, - ) -> None: - super().__init__() - - assert vllm_config.speculative_config is not None - config = vllm_config.speculative_config.draft_model_config.hf_config - self.config = config - quant_config = vllm_config.quant_config - self.rms_norm_eps = config.rms_norm_eps - - self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than - # fusing them the way V3 does with eh_proj. - self.e_proj = ReplicatedLinear( - config.hidden_size, - config.hidden_size, - bias=False, - return_bias=False, - quant_config=quant_config, - ) - self.h_proj = ReplicatedLinear( - config.hidden_size, - config.hidden_size, - bias=False, - return_bias=False, - quant_config=quant_config, - ) - - self.hc_eps = config.hc_eps - self.hc_mult = config.hc_mult - self.hc_dim = self.hc_mult * config.hidden_size - self.hc_head_fn = nn.Parameter( - torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_base = nn.Parameter( - torch.empty(self.hc_mult, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_scale = nn.Parameter( - torch.empty(1, dtype=torch.float32), - requires_grad=False, - ) - - self.shared_head = SharedHead( - config=config, prefix=prefix, quant_config=quant_config - ) - self.mtp_block = DeepseekV4DecoderLayer( - vllm_config, - prefix, - topk_indices_buffer=topk_indices_buffer, - aux_stream_list=aux_stream_list, - ) - - self.hc_head_op = HCHeadOp() - self.has_tilelang = has_tilelang() - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - previous_hidden_states: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - spec_step_index: int = 0, - ) -> torch.Tensor: - assert inputs_embeds is not None - # Target stashes pre-hc_head residual as flat (T, hc_mult * D); - # reshape to (T, hc_mult, D) — the training-time layout — before - # the fused norm pass so both inputs are 3D-friendly. - previous_hidden_states = previous_hidden_states.view( - -1, self.hc_mult, self.config.hidden_size - ) - # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm. - inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm( - inputs_embeds, - positions, - previous_hidden_states, - self.enorm.weight.data, - self.hnorm.weight.data, - self.enorm.variance_epsilon, - self.hc_mult, - ) - hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( - inputs_embeds - ).unsqueeze(-2) - hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=None - ) - if self.has_tilelang: - hidden_states = self.mtp_block.hc_post( - hidden_states, residual, post_mix, res_mix - ) - # Return the flat pre-hc_head residual so it can be re-fed as the - # next spec step's `previous_hidden_states` when - # num_speculative_tokens > 1. hc_head is deferred to compute_logits. - return hidden_states.flatten(1) - - -class DeepSeekV4MultiTokenPredictor(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - self.mtp_start_layer_idx = config.num_hidden_layers - self.num_mtp_layers = config.num_nextn_predict_layers - self.device = current_platform.device_type - - topk_tokens = config.index_topk - self.topk_indices_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - topk_tokens, - dtype=torch.int32, - device=self.device, - ) - - # Three aux streams shared across all MTP layers, mirroring - # DeepseekV4Model. ROCm runs the same work serially for now. - aux_stream_list = ( - None - if current_platform.is_rocm() - else [torch.cuda.Stream() for _ in range(3)] - ) - - # to map the exact layer index from weights - self.layers = torch.nn.ModuleDict( - { - str(idx): DeepSeekV4MultiTokenPredictorLayer( - vllm_config, - self.topk_indices_buffer, - f"{prefix}.layers.{idx}", - aux_stream_list=aux_stream_list, - ) - for idx in range( - self.mtp_start_layer_idx, - self.mtp_start_layer_idx + self.num_mtp_layers, - ) - } - ) - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "embed_tokens"), - ) - self.logits_processor = LogitsProcessor(config.vocab_size) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - previous_hidden_states: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - spec_step_idx: int = 0, - ) -> torch.Tensor: - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - current_step_idx = spec_step_idx % self.num_mtp_layers - return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( - input_ids, - positions, - previous_hidden_states, - inputs_embeds, - current_step_idx, - ) - - def compute_logits( - self, - hidden_states: torch.Tensor, - spec_step_idx: int = 0, - ) -> torch.Tensor: - current_step_idx = spec_step_idx % self.num_mtp_layers - mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] - # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply - # hc_head here so logits are computed from the dense hidden state. - hidden_states = hidden_states.view( - -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size - ) - hidden_states = mtp_layer.hc_head_op( - hidden_states, - mtp_layer.hc_head_fn, - mtp_layer.hc_head_scale, - mtp_layer.hc_head_base, - mtp_layer.rms_norm_eps, - mtp_layer.hc_eps, - ) - hidden_states = mtp_shared_head_rmsnorm( - hidden_states, - mtp_layer.shared_head.norm.weight.data, - mtp_layer.shared_head.norm.variance_epsilon, - ) - logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) - return logits - - -class DeepSeekV4MTP(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - self.config = vllm_config.model_config.hf_config - self.quant_config = vllm_config.quant_config - self.model = DeepSeekV4MultiTokenPredictor( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - hidden_states: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - spec_step_idx: int = 0, - ) -> torch.Tensor: - hidden_states = self.model( - input_ids, positions, hidden_states, inputs_embeds, spec_step_idx - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - spec_step_idx: int = 0, - ) -> torch.Tensor | None: - return self.model.compute_logits(hidden_states, spec_step_idx) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Weight name remapping for checkpoint compatibility. - # Maps checkpoint weight paths to model parameter paths. - WEIGHT_NAME_REMAPPING: dict[str, str] = { - ".emb.tok_emb.weight": ".embed_tokens.weight", - ".head.weight": ".shared_head.head.weight", - ".norm.weight": ".shared_head.norm.weight", - } - - def _remap_weight_name(name: str) -> str: - """Remap checkpoint weight names to model parameter names.""" - for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items(): - if old_pattern in name: - name = name.replace(old_pattern, new_pattern) - return name - - def _find_mtp_layer_idx(name: str) -> int: - subnames = name.split(".") - for subname in subnames: - try: - # we return the first encountered integer - return int(subname) - except ValueError: - continue - return 0 - - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("gate_up_proj", "w1", 0), - ("gate_up_proj", "w3", 1), - ("attn.fused_wqa_wkv", "attn.wq_a", 0), - ("attn.fused_wqa_wkv", "attn.wkv", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - - # Pre-compute expert mapping ONCE. - expert_mapping = FusedMoE.make_expert_params_mapping( - self, - ckpt_gate_proj_name="w1", - ckpt_down_proj_name="w2", - ckpt_up_proj_name="w3", - num_experts=self.config.n_routed_experts, - ) - - # FP8 experts register ``..._weight_scale_inv`` (block_quant) while - # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix - # for the rename below based on the model's expert dtype. - expert_scale_suffix = ( - ".weight_scale" - if getattr(self.config, "expert_dtype", "fp4") == "fp4" - else ".weight_scale_inv" - ) - - for name, loaded_weight in weights: - mtp_layer_idx = _find_mtp_layer_idx(name) - # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to - # `model.layers.{num_hidden_layers + i}.*` so that - # get_spec_layer_idx_from_weight_name can identify them. - name = name.replace( - f"mtp.{mtp_layer_idx}.", - f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", - ) - - spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) - if spec_layer is None: - continue - - name = _remap_weight_name(name) - name = self._rewrite_spec_layer_name(spec_layer, name) - - if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name: - continue - if name.endswith(".scale"): - suffix = ( - expert_scale_suffix - if _EXPERT_SCALE_RE.search(name) - else ".weight_scale_inv" - ) - name = name.removesuffix(".scale") + suffix - for param_name, weight_name, shard_id in stacked_params_mapping: - # Skip non-stacked layers and experts (experts handled below). - if ".experts." in name: - continue - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(name) - break - else: - if ".experts." in name: - # Reinterpret E8M0 scales as uint8 to preserve raw - # exponent bytes; numeric copy_() would zero them. - # Mirrors the main DeepseekV4 loader. - if ( - "weight_scale" in name - and loaded_weight.dtype == torch.float8_e8m0fnu - ): - loaded_weight = loaded_weight.view(torch.uint8) - for mapping in expert_mapping: - param_name, weight_name, expert_id, expert_shard_id = mapping - if weight_name not in name: - continue - name_mapped = name.replace(weight_name, param_name) - param = params_dict[name_mapped] - # We should ask the weight loader to return success or not - # here since otherwise we may skip experts with other - # available replicas. - weight_loader = typing.cast( - Callable[..., bool], param.weight_loader - ) - success = weight_loader( - param, - loaded_weight, - name_mapped, - shard_id=expert_shard_id, - expert_id=expert_id, - return_success=True, - ) - if success: - name = name_mapped - loaded_params.add(name_mapped) - break - continue - elif "attn_sink" in name: - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue - else: - if ".shared_experts.w2" in name: - name = name.replace( - ".shared_experts.w2", ".shared_experts.down_proj" - ) - if name.endswith(".ffn.gate.bias"): - name = name.replace(".bias", ".e_score_correction_bias") - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - continue - - loaded_layers: set[int] = set() - for param_name in loaded_params: - spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) - if spec_layer is not None: - loaded_layers.add(spec_layer) - for layer_idx in range( - self.model.mtp_start_layer_idx, - self.model.mtp_start_layer_idx + self.model.num_mtp_layers, - ): - if layer_idx not in loaded_layers: - raise ValueError( - f"MTP speculative decoding layer {layer_idx} weights " - f"missing from checkpoint. The checkpoint may have " - f"been quantized without including the MTP layers. " - f"Use a checkpoint that includes MTP layer weights, " - f"or disable speculative decoding." - ) - logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) - return loaded_params - - def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. - Add .mtp_block for modules in transformer layer block for spec layer - and rename shared layer weights to be top level. - """ - spec_layer_weight_names = [ - "embed_tokens", - "enorm", - "hnorm", - "h_proj", - "e_proj", - "shared_head", - "hc_head_fn", - "hc_head_base", - "hc_head_scale", - ] - shared_weight_names = ["embed_tokens"] - spec_layer_weight = False - shared_weight = False - for weight_name in spec_layer_weight_names: - if weight_name in name: - spec_layer_weight = True - if weight_name in shared_weight_names: - shared_weight = True - break - if not spec_layer_weight: - # treat rest weights as weights for transformer layer block - name = name.replace( - f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." - ) - elif shared_weight: - # treat shared weights as top level weights - name = name.replace(f"model.layers.{spec_layer}.", "model.") - return name diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py b/TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py deleted file mode 100644 index 2af93fba..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py +++ /dev/null @@ -1,856 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from dataclasses import dataclass -from typing import TYPE_CHECKING, cast - -import torch - -from vllm.forward_context import get_forward_context -from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache -from vllm.models.deepseek_v4.nvidia.flashmla import ( - DeepseekV4FlashMLASparseBackend, - DeepseekV4SparseMLAAttentionImpl, -) -from vllm.triton_utils import tl, triton -from vllm.v1.attention.backend import ( - CommonAttentionMetadata, -) -from vllm.v1.attention.backends.mla.flashmla_sparse import ( - FlashMLASparseMetadata, - FlashMLASparseMetadataBuilder, -) -from vllm.v1.attention.backends.mla.sparse_swa import ( - DeepseekSparseSWAMetadata, - DeepseekSparseSWAMetadataBuilder, -) -from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - build_ragged_indices_from_dense, - rocm_sparse_attn_decode, - rocm_sparse_attn_prefill, -) -from vllm.v1.worker.workspace import current_workspace_manager - -if TYPE_CHECKING: - from vllm.models.deepseek_v4.attention import ( - DeepseekV4MLAAttention, - ) - - -def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: - lengths = lengths.to(dtype=torch.int32).contiguous() - indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) - torch.cumsum(lengths, dim=0, out=indptr[1:]) - return indptr - - -# ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes -# do not touch the shared DeepSeek V4 cache utilities. -_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 - - -@triton.jit -def _combine_topk_swa_indices_kernel( - combined_indices_ptr, - combined_indices_stride, - combined_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - TOPK_WIDTH: tl.constexpr, - PADDED_TOP_K: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) - - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - gather_start = seq_len - gather_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - - topk_offset = tl.arange(0, PADDED_TOP_K) - topk_mask = topk_offset < topk_len - safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0) - topk_indices = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset, - mask=topk_mask, - other=-1, - ) - valid_topk = (topk_indices >= 0) & (topk_indices < N) - topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1) - tl.store( - combined_indices_ptr + token_idx * combined_indices_stride + topk_offset, - topk_indices, - mask=topk_mask, - ) - - swa_offset = tl.arange(0, WINDOW_SIZE) - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + swa_offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_offset < swa_len, - ) - - tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) - - -def combine_topk_swa_indices( - topk_indices: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor, - window_size: int, - compress_ratio: int, - topk: int, - M: int, - N: int, -) -> tuple[torch.Tensor, torch.Tensor]: - topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() - num_tokens = topk_indices.shape[0] - num_reqs = seq_lens.shape[0] - combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) - // _SPARSE_PREFILL_TOPK_ALIGNMENT - * _SPARSE_PREFILL_TOPK_ALIGNMENT - ) - combined_indices = torch.full( - (num_tokens, combined_topk), - fill_value=-1, - dtype=torch.int32, - device=topk_indices.device, - ) - combined_lens = torch.empty( - num_tokens, dtype=torch.int32, device=topk_indices.device - ) - - num_workers = 128 - _combine_topk_swa_indices_kernel[(num_reqs, num_workers)]( - combined_indices, - combined_indices.stride(0), - combined_lens, - topk_indices, - topk_indices.stride(0), - query_start_loc, - seq_lens, - gather_lens, - M, - N, - TOP_K=topk, - COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=window_size, - TOPK_WIDTH=topk_indices.shape[-1], - PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), - ) - return combined_indices, combined_lens - - -@triton.jit -def _compute_topk_lens_kernel( - topk_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - topk, - is_valid_token_ptr, - TRITON_BLOCK_SIZE: tl.constexpr, -): - token_idx = tl.program_id(0) - is_valid_token = tl.load(is_valid_token_ptr + token_idx) - - count = tl.zeros((), dtype=tl.int32) - for i in range(0, topk, TRITON_BLOCK_SIZE): - offset = i + tl.arange(0, TRITON_BLOCK_SIZE) - mask = offset < topk - local_idx = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=mask, - other=-1, - ) - count += tl.sum((local_idx >= 0).to(tl.int32), axis=0) - - tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) - - -@triton.jit -def _pack_global_topk_ragged_kernel( - global_topk_ragged_ptr, - topk_indptr_ptr, - topk_indices_ptr, - topk_indices_stride, - token_to_req_indices_ptr, - block_table_ptr, - block_table_stride, - block_size, - topk, - BLOCK_SIZE: tl.constexpr, -): - token_idx = tl.program_id(0) - block_idx = tl.program_id(1) - offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - - out_start = tl.load(topk_indptr_ptr + token_idx) - out_end = tl.load(topk_indptr_ptr + token_idx + 1) - out_len = out_end - out_start - if block_idx * BLOCK_SIZE >= out_len: - return - - req_idx = tl.load(token_to_req_indices_ptr + token_idx) - mask = (offset < out_len) & (offset < topk) - local_idx = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=mask, - other=-1, - ) - valid = mask & (local_idx >= 0) - block_indices = local_idx // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=valid, - other=0, - ) - block_offsets = local_idx % block_size - slot_ids = tl.where(valid, block_numbers * block_size + block_offsets, -1) - tl.store(global_topk_ragged_ptr + out_start + offset, slot_ids, mask=mask) - - -def compute_global_topk_ragged_indices_and_indptr( - topk_indices: torch.Tensor, - token_to_req_indices: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - is_valid_token: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() - num_tokens = topk_indices.shape[0] - topk = topk_indices.shape[1] - - topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) - _compute_topk_lens_kernel[(num_tokens,)]( - topk_lens, - topk_indices, - topk_indices.stride(0), - topk, - is_valid_token, - TRITON_BLOCK_SIZE=1024, - ) - - topk_indptr = _build_indptr_from_lengths(topk_lens) - global_topk_ragged = torch.empty( - num_tokens * topk, - dtype=torch.int32, - device=topk_indices.device, - ) - if global_topk_ragged.numel() > 0: - block = 128 - _pack_global_topk_ragged_kernel[(num_tokens, triton.cdiv(topk, block))]( - global_topk_ragged, - topk_indptr, - topk_indices, - topk_indices.stride(0), - token_to_req_indices, - block_table, - block_table.stride(0), - block_size, - topk, - BLOCK_SIZE=block, - ) - return global_topk_ragged, topk_indptr, topk_lens - - -@triton.jit -def _compute_combined_lens_kernel( - combined_lens_ptr, - query_start_loc_ptr, - seq_lens_ptr, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) - - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - start_pos = seq_len - query_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) - - -@triton.jit -def _combine_topk_swa_indices_ragged_kernel( - combined_ragged_ptr, - combined_indptr_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - topk_width, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - block_idx = tl.program_id(2) - num_workers = tl.num_programs(1) - - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - gather_start = seq_len - gather_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - combined_len = topk_len + swa_len - - offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - if block_idx * BLOCK_SIZE < combined_len: - out_start = tl.load(combined_indptr_ptr + token_idx) - topk_mask = (offset < topk_len) & (offset < topk_width) - topk_vals = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=topk_mask, - other=-1, - ) - tl.store( - combined_ragged_ptr + out_start + offset, - topk_vals + M * batch_idx, - mask=topk_mask, - ) - - swa_offset = offset - topk_len - swa_mask = (offset >= topk_len) & (swa_offset < swa_len) - tl.store( - combined_ragged_ptr + out_start + offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_mask, - ) - - -def combine_topk_swa_indices_ragged( - topk_indices: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor, - window_size: int, - compress_ratio: int, - topk: int, - M: int, - N: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() - num_tokens = topk_indices.shape[0] - num_reqs = seq_lens.shape[0] - combined_lens = torch.empty( - num_tokens, dtype=torch.int32, device=topk_indices.device - ) - - num_workers = 128 - _compute_combined_lens_kernel[(num_reqs, num_workers)]( - combined_lens, - query_start_loc, - seq_lens, - TOP_K=topk, - COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=window_size, - ) - - combined_indptr = _build_indptr_from_lengths(combined_lens) - combined_ragged = torch.empty( - num_tokens * (topk + window_size), - dtype=torch.int32, - device=topk_indices.device, - ) - if combined_ragged.numel() > 0: - block = 128 - _combine_topk_swa_indices_ragged_kernel[ - (num_reqs, num_workers, triton.cdiv(topk + window_size, block)) - ]( - combined_ragged, - combined_indptr, - topk_indices, - topk_indices.stride(0), - query_start_loc, - seq_lens, - gather_lens, - M, - N, - topk_indices.shape[-1], - TOP_K=topk, - COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=window_size, - BLOCK_SIZE=block, - ) - return combined_ragged, combined_indptr, combined_lens - - -def _copy_ragged_to_graph_buffers( - ragged_indices: torch.Tensor, - ragged_indptr: torch.Tensor, - ragged_indices_buffer: torch.Tensor, - ragged_indptr_buffer: torch.Tensor, - num_rows: int, - max_entries_per_row: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Copy dynamic ragged metadata into persistent CUDA graph buffers. - - FULL decode graphs capture kernel argument addresses. Keep the returned - tensors backed by stable storage, while indptr continues to bound reads. - """ - indptr_out = ragged_indptr_buffer[: num_rows + 1] - indptr_out.copy_(ragged_indptr, non_blocking=True) - - max_entries = max(num_rows * max_entries_per_row, 1) - ragged_out = ragged_indices_buffer[:max_entries] - nnz = ragged_indices.numel() - if nnz > 0: - ragged_out[:nnz].copy_(ragged_indices, non_blocking=True) - return ragged_out, indptr_out - - -@dataclass -class DeepseekV4ROCMAiterMLASparseMetadata(FlashMLASparseMetadata): - """ROCm-specific DeepSeek V4 metadata carrying ragged decode topk.""" - - c128a_decode_topk_ragged_indices: torch.Tensor | None = None - c128a_decode_topk_ragged_indptr: torch.Tensor | None = None - - -@dataclass -class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): - decode_swa_ragged_indices: torch.Tensor | None = None - decode_swa_ragged_indptr: torch.Tensor | None = None - - -class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None - self.c128a_decode_topk_ragged_indptr_buffer: torch.Tensor | None = None - if self.is_deepseek_v4 and self.compress_ratio == 128: - max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens - self.c128a_decode_topk_ragged_indices_buffer = torch.empty( - max_tokens * self.c128a_max_compressed, - dtype=torch.int32, - device=self.device, - ) - self.c128a_decode_topk_ragged_indptr_buffer = torch.empty( - max_tokens + 1, - dtype=torch.int32, - device=self.device, - ) - - def build( - self, - common_prefix_len: int, - common_attn_metadata: CommonAttentionMetadata, - fast_build: bool = False, - ) -> DeepseekV4ROCMAiterMLASparseMetadata: - base = super().build( - common_prefix_len=common_prefix_len, - common_attn_metadata=common_attn_metadata, - fast_build=fast_build, - ) - - ragged_indices = None - ragged_indptr = None - dense_decode = base.c128a_global_decode_topk_indices - decode_lens = base.c128a_decode_topk_lens - if dense_decode is not None and decode_lens is not None: - ragged_indices, ragged_indptr = build_ragged_indices_from_dense( - dense_decode.reshape(dense_decode.shape[0], -1), - decode_lens, - ) - assert self.c128a_decode_topk_ragged_indices_buffer is not None - assert self.c128a_decode_topk_ragged_indptr_buffer is not None - ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( - ragged_indices, - ragged_indptr, - self.c128a_decode_topk_ragged_indices_buffer, - self.c128a_decode_topk_ragged_indptr_buffer, - dense_decode.shape[0], - self.c128a_max_compressed, - ) - - return DeepseekV4ROCMAiterMLASparseMetadata( - **vars(base), - c128a_decode_topk_ragged_indices=ragged_indices, - c128a_decode_topk_ragged_indptr=ragged_indptr, - ) - - -class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens - self.decode_swa_ragged_indices_buffer = torch.empty( - max_tokens * self.window_size, - dtype=torch.int32, - device=self.device, - ) - self.decode_swa_ragged_indptr_buffer = torch.empty( - max_tokens + 1, - dtype=torch.int32, - device=self.device, - ) - - def build( - self, - common_prefix_len: int, - common_attn_metadata: CommonAttentionMetadata, - fast_build: bool = False, - ) -> DeepseekV4ROCMAiterSparseSWAMetadata: - base = super().build( - common_prefix_len=common_prefix_len, - common_attn_metadata=common_attn_metadata, - fast_build=fast_build, - ) - - ragged_indices = None - ragged_indptr = None - if ( - base.num_decode_tokens > 0 - and base.decode_swa_indices is not None - and base.decode_swa_lens is not None - ): - ragged_indices, ragged_indptr = build_ragged_indices_from_dense( - base.decode_swa_indices.reshape(base.num_decode_tokens, -1), - base.decode_swa_lens, - ) - ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( - ragged_indices, - ragged_indptr, - self.decode_swa_ragged_indices_buffer, - self.decode_swa_ragged_indptr_buffer, - base.num_decode_tokens, - self.window_size, - ) - - return DeepseekV4ROCMAiterSparseSWAMetadata( - **vars(base), - decode_swa_ragged_indices=ragged_indices, - decode_swa_ragged_indptr=ragged_indptr, - ) - - -class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend): - @staticmethod - def get_name() -> str: - return "ROCM_V4_FLASHMLA_SPARSE" - - @staticmethod - def get_builder_cls() -> type["DeepseekV4ROCMAiterMLASparseMetadataBuilder"]: - return DeepseekV4ROCMAiterMLASparseMetadataBuilder - - @staticmethod - def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]: - return DeepseekV4ROCMAiterMLASparseImpl - - -class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl): - """ROCm sparse MLA implementation used by DeepSeek V4's custom MLA layer.""" - - backend_cls = DeepseekV4ROCMAiterMLASparseBackend - - @classmethod - def get_padded_num_q_heads(cls, num_heads: int) -> int: - return num_heads - - @classmethod - def forward_mqa( # type: ignore[override] - cls, - layer: "DeepseekV4MLAAttention", - q: torch.Tensor, - kv: torch.Tensor, - positions: torch.Tensor, - output: torch.Tensor, - ) -> None: - assert output.shape == q.shape, ( - f"output buffer shape {output.shape} must match q shape {q.shape}" - ) - assert output.dtype == q.dtype, ( - f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" - ) - - forward_context = get_forward_context() - attn_metadata = forward_context.attn_metadata - - if attn_metadata is None: - # Warmup dummy run: no real metadata. Reserve the same bf16 - # gather workspace _forward_prefill would; the dequantize / topk - # / sparse_fwd kernels are skipped this step. - swa_only = layer.compress_ratio <= 1 - N = ( - 0 - if swa_only - else (layer.max_model_len + layer.compress_ratio - 1) - // layer.compress_ratio - ) - M = N + layer.window_size + layer.max_num_batched_tokens - current_workspace_manager().get_simultaneous( - ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), - ) - output.zero_() - return - - assert isinstance(attn_metadata, dict) - rocm_metadata = cast( - DeepseekV4ROCMAiterMLASparseMetadata | None, - attn_metadata.get(layer.prefix), - ) - swa_metadata = cast( - DeepseekV4ROCMAiterSparseSWAMetadata | None, - attn_metadata.get(layer.swa_cache_layer.prefix), - ) - assert swa_metadata is not None - - swa_only = layer.compress_ratio <= 1 - self_kv_cache = layer.kv_cache if not swa_only else None - swa_kv_cache = layer.swa_cache_layer.kv_cache - - num_decodes = swa_metadata.num_decodes - num_prefills = swa_metadata.num_prefills - num_decode_tokens = swa_metadata.num_decode_tokens - - if num_prefills > 0: - cls._forward_prefill( - layer=layer, - q=q[num_decode_tokens:], - positions=positions[num_decode_tokens:], - compressed_k_cache=self_kv_cache, - swa_k_cache=swa_kv_cache, - output=output[num_decode_tokens:], - attn_metadata=rocm_metadata, - swa_metadata=swa_metadata, - ) - if num_decodes > 0: - cls._forward_decode( - layer=layer, - q=q[:num_decode_tokens], - kv_cache=self_kv_cache, - swa_metadata=swa_metadata, - attn_metadata=rocm_metadata, - swa_only=swa_only, - output=output[:num_decode_tokens], - ) - - @classmethod - def _forward_decode( - cls, - layer: "DeepseekV4MLAAttention", - q: torch.Tensor, - kv_cache: torch.Tensor | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, - attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, - swa_only: bool, - output: torch.Tensor, - ) -> None: - num_decodes = swa_metadata.num_decodes - num_decode_tokens = swa_metadata.num_decode_tokens - - topk_indices = None - topk_lens = None - topk_ragged_indices = None - topk_ragged_indptr = None - if not swa_only: - assert attn_metadata is not None - assert swa_metadata.is_valid_token is not None - block_size = attn_metadata.block_size // layer.compress_ratio - is_valid = swa_metadata.is_valid_token[:num_decode_tokens] - if layer.compress_ratio == 4: - assert layer.topk_indices_buffer is not None - ( - topk_ragged_indices, - topk_ragged_indptr, - topk_lens, - ) = compute_global_topk_ragged_indices_and_indptr( - layer.topk_indices_buffer[:num_decode_tokens], - swa_metadata.token_to_req_indices, - attn_metadata.block_table[:num_decodes], - block_size, - is_valid, - ) - else: - topk_indices = attn_metadata.c128a_global_decode_topk_indices - topk_lens = attn_metadata.c128a_decode_topk_lens - topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices - topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr - - rocm_sparse_attn_decode( - q=q, - kv_cache=kv_cache, - swa_k_cache=layer.swa_cache_layer.kv_cache, - swa_only=swa_only, - topk_indices=topk_indices, - topk_lens=topk_lens, - swa_indices=swa_metadata.decode_swa_indices, - swa_lens=swa_metadata.decode_swa_lens, - swa_ragged_indices=swa_metadata.decode_swa_ragged_indices, - swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr, - topk_ragged_indices=topk_ragged_indices, - topk_ragged_indptr=topk_ragged_indptr, - attn_sink=layer.attn_sink, - scale=layer.scale, - head_dim=layer.head_dim, - nope_head_dim=layer.nope_head_dim, - rope_head_dim=layer.rope_head_dim, - output=output, - ) - - @classmethod - def _forward_prefill( - cls, - layer: "DeepseekV4MLAAttention", - q: torch.Tensor, - positions: torch.Tensor, - compressed_k_cache: torch.Tensor | None, - swa_k_cache: torch.Tensor, - output: torch.Tensor, - attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, - ) -> None: - swa_only = attn_metadata is None - - num_prefills = swa_metadata.num_prefills - num_prefill_tokens = swa_metadata.num_prefill_tokens - num_decodes = swa_metadata.num_decodes - num_decode_tokens = swa_metadata.num_decode_tokens - - seq_lens = swa_metadata.prefill_seq_lens - gather_lens = swa_metadata.prefill_gather_lens - assert seq_lens is not None - assert gather_lens is not None - - query_start_loc_cpu = swa_metadata.query_start_loc_cpu - query_start_loc = swa_metadata.query_start_loc - assert query_start_loc_cpu is not None - assert query_start_loc is not None - prefill_token_base = query_start_loc_cpu[num_decodes] - - if not swa_only: - if layer.compress_ratio == 4: - assert layer.topk_indices_buffer is not None - topk_indices = layer.topk_indices_buffer[num_decode_tokens:] - topk_indices = topk_indices[:num_prefill_tokens] - else: - assert attn_metadata is not None - topk_indices = attn_metadata.c128a_prefill_topk_indices - assert topk_indices is not None - top_k = topk_indices.shape[-1] - N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio - else: - assert layer.topk_indices_buffer is not None - topk_indices = layer.topk_indices_buffer[num_decode_tokens:] - top_k = 0 - N = 0 - - M = N + layer.window_size + layer.max_num_batched_tokens - num_chunks = (num_prefills + cls.PREFILL_CHUNK_SIZE - 1) // ( - cls.PREFILL_CHUNK_SIZE - ) - - workspace_manager = current_workspace_manager() - kv = workspace_manager.get_simultaneous( - ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), - )[0] - for chunk_idx in range(num_chunks): - chunk_start = chunk_idx * cls.PREFILL_CHUNK_SIZE - chunk_end = min(chunk_start + cls.PREFILL_CHUNK_SIZE, num_prefills) - chunk_size = chunk_end - chunk_start - if not swa_only: - assert attn_metadata is not None - assert compressed_k_cache is not None - block_table = attn_metadata.block_table[num_decodes:] - dequantize_and_gather_k_cache( - kv[:chunk_size], - compressed_k_cache, - seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio, - gather_lens=None, - block_table=block_table[chunk_start:chunk_end], - block_size=attn_metadata.block_size // layer.compress_ratio, - offset=0, - ) - - swa_block_table = swa_metadata.block_table[num_decodes:] - dequantize_and_gather_k_cache( - kv[:chunk_size], - swa_k_cache, - seq_lens=seq_lens[chunk_start:chunk_end], - gather_lens=gather_lens[chunk_start:chunk_end], - block_table=swa_block_table[chunk_start:chunk_end], - block_size=swa_metadata.block_size, - offset=N, - ) - - query_start = ( - query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base - ) - query_end = ( - query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base - ) - - combined_indices, combined_lens = combine_topk_swa_indices( - topk_indices[query_start:query_end], - query_start_loc[ - num_decodes + chunk_start : num_decodes + chunk_end + 1 - ], - seq_lens[chunk_start:chunk_end], - gather_lens[chunk_start:chunk_end], - layer.window_size, - layer.compress_ratio, - top_k, - M, - N, - ) - rocm_sparse_attn_prefill( - q=q[query_start:query_end], - kv=kv.view(-1, 1, q.shape[-1]), - indices=combined_indices, - topk_length=combined_lens, - scale=layer.scale, - head_dim=layer.head_dim, - nope_head_dim=layer.nope_head_dim, - rope_head_dim=layer.rope_head_dim, - attn_sink=layer.attn_sink, - output=output[query_start:query_end], - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/attention.py b/TEMP/deepseek_v4_ref/deepseek_v4/attention.py deleted file mode 100644 index 55cb3d94..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/attention.py +++ /dev/null @@ -1,806 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -DeepseekV4 MLA Attention Layer -""" - -from collections.abc import Callable -from typing import TYPE_CHECKING, Any, cast - -import torch -import torch.nn as nn -import torch.nn.functional as F -from transformers import DeepseekV2Config, DeepseekV3Config - -import vllm.envs as envs -from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.model_executor.layers.linear import ( - ReplicatedLinear, -) -from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer -from vllm.models.deepseek_v4.common.ops import ( - fused_indexer_q_rope_quant, - fused_inv_rope_fp8_quant, - fused_q_kv_rmsnorm, -) -from vllm.utils.deep_gemm import fp8_einsum -from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum - -if TYPE_CHECKING: - from vllm.v1.attention.backends.mla.sparse_swa import ( - DeepseekSparseSWAMetadata, - ) - -from vllm.config import ( - CacheConfig, - VllmConfig, - get_current_vllm_config, -) -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.forward_context import get_forward_context -from vllm.logger import init_logger -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.quantization.input_quant_fp8 import ( - QuantFP8, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - GroupShape, -) -from vllm.models.deepseek_v4.compressor import DeepseekCompressor -from vllm.platforms import current_platform -from vllm.utils.multi_stream_utils import ( - execute_in_parallel, - maybe_execute_in_parallel, -) -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata -from vllm.v1.attention.backends.mla.flashmla_sparse import ( - FlashMLASparseBackend, -) -from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV4IndexerBackend, - get_max_prefill_buffer_size, -) -from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache -from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec - -if TYPE_CHECKING: - from vllm.models.deepseek_v4.nvidia.flashmla import ( - DeepseekV4SparseMLAAttentionImpl, - ) - -logger = init_logger(__name__) - - -def _select_v4_sparse_impl() -> "type[DeepseekV4SparseMLAAttentionImpl]": - """Pick the platform-specific V4 sparse MLA impl class. Sole platform check.""" - if current_platform.is_rocm(): - from vllm.models.deepseek_v4.amd.rocm import ( - DeepseekV4ROCMAiterMLASparseImpl, - ) - - return DeepseekV4ROCMAiterMLASparseImpl - from vllm.models.deepseek_v4.nvidia.flashmla import ( - DeepseekV4FlashMLASparseImpl, - ) - - return DeepseekV4FlashMLASparseImpl - - -class DeepseekV4MLA(nn.Module): - def __init__( - self, - hidden_size: int, - num_heads: int, - head_dim: int, - scale: float, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - v_head_dim: int, - q_lora_rank: int | None, - kv_lora_rank: int, - o_lora_rank: int | None, - vllm_config: VllmConfig, - fused_wqa_wkv: torch.nn.Module, - q_norm: torch.nn.Module, - wq_b: torch.nn.Module, - kv_norm: torch.nn.Module, - wo_a: torch.nn.Module, - wo_b: torch.nn.Module, - attn_sink: torch.nn.Module, - rotary_emb: torch.nn.Module, - indexer: torch.nn.Module | None, - indexer_rotary_emb: torch.nn.Module, - topk_indices_buffer: torch.Tensor | None, - aux_stream_list: list[torch.cuda.Stream] | None, - window_size: int, - compress_ratio: int | None, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = hidden_size - self.n_local_heads = num_heads - self.head_dim = head_dim - self.scale = scale - - self.q_lora_rank = q_lora_rank - self.kv_lora_rank = kv_lora_rank - self.window_size = window_size - self.compress_ratio = compress_ratio if compress_ratio is not None else 1 - self.prefix = prefix - - # Extract config from vllm_config - config = vllm_config.model_config.hf_config - tp_size = get_tensor_model_parallel_world_size() - - # DeepseekV4-specific attributes (num_heads is already TP-adjusted) - self.eps = config.rms_norm_eps - self.rope_head_dim = config.qk_rope_head_dim - self.nope_head_dim = head_dim - self.rope_head_dim - self.n_local_groups = config.o_groups // tp_size - self.o_lora_rank = config.o_lora_rank - - # Store projection modules - self.fused_wqa_wkv = fused_wqa_wkv - self.q_norm = q_norm - self.wq_b = wq_b - - self.kv_norm = kv_norm - self.wo_a = wo_a - - self._wo_a_act_quant = QuantFP8( - static=False, - group_shape=GroupShape(1, 128), - use_ue8m0=True, - ) - # Bypass packed-for-deepgemm path — we need FP32 scales (not packed - # INT32) so fp8_einsum can handle layout transform internally. - self._wo_a_act_quant.use_deep_gemm_supported = False - self.wo_b = wo_b - - # Pick fp8_einsum recipe based on GPU arch: - # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128 - # SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1 - cap = current_platform.get_device_capability() - assert cap is not None, "DeepseekV4 attention requires a CUDA device" - self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) - self._tma_aligned_scales = cap.major >= 10 - - self.rotary_emb = rotary_emb - self.indexer_rotary_emb = indexer_rotary_emb - self.topk_indices_buffer = topk_indices_buffer - - self.indexer = indexer - - # Per-head RMS normalization for Q (no learnable weights) - self.q_head_norm = RMSNorm(head_dim, eps=self.eps, has_weight=False) - - # TODO(yifan): currently hardcoded for FP8 sparse, make it more generic - head_bytes = ( - self.nope_head_dim # 448 fp8 NoPE - + self.rope_head_dim * 2 # 64 bf16 RoPE - + self.nope_head_dim // 64 # 7B scale factors - + 1 # 1B pad - ) - - # Will be None on ROCm for now. - self.aux_stream_list = aux_stream_list - # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; - # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins - # before post-GEMM starts. - self.ln_events = [torch.cuda.Event() for _ in range(4)] - - assert cache_config is not None, "DeepseekV4 attention requires cache_config" - self.swa_cache_layer = DeepseekV4SWACache( - head_dim=self.head_dim, - window_size=self.window_size, - dtype=torch.uint8, - prefix=f"{prefix}.swa_cache", - cache_config=cache_config, - ) - - self.mla_attn = DeepseekV4MLAAttention( - num_heads=self.n_local_heads, - head_dim=self.head_dim, - scale=self.scale, - qk_nope_head_dim=self.nope_head_dim, - qk_rope_head_dim=self.rope_head_dim, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.kv_lora_rank, - compress_ratio=self.compress_ratio, - window_size=self.window_size, - head_bytes=head_bytes, - swa_cache_layer=self.swa_cache_layer, - attn_sink=attn_sink, # already padded with -inf - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - indexer=self.indexer, - topk_indices_buffer=self.topk_indices_buffer, - ) - # Mirror the inner layer's padded head count (single source of truth). - self.padded_heads = self.mla_attn.padded_heads - - # Create the compressor for layers with compress_ratio > 1; after - # creating the DeepseekV4MLAAttention layer to get its cache. - self.compressor = None - if self.compress_ratio > 1: - self.compressor = DeepseekCompressor( - vllm_config=vllm_config, - compress_ratio=self.compress_ratio, - hidden_size=self.hidden_size, - head_dim=self.head_dim, - rotate=True, - prefix=f"{prefix}.compressor", - k_cache_prefix=self.mla_attn.prefix, - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - llama_4_scaling: torch.Tensor | None = None, - ) -> torch.Tensor: - # Pre-allocate attention output with FlashMLA-padded head count. - # The op writes into `o_padded`; we slice to n_local_heads after. - num_tokens = hidden_states.shape[0] - o_padded = torch.empty( - (num_tokens, self.padded_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - - # attention_impl is wrapped with @eager_break_during_capture: this is - # where the breakable cudagraph capture breaks (the attention op runs - # eagerly between captured graph segments). - self.attention_impl(hidden_states, positions, o_padded) - o = o_padded[:, : self.n_local_heads, :] - - # Keep ROCm on the BF16 reference wo_a path util kernel ready. - if current_platform.is_rocm(): - z = rocm_inv_rope_einsum( - self.rotary_emb, - o, - positions, - self.rope_head_dim, - self.n_local_groups, - self.o_lora_rank, - self.wo_a, - ) - return self.wo_b(z.flatten(1)) - - # O projection: inverse RoPE + FP8 quant + einsum + wo_b - o_fp8, o_scale = fused_inv_rope_fp8_quant( - o, - positions, - self.rotary_emb.cos_sin_cache, - n_groups=self.n_local_groups, - heads_per_group=self.n_local_heads // self.n_local_groups, - nope_dim=self.nope_head_dim, - rope_dim=self.rope_head_dim, - tma_aligned_scales=self._tma_aligned_scales, - ) - - wo_a_fp8 = self.wo_a.weight - wo_a_scale = self.wo_a.weight_scale_inv - - z = torch.empty( - (num_tokens, self.n_local_groups, self.o_lora_rank), - device=o.device, - dtype=torch.bfloat16, - ) - fp8_einsum( - "bhr,hdr->bhd", - (o_fp8, o_scale), - (wo_a_fp8, wo_a_scale), - z, - recipe=self._einsum_recipe, - ) - - return self.wo_b(z.flatten(1)) - - def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: - aux_streams = self.aux_stream_list - if aux_streams is not None: - assert len(aux_streams) >= 3 - aux_streams = aux_streams[:3] - - # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs - # on aux streams 0..2 when their owning module exists. ln_events[0] - # is the fan-out start event; ln_events[1..3] are per-aux done events. - # On ROCm, aux_streams is None and execute_in_parallel runs serially. - aux_fns: list[Callable[[], Any] | None] = [None, None, None] - - if self.compressor is not None: - # Local ref so the closure keeps a non-None type for mypy. - compressor = self.compressor - - def compressor_kv_score() -> torch.Tensor: - return torch.mm( - hidden_states, - compressor.fused_wkv_wgate.weight.T, - out_dtype=torch.float32, - ) - - aux_fns[0] = compressor_kv_score - - if self.indexer is not None: - indexer = self.indexer - - def indexer_weights_proj() -> torch.Tensor: - # ReplicatedLinear returns (output, bias); bias is None. - weights, _ = indexer.weights_proj(hidden_states) - return weights - - def indexer_compressor_kv_score() -> torch.Tensor: - return torch.mm( - hidden_states, - indexer.compressor.fused_wkv_wgate.weight.T, - out_dtype=torch.float32, - ) - - aux_fns[1] = indexer_weights_proj - aux_fns[2] = indexer_compressor_kv_score - - def fused_wqa_wkv() -> torch.Tensor: - # MergedColumnParallelLinear returns (output, bias); bias is None. - qr_kv, _ = self.fused_wqa_wkv(hidden_states) - return qr_kv - - qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel( - fused_wqa_wkv, - aux_fns, - self.ln_events[0], - self.ln_events[1:4], - aux_streams, - enable=hidden_states.shape[0] - <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, - ) - - return qr_kv, kv_score, indexer_kv_score, indexer_weights - - @eager_break_during_capture - def attention_impl( - self, - hidden_states: torch.Tensor, - positions: torch.Tensor, - out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place - ) -> None: - forward_context = get_forward_context() - attn_metadata = forward_context.attn_metadata - - qr_kv, kv_score, indexer_kv_score, indexer_weights = ( - self.attn_gemm_parallel_execute(hidden_states) - ) - - qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) - qr, kv = fused_q_kv_rmsnorm( - qr, - kv, - self.q_norm.weight.data, - self.kv_norm.weight.data, - self.eps, - ) - - # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride - # on the default stream so q stays on its consumer stream (mla_attn - # downstream reads q on default). Indexer/compressor go on aux for - # overlap with default's GEMM + cache write. - if self.indexer is not None: - aux_streams = self.aux_stream_list - indexer = self.indexer - # Local ref so the closure keeps a non-None type for mypy. - assert self.compressor is not None - compressor = self.compressor - - def wq_b_kv_insert() -> torch.Tensor: - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) - return q - - # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs - # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the - # MLA compressor. Slot [2] is reserved for the indexer's inner - # overlap. ROCm (aux_streams is None) falls back to sequential. - q, _ = execute_in_parallel( - wq_b_kv_insert, - [ - lambda: indexer( - hidden_states, - qr, - indexer_kv_score, - indexer_weights, - positions, - self.indexer_rotary_emb, - ), - lambda: compressor(kv_score, positions, self.rotary_emb), - ], - self.ln_events[0], - [self.ln_events[1], self.ln_events[2]], - [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, - enable=aux_streams is not None, - ) - elif self.compressor is not None: - # wq_b + kv_insert on default, compressor on aux. - aux_stream = ( - self.aux_stream_list[0] if self.aux_stream_list is not None else None - ) - compressor = self.compressor - - def wq_b_kv_insert() -> torch.Tensor: - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) - return q - - q, _ = maybe_execute_in_parallel( - wq_b_kv_insert, - lambda: compressor(kv_score, positions, self.rotary_emb), - self.ln_events[0], - self.ln_events[1], - aux_stream, - ) - else: - # SWA-only layer: no compressor, no overlap. - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) - - # MLA attention writes into the pre-allocated `out` buffer - # ([num_tokens, padded_heads, head_dim]). - self.mla_attn(q, kv, positions, output=out) - - def _fused_qnorm_rope_kv_insert( - self, - q: torch.Tensor, - kv: torch.Tensor, - positions: torch.Tensor, - attn_metadata: ( - dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None - ), - ) -> torch.Tensor: - if not isinstance(attn_metadata, dict): - # Profile run: kernel doesn't fire; produce a padded tensor so - # downstream FlashMLA gets the right shape. - if self.n_local_heads < self.padded_heads: - return F.pad( - q, - (0, 0, 0, self.padded_heads - self.n_local_heads), - value=0.0, - ) - return q - - swa_metadata = cast( - "DeepseekSparseSWAMetadata | None", - attn_metadata.get(self.swa_cache_layer.prefix), - ) - assert swa_metadata is not None - - swa_kv_cache = self.swa_cache_layer.kv_cache - swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) - - # Horizontally fused: - # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE, - # with zero-fill for the padding head slots. The kernel - # allocates and returns the padded q tensor. - # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert - # kv is unchanged; mla_attn reads kv solely via swa_kv_cache. - return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - q, - kv, - swa_kv_cache_2d, - swa_metadata.slot_mapping, - positions.to(torch.int64), - self.rotary_emb.cos_sin_cache, - self.padded_heads, - self.eps, - swa_metadata.block_size, - ) - - -class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): - def __init__( - self, - num_heads: int, - head_dim: int, - scale: float, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - q_lora_rank: int | None, - kv_lora_rank: int, - compress_ratio: int, - window_size: int, - head_bytes: int, - swa_cache_layer: DeepseekV4SWACache, - attn_sink: torch.Tensor, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - # Sparse MLA Args - indexer: object | None = None, - topk_indices_buffer: torch.Tensor | None = None, - aux_stream: torch.cuda.Stream | None = None, - **extra_impl_args, - ) -> None: - super().__init__() - self.impl_cls = _select_v4_sparse_impl() - self.backend_cls = self.impl_cls.backend_cls - self.num_heads = num_heads - self.num_kv_heads = 1 - self.head_dim = head_dim - self.scale = scale - self.window_size = window_size - self.head_bytes = head_bytes - self.compress_ratio = compress_ratio - self.q_lora_rank = q_lora_rank - self.kv_lora_rank = kv_lora_rank - self.nope_head_dim = qk_nope_head_dim - self.rope_head_dim = qk_rope_head_dim - self.indexer = indexer - self.topk_indices_buffer = topk_indices_buffer - - self.prefix = prefix # Alias for compatibility with compressor - - self.aux_stream = aux_stream - self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] - - # Padded Q head count is dictated by the selected impl. - self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads) - - # Store attention sink - assert attn_sink is not None - self.attn_sink: torch.Tensor = attn_sink - # Store SWA cache - assert swa_cache_layer is not None - self.swa_cache_layer: DeepseekV4SWACache = swa_cache_layer - - # Get vllm config for cache setup - vllm_config = get_current_vllm_config() - self.max_num_batched_tokens = ( - vllm_config.scheduler_config.max_num_batched_tokens - ) - self.max_model_len = vllm_config.model_config.max_model_len - # DeepseekV4 only supports fp8 kv-cache format for now. - kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8" - - assert kv_cache_dtype.startswith("fp8"), ( - f"DeepseekV4 only supports fp8 kv-cache format for now, " - f"got {kv_cache_dtype}" - ) - assert issubclass(self.get_attn_backend(), FlashMLASparseBackend), ( - "Only FlashMLA Sparse Attention backend is supported for DeepseekV4 for now" - ) - # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format - # Automatically convert fp8 kv-cache format to "fp8_ds_mla" - if ( - issubclass(self.get_attn_backend(), FlashMLASparseBackend) - and kv_cache_dtype.startswith("fp8") - and kv_cache_dtype != "fp8_ds_mla" - ): - assert cache_config is not None - cache_config.cache_dtype = "fp8_ds_mla" - kv_cache_dtype = "fp8_ds_mla" - logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") - - self.kv_cache_dtype = kv_cache_dtype - - # Register with compilation context for metadata lookup - compilation_config = vllm_config.compilation_config - if prefix and prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - if prefix: - compilation_config.static_forward_context[prefix] = self - - self.kv_cache = torch.tensor([]) - - def get_attn_backend(self) -> type[AttentionBackend]: - return self.backend_cls - - def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: - if ( - self.compress_ratio <= 1 - ): # SWA part. Allocated separately as DeepseekV4SWACache. - return None - return MLAAttentionSpec( - block_size=vllm_config.cache_config.block_size, - num_kv_heads=1, - head_size=self.head_dim, - dtype=torch.uint8, - compress_ratio=self.compress_ratio, - cache_dtype_str=self.kv_cache_dtype, - alignment=576, # NOTE: FlashMLA requires 576B alignment - model_version="deepseek_v4", - ) - - def forward( - self, - q: torch.Tensor, - kv: torch.Tensor, - positions: torch.Tensor, - output: torch.Tensor, - ) -> None: - self.impl_cls.forward_mqa(self, q, kv, positions, output) - - -class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): - def __init__( - self, - head_dim: int, - dtype: torch.dtype, - prefix: str, - cache_config: CacheConfig, - compress_ratio: int = 1, - ): - super().__init__() - self.kv_cache = torch.tensor([]) - self.head_dim = head_dim - self.prefix = prefix - self.cache_config = cache_config - self.dtype = dtype - self.compress_ratio = compress_ratio - compilation_config = get_current_vllm_config().compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - - def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # head_dim already carries the fp8 scale padding - # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. - return MLAAttentionSpec( - block_size=self.cache_config.block_size, - num_kv_heads=1, - head_size=self.head_dim, - dtype=self.dtype, - compress_ratio=self.compress_ratio, - # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with - # the indexer's compressor state cache. V3.2 keeps the legacy layout. - alignment=576, - ) - - def forward(self): ... - - def get_attn_backend(self) -> type[AttentionBackend]: - return DeepseekV4IndexerBackend - - -class DeepseekV4Indexer(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - config: DeepseekV2Config | DeepseekV3Config, - hidden_size: int, - q_lora_rank: int, - quant_config: QuantizationConfig | None, - cache_config: CacheConfig | None, - topk_indices_buffer: torch.Tensor | None, - compress_ratio: int = 1, - prefix: str = "", - aux_stream: torch.cuda.Stream | None = None, - ): - super().__init__() - self.vllm_config = vllm_config - self.config = config - self.quant_config = quant_config - # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] - self.topk_tokens = config.index_topk - self.n_head = config.index_n_heads # 64 - self.head_dim = config.index_head_dim # 128 - self.rope_dim = config.qk_rope_head_dim # 64 - self.q_lora_rank = q_lora_rank # 1536 - self.compress_ratio = compress_ratio - self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache - logger.info_once( - "Using %s indexer cache for Lightning Indexer.", - "MXFP4" if self.use_fp4_kv else "FP8", - ) - - # no tensor parallel, just replicated - self.wq_b = ReplicatedLinear( - self.q_lora_rank, - self.head_dim * self.n_head, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.wq_b", - ) - self.weights_proj = ReplicatedLinear( - hidden_size, - self.n_head, - bias=False, - quant_config=None, - prefix=f"{prefix}.weights_proj", - ) - self.softmax_scale = self.head_dim**-0.5 - - self.scale_fmt = "ue8m0" - self.quant_block_size = 128 # TODO: get from config - self.topk_indices_buffer = topk_indices_buffer - - self.max_model_len = ( - vllm_config.model_config.max_model_len // self.compress_ratio - ) - self.prefix = prefix - - self.max_total_seq_len = ( - get_max_prefill_buffer_size(vllm_config) // self.compress_ratio - ) - - assert cache_config is not None, "Deepseek V4 indexer requires cache_config" - # NOTE(yifan): FP8 indxer cache use the same layout as V3.2: - # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. - # For FP4 indexer cache, we still allocate the same amount of memory as FP8, - # but only use the first half of the memory. - k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4 - self.k_cache = DeepseekV4IndexerCache( - head_dim=k_cache_head_dim, - dtype=torch.uint8, - prefix=f"{prefix}.k_cache", - cache_config=cache_config, - compress_ratio=self.compress_ratio, - ) - self.compressor = DeepseekCompressor( - vllm_config=vllm_config, - compress_ratio=self.compress_ratio, - hidden_size=hidden_size, - head_dim=self.head_dim, - rotate=True, - prefix=f"{prefix}.compressor", - k_cache_prefix=self.k_cache.prefix, - use_fp4_cache=self.use_fp4_kv, - ) - - self.indexer_op = SparseAttnIndexer( - self.k_cache, - self.quant_block_size, - self.scale_fmt, - self.topk_tokens, - self.head_dim, - self.max_model_len, - self.max_total_seq_len, - self.topk_indices_buffer, - skip_k_cache_insert=True, - use_fp4_cache=self.use_fp4_kv, - ) - - # None on ROCm — maybe_execute_in_parallel falls back to sequential. - self.aux_stream = aux_stream - self.ln_events: list[torch.cuda.Event] = [ - torch.cuda.Event(), - torch.cuda.Event(), - ] - - def forward( - self, - hidden_states: torch.Tensor, - qr: torch.Tensor, - compressed_kv_score: torch.Tensor, - indexer_weights: torch.Tensor, - positions: torch.Tensor, - rotary_emb: nn.Module, - ) -> torch.Tensor: - compressor = self.compressor - - def wq_b_and_q_quant(): - # ReplicatedLinear returns (output, bias); bias is None. - q, _ = self.wq_b(qr) - q = q.view(-1, self.n_head, self.head_dim) - return fused_indexer_q_rope_quant( - positions, - q, - rotary_emb.cos_sin_cache, - indexer_weights, - self.softmax_scale, - self.n_head**-0.5, - use_fp4=self.use_fp4_kv, - ) - - # compressor returns None and writes K to the indexer KV cache; the - # join orders that write before indexer_op (skip_k_cache_insert=True). - (q_quant, weights), k = maybe_execute_in_parallel( - wq_b_and_q_quant, - lambda: compressor(compressed_kv_score, positions, rotary_emb), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) - return self.indexer_op(hidden_states, q_quant, k, weights) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py deleted file mode 100644 index 208f01a7..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py deleted file mode 100644 index dc6f3c60..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from .cache_utils import ( - combine_topk_swa_indices, - compute_global_topk_indices_and_lens, - dequantize_and_gather_k_cache, - quantize_and_insert_k_cache, -) -from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant -from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant -from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm -from .fused_qk_rmsnorm import fused_q_kv_rmsnorm -from .save_partial_states import save_partial_states - -__all__ = [ - "MXFP4_BLOCK_SIZE", - "combine_topk_swa_indices", - "compute_global_topk_indices_and_lens", - "dequantize_and_gather_k_cache", - "fused_indexer_q_rope_quant", - "fused_inv_rope_fp8_quant", - "fused_mtp_input_rmsnorm", - "fused_q_kv_rmsnorm", - "mtp_shared_head_rmsnorm", - "quantize_and_insert_k_cache", - "save_partial_states", -] diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py deleted file mode 100644 index ac66751e..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py +++ /dev/null @@ -1,594 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index -preparation. - -- quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into - the paged cache. -- dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged - cache for sparse/SWA prefill. -- compute_global_topk_indices_and_lens: map local topk indices to global KV - cache slots and count valid entries. -- combine_topk_swa_indices: concatenate topk compressed indices with SWA - window indices for sparse prefill. -""" - -import torch - -from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_cutedsl - - -@triton.jit -def quantize_and_insert_k_kernel( - # Input tensors - k_ptr, # [num_tokens, 512] bf16 - slot_mapping_ptr, # [num_tokens] int64 - # Output tensor - k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) - # Dimensions - num_tokens, - input_dim: tl.constexpr, # 512 - fp8_dim: tl.constexpr, # 448 - bf16_dim: tl.constexpr, # 64 - scale_dim: tl.constexpr, # 8 - quant_block: tl.constexpr, # 64 (quantization block size) - cache_block_size: tl.constexpr, # 64 (paged cache block size) - token_data_size: tl.constexpr, # 576 bytes per token data - block_stride: tl.constexpr, # total bytes per block (padded) - fp8_max: tl.constexpr, - n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding) -): - """ - Quantize K tensor and insert into paged K cache. - - K Cache block layout (block_size=64 tokens): - - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16 - - [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales - - [64*576 + 64*8, block_stride): Padding - - One program per token. - """ - pid = tl.program_id(0) - - if pid >= num_tokens: - return - - # Get slot mapping - slot_idx = tl.load(slot_mapping_ptr + pid) - if slot_idx == -1: - return - - block_idx = slot_idx // cache_block_size - pos_in_block = slot_idx % cache_block_size - - # Input pointer for this token - input_row_ptr = k_ptr + pid * input_dim - - # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks - # (e.g. >= 57K at block_stride ~37K). Matches gather path below. - cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride - - # Token data pointer: token data is stored contiguously at start of block - # Each token's data is at offset pos_in_block * token_data_size - token_data_ptr = cache_block_ptr + pos_in_block * token_data_size - - # Scale pointer: scales are stored after ALL token data in the block - # Scale for this token is at offset (64 * 576) + pos_in_block * 8 - token_scale_ptr = ( - cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim - ) - - # Token data layout: [0:448] fp8, [448:576] bf16 - token_fp8_ptr = token_data_ptr - token_bf16_ptr = token_data_ptr + fp8_dim - - # ========== Quantize and store FP8 portion (first 448 elements) ========== - # Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent) - for qblock_idx in tl.static_range(n_quant_blocks): - qblock_start = qblock_idx * quant_block - - if qblock_start < fp8_dim: - offsets = qblock_start + tl.arange(0, quant_block) - mask = offsets < fp8_dim - - # Load bf16 input - x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0) - - # Compute absmax scale (same as CUDA kernel) - abs_x = tl.abs(x) - block_max = tl.max(abs_x, axis=0) - block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4) - - # UE8M0: Round scale UP to next power of 2 - # scale = 2^ceil(log2(block_max / fp8_max)) - raw_scale = block_max / fp8_max - log_scale = tl.log2(raw_scale) - exponent = tl.ceil(log_scale) # Round UP to next integer exponent - scale = tl.exp2(exponent) # scale = 2^exponent (power of 2) - - # Quantize to fp8: fp8_value = bf16_value / scale - x_scaled = x / scale - x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max) - - # Convert to fp8, then bitcast to uint8 for storage - x_fp8 = x_clamped.to(tl.float8e4nv) - x_uint8 = x_fp8.to(tl.uint8, bitcast=True) - - # Store as uint8 (1 byte each) - tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask) - - # UE8M0 scale encoding: stored_value = exponent + 127 (bias) - # During dequant: scale = 2^(stored_value - 127) - encoded_scale = exponent + 127.0 - encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0) - tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8)) - - # Padding scale at index 7 - tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8)) - - # ========== Store BF16 portion (last 64 elements, no quantization) ========== - bf16_input_offset = fp8_dim - - # Process bf16 in chunks of 16 - bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) - for i in tl.static_range(bf16_dim // 16): - chunk_offsets = i * 16 + tl.arange(0, 16) - bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets) - tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) - - -def quantize_and_insert_k_cache( - k: torch.Tensor, # [num_tokens, 512] bf16 - k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 - slot_mapping: torch.Tensor, # [num_tokens] int64 - block_size: int = 64, - is_ue8m0: bool = True, -): - """ - Quantize K tensor and insert into paged K cache. - - K Cache block layout (block_size=64 tokens): - - First 64 * 576 = 36864 bytes: Token data - - Each token: 448 bytes (fp8) + 128 bytes (bf16) - - Next 64 * 8 = 512 bytes: Scales - - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) - - Padded to multiple of 576 - """ - assert k.dim() == 2 and k.shape[1] == 512, ( - f"K must be [num_tokens, 512], got {k.shape}" - ) - assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}" - assert is_ue8m0, "Only support ue8m0 quantization." - - # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to - # padding. Always use slot_mapping.shape[0] as the token count. - num_tokens = slot_mapping.shape[0] - block_stride = k_cache.stride(0) # bytes per block - - TOKEN_FP8_DIM = 448 - TOKEN_BF16_DIM = 64 - TOKEN_SCALE_DIM = 8 - QUANT_BLOCK_SIZE = 64 - FP8_MAX = 448.0 - TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 - - grid = (num_tokens,) - - quantize_and_insert_k_kernel[grid]( - k, - slot_mapping, - k_cache, - num_tokens, - input_dim=512, - fp8_dim=TOKEN_FP8_DIM, - bf16_dim=TOKEN_BF16_DIM, - scale_dim=TOKEN_SCALE_DIM, - quant_block=QUANT_BLOCK_SIZE, - cache_block_size=block_size, - token_data_size=TOKEN_DATA_SIZE, - block_stride=block_stride, - fp8_max=FP8_MAX, - n_quant_blocks=8, - ) - - -@triton.jit -def _dequantize_and_gather_k_kernel( - out_ptr, - out_stride0, - out_stride1, - k_cache_ptr, - seq_lens_ptr, - block_table_ptr, - offset, - gather_lens_ptr, - # Constants - max_blocks_per_seq: tl.constexpr, - fp8_dim: tl.constexpr, # 448 - bf16_dim: tl.constexpr, # 64 - scale_dim: tl.constexpr, # 8 - quant_block: tl.constexpr, # 64 (quantization block size) - cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size) - token_data_size: tl.constexpr, # 576 bytes per token data - block_stride: tl.constexpr, # total bytes per block (padded) int32 - output_dim: tl.constexpr, # 512 - fp8_max: tl.constexpr, - n_quant_blocks: tl.constexpr, # 7 real blocks -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) - - seq_len = tl.load(seq_lens_ptr + batch_idx) - if gather_lens_ptr is not None: # noqa: SIM108 - gather_len = tl.load(gather_lens_ptr + batch_idx) - else: - # Gather all tokens - gather_len = seq_len - start_pos = seq_len - gather_len - - for i in range(worker_id, gather_len, num_workers): - # Calculate the actual token index in the sequence - pos = start_pos + i - - # Calculate which block and position within block - block_in_seq = pos // cache_block_size - pos_in_block = pos % cache_block_size - - # Get physical block index from block table - block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq - physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 - - # int64: physical_block_idx * block_stride can exceed 2^31 with many - # KV-cache blocks (e.g. >= 57K at block_stride ~37K). - cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride - - # Token data pointer - token_data_ptr = cache_block_ptr + pos_in_block * token_data_size - - # Scale pointer: after all token data - token_scale_ptr = ( - cache_block_ptr - + cache_block_size * token_data_size - + pos_in_block * scale_dim - ) - - # Token data layout: [0:448] fp8, [448:576] bf16 - token_fp8_ptr = token_data_ptr - token_bf16_ptr = token_data_ptr + fp8_dim - - # Output pointer for this token (flattened) - output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 - - # ========== Dequantize FP8 portion using UE8M0 ========== - for qblock_idx in tl.static_range(n_quant_blocks): - qblock_start = qblock_idx * quant_block - - if qblock_start < fp8_dim: - offsets = qblock_start + tl.arange(0, quant_block) - mask = offsets < fp8_dim - - # Load quantized fp8 values (stored as uint8) - x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0) - - # Bitcast uint8 back to fp8 - x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) - - # Convert fp8 to float32 for computation - x_float = x_fp8.to(tl.float32) - - # Load and decode UE8M0 scale - # UE8M0: scale = 2^(stored_value - 127) - encoded_scale = tl.load(token_scale_ptr + qblock_idx) - exponent = encoded_scale.to(tl.float32) - 127.0 - scale = tl.exp2(exponent) - - # Dequantize: bf16_value = fp8_value * scale - x_dequant = x_float * scale - - # Store as bf16 - tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask) - - # ========== Copy BF16 portion directly ========== - bf16_output_offset = fp8_dim # After 448 elements in output - - # Read bf16 from cache - bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) - - # Process in chunks of 16 - for j in tl.static_range(bf16_dim // 16): - chunk_offsets = j * 16 + tl.arange(0, 16) - bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets) - tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) - - -def dequantize_and_gather_k_cache_triton( - # [num_reqs, max_num_tokens, head_size] - out: torch.Tensor, - # [num_blocks, block_size, head_bytes] - k_cache: torch.Tensor, - # [num_reqs] - seq_lens: torch.Tensor, - # [num_reqs] - gather_lens: torch.Tensor | None, - # [num_reqs, max_blocks_per_seq] - block_table: torch.Tensor, - block_size: int, - offset: int, -) -> None: - TOKEN_FP8_DIM = 448 - TOKEN_BF16_DIM = 64 - TOKEN_SCALE_DIM = 8 - QUANT_BLOCK_SIZE = 64 - FP8_MAX = 448.0 - TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 - - num_reqs = seq_lens.shape[0] - NUM_WORKERS = 128 - _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( - out, - out.stride(0), - out.stride(1), - k_cache, - seq_lens, - block_table, - offset, - gather_lens, - max_blocks_per_seq=block_table.shape[-1], - fp8_dim=TOKEN_FP8_DIM, - bf16_dim=TOKEN_BF16_DIM, - scale_dim=TOKEN_SCALE_DIM, - quant_block=QUANT_BLOCK_SIZE, - cache_block_size=block_size, - token_data_size=TOKEN_DATA_SIZE, - block_stride=k_cache.stride(0), - output_dim=512, - fp8_max=FP8_MAX, - n_quant_blocks=7, - ) - - -def dequantize_and_gather_k_cache( - # [num_reqs, max_num_tokens, head_size] - out: torch.Tensor, - # [num_blocks, block_size, head_bytes] - k_cache: torch.Tensor, - # [num_reqs] - seq_lens: torch.Tensor, - # [num_reqs] - gather_lens: torch.Tensor | None, - # [num_reqs, max_blocks_per_seq] - block_table: torch.Tensor, - block_size: int, - offset: int, -) -> None: - if has_cutedsl(): - # lazily import, otherwise some tests fail due to CUDA driver init failure. - from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( - dequantize_and_gather_k_cache_cutedsl, - ) - - dequantize_and_gather_k_cache_cutedsl( - out, k_cache, seq_lens, gather_lens, block_table, block_size, offset - ) - return - - dequantize_and_gather_k_cache_triton( - out, k_cache, seq_lens, gather_lens, block_table, block_size, offset - ) - - -def compute_global_topk_indices_and_lens( - topk_indices: torch.Tensor, - token_to_req_indices: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - is_valid_token: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Map local topk indices to global KV cache slots and count valid entries. - - Fuses three operations into a single kernel: - 1. Block-table lookup (local index → global slot id) - 2. Valid-entry counting (topk_lens per token) - 3. Masking padding tokens to length 0 - """ - num_tokens = topk_indices.shape[0] - global_topk_indices = torch.empty_like(topk_indices) - topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) - _compute_global_topk_indices_and_lens_kernel[(num_tokens,)]( - global_topk_indices, - global_topk_indices.stride(0), - topk_lens, - topk_indices, - topk_indices.stride(0), - topk_indices.shape[-1], - token_to_req_indices, - block_table, - block_table.stride(0), - block_size, - is_valid_token, - TRITON_BLOCK_SIZE=1024, - ) - return global_topk_indices, topk_lens - - -@triton.jit -def _compute_global_topk_indices_and_lens_kernel( - global_topk_indices_ptr, - global_topk_indices_stride, - topk_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - topk, - token_to_req_indices_ptr, - block_table_ptr, - block_table_stride, - block_size, - is_valid_token_ptr, - TRITON_BLOCK_SIZE: tl.constexpr, -): - token_idx = tl.program_id(0) - is_valid_token = tl.load(is_valid_token_ptr + token_idx) - req_idx = tl.load(token_to_req_indices_ptr + token_idx) - - count = tl.zeros((), dtype=tl.int32) - for i in range(0, topk, TRITON_BLOCK_SIZE): - offset = i + tl.arange(0, TRITON_BLOCK_SIZE) - mask = offset < topk - - local_idx = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=mask, - other=-1, - ) - is_valid = local_idx >= 0 - - block_indices = local_idx // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask & is_valid, - ) - block_offsets = local_idx % block_size - - slot_ids = block_numbers * block_size + block_offsets - slot_ids = tl.where(is_valid, slot_ids, -1) - tl.store( - global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset, - slot_ids, - mask=mask, - ) - count += tl.sum(is_valid.to(tl.int32), axis=0) - - # Zero out length for padding tokens. - tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) - - -# FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see -# flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is -# 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both. -# The extra slots stay as -1 sentinels and `combined_lens` caps the valid -# range via `topk_length`, so padding is a no-op at kernel level. -_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 - - -def combine_topk_swa_indices( - topk_indices: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor, - window_size: int, - compress_ratio: int, - topk: int, - M: int, - N: int, -) -> tuple[torch.Tensor, torch.Tensor]: - num_tokens = topk_indices.shape[0] - num_reqs = seq_lens.shape[0] - combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) - // _SPARSE_PREFILL_TOPK_ALIGNMENT - * _SPARSE_PREFILL_TOPK_ALIGNMENT - ) - combined_indices = torch.full( - (num_tokens, combined_topk), - fill_value=-1, - dtype=torch.int32, - device=topk_indices.device, - ) - combined_lens = torch.empty( - num_tokens, dtype=torch.int32, device=topk_indices.device - ) - - NUM_WORKERS = 128 - _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)]( - combined_indices, - combined_indices.stride(0), - combined_lens, - topk_indices, - topk_indices.stride(0), - query_start_loc, - seq_lens, - gather_lens, - M, - N, - TOP_K=topk, - COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=window_size, - PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), - ) - return combined_indices, combined_lens - - -@triton.jit -def _combine_topk_swa_indices_kernel( - combined_indices_ptr, - combined_indices_stride, - combined_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - PADDED_TOP_K: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) - - # query_start_loc is a global tensor; rebase to chunk-local offsets - # by subtracting the chunk's starting value. - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - # The SWA portion of the gathered buffer starts from position - # (seq_len - gather_len), not position 0. We need this offset - # to correctly index into the gathered buffer. - gather_start = seq_len - gather_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - # topk_len is fully determined by the query token's absolute position: - # both the C4A indexer and the C128A metadata builder emit - # min((pos + 1) // compress_ratio, topk_tokens) valid entries. - # Caller passes TOP_K=0 for SWA-only layers to zero this out. - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - - offset = tl.arange(0, PADDED_TOP_K) - mask = offset < topk_len - topk_indices = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=mask, - ) - tl.store( - combined_indices_ptr + token_idx * combined_indices_stride + offset, - topk_indices + M * batch_idx, - mask=mask, - ) - offset = tl.arange(0, WINDOW_SIZE) - # Index into gathered buffer: N + (position - gather_start) - # For positions [pos - swa_len + 1, pos], the buffer indices are: - # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + offset, - M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, - mask=offset < swa_len, - ) - - combined_len = topk_len + swa_len - tl.store(combined_lens_ptr + token_idx, combined_len) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py deleted file mode 100644 index 9a5e478e..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ /dev/null @@ -1,666 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels. - -Three specialized kernels: - - _fused_kv_compress_norm_rope_insert_sparse_attn: - head=512, nope=448 FP8 + rope=64 bf16 - - _fused_kv_compress_norm_rope_insert_indexer_attn: - head=128, all FP8, 1 block/token - - _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn: - head=128, MXFP4 (block=32), 4 ue8m0 bytes - -RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the -even/odd halves are consumed directly for MXFP4, no interleave needed). -FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for -per-block absmax entirely in registers. MXFP4 does the same tiling on the -even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles -and N_QUANT_BLOCKS ue8m0 bytes. -""" - -from typing import Any - -import torch - -from vllm.triton_utils import tl, triton - -from .fused_indexer_q import _fp32x2_to_fp4x2 - - -def compress_norm_rope_store_triton( - state_cache: torch.Tensor, - num_actual: int, - token_to_req_indices: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - state_width: int, - cos_sin_cache: torch.Tensor, - kv_cache: torch.Tensor, - k_cache_metadata: Any, - pdl_kwargs: dict, - head_dim: int, - rope_head_dim: int, - compress_ratio: int, - overlap: bool, - use_fp4_cache: bool, - rms_norm_weight: torch.Tensor, - rms_norm_eps: float, - quant_block: int, - token_stride: int, - scale_dim: int, -) -> None: - """Shared triton launcher for the fused compress+norm+RoPE+insert path. - - Picks one of the three kernels in this module based on ``head_dim`` and - ``use_fp4_cache``. Identical launch signature for all three. - """ - if head_dim == 512: - kernel = _fused_kv_compress_norm_rope_insert_sparse_attn - num_warps = 4 - elif use_fp4_cache: - kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn - num_warps = 1 - else: - kernel = _fused_kv_compress_norm_rope_insert_indexer_attn - num_warps = 1 - - kernel[(num_actual,)]( - # state cache - state_cache, - state_cache.stride(0), - state_cache.stride(1), - # metadata - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_table.stride(0), - block_size, - # RMSNorm - rms_norm_weight, - rms_norm_eps, - # RoPE - cos_sin_cache, - cos_sin_cache.stride(0), - # KV cache - kv_cache, - k_cache_metadata.slot_mapping, - kv_cache.shape[1], # paged KV cache block size (tokens per block) - # constexprs - HEAD_SIZE=head_dim, - TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim), - STATE_WIDTH=state_width, - COMPRESS_RATIO=compress_ratio, - OVERLAP=overlap, - ROPE_HEAD_DIM=rope_head_dim, - FP8_MAX=448.0, - QUANT_BLOCK=quant_block, - TOKEN_STRIDE=token_stride, - SCALE_DIM=scale_dim, - KV_BLOCK_STRIDE=kv_cache.stride(0), - num_warps=num_warps, - **pdl_kwargs, - ) - - -# ============================================================================= -# DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16) -# ============================================================================= -@triton.jit -def _fused_kv_compress_norm_rope_insert_sparse_attn( - # ── state cache (compressor internal state) ── - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - # ── metadata ── - token_to_req_indices_ptr, - positions_ptr, - slot_mapping_ptr, - block_table_ptr, - block_table_stride, - block_size, - # ── RMSNorm ── - rms_norm_weight_ptr, - rms_norm_eps, - # ── RoPE ── - cos_sin_cache_ptr, - cos_sin_stride, - # ── KV cache output ── - k_cache_ptr, - kv_slot_mapping_ptr, - kv_cache_block_size, - # ── constexprs ── - HEAD_SIZE: tl.constexpr, - TRITON_BLOCK_SIZE: tl.constexpr, - STATE_WIDTH: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - OVERLAP: tl.constexpr, - ROPE_HEAD_DIM: tl.constexpr, - FP8_MAX: tl.constexpr, # 448.0 - QUANT_BLOCK: tl.constexpr, # 64 for DeepseekV4 - TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4 - SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad) - KV_BLOCK_STRIDE: tl.constexpr, -): - """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope). - - One program per token; early-exits for non-boundary positions. - - Cache block layout (``block_size`` tokens): - [0, bs*576): token data (448 fp8 + 128 bf16 each) - [bs*576, +bs*8): uint8 UE8M0 scales (7 real + 1 pad each) - """ - token_idx = tl.program_id(0) - - slot_id = tl.load(slot_mapping_ptr + token_idx) - if slot_id < 0: - return - - position = tl.load(positions_ptr + token_idx) - if (position + 1) % COMPRESS_RATIO != 0: - return - - req_idx = tl.load(token_to_req_indices_ptr + token_idx) - - # ── Gather state cache entries ──────────────────────────────────── - start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 - tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) - pos = start + tokens - mask_pos = pos >= 0 - - block_indices = pos // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask_pos, - other=0, - ) - block_offsets = pos % block_size - head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE - - block = tl.arange(0, TRITON_BLOCK_SIZE) - mask = block < HEAD_SIZE - block_numbers_i64 = block_numbers.to(tl.int64) - - # Precomputed row base shared by score and kv loads - row_base = ( - state_cache_ptr - + block_numbers_i64 * state_cache_stride0 - + block_offsets * state_cache_stride1 - + head_offset - ) - - combined_mask = mask_pos[:, None] & mask[None, :] - - # ── Softmax + weighted sum ─────────────────────────────────────── - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), - ) - score = tl.softmax(score, dim=0) - - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - - compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 - - # ── RMSNorm (fp32 throughout) ────────────────────────────────────── - rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) - variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE - rrms = tl.rsqrt(variance + rms_norm_eps) - normed = compressed_kv * rrms * rms_w - - # ── KV cache pointers ──────────────────────────────────────────── - kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) - if kv_slot_idx < 0: - return - kv_block_idx = kv_slot_idx // kv_cache_block_size - kv_pos_in_block = kv_slot_idx % kv_cache_block_size - - cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE - fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE - scale_ptr = ( - cache_block_ptr - + kv_cache_block_size * TOKEN_STRIDE - + kv_pos_in_block * SCALE_DIM - ) - - NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM # 448 - HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 # 32 - - # FP8 UE8M0 quant: cast fp32 → bf16 → fp32 before quant to match reference. - N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK - N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK # 7 - INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX - - quant_input = normed.to(tl.bfloat16).to(tl.float32) - quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK)) - abs_2d = tl.abs(quant_2d) - block_absmax = tl.max(abs_2d, axis=1) # [N_QUANT_BLOCKS] fp32 - block_absmax = tl.maximum(block_absmax, 1e-4) - - raw_scales = block_absmax * INV_FP8_MAX - exponents = tl.ceil(tl.log2(raw_scales)) - inv_scales = tl.exp2(-exponents) - inv_scales_col = tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1)) - x_scaled = quant_2d * inv_scales_col - x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX) - x_fp8 = x_clamped.to(tl.float8e4nv) - x_uint8 = x_fp8.to(tl.uint8, bitcast=True) - x_uint8_flat = tl.reshape(x_uint8, (TRITON_BLOCK_SIZE,)) - - nope_mask = block < NOPE_HEAD_DIM - tl.store(fp8_ptr + block, x_uint8_flat, mask=nope_mask) - - scale_idx = tl.arange(0, N_QUANT_BLOCKS) - encoded = exponents + 127.0 - encoded = tl.maximum(tl.minimum(encoded, 255.0), 0.0) - tl.store( - scale_ptr + scale_idx, - encoded.to(tl.uint8), - mask=scale_idx < N_NOPE_BLOCKS, - ) - tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8)) - - # Register-based GPT-J RoPE in fp32. - NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 - NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 - - pair_2d = tl.reshape(normed, (NUM_PAIRS, 2)) - even, odd = tl.split(pair_2d) # each [NUM_PAIRS] fp32 - - pair_idx = tl.arange(0, NUM_PAIRS) - rope_pair_local = pair_idx - NOPE_PAIRS - is_rope_pair = rope_pair_local >= 0 - cs_idx = tl.maximum(rope_pair_local, 0) - - compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO - cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride - cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) - sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) - - new_even = even * cos_v - odd * sin_v - new_odd = odd * cos_v + even * sin_v - result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32 - - # Store rotated rope portion as bf16 into the cache's bf16 area. - bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) - rope_local = block - NOPE_HEAD_DIM - is_rope = (block >= NOPE_HEAD_DIM) & mask - tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) - - -# ============================================================================= -# Indexer path (head=128, all FP8, single quant block) -# ============================================================================= -@triton.jit -def _fused_kv_compress_norm_rope_insert_indexer_attn( - # ── state cache (compressor internal state) ── - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - # ── metadata ── - token_to_req_indices_ptr, - positions_ptr, - slot_mapping_ptr, - block_table_ptr, - block_table_stride, - block_size, - # ── RMSNorm ── - rms_norm_weight_ptr, - rms_norm_eps, - # ── RoPE ── - cos_sin_cache_ptr, - cos_sin_stride, - # ── KV cache output ── - k_cache_ptr, - kv_slot_mapping_ptr, - kv_cache_block_size, - # ── constexprs ── - HEAD_SIZE: tl.constexpr, - TRITON_BLOCK_SIZE: tl.constexpr, - STATE_WIDTH: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - OVERLAP: tl.constexpr, - ROPE_HEAD_DIM: tl.constexpr, - FP8_MAX: tl.constexpr, # 448.0 - QUANT_BLOCK: tl.constexpr, # 128 for indexer - TOKEN_STRIDE: tl.constexpr, # 128 for indexer - SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32) - KV_BLOCK_STRIDE: tl.constexpr, -): - """Fused compress → RMSNorm → RoPE → FP8 quant → store. - - One program per token; early-exits for non-boundary positions. - - Cache block layout: - [0, bs*128): FP8 data (128 bytes/token) - [bs*128, +bs*4): float32 scales (4 bytes/token) - - For head_dim=128 we have exactly one quant block, so we skip the - [N_QUANT_BLOCKS, QUANT_BLOCK] reshape entirely and use a flat - ``tl.max`` reduction. - """ - token_idx = tl.program_id(0) - - slot_id = tl.load(slot_mapping_ptr + token_idx) - if slot_id < 0: - return - - position = tl.load(positions_ptr + token_idx) - if (position + 1) % COMPRESS_RATIO != 0: - return - - req_idx = tl.load(token_to_req_indices_ptr + token_idx) - - # ── Gather state cache entries ──────────────────────────────────── - start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 - tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) - pos = start + tokens - mask_pos = pos >= 0 - - block_indices = pos // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask_pos, - other=0, - ) - block_offsets = pos % block_size - head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE - - block = tl.arange(0, TRITON_BLOCK_SIZE) - mask = block < HEAD_SIZE - block_numbers_i64 = block_numbers.to(tl.int64) - - row_base = ( - state_cache_ptr - + block_numbers_i64 * state_cache_stride0 - + block_offsets * state_cache_stride1 - + head_offset - ) - - combined_mask = mask_pos[:, None] & mask[None, :] - - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), - ) - score = tl.softmax(score, dim=0) - - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - - compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 - - # ── RMSNorm (fp32 throughout) ────────────────────────────────────── - rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) - variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE - rrms = tl.rsqrt(variance + rms_norm_eps) - normed = compressed_kv * rrms * rms_w - - # ── KV cache pointers ──────────────────────────────────────────── - kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) - if kv_slot_idx < 0: - return - kv_block_idx = kv_slot_idx // kv_cache_block_size - kv_pos_in_block = kv_slot_idx % kv_cache_block_size - - cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE - fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE - scale_ptr = ( - cache_block_ptr - + kv_cache_block_size * TOKEN_STRIDE - + kv_pos_in_block * SCALE_DIM - ) - - NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM - HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 - - # ── Register-based GPT-J forward RoPE in fp32 ───────────────────── - NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 - NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 - - normed_2d = tl.reshape(normed, (NUM_PAIRS, 2)) - even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32 - - pair_idx = tl.arange(0, NUM_PAIRS) - rope_pair_local = pair_idx - NOPE_PAIRS - is_rope_pair = rope_pair_local >= 0 - cs_idx = tl.maximum(rope_pair_local, 0) - - compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO - cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride - cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) - sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) - - new_even = even * cos_v - odd * sin_v - new_odd = odd * cos_v + even * sin_v - result = tl.interleave(new_even, new_odd) # fp32 - - # ── FP8 UE8M0 quant: single block, flat reduction ──────────────── - tl.static_assert( - TRITON_BLOCK_SIZE == QUANT_BLOCK, - "Indexer expects one quant block (QUANT_BLOCK == TRITON_BLOCK_SIZE)", - ) - INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX - - result_bf16 = result.to(tl.bfloat16).to(tl.float32) - absmax = tl.max(tl.abs(result_bf16), axis=0) # scalar - absmax = tl.maximum(absmax, 1e-4) - raw_scale = absmax * INV_FP8_MAX - exponent = tl.ceil(tl.log2(raw_scale)) - inv_scale = tl.exp2(-exponent) - - x_scaled = result_bf16 * inv_scale - x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX) - x_fp8 = x_clamped.to(tl.float8e4nv) - x_uint8 = x_fp8.to(tl.uint8, bitcast=True) - - tl.store(fp8_ptr + block, x_uint8, mask=mask) - - # Single float32 scale - scale_val = tl.exp2(exponent) - tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) - - -# ============================================================================= -# Indexer path (head=128, MXFP4: 2 nibbles/byte + ue8m0 per 32-elem block) -# ============================================================================= -@triton.jit -def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( - # ── state cache (compressor internal state) ── - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - # ── metadata ── - token_to_req_indices_ptr, - positions_ptr, - slot_mapping_ptr, - block_table_ptr, - block_table_stride, - block_size, - # ── RMSNorm ── - rms_norm_weight_ptr, - rms_norm_eps, - # ── RoPE ── - cos_sin_cache_ptr, - cos_sin_stride, - # ── KV cache output ── - k_cache_ptr, - kv_slot_mapping_ptr, - kv_cache_block_size, - # ── constexprs ── - HEAD_SIZE: tl.constexpr, - TRITON_BLOCK_SIZE: tl.constexpr, - STATE_WIDTH: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - OVERLAP: tl.constexpr, - ROPE_HEAD_DIM: tl.constexpr, - FP8_MAX: tl.constexpr, # unused for MXFP4 (kept for signature parity) - QUANT_BLOCK: tl.constexpr, # 32 for MXFP4 - TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token - SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token - KV_BLOCK_STRIDE: tl.constexpr, -): - """Fused compress → RMSNorm → RoPE → MXFP4 quant → store. - - One program per token; early-exits for non-boundary positions. - - Cache block layout (``block_size`` tokens per cache block): - [0, bs*TOKEN_STRIDE): packed MXFP4 nibbles (2 values/byte) - [bs*TOKEN_STRIDE, +bs*SCALE_DIM): ue8m0 scale bytes (one per 32-elem block) - - MXFP4 format: - - E2M1 4-bit values packed two per byte (low nibble first, then high). - - Per-32-element block scale = 2^ceil(log2(amax / 6.0)), stored ue8m0 - (byte = exponent + 127). - - Max representable magnitude = 6.0. - """ - token_idx = tl.program_id(0) - - slot_id = tl.load(slot_mapping_ptr + token_idx) - if slot_id < 0: - return - - position = tl.load(positions_ptr + token_idx) - if (position + 1) % COMPRESS_RATIO != 0: - return - - req_idx = tl.load(token_to_req_indices_ptr + token_idx) - - # ── Gather state cache entries ──────────────────────────────────── - start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 - tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) - pos = start + tokens - mask_pos = pos >= 0 - - block_indices = pos // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask_pos, - other=0, - ) - block_offsets = pos % block_size - head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE - - block = tl.arange(0, TRITON_BLOCK_SIZE) - mask = block < HEAD_SIZE - block_numbers_i64 = block_numbers.to(tl.int64) - - row_base = ( - state_cache_ptr - + block_numbers_i64 * state_cache_stride0 - + block_offsets * state_cache_stride1 - + head_offset - ) - - combined_mask = mask_pos[:, None] & mask[None, :] - - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), - ) - score = tl.softmax(score, dim=0) - - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - - compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 - - # ── RMSNorm (fp32 throughout) ────────────────────────────────────── - rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) - variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE - rrms = tl.rsqrt(variance + rms_norm_eps) - normed = compressed_kv * rrms * rms_w - - # ── KV cache pointers (segregated: values first, then scales) ──── - kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) - if kv_slot_idx < 0: - return - kv_block_idx = kv_slot_idx // kv_cache_block_size - kv_pos_in_block = kv_slot_idx % kv_cache_block_size - - cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE - val_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE - scale_ptr = ( - cache_block_ptr - + kv_cache_block_size * TOKEN_STRIDE - + kv_pos_in_block * SCALE_DIM - ) - - NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM - HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 - - # ── Register-based GPT-J forward RoPE in fp32 ───────────────────── - # We keep the even/odd halves (no tl.interleave afterwards) because the - # MXFP4 per-block absmax / pack naturally operates on (even, odd) pairs. - NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 - NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 - - normed_2d = tl.reshape(normed, (NUM_PAIRS, 2)) - even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32 - - pair_idx = tl.arange(0, NUM_PAIRS) - rope_pair_local = pair_idx - NOPE_PAIRS - is_rope_pair = rope_pair_local >= 0 - cs_idx = tl.maximum(rope_pair_local, 0) - - compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO - cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride - cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) - sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) - - new_even = even * cos_v - odd * sin_v - new_odd = odd * cos_v + even * sin_v - - # bf16 roundtrip for parity with reference / Q-side kernel numerics. - new_even = new_even.to(tl.bfloat16).to(tl.float32) - new_odd = new_odd.to(tl.bfloat16).to(tl.float32) - - # ── MXFP4 quant: tile even/odd halves into (N_BLOCKS, HALF_BLOCK) ── - # Each MXFP4 block of QUANT_BLOCK elements = HALF_BLOCK consecutive pairs, - # so (N_BLOCKS, HALF_BLOCK) rows of even/odd each land exactly one block. - N_QUANT_BLOCKS: tl.constexpr = HEAD_SIZE // QUANT_BLOCK - HALF_BLOCK: tl.constexpr = QUANT_BLOCK // 2 - tl.static_assert(TRITON_BLOCK_SIZE == HEAD_SIZE) - tl.static_assert(HEAD_SIZE % QUANT_BLOCK == 0) - tl.static_assert(TOKEN_STRIDE == HEAD_SIZE // 2) - tl.static_assert(SCALE_DIM == N_QUANT_BLOCKS) - - even_2d = tl.reshape(new_even, (N_QUANT_BLOCKS, HALF_BLOCK)) - odd_2d = tl.reshape(new_odd, (N_QUANT_BLOCKS, HALF_BLOCK)) - - amax = tl.maximum( - tl.max(tl.abs(even_2d), axis=1), - tl.max(tl.abs(odd_2d), axis=1), - ) - amax = tl.maximum(amax, 6.0 * (2**-126)) - - # ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored as (exp + 127) byte. - log2_ratio = tl.ceil(tl.log2(amax * (1.0 / 6.0))) - log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0) - inv_scale = tl.exp2(-log2_ratio) - ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS] - - inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1)) - packed = _fp32x2_to_fp4x2( - even_2d * inv_scale_col, odd_2d * inv_scale_col - ) # (N_BLOCKS, HALF_BLOCK) uint8 - packed_flat = tl.reshape(packed, (TOKEN_STRIDE,)) - - tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat) - tl.store(scale_ptr + tl.arange(0, SCALE_DIM), ue8m0) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py deleted file mode 100644 index d5aaf10f..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py +++ /dev/null @@ -1,438 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_cutedsl - -# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale. -MXFP4_BLOCK_SIZE = 32 - - -@triton.jit -def _get_cos_sin( - cos_sin_cache_ptr, - cos_sin_cache_stride, - pos, - HALF_ROT_DIM: tl.constexpr, -): - block = tl.arange(0, HALF_ROT_DIM) - cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block) - cos = cos.to(tl.float32) - sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM) - sin = sin.to(tl.float32) - return cos, sin - - -@triton.jit -def _fp32x2_to_fp4x2(x_lo, x_hi): - # NOTE: $1 is high nibble, $2 is low nibble - return tl.inline_asm_elementwise( - """ - { - .reg .b8 tmp; - cvt.rn.satfinite.e2m1x2.f32 tmp, $1, $2; - cvt.u32.u8 $0, tmp; - } - """, - constraints="=r,f,f", - args=[x_hi, x_lo], - dtype=tl.uint32, - is_pure=True, - pack=1, - ).to(tl.uint8) - - -@triton.jit -def _quantize_mxfp4_pair(x_lo, x_hi): - """Quantize a block of MXFP4_BLOCK_SIZE fp32 values given as two - interleaved halves (x_lo = values at even positions in the block, - x_hi = values at odd positions). Returns: - - packed : uint8[BLOCK/2] (low nibble = quant(x_lo), high = quant(x_hi)) - - ue8m0 : scalar uint8 (block scale = 2^(ue8m0 - 127)) - """ - amax = tl.maximum(tl.max(tl.abs(x_lo)), tl.max(tl.abs(x_hi))) - # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163 - amax = tl.maximum(amax, 6.0 * (2**-126)) - # ue8m0 block scale: 2^ceil(log2(amax/6.0)). - log2_ratio = tl.math.ceil(tl.math.log2(amax * (1.0 / 6.0))) - log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0) - scale = tl.math.exp2(log2_ratio) - ue8m0 = (log2_ratio + 127.0).to(tl.uint8) - - inv_scale = 1.0 / scale - packed = _fp32x2_to_fp4x2(x_lo * inv_scale, x_hi * inv_scale) - return packed, ue8m0 - - -@triton.jit -def _fused_indexer_q_rope_quant_kernel( - pos_ptr, - # Index Q RoPE - index_q_ptr, - index_q_stride0, - index_q_stride1, - index_q_cos_sin_ptr, - index_q_cos_sin_stride, - INDEX_Q_HALF_ROT_DIM: tl.constexpr, - # Index Q Quantize - index_q_fp8_ptr, - index_q_fp8_stride0, - index_q_fp8_stride1, - INDEX_Q_HEAD_DIM: tl.constexpr, - # Index weights - index_weights_ptr, - index_weights_stride, - index_weights_softmax_scale, - index_weights_head_scale, - index_weights_out_ptr, - index_weights_out_stride, -): - # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding - # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the - # LAST rope_dim dims of each head; the leading [0, NOPE_DIM) is passed - # through unchanged. - INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM - INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM - tl.static_assert(INDEX_Q_NOPE_DIM >= 0) - - tok_idx = tl.program_id(0) - head_idx = tl.program_id(1) - - pos = tl.load(pos_ptr + tok_idx) - cos, sin = _get_cos_sin( - index_q_cos_sin_ptr, - index_q_cos_sin_stride, - pos, - INDEX_Q_HALF_ROT_DIM, - ) - half_offset = tl.arange(0, INDEX_Q_HALF_ROT_DIM) - base_ptr = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1 - - # Interleaved (GPT-J) RoPE on dims [NOPE_DIM, HEAD_DIM): - # even = q[NOPE_DIM + 2*i], odd = q[NOPE_DIM + 2*i + 1] - rot_base = base_ptr + INDEX_Q_NOPE_DIM - x_even = tl.load(rot_base + half_offset * 2).to(tl.float32) - x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32) - r_even = x_even * cos - x_odd * sin - r_odd = x_odd * cos + x_even * sin - - # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax. - # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py). - r_even = r_even.to(tl.bfloat16).to(tl.float32) - r_odd = r_odd.to(tl.bfloat16).to(tl.float32) - - amax = tl.maximum(tl.max(tl.abs(r_even)), tl.max(tl.abs(r_odd))) - if INDEX_Q_NOPE_DIM > 0: - nope_offset = tl.arange(0, INDEX_Q_NOPE_DIM) - x_nope = tl.load(base_ptr + nope_offset).to(tl.float32) - amax = tl.maximum(amax, tl.max(tl.abs(x_nope))) - index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0) - index_q_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(index_q_scale))) - - # Store quantized values to index_q_fp8 - fp8_base_ptr = ( - index_q_fp8_ptr + tok_idx * index_q_fp8_stride0 + head_idx * index_q_fp8_stride1 - ) - if INDEX_Q_NOPE_DIM > 0: - tl.store( - fp8_base_ptr + nope_offset, - tl.div_rn(x_nope, index_q_scale).to(tl.float8e4nv), - ) - fp8_rot_base = fp8_base_ptr + INDEX_Q_NOPE_DIM - tl.store( - fp8_rot_base + half_offset * 2, - tl.div_rn(r_even, index_q_scale).to(tl.float8e4nv), - ) - tl.store( - fp8_rot_base + half_offset * 2 + 1, - tl.div_rn(r_odd, index_q_scale).to(tl.float8e4nv), - ) - - # FP8 weight-fold contract: - # index_weights_out = index_weights * q_scale * softmax_scale * head_scale - # The per-token-per-head q_scale (fp32) IS folded into the output weights - # here because FP8 Q is stored WITHOUT a companion scale tensor — the - # downstream fp8_fp4_mqa_logits/fp8_fp4_paged_mqa_logits kernels use `weights` to - # apply per-token Q scale inline. See the MXFP4 kernel below for the - # contrasting convention (scales live with the Q values, weights are NOT - # q-scaled). - index_weights = tl.load( - index_weights_ptr + tok_idx * index_weights_stride + head_idx - ) - index_weights = index_weights.to(tl.float32) - index_weights *= index_q_scale - index_weights *= index_weights_softmax_scale - index_weights *= index_weights_head_scale - tl.store( - index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx, - index_weights, - ) - - -@triton.jit -def _fused_indexer_q_rope_mxfp4_kernel( - pos_ptr, - # Index Q RoPE input (fp/bf16) - index_q_ptr, - index_q_stride0, - index_q_stride1, - index_q_cos_sin_ptr, - index_q_cos_sin_stride, - INDEX_Q_HALF_ROT_DIM: tl.constexpr, - # MXFP4 Q outputs - index_q_mxfp4_ptr, # uint8, (T, H, HEAD_DIM // 2) - index_q_mxfp4_stride0, - index_q_mxfp4_stride1, - index_q_scale_ptr, # uint8 ue8m0, (T, H, HEAD_DIM // BLOCK) - index_q_scale_stride0, - index_q_scale_stride1, - INDEX_Q_HEAD_DIM: tl.constexpr, - MXFP4_BLOCK: tl.constexpr, - # Weights (NO per-token q_scale fold for MXFP4; per-block scales stay - # with the Q values in the output scale tensor). - index_weights_ptr, - index_weights_stride, - index_weights_softmax_scale, - index_weights_head_scale, - index_weights_out_ptr, - index_weights_out_stride, -): - INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM - INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM - NUM_NOPE_BLOCKS: tl.constexpr = INDEX_Q_NOPE_DIM // MXFP4_BLOCK - NUM_ROPE_BLOCKS: tl.constexpr = INDEX_Q_ROT_DIM // MXFP4_BLOCK - HALF_BLOCK: tl.constexpr = MXFP4_BLOCK // 2 - tl.static_assert(INDEX_Q_NOPE_DIM >= 0) - tl.static_assert(INDEX_Q_NOPE_DIM % MXFP4_BLOCK == 0) - tl.static_assert(INDEX_Q_ROT_DIM % MXFP4_BLOCK == 0) - tl.static_assert(MXFP4_BLOCK % 2 == 0) - - tok_idx = tl.program_id(0) - head_idx = tl.program_id(1) - - pos = tl.load(pos_ptr + tok_idx) - - q_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1 - out_base = ( - index_q_mxfp4_ptr - + tok_idx * index_q_mxfp4_stride0 - + head_idx * index_q_mxfp4_stride1 - ) - scale_base = ( - index_q_scale_ptr - + tok_idx * index_q_scale_stride0 - + head_idx * index_q_scale_stride1 - ) - - half_off = tl.arange(0, HALF_BLOCK) - - # ---- NoPE blocks: direct load, pair as (even-index, odd-index) values ---- - for b in tl.static_range(NUM_NOPE_BLOCKS): - base = b * MXFP4_BLOCK - x_lo = tl.load(q_base + base + half_off * 2).to(tl.float32) - x_hi = tl.load(q_base + base + half_off * 2 + 1).to(tl.float32) - packed, ue8m0 = _quantize_mxfp4_pair(x_lo, x_hi) - tl.store(out_base + base // 2 + half_off, packed) - tl.store(scale_base + b, ue8m0) - - # ---- RoPE blocks: apply GPT-J interleaved RoPE to the block's 16 pairs, - # then quantize. Each block covers HALF_BLOCK (=16) cos/sin pairs. ---- - rot_q_base = q_base + INDEX_Q_NOPE_DIM - for b in tl.static_range(NUM_ROPE_BLOCKS): - pair_off = b * HALF_BLOCK + half_off # indices in [0, HALF_ROT_DIM) - cos_b = tl.load( - index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + pair_off - ).to(tl.float32) - sin_b = tl.load( - index_q_cos_sin_ptr - + pos * index_q_cos_sin_stride - + pair_off - + INDEX_Q_HALF_ROT_DIM - ).to(tl.float32) - x_even = tl.load(rot_q_base + pair_off * 2).to(tl.float32) - x_odd = tl.load(rot_q_base + pair_off * 2 + 1).to(tl.float32) - r_even = x_even * cos_b - x_odd * sin_b - r_odd = x_odd * cos_b + x_even * sin_b - # bf16 roundtrip for parity with the FP8 kernel / reference numerics. - r_even = r_even.to(tl.bfloat16).to(tl.float32) - r_odd = r_odd.to(tl.bfloat16).to(tl.float32) - packed, ue8m0 = _quantize_mxfp4_pair(r_even, r_odd) - rope_byte_off = (INDEX_Q_NOPE_DIM + b * MXFP4_BLOCK) // 2 - tl.store(out_base + rope_byte_off + half_off, packed) - tl.store(scale_base + NUM_NOPE_BLOCKS + b, ue8m0) - - # MXFP4 weight-fold contract: - # index_weights_out = index_weights * softmax_scale * head_scale - # NOTE: q_scale is NOT folded here (contrast with the FP8 kernel above). - # MXFP4 Q emits a separate ue8m0 scale tensor of shape - # (T, H, HEAD_DIM // MXFP4_BLOCK) alongside the packed values, so each - # per-block scale is applied by the downstream MXFP4 logits kernel when - # dequantizing Q — there is no per-token scalar to fold into `weights`. - index_weights = tl.load( - index_weights_ptr + tok_idx * index_weights_stride + head_idx - ).to(tl.float32) - index_weights *= index_weights_softmax_scale - index_weights *= index_weights_head_scale - tl.store( - index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx, - index_weights, - ) - - -def fused_indexer_q_rope_quant( - positions: torch.Tensor, - index_q: torch.Tensor, - index_q_cos_sin_cache: torch.Tensor, - # Index weights - index_weights: torch.Tensor, - index_weights_softmax_scale: float, - index_weights_head_scale: float, - use_fp4: bool = False, -) -> tuple[ - torch.Tensor | tuple[torch.Tensor, torch.Tensor], - torch.Tensor, -]: - """Fused RoPE + quantize Q for the sparse indexer. - - Weight-fold semantics (important — the two paths differ): - - FP8 path (use_fp4=False, default): - q_fp8 : (T, H, HEAD_DIM) float8_e4m3fn, per-token-per-head - scalar scale (NOT stored — folded into weights below) - weights_out = weights * q_scale * softmax_scale * head_scale - Rationale: a single per-token q_scale is a scalar the downstream FP8 - logits kernel would otherwise multiply in. Folding it into `weights` - avoids emitting a separate tensor and is free for the logits kernel. - - MXFP4 path (use_fp4=True): - q_packed : (T, H, HEAD_DIM // 2) uint8 (2 E2M1 nibbles per byte) - q_scale : (T, H, HEAD_DIM // MXFP4_BLOCK_SIZE) uint8 ue8m0 bytes - weights_out = weights * softmax_scale * head_scale - Rationale: MXFP4 has PER-BLOCK (32-element) scales that live with - the Q values — they cannot be folded into a per-token weight - scalar, so `weights` carries only the softmax and head scales. - - Returns (q_quant, weights_out) where q_quant is either a Tensor (FP8) or - a (values, scales) tuple (MXFP4). This matches the union type accepted - by `SparseAttnIndexer.forward_*`. - """ - assert positions.ndim == 1 - assert index_q.ndim == 3 - assert index_q_cos_sin_cache.ndim == 2 - - num_tokens = positions.shape[0] - num_index_q_heads = index_q.shape[1] - index_q_head_dim = index_q.shape[2] - - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) - - if use_fp4: - assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( - f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block " - f"size {MXFP4_BLOCK_SIZE}" - ) - num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE - index_q_packed = torch.empty( - (num_tokens, num_index_q_heads, index_q_head_dim // 2), - dtype=torch.uint8, - device=index_q.device, - ) - index_q_scale = torch.empty( - (num_tokens, num_index_q_heads, num_scale_blocks), - dtype=torch.uint8, - device=index_q.device, - ) - if has_cutedsl(): - # lazily import, otherwise some tests fail due to CUDA driver init failure. - from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( - fused_indexer_q_rope_quant_mxfp4_cutedsl, - ) - - fused_indexer_q_rope_quant_mxfp4_cutedsl( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_weights_softmax_scale, - index_weights_head_scale, - index_q_packed, - index_q_scale, - index_weights_out, - ) - else: - _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)]( - positions, - index_q, - index_q.stride(0), - index_q.stride(1), - index_q_cos_sin_cache, - index_q_cos_sin_cache.stride(0), - index_q_cos_sin_cache.shape[-1] // 2, - index_q_packed, - index_q_packed.stride(0), - index_q_packed.stride(1), - index_q_scale, - index_q_scale.stride(0), - index_q_scale.stride(1), - index_q_head_dim, - MXFP4_BLOCK_SIZE, - index_weights, - index_weights.stride(0), - index_weights_softmax_scale, - index_weights_head_scale, - index_weights_out, - index_weights_out.stride(0), - num_warps=1, # TODO: Tune this - ) - - # Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0 - # bytes per (token, head) reinterpreted as one int32, then squeezed - # from (T, H, 1) to (T, H) to match DeepGEMM's expected q_sf rank - # (prefill wants 2-D (seq_len, num_heads); decode reshapes this to - # 3-D (batch, next_n, num_heads)). - return ( - index_q_packed, - index_q_scale.view(torch.int32).squeeze(-1), - ), index_weights_out - - index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn) - if has_cutedsl(): - # lazily import, otherwise some tests fail due to CUDA driver init failure. - from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( - fused_indexer_q_rope_quant_fp8_cutedsl, - ) - - fused_indexer_q_rope_quant_fp8_cutedsl( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_weights_softmax_scale, - index_weights_head_scale, - index_q_fp8, - index_weights_out, - ) - else: - _fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)]( - positions, - index_q, - index_q.stride(0), - index_q.stride(1), - index_q_cos_sin_cache, - index_q_cos_sin_cache.stride(0), - index_q_cos_sin_cache.shape[-1] // 2, - index_q_fp8, - index_q_fp8.stride(0), - index_q_fp8.stride(1), - index_q_head_dim, - index_weights, - index_weights.stride(0), - index_weights_softmax_scale, - index_weights_head_scale, - index_weights_out, - index_weights_out.stride(0), - num_warps=1, # TODO: Tune this - ) - return index_q_fp8, index_weights_out diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py deleted file mode 100644 index 97fc0962..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py +++ /dev/null @@ -1,318 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Fused inverse RoPE + block-scaled FP8 quantization kernel for DeepseekV4 attention. - -Output scale format is pre-transformed (MN-major TMA-aligned; FP32 on SM90, -INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layout. -""" - -import torch - -from vllm.platforms import current_platform -from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import direct_register_custom_op - - -@triton.jit(do_not_specialize=["num_tokens"]) -def _fused_inv_rope_fp8_quant_per_head( - o_ptr, - positions_ptr, - cos_sin_cache_ptr, - fp8_ptr, - scale_ptr, - num_tokens, - heads_per_group: tl.constexpr, - o_stride_token, - o_stride_head, - cache_stride_pos, - fp8_stride_group, - fp8_stride_token, - scale_stride_group, - scale_stride_k, - fp8_max: tl.constexpr, - eps: tl.constexpr, - QUANT_GROUP_SIZE: tl.constexpr, - CHUNKS_PER_HEAD: tl.constexpr, - ROPE_START: tl.constexpr, - HALF_ROPE: tl.constexpr, - TMA_ALIGNED_SCALES: tl.constexpr, -): - # int64: stride multiply overflows int32 past num_tokens=32768 (IMA). - pid_token = tl.program_id(0).to(tl.int64) - pid_gh = tl.program_id(1).to(tl.int64) - - g = pid_gh // heads_per_group - head_in_group = pid_gh % heads_per_group - global_head = pid_gh - qb_start = head_in_group * CHUNKS_PER_HEAD - - # Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant. - if pid_token >= num_tokens: - if TMA_ALIGNED_SCALES: - scale_addr = ( - scale_ptr - + g * scale_stride_group - + pid_token - + head_in_group * scale_stride_k - ) - tl.store(scale_addr, tl.zeros((), dtype=tl.int32)) - else: - block_offsets = tl.arange(0, CHUNKS_PER_HEAD) - qb_indices = qb_start + block_offsets - scale_addrs = ( - scale_ptr - + g * scale_stride_group - + pid_token - + qb_indices * scale_stride_k - ) - tl.store(scale_addrs, tl.zeros((CHUNKS_PER_HEAD,), dtype=tl.float32)) - return - - input_base = o_ptr + pid_token * o_stride_token + global_head * o_stride_head - - HEAD_DIM: tl.constexpr = CHUNKS_PER_HEAD * QUANT_GROUP_SIZE - offsets = tl.arange(0, HEAD_DIM) - x = tl.load(input_base + offsets).to(tl.float32) - - rope_abs_start: tl.constexpr = (CHUNKS_PER_HEAD - 1) * QUANT_GROUP_SIZE + ROPE_START - pos = tl.load(positions_ptr + pid_token) - cache_base = cos_sin_cache_ptr + pos * cache_stride_pos - is_rope = offsets >= rope_abs_start - rope_local = offsets - rope_abs_start - - x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to( - tl.float32 - ) - cs_idx = tl.maximum(rope_local >> 1, 0) - cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0) - sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0) - x_add = x * cos_v + x_partner * sin_v - x_sub = x * cos_v - x_partner * sin_v - is_even = (rope_local & 1) == 0 - rotated = tl.where(is_even, x_add, x_sub) - x = tl.where(is_rope, rotated, x) - - x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE)) - block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps) - scale_raw = block_absmax * (1.0 / fp8_max) - scales = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) - - scales_exp = tl.reshape( - tl.broadcast_to( - tl.reshape(scales, (CHUNKS_PER_HEAD, 1)), - (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE), - ), - (HEAD_DIM,), - ) - x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv) - - fp8_base = ( - fp8_ptr - + g * fp8_stride_group - + pid_token * fp8_stride_token - + qb_start * QUANT_GROUP_SIZE - ) - tl.store(fp8_base + offsets, x_quant) - - block_offsets = tl.arange(0, CHUNKS_PER_HEAD) - qb_indices = qb_start + block_offsets - if TMA_ALIGNED_SCALES: - scale_bits = scales.to(tl.int32, bitcast=True) - ue8m0_bytes = (scale_bits >> 23) & 0xFF - packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8)) - scale_addr = ( - scale_ptr - + g * scale_stride_group - + pid_token - + head_in_group * scale_stride_k - ) - tl.store(scale_addr, packed_val) - else: - scale_addrs = ( - scale_ptr + g * scale_stride_group + pid_token + qb_indices * scale_stride_k - ) - tl.store(scale_addrs, scales) - - -def fused_inv_rope_fp8_quant( - o: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - n_groups: int, - heads_per_group: int, - nope_dim: int = 448, - rope_dim: int = 64, - quant_group_size: int = 128, - tma_aligned_scales: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - """Fused inverse RoPE + block-scaled FP8 quantization. - - Args: - o: Attention output [num_tokens, num_heads, head_dim] bf16. - positions: Token positions [num_tokens] int64. - cos_sin_cache: Precomputed [max_pos, rope_dim] with cos||sin. - n_groups: Number of output groups. - heads_per_group: Heads per group. - nope_dim: Non-RoPE dimensions per head (default 448). - rope_dim: RoPE dimensions per head (default 64). - quant_group_size: FP8 quantization block size (default 128). - tma_aligned_scales: Output INT32 packed UE8M0 for SM100 (True) - or FP32 for SM90 (False). - - Returns: - o_fp8: [T, G, D] float8_e4m3fn, strides (D, T*D, 1). - o_scale: Pre-transformed scale tensor for fp8_einsum. - """ - from vllm.utils.deep_gemm import get_tma_aligned_size - - num_tokens, num_heads, head_dim = o.shape - assert num_heads == n_groups * heads_per_group - assert head_dim == nope_dim + rope_dim - assert head_dim % quant_group_size == 0 - assert nope_dim % quant_group_size == (quant_group_size - rope_dim) - assert rope_dim % 2 == 0 - assert cos_sin_cache.shape[-1] == rope_dim - assert cos_sin_cache.dtype == torch.float32 - - d = heads_per_group * head_dim - num_scale_blocks = d // quant_group_size - chunks_per_head = head_dim // quant_group_size - - fp8_dtype = torch.float8_e4m3fn - fp8_max = torch.finfo(fp8_dtype).max - - tma_aligned_T = get_tma_aligned_size(num_tokens, 4) - if tma_aligned_scales: - packed_sf_k = (num_scale_blocks + 3) // 4 - scale_inner = packed_sf_k - else: - scale_inner = num_scale_blocks - - # Run kernel through a custom op so inductor sees an opaque boundary. - # It's a pytorch bug, see https://github.com/vllm-project/vllm/issues/41106 - fp8_buf, scale_buf = torch.ops.vllm.fused_inv_rope_fp8_quant_kernel( - o, - positions, - cos_sin_cache, - heads_per_group, - quant_group_size, - chunks_per_head, - nope_dim % quant_group_size, - rope_dim // 2, - tma_aligned_scales, - fp8_max, - tma_aligned_T, - num_tokens, - n_groups, - d, - scale_inner, - ) - return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1) - - -def _fused_inv_rope_fp8_quant_kernel_impl( - o: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - heads_per_group: int, - quant_group_size: int, - chunks_per_head: int, - rope_start: int, - half_rope: int, - tma_aligned_scales: bool, - fp8_max: float, - tma_aligned_T: int, - num_tokens: int, - n_groups: int, - d: int, - scale_inner: int, -) -> tuple[torch.Tensor, torch.Tensor]: - fp8_buf = torch.empty( - (n_groups, num_tokens, d), - dtype=torch.float8_e4m3fn, - device=o.device, - ) - scale_dtype = torch.int32 if tma_aligned_scales else torch.float32 - scale_buf = torch.empty( - n_groups * scale_inner * tma_aligned_T, - dtype=scale_dtype, - device=o.device, - ).as_strided( - (n_groups, num_tokens, scale_inner), - (scale_inner * tma_aligned_T, 1, tma_aligned_T), - ) - grid = (tma_aligned_T, n_groups * heads_per_group) - pdl_kwargs = ( - {} - if current_platform.is_rocm() or current_platform.is_xpu() - else {"launch_pdl": False} - ) - _fused_inv_rope_fp8_quant_per_head[grid]( - o, - positions, - cos_sin_cache, - fp8_buf, - scale_buf, - num_tokens, - heads_per_group=heads_per_group, - o_stride_token=o.stride(0), - o_stride_head=o.stride(1), - cache_stride_pos=cos_sin_cache.stride(0), - fp8_stride_group=fp8_buf.stride(0), - fp8_stride_token=fp8_buf.stride(1), - scale_stride_group=scale_buf.stride(0), - scale_stride_k=scale_buf.stride(2), - fp8_max=fp8_max, - eps=1e-10, - QUANT_GROUP_SIZE=quant_group_size, - CHUNKS_PER_HEAD=chunks_per_head, - ROPE_START=rope_start, - HALF_ROPE=half_rope, - TMA_ALIGNED_SCALES=tma_aligned_scales, - num_stages=1, - **pdl_kwargs, - num_warps=1, - ) - return fp8_buf, scale_buf - - -def _fused_inv_rope_fp8_quant_kernel_fake( - o: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - heads_per_group: int, - quant_group_size: int, - chunks_per_head: int, - rope_start: int, - half_rope: int, - tma_aligned_scales: bool, - fp8_max: float, - tma_aligned_T: int, - num_tokens: int, - n_groups: int, - d: int, - scale_inner: int, -) -> tuple[torch.Tensor, torch.Tensor]: - fp8_buf = torch.empty( - (n_groups, num_tokens, d), - dtype=torch.float8_e4m3fn, - device=o.device, - ) - scale_dtype = torch.int32 if tma_aligned_scales else torch.float32 - scale_buf = torch.empty( - n_groups * scale_inner * tma_aligned_T, - dtype=scale_dtype, - device=o.device, - ).as_strided( - (n_groups, num_tokens, scale_inner), - (scale_inner * tma_aligned_T, 1, tma_aligned_T), - ) - return fp8_buf, scale_buf - - -direct_register_custom_op( - op_name="fused_inv_rope_fp8_quant_kernel", - op_func=_fused_inv_rope_fp8_quant_kernel_impl, - fake_impl=_fused_inv_rope_fp8_quant_kernel_fake, -) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py deleted file mode 100644 index 874b4f60..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py +++ /dev/null @@ -1,203 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused MTP-input RMSNorm: enorm (with mask-zero at position 0) + hnorm. - -Replaces the eager sequence at the top of the MTP draft forward: - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) - inputs_embeds = self.enorm(inputs_embeds) - previous_hidden_states = previous_hidden_states.view(-1, hc_mult, H) - previous_hidden_states = self.hnorm(previous_hidden_states) - -which lowers to ~6 small kernels (CompareEq, where, Fill, enorm rms_norm, -hnorm rms_norm, plus aten elementwise helpers) on the breakable-cudagraph -path. Math is preserved: positions==0 → masked row → zero RMS output -regardless of weight. - -A single grid (T, hc_mult+1) drives both norms: task 0 is enorm on -inputs_embeds[token, :], task k+1 is hnorm on previous_hidden_states[token, k, :]. -""" - -import torch - -from vllm.triton_utils import tl, triton - - -@triton.jit -def _rmsnorm_row( - x, - w_ptr, - out_row_ptr, - block, - mask, - eps, - HIDDEN: tl.constexpr, -): - x = x.to(tl.float32) - variance = tl.sum(x * x, axis=0) / HIDDEN - rrms = tl.rsqrt(variance + eps) - w = tl.load(w_ptr + block, mask=mask, other=0.0).to(tl.float32) - y = x * rrms * w - tl.store(out_row_ptr + block, y.to(out_row_ptr.dtype.element_ty), mask=mask) - - -@triton.jit -def _fused_mtp_input_rmsnorm_kernel( - inputs_embeds_ptr, - positions_ptr, - prev_hidden_ptr, - enorm_weight_ptr, - hnorm_weight_ptr, - enorm_out_ptr, - hnorm_out_ptr, - eps, - HIDDEN: tl.constexpr, - HC_MULT: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - # int64 token index so per-token offsets don't overflow int32 at - # large num_tokens (matches the convention in fused_q_kv_rmsnorm). - token_idx = tl.program_id(0).to(tl.int64) - pid_task = tl.program_id(1) - - block = tl.arange(0, BLOCK_SIZE) - mask = block < HIDDEN - - if pid_task == 0: - # enorm path: load inputs_embeds[token, :] then zero-mask at pos==0. - # Math is preserved: pos==0 → x=0 → variance=0 → RMSNorm output is 0 - # regardless of weight, matching torch.where(pos==0, 0, x) + RMSNorm. - pos = tl.load(positions_ptr + token_idx) - keep = pos != 0 - x = tl.load( - inputs_embeds_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0 - ) - x = tl.where(keep, x, 0.0) - _rmsnorm_row( - x, - enorm_weight_ptr, - enorm_out_ptr + token_idx * HIDDEN, - block, - mask, - eps, - HIDDEN, - ) - else: - # hnorm path: load prev_hidden[token, slot, :]. - slot = pid_task - 1 - row_offset = (token_idx * HC_MULT + slot) * HIDDEN - x = tl.load(prev_hidden_ptr + row_offset + block, mask=mask, other=0.0) - _rmsnorm_row( - x, - hnorm_weight_ptr, - hnorm_out_ptr + row_offset, - block, - mask, - eps, - HIDDEN, - ) - - -@triton.jit -def _mtp_shared_head_rmsnorm_kernel( - x_ptr, - weight_ptr, - out_ptr, - eps, - HIDDEN: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - token_idx = tl.program_id(0).to(tl.int64) - block = tl.arange(0, BLOCK_SIZE) - mask = block < HIDDEN - x = tl.load(x_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0) - _rmsnorm_row( - x, - weight_ptr, - out_ptr + token_idx * HIDDEN, - block, - mask, - eps, - HIDDEN, - ) - - -def mtp_shared_head_rmsnorm( - hidden_states: torch.Tensor, - weight: torch.Tensor, - eps: float, -) -> torch.Tensor: - """RMSNorm for MTP's SharedHead.norm, on (T, H) bf16 input. - - Uses the same ``_rmsnorm_row`` body as ``fused_mtp_input_rmsnorm`` so the - MTP draft path runs one consistent RMSNorm implementation end to end. - """ - assert hidden_states.ndim == 2 - assert hidden_states.is_contiguous() - assert weight.is_contiguous() - num_tokens, hidden = hidden_states.shape - out = torch.empty_like(hidden_states) - if num_tokens == 0: - return out - block_size = triton.next_power_of_2(hidden) - _mtp_shared_head_rmsnorm_kernel[(num_tokens,)]( - hidden_states, - weight, - out, - eps, - HIDDEN=hidden, - BLOCK_SIZE=block_size, - ) - return out - - -def fused_mtp_input_rmsnorm( - inputs_embeds: torch.Tensor, - positions: torch.Tensor, - previous_hidden_states: torch.Tensor, - enorm_weight: torch.Tensor, - hnorm_weight: torch.Tensor, - eps: float, - hc_mult: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Returns (enorm_out, hnorm_out). - - enorm_out has the same shape as inputs_embeds (2D, [T, H]). - hnorm_out has the same shape as previous_hidden_states (3D, [T, hc_mult, H]). - previous_hidden_states must already be reshaped to 3D. - """ - assert inputs_embeds.ndim == 2 - assert previous_hidden_states.ndim == 3 - assert previous_hidden_states.shape[1] == hc_mult - assert inputs_embeds.shape[0] == previous_hidden_states.shape[0], ( - "token dim mismatch" - ) - assert ( - inputs_embeds.shape[1] - == previous_hidden_states.shape[2] - == enorm_weight.shape[0] - == hnorm_weight.shape[0] - ) - assert inputs_embeds.is_contiguous() and previous_hidden_states.is_contiguous() - assert enorm_weight.is_contiguous() and hnorm_weight.is_contiguous() - - num_tokens, hidden = inputs_embeds.shape - enorm_out = torch.empty_like(inputs_embeds) - hnorm_out = torch.empty_like(previous_hidden_states) - if num_tokens == 0: - return enorm_out, hnorm_out - - block_size = triton.next_power_of_2(hidden) - _fused_mtp_input_rmsnorm_kernel[(num_tokens, hc_mult + 1)]( - inputs_embeds, - positions, - previous_hidden_states, - enorm_weight, - hnorm_weight, - enorm_out, - hnorm_out, - eps, - HIDDEN=hidden, - HC_MULT=hc_mult, - BLOCK_SIZE=block_size, - ) - return enorm_out, hnorm_out diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py deleted file mode 100644 index 0dd348a4..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch - -from vllm.triton_utils import tl, triton - - -@triton.jit -def _fused_q_kv_rmsnorm_kernel( - q_ptr, - q_out_ptr, - q_weight_ptr, - q_in_stride, - q_out_stride, - kv_ptr, - kv_out_ptr, - kv_weight_ptr, - kv_in_stride, - kv_out_stride, - eps, - Q_SIZE: tl.constexpr, - KV_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y. - # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes - # the launch at max-num-batched-tokens >= 65536 with "invalid argument". - # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32 - # past num_tokens ~87K under large chunked prefill. - token_idx = tl.program_id(0).to(tl.int64) - pid_task = tl.program_id(1) - - if pid_task == 0: - SIZE = Q_SIZE - row_in = q_ptr + token_idx * q_in_stride - weight_ptr = q_weight_ptr - row_out = q_out_ptr + token_idx * q_out_stride - else: - SIZE = KV_SIZE - row_in = kv_ptr + token_idx * kv_in_stride - weight_ptr = kv_weight_ptr - row_out = kv_out_ptr + token_idx * kv_out_stride - - # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's - # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which - # keep x, rrms, and w all in fp32 and perform a single cast at store. - block = tl.arange(0, BLOCK_SIZE) - mask = block < SIZE - x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32) - variance = tl.sum(x * x, axis=0) / SIZE - rrms = tl.rsqrt(variance + eps) - w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32) - y = x * rrms * w - tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask) - - -def fused_q_kv_rmsnorm( - qr: torch.Tensor, - kv: torch.Tensor, - q_weight: torch.Tensor, - kv_weight: torch.Tensor, - eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - assert qr.ndim == 2 and kv.ndim == 2 - assert qr.shape[0] == kv.shape[0], ( - f"token dim mismatch: qr={qr.shape}, kv={kv.shape}" - ) - assert qr.stride(-1) == 1 and kv.stride(-1) == 1 - assert q_weight.is_contiguous() and kv_weight.is_contiguous() - - q_size = qr.shape[1] - kv_size = kv.shape[1] - num_tokens = qr.shape[0] - qr_out = torch.empty_like(qr) - kv_out = torch.empty_like(kv) - if num_tokens == 0: - return qr_out, kv_out - - block_size = triton.next_power_of_2(max(q_size, kv_size)) - _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)]( - qr, - qr_out, - q_weight, - qr.stride(0), - qr_out.stride(0), - kv, - kv_out, - kv_weight, - kv.stride(0), - kv_out.stride(0), - eps, - Q_SIZE=q_size, - KV_SIZE=kv_size, - BLOCK_SIZE=block_size, - ) - return qr_out, kv_out diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py deleted file mode 100644 index e3d7d38f..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -from vllm.triton_utils import tl, triton - - -def save_partial_states( - kv: torch.Tensor, - score: torch.Tensor, - ape: torch.Tensor, - positions: torch.Tensor, - state_cache: torch.Tensor, - slot_mapping: torch.Tensor, - block_size: int, - state_width: int, - compress_ratio: int, - pdl_kwargs: dict | None = None, -) -> None: - """Write packed [kv, score+ape] partial states into the compressor cache. - - One program per token; pads (slot_id == -1) are skipped. - """ - num_actual = slot_mapping.shape[0] - head_size = kv.shape[-1] - _save_partial_states_kernel[(num_actual,)]( - kv, - kv.stride(0), - score, - score.stride(0), - ape, - ape.stride(0), - positions, - state_cache, - state_cache.stride(0), - state_cache.stride(1), - slot_mapping, - block_size, - HEAD_SIZE=head_size, - TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size), - STATE_WIDTH=state_width, - COMPRESS_RATIO=compress_ratio, - **(pdl_kwargs or {}), - ) - - -@triton.jit -def _save_partial_states_kernel( - kv_ptr, - kv_stride, - score_ptr, - score_stride, - ape_ptr, - ape_stride, - positions_ptr, - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - slot_mapping_ptr, - block_size, - HEAD_SIZE: tl.constexpr, - TRITON_BLOCK_SIZE: tl.constexpr, - # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide. - STATE_WIDTH: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, -): - token_idx = tl.program_id(0) - slot_id = tl.load(slot_mapping_ptr + token_idx) - - # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used - # by vLLM). During CUDA graph replay the batch may contain padding - # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an - # illegal memory access. - if slot_id < 0: - return - - block_idx = slot_id // block_size - pos_in_block = slot_id % block_size - base_ptr = ( - state_cache_ptr - + block_idx * state_cache_stride0 - + pos_in_block * state_cache_stride1 - ) - - block = tl.arange(0, TRITON_BLOCK_SIZE) - mask = block < HEAD_SIZE - - kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask) - tl.store(base_ptr + block, kv, mask=mask) - - # Fused: score += ape[position % compress_ratio] - position = tl.load(positions_ptr + token_idx) - ape_row = position % COMPRESS_RATIO - ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask) - score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask) - tl.store( - base_ptr + STATE_WIDTH + block, - score + ape, - mask=mask, - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py deleted file mode 100644 index 44ae3286..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py +++ /dev/null @@ -1,36 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepseekV4 rotary embedding initialization.""" - -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding - - -def build_deepseek_v4_rope( - config, - *, - head_dim: int, - rope_head_dim: int, - max_position_embeddings: int, - compress_ratio: int, -) -> RotaryEmbedding: - rope_parameters = config.rope_parameters - rope_parameters["rope_theta"] = ( - config.compress_rope_theta if compress_ratio > 1 else config.rope_theta - ) - if rope_parameters["rope_type"] != "default": - rope_parameters["rope_type"] = ( - "deepseek_yarn" - if rope_parameters.get("apply_yarn_scaling", True) - else "deepseek_llama_scaling" - ) - rope_parameters["mscale"] = 0 # Disable mscale - rope_parameters["mscale_all_dim"] = 0 # Disable mscale - rope_parameters["is_deepseek_v4"] = True - rope_parameters["rope_dim"] = rope_head_dim - return get_rope( - head_dim, - max_position=max_position_embeddings, - rope_parameters=rope_parameters, - is_neox_style=False, - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/compressor.py b/TEMP/deepseek_v4_ref/deepseek_v4/compressor.py deleted file mode 100644 index f36dc8f1..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/compressor.py +++ /dev/null @@ -1,380 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from dataclasses import dataclass -from typing import Any, ClassVar, cast - -import torch -from torch import nn - -from vllm.config import VllmConfig, get_current_vllm_config -from vllm.forward_context import get_forward_context -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import MergedColumnParallelLinear -from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( - compress_norm_rope_store_triton, -) -from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE -from vllm.models.deepseek_v4.common.ops.save_partial_states import ( - save_partial_states, -) -from vllm.platforms import current_platform -from vllm.v1.attention.backend import ( - AttentionBackend, - AttentionCGSupport, - AttentionMetadataBuilder, - CommonAttentionMetadata, - MultipleOf, -) -from vllm.v1.kv_cache_interface import ( - KVCacheSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, -) - - -class CompressorBackend(AttentionBackend): - def __init__(self): - super().__init__() - - @staticmethod - def get_name() -> str: - return "CompressorBackend" - - @staticmethod - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [MultipleOf(1)] - - @classmethod - def get_supported_head_sizes(cls) -> list[int]: - return [512, 1024] - - @staticmethod - def get_builder_cls() -> type["CompressorMetadataBuilder"]: - return CompressorMetadataBuilder - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - assert num_kv_heads == 1 - return (num_blocks, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (0, 1, 2, 3) - return (0, 1, 2) - - -@dataclass -class CompressorMetadata: - block_table: torch.Tensor - slot_mapping: torch.Tensor - block_size: int - - token_to_req_indices: torch.Tensor | None = None # [num_tokens] - - -class CompressorMetadataBuilder(AttentionMetadataBuilder): - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) - mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) - self.block_size = mla_spec.block_size - - self.token_to_req_indices = torch.zeros( - self.vllm_config.scheduler_config.max_num_batched_tokens, - dtype=torch.int32, - device=self.device, - ) - - def build( - self, - common_prefix_len: int, - common_attn_metadata: CommonAttentionMetadata, - fast_build: bool = False, - ) -> CompressorMetadata: - query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu - num_reqs = common_attn_metadata.num_reqs - query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() - token_to_req_indices = self.token_to_req_indices[: x.shape[0]] - token_to_req_indices.copy_(x, non_blocking=True) - return CompressorMetadata( - block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), - slot_mapping=common_attn_metadata.slot_mapping, - block_size=self.block_size, - token_to_req_indices=token_to_req_indices, - ) - - -class CompressorStateCache(torch.nn.Module, AttentionLayerBase): - def __init__( - self, - state_dim: int, - dtype: torch.dtype, - compress_ratio: int, - prefix: str, - ): - super().__init__() - self.state_dim = state_dim - self.dtype = dtype - self.prefix = prefix - self.kv_cache = torch.tensor([]) - compilation_config = get_current_vllm_config().compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - - assert self.dtype == torch.float32 - assert compress_ratio in [4, 128] - coff = 1 + (compress_ratio == 4) - self.sliding_window = coff * compress_ratio - # Block size is constrained by tensor sharing between compressor states - # and KV blocks. Since compressor states share the same physical tensor - # as KV blocks, they must use the same page size. - # The KV block shape [256//4, head_dim] = [64, 584] determines: - # - C4 compressor block shape [4, 2*512*2*4] -> block_size = 4 - # - C128 compressor block shape [8, 512*2*4] -> block_size = 8 - # TODO(yifan): make block size automatically determined and configurable. - if compress_ratio == 4: - self.block_size = 4 - elif compress_ratio == 128: - self.block_size = 8 - else: - raise ValueError(f"Invalid compress ratio: {compress_ratio}") - - def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - return SlidingWindowMLASpec( # only has one vector instead of K + V - block_size=self.block_size, - num_kv_heads=1, - head_size=self.state_dim, - dtype=self.dtype, - sliding_window=self.sliding_window, - alignment=576, # NOTE: FlashMLA requires 576B alignment - ) - - def forward(self): ... - - def get_attn_backend(self) -> type[AttentionBackend]: - return CompressorBackend - - -class DeepseekCompressor(nn.Module): - """DeepSeek V4 KV/score compressor. - - Owns the linear / norm / state-cache / ape state and the shared forward - prologue (kv/score split, save_partial_states launch). The - compress → norm → RoPE → store step is dispatched to a triton kernel - (``compress_norm_rope_store_triton``) by default, except for the NVIDIA - head_dim=128 indexer path which uses the cutedsl kernel - (``compress_norm_rope_store_cutedsl``) for better performance. - """ - - def __init__( - self, - vllm_config: VllmConfig, - compress_ratio: int, - hidden_size: int, - head_dim: int, - rotate: bool = False, - prefix: str = "", - k_cache_prefix="", - use_fp4_cache: bool = False, - ): - super().__init__() - self.compress_ratio = compress_ratio - self.hidden_size = hidden_size - self.head_dim = head_dim - self.rotate = rotate - self.prefix = prefix - self.k_cache_prefix = k_cache_prefix - self.use_fp4_cache = use_fp4_cache - - config = vllm_config.model_config.hf_config - self.rope_head_dim = config.qk_rope_head_dim - self.nope_head_dim = self.head_dim - self.rope_head_dim - self.rms_norm_eps = config.rms_norm_eps - self.device = current_platform.device_type - self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs - self.max_model_len = vllm_config.model_config.max_model_len - - self.overlap = compress_ratio == 4 - self.coff = 1 + self.overlap - - state_dtype = torch.float32 - self.ape = nn.Parameter( - torch.empty( - (compress_ratio, self.coff * self.head_dim), - dtype=state_dtype, - device=self.device, - ), - requires_grad=False, - ) - - self.fused_wkv_wgate = MergedColumnParallelLinear( - self.hidden_size, - [self.coff * self.head_dim, self.coff * self.head_dim], - bias=False, - return_bias=False, - quant_config=None, - disable_tp=True, - prefix=f"{prefix}.fused_wkv_wgate", - ) - self.norm = RMSNorm(self.head_dim, self.rms_norm_eps) - - self.state_cache = CompressorStateCache( - state_dim=2 * self.coff * self.head_dim, # kv_state + score_state - dtype=state_dtype, - compress_ratio=compress_ratio, - prefix=f"{prefix}.state_cache", - ) - - # Save reference to static_forward_context for forward-time KV cache lookup. - # get_current_vllm_config() is only available during __init__, not forward. - self._static_forward_context = ( - vllm_config.compilation_config.static_forward_context - ) - - if self.head_dim == 512: - assert not use_fp4_cache, ( - "MXFP4 cache is only supported for indexer (head=128)" - ) - self._quant_block = 64 - self._token_stride = self.nope_head_dim + self.rope_head_dim * 2 - self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad - elif self.head_dim == 128: - if use_fp4_cache: - self._quant_block = MXFP4_BLOCK_SIZE - self._token_stride = self.head_dim // 2 - self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE - else: - self._quant_block = 128 - self._token_stride = self.head_dim - self._scale_dim = 4 # single float32 scale - else: - raise ValueError( - f"Unsupported head_dim for fused quant+cache: {self.head_dim}" - ) - - def forward( - self, - # [num_tokens, 2 * self.coff * self.head_dim] - kv_score: torch.Tensor, - # [num_tokens] - positions: torch.Tensor, - rotary_emb, - ) -> None: - # Each of shape [num_tokens, coff * self.head_dim] - # input bf16, output are fp32 - kv, score = kv_score.split( - [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1 - ) - - # Get the metadata and handle dummy profiling run. - attn_metadata = get_forward_context().attn_metadata - if not isinstance(attn_metadata, dict): - return - - state_metadata = cast( - CompressorMetadata, attn_metadata[self.state_cache.prefix] - ) - token_to_req_indices = state_metadata.token_to_req_indices - slot_mapping = state_metadata.slot_mapping - num_actual = slot_mapping.shape[0] - block_table = state_metadata.block_table - block_size = state_metadata.block_size - - # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim - state_cache = self.state_cache.kv_cache - # kv_state stored in first half, score_state stored in second half - state_width = state_cache.shape[-1] // 2 - pdl_kwargs = ( - {} - if current_platform.is_rocm() or current_platform.is_xpu() - else {"launch_pdl": False} - ) - - # Store the KV and score (with fused APE addition) in the state. - # NOTE: PDL is disabled — both this kernel and the compress kernels - # below depend on preceding kernel outputs (kv/score from the cublas - # GEMM; state_cache from this kernel) but neither emits/waits on PDL - # grid dependency primitives, so launch_pdl=True caused a - # read-after-write race and non-deterministic output. - save_partial_states( - kv=kv, - score=score, - ape=self.ape, - positions=positions, - state_cache=state_cache, - slot_mapping=slot_mapping, - block_size=block_size, - state_width=state_width, - compress_ratio=self.compress_ratio, - pdl_kwargs=pdl_kwargs, - ) - - # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write. - # RoPE requirements (kernel applies forward GPT-J style rotation): - # - is_neox_style=False (interleaved pairs, NOT split-half) - # - cos_sin_cache layout: [max_pos, rope_head_dim] with first half cos, - # second half sin (per-pair, length rope_head_dim // 2 each) - # - applied to LAST rope_head_dim elements of head_dim - # - position used: (positions // compress_ratio) * compress_ratio - cos_sin_cache = rotary_emb.cos_sin_cache - k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) - kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache - - if current_platform.is_cuda(): - # NVIDIA GPUs. - if self.head_dim == 512: - from .nvidia.ops.sparse_attn_compress_cutedsl import ( - compress_norm_rope_store_cutedsl, - ) - - # Main compressor path. - # Use a cutedsl kernel for better performance. - compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl - else: - # Indexer path (head_dim == 128). - # Use a triton kernel. - compress_norm_rope_store_fn = compress_norm_rope_store_triton - else: - # AMD GPUs. - # Always use a triton kernel. - compress_norm_rope_store_fn = compress_norm_rope_store_triton - - compress_norm_rope_store_fn( - state_cache=state_cache, - num_actual=num_actual, - token_to_req_indices=token_to_req_indices, - positions=positions, - slot_mapping=slot_mapping, - block_table=block_table, - block_size=block_size, - state_width=state_width, - cos_sin_cache=cos_sin_cache, - kv_cache=kv_cache, - k_cache_metadata=k_cache_metadata, - pdl_kwargs=pdl_kwargs, - head_dim=self.head_dim, - rope_head_dim=self.rope_head_dim, - compress_ratio=self.compress_ratio, - overlap=self.overlap, - use_fp4_cache=self.use_fp4_cache, - rms_norm_weight=self.norm.weight, - rms_norm_eps=self.rms_norm_eps, - quant_block=self._quant_block, - token_stride=self._token_stride, - scale_dim=self._scale_dim, - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py deleted file mode 100644 index 208f01a7..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py deleted file mode 100644 index 5c8b08d4..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py +++ /dev/null @@ -1,424 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from abc import abstractmethod -from typing import TYPE_CHECKING, ClassVar, cast - -import torch - -from vllm.forward_context import get_forward_context -from vllm.models.deepseek_v4.common.ops import ( - combine_topk_swa_indices, - compute_global_topk_indices_and_lens, - dequantize_and_gather_k_cache, -) -from vllm.v1.attention.backend import ( - AttentionBackend, - MultipleOf, - SparseMLAAttentionImpl, -) -from vllm.v1.attention.backends.mla.flashmla_sparse import ( - FlashMLASparseBackend, - FlashMLASparseMetadata, -) -from vllm.v1.attention.ops.flashmla import ( - flash_mla_sparse_fwd, - flash_mla_with_kvcache, -) -from vllm.v1.worker.workspace import current_workspace_manager - -if TYPE_CHECKING: - from vllm.models.deepseek_v4.attention import ( - DeepseekV4MLAAttention, - ) - from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata - - -class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]): - """Abstract parent for DeepseekV4 sparse MLA impls. - - V4 sparse MLA is driven by the layer (``DeepseekV4MLAAttention.forward``) - rather than the v1 framework, so ``forward_mqa`` is overridden with a - classmethod that takes the layer as its first argument. This Liskov-broken - override is intentional: the grandparent's instance-method ``forward_mqa`` - is never called on V4 layers. - """ - - backend_cls: ClassVar[type[AttentionBackend]] - - # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather - # workspace allocated in _forward_prefill and is also read by the V4 layer's - # dummy-run path to pre-reserve that workspace. - PREFILL_CHUNK_SIZE: ClassVar[int] = 4 - - @classmethod - @abstractmethod - def forward_mqa( # type: ignore[override] - cls, - layer: "DeepseekV4MLAAttention", - q: torch.Tensor, - kv: torch.Tensor, - positions: torch.Tensor, - output: torch.Tensor, - ) -> None: - raise NotImplementedError - - @classmethod - @abstractmethod - def get_padded_num_q_heads(cls, num_heads: int) -> int: - """Q head count the backend wants q allocated at. - - The MLA wrapper allocates the q/output buffers at - ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must - satisfy ``result >= num_heads``. Backends with no padding constraint - return ``num_heads``. - """ - raise NotImplementedError - - -class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend): - @staticmethod - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [256] - - @staticmethod - def get_name() -> str: - return "V4_FLASHMLA_SPARSE" - - @staticmethod - def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]: - return DeepseekV4FlashMLASparseImpl - - @classmethod - def get_supported_head_sizes(cls) -> list[int]: - # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512 (overrides the - # V3.2 default of 576 from FlashMLASparseBackend). - return [512] - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if cache_dtype_str == "fp8_ds_mla": - # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale). - # head_size passed in is the semantic head_dim (512). - return (num_blocks, block_size, 584) - else: - return (num_blocks, block_size, head_size) - - -class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl): - """FlashMLA sparse MLA implementation for DeepSeek V4's custom MLA layer.""" - - backend_cls = DeepseekV4FlashMLASparseBackend - - @classmethod - def get_padded_num_q_heads(cls, num_heads: int) -> int: - # FP8 decode kernel only supports h_q = 64 or 128. - if num_heads > 128: - raise ValueError( - f"DeepseekV4 FlashMLA does not support {num_heads} heads " - "(FP8 decode kernel requires h_q in {64, 128})." - ) - return 64 if num_heads <= 64 else 128 - - @classmethod - def forward_mqa( # type: ignore[override] - cls, - layer: "DeepseekV4MLAAttention", - q: torch.Tensor, - kv: torch.Tensor, - positions: torch.Tensor, - output: torch.Tensor, - ) -> None: - assert output.shape == q.shape, ( - f"output buffer shape {output.shape} must match q shape {q.shape}" - ) - assert output.dtype == q.dtype, ( - f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" - ) - - # Get SWA and indexer metadata from forward context - forward_context = get_forward_context() - attn_metadata = forward_context.attn_metadata - - if attn_metadata is None: - # Warmup dummy run: no real metadata. Reserve the same bf16 - # gather workspace _forward_prefill would; the dequantize / topk - # / sparse_fwd kernels are skipped this step. - swa_only = layer.compress_ratio <= 1 - N = ( - 0 - if swa_only - else (layer.max_model_len + layer.compress_ratio - 1) - // layer.compress_ratio - ) - M = N + layer.window_size + layer.max_num_batched_tokens - current_workspace_manager().get_simultaneous( - ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), - ) - output.zero_() - return - - assert isinstance(attn_metadata, dict) - flashmla_metadata = cast( - FlashMLASparseMetadata | None, attn_metadata.get(layer.prefix) - ) - swa_metadata = cast( - "DeepseekSparseSWAMetadata | None", - attn_metadata.get(layer.swa_cache_layer.prefix), - ) - assert swa_metadata is not None - - swa_only = layer.compress_ratio <= 1 - # SWA-only layers (compress_ratio <= 1) don't have their own KV cache - # allocation, so layer.kv_cache may be empty after profiling cleanup. - self_kv_cache = layer.kv_cache if not swa_only else None - swa_kv_cache = layer.swa_cache_layer.kv_cache - - # Split prefill and decode - num_decodes = swa_metadata.num_decodes - num_prefills = swa_metadata.num_prefills - num_decode_tokens = swa_metadata.num_decode_tokens - - if num_prefills > 0: - cls._forward_prefill( - layer=layer, - q=q[num_decode_tokens:], - positions=positions[num_decode_tokens:], - compressed_k_cache=self_kv_cache, - swa_k_cache=swa_kv_cache, - output=output[num_decode_tokens:], - attn_metadata=flashmla_metadata, - swa_metadata=swa_metadata, - ) - if num_decodes > 0: - cls._forward_decode( - layer=layer, - q=q[:num_decode_tokens], - kv_cache=self_kv_cache, - swa_metadata=swa_metadata, - attn_metadata=flashmla_metadata, - swa_only=swa_only, - output=output[:num_decode_tokens], - ) - - @classmethod - def _forward_decode( - cls, - layer: "DeepseekV4MLAAttention", - q: torch.Tensor, - kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1 - swa_metadata: "DeepseekSparseSWAMetadata", - attn_metadata: FlashMLASparseMetadata | None, - swa_only: bool, - output: torch.Tensor, - ) -> None: - num_decodes = swa_metadata.num_decodes - num_decode_tokens = swa_metadata.num_decode_tokens - - topk_indices = None - topk_lens = None - if not swa_only: - assert attn_metadata is not None - assert swa_metadata.is_valid_token is not None - block_size = attn_metadata.block_size // layer.compress_ratio - is_valid = swa_metadata.is_valid_token[:num_decode_tokens] - if layer.compress_ratio == 4: - # C4A: local indices differ per layer (filled by Indexer). - assert layer.topk_indices_buffer is not None - global_indices, topk_lens = compute_global_topk_indices_and_lens( - layer.topk_indices_buffer[:num_decode_tokens], - swa_metadata.token_to_req_indices, - attn_metadata.block_table[:num_decodes], - block_size, - is_valid, - ) - topk_indices = global_indices.view(num_decode_tokens, 1, -1) - else: - # C128A: pre-computed during metadata build. - topk_indices = attn_metadata.c128a_global_decode_topk_indices - topk_lens = attn_metadata.c128a_decode_topk_lens - - swa_indices = swa_metadata.decode_swa_indices - swa_lens = swa_metadata.decode_swa_lens - - # We treat queries in the same seq as different queries - # and later we only attend by generated indices. - # q arrives pre-padded to layer.padded_heads by the outer wrapper. - q = q.unsqueeze(1) - - # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes) - # Use unsqueeze to preserve strides (handles padded blocks correctly) - swa_cache = layer.swa_cache_layer.kv_cache.unsqueeze(-2) - # Reshape KV cache to (num_blocks, block_size, 1, head_bytes) - if kv_cache is not None: - kv_cache = kv_cache.unsqueeze(-2) - - # One FlashMLASchedMeta per layer type, shared across all same-type - # layers within this decode step. The first forward call per type - # triggers the in-kernel planner (allocating tile_scheduler_metadata - # and num_splits via PyTorch's graph-aware allocator so CUDA graph - # capture reuses the same addresses on replay); subsequent same-type - # layers see have_initialized=True and skip the planner. - if layer.compress_ratio <= 1: - tile_metadata = swa_metadata.tile_sched_swaonly - elif layer.compress_ratio == 4: - tile_metadata = swa_metadata.tile_sched_c4a - elif layer.compress_ratio == 128: - tile_metadata = swa_metadata.tile_sched_c128a - else: - raise ValueError( - f"Unsupported compress_ratio={layer.compress_ratio}; " - "expected 1, 4, or 128." - ) - assert tile_metadata is not None, ( - "swa_metadata missing tile_sched entry for " - f"compress_ratio={layer.compress_ratio}; " - "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not " - "allocate one for this layer type." - ) - - out, _ = flash_mla_with_kvcache( - q=q, - k_cache=swa_cache, - block_table=None, - head_dim_v=512, - tile_scheduler_metadata=tile_metadata, - cache_seqlens=None, - is_fp8_kvcache=True, - indices=swa_indices, - topk_length=swa_lens, - softmax_scale=layer.scale, - attn_sink=layer.attn_sink, - extra_k_cache=kv_cache if not swa_only else None, - extra_indices_in_kvcache=topk_indices, - extra_topk_length=topk_lens, - out=output.unsqueeze(1), - ) - - @classmethod - def _forward_prefill( - cls, - layer: "DeepseekV4MLAAttention", - q: torch.Tensor, - positions: torch.Tensor, - compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1 - swa_k_cache: torch.Tensor, - output: torch.Tensor, - attn_metadata: FlashMLASparseMetadata | None, - swa_metadata: "DeepseekSparseSWAMetadata", - ) -> None: - swa_only = attn_metadata is None - - num_prefills = swa_metadata.num_prefills - num_prefill_tokens = swa_metadata.num_prefill_tokens - num_decodes = swa_metadata.num_decodes - num_decode_tokens = swa_metadata.num_decode_tokens - - # Use pre-computed prefill metadata. - seq_lens = swa_metadata.prefill_seq_lens - gather_lens = swa_metadata.prefill_gather_lens - assert seq_lens is not None - assert gather_lens is not None - - # Derive prefill-local token offsets from the full query_start_loc_cpu. - query_start_loc_cpu = swa_metadata.query_start_loc_cpu - query_start_loc = swa_metadata.query_start_loc - assert query_start_loc_cpu is not None - assert query_start_loc is not None - prefill_token_base = query_start_loc_cpu[num_decodes] - - if not swa_only: - if layer.compress_ratio == 4: - assert layer.topk_indices_buffer is not None - topk_indices = layer.topk_indices_buffer[num_decode_tokens:] - topk_indices = topk_indices[:num_prefill_tokens] - else: - # C128A: pre-computed during metadata build. - assert attn_metadata is not None - topk_indices = attn_metadata.c128a_prefill_topk_indices - top_k = topk_indices.shape[-1] - # Compressed region must fit the full compressed pool (seq_len // - # compress_ratio), not just top_k. top_k bounds how many indices - # the indexer selects, not the pool size it indexes into. - N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio - else: - # NOTE(woosuk): topk_indices will not be used for SWA-only layers. - assert layer.topk_indices_buffer is not None - topk_indices = layer.topk_indices_buffer[num_decode_tokens:] - top_k = 0 - N = 0 - - M = N + layer.window_size + layer.max_num_batched_tokens - chunk_size_const = cls.PREFILL_CHUNK_SIZE - num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const - - workspace_manager = current_workspace_manager() - kv = workspace_manager.get_simultaneous( - ((chunk_size_const, M, q.shape[-1]), torch.bfloat16), - )[0] - for chunk_idx in range(num_chunks): - chunk_start = chunk_idx * chunk_size_const - chunk_end = min(chunk_start + chunk_size_const, num_prefills) - chunk_size = chunk_end - chunk_start - if not swa_only: - # Gather compressed KV - assert attn_metadata is not None - block_table = attn_metadata.block_table[num_decodes:] - dequantize_and_gather_k_cache( - kv[:chunk_size], - compressed_k_cache, - seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio, - gather_lens=None, - block_table=block_table[chunk_start:chunk_end], - block_size=attn_metadata.block_size // layer.compress_ratio, - offset=0, - ) - - # Gather SWA KV - swa_block_table = swa_metadata.block_table[num_decodes:] - dequantize_and_gather_k_cache( - kv[:chunk_size], - swa_k_cache, - seq_lens=seq_lens[chunk_start:chunk_end], - gather_lens=gather_lens[chunk_start:chunk_end], - block_table=swa_block_table[chunk_start:chunk_end], - block_size=swa_metadata.block_size, - offset=N, - ) - - # Combine the topk indices and SWA indices for gathered KV cache - query_start = ( - query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base - ) - query_end = ( - query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base - ) - - combined_indices, combined_lens = combine_topk_swa_indices( - topk_indices[query_start:query_end], - query_start_loc[ - num_decodes + chunk_start : num_decodes + chunk_end + 1 - ], - seq_lens[chunk_start:chunk_end], - gather_lens[chunk_start:chunk_end], - layer.window_size, - layer.compress_ratio, - top_k, - M, - N, - ) - flash_mla_sparse_fwd( - q=q[query_start:query_end], - kv=kv.view(-1, 1, q.shape[-1]), - indices=combined_indices.unsqueeze(1), - sm_scale=layer.scale, - attn_sink=layer.attn_sink, - topk_length=combined_lens, - out=output[query_start:query_end], - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py deleted file mode 100644 index 547048ab..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py +++ /dev/null @@ -1,1476 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import typing -from collections.abc import Callable, Iterable, MutableSequence, Sequence -from itertools import islice - -import regex as re -import torch -import torch.nn as nn - -from vllm.config import VllmConfig -from vllm.distributed import ( - get_ep_group, - get_pp_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.distributed.eplb.eplb_state import EplbLayerState -from vllm.model_executor.kernels.mhc.tilelang import ( - hc_head_fused_kernel_tilelang, - mhc_fused_post_pre_tilelang, - mhc_post_tilelang, - mhc_pre_tilelang, -) -from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp -from vllm.model_executor.layers.fused_moe import FusedMoE -from vllm.model_executor.layers.fused_moe.router.base_router import ( - eplb_map_to_physical_and_record, -) -from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( - fused_topk_bias, -) -from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - MergedColumnParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - PPMissingLayer, - WeightsMapper, - extract_layer_index, - is_pp_missing_parameter, - make_layers, - maybe_prefix, -) -from vllm.model_executor.utils import set_weight_attrs -from vllm.models.deepseek_v4.attention import ( - DeepseekV4Indexer, - DeepseekV4MLA, -) -from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope -from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs -from vllm.sequence import IntermediateTensors - - -class DeepseekV4MLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - swiglu_limit: float | None = None, - quant_config: QuantizationConfig | None = None, - reduce_results: bool = True, - is_sequence_parallel: bool = False, - prefix: str = "", - ) -> None: - super().__init__() - - # If is_sequence_parallel, the input and output tensors are sharded - # across the ranks within the tp_group. In this case the weights are - # replicated and no collective ops are needed. - # Otherwise we use standard TP with an allreduce at the end. - self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, - [intermediate_size] * 2, - bias=False, - quant_config=quant_config, - disable_tp=is_sequence_parallel, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=False, - quant_config=quant_config, - reduce_results=reduce_results, - disable_tp=is_sequence_parallel, - prefix=f"{prefix}.down_proj", - ) - if hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported for now." - ) - if swiglu_limit is not None: - self.act_fn = SiluAndMulWithClamp(swiglu_limit) - else: - self.act_fn = SiluAndMul() - - def forward(self, x): - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -def make_deepseek_v4_expert_params_mapping( - num_experts: int, -) -> list[tuple[str, str, int, str]]: - return [ - ( - "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_", - f"experts.{expert_id}.{weight_name}.", - expert_id, - shard_id, - ) - for expert_id in range(num_experts) - for shard_id, weight_name in [ - ("w1", "w1"), - ("w2", "w2"), - ("w3", "w3"), - ] - ] - - -class DeepseekV4MegaMoEExperts(nn.Module): - _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} - - def __init__( - self, - vllm_config: VllmConfig, - *, - num_experts: int, - num_local_experts: int, - experts_start_idx: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - prefix: str = "", - num_logical_experts: int | None = None, - ): - super().__init__() - self.prefix = prefix - self.num_experts = num_experts - self.num_local_experts = num_local_experts - self.experts_start_idx = experts_start_idx - self.experts_end_idx = experts_start_idx + num_local_experts - self.top_k = top_k - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens - - self.num_logical_experts = ( - num_logical_experts if num_logical_experts is not None else num_experts - ) - - self.eplb_state = EplbLayerState() - - weight_attrs = {"weight_loader": self.weight_loader} - self.w13_weight = nn.Parameter( - torch.zeros( - num_local_experts, - 2 * intermediate_size, - hidden_size // 2, - dtype=torch.uint8, - ), - requires_grad=False, - ) - set_weight_attrs(self.w13_weight, weight_attrs) - - self.w13_weight_scale = nn.Parameter( - torch.zeros( - num_local_experts, - 2 * intermediate_size, - hidden_size // 32, - dtype=torch.uint8, - ), - requires_grad=False, - ) - set_weight_attrs(self.w13_weight_scale, weight_attrs) - self.w13_weight_scale.quant_method = "block" - - self.w2_weight = nn.Parameter( - torch.zeros( - num_local_experts, - hidden_size, - intermediate_size // 2, - dtype=torch.uint8, - ), - requires_grad=False, - ) - set_weight_attrs(self.w2_weight, weight_attrs) - - self.w2_weight_scale = nn.Parameter( - torch.zeros( - num_local_experts, - hidden_size, - intermediate_size // 32, - dtype=torch.uint8, - ), - requires_grad=False, - ) - set_weight_attrs(self.w2_weight_scale, weight_attrs) - self.w2_weight_scale.quant_method = "block" - - self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None - self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None - - # Register in the static forward context so the custom-op wrapper - # can look up this module by name from within a torch.compile graph. - compilation_config = vllm_config.compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - - def _map_global_expert_id(self, expert_id: int) -> list[int]: - """Return local (per-rank) slot offsets where logical expert - `expert_id` should land on this rank. - """ - physical_ids: list[int] = [] - for p in range(self.experts_start_idx, self.experts_end_idx): - if p % self.num_logical_experts == expert_id: - physical_ids.append(p - self.experts_start_idx) - return physical_ids - - def weight_loader( - self, - param: nn.Parameter, - loaded_weight: torch.Tensor, - weight_name: str, - shard_id: str, - expert_id: int, - return_success: bool = False, - ) -> bool | None: - local_expert_ids = self._map_global_expert_id(expert_id) - if not local_expert_ids: - return False if return_success else None - - loaded_any = False - for local_expert_id in local_expert_ids: - expert_data = param.data[local_expert_id] - if shard_id in ("w1", "w3"): - if "w13_" not in weight_name: - continue - shard_offset = 0 if shard_id == "w1" else self.intermediate_size - expert_data = expert_data.narrow( - 0, shard_offset, self.intermediate_size - ) - elif shard_id == "w2": - if "w2_" not in weight_name: - continue - else: - raise ValueError(f"Unsupported expert shard id: {shard_id}") - - if expert_data.shape != loaded_weight.shape: - raise ValueError( - f"DeepSeek V4 MegaMoE expert weight shape mismatch for " - f"{weight_name}: parameter shard {tuple(expert_data.shape)} " - f"vs checkpoint {tuple(loaded_weight.shape)}" - ) - expert_data.copy_(loaded_weight) - loaded_any = True - - if return_success: - return loaded_any - return None - - @staticmethod - def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: - return (sf.to(torch.int32) << 23).view(torch.float32) - - def _check_runtime_supported(self) -> None: - device = self.w13_weight.device - if torch.cuda.get_device_capability(device)[0] != 10: - raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") - if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: - raise ValueError( - "DeepGEMM MegaMoE requires hidden and intermediate sizes " - "to be multiples of 128." - ) - - def finalize_weights(self) -> None: - if self._transformed_l1_weights is not None: - return - - self._check_runtime_supported() - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - - w13_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), - 2 * self.intermediate_size, - self.hidden_size, - (1, 32), - self.num_local_experts, - ) - w2_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), - self.hidden_size, - self.intermediate_size, - (1, 32), - self.num_local_experts, - ) - self._transformed_l1_weights, self._transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe( - (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), - (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), - ) - ) - # Drop the original loader-side parameters: the MegaMoE kernels only - # consume the transformed views above. transform_weights_for_mega_moe - # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights) - # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that - # aliases the original storage, and _transformed_l2_weights still holds - # it, so the storage stays live after we drop the Parameter. - self.w13_weight = None - self.w13_weight_scale = None - self.w2_weight = None - self.w2_weight_scale = None - - def get_symm_buffer(self): - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - - group = get_ep_group().device_group - device = torch.accelerator.current_device_index() - key = ( - id(group), - device, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - ) - symm_buffer = self._symm_buffer_cache.get(key) - if symm_buffer is None: - symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( - group, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - ) - self._symm_buffer_cache[key] = symm_buffer - return symm_buffer - - def set_eplb_state( - self, - moe_layer_idx: int, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ) -> None: - self.eplb_state.set_layer_state( - moe_layer_idx, - expert_load_view, - logical_to_physical_map, - logical_replica_count, - ) - - def get_expert_weights(self) -> list[torch.Tensor]: - self.finalize_weights() - assert self._transformed_l1_weights is not None - assert self._transformed_l2_weights is not None - - def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: - """Return a (num_local_experts, -1) view with contiguous memory layout.""" - assert t.shape[0] == self.num_local_experts - if t.is_contiguous(): - return t.view(self.num_local_experts, -1) - elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: - # scales have shape (E, M, N) with memory layout (E, N, M) - back = torch.transpose(t, 1, 2) - assert back.is_contiguous() - return back.view(self.num_local_experts, -1) - - raise AssertionError( - f"DSv4 EPLB {name}: non-contiguous expert tensor with " - f"unexpected layout shape={tuple(t.shape)} " - f"stride={tuple(t.stride())} dtype={t.dtype}" - ) - - return [ - _to_eplb_view("l1_packed", self._transformed_l1_weights[0]), - _to_eplb_view("l1_scale", self._transformed_l1_weights[1]), - _to_eplb_view("l2_weight", self._transformed_l2_weights[0]), - _to_eplb_view("l2_scale", self._transformed_l2_weights[1]), - ] - - def update_expert_map(self) -> None: - pass - - def forward( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - *, - activation_clamp: float | None, - fast_math: bool = True, - ) -> torch.Tensor: - if hidden_states.shape[0] > self.max_num_tokens: - raise ValueError( - f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " - f"but the symmetric buffer was sized for {self.max_num_tokens}." - ) - y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - - symm_buffer = self.get_symm_buffer() - num_tokens = hidden_states.shape[0] - - # EPLB: map logical expert IDs to physical replicas and record load. - eplb_state = self.eplb_state - if eplb_state.logical_to_physical_map is not None: - assert eplb_state.expert_load_view is not None - assert eplb_state.logical_replica_count is not None - assert eplb_state.should_record_tensor is not None - topk_ids = eplb_map_to_physical_and_record( - topk_ids=topk_ids, - expert_load_view=eplb_state.expert_load_view, - logical_to_physical_map=eplb_state.logical_to_physical_map, - logical_replica_count=eplb_state.logical_replica_count, - record_enabled=eplb_state.should_record_tensor, - ) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens], - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - ) - - # This method must have been already called during the weight loading phase. - # We call it again here to cover the dummy weight loading case. - self.finalize_weights() - - assert self._transformed_l1_weights is not None - assert self._transformed_l2_weights is not None - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, - activation_clamp=activation_clamp, - fast_math=fast_math, - ) - return y - - -DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] - - -class DeepseekV4MoE(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - prefix: str = "", - ): - super().__init__() - - self.tp_size = get_tensor_model_parallel_world_size() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.prefix = prefix - self.use_mega_moe = ( - vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" - ) - if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: - raise NotImplementedError( - "DeepSeek V4 MegaMoE currently requires expert parallel. " - "Enable it with --enable-expert-parallel, or pick a different " - "moe backend." - ) - - self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) - self.hidden_size = config.hidden_size - - self.n_routed_experts = config.n_routed_experts - self.n_activated_experts = config.num_experts_per_tok - self.moe_intermediate_size = config.moe_intermediate_size - self.swiglu_limit = config.swiglu_limit - self.renormalize = config.norm_topk_prob - self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") - if self.use_mega_moe and self.scoring_func != "sqrtsoftplus": - raise NotImplementedError( - "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." - ) - if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": - raise NotImplementedError( - "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" - f"{config.expert_dtype!r}. Drop --kernel-config moe_backend=" - "deep_gemm_mega_moe for this checkpoint." - ) - - self.gate = GateLinear( - input_size=config.hidden_size, - output_size=config.n_routed_experts, - bias=False, - out_dtype=torch.float32, - prefix=f"{prefix}.gate", - ) - - self.gate.e_score_correction_bias = None - self.gate.tid2eid = None - is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers - self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 - if is_hash_moe: - # hash MoE doesn't use e_score_correction_bias - # Use randint instead of empty to avoid garbage values causing - # invalid memory access in dummy mode (--load-format="dummy") - self.gate.tid2eid = nn.Parameter( - torch.randint( - 0, - config.n_routed_experts, - (config.vocab_size, config.num_experts_per_tok), - dtype=self.hash_indices_dtype, - ), - requires_grad=False, - ) - elif getattr(config, "topk_method", None) == "noaux_tc": - self.gate.e_score_correction_bias = nn.Parameter( - torch.empty(config.n_routed_experts, dtype=torch.float32), - requires_grad=False, - ) - - if config.n_shared_experts is None: - self.shared_experts = None - else: - intermediate_size = config.moe_intermediate_size * config.n_shared_experts - - self.shared_experts = DeepseekV4MLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - swiglu_limit=self.swiglu_limit, - quant_config=quant_config, - reduce_results=self.use_mega_moe, - prefix=f"{prefix}.shared_experts", - ) - - if self.use_mega_moe: - self._init_mega_moe_experts(vllm_config, config, prefix) - else: - self._init_fused_moe_experts(config, quant_config, prefix) - - def _init_mega_moe_experts( - self, - vllm_config: VllmConfig, - config, - prefix: str, - ) -> None: - self.ep_group = get_ep_group() - self.ep_size = self.ep_group.world_size - self.ep_rank = self.ep_group.rank_in_group - - eplb_config = vllm_config.parallel_config.eplb_config - self.n_redundant_experts = eplb_config.num_redundant_experts - self.n_routed_experts = config.n_routed_experts - self.n_shared_experts = config.n_shared_experts or 0 - self.n_logical_experts = self.n_routed_experts - self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts - assert self.n_physical_experts % self.ep_size == 0, ( - f"n_physical_experts={self.n_physical_experts} must be divisible by " - f"ep_size={self.ep_size}. Adjust num_redundant_experts." - ) - self.n_local_physical_experts = self.n_physical_experts // self.ep_size - self.physical_expert_start = self.ep_rank * self.n_local_physical_experts - self.physical_expert_end = ( - self.physical_expert_start + self.n_local_physical_experts - ) - - self.n_local_experts = self.n_local_physical_experts - self.experts_start_idx = self.physical_expert_start - self.experts_end_idx = self.physical_expert_end - - self.experts = DeepseekV4MegaMoEExperts( - vllm_config, - num_experts=self.n_physical_experts, - num_local_experts=self.n_local_physical_experts, - experts_start_idx=self.physical_expert_start, - num_logical_experts=self.n_logical_experts, - top_k=config.num_experts_per_tok, - hidden_size=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - prefix=f"{prefix}.experts", - ) - - def _init_fused_moe_experts( - self, - config, - quant_config, - prefix: str, - ) -> None: - self.tp_rank = get_tensor_model_parallel_rank() - assert config.n_routed_experts % self.tp_size == 0 - - self.n_local_experts = config.n_routed_experts // self.tp_size - self.experts_start_idx = self.tp_rank * self.n_local_experts - self.experts_end_idx = self.experts_start_idx + self.n_local_experts - - self.n_redundant_experts = 0 - self.n_shared_experts = config.n_shared_experts or 0 - self.n_logical_experts = self.n_routed_experts - self.n_physical_experts = self.n_logical_experts - self.n_local_physical_experts = self.n_local_experts - self.physical_expert_start = self.experts_start_idx - self.physical_expert_end = self.experts_end_idx - - self.experts = FusedMoE( - shared_experts=self.shared_experts, - gate=self.gate, - num_experts=config.n_routed_experts, - top_k=config.num_experts_per_tok, - hidden_size=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - renormalize=config.norm_topk_prob, - quant_config=quant_config, - prefix=f"{prefix}.experts", - scoring_func=self.scoring_func, - routed_scaling_factor=self.routed_scaling_factor, - e_score_correction_bias=self.gate.e_score_correction_bias, - hash_indices_table=self.gate.tid2eid, - swiglu_limit=self.swiglu_limit, - router_logits_dtype=torch.float32, - ) - - def forward( - self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None - ) -> torch.Tensor: - if self.gate.tid2eid is not None and input_ids is None: - raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") - - if not self.use_mega_moe: - return self._forward_fused_moe(hidden_states, input_ids) - - org_shape = hidden_states.shape - router_logits, _ = self.gate(hidden_states) - topk_weights, topk_ids = fused_topk_bias( - hidden_states=hidden_states, - gating_output=router_logits, - scoring_func=self.scoring_func, - e_score_correction_bias=self.gate.e_score_correction_bias.data - if self.gate.e_score_correction_bias is not None - else None, - topk=self.n_activated_experts, - renormalize=self.renormalize, - indices_type=self.hash_indices_dtype, - input_tokens=input_ids, - hash_indices_table=self.gate.tid2eid, - routed_scaling_factor=self.routed_scaling_factor, - ) - activation_clamp = ( - float(self.swiglu_limit) if self.swiglu_limit is not None else None - ) - final_hidden_states = self.experts( - hidden_states, - topk_weights, - topk_ids, - activation_clamp=activation_clamp, - ) - - if self.shared_experts is not None: - shared_output = self.shared_experts(hidden_states) - final_hidden_states += shared_output - - return final_hidden_states.view(org_shape) - - def _forward_fused_moe( - self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None - ) -> torch.Tensor: - org_shape = hidden_states.shape - if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - final_hidden_states = self.experts( - hidden_states=hidden_states, - router_logits=hidden_states, - input_ids=input_ids, - ) - else: - router_logits, _ = self.gate(hidden_states) - final_hidden_states = self.experts( - hidden_states=hidden_states, - router_logits=router_logits, - input_ids=input_ids, - ) - - return final_hidden_states.view(org_shape) - - def finalize_mega_moe_weights(self) -> None: - if self.use_mega_moe: - self.experts.finalize_weights() - - -class DeepseekV4Attention(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - prefix: str, - topk_indices_buffer: torch.Tensor | None = None, - aux_stream_list: list[torch.cuda.Stream] | None = None, - ): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - layer_id = extract_layer_index(prefix) - - self.layer_id = layer_id - self.hidden_size = config.hidden_size - self.n_heads = config.num_attention_heads - tp_size = get_tensor_model_parallel_world_size() - assert self.n_heads % tp_size == 0 - - self.n_local_heads = self.n_heads // tp_size - self.q_lora_rank = config.q_lora_rank - self.o_lora_rank = config.o_lora_rank - self.head_dim = config.head_dim - self.rope_head_dim = config.qk_rope_head_dim - self.nope_head_dim = self.head_dim - self.rope_head_dim - self.n_groups = config.o_groups - self.n_local_groups = self.n_groups // tp_size - self.window_size = config.sliding_window - # NOTE(zyongye) Compress ratio can't be 0 - # we do this for because MTP layer is not included - # in the compress ratio list - if layer_id < config.num_hidden_layers: - self.compress_ratio = max(1, config.compress_ratios[layer_id]) - else: - self.compress_ratio = 1 - self.eps = config.rms_norm_eps - self.max_position_embeddings = config.max_position_embeddings - - # Padded to min 64 heads for FlashMLA, initialized to -inf - # (no sink effect). Weight loading fills the first n_local_heads slots. - padded_heads = max(self.n_local_heads, 64) - self.attn_sink = nn.Parameter( - torch.full((padded_heads,), -float("inf"), dtype=torch.float32), - requires_grad=False, - ) - - self.fused_wqa_wkv = MergedColumnParallelLinear( - self.hidden_size, - [self.q_lora_rank, self.head_dim], - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.fused_wqa_wkv", - disable_tp=True, # fused ReplicatedLinear - ) - self.q_norm = RMSNorm(self.q_lora_rank, self.eps) - self.wq_b = ColumnParallelLinear( - self.q_lora_rank, - self.n_heads * self.head_dim, - bias=False, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.wq_b", - ) - - self.kv_norm = RMSNorm(self.head_dim, self.eps) - self.wo_a = ColumnParallelLinear( - self.n_heads * self.head_dim // self.n_groups, - self.n_groups * self.o_lora_rank, - bias=False, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.wo_a", - ) - self.wo_a.is_bmm = True - self.wo_a.bmm_batch_size = self.n_local_groups - self.wo_b = RowParallelLinear( - self.n_groups * self.o_lora_rank, - self.hidden_size, - bias=False, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.wo_b", - ) - self.softmax_scale = self.head_dim**-0.5 - self.scale_fmt = config.quantization_config["scale_fmt"] - - self.rope_parameters = config.rope_scaling - - # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it) - self.rotary_emb = build_deepseek_v4_rope( - config, - head_dim=self.head_dim, - rope_head_dim=self.rope_head_dim, - max_position_embeddings=self.max_position_embeddings, - compress_ratio=self.compress_ratio, - ) - - self.indexer = None - if self.compress_ratio == 4: - # Only C4A uses sparse attention and hence has indexer. - # aux_stream_list[0] runs indexer.forward() in the wrapper; [2] is - # free here (outer GEMMs joined) for the inner overlap of - # wq_b+fused_indexer_q_rope_quant vs compressor. - indexer_aux_stream = ( - aux_stream_list[2] if aux_stream_list is not None else None - ) - self.indexer = DeepseekV4Indexer( - vllm_config, - config=config, - hidden_size=self.hidden_size, - q_lora_rank=self.q_lora_rank, - quant_config=quant_config, - cache_config=vllm_config.cache_config, - topk_indices_buffer=topk_indices_buffer, - compress_ratio=self.compress_ratio, - prefix=f"{prefix}.indexer", - aux_stream=indexer_aux_stream, - ) - - self.mla_attn = DeepseekV4MLA( - hidden_size=self.hidden_size, - num_heads=self.n_local_heads, - head_dim=self.head_dim, - scale=self.softmax_scale, - qk_nope_head_dim=self.nope_head_dim, - qk_rope_head_dim=self.rope_head_dim, - v_head_dim=self.head_dim, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.head_dim, - o_lora_rank=self.o_lora_rank, - vllm_config=vllm_config, - fused_wqa_wkv=self.fused_wqa_wkv, - q_norm=self.q_norm, - wq_b=self.wq_b, - kv_norm=self.kv_norm, - wo_a=self.wo_a, - wo_b=self.wo_b, - attn_sink=self.attn_sink, - rotary_emb=self.rotary_emb, - indexer=self.indexer, - indexer_rotary_emb=self.rotary_emb, - topk_indices_buffer=topk_indices_buffer, - aux_stream_list=aux_stream_list, - window_size=self.window_size, - compress_ratio=self.compress_ratio, - cache_config=vllm_config.cache_config, - quant_config=quant_config, - prefix=prefix, - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - llama_4_scaling: torch.Tensor | None, - ): - return self.mla_attn(positions, hidden_states, llama_4_scaling) - - -class DeepseekV4DecoderLayer(nn.Module): - def __init__( - self, - vllm_config, - prefix, - topk_indices_buffer: torch.Tensor | None = None, - aux_stream_list: list[torch.cuda.Stream] | None = None, - ): - super().__init__() - - config = vllm_config.model_config.hf_config - self.hidden_size = config.hidden_size - - self.rms_norm_eps = config.rms_norm_eps - self.attn = DeepseekV4Attention( - vllm_config, - prefix=f"{prefix}.attn", - topk_indices_buffer=topk_indices_buffer, - aux_stream_list=aux_stream_list, - ) - self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") - - self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) - self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) - self.hc_mult = config.hc_mult - self.hc_sinkhorn_iters = config.hc_sinkhorn_iters - self.hc_eps = config.hc_eps - self.hc_post_alpha = 2.0 - mix_hc = (2 + self.hc_mult) * self.hc_mult - hc_dim = self.hc_mult * self.hidden_size - self.hc_attn_fn = nn.Parameter( - torch.empty( - (mix_hc, hc_dim), - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_ffn_fn = nn.Parameter( - torch.empty( - (mix_hc, hc_dim), - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_attn_base = nn.Parameter( - torch.empty( - mix_hc, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_ffn_base = nn.Parameter( - torch.empty( - mix_hc, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_attn_scale = nn.Parameter( - torch.empty( - 3, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_ffn_scale = nn.Parameter( - torch.empty( - 3, - dtype=torch.float32, - ), - requires_grad=False, - ) - - def forward( - self, - x: torch.Tensor, - positions: torch.Tensor, - input_ids: torch.Tensor | None, - post_mix: torch.Tensor | None = None, - res_mix: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - attn_norm_weight = self.attn_norm.weight.data - attn_norm_eps = self.attn_norm.variance_epsilon - if residual is None: - # Run standalone mhc_pre on first layer - residual = x - post_mix, res_mix, x = mhc_pre_tilelang( - x, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - norm_weight=attn_norm_weight, - norm_eps=attn_norm_eps, - ) - else: - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - n_splits=1, - tile_n=1, - norm_weight=attn_norm_weight, - norm_eps=attn_norm_eps, - ) - - # attn_norm is fused into mhc_pre_tilelang / mhc_fused_post_pre above. - x = self.attn(positions, x, None) - - ffn_norm_weight = self.ffn_norm.weight.data - ffn_norm_eps = self.ffn_norm.variance_epsilon - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - n_splits=1, - tile_n=1, - norm_weight=ffn_norm_weight, - norm_eps=ffn_norm_eps, - ) - - x = self.ffn(x, input_ids) - return x, residual, post_mix, res_mix - - -class DeepseekV4Model(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.config = config - self.use_mega_moe = ( - vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" - ) - if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: - raise NotImplementedError( - "DeepSeek V4 MegaMoE currently requires expert parallel. " - "Enable it with --enable-expert-parallel, or pick a different " - "moe backend." - ) - self.vocab_size = config.vocab_size - self.hc_eps = config.hc_eps - self.hc_mult = config.hc_mult - self.hc_dim = self.hc_mult * config.hidden_size - self.rms_norm_eps = config.rms_norm_eps - - # Three aux streams: one per non-default input GEMM in - # DeepseekV4MLA.attn_gemm_parallel_execute - # (compressor kv_score, indexer.weights_proj, indexer.compressor - # kv_score). fused_wqa_wkv stays on the default stream. - aux_stream_list = [torch.cuda.Stream() for _ in range(3)] - - # Reserved topk indices buffer for all Indexer layers to reuse. - self.topk_indices_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - config.index_topk, - dtype=torch.int32, - ) - - if get_pp_group().is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - else: - self.embed_tokens = PPMissingLayer() - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: DeepseekV4DecoderLayer( - vllm_config, - prefix=prefix, - topk_indices_buffer=self.topk_indices_buffer, - aux_stream_list=aux_stream_list, - ), - prefix=f"{prefix}.layers", - ) - - if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) - else: - self.norm = PPMissingLayer() - - self.hc_head_fn = nn.Parameter( - torch.empty( - self.hc_mult, - self.hc_dim, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_head_base = nn.Parameter( - torch.empty( - self.hc_mult, - dtype=torch.float32, - ), - requires_grad=False, - ) - self.hc_head_scale = nn.Parameter( - torch.empty(1, dtype=torch.float32), - requires_grad=False, - ) - # Pre-hc_head residual stream buffer for the MTP draft. Stable - # address (outside the cudagraph pool) so the copy_ in forward() - # refreshes it correctly across captured shapes. - # refreshes it correctly across captured shapes. Only allocated on - # the last PP rank — that's where MTP target hidden states are - # produced. - if get_pp_group().is_last_rank: - self._mtp_hidden_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - self.hc_dim, - dtype=vllm_config.model_config.dtype, - ) - else: - self._mtp_hidden_buffer = None - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def make_empty_intermediate_tensors( - self, - batch_size: int, - dtype: torch.dtype, - device: torch.device, - ) -> IntermediateTensors: - # PP intermediate tensors carry the multi-stream hidden_states - # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the - # token embedding to hc_mult streams before the first decoder - # layer and keeps that shape until hc_head() collapses it. - return IntermediateTensors( - { - "hidden_states": torch.zeros( - (batch_size, self.hc_mult, self.config.hidden_size), - dtype=dtype, - device=device, - ), - } - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - - if self.use_mega_moe: - input_ids = input_ids.to(torch.int64) - - residual, post_mix, res_mix = None, None, None - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual, post_mix, res_mix = layer( - hidden_states, - positions, - input_ids, - post_mix, - res_mix, - residual, - ) - if layer is not None: - hidden_states = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix - ) - - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) - - # Stash pre-hc_head residual for the MTP draft (captured copy_). - num_tokens = hidden_states.shape[0] - self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) - - hidden_states = hc_head_fused_kernel_tilelang( - hidden_states, - self.hc_head_fn, - self.hc_head_scale, - self.hc_head_base, - self.rms_norm_eps, - self.hc_eps, - ) - hidden_states = self.norm(hidden_states) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("gate_up_proj", "w1", 0), - ("gate_up_proj", "w3", 1), - ("attn.fused_wqa_wkv", "attn.wq_a", 0), - ("attn.fused_wqa_wkv", "attn.wkv", 1), - ("compressor.fused_wkv_wgate", "compressor.wkv", 0), - ("compressor.fused_wkv_wgate", "compressor.wgate", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - - # Pre-compute expert mapping ONCE. - expert_mapping = self.get_expert_mapping() - - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - # Skip non-stacked layers and experts (experts handled below). - if ".experts." in name: - continue - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - - if is_pp_missing_parameter(name, self): - break - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(name) - break - else: - if ".experts." in name: - # E8M0 scales are stored as float8_e8m0fnu in - # checkpoints but the MoE param is uint8. copy_() - # would do a numeric conversion (e.g. 2^-7 → 0), - # destroying the raw exponent bytes. - if ( - "weight_scale" in name - and loaded_weight.dtype == torch.float8_e8m0fnu - ): - loaded_weight = loaded_weight.view(torch.uint8) - for mapping in expert_mapping: - param_name, weight_name, expert_id, expert_shard_id = mapping - if weight_name not in name: - continue - name_mapped = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name_mapped, self): - continue - param = params_dict[name_mapped] - # We should ask the weight loader to return success or not - # here since otherwise we may skip experts with other - # available replicas. - weight_loader = typing.cast( - Callable[..., bool], param.weight_loader - ) - success = weight_loader( - param, - loaded_weight, - name_mapped, - shard_id=expert_shard_id, - expert_id=expert_id, - return_success=True, - ) - if success: - name = name_mapped - break - loaded_params.add(name_mapped) - continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue - else: - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - continue - - return loaded_params - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) - if first_layer.ffn.use_mega_moe: - return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - return FusedMoE.make_expert_params_mapping( - self, - ckpt_gate_proj_name="w1", - ckpt_down_proj_name="w2", - ckpt_up_proj_name="w3", - num_experts=self.config.n_routed_experts, - ) - - def finalize_mega_moe_weights(self) -> None: - for layer in islice(self.layers, self.start_layer, self.end_layer): - layer.ffn.finalize_mega_moe_weights() - - -def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: - if expert_dtype == "fp4": - # MXFP4 experts use Mxfp4MoEMethod, which registers scales as - # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and - # shared experts use Fp8LinearMethod's block scales, which - # register as ``weight_scale_inv``. - scale_regex = { - re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", - re.compile(r"\.scale$"): ".weight_scale_inv", - } - else: - # FP8 experts use Fp8MoEMethod (block_quant=True), which registers - # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys - # there. - scale_regex = { - re.compile(r"\.scale$"): ".weight_scale_inv", - } - return WeightsMapper( - orig_to_new_prefix={ - "layers.": "model.layers.", - "embed.": "model.embed.", - "norm.": "model.norm.", - "hc_head": "model.hc_head", - "mtp.": "model.mtp.", - }, - orig_to_new_regex=scale_regex, - orig_to_new_suffix={ - "head.weight": "lm_head.weight", - "embed.weight": "embed_tokens.weight", - ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", - }, - orig_to_new_substr={ - ".attn.compressor.": ".attn.mla_attn.compressor.", - ".shared_experts.w2": ".shared_experts.down_proj", - }, - ) - - -class DeepseekV4MixtureOfExperts(MixtureOfExperts): - moe_mlp_layers: list["DeepseekV4MoE"] - - def extract_moe_parameters(self, example_moe: "DeepseekV4MoE | None") -> None: - if example_moe is None: - self.num_moe_layers = 0 - self.num_expert_groups = 0 - self.num_logical_experts = 0 - self.num_physical_experts = 0 - self.num_local_physical_experts = 0 - self.num_routed_experts = 0 - self.num_shared_experts = 0 - self.num_redundant_experts = 0 - return - self.num_logical_experts = example_moe.n_logical_experts - self.num_physical_experts = example_moe.n_physical_experts - self.num_local_physical_experts = example_moe.n_local_physical_experts - self.num_routed_experts = example_moe.n_routed_experts - self.num_shared_experts = example_moe.n_shared_experts - self.num_redundant_experts = example_moe.n_redundant_experts - - def update_physical_experts_metadata( - self, - num_physical_experts: int, - num_local_physical_experts: int, - ) -> None: - assert self.num_local_physical_experts == num_local_physical_experts - self.num_physical_experts = num_physical_experts - self.num_local_physical_experts = num_local_physical_experts - self.num_redundant_experts = num_physical_experts - self.num_logical_experts - for moe in self.moe_mlp_layers: - moe.n_local_physical_experts = num_local_physical_experts - moe.n_physical_experts = num_physical_experts - moe.n_redundant_experts = self.num_redundant_experts - moe.experts.update_expert_map() - - -class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts): - model_cls = DeepseekV4Model - - # Default mapper assumes the original FP4-expert checkpoint layout. - # Overridden per-instance in __init__ when expert_dtype != "fp4". - hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4") - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - self.config = config - expert_dtype = getattr(config, "expert_dtype", "fp4") - if expert_dtype != "fp4": - self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype) - - self.model = self.model_cls( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - if get_pp_group().is_last_rank: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) - else: - self.lm_head = PPMissingLayer() - self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] - self.model.make_empty_intermediate_tensors - ) - - self.set_moe_parameters() - - def set_moe_parameters(self) -> None: - self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = [] - self.num_expert_groups = getattr(self.config, "n_group", 1) - self.num_moe_layers = self.config.num_hidden_layers - self.moe_layers: list[nn.Module] = [] - self.moe_mlp_layers: list[DeepseekV4MoE] = [] - example_moe: DeepseekV4MoE | None = None - for layer in self.model.layers: - if isinstance(layer, PPMissingLayer): - continue - if not isinstance(layer, DeepseekV4DecoderLayer): - continue - if isinstance(layer.ffn, DeepseekV4MoE): - example_moe = layer.ffn - self.moe_mlp_layers.append(layer.ffn) - self.moe_layers.append(layer.ffn.experts) - - self.num_moe_layers = len(self.moe_layers) - self.extract_moe_parameters(example_moe) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def get_mtp_target_hidden_states(self) -> torch.Tensor | None: - """Pre-hc_head residual stream buffer (max_num_batched_tokens, - hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" - return getattr(self.model, "_mtp_hidden_buffer", None) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) - loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - self.model.finalize_mega_moe_weights() - return loaded_params - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - return self.model.get_expert_mapping() diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py deleted file mode 100644 index 133a96e3..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py +++ /dev/null @@ -1,516 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4). - -Split from ``deepseek_mtp.py`` because the V4 architecture introduces several -pieces that have no analogue in V3/V32: - * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of - the fused ``eh_proj``); - * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``; - * ``DeepseekV4DecoderLayer`` with its own aux-stream management; - * V4-specific checkpoint weight-name remapping in ``load_weights``. -""" - -import typing -from collections.abc import Callable, Iterable - -import regex as re -import torch -import torch.nn as nn - -from vllm.config import VllmConfig -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.model_executor.kernels.mhc.tilelang import ( - hc_head_fused_kernel_tilelang, - mhc_post_tilelang, -) -from vllm.model_executor.layers.fused_moe import FusedMoE -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ReplicatedLinear -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.vocab_parallel_embedding import ( - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.deepseek_mtp import SharedHead -from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name -from vllm.model_executor.models.utils import maybe_prefix -from vllm.models.deepseek_v4.common.ops import ( - fused_mtp_input_rmsnorm, - mtp_shared_head_rmsnorm, -) -from vllm.sequence import IntermediateTensors - -from .model import ( - DeepseekV4DecoderLayer, - make_deepseek_v4_expert_params_mapping, -) - -logger = init_logger(__name__) - -# MoE expert scales are fused into per-layer w13/w2 tensors. The exact -# parameter suffix depends on which FusedMoE method handles the experts: -# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``; -# - fp8 experts (Fp8MoEMethod with block_quant=True) register -# ``w{1,2,3}_weight_scale_inv``. -# Other FP8 linear scales (including shared experts) always use -# ``.weight_scale_inv``. Mirrors the per-instance mapper built by -# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py. -_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") - - -class DeepSeekV4MultiTokenPredictorLayer(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - topk_indices_buffer: torch.Tensor, - prefix: str, - aux_stream_list: list[torch.cuda.Stream] | None = None, - ) -> None: - super().__init__() - - assert vllm_config.speculative_config is not None - config = vllm_config.speculative_config.draft_model_config.hf_config - self.config = config - quant_config = vllm_config.quant_config - self.rms_norm_eps = config.rms_norm_eps - - self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than - # fusing them the way V3 does with eh_proj. - self.e_proj = ReplicatedLinear( - config.hidden_size, - config.hidden_size, - bias=False, - return_bias=False, - quant_config=quant_config, - ) - self.h_proj = ReplicatedLinear( - config.hidden_size, - config.hidden_size, - bias=False, - return_bias=False, - quant_config=quant_config, - ) - - self.hc_eps = config.hc_eps - self.hc_mult = config.hc_mult - self.hc_dim = self.hc_mult * config.hidden_size - self.hc_head_fn = nn.Parameter( - torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_base = nn.Parameter( - torch.empty(self.hc_mult, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_scale = nn.Parameter( - torch.empty(1, dtype=torch.float32), - requires_grad=False, - ) - - self.shared_head = SharedHead( - config=config, prefix=prefix, quant_config=quant_config - ) - self.mtp_block = DeepseekV4DecoderLayer( - vllm_config, - prefix, - topk_indices_buffer=topk_indices_buffer, - aux_stream_list=aux_stream_list, - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - previous_hidden_states: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - spec_step_index: int = 0, - ) -> torch.Tensor: - assert inputs_embeds is not None - # Target stashes pre-hc_head residual as flat (T, hc_mult * D); - # reshape to (T, hc_mult, D) — the training-time layout — before - # the fused norm pass so both inputs are 3D-friendly. - previous_hidden_states = previous_hidden_states.view( - -1, self.hc_mult, self.config.hidden_size - ) - # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm. - inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm( - inputs_embeds, - positions, - previous_hidden_states, - self.enorm.weight.data, - self.hnorm.weight.data, - self.enorm.variance_epsilon, - self.hc_mult, - ) - hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( - inputs_embeds - ).unsqueeze(-2) - hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=None - ) - hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) - # Return the flat pre-hc_head residual so it can be re-fed as the - # next spec step's `previous_hidden_states` when - # num_speculative_tokens > 1. hc_head is deferred to compute_logits. - return hidden_states.flatten(1) - - -class DeepSeekV4MultiTokenPredictor(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - self.mtp_start_layer_idx = config.num_hidden_layers - self.num_mtp_layers = config.num_nextn_predict_layers - - topk_tokens = config.index_topk - self.topk_indices_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - topk_tokens, - dtype=torch.int32, - ) - - # Three aux streams shared across all MTP layers, mirroring DeepseekV4Model. - aux_stream_list = [torch.cuda.Stream() for _ in range(3)] - - # to map the exact layer index from weights - self.layers = torch.nn.ModuleDict( - { - str(idx): DeepSeekV4MultiTokenPredictorLayer( - vllm_config, - self.topk_indices_buffer, - f"{prefix}.layers.{idx}", - aux_stream_list=aux_stream_list, - ) - for idx in range( - self.mtp_start_layer_idx, - self.mtp_start_layer_idx + self.num_mtp_layers, - ) - } - ) - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "embed_tokens"), - ) - self.logits_processor = LogitsProcessor(config.vocab_size) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - previous_hidden_states: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - spec_step_idx: int = 0, - ) -> torch.Tensor: - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - current_step_idx = spec_step_idx % self.num_mtp_layers - return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( - input_ids, - positions, - previous_hidden_states, - inputs_embeds, - current_step_idx, - ) - - def compute_logits( - self, - hidden_states: torch.Tensor, - spec_step_idx: int = 0, - ) -> torch.Tensor: - current_step_idx = spec_step_idx % self.num_mtp_layers - mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] - # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply - # hc_head here so logits are computed from the dense hidden state. - hidden_states = hidden_states.view( - -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size - ) - hidden_states = hc_head_fused_kernel_tilelang( - hidden_states, - mtp_layer.hc_head_fn, - mtp_layer.hc_head_scale, - mtp_layer.hc_head_base, - mtp_layer.rms_norm_eps, - mtp_layer.hc_eps, - ) - hidden_states = mtp_shared_head_rmsnorm( - hidden_states, - mtp_layer.shared_head.norm.weight.data, - mtp_layer.shared_head.norm.variance_epsilon, - ) - logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) - return logits - - -class DeepSeekV4MTP(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - self.config = vllm_config.model_config.hf_config - self.quant_config = vllm_config.quant_config - self.model = DeepSeekV4MultiTokenPredictor( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - hidden_states: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - spec_step_idx: int = 0, - ) -> torch.Tensor: - hidden_states = self.model( - input_ids, positions, hidden_states, inputs_embeds, spec_step_idx - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - spec_step_idx: int = 0, - ) -> torch.Tensor | None: - return self.model.compute_logits(hidden_states, spec_step_idx) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Weight name remapping for checkpoint compatibility. - # Maps checkpoint weight paths to model parameter paths. - WEIGHT_NAME_REMAPPING: dict[str, str] = { - ".emb.tok_emb.weight": ".embed_tokens.weight", - ".head.weight": ".shared_head.head.weight", - ".norm.weight": ".shared_head.norm.weight", - } - - def _remap_weight_name(name: str) -> str: - """Remap checkpoint weight names to model parameter names.""" - for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items(): - if old_pattern in name: - name = name.replace(old_pattern, new_pattern) - return name - - def _find_mtp_layer_idx(name: str) -> int: - subnames = name.split(".") - for subname in subnames: - try: - # we return the first encountered integer - return int(subname) - except ValueError: - continue - return 0 - - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("gate_up_proj", "w1", 0), - ("gate_up_proj", "w3", 1), - ("attn.fused_wqa_wkv", "attn.wq_a", 0), - ("attn.fused_wqa_wkv", "attn.wkv", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - - # Pre-compute expert mapping ONCE. - first_layer = next(iter(self.model.layers.values())) - if first_layer.mtp_block.ffn.use_mega_moe: - expert_mapping = make_deepseek_v4_expert_params_mapping( - self.config.n_routed_experts - ) - else: - expert_mapping = FusedMoE.make_expert_params_mapping( - self, - ckpt_gate_proj_name="w1", - ckpt_down_proj_name="w2", - ckpt_up_proj_name="w3", - num_experts=self.config.n_routed_experts, - ) - - # FP8 experts register ``..._weight_scale_inv`` (block_quant) while - # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix - # for the rename below based on the model's expert dtype. - expert_scale_suffix = ( - ".weight_scale" - if getattr(self.config, "expert_dtype", "fp4") == "fp4" - else ".weight_scale_inv" - ) - - for name, loaded_weight in weights: - mtp_layer_idx = _find_mtp_layer_idx(name) - # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to - # `model.layers.{num_hidden_layers + i}.*` so that - # get_spec_layer_idx_from_weight_name can identify them. - name = name.replace( - f"mtp.{mtp_layer_idx}.", - f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", - ) - - spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) - if spec_layer is None: - continue - - name = _remap_weight_name(name) - name = self._rewrite_spec_layer_name(spec_layer, name) - - if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name: - continue - if name.endswith(".scale"): - suffix = ( - expert_scale_suffix - if _EXPERT_SCALE_RE.search(name) - else ".weight_scale_inv" - ) - name = name.removesuffix(".scale") + suffix - for param_name, weight_name, shard_id in stacked_params_mapping: - # Skip non-stacked layers and experts (experts handled below). - if ".experts." in name: - continue - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(name) - break - else: - if ".experts." in name: - # Reinterpret E8M0 scales as uint8 to preserve raw - # exponent bytes; numeric copy_() would zero them. - # Mirrors the main DeepseekV4 loader. - if ( - "weight_scale" in name - and loaded_weight.dtype == torch.float8_e8m0fnu - ): - loaded_weight = loaded_weight.view(torch.uint8) - for mapping in expert_mapping: - param_name, weight_name, expert_id, expert_shard_id = mapping - if weight_name not in name: - continue - name_mapped = name.replace(weight_name, param_name) - param = params_dict[name_mapped] - # We should ask the weight loader to return success or not - # here since otherwise we may skip experts with other - # available replicas. - weight_loader = typing.cast( - Callable[..., bool], param.weight_loader - ) - success = weight_loader( - param, - loaded_weight, - name_mapped, - shard_id=expert_shard_id, - expert_id=expert_id, - return_success=True, - ) - if success: - name = name_mapped - loaded_params.add(name_mapped) - break - continue - elif "attn_sink" in name: - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue - else: - if ".shared_experts.w2" in name: - name = name.replace( - ".shared_experts.w2", ".shared_experts.down_proj" - ) - if name.endswith(".ffn.gate.bias"): - # ``e_score_correction_bias`` lives on the gate - # under a different attribute name. - name = name.replace( - ".ffn.gate.bias", - ".ffn.gate.e_score_correction_bias", - ) - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - continue - - loaded_layers: set[int] = set() - for param_name in loaded_params: - spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) - if spec_layer is not None: - loaded_layers.add(spec_layer) - for layer_idx in range( - self.model.mtp_start_layer_idx, - self.model.mtp_start_layer_idx + self.model.num_mtp_layers, - ): - if layer_idx not in loaded_layers: - raise ValueError( - f"MTP speculative decoding layer {layer_idx} weights " - f"missing from checkpoint. The checkpoint may have " - f"been quantized without including the MTP layers. " - f"Use a checkpoint that includes MTP layer weights, " - f"or disable speculative decoding." - ) - self.finalize_mega_moe_weights() - logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) - return loaded_params - - def finalize_mega_moe_weights(self) -> None: - for layer in self.model.layers.values(): - layer.mtp_block.ffn.finalize_mega_moe_weights() - - def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. - Add .mtp_block for modules in transformer layer block for spec layer - and rename shared layer weights to be top level. - """ - spec_layer_weight_names = [ - "embed_tokens", - "enorm", - "hnorm", - "h_proj", - "e_proj", - "shared_head", - "hc_head_fn", - "hc_head_base", - "hc_head_scale", - ] - shared_weight_names = ["embed_tokens"] - spec_layer_weight = False - shared_weight = False - for weight_name in spec_layer_weight_names: - if weight_name in name: - spec_layer_weight = True - if weight_name in shared_weight_names: - shared_weight = True - break - if not spec_layer_weight: - # treat rest weights as weights for transformer layer block - name = name.replace( - f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." - ) - elif shared_weight: - # treat shared weights as top level weights - name = name.replace(f"model.layers.{spec_layer}.", "model.") - return name diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py deleted file mode 100644 index 20752bc1..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NVIDIA-only (cutedsl/cutlass) kernels for DeepSeek V4. - -These modules import ``cutlass``/``cutedsl`` at module top level, so they must -not be imported on non-CUDA platforms. Callers should gate on -``vllm.utils.import_utils.has_cutedsl()`` before importing from here. - -This ``__init__`` deliberately imports nothing: re-exporting the cutedsl -modules here would eagerly ``import cutlass`` (initializing the CUDA driver) for -anyone who imports ``vllm.models.deepseek_v4``, breaking forked subprocesses. -Import the leaf modules directly under a ``has_cutedsl()``/``is_cuda()`` gate. -""" diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py deleted file mode 100644 index 6c443246..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ /dev/null @@ -1,331 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from functools import cache - -import cutlass -import cutlass.cute as cute -import torch -from cuda.bindings.driver import CUstream -from cutlass import BFloat16, Int32, Uint8, Uint32 -from cutlass.cute.nvgpu import cpasync -from quack.compile_utils import make_fake_tensor - -from vllm.cute_utils import _bf16x2_mul, cvt - - -def dequantize_and_gather_k_cache_cutedsl( - out: torch.Tensor, - k_cache: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor | None, - block_table: torch.Tensor, - block_size: int, - offset: int, -) -> None: - DequantGatherKCacheKernel.compile( - block_size=block_size, - has_gather_lens=gather_lens is not None, - )(out, k_cache, seq_lens, gather_lens, block_table, offset) - - -class DequantGatherKCacheKernel: - # Hard-coded for DSv4. - head_dim = 512 - group_size = 64 # 1 scale per 64 elems - - def __init__(self, fp8_dim: int = 448, block_size: int = 64): - self.fp8_dim = fp8_dim - self.bf16_dim = self.head_dim - fp8_dim - self.data_dim = fp8_dim + self.bf16_dim * 2 - self.block_size = block_size - - self.num_warps = 4 - self.tb_size = self.num_warps * 32 - self.num_stages = 4 - - @cute.jit - def __call__( - self, - out: cute.Tensor, - k_cache: cute.Tensor, - seq_lens: cute.Tensor, - gather_lens: cute.Tensor | None, - block_table: cute.Tensor, - offset: Int32, - stream: CUstream, - ): - # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] - # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. - k_data = cute.make_tensor( - k_cache.iterator, - layout=cute.make_layout( - (k_cache.shape[0], self.block_size, self.data_dim), - stride=(k_cache.stride[0], self.data_dim, 1), - ), - ) - k_scale = cute.make_tensor( - k_cache.iterator + (self.block_size * self.data_dim), - layout=cute.make_layout( - (k_cache.shape[0], self.block_size, 8), - stride=(k_cache.stride[0], 8, 1), - ), - ) - - grid = (out.shape[0], 1024, 1) - self.kernel( - out, - k_data, - k_scale, - seq_lens, - gather_lens, - block_table, - offset, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.jit - def load_g2s( - self, - k_data_slice: cute.Tensor, - k_scale: cute.Tensor, - block_table: cute.Tensor, - s_kdata_slice: cute.Tensor, - s_kscale: cute.Tensor, - req_id, - pos, - lane_id, - stage_id, - ): - # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] - # s_kdata_slice: [(4, data_dim/16), num_stages] - - op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) - cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) - cp8_atom = cute.make_copy_atom(cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64) - page_id = block_table[req_id, pos // self.block_size] - block_offset = pos % self.block_size - - # Load the first 512 bytes (32x16B). - idx = lane_id - src = k_data_slice[page_id, block_offset, (None, idx)] - cute.copy( - cp16_atom, - cute.recast_tensor(src, Uint32), - s_kdata_slice[(None, idx), stage_id], - ) - - # Load the tail 64 bytes. - idx += 32 - if idx < cutlass.const_expr(self.data_dim // 16): - src = k_data_slice[page_id, block_offset, (None, idx)] - cute.copy( - cp16_atom, - cute.recast_tensor(src, Uint32), - s_kdata_slice[(None, idx), stage_id], - ) - elif idx == cutlass.const_expr(self.data_dim // 16): - cute.copy( - cp8_atom, - k_scale[page_id, block_offset, None], - s_kscale[None, stage_id], - ) - - @cute.kernel - def kernel( - self, - out: cute.Tensor, - k_data: cute.Tensor, - k_scale: cute.Tensor, - seq_lens: cute.Tensor, - gather_lens: cute.Tensor | None, - block_table: cute.Tensor, - offset: Int32, - ): - req_id, worker_id, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - - _, num_workers, _ = cute.arch.grid_dim() - - # Prepare smem. - smem = cutlass.utils.SmemAllocator() - s_kdata = smem.allocate_tensor( - Uint32, - cute.make_layout((self.data_dim // 4, self.num_warps, self.num_stages)), - byte_alignment=16, - )[None, warp_id, None] - s_kscale = smem.allocate_tensor( - Uint8, - cute.make_layout((8, self.num_warps, self.num_stages)), - byte_alignment=8, - )[None, warp_id, None] - - # Prepare for 16B cp.async, also for BF16 smem loads later. - k_data_slice = cute.logical_divide(k_data, (None, None, 16)) - s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None)) - - # Load FP8 elems in 8B units, so once dequantized, they are 16B units. - s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None)) - - # 16B st.global. - out_slice = cute.logical_divide(out, (None, None, 8)) - - cp_op = cute.nvgpu.CopyUniversalOp() - cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) - cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - - seq_len = seq_lens[req_id] - gather_len = seq_len - if cutlass.const_expr(gather_lens is not None): - gather_len = gather_lens[req_id] # type: ignore[index] - start_pos = seq_len - gather_len - - # Start prefetch. - for i in cutlass.range_constexpr(self.num_stages - 1): - next_pos = ( - start_pos - + worker_id * self.num_warps - + warp_id - + i * num_workers * self.num_warps - ) - if next_pos < seq_len: - self.load_g2s( - k_data_slice, - k_scale, - block_table, - s_kdata_16B_slice, - s_kscale, - req_id, - next_pos, - lane_id, - i, - ) - cute.arch.cp_async_commit_group() - prefetch_stage = self.num_stages - 1 - compute_stage = 0 - - # Main loop. - for i in range( - worker_id * self.num_warps + warp_id, - gather_len, - num_workers * self.num_warps, - ): - pos = start_pos + i - - # Prefetch next stage. - next_pos = pos + num_workers * self.num_warps * (self.num_stages - 1) - if next_pos < seq_len: - self.load_g2s( - k_data_slice, - k_scale, - block_table, - s_kdata_16B_slice, - s_kscale, - req_id, - next_pos, - lane_id, - prefetch_stage, - ) - prefetch_stage = (prefetch_stage + 1) % self.num_stages - cute.arch.cp_async_commit_group() - - # Wait for gmem->smem to finish. - cute.arch.cp_async_wait_group(self.num_stages - 1) - cute.arch.sync_warp() - - # There are 512 elems per token. As a warp, data0 holds the first - # 256 elems and data1 holds the second 256 elems, i.e. each thread - # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as - # contiguous 16B global stores. On Blackwell, this might not be - # necessary as we have 32B global stores, but doing it this way - # does not seem to be slower. - data0 = cute.make_rmem_tensor((2,), Uint32) - data1 = cute.make_rmem_tensor((2,), Uint32) - cute.copy(cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0) - cute.copy( - cp8_atom, - s_kdata_8B_slice[(None, lane_id + 32), compute_stage], - data1, - ) - - # Convert to bf16x2 via bit manipulation. FP8 scales are per 64 - # elements. An 8-element chunk advances the scale index by - # chunk_id * 8 // group_size. - scale0_u32 = Uint32(s_kscale[lane_id * 8 // self.group_size, compute_stage]) - scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7)) - scale1_u32 = Uint32( - s_kscale[(lane_id + 32) * 8 // self.group_size, compute_stage] - ) - scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7)) - - # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2 - # (CUDA 13.2). - dequant0 = cute.make_rmem_tensor(4, Uint32) - dequant1 = cute.make_rmem_tensor(4, Uint32) - for j in cutlass.range_constexpr(2): - tmp0 = cvt.fp8x4_to_bf16x4(data0[j]) - tmp1 = cvt.fp8x4_to_bf16x4(data1[j]) - - # BF16 multiply is safe because the scales are exact powers of 2. - dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2) - dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2) - dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2) - dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2) - - # Last 64 elems are BF16 tail, corresponds to dequant1 of last - # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + - # 8x 16B for BF16. - if lane_id + 32 >= self.fp8_dim // 8: - idx = self.fp8_dim // 16 + (lane_id + 32) - self.fp8_dim // 8 - cute.copy( - cp16_atom, - s_kdata_16B_slice[(None, idx), compute_stage], - dequant1, - ) - - # Store two 16B BF16 chunks per lane: first half, then second half. - dst = out_slice[req_id, offset + i, (None, lane_id)] - cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32)) - - dst = out_slice[req_id, offset + i, (None, lane_id + 32)] - cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32)) - - compute_stage = (compute_stage + 1) % self.num_stages - - @cache - @staticmethod - def compile( - fp8_dim: int = 448, - block_size: int = 64, - has_gather_lens: bool = True, - ): - num_reqs = cute.sym_int() - head_dim = DequantGatherKCacheKernel.head_dim - head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 - - out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) - k_cache = cute.runtime.make_fake_tensor( - Uint8, - (cute.sym_int(), block_size, head_bytes), - stride=(cute.sym_int64(divisibility=32), head_bytes, 1), - assumed_align=32, - ) - seq_lens = make_fake_tensor(Int32, (num_reqs,)) - gather_lens = make_fake_tensor(Int32, (num_reqs,)) if has_gather_lens else None - block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int())) - - kernel = DequantGatherKCacheKernel(fp8_dim, block_size) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, - out, - k_cache, - seq_lens, - gather_lens, - block_table, - Int32(0), - stream, - options="--enable-tvm-ffi", - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py deleted file mode 100644 index 3bd5e1b5..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ /dev/null @@ -1,610 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import cache - -import cutlass -import cutlass.cute as cute -import torch -from cuda.bindings.driver import CUstream -from cutlass import BFloat16, Float32, Int64, Uint8, Uint32, const_expr -from quack.compile_utils import make_fake_tensor - -from vllm.cute_utils import ( - _bf16x2_abs, - _bf16x2_max, - cvt, - recast_val, -) -from vllm.vllm_flash_attn.cute import utils as cute_utils - -# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale. -MXFP4_BLOCK_SIZE = 32 - -_TORCH_TO_CUTE = { - torch.bfloat16: BFloat16, - torch.float32: Float32, -} - - -def fused_indexer_q_rope_quant_mxfp4_cutedsl( - positions: torch.Tensor, - index_q: torch.Tensor, - index_q_cos_sin_cache: torch.Tensor, - index_weights: torch.Tensor, - index_weights_softmax_scale: float, - index_weights_head_scale: float, - index_q_packed: torch.Tensor, - index_q_scale: torch.Tensor, - index_weights_out: torch.Tensor, -) -> None: - num_tokens, num_heads, head_dim = index_q.shape - rope_dim = index_q_cos_sin_cache.shape[-1] - rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] - - # compile all variants at first invocation - for coarsen in (1, 4): - IndexerQMxFp4Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) - - # heuristic - coarsen = 1 if num_tokens < 512 else 4 - compiled = IndexerQMxFp4Kernel.compile( - head_dim, rope_dim, num_heads, rope_type, coarsen - ) - scale = float(index_weights_softmax_scale * index_weights_head_scale) - compiled( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_q_packed, - index_q_scale, - index_weights_out, - scale, - ) - - -def fused_indexer_q_rope_quant_fp8_cutedsl( - positions: torch.Tensor, - index_q: torch.Tensor, - index_q_cos_sin_cache: torch.Tensor, - index_weights: torch.Tensor, - index_weights_softmax_scale: float, - index_weights_head_scale: float, - index_q_fp8: torch.Tensor, - index_weights_out: torch.Tensor, -) -> None: - num_tokens, num_heads, head_dim = index_q.shape - rope_dim = index_q_cos_sin_cache.shape[-1] - rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] - - for coarsen in (1, 4): - IndexerQFp8Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) - - coarsen = 1 if num_tokens < 512 else 4 - compiled = IndexerQFp8Kernel.compile( - head_dim, rope_dim, num_heads, rope_type, coarsen - ) - scale = float(index_weights_softmax_scale * index_weights_head_scale) - # The cute kernel treats the FP8 buffer as raw bytes (Uint8). - compiled( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_q_fp8.view(torch.uint8), - index_weights_out, - scale, - ) - - -class IndexerQRopeQuantKernel: - """Shared infrastructure for indexer-Q RoPE+quant fused kernels. - - Subclasses implement ``kernel`` for a particular Q quantization scheme - (MXFP4, FP8 e4m3, …). The base class owns the launch geometry and the - common preamble: thread/token addressing, the BF16 Q load, and the - interleaved-RoPE pass over the trailing ``rope_dim`` lanes. - """ - - def __init__( - self, - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - self.head_dim = head_dim - self.rope_dim = rope_dim - self.nope_dim = head_dim - rope_dim - self.num_heads = num_heads - self.cos_sin_dtype = cos_sin_dtype - - # process multiple heads at the same time to armotize RoPE load costs - assert num_heads % coarsen == 0 - self.coarsen = coarsen - - # later we will use 32B load = 16 BF16 elems - # thus, head_dim=128 requires 8 threads to handle. - # let's call subwarp = 8 threads. - self.subwarp_size = head_dim // 16 - self.tb_size = 128 - self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size - - @cute.jit - def _load_q_and_rope( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - ): - """Compute thread indices, load Q (BF16), and apply interleaved RoPE. - - Returns a tuple - (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id, - head_start, in_bounds, num_token_heads) - where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed - bf16x2 pairs covering the 16 BF16 lanes owned by this thread for - each of ``coarsen`` heads. RoPE is applied in place to the - trailing ``rope_dim`` lanes; the leading nope lanes pass through. - """ - block_id, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - - num_tokens = q.shape[0] - num_token_heads = num_tokens * self.num_heads - global_tid = block_id * self.tb_size + tid - - global_subwarp_id = global_tid // self.subwarp_size - sublane = tid % self.subwarp_size - - token_id = global_subwarp_id // (self.num_heads // self.coarsen) - head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen) - head_start = head_tile_id * self.coarsen - - # NOTE: token_id may exceed bounds, hence we need to add load/store guards - # we can't do early exit because CuteDSL doesn't support it. and we also need - # all threads in a warp to be active since we utilize warp shuffle later. - # must_in_bounds is constexpr, True when 1 threadblock fit within 1 token - # position. the compiler will remove bounds check when that happens. - must_in_bounds = cutlass.const_expr(self.tb_size % self.threads_per_token == 0) - in_bounds = must_in_bounds or (token_id < num_tokens) - - cp_op = cute.nvgpu.CopyUniversalOp() - - _layout = cute.make_layout((self.coarsen, 8), stride=(8, 1)) - q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32) - - if in_bounds: - # we can't do cute.copy() on the whole 2D tile directly because - # cute.copy() wants the 1st mode to be covered by the copy atom, - # and other modes as for loop. there is no fast way to - # "transpose" the tensor view. - q_tile = cute.local_tile( - q[token_id, None, None], - tiler=(self.coarsen, 16), - coord=(head_tile_id, sublane), - ) - cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) - for i in cutlass.range_constexpr(self.coarsen): - src = cute.recast_tensor(q_tile[i, None], Uint32) - cute.copy(cp_u32x8, src, q_bf16x2[i, None]) - - # RoPE applies only to the trailing rope_dim values. We keep the rounded - # BF16 result in q_bits so the later amax and quantization see BF16. - # cos_sin_cache layout: [max_pos, rope_dim] - if in_bounds and sublane * 16 >= self.nope_dim: - cos_vals = cute.make_rmem_tensor((8,), Float32) - sin_vals = cute.make_rmem_tensor((8,), Float32) - - pos = positions[token_id] - - # select 8 elems from cos and sin - cos_id = sublane - self.nope_dim // 16 - sin_id = cos_id + self.rope_dim // 16 - cos_src = cute.local_tile( - cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,) - ) - sin_src = cute.local_tile( - cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,) - ) - - cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256) - cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - - if const_expr(self.cos_sin_dtype is Float32): - cute.copy(cp_f32x8, cos_src, cos_vals) - cute.copy(cp_f32x8, sin_src, sin_vals) - else: - cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32) - sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32) - cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2) - cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2) - - for i in cutlass.range_constexpr(4): - cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i]) - sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i]) - cos_vals[i * 2] = cos0 - cos_vals[i * 2 + 1] = cos1 - sin_vals[i * 2] = sin0 - sin_vals[i * 2 + 1] = sin1 - - for i in cutlass.range_constexpr(self.coarsen): - for j in cutlass.range_constexpr(8): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) - rot0 = q0 * cos_vals[j] - q1 * sin_vals[j] - rot1 = q0 * sin_vals[j] + q1 * cos_vals[j] - # convert back to BF16 to match numerics - q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1) - - return ( - q_bf16x2, - tid, - global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - num_token_heads, - ) - - -class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel): - """Eight-thread subwarps process one ``(token, head)`` row.""" - - @cute.jit - def __call__( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_quant: cute.Tensor, - q_scale: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - stream: CUstream, - ): - total_threads = q.shape[0] * self.threads_per_token - grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) - self.kernel( - positions, - q, - cos_sin_cache, - weights, - q_quant, - q_scale, - weights_out, - scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_quant: cute.Tensor, - q_scale: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - ): - ( - q_bf16x2, - tid, - global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - num_token_heads, - ) = self._load_q_and_rope(positions, q, cos_sin_cache) - - cp_op = cute.nvgpu.CopyUniversalOp() - - # layout: [coarsen, 8] - q_fp4_tile = cute.local_tile( - q_quant[token_id, None, None], - tiler=(self.coarsen, 8), - coord=(head_tile_id, sublane), - ) - - for i in cutlass.range_constexpr(self.coarsen): - # compute amax in packed bf16x2 to save instructions - # Each thread holds 16 elems. Two adjacent threads form one 32-elem - # MXFP4 block, so a width-2 shuffle gives the block amax. - amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) - for j in cutlass.range_constexpr(1, 8): - amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) - amax_bf16x2 = cute_utils.warp_reduce( - amax_bf16x2, - _bf16x2_max, - width=MXFP4_BLOCK_SIZE // 16, - ) - amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) - amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - - if in_bounds: - # compute block scale with bit manipulation - # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask - # increments the exponent whenever fp4_scale is not exactly a power of 2 - eps = cutlass.const_expr(float.fromhex("0x6p-126")) - fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0) - bits = recast_val(fp4_scale, Uint32) - ue8m0 = cute_utils.shr_u32( - bits + Uint32(0x7FFFFF), Uint32(23) - ) & Uint32(0xFF) - - # Only one of the two threads in an MXFP4 block writes the shared scale. - if tid % 2 == 0: - mx_block = sublane // 2 - q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0) - - # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent - # -A + 127 = 254 - ue8m0. - inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23) - inv_fp4_scale = recast_val(inv_scale_bits, Float32) - - vals = cute.make_rmem_tensor(16, Float32) - for j in cutlass.range_constexpr(8): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) - vals[j * 2] = q0 * inv_fp4_scale - vals[j * 2 + 1] = q1 * inv_fp4_scale - - # pack to FP4 - packed = cute.make_rmem_tensor((2,), Uint32) - packed[0] = cvt.fp32x8_to_fp4x8(vals, 0) - packed[1] = cvt.fp32x8_to_fp4x8(vals, 8) - - dst = q_fp4_tile[i, None] - cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) - cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32)) - - # Weight scaling is independent of the Q subwarp work. The first - # num_tokens * num_heads logical threads cover one weight each. - if global_tid < num_token_heads: - weight_token_id = global_tid // self.num_heads - weight_head_id = global_tid % self.num_heads - weights_out[weight_token_id, weight_head_id] = ( - weights[weight_token_id, weight_head_id].to(Float32) * scale - ) - - @cache - @staticmethod - def compile( - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - num_tokens = cute.sym_int() - max_pos = cute.sym_int() - - q = make_fake_tensor( - BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 - ) - positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) - cos_sin_cache = make_fake_tensor( - cos_sin_dtype, - (max_pos, rope_dim), - divisibility=8, - ) - weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) - q_fp4 = make_fake_tensor( - Uint8, - (num_tokens, num_heads, head_dim // 2), - divisibility=16, - ) - q_scale = make_fake_tensor( - Uint8, - (num_tokens, num_heads, head_dim // MXFP4_BLOCK_SIZE), - divisibility=4, - ) - weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) - - kernel = IndexerQMxFp4Kernel( - head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, - positions, - q, - cos_sin_cache, - weights, - q_fp4, - q_scale, - weights_out, - Float32(0.0), - stream, - options="--enable-tvm-ffi", - ) - - -class IndexerQFp8Kernel(IndexerQRopeQuantKernel): - """Eight-thread subwarps process one ``(token, head)`` row and emit - float8 e4m3fn with a single per-(token, head) scalar scale folded - into the per-token weight (mirrors ``_fused_indexer_q_rope_quant_kernel``). - """ - - def __init__( - self, - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - super().__init__(head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen) - # Each subwarp owns `coarsen` heads; we use the first `coarsen` - # threads of the subwarp to write the per-head weights using the - # fp8 scale computed in the matching loop iteration. - assert self.coarsen <= self.subwarp_size, ( - f"FP8 kernel requires coarsen ({self.coarsen}) <= " - f"subwarp_size ({self.subwarp_size}) for the weight-fold step" - ) - - @cute.jit - def __call__( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_fp8: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - stream: CUstream, - ): - total_threads = q.shape[0] * self.threads_per_token - grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) - self.kernel( - positions, - q, - cos_sin_cache, - weights, - q_fp8, - weights_out, - scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_fp8: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - ): - ( - q_bf16x2, - _tid, - _global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - _num_token_heads, - ) = self._load_q_and_rope(positions, q, cos_sin_cache) - - cp_op = cute.nvgpu.CopyUniversalOp() - - # layout: [coarsen, 16] bytes (one e4m3fn per element). - q_fp8_tile = cute.local_tile( - q_fp8[token_id, None, None], - tiler=(self.coarsen, 16), - coord=(head_tile_id, sublane), - ) - - for i in cutlass.range_constexpr(self.coarsen): - # Reduce amax across the full head_dim: each thread already holds - # the max over its 16 lanes; a width=subwarp_size warp shuffle - # spreads the head-wide max to every lane in the subwarp. - amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) - for j in cutlass.range_constexpr(1, 8): - amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) - amax_bf16x2 = cute_utils.warp_reduce( - amax_bf16x2, - _bf16x2_max, - width=self.subwarp_size, - ) - amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) - amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - - # scale = max(amax, eps) / fp8_max, then rounded UP to the next - # power of two. Adding the mantissa mask before shifting out the - # mantissa bumps the exponent whenever s isn't a pure pow2. - fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0) - bits = recast_val(fp32_scale, Uint32) - scale_exp = cute_utils.shr_u32( - bits + Uint32(0x7FFFFF), Uint32(23) - ) & Uint32(0xFF) - - # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23 - fp8_scale_bits = scale_exp << Uint32(23) - fp8_scale = recast_val(fp8_scale_bits, Float32) - # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23 - inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23) - inv_fp8_scale = recast_val(inv_scale_bits, Float32) - - # Weight fold: weights_out = weights * q_scale * scale_combined. - # All threads in the subwarp share the same fp8_scale after the - # warp_reduce above, so we let thread `sublane == i` write the - # weight for head `head_start + i`. - if in_bounds and sublane == i: - head_id = head_start + i - weights_out[token_id, head_id] = ( - weights[token_id, head_id].to(Float32) * scale * fp8_scale - ) - - if in_bounds: - # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s - # (one cp.async-shaped 128-bit store per row). - packed = cute.make_rmem_tensor((4,), Uint32) - for j in cutlass.range_constexpr(4): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2]) - q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1]) - packed[j] = cvt.fp32x4_to_fp8x4( - q0 * inv_fp8_scale, - q1 * inv_fp8_scale, - q2 * inv_fp8_scale, - q3 * inv_fp8_scale, - ) - - dst = q_fp8_tile[i, None] - cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32)) - - @cache - @staticmethod - def compile( - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - num_tokens = cute.sym_int() - max_pos = cute.sym_int() - - q = make_fake_tensor( - BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 - ) - positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) - cos_sin_cache = make_fake_tensor( - cos_sin_dtype, - (max_pos, rope_dim), - divisibility=8, - ) - weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) - q_fp8 = make_fake_tensor( - Uint8, - (num_tokens, num_heads, head_dim), - divisibility=16, - ) - weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) - - kernel = IndexerQFp8Kernel( - head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, - positions, - q, - cos_sin_cache, - weights, - q_fp8, - weights_out, - Float32(0.0), - stream, - options="--enable-tvm-ffi", - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py deleted file mode 100644 index 7cdb39e9..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ /dev/null @@ -1,173 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Triton input-staging kernel for DeepSeek V4 MegaMoE. - -Quantizes hidden states to fp8 with E8M0 group scales and repacks the -routing top-k tensors into the int64/float32 layout that the DeepGEMM -MegaMoE kernels consume. -""" - -import torch - -from vllm.triton_utils import tl, triton - - -@triton.jit -def _prepare_megamoe_inputs_kernel( - hidden_states, - x_fp8, - x_sf, - topk_ids, - topk_weights, - topk_idx_out, - topk_weights_out, - hidden_stride_m: tl.constexpr, - hidden_stride_k: tl.constexpr, - x_stride_m: tl.constexpr, - x_stride_k: tl.constexpr, - x_sf_stride_m: tl.constexpr, - x_sf_stride_k: tl.constexpr, - topk_ids_stride_m: tl.constexpr, - topk_ids_stride_k: tl.constexpr, - topk_weights_stride_m: tl.constexpr, - topk_weights_stride_k: tl.constexpr, - topk_idx_stride_m: tl.constexpr, - topk_idx_stride_k: tl.constexpr, - topk_weights_out_stride_m: tl.constexpr, - topk_weights_out_stride_k: tl.constexpr, - hidden_size: tl.constexpr, - top_k: tl.constexpr, - BLOCK_K: tl.constexpr, - GROUP_K: tl.constexpr, - BLOCK_TOPK: tl.constexpr, -) -> None: - token_id = tl.program_id(0) - k_block_id = tl.program_id(1) - - k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) - k_mask = k_offsets < hidden_size - hidden = tl.load( - hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k, - mask=k_mask, - other=0.0, - ).to(tl.float32) - - num_groups: tl.constexpr = BLOCK_K // GROUP_K - hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K]) - amax = tl.max(hidden_groups, axis=1) - amax = tl.maximum(amax, 1.0e-4) - - scale = amax / 448.0 - scale_bits = scale.to(tl.uint32, bitcast=True) - scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( - tl.uint32 - ) - scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) - rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True) - - hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K]) - scaled = hidden_groups * (1.0 / rounded_scale)[:, None] - scaled = tl.reshape(scaled, [BLOCK_K]) - fp8 = scaled.to(tl.float8e4nv) - tl.store( - x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k, - fp8, - mask=k_mask, - ) - - scale_offsets = tl.arange(0, num_groups) - packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32) - tl.store( - x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, - packed_scale, - ) - - if k_block_id == 0: - topk_offsets = tl.arange(0, BLOCK_TOPK) - topk_mask = topk_offsets < top_k - - ids = tl.load( - topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k, - mask=topk_mask, - other=0, - ).to(tl.int64) - tl.store( - topk_idx_out - + token_id * topk_idx_stride_m - + topk_offsets * topk_idx_stride_k, - ids, - mask=topk_mask, - ) - - weights = tl.load( - topk_weights - + token_id * topk_weights_stride_m - + topk_offsets * topk_weights_stride_k, - mask=topk_mask, - other=0.0, - ) - tl.store( - topk_weights_out - + token_id * topk_weights_out_stride_m - + topk_offsets * topk_weights_out_stride_k, - weights, - mask=topk_mask, - ) - - -def prepare_megamoe_inputs( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - x_fp8: torch.Tensor, - x_sf: torch.Tensor, - topk_idx_out: torch.Tensor, - topk_weights_out: torch.Tensor, -) -> None: - num_tokens, hidden_size = hidden_states.shape - if num_tokens == 0: - return - if hidden_size % 128 != 0: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires hidden_size to be " - "a multiple of 128." - ) - top_k = topk_ids.shape[1] - if topk_weights.shape != topk_ids.shape: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires topk_weights and " - "topk_ids to have the same shape." - ) - - block_k = 128 - grid = (num_tokens, triton.cdiv(hidden_size, block_k)) - block_topk = triton.next_power_of_2(top_k) - _prepare_megamoe_inputs_kernel[grid]( - hidden_states, - x_fp8, - x_sf, - topk_ids, - topk_weights, - topk_idx_out, - topk_weights_out, - hidden_states.stride(0), - hidden_states.stride(1), - x_fp8.stride(0), - x_fp8.stride(1), - x_sf.stride(0), - x_sf.stride(1), - topk_ids.stride(0), - topk_ids.stride(1), - topk_weights.stride(0), - topk_weights.stride(1), - topk_idx_out.stride(0), - topk_idx_out.stride(1), - topk_weights_out.stride(0), - topk_weights_out.stride(1), - hidden_size, - top_k, - BLOCK_K=block_k, - GROUP_K=32, - BLOCK_TOPK=block_topk, - num_warps=4, - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py deleted file mode 100644 index ed16ca6d..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ /dev/null @@ -1,1380 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CuTe DSL sparse-attention compressor for DeepSeek V4. - -The public wrappers provide the C4 fused and C128 split kernels. -""" - -from __future__ import annotations - -from functools import cache -from typing import Any - -import cutlass -import cutlass.cute as cute -import torch -from cuda.bindings.driver import CUstream -from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr -from cutlass._mlir.dialects import llvm -from cutlass.cutlass_dsl import T, dsl_user_op -from quack.compile_utils import make_fake_tensor - -_TORCH_TO_CUTE = { - torch.bfloat16: BFloat16, - torch.float32: Float32, -} - - -@dsl_user_op -def _recast_val(x, dtype, *, loc=None, ip=None): - return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip))) - - -@dsl_user_op -def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32: - out = llvm.inline_asm( - T.i32(), - [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], - "cvt.rn.bf16x2.f32 $0, $2, $1;", - "=r,f,f", - has_side_effects=False, - is_align_stack=False, - ) - return Uint32(out) - - -@dsl_user_op -def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]: - out = llvm.inline_asm( - llvm.StructType.get_literal([T.f32(), T.f32()]), - [data.ir_value(loc=loc, ip=ip)], - "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n", - "=f,=f,r", - has_side_effects=False, - is_align_stack=False, - ) - return ( - Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)), - Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)), - ) - - -@dsl_user_op -def _fp32x2_to_fp8e4m3x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint16: - out = llvm.inline_asm( - T.i16(), - [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], - "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", - "=h,f,f", - has_side_effects=False, - is_align_stack=False, - ) - return Uint16(out) - - -class SparseAttnCompressNormRopeStoreC4Kernel: - min_scale = 1.0e-4 - rcp_ln2 = 1.4426950408889634 - - def __init__( - self, - head_size: int, - state_width: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - overlap: bool, - ): - self.head_dim = head_size - self.state_width = state_width - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.elems_per_lane = 8 - self.copy_elems = 4 - self.copy_chunks = self.elems_per_lane // self.copy_elems - self.lanes_per_group = quant_block // self.elems_per_lane - self.groups_per_warp = 32 // self.lanes_per_group - self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1 - self.scale_reduce_offset = self.lanes_per_group // 2 - self.num_warps = (head_size // quant_block) // self.groups_per_warp - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = self.num_warps * 32 - self.compress_ratio = compress_ratio - self.overlap = overlap - self.window = (1 + int(overlap)) * compress_ratio - - @cute.jit - def __call__( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - group_lane = lane_id % self.lanes_per_group - group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group - elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) - ) - has_req_idx = token_idx < token_to_req_indices.shape[0] - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = ( - slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0) - ) - - if active: - req_idx = token_to_req_indices[token_idx] - start = position - Int64(self.window - 1) - - smem = cutlass.utils.SmemAllocator() - s_block_numbers = smem.allocate_tensor( - Int32, cute.make_layout((self.window,)), byte_alignment=4 - ) - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 - ) - - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if tid == row: - block_number_i32 = Int32(0) - if pos >= Int64(0): - block_index = pos // block_size - block_number_i32 = block_table[req_idx, block_index] - s_block_numbers[row] = block_number_i32 - cute.arch.sync_threads() - - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) - - cp_f32x4 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 - ) - copy_layout = cute.make_layout( - (self.copy_chunks, self.copy_elems), - stride=(self.copy_elems, 1), - ) - kv_vals = cute.make_rmem_tensor(copy_layout, Float32) - score_vals = cute.make_rmem_tensor(copy_layout, Float32) - - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if pos >= Int64(0): - block_index = pos // block_size - block_offset = pos - block_index * block_size - block_number = s_block_numbers[row].to(Int64) - head_offset = Int64((row // self.compress_ratio) * self.head_dim) - row_tensor = state_cache[block_number, block_offset, None] - for chunk in cutlass.range_constexpr(self.copy_chunks): - copy_elem = const_expr(chunk * self.copy_elems) - col_tile = ( - head_offset + (elem_base + Int32(copy_elem)).to(Int64) - ) // Int64(self.copy_elems) - kv_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=(col_tile,), - ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=( - col_tile + Int64(self.state_width // self.copy_elems), - ), - ) - cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) - cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - - for e in cutlass.range_constexpr(self.elems_per_lane): - chunk = const_expr(e // self.copy_elems) - copy_elem = const_expr(e % self.copy_elems) - score = score_vals[chunk, copy_elem] - kv = kv_vals[chunk, copy_elem] - new_max = cute.arch.fmax(local_max[e], score) - old_scale = cute.math.exp2( - (local_max[e] - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - new_scale = cute.math.exp2( - (score - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - local_sum[e] = local_sum[e] * old_scale + new_scale - local_product[e] = local_product[e] * old_scale + kv * new_scale - local_max[e] = new_max - - x = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sumsq = Float32(0.0) - for e in cutlass.range_constexpr(self.elems_per_lane): - x[e] = local_product[e] / local_sum[e] - local_sumsq += x[e] * x[e] - - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - cute.arch.sync_threads() - - rrms = rrms_shared[0] - for e in cutlass.range_constexpr(self.elems_per_lane): - elem = elem_base + e - x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) - - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) - scale_base = ( - page * k_cache.stride[0] - + kv_cache_block_size * Int64(self.token_stride) - + kv_offset * Int64(self.scale_dim) - ) - - if group_idx == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x[elem] * cos_v - x[elem + 1] * sin_v - imag = x[elem] * sin_v + x[elem + 1] * cos_v - packed = _fp32x2_to_bf16x2(real, imag) - out_base = ( - value_base - + Int64(self.nope_dim) - + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64) - ) - k_cache_u32.iterator[out_base // Int64(4)] = packed - else: - q = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_absmax = Float32(0.0) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) - q0, q1 = _bf16x2_to_fp32(q_packed) - q[elem] = q0 - q[elem + 1] = q1 - local_absmax = cute.arch.fmax( - local_absmax, - cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)), - ) - absmax = local_absmax - group_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.lanes_per_group) << 8 - | (cute.arch.WARP_SIZE - 1) - ) - for step in cutlass.range_constexpr(self.scale_reduce_steps): - offset = const_expr(self.scale_reduce_offset >> step) - absmax = cute.arch.fmax( - absmax, - cute.arch.shuffle_sync_bfly( - absmax, - offset=offset, - mask_and_clamp=group_mask_and_clamp, - ), - ) - scale_raw = cute.arch.fmax( - Float32(self.min_scale), - absmax / Float32(self.fp8_max), - ) - bits = _recast_val(scale_raw, Uint32) - ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) - inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - y0 = cutlass.min( - cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (elem_base + Int32(elem)).to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - if group_lane == 0: - k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8) - if group_idx == 0: - k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( - 0 - ) - - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 1024, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 4, - overlap: bool = True, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): - if compress_ratio != 4 or not overlap: - raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.") - if head_size != 512: - raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires head_size=512." - ) - if state_width != 2 * head_size: - raise ValueError( - "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size." - ) - if quant_block != 64: - raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires quant_block=64." - ) - if rope_head_dim != 64: - raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64." - ) - if token_stride < head_size + rope_head_dim: - raise ValueError("token_stride is too small for the packed FP8/BF16 row.") - expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1 - if scale_dim < expected_scale_dim: - raise ValueError("scale_dim is too small for the UE8M0 scale row.") - - num_positions = cute.sym_int() - num_slots = cute.sym_int() - num_req_indices = cute.sym_int() - num_kv_slots = cute.sym_int() - num_state_blocks = cute.sym_int() - num_kv_blocks = cute.sym_int() - state_cache_block_size = cute.sym_int() - block_table_width = cute.sym_int() - max_pos = cute.sym_int() - state_cache_width = state_width * 2 - - state_cache = cute.runtime.make_fake_tensor( - Float32, - (num_state_blocks, state_cache_block_size, state_cache_width), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=16), - 1, - ), - assumed_align=16, - ) - token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 - ) - positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) - slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) - block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 - ) - rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 - ) - cos_sin_cache = cute.runtime.make_fake_tensor( - Float32, - (max_pos, rope_head_dim), - stride=(cute.sym_int64(divisibility=4), 1), - assumed_align=4, - ) - k_cache = cute.runtime.make_fake_tensor( - Uint8, - (num_kv_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), - assumed_align=16, - ) - kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) - - kernel = SparseAttnCompressNormRopeStoreC4Kernel( - head_size, - state_width, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - overlap, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - Int64(0), - rms_norm_weight, - Float32(0.0), - cos_sin_cache, - k_cache, - kv_slot_mapping, - Int64(0), - stream, - options="--enable-tvm-ffi", - ) - - -class SparseAttnCompressKernel: - head_tile = 64 - rows_per_warp = 16 - row_pairs_per_warp = rows_per_warp // 2 - elems_per_lane = 4 - lanes_per_row = head_tile // elems_per_lane - num_warps = 8 - stats_warp_stride = num_warps + 1 - tb_size = num_warps * 32 - rcp_ln2 = 1.4426950408889634 - - def __init__( - self, - head_size: int, - state_width: int, - compress_ratio: int, - overlap: bool, - ): - self.head_dim = head_size - self.num_splits = head_size // self.head_tile - self.state_width = state_width - self.compress_ratio = compress_ratio - self.overlap = overlap - self.window = (1 + int(overlap)) * compress_ratio - - @cute.jit - def __call__( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - compressed_kv: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0] * self.num_splits, 1, 1) - self.kernel( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - compressed_kv, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - compressed_kv: cute.Tensor, - ): - block_id, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - row_lane = lane_id // self.lanes_per_row - col_group = lane_id % self.lanes_per_row - - token_idx = block_id // self.num_splits - split_idx = block_id - token_idx * self.num_splits - col_base = split_idx * self.head_tile + col_group * self.elems_per_lane - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) - ) - has_req_idx = token_idx < token_to_req_indices.shape[0] - active = slot_id >= Int64(0) and has_req_idx and boundary - - if active: - smem = cutlass.utils.SmemAllocator() - s_max = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.lanes_per_row, - self.elems_per_lane, - self.stats_warp_stride, - ), - stride=( - self.elems_per_lane * self.stats_warp_stride, - self.stats_warp_stride, - 1, - ), - ), - byte_alignment=4, - ) - s_sum = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.lanes_per_row, - self.elems_per_lane, - self.stats_warp_stride, - ), - stride=( - self.elems_per_lane * self.stats_warp_stride, - self.stats_warp_stride, - 1, - ), - ), - byte_alignment=4, - ) - s_product = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.lanes_per_row, - self.elems_per_lane, - self.stats_warp_stride, - ), - stride=( - self.elems_per_lane * self.stats_warp_stride, - self.stats_warp_stride, - 1, - ), - ), - byte_alignment=4, - ) - - row_pair_layout = cute.make_layout( - (self.row_pairs_per_warp, self.elems_per_lane), - stride=(self.elems_per_lane, 1), - ) - kv_vals = cute.make_rmem_tensor(row_pair_layout, Float32) - score_vals = cute.make_rmem_tensor(row_pair_layout, Float32) - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) - - req_idx = token_to_req_indices[token_idx] - start = position - Int64(self.window - 1) - cp_f32x4 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 - ) - row_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.lanes_per_row) << 8 - | (cute.arch.WARP_SIZE - 1) - ) - - for i in cutlass.range_constexpr(self.row_pairs_per_warp): - row = warp_id * self.rows_per_warp + i * 2 + row_lane - pos = start + row.to(Int64) - valid = row < self.window and pos >= Int64(0) - head_offset = ((row // self.compress_ratio) * self.head_dim).to(Int64) - - for e in cutlass.range_constexpr(self.elems_per_lane): - kv = Float32(0.0) - score = -Float32.inf - kv_vals[i, e] = kv - score_vals[i, e] = score - - block_index = Int64(0) - block_offset = Int64(0) - block_number_i32 = Int32(0) - if valid: - block_index = pos // block_size - block_offset = pos - block_index * block_size - if col_group == 0: - block_number_i32 = block_table[req_idx, block_index] - block_number_i32 = cute.arch.shuffle_sync( - block_number_i32, - offset=0, - mask_and_clamp=row_mask_and_clamp, - ) - - if valid: - block_number = block_number_i32.to(Int64) - row_tensor = state_cache[block_number, block_offset, None] - col_tile = (head_offset + col_base.to(Int64)) // Int64( - self.elems_per_lane - ) - kv_src = cute.local_tile( - row_tensor, - tiler=(self.elems_per_lane,), - coord=(col_tile,), - ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.elems_per_lane,), - coord=( - col_tile + Int64(self.state_width // self.elems_per_lane), - ), - ) - cute.copy(cp_f32x4, kv_src, kv_vals[i, None]) - cute.copy(cp_f32x4, score_src, score_vals[i, None]) - - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e]) - - for e in cutlass.range_constexpr(self.elems_per_lane): - if local_max[e] > -Float32.inf: - for i in cutlass.range_constexpr(self.row_pairs_per_warp): - exp_score = cute.math.exp2( - (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2), - fastmath=True, - ) - local_sum[e] += exp_score - local_product[e] += kv_vals[i, e] * exp_score - - for e in cutlass.range_constexpr(self.elems_per_lane): - pair_max = cute.arch.shuffle_sync_bfly(local_max[e], offset=16) - pair_sum = cute.arch.shuffle_sync_bfly(local_sum[e], offset=16) - pair_product = cute.arch.shuffle_sync_bfly(local_product[e], offset=16) - warp_max = cute.arch.fmax(local_max[e], pair_max) - warp_sum = Float32(0.0) - warp_product = Float32(0.0) - if warp_max > -Float32.inf: - local_scale = cute.math.exp2( - (local_max[e] - warp_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - pair_scale = cute.math.exp2( - (pair_max - warp_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - warp_sum = local_sum[e] * local_scale + pair_sum * pair_scale - warp_product = ( - local_product[e] * local_scale + pair_product * pair_scale - ) - if lane_id < self.lanes_per_row: - s_max[col_group, e, warp_id] = warp_max - s_sum[col_group, e, warp_id] = warp_sum - s_product[col_group, e, warp_id] = warp_product - cute.arch.sync_threads() - - out_group = tid // self.num_warps - final_lane = tid % self.num_warps - final_groups_per_pass = const_expr(self.tb_size // self.num_warps) - for pass_idx in cutlass.range_constexpr( - self.head_tile // final_groups_per_pass - ): - out_idx = pass_idx * final_groups_per_pass + out_group - out_lane = out_idx // self.elems_per_lane - out_elem = out_idx % self.elems_per_lane - - local_warp_max = s_max[out_lane, out_elem, final_lane] - global_max = local_warp_max - for step in cutlass.range_constexpr(3): - offset = const_expr(4 >> step) - global_max = cute.arch.fmax( - global_max, - cute.arch.shuffle_sync_bfly( - global_max, - offset=offset, - mask_and_clamp=row_mask_and_clamp, - ), - ) - - scale = cute.math.exp2( - (local_warp_max - global_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - global_sum = s_sum[out_lane, out_elem, final_lane] * scale - global_product = s_product[out_lane, out_elem, final_lane] * scale - for step in cutlass.range_constexpr(3): - offset = const_expr(4 >> step) - global_sum += cute.arch.shuffle_sync_bfly( - global_sum, - offset=offset, - mask_and_clamp=row_mask_and_clamp, - ) - global_product += cute.arch.shuffle_sync_bfly( - global_product, - offset=offset, - mask_and_clamp=row_mask_and_clamp, - ) - - if final_lane == 0: - compressed_kv.iterator[ - token_idx.to(Int64) * compressed_kv.stride[0] - + (split_idx * self.head_tile + out_idx).to(Int64) - ] = global_product / global_sum - - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 512, - compress_ratio: int = 128, - overlap: bool = False, - ): - if head_size % SparseAttnCompressKernel.head_tile != 0: - raise ValueError("head_size must be divisible by the 64-wide head tile.") - num_positions = cute.sym_int() - num_slots = cute.sym_int() - num_req_indices = cute.sym_int() - num_blocks = cute.sym_int() - state_cache_block_size = cute.sym_int() - block_table_width = cute.sym_int() - state_cache_width = state_width * 2 - - state_cache = cute.runtime.make_fake_tensor( - Float32, - (num_blocks, state_cache_block_size, state_cache_width), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=16), - 1, - ), - assumed_align=16, - ) - token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 - ) - positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) - slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) - block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 - ) - compressed_kv = cute.runtime.make_fake_tensor( - Float32, - (num_slots, head_size), - stride=(cute.sym_int64(divisibility=4), 1), - assumed_align=4, - ) - - kernel = SparseAttnCompressKernel( - head_size, - state_width, - compress_ratio, - overlap, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - Int64(0), - compressed_kv, - stream, - options="--enable-tvm-ffi", - ) - - -class SparseAttnNormRopeStoreKernel: - min_scale = 1.0e-4 - - def __init__( - self, - head_size: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - ): - self.head_dim = head_size - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.num_warps = head_size // quant_block - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = head_size // 2 - self.compress_ratio = compress_ratio - - @cute.jit - def __call__( - self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - elem0 = tid * 2 - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) - ) - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) - - if active: - base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) - x0 = compressed_kv.iterator[base] - x1 = compressed_kv.iterator[base + Int64(1)] - - local_sumsq = x0 * x0 + x1 * x1 - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - - smem = cutlass.utils.SmemAllocator() - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 - ) - - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - cute.arch.sync_threads() - - rrms = rrms_shared[0] - x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32) - x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32) - - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) - scale_base = ( - page * k_cache.stride[0] - + kv_cache_block_size * Int64(self.token_stride) - + kv_offset * Int64(self.scale_dim) - ) - - if warp_id == self.nope_blocks: - pair_idx = lane_id - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64) - cos_v = cos_sin_cache.iterator[cs_base] - sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] - real = x0 * cos_v - x1 * sin_v - imag = x0 * sin_v + x1 * cos_v - packed = _fp32x2_to_bf16x2(real, imag) - out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64) - k_cache_u32.iterator[out_base // Int64(4)] = packed - else: - q_packed = _fp32x2_to_bf16x2(x0, x1) - q0, q1 = _bf16x2_to_fp32(q_packed) - abs0 = cute.math.absf(q0) - abs1 = cute.math.absf(q1) - local_absmax = cute.arch.fmax(abs0, abs1) - absmax = local_absmax - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - absmax = cute.arch.fmax( - absmax, cute.arch.shuffle_sync_bfly(absmax, offset) - ) - scale_raw = cute.arch.fmax( - Float32(self.min_scale), - absmax / Float32(self.fp8_max), - ) - bits = _recast_val(scale_raw, Uint32) - ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) - inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) - y0 = cutlass.min( - cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to( - Int64 - ) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - if lane_id == 0: - k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8) - if warp_id == 0: - k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( - 0 - ) - - @cache - @staticmethod - def compile( - head_size: int = 512, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 128, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): - if quant_block != 64: - raise ValueError( - "CuTe DSL sparse-attn store currently requires quant_block=64." - ) - if rope_head_dim != 64: - raise ValueError( - "CuTe DSL sparse-attn store currently requires rope_head_dim=64." - ) - if head_size % quant_block != 0: - raise ValueError("head_size must be divisible by quant_block.") - if token_stride < head_size + rope_head_dim: - raise ValueError("token_stride is too small for the packed FP8/BF16 row.") - expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1 - if scale_dim < expected_scale_dim: - raise ValueError("scale_dim is too small for the UE8M0 scale row.") - num_positions = cute.sym_int() - num_slots = cute.sym_int() - num_kv_slots = cute.sym_int() - max_pos = cute.sym_int() - num_blocks = cute.sym_int() - - compressed_kv = cute.runtime.make_fake_tensor( - Float32, - (num_slots, head_size), - stride=(cute.sym_int64(divisibility=4), 1), - assumed_align=4, - ) - positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) - slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) - rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 - ) - cos_sin_cache = cute.runtime.make_fake_tensor( - Float32, - (max_pos, rope_head_dim), - stride=(cute.sym_int64(divisibility=4), 1), - assumed_align=4, - ) - k_cache = cute.runtime.make_fake_tensor( - Uint8, - (num_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), - assumed_align=16, - ) - kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) - - kernel = SparseAttnNormRopeStoreKernel( - head_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - Float32(0.0), - cos_sin_cache, - k_cache, - kv_slot_mapping, - Int64(0), - stream, - options="--enable-tvm-ffi", - ) - - -def compress_kv_sparse_attn_cutedsl( - state_cache: torch.Tensor, - token_to_req_indices: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - compressed_kv: torch.Tensor, - head_size: int = 512, - state_width: int = 512, - compress_ratio: int = 128, - overlap: bool = False, -) -> None: - if positions.numel() == 0: - return - compiled = SparseAttnCompressKernel.compile( - head_size=head_size, - state_width=state_width, - compress_ratio=compress_ratio, - overlap=overlap, - ) - compiled( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - compressed_kv, - ) - - -def norm_rope_insert_sparse_attn_cutedsl( - compressed_kv: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - rms_norm_weight: torch.Tensor, - rms_norm_eps: float, - cos_sin_cache: torch.Tensor, - k_cache: torch.Tensor, - kv_slot_mapping: torch.Tensor, - kv_cache_block_size: int, - kv_block_stride: int, - head_size: int = 512, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - compress_ratio: int = 128, -) -> None: - if positions.numel() == 0: - return - norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype) - if norm_weight_dtype is None: - raise ValueError( - "CuTe DSL sparse-attn store supports rms_norm_weight dtype " - f"bf16/fp32, got {rms_norm_weight.dtype}." - ) - if k_cache.ndim != 3: - raise ValueError( - "CuTe DSL sparse-attn store expects the real DeepSeek V4 " - f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." - ) - compiled = SparseAttnNormRopeStoreKernel.compile( - head_size=head_size, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, - compress_ratio=compress_ratio, - norm_weight_dtype=norm_weight_dtype, - ) - compiled( - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, - ) - - -def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( - state_cache: torch.Tensor, - token_to_req_indices: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - rms_norm_weight: torch.Tensor, - rms_norm_eps: float, - cos_sin_cache: torch.Tensor, - k_cache: torch.Tensor, - kv_slot_mapping: torch.Tensor, - kv_cache_block_size: int, - kv_block_stride: int, - head_size: int = 512, - state_width: int = 1024, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - compress_ratio: int = 4, - overlap: bool = True, -) -> None: - if positions.numel() == 0: - return - norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype) - if norm_weight_dtype is None: - raise ValueError( - "CuTe DSL sparse-attn fused store supports rms_norm_weight dtype " - f"bf16/fp32, got {rms_norm_weight.dtype}." - ) - if k_cache.ndim != 3: - raise ValueError( - "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 " - f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." - ) - compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile( - head_size=head_size, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, - compress_ratio=compress_ratio, - overlap=overlap, - norm_weight_dtype=norm_weight_dtype, - ) - compiled( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, - ) - - -def compress_norm_rope_store_cutedsl( - state_cache: torch.Tensor, - num_actual: int, - token_to_req_indices: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - state_width: int, - cos_sin_cache: torch.Tensor, - kv_cache: torch.Tensor, - k_cache_metadata: Any, - pdl_kwargs: dict, - head_dim: int, - rope_head_dim: int, - compress_ratio: int, - overlap: bool, - use_fp4_cache: bool, - rms_norm_weight: torch.Tensor, - rms_norm_eps: float, - quant_block: int, - token_stride: int, - scale_dim: int, -) -> None: - if compress_ratio == 4: - # For C4A, the single fused kernel is faster than the two-kernel version. - fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - kv_cache, - k_cache_metadata.slot_mapping, - kv_cache.shape[1], # paged KV cache block size - kv_cache.stride(0), - head_size=head_dim, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=448.0, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - compress_ratio=compress_ratio, - overlap=overlap, - ) - else: - # For C128, the two-kernel version is faster than the single fused kernel. - compressed_kv = torch.empty( - (num_actual, head_dim), - dtype=torch.float32, - device=state_cache.device, - ) - compress_kv_sparse_attn_cutedsl( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - compressed_kv, - head_size=head_dim, - state_width=state_width, - compress_ratio=compress_ratio, - overlap=overlap, - ) - norm_rope_insert_sparse_attn_cutedsl( - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - kv_cache, - k_cache_metadata.slot_mapping, - kv_cache.shape[1], # paged KV cache block size - kv_cache.stride(0), - head_size=head_dim, - rope_head_dim=rope_head_dim, - fp8_max=448.0, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - compress_ratio=compress_ratio, - ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py b/TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py deleted file mode 100644 index 85a78883..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Quantization config for DeepSeek V4.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -from vllm.config import get_current_vllm_config -from vllm.model_executor.layers.fused_moe import FusedMoE -from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod -from vllm.model_executor.layers.quantization import QuantizationMethods -from vllm.model_executor.layers.quantization.fp8 import Fp8Config -from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - is_layer_skipped, -) - -_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8") - -if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4Config, - ) - - -class DeepseekV4FP8Config(Fp8Config): - """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch. - - DeepSeek V4 checkpoints always use FP8 block quantization for - linear/attention layers. The MoE expert weights vary by checkpoint: - - ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts - with ue8m0 (e8m0fnu) FP8 linear scales. - - ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block - experts with float32 FP8 linear scales. - - The dispatch and the linear scale dtype are both keyed off - ``expert_dtype`` from the model's hf_config; missing values default - to ``"fp4"`` so existing FP4 checkpoints stay unchanged. - - NOTE: ``expert_dtype`` is resolved lazily because this config is - constructed during VllmConfig setup, before ``set_current_vllm_config`` - is active. Reading hf_config eagerly in ``__init__`` would always see - the default ``"fp4"`` and silently misroute Flash-Base checkpoints. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._resolved_expert_dtype: str | None = None - self._resolved_moe_quant_algo: str | None = None - self._nvfp4_config: ModelOptNvFp4Config | None = None - # ``is_scale_e8m0`` is a property that resolves on first read, - # by which time the current vllm_config has been set. - - @property - def expert_dtype(self) -> str: - if self._resolved_expert_dtype is None: - try: - hf_config = get_current_vllm_config().model_config.hf_config - except Exception: - # vllm_config not yet set; defer the decision until a - # later call lands inside set_current_vllm_config. - return "fp4" - expert_dtype = getattr(hf_config, "expert_dtype", "fp4") - if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES: - raise ValueError( - f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; " - f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}." - ) - self._resolved_expert_dtype = expert_dtype - from vllm.logger import init_logger - - init_logger(__name__).info_once( - "DeepSeek V4 expert_dtype resolved to %r", expert_dtype - ) - return self._resolved_expert_dtype - - @property - def is_scale_e8m0(self) -> bool: - # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert - # checkpoints (Flash-Base) store them as float32. - return self.expert_dtype == "fp4" - - def _resolve_moe_overrides(self) -> None: - if self._resolved_moe_quant_algo is not None: - return - try: - hf_config = get_current_vllm_config().model_config.hf_config - except Exception: - return - quant_cfg = getattr(hf_config, "quantization_config", None) or {} - algo = (quant_cfg.get("moe_quant_algo") or "").upper() or None - self._resolved_moe_quant_algo = algo or "" - - @property - def moe_quant_algo(self) -> str: - self._resolve_moe_overrides() - return self._resolved_moe_quant_algo or "" - - def _get_nvfp4_config(self) -> ModelOptNvFp4Config: - if self._nvfp4_config is None: - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4Config, - ) - - self._nvfp4_config = ModelOptNvFp4Config( - is_checkpoint_nvfp4_serialized=True, - kv_cache_quant_algo=None, - exclude_modules=[], - group_size=16, - ) - return self._nvfp4_config - - @classmethod - def get_name(cls) -> QuantizationMethods: - return "deepseek_v4_fp8" - - @classmethod - def override_quantization_method( - cls, hf_quant_cfg, user_quant, hf_config=None - ) -> QuantizationMethods | None: - if not ( - isinstance(hf_quant_cfg, dict) - and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8") - ): - return None - model_type = getattr(hf_config, "model_type", None) - if model_type == "deepseek_v4" or user_quant == "deepseek_v4_fp8": - return "deepseek_v4_fp8" - return None - - def get_quant_method(self, layer, prefix): - if isinstance(layer, FusedMoE): - if is_layer_skipped( - prefix=prefix, - ignored_layers=self.ignored_layers, - fused_mapping=self.packed_modules_mapping, - ): - return UnquantizedFusedMoEMethod(layer.moe_config) - if self.expert_dtype == "fp4": - if self.moe_quant_algo == "NVFP4": - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4FusedMoE, - ) - - return ModelOptNvFp4FusedMoE( - quant_config=self._get_nvfp4_config(), - moe_config=layer.moe_config, - ) - return Mxfp4MoEMethod(layer.moe_config) - # expert_dtype == "fp8": fall through to Fp8Config which - # returns Fp8MoEMethod with block-wise float32 scales. - return super().get_quant_method(layer, prefix) - - def is_mxfp4_quant(self, prefix, layer): - if not isinstance(layer, FusedMoE) or self.expert_dtype != "fp4": - return False - return self.moe_quant_algo != "NVFP4" diff --git a/TEMP/dsv4thing.zip b/TEMP/dsv4thing.zip new file mode 100644 index 0000000000000000000000000000000000000000..2b1685dd27fba9b51d7a303181b34323c740079a GIT binary patch literal 20388 zcmZ^pQ?O{unyja7+qP}nwr$(Cxu$K~YudJL+jpOi6LGswbdRX48c+XA=EE0Rr63Iq zf&%cb&z5OW=f8gby+HvG0GK$tvA9@R+L_b4dO!mJf}8;Z0Q}dcssalDUiXV{)$ou0 zbt?SV90&ji0L}lHgZw|{D2WP+$cob0n*47AqyHK7FJ=Ct9%b414F-gsQ*{~2aNUa1 z@EbtYg8;O%LRtjCN-yFahB(-PRvYjrgkRrL$hqyx)fcfHj}zaiJH!#KeBm78BePh7 zm|(y{T>?8paf$O}lj1+{Dy!ZufF`OaQs_a-^4~ykyNy8*3m!l`!)HP_j4%nq0vbYtF`3= z_l+dp*NY^);`@?~DH2ObF}NtP-gziX4w`=LnJ^>A3#wd7ia?yiOr^#%&oUGWMTdac z==QYT$?I$6&5^`HQs5U;!5hX!HYae2e@{8*R8g&QqsqXH{D15@YeSPup5N2_gw9TjA=gwPruPLyHmo33= z{tOwv!0N_;;>rVomk98^4=JtJenGOL{wU$xP2yu((~dPO`xXt)s9Jydnwg>i8ri{f z09UQ=(40x*5*j8Y8#}4xdR%pGcrQun6r{l+psx{p#{dBU+Wv{p{|}1se*^{fe+0$Y z-pkO7=ey`Ir9uqhlp)yFmhkuUkqYfAz#*JWLwg-Lyp;3nMF0nDh zPDqeNl>M$l4+N$k+!fwZPAGLDG|4Y%s@H4bI|o?D#gyPX%+u%r$BT-}uoxvoo5O6S zRFRZv2Ew4&N5UVOvl|zGvbleK-rqt0W#|Ae^Ml<#hSGujcSBA8!%&p}G}O)1$%W3r z^MCU*?;juY|E&^CYx~2I`!4_Cpm3C8>3bT@4vf{a0BLMsz^tQ$H(@|QU4;Xs4dn5X z4r=LI$KM|2CLbEniH>G48%6Hh(V5^h&kCx+l-{Hl$y+!vGj+=nCp7gyu12Fj!l^wY!1Fn<;RSpX-dBEDHX>CRP|VlqRQ!+g9{E+F%&ITuEN7G1VEkC zLWuo)$vYLZ{=JKz#XtN$U;BCdjiIKosj#pZG_neizs1rDT~=vo|#Xs7Q%Svo)~k%-&gO+ZA=D z9ygosCTO_0qC{W0pt!Oh@I$Jjmc;2CeI_O`o6S&bT>s1vCE+@%b1$_JihZ zOf8e`Tz#8nFG%m9m{zWR0GepK$4q58%Lr%ElPv-$F|Z{7@{tiT`3*k24GDUm2>%7U z%KHAg*Z*4m{M_~TF?bimPLq1FPbT2YZ+h=;J%H7N$^Y(fIXZsE!7H-+(O;{irbbc7RdXK`28|6rQT8m!%_@ zLo1QjO@4w8hL}vIOAXTra6*ikXwO|_Jcfm6co$afhyPI^2)mStz~)O@MPGwF#86Ew z)xzjabGXIcX9V!ch9TU9Ry=qVc0~a6p8Q&We%8d%YngUAzRKnty@@DL`3l+Hi{$)!QKh>VjH?d$|Q1{Z^7|_n?9HLwB6*Yw3ZRt;?IdJ{lCQ z3<;|qPD3wj*5{B#7u3`tiYgyr!Oh?>N24A&wq`~i&=y40?=Z9IU~QNXnMpS+R;olo z6RHv+YqAYpRNnfhVwEe3Q9ql?W0r59xt@U=@VN*HyS-`FrU(kSbPrb)K({yWs@!yI zcB)KrJV|)fThcCUOg5@CAUMJ~%a&1ZYlhW!i$|3YxDLsdoVpPb_bl z;Jy2DY#qCs0D>#He#vDrc$X=BMrXK#qtdGHP6`+^H-fwa5!zlEuNxptEX_>^Ya5C{ zuw}u5KDPz`iSZ6xH}_TCi-9Y)b#K6nRgi2TK%sO9w@knRzw+9ft$QZ|v6 ztkw0bs_xKk0)jSlQrvZ1crwD$){-Qucw{&VX+3EV6#Z%7k|a790{~aFq0jtY>@62f zGwQkuU{zW2A&d``bgBtSGJnvQHgq|A(~K^foTQ9ie^LOTrV=Ea-(#|2yVkTi5#BdW zjF~Riki}9bYF&6YZO3jgV@t(wR)P_kso+t`gSveHhMd}kFNUt#?Yn%F5E(nvHB<7? z8)5m_gi8h}(UK9b6Pm}YC^BTsK!SHiXH3@dC_sD>%c0gW-Qmf6iR(}(-Zu1tFis@4 z-WqNYz>vqJ4n$Wi@P?(nY*vxxE=nD}_Xcf$_`%?-+ThGOFq8@QZ@Yj|`X7r^#vcNb zwyFZ!I5dfOUJN(h@@Wft0RH}dnst3c3ycFDH3)w?J^Qafykw=6Vlre#sr2<_Y|i&Z zx^0f+K`3}wm{m85nH_!;xMQSREN{R)^sa{5hIu|o+SU4$>b!MnJ->UpD8XIwY3opD zLtTHYPs|Wrs)I@;@m)H^y{tdt2elJrF0R^1H1m|b+;*LpeWwn0Q){vcWS^G+bSQd@P5 zuuUO{NE}AbfbH4^Us7q+598x4$^&8kwxsLeCQl)|`4>oTDw}L`Y)kNRcF!2*360*i zP;!&!!O8)QEJnN?twgs;_;uU1$HS^|0xjgp^DgF6mzI&L$&BiTSC7HDv1axR;!g0- zc6L8=+f$#Cw%!#*{qYdCc67Um)R%4oW>oCE0pLZxe)q?NfP=7mE73z@61g&@+UVSC{0w~l_;~O z7m$%WqEu~Y4Og`)n*Hs$UV}QBi-;A?5Yk(s0&DSjD5W;RVyp7+Tp)ATu^(vQN~d?3 z-plnGnd(dzd!?#y}!5d9}`*@DCCgw*6p{lWL+*nd>Q#L zcNE{70BmCkst5M~-o{GmkLnuv>_7M)MhhqjS$GqF^9ssDvRFQke#eF?8498nn~a;3 z>%Gp0ldKTnUXQo5Ije7t+s)exOYRAop}l}jcCQC7?6p~-is4b#pLV?F`jgAGW&0Yy zO_1$vD#77z(+6sKxjc_fp<8i<&z2z^Z^`*^IVrbSDR7j_`L25SJ3i?DMO@i7r!xqU z002eU|DCwZ{)4#wCqkQ>+L<~Tx|seS$W`zUZWaAoH5!*mo2{rnb!7tYqkCb6E$gl4 z7L_@jGp0byoCNE$PS|PjkUG|xM`DVHUt=!gRmQGE8R%VztgW+zse&} znouJL9)uS%$9pd^k2y)Kupwkenpu3ek$J<=5!FR0tRQo9szN7PLi73KA546UZGAVs zjqspVAuC@5EwUKWf$!*;r{Dcj~7gEI5s*-wK`w38@7G;DDrV0#zNqdQOt8r+lt(b{6T zh!o85)hXhF*P4H#KE;+1;ObN6CpqK_+!Ruk5rvO__cTm(A^OG1XtXSCn4J(ci?iaQ8PD5e_sJ>g8RtC!CIHg<PMD~sRk+N=qEZCO{i0q z#wsDzs>Q{DtrQ>!aM{An(@??FT3&u%XvXz8BdOA!Z)HMI2~W#WT~r14vc<9Z@|t61zyPnu@%@!R;>u^Yl(s9k9hw1*Bb^X6vP-Z z5zcM5IG@x7k*__?$4Kv)@9e6!)Yn@hBz-Pk4pVWAH%Pcqr9($8?v<04hzr$cl>7LEH5FG zEZ8$-$*u<#-0;zXJ4l(syq56gt23*zD!)C581e}PXKhP&kxSl_>xOa&Z zhxex2!PVVxW~UY}uOl+MQ(rLTNuhh$29r6~iXH&~VS5yw!);H=yngTR{&M=&besP- zm6U`F<~tlP@1cj(*I@MRPl!oXZnX-;!Dytu@9)DfwcB4HOt`BKvUKAxN8j)eslhKE za}H)~tYu9#vAA!AO5W;FS?CQfPc{DLzWm8G=PQ`*jnh9m^&vVs46Ut-_r@TX43|#1 zR)Ci+?Wve<1v>>Ea;lM1E%21WR^^lgA;Mpv)W?F05Nd4YX1(D#deOi;)=yZ@n4w93 zDuV`zCayB>u=jROZXep5wJ5mE?DBP-W&P3d)Vgo_m0*#x^mKK7?}wLDcNMR}N&x1T z02xQJ6&b_8I!RXM9Bd2F(LJ=K?0HibzF#H zOi);0+O%x`ZnGd5d`w~~$k5qW^u!AM23zM=69Nfijs%h-dL5~V@alYO=WOyPt*ch0 z2X!Todc;e-)mLY|5}WZC57aOQaYm`nP1#&6@nv610tgZJ`Kl$-hg$4FIFkJStn&W? zj>?zj_Z7tEY=C`&ZA4$=F2-@vvS$#bFFKn?>VC4-D$bvMd}N8-hsdP(=Gd}5TK+^5 z-4(suikQ7Z^W(dble9u!X^*LZweg;CWZ<&qauZf{I6ht*ltM@U{1#dsRFZGhx;;j` zlFXguMSEKoDqoHu@<2L{Wu^wT#n-<8NYDcEwjf;Kc`1>j#XG$lzVcbAQWmmCa#-+DSUXnuD=#H?-T z)o5F!p7|O{2q7DoYU?w2$Xll@(;|oKu*TzUfK=j*hg`XnKxQQMXWP`1sPepPv%$gg7(Tym#e z&&c@hFEUbYALY#cI1^}mgug8*6cT9NH#V}%<>FC%=vjDW;jnt|y5#+e@lu!`j6%W# zE-KC=mqxB$uhTD~4Btcxqt69MZv4iU%(z*EF{qlsxLZQgf8Cthz5Yx)kyFYui}`xW zqDAL;@y`l+e>sTpX13YmcW*>}2NT~|{BoOor_Zzy zuZz0!V|POYl6Z*+lP_rQ3kn1F#0w600gTk`lvuVDehE4gGigj9ivOwsQ=_G3S!>&} z&mK3i6M2Y$9K^#y>+t2MIQ#vlLel`2DxKaGLPK1B{hP6P(qmAlfw7T z@GosQxKww{Rx6n$5BWP4WFa__*A87$0@o3~B;HUWdaD*KCtFbm_6WS4M4p933qS5J zQ#Z!XR;zZ+&?M7pe39f`LL%V%Em~kajytTp{Q97D=f#en%delKv(Ll*6e#5C&Ghs7 z`SAT{^e@-vCt`;z`qH|?)Ae%0=RJPjj2~s#w|teWpD)B?8@b`taX8)eEPr!B4gKZ^ z-SzY9#_nTU!J*zZEn_Z1 zWP0P5^A0{(A!|$ghS28C>qU=T8Yd|Y?w5dYHcu!XCr39ov7Z;+yqKUrQ#}+CXa(VE zw@HU(92n*#(1+$g#%_)bfsl&?6F~fzK{{mkVrDl;%^laJD7lL)l9qNY=sKL?G`1`8 zbjT~yNWgwyiBEiMM|CFlo_;f4{aHg1>hP zfWe$c5jI=`;I0!q5~FVsh~W&O6Jhbac+OHp;bwT(*in*0a_M_pC%!oc&bej<#UR0O zl_@h<^HKTDdK?fsq$SuksFOc z6S-?%yv28Woe+a}?^p?234QTx)NZaXn>V^K0f$}D3L&l2E8CkZtn3SVsJ1N_=O9FFtZ;bF+F?b+P=SK1K3k`1Z zmM)@^yRD`Uhq+#;y1`w)&I_}1GAr4YE!nk24>rQ#(cJ4PQMA>5lp_@2!1K!u_JInf z@Aynts>SJ5c5AW-KUz%jI)o7(8o|-1MdO!ba66T&p`^-2;k|2eX{tyMGA~;6BIF3R zsHZ8FHhuZ)?zO()_dsXC=5rdn)a z#9-Zy;nlW%Nl;*j5C|f{bBxG9;a_t_?bjc9U}F01HQEg{W6Q|-2=S&x<#_u~YnwfU zg8Rp)OO|*4{cl+<`-A~A?fkrA>k8qa0pD6(yeXv;&5HDZ1p*ci)&=xp&oN!QdDsYM zi`v*i^J1NRDE{{ri+o_7OwM42cuer(U>=63h;t_eFE!?Zf%7&0=$aIs z@$c524hnf=|1fJg;a!W(DRw>2Kb439DC#;MA*_JPYBGWw(`-D6cF*Jk$}~d^F(PXe zDO3Zf&A~3nS-QY$M9cw(GR9~h?H(_`)rSu=$kdS^(1=QH42S@LCEZFEVm%MPyC4z_ zaGc=VtmP z+1tEg9G^_kHiCG&VlVv!))soM_CLj*{tMmwFw>Xc<2Nu@Hvnpu&u(s9pFO~vGVQI! z_TOe7FTWYa0|TF_YpUWIbMs6Ri>L9zb!=^u_*wePPCkzQfY+PO)osp&R07_EHT^F> zC$TccQk127LaJi(%m~Q`PV>4iJ43QEJXJzURtjBzWvP~ogq0eS(0XUr!h7m>tX6KjM7E>fC<|G05Z$!HT)FOQO#-uEz!NA7_%_d}LcVwo)~_5Z zHCG2;ZWy69`{@L(rn6D`l>hl5w~8;Q7(EYD$nA4==n)g~f9EPxKZ;S|oL=*#zQ;xY z!^9GneN&zHd8X3Gpncxr(qd%_J#@`an^u#dZWWycA+LkQuf_H2yuTVW)A1XO`Rb&| z7MD^DcHPP^4MM6Fd25A^9GlJ!PDk%=i3}$k zwO@F>k`m`H*0ajVQt&$_T!p}BwAVroBcIv27e-c}_UEC!9|%)}v^J$;<^GVR z#1D=zcfkhn#NpyO%~ZcyTOwaLVZa_1jrl%AN+qkNEHtA3>lR$J$zrVnd;O@)8W<@@ zBYP?_+(cO4@V9r^_kqtY!3W z@|gEM^3N9Re)_JRCs`=Ll>pa~v>7`N2elehka3LBnWId%g^V1D1bV{6B^XSn_B;&z9^pH%qGCdV^c^#M3x5+Ma38_GEx4vrD)OTPDalj*Q zPODxGIb#=riY;F%J#H`7T{gzScdOw9UiTN*rS(pLFnuu++;K{^`3lRFVc&d_G;9Qo zM-CWd0f%!a9hT;}GxH11+61pk74l>A;0p$hmSd%xCmu?*EGGGH(6GE;RZu#`@p3;b zHrOWCiJ&=@L0kX`0F+0}zNi2~?0FK8$9qcY8QkqaEMrj;P43Z~)Gil%kByfO%!SE2 zRWVy%DmVCOK>$_+>w}X&6xN(Lx}3fSeR>2Z-TBnmu|JXpoPq!tL*guzeyLfBoD5Ag z^HV5G0BR^T(7O+cQW9Mugeypk8#7Pi+-&d|Ji2j;;N4KPhP+t`O@x^Viieb}MQCdb zh;63&#%UpMWwUr8E*fhU&)?Ekl7O#5(y_HzHP5BT6TpTy-!Y7ZZMgMIzfpPF8nAfr zsQ%rz@2=n^%6(pG{WwF}HSx~m_&Z2er1ud$;ff2tlKM*V)U1Hy;AzA3(fU3KNK`6j z@5-l`nC%1E<8fFOa2eddn$);(wITnVWpmBnVE{7n34VJhAKgF1(utHYZC84viE6a?jF|m})i(-iI zVcEfv5kky-W`X#x+6FeePQM7~_G`;cN6{%e_A!~6DPe+;b0*Er-LbNa_bCp4W918fQQ0fCKJPP2{C1 z7Q=*ZSi%a$eMdnT0`F@znobRO6>L$f?~1jmzd(|)(9 z71(3UA0Y}&Q^lJaP3JrY;jI@bSyr<5suv2;^&Y_1R&IBM6_4{zh!1NU&F=qC7Bv%r?4x)Re%@0*s5HJz+h z>y$&ihc}bHC__Avz&IDpvC2r7U!mjlc?{~hM2TgY_IZd}vm!2eG&+@iLXDpHQ2u

@M%^@q@SWQ-21DUt3ru&)jG?p|**i<@*o$J@Mt?RpC_bK!PoY?sG@Rg|3n5lTQL==BxnD6eRBkJ6mSpVlIZ zsShBuqux}nsw|h=sGd12+dEIhp>NxK8Ki~0RbnQI=CUUgQnh~c7 zUEcR4)<7iad)Y!-6BR!vm8fWOR#37YxlAiGsJ92P%_w+LISOl$Z+14QGT5bU42yap z9G*CNz>}#5=4&~_XiD$HenRnUYBe{F7}~QYsKO|#+oAkca5-9>sN@2aeb?b3ZPcIi zO-Ngvi0`K>NpM2Y8d#bMD%#LQ68%pnPXi`}&tZ2Q5jbo||8h9zd)nxvOc6Ad+#d?NaAC(Ragt(;nMV_i`FH6ba%FSgZwyV&ak2zC$Y#=cE{dOtq8PC?&n+`#-{4~jgURj|PjhX+&BMs9riv z*luqWC7{kn#B4et<@>vR7==$R*>1EsDQj)G4Q~gNm8gFcjjNs;T=Q<60jXZP)an9V zXULx|#omClq0u{jx{zlX`UhLhEhsFb=)^-A=2H=UQpzF7lgQT!N_rRzRL+u+I%a5> zkSf>&^oA5_)eUl5Nh=Xuh81~JD#0&wVd00JiguZ{Uf1)yo~Wu7+CRZbY3ur0O6pS; z@%I(A;+Idfq>|-^SMSj}Qb=PEls-=HQK_iDcH4pdqpMa!_YvK5x80}(1_HSO;Pv(^ zLNkh#MUD6+z!NMkHk^{)D75=E z?0LZhxCKPvQpx>Jt8N{wFS!9UqJX|+1zeq4kVYtM24XZS&on}bfoZLRv%<4+hsfgx zmi>p+t$l4y&C)$uOH-*T?&n4l6pErIGRgc8+6MdbPi11C=`np}KU23QroP%+SI($! z&B9hCtlJasceh|gZRG6_w3X}UiPrYdZK0a_9CqwBK?h<)7U}|eY8Mi+AG}GC2f)Bg zMnqI~mo_pINgdRw0qAEH&Un!T58XAtb+A8q%K+KI4}Cf8kN&?ptevMf9%Cc`fNr(_ z-eFn)M~8*Drv+Q z$za3_V%^S@=AMfy@5i0fW*Sla@EBLFYU!EtP9hzmP*wOD7rUzd>sXC^Ur092T zW~nQ(X~nkMBSws}S*nR`lF8|-(Uo-YaBw%3of!kc`b><7hrwL5b^DX!1Ez|ybfjKIo^J=SlPN=MF#2f>zNB1DzvF1qZ&i2Y@LYOZco;9jA^4dghmJ_)1dv6`H zO-5pUC!7gFhN#bhe;r`}1jBaD_3wpQC!^S(Rg30@j&<{eSNcy$KpW#b&)T?UJOc#jJhF9+1-SO6aT`vx~>1bR8Zx=__Ky zyi3_O-uI7vzK#uP-~yT8>GO$v3$n6UWNeEO%DZOtXyx14{7XL8!OTeJ=PiMXC1E2* zU|jrJ9+#sR?=0-9EWOUBTzx(^7uCsmDty*6&l)F?JW`e=%bR7Qp78E!p?+lfbc?6| zQgbDWPP!eZ#;6w285B2zDM&Yi`gibZS=`06&h|FA6*=M~+m6 zi2~J{8a7m%)C;pmjFI(eslCVa?g&^@Xrf+*C0_CnI3Vv{!X~Rcern z_D*RldoU?2_DG&%C*azAfOMqfPL5L&gCd&JkJ^uvHEyb=y@etrnyS(#+bnAwj)!$e zs5*fH>v>|-w8HGs(VH*lh>LC6Xwe$Nv%|80hygQ5a^wh;4yn2zLBHLa#yscrkpo8t ze4@0sKtq;+cD2h+@9F3En-8&o(%eG5NB3#U!zgvH7+0~T&Ia;TZh*;~rq4l5WT^%E zI9lbhi@KgBts{PAnPPv)?ItlOI%oDNKn3e;s8w0}`^F2@G6KpZlitas>lREUFVG1h zv61){PIO)R6TDx__X|Z6QSZ>vH?6~nNPsBeXk2yHlO)oiw6B{{^hQJZ5r&hF1mTcS zV2nx9=qBe~j@Ep858$up-c3NkA`0KO{TJeR5sbdxr<^dsx>(cs-Us$!v zWLuv-Fl7(wh2}l|RcDOyF~_h1 z<6sCbR|o2f1Si^YF^B?~n8;ci0lH8F>QIFGKO3vu2n;u9`5PXZFb&u;@qb}Xu8L~y zLVNY{4DzX-VPLS6AURI!0=vZ=iT{egCxwd%_Xh8 zaX_#yu~3^g8Y1j8@d9j{zApnwVpu$plYSDbQ3_~U-4-FlXw(e@u)v5S3COq$OePVr znJY2BLzw6+12g(u^DwvoLZ!Co)eWzq?1M+5HgqEDgPw9C{Dy|skpnKl;7CMN#A1LI zw4=^8EJ9De-8sYe?ZiIWMX)u*4h-}*b-)LZfdpOfF>y0?vLjqp>#mlxgTn<|&>6Uw zx*qMbqwyg8o}1=7!qeGotq@a2Lu3aL)Mlzs2Ua7jX`+Rq5<#>;UED@! zfdQBdzaaKWHKtFZKRz(Sye=1oZJL<#QM5cjNL>Swg7lpx1sokB33^dSfi%l2!ufsM zeVN8c;|F`P2Gi4{5(4T2J{AG0R_HPXZ+M=pF@p!XSm$>YWd-_y#vW_9v2kl+A3i*VVDv6k zfc8Ga@e&u|_4o4v^bgtPK@)XNm7hOZa_u;{C;@*5PKt$o3;b5>_D5(kc3f3 zMiB%7B}ohX3p_a$HJ7Ke(o=i|-SY*W31pa`lkiV@FS5klhk1*={23sv> z^6P_8@1Hc>$oXKd4WZpg`|%oKE+U8zAYZ3-;KIOZP%^A*6RAG~GKVmF@w*r2CA_i-q10?a~; z!M|ZDY?+>_$!2^UwQ0jVoU3&rnKQZWBM5ydo*P&DRE)NQCU%9T`Q5lt@XTN$79%qJ z3-34z9Yrn5gSP(UrCI=v zz!2a_XfUi076?`EgT^^Wva%>7LDR_Pd})LMNC_}V-lmxuT*|$%iwq25PlVwm!0zGL z*#|vBAZxnX}C@&XSH{)tF60 zeHcg0O$H9ifbk$dIzotxjqlP&enh^lFTC?av}xc0iB>A&w=uV|q8fc8ddO&+{LH`5 z*S)$z1^Bf|M9mExf(!jcuUiv^!!@8;)8ynuPhrwr8I@8P3&Myx-B`H+3YdI@o}Ym_ zb3_NM;N@f@efd5xXPp57NVMv3s5Sf8{fD#^u5+o#sYk}a=q?P>rJn+t?j0)qAkGLWF@o1Tz{G0nvYxQ)6g6gd3=^ z?t&r&5Vjrs524wW=0LWYCn89jvq#P}U?>Mf%$2}&HwnU$4@c|V-;e|X!Z8XNkW@%E zRXh$lbqb+G=+NFS*Wj%$6Bj2YHV>mYjAh})$To=ykK1Eu0|`O$=06ZA{+S)~444FM zzytLm2y~8U~WCry%iPSfverL4SSeGI6g;APN}I0PhUx_|*`ba4C zg|FA#BwHNL#X}_^{4l4M0YL8~&)|gs7&U1r@f>Wrq|phMAz5jTrvQBuiw?*6F21Q` z83r7SfvPdaA;#$2_~iA?_!;>aRbn#4n-G4=X<;ZZljo!$cXpDpk#&}saOKyT8L0tu z9Nn2bVapF4t*#bK=FEwM@bJ(2Y&c>v77azxuOWrzhugu)#`W0ZjfuR4&SvPl(;r6z4aQ#~q- z+QvIz#Ok1Xh|SdEtShadqg7pA{v>g-yt}a5MwNl2Tq^MX=iFS-TJC-q$R!U+gF$K? zQ{otkmO^R^#os--=qK)=eZ17#;5}p53D#ztttDSFv35)bRfZvXTzbvP`bs%X35x-l zC)>Fi$FxN*Pl(N5vu09L>$La$=$%Gss*a^1ts#A?WHI(Y0ZfZVG^4DSXK6Li7{e6@ zuadt!zqb?h&k{T6i09b~eb2|orGvRLI_2Do?;fbEXC_yRLR9ZlS-{y~2-Zg-beX_J zu_nK>0A72|;_?bb+@|-GK)mkJ3#XL|zxct;N1Dwew-TX0t^8tQ6TT792qZ>yGeDF1 z%!l+0cK;5>3Z4KSaKHH$7+Hp z?ZnfQmfmiofm4gr`Mnqk@Z@A&f-~wleTKCh);6|^p`z;2b&ct}0El--?yO)N`JWkp zQ3IK3G>YCwCoN?U92q2Q(m$c1qX{|<%qV8IZhmOd%!Nua&*zODYbwUIAiSkSe5)Xs` zfpIhBf+y_P%p9F}dSr&Lexo2O$e&h60(2B$ZLwp=_+d{R%NOpl-k8!1SHV?@OaCw3 z>mHNE4=-?tjXtvO@}|D7VjnYy!x4oLMiQA}6S0SYk^p4L*u+?IPsL@RH+qrDkGcK<)l{7sF@Q zt?`QMysrk3S^Bf{vtr5-aLp{+nGTt#_g-<+?E`O@glaF~Zl&{;fc6$h*rI z0^VNBHB*Tjmy{jXXXWk#$KSP0{29^5{C$l%*yzLWqzN2tdtu=&?NPQrpp1kjr;zD= z5g{&26lYwXZLq2BeGnC26Z_Fnl)Y`?4|8>KR*J_#%I;?+mA?P3(kvR?zFs1-1@at% z%^lSUwjL5bn&s1x2AM=5U^li56eXtjfD4cYJi1J6=T?D=heR1KL|a`!>X|x{JKR|~ z4e(e8#B;iK7)FIXD06G?A_PZkj}YhY_n7#Xmr>3psRi#JZaGyD z{%h<5lBPx@7z0d3B1RMH%mK#@?P|>F)%hO+cu~7FiTMsB{E&wp`ONS_gPQb41eX?H zJx6e<9x%Yjn~Z@I2aX8sGXieS1-w2zj0QKBdj1*bs3%97k!E1sS_1iUX|b+)8Iqgq zk6-%ZidrLZsh~@ty&vv6eZ54cW%C@5%OAEW4QW-m z3_qWhTA*I<;<9M@;v4mXSLHmMqCvIIZdrYGwVhucEfm-*UL{IthsPvrwr{X~$#uI` zyWX-aewq;pObUy$SV|>DquN zV(~_EVtWp?5YLq6<>FP6V5KcQb`5{|n$Ez#Ed_u15HxLq(WTZ$tcSNt*Y=(UJ7a`u zC|(Kk(*q)6jKZWetaOQ48Lp47Wcv!{gOKcjk7SOw(9BL_(WhH^)VQ>B7>)D{V)+1N zc4k47MFnThDu2qLO1eHx&3cZvO1rqNudj{erpN=#a$v2S4xAqq~rgIBZM29U5ckrnqFnX6fu~#J)iYskRKSIqQ+S zdWa4N*QB*78^|LCZ!+qu{QXQAA3k4V=ehR^RaDtU$TuS=F^;2%0eSUtNFlw~7Lxc6 ze{Tdb`F-X zKgq1tb0ZpWH?B`dZZnb@JVK?x8+s{usaK6pvx84eo4 z(H5Xzr9~+2L``xAO3-_>+}Vl8VR^fcs?zjPRqj5h1cSwCXOEli$rj;;ER@2h<-6AF zs-OSa<2Q108bj6T86LD`m}4%e1K-!4KRUn=8``J)(=^S=)UnC zQ?HO)kAryfKNR=A8bOfafQbI3cq)0Cb za1cUoA_7tZN;fnqp>vZv_sj`*uIGL;d*=J~%-Va_`uDDfE@$|8c_=5Ddn|1516wvZ zd%_`=xpg1=j>J}~sw_Q!%P|qETs`bp^L&TV$<5;fgOjeQ=}+F}%5&d+hwc{$XB$22 zr$Q;;DoLes+&3TfrxpFV zw(dnQhN{{mbh~L*GvGfpJyfK9R~qA48Z`s1RlkY~x{gYk%X{e6ZkZLJmdUiEv`1@v zh>YN{YBa1d0Ksl=s;vR3v#!vkOb-$PJ8gCCLzH+J~M!VO%{^Tcny-_IfFR0Su``_pNE1=#EEkXt4m1tH+X(3T9xLWMD4 zDfEng{`q0qg#Gl0y;+cwBKM5hPTFOE%cAep!qXG}=^9K+)$%WT5Ze3s`2>7|nq^h9 z8pVg9!7TC!Fb)6a*X%9-zKNc-cT@YhCap_l$APmoBaJB)F}wQG13@>^lPJBTvMX&0 zdY>9*6nT&Ajon*Vdks{GYyO1V&J+5oTMV;>14?N=WU2UQHg1Zm+=`QLTsENfpLZm6 zxn=`J`|^ssGKzU1JlA!=FWuN3AErdDLc5b1m~e-ug|+;iKK zJKwdH1)$`#dzio3TiltYN$V%*Pb3cSypvP5Em?*N;vV)LYd?Q|&>T?8HcL_4{_`Z4 zde*bL+fH0p>!6#TAR-1PD@Xt8HAxVRuM&eCe8P{slkI0Je5nQG@lIc;X=X7jz;e0` zP`(BB2@o<2yT4&!Hs{=J{D{b(&$10^_zG_KGe_C~(IfG-W;Wuez9tVFVxXbZ^>%%DFp5ISM^hb83m7*KEp^jU` zR`$}ru61xE)KM)Y1Kb+y#vLszoK}M>kbV{7{%fObw6-h|0n5~HzHxe+2)Z+}`Oz80 zO4_#}*Yk4vYE87_7xl=`PC6ql8DMXRxy@{3^Ab3f&6w7JA@O5;v&y`k63sPTlL9pw zbzDSLPd=IcT;MSKQBkk!wc4=P7$ed+(L|5xJPDI|?1GMNW07#6>ZCDx>EEE{*S5DXr&w5&eob`;x_eM^EqOVQdOU) z@po|u(lHl&<_k>@SMCb+>t~nG_~vPG^91;=EzV-5glWGeL2AP{*o!>?w!Iu2SW#gX zIi^u2;46EDQDKb{em~#%@!1c%#sG26OwcSljO&V3_*^D~QdbwK{fM9q{Jw-rL4J?4tkm7FkSFwU@3i+;x2sBwcC;jxf9Rl-@`CH5Uw1Xr*RgJ; zH+8B+^2m2R!?cZk&!eug5DK1p*b;lc)(>QKEX$yfEUB}4IyKAPa7;$@TwgUzSgdhrcxD?j>2GQ$J8)-gAS9B7o`SxkNt&jEs|lM&Gfe^=oQjZu?y8^jL&BGm zsNR;rH7bWuLE^SP!e@}7vvswr!+w7nZ9L`+#73xL`Ax)Tk*>To3?atdJLUV1t^Srf zaqq*eJEx{QPpQn~ek2l!Wc9=tR_`lEJ>DISwvN#%m&V>6Xcid; z_V!TJa_<)<7HW1$p1bW3Kr6HY=1 z2{G-Mg>U3Ea=A=#T>0f-*3qdYP|mS*C8Zmsu^GzT3VOf^yq5RzmI4f!RA4Xs#5p;< zLn+DXL{8mSt)hlIG0s_(Ibvgg?8|`5j>AmB3n{gq8w88*wd@2t9{F=YMSnlc>|^CB$`)@oY0FQcHK32qmo2W;1put1;%89)+XlvZA%o%$e(vSq4t55+czXQd z3Gw`2MxSzQ4h45Aut&a6lcMc96<`DrMQ$pvu2#X7rv^L-;lY$Nt?r8!Bvjk6Enqj8IS_;IL<4~@KvkN*7{9~O*T tuple[torch.Tensor, torch.Tensor]: + """ + Casts a tensor from e2m1fn to e4m3fn losslessly. + """ + assert x.dtype == torch.int8 + assert x.ndim == 2 + out_dim, in_dim = x.size() + in_dim *= 2 + fp8_block_size = 128 + fp4_block_size = 32 + assert in_dim % fp8_block_size == 0 and out_dim % fp8_block_size == 0 + assert scale.size(0) == out_dim and scale.size(1) == in_dim // fp4_block_size + + x = x.view(torch.uint8) + low = x & 0x0F + high = (x >> 4) & 0x0F + x = torch.stack([FP4_TABLE[low.long()], FP4_TABLE[high.long()]], dim=-1).flatten(2) + + # max_fp4 (6.0) * MAX_OFFSET must fit in e4m3fn (max 448) + # 6.0 * 2^6 = 384 < 448; 6.0 * 2^7 = 768 > 448; so MAX_OFFSET_BITS = 6 + MAX_OFFSET_BITS = 6 + + bOut = out_dim // fp8_block_size + bIn = in_dim // fp8_block_size + # bOut, bIn, 128, 128 + x = x.view(bOut, fp8_block_size, bIn, fp8_block_size).transpose(1, 2) + # bOut, bIn, 128*4 + scale = scale.float().view(bOut, fp8_block_size, bIn, -1).transpose(1, 2).flatten(2) + ## bOut, bIn, 1 + scale_max_offset_bits = scale.amax(dim=-1, keepdim=True) / (2**MAX_OFFSET_BITS) + # bOut, bIn, 128*4 + offset = scale / scale_max_offset_bits + # bOut, bIn, 128, 128 + offset = offset.unflatten(-1, (fp8_block_size, -1)).repeat_interleave(fp4_block_size, dim=-1) + x = (x * offset).transpose(1, 2).reshape(out_dim, in_dim) + return x.to(torch.float8_e4m3fn), scale_max_offset_bits.squeeze(-1).to(torch.float8_e8m0fnu) + + +mapping = { + "embed_tokens": ("embed", 0), + "input_layernorm": ("attn_norm", None), + "post_attention_layernorm": ("ffn_norm", None), + "q_proj": ("wq", 0), + "q_a_proj": ("wq_a", None), + "q_a_layernorm": ("q_norm", None), + "q_b_proj": ("wq_b", 0), + "kv_a_proj_with_mqa": ("wkv_a", None), + "kv_a_layernorm": ("kv_norm", None), + "kv_b_proj": ("wkv_b", 0), + "o_proj": ("wo", 1), + "gate_proj": ("w1", 0), + "down_proj": ("w2", 1), + "up_proj": ("w3", 0), + "lm_head": ("head", 0), + + "embed": ("embed", 0), + "wq_b": ("wq_b", 0), + "wo_a": ("wo_a", 0), + "wo_b": ("wo_b", 1), + "head": ("head", 0), + "attn_sink": ("attn_sink", 0), + "weights_proj": ("weights_proj", 0), +} + + +def main(hf_ckpt_path, save_path, n_experts, mp, expert_dtype): + """ + Converts and saves model checkpoint files into a specified format. + + Args: + hf_ckpt_path (str): Path to the directory containing the input checkpoint files. + save_path (str): Path to the directory where the converted checkpoint files will be saved. + n_experts (int): Total number of experts in the model. + mp (int): Model parallelism factor. + + Returns: + None + """ + torch.set_num_threads(8) + n_local_experts = n_experts // mp + state_dicts = [{} for _ in range(mp)] + + for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))): + with safe_open(file_path, framework="pt", device="cpu") as f: + for name in f.keys(): + param: torch.Tensor = f.get_tensor(name) + if name.startswith("model."): + name = name[len("model."):] + if name.startswith("mtp.") and ("emb" in name or name.endswith("head.weight")): + continue + name = name.replace("self_attn", "attn") + name = name.replace("mlp", "ffn") + name = name.replace("weight_scale_inv", "scale") + name = name.replace("e_score_correction_bias", "bias") + if any(x in name for x in ["hc", "attn_sink", "tie2eid", "ape"]): # without .weight + key = name.split(".")[-1] + else: + key = name.split(".")[-2] + if key in mapping: + new_key, dim = mapping[key] + else: + new_key, dim = key, None + name = name.replace(key, new_key) + for i in range(mp): + new_param = param + if "experts" in name and "shared_experts" not in name: + idx = int(name.split(".")[-3]) + if idx < i * n_local_experts or idx >= (i + 1) * n_local_experts: + continue + elif dim is not None: + assert param.size(dim) % mp == 0, f"Dimension {dim} must be divisible by {mp}" + shard_size = param.size(dim) // mp + new_param = param.narrow(dim, i * shard_size, shard_size).contiguous() + state_dicts[i][name] = new_param + + os.makedirs(save_path, exist_ok=True) + + for i in trange(mp): + names = list(state_dicts[i].keys()) + for name in names: + if name.endswith("wo_a.weight"): + weight = state_dicts[i][name] + scale = state_dicts[i].pop(name.replace("weight", "scale")) + weight = weight.unflatten(0, (-1, 128)).unflatten(-1, (-1, 128)).float() * scale[:, None, :, None].float() + state_dicts[i][name] = weight.flatten(2, 3).flatten(0, 1).bfloat16() + elif "experts" in name and state_dicts[i][name].dtype == torch.int8: + if expert_dtype == "fp8": + scale_name = name.replace("weight", "scale") + weight = state_dicts[i].pop(name) + scale = state_dicts[i].pop(scale_name) + state_dicts[i][name], state_dicts[i][scale_name] = cast_e2m1fn_to_e4m3fn(weight, scale) + else: + state_dicts[i][name] = state_dicts[i][name].view(torch.float4_e2m1fn_x2) + save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors")) + + for file in ["tokenizer.json", "tokenizer_config.json"]: + old_file_path = os.path.join(hf_ckpt_path, file) + new_file_path = os.path.join(save_path, file) + if os.path.exists(old_file_path): + shutil.copyfile(old_file_path, new_file_path) + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--hf-ckpt-path", type=str, required=True) + parser.add_argument("--save-path", type=str, required=True) + parser.add_argument("--n-experts", type=int, required=True) + parser.add_argument("--model-parallel", type=int, required=True) + parser.add_argument("--expert-dtype", type=str, choices=["fp8", "fp4"], required=False, default=None) + args = parser.parse_args() + assert args.n_experts % args.model_parallel == 0, "Number of experts must be divisible by model parallelism" + main(args.hf_ckpt_path, args.save_path, args.n_experts, args.model_parallel, args.expert_dtype) diff --git a/TEMP/dsv4thing/generate.py b/TEMP/dsv4thing/generate.py new file mode 100644 index 00000000..c35c8030 --- /dev/null +++ b/TEMP/dsv4thing/generate.py @@ -0,0 +1,155 @@ +import os +import json +import sys +from argparse import ArgumentParser +from typing import List + +import torch +import torch.distributed as dist +from transformers import AutoTokenizer +from safetensors.torch import load_model + +from model import Transformer, ModelArgs +current_dir = os.path.dirname(os.path.abspath(__file__)) +encoding_dir = os.path.join(current_dir, '../encoding') +sys.path.insert(0, os.path.abspath(encoding_dir)) +from encoding_dsv4 import encode_messages, parse_message_from_completion_text + + +def sample(logits, temperature: float = 1.0): + """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU, + since it avoids the GPU-to-CPU sync in torch.multinomial.""" + logits = logits / max(temperature, 1e-5) + probs = torch.softmax(logits, dim=-1, dtype=torch.float32) + return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1) + + +@torch.inference_mode() +def generate( + model: Transformer, + prompt_tokens: List[List[int]], + max_new_tokens: int, + eos_id: int, + temperature: float = 1.0 +) -> List[List[int]]: + """Batch generation with left-padded prompts. + + The first forward pass processes [min_prompt_len:] tokens (prefill phase). + Subsequent passes generate one token at a time (decode phase). For positions + still within a prompt, the ground-truth token overrides the model's prediction. + """ + prompt_lens = [len(t) for t in prompt_tokens] + assert max(prompt_lens) <= model.max_seq_len, f"Prompt length exceeds model maximum sequence length (max_seq_len={model.max_seq_len})" + total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens)) + tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long) + for i, t in enumerate(prompt_tokens): + tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long) + prev_pos = 0 + finished = torch.tensor([False] * len(prompt_tokens)) + prompt_mask = tokens != -1 + for cur_pos in range(min(prompt_lens), total_len): + logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos) + if temperature > 0: + next_token = sample(logits, temperature) + else: + next_token = logits.argmax(dim=-1) + next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token) + tokens[:, cur_pos] = next_token + finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id) + prev_pos = cur_pos + if finished.all(): + break + completion_tokens = [] + for i, toks in enumerate(tokens.tolist()): + toks = toks[prompt_lens[i]:prompt_lens[i]+max_new_tokens] + if eos_id in toks: + toks = toks[:toks.index(eos_id)] + toks.append(eos_id) + completion_tokens.append(toks) + return completion_tokens + + +def main( + ckpt_path: str, + config: str, + input_file: str = "", + interactive: bool = True, + max_new_tokens: int = 100, + temperature: float = 1.0, +) -> None: + world_size = int(os.getenv("WORLD_SIZE", "1")) + rank = int(os.getenv("RANK", "0")) + local_rank = int(os.getenv("LOCAL_RANK", "0")) + if world_size > 1: + dist.init_process_group("nccl") + global print + if rank != 0: + print = lambda *_, **__: None + torch.cuda.set_device(local_rank) + torch.cuda.memory._set_allocator_settings("expandable_segments:True") + torch.set_default_dtype(torch.bfloat16) + torch.set_num_threads(8) + torch.manual_seed(33377335) + with open(config) as f: + args = ModelArgs(**json.load(f)) + if interactive: + args.max_batch_size = 1 + print(args) + with torch.device("cuda"): + model = Transformer(args) + tokenizer = AutoTokenizer.from_pretrained(ckpt_path) + print("load model") + load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"), strict=False) + torch.set_default_device("cuda") + print("I'm DeepSeek 👋") + + if interactive: + messages = [] + while True: + if world_size == 1: + prompt = input(">>> ") + elif rank == 0: + prompt = input(">>> ") + objects = [prompt] + dist.broadcast_object_list(objects, 0) + else: + objects = [None] + dist.broadcast_object_list(objects, 0) + prompt = objects[0] + if prompt == "/exit": + break + elif prompt == "/clear": + messages.clear() + continue + messages.append({"role": "user", "content": prompt}) + prompt_tokens = tokenizer.encode(encode_messages(messages, thinking_mode="chat")) + completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id, temperature) + completion = tokenizer.decode(completion_tokens[0]) + print(completion) + messages.append(parse_message_from_completion_text(completion, thinking_mode="chat")) + else: + with open(input_file) as f: + prompts = f.read().split("\n\n") + prompt_tokens = [tokenizer.encode(encode_messages([{"role": "user", "content": prompt}], thinking_mode="chat")) for prompt in prompts] + completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id, temperature) + completions = tokenizer.batch_decode(completion_tokens) + for prompt, completion in zip(prompts, completions): + print("Prompt:", prompt) + print("Completion:", completion) + print() + + if world_size > 1: + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--ckpt-path", type=str, required=True) + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--input-file", type=str, default="") + parser.add_argument("--interactive", action="store_true") + parser.add_argument("--max-new-tokens", type=int, default=300) + parser.add_argument("--temperature", type=float, default=0.6) + args = parser.parse_args() + assert args.input_file or args.interactive, "Either input-file or interactive mode must be specified" + main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature) diff --git a/TEMP/dsv4thing/kernel.py b/TEMP/dsv4thing/kernel.py new file mode 100644 index 00000000..ea7976fa --- /dev/null +++ b/TEMP/dsv4thing/kernel.py @@ -0,0 +1,536 @@ +import torch +import tilelang +import tilelang.language as T +from typing import Tuple, Optional + + +tilelang.set_log_level("WARNING") + +pass_configs = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, +} + +FP8 = "float8_e4m3" +FP4 = "float4_e2m1fn" +FE8M0 = "float8_e8m0fnu" +BF16 = "bfloat16" +FP32 = "float32" +INT32 = "int32" + + +def fast_log2_ceil(x): + """Compute ceil(log2(x)) via IEEE 754 bit manipulation. Avoids slow log/ceil intrinsics.""" + bits_x = T.reinterpret("uint32", x) + exp_x = (bits_x >> 23) & 0xFF + man_bits = bits_x & ((1 << 23) - 1) + return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0)) + + +def fast_pow2(x): + """Compute 2^x for integer x via IEEE 754 bit manipulation.""" + bits_x = (x + 127) << 23 + return T.reinterpret("float32", bits_x) + + +def fast_round_scale(amax, fp8_max_inv): + return fast_pow2(fast_log2_ceil(amax * fp8_max_inv)) + + +@tilelang.jit(pass_configs=pass_configs) +def act_quant_kernel( + N, block_size=128, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, + round_scale=False, inplace=False +): + """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16.""" + M = T.symbolic("M") + fp8_min = -448.0 + fp8_max = 448.0 + fp8_max_inv = 1 / fp8_max + num_stages = 0 if round_scale or inplace else 2 + blk_m = 32 + group_size = block_size + # Internal computation in FP32; scale_dtype controls output storage format. + compute_dtype = FP32 + out_dtype = in_dtype if inplace else out_dtype + + @T.prim_func + def act_quant_kernel_( + X: T.Tensor[(M, N), in_dtype], + Y: T.Tensor[(M, N), out_dtype], + S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( + pid_m, + pid_n, + ): + x_shared = T.alloc_shared((blk_m, group_size), in_dtype) + x_local = T.alloc_fragment((blk_m, group_size), in_dtype) + amax_local = T.alloc_fragment((blk_m,), compute_dtype) + s_local = T.alloc_fragment((blk_m,), compute_dtype) + y_local = T.alloc_fragment((blk_m, group_size), out_dtype) + y_shared = T.alloc_shared((blk_m, group_size), out_dtype) + + for _ in T.Pipelined(1, num_stages=num_stages): + T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) + T.copy(x_shared, x_local) + T.reduce_absmax(x_local, amax_local, dim=1) + for i in T.Parallel(blk_m): + amax_local[i] = T.max(amax_local[i], 1e-4) + if round_scale: + s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv) + else: + s_local[i] = amax_local[i] * fp8_max_inv + if inplace: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.Cast( + out_dtype, + T.Cast(compute_dtype, T.Cast(out_dtype, T.clamp( + x_local[i, j] / s_local[i], fp8_min, fp8_max + ))) * s_local[i], + ) + else: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.clamp( + x_local[i, j] / s_local[i], fp8_min, fp8_max + ) + for i in T.Parallel(blk_m): + S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i]) + T.copy(y_local, y_shared) + T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) + + return act_quant_kernel_ + + +def act_quant( + x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None, + scale_dtype: torch.dtype = torch.float32, inplace: bool = False, +) -> torch.Tensor: + """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16. + When scale_fmt is set, scales are rounded to power-of-2 (MXFP).""" + N = x.size(-1) + assert N % block_size == 0 + tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32 + z = x.contiguous() + y = torch.empty_like(z) if inplace else torch.empty_like(z, dtype=torch.float8_e4m3fn) + s = z.new_empty(*z.size()[:-1], N // block_size, dtype=scale_dtype) + kernel = act_quant_kernel( + N, block_size, scale_dtype=tl_dtype, + round_scale=scale_fmt is not None, inplace=inplace, + ) + kernel(z.view(-1, N), y.view(-1, N), s.view(-1, N // block_size)) + if inplace: + x.copy_(y) + return x + return y, s + + +@tilelang.jit(pass_configs=pass_configs) +def fp4_quant_kernel( + N, block_size=32, in_dtype=BF16, scale_dtype=FE8M0, inplace=False +): + """Block-wise FP4 quantization. Power-of-2 scale via bit ops. inplace=True does fused quant+dequant.""" + M = T.symbolic("M") + fp4_max = 6.0 + fp4_max_inv = 1.0 / fp4_max + blk_m = 32 + group_size = block_size + compute_dtype = FP32 + out_dtype = in_dtype if inplace else FP4 + + @T.prim_func + def fp4_quant_kernel_( + X: T.Tensor[(M, N), in_dtype], + Y: T.Tensor[(M, N), out_dtype], + S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( + pid_m, + pid_n, + ): + x_shared = T.alloc_shared((blk_m, group_size), in_dtype) + x_local = T.alloc_fragment((blk_m, group_size), in_dtype) + amax_local = T.alloc_fragment((blk_m,), compute_dtype) + s_local = T.alloc_fragment((blk_m,), compute_dtype) + y_local = T.alloc_fragment((blk_m, group_size), out_dtype) + y_shared = T.alloc_shared((blk_m, group_size), out_dtype) + + for _ in T.Pipelined(1, num_stages=2): + T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) + T.copy(x_shared, x_local) + T.reduce_absmax(x_local, amax_local, dim=1) + for i in T.Parallel(blk_m): + amax_local[i] = T.max(amax_local[i], 6 * (2**-126)) + s_local[i] = fast_round_scale(amax_local[i], fp4_max_inv) + if inplace: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.Cast( + out_dtype, + T.Cast(compute_dtype, T.Cast(FP4, T.clamp( + x_local[i, j] / s_local[i], -fp4_max, fp4_max + ))) * s_local[i], + ) + else: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.clamp( + x_local[i, j] / s_local[i], -fp4_max, fp4_max + ) + for i in T.Parallel(blk_m): + S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i]) + T.copy(y_local, y_shared) + T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) + + return fp4_quant_kernel_ + + +def fp4_act_quant( + x: torch.Tensor, block_size: int = 32, inplace: bool = False, +) -> torch.Tensor: + """Block-wise FP4 quantization. inplace=True does fused quant+dequant back to BF16.""" + N = x.size(-1) + assert N % block_size == 0 + z = x.contiguous() + y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2) + s = z.new_empty(*z.size()[:-1], N // block_size, dtype=torch.float8_e8m0fnu) + kernel = fp4_quant_kernel(N, block_size, inplace=inplace) + kernel(z.view(-1, N), y.view(-1, y.size(-1)), s.view(-1, N // block_size)) + if inplace: + x.copy_(y) + return x + return y, s + + +@tilelang.jit(pass_configs=pass_configs) +def fp8_gemm_kernel(N, K, out_dtype=BF16, accum_dtype=FP32, scale_dtype=FP32): + assert out_dtype in [BF16, FP32] + + M = T.symbolic("M") + group_size = 128 + block_M = 32 + block_N = 128 + block_K = 128 + + @T.prim_func + def fp8_gemm_kernel_( + A: T.Tensor[(M, K), FP8], + B: T.Tensor[(N, K), FP8], + C: T.Tensor[(M, N), out_dtype], + scales_a: T.Tensor[(M, T.ceildiv(K, group_size)), scale_dtype], + scales_b: T.Tensor[(T.ceildiv(N, group_size), T.ceildiv(K, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as ( + bx, + by, + ): + A_shared = T.alloc_shared((block_M, block_K), FP8) + B_shared = T.alloc_shared((block_N, block_K), FP8) + C_shared = T.alloc_shared((block_M, block_N), out_dtype) + Scale_C_shared = T.alloc_shared((block_M), FP32) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype) + + # Improve L2 Cache + T.use_swizzle(panel_size=10) + T.clear(C_local) + T.clear(C_local_accum) + + K_iters = T.ceildiv(K, block_K) + for k in T.Pipelined(K_iters, num_stages=4): + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[bx * block_N, k * block_K], B_shared) + # Cast scales to FP32 for computation; scales_b has one value per block_N group + Scale_B = T.Cast(FP32, scales_b[bx * block_N // group_size, k]) + for i in T.Parallel(block_M): + Scale_C_shared[i] = T.Cast(FP32, scales_a[by * block_M + i, k]) * Scale_B + + T.gemm(A_shared, B_shared, C_local, transpose_B=True) + # Separate accumulator for scale-corrected results (2x accumulation precision) + for i, j in T.Parallel(block_M, block_N): + C_local_accum[i, j] += C_local[i, j] * Scale_C_shared[i] + T.clear(C_local) + T.copy(C_local_accum, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return fp8_gemm_kernel_ + + +def fp8_gemm( + a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor, + scale_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """C[M,N] = A[M,K] @ B[N,K]^T with per-128 block FP8 scaling on both A and B.""" + assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous" + assert a_s.is_contiguous() and b_s.is_contiguous(), ( + "Scaling factor tensors must be contiguous" + ) + tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32 + K = a.size(-1) + M = a.numel() // K + N = b.size(0) + c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype()) + kernel = fp8_gemm_kernel(N, K, scale_dtype=tl_dtype) + kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s) + return c + + +@tilelang.jit(pass_configs=pass_configs) +def sparse_attn_kernel(h: int, d: int, scale=None): + """Sparse multi-head attention via index gathering + online softmax (FlashAttention-style). + For each (batch, seq_pos), gathers top-k KV positions by index, computes attention + with numerically stable running max/sum, and includes a learnable attn_sink bias.""" + b = T.symbolic("b") + m = T.symbolic("m") + n = T.symbolic("n") + topk = T.symbolic("topk") + if scale is None: + scale = (1.0 / d) ** 0.5 + + num_stages = 2 + threads = 256 + block = 64 + num_blocks = tilelang.cdiv(topk, block) + + @T.prim_func + def sparse_attn_kernel_( + q: T.Tensor[(b, m, h, d), BF16], + kv: T.Tensor[(b, n, d), BF16], + o: T.Tensor[(b, m, h, d), BF16], + attn_sink: T.Tensor[(h,), FP32], + topk_idxs: T.Tensor[(b, m, topk), INT32], + ): + with T.Kernel(m, b, threads=threads) as (bx, by): + q_shared = T.alloc_shared((h, d), BF16) + kv_shared = T.alloc_shared((block, d), BF16) + o_shared = T.alloc_shared((h, d), BF16) + acc_s_cast = T.alloc_shared((h, block), BF16) + + idxs = T.alloc_fragment(block, INT32) + acc_s = T.alloc_fragment((h, block), FP32) + acc_o = T.alloc_fragment((h, d), FP32) + scores_max = T.alloc_fragment(h, FP32) + scores_max_prev = T.alloc_fragment(h, FP32) + scores_scale = T.alloc_fragment(h, FP32) + scores_sum = T.alloc_fragment(h, FP32) + sum_exp = T.alloc_fragment(h, FP32) + + T.clear(acc_o) + T.clear(sum_exp) + T.fill(scores_max, -T.infinity(FP32)) + T.copy(q[by, bx, :, :], q_shared) + + for t in T.Pipelined(num_blocks, num_stages=num_stages): + for i in T.Parallel(block): + idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1) + for i, j in T.Parallel(block, d): + kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0) + for i, j in T.Parallel(h, block): + acc_s[i, j] = T.if_then_else(idxs[j] != -1, 0, -T.infinity(FP32)) + T.gemm(q_shared, kv_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + for i, j in T.Parallel(h, block): + acc_s[i, j] *= scale + T.copy(scores_max, scores_max_prev) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(h): + scores_scale[i] = T.exp(scores_max_prev[i] - scores_max[i]) + for i, j in T.Parallel(h, block): + acc_s[i, j] = T.exp(acc_s[i, j] - scores_max[i]) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(h): + sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, acc_s_cast) + for i, j in T.Parallel(h, d): + acc_o[i, j] *= scores_scale[i] + T.gemm(acc_s_cast, kv_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + for i in T.Parallel(h): + sum_exp[i] += T.exp(attn_sink[i] - scores_max[i]) + for i, j in T.Parallel(h, d): + acc_o[i, j] /= sum_exp[i] + T.copy(acc_o, o_shared) + T.copy(o_shared, o[by, bx, :, :]) + + return sparse_attn_kernel_ + + +def sparse_attn( + q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float +) -> torch.Tensor: + b, s, h, d = q.size() + # Pad heads to 16 for kernel efficiency (stripped after) + if h < 16: + q = torch.cat([q, q.new_zeros(b, s, 16 - h, d)], dim=2) + attn_sink = torch.cat([attn_sink, attn_sink.new_zeros(16 - h)]) + o = torch.empty_like(q) + kernel = sparse_attn_kernel(q.size(2), d, softmax_scale) + kernel(q, kv, o, attn_sink, topk_idxs) + if h < 16: + o = o.narrow(2, 0, h).contiguous() + return o + + +@tilelang.jit(pass_configs=pass_configs) +def hc_split_sinkhorn_kernel(hc: int, sinkhorn_iters: int, eps: float): + n = T.symbolic("n") + mix_hc = (2 + hc) * hc + threads = 64 + + @T.prim_func + def hc_split_sinkhorn_kernel_( + mixes: T.Tensor[(n, mix_hc), FP32], + hc_scale: T.Tensor[(3,), FP32], + hc_base: T.Tensor[(mix_hc,), FP32], + pre: T.Tensor[(n, hc), FP32], + post: T.Tensor[(n, hc), FP32], + comb: T.Tensor[(n, hc, hc), FP32], + ): + with T.Kernel(n, threads=threads) as i: + mixes_shared = T.alloc_shared(mix_hc, FP32) + comb_frag = T.alloc_fragment((hc, hc), FP32) + T.copy(mixes[i, :], mixes_shared) + + for j in T.Parallel(hc): + pre[i, j] = T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + eps + for j in T.Parallel(hc): + post[i, j] = 2 * T.sigmoid(mixes_shared[j + hc] * hc_scale[1] + hc_base[j + hc]) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = mixes_shared[j * hc + k + hc * 2] * hc_scale[2] + hc_base[j * hc + k + hc * 2] + + row_sum = T.alloc_fragment(hc, FP32) + col_sum = T.alloc_fragment(hc, FP32) + + # comb = comb.softmax(-1) + eps + row_max = T.alloc_fragment(hc, FP32) + T.reduce_max(comb_frag, row_max, dim=1) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = T.exp(comb_frag[j, k] - row_max[j]) + T.reduce_sum(comb_frag, row_sum, dim=1) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / row_sum[j] + eps + + # comb = comb / (comb.sum(-2) + eps) + T.reduce_sum(comb_frag, col_sum, dim=0) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps) + + for _ in T.serial(sinkhorn_iters - 1): + # comb = comb / (comb.sum(-1) + eps) + T.reduce_sum(comb_frag, row_sum, dim=1) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / (row_sum[j] + eps) + # comb = comb / (comb.sum(-2) + eps) + T.reduce_sum(comb_frag, col_sum, dim=0) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps) + + T.copy(comb_frag, comb[i, :, :]) + + return hc_split_sinkhorn_kernel_ + + +def hc_split_sinkhorn(mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6): + b, s, _ = mixes.size() + pre = mixes.new_empty(b, s, hc_mult) + post = mixes.new_empty(b, s, hc_mult) + comb = mixes.new_empty(b, s, hc_mult, hc_mult) + kernel = hc_split_sinkhorn_kernel(hc_mult, sinkhorn_iters, eps) + kernel(mixes.view(-1, (2 + hc_mult) * hc_mult), hc_scale, hc_base, + pre.view(-1, hc_mult), post.view(-1, hc_mult), comb.view(-1, hc_mult, hc_mult)) + return pre, post, comb + + +@tilelang.jit(pass_configs=pass_configs) +def fp4_gemm_kernel(N, K, out_dtype=BF16, accum_dtype=FP32, scale_dtype=FP32): + """FP8 act x FP4 weight GEMM kernel. + + C[M, N] = A_fp8[M, K] @ B_fp4[N, K]^T + + Act: 1x128 quant on K (reduce dim), FP8 with configurable scale dtype + Weight: 1x32 quant on K (reduce dim), FP4 with E8M0 scale + + B is stored as [N, K//2] in float4_e2m1fn_x2, logical [N, K] in fp4. + The FP4 values are packed along the K (last) dimension. + + Strategy: load FP4 sub-blocks of size [block_N, sub_K] (sub_K=32), + cast FP4 to FP8 via float, then do FP8xFP8 GEMM. + Apply act scale (per 128 on K) and weight scale (per 32 on K) to the accumulator. + """ + M = T.symbolic("M") + act_group_size = 128 + weight_group_size = 32 + block_M = 32 + block_N = 128 + block_K = 32 # matches weight_group_size for simple scale handling + n_sub = act_group_size // block_K # 4 sub-blocks per act scale group + + @T.prim_func + def fp4_gemm_kernel_( + A: T.Tensor[(M, K), FP8], + B: T.Tensor[(N, K), FP4], + C: T.Tensor[(M, N), out_dtype], + scales_a: T.Tensor[(M, T.ceildiv(K, act_group_size)), scale_dtype], + scales_b: T.Tensor[(N, T.ceildiv(K, weight_group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as ( + bx, + by, + ): + A_shared = T.alloc_shared((block_M, block_K), FP8) + B_fp4_shared = T.alloc_shared((block_N, block_K), FP4) + B_shared = T.alloc_shared((block_N, block_K), FP8) + C_shared = T.alloc_shared((block_M, block_N), out_dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype) + scale_a_frag = T.alloc_fragment((block_M,), FP32) + scale_b_frag = T.alloc_fragment((block_N,), FP32) + + T.use_swizzle(panel_size=10) + T.clear(C_local) + T.clear(C_local_accum) + + K_iters = T.ceildiv(K, block_K) + for k in T.Pipelined(K_iters, num_stages=2): + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[bx * block_N, k * block_K], B_fp4_shared) + # FP4->FP8 cast must go through FP32 to avoid ambiguous C++ overload + for i, j in T.Parallel(block_N, block_K): + B_shared[i, j] = T.Cast(FP8, T.Cast(FP32, B_fp4_shared[i, j])) + + # Weight scale: per 32 on K, indexed by k (each k is one block_K=32) + for i in T.Parallel(block_N): + scale_b_frag[i] = T.Cast(FP32, scales_b[bx * block_N + i, k]) + + # Act scale: per 128 on K, indexed by k // 4 + for i in T.Parallel(block_M): + scale_a_frag[i] = T.Cast(FP32, scales_a[by * block_M + i, k // n_sub]) + + T.gemm(A_shared, B_shared, C_local, transpose_B=True) + + for i, j in T.Parallel(block_M, block_N): + C_local_accum[i, j] += C_local[i, j] * scale_a_frag[i] * scale_b_frag[j] + T.clear(C_local) + + T.copy(C_local_accum, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return fp4_gemm_kernel_ + + +def fp4_gemm( + a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor, + scale_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T. + A has per-128 act scale; B has per-32 E8M0 weight scale. + B is stored as [N, K//2] in float4_e2m1fn_x2 (2 FP4 values per byte, packed along K).""" + assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous" + assert a_s.is_contiguous() and b_s.is_contiguous(), ( + "Scaling factor tensors must be contiguous" + ) + tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32 + K = a.size(-1) + M = a.numel() // K + N = b.size(0) + c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype()) + kernel = fp4_gemm_kernel(N, K, scale_dtype=tl_dtype) + kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s) + return c diff --git a/TEMP/dsv4thing/model.py b/TEMP/dsv4thing/model.py new file mode 100644 index 00000000..167ade8f --- /dev/null +++ b/TEMP/dsv4thing/model.py @@ -0,0 +1,827 @@ +import math +from dataclasses import dataclass +from typing import Tuple, Optional, Literal +from functools import lru_cache +from contextlib import contextmanager + +import torch +from torch import nn +import torch.nn.functional as F +import torch.distributed as dist + +from kernel import act_quant, fp4_act_quant, fp8_gemm, fp4_gemm, sparse_attn, hc_split_sinkhorn + + +world_size = 1 +rank = 0 +block_size = 128 +fp4_block_size = 32 +default_dtype = torch.bfloat16 +scale_fmt = None +scale_dtype = torch.float32 + + +@contextmanager +def set_dtype(dtype): + """Temporarily override torch default dtype, restoring it on exit (even if an exception occurs).""" + prev = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(prev) + +@dataclass +class ModelArgs: + """Model hyperparameters. Field names match the config JSON keys.""" + max_batch_size: int = 4 + max_seq_len: int = 4096 + dtype: Literal["bf16", "fp8"] = "fp8" + scale_fmt: Literal[None, "ue8m0"] = "ue8m0" + expert_dtype: Literal[None, "fp4"] = None + scale_dtype: Literal["fp32", "fp8"] = "fp8" + vocab_size: int = 129280 + dim: int = 4096 + moe_inter_dim: int = 4096 + n_layers: int = 7 + n_hash_layers: int = 0 + n_mtp_layers: int = 1 + n_heads: int = 64 + # moe + n_routed_experts: int = 8 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + route_scale: float = 1. + swiglu_limit: float = 0. + # mqa + q_lora_rank: int = 1024 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-6 + o_groups: int = 8 + o_lora_rank: int = 1024 + window_size: int = 128 + compress_ratios: Tuple[int] = (0, 0, 4, 128, 4, 128, 4, 0) + # yarn + compress_rope_theta: float = 40000.0 + original_seq_len: int = 0 + rope_theta: float = 10000.0 + rope_factor: float = 40 + beta_fast: int = 32 + beta_slow: int = 1 + # index + index_n_heads: int = 64 + index_head_dim: int = 128 + index_topk: int = 512 + # hc + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + + +class ParallelEmbedding(nn.Module): + """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows. + Out-of-range indices are zero-masked before all_reduce to combine partial embeddings.""" + def __init__(self, vocab_size: int, dim: int): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + assert vocab_size % world_size == 0, f"Vocabulary size must be divisible by world size (world_size={world_size})" + self.part_vocab_size = (vocab_size // world_size) + self.vocab_start_idx = rank * self.part_vocab_size + self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if world_size > 1: + mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx) + x = x - self.vocab_start_idx + x[mask] = 0 + y = F.embedding(x, self.weight) + if world_size > 1: + y[mask] = 0 + dist.all_reduce(y) + return y + + +def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor: + """Dispatches to fp4_gemm / fp8_gemm / F.linear based on weight dtype. + For quantized weights, x is first quantized to FP8 via act_quant.""" + assert bias is None + + if weight.dtype == torch.float4_e2m1fn_x2: + x, s = act_quant(x, block_size, scale_fmt, scale_dtype) + return fp4_gemm(x, s, weight, weight.scale, scale_dtype) + elif weight.dtype == torch.float8_e4m3fn: + x, s = act_quant(x, block_size, scale_fmt, scale_dtype) + return fp8_gemm(x, s, weight, weight.scale, scale_dtype) + else: + return F.linear(x, weight) + + +class Linear(nn.Module): + """Linear layer supporting BF16, FP8, and FP4 weight formats with per-block scaling.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + super().__init__() + self.in_features = in_features + self.out_features = out_features + dtype = dtype or default_dtype + if dtype == torch.float4_e2m1fn_x2: + # FP4: weight is [out, in//2] in float4_e2m1fn_x2, logically [out, in] in fp4 + # Scale is [out, in//32] in float8_e8m0fnu (1 scale per 32 fp4 elements along K) + self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2)) + scale_out_features = out_features + scale_in_features = in_features // fp4_block_size + self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float8_e8m0fnu)) + elif dtype == torch.float8_e4m3fn: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + scale_out_features = (out_features + block_size - 1) // block_size + scale_in_features = (in_features + block_size - 1) // block_size + self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float8_e8m0fnu)) + else: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.register_parameter("scale", None) + if bias: + self.bias = nn.Parameter(torch.empty(out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class ColumnParallelLinear(Linear): + """Shards output dim across TP ranks. No all-reduce needed on output.""" + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})" + self.part_out_features = out_features // world_size + super().__init__(in_features, self.part_out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class RowParallelLinear(Linear): + """Shards input dim across TP ranks. All-reduce on output to sum partial results.""" + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})" + self.part_in_features = in_features // world_size + super().__init__(self.part_in_features, out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = linear(x, self.weight, None) + if world_size > 1: + y = y.float() + dist.all_reduce(y) + if self.bias is not None: + y += self.bias + return y.type_as(x) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + # rmsnorm in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for convenient. + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +@lru_cache(2) +def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor: + """Precomputes complex exponentials for rotary embeddings with YaRN scaling. + When original_seq_len > 0, applies frequency interpolation with a smooth + linear ramp between beta_fast and beta_slow correction ranges.""" + + def find_correction_dim(num_rotations, dim, base, max_seq_len): + return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): + low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) + return max(low, 0), min(high, dim-1) + + def linear_ramp_factor(min, max, dim): + if min == max: + max += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len) + smooth = 1 - linear_ramp_factor(low, high, dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Applies rotary positional embeddings in-place. Uses conjugate for inverse (de-rotation).""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + """Applies randomized Hadamard rotation to spread information across dims before FP8 quant.""" + assert x.dtype == torch.bfloat16 + from fast_hadamard_transform import hadamard_transform + return hadamard_transform(x, scale=x.size(-1) ** -0.5) + + +@lru_cache(1) +def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int): + if start_pos >= window_size - 1: + start_pos %= window_size + matrix = torch.cat([torch.arange(start_pos + 1, window_size), torch.arange(0, start_pos + 1)], dim=0) + elif start_pos > 0: + matrix = F.pad(torch.arange(start_pos + 1), (0, window_size - start_pos - 1), value=-1) + else: + base = torch.arange(seqlen).unsqueeze(1) + matrix = (base - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size)) + matrix = torch.where(matrix > base, -1, matrix) + return matrix.unsqueeze(0).expand(bsz, -1, -1) + + +@lru_cache(2) +def get_compress_topk_idxs(ratio: int, bsz: int, seqlen: int, start_pos: int, offset: int): + if start_pos > 0: + matrix = torch.arange(0, (start_pos + 1) // ratio) + offset + else: + matrix = torch.arange(seqlen // ratio).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix.unsqueeze(0).expand(bsz, -1, -1) + + +class Compressor(nn.Module): + """Compresses KV cache via learned gated pooling over `compress_ratio` consecutive tokens. + When overlap=True (ratio==4), uses overlapping windows for smoother compression boundaries.""" + + def __init__(self, args: ModelArgs, compress_ratio: int = 4, head_dim: int = 512, rotate: bool = False): + super().__init__() + self.dim = args.dim + self.head_dim = head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = head_dim - args.rope_head_dim + self.compress_ratio = compress_ratio + self.overlap = compress_ratio == 4 + self.rotate = rotate + coff = 1 + self.overlap + + self.ape = nn.Parameter(torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32)) + # wkv and wgate in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for convenient. + # When overlap, the first half of dims is for overlapping compression, second half for normal. + self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.norm = RMSNorm(self.head_dim, args.norm_eps) + self.kv_cache: torch.Tensor = None # assigned lazily from Attention.kv_cache + # State buffers for decode-phase incremental compression. + # With overlap: state[:, :ratio] = overlapping window, state[:, ratio:] = current window. + self.register_buffer("kv_state", torch.zeros(args.max_batch_size, coff * compress_ratio, coff * self.head_dim, dtype=torch.float32), persistent=False) + self.register_buffer("score_state", torch.full((args.max_batch_size, coff * compress_ratio, coff * self.head_dim), float("-inf"), dtype=torch.float32), persistent=False) + self.freqs_cis: torch.Tensor = None + + def overlap_transform(self, tensor: torch.Tensor, value=0): + # tensor: [b,s,r,2d] + b, s, _, _ = tensor.size() + ratio, d = self.compress_ratio, self.head_dim + new_tensor = tensor.new_full((b, s, 2 * ratio, d), value) + new_tensor[:, :, ratio:] = tensor[:, :, :, d:] + new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor, start_pos: int): + assert self.kv_cache is not None + bsz, seqlen, _ = x.size() + ratio, overlap, d, rd = self.compress_ratio, self.overlap, self.head_dim, self.rope_head_dim + dtype = x.dtype + # compression need fp32 + x = x.float() + kv = self.wkv(x) + score = self.wgate(x) + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + offset = ratio if overlap else 0 + if overlap and cutoff >= ratio: + self.kv_state[:bsz, :ratio] = kv[:, cutoff-ratio : cutoff] + self.score_state[:bsz, :ratio] = score[:, cutoff-ratio : cutoff] + self.ape + if remainder > 0: + kv, self.kv_state[:bsz, offset : offset+remainder] = kv.split([cutoff, remainder], dim=1) + self.score_state[:bsz, offset : offset+remainder] = score[:, cutoff:] + self.ape[:remainder] + score = score[:, :cutoff] + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + self.ape + if overlap: + kv = self.overlap_transform(kv, 0) + score = self.overlap_transform(score, float("-inf")) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: + should_compress = (start_pos + 1) % self.compress_ratio == 0 + score += self.ape[start_pos % ratio] + if overlap: + self.kv_state[:bsz, ratio + start_pos % ratio] = kv.squeeze(1) + self.score_state[:bsz, ratio + start_pos % ratio] = score.squeeze(1) + if should_compress: + kv_state = torch.cat([self.kv_state[:bsz, :ratio, :d], self.kv_state[:bsz, ratio:, d:]], dim=1) + score_state = torch.cat([self.score_state[:bsz, :ratio, :d], self.score_state[:bsz, ratio:, d:]], dim=1) + kv = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True) + self.kv_state[:bsz, :ratio] = self.kv_state[:bsz, ratio:] + self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:] + else: + self.kv_state[:bsz, start_pos % ratio] = kv.squeeze(1) + self.score_state[:bsz, start_pos % ratio] = score.squeeze(1) + if should_compress: + kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True) + if not should_compress: + return + kv = self.norm(kv.to(dtype)) + if start_pos == 0: + freqs_cis = self.freqs_cis[:cutoff:ratio] + else: + freqs_cis = self.freqs_cis[start_pos + 1 - self.compress_ratio].unsqueeze(0) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + if self.rotate: + kv = rotate_activation(kv) + fp4_act_quant(kv, fp4_block_size, True) + else: + act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True) + if start_pos == 0: + self.kv_cache[:bsz, :seqlen // ratio] = kv + else: + self.kv_cache[:bsz, start_pos // ratio] = kv.squeeze(1) + return kv + + +class Indexer(torch.nn.Module): + """Selects top-k compressed KV positions for sparse attention via learned scoring. + Has its own Compressor (with Hadamard rotation) to build compressed KV for scoring.""" + + def __init__(self, args: ModelArgs, compress_ratio: int = 4): + super().__init__() + self.dim = args.dim + self.n_heads = args.index_n_heads + self.n_local_heads = args.index_n_heads // world_size + self.head_dim = args.index_head_dim + self.rope_head_dim = args.rope_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16) + self.softmax_scale = self.head_dim ** -0.5 + self.compress_ratio = compress_ratio + + self.compressor = Compressor(args, compress_ratio, self.head_dim, True) + self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len // compress_ratio, self.head_dim), persistent=False) + self.freqs_cis = None + + def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] + ratio = self.compress_ratio + rd = self.rope_head_dim + end_pos = start_pos + seqlen + if self.compressor.kv_cache is None: + self.compressor.kv_cache = self.kv_cache + self.compressor.freqs_cis = self.freqs_cis + q = self.wq_b(qr) + q = q.unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + q = rotate_activation(q) + # use fp4 simulation for q and kv in indexer + fp4_act_quant(q, fp4_block_size, True) + self.compressor(x, start_pos) + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) + # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16 + index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio]) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + if world_size > 1: + dist.all_reduce(index_score) + if start_pos == 0: + mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + index_score += torch.where(mask, float("-inf"), 0) + topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1] + if start_pos == 0: + mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + topk_idxs = torch.where(mask, -1, topk_idxs + offset) + else: + topk_idxs += offset + return topk_idxs + + +class Attention(nn.Module): + """Multi-head Latent Attention (MLA) with sliding window + optional KV compression. + Uses low-rank Q projection (wq_a -> q_norm -> wq_b) and grouped low-rank O projection.""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.dim = args.dim + self.n_heads = args.n_heads + self.n_local_heads = args.n_heads // world_size + self.q_lora_rank = args.q_lora_rank + self.o_lora_rank = args.o_lora_rank + self.head_dim = args.head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = args.head_dim - args.rope_head_dim + self.n_groups = args.o_groups + self.n_local_groups = self.n_groups // world_size + self.window_size = args.window_size + self.compress_ratio = args.compress_ratios[layer_id] + self.eps = args.norm_eps + + self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32)) + self.wq_a = Linear(self.dim, self.q_lora_rank) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.wkv = Linear(self.dim, self.head_dim) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear(self.n_heads * self.head_dim // self.n_groups, self.n_groups * args.o_lora_rank, dtype=torch.bfloat16) + self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim) + self.softmax_scale = self.head_dim ** -0.5 + + if self.compress_ratio: + self.compressor = Compressor(args, self.compress_ratio, self.head_dim) + if self.compress_ratio == 4: + self.indexer = Indexer(args, self.compress_ratio) + else: + self.indexer = None + + kv_cache_size = args.window_size + (args.max_seq_len // self.compress_ratio if self.compress_ratio else 0) + self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, kv_cache_size, self.head_dim), persistent=False) + if self.compress_ratio: + original_seq_len, rope_theta = args.original_seq_len, args.compress_rope_theta + else: + # disable YaRN and use base rope_theta in pure sliding-window attention + original_seq_len, rope_theta = 0, args.rope_theta + freqs_cis = precompute_freqs_cis(self.rope_head_dim, args.max_seq_len, original_seq_len, + rope_theta, args.rope_factor, args.beta_fast, args.beta_slow) + self.register_buffer("freqs_cis", freqs_cis, persistent=False) + + def forward(self, x: torch.Tensor, start_pos: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] + win = self.window_size + ratio = self.compress_ratio + rd = self.rope_head_dim + if self.compress_ratio and self.compressor.kv_cache is None: + self.compressor.kv_cache = self.kv_cache[:, win:] + self.compressor.freqs_cis = self.freqs_cis + if self.indexer is not None: + self.indexer.freqs_cis = self.freqs_cis + # q + qr = q = self.q_norm(self.wq_a(x)) + q = self.wq_b(q).unflatten(-1, (self.n_local_heads, self.head_dim)) + q *= torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + # win kv & topk_idxs + kv = self.wkv(x) + kv = self.kv_norm(kv) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + # FP8-simulate non-rope dims to match QAT; rope dims stay bf16 for positional precision + act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True) + topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos) + if self.compress_ratio: + offset = kv.size(1) if start_pos == 0 else win + if self.indexer is not None: + compress_topk_idxs = self.indexer(x, qr, start_pos, offset) + else: + compress_topk_idxs = get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset) + topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1) + topk_idxs = topk_idxs.int() + + # compress kv & attn + if start_pos == 0: + if seqlen <= win: + self.kv_cache[:bsz, :seqlen] = kv + else: + cutoff = seqlen % win + self.kv_cache[:bsz, cutoff: win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split([win - cutoff, cutoff], dim=1) + if self.compress_ratio: + if (kv_compress := self.compressor(x, start_pos)) is not None: + kv = torch.cat([kv, kv_compress], dim=1) + # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16 + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + else: + self.kv_cache[:bsz, start_pos % win] = kv.squeeze(1) + if self.compress_ratio: + self.compressor(x, start_pos) + o = sparse_attn(q, self.kv_cache[:bsz], self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + # o + o = o.view(bsz, seqlen, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + # NOTE: wo_a is FP8 in checkpoint; could do FP8 einsum here for better perf, + # but using BF16 for simplicity. + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + +class Gate(nn.Module): + """MoE gating: computes expert routing scores and selects top-k experts. + Supports hash-based routing (first n_hash_layers) where expert indices are + predetermined per token ID, and score-based routing (remaining layers).""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.dim = args.dim + self.topk = args.n_activated_experts + self.score_func = args.score_func + self.route_scale = args.route_scale + self.hash = layer_id < args.n_hash_layers + self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim)) + if self.hash: + self.tid2eid = nn.Parameter(torch.empty(args.vocab_size, args.n_activated_experts, dtype=torch.int32), requires_grad=False) + self.bias = None + else: + self.bias = nn.Parameter(torch.empty(args.n_routed_experts, dtype=torch.float32)) + + def forward(self, x: torch.Tensor, input_ids: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: + scores = linear(x.float(), self.weight.float()) + if self.score_func == "softmax": + scores = scores.softmax(dim=-1) + elif self.score_func == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + original_scores = scores + # Bias shifts scores for expert selection (topk) but does not affect routing weights. + if self.bias is not None: + scores = scores + self.bias + if self.hash: + indices = self.tid2eid[input_ids] + else: + indices = scores.topk(self.topk, dim=-1)[1] + weights = original_scores.gather(1, indices) + if self.score_func != "softmax": + weights /= weights.sum(dim=-1, keepdim=True) + weights *= self.route_scale + return weights, indices + + +class Expert(nn.Module): + """Single MoE expert: SwiGLU FFN (w1, w2, w3). Computation in float32 for stability.""" + def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0): + super().__init__() + self.w1 = Linear(dim, inter_dim, dtype=dtype) + self.w2 = Linear(inter_dim, dim, dtype=dtype) + self.w3 = Linear(dim, inter_dim, dtype=dtype) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor] = None) -> torch.Tensor: + dtype = x.dtype + gate = self.w1(x).float() + up = self.w3(x).float() + if self.swiglu_limit > 0: + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + gate = torch.clamp(gate, max=self.swiglu_limit) + x = F.silu(gate) * up + if weights is not None: + x = weights * x + return self.w2(x.to(dtype)) + + +class MoE(nn.Module): + """Mixture-of-Experts: gate routes each token to top-k routed experts + 1 shared expert. + Experts are sharded across TP ranks; each rank handles n_routed_experts // world_size experts.""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.dim = args.dim + assert args.n_routed_experts % world_size == 0, f"Number of experts must be divisible by world size (world_size={world_size})" + self.n_routed_experts = args.n_routed_experts + self.n_local_experts = args.n_routed_experts // world_size + self.n_activated_experts = args.n_activated_experts + self.experts_start_idx = rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.gate = Gate(layer_id, args) + expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None + self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim, dtype=expert_dtype, swiglu_limit=args.swiglu_limit) if self.experts_start_idx <= i < self.experts_end_idx else None + for i in range(self.n_routed_experts)]) + assert args.n_shared_experts == 1 + self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit) + + def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + shape = x.size() + x = x.view(-1, self.dim) + weights, indices = self.gate(x, input_ids.flatten()) + y = torch.zeros_like(x, dtype=torch.float32) + counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() + for i in range(self.experts_start_idx, self.experts_end_idx): + if counts[i] == 0: + continue + expert = self.experts[i] + idx, top = torch.where(indices == i) + y[idx] += expert(x[idx], weights[idx, top, None]) + if world_size > 1: + dist.all_reduce(y) + y += self.shared_experts(x) + return y.type_as(x).view(shape) + + +class Block(nn.Module): + """Transformer block with Hyper-Connections (HC) mixing. + Instead of a simple residual, HC maintains `hc_mult` copies of the hidden state. + hc_pre: reduces hc copies -> 1 via learned weighted sum (pre-weights from Sinkhorn). + hc_post: expands 1 -> hc copies via learned post-weights + combination matrix.""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.norm_eps = args.norm_eps + self.attn = Attention(layer_id, args) + self.ffn = MoE(layer_id, args) + self.attn_norm = RMSNorm(args.dim, self.norm_eps) + self.ffn_norm = RMSNorm(args.dim, self.norm_eps) + self.hc_mult = hc_mult = args.hc_mult + self.hc_sinkhorn_iters = args.hc_sinkhorn_iters + self.hc_eps = args.hc_eps + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_attn_scale = nn.Parameter(torch.empty(3)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3)) + + def hc_pre(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + # x: [b,s,hc,d], hc_fn: [mix_hc,hc*d], hc_scale: [3], hc_base: [mix_hc], y: [b,s,hc,d] + shape, dtype = x.size(), x.dtype + x = x.flatten(2).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + pre, post, comb = hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps) + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2) + return y.to(dtype), post, comb + + def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor): + # x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc], y: [b,s,hc,d] + y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) + return y.type_as(x) + + def forward(self, x: torch.Tensor, start_pos: int, input_ids: Optional[torch.Tensor]) -> torch.Tensor: + residual = x + x, post, comb = self.hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x = self.attn_norm(x) + x = self.attn(x, start_pos) + x = self.hc_post(x, residual, post, comb) + + residual = x + x, post, comb = self.hc_pre(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + x = self.hc_post(x, residual, post, comb) + return x + + +class ParallelHead(nn.Module): + + def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + self.norm_eps = norm_eps + self.hc_eps = hc_eps + self.part_vocab_size = (vocab_size // world_size) + # lm_head in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for easier computation of logits later. + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32)) + + def get_logits(self, x): + return F.linear(x[:, -1].float(), self.weight) + + def forward(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, norm: RMSNorm): + # x: [b,s,hc,d] + x = self.hc_head(x, hc_fn, hc_scale, hc_base) + logits = self.get_logits(norm(x)) + if world_size > 1: + all_logits = [torch.empty_like(logits) for _ in range(world_size)] + dist.all_gather(all_logits, logits) + logits = torch.cat(all_logits, dim=-1) + return logits + + def hc_head(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + shape, dtype = x.size(), x.dtype + x = x.flatten(2).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + pre = torch.sigmoid(mixes * hc_scale + hc_base) + self.hc_eps + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2) + return y.to(dtype) + + +class MTPBlock(Block): + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__(layer_id, args) + self.e_proj = Linear(args.dim, args.dim) + self.h_proj = Linear(args.dim, args.dim) + self.enorm = RMSNorm(args.dim, args.norm_eps) + self.hnorm = RMSNorm(args.dim, args.norm_eps) + self.norm = RMSNorm(args.dim, args.norm_eps) + self.hc_mult = hc_mult = args.hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim)) + self.hc_head_base = nn.Parameter(torch.empty(hc_mult)) + self.hc_head_scale = nn.Parameter(torch.empty(1)) + self.embed: ParallelEmbedding = None + self.head: ParallelHead = None + + @torch.inference_mode() + def forward(self, x: torch.Tensor, start_pos: int, input_ids: torch.Tensor) -> torch.Tensor: + # x: [b,s,hc,d] + assert self.embed is not None and self.head is not None + e = self.embed(input_ids) + e = self.enorm(e) + x = self.hnorm(x) + x = self.e_proj(e).unsqueeze(2) + self.h_proj(x) + x = super().forward(x, start_pos, input_ids) + logits = self.head(x, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm) + return logits + + +class Transformer(nn.Module): + """Full DeepSeek-V4 model: embed -> HC-expand -> N blocks -> HC-head -> logits. + Sets global state (world_size, rank, default_dtype, scale_fmt, scale_dtype) in __init__.""" + def __init__(self, args: ModelArgs): + global world_size, rank, default_dtype, scale_fmt, scale_dtype + world_size = dist.get_world_size() if dist.is_initialized() else 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16 + scale_fmt = "ue8m0" if args.scale_dtype == "fp8" else args.scale_fmt + scale_dtype = torch.float8_e8m0fnu if args.scale_dtype == "fp8" else torch.float32 + super().__init__() + self.max_seq_len = args.max_seq_len + self.norm_eps = args.norm_eps + self.hc_eps = args.hc_eps + self.embed = ParallelEmbedding(args.vocab_size, args.dim) + self.layers = torch.nn.ModuleList() + for layer_id in range(args.n_layers): + self.layers.append(Block(layer_id, args)) + self.norm = RMSNorm(args.dim, self.norm_eps) + self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps) + self.mtp = torch.nn.ModuleList() + for layer_id in range(args.n_mtp_layers): + self.mtp.append(MTPBlock(args.n_layers + layer_id, args)) + self.mtp[-1].embed = self.embed + self.mtp[-1].head = self.head + self.hc_mult = hc_mult = args.hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim)) + self.hc_head_base = nn.Parameter(torch.empty(hc_mult)) + self.hc_head_scale = nn.Parameter(torch.empty(1)) + + @torch.inference_mode() + def forward(self, input_ids: torch.Tensor, start_pos: int = 0): + h = self.embed(input_ids) + # Expand to hc_mult copies for Hyper-Connections + h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + for layer in self.layers: + h = layer(h, start_pos, input_ids) + logits = self.head(h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm) + return logits + + +if __name__ == "__main__": + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.manual_seed(0) + args = ModelArgs(n_hash_layers=0) + x = torch.randint(0, args.vocab_size, (2, 128)) + model = Transformer(args) + + print(model(x).size()) + for i in range(128, 150): + print(i, model(x[:, 0:1], i).size()) + + h = torch.randn(2, 128, args.hc_mult, args.dim) + mtp = model.mtp[0] + print(mtp(h, 0, x).size()) + print(mtp(h[:, 0:1], 1, x[:, 0:1]).size()) diff --git a/TEMP/dsv4thing/requirements.txt b/TEMP/dsv4thing/requirements.txt new file mode 100644 index 00000000..7e1cc78c --- /dev/null +++ b/TEMP/dsv4thing/requirements.txt @@ -0,0 +1,5 @@ +torch>=2.10.0 +transformers>=5.0.0 +safetensors>=0.7.0 +fast_hadamard_transform +tilelang==0.1.8 \ No newline at end of file diff --git a/single_shot_inference.py b/single_shot_inference.py index 7f445d19..369b1d50 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -47,8 +47,16 @@ PROMPT = _args.prompt or "The capital of France is" NUM_GPUS = _args.num_gpus SEED = _args.seed VERBOSE = _args.verbose -THINK_START, THINK_END = 128821, 128822 -USER_TOKEN, ASSISTANT_TOKEN = 128803, 128804 +# Special token IDs — derived from official encoding module strings + tokenizer. +# Do NOT hardcode these; the encoding module defines the canonical token strings. +from encoding.deepseek_v4_encoding import ( + thinking_start_token as _THINK_START_STR, + thinking_end_token as _THINK_END_STR, + USER_SP_TOKEN as _USER_STR, + ASSISTANT_SP_TOKEN as _ASSISTANT_STR, + eos_token as _EOS_STR, + bos_token as _BOS_STR, +) FP4_LUT = torch.tensor([0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) # ===================================================================== @@ -1461,6 +1469,14 @@ def main(): from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR) + # Derive special token IDs from official encoding strings + tokenizer. + # This is the ONLY source of truth — never hardcode these IDs. + THINK_START = tokenizer.convert_tokens_to_ids(_THINK_START_STR) + THINK_END = tokenizer.convert_tokens_to_ids(_THINK_END_STR) + USER_TOKEN = tokenizer.convert_tokens_to_ids(_USER_STR) + ASSISTANT_TOKEN = tokenizer.convert_tokens_to_ids(_ASSISTANT_STR) + bos = tokenizer.bos_token_id or 0 + # A1: Build explicit stop set — DSV4 uses special turn-end tokens beyond eos STOP_IDS = set() eos_id = tokenizer.eos_token_id @@ -1476,7 +1492,6 @@ def main(): print(f" Special tokens: {tokenizer.special_tokens_map}") print(f" THINK_START={THINK_START} THINK_END={THINK_END} USER={USER_TOKEN} ASST={ASSISTANT_TOKEN}") - bos = tokenizer.bos_token_id or 0 if _args.prefill_tokens: generated = [int(x) for x in _args.prefill_tokens.split(',')] else: @@ -1749,12 +1764,35 @@ def main(): print(f" L{li} {tag}: {dt_ms:.2f}ms") prev_t = t - out = tokenizer.decode(all_tokens, skip_special_tokens=True) - print(f"\n{'='*70}") - print(f"Input: '{PROMPT}'") - print(f"Output: '{out}'") - print(f"Total: {time.time()-t0:.1f}s") - print(f"{'='*70}") + out_raw = tokenizer.decode(all_tokens, skip_special_tokens=False) + # Use official DSV4 parser for structured output + try: + from encoding.deepseek_v4_encoding import parse_message_from_completion_text + # Find the assistant portion — after the last ASSISTANT token + assistant_start = out_raw.find(_ASSISTANT_STR) + if assistant_start >= 0: + assistant_text = out_raw[assistant_start + len(_ASSISTANT_STR):] + else: + assistant_text = out_raw + parsed = parse_message_from_completion_text(assistant_text, thinking_mode=_args.thinking_mode) + reasoning = parsed.get('reasoning', '') + content = parsed.get('content', '') + print(f"\n{'='*70}") + print(f"Input: '{PROMPT}'") + if reasoning: + print(f"Reasoning: {reasoning[:500]}{'...' if len(reasoning) > 500 else ''}") + print(f"Content: {content}") + print(f"Total: {time.time()-t0:.1f}s") + print(f"{'='*70}") + except Exception as e: + # Fallback: raw decode (shouldn't happen with correct output) + out = tokenizer.decode(all_tokens, skip_special_tokens=True) + print(f"\n{'='*70}") + print(f"Input: '{PROMPT}'") + print(f"Output (raw): '{out}'") + print(f"Parse error: {e}") + print(f"Total: {time.time()-t0:.1f}s") + print(f"{'='*70}") if __name__ == "__main__": main()