DSV4-Flash-Vision EXL3 MixedK D2-K2x3-Dense6 + lna-lab recipe (2026-09-03)
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +2 -0
- README.md +79 -0
- config.json +1178 -0
- generation_config.json +9 -0
- lna-lab/Dockerfile +14 -0
- lna-lab/PLAN.md +74 -0
- lna-lab/README-dsv4-carve.md +65 -0
- lna-lab/UPSTREAM-vcruz305-README.md +51 -0
- lna-lab/bench-dsv4.py +37 -0
- lna-lab/bench-streams.py +29 -0
- lna-lab/exllamav3_ext.cpython-312-x86_64-linux-gnu.so +3 -0
- lna-lab/patches/exl3.py.lna +389 -0
- lna-lab/patches/exl3_gemm.cu.lna +703 -0
- lna-lab/patches/patch_dsv4_attention_compressor_exl3.py +62 -0
- lna-lab/patches/patch_dsv4_aux_streams_env.py +19 -0
- lna-lab/patches/patch_dsv4_dense_exl3.py +238 -0
- lna-lab/patches/vllm_exl3_exl3.py.lna +1528 -0
- lna-lab/ppl-vllm.py +23 -0
- lna-lab/serve-dsv4-tp8.sh +20 -0
- lna-lab/tools/dense_bake.py +924 -0
- lna-lab/tools/layer_overlay.py +957 -0
- lna-lab/tools/mtp_overlay.py +1032 -0
- model-00001-of-00048.safetensors +3 -0
- model-00002-of-00048.safetensors +3 -0
- model-00003-of-00048.safetensors +3 -0
- model-00004-of-00048.safetensors +3 -0
- model-00005-of-00048.safetensors +3 -0
- model-00006-of-00048.safetensors +3 -0
- model-00007-of-00048.safetensors +3 -0
- model-00008-of-00048.safetensors +3 -0
- model-00009-of-00048.safetensors +3 -0
- model-00010-of-00048.safetensors +3 -0
- model-00011-of-00048.safetensors +3 -0
- model-00012-of-00048.safetensors +3 -0
- model-00013-of-00048.safetensors +3 -0
- model-00014-of-00048.safetensors +3 -0
- model-00015-of-00048.safetensors +3 -0
- model-00016-of-00048.safetensors +3 -0
- model-00017-of-00048.safetensors +3 -0
- model-00018-of-00048.safetensors +3 -0
- model-00019-of-00048.safetensors +3 -0
- model-00020-of-00048.safetensors +3 -0
- model-00021-of-00048.safetensors +3 -0
- model-00022-of-00048.safetensors +3 -0
- model-00023-of-00048.safetensors +3 -0
- model-00024-of-00048.safetensors +3 -0
- model-00025-of-00048.safetensors +3 -0
- model-00026-of-00048.safetensors +3 -0
- model-00027-of-00048.safetensors +3 -0
- model-00028-of-00048.safetensors +3 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
lna-lab/exllamav3_ext.cpython-312-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: other
|
| 3 |
+
license_name: deepseek
|
| 4 |
+
license_link: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp
|
| 5 |
+
base_model: deepseek-ai/DeepSeek-V4-Flash-Vision-Exp
|
| 6 |
+
base_model_relation: quantized
|
| 7 |
+
library_name: exllamav3
|
| 8 |
+
pipeline_tag: image-text-to-text
|
| 9 |
+
tags:
|
| 10 |
+
- exl3
|
| 11 |
+
- exllamav3
|
| 12 |
+
- deepseek_v4
|
| 13 |
+
- moe
|
| 14 |
+
- mixed-precision
|
| 15 |
+
- vllm
|
| 16 |
+
- blackwell
|
| 17 |
+
- lna-lab
|
| 18 |
+
language:
|
| 19 |
+
- en
|
| 20 |
+
- ja
|
| 21 |
+
- zh
|
| 22 |
+
---
|
| 23 |
+
|
| 24 |
+
# DSV4-Flash-Vision — EXL3 MixedK · D2 · K2×3 · Dense6 (380K context on 8 × 16 GB)
|
| 25 |
+
|
| 26 |
+
A serving-tuned derivative of [vcruz305/DSV4-Flash-Vision-EXL3-MixedK](https://huggingface.co/vcruz305/DSV4-Flash-Vision-EXL3-MixedK)
|
| 27 |
+
(DeepSeek-V4-Flash-Vision-Exp, 305B MoE, vision). Everything that was still BF16 or fp8 in the original has been carved into EXL3
|
| 28 |
+
so that a **380K-context, 4-stream, DSpark3-speculative** seat fits on eight RTX PRO 2000 16 GB under vLLM.
|
| 29 |
+
Baked and measured by **Lna-Lab** (YUKI) on 2026-09-03; house model of SAZANAMI ("オオタニ").
|
| 30 |
+
|
| 31 |
+
Recipe, patches and the serve script: **github.com/lna-lab/dsv4-carve** (mirrored in `lna-lab/` here).
|
| 32 |
+
|
| 33 |
+
## Composition
|
| 34 |
+
|
| 35 |
+
| Part | Here | Original MixedK |
|
| 36 |
+
|---|---|---|
|
| 37 |
+
| Routed experts | EXL3 2/3-bit mixed; layers 3, 21, 41 → K2 (wrldsuksgo2mars K2-v1) | EXL3 2/3-bit |
|
| 38 |
+
| MTP draft experts (3 layers, DSpark) | EXL3 K2 (wrldsuksgo2mars K2.2-D2-v1) | fp8 |
|
| 39 |
+
| Attention + shared experts (790 linears) | **EXL3 6-bit**, calibrated (250 × 2048) | BF16 |
|
| 40 |
+
| lm_head / embeddings / indexer weights_proj | BF16 (unchanged) | BF16 |
|
| 41 |
+
|
| 42 |
+
86 GB on disk, 51 shards. Per GPU at TP8: 12.0 GiB weights → **396,656 tokens** of fp8 KV at `--max-model-len 389120`.
|
| 43 |
+
|
| 44 |
+
## Measured (vLLM 0.28.1rc1.dev337 + vllm-exl3 0.2.3, TP8, CUDA graphs, RTX PRO 2000 × 8)
|
| 45 |
+
|
| 46 |
+
| | |
|
| 47 |
+
|---|---|
|
| 48 |
+
| ppl wikitext-2 (512 × 16) | **6.7159** — original MixedK 6.6271 (+1.34 %); 4-bit dense was +3.8 % and rejected |
|
| 49 |
+
| single stream, no speculation | en 42.1 / ja 42.0 / code 41.9 tok/s |
|
| 50 |
+
| single stream, DSpark3 (3 MTP tokens) | en 62 / ja 55 / code 86 tok/s |
|
| 51 |
+
| 4 streams, DSpark3, aggregate | en 136 / ja 116 / code 186 tok/s |
|
| 52 |
+
| needle at 166K tokens | found; TTFT 167 s (prefill ≈ 1.0k tok/s), decode at that depth 37 tok/s |
|
| 53 |
+
| vision | works (shapes, colours, text read correctly) |
|
| 54 |
+
|
| 55 |
+
## You need the patches
|
| 56 |
+
|
| 57 |
+
Stock exllamav3 1.4.5 / vllm-exl3 0.2.3 will **not** serve this pack: dense EXL3 linears in attention trigger (a) the cooperative-GEMM
|
| 58 |
+
autotuner inside CUDA-graph capture and (b) a two-stream deadlock between vLLM's shared-experts stream and the routed EXL3 kernels
|
| 59 |
+
(both use exllamav3's per-device lock buffer). The `lna-lab/` folder carries the patched sources, the vLLM patches, the prebuilt
|
| 60 |
+
extension for sm_120 (`exllamav3_ext.cpython-312-x86_64-linux-gnu.so`, torch cu13) and the serve script; the README there explains each.
|
| 61 |
+
Minimum: `VLLM_DISABLE_SHARED_EXPERTS_STREAM=1`, `VLLM_SPARSE_INDEXER_MAX_LOGITS_MB=128`, aux streams off, `--kv-cache-dtype fp8`.
|
| 62 |
+
|
| 63 |
+
```bash
|
| 64 |
+
AUX_STREAMS=0 MODEL=. UTIL=0.97 MAXLEN=389120 BT=512 SEQS=4 \
|
| 65 |
+
SPEC='{"method":"dspark","num_speculative_tokens":3}' bash lna-lab/serve-dsv4-tp8.sh
|
| 66 |
+
```
|
| 67 |
+
|
| 68 |
+
Thinking is on by default; pass `chat_template_kwargs: {"thinking": false}` or a large `max_tokens`.
|
| 69 |
+
|
| 70 |
+
## Files
|
| 71 |
+
|
| 72 |
+
- `model-*.safetensors`, `model.safetensors.index.json`, `config.json` (`quantization_config.non_routed_exl3` lists the 6-bit dense modules), tokenizer, `generation_config.json`
|
| 73 |
+
- `lna-lab/` — Dockerfile, serve/bench/ppl scripts, `tools/` (overlay + dense bake), `patches/`, prebuilt extension, upstream model card
|
| 74 |
+
|
| 75 |
+
## Credits and license
|
| 76 |
+
|
| 77 |
+
- Base weights: DeepSeek (MIT). This pack follows the DeepSeek model license like the upstream MixedK.
|
| 78 |
+
- MixedK, vllm-exl3 and its DSV4 patches: [vcruz305](https://huggingface.co/vcruz305) · K2 packs: [wrldsuksgo2mars](https://huggingface.co/wrldsuksgo2mars) · ExLlamaV3: turboderp · vLLM
|
| 79 |
+
- Bake, patches, measurements: Lna-Lab / YUKI with Ken ([@Tono_Ken3](https://x.com/Tono_Ken3))
|
config.json
ADDED
|
@@ -0,0 +1,1178 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"DeepseekV4ForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"attention_bias": false,
|
| 6 |
+
"attention_dropout": 0.0,
|
| 7 |
+
"bos_token_id": 0,
|
| 8 |
+
"eos_token_id": 1,
|
| 9 |
+
"expert_dtype": "fp4",
|
| 10 |
+
"hc_eps": 1e-06,
|
| 11 |
+
"hc_mult": 4,
|
| 12 |
+
"hc_sinkhorn_iters": 20,
|
| 13 |
+
"head_dim": 512,
|
| 14 |
+
"hidden_act": "silu",
|
| 15 |
+
"hidden_size": 4096,
|
| 16 |
+
"index_head_dim": 128,
|
| 17 |
+
"index_n_heads": 64,
|
| 18 |
+
"index_topk": 512,
|
| 19 |
+
"initializer_range": 0.02,
|
| 20 |
+
"max_position_embeddings": 1048576,
|
| 21 |
+
"model_type": "deepseek_v4",
|
| 22 |
+
"moe_intermediate_size": 2048,
|
| 23 |
+
"n_routed_experts": 256,
|
| 24 |
+
"n_shared_experts": 1,
|
| 25 |
+
"norm_topk_prob": true,
|
| 26 |
+
"num_attention_heads": 64,
|
| 27 |
+
"num_experts_per_tok": 6,
|
| 28 |
+
"num_hidden_layers": 43,
|
| 29 |
+
"num_hash_layers": 3,
|
| 30 |
+
"num_key_value_heads": 1,
|
| 31 |
+
"num_nextn_predict_layers": 3,
|
| 32 |
+
"o_groups": 8,
|
| 33 |
+
"o_lora_rank": 1024,
|
| 34 |
+
"q_lora_rank": 1024,
|
| 35 |
+
"qk_rope_head_dim": 64,
|
| 36 |
+
"quantization_config": {
|
| 37 |
+
"quant_method": "exl3",
|
| 38 |
+
"bits": 2,
|
| 39 |
+
"codebook": "mcg",
|
| 40 |
+
"head_bits": 16,
|
| 41 |
+
"non_routed_dtype_policy": "bf16_as_stored",
|
| 42 |
+
"scope": "dsv4_routed_experts_only",
|
| 43 |
+
"serving_reader_qualified": false,
|
| 44 |
+
"version": "0.0.43",
|
| 45 |
+
"layer_bits": {
|
| 46 |
+
"13": 3,
|
| 47 |
+
"22": 3,
|
| 48 |
+
"28": 3,
|
| 49 |
+
"43": 2,
|
| 50 |
+
"44": 2,
|
| 51 |
+
"45": 2
|
| 52 |
+
},
|
| 53 |
+
"non_routed_quantization": {
|
| 54 |
+
"activation_scheme": "dynamic",
|
| 55 |
+
"fmt": "e4m3",
|
| 56 |
+
"quant_method": "fp8",
|
| 57 |
+
"scale_fmt": "ue8m0",
|
| 58 |
+
"weight_block_size": [
|
| 59 |
+
128,
|
| 60 |
+
128
|
| 61 |
+
]
|
| 62 |
+
},
|
| 63 |
+
"non_routed_exl3": {
|
| 64 |
+
"codebook": "mcg",
|
| 65 |
+
"layers": {
|
| 66 |
+
"language_model.model.layers.0.attn.fused_wqa_wkv": {
|
| 67 |
+
"bits": 6
|
| 68 |
+
},
|
| 69 |
+
"language_model.model.layers.0.attn.wo_a": {
|
| 70 |
+
"bits": 6
|
| 71 |
+
},
|
| 72 |
+
"language_model.model.layers.0.attn.wo_b": {
|
| 73 |
+
"bits": 6
|
| 74 |
+
},
|
| 75 |
+
"language_model.model.layers.0.attn.wq_b": {
|
| 76 |
+
"bits": 6
|
| 77 |
+
},
|
| 78 |
+
"language_model.model.layers.0.ffn.shared_experts.down_proj": {
|
| 79 |
+
"bits": 6
|
| 80 |
+
},
|
| 81 |
+
"language_model.model.layers.0.ffn.shared_experts.gate_up_proj": {
|
| 82 |
+
"bits": 6
|
| 83 |
+
},
|
| 84 |
+
"language_model.model.layers.1.attn.fused_wqa_wkv": {
|
| 85 |
+
"bits": 6
|
| 86 |
+
},
|
| 87 |
+
"language_model.model.layers.1.attn.wo_a": {
|
| 88 |
+
"bits": 6
|
| 89 |
+
},
|
| 90 |
+
"language_model.model.layers.1.attn.wo_b": {
|
| 91 |
+
"bits": 6
|
| 92 |
+
},
|
| 93 |
+
"language_model.model.layers.1.attn.wq_b": {
|
| 94 |
+
"bits": 6
|
| 95 |
+
},
|
| 96 |
+
"language_model.model.layers.1.ffn.shared_experts.down_proj": {
|
| 97 |
+
"bits": 6
|
| 98 |
+
},
|
| 99 |
+
"language_model.model.layers.1.ffn.shared_experts.gate_up_proj": {
|
| 100 |
+
"bits": 6
|
| 101 |
+
},
|
| 102 |
+
"language_model.model.layers.10.attn.compressor.fused_wkv_wgate": {
|
| 103 |
+
"bits": 6
|
| 104 |
+
},
|
| 105 |
+
"language_model.model.layers.10.attn.fused_wqa_wkv": {
|
| 106 |
+
"bits": 6
|
| 107 |
+
},
|
| 108 |
+
"language_model.model.layers.10.attn.indexer.compressor.fused_wkv_wgate": {
|
| 109 |
+
"bits": 6
|
| 110 |
+
},
|
| 111 |
+
"language_model.model.layers.10.attn.indexer.wq_b": {
|
| 112 |
+
"bits": 6
|
| 113 |
+
},
|
| 114 |
+
"language_model.model.layers.10.attn.wo_a": {
|
| 115 |
+
"bits": 6
|
| 116 |
+
},
|
| 117 |
+
"language_model.model.layers.10.attn.wo_b": {
|
| 118 |
+
"bits": 6
|
| 119 |
+
},
|
| 120 |
+
"language_model.model.layers.10.attn.wq_b": {
|
| 121 |
+
"bits": 6
|
| 122 |
+
},
|
| 123 |
+
"language_model.model.layers.10.ffn.shared_experts.down_proj": {
|
| 124 |
+
"bits": 6
|
| 125 |
+
},
|
| 126 |
+
"language_model.model.layers.10.ffn.shared_experts.gate_up_proj": {
|
| 127 |
+
"bits": 6
|
| 128 |
+
},
|
| 129 |
+
"language_model.model.layers.11.attn.compressor.fused_wkv_wgate": {
|
| 130 |
+
"bits": 6
|
| 131 |
+
},
|
| 132 |
+
"language_model.model.layers.11.attn.fused_wqa_wkv": {
|
| 133 |
+
"bits": 6
|
| 134 |
+
},
|
| 135 |
+
"language_model.model.layers.11.attn.wo_a": {
|
| 136 |
+
"bits": 6
|
| 137 |
+
},
|
| 138 |
+
"language_model.model.layers.11.attn.wo_b": {
|
| 139 |
+
"bits": 6
|
| 140 |
+
},
|
| 141 |
+
"language_model.model.layers.11.attn.wq_b": {
|
| 142 |
+
"bits": 6
|
| 143 |
+
},
|
| 144 |
+
"language_model.model.layers.11.ffn.shared_experts.down_proj": {
|
| 145 |
+
"bits": 6
|
| 146 |
+
},
|
| 147 |
+
"language_model.model.layers.11.ffn.shared_experts.gate_up_proj": {
|
| 148 |
+
"bits": 6
|
| 149 |
+
},
|
| 150 |
+
"language_model.model.layers.12.attn.compressor.fused_wkv_wgate": {
|
| 151 |
+
"bits": 6
|
| 152 |
+
},
|
| 153 |
+
"language_model.model.layers.12.attn.fused_wqa_wkv": {
|
| 154 |
+
"bits": 6
|
| 155 |
+
},
|
| 156 |
+
"language_model.model.layers.12.attn.indexer.compressor.fused_wkv_wgate": {
|
| 157 |
+
"bits": 6
|
| 158 |
+
},
|
| 159 |
+
"language_model.model.layers.12.attn.indexer.wq_b": {
|
| 160 |
+
"bits": 6
|
| 161 |
+
},
|
| 162 |
+
"language_model.model.layers.12.attn.wo_a": {
|
| 163 |
+
"bits": 6
|
| 164 |
+
},
|
| 165 |
+
"language_model.model.layers.12.attn.wo_b": {
|
| 166 |
+
"bits": 6
|
| 167 |
+
},
|
| 168 |
+
"language_model.model.layers.12.attn.wq_b": {
|
| 169 |
+
"bits": 6
|
| 170 |
+
},
|
| 171 |
+
"language_model.model.layers.12.ffn.shared_experts.down_proj": {
|
| 172 |
+
"bits": 6
|
| 173 |
+
},
|
| 174 |
+
"language_model.model.layers.12.ffn.shared_experts.gate_up_proj": {
|
| 175 |
+
"bits": 6
|
| 176 |
+
},
|
| 177 |
+
"language_model.model.layers.13.attn.compressor.fused_wkv_wgate": {
|
| 178 |
+
"bits": 6
|
| 179 |
+
},
|
| 180 |
+
"language_model.model.layers.13.attn.fused_wqa_wkv": {
|
| 181 |
+
"bits": 6
|
| 182 |
+
},
|
| 183 |
+
"language_model.model.layers.13.attn.wo_a": {
|
| 184 |
+
"bits": 6
|
| 185 |
+
},
|
| 186 |
+
"language_model.model.layers.13.attn.wo_b": {
|
| 187 |
+
"bits": 6
|
| 188 |
+
},
|
| 189 |
+
"language_model.model.layers.13.attn.wq_b": {
|
| 190 |
+
"bits": 6
|
| 191 |
+
},
|
| 192 |
+
"language_model.model.layers.13.ffn.shared_experts.down_proj": {
|
| 193 |
+
"bits": 6
|
| 194 |
+
},
|
| 195 |
+
"language_model.model.layers.13.ffn.shared_experts.gate_up_proj": {
|
| 196 |
+
"bits": 6
|
| 197 |
+
},
|
| 198 |
+
"language_model.model.layers.14.attn.compressor.fused_wkv_wgate": {
|
| 199 |
+
"bits": 6
|
| 200 |
+
},
|
| 201 |
+
"language_model.model.layers.14.attn.fused_wqa_wkv": {
|
| 202 |
+
"bits": 6
|
| 203 |
+
},
|
| 204 |
+
"language_model.model.layers.14.attn.indexer.compressor.fused_wkv_wgate": {
|
| 205 |
+
"bits": 6
|
| 206 |
+
},
|
| 207 |
+
"language_model.model.layers.14.attn.indexer.wq_b": {
|
| 208 |
+
"bits": 6
|
| 209 |
+
},
|
| 210 |
+
"language_model.model.layers.14.attn.wo_a": {
|
| 211 |
+
"bits": 6
|
| 212 |
+
},
|
| 213 |
+
"language_model.model.layers.14.attn.wo_b": {
|
| 214 |
+
"bits": 6
|
| 215 |
+
},
|
| 216 |
+
"language_model.model.layers.14.attn.wq_b": {
|
| 217 |
+
"bits": 6
|
| 218 |
+
},
|
| 219 |
+
"language_model.model.layers.14.ffn.shared_experts.down_proj": {
|
| 220 |
+
"bits": 6
|
| 221 |
+
},
|
| 222 |
+
"language_model.model.layers.14.ffn.shared_experts.gate_up_proj": {
|
| 223 |
+
"bits": 6
|
| 224 |
+
},
|
| 225 |
+
"language_model.model.layers.15.attn.compressor.fused_wkv_wgate": {
|
| 226 |
+
"bits": 6
|
| 227 |
+
},
|
| 228 |
+
"language_model.model.layers.15.attn.fused_wqa_wkv": {
|
| 229 |
+
"bits": 6
|
| 230 |
+
},
|
| 231 |
+
"language_model.model.layers.15.attn.wo_a": {
|
| 232 |
+
"bits": 6
|
| 233 |
+
},
|
| 234 |
+
"language_model.model.layers.15.attn.wo_b": {
|
| 235 |
+
"bits": 6
|
| 236 |
+
},
|
| 237 |
+
"language_model.model.layers.15.attn.wq_b": {
|
| 238 |
+
"bits": 6
|
| 239 |
+
},
|
| 240 |
+
"language_model.model.layers.15.ffn.shared_experts.down_proj": {
|
| 241 |
+
"bits": 6
|
| 242 |
+
},
|
| 243 |
+
"language_model.model.layers.15.ffn.shared_experts.gate_up_proj": {
|
| 244 |
+
"bits": 6
|
| 245 |
+
},
|
| 246 |
+
"language_model.model.layers.16.attn.compressor.fused_wkv_wgate": {
|
| 247 |
+
"bits": 6
|
| 248 |
+
},
|
| 249 |
+
"language_model.model.layers.16.attn.fused_wqa_wkv": {
|
| 250 |
+
"bits": 6
|
| 251 |
+
},
|
| 252 |
+
"language_model.model.layers.16.attn.indexer.compressor.fused_wkv_wgate": {
|
| 253 |
+
"bits": 6
|
| 254 |
+
},
|
| 255 |
+
"language_model.model.layers.16.attn.indexer.wq_b": {
|
| 256 |
+
"bits": 6
|
| 257 |
+
},
|
| 258 |
+
"language_model.model.layers.16.attn.wo_a": {
|
| 259 |
+
"bits": 6
|
| 260 |
+
},
|
| 261 |
+
"language_model.model.layers.16.attn.wo_b": {
|
| 262 |
+
"bits": 6
|
| 263 |
+
},
|
| 264 |
+
"language_model.model.layers.16.attn.wq_b": {
|
| 265 |
+
"bits": 6
|
| 266 |
+
},
|
| 267 |
+
"language_model.model.layers.16.ffn.shared_experts.down_proj": {
|
| 268 |
+
"bits": 6
|
| 269 |
+
},
|
| 270 |
+
"language_model.model.layers.16.ffn.shared_experts.gate_up_proj": {
|
| 271 |
+
"bits": 6
|
| 272 |
+
},
|
| 273 |
+
"language_model.model.layers.17.attn.compressor.fused_wkv_wgate": {
|
| 274 |
+
"bits": 6
|
| 275 |
+
},
|
| 276 |
+
"language_model.model.layers.17.attn.fused_wqa_wkv": {
|
| 277 |
+
"bits": 6
|
| 278 |
+
},
|
| 279 |
+
"language_model.model.layers.17.attn.wo_a": {
|
| 280 |
+
"bits": 6
|
| 281 |
+
},
|
| 282 |
+
"language_model.model.layers.17.attn.wo_b": {
|
| 283 |
+
"bits": 6
|
| 284 |
+
},
|
| 285 |
+
"language_model.model.layers.17.attn.wq_b": {
|
| 286 |
+
"bits": 6
|
| 287 |
+
},
|
| 288 |
+
"language_model.model.layers.17.ffn.shared_experts.down_proj": {
|
| 289 |
+
"bits": 6
|
| 290 |
+
},
|
| 291 |
+
"language_model.model.layers.17.ffn.shared_experts.gate_up_proj": {
|
| 292 |
+
"bits": 6
|
| 293 |
+
},
|
| 294 |
+
"language_model.model.layers.18.attn.compressor.fused_wkv_wgate": {
|
| 295 |
+
"bits": 6
|
| 296 |
+
},
|
| 297 |
+
"language_model.model.layers.18.attn.fused_wqa_wkv": {
|
| 298 |
+
"bits": 6
|
| 299 |
+
},
|
| 300 |
+
"language_model.model.layers.18.attn.indexer.compressor.fused_wkv_wgate": {
|
| 301 |
+
"bits": 6
|
| 302 |
+
},
|
| 303 |
+
"language_model.model.layers.18.attn.indexer.wq_b": {
|
| 304 |
+
"bits": 6
|
| 305 |
+
},
|
| 306 |
+
"language_model.model.layers.18.attn.wo_a": {
|
| 307 |
+
"bits": 6
|
| 308 |
+
},
|
| 309 |
+
"language_model.model.layers.18.attn.wo_b": {
|
| 310 |
+
"bits": 6
|
| 311 |
+
},
|
| 312 |
+
"language_model.model.layers.18.attn.wq_b": {
|
| 313 |
+
"bits": 6
|
| 314 |
+
},
|
| 315 |
+
"language_model.model.layers.18.ffn.shared_experts.down_proj": {
|
| 316 |
+
"bits": 6
|
| 317 |
+
},
|
| 318 |
+
"language_model.model.layers.18.ffn.shared_experts.gate_up_proj": {
|
| 319 |
+
"bits": 6
|
| 320 |
+
},
|
| 321 |
+
"language_model.model.layers.19.attn.compressor.fused_wkv_wgate": {
|
| 322 |
+
"bits": 6
|
| 323 |
+
},
|
| 324 |
+
"language_model.model.layers.19.attn.fused_wqa_wkv": {
|
| 325 |
+
"bits": 6
|
| 326 |
+
},
|
| 327 |
+
"language_model.model.layers.19.attn.wo_a": {
|
| 328 |
+
"bits": 6
|
| 329 |
+
},
|
| 330 |
+
"language_model.model.layers.19.attn.wo_b": {
|
| 331 |
+
"bits": 6
|
| 332 |
+
},
|
| 333 |
+
"language_model.model.layers.19.attn.wq_b": {
|
| 334 |
+
"bits": 6
|
| 335 |
+
},
|
| 336 |
+
"language_model.model.layers.19.ffn.shared_experts.down_proj": {
|
| 337 |
+
"bits": 6
|
| 338 |
+
},
|
| 339 |
+
"language_model.model.layers.19.ffn.shared_experts.gate_up_proj": {
|
| 340 |
+
"bits": 6
|
| 341 |
+
},
|
| 342 |
+
"language_model.model.layers.2.attn.compressor.fused_wkv_wgate": {
|
| 343 |
+
"bits": 6
|
| 344 |
+
},
|
| 345 |
+
"language_model.model.layers.2.attn.fused_wqa_wkv": {
|
| 346 |
+
"bits": 6
|
| 347 |
+
},
|
| 348 |
+
"language_model.model.layers.2.attn.indexer.compressor.fused_wkv_wgate": {
|
| 349 |
+
"bits": 6
|
| 350 |
+
},
|
| 351 |
+
"language_model.model.layers.2.attn.indexer.wq_b": {
|
| 352 |
+
"bits": 6
|
| 353 |
+
},
|
| 354 |
+
"language_model.model.layers.2.attn.wo_a": {
|
| 355 |
+
"bits": 6
|
| 356 |
+
},
|
| 357 |
+
"language_model.model.layers.2.attn.wo_b": {
|
| 358 |
+
"bits": 6
|
| 359 |
+
},
|
| 360 |
+
"language_model.model.layers.2.attn.wq_b": {
|
| 361 |
+
"bits": 6
|
| 362 |
+
},
|
| 363 |
+
"language_model.model.layers.2.ffn.shared_experts.down_proj": {
|
| 364 |
+
"bits": 6
|
| 365 |
+
},
|
| 366 |
+
"language_model.model.layers.2.ffn.shared_experts.gate_up_proj": {
|
| 367 |
+
"bits": 6
|
| 368 |
+
},
|
| 369 |
+
"language_model.model.layers.20.attn.compressor.fused_wkv_wgate": {
|
| 370 |
+
"bits": 6
|
| 371 |
+
},
|
| 372 |
+
"language_model.model.layers.20.attn.fused_wqa_wkv": {
|
| 373 |
+
"bits": 6
|
| 374 |
+
},
|
| 375 |
+
"language_model.model.layers.20.attn.indexer.compressor.fused_wkv_wgate": {
|
| 376 |
+
"bits": 6
|
| 377 |
+
},
|
| 378 |
+
"language_model.model.layers.20.attn.indexer.wq_b": {
|
| 379 |
+
"bits": 6
|
| 380 |
+
},
|
| 381 |
+
"language_model.model.layers.20.attn.wo_a": {
|
| 382 |
+
"bits": 6
|
| 383 |
+
},
|
| 384 |
+
"language_model.model.layers.20.attn.wo_b": {
|
| 385 |
+
"bits": 6
|
| 386 |
+
},
|
| 387 |
+
"language_model.model.layers.20.attn.wq_b": {
|
| 388 |
+
"bits": 6
|
| 389 |
+
},
|
| 390 |
+
"language_model.model.layers.20.ffn.shared_experts.down_proj": {
|
| 391 |
+
"bits": 6
|
| 392 |
+
},
|
| 393 |
+
"language_model.model.layers.20.ffn.shared_experts.gate_up_proj": {
|
| 394 |
+
"bits": 6
|
| 395 |
+
},
|
| 396 |
+
"language_model.model.layers.21.attn.compressor.fused_wkv_wgate": {
|
| 397 |
+
"bits": 6
|
| 398 |
+
},
|
| 399 |
+
"language_model.model.layers.21.attn.fused_wqa_wkv": {
|
| 400 |
+
"bits": 6
|
| 401 |
+
},
|
| 402 |
+
"language_model.model.layers.21.attn.wo_a": {
|
| 403 |
+
"bits": 6
|
| 404 |
+
},
|
| 405 |
+
"language_model.model.layers.21.attn.wo_b": {
|
| 406 |
+
"bits": 6
|
| 407 |
+
},
|
| 408 |
+
"language_model.model.layers.21.attn.wq_b": {
|
| 409 |
+
"bits": 6
|
| 410 |
+
},
|
| 411 |
+
"language_model.model.layers.21.ffn.shared_experts.down_proj": {
|
| 412 |
+
"bits": 6
|
| 413 |
+
},
|
| 414 |
+
"language_model.model.layers.21.ffn.shared_experts.gate_up_proj": {
|
| 415 |
+
"bits": 6
|
| 416 |
+
},
|
| 417 |
+
"language_model.model.layers.22.attn.compressor.fused_wkv_wgate": {
|
| 418 |
+
"bits": 6
|
| 419 |
+
},
|
| 420 |
+
"language_model.model.layers.22.attn.fused_wqa_wkv": {
|
| 421 |
+
"bits": 6
|
| 422 |
+
},
|
| 423 |
+
"language_model.model.layers.22.attn.indexer.compressor.fused_wkv_wgate": {
|
| 424 |
+
"bits": 6
|
| 425 |
+
},
|
| 426 |
+
"language_model.model.layers.22.attn.indexer.wq_b": {
|
| 427 |
+
"bits": 6
|
| 428 |
+
},
|
| 429 |
+
"language_model.model.layers.22.attn.wo_a": {
|
| 430 |
+
"bits": 6
|
| 431 |
+
},
|
| 432 |
+
"language_model.model.layers.22.attn.wo_b": {
|
| 433 |
+
"bits": 6
|
| 434 |
+
},
|
| 435 |
+
"language_model.model.layers.22.attn.wq_b": {
|
| 436 |
+
"bits": 6
|
| 437 |
+
},
|
| 438 |
+
"language_model.model.layers.22.ffn.shared_experts.down_proj": {
|
| 439 |
+
"bits": 6
|
| 440 |
+
},
|
| 441 |
+
"language_model.model.layers.22.ffn.shared_experts.gate_up_proj": {
|
| 442 |
+
"bits": 6
|
| 443 |
+
},
|
| 444 |
+
"language_model.model.layers.23.attn.compressor.fused_wkv_wgate": {
|
| 445 |
+
"bits": 6
|
| 446 |
+
},
|
| 447 |
+
"language_model.model.layers.23.attn.fused_wqa_wkv": {
|
| 448 |
+
"bits": 6
|
| 449 |
+
},
|
| 450 |
+
"language_model.model.layers.23.attn.wo_a": {
|
| 451 |
+
"bits": 6
|
| 452 |
+
},
|
| 453 |
+
"language_model.model.layers.23.attn.wo_b": {
|
| 454 |
+
"bits": 6
|
| 455 |
+
},
|
| 456 |
+
"language_model.model.layers.23.attn.wq_b": {
|
| 457 |
+
"bits": 6
|
| 458 |
+
},
|
| 459 |
+
"language_model.model.layers.23.ffn.shared_experts.down_proj": {
|
| 460 |
+
"bits": 6
|
| 461 |
+
},
|
| 462 |
+
"language_model.model.layers.23.ffn.shared_experts.gate_up_proj": {
|
| 463 |
+
"bits": 6
|
| 464 |
+
},
|
| 465 |
+
"language_model.model.layers.24.attn.compressor.fused_wkv_wgate": {
|
| 466 |
+
"bits": 6
|
| 467 |
+
},
|
| 468 |
+
"language_model.model.layers.24.attn.fused_wqa_wkv": {
|
| 469 |
+
"bits": 6
|
| 470 |
+
},
|
| 471 |
+
"language_model.model.layers.24.attn.indexer.compressor.fused_wkv_wgate": {
|
| 472 |
+
"bits": 6
|
| 473 |
+
},
|
| 474 |
+
"language_model.model.layers.24.attn.indexer.wq_b": {
|
| 475 |
+
"bits": 6
|
| 476 |
+
},
|
| 477 |
+
"language_model.model.layers.24.attn.wo_a": {
|
| 478 |
+
"bits": 6
|
| 479 |
+
},
|
| 480 |
+
"language_model.model.layers.24.attn.wo_b": {
|
| 481 |
+
"bits": 6
|
| 482 |
+
},
|
| 483 |
+
"language_model.model.layers.24.attn.wq_b": {
|
| 484 |
+
"bits": 6
|
| 485 |
+
},
|
| 486 |
+
"language_model.model.layers.24.ffn.shared_experts.down_proj": {
|
| 487 |
+
"bits": 6
|
| 488 |
+
},
|
| 489 |
+
"language_model.model.layers.24.ffn.shared_experts.gate_up_proj": {
|
| 490 |
+
"bits": 6
|
| 491 |
+
},
|
| 492 |
+
"language_model.model.layers.25.attn.compressor.fused_wkv_wgate": {
|
| 493 |
+
"bits": 6
|
| 494 |
+
},
|
| 495 |
+
"language_model.model.layers.25.attn.fused_wqa_wkv": {
|
| 496 |
+
"bits": 6
|
| 497 |
+
},
|
| 498 |
+
"language_model.model.layers.25.attn.wo_a": {
|
| 499 |
+
"bits": 6
|
| 500 |
+
},
|
| 501 |
+
"language_model.model.layers.25.attn.wo_b": {
|
| 502 |
+
"bits": 6
|
| 503 |
+
},
|
| 504 |
+
"language_model.model.layers.25.attn.wq_b": {
|
| 505 |
+
"bits": 6
|
| 506 |
+
},
|
| 507 |
+
"language_model.model.layers.25.ffn.shared_experts.down_proj": {
|
| 508 |
+
"bits": 6
|
| 509 |
+
},
|
| 510 |
+
"language_model.model.layers.25.ffn.shared_experts.gate_up_proj": {
|
| 511 |
+
"bits": 6
|
| 512 |
+
},
|
| 513 |
+
"language_model.model.layers.26.attn.compressor.fused_wkv_wgate": {
|
| 514 |
+
"bits": 6
|
| 515 |
+
},
|
| 516 |
+
"language_model.model.layers.26.attn.fused_wqa_wkv": {
|
| 517 |
+
"bits": 6
|
| 518 |
+
},
|
| 519 |
+
"language_model.model.layers.26.attn.indexer.compressor.fused_wkv_wgate": {
|
| 520 |
+
"bits": 6
|
| 521 |
+
},
|
| 522 |
+
"language_model.model.layers.26.attn.indexer.wq_b": {
|
| 523 |
+
"bits": 6
|
| 524 |
+
},
|
| 525 |
+
"language_model.model.layers.26.attn.wo_a": {
|
| 526 |
+
"bits": 6
|
| 527 |
+
},
|
| 528 |
+
"language_model.model.layers.26.attn.wo_b": {
|
| 529 |
+
"bits": 6
|
| 530 |
+
},
|
| 531 |
+
"language_model.model.layers.26.attn.wq_b": {
|
| 532 |
+
"bits": 6
|
| 533 |
+
},
|
| 534 |
+
"language_model.model.layers.26.ffn.shared_experts.down_proj": {
|
| 535 |
+
"bits": 6
|
| 536 |
+
},
|
| 537 |
+
"language_model.model.layers.26.ffn.shared_experts.gate_up_proj": {
|
| 538 |
+
"bits": 6
|
| 539 |
+
},
|
| 540 |
+
"language_model.model.layers.27.attn.compressor.fused_wkv_wgate": {
|
| 541 |
+
"bits": 6
|
| 542 |
+
},
|
| 543 |
+
"language_model.model.layers.27.attn.fused_wqa_wkv": {
|
| 544 |
+
"bits": 6
|
| 545 |
+
},
|
| 546 |
+
"language_model.model.layers.27.attn.wo_a": {
|
| 547 |
+
"bits": 6
|
| 548 |
+
},
|
| 549 |
+
"language_model.model.layers.27.attn.wo_b": {
|
| 550 |
+
"bits": 6
|
| 551 |
+
},
|
| 552 |
+
"language_model.model.layers.27.attn.wq_b": {
|
| 553 |
+
"bits": 6
|
| 554 |
+
},
|
| 555 |
+
"language_model.model.layers.27.ffn.shared_experts.down_proj": {
|
| 556 |
+
"bits": 6
|
| 557 |
+
},
|
| 558 |
+
"language_model.model.layers.27.ffn.shared_experts.gate_up_proj": {
|
| 559 |
+
"bits": 6
|
| 560 |
+
},
|
| 561 |
+
"language_model.model.layers.28.attn.compressor.fused_wkv_wgate": {
|
| 562 |
+
"bits": 6
|
| 563 |
+
},
|
| 564 |
+
"language_model.model.layers.28.attn.fused_wqa_wkv": {
|
| 565 |
+
"bits": 6
|
| 566 |
+
},
|
| 567 |
+
"language_model.model.layers.28.attn.indexer.compressor.fused_wkv_wgate": {
|
| 568 |
+
"bits": 6
|
| 569 |
+
},
|
| 570 |
+
"language_model.model.layers.28.attn.indexer.wq_b": {
|
| 571 |
+
"bits": 6
|
| 572 |
+
},
|
| 573 |
+
"language_model.model.layers.28.attn.wo_a": {
|
| 574 |
+
"bits": 6
|
| 575 |
+
},
|
| 576 |
+
"language_model.model.layers.28.attn.wo_b": {
|
| 577 |
+
"bits": 6
|
| 578 |
+
},
|
| 579 |
+
"language_model.model.layers.28.attn.wq_b": {
|
| 580 |
+
"bits": 6
|
| 581 |
+
},
|
| 582 |
+
"language_model.model.layers.28.ffn.shared_experts.down_proj": {
|
| 583 |
+
"bits": 6
|
| 584 |
+
},
|
| 585 |
+
"language_model.model.layers.28.ffn.shared_experts.gate_up_proj": {
|
| 586 |
+
"bits": 6
|
| 587 |
+
},
|
| 588 |
+
"language_model.model.layers.29.attn.compressor.fused_wkv_wgate": {
|
| 589 |
+
"bits": 6
|
| 590 |
+
},
|
| 591 |
+
"language_model.model.layers.29.attn.fused_wqa_wkv": {
|
| 592 |
+
"bits": 6
|
| 593 |
+
},
|
| 594 |
+
"language_model.model.layers.29.attn.wo_a": {
|
| 595 |
+
"bits": 6
|
| 596 |
+
},
|
| 597 |
+
"language_model.model.layers.29.attn.wo_b": {
|
| 598 |
+
"bits": 6
|
| 599 |
+
},
|
| 600 |
+
"language_model.model.layers.29.attn.wq_b": {
|
| 601 |
+
"bits": 6
|
| 602 |
+
},
|
| 603 |
+
"language_model.model.layers.29.ffn.shared_experts.down_proj": {
|
| 604 |
+
"bits": 6
|
| 605 |
+
},
|
| 606 |
+
"language_model.model.layers.29.ffn.shared_experts.gate_up_proj": {
|
| 607 |
+
"bits": 6
|
| 608 |
+
},
|
| 609 |
+
"language_model.model.layers.3.attn.compressor.fused_wkv_wgate": {
|
| 610 |
+
"bits": 6
|
| 611 |
+
},
|
| 612 |
+
"language_model.model.layers.3.attn.fused_wqa_wkv": {
|
| 613 |
+
"bits": 6
|
| 614 |
+
},
|
| 615 |
+
"language_model.model.layers.3.attn.wo_a": {
|
| 616 |
+
"bits": 6
|
| 617 |
+
},
|
| 618 |
+
"language_model.model.layers.3.attn.wo_b": {
|
| 619 |
+
"bits": 6
|
| 620 |
+
},
|
| 621 |
+
"language_model.model.layers.3.attn.wq_b": {
|
| 622 |
+
"bits": 6
|
| 623 |
+
},
|
| 624 |
+
"language_model.model.layers.3.ffn.shared_experts.down_proj": {
|
| 625 |
+
"bits": 6
|
| 626 |
+
},
|
| 627 |
+
"language_model.model.layers.3.ffn.shared_experts.gate_up_proj": {
|
| 628 |
+
"bits": 6
|
| 629 |
+
},
|
| 630 |
+
"language_model.model.layers.30.attn.compressor.fused_wkv_wgate": {
|
| 631 |
+
"bits": 6
|
| 632 |
+
},
|
| 633 |
+
"language_model.model.layers.30.attn.fused_wqa_wkv": {
|
| 634 |
+
"bits": 6
|
| 635 |
+
},
|
| 636 |
+
"language_model.model.layers.30.attn.indexer.compressor.fused_wkv_wgate": {
|
| 637 |
+
"bits": 6
|
| 638 |
+
},
|
| 639 |
+
"language_model.model.layers.30.attn.indexer.wq_b": {
|
| 640 |
+
"bits": 6
|
| 641 |
+
},
|
| 642 |
+
"language_model.model.layers.30.attn.wo_a": {
|
| 643 |
+
"bits": 6
|
| 644 |
+
},
|
| 645 |
+
"language_model.model.layers.30.attn.wo_b": {
|
| 646 |
+
"bits": 6
|
| 647 |
+
},
|
| 648 |
+
"language_model.model.layers.30.attn.wq_b": {
|
| 649 |
+
"bits": 6
|
| 650 |
+
},
|
| 651 |
+
"language_model.model.layers.30.ffn.shared_experts.down_proj": {
|
| 652 |
+
"bits": 6
|
| 653 |
+
},
|
| 654 |
+
"language_model.model.layers.30.ffn.shared_experts.gate_up_proj": {
|
| 655 |
+
"bits": 6
|
| 656 |
+
},
|
| 657 |
+
"language_model.model.layers.31.attn.compressor.fused_wkv_wgate": {
|
| 658 |
+
"bits": 6
|
| 659 |
+
},
|
| 660 |
+
"language_model.model.layers.31.attn.fused_wqa_wkv": {
|
| 661 |
+
"bits": 6
|
| 662 |
+
},
|
| 663 |
+
"language_model.model.layers.31.attn.wo_a": {
|
| 664 |
+
"bits": 6
|
| 665 |
+
},
|
| 666 |
+
"language_model.model.layers.31.attn.wo_b": {
|
| 667 |
+
"bits": 6
|
| 668 |
+
},
|
| 669 |
+
"language_model.model.layers.31.attn.wq_b": {
|
| 670 |
+
"bits": 6
|
| 671 |
+
},
|
| 672 |
+
"language_model.model.layers.31.ffn.shared_experts.down_proj": {
|
| 673 |
+
"bits": 6
|
| 674 |
+
},
|
| 675 |
+
"language_model.model.layers.31.ffn.shared_experts.gate_up_proj": {
|
| 676 |
+
"bits": 6
|
| 677 |
+
},
|
| 678 |
+
"language_model.model.layers.32.attn.compressor.fused_wkv_wgate": {
|
| 679 |
+
"bits": 6
|
| 680 |
+
},
|
| 681 |
+
"language_model.model.layers.32.attn.fused_wqa_wkv": {
|
| 682 |
+
"bits": 6
|
| 683 |
+
},
|
| 684 |
+
"language_model.model.layers.32.attn.indexer.compressor.fused_wkv_wgate": {
|
| 685 |
+
"bits": 6
|
| 686 |
+
},
|
| 687 |
+
"language_model.model.layers.32.attn.indexer.wq_b": {
|
| 688 |
+
"bits": 6
|
| 689 |
+
},
|
| 690 |
+
"language_model.model.layers.32.attn.wo_a": {
|
| 691 |
+
"bits": 6
|
| 692 |
+
},
|
| 693 |
+
"language_model.model.layers.32.attn.wo_b": {
|
| 694 |
+
"bits": 6
|
| 695 |
+
},
|
| 696 |
+
"language_model.model.layers.32.attn.wq_b": {
|
| 697 |
+
"bits": 6
|
| 698 |
+
},
|
| 699 |
+
"language_model.model.layers.32.ffn.shared_experts.down_proj": {
|
| 700 |
+
"bits": 6
|
| 701 |
+
},
|
| 702 |
+
"language_model.model.layers.32.ffn.shared_experts.gate_up_proj": {
|
| 703 |
+
"bits": 6
|
| 704 |
+
},
|
| 705 |
+
"language_model.model.layers.33.attn.compressor.fused_wkv_wgate": {
|
| 706 |
+
"bits": 6
|
| 707 |
+
},
|
| 708 |
+
"language_model.model.layers.33.attn.fused_wqa_wkv": {
|
| 709 |
+
"bits": 6
|
| 710 |
+
},
|
| 711 |
+
"language_model.model.layers.33.attn.wo_a": {
|
| 712 |
+
"bits": 6
|
| 713 |
+
},
|
| 714 |
+
"language_model.model.layers.33.attn.wo_b": {
|
| 715 |
+
"bits": 6
|
| 716 |
+
},
|
| 717 |
+
"language_model.model.layers.33.attn.wq_b": {
|
| 718 |
+
"bits": 6
|
| 719 |
+
},
|
| 720 |
+
"language_model.model.layers.33.ffn.shared_experts.down_proj": {
|
| 721 |
+
"bits": 6
|
| 722 |
+
},
|
| 723 |
+
"language_model.model.layers.33.ffn.shared_experts.gate_up_proj": {
|
| 724 |
+
"bits": 6
|
| 725 |
+
},
|
| 726 |
+
"language_model.model.layers.34.attn.compressor.fused_wkv_wgate": {
|
| 727 |
+
"bits": 6
|
| 728 |
+
},
|
| 729 |
+
"language_model.model.layers.34.attn.fused_wqa_wkv": {
|
| 730 |
+
"bits": 6
|
| 731 |
+
},
|
| 732 |
+
"language_model.model.layers.34.attn.indexer.compressor.fused_wkv_wgate": {
|
| 733 |
+
"bits": 6
|
| 734 |
+
},
|
| 735 |
+
"language_model.model.layers.34.attn.indexer.wq_b": {
|
| 736 |
+
"bits": 6
|
| 737 |
+
},
|
| 738 |
+
"language_model.model.layers.34.attn.wo_a": {
|
| 739 |
+
"bits": 6
|
| 740 |
+
},
|
| 741 |
+
"language_model.model.layers.34.attn.wo_b": {
|
| 742 |
+
"bits": 6
|
| 743 |
+
},
|
| 744 |
+
"language_model.model.layers.34.attn.wq_b": {
|
| 745 |
+
"bits": 6
|
| 746 |
+
},
|
| 747 |
+
"language_model.model.layers.34.ffn.shared_experts.down_proj": {
|
| 748 |
+
"bits": 6
|
| 749 |
+
},
|
| 750 |
+
"language_model.model.layers.34.ffn.shared_experts.gate_up_proj": {
|
| 751 |
+
"bits": 6
|
| 752 |
+
},
|
| 753 |
+
"language_model.model.layers.35.attn.compressor.fused_wkv_wgate": {
|
| 754 |
+
"bits": 6
|
| 755 |
+
},
|
| 756 |
+
"language_model.model.layers.35.attn.fused_wqa_wkv": {
|
| 757 |
+
"bits": 6
|
| 758 |
+
},
|
| 759 |
+
"language_model.model.layers.35.attn.wo_a": {
|
| 760 |
+
"bits": 6
|
| 761 |
+
},
|
| 762 |
+
"language_model.model.layers.35.attn.wo_b": {
|
| 763 |
+
"bits": 6
|
| 764 |
+
},
|
| 765 |
+
"language_model.model.layers.35.attn.wq_b": {
|
| 766 |
+
"bits": 6
|
| 767 |
+
},
|
| 768 |
+
"language_model.model.layers.35.ffn.shared_experts.down_proj": {
|
| 769 |
+
"bits": 6
|
| 770 |
+
},
|
| 771 |
+
"language_model.model.layers.35.ffn.shared_experts.gate_up_proj": {
|
| 772 |
+
"bits": 6
|
| 773 |
+
},
|
| 774 |
+
"language_model.model.layers.36.attn.compressor.fused_wkv_wgate": {
|
| 775 |
+
"bits": 6
|
| 776 |
+
},
|
| 777 |
+
"language_model.model.layers.36.attn.fused_wqa_wkv": {
|
| 778 |
+
"bits": 6
|
| 779 |
+
},
|
| 780 |
+
"language_model.model.layers.36.attn.indexer.compressor.fused_wkv_wgate": {
|
| 781 |
+
"bits": 6
|
| 782 |
+
},
|
| 783 |
+
"language_model.model.layers.36.attn.indexer.wq_b": {
|
| 784 |
+
"bits": 6
|
| 785 |
+
},
|
| 786 |
+
"language_model.model.layers.36.attn.wo_a": {
|
| 787 |
+
"bits": 6
|
| 788 |
+
},
|
| 789 |
+
"language_model.model.layers.36.attn.wo_b": {
|
| 790 |
+
"bits": 6
|
| 791 |
+
},
|
| 792 |
+
"language_model.model.layers.36.attn.wq_b": {
|
| 793 |
+
"bits": 6
|
| 794 |
+
},
|
| 795 |
+
"language_model.model.layers.36.ffn.shared_experts.down_proj": {
|
| 796 |
+
"bits": 6
|
| 797 |
+
},
|
| 798 |
+
"language_model.model.layers.36.ffn.shared_experts.gate_up_proj": {
|
| 799 |
+
"bits": 6
|
| 800 |
+
},
|
| 801 |
+
"language_model.model.layers.37.attn.compressor.fused_wkv_wgate": {
|
| 802 |
+
"bits": 6
|
| 803 |
+
},
|
| 804 |
+
"language_model.model.layers.37.attn.fused_wqa_wkv": {
|
| 805 |
+
"bits": 6
|
| 806 |
+
},
|
| 807 |
+
"language_model.model.layers.37.attn.wo_a": {
|
| 808 |
+
"bits": 6
|
| 809 |
+
},
|
| 810 |
+
"language_model.model.layers.37.attn.wo_b": {
|
| 811 |
+
"bits": 6
|
| 812 |
+
},
|
| 813 |
+
"language_model.model.layers.37.attn.wq_b": {
|
| 814 |
+
"bits": 6
|
| 815 |
+
},
|
| 816 |
+
"language_model.model.layers.37.ffn.shared_experts.down_proj": {
|
| 817 |
+
"bits": 6
|
| 818 |
+
},
|
| 819 |
+
"language_model.model.layers.37.ffn.shared_experts.gate_up_proj": {
|
| 820 |
+
"bits": 6
|
| 821 |
+
},
|
| 822 |
+
"language_model.model.layers.38.attn.compressor.fused_wkv_wgate": {
|
| 823 |
+
"bits": 6
|
| 824 |
+
},
|
| 825 |
+
"language_model.model.layers.38.attn.fused_wqa_wkv": {
|
| 826 |
+
"bits": 6
|
| 827 |
+
},
|
| 828 |
+
"language_model.model.layers.38.attn.indexer.compressor.fused_wkv_wgate": {
|
| 829 |
+
"bits": 6
|
| 830 |
+
},
|
| 831 |
+
"language_model.model.layers.38.attn.indexer.wq_b": {
|
| 832 |
+
"bits": 6
|
| 833 |
+
},
|
| 834 |
+
"language_model.model.layers.38.attn.wo_a": {
|
| 835 |
+
"bits": 6
|
| 836 |
+
},
|
| 837 |
+
"language_model.model.layers.38.attn.wo_b": {
|
| 838 |
+
"bits": 6
|
| 839 |
+
},
|
| 840 |
+
"language_model.model.layers.38.attn.wq_b": {
|
| 841 |
+
"bits": 6
|
| 842 |
+
},
|
| 843 |
+
"language_model.model.layers.38.ffn.shared_experts.down_proj": {
|
| 844 |
+
"bits": 6
|
| 845 |
+
},
|
| 846 |
+
"language_model.model.layers.38.ffn.shared_experts.gate_up_proj": {
|
| 847 |
+
"bits": 6
|
| 848 |
+
},
|
| 849 |
+
"language_model.model.layers.39.attn.compressor.fused_wkv_wgate": {
|
| 850 |
+
"bits": 6
|
| 851 |
+
},
|
| 852 |
+
"language_model.model.layers.39.attn.fused_wqa_wkv": {
|
| 853 |
+
"bits": 6
|
| 854 |
+
},
|
| 855 |
+
"language_model.model.layers.39.attn.wo_a": {
|
| 856 |
+
"bits": 6
|
| 857 |
+
},
|
| 858 |
+
"language_model.model.layers.39.attn.wo_b": {
|
| 859 |
+
"bits": 6
|
| 860 |
+
},
|
| 861 |
+
"language_model.model.layers.39.attn.wq_b": {
|
| 862 |
+
"bits": 6
|
| 863 |
+
},
|
| 864 |
+
"language_model.model.layers.39.ffn.shared_experts.down_proj": {
|
| 865 |
+
"bits": 6
|
| 866 |
+
},
|
| 867 |
+
"language_model.model.layers.39.ffn.shared_experts.gate_up_proj": {
|
| 868 |
+
"bits": 6
|
| 869 |
+
},
|
| 870 |
+
"language_model.model.layers.4.attn.compressor.fused_wkv_wgate": {
|
| 871 |
+
"bits": 6
|
| 872 |
+
},
|
| 873 |
+
"language_model.model.layers.4.attn.fused_wqa_wkv": {
|
| 874 |
+
"bits": 6
|
| 875 |
+
},
|
| 876 |
+
"language_model.model.layers.4.attn.indexer.compressor.fused_wkv_wgate": {
|
| 877 |
+
"bits": 6
|
| 878 |
+
},
|
| 879 |
+
"language_model.model.layers.4.attn.indexer.wq_b": {
|
| 880 |
+
"bits": 6
|
| 881 |
+
},
|
| 882 |
+
"language_model.model.layers.4.attn.wo_a": {
|
| 883 |
+
"bits": 6
|
| 884 |
+
},
|
| 885 |
+
"language_model.model.layers.4.attn.wo_b": {
|
| 886 |
+
"bits": 6
|
| 887 |
+
},
|
| 888 |
+
"language_model.model.layers.4.attn.wq_b": {
|
| 889 |
+
"bits": 6
|
| 890 |
+
},
|
| 891 |
+
"language_model.model.layers.4.ffn.shared_experts.down_proj": {
|
| 892 |
+
"bits": 6
|
| 893 |
+
},
|
| 894 |
+
"language_model.model.layers.4.ffn.shared_experts.gate_up_proj": {
|
| 895 |
+
"bits": 6
|
| 896 |
+
},
|
| 897 |
+
"language_model.model.layers.40.attn.compressor.fused_wkv_wgate": {
|
| 898 |
+
"bits": 6
|
| 899 |
+
},
|
| 900 |
+
"language_model.model.layers.40.attn.fused_wqa_wkv": {
|
| 901 |
+
"bits": 6
|
| 902 |
+
},
|
| 903 |
+
"language_model.model.layers.40.attn.indexer.compressor.fused_wkv_wgate": {
|
| 904 |
+
"bits": 6
|
| 905 |
+
},
|
| 906 |
+
"language_model.model.layers.40.attn.indexer.wq_b": {
|
| 907 |
+
"bits": 6
|
| 908 |
+
},
|
| 909 |
+
"language_model.model.layers.40.attn.wo_a": {
|
| 910 |
+
"bits": 6
|
| 911 |
+
},
|
| 912 |
+
"language_model.model.layers.40.attn.wo_b": {
|
| 913 |
+
"bits": 6
|
| 914 |
+
},
|
| 915 |
+
"language_model.model.layers.40.attn.wq_b": {
|
| 916 |
+
"bits": 6
|
| 917 |
+
},
|
| 918 |
+
"language_model.model.layers.40.ffn.shared_experts.down_proj": {
|
| 919 |
+
"bits": 6
|
| 920 |
+
},
|
| 921 |
+
"language_model.model.layers.40.ffn.shared_experts.gate_up_proj": {
|
| 922 |
+
"bits": 6
|
| 923 |
+
},
|
| 924 |
+
"language_model.model.layers.41.attn.compressor.fused_wkv_wgate": {
|
| 925 |
+
"bits": 6
|
| 926 |
+
},
|
| 927 |
+
"language_model.model.layers.41.attn.fused_wqa_wkv": {
|
| 928 |
+
"bits": 6
|
| 929 |
+
},
|
| 930 |
+
"language_model.model.layers.41.attn.wo_a": {
|
| 931 |
+
"bits": 6
|
| 932 |
+
},
|
| 933 |
+
"language_model.model.layers.41.attn.wo_b": {
|
| 934 |
+
"bits": 6
|
| 935 |
+
},
|
| 936 |
+
"language_model.model.layers.41.attn.wq_b": {
|
| 937 |
+
"bits": 6
|
| 938 |
+
},
|
| 939 |
+
"language_model.model.layers.41.ffn.shared_experts.down_proj": {
|
| 940 |
+
"bits": 6
|
| 941 |
+
},
|
| 942 |
+
"language_model.model.layers.41.ffn.shared_experts.gate_up_proj": {
|
| 943 |
+
"bits": 6
|
| 944 |
+
},
|
| 945 |
+
"language_model.model.layers.42.attn.compressor.fused_wkv_wgate": {
|
| 946 |
+
"bits": 6
|
| 947 |
+
},
|
| 948 |
+
"language_model.model.layers.42.attn.fused_wqa_wkv": {
|
| 949 |
+
"bits": 6
|
| 950 |
+
},
|
| 951 |
+
"language_model.model.layers.42.attn.indexer.compressor.fused_wkv_wgate": {
|
| 952 |
+
"bits": 6
|
| 953 |
+
},
|
| 954 |
+
"language_model.model.layers.42.attn.indexer.wq_b": {
|
| 955 |
+
"bits": 6
|
| 956 |
+
},
|
| 957 |
+
"language_model.model.layers.42.attn.wo_a": {
|
| 958 |
+
"bits": 6
|
| 959 |
+
},
|
| 960 |
+
"language_model.model.layers.42.attn.wo_b": {
|
| 961 |
+
"bits": 6
|
| 962 |
+
},
|
| 963 |
+
"language_model.model.layers.42.attn.wq_b": {
|
| 964 |
+
"bits": 6
|
| 965 |
+
},
|
| 966 |
+
"language_model.model.layers.42.ffn.shared_experts.down_proj": {
|
| 967 |
+
"bits": 6
|
| 968 |
+
},
|
| 969 |
+
"language_model.model.layers.42.ffn.shared_experts.gate_up_proj": {
|
| 970 |
+
"bits": 6
|
| 971 |
+
},
|
| 972 |
+
"language_model.model.layers.5.attn.compressor.fused_wkv_wgate": {
|
| 973 |
+
"bits": 6
|
| 974 |
+
},
|
| 975 |
+
"language_model.model.layers.5.attn.fused_wqa_wkv": {
|
| 976 |
+
"bits": 6
|
| 977 |
+
},
|
| 978 |
+
"language_model.model.layers.5.attn.wo_a": {
|
| 979 |
+
"bits": 6
|
| 980 |
+
},
|
| 981 |
+
"language_model.model.layers.5.attn.wo_b": {
|
| 982 |
+
"bits": 6
|
| 983 |
+
},
|
| 984 |
+
"language_model.model.layers.5.attn.wq_b": {
|
| 985 |
+
"bits": 6
|
| 986 |
+
},
|
| 987 |
+
"language_model.model.layers.5.ffn.shared_experts.down_proj": {
|
| 988 |
+
"bits": 6
|
| 989 |
+
},
|
| 990 |
+
"language_model.model.layers.5.ffn.shared_experts.gate_up_proj": {
|
| 991 |
+
"bits": 6
|
| 992 |
+
},
|
| 993 |
+
"language_model.model.layers.6.attn.compressor.fused_wkv_wgate": {
|
| 994 |
+
"bits": 6
|
| 995 |
+
},
|
| 996 |
+
"language_model.model.layers.6.attn.fused_wqa_wkv": {
|
| 997 |
+
"bits": 6
|
| 998 |
+
},
|
| 999 |
+
"language_model.model.layers.6.attn.indexer.compressor.fused_wkv_wgate": {
|
| 1000 |
+
"bits": 6
|
| 1001 |
+
},
|
| 1002 |
+
"language_model.model.layers.6.attn.indexer.wq_b": {
|
| 1003 |
+
"bits": 6
|
| 1004 |
+
},
|
| 1005 |
+
"language_model.model.layers.6.attn.wo_a": {
|
| 1006 |
+
"bits": 6
|
| 1007 |
+
},
|
| 1008 |
+
"language_model.model.layers.6.attn.wo_b": {
|
| 1009 |
+
"bits": 6
|
| 1010 |
+
},
|
| 1011 |
+
"language_model.model.layers.6.attn.wq_b": {
|
| 1012 |
+
"bits": 6
|
| 1013 |
+
},
|
| 1014 |
+
"language_model.model.layers.6.ffn.shared_experts.down_proj": {
|
| 1015 |
+
"bits": 6
|
| 1016 |
+
},
|
| 1017 |
+
"language_model.model.layers.6.ffn.shared_experts.gate_up_proj": {
|
| 1018 |
+
"bits": 6
|
| 1019 |
+
},
|
| 1020 |
+
"language_model.model.layers.7.attn.compressor.fused_wkv_wgate": {
|
| 1021 |
+
"bits": 6
|
| 1022 |
+
},
|
| 1023 |
+
"language_model.model.layers.7.attn.fused_wqa_wkv": {
|
| 1024 |
+
"bits": 6
|
| 1025 |
+
},
|
| 1026 |
+
"language_model.model.layers.7.attn.wo_a": {
|
| 1027 |
+
"bits": 6
|
| 1028 |
+
},
|
| 1029 |
+
"language_model.model.layers.7.attn.wo_b": {
|
| 1030 |
+
"bits": 6
|
| 1031 |
+
},
|
| 1032 |
+
"language_model.model.layers.7.attn.wq_b": {
|
| 1033 |
+
"bits": 6
|
| 1034 |
+
},
|
| 1035 |
+
"language_model.model.layers.7.ffn.shared_experts.down_proj": {
|
| 1036 |
+
"bits": 6
|
| 1037 |
+
},
|
| 1038 |
+
"language_model.model.layers.7.ffn.shared_experts.gate_up_proj": {
|
| 1039 |
+
"bits": 6
|
| 1040 |
+
},
|
| 1041 |
+
"language_model.model.layers.8.attn.compressor.fused_wkv_wgate": {
|
| 1042 |
+
"bits": 6
|
| 1043 |
+
},
|
| 1044 |
+
"language_model.model.layers.8.attn.fused_wqa_wkv": {
|
| 1045 |
+
"bits": 6
|
| 1046 |
+
},
|
| 1047 |
+
"language_model.model.layers.8.attn.indexer.compressor.fused_wkv_wgate": {
|
| 1048 |
+
"bits": 6
|
| 1049 |
+
},
|
| 1050 |
+
"language_model.model.layers.8.attn.indexer.wq_b": {
|
| 1051 |
+
"bits": 6
|
| 1052 |
+
},
|
| 1053 |
+
"language_model.model.layers.8.attn.wo_a": {
|
| 1054 |
+
"bits": 6
|
| 1055 |
+
},
|
| 1056 |
+
"language_model.model.layers.8.attn.wo_b": {
|
| 1057 |
+
"bits": 6
|
| 1058 |
+
},
|
| 1059 |
+
"language_model.model.layers.8.attn.wq_b": {
|
| 1060 |
+
"bits": 6
|
| 1061 |
+
},
|
| 1062 |
+
"language_model.model.layers.8.ffn.shared_experts.down_proj": {
|
| 1063 |
+
"bits": 6
|
| 1064 |
+
},
|
| 1065 |
+
"language_model.model.layers.8.ffn.shared_experts.gate_up_proj": {
|
| 1066 |
+
"bits": 6
|
| 1067 |
+
},
|
| 1068 |
+
"language_model.model.layers.9.attn.compressor.fused_wkv_wgate": {
|
| 1069 |
+
"bits": 6
|
| 1070 |
+
},
|
| 1071 |
+
"language_model.model.layers.9.attn.fused_wqa_wkv": {
|
| 1072 |
+
"bits": 6
|
| 1073 |
+
},
|
| 1074 |
+
"language_model.model.layers.9.attn.wo_a": {
|
| 1075 |
+
"bits": 6
|
| 1076 |
+
},
|
| 1077 |
+
"language_model.model.layers.9.attn.wo_b": {
|
| 1078 |
+
"bits": 6
|
| 1079 |
+
},
|
| 1080 |
+
"language_model.model.layers.9.attn.wq_b": {
|
| 1081 |
+
"bits": 6
|
| 1082 |
+
},
|
| 1083 |
+
"language_model.model.layers.9.ffn.shared_experts.down_proj": {
|
| 1084 |
+
"bits": 6
|
| 1085 |
+
},
|
| 1086 |
+
"language_model.model.layers.9.ffn.shared_experts.gate_up_proj": {
|
| 1087 |
+
"bits": 6
|
| 1088 |
+
}
|
| 1089 |
+
}
|
| 1090 |
+
}
|
| 1091 |
+
},
|
| 1092 |
+
"rms_norm_eps": 1e-20,
|
| 1093 |
+
"rope_scaling": {
|
| 1094 |
+
"beta_fast": 32,
|
| 1095 |
+
"beta_slow": 1,
|
| 1096 |
+
"factor": 16,
|
| 1097 |
+
"original_max_position_embeddings": 65536,
|
| 1098 |
+
"type": "yarn"
|
| 1099 |
+
},
|
| 1100 |
+
"rope_theta": 10000,
|
| 1101 |
+
"routed_scaling_factor": 1.5,
|
| 1102 |
+
"scoring_func": "sqrtsoftplus",
|
| 1103 |
+
"sliding_window": 128,
|
| 1104 |
+
"swiglu_limit": 10.0,
|
| 1105 |
+
"tie_word_embeddings": false,
|
| 1106 |
+
"topk_method": "noaux_tc",
|
| 1107 |
+
"torch_dtype": "bfloat16",
|
| 1108 |
+
"transformers_version": "5.0.0",
|
| 1109 |
+
"use_cache": true,
|
| 1110 |
+
"vocab_size": 129280,
|
| 1111 |
+
"compress_rope_theta": 160000,
|
| 1112 |
+
"compress_ratios": [
|
| 1113 |
+
0,
|
| 1114 |
+
0,
|
| 1115 |
+
4,
|
| 1116 |
+
128,
|
| 1117 |
+
4,
|
| 1118 |
+
128,
|
| 1119 |
+
4,
|
| 1120 |
+
128,
|
| 1121 |
+
4,
|
| 1122 |
+
128,
|
| 1123 |
+
4,
|
| 1124 |
+
128,
|
| 1125 |
+
4,
|
| 1126 |
+
128,
|
| 1127 |
+
4,
|
| 1128 |
+
128,
|
| 1129 |
+
4,
|
| 1130 |
+
128,
|
| 1131 |
+
4,
|
| 1132 |
+
128,
|
| 1133 |
+
4,
|
| 1134 |
+
128,
|
| 1135 |
+
4,
|
| 1136 |
+
128,
|
| 1137 |
+
4,
|
| 1138 |
+
128,
|
| 1139 |
+
4,
|
| 1140 |
+
128,
|
| 1141 |
+
4,
|
| 1142 |
+
128,
|
| 1143 |
+
4,
|
| 1144 |
+
128,
|
| 1145 |
+
4,
|
| 1146 |
+
128,
|
| 1147 |
+
4,
|
| 1148 |
+
128,
|
| 1149 |
+
4,
|
| 1150 |
+
128,
|
| 1151 |
+
4,
|
| 1152 |
+
128,
|
| 1153 |
+
4,
|
| 1154 |
+
128,
|
| 1155 |
+
4,
|
| 1156 |
+
0,
|
| 1157 |
+
0,
|
| 1158 |
+
0
|
| 1159 |
+
],
|
| 1160 |
+
"dspark_block_size": 5,
|
| 1161 |
+
"dspark_noise_token_id": 128799,
|
| 1162 |
+
"dspark_target_layer_ids": [
|
| 1163 |
+
40,
|
| 1164 |
+
41,
|
| 1165 |
+
42
|
| 1166 |
+
],
|
| 1167 |
+
"dspark_markov_rank": 256,
|
| 1168 |
+
"vision_n_layers": 32,
|
| 1169 |
+
"vision_dim": 1024,
|
| 1170 |
+
"vision_n_heads": 16,
|
| 1171 |
+
"vision_inter_dim": 2816,
|
| 1172 |
+
"vision_patch_size": 14,
|
| 1173 |
+
"vision_rope_theta": 10000.0,
|
| 1174 |
+
"vision_downsample_ratio": 3,
|
| 1175 |
+
"vision_max_n_token": 384,
|
| 1176 |
+
"vision_min_pixels": 147456,
|
| 1177 |
+
"vision_max_wh_ratio": 8
|
| 1178 |
+
}
|
generation_config.json
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_from_model_config": true,
|
| 3 |
+
"bos_token_id": 0,
|
| 4 |
+
"eos_token_id": 1,
|
| 5 |
+
"do_sample": true,
|
| 6 |
+
"temperature": 1.0,
|
| 7 |
+
"top_p": 1.0,
|
| 8 |
+
"transformers_version": "4.46.3"
|
| 9 |
+
}
|
lna-lab/Dockerfile
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# lna-lab/vllm-exl3:dsv4 — vLLM nightly dev337 (vision class) + exllamav3 1.4.5 (compiled ext, sm_120) + vllm-exl3 + recipe patches
|
| 2 |
+
FROM vllm/vllm-openai:nightly
|
| 3 |
+
ENV TORCH_CUDA_ARCH_LIST="12.0" MAX_JOBS=32 CPATH=/usr/local/lib/python3.12/dist-packages/nvidia/cu13/include LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib
|
| 4 |
+
RUN pip install --no-deps "https://wheels.vllm.ai/27a94d1ce4e3fc100c4732439ccec10f8246a804/vllm-0.28.1rc1.dev337%2Bg27a94d1ce-cp38-abi3-manylinux_2_28_x86_64.whl" 2>&1 | tail -2 \
|
| 5 |
+
&& python3 -c "import vllm,os;print(vllm.__version__); print(os.listdir(os.path.join(os.path.dirname(vllm.__file__),'models','deepseek_v4','nvidia')))"
|
| 6 |
+
COPY exllamav3-src /opt/exllamav3
|
| 7 |
+
RUN pip install ninja 2>&1 | tail -1 && pip install --no-build-isolation --no-deps -v /opt/exllamav3 2>&1 | grep -i -E "error|Successfully" | tail -5 \
|
| 8 |
+
&& python3 -c "import torch, exllamav3_ext; print('exllamav3_ext OK')"
|
| 9 |
+
COPY vllm-exl3 /opt/vllm-exl3
|
| 10 |
+
RUN pip install --no-deps /opt/vllm-exl3
|
| 11 |
+
COPY recipe/scripts /opt/recipe
|
| 12 |
+
COPY recipe-lna /opt/recipe-lna
|
| 13 |
+
RUN python3 /opt/recipe/patch_dsv4_stock028.py && python3 /opt/recipe/patch_dsv4_vl_stream_load.py && python3 /opt/recipe/patch_dsv4_vl_sm120_wide_swa.py \
|
| 14 |
+
&& python3 /opt/recipe-lna/patch_dsv4_dense_exl3.py && python3 /opt/recipe-lna/patch_dsv4_attention_compressor_exl3.py && python3 /opt/recipe-lna/patch_dsv4_aux_streams_env.py
|
lna-lab/PLAN.md
ADDED
|
@@ -0,0 +1,74 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 鉱脈 dsv4-carve — DSV4-Flash-Vision EXL3 MixedK を TP=8 (16GB×8) で DSpark3 + 長文脈が同居する形に彫る
|
| 2 |
+
総監督: ユキ(雲)。開戦 2026-09-03。ケン「DSpark3 が入る隙間を削りで作る。-c128 ではなく(笑)」。
|
| 3 |
+
職人: GLM(Z.ai)/Luna(codex gpt-5.6-luna xhigh)。
|
| 4 |
+
|
| 5 |
+
## 北極星
|
| 6 |
+
TP=8 **単流**(ケン 09-03 18:10「単流は目をつぶる」)で **DSpark3 が乗り、文脈 128k**(ケン決定)、decode は素を割らない。
|
| 7 |
+
画像(vision)は落とさない。無検閲版(ablit)は素で型が固まってから同じ型で。
|
| 8 |
+
|
| 9 |
+
## 実測の基点(2026-09-03、像 lna-lab/vllm-exl3:dsv4 = vLLM dev337 + exllamav3 1.4.5 + vllm-exl3 0.2.3)
|
| 10 |
+
- 札 15.48 GiB、util 0.96 → 14.86 使用可。重み+非torch 13.18、活性の山 0.73(BT 512)、graph 0.04、KV 1.12 = 69k tok(32k×2.1)
|
| 11 |
+
- decode 単流 en/ja/code 40 tok/s(CUDA graph sizes [1,2,4])。eager 12-16。プリフィル 1.3k tok/s。画像 OK。
|
| 12 |
+
- draft 3 層を足すと 15.2/枚 → KV 0 → OOM。NCCL_BUFFSIZE 絞りは効かない(13.18 のまま)。
|
| 13 |
+
- ディスク家計簿(/枚): routed 8.67 / attention BF16 1.18 / draft experts fp8 1.20 / vision 0.87(複製疑い) / shared 0.25 / embed 0.12 / head 0.12 / draft他 0.12
|
| 14 |
+
|
| 15 |
+
## 職人と発注
|
| 16 |
+
- T2 draft overlay(Luna/codex): DATA1/vllm-exl3-lab/orders/T2-mtp-overlay.md。出所 = wrldsuksgo2mars K2.2-D2(draft が EXL3 K2、shard 2/11 に集中、≈2.4GB を range-read)
|
| 17 |
+
- T3 attention 自家焼き(GLM、第二段): 元 checkpoint の置き場が先
|
| 18 |
+
|
| 19 |
+
## 必要な隙間
|
| 20 |
+
64k+draft: 2.9 GiB/128k+draft: 4.7 GiB(KV fp8 ≈ 28 KB/token、1.82 GiB/64k)
|
| 21 |
+
|
| 22 |
+
## 彫る順(効き × 易しさ)
|
| 23 |
+
1. ~~vision tower を TP 分割~~ → dev337 で既に分割済(QKV/Row parallel)。家計簿の 0.87 は誤り、実は ≈0.11/枚。削れない
|
| 24 |
+
2. **draft experts fp8 → EXL3 2bit**(−0.9)— 出所: 自家焼き(exllamav3 convert -mb 2)or 外の EXL3 pack から overlay。plugin は mtp_experts="exl3" 既定で読める
|
| 25 |
+
3. **attention/shared BF16 → EXL3 4bit overlay**(−1.0)— 世界に EXL3 版 attention の pack は無い(D2 も BF16)→ **自家焼き**が要る(元 checkpoint ≈300GB、置き場: DATA1 115G/stripe 101G では足りず→NAS 経由か掃除)。exllamav3 の計測パスを非 expert 線形だけに絞る道具が要る。第二段。⚠ bf16_shards は TP>1 不可 → 融合線形は全 shard EXL3 に
|
| 26 |
+
4. 活性 BT 512→384(−0.2)※画像 1 枚 387 token が下限
|
| 27 |
+
5. 3bit 層 6 つ → 2bit(−0.4)— 最後。ppl の門つき
|
| 28 |
+
6. 1トークンの家計簿(段ごとの issue_ms)で decode 40 を守る/伸ばす
|
| 29 |
+
|
| 30 |
+
## 門
|
| 31 |
+
- 各段: 正気三言語 + 画像テストカード + decode ≥ 38 tok/s + KV トークン数(ログ行)を記録。
|
| 32 |
+
- 投機: spec-OFF と同一 greedy 出力 + 受理長(/metrics)。[[speculation-gate-doctrine]]
|
| 33 |
+
- 化粧合格禁止。実測と ESTIMATE を書き分ける。
|
| 34 |
+
|
| 35 |
+
## 台本
|
| 36 |
+
DATA1/vllm-exl3-lab/{Dockerfile, serve-dsv4-tp8.sh, bench-dsv4.py, recipe/, vllm-exl3/}。口 :8899。GPU 0,1,2,3,5,7,8,9。GPU4/6 は触らない。
|
| 37 |
+
|
| 38 |
+
## 09-03 18:10 中間実測
|
| 39 |
+
- 🎖️ DSpark3 が TP8 に乗った(D2 overlay + VLLM_SPARSE_INDEXER_MAX_LOGITS_MB=128、16k): en 67/ja 55/code 70 tok/s、素 45。受理長 ≈2.1。画像 OK。
|
| 40 |
+
- 128k への残り: KV ≈2.1 GiB 必要、空き 0.51 → **+1.6 GiB/枚**。削り代: T3 3bit→K2 −0.4/attention EXL3 自家焼き −0.85〜1.0/BT・util −0.3/(未確認) fp4 KV で KV 半減。
|
| 41 |
+
- 決定性: この建て付けは greedy でも走行間不一致 → ビット一致は門に使えない。品質は ppl/実食。
|
| 42 |
+
|
| 43 |
+
## 一週間の段取り(ケン 09-03 18:15「一週間かけてもいい。精度を保ったまま実用域へ」「コードを読んで無駄時間を詰めるスタイル」)
|
| 44 |
+
北極星: **TP=8 単流 DSpark3 付き 128k、精度は素の pack と同等(ppl 門)、decode ≥ 素**。
|
| 45 |
+
- **段 A(〜09-04): 32k**。T3 layer_overlay で 3bit 層 6 つ → K2(−0.4 GiB)。donor shard 7 本を DATA1/.tmp/k2v1-donor に落として local 経路で組む。門: ppl(vLLM prompt_logprobs で自前計器を作る、同一文・同一 ctx で pack 間比較)+三言語正気+画像。
|
| 46 |
+
- **段 B(09-05〜07): attention/shared の EXL3 化(−0.85〜1.0)→ 64k〜128k**。世界に無いので自家焼き。元 = NAS models-cold/DeepSeek-V4-Flash-Vision-Exp(157G)。置き場は掃除待ち(ケン決裁)。方式は exllamav3 の量子化器(quantize_linears_single)を非 expert 線形だけに当て、前段の forward は D2 pack の EXL3 experts で回す「部分焼き」。⚠ plugin の non_routed_exl3 は TP>1 で bf16_shards 不可 → 融合線形は全 shard EXL3。lm_head は plugin 未対応(0.12 だけなので後回し)。
|
| 47 |
+
- **段 C(並行): 1 トークンの家計簿**。vLLM の torch profiler で段ごとの時間(attention/indexer/MoE/all-reduce/lm_head/draft)を取り、CUDA graph の外に漏れている起動律速を潰す。目安: 素 45 → 50+、DSpark3 70 → 80+。
|
| 48 |
+
- **段 D: 128k 実走**。長文プロンプト(100k)で TTFT・decode・受理長。KV 2.1 GiB の確保を確認。
|
| 49 |
+
- 撤退線: ppl が素 pack より 2% 超悪化する削りは採らない(数字は段 A の計器で確定)。
|
| 50 |
+
- **ppl 計器(09-03 18:50 確定)**: `DATA1/vllm-exl3-lab/ppl-vllm.py <port> wiki.test.raw 512 16`、計測席 = draft 無し pack・spec OFF・util 0.95・max-len 4096・BT 1024(prompt_logprobs は語彙 129k の logits を rank0 に要求するので満杯席では OOM)。**基準 = 元 pack 6.6271**(7,215 tok)。撤退線 2% → 6.76。
|
| 51 |
+
- **家計簿(09-03 19:00、torch profiler、rank0、DSpark3 付き 64tok ≈29 手)**: 密 BF16 GEMM **51%**(cuBLAS が sm120 核を持たず sm80 WMMA 16×16 核に落ちる、11,686 発 × 43-50µs、一手 ≈21.8 ms、帯域計算なら ≈3 ms)/EXL3 expert GEMM 29.5%(exl3_moe_kernel 282µs × 1120)/NCCL 7.3%/indexer・norm・attention 各 ≤1%。→ **段 B(attention EXL3 化)は容量と速度の両方の本丸**。計器: prof-summary.py、席は PROFILE='{"profiler":"torch","torch_profiler_dir":"/lab/prof"}' で立て /start_profile /stop_profile。
|
| 52 |
+
- **段 A 実測(09-03 19:25)**: K2×6(3,13,21,22,28,41 を K2)= ppl **6.8845(+3.9%)** vs 元 6.6271 → 撤退線 2% 超、**丸ごとは不採用**。容量は 13.78→13.22 GiB/枚、KV 0.49→**1.05 GiB(64.5k tok)**、DSpark3 decode en 66/ja 58/code 69、受理長 2.2/1.8/2.3。→ 層ごとの感度掃引(t3-sens.sh、一層ずつ K2 → ppl)で痛くない層だけ採る。
|
| 53 |
+
- 道具の傷(直した): layer_overlay の shard 書き換えに (名前,meta) のリストを渡していて何も落ちていなかった(verify も見逃し)。verify に「書き換え shard に落とした名前が残っていない」を追加。教訓: **verify は「残したものが同じ」だけでなく「消したものが無い」も見る**。
|
| 54 |
+
- 段 B 設計メモ: docs/T4-dense-partial-bake-design.md(VariantSafetensorsCollection で EXL3 experts を差し、密線形だけ校正付きで焼く)。
|
| 55 |
+
- 🎖️ **段 A 着地(09-03 20:45)**: 層別感度 3/13/21/22/28/41 = +0.41/+1.06/+0.67/+0.89/+1.17/+0.66%。**K2×3(3,21,41)= ppl 6.7041(+1.16%)**、重み 13.50 GiB、KV 0.77 GiB = 47k tok → **32k + DSpark3**(en 60/ja 59/code 68、受理長 2.1/1.9/2.2)。給仕 pack = `DATA1/DSV4-Flash-Vision-EXL3-MixedK-D2-K2x3`、計測 pack = `…-MixedK-K2x3`。六層版は不採用(+3.9%)。
|
| 56 |
+
- 128k までの残り(ESTIMATE): KV 2.1 GiB 必要 vs 0.77 → +1.35 GiB。段 B の attention EXL3(−1.1、wo_a 込み)+ BT/util(−0.3)でちょうど。段 B は T4(Luna 製作中)+ plugin/vLLM 改変(docs/T5)。
|
| 57 |
+
- **段 B 着手(09-03 夜)**: 元 checkpoint を DATA1 に写し(157G)。T4 `tools/dense_bake.py`(Luna)= exllamav3 の校正付き量子化で密線形 790 本(43 層 × 15〜20)だけを 4bit EXL3 に。GPU 10,11 で焼き中(作業場 DATA1/.tmp/t4-dense-work、層ごとに ckpt、`--resume`)。私が直した傷: 層鍵の正規表現・target の順序比較・inference_mode。
|
| 58 |
+
- T6(Luna): plugin に実効 TP(disable_tp/Replicated)と rank 局所 `wo_a.slice.{rank}`、vLLM 側は recipe-lna/patch_dsv4_dense_exl3.py(compressor quant_config・_o_proj EXL3 分岐)。像 **lna-lab/vllm-exl3:dsv4-dense**(旧 :dsv4 は温存)。給仕台本は IMAGE= で切替。
|
| 59 |
+
- 教訓(今日三度目): `pgrep -f`/`pkill -f` のパターンが自分の bash -c に入る → 自殺(exit 144)。**pid で殺す**。docker build のタグは最初から別名で。
|
| 60 |
+
- 段 B 焼き(09-03 16:46〜): 8 枚並列(10 枚は peer 上限で不可 → [[exllamav3-tp-peer-limit]])、層 0→1 が 77 秒、1 層 66 MB、wo_a 4bit proxy_err ≈0.001。完了後は t4-after.sh が自動で merge(K2x3 計測 pack / D2-K2x3 給仕 pack)→ ppl → 64k/128k DSpark3 席まで回す(ログ DATA1/.tmp/t4-after.log)。
|
| 61 |
+
- **段 B 初点灯(09-03 18:30)**: 焼き 43 層 68 分(8 枚)。merge は 43 shard 書き直しで 66G(元 checkpoint は NAS へ退去して場所を作った)。**密 EXL3 で重み 13.50 → 10.56 GiB/枚(−2.9)**。vLLM 側で三つ直した: (1) config の prefix は `language_model.model.layers.N`/`ffn.shared_experts`(Vision クラス)、(2) wo_a EXL3 分岐の出力は fp32→bf16、(3) compressor の kv_score は `torch.mm(weight.T)` 直叩き → EXL3 shard 分岐(recipe-lna/patch_dsv4_attention_compressor_exl3.py)。CUDA graph + 多ストリーム有効だと warmup で無言ハング(GPU 100% spin)、eager + `VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD=0` で正気 OK。切り分けと 64k/128k を t4-after4.sh で実施中。
|
| 62 |
+
- 段 B ppl(09-03 18:29、eager): 密 4bit = **6.8789(基準 +3.8%、K2×3 比 +2.6%)** → 線越え。余裕 2.9 GiB があるので attention/shared を 6bit で���き直す(元 checkpoint を NAS から再写し)。graph ON + 多ストリーム OFF は起動するが生成で 500(要 traceback)。
|
| 63 |
+
- **128k の席が立った(09-03 19:00、密 4bit pack、graph ON・多ストリーム OFF・DSpark3)**: 重み 11.51 GiB/枚(draft 込み)、**KV 396k〜420k トークン(128k の 3 倍)**。ただし生成が「無言→EngineCore TimeoutError」で落ちる(graph ON でも eager でも。12 トークンの補完は通る、思考付きの長い生成で落ちる)。多ストリーム ON は warmup で固まる。容疑: 小バッチ経路 `LinearEXL3.forward` → `ext.BC_LinearEXL3.run_alloc`(rows ≤ 144)— TP8 の 8 プロセス×多モジュールでの静的状態か同期。exllamav3_ext を読んで特定する。
|
| 64 |
+
- 6bit 焼き直し開始(09-03 19:05、8 枚、work=.tmp/t4-dense-work-b6)。4bit の ppl は 6.8789(+3.8%)なので不採用、6bit で線内を狙う。
|
| 65 |
+
- 6bit 焼き完了(09-03 20:22、43 層 RC=0、trellis 幅 96)。merge が config に bits=4 を書いていて(--bits 未指定=既定)読み込みで形不一致 → 設定を 6 に直し、dense_bake merge は work/args.json の bits を既定にするよう修正(cd3b686)。席の連鎖は .tmp/t4-after6.sh / t4-after6.log で張り直し(20:33)。
|
| 66 |
+
- 6bit 席の二つ目の傷(09-03 20:45): 密 EXL3 の小バッチ GEMM が CUDA graph 捕捉中に**未チューンの形**で自動チューナへ入り `coop_autotune.cu:464` "operation not permitted when stream is capturing" で全 rank 死亡。4bit で通ったのは密が BF16 だった頃のキャッシュ形と重なっていたため。直し=①exl3_gemm.cu: 捕捉中は `cudaStreamIsCapturing` を見てチューンせず静的ヒューリスティックへ(mgemm も同条件)②プラグイン: 重み読込直後に行数 1,2,4,8,16 で事前実行(LNA_EXL3_PREWARM_ROWS)③像は焼き直さず .so と plugin を DATA1 から差し込む(serve の EXT_SO / PLUGIN_SRC)。a131603。
|
| 67 |
+
- ⚠️ルート盤事故(09-03 20:50): 落ちた席のコンテナ書込層が 27.6GB(コアダンプ/キャッシュ)でルート 0 バイト。docker data-root がルートに在る構造問題。対策=serve に --ulimit core=0・/root/.cache と .tilelang を DATA1 へ。**data-root の DATA1 移設はケン決裁待ち**。
|
| 68 |
+
- ★decode ハングの真因(09-03 21:40、py-spy+CUDA_LAUNCH_BLOCKING=1 で確定): CUDA_LAUNCH_BLOCKING=1 では止まらない=同時走行の衝突。vLLM FusedMoE は ≤256 トークンで **shared experts を別ストリームに重ねる**(VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD=256)。密焼きで shared experts も EXL3 協調カーネルになり、routed の mgemm と同じ DevCtx ロック領域を別ストリームから同時に叩いて永久 spin(GPU 100%)。プリフィル(1024行)で平気・decode で止まる症状と一致。対策= VLLM_DISABLE_SHARED_EXPERTS_STREAM=1(serve 既定)。
|
| 69 |
+
- 🎖️**128k 単流 DSpark3 到達(09-03 21:30)**: pack=`DSV4-Flash-Vision-EXL3-MixedK-D2-K2x3-Dense6`(68G)。KV **324,519 tok**(128k×2.48)、重み 12.0 GiB/枚。ppl **6.7159(+1.34%)**。投機なし 4k: en 42.1/ja 42.0/code 41.9。**DSpark3 128k: en 54-56 / ja 53 / code 62 tok/s**(受理 2.0/1.8/2.2)。32k BF16 密より約 1 割低い(shared 別ストリーム OFF+密 EXL3 の小バッチ経路、要 profiler)。席の起動= serve-dsv4-tp8.sh(EXT_SO/PLUGIN_SRC 差し込み・AUX_STREAMS=0・SHARED_STREAM_OFF 既定・UTIL 0.97・MAXLEN 131072・BT 512・SEQS 1・SPEC dspark 3)。Dense4 と 4bit qtensors は削除。
|
| 70 |
+
- **-c 307,200 実測(09-03 21:45)**: 同 pack・同席で MAXLEN=307200 → KV **394,570 tok**(1.28×)、短文 DSpark3 en 58.6/ja 53.1/code 62.6。**針テスト 166k tok**(実測 prompt_tokens、84 万字、針は 37% 位置): 正答 `sazanami-hotaru-2026`、TTFT **167 s**(プリフィル ≈1.0k tok/s、実測から算出)、その深さの decode **37.4 tok/s**。作文テスト 7 本(思考オフ)品質良、誤字 1(冼)。思考オンは max_tokens を大きく。Vision OK。
|
| 71 |
+
- **多流実測(09-03 21:55、-c 307,200・SEQS=4・DSpark3、bench-streams.py=usage 計数・metrics 突合一致)**: en 1流 62.4 / 2流 35×2=69 / 4流 35×4=**136** ・ ja 55 / 62 / **116** ・ code 86 / 96 / **186** tok/s。★2 流と 4 流の壁時計がほぼ同じ=2 流の一手が 4 流と同コスト(graph 捕捉サイズと EXL3 小バッチの丸め。2 流は損、4 流が甘い)。⚠️SSE チャンク数で数えると投機で過小(1流 28 に見えた)。
|
| 72 |
+
- 🎖️**制式採用(09-03 22:00 ケン決裁「俺達のショウヘイオオタニがやってきた」)**: 制式構成 = -c 307,200・SEQS 4・DSpark3。レシピは Models/RECIPES.md。残: docker data-root 移設(先)→ 常設化(systemd/DSH)→ HF+GitHub lna-lab 公開(private 先行)→ プリフィル彫刻・8 流。
|
| 73 |
+
- -c 389,120(380K)・SEQS 4 でも起動: KV **396,656 tok**(1.02×)、4 流 en 132 / code 141 合計(128 tok 短走)。制式は **-c 389,120・4 流**に更新(KV は共用プール、1 本 380k or 4 本 ≈99k ずつ)。
|
| 74 |
+
- DATA1: ds4 222G と焼き作業場 4 つ(ケン決裁)を削除、533G 空き。ds4 の履歴は DATA2/Lna-Lab/archive。
|
lna-lab/README-dsv4-carve.md
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# dsv4-carve — DeepSeek-V4-Flash-Vision (305B) EXL3 on 8 × 16 GB, 380K context, DSpark3
|
| 2 |
+
|
| 3 |
+
Lna-Lab's serving recipe for [vcruz305/DSV4-Flash-Vision-EXL3-MixedK](https://huggingface.co/vcruz305/DSV4-Flash-Vision-EXL3-MixedK)
|
| 4 |
+
carved down until a **380K-context, 4-stream, speculative (DSpark3) seat** fits on eight RTX PRO 2000 (16 GB, Blackwell sm_120)
|
| 5 |
+
under vLLM. Adopted as the house model of SAZANAMI on 2026-09-03 ("オオタニ").
|
| 6 |
+
|
| 7 |
+
Weights (private mirror, 86 GB): `sakamakismile/DSV4-Flash-Vision-EXL3-MixedK-D2-K2x3-Dense6`.
|
| 8 |
+
|
| 9 |
+
## What is in the pack
|
| 10 |
+
|
| 11 |
+
| Part | Format | Source |
|
| 12 |
+
|---|---|---|
|
| 13 |
+
| Routed experts (43 layers) | EXL3 2/3 bit mixed, layers 3/21/41 replaced by K2 | vcruz305 MixedK + wrldsuksgo2mars K2-v1 (`tools/layer_overlay.py`) |
|
| 14 |
+
| MTP draft experts (3 layers) | EXL3 K2 (was fp8) | wrldsuksgo2mars K2.2-D2-v1 (`tools/mtp_overlay.py`) |
|
| 15 |
+
| Attention + shared experts (790 dense linears) | **EXL3 6 bit** (was BF16), calibrated 250 × 2048 | baked here (`tools/dense_bake.py`) |
|
| 16 |
+
| KV cache | fp8 (≈17 KB/token incl. draft) | vLLM |
|
| 17 |
+
|
| 18 |
+
Per-GPU (TP8): weights 12.0 GiB, KV pool **396,656 tokens** at `--max-model-len 389120`.
|
| 19 |
+
|
| 20 |
+
## Measured (2026-09-03, TP8, CUDA graphs on)
|
| 21 |
+
|
| 22 |
+
| | tok/s |
|
| 23 |
+
|---|---|
|
| 24 |
+
| ppl (wikitext-2, 512 × 16) | **6.7159** (Cruz original 6.6271, +1.34 %) |
|
| 25 |
+
| single stream, no speculation | en 42.1 / ja 42.0 / code 41.9 |
|
| 26 |
+
| single stream, DSpark3 | en 62 / ja 55 / code 86 |
|
| 27 |
+
| 4 streams, DSpark3 (aggregate) | en 136 / ja 116 / code 186 |
|
| 28 |
+
| needle at 166K tokens | found; TTFT 167 s (prefill ≈ 1.0k tok/s), decode at depth 37 |
|
| 29 |
+
|
| 30 |
+
2 streams cost the same wall time as 4 (graph capture sizes + EXL3 small-batch bucketing): run 1 or 4.
|
| 31 |
+
|
| 32 |
+
## The three traps (all fixed in this repo)
|
| 33 |
+
|
| 34 |
+
1. **`dense_bake.py --merge` stamped `bits: 4` into `config.json` for a 6-bit bake** → shape mismatch at load. Merge now reads bits from `work/args.json`.
|
| 35 |
+
2. **exllamav3's cooperative-GEMM autotuner runs inside CUDA-graph capture** (`coop_autotune.cu:464`, "operation not permitted when stream is capturing"). Patch in `recipe-lna/exllamav3/exl3_gemm.cu.lna`: skip tuning while `cudaStreamIsCapturing`, fall back to the static heuristic; plus the plugin pre-tunes decode row counts (1,2,4,8,16) right after weight load (`LNA_EXL3_PREWARM_ROWS`).
|
| 36 |
+
3. **Decode deadlock, GPU 100 % on all ranks.** vLLM's FusedMoE overlaps *shared experts on a second stream* for ≤ 256 tokens. With shared experts also EXL3, two cooperative kernels share exllamav3's per-device lock buffer (`DevCtx::get_locks`) from two streams and spin forever. Diagnosed with `CUDA_LAUNCH_BLOCKING=1` (no hang → concurrency) and py-spy. Fix: **`VLLM_DISABLE_SHARED_EXPERTS_STREAM=1`** (default in `serve-dsv4-tp8.sh`).
|
| 37 |
+
|
| 38 |
+
## Serving
|
| 39 |
+
|
| 40 |
+
Image: `Dockerfile` (vllm/vllm-openai:nightly + vLLM 0.28.1rc1.dev337 wheel + exllamav3 1.4.5 built from source with the house patches + vllm-exl3 0.2.3 + `recipe-lna/patch_*.py`). The rebuilt extension and plugin are mounted over the image at run time (`EXT_SO`, `PLUGIN_SRC`) so the image never needs re-baking.
|
| 41 |
+
|
| 42 |
+
```bash
|
| 43 |
+
export EXT_SO=/path/to/exllamav3_ext.cpython-312-x86_64-linux-gnu.so # built from recipe-lna/exllamav3/*.lna
|
| 44 |
+
export PLUGIN_SRC=/path/to/vllm_exl3 # vllm-exl3 0.2.3 + recipe-lna/vllm_exl3_exl3.py.lna
|
| 45 |
+
AUX_STREAMS=0 IMAGE=dsv4-dense \
|
| 46 |
+
NCCL_EXTRA="-e VLLM_SPARSE_INDEXER_MAX_LOGITS_MB=128 -e VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION=1" \
|
| 47 |
+
MODEL=/path/to/DSV4-Flash-Vision-EXL3-MixedK-D2-K2x3-Dense6 \
|
| 48 |
+
UTIL=0.97 MAXLEN=389120 BT=512 SEQS=4 SPEC='{"method":"dspark","num_speculative_tokens":3}' \
|
| 49 |
+
bash serve-dsv4-tp8.sh # OpenAI API on 127.0.0.1:8899, model name DSV4-Flash
|
| 50 |
+
```
|
| 51 |
+
|
| 52 |
+
Notes: thinking is on by default and eats the budget (700 tokens → empty answer); send `chat_template_kwargs: {"thinking": false}` or a large `max_tokens`. Vision works. `--ulimit core=0` and cache mounts are set because a crashed seat once wrote 27 GB of core dumps into the container layer.
|
| 53 |
+
|
| 54 |
+
## Layout
|
| 55 |
+
|
| 56 |
+
- `serve-dsv4-tp8.sh` — the seat. `bench-dsv4.py` / `bench-streams.py` (usage-based token count; SSE chunk counting undercounts with speculation) / `ppl-vllm.py` / `prof-summary.py`.
|
| 57 |
+
- `tools/` — `mtp_overlay.py`, `layer_overlay.py`, `dense_bake.py` (+ `README-mtp-overlay.md`).
|
| 58 |
+
- `recipe-lna/` — vLLM/exllamav3/plugin patches applied in the image; `exllamav3/*.lna` are the patched source files.
|
| 59 |
+
- `orders/` — the work orders given to the craftspeople (codex "Luna", Explore agents). `docs/` — campaign canon (`PLAN.md`) and design notes.
|
| 60 |
+
- `exl3-tune-cache/` — autotuner disk cache for the 8 × RTX PRO 2000.
|
| 61 |
+
|
| 62 |
+
## Credits
|
| 63 |
+
|
| 64 |
+
vcruz305 (MixedK, vllm-exl3 patches), wrldsuksgo2mars (K2 packs), turboderp (exllamav3), DeepSeek (weights, MIT), vLLM.
|
| 65 |
+
Bake, patches and measurements: Lna-Lab / YUKI with Ken ([@Tono_Ken3](https://x.com/Tono_Ken3)). License for this repo: MIT; the weights follow the DeepSeek model license.
|
lna-lab/UPSTREAM-vcruz305-README.md
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: other
|
| 3 |
+
license_name: deepseek
|
| 4 |
+
license_link: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp
|
| 5 |
+
base_model: deepseek-ai/DeepSeek-V4-Flash-Vision-Exp
|
| 6 |
+
base_model_relation: quantized
|
| 7 |
+
library_name: exllamav3
|
| 8 |
+
tags: [exl3, trellis, deepseek_v4, moe, mixed-precision, dgx-spark]
|
| 9 |
+
---
|
| 10 |
+
|
| 11 |
+
# DeepSeek-V4-Flash-Vision-Exp — EXL3 MixedK
|
| 12 |
+
|
| 13 |
+
Sensitivity-allocated mixed-precision EXL3 quantization of
|
| 14 |
+
[deepseek-ai/DeepSeek-V4-Flash-Vision-Exp](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp):
|
| 15 |
+
a K2 base with the most quantization-sensitive MoE layers upgraded to K3.
|
| 16 |
+
~95 GiB — fits one DGX Spark.
|
| 17 |
+
|
| 18 |
+
## Strategy
|
| 19 |
+
Full K2 base (~2.0 bpw experts) with **K3 (~3.0 bpw) on layers 21, 22, 13,
|
| 20 |
+
41, 28, 3**, selected by a per-layer K2-vs-K3 proxy-error scan. K is
|
| 21 |
+
self-describing per tensor via trellis shape, so the pack needs no metadata
|
| 22 |
+
changes — kernels dispatch per layer at load.
|
| 23 |
+
|
| 24 |
+
| Layer | K2->K3 gain |
|
| 25 |
+
|---|---|
|
| 26 |
+
| 21 | 0.081 |
|
| 27 |
+
| 22 | 0.081 |
|
| 28 |
+
| 13 | 0.081 |
|
| 29 |
+
| 41 | 0.081 |
|
| 30 |
+
| 28 | 0.081 |
|
| 31 |
+
| 3 | 0.081 |
|
| 32 |
+
|
| 33 |
+
The landscape is **near-flat** (~0.3% spread across all 43 layers), so this
|
| 34 |
+
top-6 is a noise-level ordering of near-ties, not a set of distinct outliers.
|
| 35 |
+
Whether mixed-K measurably beats uniform K2 is pending PPL/KLD.
|
| 36 |
+
|
| 37 |
+
## Source, format, serving
|
| 38 |
+
Identical to the [K2 base](https://huggingface.co/vcruz305/DSV4-Flash-Vision-EXL3-K2):
|
| 39 |
+
mixed-format source (MXFP4 experts + FP8 attention) dequanted to BF16 before
|
| 40 |
+
trellis (4-bit -> 2-bit double quant on experts); MTP/DSpark drafter kept at
|
| 41 |
+
source MXFP4 verbatim. vLLM EXL3 overlay, one Spark, text-only serve.
|
| 42 |
+
|
| 43 |
+
## Limitations
|
| 44 |
+
Not boot-tested; mixed-format load (EXL3 + MXFP4 drafter) unverified; quality
|
| 45 |
+
gain over K2 not yet measured (flat landscape suggests it may be small);
|
| 46 |
+
DeepSeek license (`other`). See the K2 card for full detail.
|
| 47 |
+
|
| 48 |
+
## Related
|
| 49 |
+
- [K2 base](https://huggingface.co/vcruz305/DSV4-Flash-Vision-EXL3-K2)
|
| 50 |
+
- [k3-delta + evidence](https://huggingface.co/vcruz305/dsv4-flash-vision-exl3-k3-delta)
|
| 51 |
+
- [ablit MixedK](https://huggingface.co/vcruz305/DSV4-Flash-Vision-ablit-EXL3-MixedK) — uncensored variant
|
lna-lab/bench-dsv4.py
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""tok/s + DSpark acceptance meter for the vLLM DSV4 seat. usage: bench-dsv4.py [port] [max_tokens]"""
|
| 3 |
+
import json, sys, time, urllib.request, re
|
| 4 |
+
port = sys.argv[1] if len(sys.argv) > 1 else "8899"; N = int(sys.argv[2]) if len(sys.argv) > 2 else 256
|
| 5 |
+
U = f"http://127.0.0.1:{port}"
|
| 6 |
+
P = {"en": "Explain in detail how a transformer language model generates text, step by step.",
|
| 7 |
+
"ja": "日本の四季それぞれの特徴と、季節ごとの代表的な行事や食べ物について詳しく説明してください。",
|
| 8 |
+
"code": "Write a Python class implementing an LRU cache with get/put and O(1) operations, with docstrings and a small test."}
|
| 9 |
+
def metrics():
|
| 10 |
+
try: t = urllib.request.urlopen(U + "/metrics", timeout=10).read().decode()
|
| 11 |
+
except Exception: return {}
|
| 12 |
+
out = {}
|
| 13 |
+
for k in ("vllm:spec_decode_num_accepted_tokens_total", "vllm:spec_decode_num_drafts_total", "vllm:spec_decode_num_draft_tokens_total"):
|
| 14 |
+
m = re.findall(r"^" + re.escape(k) + r"(?:\{[^}]*\})?\s+([\d.e+]+)", t, re.M)
|
| 15 |
+
if m: out[k.split("num_")[1].replace("_total","")] = sum(float(x) for x in m)
|
| 16 |
+
return out
|
| 17 |
+
for name, prompt in P.items():
|
| 18 |
+
body = {"model": "DSV4-Flash", "messages": [{"role": "user", "content": prompt}], "max_tokens": N, "temperature": 0, "stream": True, "stream_options": {"include_usage": True}}
|
| 19 |
+
m0 = metrics(); t0 = time.time(); first = None; toks = 0; text = ""
|
| 20 |
+
req = urllib.request.Request(U + "/v1/chat/completions", data=json.dumps(body).encode(), headers={"content-type": "application/json"})
|
| 21 |
+
with urllib.request.urlopen(req, timeout=600) as r:
|
| 22 |
+
for line in r:
|
| 23 |
+
line = line.decode().strip()
|
| 24 |
+
if not line.startswith("data:") or line.endswith("[DONE]"): continue
|
| 25 |
+
d = json.loads(line[5:])
|
| 26 |
+
if d.get("usage"): toks = d["usage"]["completion_tokens"]
|
| 27 |
+
ch = d.get("choices") or []
|
| 28 |
+
if ch and (ch[0]["delta"].get("content") or (ch[0]["delta"].get("reasoning_content") or ch[0]["delta"].get("reasoning"))):
|
| 29 |
+
if first is None: first = time.time()
|
| 30 |
+
text += ch[0]["delta"].get("content") or ""
|
| 31 |
+
wall = time.time() - t0; m1 = metrics()
|
| 32 |
+
acc = ""
|
| 33 |
+
if m1.get("drafts") and m0 is not None:
|
| 34 |
+
dr = m1["drafts"] - m0.get("drafts", 0); ac = m1["accepted_tokens"] - m0.get("accepted_tokens", 0)
|
| 35 |
+
if dr: acc = f" accept_len={1 + ac / dr:.2f} (accepted {ac:.0f}/{m1['draft_tokens'] - m0.get('draft_tokens', 0):.0f} drafted)"
|
| 36 |
+
print(f"{name}: tokens={toks} wall={wall:.2f}s TTFT={(first or t0) - t0:.2f}s decode={toks / (wall - ((first or t0) - t0) + 1e-9):.2f} tok/s{acc}")
|
| 37 |
+
print(" ", repr(text[:100]))
|
lna-lab/bench-streams.py
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Multi-stream decode meter for the DSV4 seat: per-stream and aggregate tok/s (streaming, wall-clock)."""
|
| 3 |
+
import json, sys, time, threading, urllib.request
|
| 4 |
+
port = sys.argv[1] if len(sys.argv) > 1 else "8899"; N = int(sys.argv[2]) if len(sys.argv) > 2 else 256
|
| 5 |
+
streams_list = [int(x) for x in (sys.argv[3] if len(sys.argv) > 3 else "1,2,4").split(",")]
|
| 6 |
+
P = {"en": "Write a detailed essay about the history of the printing press and its effects on European society.",
|
| 7 |
+
"ja": "日本の四季それぞれの風物詩と、それが文学に与えた影響について詳しく説明してください。",
|
| 8 |
+
"code": "Write a complete Python implementation of an LRU cache with tests and explanations."}
|
| 9 |
+
def one(prompt, out):
|
| 10 |
+
body = {"model": "DSV4-Flash", "messages": [{"role": "user", "content": prompt}], "max_tokens": N, "temperature": 0.0, "stream": True, "stream_options": {"include_usage": True}, "chat_template_kwargs": {"thinking": False}}
|
| 11 |
+
t0 = time.time(); first = None; n = 0
|
| 12 |
+
r = urllib.request.urlopen(urllib.request.Request(f"http://127.0.0.1:{port}/v1/chat/completions", data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}), timeout=3600)
|
| 13 |
+
for line in r:
|
| 14 |
+
line = line.decode().strip()
|
| 15 |
+
if not line.startswith("data:") or line.endswith("[DONE]"): continue
|
| 16 |
+
j = json.loads(line[5:])
|
| 17 |
+
if j.get("usage") and j["usage"].get("completion_tokens"): n = j["usage"]["completion_tokens"]; last = time.time(); continue
|
| 18 |
+
ch = j.get("choices") or []
|
| 19 |
+
d = ch[0].get("delta", {}) if ch else {}
|
| 20 |
+
if d.get("content") or d.get("reasoning_content") or d.get("reasoning"):
|
| 21 |
+
if first is None: first = time.time()
|
| 22 |
+
last = time.time()
|
| 23 |
+
out.append((first - t0, n, last - first))
|
| 24 |
+
for lang, prompt in P.items():
|
| 25 |
+
for s in streams_list:
|
| 26 |
+
res = []; th = [threading.Thread(target=one, args=(prompt, res)) for _ in range(s)]
|
| 27 |
+
t = time.time(); [x.start() for x in th]; [x.join() for x in th]; wall = time.time() - t
|
| 28 |
+
tot = sum(n for _, n, _ in res); per = [n / dt for _, n, dt in res]
|
| 29 |
+
print(f"{lang}: streams={s} tokens={tot} wall={wall:.2f}s TTFT={max(f for f,_,_ in res):.2f}s agg={tot/wall:.1f} tok/s per_stream={min(per):.1f}-{max(per):.1f}", flush=True)
|
lna-lab/exllamav3_ext.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:175c099f07950496da3f5faa40e32d78eab94228ba5a5961770f8d4605691bf4
|
| 3 |
+
size 356940600
|
lna-lab/patches/exl3.py.lna
ADDED
|
@@ -0,0 +1,389 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
import torch
|
| 3 |
+
from ...model.config import Config
|
| 4 |
+
from .exl3_lib.quantize import preapply_had_l, preapply_had_r, had_k, had_n
|
| 5 |
+
from ...ext import exllamav3_ext as ext
|
| 6 |
+
from ...util.tensor import g_tensor_cache
|
| 7 |
+
import os
|
| 8 |
+
from ...util import profile_opt
|
| 9 |
+
|
| 10 |
+
AUTO_RECONSTRUCT_THRESHOLD = 144
|
| 11 |
+
MAX_RECONSTRUCT_SLICE_N = 32768
|
| 12 |
+
RECONSTRUCT_SLICE_GRANULARITY_N = 128
|
| 13 |
+
|
| 14 |
+
no_fused_reconstruct = os.environ.get("EXL3_NO_FUSED_RECONSTRUCT", "0") != "0"
|
| 15 |
+
|
| 16 |
+
class LinearEXL3:
|
| 17 |
+
|
| 18 |
+
quant_type: str = "exl3"
|
| 19 |
+
|
| 20 |
+
def __init__(
|
| 21 |
+
self,
|
| 22 |
+
config: Config | None,
|
| 23 |
+
in_features: int,
|
| 24 |
+
out_features: int,
|
| 25 |
+
scale: torch.Tensor | None = None,
|
| 26 |
+
su: torch.Tensor | None = None,
|
| 27 |
+
sv: torch.Tensor | None = None,
|
| 28 |
+
suh: torch.Tensor | None = None,
|
| 29 |
+
svh: torch.Tensor | None = None,
|
| 30 |
+
trellis: torch.Tensor | None = None,
|
| 31 |
+
mcg: torch.Tensor | None = None,
|
| 32 |
+
mul1: torch.Tensor | None = None,
|
| 33 |
+
bias: torch.Tensor | None = None,
|
| 34 |
+
out_dtype: torch.dtype | None = None,
|
| 35 |
+
transformers_fix: bool = False,
|
| 36 |
+
key: str | None = None
|
| 37 |
+
):
|
| 38 |
+
assert scale is None, "scale is no longer used"
|
| 39 |
+
assert su is not None or suh is not None, "either su (packed) or suh (unpacked) is required"
|
| 40 |
+
assert sv is not None or svh is not None, "either sv (packed) or svh (unpacked) is required"
|
| 41 |
+
assert trellis is not None, "trellis is required"
|
| 42 |
+
if su is not None: assert su.dtype == torch.int16, "su is wrong datatype"
|
| 43 |
+
if sv is not None: assert sv.dtype == torch.int16, "sv is wrong datatype"
|
| 44 |
+
if suh is not None: assert suh.dtype == torch.half, "suh is wrong datatype"
|
| 45 |
+
if svh is not None: assert svh.dtype == torch.half, "svh is wrong datatype"
|
| 46 |
+
assert trellis.dtype == torch.int16, "trellis is wrong datatype"
|
| 47 |
+
assert len(trellis.shape) == 3, "trellis must have dim = 3"
|
| 48 |
+
|
| 49 |
+
if bias is not None and bias.dtype == torch.float: bias = bias.to(torch.half)
|
| 50 |
+
|
| 51 |
+
# Not a Module subclass, so the config-or-NullConfig default doesn't apply here; TP imports pass
|
| 52 |
+
# config=None and forward() reads config.infer_params
|
| 53 |
+
if config is None:
|
| 54 |
+
from ...model.config import NullConfig
|
| 55 |
+
config = NullConfig()
|
| 56 |
+
self.config = config
|
| 57 |
+
self.transformers_fix = transformers_fix
|
| 58 |
+
self.key = key
|
| 59 |
+
|
| 60 |
+
# self.scale = scale.item()
|
| 61 |
+
self.su = None
|
| 62 |
+
self.sv = None
|
| 63 |
+
self.suh = suh if suh is not None else self.unpack_bf(su)
|
| 64 |
+
self.svh = svh if svh is not None else self.unpack_bf(sv)
|
| 65 |
+
self.trellis = trellis
|
| 66 |
+
self.K = trellis.shape[-1] // 16
|
| 67 |
+
self.in_features = in_features
|
| 68 |
+
self.out_features = out_features
|
| 69 |
+
self.bias = bias
|
| 70 |
+
self.swap_device = None
|
| 71 |
+
self.out_dtype = out_dtype
|
| 72 |
+
self.default_out_dtype = out_dtype or torch.half
|
| 73 |
+
|
| 74 |
+
self.mcg_tensor = mcg
|
| 75 |
+
self.mul1_tensor = mul1
|
| 76 |
+
self.mcg = self.mcg_tensor is not None
|
| 77 |
+
self.mul1 = self.mul1_tensor is not None
|
| 78 |
+
|
| 79 |
+
self._fused_reconstruct = None
|
| 80 |
+
self.bsz1_xh_args = (self.trellis.device, (1, self.in_features), self.out_dtype)
|
| 81 |
+
self.bc = ext.BC_LinearEXL3(
|
| 82 |
+
self.trellis,
|
| 83 |
+
self.suh,
|
| 84 |
+
self.svh,
|
| 85 |
+
self.K,
|
| 86 |
+
self.bias,
|
| 87 |
+
self.mcg,
|
| 88 |
+
self.mul1,
|
| 89 |
+
torch.empty((1, self.in_features), dtype = self.out_dtype, device = self.trellis.device) # LNA-LAB: per-linear scratch (was a shared cache entry → races across streams)
|
| 90 |
+
)
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def unload(self):
|
| 94 |
+
# g_tensor_cache.drop(*self.bsz1_xh_args)
|
| 95 |
+
pass
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def get_tensors(self, key: str):
|
| 99 |
+
return {
|
| 100 |
+
f"{key}.{subkey}": tensor.contiguous()
|
| 101 |
+
for subkey, tensor in [
|
| 102 |
+
("su", self.su),
|
| 103 |
+
("sv", self.sv),
|
| 104 |
+
("suh", self.suh),
|
| 105 |
+
("svh", self.svh),
|
| 106 |
+
("trellis", self.trellis),
|
| 107 |
+
("bias", self.bias),
|
| 108 |
+
("mcg", self.mcg_tensor),
|
| 109 |
+
("mul1", self.mul1_tensor),
|
| 110 |
+
] if tensor is not None
|
| 111 |
+
}
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def forward(
|
| 115 |
+
self,
|
| 116 |
+
x: torch.Tensor,
|
| 117 |
+
params: dict,
|
| 118 |
+
out_dtype: torch.dtype | None = None,
|
| 119 |
+
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
| 120 |
+
|
| 121 |
+
if "ovr" in params:
|
| 122 |
+
ovr = params["ovr"]
|
| 123 |
+
if self.key in ovr and ovr[self.key].inner is not self:
|
| 124 |
+
return ovr[self.key].forward(x, params, out_dtype)
|
| 125 |
+
|
| 126 |
+
# The EXL3 kernels read x as contiguous rows; a strided view (e.g. a head-group slice
|
| 127 |
+
# of a wider tensor) would be silently misread as interleaved garbage. Producers are
|
| 128 |
+
# responsible for contiguity (a silent copy here would hide a hot-path inefficiency
|
| 129 |
+
# and break CUDA-graph address stability)
|
| 130 |
+
assert x.is_contiguous(), f"LinearEXL3 {self.key}: non-contiguous input {tuple(x.shape)}"
|
| 131 |
+
|
| 132 |
+
reconstruct = params.get("reconstruct")
|
| 133 |
+
if not reconstruct:
|
| 134 |
+
rows = x.numel() // x.shape[-1]
|
| 135 |
+
if rows <= AUTO_RECONSTRUCT_THRESHOLD or self.config.infer_params.no_reconstruct:
|
| 136 |
+
dtype = out_dtype or self.default_out_dtype
|
| 137 |
+
return self.bc.run_alloc(x, self.out_features, dtype == torch.float)
|
| 138 |
+
|
| 139 |
+
return self.reconstruct_hgemm(x, out_dtype)
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
def unpack_bf(self, bitfield: torch.Tensor):
|
| 143 |
+
# For some reason this operation causes a GPU assert on Transformers. Running on CPU seems to fix it
|
| 144 |
+
device = bitfield.device
|
| 145 |
+
if self.transformers_fix:
|
| 146 |
+
bitfield = bitfield.cpu()
|
| 147 |
+
|
| 148 |
+
# (Only used for full reconstruct and loading old models, not during inference)
|
| 149 |
+
bitfield = bitfield.view(torch.uint16).to(torch.int)
|
| 150 |
+
masks = (1 << torch.arange(16)).to(bitfield.device)
|
| 151 |
+
expanded = (bitfield.unsqueeze(-1) & masks) > 0
|
| 152 |
+
expanded = expanded.flatten()
|
| 153 |
+
# NOT torch.where with CPU scalar tensors: that path misses the device guard when the
|
| 154 |
+
# condition lives on a non-current device (observed on torch 2.11.0+cu130) — the kernel
|
| 155 |
+
# launches on the current device's context, faults there, silently zero-fills the output
|
| 156 |
+
# and leaves every other device in the process unusable. Map bool -> {-1, +1} arithmetically
|
| 157 |
+
expanded = 1.0 - expanded.to(torch.float16) * 2.0
|
| 158 |
+
return expanded.contiguous().to(device)
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
def reconstruct_hgemm(self, x: torch.Tensor, out_dtype):
|
| 162 |
+
|
| 163 |
+
shape = x.shape
|
| 164 |
+
rows = x.numel() // shape[-1]
|
| 165 |
+
out_shape = shape[:-1] + (self.out_features,)
|
| 166 |
+
x = x.view(rows, self.in_features)
|
| 167 |
+
y = torch.empty(out_shape, dtype = out_dtype or self.default_out_dtype, device = x.device)
|
| 168 |
+
|
| 169 |
+
y_ = y.view(rows, self.out_features)
|
| 170 |
+
|
| 171 |
+
# Fused path: reconstruct emits ORIGINAL-basis weights (both Hadamards + sign
|
| 172 |
+
# vectors folded into the memory-bound reconstruct kernel), so the gemm runs on the
|
| 173 |
+
# raw input and the standalone input/output had_r_128 launches disappear (~14% of
|
| 174 |
+
# long-chunk prefill GPU time). Requires 128-divisible dims (always true for EXL3
|
| 175 |
+
# tensors: both sides are had-transformed at quant time)
|
| 176 |
+
if self._fused_reconstruct is None:
|
| 177 |
+
self._fused_reconstruct = (
|
| 178 |
+
self.in_features % 128 == 0 and self.out_features % 128 == 0
|
| 179 |
+
and not no_fused_reconstruct
|
| 180 |
+
)
|
| 181 |
+
|
| 182 |
+
# The fused kernel costs ~4x plain reconstruct (k*n-proportional) while the saved
|
| 183 |
+
# had launches scale with rows*(k+n); breakeven is rows ~400-900 across shapes
|
| 184 |
+
use_fused = self._fused_reconstruct and rows >= 1024
|
| 185 |
+
|
| 186 |
+
if use_fused:
|
| 187 |
+
xh = x
|
| 188 |
+
else:
|
| 189 |
+
xh = torch.empty_like(x)
|
| 190 |
+
ext.had_r_128(x, xh, self.suh, None, 1.0)
|
| 191 |
+
|
| 192 |
+
if self.out_features <= MAX_RECONSTRUCT_SLICE_N:
|
| 193 |
+
w = torch.empty((self.in_features, self.out_features), dtype = torch.half, device = self.trellis.device)
|
| 194 |
+
if use_fused:
|
| 195 |
+
ext.reconstruct_had_slice(w, self.trellis, self.suh, self.svh, self.K, self.mcg, self.mul1, 0)
|
| 196 |
+
else:
|
| 197 |
+
ext.reconstruct(w, self.trellis, self.K, self.mcg, self.mul1)
|
| 198 |
+
ext.hgemm(xh, w, y_)
|
| 199 |
+
else:
|
| 200 |
+
numel_ = self.in_features * MAX_RECONSTRUCT_SLICE_N
|
| 201 |
+
w_ = torch.empty((numel_,), dtype = torch.half, device = self.trellis.device)
|
| 202 |
+
for n_start in range(0, self.out_features, MAX_RECONSTRUCT_SLICE_N):
|
| 203 |
+
n_end = min(n_start + MAX_RECONSTRUCT_SLICE_N, self.out_features)
|
| 204 |
+
numel = self.in_features * (n_end - n_start)
|
| 205 |
+
w = w_[:numel].view(self.in_features, n_end - n_start)
|
| 206 |
+
if use_fused:
|
| 207 |
+
ext.reconstruct_had_slice(
|
| 208 |
+
w, self.trellis, self.suh, self.svh[n_start:], self.K, self.mcg, self.mul1, n_start)
|
| 209 |
+
else:
|
| 210 |
+
ext.reconstruct_slice(w, self.trellis, self.K, self.mcg, self.mul1, n_start)
|
| 211 |
+
ext.hgemm(xh, w, y_[:, n_start:n_end])
|
| 212 |
+
|
| 213 |
+
if not use_fused:
|
| 214 |
+
ext.had_r_128(y_, y_, None, self.svh, 1.0)
|
| 215 |
+
|
| 216 |
+
if self.bias is not None:
|
| 217 |
+
y += self.bias
|
| 218 |
+
return y
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def get_inner_weight_tensor(self, n_offset: int = 0, n_features: int | None = None):
|
| 222 |
+
w = torch.empty((self.in_features, self.out_features), dtype = torch.half, device = self.trellis.device)
|
| 223 |
+
ext.reconstruct(w, self.trellis, self.K, self.mcg, self.mul1)
|
| 224 |
+
return w
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
def get_weight_tensor(self):
|
| 228 |
+
# suh = self.unpack_bf(self.su).unsqueeze(1)
|
| 229 |
+
suh = self.unpack_bf(self.su).unsqueeze(1) if self.su else self.suh.unsqueeze(1)
|
| 230 |
+
svh = self.unpack_bf(self.sv).unsqueeze(0) if self.sv else self.svh.unsqueeze(0)
|
| 231 |
+
w = self.get_inner_weight_tensor()
|
| 232 |
+
w = preapply_had_l(w, had_k)
|
| 233 |
+
w *= suh
|
| 234 |
+
w = preapply_had_r(w, had_n)
|
| 235 |
+
w *= svh
|
| 236 |
+
# w *= self.scale
|
| 237 |
+
return w
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def get_bias_tensor(self) -> torch.Tensor | None:
|
| 241 |
+
return self.bias
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
# Swap tensors to CPU (to free some space while quantizing)
|
| 245 |
+
def swap_cpu(self):
|
| 246 |
+
if self.swap_device is not None:
|
| 247 |
+
return
|
| 248 |
+
self.swap_device = self.trellis.device
|
| 249 |
+
if self.su is not None: self.su = self.su.cpu()
|
| 250 |
+
if self.sv is not None: self.sv = self.sv.cpu()
|
| 251 |
+
if self.suh is not None: self.suh = self.suh.cpu()
|
| 252 |
+
if self.svh is not None: self.svh = self.svh.cpu()
|
| 253 |
+
if self.trellis is not None: self.trellis = self.trellis.cpu()
|
| 254 |
+
if self.bias is not None: self.bias = self.bias.cpu()
|
| 255 |
+
|
| 256 |
+
|
| 257 |
+
def unswap_cpu(self):
|
| 258 |
+
if self.swap_device is None:
|
| 259 |
+
return
|
| 260 |
+
if self.su is not None: self.su = self.su.to(self.swap_device)
|
| 261 |
+
if self.sv is not None: self.sv = self.sv.to(self.swap_device)
|
| 262 |
+
if self.suh is not None: self.suh = self.suh.to(self.swap_device)
|
| 263 |
+
if self.svh is not None: self.svh = self.svh.to(self.swap_device)
|
| 264 |
+
if self.trellis is not None: self.trellis = self.trellis.to(self.swap_device)
|
| 265 |
+
if self.bias is not None: self.bias = self.bias.to(self.swap_device)
|
| 266 |
+
self.swap_device = None
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
def tp_export(self, plan, producer):
|
| 270 |
+
return {
|
| 271 |
+
"cls": LinearEXL3,
|
| 272 |
+
"in_features": self.in_features,
|
| 273 |
+
"out_features": self.out_features,
|
| 274 |
+
"suh": producer.send(self.suh),
|
| 275 |
+
"svh": producer.send(self.svh),
|
| 276 |
+
"trellis": producer.send(self.trellis),
|
| 277 |
+
"bias": producer.send(self.bias),
|
| 278 |
+
"mcg": producer.send(self.mcg_tensor),
|
| 279 |
+
"mul1": producer.send(self.mul1_tensor),
|
| 280 |
+
"out_dtype": self.out_dtype,
|
| 281 |
+
}
|
| 282 |
+
|
| 283 |
+
|
| 284 |
+
@staticmethod
|
| 285 |
+
def tp_import_split(local_context, exported, plan, split):
|
| 286 |
+
consumer = local_context["consumer"]
|
| 287 |
+
device = local_context["device"]
|
| 288 |
+
id_suh = exported["suh"]
|
| 289 |
+
id_svh = exported["svh"]
|
| 290 |
+
id_trellis = exported["trellis"]
|
| 291 |
+
id_bias = exported["bias"]
|
| 292 |
+
mcg = consumer.recv(exported["mcg"], cuda = True)
|
| 293 |
+
mul1 = consumer.recv(exported["mul1"], cuda = True)
|
| 294 |
+
|
| 295 |
+
if split is not None:
|
| 296 |
+
split_out, first, last = split
|
| 297 |
+
else:
|
| 298 |
+
split_out, first, last = True, 0, exported["out_features"]
|
| 299 |
+
|
| 300 |
+
if split_out:
|
| 301 |
+
suh = consumer.recv(id_suh, cuda = True)
|
| 302 |
+
svh = consumer.recv(id_svh, cuda = True, slice_dim = 0, first = first, last = last)
|
| 303 |
+
trellis = consumer.recv(id_trellis, cuda = True, slice_dim = 1, first = first // 16, last = last // 16)
|
| 304 |
+
bias = consumer.recv(id_bias, cuda = True, slice_dim = 0, first = first, last = last)
|
| 305 |
+
in_features = exported["in_features"]
|
| 306 |
+
out_features = last - first
|
| 307 |
+
else:
|
| 308 |
+
suh = consumer.recv(id_suh, cuda = True, slice_dim = 0, first = first, last = last)
|
| 309 |
+
svh = consumer.recv(id_svh, cuda = True)
|
| 310 |
+
trellis = consumer.recv(id_trellis, cuda = True, slice_dim = 0, first = first // 16, last = last // 16)
|
| 311 |
+
bias = consumer.recv(id_bias, cuda = True) if first == 0 else None
|
| 312 |
+
in_features = last - first
|
| 313 |
+
out_features = exported["out_features"]
|
| 314 |
+
|
| 315 |
+
module = LinearEXL3(
|
| 316 |
+
config = None,
|
| 317 |
+
in_features = in_features,
|
| 318 |
+
out_features = out_features,
|
| 319 |
+
scale = None,
|
| 320 |
+
su = None,
|
| 321 |
+
sv = None,
|
| 322 |
+
suh = suh,
|
| 323 |
+
svh = svh,
|
| 324 |
+
trellis = trellis,
|
| 325 |
+
mcg = mcg,
|
| 326 |
+
mul1 = mul1,
|
| 327 |
+
bias = bias,
|
| 328 |
+
out_dtype = exported["out_dtype"],
|
| 329 |
+
)
|
| 330 |
+
return module
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
@staticmethod
|
| 334 |
+
def tp_import_split_3(local_context, exported, plan, split_0, split_1, split_2, dbg = False):
|
| 335 |
+
return LinearEXL3.tp_import_split_n(local_context, exported, plan, [split_0, split_1, split_2], dbg)
|
| 336 |
+
|
| 337 |
+
|
| 338 |
+
@staticmethod
|
| 339 |
+
def tp_import_split_n(local_context, exported, plan, splits, dbg = False):
|
| 340 |
+
consumer = local_context["consumer"]
|
| 341 |
+
device = local_context["device"]
|
| 342 |
+
id_suh = exported["suh"]
|
| 343 |
+
id_svh = exported["svh"]
|
| 344 |
+
id_trellis = exported["trellis"]
|
| 345 |
+
id_bias = exported["bias"]
|
| 346 |
+
mcg = consumer.recv(exported["mcg"], cuda = True)
|
| 347 |
+
mul1 = consumer.recv(exported["mul1"], cuda = True)
|
| 348 |
+
|
| 349 |
+
svh_ = []
|
| 350 |
+
trellis_ = []
|
| 351 |
+
bias_ = []
|
| 352 |
+
in_features = 0
|
| 353 |
+
out_features = 0
|
| 354 |
+
|
| 355 |
+
for split in splits:
|
| 356 |
+
assert split is not None
|
| 357 |
+
split_out, first, last = split
|
| 358 |
+
assert split_out
|
| 359 |
+
|
| 360 |
+
suh = consumer.recv(id_suh, cuda = True)
|
| 361 |
+
svh = consumer.recv(id_svh, cuda = True, slice_dim = 0, first = first, last = last)
|
| 362 |
+
trellis = consumer.recv(id_trellis, cuda = True, slice_dim = 1, first = first // 16, last = last // 16)
|
| 363 |
+
bias = consumer.recv(id_bias, cuda = True, slice_dim = 0, first = first, last = last)
|
| 364 |
+
in_features = exported["in_features"]
|
| 365 |
+
out_features += last - first
|
| 366 |
+
svh_.append(svh)
|
| 367 |
+
trellis_.append(trellis)
|
| 368 |
+
bias_.append(bias)
|
| 369 |
+
|
| 370 |
+
svh = torch.cat(svh_, dim = 0)
|
| 371 |
+
trellis = torch.cat(trellis_, dim = 1)
|
| 372 |
+
bias = torch.cat(bias_, dim = 0) if bias_[0] is not None else None
|
| 373 |
+
|
| 374 |
+
module = LinearEXL3(
|
| 375 |
+
config = None,
|
| 376 |
+
in_features = in_features,
|
| 377 |
+
out_features = out_features,
|
| 378 |
+
scale = None,
|
| 379 |
+
su = None,
|
| 380 |
+
sv = None,
|
| 381 |
+
suh = suh,
|
| 382 |
+
svh = svh,
|
| 383 |
+
trellis = trellis,
|
| 384 |
+
mcg = mcg,
|
| 385 |
+
mul1 = mul1,
|
| 386 |
+
bias = bias,
|
| 387 |
+
out_dtype = exported["out_dtype"],
|
| 388 |
+
)
|
| 389 |
+
return module
|
lna-lab/patches/exl3_gemm.cu.lna
ADDED
|
@@ -0,0 +1,703 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#include <cuda_fp16.h>
|
| 2 |
+
#include "exl3_gemm.cuh"
|
| 3 |
+
|
| 4 |
+
#include <c10/cuda/CUDAGuard.h>
|
| 5 |
+
#include <ATen/cuda/CUDAContext.h>
|
| 6 |
+
#include <cooperative_groups.h>
|
| 7 |
+
namespace cg = cooperative_groups;
|
| 8 |
+
#include "../util.h"
|
| 9 |
+
#include "../util.cuh"
|
| 10 |
+
#include "exl3_gemm_kernel.cuh"
|
| 11 |
+
#include "exl3_kernel_map.cuh"
|
| 12 |
+
#include "exl3_devctx.cuh"
|
| 13 |
+
#include "exl3_gemv.cuh"
|
| 14 |
+
#include "exl3_gemv_int8.cuh"
|
| 15 |
+
#include "coop_autotune.cuh"
|
| 16 |
+
#include <atomic>
|
| 17 |
+
#include <cstdio>
|
| 18 |
+
#include <set>
|
| 19 |
+
#include <vector>
|
| 20 |
+
|
| 21 |
+
int exl3_gemm_tilesize_k_g[] = {EXL3_GEMM_TILESIZE_K};
|
| 22 |
+
int exl3_gemm_tilesize_n_g[] = {EXL3_GEMM_TILESIZE_N};
|
| 23 |
+
int exl3_gemm_blockdim_g[] = {EXL3_GEMM_BLOCKDIM};
|
| 24 |
+
|
| 25 |
+
/*
|
| 26 |
+
EXL3 matmul, A @ B -> C
|
| 27 |
+
|
| 28 |
+
- A: row-major A tensor, shape (m, k), dtype float16, contiguous
|
| 29 |
+
- B: EXL3-quantized B tensor, shape (k//16, n//16, 16*K), dtype uint16
|
| 30 |
+
- C: empty row-major C tensor, shape (m, n), dtype float16 or float32, contiguous. Does not need to be zero-initialized
|
| 31 |
+
- suh: optional, packed input scales/flips, shape (k//16), dtype float16
|
| 32 |
+
- A_had: required if suh given, may be reference to A, temporary storage for input transform, size and dtype as A
|
| 33 |
+
- svh: optional, packed output scales/flips, shape (n//16), dtype float16
|
| 34 |
+
|
| 35 |
+
limitations:
|
| 36 |
+
- k % 16 == 0
|
| 37 |
+
- n % 128 == 0
|
| 38 |
+
*/
|
| 39 |
+
|
| 40 |
+
std::set<void*> kernel_attr_set[MAX_DEVICES] = {};
|
| 41 |
+
|
| 42 |
+
uint64_t roundup_pow2(uint64_t x)
|
| 43 |
+
{
|
| 44 |
+
if (x == 0) return 1;
|
| 45 |
+
x--;
|
| 46 |
+
x |= x >> 1;
|
| 47 |
+
x |= x >> 2;
|
| 48 |
+
x |= x >> 4;
|
| 49 |
+
x |= x >> 8;
|
| 50 |
+
x |= x >> 16;
|
| 51 |
+
x |= x >> 32;
|
| 52 |
+
return x + 1;
|
| 53 |
+
}
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
// LNA-LAB: the autotuner times candidates with cudaEvent/cudaStreamSynchronize, which is
|
| 57 |
+
// forbidden while the stream is being captured into a CUDA graph. If a shape reaches here
|
| 58 |
+
// untuned during capture, fall back to the static heuristic instead of asserting.
|
| 59 |
+
static inline bool lna_stream_is_capturing(cudaStream_t stream)
|
| 60 |
+
{
|
| 61 |
+
cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone;
|
| 62 |
+
if (cudaStreamIsCapturing(stream, &st) != cudaSuccess) { (void) cudaGetLastError(); return false; }
|
| 63 |
+
return st != cudaStreamCaptureStatusNone;
|
| 64 |
+
}
|
| 65 |
+
|
| 66 |
+
uint64_t gemm_autotune_hash
|
| 67 |
+
(
|
| 68 |
+
int size_m,
|
| 69 |
+
int size_k,
|
| 70 |
+
int size_n,
|
| 71 |
+
int K,
|
| 72 |
+
bool c_fp32,
|
| 73 |
+
int device,
|
| 74 |
+
int cc,
|
| 75 |
+
int max_num_sms,
|
| 76 |
+
int cb
|
| 77 |
+
)
|
| 78 |
+
{
|
| 79 |
+
uint64_t h = 1469598103934665603ull;
|
| 80 |
+
auto mix = [&] (uint64_t v)
|
| 81 |
+
{
|
| 82 |
+
h ^= v;
|
| 83 |
+
h *= 1099511628211ull;
|
| 84 |
+
};
|
| 85 |
+
mix((uint64_t) MIN(roundup_pow2(size_m), 16));
|
| 86 |
+
mix((uint64_t) size_k);
|
| 87 |
+
mix((uint64_t) size_n);
|
| 88 |
+
mix((uint64_t) K);
|
| 89 |
+
mix(c_fp32 ? 1ull : 0ull);
|
| 90 |
+
mix((uint64_t) device);
|
| 91 |
+
mix((uint64_t) cc);
|
| 92 |
+
mix((uint64_t) max_num_sms);
|
| 93 |
+
mix((uint64_t) cb);
|
| 94 |
+
return h;
|
| 95 |
+
}
|
| 96 |
+
|
| 97 |
+
uint64_t mgemm_autotune_hash
|
| 98 |
+
(
|
| 99 |
+
int size_m,
|
| 100 |
+
int size_k,
|
| 101 |
+
int size_n,
|
| 102 |
+
int K,
|
| 103 |
+
bool c_fp32,
|
| 104 |
+
int device,
|
| 105 |
+
int cc,
|
| 106 |
+
int max_num_sms,
|
| 107 |
+
int cb,
|
| 108 |
+
int bszm_in,
|
| 109 |
+
int bszm_out
|
| 110 |
+
)
|
| 111 |
+
{
|
| 112 |
+
uint64_t h = gemm_autotune_hash(size_m, size_k, size_n, K, c_fp32, device, cc, max_num_sms, cb);
|
| 113 |
+
auto mix = [&] (uint64_t v)
|
| 114 |
+
{
|
| 115 |
+
h ^= v;
|
| 116 |
+
h *= 1099511628211ull;
|
| 117 |
+
};
|
| 118 |
+
mix((uint64_t) MIN(bszm_in, 24));
|
| 119 |
+
mix((uint64_t) MIN(bszm_out, 24));
|
| 120 |
+
return h;
|
| 121 |
+
}
|
| 122 |
+
|
| 123 |
+
int exl3_gemm_gr
|
| 124 |
+
(
|
| 125 |
+
const at::Tensor& A,
|
| 126 |
+
const at::Tensor& B,
|
| 127 |
+
at::Tensor& C,
|
| 128 |
+
const c10::optional<at::Tensor>& suh,
|
| 129 |
+
const c10::optional<at::Tensor>& A_had,
|
| 130 |
+
const c10::optional<at::Tensor>& svh,
|
| 131 |
+
int force_shape_idx,
|
| 132 |
+
bool mcg,
|
| 133 |
+
bool mul1,
|
| 134 |
+
int force_num_sms,
|
| 135 |
+
Graph* graph
|
| 136 |
+
)
|
| 137 |
+
{
|
| 138 |
+
const at::cuda::OptionalCUDAGuard device_guard(A.device());
|
| 139 |
+
cudaStream_t stream = graph ? graph->capture_stream : at::cuda::getCurrentCUDAStream().stream();
|
| 140 |
+
|
| 141 |
+
TORCH_CHECK_DIM(B, 3);
|
| 142 |
+
TORCH_CHECK_SHAPES(A, -1, B, 0, 16);
|
| 143 |
+
TORCH_CHECK_SHAPES(C, -1, B, 1, 16);
|
| 144 |
+
// TORCH_CHECK_SHAPES(A, 0, C, 0, 1);
|
| 145 |
+
TORCH_CHECK_DTYPE(A, kHalf);
|
| 146 |
+
TORCH_CHECK_DTYPE(B, kShort);
|
| 147 |
+
bool c_fp32 = C.dtype() == at::kFloat;
|
| 148 |
+
if (!c_fp32) TORCH_CHECK_DTYPE(C, kHalf);
|
| 149 |
+
|
| 150 |
+
// Get SU, optionally
|
| 151 |
+
const half* suh_ptr = (const half*) OPTPTR(suh);
|
| 152 |
+
half* A_had_ptr = nullptr;
|
| 153 |
+
if (suh_ptr)
|
| 154 |
+
{
|
| 155 |
+
// TORCH_CHECK_SHAPES(suh.value(), 0, A, 1, 1);
|
| 156 |
+
A_had_ptr = (half*) OPTPTR(A_had);
|
| 157 |
+
// TORCH_CHECK(A_had_ptr, "Must supply A_had with suh");
|
| 158 |
+
// TORCH_CHECK_SHAPES_FULL(A_had.value(), A);
|
| 159 |
+
}
|
| 160 |
+
|
| 161 |
+
// Get SV, optionally
|
| 162 |
+
const half* svh_ptr = (const half*) OPTPTR(svh);
|
| 163 |
+
// if (svh_ptr)
|
| 164 |
+
// TORCH_CHECK_SHAPES(svh.value(), 0, B, 1, 16);
|
| 165 |
+
|
| 166 |
+
// Device properties
|
| 167 |
+
int device;
|
| 168 |
+
cudaGetDevice(&device);
|
| 169 |
+
int num_sms = force_num_sms ? force_num_sms : DevCtx::instance().get_num_sms(device);
|
| 170 |
+
int cc = DevCtx::instance().get_cc(device);
|
| 171 |
+
int* locks = DevCtx::instance().get_locks(device);
|
| 172 |
+
|
| 173 |
+
// Dispatch
|
| 174 |
+
int K = B.size(2) / 16;
|
| 175 |
+
const half* A_ptr = (const half*) A.data_ptr();
|
| 176 |
+
const uint16_t* B_ptr = (const uint16_t*) B.data_ptr();
|
| 177 |
+
void* C_ptr = (void*) C.data_ptr();
|
| 178 |
+
|
| 179 |
+
int size_m = 1;
|
| 180 |
+
int dim = A.dim();
|
| 181 |
+
for (int d = 0; d < dim - 1; ++d) size_m *= A.size(d);
|
| 182 |
+
int size_k = A.size(-1);
|
| 183 |
+
int size_n = B.size(1) * 16;
|
| 184 |
+
|
| 185 |
+
// Select kernel
|
| 186 |
+
TORCH_CHECK(!(mcg && mul1), "Specified both mcg and mul1")
|
| 187 |
+
int cb = 0;
|
| 188 |
+
if (mcg) cb = 1;
|
| 189 |
+
if (mul1) cb = 2;
|
| 190 |
+
|
| 191 |
+
// Experimental fused int8-activation GEMV path (EXL3_INT8_GEMV=1) for mul1 tensors. Rows are
|
| 192 |
+
// processed as successive GEMV launches, so this is only sensible for small m (the reconstruct
|
| 193 |
+
// threshold keeps m <= 144 in practice). Not graph-capturable yet; graphed callers fall through
|
| 194 |
+
// to the regular kernel.
|
| 195 |
+
if (mul1 && exl3_gemv_int8_enabled())
|
| 196 |
+
{
|
| 197 |
+
if (exl3_gemv_int8(A, B, C, suh, A_had, svh, stream, graph))
|
| 198 |
+
return 0;
|
| 199 |
+
}
|
| 200 |
+
|
| 201 |
+
int block_dim;
|
| 202 |
+
int shape_idx;
|
| 203 |
+
fp_exl3_gemm_kernel kernel;
|
| 204 |
+
|
| 205 |
+
void* kernelArgs[] =
|
| 206 |
+
{
|
| 207 |
+
(void*)& A_ptr,
|
| 208 |
+
(void*)& B_ptr,
|
| 209 |
+
(void*)& C_ptr,
|
| 210 |
+
(void*)& size_m,
|
| 211 |
+
(void*)& size_k,
|
| 212 |
+
(void*)& size_n,
|
| 213 |
+
(void*)& locks,
|
| 214 |
+
(void*)& suh_ptr,
|
| 215 |
+
(void*)& A_had_ptr,
|
| 216 |
+
(void*)& svh_ptr
|
| 217 |
+
};
|
| 218 |
+
|
| 219 |
+
auto add_graph_args = [&](void* kernel_ptr)
|
| 220 |
+
{
|
| 221 |
+
if (graph)
|
| 222 |
+
{
|
| 223 |
+
graph->record_param(kernel_ptr, GP_gemm_A, 0);
|
| 224 |
+
graph->record_param(kernel_ptr, GP_gemm_B_trellis, 1);
|
| 225 |
+
graph->record_param(kernel_ptr, GP_gemm_C, 2);
|
| 226 |
+
graph->record_param(kernel_ptr, GP_gemm_B_suh, 7);
|
| 227 |
+
graph->record_param(kernel_ptr, GP_gemm_A_had, 8);
|
| 228 |
+
graph->record_param(kernel_ptr, GP_gemm_B_svh, 9);
|
| 229 |
+
graph->record_param(kernel_ptr, GP_end, 0);
|
| 230 |
+
}
|
| 231 |
+
};
|
| 232 |
+
|
| 233 |
+
// QTIP-style GEMV path for small m (exl3_gemv_kernel.cuh). Same kernel arguments, so graph
|
| 234 |
+
// recording is identical; falls through to the regular kernel when the heuristic declines
|
| 235 |
+
if (force_shape_idx <= 0 && force_num_sms <= 0)
|
| 236 |
+
{
|
| 237 |
+
void* gemv_kernel = nullptr;
|
| 238 |
+
if (exl3_gemv_try_launch
|
| 239 |
+
(
|
| 240 |
+
kernelArgs, size_m, size_k, size_n, K, cb, c_fp32,
|
| 241 |
+
suh_ptr && A_had_ptr && svh_ptr,
|
| 242 |
+
device, stream, &gemv_kernel, false
|
| 243 |
+
))
|
| 244 |
+
{
|
| 245 |
+
add_graph_args(gemv_kernel);
|
| 246 |
+
cuda_check(cudaPeekAtLastError());
|
| 247 |
+
return 90;
|
| 248 |
+
}
|
| 249 |
+
}
|
| 250 |
+
|
| 251 |
+
bool autotune = force_shape_idx <= 0 && force_num_sms <= 0;
|
| 252 |
+
if (autotune)
|
| 253 |
+
{
|
| 254 |
+
uint64_t autotune_key = gemm_autotune_hash(MAX(size_m, 2), size_k, size_n, K, c_fp32, device, cc, num_sms, cb);
|
| 255 |
+
CoopAutotuneLaunch tuned;
|
| 256 |
+
if (CoopKernelAutotuner::launch_locked(autotune_key, kernelArgs, SMEM_MAX, stream, &tuned))
|
| 257 |
+
{
|
| 258 |
+
add_graph_args((void*) tuned.kernel);
|
| 259 |
+
cuda_check(cudaPeekAtLastError());
|
| 260 |
+
return tuned.tag;
|
| 261 |
+
}
|
| 262 |
+
if (lna_stream_is_capturing(stream))
|
| 263 |
+
{
|
| 264 |
+
static std::atomic<int> warned{0};
|
| 265 |
+
if (warned.fetch_add(1) == 0)
|
| 266 |
+
fprintf(stderr, "[exl3_gemm] LNA-LAB: untuned shape m=%d k=%d n=%d K=%d during CUDA graph capture; using static heuristic (pre-warm to tune)\n", size_m, size_k, size_n, K);
|
| 267 |
+
autotune = false;
|
| 268 |
+
}
|
| 269 |
+
}
|
| 270 |
+
if (autotune)
|
| 271 |
+
{
|
| 272 |
+
uint64_t autotune_key = gemm_autotune_hash(MAX(size_m, 2), size_k, size_n, K, c_fp32, device, cc, num_sms, cb);
|
| 273 |
+
CoopAutotuneLaunch tuned;
|
| 274 |
+
std::vector<CoopAutotuneCandidate> candidates;
|
| 275 |
+
for (int candidate_shape_idx = 1; candidate_shape_idx <= EXL3_GEMM_NUM_SHAPES; ++candidate_shape_idx)
|
| 276 |
+
{
|
| 277 |
+
if (!exl3_gemm_shape_compat(candidate_shape_idx, size_m, size_k, size_n, K)) continue;
|
| 278 |
+
|
| 279 |
+
fp_exl3_gemm_kernel candidate_kernel = get_gemm_kernel_ptr(K, candidate_shape_idx, c_fp32, cb);
|
| 280 |
+
if (!candidate_kernel) continue;
|
| 281 |
+
|
| 282 |
+
int tilesize_k = exl3_gemm_tilesize_k_g[candidate_shape_idx];
|
| 283 |
+
int tilesize_n = exl3_gemm_tilesize_n_g[candidate_shape_idx];
|
| 284 |
+
int max_slices = MAX(size_k / tilesize_k * size_n / tilesize_n, 1);
|
| 285 |
+
int max_candidate_sms = MAX(MIN(max_slices, num_sms), 1);
|
| 286 |
+
|
| 287 |
+
candidates.push_back
|
| 288 |
+
({
|
| 289 |
+
(void*) candidate_kernel,
|
| 290 |
+
exl3_gemm_blockdim_g[candidate_shape_idx],
|
| 291 |
+
max_candidate_sms,
|
| 292 |
+
1,
|
| 293 |
+
max_candidate_sms,
|
| 294 |
+
candidate_shape_idx
|
| 295 |
+
});
|
| 296 |
+
}
|
| 297 |
+
TORCH_CHECK(!candidates.empty(), "exl3_gemm autotune: no compatible kernel shapes");
|
| 298 |
+
|
| 299 |
+
tuned = CoopKernelAutotuner::launch(autotune_key, candidates, kernelArgs, SMEM_MAX, stream, (size_t) size_k * size_n);
|
| 300 |
+
if (graph)
|
| 301 |
+
add_graph_args((void*) tuned.kernel);
|
| 302 |
+
cuda_check(cudaPeekAtLastError());
|
| 303 |
+
return tuned.tag;
|
| 304 |
+
}
|
| 305 |
+
|
| 306 |
+
kernel = select_exl3_gemm_kernel
|
| 307 |
+
(
|
| 308 |
+
cc, size_m, size_k, size_n, K, c_fp32,
|
| 309 |
+
force_shape_idx, &block_dim, &shape_idx,
|
| 310 |
+
&num_sms, cb
|
| 311 |
+
);
|
| 312 |
+
if (!kernel) return 0;
|
| 313 |
+
|
| 314 |
+
// Launch
|
| 315 |
+
if (kernel_attr_set[device].find((void*) kernel) == kernel_attr_set[device].end())
|
| 316 |
+
{
|
| 317 |
+
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_MAX);
|
| 318 |
+
kernel_attr_set[device].insert((void*) kernel);
|
| 319 |
+
cuda_check(cudaPeekAtLastError());
|
| 320 |
+
}
|
| 321 |
+
cudaLaunchCooperativeKernel
|
| 322 |
+
(
|
| 323 |
+
(void*) kernel,
|
| 324 |
+
num_sms,
|
| 325 |
+
block_dim,
|
| 326 |
+
kernelArgs,
|
| 327 |
+
SMEM_MAX,
|
| 328 |
+
stream
|
| 329 |
+
);
|
| 330 |
+
add_graph_args((void*) kernel);
|
| 331 |
+
|
| 332 |
+
cuda_check(cudaPeekAtLastError());
|
| 333 |
+
return shape_idx;
|
| 334 |
+
}
|
| 335 |
+
|
| 336 |
+
int exl3_gemm
|
| 337 |
+
(
|
| 338 |
+
const at::Tensor& A,
|
| 339 |
+
const at::Tensor& B,
|
| 340 |
+
at::Tensor& C,
|
| 341 |
+
const c10::optional<at::Tensor>& suh,
|
| 342 |
+
const c10::optional<at::Tensor>& A_had,
|
| 343 |
+
const c10::optional<at::Tensor>& svh,
|
| 344 |
+
int force_shape_idx,
|
| 345 |
+
bool mcg,
|
| 346 |
+
bool mul1,
|
| 347 |
+
int force_num_sms
|
| 348 |
+
)
|
| 349 |
+
{
|
| 350 |
+
return exl3_gemm_gr
|
| 351 |
+
(
|
| 352 |
+
A,
|
| 353 |
+
B,
|
| 354 |
+
C,
|
| 355 |
+
suh,
|
| 356 |
+
A_had,
|
| 357 |
+
svh,
|
| 358 |
+
force_shape_idx,
|
| 359 |
+
mcg,
|
| 360 |
+
mul1,
|
| 361 |
+
force_num_sms,
|
| 362 |
+
nullptr
|
| 363 |
+
);
|
| 364 |
+
}
|
| 365 |
+
|
| 366 |
+
/*
|
| 367 |
+
EXL3 batched/multi-matrix matmul.
|
| 368 |
+
|
| 369 |
+
This is not a conventional batched A @ B. B, suh and svh are CUDA int64
|
| 370 |
+
tensors containing device addresses (one address per quantized matrix), rather
|
| 371 |
+
than the matrix data themselves. Entry q of each table describes one linear:
|
| 372 |
+
|
| 373 |
+
B[q] -> EXL3 trellis, logically (k / 16, n / 16, 16 * K) uint16
|
| 374 |
+
suh[q] -> packed input scales/flips, logically (k / 16) float16
|
| 375 |
+
svh[q] -> packed output scales/flips, logically (n / 16) float16
|
| 376 |
+
|
| 377 |
+
A is contiguous float16 [a_batches, m, k], C is contiguous float16 or
|
| 378 |
+
float32 [c_batches, m, n], and A_had is float16 scratch with room for every
|
| 379 |
+
active matrix. The kernel applies the input Hadamard transform into A_had,
|
| 380 |
+
performs the selected EXL3 matmul, then applies the output transform.
|
| 381 |
+
|
| 382 |
+
The active matrix/output slot j selects q = indices[j] when indices is given,
|
| 383 |
+
or q = j otherwise. This supports the following modes:
|
| 384 |
+
|
| 385 |
+
- Multiple inputs and outputs: A[j] @ B[q] -> C[j].
|
| 386 |
+
- One input, multiple outputs: when a_batches == 1, A[0] is broadcast and
|
| 387 |
+
transformed separately for each selected B[q], producing C[j]. This is used
|
| 388 |
+
for e.g. fused gate/up projections and MoE expert fan-out.
|
| 389 |
+
- Indexed matrices: indices is a contiguous int64 [*, num_indices] tensor;
|
| 390 |
+
the kernel reads its first num_indices entries as q values. Negative indices
|
| 391 |
+
skip that slot.
|
| 392 |
+
- Weighted MoE reduction: weights is a float16 tensor parallel to indices.
|
| 393 |
+
Each transformed result is multiplied by weights[j], then all active C[j]
|
| 394 |
+
are summed into C[0]. C therefore also serves as per-expert scratch; only
|
| 395 |
+
C[0] is the reduced result.
|
| 396 |
+
- Expert-range filtering: with min_index >= 0, selections outside
|
| 397 |
+
[min_index, max_index) are removed and retained indices are rebased by
|
| 398 |
+
min_index. This allows B/suh/svh to be local pointer tables for an expert
|
| 399 |
+
shard. At num_tokens == 1 the retained indices (and their weights) are
|
| 400 |
+
compacted; at num_tokens > 1 out-of-range slots are instead masked to -1 in
|
| 401 |
+
place, preserving the per-token slot groups the final reduction depends on
|
| 402 |
+
(and, with bszm_in > 1, the slot -> input-row correspondence).
|
| 403 |
+
|
| 404 |
+
Without weights, every active C[j] is a separate output. The active slot count
|
| 405 |
+
is max(a_batches, c_batches), capped to num_indices when indices is present.
|
| 406 |
+
|
| 407 |
+
Limitations: k must be divisible by 16 and n by 128. Range filtering supports
|
| 408 |
+
at most 128 slots (the kernel's index-compaction capacity).
|
| 409 |
+
*/
|
| 410 |
+
|
| 411 |
+
int exl3_mgemm_gr
|
| 412 |
+
(
|
| 413 |
+
const at::Tensor& A,
|
| 414 |
+
const at::Tensor& B,
|
| 415 |
+
at::Tensor& C,
|
| 416 |
+
const at::Tensor& suh,
|
| 417 |
+
const at::Tensor& A_had,
|
| 418 |
+
const at::Tensor& svh,
|
| 419 |
+
const c10::optional<at::Tensor>& indices,
|
| 420 |
+
const c10::optional<at::Tensor>& weights,
|
| 421 |
+
int K,
|
| 422 |
+
int force_shape_idx,
|
| 423 |
+
bool mcg,
|
| 424 |
+
bool mul1,
|
| 425 |
+
int min_index,
|
| 426 |
+
int max_index,
|
| 427 |
+
int force_num_sms,
|
| 428 |
+
Graph* graph,
|
| 429 |
+
int num_tokens,
|
| 430 |
+
const c10::optional<at::Tensor>& size_n_list,
|
| 431 |
+
const c10::optional<at::Tensor>& c_ptrs
|
| 432 |
+
)
|
| 433 |
+
{
|
| 434 |
+
const at::cuda::OptionalCUDAGuard device_guard(A.device());
|
| 435 |
+
cudaStream_t stream = graph ? graph->capture_stream : at::cuda::getCurrentCUDAStream().stream();
|
| 436 |
+
|
| 437 |
+
// num_tokens > 1 with expert-range filtering (min_index >= 0) uses position-preserving
|
| 438 |
+
// masking in the kernel (out-of-range slots marked -1 in place) instead of index
|
| 439 |
+
// compaction, so the grouped reduction's fixed per-token slot runs stay intact
|
| 440 |
+
|
| 441 |
+
TORCH_CHECK_DTYPE(A, kHalf);
|
| 442 |
+
TORCH_CHECK_DTYPE(B, kLong);
|
| 443 |
+
TORCH_CHECK_DTYPE(suh, kLong);
|
| 444 |
+
TORCH_CHECK_DTYPE(svh, kLong);
|
| 445 |
+
bool c_fp32 = C.dtype() == at::kFloat;
|
| 446 |
+
if (!c_fp32) TORCH_CHECK_DTYPE(C, kHalf);
|
| 447 |
+
TORCH_CHECK_DIM(A, 3);
|
| 448 |
+
TORCH_CHECK_DIM(B, 1);
|
| 449 |
+
TORCH_CHECK_DIM(suh, 1);
|
| 450 |
+
TORCH_CHECK_DIM(svh, 1);
|
| 451 |
+
TORCH_CHECK_DIM(C, 3);
|
| 452 |
+
|
| 453 |
+
TORCH_CHECK_SHAPES(A, 1, C, 1, 1);
|
| 454 |
+
TORCH_CHECK_SHAPES(B, 0, suh, 0, 1);
|
| 455 |
+
TORCH_CHECK_SHAPES(B, 0, svh, 0, 1);
|
| 456 |
+
|
| 457 |
+
int bsz = A.size(1);
|
| 458 |
+
int bszm_in = A.size(0);
|
| 459 |
+
int bszm_out = C.size(0);
|
| 460 |
+
|
| 461 |
+
// Per-matrix output widths/pointers (uniform-width callers pass neither): C then only
|
| 462 |
+
// provides the dtype and the max width (locks/shape sizing); outputs go to c_ptrs
|
| 463 |
+
const int* size_n_list_ptr = nullptr;
|
| 464 |
+
void** c_list_ptr = nullptr;
|
| 465 |
+
if (size_n_list)
|
| 466 |
+
{
|
| 467 |
+
TORCH_CHECK(c_ptrs, "exl3_mgemm: size_n_list requires c_ptrs");
|
| 468 |
+
TORCH_CHECK_DTYPE(size_n_list.value(), kInt);
|
| 469 |
+
TORCH_CHECK_DTYPE(c_ptrs.value(), kLong);
|
| 470 |
+
TORCH_CHECK(num_tokens == 1 && min_index < 0 && !weights,
|
| 471 |
+
"exl3_mgemm: per-matrix widths incompatible with multi-token/filtering/weights");
|
| 472 |
+
size_n_list_ptr = (const int*) size_n_list.value().data_ptr();
|
| 473 |
+
c_list_ptr = (void**) c_ptrs.value().data_ptr();
|
| 474 |
+
bszm_out = (int) c_ptrs.value().size(0);
|
| 475 |
+
}
|
| 476 |
+
int bszm = MAX(bszm_in, bszm_out);
|
| 477 |
+
|
| 478 |
+
// The kernel writes one hadamard-transformed input slab PER MATRIX (A_had + j * m * k);
|
| 479 |
+
// an undersized scratch is silent OOB corruption (found the hard way)
|
| 480 |
+
TORCH_CHECK(A_had.numel() >= (int64_t) bszm * A.size(1) * A.size(2),
|
| 481 |
+
"exl3_mgemm: A_had must hold bszm * m * k elements");
|
| 482 |
+
|
| 483 |
+
const int64_t* indices_ptr = (const int64_t*) OPTPTR(indices);
|
| 484 |
+
const half* weights_ptr = (const half*) OPTPTR(weights);
|
| 485 |
+
|
| 486 |
+
if (indices)
|
| 487 |
+
{
|
| 488 |
+
TORCH_CHECK_DIM(indices.value(), 2);
|
| 489 |
+
int num_indices = indices.value().size(1);
|
| 490 |
+
TORCH_CHECK(num_indices <= bszm_in || num_indices <= bszm_out, "mgemm: too many indices for tensor batch");
|
| 491 |
+
if (bszm_in > num_indices) bszm_in = num_indices;
|
| 492 |
+
if (bszm_out > num_indices) bszm_out = num_indices;
|
| 493 |
+
}
|
| 494 |
+
|
| 495 |
+
if (weights)
|
| 496 |
+
{
|
| 497 |
+
TORCH_CHECK_DIM(weights.value(), 2);
|
| 498 |
+
}
|
| 499 |
+
|
| 500 |
+
int size_m = A.size(1);
|
| 501 |
+
int size_k = A.size(2);
|
| 502 |
+
int size_n = C.size(2);
|
| 503 |
+
|
| 504 |
+
// Device properties
|
| 505 |
+
int device;
|
| 506 |
+
cudaGetDevice(&device);
|
| 507 |
+
int total_sms = DevCtx::instance().get_num_sms(device);
|
| 508 |
+
int num_sms = force_num_sms ? force_num_sms : total_sms;
|
| 509 |
+
int cc = DevCtx::instance().get_cc(device);
|
| 510 |
+
int* locks = DevCtx::instance().get_locks(device);
|
| 511 |
+
|
| 512 |
+
// Dispatch
|
| 513 |
+
const half* A_ptr = (const half*) A.data_ptr();
|
| 514 |
+
const uintptr_t* B_ptr_ptr = (const uintptr_t*) B.data_ptr();
|
| 515 |
+
void* C_ptr = (void*) C.data_ptr();
|
| 516 |
+
const half* A_had_ptr = (const half*) A_had.data_ptr();
|
| 517 |
+
const uintptr_t* suh_ptr_ptr = (const uintptr_t*) suh.data_ptr();
|
| 518 |
+
const uintptr_t* svh_ptr_ptr = (const uintptr_t*) svh.data_ptr();
|
| 519 |
+
|
| 520 |
+
// Select kernel
|
| 521 |
+
TORCH_CHECK(!(mcg && mul1), "Specified both mcg and mul1")
|
| 522 |
+
int cb = 0;
|
| 523 |
+
if (mcg) cb = 1;
|
| 524 |
+
if (mul1) cb = 2;
|
| 525 |
+
|
| 526 |
+
int shape_idx;
|
| 527 |
+
int block_dim;
|
| 528 |
+
fp_exl3_mgemm_kernel kernel;
|
| 529 |
+
int concurrency;
|
| 530 |
+
|
| 531 |
+
void* kernelArgs[] =
|
| 532 |
+
{
|
| 533 |
+
(void*)& A_ptr,
|
| 534 |
+
(void*)& B_ptr_ptr,
|
| 535 |
+
(void*)& C_ptr,
|
| 536 |
+
(void*)& size_m,
|
| 537 |
+
(void*)& size_k,
|
| 538 |
+
(void*)& size_n,
|
| 539 |
+
(void*)& locks,
|
| 540 |
+
(void*)& suh_ptr_ptr,
|
| 541 |
+
(void*)& A_had_ptr,
|
| 542 |
+
(void*)& svh_ptr_ptr,
|
| 543 |
+
(void*)& indices_ptr,
|
| 544 |
+
(void*)& weights_ptr,
|
| 545 |
+
(void*)& bszm_in,
|
| 546 |
+
(void*)& bszm_out,
|
| 547 |
+
(void*)& min_index,
|
| 548 |
+
(void*)& max_index,
|
| 549 |
+
(void*)& num_tokens,
|
| 550 |
+
(void*)& size_n_list_ptr,
|
| 551 |
+
(void*)& c_list_ptr
|
| 552 |
+
};
|
| 553 |
+
|
| 554 |
+
auto add_graph_args = [&](void* kernel_ptr)
|
| 555 |
+
{
|
| 556 |
+
if (graph)
|
| 557 |
+
{
|
| 558 |
+
graph->record_param(kernel_ptr, GP_mgemm_A, 0);
|
| 559 |
+
graph->record_param(kernel_ptr, GP_mgemm_C, 2);
|
| 560 |
+
graph->record_param(kernel_ptr, GP_mgemm_indices, 10);
|
| 561 |
+
graph->record_param(kernel_ptr, GP_mgemm_weights, 11);
|
| 562 |
+
graph->record_param(kernel_ptr, GP_end, 0);
|
| 563 |
+
}
|
| 564 |
+
};
|
| 565 |
+
|
| 566 |
+
bool autotune = force_shape_idx <= 0 && force_num_sms <= 0;
|
| 567 |
+
if (autotune)
|
| 568 |
+
{
|
| 569 |
+
uint64_t autotune_key = mgemm_autotune_hash
|
| 570 |
+
(
|
| 571 |
+
size_m, size_k, size_n, K, c_fp32, device, cc, total_sms, cb, bszm_in, bszm_out
|
| 572 |
+
);
|
| 573 |
+
|
| 574 |
+
CoopAutotuneLaunch tuned;
|
| 575 |
+
if (CoopKernelAutotuner::launch_locked(autotune_key, kernelArgs, SMEM_MAX, stream, &tuned))
|
| 576 |
+
{
|
| 577 |
+
add_graph_args((void*) tuned.kernel);
|
| 578 |
+
cuda_check(cudaPeekAtLastError());
|
| 579 |
+
return tuned.tag;
|
| 580 |
+
}
|
| 581 |
+
if (!graph && !lna_stream_is_capturing(stream))
|
| 582 |
+
{
|
| 583 |
+
std::vector<CoopAutotuneCandidate> candidates;
|
| 584 |
+
for (int candidate_shape_idx = 1; candidate_shape_idx <= EXL3_GEMM_NUM_SHAPES; ++candidate_shape_idx)
|
| 585 |
+
{
|
| 586 |
+
if (!exl3_gemm_shape_compat(candidate_shape_idx, size_m, size_k, size_n, K)) continue;
|
| 587 |
+
|
| 588 |
+
fp_exl3_mgemm_kernel candidate_kernel = get_mgemm_kernel_ptr(K, candidate_shape_idx, c_fp32, cb);
|
| 589 |
+
if (!candidate_kernel) continue;
|
| 590 |
+
|
| 591 |
+
int tilesize_k = exl3_gemm_tilesize_k_g[candidate_shape_idx];
|
| 592 |
+
int tilesize_n = exl3_gemm_tilesize_n_g[candidate_shape_idx];
|
| 593 |
+
int max_slices = MAX(size_k / tilesize_k * size_n / tilesize_n, 1);
|
| 594 |
+
int max_candidate_sms = MAX(MIN(max_slices, total_sms), 1);
|
| 595 |
+
|
| 596 |
+
candidates.push_back
|
| 597 |
+
({
|
| 598 |
+
(void*) candidate_kernel,
|
| 599 |
+
exl3_gemm_blockdim_g[candidate_shape_idx],
|
| 600 |
+
max_candidate_sms,
|
| 601 |
+
bszm,
|
| 602 |
+
total_sms,
|
| 603 |
+
candidate_shape_idx
|
| 604 |
+
});
|
| 605 |
+
}
|
| 606 |
+
TORCH_CHECK(!candidates.empty(), "exl3_mgemm autotune: no compatible kernel shapes");
|
| 607 |
+
|
| 608 |
+
tuned = CoopKernelAutotuner::launch(autotune_key, candidates, kernelArgs, SMEM_MAX, stream, (size_t) size_k * size_n * bszm);
|
| 609 |
+
add_graph_args((void*) tuned.kernel);
|
| 610 |
+
|
| 611 |
+
// DBGI10(size_m, size_k, size_n, K, bszm_in, bszm_out, tuned.tag, tuned.block_dim, tuned.num_sms, tuned.concurrency);
|
| 612 |
+
|
| 613 |
+
cuda_check(cudaPeekAtLastError());
|
| 614 |
+
return tuned.tag;
|
| 615 |
+
}
|
| 616 |
+
}
|
| 617 |
+
|
| 618 |
+
kernel = select_exl3_mgemm_kernel
|
| 619 |
+
(
|
| 620 |
+
cc, size_m, size_k, size_n, K, c_fp32,
|
| 621 |
+
force_shape_idx, &block_dim, &shape_idx,
|
| 622 |
+
&num_sms, cb, bszm_in, bszm_out
|
| 623 |
+
);
|
| 624 |
+
int tilesize_k = exl3_gemm_tilesize_k_g[shape_idx];
|
| 625 |
+
int tilesize_n = exl3_gemm_tilesize_n_g[shape_idx];
|
| 626 |
+
int tiles = MAX(size_k / tilesize_k * size_n / tilesize_n, 1);
|
| 627 |
+
num_sms = tiles;
|
| 628 |
+
if (num_sms * bszm > total_sms) num_sms = MAX(total_sms / bszm, 1);
|
| 629 |
+
if (num_sms <= total_sms && tiles / num_sms > 48) num_sms = MIN(total_sms, num_sms * 2);
|
| 630 |
+
concurrency = MIN(total_sms / num_sms, bszm);
|
| 631 |
+
|
| 632 |
+
// DBGI10(size_m, size_k, size_n, K, bszm_in, bszm_out, shape_idx, block_dim, num_sms, concurrency);
|
| 633 |
+
|
| 634 |
+
// Launch bigger grid if possible
|
| 635 |
+
dim3 block_grid(num_sms, 1, concurrency);
|
| 636 |
+
|
| 637 |
+
// Launch
|
| 638 |
+
if (kernel_attr_set[device].find((void*) kernel) == kernel_attr_set[device].end())
|
| 639 |
+
{
|
| 640 |
+
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_MAX);
|
| 641 |
+
kernel_attr_set[device].insert((void*) kernel);
|
| 642 |
+
}
|
| 643 |
+
|
| 644 |
+
cudaLaunchCooperativeKernel
|
| 645 |
+
(
|
| 646 |
+
(void*) kernel,
|
| 647 |
+
block_grid,
|
| 648 |
+
block_dim,
|
| 649 |
+
kernelArgs,
|
| 650 |
+
SMEM_MAX,
|
| 651 |
+
stream
|
| 652 |
+
);
|
| 653 |
+
add_graph_args((void*) kernel);
|
| 654 |
+
|
| 655 |
+
cuda_check(cudaPeekAtLastError());
|
| 656 |
+
return shape_idx;
|
| 657 |
+
}
|
| 658 |
+
|
| 659 |
+
int exl3_mgemm
|
| 660 |
+
(
|
| 661 |
+
const at::Tensor& A,
|
| 662 |
+
const at::Tensor& B,
|
| 663 |
+
at::Tensor& C,
|
| 664 |
+
const at::Tensor& suh,
|
| 665 |
+
const at::Tensor& A_had,
|
| 666 |
+
const at::Tensor& svh,
|
| 667 |
+
const c10::optional<at::Tensor>& indices,
|
| 668 |
+
const c10::optional<at::Tensor>& weights,
|
| 669 |
+
int K,
|
| 670 |
+
int force_shape_idx,
|
| 671 |
+
uint32_t mcg_mult,
|
| 672 |
+
uint32_t mul1_mult,
|
| 673 |
+
int min_index,
|
| 674 |
+
int max_index,
|
| 675 |
+
int force_num_sms,
|
| 676 |
+
int num_tokens,
|
| 677 |
+
const c10::optional<at::Tensor>& size_n_list,
|
| 678 |
+
const c10::optional<at::Tensor>& c_ptrs
|
| 679 |
+
)
|
| 680 |
+
{
|
| 681 |
+
return exl3_mgemm_gr
|
| 682 |
+
(
|
| 683 |
+
A,
|
| 684 |
+
B,
|
| 685 |
+
C,
|
| 686 |
+
suh,
|
| 687 |
+
A_had,
|
| 688 |
+
svh,
|
| 689 |
+
indices,
|
| 690 |
+
weights,
|
| 691 |
+
K,
|
| 692 |
+
force_shape_idx,
|
| 693 |
+
mcg_mult,
|
| 694 |
+
mul1_mult,
|
| 695 |
+
min_index,
|
| 696 |
+
max_index,
|
| 697 |
+
force_num_sms,
|
| 698 |
+
nullptr,
|
| 699 |
+
num_tokens,
|
| 700 |
+
size_n_list,
|
| 701 |
+
c_ptrs
|
| 702 |
+
);
|
| 703 |
+
}
|
lna-lab/patches/patch_dsv4_attention_compressor_exl3.py
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""LNA-LAB: DSV4 attention.py reads compressor.fused_wkv_wgate.weight directly (torch.mm). When the
|
| 3 |
+
module was built by vllm-exl3 (no .weight, has _exl3_linears), run the EXL3 shards instead.
|
| 4 |
+
Idempotent, exact anchors, backup attention.py.orig-lna2."""
|
| 5 |
+
import shutil, sys
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
HELPER = '''
|
| 9 |
+
def _lna_kv_score(module, hidden_states):
|
| 10 |
+
"""LNA-LAB: fused_wkv_wgate as bf16 weight (torch.mm) or as EXL3 shards (vllm-exl3)."""
|
| 11 |
+
linears = getattr(module, "_exl3_linears", None)
|
| 12 |
+
if linears:
|
| 13 |
+
x = hidden_states.reshape(-1, hidden_states.shape[-1]).to(torch.float16).contiguous()
|
| 14 |
+
outs = [lin.forward(x, {}, out_dtype=torch.float32) for lin in linears]
|
| 15 |
+
out = outs[0] if len(outs) == 1 else torch.cat(outs, dim=-1)
|
| 16 |
+
return out.reshape(*hidden_states.shape[:-1], out.shape[-1])
|
| 17 |
+
return torch.mm(hidden_states, module.fused_wkv_wgate.weight.T, out_dtype=torch.float32)
|
| 18 |
+
|
| 19 |
+
'''
|
| 20 |
+
EDITS = [
|
| 21 |
+
(''' return torch.mm(
|
| 22 |
+
hidden_states,
|
| 23 |
+
compressor.fused_wkv_wgate.weight.T,
|
| 24 |
+
out_dtype=torch.float32,
|
| 25 |
+
)
|
| 26 |
+
''', ''' return _lna_kv_score(compressor.fused_wkv_wgate, hidden_states) # LNA-LAB
|
| 27 |
+
'''),
|
| 28 |
+
(''' return torch.mm(
|
| 29 |
+
hidden_states,
|
| 30 |
+
indexer.compressor.fused_wkv_wgate.weight.T,
|
| 31 |
+
out_dtype=torch.float32,
|
| 32 |
+
)
|
| 33 |
+
''', ''' return _lna_kv_score(indexer.compressor.fused_wkv_wgate, hidden_states) # LNA-LAB
|
| 34 |
+
'''),
|
| 35 |
+
]
|
| 36 |
+
|
| 37 |
+
def main():
|
| 38 |
+
root = Path(sys.argv[1]) if len(sys.argv) > 1 else None
|
| 39 |
+
if root is None:
|
| 40 |
+
import vllm
|
| 41 |
+
root = Path(vllm.__file__).resolve().parent
|
| 42 |
+
f = root / "models" / "deepseek_v4" / "attention.py"
|
| 43 |
+
text = f.read_text()
|
| 44 |
+
if "_lna_kv_score" in text:
|
| 45 |
+
print(f"{f}: already patched"); return 0
|
| 46 |
+
for old, new in EDITS:
|
| 47 |
+
if text.count(old) != 1:
|
| 48 |
+
raise SystemExit(f"{f}: expected 1 anchor, found {text.count(old)}: {old[:60]!r}")
|
| 49 |
+
text = text.replace(old, new)
|
| 50 |
+
# helper: insert before the first class definition
|
| 51 |
+
idx = text.index("\nclass ")
|
| 52 |
+
text = text[:idx] + "\n" + HELPER + text[idx:]
|
| 53 |
+
# the helper only reads .weight in the fallback path; keep it after fused_wkv_wgate module refs
|
| 54 |
+
text = text.replace("return torch.mm(hidden_states, module.fused_wkv_wgate.weight.T, out_dtype=torch.float32)",
|
| 55 |
+
"return torch.mm(hidden_states, module.weight.T, out_dtype=torch.float32)")
|
| 56 |
+
bak = f.with_suffix(f.suffix + ".orig-lna2")
|
| 57 |
+
if not bak.exists(): shutil.copy2(f, bak)
|
| 58 |
+
f.write_text(text); compile(text, str(f), "exec"); print(f"{f}: patched")
|
| 59 |
+
return 0
|
| 60 |
+
|
| 61 |
+
if __name__ == "__main__":
|
| 62 |
+
sys.exit(main())
|
lna-lab/patches/patch_dsv4_aux_streams_env.py
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""LNA-LAB: let LNA_DSV4_AUX_STREAMS=0 disable DSV4's attention aux CUDA streams (EXL3 dense GEMMs are
|
| 3 |
+
cooperative kernels sharing one device lock buffer; running them concurrently on aux streams deadlocks).
|
| 4 |
+
Idempotent exact-anchor patch of nvidia/model.py, backup model.py.orig-lna3."""
|
| 5 |
+
import shutil, sys
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
OLD = " aux_stream_list = [torch.cuda.Stream() for _ in range(3)]\n"
|
| 8 |
+
NEW = (" import os as _lna_os # LNA-LAB\n"
|
| 9 |
+
" aux_stream_list = (None if _lna_os.environ.get(\"LNA_DSV4_AUX_STREAMS\", \"1\") == \"0\"\n"
|
| 10 |
+
" else [torch.cuda.Stream() for _ in range(3)]) # LNA-LAB: EXL3 dense needs sequential GEMMs\n")
|
| 11 |
+
def main():
|
| 12 |
+
import vllm
|
| 13 |
+
f = Path(vllm.__file__).resolve().parent / "models" / "deepseek_v4" / "nvidia" / "model.py"
|
| 14 |
+
t = f.read_text()
|
| 15 |
+
if "LNA_DSV4_AUX_STREAMS" in t: print(f"{f}: already patched"); return 0
|
| 16 |
+
if t.count(OLD) != 1: raise SystemExit(f"{f}: expected 1 anchor, found {t.count(OLD)}")
|
| 17 |
+
b = f.with_suffix(f.suffix + ".orig-lna3"); b.exists() or shutil.copy2(f, b)
|
| 18 |
+
t = t.replace(OLD, NEW); compile(t, str(f), "exec"); f.write_text(t); print(f"{f}: patched"); return 0
|
| 19 |
+
if __name__ == "__main__": sys.exit(main())
|
lna-lab/patches/patch_dsv4_dense_exl3.py
ADDED
|
@@ -0,0 +1,238 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Patch the DSV4 vLLM fork to consume the dense EXL3 overlay at TP=8.
|
| 3 |
+
|
| 4 |
+
The image keeps vLLM's model files private, so this recipe is deliberately an
|
| 5 |
+
exact-anchor patch. Every edit is idempotent, syntax-checked, and backed up as
|
| 6 |
+
``*.orig-lna`` before it is written.
|
| 7 |
+
|
| 8 |
+
Usage: python recipe-lna/patch_dsv4_dense_exl3.py [path/to/site-packages/vllm]
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import ast
|
| 14 |
+
import shutil
|
| 15 |
+
import sys
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
# LNA-LAB: the common compressor class is instantiated for both attention and
|
| 20 |
+
# indexer compressors, so this one constructor anchor covers both instances.
|
| 21 |
+
COMPRESSOR_ANCHOR = (
|
| 22 |
+
" bias=False,\n"
|
| 23 |
+
" return_bias=False,\n"
|
| 24 |
+
" quant_config=None,\n"
|
| 25 |
+
" disable_tp=True,\n"
|
| 26 |
+
)
|
| 27 |
+
COMPRESSOR_PATCH = (
|
| 28 |
+
" bias=False,\n"
|
| 29 |
+
" return_bias=False,\n"
|
| 30 |
+
" # LNA-LAB: dense EXL3 must receive the real quant config.\n"
|
| 31 |
+
" quant_config=vllm_config.quant_config,\n"
|
| 32 |
+
" disable_tp=True,\n"
|
| 33 |
+
)
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
OPROJ_SIGNATURE = (
|
| 37 |
+
" def _o_proj(self, o: torch.Tensor, positions: torch.Tensor)"
|
| 38 |
+
" -> torch.Tensor:\n"
|
| 39 |
+
)
|
| 40 |
+
OPROJ_BASELINE_ANCHOR = OPROJ_SIGNATURE + " return deep_gemm_fp8_o_proj(\n"
|
| 41 |
+
|
| 42 |
+
# This is the exact prefix emitted by recipe/scripts/patch_dsv4_stock028.py;
|
| 43 |
+
# accepting it keeps this patch composable with the existing serving recipe.
|
| 44 |
+
OPROJ_STOCK028_ANCHOR = (
|
| 45 |
+
OPROJ_SIGNATURE
|
| 46 |
+
+ " if self.wo_a.weight.dtype != torch.float8_e4m3fn:\n"
|
| 47 |
+
+ " # bf16 wo_a (packs that keep non-routed weights unquantized):\n"
|
| 48 |
+
+ " # the fp8 einsum path needs block scales that do not exist, so\n"
|
| 49 |
+
+ " # use the Triton inverse-RoPE + bf16 einsum reference instead.\n"
|
| 50 |
+
+ " from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (\n"
|
| 51 |
+
+ " rocm_inv_rope_einsum,\n"
|
| 52 |
+
+ " )\n"
|
| 53 |
+
+ "\n"
|
| 54 |
+
+ " z = rocm_inv_rope_einsum(\n"
|
| 55 |
+
+ " self.rotary_emb,\n"
|
| 56 |
+
+ " o,\n"
|
| 57 |
+
+ " positions,\n"
|
| 58 |
+
+ " self.rope_head_dim,\n"
|
| 59 |
+
+ " self.n_local_groups,\n"
|
| 60 |
+
+ " self.o_lora_rank,\n"
|
| 61 |
+
+ " self.wo_a,\n"
|
| 62 |
+
+ " )\n"
|
| 63 |
+
+ " return self.wo_b(z.flatten(1))\n"
|
| 64 |
+
+ " return deep_gemm_fp8_o_proj(\n"
|
| 65 |
+
)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
OPROJ_COMMON = (
|
| 69 |
+
OPROJ_SIGNATURE
|
| 70 |
+
# LNA-LAB: rank-local EXL3 wo_a is a regular one-group LinearEXL3 call.
|
| 71 |
+
+ " exl3_linears = getattr(self.wo_a, \"_exl3_linears\", None)\n"
|
| 72 |
+
+ " if exl3_linears is not None:\n"
|
| 73 |
+
+ " if self.n_local_groups != 1:\n"
|
| 74 |
+
+ " raise NotImplementedError(\n"
|
| 75 |
+
+ " \"EXL3 wo_a requires TP=8 (one local group); \"\n"
|
| 76 |
+
+ " f\"got n_local_groups={self.n_local_groups}\"\n"
|
| 77 |
+
+ " )\n"
|
| 78 |
+
+ " if len(exl3_linears) != 1 or exl3_linears[0] is None:\n"
|
| 79 |
+
+ " raise RuntimeError(\n"
|
| 80 |
+
+ " \"EXL3 wo_a has no rank-local LinearEXL3 instance\"\n"
|
| 81 |
+
+ " )\n"
|
| 82 |
+
+ " # LNA-LAB: use the same inverse GPT-J RoPE convention as the\n"
|
| 83 |
+
+ " # Triton rocm_inv_rope_einsum reference before the EXL3 GEMM.\n"
|
| 84 |
+
+ " try:\n"
|
| 85 |
+
+ " from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (\n"
|
| 86 |
+
+ " _fused_inverse_rope_gptj,\n"
|
| 87 |
+
+ " )\n"
|
| 88 |
+
+ " o_ref = _fused_inverse_rope_gptj(\n"
|
| 89 |
+
+ " o, positions, self.rotary_emb.cos_sin_cache,\n"
|
| 90 |
+
+ " self.rope_head_dim,\n"
|
| 91 |
+
+ " )\n"
|
| 92 |
+
+ " except ImportError:\n"
|
| 93 |
+
+ " from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (\n"
|
| 94 |
+
+ " _apply_inv_rope_ref,\n"
|
| 95 |
+
+ " )\n"
|
| 96 |
+
+ " o_ref = _apply_inv_rope_ref(\n"
|
| 97 |
+
+ " self.rotary_emb, o, positions, self.rope_head_dim\n"
|
| 98 |
+
+ " )\n"
|
| 99 |
+
+ " b = o_ref.shape[0]\n"
|
| 100 |
+
+ " exl3_linear = exl3_linears[0]\n"
|
| 101 |
+
+ " z = exl3_linear.forward(\n"
|
| 102 |
+
+ " o_ref.reshape(b, -1).contiguous().half(),\n"
|
| 103 |
+
+ " {},\n"
|
| 104 |
+
+ " out_dtype=torch.float32,\n"
|
| 105 |
+
+ " )\n"
|
| 106 |
+
+ " z = z.to(torch.bfloat16).reshape(b, 1, self.o_lora_rank) # LNA-LAB: EXL3 GEMM emits fp32/fp16\n"
|
| 107 |
+
+ " return self.wo_b(z.flatten(1))\n"
|
| 108 |
+
+ " wo_a_weight = getattr(self.wo_a, \"weight\", None)\n"
|
| 109 |
+
+ " if wo_a_weight is None:\n"
|
| 110 |
+
+ " raise RuntimeError(\n"
|
| 111 |
+
+ " \"DSV4 wo_a has neither dense weight nor EXL3 linears\"\n"
|
| 112 |
+
+ " )\n"
|
| 113 |
+
+ " if wo_a_weight.dtype != torch.float8_e4m3fn:\n"
|
| 114 |
+
+ " # LNA-LAB: preserve the existing bf16 reference path.\n"
|
| 115 |
+
+ " from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (\n"
|
| 116 |
+
+ " rocm_inv_rope_einsum,\n"
|
| 117 |
+
+ " )\n"
|
| 118 |
+
+ "\n"
|
| 119 |
+
+ " z = rocm_inv_rope_einsum(\n"
|
| 120 |
+
+ " self.rotary_emb,\n"
|
| 121 |
+
+ " o,\n"
|
| 122 |
+
+ " positions,\n"
|
| 123 |
+
+ " self.rope_head_dim,\n"
|
| 124 |
+
+ " self.n_local_groups,\n"
|
| 125 |
+
+ " self.o_lora_rank,\n"
|
| 126 |
+
+ " self.wo_a,\n"
|
| 127 |
+
+ " )\n"
|
| 128 |
+
+ " return self.wo_b(z.flatten(1))\n"
|
| 129 |
+
+ " return deep_gemm_fp8_o_proj(\n"
|
| 130 |
+
)
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
# LNA-LAB: AutoWeightsLoader maps the pre-fusion bake names to wo_a.tensors;
|
| 134 |
+
# integer ids are then consumed by Exl3LinearMethod as the physical TP rank.
|
| 135 |
+
WO_A_MAPPING_ANCHOR = (
|
| 136 |
+
' ("compressor.fused_wkv_wgate", "compressor.wkv", 0),\n'
|
| 137 |
+
' ("compressor.fused_wkv_wgate", "compressor.wgate", 1),\n'
|
| 138 |
+
" ]\n"
|
| 139 |
+
)
|
| 140 |
+
WO_A_MAPPING_PATCH = (
|
| 141 |
+
' ("compressor.fused_wkv_wgate", "compressor.wkv", 0),\n'
|
| 142 |
+
' ("compressor.fused_wkv_wgate", "compressor.wgate", 1),\n'
|
| 143 |
+
" # LNA-LAB: each pre-fusion wo_a slice is one rank-local shard.\n"
|
| 144 |
+
' ("attn.wo_a", "attn.wo_a.slice.0", 0),\n'
|
| 145 |
+
' ("attn.wo_a", "attn.wo_a.slice.1", 1),\n'
|
| 146 |
+
' ("attn.wo_a", "attn.wo_a.slice.2", 2),\n'
|
| 147 |
+
' ("attn.wo_a", "attn.wo_a.slice.3", 3),\n'
|
| 148 |
+
' ("attn.wo_a", "attn.wo_a.slice.4", 4),\n'
|
| 149 |
+
' ("attn.wo_a", "attn.wo_a.slice.5", 5),\n'
|
| 150 |
+
' ("attn.wo_a", "attn.wo_a.slice.6", 6),\n'
|
| 151 |
+
' ("attn.wo_a", "attn.wo_a.slice.7", 7),\n'
|
| 152 |
+
" ]\n"
|
| 153 |
+
)
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def patch_file(path: Path, anchor: str, patched: str, expect: int) -> str:
|
| 157 |
+
text = path.read_text(encoding="utf-8")
|
| 158 |
+
have = text.count(patched)
|
| 159 |
+
if have == expect:
|
| 160 |
+
return "already patched"
|
| 161 |
+
if have:
|
| 162 |
+
raise SystemExit(f"{path}: partially patched ({have}/{expect}); refusing to guess")
|
| 163 |
+
count = text.count(anchor)
|
| 164 |
+
if count != expect:
|
| 165 |
+
raise SystemExit(f"{path}: expected {expect} anchor(s), found {count}")
|
| 166 |
+
backup = path.with_name(path.name + ".orig-lna")
|
| 167 |
+
if not backup.exists():
|
| 168 |
+
shutil.copy2(path, backup)
|
| 169 |
+
new_text = text.replace(anchor, patched)
|
| 170 |
+
ast.parse(new_text)
|
| 171 |
+
path.write_text(new_text, encoding="utf-8")
|
| 172 |
+
return f"patched x{expect} (backup {backup.name})"
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
def patch_flashinfer(path: Path) -> str:
|
| 176 |
+
text = path.read_text(encoding="utf-8")
|
| 177 |
+
marker = "exl3_linears = getattr(self.wo_a, \"_exl3_linears\", None)"
|
| 178 |
+
if text.count(marker) == 2:
|
| 179 |
+
return "already patched"
|
| 180 |
+
if marker in text:
|
| 181 |
+
raise SystemExit(f"{path}: partially patched; refusing to guess")
|
| 182 |
+
baseline_count = text.count(OPROJ_BASELINE_ANCHOR)
|
| 183 |
+
stock_count = text.count(OPROJ_STOCK028_ANCHOR)
|
| 184 |
+
if baseline_count == 2 and stock_count == 0:
|
| 185 |
+
new_text = text.replace(OPROJ_BASELINE_ANCHOR, OPROJ_COMMON)
|
| 186 |
+
expected = "baseline"
|
| 187 |
+
elif stock_count == 2 and baseline_count == 0:
|
| 188 |
+
new_text = text.replace(OPROJ_STOCK028_ANCHOR, OPROJ_COMMON)
|
| 189 |
+
expected = "stock028-patched"
|
| 190 |
+
else:
|
| 191 |
+
raise SystemExit(
|
| 192 |
+
f"{path}: expected two matching _o_proj methods; "
|
| 193 |
+
f"baseline={baseline_count}, stock028={stock_count}"
|
| 194 |
+
)
|
| 195 |
+
backup = path.with_name(path.name + ".orig-lna")
|
| 196 |
+
if not backup.exists():
|
| 197 |
+
shutil.copy2(path, backup)
|
| 198 |
+
ast.parse(new_text)
|
| 199 |
+
path.write_text(new_text, encoding="utf-8")
|
| 200 |
+
return f"patched x2 ({expected}; backup {backup.name})"
|
| 201 |
+
|
| 202 |
+
|
| 203 |
+
def main() -> int:
|
| 204 |
+
if len(sys.argv) > 1:
|
| 205 |
+
root = Path(sys.argv[1]).expanduser().resolve()
|
| 206 |
+
else:
|
| 207 |
+
import vllm
|
| 208 |
+
|
| 209 |
+
root = Path(vllm.__file__).resolve().parent
|
| 210 |
+
model_dir = root / "models" / "deepseek_v4"
|
| 211 |
+
nvidia_dir = model_dir / "nvidia"
|
| 212 |
+
targets = {
|
| 213 |
+
"compressor.py": model_dir / "compressor.py",
|
| 214 |
+
"flashinfer_sparse.py": nvidia_dir / "flashinfer_sparse.py",
|
| 215 |
+
"model.py": nvidia_dir / "model.py",
|
| 216 |
+
}
|
| 217 |
+
missing = [str(path) for path in targets.values() if not path.is_file()]
|
| 218 |
+
if missing:
|
| 219 |
+
print("missing DSV4 vLLM file(s):", ", ".join(missing), file=sys.stderr)
|
| 220 |
+
return 2
|
| 221 |
+
|
| 222 |
+
try:
|
| 223 |
+
print("compressor.py:", patch_file(
|
| 224 |
+
targets["compressor.py"], COMPRESSOR_ANCHOR, COMPRESSOR_PATCH, 1
|
| 225 |
+
))
|
| 226 |
+
print("flashinfer_sparse.py:", patch_flashinfer(targets["flashinfer_sparse.py"]))
|
| 227 |
+
print("model.py:", patch_file(
|
| 228 |
+
targets["model.py"], WO_A_MAPPING_ANCHOR, WO_A_MAPPING_PATCH, 1
|
| 229 |
+
))
|
| 230 |
+
except (OSError, SystemExit) as exc:
|
| 231 |
+
print(exc, file=sys.stderr)
|
| 232 |
+
return 1
|
| 233 |
+
print("DSV4_DENSE_EXL3_PATCH_OK")
|
| 234 |
+
return 0
|
| 235 |
+
|
| 236 |
+
|
| 237 |
+
if __name__ == "__main__":
|
| 238 |
+
raise SystemExit(main())
|
lna-lab/patches/vllm_exl3_exl3.py.lna
ADDED
|
@@ -0,0 +1,1528 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""EXL3/MCG trellis quantization for GLM-5.3-Flash routed experts.
|
| 3 |
+
|
| 4 |
+
Checkpoint ABI used by this pack:
|
| 5 |
+
quant_method=exl3, codebook=mcg, scope=glm53_routed_experts_only
|
| 6 |
+
per expert matrix: trellis (int16) + suh/svh (fp16) + mcg (int32 marker)
|
| 7 |
+
|
| 8 |
+
Non-routed tensors stay native (UnquantizedLinearMethod). Experts never
|
| 9 |
+
expand to a persistent BF16 weight; LinearEXL3 / exllamav3_ext runs the
|
| 10 |
+
trellis GEMM. TP=2 shards gate/up column-wise and down row-wise; the MoE
|
| 11 |
+
runner all-reduces the combined output.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import importlib
|
| 17 |
+
import os
|
| 18 |
+
from typing import TYPE_CHECKING, Any
|
| 19 |
+
|
| 20 |
+
import re
|
| 21 |
+
import torch
|
| 22 |
+
import torch.nn.functional as F
|
| 23 |
+
from torch.nn.parameter import Parameter
|
| 24 |
+
|
| 25 |
+
from vllm.logger import init_logger
|
| 26 |
+
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
|
| 27 |
+
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
|
| 28 |
+
FusedMoEMethodBase,
|
| 29 |
+
)
|
| 30 |
+
from vllm.model_executor.layers.linear import (
|
| 31 |
+
LinearBase,
|
| 32 |
+
LinearMethodBase,
|
| 33 |
+
UnquantizedLinearMethod,
|
| 34 |
+
)
|
| 35 |
+
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
| 36 |
+
from vllm.model_executor.layers.quantization import register_quantization_config
|
| 37 |
+
from vllm.model_executor.utils import set_weight_attrs
|
| 38 |
+
|
| 39 |
+
if TYPE_CHECKING:
|
| 40 |
+
from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
|
| 41 |
+
from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
|
| 42 |
+
SharedExperts,
|
| 43 |
+
)
|
| 44 |
+
|
| 45 |
+
# Under the "vllm." hierarchy so vLLM's logging config actually emits these
|
| 46 |
+
# INFO lines; a bare module name is dropped and the load log shows nothing.
|
| 47 |
+
logger = init_logger("vllm." + __name__)
|
| 48 |
+
|
| 49 |
+
MCG_MULTIPLIER = 0xCBAC1FED
|
| 50 |
+
MCG_MARKER_SIGNED_INT32 = -877912083
|
| 51 |
+
MUL1_MULTIPLIER = 0x83DCD12D
|
| 52 |
+
MUL1_MARKER_SIGNED_INT32 = -2082680531
|
| 53 |
+
EXL3_SUFFIXES = ("trellis", "suh", "svh", "mcg", "mul1")
|
| 54 |
+
_WO_A_SLICE_RE = re.compile(r"(?:^|\.)slice\.(?P<rank>[0-9]+)(?:\.|$)")
|
| 55 |
+
SWIGLU_LIMIT_DEFAULT = 10.0
|
| 56 |
+
TEMP_ROWS_FUSED = 2048
|
| 57 |
+
MOE_ACT_SILU = 0
|
| 58 |
+
# Shared fused scratch: decode is sequential across layers.
|
| 59 |
+
_FUSED_TEMP_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = {}
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def _narrow_tp(tensor: torch.Tensor, dim: int, tp_rank: int, tp_size: int) -> torch.Tensor:
|
| 63 |
+
if tp_size <= 1:
|
| 64 |
+
return tensor
|
| 65 |
+
size = int(tensor.shape[dim])
|
| 66 |
+
if size % tp_size:
|
| 67 |
+
raise ValueError(
|
| 68 |
+
f"EXL3 TP shard: dim {dim} size {size} is not divisible by tp={tp_size}"
|
| 69 |
+
)
|
| 70 |
+
chunk = size // tp_size
|
| 71 |
+
return tensor.narrow(dim, chunk * tp_rank, chunk).contiguous()
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def _effective_tp_for_layer(
|
| 75 |
+
layer: torch.nn.Module, global_tp_size: int, global_tp_rank: int
|
| 76 |
+
) -> tuple[int, int]:
|
| 77 |
+
"""Return the TP geometry that the checkpoint tensor actually uses."""
|
| 78 |
+
# LNA-LAB: replicated/disable_tp layers must not be narrowed by global TP.
|
| 79 |
+
try:
|
| 80 |
+
from vllm.model_executor.layers.linear import ReplicatedLinear
|
| 81 |
+
except ImportError: # pragma: no cover - only old vLLM builds lack this class.
|
| 82 |
+
is_replicated = layer.__class__.__name__ == "ReplicatedLinear"
|
| 83 |
+
else:
|
| 84 |
+
is_replicated = isinstance(layer, ReplicatedLinear)
|
| 85 |
+
if getattr(layer, "disable_tp", False) or is_replicated:
|
| 86 |
+
return 1, 0
|
| 87 |
+
return int(global_tp_size), int(global_tp_rank)
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def _wo_a_slice_rank(loaded_shard_id: str | int | None) -> int | None:
|
| 91 |
+
"""Decode a model-loader shard id or a literal ``wo_a.slice.N`` name."""
|
| 92 |
+
# LNA-LAB: dense-bake names carry the group rank; model.py maps them to ints.
|
| 93 |
+
if isinstance(loaded_shard_id, int):
|
| 94 |
+
return loaded_shard_id
|
| 95 |
+
if loaded_shard_id is None:
|
| 96 |
+
return None
|
| 97 |
+
if str(loaded_shard_id).isdigit():
|
| 98 |
+
return int(str(loaded_shard_id))
|
| 99 |
+
match = _WO_A_SLICE_RE.search(str(loaded_shard_id))
|
| 100 |
+
return int(match.group("rank")) if match else None
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def _is_wo_a_layer(layer: torch.nn.Module) -> bool:
|
| 104 |
+
# LNA-LAB: vLLM sets is_bmm after construction, so prefix is the fallback.
|
| 105 |
+
prefix = str(getattr(layer, "prefix", ""))
|
| 106 |
+
return bool(
|
| 107 |
+
_prefix_has_suffix(prefix, "attn.wo_a")
|
| 108 |
+
or (getattr(layer, "is_bmm", False) and prefix.endswith(".wo_a"))
|
| 109 |
+
)
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def shard_exl3_col(loaded: torch.Tensor, suffix: str, tp_rank: int, tp_size: int) -> torch.Tensor:
|
| 113 |
+
"""Gate/up: trellis dim 1 and svh dim 0 are column-parallel."""
|
| 114 |
+
if suffix == "trellis":
|
| 115 |
+
return _narrow_tp(loaded, 1, tp_rank, tp_size)
|
| 116 |
+
if suffix == "svh":
|
| 117 |
+
return _narrow_tp(loaded, 0, tp_rank, tp_size)
|
| 118 |
+
return loaded.contiguous()
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def shard_exl3_row(loaded: torch.Tensor, suffix: str, tp_rank: int, tp_size: int) -> torch.Tensor:
|
| 122 |
+
"""Down: trellis dim 0 and suh dim 0 are row-parallel."""
|
| 123 |
+
if suffix == "trellis":
|
| 124 |
+
return _narrow_tp(loaded, 0, tp_rank, tp_size)
|
| 125 |
+
if suffix == "suh":
|
| 126 |
+
return _narrow_tp(loaded, 0, tp_rank, tp_size)
|
| 127 |
+
return loaded.contiguous()
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
def _install_exllamav3_namespace() -> None:
|
| 131 |
+
"""Validate that the native ExLlamaV3 package and extension are importable."""
|
| 132 |
+
import exllamav3_ext # noqa: F401 — compiled extension must exist
|
| 133 |
+
|
| 134 |
+
# ExLlamaV3 1.4+ imports cleanly as a regular package and its LinearEXL3
|
| 135 |
+
# constructor relies on the real NullConfig/InferParams implementation.
|
| 136 |
+
# Namespace stubs used by much older builds hide those classes and fail only
|
| 137 |
+
# after a full checkpoint load, so deliberately exercise the normal import.
|
| 138 |
+
importlib.import_module("exllamav3.modules.quant.exl3")
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def load_linear_exl3_cls():
|
| 142 |
+
_install_exllamav3_namespace()
|
| 143 |
+
return importlib.import_module("exllamav3.modules.quant.exl3").LinearEXL3
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def make_linear_exl3(
|
| 147 |
+
trellis: torch.Tensor,
|
| 148 |
+
suh: torch.Tensor,
|
| 149 |
+
svh: torch.Tensor,
|
| 150 |
+
mcg: torch.Tensor | None = None,
|
| 151 |
+
mul1: torch.Tensor | None = None,
|
| 152 |
+
*,
|
| 153 |
+
out_dtype: torch.dtype = torch.float16,
|
| 154 |
+
):
|
| 155 |
+
"""Build a LinearEXL3 over already-sharded packed tensors. No BF16 expand."""
|
| 156 |
+
cls = load_linear_exl3_cls()
|
| 157 |
+
return cls(
|
| 158 |
+
config=None,
|
| 159 |
+
in_features=int(suh.numel()),
|
| 160 |
+
out_features=int(svh.numel()),
|
| 161 |
+
trellis=trellis.contiguous(),
|
| 162 |
+
suh=suh.contiguous(),
|
| 163 |
+
svh=svh.contiguous(),
|
| 164 |
+
mcg=mcg.contiguous() if mcg is not None else None,
|
| 165 |
+
mul1=mul1.contiguous() if mul1 is not None else None,
|
| 166 |
+
out_dtype=out_dtype,
|
| 167 |
+
transformers_fix=True,
|
| 168 |
+
)
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def execute_exl3_linear(
|
| 172 |
+
x: torch.Tensor,
|
| 173 |
+
trellis: torch.Tensor,
|
| 174 |
+
suh: torch.Tensor,
|
| 175 |
+
svh: torch.Tensor,
|
| 176 |
+
mcg: torch.Tensor | None = None,
|
| 177 |
+
mul1: torch.Tensor | None = None,
|
| 178 |
+
*,
|
| 179 |
+
out_dtype: torch.dtype = torch.float32,
|
| 180 |
+
) -> torch.Tensor:
|
| 181 |
+
"""Real EXL3 expert GEMM entry (LinearEXL3 / exllamav3_ext)."""
|
| 182 |
+
inner = make_linear_exl3(trellis, suh, svh, mcg, mul1, out_dtype=torch.float16)
|
| 183 |
+
return inner.forward(x.contiguous().half(), {}, out_dtype=out_dtype)
|
| 184 |
+
|
| 185 |
+
|
| 186 |
+
def fused_moe_enabled() -> bool:
|
| 187 |
+
return os.environ.get("EXL3_FUSED_MOE", "1") != "0"
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def load_exllamav3_ext():
|
| 191 |
+
import exllamav3_ext
|
| 192 |
+
|
| 193 |
+
return exllamav3_ext
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def _exl3_moe_accepts_num_active(fn) -> bool:
|
| 197 |
+
try:
|
| 198 |
+
import inspect
|
| 199 |
+
|
| 200 |
+
if "num_active" in inspect.signature(fn).parameters:
|
| 201 |
+
return True
|
| 202 |
+
except (TypeError, ValueError):
|
| 203 |
+
pass
|
| 204 |
+
doc = getattr(fn, "__doc__", None) or ""
|
| 205 |
+
return "num_active" in doc or "arg29" in doc or doc.count("arg") >= 30
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def pin_exl3_expert_map(
|
| 209 |
+
layer: torch.nn.Module, device: torch.device
|
| 210 |
+
) -> torch.Tensor | None:
|
| 211 |
+
"""Move expert_map onto `device` once. CUDA graph capture forbids a CPU→GPU copy."""
|
| 212 |
+
emap = getattr(layer, "expert_map", None)
|
| 213 |
+
if emap is None:
|
| 214 |
+
return None
|
| 215 |
+
if emap.device != device or emap.dtype != torch.long:
|
| 216 |
+
layer.expert_map = emap.to(device=device, dtype=torch.long)
|
| 217 |
+
return layer.expert_map
|
| 218 |
+
|
| 219 |
+
|
| 220 |
+
def map_topk_to_local(
|
| 221 |
+
ids: torch.Tensor,
|
| 222 |
+
n_local: int,
|
| 223 |
+
expert_map: torch.Tensor | None,
|
| 224 |
+
) -> torch.Tensor:
|
| 225 |
+
"""ids (T, K) global expert ids → local ids, invalid/non-local → n_local sentinel.
|
| 226 |
+
|
| 227 |
+
`expert_map` must already live on `ids.device` (see pin_exl3_expert_map).
|
| 228 |
+
"""
|
| 229 |
+
flat = ids.reshape(-1)
|
| 230 |
+
if expert_map is None:
|
| 231 |
+
invalid = (flat < 0) | (flat >= n_local)
|
| 232 |
+
return torch.where(invalid, flat.new_full(flat.shape, n_local), flat)
|
| 233 |
+
if expert_map.device != flat.device or expert_map.dtype != torch.long:
|
| 234 |
+
raise RuntimeError(
|
| 235 |
+
"EXL3 expert_map is not pinned to the hidden-state device; "
|
| 236 |
+
"call pin_exl3_expert_map before fused apply (CUDA graphs forbid the copy)"
|
| 237 |
+
)
|
| 238 |
+
n_global = int(expert_map.numel())
|
| 239 |
+
safe = flat.clamp(min=0, max=max(n_global - 1, 0))
|
| 240 |
+
mapped = expert_map[safe] if n_global else flat.new_full(flat.shape, n_local)
|
| 241 |
+
invalid = (flat < 0) | (flat >= n_global) | (mapped < 0) | (mapped >= n_local)
|
| 242 |
+
return torch.where(invalid, flat.new_full(flat.shape, n_local), mapped)
|
| 243 |
+
|
| 244 |
+
|
| 245 |
+
def apply_exl3_python_loop(
|
| 246 |
+
x2d: torch.Tensor,
|
| 247 |
+
ids: torch.Tensor,
|
| 248 |
+
weights: torch.Tensor,
|
| 249 |
+
inners: list[dict[str, Any]],
|
| 250 |
+
expert_map: torch.Tensor | None,
|
| 251 |
+
limit: float,
|
| 252 |
+
*,
|
| 253 |
+
only_experts: set[int] | None = None,
|
| 254 |
+
out: torch.Tensor | None = None,
|
| 255 |
+
) -> torch.Tensor:
|
| 256 |
+
"""Unique-expert LinearEXL3 loop. `only_experts` is local ids (fat-expert fallback)."""
|
| 257 |
+
tokens, hidden = x2d.shape
|
| 258 |
+
if out is None:
|
| 259 |
+
out = torch.zeros(tokens, hidden, dtype=torch.float32, device=x2d.device)
|
| 260 |
+
unique = torch.unique(ids)
|
| 261 |
+
for raw in unique.tolist():
|
| 262 |
+
e_raw = int(raw)
|
| 263 |
+
if e_raw < 0:
|
| 264 |
+
continue
|
| 265 |
+
e = e_raw
|
| 266 |
+
if expert_map is not None:
|
| 267 |
+
mapped = int(expert_map[e].item()) if expert_map.numel() > e else e
|
| 268 |
+
if mapped < 0:
|
| 269 |
+
continue
|
| 270 |
+
e = mapped
|
| 271 |
+
if e >= len(inners):
|
| 272 |
+
continue
|
| 273 |
+
if only_experts is not None and e not in only_experts:
|
| 274 |
+
continue
|
| 275 |
+
token_idx, k_pos = (ids == int(raw)).nonzero(as_tuple=True)
|
| 276 |
+
h = x2d.index_select(0, token_idx)
|
| 277 |
+
pack = inners[e]
|
| 278 |
+
gate = pack["gate"].forward(h.contiguous().half(), {}, out_dtype=torch.float32)
|
| 279 |
+
up = pack["up"].forward(h.contiguous().half(), {}, out_dtype=torch.float32)
|
| 280 |
+
act = F.silu(gate.clamp(max=limit)) * up.clamp(min=-limit, max=limit)
|
| 281 |
+
down = pack["down"].forward(act.contiguous().half(), {}, out_dtype=torch.float32)
|
| 282 |
+
scale = weights[token_idx, k_pos].unsqueeze(-1).to(dtype=torch.float32)
|
| 283 |
+
out.index_add_(0, token_idx, down * scale)
|
| 284 |
+
return out
|
| 285 |
+
|
| 286 |
+
|
| 287 |
+
def build_exl3_fused_state(layer: torch.nn.Module, inners: list[dict[str, Any]]) -> None:
|
| 288 |
+
"""Pointer tables + fused temps, once after load. No per-token alloc."""
|
| 289 |
+
import exllamav3_ext
|
| 290 |
+
|
| 291 |
+
device = layer.w13_trellis.device
|
| 292 |
+
n_exp = len(inners)
|
| 293 |
+
hidden = int(layer._exl3_hidden_size)
|
| 294 |
+
intermediate = int(layer._exl3_intermediate_local)
|
| 295 |
+
|
| 296 |
+
def _ptrs(which: str, attr: str) -> torch.Tensor:
|
| 297 |
+
return torch.tensor(
|
| 298 |
+
[int(getattr(pack[which], attr).data_ptr()) for pack in inners],
|
| 299 |
+
dtype=torch.int64,
|
| 300 |
+
device=device,
|
| 301 |
+
)
|
| 302 |
+
|
| 303 |
+
layer._exl3_ptrs = {
|
| 304 |
+
"gate_trellis": _ptrs("gate", "trellis"),
|
| 305 |
+
"gate_suh": _ptrs("gate", "suh"),
|
| 306 |
+
"gate_svh": _ptrs("gate", "svh"),
|
| 307 |
+
"up_trellis": _ptrs("up", "trellis"),
|
| 308 |
+
"up_suh": _ptrs("up", "suh"),
|
| 309 |
+
"up_svh": _ptrs("up", "svh"),
|
| 310 |
+
"down_trellis": _ptrs("down", "trellis"),
|
| 311 |
+
"down_suh": _ptrs("down", "suh"),
|
| 312 |
+
"down_svh": _ptrs("down", "svh"),
|
| 313 |
+
}
|
| 314 |
+
idx = int(device.index) if device.index is not None else 0
|
| 315 |
+
concurrency = int(exllamav3_ext.exl3_moe_max_concurrency(idx))
|
| 316 |
+
if concurrency < 1:
|
| 317 |
+
concurrency = 1
|
| 318 |
+
key = (str(device), hidden, intermediate, concurrency)
|
| 319 |
+
temps = _FUSED_TEMP_CACHE.get(key)
|
| 320 |
+
if temps is None:
|
| 321 |
+
temps = (
|
| 322 |
+
torch.empty((concurrency, TEMP_ROWS_FUSED, hidden), dtype=torch.float16, device=device),
|
| 323 |
+
torch.empty((concurrency, TEMP_ROWS_FUSED, hidden), dtype=torch.float16, device=device),
|
| 324 |
+
torch.empty((concurrency, TEMP_ROWS_FUSED, intermediate), dtype=torch.float16, device=device),
|
| 325 |
+
torch.empty((concurrency, TEMP_ROWS_FUSED, intermediate), dtype=torch.float16, device=device),
|
| 326 |
+
)
|
| 327 |
+
_FUSED_TEMP_CACHE[key] = temps
|
| 328 |
+
layer._exl3_fused_temps = temps
|
| 329 |
+
layer._exl3_fused_concurrency = concurrency
|
| 330 |
+
layer._exl3_k = int(layer._exl3_bits)
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
def apply_exl3_fused_moe(
|
| 334 |
+
x2d: torch.Tensor,
|
| 335 |
+
ids: torch.Tensor,
|
| 336 |
+
weights: torch.Tensor,
|
| 337 |
+
layer: torch.nn.Module,
|
| 338 |
+
inners: list[dict[str, Any]],
|
| 339 |
+
expert_map: torch.Tensor | None,
|
| 340 |
+
limit: float,
|
| 341 |
+
) -> torch.Tensor:
|
| 342 |
+
"""One exl3_moe launch per layer. Experts with count > 128 fall back to LinearEXL3."""
|
| 343 |
+
import exllamav3_ext
|
| 344 |
+
|
| 345 |
+
tokens, hidden = x2d.shape
|
| 346 |
+
n_exp = len(inners)
|
| 347 |
+
ptrs = getattr(layer, "_exl3_ptrs", None)
|
| 348 |
+
temps = getattr(layer, "_exl3_fused_temps", None)
|
| 349 |
+
if not ptrs or temps is None:
|
| 350 |
+
raise RuntimeError("EXL3 fused pointer tables were not built after weight load")
|
| 351 |
+
|
| 352 |
+
local = map_topk_to_local(ids, n_exp, expert_map)
|
| 353 |
+
topk = int(ids.shape[-1])
|
| 354 |
+
flat_token = torch.arange(tokens, device=x2d.device, dtype=torch.long).repeat_interleave(topk)
|
| 355 |
+
flat_weight = weights.reshape(-1).to(dtype=torch.float16)
|
| 356 |
+
order = local.argsort()
|
| 357 |
+
token_sorted = flat_token[order]
|
| 358 |
+
weight_sorted = flat_weight[order]
|
| 359 |
+
# scatter_add stays on GPU. torch.bincount can host-stage and break CUDA graphs.
|
| 360 |
+
expert_count = torch.zeros(n_exp + 1, dtype=torch.long, device=local.device)
|
| 361 |
+
expert_count.scatter_add_(
|
| 362 |
+
0, local.long(), torch.ones(local.shape, dtype=torch.long, device=local.device)
|
| 363 |
+
)
|
| 364 |
+
out = torch.zeros(tokens, hidden, dtype=torch.float32, device=x2d.device)
|
| 365 |
+
xh = x2d.contiguous().half()
|
| 366 |
+
|
| 367 |
+
counts = expert_count[:n_exp]
|
| 368 |
+
|
| 369 |
+
if tokens > TEMP_ROWS_FUSED and bool((counts > TEMP_ROWS_FUSED).any().item()):
|
| 370 |
+
logger.info_once("EXL3 fat-chunk slicing ACTIVE (tokens=%d)" % tokens)
|
| 371 |
+
# Deep-context prefill chunks can route more than TEMP_ROWS_FUSED rows
|
| 372 |
+
# to a single expert. The fused kernel covers at most TEMP_ROWS_FUSED
|
| 373 |
+
# rows per expert, and the old fallback reconstructed whole experts
|
| 374 |
+
# per chunk, stalling prefill by orders of magnitude past ~160k
|
| 375 |
+
# context (the ">163k hang"). Within a slice of <= TEMP_ROWS_FUSED
|
| 376 |
+
# tokens no expert can exceed TEMP_ROWS_FUSED rows (each token adds at
|
| 377 |
+
# most one row per expert), so re-run the fused path per slice.
|
| 378 |
+
# Prefill-only: decode batches are at most the largest capture size,
|
| 379 |
+
# far below TEMP_ROWS_FUSED, and never reach this host sync.
|
| 380 |
+
for s in range(0, tokens, TEMP_ROWS_FUSED):
|
| 381 |
+
e = min(s + TEMP_ROWS_FUSED, tokens)
|
| 382 |
+
out[s:e] = apply_exl3_fused_moe(
|
| 383 |
+
x2d[s:e], ids[s:e], weights[s:e], layer, inners, expert_map, limit
|
| 384 |
+
)
|
| 385 |
+
return out
|
| 386 |
+
fn = exllamav3_ext.exl3_moe
|
| 387 |
+
# -1 = unknown active count: max-concurrency grid, no .item() host sync.
|
| 388 |
+
n_active_host = -1 if _exl3_moe_accepts_num_active(fn) else None
|
| 389 |
+
|
| 390 |
+
k = int(getattr(layer, "_exl3_k", 4))
|
| 391 |
+
args = (
|
| 392 |
+
xh,
|
| 393 |
+
out,
|
| 394 |
+
expert_count,
|
| 395 |
+
token_sorted,
|
| 396 |
+
weight_sorted,
|
| 397 |
+
temps[0],
|
| 398 |
+
temps[1],
|
| 399 |
+
temps[2],
|
| 400 |
+
temps[3],
|
| 401 |
+
MOE_ACT_SILU,
|
| 402 |
+
k,
|
| 403 |
+
k,
|
| 404 |
+
k,
|
| 405 |
+
ptrs["gate_trellis"],
|
| 406 |
+
ptrs["gate_suh"],
|
| 407 |
+
ptrs["gate_svh"],
|
| 408 |
+
ptrs["up_trellis"],
|
| 409 |
+
ptrs["up_suh"],
|
| 410 |
+
ptrs["up_svh"],
|
| 411 |
+
ptrs["down_trellis"],
|
| 412 |
+
ptrs["down_suh"],
|
| 413 |
+
ptrs["down_svh"],
|
| 414 |
+
True,
|
| 415 |
+
False,
|
| 416 |
+
True,
|
| 417 |
+
False,
|
| 418 |
+
True,
|
| 419 |
+
False,
|
| 420 |
+
float(limit),
|
| 421 |
+
)
|
| 422 |
+
if n_active_host is not None:
|
| 423 |
+
fn(*args, n_active_host)
|
| 424 |
+
else:
|
| 425 |
+
fn(*args)
|
| 426 |
+
|
| 427 |
+
if tokens > TEMP_ROWS_FUSED:
|
| 428 |
+
fat = (counts > TEMP_ROWS_FUSED).nonzero(as_tuple=False).view(-1)
|
| 429 |
+
if fat.numel():
|
| 430 |
+
apply_exl3_python_loop(
|
| 431 |
+
x2d,
|
| 432 |
+
ids,
|
| 433 |
+
weights,
|
| 434 |
+
inners,
|
| 435 |
+
expert_map,
|
| 436 |
+
limit,
|
| 437 |
+
only_experts=set(int(i) for i in fat.tolist()),
|
| 438 |
+
out=out,
|
| 439 |
+
)
|
| 440 |
+
return out
|
| 441 |
+
|
| 442 |
+
|
| 443 |
+
def apply_exl3_experts(
|
| 444 |
+
x: torch.Tensor,
|
| 445 |
+
topk_ids: torch.Tensor,
|
| 446 |
+
topk_weights: torch.Tensor,
|
| 447 |
+
layer: torch.nn.Module,
|
| 448 |
+
*,
|
| 449 |
+
limit: float = SWIGLU_LIMIT_DEFAULT,
|
| 450 |
+
fused: bool | None = None,
|
| 451 |
+
) -> torch.Tensor:
|
| 452 |
+
"""Shipped routed-expert apply. `fused=None` honors EXL3_FUSED_MOE."""
|
| 453 |
+
inners = getattr(layer, "_exl3_inners", None)
|
| 454 |
+
if not inners:
|
| 455 |
+
raise RuntimeError("EXL3 experts were not built after weight load")
|
| 456 |
+
tokens, hidden = x.shape[-2], x.shape[-1]
|
| 457 |
+
x2d = x.reshape(tokens, hidden)
|
| 458 |
+
ids = topk_ids.reshape(tokens, -1).to(torch.long)
|
| 459 |
+
weights = topk_weights.reshape(tokens, -1)
|
| 460 |
+
expert_map = pin_exl3_expert_map(layer, x2d.device)
|
| 461 |
+
have_ptrs = bool(getattr(layer, "_exl3_ptrs", None))
|
| 462 |
+
if fused is True and not have_ptrs:
|
| 463 |
+
raise RuntimeError("EXL3 fused apply requested but pointer tables are missing")
|
| 464 |
+
use_fused = (fused_moe_enabled() if fused is None else bool(fused)) and have_ptrs
|
| 465 |
+
if use_fused:
|
| 466 |
+
try:
|
| 467 |
+
import exllamav3_ext
|
| 468 |
+
|
| 469 |
+
use_fused = hasattr(exllamav3_ext, "exl3_moe")
|
| 470 |
+
except Exception:
|
| 471 |
+
use_fused = False
|
| 472 |
+
if use_fused:
|
| 473 |
+
out = apply_exl3_fused_moe(x2d, ids, weights, layer, inners, expert_map, limit)
|
| 474 |
+
layer._exl3_last_apply = "fused"
|
| 475 |
+
else:
|
| 476 |
+
out = apply_exl3_python_loop(x2d, ids, weights, inners, expert_map, limit)
|
| 477 |
+
layer._exl3_last_apply = "loop"
|
| 478 |
+
return out.to(dtype=x.dtype)
|
| 479 |
+
|
| 480 |
+
|
| 481 |
+
def _suffix_from_mapped_name(weight_name: str) -> str:
|
| 482 |
+
tail = weight_name.rsplit(".", 1)[-1]
|
| 483 |
+
for suffix in EXL3_SUFFIXES:
|
| 484 |
+
if tail == suffix or tail.endswith("_" + suffix):
|
| 485 |
+
return suffix
|
| 486 |
+
raise ValueError(f"not an EXL3 packed name: {weight_name}")
|
| 487 |
+
|
| 488 |
+
|
| 489 |
+
def _prefix_has_suffix(prefix: str, suffix: str) -> bool:
|
| 490 |
+
"""Module-path suffix match: "self_attn.o_proj" matches
|
| 491 |
+
"model.layers.3.self_attn.o_proj" but not "...cross_attn.o_proj_x"."""
|
| 492 |
+
return prefix == suffix or prefix.endswith("." + suffix)
|
| 493 |
+
|
| 494 |
+
|
| 495 |
+
@register_quantization_config("exl3")
|
| 496 |
+
class Exl3Config(QuantizationConfig):
|
| 497 |
+
"""Routed-experts-only EXL3/MCG. Dense / shared / attention stay native."""
|
| 498 |
+
|
| 499 |
+
def __init__(
|
| 500 |
+
self,
|
| 501 |
+
bits: int = 4,
|
| 502 |
+
codebook: str = "mcg",
|
| 503 |
+
scope: str = "glm53_routed_experts_only",
|
| 504 |
+
**kwargs: Any,
|
| 505 |
+
) -> None:
|
| 506 |
+
super().__init__()
|
| 507 |
+
self.bits = int(bits)
|
| 508 |
+
self.codebook = str(codebook)
|
| 509 |
+
self.scope = str(scope)
|
| 510 |
+
# Optional per-layer override, e.g. {"42": 3, "27": 3}. Layers absent
|
| 511 |
+
# from the map use `bits`. This is how a mixed-K checkpoint (K2 base
|
| 512 |
+
# with K3 delta layers) declares itself; the trellis tensors for those
|
| 513 |
+
# layers are shaped for their own K and would fail the load shape
|
| 514 |
+
# check under the base K.
|
| 515 |
+
raw_layer_bits = kwargs.pop("layer_bits", None) or {}
|
| 516 |
+
self.layer_bits: dict[int, int] = {
|
| 517 |
+
int(k): int(v) for k, v in dict(raw_layer_bits).items()
|
| 518 |
+
}
|
| 519 |
+
for layer_idx, layer_k in self.layer_bits.items():
|
| 520 |
+
if layer_k not in (2, 3, 4, 5, 6):
|
| 521 |
+
raise ValueError(
|
| 522 |
+
f"unsupported EXL3 bits={layer_k} for layer {layer_idx}"
|
| 523 |
+
)
|
| 524 |
+
# Non-routed dense linear config: optional {"modules": [...], "bits": K, "layer_bits": {...},
|
| 525 |
+
# "codebook": "mcg"|"mul1", "layers": {prefix: {"bits": K, "bf16_shards": [...]}, ...}}
|
| 526 |
+
raw_nr_exl3 = kwargs.pop("non_routed_exl3", None) or {}
|
| 527 |
+
self.non_routed_exl3: dict[str, Any] = dict(raw_nr_exl3) if raw_nr_exl3 else {}
|
| 528 |
+
# Validate non-routed bits if present
|
| 529 |
+
nr_bits = self.non_routed_exl3.get("bits")
|
| 530 |
+
if nr_bits is not None and nr_bits not in (2, 3, 4, 5, 6):
|
| 531 |
+
raise ValueError(f"unsupported non_routed_exl3 bits={nr_bits}")
|
| 532 |
+
nr_layer_bits = self.non_routed_exl3.get("layer_bits", {})
|
| 533 |
+
for suffix, k in (nr_layer_bits or {}).items():
|
| 534 |
+
if k not in (2, 3, 4, 5, 6):
|
| 535 |
+
raise ValueError(
|
| 536 |
+
f"unsupported non_routed_exl3 bits={k} for suffix {suffix}"
|
| 537 |
+
)
|
| 538 |
+
# Validate non-routed layers dict: each value is {"bits": K[, "bf16_shards": [...]]}
|
| 539 |
+
nr_layers = self.non_routed_exl3.get("layers", {})
|
| 540 |
+
for prefix, layer_cfg in (nr_layers or {}).items():
|
| 541 |
+
if not isinstance(layer_cfg, dict):
|
| 542 |
+
raise ValueError(
|
| 543 |
+
f"non_routed_exl3 layers[{prefix}] must be a dict, got {type(layer_cfg)}"
|
| 544 |
+
)
|
| 545 |
+
layer_bits = layer_cfg.get("bits")
|
| 546 |
+
if layer_bits is not None and layer_bits not in (2, 3, 4, 5, 6):
|
| 547 |
+
raise ValueError(
|
| 548 |
+
f"unsupported non_routed_exl3 layers[{prefix}] bits={layer_bits}"
|
| 549 |
+
)
|
| 550 |
+
# Validate non-routed codebook
|
| 551 |
+
nr_codebook = self.non_routed_exl3.get("codebook", "mcg")
|
| 552 |
+
if nr_codebook not in ("mcg", "mul1"):
|
| 553 |
+
raise ValueError(
|
| 554 |
+
f"unsupported non_routed_exl3 codebook={nr_codebook!r}; must be 'mcg' or 'mul1'"
|
| 555 |
+
)
|
| 556 |
+
self.raw_config = dict(kwargs)
|
| 557 |
+
if self.codebook not in ("mcg", "mul1"):
|
| 558 |
+
raise ValueError(
|
| 559 |
+
f"unsupported codebook={self.codebook!r}; must be 'mcg' or 'mul1'"
|
| 560 |
+
)
|
| 561 |
+
if self.bits not in (2, 3, 4, 5, 6):
|
| 562 |
+
raise ValueError(f"unsupported EXL3 bits={self.bits}")
|
| 563 |
+
|
| 564 |
+
def get_name(self) -> str:
|
| 565 |
+
return "exl3"
|
| 566 |
+
|
| 567 |
+
_LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")
|
| 568 |
+
|
| 569 |
+
def bits_for_prefix(self, prefix: str) -> int:
|
| 570 |
+
"""Per-layer K: `layer_bits` entry for this layer, else the base K."""
|
| 571 |
+
if not self.layer_bits:
|
| 572 |
+
return self.bits
|
| 573 |
+
m = self._LAYER_RE.search(prefix or "")
|
| 574 |
+
if m is None:
|
| 575 |
+
return self.bits
|
| 576 |
+
return self.layer_bits.get(int(m.group(1)), self.bits)
|
| 577 |
+
|
| 578 |
+
def _matches_non_routed_exl3(self, prefix: str) -> bool:
|
| 579 |
+
"""Check if prefix matches non_routed_exl3: either layers dict keys or modules list."""
|
| 580 |
+
if not self.non_routed_exl3:
|
| 581 |
+
return False
|
| 582 |
+
# Check if prefix is a key in the layers dict
|
| 583 |
+
layers = self.non_routed_exl3.get("layers", {})
|
| 584 |
+
if layers and prefix in layers:
|
| 585 |
+
return True
|
| 586 |
+
# Fall back to suffix matching on modules list
|
| 587 |
+
modules = self.non_routed_exl3.get("modules", [])
|
| 588 |
+
if not modules:
|
| 589 |
+
return False
|
| 590 |
+
return any(_prefix_has_suffix(prefix, m) for m in modules)
|
| 591 |
+
|
| 592 |
+
def _bits_for_non_routed(self, prefix: str) -> int:
|
| 593 |
+
"""Get K bits for non_routed_exl3 layer, checking layers dict first, then suffix form."""
|
| 594 |
+
if not self.non_routed_exl3:
|
| 595 |
+
return self.bits
|
| 596 |
+
# Check layers dict first
|
| 597 |
+
layers = self.non_routed_exl3.get("layers", {})
|
| 598 |
+
if layers and prefix in layers:
|
| 599 |
+
layer_cfg = layers[prefix]
|
| 600 |
+
if "bits" in layer_cfg:
|
| 601 |
+
return int(layer_cfg["bits"])
|
| 602 |
+
return int(self.non_routed_exl3.get("bits", self.bits))
|
| 603 |
+
# Fall back to suffix matching
|
| 604 |
+
modules = self.non_routed_exl3.get("modules", [])
|
| 605 |
+
matched_suffix = None
|
| 606 |
+
for suffix in modules:
|
| 607 |
+
if _prefix_has_suffix(prefix, suffix):
|
| 608 |
+
matched_suffix = suffix
|
| 609 |
+
break
|
| 610 |
+
if matched_suffix is None:
|
| 611 |
+
return self.bits
|
| 612 |
+
# Check layer_bits override for this suffix
|
| 613 |
+
layer_bits = self.non_routed_exl3.get("layer_bits", {})
|
| 614 |
+
if matched_suffix in layer_bits:
|
| 615 |
+
return int(layer_bits[matched_suffix])
|
| 616 |
+
# Fall back to non_routed_exl3 bits or base bits
|
| 617 |
+
return int(self.non_routed_exl3.get("bits", self.bits))
|
| 618 |
+
|
| 619 |
+
def _bf16_shards_for(self, prefix: str) -> list[int]:
|
| 620 |
+
"""Get bf16 shard indices for a non_routed_exl3 layer from the layers dict."""
|
| 621 |
+
if not self.non_routed_exl3:
|
| 622 |
+
return []
|
| 623 |
+
layers = self.non_routed_exl3.get("layers", {})
|
| 624 |
+
if layers and prefix in layers:
|
| 625 |
+
layer_cfg = layers[prefix]
|
| 626 |
+
return list(layer_cfg.get("bf16_shards", []))
|
| 627 |
+
return []
|
| 628 |
+
|
| 629 |
+
def get_supported_act_dtypes(self) -> list[torch.dtype]:
|
| 630 |
+
return [torch.bfloat16, torch.float16, torch.float32]
|
| 631 |
+
|
| 632 |
+
@classmethod
|
| 633 |
+
def get_min_capability(cls) -> int:
|
| 634 |
+
# LinearEXL3 uses CUDA >= Ampere; GB10 is SM121.
|
| 635 |
+
return 80
|
| 636 |
+
|
| 637 |
+
@staticmethod
|
| 638 |
+
def get_config_filenames() -> list[str]:
|
| 639 |
+
return ["quantization_config.json"]
|
| 640 |
+
|
| 641 |
+
@classmethod
|
| 642 |
+
def from_config(cls, config: dict[str, Any]) -> "Exl3Config":
|
| 643 |
+
skip = {
|
| 644 |
+
"bits",
|
| 645 |
+
"codebook",
|
| 646 |
+
"scope",
|
| 647 |
+
"quant_method",
|
| 648 |
+
# Some packs ship a large per-tensor ledger here; keep it off the config object.
|
| 649 |
+
"tensor_storage",
|
| 650 |
+
"non_routed_exl3",
|
| 651 |
+
"non_routed_quantization",
|
| 652 |
+
"mtp_experts",
|
| 653 |
+
"mtp_experts_start_layer",
|
| 654 |
+
}
|
| 655 |
+
inst = cls(
|
| 656 |
+
bits=int(config.get("bits", 4)),
|
| 657 |
+
codebook=str(config.get("codebook", "mcg")),
|
| 658 |
+
scope=str(config.get("scope", "glm53_routed_experts_only")),
|
| 659 |
+
non_routed_exl3=config.get("non_routed_exl3"),
|
| 660 |
+
**{k: v for k, v in config.items() if k not in skip},
|
| 661 |
+
)
|
| 662 |
+
# __init__ swallows unknown kwargs; stash the delegation dict explicitly.
|
| 663 |
+
inst.non_routed_quantization = config.get("non_routed_quantization")
|
| 664 |
+
# "bf16_as_stored": dense linears are BF16 tensors; never delegate them
|
| 665 |
+
# (the delegate still serves source-format MTP experts).
|
| 666 |
+
inst.non_routed_dtype_policy = str(config.get("non_routed_dtype_policy", ""))
|
| 667 |
+
# Mixed-format packs: draft/MTP blocks appended past the main stack can
|
| 668 |
+
# keep their experts in the source format (e.g. MXFP4). Declare
|
| 669 |
+
# mtp_experts: "source" plus mtp_experts_start_layer: <first draft
|
| 670 |
+
# layer index>; those layers delegate to non_routed_quantization.
|
| 671 |
+
inst.mtp_experts = str(config.get("mtp_experts", "exl3"))
|
| 672 |
+
inst.mtp_experts_start_layer = config.get("mtp_experts_start_layer")
|
| 673 |
+
return inst
|
| 674 |
+
|
| 675 |
+
@classmethod
|
| 676 |
+
def override_quantization_method(
|
| 677 |
+
cls,
|
| 678 |
+
hf_quant_cfg: dict[str, Any],
|
| 679 |
+
user_quant: str | None,
|
| 680 |
+
hf_config: Any = None,
|
| 681 |
+
) -> str | None:
|
| 682 |
+
method = str((hf_quant_cfg or {}).get("quant_method", "")).lower()
|
| 683 |
+
if method == "exl3":
|
| 684 |
+
return "exl3"
|
| 685 |
+
return None
|
| 686 |
+
|
| 687 |
+
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
|
| 688 |
+
from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
|
| 689 |
+
|
| 690 |
+
if isinstance(layer, RoutedExperts):
|
| 691 |
+
# Draft/MTP blocks construct with plain layers.N prefixes (the
|
| 692 |
+
# mtp_block name appears only in parameter paths), so gate by
|
| 693 |
+
# declared layer index, never by name.
|
| 694 |
+
if getattr(self, "mtp_experts", "exl3") == "source":
|
| 695 |
+
_start = getattr(self, "mtp_experts_start_layer", None)
|
| 696 |
+
_lm = re.search(r"layers\.(\d+)\.", prefix)
|
| 697 |
+
if _start is not None and _lm and int(_lm.group(1)) >= int(_start):
|
| 698 |
+
d = self._non_routed_delegate()
|
| 699 |
+
if d is not None:
|
| 700 |
+
dm = d.get_quant_method(layer, prefix)
|
| 701 |
+
if dm is not None:
|
| 702 |
+
return dm
|
| 703 |
+
return Exl3MoEMethod(
|
| 704 |
+
layer.moe_config, self, bits=self.bits_for_prefix(prefix)
|
| 705 |
+
)
|
| 706 |
+
if isinstance(layer, LinearBase):
|
| 707 |
+
# Check if this LinearBase should use non_routed_exl3
|
| 708 |
+
if self._matches_non_routed_exl3(prefix):
|
| 709 |
+
bits = self._bits_for_non_routed(prefix)
|
| 710 |
+
return Exl3LinearMethod(self, bits=bits)
|
| 711 |
+
if getattr(self, "non_routed_dtype_policy", "") == "bf16_as_stored":
|
| 712 |
+
return UnquantizedLinearMethod()
|
| 713 |
+
d = self._non_routed_delegate()
|
| 714 |
+
if d is not None:
|
| 715 |
+
m = d.get_quant_method(layer, prefix)
|
| 716 |
+
if m is not None:
|
| 717 |
+
return m
|
| 718 |
+
return UnquantizedLinearMethod()
|
| 719 |
+
return None
|
| 720 |
+
|
| 721 |
+
def _non_routed_delegate(self):
|
| 722 |
+
# Packs that keep non-routed weights in the official source format
|
| 723 |
+
# (e.g. DeepSeek block-FP8) declare it under
|
| 724 |
+
# ``quantization_config.non_routed_quantization``; delegate those
|
| 725 |
+
# layers to the matching quant method so arch-specific fp8 forward
|
| 726 |
+
# paths get real scale tensors. Absent key = unquantized (GLM).
|
| 727 |
+
if not hasattr(self, "_nr_delegate_cached"):
|
| 728 |
+
self._nr_delegate_cached = None
|
| 729 |
+
nrq = getattr(self, "non_routed_quantization", None)
|
| 730 |
+
if isinstance(nrq, dict) and nrq.get("quant_method"):
|
| 731 |
+
from vllm.model_executor.layers.quantization import (
|
| 732 |
+
get_quantization_config,
|
| 733 |
+
)
|
| 734 |
+
for name in ("deepseek_v4_fp8", str(nrq.get("quant_method"))):
|
| 735 |
+
try:
|
| 736 |
+
cls = get_quantization_config(name)
|
| 737 |
+
self._nr_delegate_cached = cls.from_config(dict(nrq))
|
| 738 |
+
break
|
| 739 |
+
except Exception:
|
| 740 |
+
continue
|
| 741 |
+
return self._nr_delegate_cached
|
| 742 |
+
|
| 743 |
+
|
| 744 |
+
class Exl3MoEMethod(FusedMoEMethodBase):
|
| 745 |
+
"""Packed MCG trellis experts: create/load packed tensors, LinearEXL3 apply."""
|
| 746 |
+
|
| 747 |
+
def __init__(
|
| 748 |
+
self, moe, quant_config: Exl3Config, bits: int | None = None
|
| 749 |
+
) -> None:
|
| 750 |
+
super().__init__(moe)
|
| 751 |
+
self.quant_config = quant_config
|
| 752 |
+
# One method instance per RoutedExperts layer, so this is per-layer K.
|
| 753 |
+
self.bits = int(bits) if bits is not None else quant_config.bits
|
| 754 |
+
self._logged = False
|
| 755 |
+
|
| 756 |
+
def get_fused_moe_quant_config(self, layer: "RoutedExperts") -> FusedMoEQuantConfig | None:
|
| 757 |
+
return None
|
| 758 |
+
|
| 759 |
+
def create_weights(
|
| 760 |
+
self,
|
| 761 |
+
layer: "RoutedExperts",
|
| 762 |
+
num_experts: int,
|
| 763 |
+
hidden_size: int,
|
| 764 |
+
intermediate_size_per_partition: int,
|
| 765 |
+
params_dtype: torch.dtype,
|
| 766 |
+
**extra_weight_attrs,
|
| 767 |
+
) -> None:
|
| 768 |
+
del params_dtype
|
| 769 |
+
if hidden_size % 16 or intermediate_size_per_partition % 16:
|
| 770 |
+
raise ValueError(
|
| 771 |
+
"EXL3 trellis tiles are 16-wide; "
|
| 772 |
+
f"hidden={hidden_size} intermediate_local={intermediate_size_per_partition}"
|
| 773 |
+
)
|
| 774 |
+
k_words = self.bits * 16
|
| 775 |
+
in_tiles = hidden_size // 16
|
| 776 |
+
out_tiles = intermediate_size_per_partition // 16
|
| 777 |
+
|
| 778 |
+
extra = {k: v for k, v in extra_weight_attrs.items() if k != "weight_loader"}
|
| 779 |
+
|
| 780 |
+
# w13_* : stacked [expert, {gate=0, up=1}, ...] so the stock
|
| 781 |
+
# expert_params_mapping (experts.w13_ + suffix) hits these names.
|
| 782 |
+
w13_trellis = Parameter(
|
| 783 |
+
torch.empty(
|
| 784 |
+
num_experts, 2, in_tiles, out_tiles, k_words, dtype=torch.int16
|
| 785 |
+
),
|
| 786 |
+
requires_grad=False,
|
| 787 |
+
)
|
| 788 |
+
w13_suh = Parameter(
|
| 789 |
+
torch.empty(num_experts, 2, hidden_size, dtype=torch.float16),
|
| 790 |
+
requires_grad=False,
|
| 791 |
+
)
|
| 792 |
+
w13_svh = Parameter(
|
| 793 |
+
torch.empty(
|
| 794 |
+
num_experts, 2, intermediate_size_per_partition, dtype=torch.float16
|
| 795 |
+
),
|
| 796 |
+
requires_grad=False,
|
| 797 |
+
)
|
| 798 |
+
w13_mcg = Parameter(
|
| 799 |
+
torch.empty(num_experts, 2, 1, dtype=torch.int32),
|
| 800 |
+
requires_grad=False,
|
| 801 |
+
)
|
| 802 |
+
w2_trellis = Parameter(
|
| 803 |
+
torch.empty(
|
| 804 |
+
num_experts, out_tiles, in_tiles, k_words, dtype=torch.int16
|
| 805 |
+
),
|
| 806 |
+
requires_grad=False,
|
| 807 |
+
)
|
| 808 |
+
w2_suh = Parameter(
|
| 809 |
+
torch.empty(
|
| 810 |
+
num_experts, intermediate_size_per_partition, dtype=torch.float16
|
| 811 |
+
),
|
| 812 |
+
requires_grad=False,
|
| 813 |
+
)
|
| 814 |
+
w2_svh = Parameter(
|
| 815 |
+
torch.empty(num_experts, hidden_size, dtype=torch.float16),
|
| 816 |
+
requires_grad=False,
|
| 817 |
+
)
|
| 818 |
+
w2_mcg = Parameter(
|
| 819 |
+
torch.empty(num_experts, 1, dtype=torch.int32),
|
| 820 |
+
requires_grad=False,
|
| 821 |
+
)
|
| 822 |
+
|
| 823 |
+
packed = {
|
| 824 |
+
"w13_trellis": w13_trellis,
|
| 825 |
+
"w13_suh": w13_suh,
|
| 826 |
+
"w13_svh": w13_svh,
|
| 827 |
+
"w13_mcg": w13_mcg,
|
| 828 |
+
"w2_trellis": w2_trellis,
|
| 829 |
+
"w2_suh": w2_suh,
|
| 830 |
+
"w2_svh": w2_svh,
|
| 831 |
+
"w2_mcg": w2_mcg,
|
| 832 |
+
}
|
| 833 |
+
for name, param in packed.items():
|
| 834 |
+
layer.register_parameter(name, param)
|
| 835 |
+
set_weight_attrs(param, extra)
|
| 836 |
+
param.weight_loader = self._load_exl3
|
| 837 |
+
param._exl3_owner = layer
|
| 838 |
+
if hasattr(layer, "w13_weight") or hasattr(layer, "w2_weight"):
|
| 839 |
+
raise RuntimeError("EXL3 create_weights must not allocate dense expert weights")
|
| 840 |
+
|
| 841 |
+
layer._exl3_hidden_size = hidden_size
|
| 842 |
+
layer._exl3_intermediate_local = intermediate_size_per_partition
|
| 843 |
+
layer._exl3_k_words = k_words
|
| 844 |
+
layer._exl3_bits = self.bits
|
| 845 |
+
|
| 846 |
+
def _load_exl3(
|
| 847 |
+
self,
|
| 848 |
+
param: Parameter,
|
| 849 |
+
loaded_weight: torch.Tensor,
|
| 850 |
+
weight_name: str,
|
| 851 |
+
shard_id: str = "w1",
|
| 852 |
+
expert_id: int = 0,
|
| 853 |
+
return_success: bool = False,
|
| 854 |
+
) -> bool | None:
|
| 855 |
+
from vllm.distributed import (
|
| 856 |
+
get_tensor_model_parallel_rank,
|
| 857 |
+
get_tensor_model_parallel_world_size,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
layer = param
|
| 861 |
+
# param is the Parameter; expert_id is already physical. Map to local
|
| 862 |
+
# via the owning module if present on the weight_loader closure... we
|
| 863 |
+
# look up from param's __dict__ after register. RoutedExperts.weight_loader
|
| 864 |
+
# maps global→local; glm5next calls *our* loader, so map here.
|
| 865 |
+
owner = getattr(param, "_exl3_owner", None)
|
| 866 |
+
if owner is not None:
|
| 867 |
+
local_id = owner._map_global_expert_id_to_local_expert_id(expert_id)
|
| 868 |
+
if local_id == -1:
|
| 869 |
+
return False if return_success else None
|
| 870 |
+
expert_id = local_id
|
| 871 |
+
|
| 872 |
+
tp_rank = get_tensor_model_parallel_rank()
|
| 873 |
+
tp_size = get_tensor_model_parallel_world_size()
|
| 874 |
+
suffix = _suffix_from_mapped_name(weight_name)
|
| 875 |
+
loaded = loaded_weight.detach().contiguous()
|
| 876 |
+
|
| 877 |
+
if shard_id in ("w1", "w3"):
|
| 878 |
+
shard_idx = 0 if shard_id == "w1" else 1
|
| 879 |
+
sharded = shard_exl3_col(loaded, suffix, tp_rank, tp_size)
|
| 880 |
+
dest = param.data[expert_id, shard_idx]
|
| 881 |
+
elif shard_id == "w2":
|
| 882 |
+
sharded = shard_exl3_row(loaded, suffix, tp_rank, tp_size)
|
| 883 |
+
dest = param.data[expert_id]
|
| 884 |
+
else:
|
| 885 |
+
raise ValueError(f"unknown EXL3 shard_id={shard_id}")
|
| 886 |
+
|
| 887 |
+
if tuple(dest.shape) != tuple(sharded.shape):
|
| 888 |
+
raise RuntimeError(
|
| 889 |
+
f"EXL3 load shape mismatch {weight_name} shard={shard_id} "
|
| 890 |
+
f"expert={expert_id}: dest {tuple(dest.shape)} != "
|
| 891 |
+
f"loaded {tuple(sharded.shape)}"
|
| 892 |
+
)
|
| 893 |
+
dest.copy_(sharded)
|
| 894 |
+
return True if return_success else None
|
| 895 |
+
|
| 896 |
+
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
| 897 |
+
if not hasattr(layer, "w13_trellis"):
|
| 898 |
+
return
|
| 899 |
+
# Bind owner for any late loads; stitch LinearEXL3 handles.
|
| 900 |
+
for name in (
|
| 901 |
+
"w13_trellis",
|
| 902 |
+
"w13_suh",
|
| 903 |
+
"w13_svh",
|
| 904 |
+
"w13_mcg",
|
| 905 |
+
"w2_trellis",
|
| 906 |
+
"w2_suh",
|
| 907 |
+
"w2_svh",
|
| 908 |
+
"w2_mcg",
|
| 909 |
+
):
|
| 910 |
+
getattr(layer, name)._exl3_owner = layer
|
| 911 |
+
|
| 912 |
+
mcg13 = layer.w13_mcg.reshape(-1)
|
| 913 |
+
mcg2 = layer.w2_mcg.reshape(-1)
|
| 914 |
+
if not torch.all(mcg13 == MCG_MARKER_SIGNED_INT32) or not torch.all(
|
| 915 |
+
mcg2 == MCG_MARKER_SIGNED_INT32
|
| 916 |
+
):
|
| 917 |
+
raise RuntimeError(
|
| 918 |
+
"EXL3 mcg marker is not the MCG int32 0xCBAC1FED / "
|
| 919 |
+
f"{MCG_MARKER_SIGNED_INT32}; packed ABI mismatch"
|
| 920 |
+
)
|
| 921 |
+
|
| 922 |
+
n_exp = int(layer.w13_trellis.shape[0])
|
| 923 |
+
inners: list[dict[str, Any]] = []
|
| 924 |
+
for e in range(n_exp):
|
| 925 |
+
gate = make_linear_exl3(
|
| 926 |
+
layer.w13_trellis[e, 0],
|
| 927 |
+
layer.w13_suh[e, 0],
|
| 928 |
+
layer.w13_svh[e, 0],
|
| 929 |
+
layer.w13_mcg[e, 0],
|
| 930 |
+
)
|
| 931 |
+
up = make_linear_exl3(
|
| 932 |
+
layer.w13_trellis[e, 1],
|
| 933 |
+
layer.w13_suh[e, 1],
|
| 934 |
+
layer.w13_svh[e, 1],
|
| 935 |
+
layer.w13_mcg[e, 1],
|
| 936 |
+
)
|
| 937 |
+
down = make_linear_exl3(
|
| 938 |
+
layer.w2_trellis[e],
|
| 939 |
+
layer.w2_suh[e],
|
| 940 |
+
layer.w2_svh[e],
|
| 941 |
+
layer.w2_mcg[e],
|
| 942 |
+
)
|
| 943 |
+
inners.append({"gate": gate, "up": up, "down": down})
|
| 944 |
+
layer._exl3_inners = inners
|
| 945 |
+
fused_ok = False
|
| 946 |
+
fused_err = None
|
| 947 |
+
if fused_moe_enabled():
|
| 948 |
+
try:
|
| 949 |
+
import exllamav3_ext
|
| 950 |
+
|
| 951 |
+
if hasattr(exllamav3_ext, "exl3_moe"):
|
| 952 |
+
build_exl3_fused_state(layer, inners)
|
| 953 |
+
fused_ok = True
|
| 954 |
+
else:
|
| 955 |
+
fused_err = "exllamav3_ext.exl3_moe missing"
|
| 956 |
+
except Exception as exc:
|
| 957 |
+
fused_err = repr(exc)
|
| 958 |
+
layer._exl3_ptrs = None
|
| 959 |
+
if not self._logged and self.bits != self.quant_config.bits:
|
| 960 |
+
logger.info(
|
| 961 |
+
"EXL3 per-layer K override: layer prefix %s uses bits=%d (base %d)",
|
| 962 |
+
getattr(layer, "layer_name", None) or getattr(layer, "prefix", "?"),
|
| 963 |
+
self.bits,
|
| 964 |
+
self.quant_config.bits,
|
| 965 |
+
)
|
| 966 |
+
if not self._logged:
|
| 967 |
+
if fused_ok:
|
| 968 |
+
logger.info(
|
| 969 |
+
"EXL3 MCG trellis engaged for routed experts: bits=%s "
|
| 970 |
+
"experts_local=%s hidden=%s intermediate_local=%s "
|
| 971 |
+
"fused_moe=exl3_moe concurrency=%s "
|
| 972 |
+
"(no BF16 expert reconstruct at load)",
|
| 973 |
+
self.bits,
|
| 974 |
+
n_exp,
|
| 975 |
+
layer._exl3_hidden_size,
|
| 976 |
+
layer._exl3_intermediate_local,
|
| 977 |
+
getattr(layer, "_exl3_fused_concurrency", "?"),
|
| 978 |
+
)
|
| 979 |
+
else:
|
| 980 |
+
logger.info(
|
| 981 |
+
"EXL3 MCG trellis engaged for routed experts: bits=%s "
|
| 982 |
+
"experts_local=%s hidden=%s intermediate_local=%s "
|
| 983 |
+
"fused_moe=python_loop (%s) "
|
| 984 |
+
"(no BF16 expert reconstruct at load)",
|
| 985 |
+
self.bits,
|
| 986 |
+
n_exp,
|
| 987 |
+
layer._exl3_hidden_size,
|
| 988 |
+
layer._exl3_intermediate_local,
|
| 989 |
+
fused_err or "EXL3_FUSED_MOE=0",
|
| 990 |
+
)
|
| 991 |
+
self._logged = True
|
| 992 |
+
|
| 993 |
+
def apply(
|
| 994 |
+
self,
|
| 995 |
+
layer: "RoutedExperts",
|
| 996 |
+
x: torch.Tensor,
|
| 997 |
+
topk_weights: torch.Tensor,
|
| 998 |
+
topk_ids: torch.Tensor,
|
| 999 |
+
shared_experts: "SharedExperts | None",
|
| 1000 |
+
shared_experts_input: torch.Tensor | None,
|
| 1001 |
+
) -> torch.Tensor:
|
| 1002 |
+
del shared_experts, shared_experts_input
|
| 1003 |
+
limit = getattr(self.moe, "swiglu_limit", None) or SWIGLU_LIMIT_DEFAULT
|
| 1004 |
+
return apply_exl3_experts(
|
| 1005 |
+
x, topk_ids, topk_weights, layer, limit=float(limit)
|
| 1006 |
+
)
|
| 1007 |
+
|
| 1008 |
+
|
| 1009 |
+
class Exl3LinearMethod(LinearMethodBase):
|
| 1010 |
+
"""Non-routed (dense) EXL3 linear method for QKV/MLP dense projections.
|
| 1011 |
+
|
| 1012 |
+
This method handles trellis/suh/svh/mcg parameters for non-routed dense
|
| 1013 |
+
linear layers, building LinearEXL3 objects after weight loading and applying
|
| 1014 |
+
them with proper TP slicing and shard concatenation.
|
| 1015 |
+
"""
|
| 1016 |
+
|
| 1017 |
+
def __init__(self, quant_config: Exl3Config, bits: int | None = None) -> None:
|
| 1018 |
+
self.quant_config = quant_config
|
| 1019 |
+
self.bits = int(bits) if bits is not None else quant_config.bits
|
| 1020 |
+
self._logged = False
|
| 1021 |
+
|
| 1022 |
+
def create_weights(
|
| 1023 |
+
self,
|
| 1024 |
+
layer,
|
| 1025 |
+
input_size_per_partition: int,
|
| 1026 |
+
output_partition_sizes: list[int],
|
| 1027 |
+
input_size: int,
|
| 1028 |
+
output_size: int,
|
| 1029 |
+
params_dtype: torch.dtype,
|
| 1030 |
+
**extra_weight_attrs,
|
| 1031 |
+
) -> None:
|
| 1032 |
+
from vllm.model_executor.layers.linear import (
|
| 1033 |
+
ColumnParallelLinear,
|
| 1034 |
+
RowParallelLinear,
|
| 1035 |
+
QKVParallelLinear,
|
| 1036 |
+
MergedColumnParallelLinear,
|
| 1037 |
+
)
|
| 1038 |
+
from vllm.distributed import (
|
| 1039 |
+
get_tensor_model_parallel_rank,
|
| 1040 |
+
get_tensor_model_parallel_world_size,
|
| 1041 |
+
)
|
| 1042 |
+
|
| 1043 |
+
# Determine layer type and shard behavior
|
| 1044 |
+
n_shards = len(output_partition_sizes)
|
| 1045 |
+
is_row_parallel = isinstance(layer, RowParallelLinear)
|
| 1046 |
+
is_col_parallel = isinstance(layer, ColumnParallelLinear)
|
| 1047 |
+
is_qkv_parallel = isinstance(layer, QKVParallelLinear)
|
| 1048 |
+
is_merged_col_parallel = isinstance(layer, MergedColumnParallelLinear)
|
| 1049 |
+
global_tp_rank = get_tensor_model_parallel_rank()
|
| 1050 |
+
global_tp_size = get_tensor_model_parallel_world_size()
|
| 1051 |
+
effective_tp_size, effective_tp_rank = _effective_tp_for_layer(
|
| 1052 |
+
layer, global_tp_size, global_tp_rank
|
| 1053 |
+
)
|
| 1054 |
+
is_wo_a = _is_wo_a_layer(layer)
|
| 1055 |
+
# LNA-LAB: EXL3 wo_a is one rank-local group only; TP<8 needs a future
|
| 1056 |
+
# multi-group path and must fail before an ambiguous partial load.
|
| 1057 |
+
if is_wo_a and effective_tp_size != 8:
|
| 1058 |
+
raise NotImplementedError(
|
| 1059 |
+
"EXL3 wo_a requires TP=8 (one local group per rank); "
|
| 1060 |
+
f"got effective tp_size={effective_tp_size}"
|
| 1061 |
+
)
|
| 1062 |
+
|
| 1063 |
+
# For column-parallel: output dimension is sharded (each shard has different outputs)
|
| 1064 |
+
# For row-parallel: input dimension is sharded (each shard has same input, different outputs)
|
| 1065 |
+
if is_row_parallel:
|
| 1066 |
+
# Input is partitioned across ranks, each rank gets full input height
|
| 1067 |
+
in_per_partition = input_size_per_partition
|
| 1068 |
+
else:
|
| 1069 |
+
# Column-parallel or unsharded: each rank gets full input
|
| 1070 |
+
in_per_partition = input_size_per_partition
|
| 1071 |
+
|
| 1072 |
+
# Get bf16_shards from config (may be empty)
|
| 1073 |
+
bf16_shards = self.quant_config._bf16_shards_for(getattr(layer, "prefix", ""))
|
| 1074 |
+
if bf16_shards:
|
| 1075 |
+
# LNA-LAB: the guard follows effective, not process-global, TP.
|
| 1076 |
+
if effective_tp_size > 1:
|
| 1077 |
+
raise RuntimeError(
|
| 1078 |
+
"EXL3 bf16 shards are not supported with effective TP size > 1; "
|
| 1079 |
+
f"tp_size={effective_tp_size}"
|
| 1080 |
+
)
|
| 1081 |
+
|
| 1082 |
+
# K words per shard
|
| 1083 |
+
k_words = self.bits * 16
|
| 1084 |
+
|
| 1085 |
+
# Validate tile alignment for all shards
|
| 1086 |
+
for i, out_size in enumerate(output_partition_sizes):
|
| 1087 |
+
if in_per_partition % 16 or out_size % 16:
|
| 1088 |
+
raise ValueError(
|
| 1089 |
+
f"EXL3 trellis tiles are 16-wide; "
|
| 1090 |
+
f"shard {i}: in={in_per_partition} out={out_size}"
|
| 1091 |
+
)
|
| 1092 |
+
|
| 1093 |
+
in_tiles = in_per_partition // 16
|
| 1094 |
+
out_tiles_list = [s // 16 for s in output_partition_sizes]
|
| 1095 |
+
total_out_tiles = sum(out_tiles_list)
|
| 1096 |
+
|
| 1097 |
+
# Allocate fused trellis covering all shards (dim1 will be narrow per-shard)
|
| 1098 |
+
trellis_param = Parameter(
|
| 1099 |
+
torch.empty(in_tiles, total_out_tiles, k_words, dtype=torch.int16),
|
| 1100 |
+
requires_grad=False,
|
| 1101 |
+
)
|
| 1102 |
+
# Per-shard suh (one per shard, each covers this rank's input partition)
|
| 1103 |
+
suh_param = Parameter(
|
| 1104 |
+
torch.empty(n_shards, in_per_partition, dtype=torch.float16),
|
| 1105 |
+
requires_grad=False,
|
| 1106 |
+
)
|
| 1107 |
+
# Per-shard svh (one per shard, concatenated)
|
| 1108 |
+
svh_param = Parameter(
|
| 1109 |
+
torch.empty(sum(output_partition_sizes), dtype=torch.float16),
|
| 1110 |
+
requires_grad=False,
|
| 1111 |
+
)
|
| 1112 |
+
# Per-shard mcg and mul1 markers (both registered, one will be nonzero)
|
| 1113 |
+
mcg_param = Parameter(
|
| 1114 |
+
torch.zeros(n_shards, 1, dtype=torch.int32),
|
| 1115 |
+
requires_grad=False,
|
| 1116 |
+
)
|
| 1117 |
+
mul1_param = Parameter(
|
| 1118 |
+
torch.zeros(n_shards, 1, dtype=torch.int32),
|
| 1119 |
+
requires_grad=False,
|
| 1120 |
+
)
|
| 1121 |
+
|
| 1122 |
+
# LNA-LAB: wo_a has no dense weight shim; _o_proj consumes _exl3_linears.
|
| 1123 |
+
weight_param = None
|
| 1124 |
+
if not is_wo_a:
|
| 1125 |
+
# Staging parameter for bf16 shards: rows are concatenated bf16 weights
|
| 1126 |
+
bf16_rows = sum(output_partition_sizes[i] for i in bf16_shards)
|
| 1127 |
+
weight_param = Parameter(
|
| 1128 |
+
torch.empty(bf16_rows, in_per_partition, dtype=params_dtype),
|
| 1129 |
+
requires_grad=False,
|
| 1130 |
+
)
|
| 1131 |
+
|
| 1132 |
+
layer.register_parameter("trellis", trellis_param)
|
| 1133 |
+
layer.register_parameter("suh", suh_param)
|
| 1134 |
+
layer.register_parameter("svh", svh_param)
|
| 1135 |
+
layer.register_parameter("mcg", mcg_param)
|
| 1136 |
+
layer.register_parameter("mul1", mul1_param)
|
| 1137 |
+
if weight_param is not None:
|
| 1138 |
+
layer.register_parameter("weight", weight_param)
|
| 1139 |
+
|
| 1140 |
+
# Custom weight loader
|
| 1141 |
+
extra = {k: v for k, v in extra_weight_attrs.items() if k != "weight_loader"}
|
| 1142 |
+
set_weight_attrs(trellis_param, extra)
|
| 1143 |
+
set_weight_attrs(suh_param, extra)
|
| 1144 |
+
set_weight_attrs(svh_param, extra)
|
| 1145 |
+
set_weight_attrs(mcg_param, extra)
|
| 1146 |
+
set_weight_attrs(mul1_param, extra)
|
| 1147 |
+
if weight_param is not None:
|
| 1148 |
+
set_weight_attrs(weight_param, extra)
|
| 1149 |
+
|
| 1150 |
+
# vLLM calls ``weight_loader(param, loaded_weight[, shard_id])`` and
|
| 1151 |
+
# never passes the checkpoint name, so bind the tensor kind per param.
|
| 1152 |
+
for suffix, p in (
|
| 1153 |
+
("trellis", trellis_param),
|
| 1154 |
+
("suh", suh_param),
|
| 1155 |
+
("svh", svh_param),
|
| 1156 |
+
("mcg", mcg_param),
|
| 1157 |
+
("mul1", mul1_param),
|
| 1158 |
+
):
|
| 1159 |
+
p.weight_loader = self._make_weight_loader(
|
| 1160 |
+
suffix,
|
| 1161 |
+
n_shards,
|
| 1162 |
+
output_partition_sizes,
|
| 1163 |
+
is_row_parallel,
|
| 1164 |
+
bf16_shards,
|
| 1165 |
+
effective_tp_rank,
|
| 1166 |
+
effective_tp_size,
|
| 1167 |
+
is_wo_a,
|
| 1168 |
+
)
|
| 1169 |
+
if weight_param is not None:
|
| 1170 |
+
weight_param.weight_loader = self._make_weight_loader(
|
| 1171 |
+
"weight",
|
| 1172 |
+
n_shards,
|
| 1173 |
+
output_partition_sizes,
|
| 1174 |
+
is_row_parallel,
|
| 1175 |
+
bf16_shards,
|
| 1176 |
+
effective_tp_rank,
|
| 1177 |
+
effective_tp_size,
|
| 1178 |
+
is_wo_a,
|
| 1179 |
+
)
|
| 1180 |
+
|
| 1181 |
+
# Store metadata
|
| 1182 |
+
layer._exl3_linear_n_shards = n_shards
|
| 1183 |
+
layer._exl3_linear_output_partition_sizes = output_partition_sizes
|
| 1184 |
+
layer._exl3_linear_input_size_per_partition = in_per_partition
|
| 1185 |
+
layer._exl3_linear_is_row_parallel = is_row_parallel
|
| 1186 |
+
layer._exl3_linear_is_qkv = is_qkv_parallel
|
| 1187 |
+
layer._exl3_linear_is_merged = is_merged_col_parallel
|
| 1188 |
+
layer._exl3_linear_bf16_shards = bf16_shards
|
| 1189 |
+
# LNA-LAB: retain both process and effective geometry for diagnostics/loaders.
|
| 1190 |
+
layer._exl3_linear_tp_size = effective_tp_size
|
| 1191 |
+
layer._exl3_linear_tp_rank = effective_tp_rank
|
| 1192 |
+
layer._exl3_linear_is_wo_a = is_wo_a
|
| 1193 |
+
|
| 1194 |
+
def _make_weight_loader(
|
| 1195 |
+
self,
|
| 1196 |
+
suffix,
|
| 1197 |
+
n_shards,
|
| 1198 |
+
output_partition_sizes,
|
| 1199 |
+
is_row_parallel,
|
| 1200 |
+
bf16_shards,
|
| 1201 |
+
tp_rank,
|
| 1202 |
+
tp_size,
|
| 1203 |
+
is_wo_a,
|
| 1204 |
+
):
|
| 1205 |
+
"""Create a weight_loader closure for EXL3 linear parameters."""
|
| 1206 |
+
|
| 1207 |
+
def weight_loader(
|
| 1208 |
+
param: Parameter,
|
| 1209 |
+
loaded_weight: torch.Tensor,
|
| 1210 |
+
loaded_shard_id: str | int | None = None,
|
| 1211 |
+
) -> None:
|
| 1212 |
+
# Map shard_id to shard index
|
| 1213 |
+
shard_idx = 0
|
| 1214 |
+
if loaded_shard_id is not None:
|
| 1215 |
+
if isinstance(loaded_shard_id, str):
|
| 1216 |
+
# "q", "k", "v" for QKV layers
|
| 1217 |
+
shard_map = {"q": 0, "k": 1, "v": 2}
|
| 1218 |
+
if loaded_shard_id in shard_map:
|
| 1219 |
+
shard_idx = shard_map[loaded_shard_id]
|
| 1220 |
+
elif is_wo_a and _wo_a_slice_rank(loaded_shard_id) is not None:
|
| 1221 |
+
# LNA-LAB: model.py supplies the wo_a group rank as the
|
| 1222 |
+
# shard id; the special loader below selects this rank.
|
| 1223 |
+
shard_idx = 0
|
| 1224 |
+
else:
|
| 1225 |
+
raise ValueError(
|
| 1226 |
+
f"unknown shard_id={loaded_shard_id} for EXL3 linear"
|
| 1227 |
+
)
|
| 1228 |
+
elif isinstance(loaded_shard_id, int):
|
| 1229 |
+
# LNA-LAB: wo_a integer ids are group ranks, not local
|
| 1230 |
+
# fused-shard indices (there is exactly one local shard).
|
| 1231 |
+
shard_idx = 0 if is_wo_a else loaded_shard_id
|
| 1232 |
+
if not is_wo_a and shard_idx >= n_shards:
|
| 1233 |
+
raise ValueError(
|
| 1234 |
+
f"shard_idx={shard_idx} out of range for n_shards={n_shards}"
|
| 1235 |
+
)
|
| 1236 |
+
|
| 1237 |
+
if is_wo_a:
|
| 1238 |
+
# LNA-LAB: each wo_a.slice.N is already a complete local
|
| 1239 |
+
# [group_width, o_lora_rank] tensor; never narrow it again.
|
| 1240 |
+
slice_rank = _wo_a_slice_rank(loaded_shard_id)
|
| 1241 |
+
if slice_rank is None:
|
| 1242 |
+
raise ValueError(
|
| 1243 |
+
"EXL3 wo_a loader requires a wo_a.slice.N shard id"
|
| 1244 |
+
)
|
| 1245 |
+
if slice_rank != tp_rank:
|
| 1246 |
+
return
|
| 1247 |
+
loaded = loaded_weight.detach().contiguous()
|
| 1248 |
+
if suffix == "trellis":
|
| 1249 |
+
dest = param.data
|
| 1250 |
+
elif suffix == "suh":
|
| 1251 |
+
dest = param.data[0]
|
| 1252 |
+
elif suffix == "svh":
|
| 1253 |
+
dest = param.data
|
| 1254 |
+
elif suffix in ("mcg", "mul1"):
|
| 1255 |
+
dest = param.data[0]
|
| 1256 |
+
if tuple(dest.shape) != (1,):
|
| 1257 |
+
raise RuntimeError(
|
| 1258 |
+
f"EXL3 wo_a {suffix} marker shape mismatch: "
|
| 1259 |
+
f"expected (1,) got {tuple(dest.shape)}"
|
| 1260 |
+
)
|
| 1261 |
+
if loaded.numel():
|
| 1262 |
+
dest[0] = int(loaded.item())
|
| 1263 |
+
return
|
| 1264 |
+
else:
|
| 1265 |
+
raise ValueError(f"unknown EXL3 wo_a suffix={suffix}")
|
| 1266 |
+
if tuple(dest.shape) != tuple(loaded.shape):
|
| 1267 |
+
raise RuntimeError(
|
| 1268 |
+
f"EXL3 wo_a load shape mismatch slice={slice_rank} "
|
| 1269 |
+
f"suffix={suffix}: dest {tuple(dest.shape)} != "
|
| 1270 |
+
f"loaded {tuple(loaded.shape)}"
|
| 1271 |
+
)
|
| 1272 |
+
dest.copy_(loaded)
|
| 1273 |
+
return
|
| 1274 |
+
|
| 1275 |
+
# Special handling for weight (bf16 staging) and markers
|
| 1276 |
+
if suffix in ("weight", "mcg", "mul1"):
|
| 1277 |
+
# Weight parameter: only load bf16 shards, discard EXL3 shards
|
| 1278 |
+
if suffix == "weight":
|
| 1279 |
+
# Check shape matches the expected shard size
|
| 1280 |
+
expected_out = output_partition_sizes[shard_idx]
|
| 1281 |
+
expected_in = param.shape[1]
|
| 1282 |
+
loaded_shape = loaded_weight.shape
|
| 1283 |
+
if loaded_shape[0] != expected_out or (len(loaded_shape) > 1 and loaded_shape[1] != expected_in):
|
| 1284 |
+
# After TP narrowing, expect (shard_out, in)
|
| 1285 |
+
tp_sharded = _narrow_tp(loaded_weight, 0, tp_rank, tp_size) if is_row_parallel else _narrow_tp(loaded_weight, 1, tp_rank, tp_size) if not is_row_parallel else loaded_weight
|
| 1286 |
+
if tuple(tp_sharded.shape) != (expected_out, expected_in):
|
| 1287 |
+
raise RuntimeError(
|
| 1288 |
+
f"EXL3 weight load shape mismatch shard={shard_idx}: "
|
| 1289 |
+
f"expected ({expected_out},{expected_in}) but got {tuple(loaded_weight.shape)} "
|
| 1290 |
+
f"(after TP: {tuple(tp_sharded.shape)})"
|
| 1291 |
+
)
|
| 1292 |
+
# If this shard is in bf16_shards, copy; otherwise discard
|
| 1293 |
+
if shard_idx in bf16_shards:
|
| 1294 |
+
bf16_idx = bf16_shards.index(shard_idx)
|
| 1295 |
+
bf16_row_start = sum(output_partition_sizes[i] for i in bf16_shards[:bf16_idx])
|
| 1296 |
+
bf16_row_end = bf16_row_start + expected_out
|
| 1297 |
+
param.data[bf16_row_start:bf16_row_end].copy_(loaded_weight.detach())
|
| 1298 |
+
# else: discard this EXL3 shard's stale BF16 weight
|
| 1299 |
+
return
|
| 1300 |
+
else:
|
| 1301 |
+
# Marker (mcg or mul1): store the value (will be 0 if marker not present)
|
| 1302 |
+
dest = param.data[shard_idx]
|
| 1303 |
+
if tuple(dest.shape) != (1,):
|
| 1304 |
+
raise RuntimeError(
|
| 1305 |
+
f"EXL3 {suffix} marker shape mismatch: expected (1,) got {tuple(dest.shape)}"
|
| 1306 |
+
)
|
| 1307 |
+
loaded_val = loaded_weight.detach().item() if loaded_weight.numel() > 0 else 0
|
| 1308 |
+
dest[0] = int(loaded_val)
|
| 1309 |
+
return
|
| 1310 |
+
|
| 1311 |
+
# Normal EXL3 suffix handling (trellis, suh, svh)
|
| 1312 |
+
loaded = loaded_weight.detach().contiguous()
|
| 1313 |
+
|
| 1314 |
+
# Apply TP slicing based on layer type
|
| 1315 |
+
if is_row_parallel:
|
| 1316 |
+
# Row-parallel: input is sharded, trellis dim 0 and suh dim 0
|
| 1317 |
+
sharded = shard_exl3_row(loaded, suffix, tp_rank, tp_size)
|
| 1318 |
+
else:
|
| 1319 |
+
# Column-parallel: output is sharded, trellis dim 1 and svh dim 0
|
| 1320 |
+
sharded = shard_exl3_col(loaded, suffix, tp_rank, tp_size)
|
| 1321 |
+
|
| 1322 |
+
# Copy into the right location
|
| 1323 |
+
if suffix == "trellis":
|
| 1324 |
+
# Trellis is fused; narrow dim1 for this shard
|
| 1325 |
+
out_tiles_start = sum(s // 16 for s in output_partition_sizes[:shard_idx])
|
| 1326 |
+
out_tiles_end = out_tiles_start + output_partition_sizes[shard_idx] // 16
|
| 1327 |
+
dest = param.data[:, out_tiles_start:out_tiles_end, :]
|
| 1328 |
+
elif suffix == "suh":
|
| 1329 |
+
# Suh per-shard
|
| 1330 |
+
dest = param.data[shard_idx]
|
| 1331 |
+
elif suffix == "svh":
|
| 1332 |
+
# Svh is concatenated; slice for this shard
|
| 1333 |
+
out_start = sum(output_partition_sizes[:shard_idx])
|
| 1334 |
+
out_end = out_start + output_partition_sizes[shard_idx]
|
| 1335 |
+
dest = param.data[out_start:out_end]
|
| 1336 |
+
else:
|
| 1337 |
+
raise ValueError(f"unknown EXL3 suffix={suffix}")
|
| 1338 |
+
|
| 1339 |
+
if tuple(dest.shape) != tuple(sharded.shape):
|
| 1340 |
+
raise RuntimeError(
|
| 1341 |
+
f"EXL3 linear load shape mismatch shard={shard_idx} "
|
| 1342 |
+
f"suffix={suffix}: dest {tuple(dest.shape)} != "
|
| 1343 |
+
f"loaded {tuple(sharded.shape)}"
|
| 1344 |
+
)
|
| 1345 |
+
dest.copy_(sharded)
|
| 1346 |
+
|
| 1347 |
+
return weight_loader
|
| 1348 |
+
|
| 1349 |
+
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
| 1350 |
+
if not hasattr(layer, "trellis"):
|
| 1351 |
+
return
|
| 1352 |
+
|
| 1353 |
+
# Get bf16 shards and verify exactly one marker per EXL3 shard
|
| 1354 |
+
n_shards = int(layer._exl3_linear_n_shards)
|
| 1355 |
+
output_sizes = layer._exl3_linear_output_partition_sizes
|
| 1356 |
+
bf16_shards = getattr(layer, "_exl3_linear_bf16_shards", [])
|
| 1357 |
+
|
| 1358 |
+
mcg_vals = layer.mcg.reshape(-1)
|
| 1359 |
+
mul1_vals = layer.mul1.reshape(-1)
|
| 1360 |
+
|
| 1361 |
+
for i in range(n_shards):
|
| 1362 |
+
# Skip marker checks for bf16 shards - they don't use LinearEXL3
|
| 1363 |
+
if i in bf16_shards:
|
| 1364 |
+
continue
|
| 1365 |
+
mcg_is_set = mcg_vals[i].item() != 0
|
| 1366 |
+
mul1_is_set = mul1_vals[i].item() != 0
|
| 1367 |
+
if mcg_is_set and mul1_is_set:
|
| 1368 |
+
raise RuntimeError(
|
| 1369 |
+
f"EXL3 linear shard {i}: both mcg and mul1 markers are set; "
|
| 1370 |
+
f"exactly one codebook marker must be present"
|
| 1371 |
+
)
|
| 1372 |
+
if not mcg_is_set and not mul1_is_set:
|
| 1373 |
+
raise RuntimeError(
|
| 1374 |
+
f"EXL3 linear shard {i}: neither mcg nor mul1 marker is set; "
|
| 1375 |
+
f"exactly one codebook marker must be present"
|
| 1376 |
+
)
|
| 1377 |
+
# Verify marker value
|
| 1378 |
+
if mcg_is_set and mcg_vals[i].item() != MCG_MARKER_SIGNED_INT32:
|
| 1379 |
+
raise RuntimeError(
|
| 1380 |
+
f"EXL3 linear shard {i}: mcg marker is {mcg_vals[i].item()}, "
|
| 1381 |
+
f"expected {MCG_MARKER_SIGNED_INT32}"
|
| 1382 |
+
)
|
| 1383 |
+
if mul1_is_set and mul1_vals[i].item() != MUL1_MARKER_SIGNED_INT32:
|
| 1384 |
+
raise RuntimeError(
|
| 1385 |
+
f"EXL3 linear shard {i}: mul1 marker is {mul1_vals[i].item()}, "
|
| 1386 |
+
f"expected {MUL1_MARKER_SIGNED_INT32}"
|
| 1387 |
+
)
|
| 1388 |
+
|
| 1389 |
+
# Build LinearEXL3 objects for EXL3 shards only (skip bf16 shards)
|
| 1390 |
+
linears = []
|
| 1391 |
+
for i in range(n_shards):
|
| 1392 |
+
if i in bf16_shards:
|
| 1393 |
+
# bf16 shards don't use LinearEXL3; store None as placeholder
|
| 1394 |
+
linears.append(None)
|
| 1395 |
+
continue
|
| 1396 |
+
out_tiles_start = sum(s // 16 for s in output_sizes[:i])
|
| 1397 |
+
out_tiles_end = out_tiles_start + output_sizes[i] // 16
|
| 1398 |
+
trellis_shard = layer.trellis[:, out_tiles_start:out_tiles_end, :].contiguous()
|
| 1399 |
+
suh_shard = layer.suh[i].contiguous()
|
| 1400 |
+
svh_shard = layer.svh[
|
| 1401 |
+
sum(output_sizes[:i]) : sum(output_sizes[: i + 1])
|
| 1402 |
+
].contiguous()
|
| 1403 |
+
mcg_shard = layer.mcg[i].contiguous() if mcg_vals[i].item() != 0 else None
|
| 1404 |
+
mul1_shard = layer.mul1[i].contiguous() if mul1_vals[i].item() != 0 else None
|
| 1405 |
+
|
| 1406 |
+
linear = make_linear_exl3(
|
| 1407 |
+
trellis_shard, suh_shard, svh_shard, mcg_shard, mul1_shard, out_dtype=torch.float16
|
| 1408 |
+
)
|
| 1409 |
+
linears.append(linear)
|
| 1410 |
+
|
| 1411 |
+
layer._exl3_linears = linears
|
| 1412 |
+
|
| 1413 |
+
# LNA-LAB: pre-tune the EXL3 cooperative GEMM for the decode row counts before vLLM
|
| 1414 |
+
# captures CUDA graphs. The exllamav3 autotuner times candidates with events and
|
| 1415 |
+
# stream syncs, which is illegal during capture; warming the launch/disk cache here
|
| 1416 |
+
# keeps capture on the tuned kernel instead of the static heuristic.
|
| 1417 |
+
prewarm = os.environ.get("LNA_EXL3_PREWARM_ROWS", "1,2,4,8,16")
|
| 1418 |
+
if prewarm:
|
| 1419 |
+
try:
|
| 1420 |
+
dev = linears[0].trellis.device if linears and linears[0] is not None else None
|
| 1421 |
+
if dev is not None:
|
| 1422 |
+
for r in (int(v) for v in prewarm.split(",") if v.strip()):
|
| 1423 |
+
for lin in linears:
|
| 1424 |
+
if lin is not None:
|
| 1425 |
+
x = torch.zeros((r, int(lin.in_features)), dtype=torch.float16, device=dev)
|
| 1426 |
+
lin.forward(x, {}, out_dtype=torch.float32)
|
| 1427 |
+
torch.cuda.synchronize(dev)
|
| 1428 |
+
except Exception as exc: # pragma: no cover - diagnostic only
|
| 1429 |
+
logger.warning("EXL3 prewarm skipped for %s: %s", getattr(layer, "prefix", "?"), exc)
|
| 1430 |
+
|
| 1431 |
+
# Keep bf16 weights if present, remove weight staging param if all loaded
|
| 1432 |
+
if bf16_shards and hasattr(layer, "weight"):
|
| 1433 |
+
bf16_rows = sum(output_sizes[i] for i in bf16_shards)
|
| 1434 |
+
if bf16_rows > 0:
|
| 1435 |
+
layer._exl3_bf16_weight = layer.weight.data.clone()
|
| 1436 |
+
# Delete weight staging param only if it has rows; empty param stays as placeholder
|
| 1437 |
+
if layer.weight.data.shape[0] > 0:
|
| 1438 |
+
try:
|
| 1439 |
+
delattr(layer, "weight")
|
| 1440 |
+
except Exception:
|
| 1441 |
+
pass
|
| 1442 |
+
elif hasattr(layer, "weight"):
|
| 1443 |
+
# No bf16 shards, delete the staging param
|
| 1444 |
+
try:
|
| 1445 |
+
delattr(layer, "weight")
|
| 1446 |
+
except Exception:
|
| 1447 |
+
pass
|
| 1448 |
+
|
| 1449 |
+
# Free fused parameters to avoid memory doubling
|
| 1450 |
+
for param_name in ("trellis", "suh", "svh", "mcg", "mul1"):
|
| 1451 |
+
if hasattr(layer, param_name):
|
| 1452 |
+
try:
|
| 1453 |
+
delattr(layer, param_name)
|
| 1454 |
+
except Exception:
|
| 1455 |
+
pass
|
| 1456 |
+
|
| 1457 |
+
def apply(
|
| 1458 |
+
self,
|
| 1459 |
+
layer,
|
| 1460 |
+
x: torch.Tensor,
|
| 1461 |
+
bias: torch.Tensor | None = None,
|
| 1462 |
+
) -> torch.Tensor:
|
| 1463 |
+
linears = getattr(layer, "_exl3_linears", None)
|
| 1464 |
+
if not linears:
|
| 1465 |
+
raise RuntimeError("EXL3 linear layers were not built after weight load")
|
| 1466 |
+
|
| 1467 |
+
# x shape: (batch, in_features) or (batch, ..., in_features)
|
| 1468 |
+
# Flatten to 2D: (rows, in_features)
|
| 1469 |
+
orig_shape = x.shape
|
| 1470 |
+
if len(orig_shape) > 2:
|
| 1471 |
+
# Multi-dim input: flatten to (rows, in)
|
| 1472 |
+
rows = 1
|
| 1473 |
+
for d in orig_shape[:-1]:
|
| 1474 |
+
rows *= d
|
| 1475 |
+
x_2d = x.reshape(rows, orig_shape[-1])
|
| 1476 |
+
else:
|
| 1477 |
+
x_2d = x
|
| 1478 |
+
|
| 1479 |
+
# Cast to contiguous fp16 for EXL3 shards
|
| 1480 |
+
x_fp16 = x_2d.to(torch.float16).contiguous()
|
| 1481 |
+
|
| 1482 |
+
# Get bf16 shards and weight if present
|
| 1483 |
+
bf16_shards = getattr(layer, "_exl3_linear_bf16_shards", [])
|
| 1484 |
+
bf16_weight = getattr(layer, "_exl3_bf16_weight", None)
|
| 1485 |
+
output_sizes = layer._exl3_linear_output_partition_sizes
|
| 1486 |
+
n_shards = len(linears)
|
| 1487 |
+
|
| 1488 |
+
# Run each shard in declared order
|
| 1489 |
+
outputs = []
|
| 1490 |
+
for i in range(n_shards):
|
| 1491 |
+
if i in bf16_shards:
|
| 1492 |
+
# BF16 shard: use dense linear
|
| 1493 |
+
if bf16_weight is None:
|
| 1494 |
+
raise RuntimeError(
|
| 1495 |
+
f"EXL3 bf16 shard {i} but _exl3_bf16_weight is missing"
|
| 1496 |
+
)
|
| 1497 |
+
bf16_idx = bf16_shards.index(i)
|
| 1498 |
+
out_start = sum(output_sizes[j] for j in bf16_shards[:bf16_idx])
|
| 1499 |
+
out_end = out_start + output_sizes[i]
|
| 1500 |
+
w_shard = bf16_weight[out_start:out_end]
|
| 1501 |
+
out = F.linear(x_2d, w_shard).to(dtype=torch.float32)
|
| 1502 |
+
outputs.append(out)
|
| 1503 |
+
else:
|
| 1504 |
+
# EXL3 shard
|
| 1505 |
+
linear = linears[i]
|
| 1506 |
+
if linear is None:
|
| 1507 |
+
raise RuntimeError(f"EXL3 linear shard {i} is None")
|
| 1508 |
+
out = linear.forward(x_fp16, {}, out_dtype=torch.float32)
|
| 1509 |
+
outputs.append(out)
|
| 1510 |
+
|
| 1511 |
+
# Concatenate shards along output dimension
|
| 1512 |
+
if len(outputs) > 1:
|
| 1513 |
+
y = torch.cat(outputs, dim=1)
|
| 1514 |
+
else:
|
| 1515 |
+
y = outputs[0]
|
| 1516 |
+
|
| 1517 |
+
# Cast back to input dtype
|
| 1518 |
+
y = y.to(dtype=x.dtype)
|
| 1519 |
+
|
| 1520 |
+
# Add bias if provided
|
| 1521 |
+
if bias is not None:
|
| 1522 |
+
y = y + bias
|
| 1523 |
+
|
| 1524 |
+
# Restore original shape
|
| 1525 |
+
if len(orig_shape) > 2:
|
| 1526 |
+
y = y.reshape(*orig_shape[:-1], y.shape[-1])
|
| 1527 |
+
|
| 1528 |
+
return y
|
lna-lab/ppl-vllm.py
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""ppl meter through a vLLM OpenAI seat: prompt_logprobs over fixed text chunks. usage: ppl-vllm.py [port] [text_file] [ctx_tokens] [n_chunks]"""
|
| 3 |
+
import json, math, sys, urllib.request
|
| 4 |
+
port = sys.argv[1] if len(sys.argv) > 1 else "8899"
|
| 5 |
+
path = sys.argv[2] if len(sys.argv) > 2 else "/run/media/tonoken3/DATA1/Models/wikitext-2-raw/wiki.test.raw"
|
| 6 |
+
ctx = int(sys.argv[3]) if len(sys.argv) > 3 else 2048; nch = int(sys.argv[4]) if len(sys.argv) > 4 else 8
|
| 7 |
+
U = f"http://127.0.0.1:{port}/v1/completions"
|
| 8 |
+
text = open(path, encoding="utf-8", errors="ignore").read()
|
| 9 |
+
chars = ctx * 4 # rough char budget per chunk; server truncates nothing, we just pick disjoint windows
|
| 10 |
+
tot_nll = 0.0; tot_n = 0
|
| 11 |
+
for i in range(nch):
|
| 12 |
+
chunk = text[i * chars:(i + 1) * chars]
|
| 13 |
+
body = {"model": "DSV4-Flash", "prompt": chunk, "max_tokens": 1, "temperature": 0, "prompt_logprobs": 0, "echo": False}
|
| 14 |
+
r = json.load(urllib.request.urlopen(urllib.request.Request(U, data=json.dumps(body).encode(), headers={"content-type": "application/json"}), timeout=900))
|
| 15 |
+
pl = r["choices"][0].get("prompt_logprobs") or []
|
| 16 |
+
lps = []
|
| 17 |
+
for tok in pl[1:ctx]: # skip first token (no context); cap at ctx tokens
|
| 18 |
+
if tok:
|
| 19 |
+
lps.append(max(v["logprob"] if isinstance(v, dict) else v for v in tok.values()))
|
| 20 |
+
n = len(lps); nll = -sum(lps)
|
| 21 |
+
tot_nll += nll; tot_n += n
|
| 22 |
+
print(f"chunk {i}: tokens={n} ppl={math.exp(nll / max(n, 1)):.4f}", flush=True)
|
| 23 |
+
print(f"PPL={math.exp(tot_nll / max(tot_n, 1)):.4f} tokens={tot_n} ctx={ctx} chunks={nch} file={path}")
|
lna-lab/serve-dsv4-tp8.sh
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# DSV4-Flash-Vision EXL3 MixedK on 8 experiment GPUs (TP=8) via lna-lab/vllm-exl3:dsv4. YUKI 2026-09-03
|
| 3 |
+
# env: GPUS (default 0,1,2,3,5,7,8,9) PORT(8899) MAXLEN(65536) SPEC ('' or json) UTIL(0.90) NAME(dsv4)
|
| 4 |
+
set -u
|
| 5 |
+
GPUS=${GPUS:-0,1,2,3,5,7,8,9}; PORT=${PORT:-8899}; MAXLEN=${MAXLEN:-65536}; UTIL=${UTIL:-0.90}; NAME=${NAME:-dsv4}
|
| 6 |
+
TP=$(echo $GPUS | tr ',' '\n' | wc -l)
|
| 7 |
+
M=${MODEL:-/run/media/tonoken3/DATA1/DSV4-Flash-Vision-EXL3-MixedK}
|
| 8 |
+
EXTRA=()
|
| 9 |
+
[[ -n "${SPEC:-}" ]] && EXTRA+=(--speculative-config "$SPEC")
|
| 10 |
+
docker rm -f $NAME >/dev/null 2>&1
|
| 11 |
+
docker run --ulimit core=0 --cap-add=SYS_PTRACE -e VLLM_DISABLE_SHARED_EXPERTS_STREAM=${SHARED_STREAM_OFF:-1} -v /run/media/tonoken3/DATA1/.tmp/dsv4-cache/root-cache:/root/.cache -v /run/media/tonoken3/DATA1/.tmp/dsv4-cache/tilelang:/root/.tilelang ${EXT_SO:+-v $EXT_SO:/usr/local/lib/python3.12/dist-packages/$(basename ${EXT_SO:-x})} ${PLUGIN_SRC:+-v $PLUGIN_SRC:/usr/local/lib/python3.12/dist-packages/vllm_exl3} -d --name $NAME --gpus "\"device=$GPUS\"" --shm-size=16g --ipc=host \
|
| 12 |
+
-e NCCL_P2P_DISABLE=1 -e NCCL_CUMEM_ENABLE=0 -e EXLLAMAV3_TUNE_CACHE=/lab/exl3-tune-cache ${AUX_STREAMS:+-e LNA_DSV4_AUX_STREAMS=$AUX_STREAMS} ${NCCL_EXTRA:-} -e VLLM_NO_USAGE_STATS=1 -e DO_NOT_TRACK=1 \
|
| 13 |
+
-p 127.0.0.1:$PORT:8000 -v /run/media/tonoken3/DATA1:/run/media/tonoken3/DATA1 -v /run/media/tonoken3/DATA1/vllm-exl3-lab:/lab \
|
| 14 |
+
lna-lab/vllm-exl3:${IMAGE:-dsv4} \
|
| 15 |
+
$M --served-model-name DSV4-Flash --tensor-parallel-size $TP --quantization exl3 \
|
| 16 |
+
--max-model-len $MAXLEN --max-num-seqs ${SEQS:-4} --max-num-batched-tokens ${BT:-2048} \
|
| 17 |
+
--kv-cache-dtype ${KVDT:-fp8} --gpu-memory-utilization $UTIL ${EAGER:+--enforce-eager} --compilation-config "${COMPILE:-{\"cudagraph_capture_sizes\":[1,2,4]\}}" --disable-custom-all-reduce \
|
| 18 |
+
--no-enable-prefix-caching --trust-remote-code ${PROFILE:+--profiler-config "$PROFILE"} \
|
| 19 |
+
--enable-auto-tool-choice --tool-call-parser deepseek_v4 --reasoning-parser deepseek_v4 "${EXTRA[@]}"
|
| 20 |
+
echo "container $NAME on :$PORT (TP=$TP, spec=${SPEC:-off}); docker logs -f $NAME"
|
lna-lab/tools/dense_bake.py
ADDED
|
@@ -0,0 +1,924 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Calibrate and bake only the dense DSV4 linears into EXL3.
|
| 3 |
+
|
| 4 |
+
The bake deliberately keeps the routed-expert tensors in the existing EXL3
|
| 5 |
+
pack. It uses the read-only exllamav3 checkout as a runtime dependency and
|
| 6 |
+
does not copy or modify that checkout. ``--dry-run`` only inspects config,
|
| 7 |
+
model layout, and (when ``--merge`` is supplied) local safetensors headers.
|
| 8 |
+
It never imports CUDA conversion kernels or writes a model pack.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import argparse
|
| 14 |
+
import copy
|
| 15 |
+
import json
|
| 16 |
+
import os
|
| 17 |
+
import re
|
| 18 |
+
import struct
|
| 19 |
+
import sys
|
| 20 |
+
from pathlib import Path
|
| 21 |
+
from types import SimpleNamespace
|
| 22 |
+
from typing import Any, Iterable
|
| 23 |
+
|
| 24 |
+
try:
|
| 25 |
+
from . import layer_overlay as overlay
|
| 26 |
+
except ImportError: # Running as ``python tools/dense_bake.py``.
|
| 27 |
+
import layer_overlay as overlay
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
CONFIG_NAME = "config.json"
|
| 31 |
+
INDEX_NAME = "model.safetensors.index.json"
|
| 32 |
+
DENSE_SHARD = "model-dense-exl3.safetensors"
|
| 33 |
+
EXL3_SUFFIXES = ("trellis", "suh", "svh", "mcg")
|
| 34 |
+
EXL3_DTYPES = {"trellis": "I16", "suh": "F16", "svh": "F16", "mcg": "I32"}
|
| 35 |
+
MCG_MARKER = -877912083
|
| 36 |
+
DEFAULT_EXL3_SOURCE = "/run/media/tonoken3/DATA1/.tmp/exl3src"
|
| 37 |
+
|
| 38 |
+
_LAYER_FILE_RE = re.compile(r"^layers\.(?P<layer>[0-9]+)\.safetensors$")
|
| 39 |
+
_WO_A_RE = re.compile(r"^attn\.wo_a\.slice\.(?P<slice>[0-7])$")
|
| 40 |
+
_DENSE_TAILS = {
|
| 41 |
+
"attn.wq_a",
|
| 42 |
+
"attn.wq_b",
|
| 43 |
+
"attn.wkv",
|
| 44 |
+
"attn.wo_b",
|
| 45 |
+
"attn.compressor.wkv",
|
| 46 |
+
"attn.compressor.wgate",
|
| 47 |
+
"attn.indexer.compressor.wkv",
|
| 48 |
+
"attn.indexer.compressor.wgate",
|
| 49 |
+
"attn.indexer.wq_b",
|
| 50 |
+
"ffn.shared_experts.w1",
|
| 51 |
+
"ffn.shared_experts.w2",
|
| 52 |
+
"ffn.shared_experts.w3",
|
| 53 |
+
}
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def log(message: str) -> None:
|
| 57 |
+
print(message, flush=True)
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def dense_kind(key: str) -> str | None:
|
| 61 |
+
"""Return ``attn`` or ``shared`` only for the ordered T4 target list."""
|
| 62 |
+
match = re.match(r"^layers\.[0-9]+\.(?P<tail>.+)$", key)
|
| 63 |
+
if match is None:
|
| 64 |
+
return None
|
| 65 |
+
tail = match.group("tail")
|
| 66 |
+
if tail in _DENSE_TAILS or _WO_A_RE.fullmatch(tail):
|
| 67 |
+
return "shared" if tail.startswith("ffn.shared_experts.") else "attn"
|
| 68 |
+
return None
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def is_dense_key(key: str) -> bool:
|
| 72 |
+
return dense_kind(key) is not None
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
def dense_keys_for_layer(layer: int, keys: Iterable[str]) -> list[str]:
|
| 76 |
+
"""Filter a module's recursive Linear keys, excluding experts, MTP, and non-linears."""
|
| 77 |
+
result = [key for key in keys if dense_kind(key) is not None and key.startswith(f"layers.{layer}.")]
|
| 78 |
+
result.sort()
|
| 79 |
+
wo_a = [key for key in result if ".attn.wo_a.slice." in key]
|
| 80 |
+
if wo_a and {int(key.rsplit(".", 1)[1]) for key in wo_a} != set(range(8)):
|
| 81 |
+
raise RuntimeError(f"layers.{layer}: wo_a must contain exactly slices 0..7, got {wo_a}")
|
| 82 |
+
return result
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def bits_for_key(key: str, bits: float, attn_bits: float | None, shared_bits: float | None) -> float:
|
| 86 |
+
kind = dense_kind(key)
|
| 87 |
+
if kind is None:
|
| 88 |
+
raise ValueError(f"not a T4 dense key: {key}")
|
| 89 |
+
value = float(attn_bits if kind == "attn" and attn_bits is not None else
|
| 90 |
+
shared_bits if kind == "shared" and shared_bits is not None else bits)
|
| 91 |
+
return int(value) if value.is_integer() else value # LNA-LAB: the quantizer shifts by K (must be int)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def _schema_bits(value: float) -> int:
|
| 95 |
+
if not float(value).is_integer() or int(value) not in (2, 3, 4, 5, 6):
|
| 96 |
+
raise ValueError(f"vLLM non_routed_exl3 requires integer K in 2..6, got {value}")
|
| 97 |
+
return int(value)
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def _layer_number(key: str) -> int:
|
| 101 |
+
match = re.match(r"^layers\.([0-9]+)(?:\.|$)", key) # LNA-LAB: module keys are bare "layers.N"
|
| 102 |
+
if match is None:
|
| 103 |
+
raise ValueError(f"not a language-layer key: {key}")
|
| 104 |
+
return int(match.group(1))
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def _parse_devices(value: str) -> list[int]:
|
| 108 |
+
try:
|
| 109 |
+
devices = [int(part.strip()) for part in value.split(",") if part.strip()]
|
| 110 |
+
except ValueError as exc:
|
| 111 |
+
raise ValueError(f"invalid --devices value: {value!r}") from exc
|
| 112 |
+
if not devices or any(device < 0 for device in devices) or len(set(devices)) != len(devices):
|
| 113 |
+
raise ValueError(f"--devices must be unique non-negative integers: {value!r}")
|
| 114 |
+
return devices
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
def _runtime_import(*, dry_run: bool = False):
|
| 118 |
+
"""Import the exact exllamav3 classes/functions used by the bake.
|
| 119 |
+
|
| 120 |
+
A dry-run can construct the Python model graph without the native extension.
|
| 121 |
+
The extension stub is never used by a real bake; it exists solely to make
|
| 122 |
+
the requested offline layout check possible on a CPU-only host.
|
| 123 |
+
"""
|
| 124 |
+
source_root = os.environ.get("EXL3_SOURCE", DEFAULT_EXL3_SOURCE)
|
| 125 |
+
if not Path(source_root).is_dir():
|
| 126 |
+
raise RuntimeError(f"exllamav3 source checkout is missing: {source_root}")
|
| 127 |
+
if source_root not in sys.path:
|
| 128 |
+
sys.path.insert(0, source_root)
|
| 129 |
+
|
| 130 |
+
if dry_run and "exllamav3.ext" not in sys.modules:
|
| 131 |
+
class _ExtensionStub:
|
| 132 |
+
def __getattr__(self, _name):
|
| 133 |
+
return lambda *args, **kwargs: None
|
| 134 |
+
|
| 135 |
+
import types
|
| 136 |
+
ext_module = types.ModuleType("exllamav3.ext")
|
| 137 |
+
ext_module.exllamav3_ext = _ExtensionStub()
|
| 138 |
+
sys.modules["exllamav3.ext"] = ext_module
|
| 139 |
+
|
| 140 |
+
from exllamav3 import Config, Model, Tokenizer
|
| 141 |
+
from exllamav3.modules import Linear
|
| 142 |
+
from exllamav3.loader.safetensors import SafetensorsCollection, VariantSafetensorsCollection
|
| 143 |
+
from exllamav3.conversion.calibration_data import get_default_calibration
|
| 144 |
+
from exllamav3.conversion.convert_model import (
|
| 145 |
+
advance_state_parallel,
|
| 146 |
+
capture_module_parallel,
|
| 147 |
+
load_parallel_calib_modules,
|
| 148 |
+
quantize_linears_parallel,
|
| 149 |
+
quantize_linears_single,
|
| 150 |
+
load_tensor,
|
| 151 |
+
save_tensor,
|
| 152 |
+
)
|
| 153 |
+
return SimpleNamespace(
|
| 154 |
+
Config=Config,
|
| 155 |
+
Model=Model,
|
| 156 |
+
Tokenizer=Tokenizer,
|
| 157 |
+
Linear=Linear,
|
| 158 |
+
SafetensorsCollection=SafetensorsCollection,
|
| 159 |
+
VariantSafetensorsCollection=VariantSafetensorsCollection,
|
| 160 |
+
get_default_calibration=get_default_calibration,
|
| 161 |
+
advance_state_parallel=advance_state_parallel,
|
| 162 |
+
capture_module_parallel=capture_module_parallel,
|
| 163 |
+
load_parallel_calib_modules=load_parallel_calib_modules,
|
| 164 |
+
quantize_linears_parallel=quantize_linears_parallel,
|
| 165 |
+
quantize_linears_single=quantize_linears_single,
|
| 166 |
+
load_tensor=load_tensor,
|
| 167 |
+
save_tensor=save_tensor,
|
| 168 |
+
)
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def _load_variant_model(src: Path, experts: Path, runtime):
|
| 172 |
+
config = runtime.Config.from_directory(str(src))
|
| 173 |
+
expert_stc = runtime.SafetensorsCollection(str(experts))
|
| 174 |
+
variant = runtime.VariantSafetensorsCollection(config.stc)
|
| 175 |
+
# This is the T4 contract: only routed expert keys are overridden.
|
| 176 |
+
variant.add_stc(["layers.*.ffn.experts.*"], expert_stc)
|
| 177 |
+
config.stc = variant
|
| 178 |
+
model = runtime.Model.from_config(config)
|
| 179 |
+
# DSV4's architecture exposes this flag for calibration; make the T4
|
| 180 |
+
# invariant explicit in case a future architecture default changes.
|
| 181 |
+
model.calibration_all_experts = True
|
| 182 |
+
return config, model
|
| 183 |
+
|
| 184 |
+
|
| 185 |
+
def _module_plan(model, linear_type) -> list[tuple[int, Any, list[Any]]]:
|
| 186 |
+
plan = []
|
| 187 |
+
for idx, module in enumerate(model.modules):
|
| 188 |
+
layer_match = re.fullmatch(r"layers\.([0-9]+)", getattr(module, "key", ""))
|
| 189 |
+
if layer_match is None:
|
| 190 |
+
continue
|
| 191 |
+
layer = int(layer_match.group(1))
|
| 192 |
+
linears = [m for m in module if isinstance(m, linear_type)]
|
| 193 |
+
keys = dense_keys_for_layer(layer, [m.key for m in linears])
|
| 194 |
+
by_key = {m.key: m for m in linears}
|
| 195 |
+
targets = [by_key[key] for key in keys]
|
| 196 |
+
if len(targets) == 0:
|
| 197 |
+
raise RuntimeError(f"{module.key}: no dense linears found")
|
| 198 |
+
plan.append((idx, module, targets))
|
| 199 |
+
if not plan:
|
| 200 |
+
raise RuntimeError("model has no layers.N modules")
|
| 201 |
+
return plan
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def dry_run_layout(src: Path, experts: Path) -> dict[str, Any]:
|
| 205 |
+
if not src.is_dir() or not experts.is_dir():
|
| 206 |
+
raise RuntimeError(f"dry-run inputs must be directories: {src}, {experts}")
|
| 207 |
+
runtime = _runtime_import(dry_run=True)
|
| 208 |
+
# Construct the real source layout with the D2 expert override, and also
|
| 209 |
+
# construct Config.from_directory(D2) explicitly: the latter is the
|
| 210 |
+
# complete experts-only config used for the offline pack sanity check.
|
| 211 |
+
config, model = _load_variant_model(src, experts, runtime)
|
| 212 |
+
d2_config = runtime.Config.from_directory(str(experts))
|
| 213 |
+
plan = _module_plan(model, runtime.Linear)
|
| 214 |
+
counts = {"attn": 0, "shared": 0}
|
| 215 |
+
for _, _, targets in plan:
|
| 216 |
+
for linear in targets:
|
| 217 |
+
counts[dense_kind(linear.key)] += 1
|
| 218 |
+
log(
|
| 219 |
+
f"DRY_RUN_LAYOUT architecture={getattr(config, 'architecture', 'unknown')} "
|
| 220 |
+
f"modules={len(model.modules)} layers={len(plan)} dense_linears={sum(counts.values())} "
|
| 221 |
+
f"attn={counts['attn']} shared={counts['shared']} routed_override=layers.*.ffn.experts.* "
|
| 222 |
+
f"d2_config={getattr(d2_config, 'architecture', 'unknown')}"
|
| 223 |
+
)
|
| 224 |
+
for _, module, targets in plan[:2]:
|
| 225 |
+
log(f"DRY_RUN_LAYER key={module.key} dense_linears={len(targets)}")
|
| 226 |
+
log(f"DRY_RUN_LAYER key={plan[-1][1].key} dense_linears={len(plan[-1][2])}")
|
| 227 |
+
return {"config": config, "model": model, "runtime": runtime, "plan": plan}
|
| 228 |
+
|
| 229 |
+
|
| 230 |
+
def _load_module(module, device, *, load_slice=None, source=None, close=True):
|
| 231 |
+
defer = source is None and module.can_defer_load()
|
| 232 |
+
if defer:
|
| 233 |
+
module.config.stc.begin_deferred_load()
|
| 234 |
+
try:
|
| 235 |
+
kwargs = {}
|
| 236 |
+
if load_slice is not None:
|
| 237 |
+
kwargs["load_slice"] = load_slice
|
| 238 |
+
if source is not None:
|
| 239 |
+
kwargs["source"] = source
|
| 240 |
+
module.load(device, **kwargs)
|
| 241 |
+
finally:
|
| 242 |
+
if defer:
|
| 243 |
+
module.config.stc.end_deferred_load()
|
| 244 |
+
if close:
|
| 245 |
+
module.config.stc.close()
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
def _capture_serial(model, module, state, original_ids, quant_preserves, bad_rows, num_ref_states=5):
|
| 249 |
+
capture_h = {}
|
| 250 |
+
ref_states = {}
|
| 251 |
+
|
| 252 |
+
def get_preserve(i, params):
|
| 253 |
+
params.update(quant_preserves[i])
|
| 254 |
+
params["quant_preserve"] = quant_preserves[i]
|
| 255 |
+
|
| 256 |
+
def put_preserve(i, params):
|
| 257 |
+
quant_preserves[i] = params["quant_preserve"]
|
| 258 |
+
|
| 259 |
+
for i in range(len(state)):
|
| 260 |
+
if i in bad_rows:
|
| 261 |
+
continue
|
| 262 |
+
params = {
|
| 263 |
+
"attn_mode": "flash_attn_nc",
|
| 264 |
+
"capture": capture_h,
|
| 265 |
+
"activate_all_experts": model.calibration_all_experts,
|
| 266 |
+
"input_ids": original_ids[i],
|
| 267 |
+
}
|
| 268 |
+
get_preserve(i, params)
|
| 269 |
+
result = module.forward(module.prepare_for_device(state[i], params), params)
|
| 270 |
+
put_preserve(i, params)
|
| 271 |
+
if i < num_ref_states:
|
| 272 |
+
if model.calibration_all_experts:
|
| 273 |
+
params = {"attn_mode": "flash_attn_nc", "input_ids": original_ids[i]}
|
| 274 |
+
get_preserve(i, params)
|
| 275 |
+
result = module.forward(module.prepare_for_device(state[i], params), params)
|
| 276 |
+
put_preserve(i, params)
|
| 277 |
+
if result.isfinite().all().item():
|
| 278 |
+
ref_states[i] = result.cpu()
|
| 279 |
+
else:
|
| 280 |
+
bad_rows.add(i)
|
| 281 |
+
log(f"WARNING non-finite reference state row={i}; excluded")
|
| 282 |
+
for item in capture_h.values():
|
| 283 |
+
item["H_swap_device"] = item["H"].device
|
| 284 |
+
item["H"] = item["H"].cpu()
|
| 285 |
+
return capture_h, ref_states, get_preserve, put_preserve
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
def _advance_serial(model, module, state, original_ids, quant_preserves, ref_states, bad_rows, have_linears):
|
| 289 |
+
error = 0.0
|
| 290 |
+
measured = 0
|
| 291 |
+
|
| 292 |
+
def get_preserve(i, params):
|
| 293 |
+
params.update(quant_preserves[i])
|
| 294 |
+
params["quant_preserve"] = quant_preserves[i]
|
| 295 |
+
|
| 296 |
+
def put_preserve(i, params):
|
| 297 |
+
quant_preserves[i] = params["quant_preserve"]
|
| 298 |
+
|
| 299 |
+
for i in range(len(state)):
|
| 300 |
+
if i in bad_rows:
|
| 301 |
+
continue
|
| 302 |
+
params = {"attn_mode": "flash_attn_nc", "input_ids": original_ids[i]}
|
| 303 |
+
state[i] = module.prepare_for_device(state[i], params)
|
| 304 |
+
if i < 5:
|
| 305 |
+
get_preserve(i, params)
|
| 306 |
+
result = module.forward(state[i], params)
|
| 307 |
+
put_preserve(i, params)
|
| 308 |
+
if not result.isfinite().all().item():
|
| 309 |
+
bad_rows.add(i)
|
| 310 |
+
continue
|
| 311 |
+
state[i] = result.cpu()
|
| 312 |
+
else:
|
| 313 |
+
# We stop after the final transformer block, so every row must be
|
| 314 |
+
# advanced; the reference/error pass remains limited to five rows.
|
| 315 |
+
get_preserve(i, params)
|
| 316 |
+
state[i] = module.forward(state[i], params).cpu()
|
| 317 |
+
put_preserve(i, params)
|
| 318 |
+
ref = ref_states.get(i)
|
| 319 |
+
if ref is not None and have_linears and i not in bad_rows:
|
| 320 |
+
x = state[i].view(-1, state[i].shape[-1]).float()
|
| 321 |
+
y = ref.view(-1, ref.shape[-1]).float()
|
| 322 |
+
error += (torch_norm(x - y) / torch_norm(y)).item()
|
| 323 |
+
measured += 1
|
| 324 |
+
ref_states[i] = None
|
| 325 |
+
return error / max(measured, 1), measured
|
| 326 |
+
|
| 327 |
+
|
| 328 |
+
def _set_new_tensors(stc, tensors) -> None:
|
| 329 |
+
"""Set the in-memory reload layer on either a plain or variant STC.
|
| 330 |
+
|
| 331 |
+
VariantSafetensorsCollection intentionally leaves ``set_new_tensors``
|
| 332 |
+
unimplemented. Dense keys resolve to its ``main`` collection, so the
|
| 333 |
+
converter's reload operation is applied to that collection only; the
|
| 334 |
+
routed-expert override remains the D2 collection.
|
| 335 |
+
"""
|
| 336 |
+
target = getattr(stc, "main", stc)
|
| 337 |
+
target.set_new_tensors(tensors)
|
| 338 |
+
|
| 339 |
+
|
| 340 |
+
def _inject_reload_keys(stc, tensors) -> list[str]:
|
| 341 |
+
"""Make variant ``has_tensor_group`` see in-memory EXL3 keys.
|
| 342 |
+
|
| 343 |
+
The audited VariantSafetensorsCollection delegates ``get_tensor`` but its
|
| 344 |
+
``has_tensor_group`` checks only ``tensor_file_map``. Adding temporary
|
| 345 |
+
sentinels lets Linear.load choose its EXL3 branch; get_tensor still takes
|
| 346 |
+
the actual values from ``new_tensors``. The sentinels are removed before
|
| 347 |
+
the collection is closed.
|
| 348 |
+
"""
|
| 349 |
+
main = getattr(stc, "main", stc)
|
| 350 |
+
added = []
|
| 351 |
+
for key in tensors:
|
| 352 |
+
if key not in main.tensor_file_map:
|
| 353 |
+
main.tensor_file_map[key] = "__dense_bake_memory__"
|
| 354 |
+
added.append(key)
|
| 355 |
+
return added
|
| 356 |
+
|
| 357 |
+
|
| 358 |
+
def _remove_reload_keys(stc, added: list[str]) -> None:
|
| 359 |
+
main = getattr(stc, "main", stc)
|
| 360 |
+
for key in added:
|
| 361 |
+
main.tensor_file_map.pop(key, None)
|
| 362 |
+
|
| 363 |
+
|
| 364 |
+
def torch_norm(tensor):
|
| 365 |
+
# Kept as a tiny late-bound helper so importing this tool for header-only
|
| 366 |
+
# merge planning does not initialize torch/CUDA.
|
| 367 |
+
import torch
|
| 368 |
+
return torch.linalg.norm(tensor, "fro")
|
| 369 |
+
|
| 370 |
+
|
| 371 |
+
def _checkpoint(runtime, work: Path, job: dict, state, original_ids, args_dict: dict):
|
| 372 |
+
ckpt = work / "ckpt"
|
| 373 |
+
ckpt.mkdir(parents=True, exist_ok=True)
|
| 374 |
+
runtime.save_tensor(state, "ckpt/state.safetensors", args_dict)
|
| 375 |
+
runtime.save_tensor(original_ids, "ckpt/original_input_ids.safetensors", args_dict)
|
| 376 |
+
(ckpt / "job.json").write_text(json.dumps(job, indent=2) + "\n", encoding="utf-8")
|
| 377 |
+
|
| 378 |
+
|
| 379 |
+
def _read_job(work: Path) -> dict:
|
| 380 |
+
path = work / "ckpt" / "job.json"
|
| 381 |
+
if not path.is_file():
|
| 382 |
+
return {"next_module_idx": 0, "bad_rows": []}
|
| 383 |
+
try:
|
| 384 |
+
data = json.loads(path.read_text(encoding="utf-8"))
|
| 385 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 386 |
+
raise RuntimeError(f"invalid checkpoint {path}: {exc}") from exc
|
| 387 |
+
if not isinstance(data, dict) or not isinstance(data.get("next_module_idx"), int):
|
| 388 |
+
raise RuntimeError(f"checkpoint {path} has no integer next_module_idx")
|
| 389 |
+
return data
|
| 390 |
+
|
| 391 |
+
|
| 392 |
+
def _prepare_real_state(runtime, experts: Path, args, model, job, work: Path):
|
| 393 |
+
import torch
|
| 394 |
+
args_dict = vars(args).copy()
|
| 395 |
+
args_dict["work_dir"] = str(work)
|
| 396 |
+
args_dict["apply_out_scales"] = True
|
| 397 |
+
args_dict["image_dump"] = False
|
| 398 |
+
args_dict["codebook"] = "mcg"
|
| 399 |
+
args_dict["verbose"] = args.verbose
|
| 400 |
+
quant_preserves = []
|
| 401 |
+
bad_rows = set(job.get("bad_rows") or [])
|
| 402 |
+
if job["next_module_idx"] > 0:
|
| 403 |
+
state = runtime.load_tensor("ckpt/state.safetensors", args_dict)
|
| 404 |
+
original_ids = runtime.load_tensor("ckpt/original_input_ids.safetensors", args_dict)
|
| 405 |
+
if not isinstance(state, list) or not isinstance(original_ids, list):
|
| 406 |
+
raise RuntimeError("checkpoint state/original_input_ids must both be tensor lists")
|
| 407 |
+
quant_preserves = [{} for _ in state]
|
| 408 |
+
return state, original_ids, quant_preserves, bad_rows, args_dict
|
| 409 |
+
|
| 410 |
+
tokenizer_config = runtime.Config.from_directory(str(experts))
|
| 411 |
+
tokenizer = runtime.Tokenizer.from_config(tokenizer_config)
|
| 412 |
+
original_ids = runtime.get_default_calibration(
|
| 413 |
+
{"cal_rows": args.cal_rows, "cal_cols": args.cal_cols}, tokenizer
|
| 414 |
+
)
|
| 415 |
+
state = list(original_ids)
|
| 416 |
+
quant_preserves = [{} for _ in state]
|
| 417 |
+
# The converter starts with token rows and advances them through embed and
|
| 418 |
+
# hc_expand before entering layers.N. These modules are intentionally not
|
| 419 |
+
# quantized or written to qtensors.
|
| 420 |
+
for idx, module in enumerate(model.modules[: model.first_block_idx]):
|
| 421 |
+
log(f"LOAD_PREFIX module={module.key}")
|
| 422 |
+
_load_module(module, torch.device(f"cuda:{args.devices[0]}"))
|
| 423 |
+
for i in range(len(state)):
|
| 424 |
+
params = {"attn_mode": "flash_attn_nc", "input_ids": original_ids[i]}
|
| 425 |
+
state[i] = module.forward(module.prepare_for_device(state[i], params), params).cpu()
|
| 426 |
+
module.unload()
|
| 427 |
+
return state, original_ids, quant_preserves, bad_rows, args_dict
|
| 428 |
+
|
| 429 |
+
|
| 430 |
+
def bake(args: argparse.Namespace) -> int:
|
| 431 |
+
src = Path(args.src).expanduser().resolve()
|
| 432 |
+
experts = Path(args.experts).expanduser().resolve()
|
| 433 |
+
work = Path(args.work).expanduser().resolve()
|
| 434 |
+
if not src.is_dir() or not experts.is_dir():
|
| 435 |
+
raise RuntimeError(f"--src and --experts must be directories: {src}, {experts}")
|
| 436 |
+
if args.bits < 1 or args.bits > 8 or any(
|
| 437 |
+
value is not None and (value < 1 or value > 8)
|
| 438 |
+
for value in (args.attn_bits, args.shared_bits)
|
| 439 |
+
):
|
| 440 |
+
raise ValueError("bits must be in the range 1..8")
|
| 441 |
+
devices = _parse_devices(args.devices)
|
| 442 |
+
if not args.dry_run:
|
| 443 |
+
import torch
|
| 444 |
+
if not torch.cuda.is_available():
|
| 445 |
+
raise RuntimeError("real bake requires CUDA; use --dry-run for the offline checks")
|
| 446 |
+
count = torch.cuda.device_count()
|
| 447 |
+
if any(device >= count for device in devices):
|
| 448 |
+
raise RuntimeError(f"requested CUDA devices {devices}, but only {count} are visible")
|
| 449 |
+
|
| 450 |
+
runtime = _runtime_import(dry_run=args.dry_run)
|
| 451 |
+
config, model = _load_variant_model(src, experts, runtime)
|
| 452 |
+
plan = _module_plan(model, runtime.Linear)
|
| 453 |
+
log(
|
| 454 |
+
f"BAKE_PLAN layers={len(plan)} dense_linears={sum(len(t) for _, _, t in plan)} "
|
| 455 |
+
f"bits={args.bits:g} attn_bits={(args.attn_bits if args.attn_bits is not None else args.bits):g} "
|
| 456 |
+
f"shared_bits={(args.shared_bits if args.shared_bits is not None else args.bits):g} codebook=mcg"
|
| 457 |
+
)
|
| 458 |
+
if args.dry_run:
|
| 459 |
+
for _, module, targets in plan[:2]:
|
| 460 |
+
log(f"DRY_RUN_BAKE module={module.key} targets={len(targets)}")
|
| 461 |
+
log(f"DRY_RUN_BAKE module={plan[-1][1].key} targets={len(plan[-1][2])}")
|
| 462 |
+
if args.merge:
|
| 463 |
+
merge_plan = plan_merge(work, experts, args)
|
| 464 |
+
log(f"DRY_RUN_MERGE dropped={len(merge_plan['dropped'])} new={len(merge_plan['new_plan'])}")
|
| 465 |
+
return 0
|
| 466 |
+
|
| 467 |
+
work.mkdir(parents=True, exist_ok=True)
|
| 468 |
+
(work / "qtensors").mkdir(exist_ok=True)
|
| 469 |
+
(work / "debug").mkdir(exist_ok=True)
|
| 470 |
+
signature = {
|
| 471 |
+
"src": str(src), "experts": str(experts), "bits": args.bits,
|
| 472 |
+
"attn_bits": args.attn_bits, "shared_bits": args.shared_bits,
|
| 473 |
+
"cal_rows": args.cal_rows, "cal_cols": args.cal_cols, "codebook": "mcg",
|
| 474 |
+
}
|
| 475 |
+
signature_path = work / "args.json"
|
| 476 |
+
if args.resume:
|
| 477 |
+
if not signature_path.is_file():
|
| 478 |
+
raise RuntimeError(f"--resume requested but {signature_path} is missing")
|
| 479 |
+
if json.loads(signature_path.read_text(encoding="utf-8")) != signature:
|
| 480 |
+
raise RuntimeError("resume settings differ from work/args.json")
|
| 481 |
+
else:
|
| 482 |
+
qtensors_dir = work / "qtensors"
|
| 483 |
+
existing_qtensors = qtensors_dir.is_dir() and any(qtensors_dir.glob("layers.*.safetensors"))
|
| 484 |
+
if signature_path.exists() or (work / "ckpt" / "job.json").exists() or existing_qtensors:
|
| 485 |
+
raise RuntimeError("work directory already contains a bake; use --resume or choose a new --work path")
|
| 486 |
+
signature_path.write_text(json.dumps(signature, indent=2) + "\n", encoding="utf-8")
|
| 487 |
+
|
| 488 |
+
job = _read_job(work) if args.resume else {"next_module_idx": 0, "bad_rows": []}
|
| 489 |
+
state, original_ids, quant_preserves, bad_rows, args_dict = _prepare_real_state(
|
| 490 |
+
runtime, experts, args, model, job, work
|
| 491 |
+
)
|
| 492 |
+
job["cal_rows"] = args.cal_rows
|
| 493 |
+
job["cal_cols"] = args.cal_cols
|
| 494 |
+
job["codebook"] = "mcg"
|
| 495 |
+
job["bad_rows"] = sorted(bad_rows)
|
| 496 |
+
parallel = len(devices) > 1
|
| 497 |
+
replicas = [runtime.Model.from_config(config) for _ in devices[1:]] if parallel else []
|
| 498 |
+
for replica in replicas:
|
| 499 |
+
replica.calibration_all_experts = True
|
| 500 |
+
import torch
|
| 501 |
+
|
| 502 |
+
for module_idx, module, targets in plan:
|
| 503 |
+
if module_idx < job["next_module_idx"]:
|
| 504 |
+
continue
|
| 505 |
+
layer = _layer_number(module.key)
|
| 506 |
+
log(f"BAKE_LAYER start={module.key} targets={len(targets)}")
|
| 507 |
+
device = torch.device(f"cuda:{devices[0]}")
|
| 508 |
+
_load_module(module, device)
|
| 509 |
+
# Recalculate recursive linears after load; dense targets must be FP16
|
| 510 |
+
# here, while the routed expert variant is already EXL3.
|
| 511 |
+
targets = [m for m in module if isinstance(m, runtime.Linear) and is_dense_key(m.key)]
|
| 512 |
+
if len(targets) == 0 or any(getattr(m, "inner", None).__class__.__name__ != "LinearFP16" for m in targets):
|
| 513 |
+
raise RuntimeError(f"{module.key}: a dense target did not load as LinearFP16")
|
| 514 |
+
if any(m.qmap is None for m in targets):
|
| 515 |
+
raise RuntimeError(f"{module.key}: a dense target has no Hessian qmap")
|
| 516 |
+
keys = [m.key for m in targets]
|
| 517 |
+
if sorted(keys) != sorted(dense_keys_for_layer(layer, keys)): # LNA-LAB: module order is not lexical
|
| 518 |
+
raise RuntimeError(f"{module.key}: dense target ordering/filter changed after load; got={keys} expected={dense_keys_for_layer(layer, keys)}")
|
| 519 |
+
|
| 520 |
+
capture_replicas = None
|
| 521 |
+
if parallel:
|
| 522 |
+
capture_replicas = runtime.load_parallel_calib_modules(
|
| 523 |
+
replicas, module_idx, devices, None
|
| 524 |
+
)
|
| 525 |
+
if capture_replicas is not None:
|
| 526 |
+
def get_preserve(i, params):
|
| 527 |
+
params.update(quant_preserves[i])
|
| 528 |
+
params["quant_preserve"] = quant_preserves[i]
|
| 529 |
+
def put_preserve(i, params):
|
| 530 |
+
quant_preserves[i] = params["quant_preserve"]
|
| 531 |
+
capture_h, ref_states = runtime.capture_module_parallel(
|
| 532 |
+
model, [module] + capture_replicas, devices, None, state,
|
| 533 |
+
original_ids, get_preserve, put_preserve, False, 0,
|
| 534 |
+
f"CAPTURE {module.key}", bad_rows,
|
| 535 |
+
)
|
| 536 |
+
for replica in capture_replicas:
|
| 537 |
+
replica.unload()
|
| 538 |
+
else:
|
| 539 |
+
capture_h, ref_states, get_preserve, put_preserve = _capture_serial(
|
| 540 |
+
model, module, state, original_ids, quant_preserves, bad_rows
|
| 541 |
+
)
|
| 542 |
+
|
| 543 |
+
strategy = {m.key: bits_for_key(m.key, args.bits, args.attn_bits, args.shared_bits) for m in targets}
|
| 544 |
+
for linear in targets:
|
| 545 |
+
linear.inner.swap_cpu()
|
| 546 |
+
if len(targets) >= len(devices) and all(bits <= 8 for bits in strategy.values()):
|
| 547 |
+
runtime.quantize_linears_parallel(
|
| 548 |
+
args_dict, targets, config, strategy, module_idx, devices, None, capture_h, state
|
| 549 |
+
)
|
| 550 |
+
else:
|
| 551 |
+
runtime.quantize_linears_single(
|
| 552 |
+
args_dict, targets, config, strategy, module_idx, devices, None, capture_h, state
|
| 553 |
+
)
|
| 554 |
+
qtensors = {}
|
| 555 |
+
for linear in targets:
|
| 556 |
+
qtensors.update(linear.get_tensors())
|
| 557 |
+
expected = {f"{m.key}.{suffix}" for m in targets for suffix in EXL3_SUFFIXES}
|
| 558 |
+
if set(qtensors) != expected:
|
| 559 |
+
raise RuntimeError(f"{module.key}: converter emitted unexpected tensors")
|
| 560 |
+
runtime.save_tensor(qtensors, f"qtensors/layers.{layer}.safetensors", args_dict)
|
| 561 |
+
module.unload()
|
| 562 |
+
reload_keys = _inject_reload_keys(config.stc, qtensors)
|
| 563 |
+
_set_new_tensors(config.stc, qtensors)
|
| 564 |
+
_load_module(module, device, source=qtensors, close=False)
|
| 565 |
+
advance_replicas = None
|
| 566 |
+
if parallel:
|
| 567 |
+
advance_replicas = runtime.load_parallel_calib_modules(
|
| 568 |
+
replicas, module_idx, devices, None, source=qtensors
|
| 569 |
+
)
|
| 570 |
+
_set_new_tensors(config.stc, None)
|
| 571 |
+
_remove_reload_keys(config.stc, reload_keys)
|
| 572 |
+
if advance_replicas is not None:
|
| 573 |
+
def get_preserve(i, params):
|
| 574 |
+
params.update(quant_preserves[i])
|
| 575 |
+
params["quant_preserve"] = quant_preserves[i]
|
| 576 |
+
def put_preserve(i, params):
|
| 577 |
+
quant_preserves[i] = params["quant_preserve"]
|
| 578 |
+
runtime.advance_state_parallel(
|
| 579 |
+
model, [module] + advance_replicas, devices, None, state,
|
| 580 |
+
original_ids, get_preserve, put_preserve, ref_states, True,
|
| 581 |
+
module_idx == plan[-1][0], f"ADVANCE {module.key}", bad_rows,
|
| 582 |
+
)
|
| 583 |
+
for replica in advance_replicas:
|
| 584 |
+
replica.unload()
|
| 585 |
+
else:
|
| 586 |
+
_advance_serial(model, module, state, original_ids, quant_preserves, ref_states, bad_rows, True)
|
| 587 |
+
module.unload()
|
| 588 |
+
job["next_module_idx"] = module_idx + 1
|
| 589 |
+
job["bad_rows"] = sorted(bad_rows)
|
| 590 |
+
job["completed_layers"] = sorted(set(job.get("completed_layers", [])) | {layer})
|
| 591 |
+
_checkpoint(runtime, work, job, state, original_ids, args_dict)
|
| 592 |
+
log(f"BAKE_LAYER done={module.key} checkpoint=after_layer_{layer}")
|
| 593 |
+
log(f"BAKE_COMPLETE work={work} layers={len(plan)}")
|
| 594 |
+
return 0
|
| 595 |
+
|
| 596 |
+
|
| 597 |
+
def _read_work_tensors(work: Path) -> list[dict[str, Any]]:
|
| 598 |
+
qtensors = work / "qtensors"
|
| 599 |
+
if not qtensors.is_dir():
|
| 600 |
+
raise RuntimeError(f"missing work qtensors directory: {qtensors}")
|
| 601 |
+
result = []
|
| 602 |
+
seen = set()
|
| 603 |
+
for path in sorted(qtensors.glob("layers.*.safetensors"), key=lambda p: p.name):
|
| 604 |
+
match = _LAYER_FILE_RE.fullmatch(path.name)
|
| 605 |
+
if match is None:
|
| 606 |
+
continue
|
| 607 |
+
header_len, header = overlay.read_header(path)
|
| 608 |
+
for name, meta in header.items():
|
| 609 |
+
if name == "__metadata__":
|
| 610 |
+
continue
|
| 611 |
+
if not is_dense_key(name.rsplit(".", 1)[0]) or name.rsplit(".", 1)[1] not in EXL3_SUFFIXES:
|
| 612 |
+
raise RuntimeError(f"work file contains non-dense/non-EXL3 tensor: {path.name}:{name}")
|
| 613 |
+
base, suffix = name.rsplit(".", 1)
|
| 614 |
+
if _layer_number(base) != int(match.group("layer")):
|
| 615 |
+
raise RuntimeError(f"tensor {name} is in the wrong work layer file {path.name}")
|
| 616 |
+
if name in seen:
|
| 617 |
+
raise RuntimeError(f"duplicate work tensor: {name}")
|
| 618 |
+
expected_dtype = EXL3_DTYPES[suffix]
|
| 619 |
+
if meta.get("dtype") != expected_dtype:
|
| 620 |
+
raise RuntimeError(f"{name}: expected dtype {expected_dtype}, got {meta.get('dtype')}")
|
| 621 |
+
seen.add(name)
|
| 622 |
+
result.append({"target": name, "base": base, "suffix": suffix,
|
| 623 |
+
"meta": meta, "path": path, "header_len": header_len})
|
| 624 |
+
if not result:
|
| 625 |
+
raise RuntimeError(f"no complete dense EXL3 tensors under {qtensors}")
|
| 626 |
+
grouped = {}
|
| 627 |
+
for item in result:
|
| 628 |
+
grouped.setdefault(item["base"], set()).add(item["suffix"])
|
| 629 |
+
incomplete = [base for base, suffixes in grouped.items() if suffixes != set(EXL3_SUFFIXES)]
|
| 630 |
+
if incomplete:
|
| 631 |
+
raise RuntimeError(f"work tensors have incomplete EXL3 groups, e.g. {incomplete[0]}")
|
| 632 |
+
for layer in sorted({_layer_number(base) for base in grouped}):
|
| 633 |
+
wo_a = {
|
| 634 |
+
int(base.rsplit(".", 1)[1])
|
| 635 |
+
for base in grouped
|
| 636 |
+
if _layer_number(base) == layer and _WO_A_RE.fullmatch(base.split(".", 2)[2])
|
| 637 |
+
}
|
| 638 |
+
if wo_a and wo_a != set(range(8)):
|
| 639 |
+
raise RuntimeError(f"work layer {layer} has incomplete wo_a slices: {sorted(wo_a)}")
|
| 640 |
+
return sorted(result, key=lambda item: item["target"])
|
| 641 |
+
|
| 642 |
+
|
| 643 |
+
def _vllm_prefixes(layer: int, root: str, shared_prefix: str) -> dict[str, str]:
|
| 644 |
+
base = f"{root}.layers.{layer}"
|
| 645 |
+
return {
|
| 646 |
+
"attn.wq_a": f"{base}.attn.fused_wqa_wkv",
|
| 647 |
+
"attn.wkv": f"{base}.attn.fused_wqa_wkv",
|
| 648 |
+
"attn.wq_b": f"{base}.attn.wq_b",
|
| 649 |
+
"attn.wo_b": f"{base}.attn.wo_b",
|
| 650 |
+
"attn.compressor.wkv": f"{base}.attn.compressor.fused_wkv_wgate",
|
| 651 |
+
"attn.compressor.wgate": f"{base}.attn.compressor.fused_wkv_wgate",
|
| 652 |
+
"attn.indexer.compressor.wkv": f"{base}.attn.indexer.compressor.fused_wkv_wgate",
|
| 653 |
+
"attn.indexer.compressor.wgate": f"{base}.attn.indexer.compressor.fused_wkv_wgate",
|
| 654 |
+
"attn.indexer.wq_b": f"{base}.attn.indexer.wq_b",
|
| 655 |
+
"ffn.shared_experts.w1": f"{base}.{shared_prefix}.gate_up_proj",
|
| 656 |
+
"ffn.shared_experts.w3": f"{base}.{shared_prefix}.gate_up_proj",
|
| 657 |
+
"ffn.shared_experts.w2": f"{base}.{shared_prefix}.down_proj",
|
| 658 |
+
}
|
| 659 |
+
|
| 660 |
+
|
| 661 |
+
def _config_with_vllm_block(base_config: dict, work_items: list[dict], args) -> dict:
|
| 662 |
+
config = copy.deepcopy(base_config)
|
| 663 |
+
q = config.setdefault("quantization_config", {})
|
| 664 |
+
layers = {}
|
| 665 |
+
for item in work_items:
|
| 666 |
+
base = item["base"]
|
| 667 |
+
tail = base.split(".", 2)[2]
|
| 668 |
+
if _WO_A_RE.fullmatch(tail):
|
| 669 |
+
# LNA-LAB: the plugin consumes one already-sliced wo_a per TP rank.
|
| 670 |
+
vllm = f"{args.vllm_root}.layers.{_layer_number(base)}.attn.wo_a"
|
| 671 |
+
else:
|
| 672 |
+
vllm = _vllm_prefixes(
|
| 673 |
+
_layer_number(base), args.vllm_root, args.vllm_shared_prefix
|
| 674 |
+
).get(tail)
|
| 675 |
+
if vllm is None:
|
| 676 |
+
raise RuntimeError(f"no vLLM prefix mapping for {base}")
|
| 677 |
+
bits = args.attn_bits if dense_kind(base) == "attn" and args.attn_bits is not None else \
|
| 678 |
+
args.shared_bits if dense_kind(base) == "shared" and args.shared_bits is not None else args.bits
|
| 679 |
+
bits = _schema_bits(bits)
|
| 680 |
+
old = layers.setdefault(vllm, {"bits": bits})
|
| 681 |
+
if old["bits"] != bits:
|
| 682 |
+
raise RuntimeError(f"fused vLLM module has mixed bits: {vllm}")
|
| 683 |
+
# LNA-LAB: unmatched dense modules (notably indexer.weights_proj) stay BF16.
|
| 684 |
+
q.setdefault("non_routed_dtype_policy", "bf16_as_stored")
|
| 685 |
+
q["non_routed_exl3"] = {
|
| 686 |
+
"codebook": "mcg",
|
| 687 |
+
"layers": dict(sorted(layers.items())),
|
| 688 |
+
}
|
| 689 |
+
return config
|
| 690 |
+
|
| 691 |
+
|
| 692 |
+
def _write_dense_shard(out: Path, items: list[dict]) -> dict[str, Any]:
|
| 693 |
+
destination = out / DENSE_SHARD
|
| 694 |
+
partial = out / (DENSE_SHARD + ".partial")
|
| 695 |
+
entries = [(item["target"], item["meta"]) for item in items]
|
| 696 |
+
with partial.open("wb") as dst:
|
| 697 |
+
header_len, _ = overlay.common.write_safetensors_header(
|
| 698 |
+
dst, entries, {"format": "pt", "source": "exllamav3-dense-bake", "codebook": "mcg"}
|
| 699 |
+
)
|
| 700 |
+
for item in items:
|
| 701 |
+
with item["path"].open("rb") as src:
|
| 702 |
+
overlay.common.copy_payload(src, dst, 8 + item["header_len"], item["meta"])
|
| 703 |
+
os.replace(partial, destination)
|
| 704 |
+
return {"name": DENSE_SHARD, "header_len": header_len, "entries": entries}
|
| 705 |
+
|
| 706 |
+
|
| 707 |
+
def _link_base_files(experts: Path, out: Path, rewrite_files: set[str]) -> None:
|
| 708 |
+
excluded = rewrite_files | {CONFIG_NAME, INDEX_NAME, DENSE_SHARD, DENSE_SHARD + ".partial"}
|
| 709 |
+
for entry in sorted(experts.iterdir(), key=lambda p: p.name):
|
| 710 |
+
if entry.name in excluded or not (entry.is_file() or entry.is_symlink()):
|
| 711 |
+
continue
|
| 712 |
+
destination = out / entry.name
|
| 713 |
+
if destination.exists() or destination.is_symlink():
|
| 714 |
+
if destination.is_symlink() and destination.resolve() == entry.resolve():
|
| 715 |
+
continue
|
| 716 |
+
raise RuntimeError(f"output path already exists: {destination}")
|
| 717 |
+
destination.symlink_to(entry.resolve())
|
| 718 |
+
|
| 719 |
+
|
| 720 |
+
def plan_merge(work: Path, experts: Path, args, *, require_complete: bool = True) -> dict[str, Any]:
|
| 721 |
+
work_items = _read_work_tensors(work)
|
| 722 |
+
source, source_map = overlay.source_headers(experts)
|
| 723 |
+
overlay.validate_source_index(experts, source_map)
|
| 724 |
+
dense_bases = sorted({item["base"] for item in work_items})
|
| 725 |
+
expected_bases = set()
|
| 726 |
+
for name in source:
|
| 727 |
+
if not name.endswith(".weight"):
|
| 728 |
+
continue
|
| 729 |
+
base = name[:-len(".weight")]
|
| 730 |
+
if base.startswith("layers.") and base.endswith(".attn.wo_a"):
|
| 731 |
+
layer = _layer_number(base)
|
| 732 |
+
expected_bases.update(f"layers.{layer}.attn.wo_a.slice.{i}" for i in range(8))
|
| 733 |
+
elif is_dense_key(base):
|
| 734 |
+
expected_bases.add(base)
|
| 735 |
+
missing_work = sorted(expected_bases - set(dense_bases))
|
| 736 |
+
extra_work = sorted(set(dense_bases) - expected_bases)
|
| 737 |
+
if (missing_work or extra_work) and require_complete:
|
| 738 |
+
raise RuntimeError(
|
| 739 |
+
"work directory is not a complete dense bake: "
|
| 740 |
+
f"missing={missing_work[:3]} extra={extra_work[:3]}"
|
| 741 |
+
)
|
| 742 |
+
if missing_work or extra_work:
|
| 743 |
+
log(
|
| 744 |
+
"WARNING dry-run merge is partial: "
|
| 745 |
+
f"missing={missing_work[:3]} extra={extra_work[:3]}"
|
| 746 |
+
)
|
| 747 |
+
dropped = set() # LNA-LAB: the 8 wo_a slices replace the single source wo_a.weight
|
| 748 |
+
for base in dense_bases:
|
| 749 |
+
m = re.match(r"^(layers\.[0-9]+\.attn\.wo_a)\.slice\.[0-7]$", base)
|
| 750 |
+
dropped.add(f"{m.group(1)}.weight" if m else f"{base}.weight")
|
| 751 |
+
missing = sorted(name for name in dropped if name not in source)
|
| 752 |
+
if missing:
|
| 753 |
+
raise RuntimeError(f"dense work tensor has no BF16/fp8 source weight to replace: {missing[0]}")
|
| 754 |
+
by_file: dict[str, set[str]] = {}
|
| 755 |
+
for name in dropped:
|
| 756 |
+
by_file.setdefault(source[name][0].name, set()).add(name)
|
| 757 |
+
rewrite_files = set(by_file)
|
| 758 |
+
new_plan = [{"target": item["target"], "meta": item["meta"]} for item in work_items]
|
| 759 |
+
base_config_path = experts / CONFIG_NAME
|
| 760 |
+
try:
|
| 761 |
+
base_config = json.loads(base_config_path.read_text(encoding="utf-8"))
|
| 762 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 763 |
+
raise RuntimeError(f"cannot read {base_config_path}: {exc}") from exc
|
| 764 |
+
config = _config_with_vllm_block(base_config, work_items, args)
|
| 765 |
+
index = overlay.output_index(source, dropped, new_plan)
|
| 766 |
+
# layer_overlay's generic planner uses its historical shard name; T4 has
|
| 767 |
+
# a distinct dense shard so it cannot be confused with a layer overlay.
|
| 768 |
+
new_targets = {item["target"] for item in new_plan} # LNA-LAB: only the dense tensors go to the dense shard;
|
| 769 |
+
index["weight_map"] = { # the base pack may itself carry a layer_overlay shard of that name
|
| 770 |
+
name: (DENSE_SHARD if (filename == overlay.NEW_SHARD and name in new_targets) else filename)
|
| 771 |
+
for name, filename in index["weight_map"].items()
|
| 772 |
+
}
|
| 773 |
+
return {
|
| 774 |
+
"source": source,
|
| 775 |
+
"source_map": source_map,
|
| 776 |
+
"items": work_items,
|
| 777 |
+
"dropped": dropped,
|
| 778 |
+
"rewrite_files": rewrite_files,
|
| 779 |
+
"new_plan": new_plan,
|
| 780 |
+
"config": config,
|
| 781 |
+
"index": index,
|
| 782 |
+
}
|
| 783 |
+
|
| 784 |
+
|
| 785 |
+
def _verify_merge(experts: Path, out: Path, plan: dict[str, Any]) -> None:
|
| 786 |
+
index_path = out / INDEX_NAME
|
| 787 |
+
if json.loads(index_path.read_text(encoding="utf-8")) != plan["index"]:
|
| 788 |
+
raise RuntimeError("merge index differs from plan")
|
| 789 |
+
output_headers = {}
|
| 790 |
+
for filename in sorted(set(plan["index"]["weight_map"].values())):
|
| 791 |
+
path = out / filename
|
| 792 |
+
if not path.is_file():
|
| 793 |
+
raise RuntimeError(f"indexed output shard missing: {path}")
|
| 794 |
+
header_len, header = overlay.read_header(path)
|
| 795 |
+
expected_names = {
|
| 796 |
+
name for name, mapped in plan["index"]["weight_map"].items() if mapped == filename
|
| 797 |
+
}
|
| 798 |
+
actual_names = {name for name in header if name != "__metadata__"}
|
| 799 |
+
if actual_names != expected_names:
|
| 800 |
+
raise RuntimeError(
|
| 801 |
+
f"output shard {filename} has unexpected tensor set: "
|
| 802 |
+
f"missing={sorted(expected_names - actual_names)[:2]} "
|
| 803 |
+
f"extra={sorted(actual_names - expected_names)[:2]}"
|
| 804 |
+
)
|
| 805 |
+
for name, mapped in plan["index"]["weight_map"].items():
|
| 806 |
+
if mapped == filename:
|
| 807 |
+
if name not in header:
|
| 808 |
+
raise RuntimeError(f"indexed tensor missing from {filename}: {name}")
|
| 809 |
+
output_headers[name] = (path, header_len, header)
|
| 810 |
+
if any(name in plan["dropped"] for name in output_headers):
|
| 811 |
+
raise RuntimeError("dropped source tensor remains in merged output")
|
| 812 |
+
for filename in plan["rewrite_files"]:
|
| 813 |
+
path = out / filename
|
| 814 |
+
if path.is_symlink():
|
| 815 |
+
raise RuntimeError(f"rewritten shard is still a symlink: {filename}")
|
| 816 |
+
_, header = overlay.read_header(path)
|
| 817 |
+
if any(name in plan["dropped"] for name in header):
|
| 818 |
+
raise RuntimeError(f"dropped tensor remains in rewritten shard: {filename}")
|
| 819 |
+
for name, (source_path, source_header_len, source_header) in plan["source"].items():
|
| 820 |
+
if name in plan["dropped"] or source_path.name not in plan["rewrite_files"]:
|
| 821 |
+
continue
|
| 822 |
+
output_path, output_header_len, output_header = output_headers[name]
|
| 823 |
+
if source_header[name].get("dtype") != output_header[name].get("dtype") or \
|
| 824 |
+
source_header[name].get("shape") != output_header[name].get("shape"):
|
| 825 |
+
raise RuntimeError(f"retained tensor metadata changed: {name}")
|
| 826 |
+
overlay.compare_payloads(source_path, output_path, source_header[name], output_header[name],
|
| 827 |
+
source_header_len, output_header_len)
|
| 828 |
+
for item in plan["items"]:
|
| 829 |
+
path, header_len, header = output_headers[item["target"]]
|
| 830 |
+
meta = header[item["target"]]
|
| 831 |
+
if meta.get("dtype") != item["meta"].get("dtype") or meta.get("shape") != item["meta"].get("shape"):
|
| 832 |
+
raise RuntimeError(f"new tensor metadata changed: {item['target']}")
|
| 833 |
+
if item["target"].endswith(".mcg"):
|
| 834 |
+
start, end = meta["data_offsets"]
|
| 835 |
+
with path.open("rb") as stream:
|
| 836 |
+
stream.seek(8 + header_len + start)
|
| 837 |
+
if struct.unpack("<i", stream.read(end - start))[0] != MCG_MARKER:
|
| 838 |
+
raise RuntimeError(f"MCG marker mismatch: {item['target']}")
|
| 839 |
+
if json.loads((out / CONFIG_NAME).read_text(encoding="utf-8")) != plan["config"]:
|
| 840 |
+
raise RuntimeError("merge config differs from plan")
|
| 841 |
+
for entry in experts.iterdir():
|
| 842 |
+
if entry.name in plan["rewrite_files"] or entry.name in {CONFIG_NAME, INDEX_NAME, DENSE_SHARD}:
|
| 843 |
+
continue
|
| 844 |
+
if not (entry.is_file() or entry.is_symlink()):
|
| 845 |
+
continue
|
| 846 |
+
linked = out / entry.name
|
| 847 |
+
if not linked.is_symlink() or linked.resolve() != entry.resolve():
|
| 848 |
+
raise RuntimeError(f"unchanged source file is not linked unchanged: {entry.name}")
|
| 849 |
+
log(f"MERGE_VERIFY_OK indexed_tensors={len(output_headers)} dropped_absent=true retained_payloads_byte_identical=true")
|
| 850 |
+
|
| 851 |
+
|
| 852 |
+
def merge(args: argparse.Namespace) -> int:
|
| 853 |
+
work = Path(args.work).expanduser().resolve()
|
| 854 |
+
# LNA-LAB: the merged config must describe the bits that were actually baked.
|
| 855 |
+
# Unless --bits was given explicitly, take bits/attn_bits/shared_bits from work/args.json.
|
| 856 |
+
if not getattr(args, "_bits_explicit", False):
|
| 857 |
+
recorded = work / "args.json"
|
| 858 |
+
if recorded.is_file():
|
| 859 |
+
rec = json.loads(recorded.read_text())
|
| 860 |
+
args.bits = rec.get("bits", args.bits)
|
| 861 |
+
args.attn_bits = rec.get("attn_bits", args.attn_bits)
|
| 862 |
+
args.shared_bits = rec.get("shared_bits", args.shared_bits)
|
| 863 |
+
log(f"MERGE_BITS from {recorded}: bits={args.bits:g} attn_bits={args.attn_bits} shared_bits={args.shared_bits}")
|
| 864 |
+
experts = Path(args.experts).expanduser().resolve()
|
| 865 |
+
out = Path(args.merge).expanduser().resolve()
|
| 866 |
+
overlay.safe_output_path(experts, out)
|
| 867 |
+
plan = plan_merge(work, experts, args, require_complete=not args.dry_run)
|
| 868 |
+
log(f"MERGE_PLAN rewrite_files={len(plan['rewrite_files'])} dropped={len(plan['dropped'])} new={len(plan['new_plan'])}")
|
| 869 |
+
if args.dry_run:
|
| 870 |
+
log(f"DRY_RUN_MERGE output={out} symlinks=all_unchanged_source_files")
|
| 871 |
+
return 0
|
| 872 |
+
if out.exists() and any(out.iterdir()):
|
| 873 |
+
raise RuntimeError(f"refusing to merge into non-empty directory: {out}")
|
| 874 |
+
out.mkdir(parents=True, exist_ok=True)
|
| 875 |
+
_link_base_files(experts, out, plan["rewrite_files"])
|
| 876 |
+
for filename in sorted(plan["rewrite_files"]):
|
| 877 |
+
overlay.rewrite_source_shard(experts / filename, out / filename, plan["dropped"])
|
| 878 |
+
_write_dense_shard(out, plan["items"])
|
| 879 |
+
overlay.write_output_metadata(out, plan["index"], plan["config"])
|
| 880 |
+
_verify_merge(experts, out, plan)
|
| 881 |
+
return 0
|
| 882 |
+
|
| 883 |
+
|
| 884 |
+
def make_parser() -> argparse.ArgumentParser:
|
| 885 |
+
parser = argparse.ArgumentParser(description=__doc__)
|
| 886 |
+
parser.add_argument("--src", required=True, help="bf16/fp8 DeepSeek-V4 source checkpoint")
|
| 887 |
+
parser.add_argument("--experts", required=True, help="existing routed-expert EXL3 pack")
|
| 888 |
+
parser.add_argument("--work", required=True, help="resumable qtensors/checkpoint directory")
|
| 889 |
+
parser.add_argument("--merge", help="write merged overlay pack to this directory")
|
| 890 |
+
parser.add_argument("--bits", type=float, default=4, help="dense default EXL3 K (default: 4)")
|
| 891 |
+
parser.add_argument("--attn-bits", type=float, default=None)
|
| 892 |
+
parser.add_argument("--shared-bits", type=float, default=None)
|
| 893 |
+
parser.add_argument("--cal-rows", type=int, default=250)
|
| 894 |
+
parser.add_argument("--cal-cols", type=int, default=2048)
|
| 895 |
+
parser.add_argument("--devices", default="0,1", help="CUDA device IDs, e.g. 0,1")
|
| 896 |
+
parser.add_argument("--resume", action="store_true", help="resume completed module checkpoints")
|
| 897 |
+
parser.add_argument("--dry-run", action="store_true", help="offline layout/merge planning; no writes")
|
| 898 |
+
parser.add_argument("--vllm-root", default="language_model.model", help="vLLM model root prefix (LNA-LAB: vision class roots the text model under language_model.model)")
|
| 899 |
+
parser.add_argument("--vllm-shared-prefix", default="ffn.shared_experts",
|
| 900 |
+
help="vLLM shared-expert child prefix under model.layers.N")
|
| 901 |
+
parser.add_argument("--verbose", action="store_true")
|
| 902 |
+
return parser
|
| 903 |
+
|
| 904 |
+
|
| 905 |
+
def main(argv=None) -> int:
|
| 906 |
+
args = make_parser().parse_args(argv)
|
| 907 |
+
args._bits_explicit = any(a == "--bits" or a.startswith("--bits=") for a in (argv if argv is not None else sys.argv[1:]))
|
| 908 |
+
try:
|
| 909 |
+
if args.merge:
|
| 910 |
+
# A merge is independent of CUDA; --dry-run additionally performs
|
| 911 |
+
# the requested source/model layout check before header planning.
|
| 912 |
+
if args.dry_run:
|
| 913 |
+
dry_run_layout(Path(args.src).expanduser().resolve(), Path(args.experts).expanduser().resolve())
|
| 914 |
+
return merge(args)
|
| 915 |
+
import torch # LNA-LAB: the converter runs under inference_mode; loaded tensors are inference tensors
|
| 916 |
+
with torch.inference_mode():
|
| 917 |
+
return bake(args)
|
| 918 |
+
except (RuntimeError, ValueError, OSError) as exc:
|
| 919 |
+
log(f"ERROR {exc}")
|
| 920 |
+
return 2
|
| 921 |
+
|
| 922 |
+
|
| 923 |
+
if __name__ == "__main__":
|
| 924 |
+
raise SystemExit(main())
|
lna-lab/tools/layer_overlay.py
ADDED
|
@@ -0,0 +1,957 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Build a selected-main-layer K2 EXL3 overlay.
|
| 3 |
+
|
| 4 |
+
The source pack is never changed. Unchanged source files are linked into the
|
| 5 |
+
output, source shards containing the selected K3 routed experts are rewritten
|
| 6 |
+
tensor-by-tensor, and the replacement K2 tensors are written to one new shard.
|
| 7 |
+
Donor tensors are fetched with HTTP Range requests one tensor at a time; a
|
| 8 |
+
local donor-shard directory can be supplied through ``LAYER_OVERLAY_LOCAL_DIR``.
|
| 9 |
+
|
| 10 |
+
``--dry-run`` is deliberately offline: it scans the local source headers and
|
| 11 |
+
donor index and plans the replacement using a K2 ABI exemplar already in the
|
| 12 |
+
source pack. A donor header cache, when supplied, adds header-level shape and
|
| 13 |
+
dtype validation without network access.
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
from __future__ import annotations
|
| 17 |
+
|
| 18 |
+
import argparse
|
| 19 |
+
import copy
|
| 20 |
+
import json
|
| 21 |
+
import os
|
| 22 |
+
import re
|
| 23 |
+
import struct
|
| 24 |
+
import sys
|
| 25 |
+
import time
|
| 26 |
+
import urllib.error
|
| 27 |
+
import urllib.request
|
| 28 |
+
from pathlib import Path
|
| 29 |
+
from typing import BinaryIO, Iterable
|
| 30 |
+
|
| 31 |
+
try:
|
| 32 |
+
from . import mtp_overlay as common
|
| 33 |
+
except ImportError: # Running as ``python tools/layer_overlay.py``.
|
| 34 |
+
import mtp_overlay as common
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
CONFIG_NAME = "config.json"
|
| 38 |
+
INDEX_NAME = "model.safetensors.index.json"
|
| 39 |
+
NEW_SHARD = "model-layers-k2.safetensors"
|
| 40 |
+
DEFAULT_DONOR_INDEX = "/run/media/tonoken3/DATA1/.tmp/k2v1-index.json"
|
| 41 |
+
DEFAULT_HEADER_CACHE = "/run/media/tonoken3/DATA1/.tmp/layer-overlay-donor-headers.json"
|
| 42 |
+
CHUNK_SIZE = common.CHUNK_SIZE
|
| 43 |
+
RETRY_COUNT = common.RETRY_COUNT
|
| 44 |
+
|
| 45 |
+
DTYPE_BYTES = common.DTYPE_BYTES
|
| 46 |
+
EXL3_SUFFIXES = ("trellis", "suh", "svh", "mcg")
|
| 47 |
+
EXL3_DTYPE = {"trellis": "I16", "suh": "F16", "svh": "F16", "mcg": "I32"}
|
| 48 |
+
PROJECTION_MAP = {"gate_proj": "w1", "up_proj": "w3", "down_proj": "w2"}
|
| 49 |
+
PROJECTION_ORDER = {"gate_proj": 0, "up_proj": 1, "down_proj": 2}
|
| 50 |
+
|
| 51 |
+
SOURCE_EXPERT_RE = re.compile(
|
| 52 |
+
r"^layers\.(?P<layer>[0-9]+)\.ffn\.experts\."
|
| 53 |
+
r"(?P<expert>[0-9]+)\.(?P<projection>w[123])\."
|
| 54 |
+
r"(?P<suffix>trellis|suh|svh|mcg)$"
|
| 55 |
+
)
|
| 56 |
+
DONOR_EXPERT_RE = re.compile(
|
| 57 |
+
r"^model\.layers\.(?P<layer>[0-9]+)\.mlp\.experts\."
|
| 58 |
+
r"(?P<expert>[0-9]+)\.(?P<projection>gate_proj|up_proj|down_proj)\."
|
| 59 |
+
r"(?P<suffix>trellis|suh|svh|mcg)$"
|
| 60 |
+
)
|
| 61 |
+
|
| 62 |
+
# These helpers are intentionally re-exported for small offline tests and for
|
| 63 |
+
# callers that used the corresponding mtp_overlay helpers.
|
| 64 |
+
parse_header = common.parse_header
|
| 65 |
+
validate_header = common.validate_header
|
| 66 |
+
read_header = common.read_header
|
| 67 |
+
tensor_nbytes = common.tensor_nbytes
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def log(message: str) -> None:
|
| 71 |
+
print(message, flush=True)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def parse_layers(value: str) -> tuple[int, ...]:
|
| 75 |
+
try:
|
| 76 |
+
values = [int(part.strip()) for part in value.split(",") if part.strip()]
|
| 77 |
+
except ValueError as exc:
|
| 78 |
+
raise ValueError(f"invalid --layers value {value!r}") from exc
|
| 79 |
+
if not values or any(layer < 0 for layer in values):
|
| 80 |
+
raise ValueError(f"--layers must contain non-negative layer numbers: {value!r}")
|
| 81 |
+
if len(set(values)) != len(values):
|
| 82 |
+
raise ValueError(f"--layers contains duplicates: {value!r}")
|
| 83 |
+
return tuple(sorted(values))
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def source_layer(name: str) -> int | None:
|
| 87 |
+
match = SOURCE_EXPERT_RE.fullmatch(name)
|
| 88 |
+
return None if match is None else int(match.group("layer"))
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def donor_layer(name: str) -> int | None:
|
| 92 |
+
match = DONOR_EXPERT_RE.fullmatch(name)
|
| 93 |
+
return None if match is None else int(match.group("layer"))
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
def map_donor_name(name: str) -> str:
|
| 97 |
+
"""Map one donor main-layer expert name to the target pack ABI."""
|
| 98 |
+
match = DONOR_EXPERT_RE.fullmatch(name)
|
| 99 |
+
if match is None:
|
| 100 |
+
raise ValueError(f"not a donor main-layer EXL3 expert tensor: {name}")
|
| 101 |
+
return (
|
| 102 |
+
f"layers.{match.group('layer')}.ffn.experts.{int(match.group('expert'))}."
|
| 103 |
+
f"{PROJECTION_MAP[match.group('projection')]}.{match.group('suffix')}"
|
| 104 |
+
)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
# Short alias used by offline tests and useful to importers.
|
| 108 |
+
map_name = map_donor_name
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def is_donor_expert_name(name: str) -> bool:
|
| 112 |
+
return DONOR_EXPERT_RE.fullmatch(name) is not None
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def is_source_expert_name(name: str) -> bool:
|
| 116 |
+
return SOURCE_EXPERT_RE.fullmatch(name) is not None
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def donor_base_url(donor: str) -> str:
|
| 120 |
+
return donor.rstrip("/") if donor.startswith("http") else f"https://huggingface.co/{donor}/resolve/main"
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def local_or_remote_bytes(url: str, byte_range: tuple[int, int] | None = None) -> bytes:
|
| 124 |
+
"""Read a range from a local donor shard or a remote URL.
|
| 125 |
+
|
| 126 |
+
The real build calls this only for header ranges and selected tensor
|
| 127 |
+
ranges. It never requests a complete donor shard.
|
| 128 |
+
"""
|
| 129 |
+
local_dir = os.environ.get("LAYER_OVERLAY_LOCAL_DIR")
|
| 130 |
+
filename = url.split("?", 1)[0].rsplit("/", 1)[-1]
|
| 131 |
+
if local_dir:
|
| 132 |
+
local_path = Path(local_dir).expanduser() / filename
|
| 133 |
+
if local_path.is_file():
|
| 134 |
+
with local_path.open("rb") as stream:
|
| 135 |
+
if byte_range is None:
|
| 136 |
+
return stream.read()
|
| 137 |
+
stream.seek(byte_range[0])
|
| 138 |
+
data = stream.read(byte_range[1] - byte_range[0] + 1)
|
| 139 |
+
expected = byte_range[1] - byte_range[0] + 1
|
| 140 |
+
if len(data) != expected:
|
| 141 |
+
raise RuntimeError(
|
| 142 |
+
f"short local range read {local_path}: got {len(data)}, expected {expected}"
|
| 143 |
+
)
|
| 144 |
+
return data
|
| 145 |
+
|
| 146 |
+
for attempt in range(RETRY_COUNT):
|
| 147 |
+
try:
|
| 148 |
+
request = urllib.request.Request(url, headers={"User-Agent": "layer_overlay/1"})
|
| 149 |
+
if byte_range is not None:
|
| 150 |
+
request.add_header("Range", f"bytes={byte_range[0]}-{byte_range[1]}")
|
| 151 |
+
with urllib.request.urlopen(request, timeout=60) as response:
|
| 152 |
+
status = getattr(response, "status", response.getcode())
|
| 153 |
+
if byte_range is not None and status != 206:
|
| 154 |
+
raise RuntimeError(f"expected HTTP 206 for Range request, got {status}")
|
| 155 |
+
data = response.read()
|
| 156 |
+
expected = None if byte_range is None else byte_range[1] - byte_range[0] + 1
|
| 157 |
+
if expected is not None and len(data) != expected:
|
| 158 |
+
raise RuntimeError(f"short range read: got {len(data)}, expected {expected}")
|
| 159 |
+
return data
|
| 160 |
+
except (OSError, TimeoutError, urllib.error.URLError, urllib.error.HTTPError, RuntimeError) as exc:
|
| 161 |
+
if attempt == RETRY_COUNT - 1:
|
| 162 |
+
raise RuntimeError(f"failed fetching {url} range={byte_range}: {exc}") from exc
|
| 163 |
+
delay = 2**attempt
|
| 164 |
+
log(f"RETRY attempt={attempt + 1}/{RETRY_COUNT} range={byte_range}: {exc}")
|
| 165 |
+
time.sleep(delay)
|
| 166 |
+
raise AssertionError("unreachable")
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
def donor_index_from_file(path: Path) -> dict[str, str]:
|
| 170 |
+
try:
|
| 171 |
+
data = json.loads(path.read_text())
|
| 172 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 173 |
+
raise RuntimeError(f"cannot read donor index {path}: {exc}") from exc
|
| 174 |
+
weight_map = data.get("weight_map")
|
| 175 |
+
if not isinstance(weight_map, dict) or not all(
|
| 176 |
+
isinstance(name, str) and isinstance(filename, str)
|
| 177 |
+
for name, filename in weight_map.items()
|
| 178 |
+
):
|
| 179 |
+
raise RuntimeError(f"donor index {path} has no valid weight_map")
|
| 180 |
+
return weight_map
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def source_headers(src: Path) -> tuple[dict[str, tuple[Path, int, dict]], dict[str, str]]:
|
| 184 |
+
"""Scan every local safetensors file, including an existing overlay shard."""
|
| 185 |
+
candidates = sorted(
|
| 186 |
+
(path for path in src.iterdir() if path.name.endswith(".safetensors") and path.is_file()),
|
| 187 |
+
key=lambda path: path.name,
|
| 188 |
+
)
|
| 189 |
+
if not candidates:
|
| 190 |
+
raise RuntimeError(f"no safetensors shards found under {src}")
|
| 191 |
+
headers: dict[str, tuple[Path, int, dict]] = {}
|
| 192 |
+
weight_map: dict[str, str] = {}
|
| 193 |
+
for path in candidates:
|
| 194 |
+
header_len, header = read_header(path)
|
| 195 |
+
for name in header:
|
| 196 |
+
if name == "__metadata__":
|
| 197 |
+
continue
|
| 198 |
+
if name in weight_map:
|
| 199 |
+
raise RuntimeError(f"duplicate source tensor {name!r}")
|
| 200 |
+
headers[name] = (path, header_len, header)
|
| 201 |
+
weight_map[name] = path.name
|
| 202 |
+
log(f"SOURCE_HEADERS files={len(candidates)} tensors={len(headers)}")
|
| 203 |
+
return headers, weight_map
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
def validate_source_index(src: Path, source_map: dict[str, str]) -> None:
|
| 207 |
+
"""Cross-check the numbered source headers against an existing index."""
|
| 208 |
+
path = src / INDEX_NAME
|
| 209 |
+
if not path.is_file():
|
| 210 |
+
log("SOURCE_INDEX absent=allowed; numbered headers are authoritative")
|
| 211 |
+
return
|
| 212 |
+
try:
|
| 213 |
+
data = json.loads(path.read_text())
|
| 214 |
+
index_map = data["weight_map"]
|
| 215 |
+
except (OSError, json.JSONDecodeError, KeyError, TypeError) as exc:
|
| 216 |
+
raise RuntimeError(f"invalid source index {path}: {exc}") from exc
|
| 217 |
+
if not isinstance(index_map, dict):
|
| 218 |
+
raise RuntimeError(f"source index {path} has no object weight_map")
|
| 219 |
+
if set(index_map) != set(source_map):
|
| 220 |
+
raise RuntimeError(
|
| 221 |
+
f"source index tensor set differs from local headers: "
|
| 222 |
+
f"missing={sorted(set(source_map) - set(index_map))[:5]} "
|
| 223 |
+
f"extra={sorted(set(index_map) - set(source_map))[:5]}"
|
| 224 |
+
)
|
| 225 |
+
for name, filename in source_map.items():
|
| 226 |
+
if index_map.get(name) != filename:
|
| 227 |
+
raise RuntimeError(f"source index disagrees with header for {name}: {index_map.get(name)!r} != {filename!r}")
|
| 228 |
+
log(f"SOURCE_INDEX checked=local_headers tensors={len(source_map)} indexed_total={len(index_map)}")
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
def _abi_meta(meta: dict) -> dict:
|
| 232 |
+
return {"dtype": meta.get("dtype"), "shape": list(meta.get("shape", []))}
|
| 233 |
+
|
| 234 |
+
|
| 235 |
+
def validate_abi(abi: dict[str, dict[str, dict]], label: str) -> None:
|
| 236 |
+
for projection in ("w1", "w2", "w3"):
|
| 237 |
+
if projection not in abi:
|
| 238 |
+
raise RuntimeError(f"{label}: missing projection {projection}")
|
| 239 |
+
for suffix in EXL3_SUFFIXES:
|
| 240 |
+
meta = abi[projection].get(suffix)
|
| 241 |
+
if not isinstance(meta, dict) or meta.get("dtype") != EXL3_DTYPE[suffix]:
|
| 242 |
+
raise RuntimeError(f"{label}: unexpected {projection}.{suffix} metadata: {meta}")
|
| 243 |
+
shape = meta.get("shape")
|
| 244 |
+
if not isinstance(shape, list) or not all(isinstance(dim, int) and dim >= 0 for dim in shape):
|
| 245 |
+
raise RuntimeError(f"{label}: invalid {projection}.{suffix} shape: {shape}")
|
| 246 |
+
if abi[projection]["trellis"]["shape"][-1:] != [32]:
|
| 247 |
+
raise RuntimeError(f"{label}: {projection}.trellis is not K2: {abi[projection]['trellis']}")
|
| 248 |
+
|
| 249 |
+
|
| 250 |
+
def find_source_k2_abi(
|
| 251 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 252 |
+
selected_layers: Iterable[int],
|
| 253 |
+
) -> tuple[dict[str, dict[str, dict]], int]:
|
| 254 |
+
"""Use one unselected local main layer as the expected K2 ABI."""
|
| 255 |
+
selected = set(selected_layers)
|
| 256 |
+
available = sorted({layer for name in source if (layer := source_layer(name)) is not None})
|
| 257 |
+
for layer in available:
|
| 258 |
+
if layer in selected:
|
| 259 |
+
continue
|
| 260 |
+
abi: dict[str, dict[str, dict]] = {}
|
| 261 |
+
complete = True
|
| 262 |
+
for projection in ("w1", "w2", "w3"):
|
| 263 |
+
abi[projection] = {}
|
| 264 |
+
for suffix in EXL3_SUFFIXES:
|
| 265 |
+
name = f"layers.{layer}.ffn.experts.0.{projection}.{suffix}"
|
| 266 |
+
if name not in source:
|
| 267 |
+
complete = False
|
| 268 |
+
break
|
| 269 |
+
abi[projection][suffix] = _abi_meta(source[name][2][name])
|
| 270 |
+
if not complete:
|
| 271 |
+
break
|
| 272 |
+
if complete:
|
| 273 |
+
try:
|
| 274 |
+
validate_abi(abi, f"source layer {layer} K2 exemplar")
|
| 275 |
+
except RuntimeError:
|
| 276 |
+
continue
|
| 277 |
+
log(f"MAIN_K2_ABI exemplar_layer={layer} validated=true")
|
| 278 |
+
return abi, layer
|
| 279 |
+
raise RuntimeError("could not find an unselected K2 main-layer ABI exemplar in the source pack")
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
def expected_names(layers: Iterable[int], *, donor: bool) -> set[str]:
|
| 283 |
+
names = set()
|
| 284 |
+
for layer in layers:
|
| 285 |
+
for expert in range(256):
|
| 286 |
+
for projection in ("gate_proj", "up_proj", "down_proj") if donor else ("w1", "w2", "w3"):
|
| 287 |
+
for suffix in EXL3_SUFFIXES:
|
| 288 |
+
if donor:
|
| 289 |
+
names.add(f"model.layers.{layer}.mlp.experts.{expert}.{projection}.{suffix}")
|
| 290 |
+
else:
|
| 291 |
+
names.add(f"layers.{layer}.ffn.experts.{expert}.{projection}.{suffix}")
|
| 292 |
+
return names
|
| 293 |
+
|
| 294 |
+
|
| 295 |
+
def validate_source_selection(
|
| 296 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 297 |
+
selected_layers: tuple[int, ...],
|
| 298 |
+
k2_abi: dict[str, dict[str, dict]],
|
| 299 |
+
) -> tuple[set[str], dict[str, list[tuple[str, dict]]], int]:
|
| 300 |
+
"""Validate selected source layers are complete K3 sets and plan drops."""
|
| 301 |
+
selected = set(selected_layers)
|
| 302 |
+
expected = expected_names(selected_layers, donor=False)
|
| 303 |
+
actual = {name for name in source if source_layer(name) in selected}
|
| 304 |
+
if actual != expected:
|
| 305 |
+
raise RuntimeError(
|
| 306 |
+
f"source selected expert set mismatch: actual={len(actual)} "
|
| 307 |
+
f"expected={len(expected)} missing={sorted(expected - actual)[:5]} "
|
| 308 |
+
f"extra={sorted(actual - expected)[:5]}"
|
| 309 |
+
)
|
| 310 |
+
for name in source:
|
| 311 |
+
raw_layer = re.match(r"^layers\.(?P<layer>[0-9]+)\.ffn\.experts\.", name)
|
| 312 |
+
if (
|
| 313 |
+
raw_layer is not None
|
| 314 |
+
and int(raw_layer.group("layer")) in selected
|
| 315 |
+
and not is_source_expert_name(name)
|
| 316 |
+
):
|
| 317 |
+
raise RuntimeError(f"malformed selected source expert name: {name}")
|
| 318 |
+
|
| 319 |
+
dropped: set[str] = set()
|
| 320 |
+
by_file: dict[str, list[tuple[str, dict]]] = {}
|
| 321 |
+
for name in sorted(actual):
|
| 322 |
+
path, _, header = source[name]
|
| 323 |
+
meta = header[name]
|
| 324 |
+
match = SOURCE_EXPERT_RE.fullmatch(name)
|
| 325 |
+
assert match is not None
|
| 326 |
+
expert = int(match.group("expert"))
|
| 327 |
+
projection = match.group("projection")
|
| 328 |
+
suffix = match.group("suffix")
|
| 329 |
+
if expert >= 256:
|
| 330 |
+
raise RuntimeError(f"selected source expert id out of range: {name}")
|
| 331 |
+
if meta.get("dtype") != EXL3_DTYPE[suffix]:
|
| 332 |
+
raise RuntimeError(f"selected source dtype mismatch for {name}: {meta}")
|
| 333 |
+
shape = meta.get("shape")
|
| 334 |
+
expected_shape = list(k2_abi[projection][suffix]["shape"])
|
| 335 |
+
if suffix == "trellis":
|
| 336 |
+
if not isinstance(shape, list) or shape[:-1] != expected_shape[:-1] or shape[-1:] != [48]:
|
| 337 |
+
raise RuntimeError(f"selected source is not K3 for {name}: {meta}")
|
| 338 |
+
elif shape != expected_shape:
|
| 339 |
+
raise RuntimeError(f"selected source ABI mismatch for {name}: {meta}, expected {k2_abi[projection][suffix]}")
|
| 340 |
+
dropped.add(name)
|
| 341 |
+
by_file.setdefault(path.name, []).append((name, meta))
|
| 342 |
+
|
| 343 |
+
expected_per_layer = 256 * 3 * 4
|
| 344 |
+
for layer in selected_layers:
|
| 345 |
+
count = sum(1 for name in actual if source_layer(name) == layer)
|
| 346 |
+
if count != expected_per_layer:
|
| 347 |
+
raise RuntimeError(f"source layer {layer}: found {count} expert tensors, expected {expected_per_layer}")
|
| 348 |
+
for filename, items in sorted(by_file.items()):
|
| 349 |
+
# A shard may contain one or more selected layers; every complete layer
|
| 350 |
+
# contributes exactly 256 * 3 * 4 expert tensors.
|
| 351 |
+
layers_here = {source_layer(name) for name, _ in items}
|
| 352 |
+
expected_count = expected_per_layer * len(layers_here)
|
| 353 |
+
if len(items) != expected_count:
|
| 354 |
+
raise RuntimeError(f"{filename}: found {len(items)} selected tensors, expected {expected_count}")
|
| 355 |
+
drop_bytes = sum(tensor_nbytes(meta) for items in by_file.values() for _, meta in items)
|
| 356 |
+
log(
|
| 357 |
+
f"SOURCE_DROP tensors={len(dropped)} bytes={drop_bytes} "
|
| 358 |
+
f"files={sorted(by_file)} layers={list(selected_layers)}"
|
| 359 |
+
)
|
| 360 |
+
return dropped, by_file, drop_bytes
|
| 361 |
+
|
| 362 |
+
|
| 363 |
+
def validate_donor_selection(weight_map: dict[str, str], layers: tuple[int, ...]) -> list[str]:
|
| 364 |
+
selected = set(layers)
|
| 365 |
+
candidates = [
|
| 366 |
+
name for name in weight_map
|
| 367 |
+
if name.startswith("model.layers.") and ".mlp.experts." in name
|
| 368 |
+
]
|
| 369 |
+
malformed = sorted(name for name in candidates if not is_donor_expert_name(name))
|
| 370 |
+
if malformed:
|
| 371 |
+
raise RuntimeError(f"malformed donor main expert names (first 5): {malformed[:5]}")
|
| 372 |
+
selected_names = [name for name in candidates if donor_layer(name) in selected]
|
| 373 |
+
expected = expected_names(layers, donor=True)
|
| 374 |
+
actual = set(selected_names)
|
| 375 |
+
if actual != expected:
|
| 376 |
+
raise RuntimeError(
|
| 377 |
+
f"donor selected expert set mismatch: actual={len(actual)} expected={len(expected)} "
|
| 378 |
+
f"missing={sorted(expected - actual)[:5]} extra={sorted(actual - expected)[:5]}"
|
| 379 |
+
)
|
| 380 |
+
selected_names.sort(
|
| 381 |
+
key=lambda name: (
|
| 382 |
+
int(name.split(".")[2]),
|
| 383 |
+
int(name.split(".")[5]),
|
| 384 |
+
PROJECTION_ORDER[name.split(".")[6]],
|
| 385 |
+
EXL3_SUFFIXES.index(name.split(".")[7]),
|
| 386 |
+
)
|
| 387 |
+
)
|
| 388 |
+
log(
|
| 389 |
+
f"DONOR_INDEX selected_tensors={len(selected_names)} layers={list(layers)} "
|
| 390 |
+
f"files={sorted({weight_map[name] for name in selected_names})}"
|
| 391 |
+
)
|
| 392 |
+
return selected_names
|
| 393 |
+
|
| 394 |
+
|
| 395 |
+
def donor_headers_from_remote(
|
| 396 |
+
donor: str,
|
| 397 |
+
weight_map: dict[str, str],
|
| 398 |
+
selected_names: Iterable[str],
|
| 399 |
+
cache_path: Path | None,
|
| 400 |
+
*,
|
| 401 |
+
write_cache: bool = True,
|
| 402 |
+
allow_network: bool = True,
|
| 403 |
+
) -> dict[str, tuple[int, dict]]:
|
| 404 |
+
"""Read one header per required donor shard, locally or with ranges."""
|
| 405 |
+
base = donor_base_url(donor)
|
| 406 |
+
files = sorted({weight_map[name] for name in selected_names})
|
| 407 |
+
cached: dict = {}
|
| 408 |
+
if cache_path is not None and cache_path.exists():
|
| 409 |
+
try:
|
| 410 |
+
cached = json.loads(cache_path.read_text()).get("headers", {})
|
| 411 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 412 |
+
raise RuntimeError(f"cannot read donor header cache {cache_path}: {exc}") from exc
|
| 413 |
+
|
| 414 |
+
result: dict[str, tuple[int, dict]] = {}
|
| 415 |
+
for filename in files:
|
| 416 |
+
item = cached.get(filename)
|
| 417 |
+
if isinstance(item, list) and len(item) == 2:
|
| 418 |
+
header_len, header = int(item[0]), item[1]
|
| 419 |
+
validate_header(header)
|
| 420 |
+
result[filename] = (header_len, header)
|
| 421 |
+
log(f"DONOR_HEADER cache file={filename} tensors={len(header) - 1}")
|
| 422 |
+
continue
|
| 423 |
+
if not allow_network:
|
| 424 |
+
raise RuntimeError(f"donor header cache has no entry for {filename}")
|
| 425 |
+
url = f"{base}/{filename}"
|
| 426 |
+
length_blob = local_or_remote_bytes(url, (0, 7))
|
| 427 |
+
if len(length_blob) != 8:
|
| 428 |
+
raise RuntimeError(f"donor header length read for {filename} was not 8 bytes")
|
| 429 |
+
header_len = struct.unpack("<Q", length_blob)[0]
|
| 430 |
+
header_blob = local_or_remote_bytes(url, (8, 8 + header_len - 1))
|
| 431 |
+
parsed_len, header = parse_header(length_blob + header_blob)
|
| 432 |
+
validate_header(header)
|
| 433 |
+
if parsed_len != header_len:
|
| 434 |
+
raise RuntimeError(f"donor header length changed for {filename}")
|
| 435 |
+
result[filename] = (header_len, header)
|
| 436 |
+
log(f"DONOR_HEADER remote_or_local file={filename} tensors={len(header) - 1}")
|
| 437 |
+
|
| 438 |
+
if cache_path is not None and write_cache:
|
| 439 |
+
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
| 440 |
+
cache_path.write_text(
|
| 441 |
+
json.dumps({"headers": {name: [hlen, header] for name, (hlen, header) in result.items()}}, indent=2)
|
| 442 |
+
)
|
| 443 |
+
log(f"DONOR_HEADER_CACHE wrote={cache_path}")
|
| 444 |
+
return result
|
| 445 |
+
|
| 446 |
+
|
| 447 |
+
def expected_output_meta(
|
| 448 |
+
donor_name: str,
|
| 449 |
+
k2_abi: dict[str, dict[str, dict]],
|
| 450 |
+
) -> dict:
|
| 451 |
+
match = DONOR_EXPERT_RE.fullmatch(donor_name)
|
| 452 |
+
assert match is not None
|
| 453 |
+
projection = PROJECTION_MAP[match.group("projection")]
|
| 454 |
+
suffix = match.group("suffix")
|
| 455 |
+
if suffix == "mcg":
|
| 456 |
+
# Donor markers are scalar [], while the destination parameter is [1].
|
| 457 |
+
return {"dtype": "I32", "shape": [1]}
|
| 458 |
+
return copy.deepcopy(k2_abi[projection][suffix])
|
| 459 |
+
|
| 460 |
+
|
| 461 |
+
def validate_and_plan_donor(
|
| 462 |
+
selected: list[str],
|
| 463 |
+
weight_map: dict[str, str],
|
| 464 |
+
donor_headers: dict[str, tuple[int, dict]] | None,
|
| 465 |
+
k2_abi: dict[str, dict[str, dict]],
|
| 466 |
+
) -> list[dict]:
|
| 467 |
+
plan: list[dict] = []
|
| 468 |
+
seen_targets: set[str] = set()
|
| 469 |
+
for donor_name in selected:
|
| 470 |
+
target_name = map_donor_name(donor_name)
|
| 471 |
+
if target_name in seen_targets:
|
| 472 |
+
raise RuntimeError(f"donor name map collision at {target_name}")
|
| 473 |
+
seen_targets.add(target_name)
|
| 474 |
+
expected = expected_output_meta(donor_name, k2_abi)
|
| 475 |
+
source_file = weight_map[donor_name]
|
| 476 |
+
actual = copy.deepcopy(expected)
|
| 477 |
+
donor_meta = None
|
| 478 |
+
if donor_headers is not None:
|
| 479 |
+
if source_file not in donor_headers:
|
| 480 |
+
raise RuntimeError(f"missing donor header for {source_file}")
|
| 481 |
+
_, header = donor_headers[source_file]
|
| 482 |
+
if donor_name not in header:
|
| 483 |
+
raise RuntimeError(f"donor header lacks indexed tensor {donor_name}")
|
| 484 |
+
donor_meta = copy.deepcopy(header[donor_name])
|
| 485 |
+
validate_header({donor_name: donor_meta})
|
| 486 |
+
actual = _abi_meta(donor_meta)
|
| 487 |
+
if donor_name.endswith(".mcg") and actual.get("shape") == []:
|
| 488 |
+
actual["shape"] = [1]
|
| 489 |
+
if actual != expected:
|
| 490 |
+
raise RuntimeError(
|
| 491 |
+
f"donor shape/dtype mismatch for {donor_name}: got {actual}, expected {expected}"
|
| 492 |
+
)
|
| 493 |
+
plan.append(
|
| 494 |
+
{
|
| 495 |
+
"donor": donor_name,
|
| 496 |
+
"target": target_name,
|
| 497 |
+
"file": source_file,
|
| 498 |
+
"meta": actual,
|
| 499 |
+
"expected": expected,
|
| 500 |
+
"donor_meta": donor_meta,
|
| 501 |
+
}
|
| 502 |
+
)
|
| 503 |
+
expected_count = len(selected)
|
| 504 |
+
if len(plan) != expected_count:
|
| 505 |
+
raise RuntimeError(f"planned donor tensor count is {len(plan)}, expected {expected_count}")
|
| 506 |
+
return plan
|
| 507 |
+
|
| 508 |
+
|
| 509 |
+
def tensor_range(header_len: int, meta: dict) -> tuple[int, int]:
|
| 510 |
+
start, end = meta["data_offsets"]
|
| 511 |
+
return 8 + header_len + start, 8 + header_len + end - 1
|
| 512 |
+
|
| 513 |
+
|
| 514 |
+
def build_new_header(plan: list[dict], donor: str, layers: tuple[int, ...]) -> tuple[list[tuple[str, dict]], int, dict]:
|
| 515 |
+
entries = [(item["target"], item["meta"]) for item in plan]
|
| 516 |
+
offset = 0
|
| 517 |
+
header: dict = {
|
| 518 |
+
"__metadata__": {
|
| 519 |
+
"format": "pt",
|
| 520 |
+
"source": donor,
|
| 521 |
+
"overlay": "main-layer-k2",
|
| 522 |
+
"layers": ",".join(str(layer) for layer in layers),
|
| 523 |
+
}
|
| 524 |
+
}
|
| 525 |
+
for name, meta in entries:
|
| 526 |
+
size = tensor_nbytes(meta)
|
| 527 |
+
header[name] = {
|
| 528 |
+
"dtype": meta["dtype"],
|
| 529 |
+
"shape": list(meta["shape"]),
|
| 530 |
+
"data_offsets": [offset, offset + size],
|
| 531 |
+
}
|
| 532 |
+
offset += size
|
| 533 |
+
return entries, offset, header
|
| 534 |
+
|
| 535 |
+
|
| 536 |
+
def header_without_offsets(header: dict) -> dict:
|
| 537 |
+
return {
|
| 538 |
+
name: {"dtype": meta["dtype"], "shape": list(meta["shape"])}
|
| 539 |
+
for name, meta in header.items()
|
| 540 |
+
if name != "__metadata__"
|
| 541 |
+
}
|
| 542 |
+
|
| 543 |
+
|
| 544 |
+
def write_new_shard(
|
| 545 |
+
out: Path,
|
| 546 |
+
plan: list[dict],
|
| 547 |
+
donor: str,
|
| 548 |
+
layers: tuple[int, ...],
|
| 549 |
+
donor_headers: dict[str, tuple[int, dict]],
|
| 550 |
+
) -> int:
|
| 551 |
+
entries, payload_bytes, expected_header = build_new_header(plan, donor, layers)
|
| 552 |
+
destination = out / NEW_SHARD
|
| 553 |
+
partial = out / (NEW_SHARD + ".partial")
|
| 554 |
+
header_len = None
|
| 555 |
+
actual_header = None
|
| 556 |
+
resume_at = 0
|
| 557 |
+
|
| 558 |
+
if destination.exists():
|
| 559 |
+
existing_len, existing_header = read_header(destination)
|
| 560 |
+
if header_without_offsets(existing_header) != header_without_offsets(expected_header):
|
| 561 |
+
raise RuntimeError(f"existing {destination} has a different tensor plan")
|
| 562 |
+
if destination.stat().st_size != 8 + existing_len + payload_bytes:
|
| 563 |
+
raise RuntimeError(f"existing {destination} is incomplete or has trailing bytes")
|
| 564 |
+
log(f"NEW_SHARD existing_complete={destination}")
|
| 565 |
+
return payload_bytes
|
| 566 |
+
|
| 567 |
+
if partial.exists():
|
| 568 |
+
existing_len, existing_header = read_header(partial, check_data=False)
|
| 569 |
+
if header_without_offsets(existing_header) == header_without_offsets(expected_header):
|
| 570 |
+
data_start = 8 + existing_len
|
| 571 |
+
data_bytes = partial.stat().st_size - data_start
|
| 572 |
+
boundaries = {0} | {
|
| 573 |
+
meta["data_offsets"][1]
|
| 574 |
+
for name, meta in existing_header.items()
|
| 575 |
+
if name != "__metadata__"
|
| 576 |
+
}
|
| 577 |
+
if data_bytes in boundaries and 0 <= data_bytes <= payload_bytes:
|
| 578 |
+
header_len, actual_header, resume_at = existing_len, existing_header, data_bytes
|
| 579 |
+
log(f"NEW_SHARD resume={partial} completed_bytes={resume_at}")
|
| 580 |
+
else:
|
| 581 |
+
log(f"NEW_SHARD discard_invalid_partial={partial}")
|
| 582 |
+
partial.unlink()
|
| 583 |
+
else:
|
| 584 |
+
log(f"NEW_SHARD discard_stale_partial={partial}")
|
| 585 |
+
partial.unlink()
|
| 586 |
+
|
| 587 |
+
if header_len is None:
|
| 588 |
+
with partial.open("wb") as stream:
|
| 589 |
+
header_len, actual_header = common.write_safetensors_header(
|
| 590 |
+
stream,
|
| 591 |
+
entries,
|
| 592 |
+
{"format": "pt", "source": donor, "overlay": "main-layer-k2", "layers": ",".join(map(str, layers))},
|
| 593 |
+
)
|
| 594 |
+
resume_at = 0
|
| 595 |
+
|
| 596 |
+
assert header_len is not None and actual_header is not None
|
| 597 |
+
with partial.open("r+b") as stream:
|
| 598 |
+
data_start = 8 + header_len
|
| 599 |
+
stream.truncate(data_start + resume_at)
|
| 600 |
+
stream.seek(data_start + resume_at)
|
| 601 |
+
for index, item in enumerate(plan):
|
| 602 |
+
name = item["target"]
|
| 603 |
+
meta = actual_header[name]
|
| 604 |
+
start, end = meta["data_offsets"]
|
| 605 |
+
if end <= resume_at:
|
| 606 |
+
log(f"FETCH skip_existing tensor={index + 1}/{len(plan)} name={name}")
|
| 607 |
+
continue
|
| 608 |
+
if start != resume_at:
|
| 609 |
+
raise RuntimeError(f"partial new shard boundary mismatch before {name}")
|
| 610 |
+
donor_header_len, donor_header = donor_headers[item["file"]]
|
| 611 |
+
donor_meta = donor_header[item["donor"]]
|
| 612 |
+
absolute = tensor_range(donor_header_len, donor_meta)
|
| 613 |
+
log(f"FETCH tensor={index + 1}/{len(plan)} name={name} bytes={end - start}")
|
| 614 |
+
data = local_or_remote_bytes(f"{donor_base_url(donor)}/{item['file']}", absolute)
|
| 615 |
+
if len(data) != end - start:
|
| 616 |
+
raise RuntimeError(f"{name}: fetched {len(data)} bytes, expected {end - start}")
|
| 617 |
+
stream.write(data)
|
| 618 |
+
resume_at = end
|
| 619 |
+
if resume_at != payload_bytes:
|
| 620 |
+
raise RuntimeError(f"new shard ended at {resume_at} bytes, expected {payload_bytes}")
|
| 621 |
+
os.replace(partial, destination)
|
| 622 |
+
log(f"NEW_SHARD wrote={destination} payload_bytes={payload_bytes}")
|
| 623 |
+
return payload_bytes
|
| 624 |
+
|
| 625 |
+
|
| 626 |
+
def rewrite_source_shard(source_path: Path, destination_path: Path, dropped: set[str]) -> int:
|
| 627 |
+
"""Copy one source shard while streaming every retained tensor payload."""
|
| 628 |
+
header_len, header = read_header(source_path)
|
| 629 |
+
entries = [
|
| 630 |
+
(name, meta)
|
| 631 |
+
for name, meta in header.items()
|
| 632 |
+
if name != "__metadata__" and name not in dropped
|
| 633 |
+
]
|
| 634 |
+
copied = sum(tensor_nbytes(meta) for _, meta in entries)
|
| 635 |
+
partial = destination_path.with_name(destination_path.name + ".partial")
|
| 636 |
+
with source_path.open("rb") as source_stream, partial.open("wb") as destination_stream:
|
| 637 |
+
common.write_safetensors_header(destination_stream, entries, header.get("__metadata__"))
|
| 638 |
+
source_data_start = 8 + header_len
|
| 639 |
+
for name, meta in entries:
|
| 640 |
+
common.copy_payload(source_stream, destination_stream, source_data_start, meta)
|
| 641 |
+
os.replace(partial, destination_path)
|
| 642 |
+
log(
|
| 643 |
+
f"REWRITE file={source_path.name} kept_tensors={len(entries)} copied_bytes={copied} "
|
| 644 |
+
f"dropped_tensors={sum(name in dropped for name in header if name != '__metadata__')}"
|
| 645 |
+
)
|
| 646 |
+
return copied
|
| 647 |
+
|
| 648 |
+
|
| 649 |
+
def safe_output_path(src: Path, out: Path) -> None:
|
| 650 |
+
src_real = src.resolve()
|
| 651 |
+
out_real = out.resolve(strict=False)
|
| 652 |
+
if out_real == src_real or src_real in out_real.parents:
|
| 653 |
+
raise RuntimeError(f"refusing output inside source pack: {out}")
|
| 654 |
+
|
| 655 |
+
|
| 656 |
+
def link_source_files(src: Path, out: Path, rewrite_files: set[str]) -> int:
|
| 657 |
+
excluded = rewrite_files | {CONFIG_NAME, INDEX_NAME, NEW_SHARD}
|
| 658 |
+
linked = 0
|
| 659 |
+
for entry in sorted(src.iterdir(), key=lambda path: path.name):
|
| 660 |
+
if entry.name in excluded or not (entry.is_file() or entry.is_symlink()):
|
| 661 |
+
continue
|
| 662 |
+
destination = out / entry.name
|
| 663 |
+
if os.path.lexists(destination):
|
| 664 |
+
if destination.is_symlink() and destination.resolve() == entry.resolve():
|
| 665 |
+
continue
|
| 666 |
+
raise RuntimeError(f"output path already exists and is not the source link: {destination}")
|
| 667 |
+
# Resolve source symlinks so the overlay remains usable independently
|
| 668 |
+
# of a symlink chain in the source overlay.
|
| 669 |
+
destination.symlink_to(entry.resolve())
|
| 670 |
+
linked += 1
|
| 671 |
+
log(f"LINKED source_files={linked}")
|
| 672 |
+
return linked
|
| 673 |
+
|
| 674 |
+
|
| 675 |
+
def output_index(
|
| 676 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 677 |
+
dropped: set[str],
|
| 678 |
+
new_plan: list[dict],
|
| 679 |
+
) -> dict:
|
| 680 |
+
weight_map: dict[str, str] = {}
|
| 681 |
+
total_size = 0
|
| 682 |
+
for name, (path, _, header) in sorted(source.items()):
|
| 683 |
+
if name in dropped:
|
| 684 |
+
continue
|
| 685 |
+
weight_map[name] = path.name
|
| 686 |
+
total_size += tensor_nbytes(header[name])
|
| 687 |
+
for item in new_plan:
|
| 688 |
+
if item["target"] in weight_map:
|
| 689 |
+
raise RuntimeError(f"new tensor collides with retained source tensor: {item['target']}")
|
| 690 |
+
weight_map[item["target"]] = NEW_SHARD
|
| 691 |
+
total_size += tensor_nbytes(item["meta"])
|
| 692 |
+
return {"metadata": {"total_size": total_size}, "weight_map": dict(sorted(weight_map.items()))}
|
| 693 |
+
|
| 694 |
+
|
| 695 |
+
def edited_config(src: Path, layers: tuple[int, ...]) -> dict:
|
| 696 |
+
config_path = src / CONFIG_NAME
|
| 697 |
+
try:
|
| 698 |
+
config = json.loads(config_path.read_text())
|
| 699 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 700 |
+
raise RuntimeError(f"cannot read {config_path}: {exc}") from exc
|
| 701 |
+
q = config.get("quantization_config")
|
| 702 |
+
if not isinstance(q, dict):
|
| 703 |
+
raise RuntimeError("source config has no quantization_config object")
|
| 704 |
+
if q.get("bits") != 2:
|
| 705 |
+
raise RuntimeError(f"source base quantization bits must be 2, got {q.get('bits')!r}")
|
| 706 |
+
layer_bits = q.get("layer_bits")
|
| 707 |
+
if not isinstance(layer_bits, dict):
|
| 708 |
+
raise RuntimeError("source quantization_config.layer_bits is not an object")
|
| 709 |
+
layer_bits = dict(layer_bits)
|
| 710 |
+
for layer in layers:
|
| 711 |
+
key = str(layer)
|
| 712 |
+
if layer_bits.get(key) != 3:
|
| 713 |
+
raise RuntimeError(f"source config layer_bits[{key!r}] must be 3 before replacement, got {layer_bits.get(key)!r}")
|
| 714 |
+
layer_bits.pop(key)
|
| 715 |
+
q["layer_bits"] = layer_bits
|
| 716 |
+
return config
|
| 717 |
+
|
| 718 |
+
|
| 719 |
+
def write_output_metadata(out: Path, index: dict, config: dict) -> None:
|
| 720 |
+
(out / INDEX_NAME).write_text(json.dumps(index, indent=2) + "\n")
|
| 721 |
+
(out / CONFIG_NAME).write_text(json.dumps(config, indent=2) + "\n")
|
| 722 |
+
log(f"WROTE metadata files={INDEX_NAME},{CONFIG_NAME}")
|
| 723 |
+
|
| 724 |
+
|
| 725 |
+
def compare_payloads(
|
| 726 |
+
source_path: Path,
|
| 727 |
+
output_path: Path,
|
| 728 |
+
source_meta: dict,
|
| 729 |
+
output_meta: dict,
|
| 730 |
+
source_header_len: int,
|
| 731 |
+
output_header_len: int,
|
| 732 |
+
) -> None:
|
| 733 |
+
source_start, source_end = source_meta["data_offsets"]
|
| 734 |
+
output_start, output_end = output_meta["data_offsets"]
|
| 735 |
+
if source_end - source_start != output_end - output_start:
|
| 736 |
+
raise RuntimeError("retained tensor payload length changed")
|
| 737 |
+
with source_path.open("rb") as source_stream, output_path.open("rb") as output_stream:
|
| 738 |
+
source_stream.seek(8 + source_header_len + source_start)
|
| 739 |
+
output_stream.seek(8 + output_header_len + output_start)
|
| 740 |
+
remaining = source_end - source_start
|
| 741 |
+
while remaining:
|
| 742 |
+
size = min(CHUNK_SIZE, remaining)
|
| 743 |
+
source_data = source_stream.read(size)
|
| 744 |
+
output_data = output_stream.read(size)
|
| 745 |
+
if len(source_data) != size or source_data != output_data:
|
| 746 |
+
raise RuntimeError(f"retained tensor payload differs: {source_path.name}")
|
| 747 |
+
remaining -= size
|
| 748 |
+
|
| 749 |
+
|
| 750 |
+
def verify_output(
|
| 751 |
+
src: Path,
|
| 752 |
+
out: Path,
|
| 753 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 754 |
+
dropped: set[str],
|
| 755 |
+
plan: list[dict],
|
| 756 |
+
expected_index: dict,
|
| 757 |
+
expected_config: dict,
|
| 758 |
+
rewrite_files: set[str],
|
| 759 |
+
) -> None:
|
| 760 |
+
index_path = out / INDEX_NAME
|
| 761 |
+
config_path = out / CONFIG_NAME
|
| 762 |
+
if not index_path.is_file() or not config_path.is_file():
|
| 763 |
+
raise RuntimeError("overlay index/config is missing")
|
| 764 |
+
actual_index = json.loads(index_path.read_text())
|
| 765 |
+
if actual_index != expected_index:
|
| 766 |
+
raise RuntimeError("overlay index differs from the planned index")
|
| 767 |
+
|
| 768 |
+
output_headers: dict[str, tuple[Path, int, dict]] = {}
|
| 769 |
+
for filename in sorted(set(actual_index["weight_map"].values())):
|
| 770 |
+
path = out / filename
|
| 771 |
+
if not path.is_file():
|
| 772 |
+
raise RuntimeError(f"indexed shard is missing: {path}")
|
| 773 |
+
header_len, header = read_header(path)
|
| 774 |
+
for name, mapped_file in actual_index["weight_map"].items():
|
| 775 |
+
if mapped_file == filename:
|
| 776 |
+
if name not in header:
|
| 777 |
+
raise RuntimeError(f"indexed tensor is missing from {filename}: {name}")
|
| 778 |
+
output_headers[name] = (path, header_len, header)
|
| 779 |
+
if set(output_headers) != set(actual_index["weight_map"]):
|
| 780 |
+
raise RuntimeError("not every indexed tensor was opened")
|
| 781 |
+
|
| 782 |
+
for entry in src.iterdir():
|
| 783 |
+
if not (entry.is_file() or entry.is_symlink()):
|
| 784 |
+
continue
|
| 785 |
+
if entry.name in rewrite_files or entry.name in (CONFIG_NAME, INDEX_NAME, NEW_SHARD):
|
| 786 |
+
continue
|
| 787 |
+
linked = out / entry.name
|
| 788 |
+
if not linked.is_symlink() or linked.resolve() != entry.resolve():
|
| 789 |
+
raise RuntimeError(f"source file is not linked unchanged: {entry.name}")
|
| 790 |
+
for filename in rewrite_files:
|
| 791 |
+
if (out / filename).is_symlink():
|
| 792 |
+
raise RuntimeError(f"rewritten source shard is still a symlink: {filename}")
|
| 793 |
+
_, rewritten_header = read_header(out / filename)
|
| 794 |
+
leftover = [n for n in rewritten_header if n != "__metadata__" and n in dropped]
|
| 795 |
+
if leftover:
|
| 796 |
+
raise RuntimeError(f"rewritten shard {filename} still carries {len(leftover)} dropped tensors, e.g. {leftover[0]}")
|
| 797 |
+
|
| 798 |
+
for name, (source_path, source_header_len, source_header) in source.items():
|
| 799 |
+
if name in dropped:
|
| 800 |
+
# Replaced in place: the same name must now point at the new K2 shard.
|
| 801 |
+
where = actual_index["weight_map"].get(name)
|
| 802 |
+
if where is not None and where != NEW_SHARD:
|
| 803 |
+
raise RuntimeError(f"dropped source tensor remains indexed in {where}: {name}")
|
| 804 |
+
continue
|
| 805 |
+
if source_path.name not in rewrite_files:
|
| 806 |
+
continue
|
| 807 |
+
output_path, output_header_len, output_header = output_headers[name]
|
| 808 |
+
if output_header[name]["dtype"] != source_header[name]["dtype"] or output_header[name]["shape"] != source_header[name]["shape"]:
|
| 809 |
+
raise RuntimeError(f"retained tensor metadata changed: {name}")
|
| 810 |
+
compare_payloads(
|
| 811 |
+
source_path,
|
| 812 |
+
output_path,
|
| 813 |
+
source_header[name],
|
| 814 |
+
output_header[name],
|
| 815 |
+
source_header_len,
|
| 816 |
+
output_header_len,
|
| 817 |
+
)
|
| 818 |
+
|
| 819 |
+
for item in plan:
|
| 820 |
+
path, header_len, header = output_headers[item["target"]]
|
| 821 |
+
actual = header[item["target"]]
|
| 822 |
+
if actual["dtype"] != item["meta"]["dtype"] or actual["shape"] != item["meta"]["shape"]:
|
| 823 |
+
raise RuntimeError(f"new tensor metadata differs from donor plan: {item['target']}")
|
| 824 |
+
if item["target"].endswith(".mcg"):
|
| 825 |
+
with path.open("rb") as stream:
|
| 826 |
+
start, end = actual["data_offsets"]
|
| 827 |
+
stream.seek(8 + header_len + start)
|
| 828 |
+
payload = stream.read(end - start)
|
| 829 |
+
if len(payload) != 4 or struct.unpack("<i", payload)[0] != -877912083:
|
| 830 |
+
raise RuntimeError(f"MCG marker mismatch for {item['target']}")
|
| 831 |
+
|
| 832 |
+
actual_config = json.loads(config_path.read_text())
|
| 833 |
+
if actual_config != expected_config:
|
| 834 |
+
raise RuntimeError("overlay config differs from the planned config")
|
| 835 |
+
log(f"VERIFY_OK indexed_tensors={len(output_headers)} retained_payloads_byte_identical=true")
|
| 836 |
+
|
| 837 |
+
|
| 838 |
+
def build(args: argparse.Namespace) -> int:
|
| 839 |
+
src = Path(args.src).expanduser().resolve()
|
| 840 |
+
out = Path(args.out).expanduser()
|
| 841 |
+
donor = args.donor
|
| 842 |
+
layers = parse_layers(args.layers)
|
| 843 |
+
if not src.is_dir():
|
| 844 |
+
raise RuntimeError(f"source directory does not exist: {src}")
|
| 845 |
+
safe_output_path(src, out)
|
| 846 |
+
|
| 847 |
+
source, source_map = source_headers(src)
|
| 848 |
+
validate_source_index(src, source_map)
|
| 849 |
+
k2_abi, exemplar_layer = find_source_k2_abi(source, layers)
|
| 850 |
+
dropped, dropped_by_file, dropped_bytes = validate_source_selection(source, layers, k2_abi)
|
| 851 |
+
rewrite_files = set(dropped_by_file)
|
| 852 |
+
expected_config = edited_config(src, layers)
|
| 853 |
+
log(f"CONFIG_PLAN remove_layer_bits={[str(layer) for layer in layers]} exemplar_layer={exemplar_layer}")
|
| 854 |
+
|
| 855 |
+
donor_index_path = Path(args.donor_index).expanduser() if args.donor_index else None
|
| 856 |
+
if donor_index_path is not None and donor_index_path.exists():
|
| 857 |
+
donor_map = donor_index_from_file(donor_index_path)
|
| 858 |
+
log(f"DONOR_INDEX read={donor_index_path}")
|
| 859 |
+
elif args.dry_run:
|
| 860 |
+
raise RuntimeError(f"--dry-run requires the local donor index; not found: {donor_index_path}")
|
| 861 |
+
else:
|
| 862 |
+
index_blob = local_or_remote_bytes(f"{donor_base_url(donor)}/{INDEX_NAME}")
|
| 863 |
+
try:
|
| 864 |
+
donor_map = json.loads(index_blob)["weight_map"]
|
| 865 |
+
except (json.JSONDecodeError, KeyError, TypeError) as exc:
|
| 866 |
+
raise RuntimeError(f"remote donor index is invalid: {exc}") from exc
|
| 867 |
+
log(f"DONOR_INDEX remote tensors={len(donor_map)}")
|
| 868 |
+
selected = validate_donor_selection(donor_map, layers)
|
| 869 |
+
|
| 870 |
+
cache_path = Path(args.donor_headers).expanduser() if args.donor_headers else None
|
| 871 |
+
donor_headers: dict[str, tuple[int, dict]] | None = None
|
| 872 |
+
if args.dry_run:
|
| 873 |
+
if cache_path is not None and cache_path.exists():
|
| 874 |
+
donor_headers = donor_headers_from_remote(
|
| 875 |
+
donor,
|
| 876 |
+
donor_map,
|
| 877 |
+
selected,
|
| 878 |
+
cache_path,
|
| 879 |
+
write_cache=False,
|
| 880 |
+
allow_network=False,
|
| 881 |
+
)
|
| 882 |
+
else:
|
| 883 |
+
log("DONOR_HEADERS offline=not-read; dry-run uses local K2 ABI")
|
| 884 |
+
else:
|
| 885 |
+
donor_headers = donor_headers_from_remote(donor, donor_map, selected, cache_path)
|
| 886 |
+
plan = validate_and_plan_donor(selected, donor_map, donor_headers, k2_abi)
|
| 887 |
+
|
| 888 |
+
fetched_bytes = sum(tensor_nbytes(item["meta"]) for item in plan)
|
| 889 |
+
rewritten_bytes = sum(
|
| 890 |
+
tensor_nbytes(header[name])
|
| 891 |
+
for name, (path, _, header) in source.items()
|
| 892 |
+
if name not in dropped and path.name in rewrite_files
|
| 893 |
+
)
|
| 894 |
+
expected_index = output_index(source, dropped, plan)
|
| 895 |
+
log(
|
| 896 |
+
f"PLAN layers={list(layers)} donor_tensors={len(plan)} bytes_to_fetch={fetched_bytes} "
|
| 897 |
+
f"({fetched_bytes / (1 << 30):.3f} GiB)"
|
| 898 |
+
)
|
| 899 |
+
log(
|
| 900 |
+
f"PLAN rewrite_files={sorted(rewrite_files)} bytes_to_rewrite={rewritten_bytes} "
|
| 901 |
+
f"bytes_removed={dropped_bytes} dropped_tensors={len(dropped)}"
|
| 902 |
+
)
|
| 903 |
+
log(
|
| 904 |
+
f"PLAN output symlinks=source_files_except_rewritten_config_index "
|
| 905 |
+
f"new_shard={NEW_SHARD} index={INDEX_NAME} config=remove_selected_layer_bits"
|
| 906 |
+
)
|
| 907 |
+
if args.dry_run:
|
| 908 |
+
log("DRY_RUN_OK network=false writes=false")
|
| 909 |
+
return 0
|
| 910 |
+
|
| 911 |
+
out.mkdir(parents=True, exist_ok=True)
|
| 912 |
+
link_source_files(src, out, rewrite_files)
|
| 913 |
+
if donor_headers is None:
|
| 914 |
+
raise RuntimeError("donor headers were not loaded for a real build")
|
| 915 |
+
write_new_shard(out, plan, donor, layers, donor_headers)
|
| 916 |
+
for filename in sorted(rewrite_files):
|
| 917 |
+
rewrite_source_shard(src / filename, out / filename, {n for n, _ in dropped_by_file[filename]})
|
| 918 |
+
write_output_metadata(out, expected_index, expected_config)
|
| 919 |
+
log(f"BUILD_OK out={out}")
|
| 920 |
+
if args.verify:
|
| 921 |
+
verify_output(src, out, source, dropped, plan, expected_index, expected_config, rewrite_files)
|
| 922 |
+
return 0
|
| 923 |
+
|
| 924 |
+
|
| 925 |
+
def main() -> int:
|
| 926 |
+
parser = argparse.ArgumentParser(description=__doc__)
|
| 927 |
+
parser.add_argument("--src", required=True, help="source pack directory")
|
| 928 |
+
parser.add_argument("--donor", required=True, help="Hugging Face repo id or base URL")
|
| 929 |
+
parser.add_argument("--out", required=True, help="new overlay directory")
|
| 930 |
+
parser.add_argument("--layers", required=True, help="comma-separated main layer numbers to replace")
|
| 931 |
+
parser.add_argument(
|
| 932 |
+
"--donor-index",
|
| 933 |
+
default=DEFAULT_DONOR_INDEX,
|
| 934 |
+
help=f"offline donor weight_map JSON (default: {DEFAULT_DONOR_INDEX})",
|
| 935 |
+
)
|
| 936 |
+
parser.add_argument(
|
| 937 |
+
"--donor-headers",
|
| 938 |
+
default=DEFAULT_HEADER_CACHE,
|
| 939 |
+
help=f"donor header cache (default: {DEFAULT_HEADER_CACHE})",
|
| 940 |
+
)
|
| 941 |
+
parser.add_argument("--dry-run", action="store_true", help="plan only; no network and no writes")
|
| 942 |
+
parser.add_argument("--verify", action="store_true", help="verify every indexed tensor after building")
|
| 943 |
+
args = parser.parse_args()
|
| 944 |
+
if args.dry_run and args.verify:
|
| 945 |
+
parser.error("--dry-run and --verify cannot be combined")
|
| 946 |
+
try:
|
| 947 |
+
return build(args)
|
| 948 |
+
except KeyboardInterrupt:
|
| 949 |
+
log("ABORT interrupted")
|
| 950 |
+
return 130
|
| 951 |
+
except Exception as exc:
|
| 952 |
+
log(f"ABORT {type(exc).__name__}: {exc}")
|
| 953 |
+
return 1
|
| 954 |
+
|
| 955 |
+
|
| 956 |
+
if __name__ == "__main__":
|
| 957 |
+
sys.exit(main())
|
lna-lab/tools/mtp_overlay.py
ADDED
|
@@ -0,0 +1,1032 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Build an EXL3-routed-expert overlay for the three DSpark MTP layers.
|
| 3 |
+
|
| 4 |
+
The source pack is never changed. The output directory links every source file
|
| 5 |
+
except ``config.json`` and source shards 46--48, rewrites those three shards with
|
| 6 |
+
only their non-draft tensors, and adds ``model-mtp-exl3.safetensors`` containing
|
| 7 |
+
the donor's EXL3 MTP experts. Donor tensor payloads are fetched with HTTP Range
|
| 8 |
+
requests one tensor at a time; whole donor shards are never downloaded.
|
| 9 |
+
|
| 10 |
+
Offline planning is intentional. ``--dry-run`` reads the source headers and the
|
| 11 |
+
local donor weight map but never opens a network connection or writes an output.
|
| 12 |
+
The default local donor index is the one named in orders/T2-mtp-overlay.md.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import argparse
|
| 18 |
+
import copy
|
| 19 |
+
import json
|
| 20 |
+
import os
|
| 21 |
+
import re
|
| 22 |
+
import struct
|
| 23 |
+
import sys
|
| 24 |
+
import time
|
| 25 |
+
import urllib.error
|
| 26 |
+
import urllib.request
|
| 27 |
+
from pathlib import Path
|
| 28 |
+
from typing import BinaryIO, Iterable
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
SOURCE_REWRITE_FILES = tuple(
|
| 32 |
+
f"model-{i:05d}-of-00048.safetensors" for i in (46, 47, 48)
|
| 33 |
+
)
|
| 34 |
+
NEW_SHARD = "model-mtp-exl3.safetensors"
|
| 35 |
+
INDEX_NAME = "model.safetensors.index.json"
|
| 36 |
+
CONFIG_NAME = "config.json"
|
| 37 |
+
DEFAULT_DONOR_INDEX = "/run/media/tonoken3/DATA1/.tmp/d2-index.json"
|
| 38 |
+
DEFAULT_HEADER_CACHE = "/run/media/tonoken3/DATA1/.tmp/mtp-overlay-donor-headers.json"
|
| 39 |
+
CHUNK_SIZE = 8 << 20
|
| 40 |
+
RETRY_COUNT = 5
|
| 41 |
+
|
| 42 |
+
DTYPE_BYTES = {
|
| 43 |
+
"BOOL": 1,
|
| 44 |
+
"U8": 1,
|
| 45 |
+
"I8": 1,
|
| 46 |
+
"F8_E4M3": 1,
|
| 47 |
+
"F8_E4M3FN": 1,
|
| 48 |
+
"F8_E5M2": 1,
|
| 49 |
+
"F8_E8M0": 1,
|
| 50 |
+
"U16": 2,
|
| 51 |
+
"I16": 2,
|
| 52 |
+
"F16": 2,
|
| 53 |
+
"BF16": 2,
|
| 54 |
+
"U32": 4,
|
| 55 |
+
"I32": 4,
|
| 56 |
+
"F32": 4,
|
| 57 |
+
"U64": 8,
|
| 58 |
+
"I64": 8,
|
| 59 |
+
"F64": 8,
|
| 60 |
+
}
|
| 61 |
+
|
| 62 |
+
DONOR_EXPERT_RE = re.compile(
|
| 63 |
+
r"^mtp\.(?P<layer>[0-2])\.mlp\.experts\."
|
| 64 |
+
r"(?P<expert>[0-9]+)\.(?P<projection>gate_proj|up_proj|down_proj)\."
|
| 65 |
+
r"(?P<suffix>trellis|suh|svh|mcg)$"
|
| 66 |
+
)
|
| 67 |
+
SOURCE_EXPERT_RE = re.compile(
|
| 68 |
+
r"^mtp\.(?P<layer>[0-2])\.ffn\.experts\."
|
| 69 |
+
r"(?P<expert>[0-9]+)\.(?P<projection>w[123])\."
|
| 70 |
+
r"(?P<suffix>weight|scale)$"
|
| 71 |
+
)
|
| 72 |
+
SOURCE_SHARD_RE = re.compile(r"^model-(?P<number>[0-9]+)-of-(?P<total>[0-9]+)\.safetensors$")
|
| 73 |
+
|
| 74 |
+
PROJECTION_MAP = {"gate_proj": "w1", "up_proj": "w3", "down_proj": "w2"}
|
| 75 |
+
PROJECTION_ORDER = {"gate_proj": 0, "up_proj": 1, "down_proj": 2}
|
| 76 |
+
EXL3_SUFFIXES = ("trellis", "suh", "svh", "mcg")
|
| 77 |
+
EXL3_DTYPE = {"trellis": "I16", "suh": "F16", "svh": "F16", "mcg": "I32"}
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def log(message: str) -> None:
|
| 81 |
+
print(message, flush=True)
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def tensor_nbytes(meta: dict) -> int:
|
| 85 |
+
"""Return the payload size and reject malformed safetensors metadata."""
|
| 86 |
+
dtype = meta.get("dtype")
|
| 87 |
+
if dtype not in DTYPE_BYTES:
|
| 88 |
+
raise ValueError(f"unsupported safetensors dtype {dtype!r}")
|
| 89 |
+
shape = meta.get("shape")
|
| 90 |
+
if not isinstance(shape, list) or any(not isinstance(x, int) or x < 0 for x in shape):
|
| 91 |
+
raise ValueError(f"invalid safetensors shape {shape!r}")
|
| 92 |
+
size = DTYPE_BYTES[dtype]
|
| 93 |
+
for dim in shape:
|
| 94 |
+
size *= dim
|
| 95 |
+
return size
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def parse_header(blob: bytes) -> tuple[int, dict]:
|
| 99 |
+
"""Parse a safetensors header from a blob containing at least its header."""
|
| 100 |
+
if len(blob) < 8:
|
| 101 |
+
raise ValueError("safetensors header is shorter than the 8-byte length")
|
| 102 |
+
header_len = struct.unpack("<Q", blob[:8])[0]
|
| 103 |
+
end = 8 + header_len
|
| 104 |
+
if end > len(blob):
|
| 105 |
+
raise ValueError(f"truncated safetensors header: need {end}, have {len(blob)}")
|
| 106 |
+
try:
|
| 107 |
+
header = json.loads(blob[8:end].decode("utf-8"))
|
| 108 |
+
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
| 109 |
+
raise ValueError(f"invalid safetensors JSON header: {exc}") from exc
|
| 110 |
+
if not isinstance(header, dict):
|
| 111 |
+
raise ValueError("safetensors header JSON must be an object")
|
| 112 |
+
return header_len, header
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def validate_header(header: dict, data_bytes: int | None = None) -> None:
|
| 116 |
+
"""Validate the structural parts needed by the range copier."""
|
| 117 |
+
for name, meta in header.items():
|
| 118 |
+
if name == "__metadata__":
|
| 119 |
+
if not isinstance(meta, dict):
|
| 120 |
+
raise ValueError("safetensors __metadata__ must be an object")
|
| 121 |
+
continue
|
| 122 |
+
if not isinstance(meta, dict):
|
| 123 |
+
raise ValueError(f"tensor {name!r} metadata is not an object")
|
| 124 |
+
offsets = meta.get("data_offsets")
|
| 125 |
+
if (
|
| 126 |
+
not isinstance(offsets, list)
|
| 127 |
+
or len(offsets) != 2
|
| 128 |
+
or any(not isinstance(x, int) for x in offsets)
|
| 129 |
+
or offsets[0] < 0
|
| 130 |
+
or offsets[1] < offsets[0]
|
| 131 |
+
):
|
| 132 |
+
raise ValueError(f"invalid data_offsets for {name!r}: {offsets!r}")
|
| 133 |
+
expected = tensor_nbytes(meta)
|
| 134 |
+
actual = offsets[1] - offsets[0]
|
| 135 |
+
if expected != actual:
|
| 136 |
+
raise ValueError(
|
| 137 |
+
f"payload size mismatch for {name!r}: metadata={expected}, offsets={actual}"
|
| 138 |
+
)
|
| 139 |
+
if data_bytes is not None and offsets[1] > data_bytes:
|
| 140 |
+
raise ValueError(
|
| 141 |
+
f"data_offsets for {name!r} exceed data region: {offsets[1]} > {data_bytes}"
|
| 142 |
+
)
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def read_header(path: Path, *, check_data: bool = True) -> tuple[int, dict]:
|
| 146 |
+
"""Read and validate only a local file's 8-byte length and JSON header."""
|
| 147 |
+
file_size = path.stat().st_size
|
| 148 |
+
with path.open("rb") as stream:
|
| 149 |
+
prefix = stream.read(8)
|
| 150 |
+
if len(prefix) != 8:
|
| 151 |
+
raise ValueError(f"{path}: missing 8-byte safetensors header length")
|
| 152 |
+
header_len = struct.unpack("<Q", prefix)[0]
|
| 153 |
+
header_blob = prefix + stream.read(header_len)
|
| 154 |
+
if len(header_blob) != 8 + header_len:
|
| 155 |
+
raise ValueError(f"{path}: truncated safetensors JSON header")
|
| 156 |
+
parsed_len, header = parse_header(header_blob)
|
| 157 |
+
if parsed_len != header_len:
|
| 158 |
+
raise AssertionError("header length changed during parsing")
|
| 159 |
+
validate_header(header, file_size - 8 - header_len if check_data else None)
|
| 160 |
+
return header_len, header
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def map_donor_name(name: str) -> str:
|
| 164 |
+
"""Map one donor EXL3 expert name to the target pack's expert ABI."""
|
| 165 |
+
match = DONOR_EXPERT_RE.fullmatch(name)
|
| 166 |
+
if match is None:
|
| 167 |
+
raise ValueError(f"not a donor MTP EXL3 expert tensor: {name}")
|
| 168 |
+
projection = match.group("projection")
|
| 169 |
+
return (
|
| 170 |
+
f"mtp.{match.group('layer')}.ffn.experts.{int(match.group('expert'))}."
|
| 171 |
+
f"{PROJECTION_MAP[projection]}.{match.group('suffix')}"
|
| 172 |
+
)
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
# Short alias used by the offline tests and useful to callers importing this tool.
|
| 176 |
+
map_name = map_donor_name
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def is_donor_expert_name(name: str) -> bool:
|
| 180 |
+
return DONOR_EXPERT_RE.fullmatch(name) is not None
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def tensor_range(header_len: int, meta: dict) -> tuple[int, int]:
|
| 184 |
+
"""Return the inclusive absolute byte range for one tensor."""
|
| 185 |
+
start, end = meta["data_offsets"]
|
| 186 |
+
return 8 + header_len + start, 8 + header_len + end - 1
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
def source_headers(src: Path) -> tuple[dict[str, tuple[Path, int, dict]], dict[str, str]]:
|
| 190 |
+
"""Scan all numbered source shards; the source intentionally has no index."""
|
| 191 |
+
candidates = []
|
| 192 |
+
for path in src.iterdir():
|
| 193 |
+
match = SOURCE_SHARD_RE.fullmatch(path.name)
|
| 194 |
+
if match and path.is_file():
|
| 195 |
+
candidates.append((int(match.group("number")), path))
|
| 196 |
+
if not candidates:
|
| 197 |
+
raise RuntimeError(f"no numbered safetensors shards found under {src}")
|
| 198 |
+
|
| 199 |
+
headers: dict[str, tuple[Path, int, dict]] = {}
|
| 200 |
+
weight_map: dict[str, str] = {}
|
| 201 |
+
for _, path in sorted(candidates):
|
| 202 |
+
header_len, header = read_header(path)
|
| 203 |
+
for name in header:
|
| 204 |
+
if name == "__metadata__":
|
| 205 |
+
continue
|
| 206 |
+
if name in weight_map:
|
| 207 |
+
raise RuntimeError(f"duplicate source tensor {name!r}")
|
| 208 |
+
headers[name] = (path, header_len, header)
|
| 209 |
+
weight_map[name] = path.name
|
| 210 |
+
log(f"SOURCE_HEADERS files={len(candidates)} tensors={len(headers)}")
|
| 211 |
+
return headers, weight_map
|
| 212 |
+
|
| 213 |
+
|
| 214 |
+
def donor_index_from_file(path: Path) -> dict[str, str]:
|
| 215 |
+
try:
|
| 216 |
+
data = json.loads(path.read_text())
|
| 217 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 218 |
+
raise RuntimeError(f"cannot read donor index {path}: {exc}") from exc
|
| 219 |
+
weight_map = data.get("weight_map")
|
| 220 |
+
if not isinstance(weight_map, dict) or not all(
|
| 221 |
+
isinstance(name, str) and isinstance(filename, str)
|
| 222 |
+
for name, filename in weight_map.items()
|
| 223 |
+
):
|
| 224 |
+
raise RuntimeError(f"donor index {path} has no valid weight_map")
|
| 225 |
+
return weight_map
|
| 226 |
+
|
| 227 |
+
|
| 228 |
+
def http_bytes(url: str, byte_range: tuple[int, int] | None = None) -> bytes:
|
| 229 |
+
"""Fetch one complete response, retrying transient/range failures.
|
| 230 |
+
|
| 231 |
+
Lna-Lab: if MTP_OVERLAY_LOCAL_DIR holds a file with the URL's basename (a
|
| 232 |
+
pre-downloaded donor shard), serve the bytes from it instead of the network.
|
| 233 |
+
"""
|
| 234 |
+
local_dir = os.environ.get("MTP_OVERLAY_LOCAL_DIR")
|
| 235 |
+
if local_dir:
|
| 236 |
+
local_path = Path(local_dir) / url.rsplit("/", 1)[-1].split("?")[0]
|
| 237 |
+
if local_path.is_file():
|
| 238 |
+
with local_path.open("rb") as fh:
|
| 239 |
+
if byte_range is None:
|
| 240 |
+
return fh.read()
|
| 241 |
+
fh.seek(byte_range[0])
|
| 242 |
+
data = fh.read(byte_range[1] - byte_range[0] + 1)
|
| 243 |
+
expected = byte_range[1] - byte_range[0] + 1
|
| 244 |
+
if len(data) != expected:
|
| 245 |
+
raise RuntimeError(f"short local range read {local_path}: got {len(data)}, expected {expected}")
|
| 246 |
+
return data
|
| 247 |
+
for attempt in range(RETRY_COUNT):
|
| 248 |
+
try:
|
| 249 |
+
request = urllib.request.Request(url, headers={"User-Agent": "mtp_overlay/1"})
|
| 250 |
+
if byte_range is not None:
|
| 251 |
+
request.add_header("Range", f"bytes={byte_range[0]}-{byte_range[1]}")
|
| 252 |
+
with urllib.request.urlopen(request, timeout=60) as response:
|
| 253 |
+
status = getattr(response, "status", response.getcode())
|
| 254 |
+
if byte_range is not None and status != 206:
|
| 255 |
+
raise RuntimeError(f"expected HTTP 206 for Range request, got {status}")
|
| 256 |
+
data = response.read()
|
| 257 |
+
expected = None if byte_range is None else byte_range[1] - byte_range[0] + 1
|
| 258 |
+
if expected is not None and len(data) != expected:
|
| 259 |
+
raise RuntimeError(f"short range read: got {len(data)}, expected {expected}")
|
| 260 |
+
return data
|
| 261 |
+
except (OSError, TimeoutError, urllib.error.URLError, urllib.error.HTTPError, RuntimeError) as exc:
|
| 262 |
+
if attempt == RETRY_COUNT - 1:
|
| 263 |
+
raise RuntimeError(f"failed fetching {url} range={byte_range}: {exc}") from exc
|
| 264 |
+
delay = 2**attempt
|
| 265 |
+
log(f"RETRY attempt={attempt + 1}/{RETRY_COUNT} range={byte_range}: {exc}")
|
| 266 |
+
time.sleep(delay)
|
| 267 |
+
raise AssertionError("unreachable")
|
| 268 |
+
|
| 269 |
+
|
| 270 |
+
def donor_headers_from_remote(
|
| 271 |
+
donor: str,
|
| 272 |
+
weight_map: dict[str, str],
|
| 273 |
+
selected_names: Iterable[str],
|
| 274 |
+
cache_path: Path | None,
|
| 275 |
+
*,
|
| 276 |
+
write_cache: bool = True,
|
| 277 |
+
allow_network: bool = True,
|
| 278 |
+
) -> dict[str, tuple[int, dict]]:
|
| 279 |
+
"""Read one remote safetensors header per selected donor shard."""
|
| 280 |
+
base = donor.rstrip("/") if donor.startswith("http") else f"https://huggingface.co/{donor}/resolve/main"
|
| 281 |
+
files = sorted({weight_map[name] for name in selected_names})
|
| 282 |
+
cached: dict = {}
|
| 283 |
+
if cache_path is not None and cache_path.exists():
|
| 284 |
+
try:
|
| 285 |
+
cached = json.loads(cache_path.read_text()).get("headers", {})
|
| 286 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 287 |
+
raise RuntimeError(f"cannot read donor header cache {cache_path}: {exc}") from exc
|
| 288 |
+
|
| 289 |
+
result: dict[str, tuple[int, dict]] = {}
|
| 290 |
+
for filename in files:
|
| 291 |
+
item = cached.get(filename)
|
| 292 |
+
if isinstance(item, list) and len(item) == 2:
|
| 293 |
+
header_len, header = int(item[0]), item[1]
|
| 294 |
+
validate_header(header)
|
| 295 |
+
result[filename] = (header_len, header)
|
| 296 |
+
log(f"DONOR_HEADER cache file={filename} tensors={len(header) - 1}")
|
| 297 |
+
continue
|
| 298 |
+
if not allow_network:
|
| 299 |
+
raise RuntimeError(f"donor header cache has no entry for {filename}")
|
| 300 |
+
url = f"{base}/{filename}"
|
| 301 |
+
length_blob = http_bytes(url, (0, 7))
|
| 302 |
+
if len(length_blob) != 8:
|
| 303 |
+
raise RuntimeError(f"donor header length read for {filename} was not 8 bytes")
|
| 304 |
+
header_len = struct.unpack("<Q", length_blob)[0]
|
| 305 |
+
header_blob = http_bytes(url, (8, 8 + header_len - 1))
|
| 306 |
+
parsed_len, header = parse_header(length_blob + header_blob)
|
| 307 |
+
validate_header(header)
|
| 308 |
+
if parsed_len != header_len:
|
| 309 |
+
raise RuntimeError(f"donor header length changed for {filename}")
|
| 310 |
+
result[filename] = (header_len, header)
|
| 311 |
+
log(f"DONOR_HEADER remote file={filename} tensors={len(header) - 1}")
|
| 312 |
+
|
| 313 |
+
if cache_path is not None and write_cache:
|
| 314 |
+
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
| 315 |
+
cache_path.write_text(
|
| 316 |
+
json.dumps(
|
| 317 |
+
{"headers": {name: [hlen, header] for name, (hlen, header) in result.items()}},
|
| 318 |
+
indent=2,
|
| 319 |
+
)
|
| 320 |
+
)
|
| 321 |
+
log(f"DONOR_HEADER_CACHE wrote={cache_path}")
|
| 322 |
+
return result
|
| 323 |
+
|
| 324 |
+
|
| 325 |
+
def expected_donor_names() -> set[str]:
|
| 326 |
+
return {
|
| 327 |
+
f"mtp.{layer}.mlp.experts.{expert}.{projection}.{suffix}"
|
| 328 |
+
for layer in range(3)
|
| 329 |
+
for expert in range(256)
|
| 330 |
+
for projection in PROJECTION_MAP
|
| 331 |
+
for suffix in EXL3_SUFFIXES
|
| 332 |
+
}
|
| 333 |
+
|
| 334 |
+
|
| 335 |
+
def source_draft_shapes(
|
| 336 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 337 |
+
) -> dict[str, tuple[int, int]]:
|
| 338 |
+
"""Validate the measured source FP8 draft weight shapes and return out/in dims."""
|
| 339 |
+
matches = [
|
| 340 |
+
(name, meta[2][name])
|
| 341 |
+
for name, meta in source.items()
|
| 342 |
+
if SOURCE_EXPERT_RE.fullmatch(name)
|
| 343 |
+
]
|
| 344 |
+
expected = {
|
| 345 |
+
f"mtp.{layer}.ffn.experts.{expert}.{projection}.{suffix}"
|
| 346 |
+
for layer in range(3)
|
| 347 |
+
for expert in range(256)
|
| 348 |
+
for projection in ("w1", "w2", "w3")
|
| 349 |
+
for suffix in ("weight", "scale")
|
| 350 |
+
}
|
| 351 |
+
actual = {name for name, _ in matches}
|
| 352 |
+
if actual != expected:
|
| 353 |
+
missing = sorted(expected - actual)[:5]
|
| 354 |
+
extra = sorted(actual - expected)[:5]
|
| 355 |
+
raise RuntimeError(f"source draft expert set mismatch missing={missing} extra={extra}")
|
| 356 |
+
|
| 357 |
+
shapes: dict[str, tuple[int, int]] = {}
|
| 358 |
+
for projection in ("w1", "w2", "w3"):
|
| 359 |
+
name = f"mtp.0.ffn.experts.0.{projection}.weight"
|
| 360 |
+
meta = source[name][2][name]
|
| 361 |
+
if meta.get("dtype") != "I8" or len(meta.get("shape", [])) != 2:
|
| 362 |
+
raise RuntimeError(f"unexpected source draft metadata for {name}: {meta}")
|
| 363 |
+
out_dim, in_dim = meta["shape"]
|
| 364 |
+
if out_dim % 16 or in_dim % 16:
|
| 365 |
+
raise RuntimeError(f"source draft dimensions are not tile-aligned for {name}: {meta['shape']}")
|
| 366 |
+
shapes[projection] = (out_dim, in_dim)
|
| 367 |
+
scale_name = f"mtp.0.ffn.experts.0.{projection}.scale"
|
| 368 |
+
scale = source[scale_name][2][scale_name]
|
| 369 |
+
if scale.get("dtype") != "F8_E8M0":
|
| 370 |
+
raise RuntimeError(f"unexpected source draft scale dtype for {scale_name}: {scale}")
|
| 371 |
+
|
| 372 |
+
for name, (_, _, header) in source.items():
|
| 373 |
+
match = SOURCE_EXPERT_RE.fullmatch(name)
|
| 374 |
+
if match is None:
|
| 375 |
+
continue
|
| 376 |
+
meta = header[name]
|
| 377 |
+
projection = match.group("projection")
|
| 378 |
+
suffix = match.group("suffix")
|
| 379 |
+
if suffix == "weight":
|
| 380 |
+
if meta["dtype"] != "I8" or tuple(meta["shape"]) != shapes[projection]:
|
| 381 |
+
raise RuntimeError(f"inconsistent source draft weight {name}: {meta}")
|
| 382 |
+
elif meta["dtype"] != "F8_E8M0":
|
| 383 |
+
raise RuntimeError(f"inconsistent source draft scale {name}: {meta}")
|
| 384 |
+
|
| 385 |
+
log(
|
| 386 |
+
"SOURCE_DRAFT_SHAPES "
|
| 387 |
+
+ " ".join(
|
| 388 |
+
f"{projection}.weight=I8{list(shapes[projection])}"
|
| 389 |
+
for projection in ("w1", "w2", "w3")
|
| 390 |
+
)
|
| 391 |
+
+ " scale=F8_E8M0"
|
| 392 |
+
)
|
| 393 |
+
return shapes
|
| 394 |
+
|
| 395 |
+
|
| 396 |
+
def main_k2_shapes(
|
| 397 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 398 |
+
) -> dict[str, dict[str, dict]]:
|
| 399 |
+
"""Read the concrete K2 EXL3 ABI from the requested main-layer exemplar."""
|
| 400 |
+
result: dict[str, dict[str, dict]] = {}
|
| 401 |
+
for projection in ("w1", "w2", "w3"):
|
| 402 |
+
result[projection] = {}
|
| 403 |
+
for suffix in EXL3_SUFFIXES:
|
| 404 |
+
name = f"layers.5.ffn.experts.0.{projection}.{suffix}"
|
| 405 |
+
if name not in source:
|
| 406 |
+
raise RuntimeError(f"required main-layer K2 exemplar is missing: {name}")
|
| 407 |
+
result[projection][suffix] = copy.deepcopy(source[name][2][name])
|
| 408 |
+
for projection, parts in result.items():
|
| 409 |
+
trellis = parts["trellis"]
|
| 410 |
+
if trellis.get("dtype") != "I16" or trellis.get("shape", [])[-1:] != [32]:
|
| 411 |
+
raise RuntimeError(f"main-layer exemplar is not K2 EXL3 for {projection}: {trellis}")
|
| 412 |
+
for suffix in EXL3_SUFFIXES:
|
| 413 |
+
if parts[suffix].get("dtype") != EXL3_DTYPE[suffix]:
|
| 414 |
+
raise RuntimeError(f"main-layer exemplar dtype mismatch for {projection}.{suffix}")
|
| 415 |
+
log(
|
| 416 |
+
"MAIN_K2_ABI "
|
| 417 |
+
+ " ".join(
|
| 418 |
+
f"{projection}="
|
| 419 |
+
+ ",".join(
|
| 420 |
+
f"{suffix}:{result[projection][suffix]['dtype']}{result[projection][suffix]['shape']}"
|
| 421 |
+
for suffix in EXL3_SUFFIXES
|
| 422 |
+
)
|
| 423 |
+
for projection in ("w1", "w2", "w3")
|
| 424 |
+
)
|
| 425 |
+
)
|
| 426 |
+
return result
|
| 427 |
+
|
| 428 |
+
|
| 429 |
+
def expected_mtp_meta(
|
| 430 |
+
projection: str,
|
| 431 |
+
suffix: str,
|
| 432 |
+
draft_shapes: dict[str, tuple[int, int]],
|
| 433 |
+
main_shapes: dict[str, dict[str, dict]],
|
| 434 |
+
) -> dict:
|
| 435 |
+
"""Build the offline expected MTP ABI from the concrete main K2 ABI.
|
| 436 |
+
|
| 437 |
+
The source FP8 MTP weights are TP-local (w1/w3 [2048, 2048], w2
|
| 438 |
+
[4096, 1024]), while the donor EXL3 tensors are full-width and match the
|
| 439 |
+
main-layer K2 tensors. ``draft_shapes`` is still passed here so the
|
| 440 |
+
source dimensions remain part of the validated plan and cannot silently
|
| 441 |
+
disappear from the offline check. EXL3 donor MCG markers are scalar
|
| 442 |
+
tensors (shape []) although the vLLM destination parameter is [1].
|
| 443 |
+
"""
|
| 444 |
+
del draft_shapes
|
| 445 |
+
if suffix == "mcg":
|
| 446 |
+
# Lna-Lab: donor markers are scalar; the vLLM destination parameter is [1].
|
| 447 |
+
return {"dtype": EXL3_DTYPE[suffix], "shape": [1]}
|
| 448 |
+
reference = main_shapes[projection][suffix]
|
| 449 |
+
return {"dtype": reference["dtype"], "shape": list(reference["shape"])}
|
| 450 |
+
|
| 451 |
+
|
| 452 |
+
def validate_donor_selection(
|
| 453 |
+
weight_map: dict[str, str],
|
| 454 |
+
) -> list[str]:
|
| 455 |
+
names = [name for name in weight_map if name.startswith("mtp.") and ".experts." in name]
|
| 456 |
+
malformed = sorted(name for name in names if not is_donor_expert_name(name))
|
| 457 |
+
if malformed:
|
| 458 |
+
raise RuntimeError(f"malformed donor MTP expert names (first 5): {malformed[:5]}")
|
| 459 |
+
expected = expected_donor_names()
|
| 460 |
+
actual = set(names)
|
| 461 |
+
if actual != expected:
|
| 462 |
+
raise RuntimeError(
|
| 463 |
+
f"donor MTP expert set mismatch: selected={len(actual)} "
|
| 464 |
+
f"missing={sorted(expected - actual)[:5]} extra={sorted(actual - expected)[:5]}"
|
| 465 |
+
)
|
| 466 |
+
selected = sorted(
|
| 467 |
+
names,
|
| 468 |
+
key=lambda name: (
|
| 469 |
+
int(name.split(".")[1]),
|
| 470 |
+
int(name.split(".")[4]),
|
| 471 |
+
PROJECTION_ORDER[name.split(".")[5]],
|
| 472 |
+
EXL3_SUFFIXES.index(name.split(".")[6]),
|
| 473 |
+
),
|
| 474 |
+
)
|
| 475 |
+
log(
|
| 476 |
+
f"DONOR_INDEX experts={len(selected)} non_expert_mtp={sum(name.startswith('mtp.') for name in weight_map) - len(selected)} "
|
| 477 |
+
f"files={sorted({weight_map[name] for name in selected})}"
|
| 478 |
+
)
|
| 479 |
+
return selected
|
| 480 |
+
|
| 481 |
+
|
| 482 |
+
def validate_and_plan_donor(
|
| 483 |
+
selected: list[str],
|
| 484 |
+
weight_map: dict[str, str],
|
| 485 |
+
donor_headers: dict[str, tuple[int, dict]] | None,
|
| 486 |
+
draft_shapes: dict[str, tuple[int, int]],
|
| 487 |
+
main_shapes: dict[str, dict[str, dict]],
|
| 488 |
+
) -> list[dict]:
|
| 489 |
+
plan = []
|
| 490 |
+
seen_targets: set[str] = set()
|
| 491 |
+
for donor_name in selected:
|
| 492 |
+
target_name = map_donor_name(donor_name)
|
| 493 |
+
if target_name in seen_targets:
|
| 494 |
+
raise RuntimeError(f"donor name map collision at {target_name}")
|
| 495 |
+
seen_targets.add(target_name)
|
| 496 |
+
projection = PROJECTION_MAP[donor_name.split(".")[5]]
|
| 497 |
+
suffix = donor_name.split(".")[6]
|
| 498 |
+
expected = expected_mtp_meta(projection, suffix, draft_shapes, main_shapes)
|
| 499 |
+
source_file = weight_map[donor_name]
|
| 500 |
+
actual = expected
|
| 501 |
+
if donor_headers is not None:
|
| 502 |
+
if source_file not in donor_headers:
|
| 503 |
+
raise RuntimeError(f"missing donor header for {source_file}")
|
| 504 |
+
_, header = donor_headers[source_file]
|
| 505 |
+
if donor_name not in header:
|
| 506 |
+
raise RuntimeError(f"donor header lacks indexed tensor {donor_name}")
|
| 507 |
+
actual = copy.deepcopy(header[donor_name])
|
| 508 |
+
if donor_name.endswith(".mcg") and actual.get("shape") == []:
|
| 509 |
+
actual["shape"] = [1] # same 4 bytes; vLLM wants a (1,) parameter
|
| 510 |
+
# Dtype and shape must agree with the concrete main K2 EXL3 ABI.
|
| 511 |
+
# The source MTP FP8 tensors are TP-local; the donor EXL3 tensors
|
| 512 |
+
# are full-width, as the local donor shard header confirms.
|
| 513 |
+
if actual.get("dtype") != expected["dtype"] or actual.get("shape") != expected["shape"]:
|
| 514 |
+
raise RuntimeError(
|
| 515 |
+
f"donor shape/dtype mismatch for {donor_name}: got "
|
| 516 |
+
f"{actual.get('dtype')}{actual.get('shape')}, expected "
|
| 517 |
+
f"{expected['dtype']}{expected['shape']} from local MTP dimensions"
|
| 518 |
+
)
|
| 519 |
+
validate_header({donor_name: actual})
|
| 520 |
+
plan.append(
|
| 521 |
+
{
|
| 522 |
+
"donor": donor_name,
|
| 523 |
+
"target": target_name,
|
| 524 |
+
"file": source_file,
|
| 525 |
+
"meta": actual,
|
| 526 |
+
"expected": expected,
|
| 527 |
+
}
|
| 528 |
+
)
|
| 529 |
+
if len(plan) != 9216:
|
| 530 |
+
raise RuntimeError(f"planned donor tensor count is {len(plan)}, expected 9216")
|
| 531 |
+
return plan
|
| 532 |
+
|
| 533 |
+
|
| 534 |
+
def source_drop_plan(
|
| 535 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 536 |
+
) -> tuple[set[str], dict[str, list[tuple[str, dict]]], int, int]:
|
| 537 |
+
dropped: set[str] = set()
|
| 538 |
+
by_file: dict[str, list[tuple[str, dict]]] = {name: [] for name in SOURCE_REWRITE_FILES}
|
| 539 |
+
for name, (path, _, header) in source.items():
|
| 540 |
+
if name.startswith("mtp.") and ".ffn.experts." in name and SOURCE_EXPERT_RE.fullmatch(name) is None:
|
| 541 |
+
raise RuntimeError(f"malformed source MTP expert tensor name: {name}")
|
| 542 |
+
match = SOURCE_EXPERT_RE.fullmatch(name)
|
| 543 |
+
if match is None:
|
| 544 |
+
continue
|
| 545 |
+
if path.name not in by_file:
|
| 546 |
+
raise RuntimeError(f"draft expert tensor is outside source shards 46-48: {name} in {path.name}")
|
| 547 |
+
if int(match.group("expert")) >= 256:
|
| 548 |
+
raise RuntimeError(f"source draft expert id out of range: {name}")
|
| 549 |
+
dropped.add(name)
|
| 550 |
+
by_file[path.name].append((name, header[name]))
|
| 551 |
+
expected_per_file = 256 * 3 * 2
|
| 552 |
+
for filename, items in by_file.items():
|
| 553 |
+
if len(items) != expected_per_file:
|
| 554 |
+
raise RuntimeError(f"{filename}: found {len(items)} draft tensors, expected {expected_per_file}")
|
| 555 |
+
drop_bytes = sum(tensor_nbytes(meta) for items in by_file.values() for _, meta in items)
|
| 556 |
+
log(f"SOURCE_DROP tensors={len(dropped)} bytes={drop_bytes} files={list(SOURCE_REWRITE_FILES)}")
|
| 557 |
+
return dropped, by_file, len(dropped), drop_bytes
|
| 558 |
+
|
| 559 |
+
|
| 560 |
+
def write_safetensors_header(
|
| 561 |
+
stream: BinaryIO,
|
| 562 |
+
entries: list[tuple[str, dict]],
|
| 563 |
+
metadata: dict | None = None,
|
| 564 |
+
) -> tuple[int, dict]:
|
| 565 |
+
header: dict = {}
|
| 566 |
+
if metadata is not None:
|
| 567 |
+
header["__metadata__"] = metadata
|
| 568 |
+
offset = 0
|
| 569 |
+
for name, meta in entries:
|
| 570 |
+
size = tensor_nbytes(meta)
|
| 571 |
+
header[name] = {
|
| 572 |
+
"dtype": meta["dtype"],
|
| 573 |
+
"shape": list(meta["shape"]),
|
| 574 |
+
"data_offsets": [offset, offset + size],
|
| 575 |
+
}
|
| 576 |
+
offset += size
|
| 577 |
+
blob = json.dumps(header, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
| 578 |
+
blob += b" " * (-len(blob) % 8)
|
| 579 |
+
stream.write(struct.pack("<Q", len(blob)))
|
| 580 |
+
stream.write(blob)
|
| 581 |
+
return len(blob), header
|
| 582 |
+
|
| 583 |
+
|
| 584 |
+
def copy_payload(
|
| 585 |
+
source: BinaryIO,
|
| 586 |
+
destination: BinaryIO,
|
| 587 |
+
source_data_start: int,
|
| 588 |
+
meta: dict,
|
| 589 |
+
) -> None:
|
| 590 |
+
start, end = meta["data_offsets"]
|
| 591 |
+
source.seek(source_data_start + start)
|
| 592 |
+
remaining = end - start
|
| 593 |
+
while remaining:
|
| 594 |
+
block = source.read(min(CHUNK_SIZE, remaining))
|
| 595 |
+
if not block:
|
| 596 |
+
raise RuntimeError("short source tensor payload while rewriting shard")
|
| 597 |
+
destination.write(block)
|
| 598 |
+
remaining -= len(block)
|
| 599 |
+
|
| 600 |
+
|
| 601 |
+
def rewrite_source_shard(
|
| 602 |
+
source_path: Path,
|
| 603 |
+
destination_path: Path,
|
| 604 |
+
dropped: set[str],
|
| 605 |
+
) -> int:
|
| 606 |
+
header_len, header = read_header(source_path)
|
| 607 |
+
entries = [(name, meta) for name, meta in header.items() if name != "__metadata__" and name not in dropped]
|
| 608 |
+
copied = sum(tensor_nbytes(meta) for _, meta in entries)
|
| 609 |
+
partial = destination_path.with_name(destination_path.name + ".partial")
|
| 610 |
+
with source_path.open("rb") as source_stream, partial.open("wb") as destination_stream:
|
| 611 |
+
_, new_header = write_safetensors_header(
|
| 612 |
+
destination_stream,
|
| 613 |
+
entries,
|
| 614 |
+
header.get("__metadata__"),
|
| 615 |
+
)
|
| 616 |
+
del new_header
|
| 617 |
+
source_data_start = 8 + header_len
|
| 618 |
+
for name, meta in entries:
|
| 619 |
+
copy_payload(source_stream, destination_stream, source_data_start, meta)
|
| 620 |
+
os.replace(partial, destination_path)
|
| 621 |
+
log(
|
| 622 |
+
f"REWRITE file={source_path.name} kept_tensors={len(entries)} "
|
| 623 |
+
f"copied_bytes={copied} dropped_tensors={sum(name in dropped for name, _ in header.items())}"
|
| 624 |
+
)
|
| 625 |
+
return copied
|
| 626 |
+
|
| 627 |
+
|
| 628 |
+
def safe_output_path(src: Path, out: Path) -> None:
|
| 629 |
+
src_real = src.resolve()
|
| 630 |
+
out_real = out.resolve(strict=False)
|
| 631 |
+
if out_real == src_real or src_real in out_real.parents:
|
| 632 |
+
raise RuntimeError(f"refusing output inside source pack: {out}")
|
| 633 |
+
|
| 634 |
+
|
| 635 |
+
def link_source_files(src: Path, out: Path) -> int:
|
| 636 |
+
excluded = set(SOURCE_REWRITE_FILES) | {CONFIG_NAME, INDEX_NAME}
|
| 637 |
+
linked = 0
|
| 638 |
+
for entry in sorted(src.iterdir(), key=lambda p: p.name):
|
| 639 |
+
if entry.name in excluded or not (entry.is_file() or entry.is_symlink()):
|
| 640 |
+
continue
|
| 641 |
+
destination = out / entry.name
|
| 642 |
+
if os.path.lexists(destination):
|
| 643 |
+
if destination.is_symlink() and destination.resolve() == entry.resolve():
|
| 644 |
+
continue
|
| 645 |
+
raise RuntimeError(f"output path already exists and is not the source link: {destination}")
|
| 646 |
+
destination.symlink_to(entry.resolve())
|
| 647 |
+
linked += 1
|
| 648 |
+
log(f"LINKED source_files={linked}")
|
| 649 |
+
return linked
|
| 650 |
+
|
| 651 |
+
|
| 652 |
+
def build_new_header(plan: list[dict], donor: str) -> tuple[list[tuple[str, dict]], int, dict]:
|
| 653 |
+
entries = [(item["target"], item["meta"]) for item in plan]
|
| 654 |
+
offset = 0
|
| 655 |
+
header: dict = {"__metadata__": {"format": "pt", "source": donor, "overlay": "mtp-exl3"}}
|
| 656 |
+
for name, meta in entries:
|
| 657 |
+
size = tensor_nbytes(meta)
|
| 658 |
+
header[name] = {
|
| 659 |
+
"dtype": meta["dtype"],
|
| 660 |
+
"shape": list(meta["shape"]),
|
| 661 |
+
"data_offsets": [offset, offset + size],
|
| 662 |
+
}
|
| 663 |
+
offset += size
|
| 664 |
+
return entries, offset, header
|
| 665 |
+
|
| 666 |
+
|
| 667 |
+
def header_without_offsets(header: dict) -> dict:
|
| 668 |
+
return {
|
| 669 |
+
name: {"dtype": meta["dtype"], "shape": list(meta["shape"])}
|
| 670 |
+
for name, meta in header.items()
|
| 671 |
+
if name != "__metadata__"
|
| 672 |
+
}
|
| 673 |
+
|
| 674 |
+
|
| 675 |
+
def write_new_shard(
|
| 676 |
+
out: Path,
|
| 677 |
+
plan: list[dict],
|
| 678 |
+
donor: str,
|
| 679 |
+
donor_headers: dict[str, tuple[int, dict]],
|
| 680 |
+
) -> int:
|
| 681 |
+
entries, payload_bytes, expected_header = build_new_header(plan, donor)
|
| 682 |
+
destination = out / NEW_SHARD
|
| 683 |
+
partial = out / (NEW_SHARD + ".partial")
|
| 684 |
+
header_len = None
|
| 685 |
+
actual_header = None
|
| 686 |
+
resume_at = 0
|
| 687 |
+
|
| 688 |
+
if destination.exists():
|
| 689 |
+
existing_len, existing_header = read_header(destination)
|
| 690 |
+
if header_without_offsets(existing_header) != header_without_offsets(expected_header):
|
| 691 |
+
raise RuntimeError(f"existing {destination} has a different tensor plan")
|
| 692 |
+
if destination.stat().st_size != 8 + existing_len + payload_bytes:
|
| 693 |
+
raise RuntimeError(f"existing {destination} is incomplete or has trailing bytes")
|
| 694 |
+
log(f"NEW_SHARD existing_complete={destination}")
|
| 695 |
+
return payload_bytes
|
| 696 |
+
|
| 697 |
+
if partial.exists():
|
| 698 |
+
# A process can stop in the middle of the current tensor. Read and
|
| 699 |
+
# validate the JSON header without requiring the partial data region to
|
| 700 |
+
# reach its final offset, then resume only at a completed boundary.
|
| 701 |
+
existing_len, existing_header = read_header(partial, check_data=False)
|
| 702 |
+
if header_without_offsets(existing_header) == header_without_offsets(expected_header):
|
| 703 |
+
data_start = 8 + existing_len
|
| 704 |
+
data_bytes = partial.stat().st_size - data_start
|
| 705 |
+
boundaries = {0} | {meta["data_offsets"][1] for name, meta in existing_header.items() if name != "__metadata__"}
|
| 706 |
+
if data_bytes in boundaries and 0 <= data_bytes <= payload_bytes:
|
| 707 |
+
header_len, actual_header, resume_at = existing_len, existing_header, data_bytes
|
| 708 |
+
log(f"NEW_SHARD resume={partial} completed_bytes={resume_at}")
|
| 709 |
+
else:
|
| 710 |
+
log(f"NEW_SHARD discard_invalid_partial={partial}")
|
| 711 |
+
partial.unlink()
|
| 712 |
+
else:
|
| 713 |
+
log(f"NEW_SHARD discard_stale_partial={partial}")
|
| 714 |
+
partial.unlink()
|
| 715 |
+
|
| 716 |
+
if header_len is None:
|
| 717 |
+
with partial.open("wb") as stream:
|
| 718 |
+
header_len, actual_header = write_safetensors_header(
|
| 719 |
+
stream, entries, {"format": "pt", "source": donor, "overlay": "mtp-exl3"}
|
| 720 |
+
)
|
| 721 |
+
resume_at = 0
|
| 722 |
+
|
| 723 |
+
assert header_len is not None and actual_header is not None
|
| 724 |
+
with partial.open("r+b") as stream:
|
| 725 |
+
data_start = 8 + header_len
|
| 726 |
+
stream.truncate(data_start + resume_at)
|
| 727 |
+
stream.seek(data_start + resume_at)
|
| 728 |
+
for index, item in enumerate(plan):
|
| 729 |
+
name = item["target"]
|
| 730 |
+
meta = actual_header[name]
|
| 731 |
+
start, end = meta["data_offsets"]
|
| 732 |
+
if end <= resume_at:
|
| 733 |
+
log(f"FETCH skip_existing tensor={index + 1}/{len(plan)} name={name}")
|
| 734 |
+
continue
|
| 735 |
+
if start != resume_at:
|
| 736 |
+
raise RuntimeError(f"partial new shard boundary mismatch before {name}")
|
| 737 |
+
remote_header_len, remote_header = donor_headers[item["file"]]
|
| 738 |
+
remote_meta = remote_header[item["donor"]]
|
| 739 |
+
absolute = tensor_range(remote_header_len, remote_meta)
|
| 740 |
+
base = donor.rstrip("/") if donor.startswith("http") else f"https://huggingface.co/{donor}/resolve/main"
|
| 741 |
+
log(f"FETCH tensor={index + 1}/{len(plan)} name={name} bytes={end - start}")
|
| 742 |
+
data = http_bytes(f"{base}/{item['file']}", absolute)
|
| 743 |
+
if len(data) != end - start:
|
| 744 |
+
raise RuntimeError(f"{name}: fetched {len(data)} bytes, expected {end - start}")
|
| 745 |
+
stream.write(data)
|
| 746 |
+
resume_at = end
|
| 747 |
+
if resume_at != payload_bytes:
|
| 748 |
+
raise RuntimeError(f"new shard ended at {resume_at} bytes, expected {payload_bytes}")
|
| 749 |
+
os.replace(partial, destination)
|
| 750 |
+
log(f"NEW_SHARD wrote={destination} payload_bytes={payload_bytes}")
|
| 751 |
+
return payload_bytes
|
| 752 |
+
|
| 753 |
+
|
| 754 |
+
def output_index(
|
| 755 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 756 |
+
dropped: set[str],
|
| 757 |
+
new_plan: list[dict],
|
| 758 |
+
) -> dict:
|
| 759 |
+
weight_map: dict[str, str] = {}
|
| 760 |
+
total_size = 0
|
| 761 |
+
for name, (path, _, header) in sorted(source.items()):
|
| 762 |
+
if name in dropped:
|
| 763 |
+
continue
|
| 764 |
+
output_file = path.name if path.name not in SOURCE_REWRITE_FILES else path.name
|
| 765 |
+
weight_map[name] = output_file
|
| 766 |
+
total_size += tensor_nbytes(header[name])
|
| 767 |
+
for item in new_plan:
|
| 768 |
+
if item["target"] in weight_map:
|
| 769 |
+
raise RuntimeError(f"new tensor collides with retained source tensor: {item['target']}")
|
| 770 |
+
weight_map[item["target"]] = NEW_SHARD
|
| 771 |
+
total_size += tensor_nbytes(item["meta"])
|
| 772 |
+
return {"metadata": {"total_size": total_size}, "weight_map": dict(sorted(weight_map.items()))}
|
| 773 |
+
|
| 774 |
+
|
| 775 |
+
def edited_config(src: Path) -> dict:
|
| 776 |
+
config_path = src / CONFIG_NAME
|
| 777 |
+
try:
|
| 778 |
+
config = json.loads(config_path.read_text())
|
| 779 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 780 |
+
raise RuntimeError(f"cannot read {config_path}: {exc}") from exc
|
| 781 |
+
q = config.get("quantization_config")
|
| 782 |
+
if not isinstance(q, dict):
|
| 783 |
+
raise RuntimeError("source config has no quantization_config object")
|
| 784 |
+
q.pop("mtp_experts", None)
|
| 785 |
+
q.pop("mtp_experts_start_layer", None)
|
| 786 |
+
layer_bits = q.get("layer_bits")
|
| 787 |
+
if not isinstance(layer_bits, dict):
|
| 788 |
+
raise RuntimeError("source quantization_config.layer_bits is not an object")
|
| 789 |
+
layer_bits = dict(layer_bits)
|
| 790 |
+
layer_bits.update({"43": 2, "44": 2, "45": 2})
|
| 791 |
+
q["layer_bits"] = layer_bits
|
| 792 |
+
return config
|
| 793 |
+
|
| 794 |
+
|
| 795 |
+
def write_output_metadata(out: Path, index: dict, config: dict) -> None:
|
| 796 |
+
(out / INDEX_NAME).write_text(json.dumps(index, indent=2) + "\n")
|
| 797 |
+
(out / CONFIG_NAME).write_text(json.dumps(config, indent=2) + "\n")
|
| 798 |
+
log(f"WROTE metadata files={INDEX_NAME},{CONFIG_NAME}")
|
| 799 |
+
|
| 800 |
+
|
| 801 |
+
def compare_payloads(source_path: Path, output_path: Path, source_meta: dict, output_meta: dict, source_header_len: int, output_header_len: int) -> None:
|
| 802 |
+
source_start, source_end = source_meta["data_offsets"]
|
| 803 |
+
output_start, output_end = output_meta["data_offsets"]
|
| 804 |
+
if source_end - source_start != output_end - output_start:
|
| 805 |
+
raise RuntimeError("retained tensor payload length changed")
|
| 806 |
+
with source_path.open("rb") as source_stream, output_path.open("rb") as output_stream:
|
| 807 |
+
source_stream.seek(8 + source_header_len + source_start)
|
| 808 |
+
output_stream.seek(8 + output_header_len + output_start)
|
| 809 |
+
remaining = source_end - source_start
|
| 810 |
+
while remaining:
|
| 811 |
+
size = min(CHUNK_SIZE, remaining)
|
| 812 |
+
a = source_stream.read(size)
|
| 813 |
+
b = output_stream.read(size)
|
| 814 |
+
if a != b:
|
| 815 |
+
raise RuntimeError(f"retained tensor payload differs: {source_path.name}")
|
| 816 |
+
remaining -= len(a)
|
| 817 |
+
|
| 818 |
+
|
| 819 |
+
def verify_output(
|
| 820 |
+
src: Path,
|
| 821 |
+
out: Path,
|
| 822 |
+
source: dict[str, tuple[Path, int, dict]],
|
| 823 |
+
dropped: set[str],
|
| 824 |
+
plan: list[dict],
|
| 825 |
+
expected_index: dict,
|
| 826 |
+
source_config: dict,
|
| 827 |
+
) -> None:
|
| 828 |
+
index_path = out / INDEX_NAME
|
| 829 |
+
config_path = out / CONFIG_NAME
|
| 830 |
+
if not index_path.is_file() or not config_path.is_file():
|
| 831 |
+
raise RuntimeError("overlay index/config is missing")
|
| 832 |
+
actual_index = json.loads(index_path.read_text())
|
| 833 |
+
if actual_index.get("weight_map") != expected_index.get("weight_map"):
|
| 834 |
+
raise RuntimeError("overlay index weight_map differs from the planned map")
|
| 835 |
+
|
| 836 |
+
# Every indexed tensor is opened through its file header and its declared
|
| 837 |
+
# shape/dtype/byte length is checked against the index.
|
| 838 |
+
output_headers: dict[str, tuple[Path, int, dict]] = {}
|
| 839 |
+
for filename in sorted(set(actual_index["weight_map"].values())):
|
| 840 |
+
path = out / filename
|
| 841 |
+
if not path.is_file():
|
| 842 |
+
raise RuntimeError(f"indexed shard is missing: {path}")
|
| 843 |
+
header_len, header = read_header(path)
|
| 844 |
+
for name, mapped_file in actual_index["weight_map"].items():
|
| 845 |
+
if mapped_file != filename:
|
| 846 |
+
continue
|
| 847 |
+
if name not in header:
|
| 848 |
+
raise RuntimeError(f"indexed tensor is missing from {filename}: {name}")
|
| 849 |
+
meta = header[name]
|
| 850 |
+
if tensor_nbytes(meta) != meta["data_offsets"][1] - meta["data_offsets"][0]:
|
| 851 |
+
raise RuntimeError(f"invalid indexed tensor payload: {name}")
|
| 852 |
+
output_headers[name] = (path, header_len, header)
|
| 853 |
+
if set(output_headers) != set(actual_index["weight_map"]):
|
| 854 |
+
raise RuntimeError("not every indexed tensor was opened")
|
| 855 |
+
|
| 856 |
+
# Confirm source links and byte identity of every retained tensor in the
|
| 857 |
+
# three rewritten shards. Other source shards are links and need no copy.
|
| 858 |
+
for entry in src.iterdir():
|
| 859 |
+
if not (entry.is_file() or entry.is_symlink()) or entry.name in SOURCE_REWRITE_FILES or entry.name in (CONFIG_NAME, INDEX_NAME):
|
| 860 |
+
continue
|
| 861 |
+
linked = out / entry.name
|
| 862 |
+
if not linked.is_symlink() or linked.resolve() != entry.resolve():
|
| 863 |
+
raise RuntimeError(f"source file is not linked unchanged: {entry.name}")
|
| 864 |
+
for name, (source_path, source_header_len, source_header) in source.items():
|
| 865 |
+
if name in dropped:
|
| 866 |
+
if name in actual_index["weight_map"]:
|
| 867 |
+
raise RuntimeError(f"dropped source tensor remains indexed: {name}")
|
| 868 |
+
continue
|
| 869 |
+
if source_path.name not in SOURCE_REWRITE_FILES:
|
| 870 |
+
continue
|
| 871 |
+
output_path, output_header_len, output_header = output_headers[name]
|
| 872 |
+
if output_header[name]["dtype"] != source_header[name]["dtype"] or output_header[name]["shape"] != source_header[name]["shape"]:
|
| 873 |
+
raise RuntimeError(f"retained tensor metadata changed: {name}")
|
| 874 |
+
compare_payloads(
|
| 875 |
+
source_path,
|
| 876 |
+
output_path,
|
| 877 |
+
source_header[name],
|
| 878 |
+
output_header[name],
|
| 879 |
+
source_header_len,
|
| 880 |
+
output_header_len,
|
| 881 |
+
)
|
| 882 |
+
|
| 883 |
+
# New names, shapes, dtypes and MCG markers must match the donor plan.
|
| 884 |
+
for item in plan:
|
| 885 |
+
path, header_len, header = output_headers[item["target"]]
|
| 886 |
+
del path
|
| 887 |
+
actual = header[item["target"]]
|
| 888 |
+
if actual["dtype"] != item["meta"]["dtype"] or actual["shape"] != item["meta"]["shape"]:
|
| 889 |
+
raise RuntimeError(f"new tensor metadata differs from donor plan: {item['target']}")
|
| 890 |
+
if item["target"].endswith(".mcg"):
|
| 891 |
+
with (out / NEW_SHARD).open("rb") as stream:
|
| 892 |
+
start, end = actual["data_offsets"]
|
| 893 |
+
stream.seek(8 + header_len + start)
|
| 894 |
+
value = struct.unpack("<i", stream.read(end - start))[0]
|
| 895 |
+
if value != -877912083:
|
| 896 |
+
raise RuntimeError(f"MCG marker mismatch for {item['target']}: {value}")
|
| 897 |
+
|
| 898 |
+
config = json.loads(config_path.read_text())
|
| 899 |
+
q = config.get("quantization_config", {})
|
| 900 |
+
if "mtp_experts" in q or "mtp_experts_start_layer" in q:
|
| 901 |
+
raise RuntimeError("MTP source-format config keys were not removed")
|
| 902 |
+
if q.get("non_routed_quantization") != source_config["quantization_config"].get("non_routed_quantization"):
|
| 903 |
+
raise RuntimeError("non_routed_quantization changed")
|
| 904 |
+
if q.get("non_routed_dtype_policy") != source_config["quantization_config"].get("non_routed_dtype_policy"):
|
| 905 |
+
raise RuntimeError("non_routed_dtype_policy changed")
|
| 906 |
+
for key in ("43", "44", "45"):
|
| 907 |
+
if q.get("layer_bits", {}).get(key) != 2:
|
| 908 |
+
raise RuntimeError(f"layer_bits[{key}] is not explicitly 2")
|
| 909 |
+
log(f"VERIFY_OK indexed_tensors={len(output_headers)} retained_payloads_byte_identical=true")
|
| 910 |
+
|
| 911 |
+
|
| 912 |
+
def build(args: argparse.Namespace) -> int:
|
| 913 |
+
src = Path(args.src).expanduser().resolve()
|
| 914 |
+
out = Path(args.out).expanduser()
|
| 915 |
+
donor = args.donor
|
| 916 |
+
if not src.is_dir():
|
| 917 |
+
raise RuntimeError(f"source directory does not exist: {src}")
|
| 918 |
+
safe_output_path(src, out)
|
| 919 |
+
|
| 920 |
+
source, source_map = source_headers(src)
|
| 921 |
+
dropped, by_file, dropped_count, dropped_bytes = source_drop_plan(source)
|
| 922 |
+
del by_file
|
| 923 |
+
draft_shapes = source_draft_shapes(source)
|
| 924 |
+
main_shapes = main_k2_shapes(source)
|
| 925 |
+
|
| 926 |
+
donor_index_path = Path(args.donor_index).expanduser() if args.donor_index else None
|
| 927 |
+
if donor_index_path is not None and donor_index_path.exists():
|
| 928 |
+
donor_map = donor_index_from_file(donor_index_path)
|
| 929 |
+
log(f"DONOR_INDEX read={donor_index_path}")
|
| 930 |
+
elif args.dry_run:
|
| 931 |
+
raise RuntimeError(
|
| 932 |
+
f"--dry-run requires the local donor index; not found: {donor_index_path}"
|
| 933 |
+
)
|
| 934 |
+
else:
|
| 935 |
+
base = donor.rstrip("/") if donor.startswith("http") else f"https://huggingface.co/{donor}/resolve/main"
|
| 936 |
+
donor_index_blob = http_bytes(f"{base}/model.safetensors.index.json")
|
| 937 |
+
try:
|
| 938 |
+
donor_map = json.loads(donor_index_blob)["weight_map"]
|
| 939 |
+
except (json.JSONDecodeError, KeyError, TypeError) as exc:
|
| 940 |
+
raise RuntimeError(f"remote donor index is invalid: {exc}") from exc
|
| 941 |
+
log(f"DONOR_INDEX remote tensors={len(donor_map)}")
|
| 942 |
+
selected = validate_donor_selection(donor_map)
|
| 943 |
+
|
| 944 |
+
cache_path = Path(args.donor_headers).expanduser() if args.donor_headers else None
|
| 945 |
+
donor_headers: dict[str, tuple[int, dict]] | None = None
|
| 946 |
+
if args.dry_run:
|
| 947 |
+
if cache_path is not None and cache_path.exists():
|
| 948 |
+
donor_headers = donor_headers_from_remote(
|
| 949 |
+
donor,
|
| 950 |
+
donor_map,
|
| 951 |
+
selected,
|
| 952 |
+
cache_path,
|
| 953 |
+
write_cache=False,
|
| 954 |
+
allow_network=False,
|
| 955 |
+
)
|
| 956 |
+
else:
|
| 957 |
+
log("DONOR_HEADERS offline=not-read; dry-run uses local MTP dimensions")
|
| 958 |
+
else:
|
| 959 |
+
donor_headers = donor_headers_from_remote(donor, donor_map, selected, cache_path)
|
| 960 |
+
plan = validate_and_plan_donor(selected, donor_map, donor_headers, draft_shapes, main_shapes)
|
| 961 |
+
|
| 962 |
+
fetched_bytes = sum(tensor_nbytes(item["meta"]) for item in plan)
|
| 963 |
+
rewritten_bytes = sum(
|
| 964 |
+
tensor_nbytes(meta)
|
| 965 |
+
for name, (_, _, header) in source.items()
|
| 966 |
+
if name not in dropped and Path(source_map[name]).name in SOURCE_REWRITE_FILES
|
| 967 |
+
for meta in [header[name]]
|
| 968 |
+
)
|
| 969 |
+
source_config = edited_config(src)
|
| 970 |
+
expected_index = output_index(source, dropped, plan)
|
| 971 |
+
log(
|
| 972 |
+
f"PLAN donor_tensors={len(plan)} bytes_to_fetch={fetched_bytes} "
|
| 973 |
+
f"({fetched_bytes / (1 << 30):.3f} GiB)"
|
| 974 |
+
)
|
| 975 |
+
log(
|
| 976 |
+
f"PLAN rewrite_files=3 bytes_to_rewrite={rewritten_bytes} "
|
| 977 |
+
f"bytes_removed={dropped_bytes} dropped_tensors={dropped_count}"
|
| 978 |
+
)
|
| 979 |
+
log(
|
| 980 |
+
f"PLAN output symlinks=source_files_except_46_48_and_config "
|
| 981 |
+
f"new_shard={NEW_SHARD} index={INDEX_NAME} config_edits=remove_mtp_source_keys+layer_bits_43_44_45=2"
|
| 982 |
+
)
|
| 983 |
+
if args.dry_run:
|
| 984 |
+
log("DRY_RUN_OK network=false writes=false")
|
| 985 |
+
return 0
|
| 986 |
+
|
| 987 |
+
out.mkdir(parents=True, exist_ok=True)
|
| 988 |
+
link_source_files(src, out)
|
| 989 |
+
if donor_headers is None:
|
| 990 |
+
raise RuntimeError("donor headers were not loaded for a real build")
|
| 991 |
+
write_new_shard(out, plan, donor, donor_headers)
|
| 992 |
+
for filename in SOURCE_REWRITE_FILES:
|
| 993 |
+
rewrite_source_shard(src / filename, out / filename, dropped)
|
| 994 |
+
write_output_metadata(out, expected_index, source_config)
|
| 995 |
+
log(f"BUILD_OK out={out}")
|
| 996 |
+
if args.verify:
|
| 997 |
+
verify_output(src, out, source, dropped, plan, expected_index, json.loads((src / CONFIG_NAME).read_text()))
|
| 998 |
+
return 0
|
| 999 |
+
|
| 1000 |
+
|
| 1001 |
+
def main() -> int:
|
| 1002 |
+
parser = argparse.ArgumentParser(description=__doc__)
|
| 1003 |
+
parser.add_argument("--src", required=True, help="source pack directory")
|
| 1004 |
+
parser.add_argument("--donor", required=True, help="Hugging Face repo id or base URL")
|
| 1005 |
+
parser.add_argument("--out", required=True, help="new overlay directory")
|
| 1006 |
+
parser.add_argument(
|
| 1007 |
+
"--donor-index",
|
| 1008 |
+
default=DEFAULT_DONOR_INDEX,
|
| 1009 |
+
help=f"offline donor weight_map JSON (default: {DEFAULT_DONOR_INDEX})",
|
| 1010 |
+
)
|
| 1011 |
+
parser.add_argument(
|
| 1012 |
+
"--donor-headers",
|
| 1013 |
+
default=DEFAULT_HEADER_CACHE,
|
| 1014 |
+
help=f"remote donor header cache (default: {DEFAULT_HEADER_CACHE})",
|
| 1015 |
+
)
|
| 1016 |
+
parser.add_argument("--dry-run", action="store_true", help="plan only; no network and no writes")
|
| 1017 |
+
parser.add_argument("--verify", action="store_true", help="verify every indexed tensor after building")
|
| 1018 |
+
args = parser.parse_args()
|
| 1019 |
+
if args.dry_run and args.verify:
|
| 1020 |
+
parser.error("--dry-run and --verify cannot be combined")
|
| 1021 |
+
try:
|
| 1022 |
+
return build(args)
|
| 1023 |
+
except KeyboardInterrupt:
|
| 1024 |
+
log("ABORT interrupted")
|
| 1025 |
+
return 130
|
| 1026 |
+
except Exception as exc:
|
| 1027 |
+
log(f"ABORT {type(exc).__name__}: {exc}")
|
| 1028 |
+
return 1
|
| 1029 |
+
|
| 1030 |
+
|
| 1031 |
+
if __name__ == "__main__":
|
| 1032 |
+
sys.exit(main())
|
model-00001-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:367c971dc3cd6a042a9bec1caff508e77eabfaef1df2de3a827397ef8bbc6af3
|
| 3 |
+
size 1991876272
|
model-00002-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d654a0be1cc7be7222d79d1e3971ee64fa78867e01fbcae815032e71951de193
|
| 3 |
+
size 1631836440
|
model-00003-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:3443031c34ff01bf4ef0b23fb36d60cd6cbf45a4e01128c816a3389adf7ac46e
|
| 3 |
+
size 1631836440
|
model-00004-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:34687841798a15ce9b46052ff8770237d8e8343b08b3b8cd4e985de2913ac512
|
| 3 |
+
size 1632383008
|
model-00005-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:aa3b869fe2e277e6f46d2618d4f5e8625d4ba359d2ebaa5fda303e5286c49440
|
| 3 |
+
size 5529472
|
model-00006-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0e97556b8989a374dec330c3c886dcd6e19a56001c7daf9e835cc306587a2e1f
|
| 3 |
+
size 1626176072
|
model-00007-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:dfecfd503aabac065be311b383150950786df03bbce177b881c0bd37b64b92a9
|
| 3 |
+
size 1625892784
|
model-00008-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:cc8bf1d81af397ef98daf0fb1653ae08454236a6e5e57d65c425fd0082dd7f51
|
| 3 |
+
size 1626176072
|
model-00009-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fad91e6c1eaa4ca3e48d4e3f2f98c69c86b06ad1acd9cf407dd46728556be008
|
| 3 |
+
size 1625892784
|
model-00010-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:1c1e892b046ab791432bddf923111f705ea762638d52ce9d0499401b5c623d80
|
| 3 |
+
size 1626176072
|
model-00011-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e7d4ca138c51b71490a0d18daa61e38fe66d001f09e3365ce52146f5228091bf
|
| 3 |
+
size 1625892784
|
model-00012-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b3a4d98b8205cf0f531abba739bfd513266cb29f1420bddc6ddc5068b864133a
|
| 3 |
+
size 1626179160
|
model-00013-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ac720c288d0664776005ce618afb0f225916cf0a017549129913ddbe1f68c676
|
| 3 |
+
size 1625895872
|
model-00014-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:cac7ec09d5cedda548777f9a1cabea3427df3a12261e8fbc44780c88a9ddfbf8
|
| 3 |
+
size 1626179160
|
model-00015-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5716a6ecfde594977d45ddc7120dec61367503d9b47714f9019973a2f8360ebc
|
| 3 |
+
size 2431202576
|
model-00016-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:1ee66c56843d7343a08f5ead3acd6ac4f2eb26edb6d6067e015bc51a90428d84
|
| 3 |
+
size 1626179160
|
model-00017-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2b9fffdfcecaa4924223fd6861a3281e331f52610e4406077f19d854a6d02ee5
|
| 3 |
+
size 1625895872
|
model-00018-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a3e3be162b08f58e7ab5fcabeed8c9933c416a9be7798706f4b4a382d5cddaee
|
| 3 |
+
size 1626179160
|
model-00019-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:67283c44bf0daababc353a34bf14d3858cb71561ce0cd1189e2012e275964d05
|
| 3 |
+
size 1625895872
|
model-00020-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4ae488fb14a99df7c32f0d4547d6ec7db03a120c14b4251e309460ddb82b1dad
|
| 3 |
+
size 1626179160
|
model-00021-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e080a0e3dc4607142b37a9b85ae3b1802fcd0c8b872815ea7c51a227e4e1ea3d
|
| 3 |
+
size 1625895872
|
model-00022-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fb567d322e29b3ce20ca332de7ac3d451c91836518bfdf6bbafb097e9d56ef3c
|
| 3 |
+
size 1626179160
|
model-00023-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ae4fe98e86714fcfdece132565b3240d88567815c3c51ab6246fa9789c937326
|
| 3 |
+
size 5529488
|
model-00024-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e8e884e9ad843f62dcd625763e4750b5aabc9a0822b4db786c722c4284389756
|
| 3 |
+
size 2431485872
|
model-00025-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:98804432b6124d2a0d7fe527e9a9e4c7d691d78ef8489f3c20e9d30a406af022
|
| 3 |
+
size 1625895872
|
model-00026-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b3e1295c53d8f9b711d57b10112929685927590327a257372288ba28185c0906
|
| 3 |
+
size 1626179160
|
model-00027-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a8bdc5ef60398f85f6ea7d5828915ad2fea0b0415adf8b1615e19cb423737632
|
| 3 |
+
size 1625895872
|
model-00028-of-00048.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d79711081aee89cb63a88758c625ae9998fc49d75da6e520a65dfb34531af68d
|
| 3 |
+
size 1626179160
|