sakamakismile commited on
Commit
03afa3a
·
verified ·
1 Parent(s): 071d1f6

DSV4-Flash-Vision EXL3 MixedK D2-K2x3-Dense6 + lna-lab recipe (2026-09-03)

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +2 -0
  2. README.md +79 -0
  3. config.json +1178 -0
  4. generation_config.json +9 -0
  5. lna-lab/Dockerfile +14 -0
  6. lna-lab/PLAN.md +74 -0
  7. lna-lab/README-dsv4-carve.md +65 -0
  8. lna-lab/UPSTREAM-vcruz305-README.md +51 -0
  9. lna-lab/bench-dsv4.py +37 -0
  10. lna-lab/bench-streams.py +29 -0
  11. lna-lab/exllamav3_ext.cpython-312-x86_64-linux-gnu.so +3 -0
  12. lna-lab/patches/exl3.py.lna +389 -0
  13. lna-lab/patches/exl3_gemm.cu.lna +703 -0
  14. lna-lab/patches/patch_dsv4_attention_compressor_exl3.py +62 -0
  15. lna-lab/patches/patch_dsv4_aux_streams_env.py +19 -0
  16. lna-lab/patches/patch_dsv4_dense_exl3.py +238 -0
  17. lna-lab/patches/vllm_exl3_exl3.py.lna +1528 -0
  18. lna-lab/ppl-vllm.py +23 -0
  19. lna-lab/serve-dsv4-tp8.sh +20 -0
  20. lna-lab/tools/dense_bake.py +924 -0
  21. lna-lab/tools/layer_overlay.py +957 -0
  22. lna-lab/tools/mtp_overlay.py +1032 -0
  23. model-00001-of-00048.safetensors +3 -0
  24. model-00002-of-00048.safetensors +3 -0
  25. model-00003-of-00048.safetensors +3 -0
  26. model-00004-of-00048.safetensors +3 -0
  27. model-00005-of-00048.safetensors +3 -0
  28. model-00006-of-00048.safetensors +3 -0
  29. model-00007-of-00048.safetensors +3 -0
  30. model-00008-of-00048.safetensors +3 -0
  31. model-00009-of-00048.safetensors +3 -0
  32. model-00010-of-00048.safetensors +3 -0
  33. model-00011-of-00048.safetensors +3 -0
  34. model-00012-of-00048.safetensors +3 -0
  35. model-00013-of-00048.safetensors +3 -0
  36. model-00014-of-00048.safetensors +3 -0
  37. model-00015-of-00048.safetensors +3 -0
  38. model-00016-of-00048.safetensors +3 -0
  39. model-00017-of-00048.safetensors +3 -0
  40. model-00018-of-00048.safetensors +3 -0
  41. model-00019-of-00048.safetensors +3 -0
  42. model-00020-of-00048.safetensors +3 -0
  43. model-00021-of-00048.safetensors +3 -0
  44. model-00022-of-00048.safetensors +3 -0
  45. model-00023-of-00048.safetensors +3 -0
  46. model-00024-of-00048.safetensors +3 -0
  47. model-00025-of-00048.safetensors +3 -0
  48. model-00026-of-00048.safetensors +3 -0
  49. model-00027-of-00048.safetensors +3 -0
  50. 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