Add files using upload-large-folder tool

README.md

@@ -0,0 +1,152 @@
+---
+license: other
+base_model:
+- thinkingmachines/Inkling
+tags:
+- moe
+- nvfp4
+- aqlm
+- quantized
+- inkling
+---
+
+# Inkling-512k-NVFP4-AQLM-hybrid
+
+A memory-optimized **hybrid-quantized** build of
+[thinkingmachines/Inkling](https://huggingface.co/thinkingmachines/Inkling)
+(MoE, 256 routed experts/layer, 6 active). Each MoE layer's experts are split
+into a small **hot** set kept at NVFP4 (~4 bpw) and a large **cold** set stored
+as **AQLM 2-bit**. This is the **512k tier**: **14.8% of expert slots are
+NVFP4 (hot)** and **85.2% are AQLM-2bit (cold)**, averaging **~2.66 bits
+per weight** over the routed experts. It targets serving at **512k context on
+4x96GB GPUs**.
+
+> **This checkpoint needs a CUSTOM FUSED MoE KERNEL** that dispatches NVFP4 hot
+> experts and AQLM cold experts within the same fused grouped-GEMM. **It does
+> NOT load on stock vLLM / Transformers.** The on-disk format and dequant math
+> are fully specified below so a kernel can consume it directly.
+
+## Serving (512k on 4x96GB)
+This is the **higher-quality 14.8%-hot** tier: **~358 GB weights** + KV @512k
+(fp8, 11 global layers) ~11.5 GB leaves only **~14.6 GB** runtime headroom on a
+384 GB box. That is **too tight for CUDA graphs + MTP at 512k** — serve with
+**`--enforce-eager`** at 512k (with MTP), or use **256k** context (KV ~5.8 GB,
+~20 GB headroom) for non-eager. For non-eager + MTP + CUDA-graphs at full 512k,
+use the **10%-hot** tier instead (~346 GB weights, ~27 GB headroom). TP4,
+`--kv-cache-dtype fp8`, modest `--max-num-seqs`.
+
+## Quality
+Held-out perplexity (teacher-forced, 5-domain mix disjoint from calibration;
+gate-proxy = cold experts requantized to NVFP4 so it loads on stock vLLM, whose
+per-expert error tracks the true hybrid). Lower is better.
+
+| domain | Inkling-NVFP4 | prior hybrid | **this build (v3, 14.8% hot)** |
+|---|---|---|---|
+| code | 2.222 | 2.250 | 2.244 |
+| agentic | 1.930 | 1.938 | 1.934 |
+| instruction | 4.674 | 5.169 | **4.892** |
+| medical | 5.464 | 5.964 | **5.817** |
+| prose | 5.399 | 6.691 | **6.563** |
+| **aggregate** | **3.406** | **3.678** | **3.608** |
+
+This **v3** build re-tiers and re-calibrates on a broad **multimodal** corpus
+(code / agentic / instruction / medical / prose, plus image + audio routing).
+It beats the prior text-only-calibrated hybrid by **−1.9% aggregate**, with the
+gains concentrated where the prior build was starved (**instruction −5.4%,
+medical −2.5%, prose −1.9%**); code/agentic are unchanged within noise. This
+14.8% tier edges the 10%-hot tier (−1.1%) at the cost of the tighter 512k fit
+noted above.
+
+## Provenance
+- **Hot experts (NVFP4):** sliced verbatim from the ModelOpt NVFP4 base
+  (`thinkingmachines/Inkling-NVFP4`), no requantization.
+- **Cold experts (AQLM):** quantized from the **Inkling BF16 teacher**.
+  Both `w13` and `w2` cold books are **activation-aware converged** then
+  **output-preserving PV-tuned** against the BF16 teacher on **multimodal**
+  calibration activations (text + image + audio routing), with a disjoint
+  held-out early-stop.
+- **Assignment** of hot vs cold per layer is in
+  `hf_quant_config.json['aqlm_hybrid']['aqlm_layer_books']` and reproduced from
+  the multimodal REAP importance study (`assign_512k_v3.json`, 14.8% hot). Reproduction scripts are in
+  `code/`.
+
+## Layout summary
+- **63 NVFP4-packed MoE layers** carry the hybrid expert format below
+  (NVFP4 hot slabs + AQLM cold).
+- **1 bf16-special MoE layer(s)** (layer **2**) have **no NVFP4
+  base**, so they use the SAME hybrid format EXCEPT the hot slabs are stored as
+  **bf16** (`experts.{proj}_hot_bf16`, no `.scale`/`.scale2`); their cold experts
+  are AQLM like every other layer. The fused uncompressed bf16 `experts.w13_weight`
+  / `experts.w2_weight` are **dropped**. Layers 0-1 are dense MLP (not experts).
+- All **non-expert** tensors (attention, embeddings, norms, `shared_experts`,
+  `gate`, vision, audio, MTP) are copied verbatim from the NVFP4 base.
+- Total hot slots 2417 / 16384 expert slots.
+
+## On-disk tensor format (per NVFP4-packed MoE layer L)
+`prefix = model.llm.layers.{L}.mlp.experts`. Routed experts (256 total) are
+partitioned into `n_hot` hot + `n_cold` cold. For each projection
+`proj in {w13, w2}` (w13 = fused gate|up, out=6144, in=6144; w2 = down,
+out=6144, in=3072):
+
+**Hot — NVFP4-packed layers (ModelOpt, `hot_format=nvfp4`):**
+| tensor | dtype | shape |
+|---|---|---|
+| `{prefix}.{proj}_hot_weight`        | uint8       | `[n_hot, out, in/2]`  (2 fp4 codes/byte) |
+| `{prefix}.{proj}_hot_weight.scale`  | float8_e4m3 | `[n_hot, out, in/16]` (per-block-16 scale) |
+| `{prefix}.{proj}_hot_weight.scale2` | float32     | `[n_hot]`             (per-expert global scale) |
+
+**Hot — bf16-special layer(s) (`hot_format=bf16`, e.g. layer 2):**
+| tensor | dtype | shape |
+|---|---|---|
+| `{prefix}.{proj}_hot_bf16` | bfloat16 | `[n_hot, out, in]` (dense bf16, no scales) |
+
+**Cold (AQLM) — same for every layer. One code tensor PER BOOK, smallest lossless dtype:**
+| tensor | dtype | shape |
+|---|---|---|
+| `{prefix}.{proj}_cold_codes.{b}`    | uint8 if book has <=256 entries else int16 | `[n_cold, out, in/8]` (group_size=8) |
+| `{prefix}.{proj}_cold_codebook.{b}` | float16 | `[entries_b, 8]` |
+| `{prefix}.{proj}_cold_scales`     | float16 | `[n_cold, out]` (per-output-row scale) |
+
+Per-book code dtypes: `w13` book0 = **int16** (65536 entries); `w2` book0 =
+**int16** (65536), `w2` book1 = **uint8** (256 entries, codes 0..255 — lossless).
+`{proj}_cold_codes.{b}` replaces the old single `{proj}_cold_codes` tensor
+(the book axis is now the trailing `.{b}` suffix), saving ~38 GB losslessly.
+
+**Index maps:**
+| tensor | dtype | shape | meaning |
+|---|---|---|---|
+| `{prefix}.hot_ids`  | int32 | `[n_hot]`  | routed-expert index (0..255) of each hot slab |
+| `{prefix}.cold_ids` | int32 | `[n_cold]` | routed-expert index (0..255) of each cold slab |
+
+`hot_ids[i]` gives the original routed-expert id for hot slab `i` (and likewise
+`cold_ids`); together they partition `0..255`. **group_size = 8.**
+- `w13`: **1 book**, `entries = [65536]`.
+- `w2`:  **2 books**, `entries = [65536, 256]`.
+
+## Dequant math
+**Cold (AQLM), expert e, output row r, group g (8 weights):**
+```
+acc[8] = 0
+for b in range(n_books):
+    acc += codebook_b[ codes_b[e, r, g] ]           # codes_b = cold_codes.{b}; gather an 8-vector
+W[e, r, g*8:(g+1)*8] = acc * cold_scales[e, r]        # per-row scale
+```
+Codes are read from the per-book tensor `{proj}_cold_codes.{b}`. uint8 books
+(<=256 entries) index the codebook directly. For the int16 65536-entry book,
+negative int16 values wrap to the intended unsigned row (`idx & 0xFFFF`).
+
+**Hot (NVFP4, standard ModelOpt e2m1):**
+```
+codes = unpack_nibbles(hot_weight)                    # 0..15 -> e2m1 LUT value
+W = LUT[codes] * scale.repeat_interleave(16, dim=-1) * scale2
+```
+e2m1 LUT = `[0, .5, 1, 1.5, 2, 3, 4, 6, -0, -.5, -1, -1.5, -2, -3, -4, -6]`;
+low nibble is the first weight of each pair.
+
+## Files
+- `model-*.safetensors` + `model.safetensors.index.json` — the sharded weights.
+- `hf_quant_config.json` — original NVFP4 config **plus** the `aqlm_hybrid` block
+  (`quant_method`, per-layer `{n_nvfp4, n_cold, packed}`, book entry counts).
+- `config.json` — base config **plus** a `quantization_config` block advertising
+  `quant_method = inkling_nvfp4_aqlm_hybrid`.
+- `code/` — the exact repackaging scripts (`build_hybrid.py`, `ink_common.py`, ...).

chat_template.jinja

@@ -0,0 +1,129 @@
+{%- set effort_map = {"none": 0.0, "minimal": 0.1, "low": 0.2, "medium": 0.7, "high": 0.9, "max": 0.99} -%}
+{%- set role_token = {"user": "<|message_user|>", "assistant": "<|message_model|>", "system": "<|message_system|>", "tool": "<|message_tool|>"} -%}
+
+{%- macro emit_thinking_effort() -%}
+    {%- set eff = reasoning_effort if reasoning_effort is defined and reasoning_effort is not none else 0.9 -%}
+    {%- if eff is string -%}
+        {%- set key = eff | trim -%}
+        {%- if key not in effort_map -%}
+            {{- raise_exception("Unknown reasoning_effort: " ~ eff) -}}
+        {%- endif -%}
+        {%- set num = effort_map[key] -%}
+    {%- else -%}
+        {%- set num = eff | float -%}
+    {%- endif -%}
+    {%- if num < 0.0 or num > 0.99 -%}
+        {{- raise_exception("reasoning_effort must be in [0.0, 0.99]") -}}
+    {%- endif -%}
+    {{- "<|message_system|><|content_text|>Thinking effort level: " -}}
+    {%- if num == 0.0 -%}0{%- else -%}{{ num }}{%- endif -%}
+    {{- "<|end_message|>" -}}
+{%- endmacro -%}
+
+{%- if tools -%}
+    {%- set tool_state = namespace(specs=[]) -%}
+    {%- for tool in tools -%}
+        {%- set fn = tool.function if tool.function is defined else tool -%}
+        {%- set spec = {
+            "description": (fn.description if fn.description is defined and fn.description else ""),
+            "name": fn.name,
+            "parameters": (fn.parameters if fn.parameters is defined and fn.parameters else {}),
+            "type": (tool.type if tool.type is defined and tool.type else "function"),
+        } -%}
+        {%- set tool_state.specs = tool_state.specs + [spec] -%}
+    {%- endfor -%}
+    {{- "<|message_system|>tool_declare<|content_xml|>" -}}
+    {{- tool_state.specs | tojson(sort_keys=true, separators=(",", ":")) -}}
+    {{- "<|end_message|>" -}}
+{%- endif -%}
+
+{%- set state = namespace(effort_emitted=false) -%}
+{%- for message in messages -%}
+    {%- if message.role not in role_token -%}
+        {{- raise_exception("Unknown message role: " ~ message.role) -}}
+    {%- endif -%}
+    {%- if not state.effort_emitted and message.role != "system" -%}
+        {{- emit_thinking_effort() -}}
+        {%- set state.effort_emitted = true -%}
+    {%- endif -%}
+
+    {%- set rtok = role_token[message.role] -%}
+
+    {%- if message.role == "tool" -%}
+        {%- set tool_name_state = namespace(name="") -%}
+        {%- if message.name is defined and message.name -%}
+            {%- set tool_name_state.name = message.name -%}
+        {%- elif message.tool_call_id is defined and message.tool_call_id -%}
+            {%- for prev in messages -%}
+                {%- if prev.role == "assistant" and prev.tool_calls -%}
+                    {%- for tc in prev.tool_calls -%}
+                        {%- if tc.id is defined and tc.id == message.tool_call_id and tc.function.name is defined -%}
+                            {%- set tool_name_state.name = tc.function.name -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+        {{- rtok -}}
+        {%- if tool_name_state.name -%}{{- tool_name_state.name -}}{%- endif -%}
+        {{- "<|content_text|>" -}}
+        {%- if message.content is string -%}{{- message.content -}}{%- endif -%}
+        {{- "<|end_message|>" -}}
+
+    {%- else -%}
+        {%- if message.role == "assistant" and message.reasoning_content is defined and message.reasoning_content -%}
+            {{- "<|message_model|><|content_thinking|>" ~ message.reasoning_content ~ "<|end_message|>" -}}
+        {%- endif -%}
+
+        {%- if message.content is string -%}
+            {{- rtok ~ "<|content_text|>" ~ message.content ~ "<|end_message|>" -}}
+        {%- elif message.content -%}
+            {%- for part in message.content -%}
+                {%- if part is string -%}
+                    {{- rtok ~ "<|content_text|>" ~ part ~ "<|end_message|>" -}}
+                {%- elif part.type is not defined or part.type in ("text", "input_text") -%}
+                    {%- set text_part = (part.text if part.text is defined and part.text is string else "") -%}
+                    {{- rtok ~ "<|content_text|>" ~ text_part ~ "<|end_message|>" -}}
+                {%- elif part.type in ("image", "input_image", "image_url") -%}
+                    {{- rtok ~ "<|content_image|><|unused_200054|><|end_message|>" -}}
+                {%- elif part.type in ("audio", "input_audio", "audio_url") -%}
+                    {{- rtok ~ "<|content_audio_input|><|unused_200053|><|audio_end|><|end_message|>" -}}
+                {%- else -%}
+                    {{- raise_exception("Unsupported content part type: " ~ part.type) -}}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+
+        {%- if message.role == "assistant" and message.tool_calls -%}
+            {%- for tc in message.tool_calls -%}
+                {%- set fn = tc.function -%}
+                {%- if fn.name is not defined or fn.name is not string -%}
+                    {{- raise_exception("tool call function name must be a string") -}}
+                {%- endif -%}
+                {%- set args = fn.arguments if fn.arguments is defined and fn.arguments else {} -%}
+                {%- if args is string -%}
+                    {{- raise_exception("tool call arguments must be a parsed object, not a JSON string; canonicalize upstream") -}}
+                {%- endif -%}
+                {%- if args is not mapping -%}
+                    {{- raise_exception("tool call arguments must be an object") -}}
+                {%- endif -%}
+                {{- "<|message_model|>" ~ fn.name ~ "<|content_invoke_tool_json|>" -}}
+                {{- '{"name":' ~ (fn.name | tojson(sort_keys=true, separators=(",", ":"))) ~ ',"args":' -}}
+                {{- (args | tojson(sort_keys=true, separators=(",", ":"))) -}}
+                {{- "}<|end_message|>" -}}
+            {%- endfor -%}
+        {%- endif -%}
+
+        {%- if message.role == "assistant" -%}
+            {{- "<|content_model_end_sampling|>" -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if not state.effort_emitted -%}
+    {{- emit_thinking_effort() -}}
+{%- endif -%}
+
+{%- if add_generation_prompt -%}
+    {{- "<|message_model|>" -}}
+{%- endif -%}

code/analyze_mm_routing.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""Diagnostic: do image/audio tokens route to experts that v2's (text-only) tiering
+left COLD? If yes, PV alone can't fully protect multimodal -> a v4 re-tier with
+multimodal routing is warranted. Reads the captured mm acts (ids) — free, no GPU.
+
+Reports, per modality: fraction of routed token-slots landing on v2-COLD experts,
+aggregate + worst layers. Writes $RUN/mm_routing.txt.
+"""
+import json, os, sys, glob
+import torch
+
+Q = "/data/inkling-quant"
+ASSIGN = f"{Q}/assign_512k_v2.json"
+OUT = sys.argv[1] if len(sys.argv) > 1 else f"{Q}/v3_run/mm_routing.txt"
+
+def cold_frac(acts_dir, hot):
+    """Over all layers with acts in acts_dir, fraction of (token,slot) routings
+    that hit a COLD expert. Returns (agg_frac, [(layer,frac,ntok)])."""
+    per = []
+    tot_cold = tot = 0
+    for p in sorted(glob.glob(f"{acts_dir}/acts_layer_*.pt")):
+        L = int(p.split("_")[-1].split(".")[0])
+        ids = torch.load(p, map_location="cpu").get("ids")
+        if ids is None or ids.numel() == 0:
+            continue
+        hotL = torch.tensor(sorted(hot.get(str(L), [])))
+        is_hot = torch.isin(ids, hotL)          # [N,6] bool
+        c = int((~is_hot).sum()); n = int(is_hot.numel())
+        tot_cold += c; tot += n
+        per.append((L, c / max(1, n), ids.shape[0]))
+    per.sort(key=lambda x: -x[1])
+    return (tot_cold / max(1, tot)), per
+
+def main():
+    hot = json.load(open(ASSIGN))["hot_experts"]
+    lines = ["=== multimodal routing vs v2 (text-only) tiering ===",
+             "fraction of image/audio routings landing on v2-COLD experts:"]
+    rec = False
+    for mm, d in [("image", f"{Q}/acts_v3_img"), ("audio", f"{Q}/acts_v3_audio")]:
+        if not os.path.isdir(d):
+            lines.append(f"  {mm}: (no acts at {d})"); continue
+        frac, per = cold_frac(d, hot)
+        lines.append(f"  {mm}: {100*frac:.1f}% of routings on cold experts "
+                     f"(worst layers: " + ", ".join(f"L{L}={100*f:.0f}%" for L,f,_ in per[:5]) + ")")
+        if frac > 0.55:   # cold experts get ~85% of text mass by design; >55% mm-cold => mm under-served
+            rec = True
+    # reference: text cold fraction (baseline expectation ~85% since 85% experts are cold)
+    if os.path.isdir(f"{Q}/acts_v3"):
+        tf, _ = cold_frac(f"{Q}/acts_v3", hot)
+        lines.append(f"  [ref] text: {100*tf:.1f}% on cold (expected high; ~85% of experts are cold)")
+    lines.append("")
+    lines.append("RE-TIER RECOMMENDED (v4): multimodal routing significantly under-served by v2 tiering."
+                 if rec else
+                 "NO re-tier needed: multimodal routing overlaps v2 hot set adequately; PV suffices.")
+    open(OUT, "w").write("\n".join(lines) + "\n")
+    print("\n".join(lines))
+
+if __name__ == "__main__":
+    main()

code/build_calib.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""Assemble a diverse local calibration corpus for routing-stats / tiering.
+Mixes code (vLLM src), prose/docs (markdown), and synthetic chat/instruction
+prompts. Not distribution-perfect, but exercises diverse expert routing.
+Output: /data/inkling-quant/calib.jsonl (one {"text": ...} per line, ~6k chars).
+"""
+import json, os, glob, random
+random.seed(42)
+OUT = "/data/inkling-quant/calib.jsonl"
+CHUNK = 6000  # chars (~1.5-2k tokens)
+MAX_CHUNKS = 700
+chunks = []
+
+def add_files(patterns, cap):
+    got = 0
+    files = []
+    for p in patterns:
+        files += glob.glob(p, recursive=True)
+    random.shuffle(files)
+    for f in files:
+        if got >= cap: break
+        try:
+            t = open(f, errors="ignore").read()
+        except Exception:
+            continue
+        if len(t) < 500: continue
+        for i in range(0, min(len(t), CHUNK*3), CHUNK):
+            c = t[i:i+CHUNK]
+            if len(c) >= 1000:
+                chunks.append(c); got += 1
+            if got >= cap: break
+
+# code (python + rust)
+add_files(["/data/vllm-inkling/vllm/**/*.py", "/data/vllm-inkling/rust/**/*.rs"], 300)
+# prose / docs
+add_files(["/data/vllm-inkling/docs/**/*.md", "/data/glm52/code/**/*.md"], 150)
+# synthetic instruction/chat/agentic (varied domains)
+prompts = [
+ "Explain how photosynthesis converts sunlight into chemical energy, step by step.",
+ "Write a haiku about the ocean at dawn, then explain the imagery.",
+ "You are a helpful assistant. A user asks: how do I safely defrost chicken? Answer thoroughly.",
+ "Summarize the causes of the French Revolution in three paragraphs.",
+ "Translate 'the quick brown fox jumps over the lazy dog' into French and German.",
+ "Given a list of integers, describe an algorithm to find the two numbers summing to a target.",
+ "A patient reports fatigue and joint pain. List a differential diagnosis reasoning process (educational).",
+ "Write a short story about a lighthouse keeper who discovers a message in a bottle.",
+ "Explain the difference between TCP and UDP and when to use each.",
+ "Describe how a mixture-of-experts transformer routes tokens to experts.",
+]
+for _ in range(12):
+    for p in prompts:
+        chunks.append(p + "\n\n" + random.choice(prompts) + "\n\n" + random.choice(prompts))
+
+random.shuffle(chunks)
+chunks = chunks[:MAX_CHUNKS]
+with open(OUT, "w") as fh:
+    for c in chunks:
+        fh.write(json.dumps({"text": c}) + "\n")
+print(f"wrote {len(chunks)} chunks to {OUT}  (~{sum(len(c) for c in chunks)//4} tokens est)")

code/build_calib_v3.py

@@ -0,0 +1,232 @@
+#!/usr/bin/env python3
+"""v3 corpus — BROAD GLM-5.2-style mix, with per-category disjoint held-out.
+
+Mirrors GLM-5.2's calib composition (code 40 / agentic 25 / instruction 15 /
+medical 10 / prose 10) instead of narrow wikitext, so routing coverage matches
+what Inkling actually serves. Each category is split 3 ways by content hash:
+  train -> calib_v3.jsonl   (PV trains; big, for cold-expert coverage)
+  pv    -> heldout_v3.jsonl (PV early-stop signal; disjoint)
+  ppl   -> ppl_v3.jsonl     (pre/post ppl; disjoint; {"text","cat"} for per-domain)
+
+Medical = PUBLIC LITERATURE / EXAM text only (PubMedQA abstracts, MedMCQA
+questions+explanations, general medical-education prompts). NO patient data / PHI.
+
+Run with the datasets-enabled venv:
+  /home/coder/git/glm52/.venv/bin/python build_calib_v3.py
+"""
+import json, os, glob, random, hashlib, re
+random.seed(1337)
+
+RAW = "/data/inkling-quant/corpus_v3/raw"
+OUT = "/data/inkling-quant"
+CHUNK = 6000; MIN_CHUNK = 1000
+# token budgets (approx via chars/4). GLM used ~15M; we aim ~9M train.
+TRAIN_TOK = 9_000_000
+MIX = {"code":0.40, "agentic":0.25, "instruction":0.15, "medical":0.10, "prose":0.10}
+# per-category held-out caps (chunks). ppl small+fast (~GLM's 46.7k tok); pv larger.
+PPL_CHUNKS_PER_CAT = 20      # ~20*1.5k = ~30k tok/cat -> ~50k+ ppl tok (~GLM's 46.7k)
+PV_CHUNKS_PER_CAT  = 60
+
+def chunks_of(text, cap=10**9):
+    out=[]; text=text.strip()
+    for i in range(0, len(text), CHUNK):
+        c=text[i:i+CHUNK]
+        if len(c)>=MIN_CHUNK: out.append(c)
+        if len(out)>=cap: break
+    return out
+
+def pack(docs, cap_tok=10**12):
+    """Concatenate a stream of (possibly short) docs into a buffer and emit
+    CHUNK-sized pieces, so short examples aren't dropped (GLM's approach)."""
+    out=[]; buf=""; tok=0
+    for d in docs:
+        if not d: continue
+        buf += d.rstrip()+"\n\n"
+        while len(buf)>=CHUNK:
+            out.append(buf[:CHUNK]); buf=buf[CHUNK:]; tok+=CHUNK//4
+            if tok>=cap_tok: return out
+    if len(buf)>=MIN_CHUNK: out.append(buf)
+    return out
+
+# ---------------- sources ----------------
+def src_code_local(cap_tok):
+    pats=["/data/vllm-inkling/vllm/**/*.py","/data/vllm-inkling/csrc/**/*.*",
+          "/data/vllm-inkling/rust/**/*.rs","/data/vllm-glm52-sm120/vllm/**/*.py",
+          "/data/vllm-glm52-sm120/csrc/**/*.cu","/home/coder/Antidoom/**/*.py"]
+    files=[]; [files.extend(glob.glob(p,recursive=True)) for p in pats]
+    files=sorted(set(files)); random.shuffle(files)
+    docs=[]; tok=0
+    for f in files:
+        if tok>=cap_tok: break
+        try: t=open(f,errors="ignore").read()
+        except: continue
+        if len(t)<400: continue
+        docs.append(f"# ==== {os.path.basename(f)} ====\n{t[:60000]}"); tok+=len(t[:60000])//4
+    return docs
+
+def _hf(name, cfg=None):
+    from datasets import load_dataset
+    return load_dataset(name, cfg, split="train", streaming=True) if cfg else \
+           load_dataset(name, split="train", streaming=True)
+
+def src_code_hf(cap_tok):
+    docs=[]; tok=0
+    try:
+        for r in _hf("m-a-p/CodeFeedback-Filtered-Instruction"):
+            docs.append(f"### Task\n{r['query']}\n\n### Solution\n{r['answer']}"); tok+=len(docs[-1])//4
+            if tok>=cap_tok*0.6: break
+    except Exception as e: print("  code_hf CodeFeedback:",str(e)[:80])
+    try:
+        for r in _hf("jtatman/python-code-dataset-500k"):
+            docs.append(f"# {r.get('instruction','')}\n{r.get('output','')}"); tok+=len(docs[-1])//4
+            if tok>=cap_tok: break
+    except Exception as e: print("  code_hf py500k:",str(e)[:80])
+    return docs
+
+TOOLS=["web_search","python","bash","read_file","sql_query","http_get","calculator","grep","git_log"]
+AG_TASKS=["Find the latest release version of numpy and summarize the changelog.",
+ "Compute the 20th Fibonacci number and verify it two ways.",
+ "Read config.yaml and report the database connection settings.",
+ "Query the orders table for total revenue in Q3 grouped by region.",
+ "Check whether the service on port 8080 is healthy and report latency.",
+ "Refactor this function to be O(n) and run the tests.",
+ "Download the CSV, parse it, and report the top 5 rows by value.",
+ "Diagnose why the CI build failed and propose a minimal fix.",
+ "Search the codebase for all callers of parse_config and summarize them.",
+ "Given a stack trace, locate the offending line and explain the bug."]
+def src_agentic(cap_tok):
+    """Varied synthetic tool-call sessions (content randomized so chunks are unique)."""
+    docs=[]; tok=0; rng=random.Random(7)
+    while tok<cap_tok:
+        task=rng.choice(AG_TASKS); n=rng.randint(2,5)
+        s=[f"<|user|>\n{task}\n<|assistant|>\nLet me work through this step by step.\n"]
+        for i in range(n):
+            tool=rng.choice(TOOLS); rid=rng.randint(1000,9999)
+            args=json.dumps({"q":f"{task[:30]}","step":i,"rid":rid,"limit":rng.randint(5,50)})
+            res=json.dumps({"rid":rid,"rows":rng.randint(0,200),"ok":rng.random()>0.1,
+                            "note":rng.choice(["cache miss","200 OK","partial","retry","done"])})
+            s.append(f'<tool_call>\n{{"name":"{tool}","arguments":{args}}}\n</tool_call>\n')
+            s.append(f'<tool_response>\n{res}\n</tool_response>\n')
+            s.append(rng.choice(["Observed the result; continuing.","That narrows it down.",
+                                 "Now I'll verify.","Proceeding to the next step."])+f" (step {i})\n")
+        s.append(f"Final answer: {task[:40]}... resolved after {n} tool calls (session {rng.randint(10**4,10**6)}).\n")
+        c="".join(s); docs.append(c); tok+=len(c)//4
+    return docs
+
+def src_instruction(cap_tok):
+    docs=[]; tok=0
+    try:
+        for r in _hf("tatsu-lab/alpaca"):
+            instr=r["instruction"]+(("\n\n"+r["input"]) if r.get("input") else "")
+            docs.append(f"<|user|>\n{instr}\n<|assistant|>\n{r['output']}"); tok+=len(docs[-1])//4
+            if tok>=cap_tok: break
+    except Exception as e: print("  instruction alpaca:",str(e)[:80])
+    return docs
+
+MEDICAL_QA=[
+ "Explain the differential diagnosis approach for acute chest pain across cardiac, pulmonary, GI, and musculoskeletal causes.",
+ "Describe how to systematically read a chest X-ray and common findings and pitfalls.",
+ "Explain the pharmacology of beta-blockers: mechanism, selectivity, indications, contraindications.",
+ "Walk through the pathophysiology of type 2 diabetes and the mechanism of each major drug class.",
+ "Explain how CT and MRI work physically and when each is preferred clinically.",
+ "Describe interpretation of a full blood count and common patterns and differentials.",
+ "Explain TNM staging of solid tumours and how imaging and histopathology contribute.",
+ "Describe the cardiac cycle and how it maps to ECG waveforms and heart sounds.",
+ "Explain sepsis: definitions, pathophysiology, recognition criteria, initial management.",
+ "Explain acid-base disturbances and step-by-step arterial blood gas interpretation."]
+RADIOLOGY_CSV="/data/inkling-quant/corpus_v3/raw/radiology/ReportsDATASET.csv"
+def src_medical(cap_tok):
+    """PUBLIC LITERATURE / EXAM text only (no patient data)."""
+    docs=[]; tok=0
+    # local de-identified radiology reports (harrison.ai domain), if present
+    if os.path.exists(RADIOLOGY_CSV):
+        import csv as _csv
+        for r in _csv.DictReader(open(RADIOLOGY_CSV,errors="ignore")):
+            t=next(iter(r.values()),"") or ""
+            if len(t)>40: docs.append(t); tok+=len(t)//4
+            if tok>=cap_tok*0.3: break
+    try:
+        for r in _hf("qiaojin/PubMedQA","pqa_labeled"):
+            ctx=" ".join(r["context"]["contexts"]) if isinstance(r.get("context"),dict) else str(r.get("context",""))
+            docs.append(f"{r['question']}\n{ctx}\n{r.get('long_answer','')}"); tok+=len(docs[-1])//4
+            if tok>=cap_tok*0.5: break
+    except Exception as e: print("  medical pubmedqa:",str(e)[:80])
+    try:
+        for r in _hf("openlifescienceai/medmcqa"):
+            opts=f"A){r['opa']} B){r['opb']} C){r['opc']} D){r['opd']}"
+            docs.append(f"{r['question']}\n{opts}\nExplanation: {r.get('exp') or ''}"); tok+=len(docs[-1])//4
+            if tok>=cap_tok: break
+    except Exception as e: print("  medical medmcqa:",str(e)[:80])
+    for q in MEDICAL_QA: docs.append(f"<|user|>\n{q} Be thorough.\n<|assistant|>")
+    return docs
+
+def src_prose(cap_tok):
+    docs=[]; tok=0
+    p=f"{RAW}/wikitext2-train.txt"           # encyclopedic
+    if os.path.exists(p):
+        raw=open(p,errors="ignore").read()
+        for para in re.split(r"\n = [^=]", raw):
+            docs.append(para); tok+=len(para)//4
+            if tok>=cap_tok*0.5: break
+    for f in sorted(glob.glob("/data/vllm-inkling/docs/**/*.md",recursive=True)):  # docs
+        if tok>=cap_tok*0.7: break
+        try: t=open(f,errors="ignore").read()
+        except: continue
+        docs.append(t); tok+=len(t)//4
+    try:                                      # dolly general knowledge
+        for r in _hf("databricks/databricks-dolly-15k"):
+            docs.append("\n\n".join(x for x in [r.get("instruction",""),r.get("context",""),r.get("response","")] if x))
+            tok+=len(docs[-1])//4
+            if tok>=cap_tok: break
+    except Exception as e: print("  prose dolly:",str(e)[:80])
+    return docs
+
+SOURCES={"code":lambda b:src_code_local(int(b*0.35))+src_code_hf(int(b*0.65)),
+         "agentic":src_agentic,"instruction":src_instruction,
+         "medical":src_medical,"prose":src_prose}
+
+# ---------------- split + write ----------------
+def _key(c): return hashlib.md5(c[:512].encode()).hexdigest()
+
+def build():
+    os.makedirs(RAW,exist_ok=True)
+    train,pv,ppl = [],[],[]      # (text,cat)
+    seen=set()
+    for cat,frac in MIX.items():
+        budget=int(TRAIN_TOK*frac / 0.8)     # +25% headroom for the held-out slices
+        docs=SOURCES[cat](budget)
+        chunks=pack(docs, budget)            # pack short docs; split long ones
+        # dedup within category + globally
+        uniq=[]
+        for c in chunks:
+            k=_key(c)
+            if k in seen: continue
+            seen.add(k); uniq.append(c)
+        random.Random(hash(cat)&0xffff).shuffle(uniq)
+        # allocate: ppl first, then pv, rest train (disjoint by slice)
+        cppl=uniq[:PPL_CHUNKS_PER_CAT]
+        cpv =uniq[PPL_CHUNKS_PER_CAT:PPL_CHUNKS_PER_CAT+PV_CHUNKS_PER_CAT]
+        ctr =uniq[PPL_CHUNKS_PER_CAT+PV_CHUNKS_PER_CAT:]
+        ppl += [(c,cat) for c in cppl]; pv += [(c,cat) for c in cpv]; train += [(c,cat) for c in ctr]
+        tt=sum(len(c)//4 for c in ctr)
+        print(f"  {cat:12s}: train {len(ctr)} (~{tt:,} tok) | pv {len(cpv)} | ppl {len(cppl)}")
+    random.shuffle(train); random.shuffle(pv)
+    def w(path, items, with_cat=False):
+        with open(path,"w") as f:
+            for c,cat in items:
+                f.write(json.dumps({"text":c,**({"cat":cat} if with_cat else {})})+"\n")
+        print(f"  wrote {path}: {len(items)} chunks (~{sum(len(c)//4 for c,_ in items):,} tok)")
+    w(f"{OUT}/calib_v3.jsonl", train)
+    w(f"{OUT}/heldout_v3.jsonl", pv)
+    w(f"{OUT}/ppl_v3.jsonl", ppl, with_cat=True)
+    # disjointness assertion
+    def keys(p): return {_key(json.loads(l)["text"]) for l in open(p)}
+    kC,kH,kP=keys(f"{OUT}/calib_v3.jsonl"),keys(f"{OUT}/heldout_v3.jsonl"),keys(f"{OUT}/ppl_v3.jsonl")
+    print(f"\n  overlap calib∩pv={len(kC&kH)} calib∩ppl={len(kC&kP)} pv∩ppl={len(kH&kP)} (all must be 0)")
+    from collections import Counter
+    print("  ppl per-category:", dict(Counter(cat for _,cat in ppl)))
+
+if __name__=="__main__":
+    print(">>> building BROAD v3 corpus (GLM-style mix) ...")
+    build()
+    print(">>> done.")

code/build_gate_ckpt.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python3
+"""Build the 512k QUALITY-GATE checkpoint: a standard ModelOpt-NVFP4 checkpoint
+identical to /data/inkling-nvfp4 EXCEPT COLD experts (per assign_512k.json) have
+their fused w13/w2 weights REPLACED by re-quantizing the AQLM parts to NVFP4
+(dequant AQLM->bf16->NVFP4). HOT experts keep their original Inkling-NVFP4 values.
+Loads on the existing vLLM NVFP4 path (no custom kernel).
+  w13 <- converged parts (parts_conv), w2 <- encode parts (parts).
+
+The three tensors of each fused expert group -- packed `...experts.{w13,w2}_weight`,
+`.scale`, `.scale2` -- live in DIFFERENT safetensors shards (per weight_map).
+Two phases:
+  phase1 (GPU 4-7, parallel): requant every cold expert ONCE per (layer,proj);
+          write compact per-group cache files (cold rows only). Resumable.
+  phase2 (CPU): stream each source shard, overwrite the cold-expert slices of any
+          expert packed/.scale/.scale2 tensor present, copy everything else
+          verbatim, keep the exact same shard layout. Resumable per shard.
+Out: /data/inkling-512k-gate/  (~592GB).
+"""
+import argparse, json, os, re, shutil
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+import ink_common as ic
+
+SRC   = ic.NVFP4_DIR
+# env-overridable so the SAME builder makes v1 / v2 / v3 gate proxies:
+#   v1: W13_PARTS=parts_conv W2_PARTS=parts       ASSIGN=assign_512k.json
+#   v2: W13_PARTS=parts_conv W2_PARTS=parts_conv  ASSIGN=assign_512k_v2.json GATE_OUT=/data/inkling-512k-gate-v2
+#   v3: W13_PARTS=parts_pv_v3 W2_PARTS=parts_pv_v3 ASSIGN=assign_512k_v2.json GATE_OUT=/data/inkling-512k-gate-v3
+OUT   = os.environ.get("GATE_OUT", "/data/inkling-512k-gate")
+W13_PARTS = os.environ.get("W13_PARTS", "/data/inkling-quant/parts_conv")   # w13 source
+W2_PARTS  = os.environ.get("W2_PARTS",  "/data/inkling-quant/parts")        # w2 source
+CONV = W13_PARTS; PARTS = W2_PARTS   # back-compat aliases
+ASSIGN = os.environ.get("ASSIGN", "/data/inkling-quant/assign_512k.json")
+CACHE = os.environ.get("REQUANT_CACHE", "/data/inkling-quant/tools/requant_cache")
+
+# packed `..._weight`, or `.scale`, or `.scale2` of a fused MoE expert group
+EXP_RE = re.compile(r"model\.llm\.layers\.(\d+)\.mlp\.experts\.(w13|w2)_weight(\.scale2|\.scale)?$")
+SUF2KEY = {"": "packed", ".scale": "scale", ".scale2": "scale2"}
+
+
+def cold_map():
+    a = json.load(open(ASSIGN))["hot_experts"]
+    return {int(L): [e for e in range(ic.N_EXPERTS) if e not in set(hot)]
+            for L, hot in a.items()}
+
+
+def full_layers():
+    """MoE layers stored full-precision (bf16, NOT NVFP4-packed) in the base
+    checkpoint -> identified by the absence of a `...w13_weight.scale` tensor.
+    In Inkling-NVFP4 this is layer 2 (the first MoE layer kept high precision)."""
+    wm = json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"]
+    fl = set()
+    for L in ic.moe_layers():
+        if f"model.llm.layers.{L}.mlp.experts.w13_weight.scale" not in wm:
+            fl.add(L)
+    return fl
+
+
+def cache_path(L, proj):
+    return f"{CACHE}/L{L}_{proj}.pt"
+
+
+# ---------------------------------------------------------------- phase 1 (GPU)
+def phase1(layers, dev="cuda"):
+    cold = cold_map()
+    full = full_layers()
+    os.makedirs(CACHE, exist_ok=True)
+    for L in layers:
+        for proj in ("w13", "w2"):
+            cp = cache_path(L, proj)
+            if os.path.exists(cp):
+                print(f"[gpu] skip cache L{L} {proj} (exists)", flush=True)
+                continue
+            d = CONV if proj == "w13" else PARTS
+            part = torch.load(f"{d}/layer_{L}_{proj}.pt", map_location="cpu")
+            ce = cold[L]
+            if L in full:
+                # full-precision (bf16) layer: cache AQLM dequant as bf16, no NVFP4
+                wl = [ic.aqlm_dequant_expert(part, proj, e, device=dev).to(torch.bfloat16).cpu()
+                      for e in ce]
+                out = {"cold": torch.tensor(ce, dtype=torch.long),
+                       "weight": torch.stack(wl)}
+                del wl
+                tag = "bf16"
+            else:
+                pk_l, sf_l, s2_l = [], [], []
+                for e in ce:
+                    W = ic.aqlm_dequant_expert(part, proj, e, device=dev).float()
+                    pk, sf8, s2 = ic.quantize_nvfp4(W)
+                    pk_l.append(pk.cpu()); sf_l.append(sf8.cpu()); s2_l.append(s2.cpu().reshape(()))
+                out = {"cold":  torch.tensor(ce, dtype=torch.long),
+                       "packed": torch.stack(pk_l),
+                       "scale":  torch.stack(sf_l),
+                       "scale2": torch.stack(s2_l)}
+                del pk_l, sf_l, s2_l
+                tag = "nvfp4"
+            tmp = cp + ".tmp"
+            torch.save(out, tmp); os.replace(tmp, cp)
+            del part, out
+            print(f"[gpu] cached L{L} {proj} ({len(ce)} cold experts, {tag})", flush=True)
+
+
+# ---------------------------------------------------------------- phase 2 (CPU)
+def phase2():
+    os.makedirs(OUT, exist_ok=True)
+    cold = cold_map()
+    idx = json.load(open(f"{SRC}/model.safetensors.index.json"))
+    wm = idx["weight_map"]
+    shards = sorted(set(wm.values()))
+    for si, shard in enumerate(shards):
+        outp = f"{OUT}/{shard}"
+        if os.path.exists(outp):
+            print(f"[{si+1}/{len(shards)}] skip {shard} (exists)", flush=True); continue
+        tensors, meta = {}, {}
+        with safe_open(f"{SRC}/{shard}", framework="pt") as f:
+            meta = f.metadata() or {}
+            for k in f.keys():
+                tensors[k] = f.get_tensor(k)
+        gcache = {}  # (L,proj) -> loaded cache dict; scoped to this shard
+
+        def gload(L, proj):
+            key = (L, proj)
+            if key not in gcache:
+                gcache[key] = torch.load(cache_path(L, proj), map_location="cpu")
+            return gcache[key]
+
+        changed = 0
+        for k in list(tensors.keys()):
+            m = EXP_RE.match(k)
+            if not m:
+                continue
+            L = int(m.group(1)); proj = m.group(2); suf = m.group(3) or ""
+            if L not in cold:
+                continue
+            g = gload(L, proj)
+            ce = g["cold"]
+            if "weight" in g:               # full-precision (bf16) layer
+                assert suf == "", f"unexpected aux tensor {k} on a bf16 layer"
+                src = g["weight"]
+            else:                            # NVFP4-packed layer
+                src = g[SUF2KEY[suf]]
+            t = tensors[k]
+            t[ce] = src.to(t.dtype)   # overwrite cold rows (cold-ordered) in place
+            tensors[k] = t
+            changed += 1
+        tmp = outp + ".tmp"
+        save_file(tensors, tmp, metadata=meta or {"format": "pt"})
+        os.replace(tmp, outp)
+        del tensors, gcache
+        print(f"[{si+1}/{len(shards)}] {shard}: {changed} expert tensors patched -> saved", flush=True)
+
+    # copy every non-safetensors file (config, tokenizer, index, hf_quant_config, ...)
+    # incl. dirs like tiktoken; skip the HF download cache dir.
+    SKIP = {".cache"}
+    for fn in os.listdir(SRC):
+        if fn.endswith(".safetensors") or fn in SKIP:
+            continue
+        src, dst = f"{SRC}/{fn}", f"{OUT}/{fn}"
+        if os.path.exists(dst):
+            continue
+        if os.path.isdir(src):
+            shutil.copytree(src, dst)
+        else:
+            shutil.copy2(src, dst)
+    print("PHASE2 DONE", flush=True)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--mode", required=True, choices=["phase1", "phase2"])
+    ap.add_argument("--worker", type=int, default=0)
+    ap.add_argument("--nworkers", type=int, default=1)
+    ap.add_argument("--device", default="cuda")
+    a = ap.parse_args()
+    if a.mode == "phase1":
+        layers = ic.moe_layers()[a.worker::a.nworkers]
+        print(f"[gpu] worker {a.worker}/{a.nworkers} layers: {layers}", flush=True)
+        phase1(layers, a.device)
+    else:
+        phase2()
+
+
+if __name__ == "__main__":
+    main()

code/build_heldout.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+"""Build a HELD-OUT text set DISJOINT from calib.jsonl for PV-generalization eval.
+
+Sources are deliberately DIFFERENT from calib's (calib = /data/vllm-inkling vllm
+python/rust src, /data/vllm-inkling docs md, /data/glm52/code md, + 10 synthetic
+prompts).  Held-out uses:
+  - code:  /home/coder/Antidoom/antidoom/src/**/*.py  (different repo, NOT vllm)
+  - prose: /data/inkling-recon/RECON.md, /data/inkling-nvfp4/README.md,
+           /data/inkling-512k-hybrid/README.md, /data/glm52-*/README.md
+  - fresh synthetic chat/instruction prompts (distinct topics from calib's 10)
+Then VERIFIES every held-out chunk is disjoint from calib (no held-out chunk is a
+substring of any calib chunk, and no long shared window) and drops any that overlap.
+Output: /data/inkling-quant/heldout.jsonl (~48 {"text": ...} chunks).
+"""
+import json, glob, random
+random.seed(7)
+OUT = "/data/inkling-quant/heldout.jsonl"
+CALIB = "/data/inkling-quant/calib.jsonl"
+CHUNK = 6000
+chunks = []
+
+def add_files(patterns, cap, per_file_cap=3):
+    got = 0; files = []
+    for p in patterns:
+        files += glob.glob(p, recursive=True)
+    random.shuffle(files)
+    for f in files:
+        if got >= cap: break
+        try:
+            t = open(f, errors="ignore").read()
+        except Exception:
+            continue
+        if len(t) < 500: continue
+        n_here = 0
+        for i in range(0, len(t), CHUNK):
+            c = t[i:i+CHUNK]
+            if len(c) >= 1000:
+                chunks.append(c); got += 1; n_here += 1
+            if got >= cap or n_here >= per_file_cap: break
+
+# code from a DIFFERENT repo than calib (antidoom, not vllm)
+add_files(["/home/coder/Antidoom/antidoom/src/**/*.py"], 22, per_file_cap=3)
+# prose/docs NOT in calib globs
+add_files(["/data/inkling-recon/RECON.md", "/data/inkling-nvfp4/README.md",
+           "/data/inkling-512k-hybrid/README.md", "/data/glm52-*/README.md",
+           "/data/glm52-fp8o/README.md"], 14, per_file_cap=2)
+# fresh synthetic prompts (distinct domains from calib's 10)
+syn = [
+ "Describe the lifecycle of a star from nebula to supernova, and what determines its final state.",
+ "Write a Python function that merges two sorted linked lists into one sorted list, with comments.",
+ "A customer emails that their package arrived damaged. Draft a professional, empathetic reply offering a resolution.",
+ "Explain the Byzantine Generals Problem and how blockchain consensus addresses it.",
+ "Compose a limerick about a cat who refuses to use its expensive bed.",
+ "What are the trade-offs between microservices and a monolithic architecture for a small startup?",
+ "Outline a 30-minute beginner yoga routine, naming each pose and its benefit.",
+ "Explain how vaccines train the adaptive immune system, appropriate for a curious teenager.",
+ "Given a noisy time series, describe how a Kalman filter estimates the underlying state.",
+ "Write a SQL query to find the second-highest salary in an employees table, and explain it.",
+ "Discuss the ecological role of mycorrhizal fungi in forest ecosystems.",
+ "Draft a short cover letter for a data-analyst role emphasizing curiosity and rigor.",
+]
+for _ in range(3):
+    for p in syn:
+        chunks.append(p + "\n\n" + random.choice(syn) + "\n\n" + random.choice(syn))
+
+# --- disjointness verification vs calib ---
+calib_chunks = [json.loads(l)["text"] for l in open(CALIB)]
+calib_blob = "\n".join(calib_chunks)
+
+def overlaps_calib(c, win=200, step=100):
+    # substring containment either way
+    if c in calib_blob:
+        return True
+    for cc in calib_chunks:
+        if cc in c:
+            return True
+    # long shared window
+    for i in range(0, max(1, len(c) - win), step):
+        if c[i:i+win] in calib_blob:
+            return True
+    return False
+
+random.shuffle(chunks)
+kept, dropped = [], 0
+for c in chunks:
+    if overlaps_calib(c):
+        dropped += 1; continue
+    kept.append(c)
+    if len(kept) >= 48:
+        break
+
+with open(OUT, "w") as fh:
+    for c in kept:
+        fh.write(json.dumps({"text": c}) + "\n")
+print(f"wrote {len(kept)} held-out chunks (dropped {dropped} for calib overlap) -> {OUT}")
+print(f"~{sum(len(c) for c in kept)//4} tokens est; calib has {len(calib_chunks)} chunks")

code/build_hybrid.py

@@ -0,0 +1,448 @@
+#!/usr/bin/env python3
+"""Assemble the COMPRESSED 512k hybrid checkpoint (the real memory-efficient
+artifact the custom kernel will run). NO requant: cold experts are already AQLM
+(parts_conv w13 / parts w2), hot are already NVFP4 (inkling-nvfp4). Pure repackage.
+
+FORMAT (per NVFP4-packed MoE layer L; layer 2 is bf16-special -> copied as-is):
+  experts.{w13,w2}_hot_weight         u8   [n_hot, out, in/2]      (NVFP4 packed)
+  experts.{w13,w2}_hot_weight.scale   f8e4m3 [n_hot, out, in/16]
+  experts.{w13,w2}_hot_weight.scale2  f32  [n_hot]
+  experts.{w13,w2}_cold_codes         int16 [n_cold, books, out, in/8]   (AQLM)
+  experts.{w13,w2}_cold_codebook.{b}  f16  [entries_b, 8]   (w13: 1 book@65536; w2: 65536 + 256)
+  experts.{w13,w2}_cold_scales        f16  [n_cold, out]
+  experts.hot_ids / experts.cold_ids  int32
+All NON-expert tensors (attn, embed, norms, shared_experts, gate, vision, audio,
+mtp, layer-2 experts) copied verbatim from inkling-nvfp4.
+Config: hf_quant_config.json extended with an 'aqlm_hybrid' block giving per-layer
+{n_nvfp4, n_cold, w13_books, w2_books, entries} + the assignment. See FORMAT in README.
+"""
+import argparse, json, os, re
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+import ink_common as ic
+
+SRC = ic.NVFP4_DIR
+# env-overridable for v3: HYBRID_OUT + COLD_PARTS (cold w13+w2 source) + ASSIGN
+#   v2: COLD_PARTS=parts_conv   HYBRID_OUT=/data/inkling-512k-hybrid-v2
+#   v3: COLD_PARTS=parts_pv_v3  HYBRID_OUT=/data/inkling-512k-hybrid-v3
+OUT = os.environ.get("HYBRID_OUT", "/data/inkling-512k-hybrid-v2")
+CONV = os.environ.get("COLD_PARTS", "/data/inkling-quant/parts_conv")
+PARTS = CONV
+ASSIGN = os.environ.get("ASSIGN", "/data/inkling-quant/assign_512k_v2.json")
+SHARD_BYTES = 8 * 2**30
+
+def hot_map():
+    a = json.load(open(ASSIGN))["hot_experts"]
+    return {int(L): sorted(v) for L, v in a.items()}
+
+def packed_layers():
+    idx = json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"]
+    ls = set()
+    for k in idx:
+        m = re.match(r"model\.llm\.layers\.(\d+)\.mlp\.experts\.w13_weight\.scale$", k)
+        if m: ls.add(int(m.group(1)))
+    return ls
+
+EXP_ANY = re.compile(r"model\.llm\.layers\.(\d+)\.mlp\.experts\.(w13|w2)_weight(\.scale2|\.scale)?$")
+
+def book_code_dtype(entries):
+    """Smallest int dtype that losslessly holds this book's codes.
+    <=256 entries -> uint8 (codes 0..255); otherwise int16 (65536-entry books
+    use int16 with the negative-wrap convention for rows >= 32768)."""
+    return torch.uint8 if entries <= 256 else torch.int16
+
+def build_layer_tensors(L, hot, nv):
+    """Hybrid expert tensors for one MoE layer.
+    Branches on the source format: NVFP4-packed layers keep packed uint8 hot
+    slabs (+scale/scale2); the bf16-special layer (layer 2, no NVFP4 base) keeps
+    bf16 hot slabs (no scales). Cold experts are AQLM in both cases, stored as
+    ONE code tensor per book with the smallest lossless dtype."""
+    out = {}
+    hids = hot[L]; cids = [e for e in range(ic.N_EXPERTS) if e not in set(hids)]
+    hidx = torch.tensor(hids); cidx = torch.tensor(cids)
+    for proj in ("w13", "w2"):
+        base = f"model.llm.layers.{L}.mlp.experts.{proj}"
+        is_nvfp4 = nv.has(f"{base}_weight.scale")
+        if is_nvfp4:
+            # HOT: NVFP4 packed slabs sliced verbatim.
+            packed = nv.get(f"{base}_weight"); scale = nv.get(f"{base}_weight.scale"); scale2 = nv.get(f"{base}_weight.scale2")
+            out[f"{base}_hot_weight"] = packed[hidx].contiguous()
+            out[f"{base}_hot_weight.scale"] = scale[hidx].contiguous()
+            out[f"{base}_hot_weight.scale2"] = scale2[hidx].contiguous()
+        else:
+            # HOT (layer 2): bf16 slabs gathered from the fused bf16 tensor, NO scales.
+            fused = nv.get(f"{base}_weight")  # bf16 [E, out, in]
+            out[f"{base}_hot_bf16"] = fused[hidx].contiguous().to(torch.bfloat16)
+        # COLD: AQLM, one code tensor per book with minimal lossless dtype.
+        part = torch.load(f"{CONV}/layer_{L}_{proj}.pt", map_location="cpu")  # v2: both w13+w2 converged
+        codes = part[f"{proj}_codes"][cidx]          # [n_cold, books, out, ng] int16
+        for b, cb in enumerate(part[f"{proj}_codebooks"]):
+            dt = book_code_dtype(cb.shape[0])
+            out[f"{base}_cold_codes.{b}"] = codes[:, b].to(dt).contiguous()
+            out[f"{base}_cold_codebook.{b}"] = cb.contiguous()
+        out[f"{base}_cold_scales"] = part[f"{proj}_scales"][cidx].contiguous()
+    out[f"model.llm.layers.{L}.mlp.experts.hot_ids"] = torch.tensor(hids, dtype=torch.int32)
+    out[f"model.llm.layers.{L}.mlp.experts.cold_ids"] = torch.tensor(cids, dtype=torch.int32)
+    return out
+
+def _dtype_str(dt):
+    return str(dt).split(".")[-1]
+
+def make_hybrid_config(hot, packed):
+    part_w13 = torch.load(f"{CONV}/layer_{min(packed)}_w13.pt", map_location="cpu")
+    part_w2  = torch.load(f"{CONV}/layer_{min(packed)}_w2.pt", map_location="cpu")
+    books = {"w13": [c.shape[0] for c in part_w13["w13_codebooks"]],
+             "w2":  [c.shape[0] for c in part_w2["w2_codebooks"]]}
+    code_dtypes = {p: [_dtype_str(book_code_dtype(e)) for e in books[p]] for p in ("w13", "w2")}
+    layer_books = {}
+    for L in sorted(hot):
+        nhot = len(hot[L]); ncold = ic.N_EXPERTS - nhot
+        packed_flag = L in packed
+        layer_books[str(L)] = {"n_hot": nhot, "n_nvfp4": nhot, "n_cold": ncold,
+                               "packed": packed_flag,
+                               "hot_format": "nvfp4" if packed_flag else "bf16"}
+    return {
+        "quant_method": "inkling_nvfp4_aqlm_hybrid",
+        "base_nvfp4": "thinkingmachines/Inkling-NVFP4",
+        "group_size": 8,
+        "w13_book_entries": books["w13"], "w2_book_entries": books["w2"],
+        "w13_code_dtypes": code_dtypes["w13"], "w2_code_dtypes": code_dtypes["w2"],
+        "cold_code_layout": "per-book tensors {proj}_cold_codes.{b}, dtype uint8 if book<=256 entries else int16",
+        "aqlm_layer_books": layer_books,
+        "note": "cold experts: AQLM, per-book code tensors {proj}_cold_codes.{b} (w2 book0=int16@65536, "
+                "book1=uint8@256; w13 book0=int16@65536) + per-book codebooks + per-row scales. "
+                "hot experts: NVFP4 (hot_format=nvfp4) except layer 2 which has NO NVFP4 base and stores "
+                "bf16 hot slabs {proj}_hot_bf16 (no scale/scale2, hot_format=bf16). Its cold experts are "
+                "AQLM like every other layer. hot_ids/cold_ids give the routed-expert index for each slab.",
+    }
+
+# ---------------------------------------------------------------------------
+# Full-checkpoint assembly (streaming, resumable, CPU/IO only).
+# ---------------------------------------------------------------------------
+import struct, shutil
+
+def is_dropped(name, built):
+    """True iff `name` is a fused MoE-expert weight/scale/scale2 of a layer that is
+    rebuilt as a hybrid layer (every MoE layer, incl. layer 2's bf16 experts).
+    input_amax/original_shape are NOT matched by EXP_ANY, so they are kept."""
+    m = EXP_ANY.match(name)
+    return bool(m) and int(m.group(1)) in built
+
+def _read_st_header(path):
+    with open(path, "rb") as f:
+        n = struct.unpack("<Q", f.read(8))[0]
+        h = json.loads(f.read(n))
+    h.pop("__metadata__", None)
+    return h
+
+class LayerCache:
+    """Builds one packed layer's hybrid tensors on demand, caching the most recent."""
+    def __init__(self, hot, nv):
+        self.hot = hot; self.nv = nv; self._L = None; self._t = None
+    def get(self, L):
+        if self._L != L:
+            self._t = build_layer_tensors(L, self.hot, self.nv)
+            self._L = L
+        return self._t
+
+def plan_tensors(hot, built):
+    """Return ordered list of plan entries + weight_map skeleton.
+    entry = {name, nbytes, kind:'src'|'hyb', src_shard?, layer?}.
+    `built` = every MoE layer rebuilt as a hybrid layer (packed NVFP4 + layer 2)."""
+    src_idx = json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"]
+    # per-shard header cache for exact source byte sizes
+    hdr_cache = {}
+    def src_nbytes(name, shard):
+        if shard not in hdr_cache:
+            hdr_cache[shard] = _read_st_header(f"{SRC}/{shard}")
+        o = hdr_cache[shard][name]["data_offsets"]
+        return o[1] - o[0]
+
+    plan = []
+    # 1) all kept source tensors, in source weight_map order
+    for name, shard in src_idx.items():
+        if is_dropped(name, built):
+            continue
+        plan.append({"name": name, "nbytes": src_nbytes(name, shard),
+                     "kind": "src", "src_shard": shard})
+    # 2) hybrid expert tensors per rebuilt layer (build each once for exact sizes)
+    nv = ic.ShardReader(SRC)
+    for L in sorted(built):
+        t = build_layer_tensors(L, hot, nv)
+        for name, v in t.items():
+            plan.append({"name": name, "nbytes": v.numel() * v.element_size(),
+                         "kind": "hyb", "layer": L})
+        del t
+    return plan
+
+def bin_pack(plan, shard_bytes):
+    """Greedy pack preserving order; a tensor larger than shard_bytes gets its own shard."""
+    shards = []; cur = []; cur_b = 0
+    for e in plan:
+        if cur and cur_b + e["nbytes"] > shard_bytes:
+            shards.append(cur); cur = []; cur_b = 0
+        cur.append(e); cur_b += e["nbytes"]
+    if cur:
+        shards.append(cur)
+    return shards
+
+def assemble():
+    os.makedirs(OUT, exist_ok=True)
+    hot = hot_map(); packed = packed_layers()
+    built = set(hot)              # every MoE layer is now rebuilt as a hybrid layer
+    bf16 = sorted(built - packed) # layer(s) with no NVFP4 base -> bf16 hot slabs
+    print(f"[plan] hybrid MoE layers: {len(built)}  (NVFP4-packed: {len(packed)}, bf16-hot: {bf16})")
+    plan = plan_tensors(hot, built)
+    total = sum(e["nbytes"] for e in plan)
+    print(f"[plan] {len(plan)} output tensors, total {total/1e9:.1f} GB")
+    shards = bin_pack(plan, SHARD_BYTES)
+    NN = len(shards)
+    print(f"[plan] {NN} output shards (~{SHARD_BYTES/2**30:.0f}GB target)")
+
+    def shard_name(i):
+        return f"model-{i:05d}-of-{NN:05d}.safetensors"
+
+    # index.json (deterministic from plan)
+    weight_map = {}
+    for i, sh in enumerate(shards, 1):
+        for e in sh:
+            weight_map[e["name"]] = shard_name(i)
+    idx = {"metadata": {"total_size": total}, "weight_map": weight_map}
+    with open(f"{OUT}/model.safetensors.index.json", "w") as f:
+        json.dump(idx, f, indent=2)
+    print(f"[index] wrote model.safetensors.index.json ({len(weight_map)} tensors)")
+
+    # write shards (resumable: fully-written shards get renamed from .tmp)
+    src_readers = {}
+    def src_get(shard, name):
+        if shard not in src_readers:
+            src_readers[shard] = safe_open(f"{SRC}/{shard}", framework="pt")
+        return src_readers[shard].get_tensor(name)
+
+    lc = LayerCache(hot, ic.ShardReader(SRC))
+    for i, sh in enumerate(shards, 1):
+        fname = shard_name(i); fpath = f"{OUT}/{fname}"
+        exp_b = sum(e["nbytes"] for e in sh)
+        if os.path.exists(fpath):
+            print(f"[shard {i}/{NN}] {fname} exists, skip")
+            continue
+        tmp = fpath + ".tmp"
+        if os.path.exists(tmp):
+            os.remove(tmp)
+        tensors = {}
+        for e in sh:
+            if e["kind"] == "src":
+                tensors[e["name"]] = src_get(e["src_shard"], e["name"])
+            else:
+                t = lc.get(e["layer"])[e["name"]]
+                assert t.numel() * t.element_size() == e["nbytes"], \
+                    f"size mismatch {e['name']}: built {t.numel()*t.element_size()} vs planned {e['nbytes']}"
+                tensors[e["name"]] = t
+        save_file(tensors, tmp, metadata={"format": "pt"})
+        os.replace(tmp, fpath)
+        del tensors
+        print(f"[shard {i}/{NN}] wrote {fname}  {exp_b/1e9:.2f} GB  ({len(sh)} tensors)")
+
+    write_side_files(hot, packed)
+    print(f"[done] assembled {NN} shards, {total/1e9:.1f} GB -> {OUT}")
+    return NN, total
+
+def write_side_files(hot, packed):
+    # verbatim copies
+    for fn in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja",
+               "processor_config.json", "special_tokens_map.json"):
+        src = f"{SRC}/{fn}"
+        if os.path.exists(src):
+            shutil.copy2(src, f"{OUT}/{fn}")
+    if os.path.isdir(f"{SRC}/tiktoken"):
+        shutil.copytree(f"{SRC}/tiktoken", f"{OUT}/tiktoken", dirs_exist_ok=True)
+
+    # config.json + quantization_config block
+    cfg = json.load(open(f"{SRC}/config.json"))
+    cfg["quantization_config"] = {
+        "quant_method": "inkling_nvfp4_aqlm_hybrid",
+        "base_nvfp4": "thinkingmachines/Inkling-NVFP4",
+        "group_size": 8,
+        "note": "Mixed per-expert quant: hot experts ModelOpt-NVFP4, cold experts AQLM-2bit. "
+                "Full spec in hf_quant_config.json['aqlm_hybrid'] and README.md. "
+                "Requires a custom fused-MoE kernel; does NOT load on stock vLLM.",
+    }
+    with open(f"{OUT}/config.json", "w") as f:
+        json.dump(cfg, f, indent=2)
+
+    # hf_quant_config.json = original NVFP4 + aqlm_hybrid block
+    hq = json.load(open(f"{SRC}/hf_quant_config.json"))
+    hq["aqlm_hybrid"] = make_hybrid_config(hot, packed)
+    with open(f"{OUT}/hf_quant_config.json", "w") as f:
+        json.dump(hq, f, indent=2)
+
+    # reproducibility: tools as code/
+    os.makedirs(f"{OUT}/code", exist_ok=True)
+    for py in os.listdir(os.path.dirname(os.path.abspath(__file__))):
+        if py.endswith(".py"):
+            shutil.copy2(os.path.join(os.path.dirname(os.path.abspath(__file__)), py),
+                         f"{OUT}/code/{py}")
+
+    write_readme(hot, packed)
+    print("[side] wrote config.json, hf_quant_config.json, tokenizer/processor files, code/, README.md")
+
+def write_readme(hot, packed):
+    built = set(hot); bf16 = sorted(built - packed)
+    tot_hot = sum(len(hot[L]) for L in built); tot_slots = len(built) * ic.N_EXPERTS
+    pct_hot = 100.0 * tot_hot / tot_slots
+    readme = README_TEMPLATE.format(
+        n_packed=len(packed), n_bf16=len(bf16), bf16_layers=", ".join(map(str, bf16)),
+        tot_hot=tot_hot, tot_slots=tot_slots,
+        pct_hot=f"{pct_hot:.1f}", pct_cold=f"{100-pct_hot:.1f}")
+    with open(f"{OUT}/README.md", "w") as f:
+        f.write(readme)
+
+README_TEMPLATE = r"""---
+license: other
+base_model:
+- thinkingmachines/Inkling
+tags:
+- moe
+- nvfp4
+- aqlm
+- quantized
+- inkling
+---
+
+# Inkling-512k-NVFP4-AQLM-hybrid
+
+A memory-optimized **hybrid-quantized** build of
+[thinkingmachines/Inkling](https://huggingface.co/thinkingmachines/Inkling)
+(MoE, 256 routed experts/layer, 6 active). Each MoE layer's experts are split
+into a small **hot** set kept at NVFP4 (~4 bpw) and a large **cold** set stored
+as **AQLM 2-bit**. This is the **512k tier**: **{pct_hot}% of expert slots are
+NVFP4 (hot)** and **{pct_cold}% are AQLM-2bit (cold)**, averaging **~2.66 bits
+per weight** over the routed experts. It targets serving at **512k context on
+4x96GB GPUs**.
+
+> **This checkpoint needs a CUSTOM FUSED MoE KERNEL** that dispatches NVFP4 hot
+> experts and AQLM cold experts within the same fused grouped-GEMM. **It does
+> NOT load on stock vLLM / Transformers.** The on-disk format and dequant math
+> are fully specified below so a kernel can consume it directly.
+
+## Quality
+Measured **+1.0% perplexity vs Inkling-NVFP4** on the gate-proxy eval — i.e. a
+1% relative degradation for a large drop in expert memory.
+
+## Provenance
+- **Hot experts (NVFP4):** sliced verbatim from the ModelOpt NVFP4 base
+  (`thinkingmachines/Inkling-NVFP4`), no requantization.
+- **Cold experts (AQLM):** quantized from the **Inkling BF16 teacher**. The
+  `w13` cold books are **activation-aware converged** (calibrated on captured
+  activations); `w2` cold books are AQLM-encoded.
+- **Assignment** of hot vs cold per layer is in
+  `hf_quant_config.json['aqlm_hybrid']['aqlm_layer_books']` and reproduced from
+  the routing/importance study (`assign_512k.json`). Reproduction scripts are in
+  `code/`.
+
+## Layout summary
+- **{n_packed} NVFP4-packed MoE layers** carry the hybrid expert format below
+  (NVFP4 hot slabs + AQLM cold).
+- **{n_bf16} bf16-special MoE layer(s)** (layer **{bf16_layers}**) have **no NVFP4
+  base**, so they use the SAME hybrid format EXCEPT the hot slabs are stored as
+  **bf16** (`experts.{{proj}}_hot_bf16`, no `.scale`/`.scale2`); their cold experts
+  are AQLM like every other layer. The fused uncompressed bf16 `experts.w13_weight`
+  / `experts.w2_weight` are **dropped**. Layers 0-1 are dense MLP (not experts).
+- All **non-expert** tensors (attention, embeddings, norms, `shared_experts`,
+  `gate`, vision, audio, MTP) are copied verbatim from the NVFP4 base.
+- Total hot slots {tot_hot} / {tot_slots} expert slots.
+
+## On-disk tensor format (per NVFP4-packed MoE layer L)
+`prefix = model.llm.layers.{{L}}.mlp.experts`. Routed experts (256 total) are
+partitioned into `n_hot` hot + `n_cold` cold. For each projection
+`proj in {{w13, w2}}` (w13 = fused gate|up, out=6144, in=6144; w2 = down,
+out=6144, in=3072):
+
+**Hot — NVFP4-packed layers (ModelOpt, `hot_format=nvfp4`):**
+| tensor | dtype | shape |
+|---|---|---|
+| `{{prefix}}.{{proj}}_hot_weight`        | uint8       | `[n_hot, out, in/2]`  (2 fp4 codes/byte) |
+| `{{prefix}}.{{proj}}_hot_weight.scale`  | float8_e4m3 | `[n_hot, out, in/16]` (per-block-16 scale) |
+| `{{prefix}}.{{proj}}_hot_weight.scale2` | float32     | `[n_hot]`             (per-expert global scale) |
+
+**Hot — bf16-special layer(s) (`hot_format=bf16`, e.g. layer 2):**
+| tensor | dtype | shape |
+|---|---|---|
+| `{{prefix}}.{{proj}}_hot_bf16` | bfloat16 | `[n_hot, out, in]` (dense bf16, no scales) |
+
+**Cold (AQLM) — same for every layer. One code tensor PER BOOK, smallest lossless dtype:**
+| tensor | dtype | shape |
+|---|---|---|
+| `{{prefix}}.{{proj}}_cold_codes.{{b}}`    | uint8 if book has <=256 entries else int16 | `[n_cold, out, in/8]` (group_size=8) |
+| `{{prefix}}.{{proj}}_cold_codebook.{{b}}` | float16 | `[entries_b, 8]` |
+| `{{prefix}}.{{proj}}_cold_scales`     | float16 | `[n_cold, out]` (per-output-row scale) |
+
+Per-book code dtypes: `w13` book0 = **int16** (65536 entries); `w2` book0 =
+**int16** (65536), `w2` book1 = **uint8** (256 entries, codes 0..255 — lossless).
+`{{proj}}_cold_codes.{{b}}` replaces the old single `{{proj}}_cold_codes` tensor
+(the book axis is now the trailing `.{{b}}` suffix), saving ~38 GB losslessly.
+
+**Index maps:**
+| tensor | dtype | shape | meaning |
+|---|---|---|---|
+| `{{prefix}}.hot_ids`  | int32 | `[n_hot]`  | routed-expert index (0..255) of each hot slab |
+| `{{prefix}}.cold_ids` | int32 | `[n_cold]` | routed-expert index (0..255) of each cold slab |
+
+`hot_ids[i]` gives the original routed-expert id for hot slab `i` (and likewise
+`cold_ids`); together they partition `0..255`. **group_size = 8.**
+- `w13`: **1 book**, `entries = [65536]`.
+- `w2`:  **2 books**, `entries = [65536, 256]`.
+
+## Dequant math
+**Cold (AQLM), expert e, output row r, group g (8 weights):**
+```
+acc[8] = 0
+for b in range(n_books):
+    acc += codebook_b[ codes_b[e, r, g] ]           # codes_b = cold_codes.{{b}}; gather an 8-vector
+W[e, r, g*8:(g+1)*8] = acc * cold_scales[e, r]        # per-row scale
+```
+Codes are read from the per-book tensor `{{proj}}_cold_codes.{{b}}`. uint8 books
+(<=256 entries) index the codebook directly. For the int16 65536-entry book,
+negative int16 values wrap to the intended unsigned row (`idx & 0xFFFF`).
+
+**Hot (NVFP4, standard ModelOpt e2m1):**
+```
+codes = unpack_nibbles(hot_weight)                    # 0..15 -> e2m1 LUT value
+W = LUT[codes] * scale.repeat_interleave(16, dim=-1) * scale2
+```
+e2m1 LUT = `[0, .5, 1, 1.5, 2, 3, 4, 6, -0, -.5, -1, -1.5, -2, -3, -4, -6]`;
+low nibble is the first weight of each pair.
+
+## Files
+- `model-*.safetensors` + `model.safetensors.index.json` — the sharded weights.
+- `hf_quant_config.json` — original NVFP4 config **plus** the `aqlm_hybrid` block
+  (`quant_method`, per-layer `{{n_nvfp4, n_cold, packed}}`, book entry counts).
+- `config.json` — base config **plus** a `quantization_config` block advertising
+  `quant_method = inkling_nvfp4_aqlm_hybrid`.
+- `code/` — the exact repackaging scripts (`build_hybrid.py`, `ink_common.py`, ...).
+"""
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layer", type=int, default=-1, help="build+inspect one layer (test)")
+    ap.add_argument("--assemble", action="store_true", help="assemble full checkpoint -> OUT")
+    ap.add_argument("--validate", action="store_true", help="validate assembled checkpoint")
+    a = ap.parse_args()
+    hot = hot_map(); packed = packed_layers()
+    if a.assemble:
+        assemble()
+    elif a.validate:
+        from validate_hybrid import validate
+        validate()
+    elif a.layer >= 0:
+        nv = ic.ShardReader(SRC)
+        t = build_layer_tensors(a.layer, hot, nv)
+        print(f"layer {a.layer} hybrid tensors ({len(t)}):")
+        for k, v in t.items():
+            print(f"  {k.split('.mlp.')[-1]}: {v.dtype} {tuple(v.shape)}")
+        tot = sum(v.numel()*v.element_size() for v in t.values())
+        print(f"layer bytes: {tot/1e6:.1f} MB")
+    else:
+        print("packed MoE layers:", len(packed), "| bf16-special layers:", sorted(set(hot)-packed))
+        print(json.dumps(make_hybrid_config(hot, packed)["aqlm_layer_books"]["3"]))

code/build_mm_manifest.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""Assemble the v3 multimodal IMAGE manifest from all staged image pools, with a
+disjoint train/held split. Emits jsonl of {"path","domain","prompt"} that
+run_capture_acts_mm.py consumes (image+text -> vLLM -> MoE acts).
+
+Domains (each exercises different vision routing):
+  chest_xray  medical radiograph      (NIH sample)
+  fracatlas   fracture radiograph
+  natural     general/natural imagery
+  ocr         receipts / scanned text
+  screenshot  desktop/web/mobile GUI  (computer-use)
+"""
+import json, glob, os, random
+random.seed(1337)
+MM = "/data/inkling-mm"
+OUT = "/data/inkling-quant"
+PROMPTS = {
+  "chest_xray": "<|content_image|>\nDescribe the findings in this chest radiograph.",
+  "fracatlas":  "<|content_image|>\nIs there a fracture in this radiograph? Describe it.",
+  "natural":    "<|content_image|>\nDescribe this image in detail.",
+  "ocr":        "<|content_image|>\nTranscribe all text visible in this image.",
+  "screenshot": "<|content_image|>\nDescribe this UI screenshot and the actions available.",
+}
+POOLS = {
+  "chest_xray": [f"{MM}/chest_xray/**/*.png", f"{MM}/chest_xray/**/*.jpeg", f"{MM}/chest_xray/**/*.jpg"],
+  "fracatlas":  [f"{MM}/fracatlas/**/*.jpg", f"{MM}/fracatlas/**/*.png"],
+  "natural":    [f"{MM}/natural/**/*.jpg", f"{MM}/natural/**/*.png"],
+  "ocr":        [f"{MM}/ocr/**/*.jpg", f"{MM}/ocr/**/*.png", f"{MM}/ocr_hf/**/*.jpg"],
+  "screenshot": [f"{MM}/screenshots/**/*.jpg", f"{MM}/screenshots/**/*.png"],
+}
+PER_DOMAIN = 600      # cap per domain for train (balanced; not ham)
+HELD = 40             # per domain held-out for mm ppl
+
+def files(pats):
+    fs=[]
+    for p in pats: fs+=glob.glob(p, recursive=True)
+    return sorted(set(fs))
+
+def main():
+    train, held = [], []
+    for dom, pats in POOLS.items():
+        fs = files(pats); random.shuffle(fs)
+        if not fs:
+            print(f"  {dom:11s}: 0 (missing)"); continue
+        h = fs[:HELD]; t = fs[HELD:HELD+PER_DOMAIN]
+        for f in t: train.append({"path":f,"domain":dom,"prompt":PROMPTS[dom]})
+        for f in h: held.append({"path":f,"domain":dom,"prompt":PROMPTS[dom]})
+        print(f"  {dom:11s}: {len(fs)} avail -> train {len(t)} held {len(h)}")
+    random.shuffle(train)
+    for name, rows in (("mm_img_train.jsonl",train),("mm_img_held.jsonl",held)):
+        with open(f"{OUT}/{name}","w") as fh:
+            for r in rows: fh.write(json.dumps(r)+"\n")
+        print(f"  wrote {OUT}/{name}: {len(rows)}")
+
+if __name__=="__main__":
+    print(">>> building multimodal image manifest ..."); main()

code/converge.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Activation-aware AQLM convergence (Phase E1) for Inkling cold experts, w13.
+
+Objective per layer over cold experts (shared codebook):
+  min  sum_{e,o,g}  s[e,o]^2 * || (v[e,o,g] - cb[code]) * sqrt(h_e[g]) ||^2
+where v = W/s (row-normalized teacher), h_e = diag input Hessian = E[x^2] over
+tokens routed to expert e (from /data/inkling-quant/acts). Warm-starts from the
+encode part; iterates weighted re-encode -> weighted codebook update -> scale
+refit. w13 input is the MoE hidden x directly (no gate/up interleaving needed).
+
+Writes /data/inkling-quant/parts_conv/layer_{L}_w13.pt (same schema as encode
+part) + records rel_err before/after (unweighted, weight-space).
+"""
+from __future__ import annotations
+import argparse, os, time, json
+import torch
+import ink_common as ic
+
+import os as _os
+PARTS = _os.environ.get("INK_PARTS", "/data/inkling-quant/parts")
+OUT = _os.environ.get("INK_OUT", "/data/inkling-quant/parts_conv")
+ACTS = _os.environ.get("INK_ACTS", "/data/inkling-quant/acts")
+ASSIGN = _os.environ.get("INK_ASSIGN", "/data/inkling-quant/assign_512k.json")
+G = 8; CHUNK = 16384
+
+def cold_experts(layer, assign):
+    hot = set(assign["hot_experts"].get(str(layer), []))
+    return [e for e in range(ic.N_EXPERTS) if e not in hot]
+
+def w13_hessians(layer, cold, device):
+    a = torch.load(f"{ACTS}/acts_layer_{layer}.pt")
+    x = a["x"].float().to(device); ids = a["ids"].to(device)   # x[N,6144], ids[N,6]
+    H = torch.ones(ic.N_EXPERTS, 6144, device=device)
+    for e in cold:
+        m = (ids == e).any(1)
+        if int(m.sum()) >= 8:
+            xe = x[m]
+            H[e] = (xe * xe).mean(0)
+    H.clamp_min_(1e-8)
+    return H
+
+def chunk_weight(st, n, out_dim, ng, hsqrt):
+    """Per-chunk weight[n,8] = sqrt(h_e[group cols]); depends only on (e,g)."""
+    idx = torch.arange(st, st + n, device=hsqrt.device)
+    e = idx // (out_dim * ng); g = idx % ng
+    return hsqrt[e, g]                                   # [n,8]
+
+def wencode(vecs, cb, out_dim, ng, hsqrt):
+    """vecs[N,8], cb[K,8] -> codes[N] int32 (weighted nearest), weights on the fly."""
+    cbh = cb.half(); cb2h = (cb * cb).half()
+    out = torch.empty(vecs.shape[0], dtype=torch.int32, device=vecs.device)
+    for s in range(0, vecs.shape[0], CHUNK):
+        v = vecs[s:s+CHUNK]; w = chunk_weight(s, v.shape[0], out_dim, ng, hsqrt)
+        score = 2.0 * ((v * w) @ cbh.t()) - (w @ cb2h.t())   # [chunk,K]
+        out[s:s+CHUNK] = score.argmax(-1).to(torch.int32)
+    return out
+
+def rel_err(vecs, recon, roww):
+    # weight-space unweighted rel err uses row scale only (compare vs encode part metric)
+    num = torch.zeros((), device=vecs.device); den = torch.zeros((), device=vecs.device)
+    for s in range(0, vecs.shape[0], CHUNK*8):
+        v = vecs[s:s+CHUNK*8].float(); r = recon[s:s+CHUNK*8].float()
+        num += (v - r).pow(2).sum(); den += v.pow(2).sum()
+    return float((num.sqrt() / den.sqrt()))
+
+def werr(vecs, codes, cb, out_dim, ng, hsqrt, s2):
+    """Weighted (activation-aware) rel err = sqrt(sum w(v-r)^2)/sqrt(sum w v^2),
+    w = h(e,g)*s^2(e,o) — the actual converge objective."""
+    num = torch.zeros((), device=vecs.device); den = torch.zeros((), device=vecs.device)
+    for st in range(0, vecs.shape[0], CHUNK*8):
+        n = min(CHUNK*8, vecs.shape[0]-st)
+        idx = torch.arange(st, st+n, device=vecs.device)
+        e = idx // (out_dim*ng); o = (idx // ng) % out_dim; g = idx % ng
+        w = (hsqrt[e, g].float()**2) * s2[e, o].unsqueeze(1)   # [n,8]
+        v = vecs[st:st+n].float(); r = cb[codes[st:st+n].long()].float()
+        num += (w*(v-r).pow(2)).sum(); den += (w*v.pow(2)).sum()
+    return float((num.sqrt()/den.sqrt()))
+
+def converge_layer(layer, iters, device="cuda"):
+    assign = json.load(open(ASSIGN))
+    cold = cold_experts(layer, assign)
+    part = torch.load(f"{PARTS}/layer_{layer}_w13.pt")
+    cb = part["w13_codebooks"][0].float().to(device)          # [K,8]
+    K = cb.shape[0]
+    tc = ic.ShardReader(ic.BF16_DIR)
+    # teacher cold experts -> normalized vecs + per-row scale, built PER-EXPERT to
+    # avoid a full [C,6144,6144] float copy (OOM). vecs half, s fp32.
+    C = len(cold); out = 6144; ind = 6144; ng = ind // G
+    s = torch.empty(C, out, device=device)
+    vecs = torch.empty(C * out * ng, G, dtype=torch.float16, device=device)
+    for ci, e in enumerate(cold):
+        we = ic.teacher_expert(tc, layer, "w13", e, device=device).float()   # [6144,6144]
+        se = we.abs().mean(-1).clamp_min(1e-8)                                # [6144]
+        s[ci] = se
+        base = ci * out * ng
+        vecs[base:base + out*ng] = (we / se.unsqueeze(-1)).reshape(out*ng, G).half()
+        del we
+    torch.cuda.empty_cache()
+    # compact hessian tables (no full [Nv,8] materialization)
+    H = w13_hessians(layer, cold, device)                      # [256,6144]
+    hcold = H[cold].reshape(C, ng, G)                          # [C,ng,8]  (= h per e,group)
+    hsqrt = torch.sqrt(hcold).half()                           # [C,ng,8]
+    s2 = (s * s)                                               # [C,out]  row weight
+    # init recon rel err (warm start = encode part codes)
+    codes0 = part["w13_codes"][:, 0][torch.tensor(cold)].reshape(-1).to(device).long()
+    err_before = rel_err(vecs, cb[codes0].half(), None)
+    werr_before = werr(vecs, codes0, cb, out, ng, hsqrt, s2)
+    print(f"  L{layer} w13 warm-start: unweighted {err_before:.4f} | WEIGHTED {werr_before:.4f}", flush=True)
+    codes = codes0; err = err_before
+    for it in range(iters):
+        codes = wencode(vecs, cb, out, ng, hsqrt)              # weighted re-encode
+        # weighted codebook update: c_k = sum(rw*v)/sum(rw), rw = h(e,g) * s^2(e,o), per chunk
+        newcb = torch.zeros(K, G, device=device); wsum = torch.zeros(K, G, device=device)
+        for st in range(0, vecs.shape[0], CHUNK*8):
+            n = min(CHUNK*8, vecs.shape[0]-st)
+            idx = torch.arange(st, st+n, device=device)
+            e = idx // (out*ng); o = (idx // ng) % out; g = idx % ng
+            r = (hsqrt[e, g].float() ** 2) * s2[e, o].unsqueeze(1)   # [n,8] = h * s^2
+            c = codes[st:st+n].long(); v = vecs[st:st+n].float()
+            newcb.index_add_(0, c, v * r); wsum.index_add_(0, c, r)
+        nz = wsum[:, 0] > 0
+        cb[nz] = (newcb[nz] / wsum[nz].clamp_min(1e-12))
+        err = rel_err(vecs, cb[codes].half(), None)
+        we = werr(vecs, codes, cb, out, ng, hsqrt, s2)
+        print(f"  L{layer} w13 iter{it}: unweighted {err:.4f} | WEIGHTED {we:.4f}", flush=True)
+    # refit per-row scale s (closed form, h-weighted): s = <W, s0*recon_h>/<recon,recon>_h ... approx: keep s (rows already normalized). recompute scales via weighted lsq on original W
+    # (scale refit folded into part's existing per-row scale; recon already in normalized space)
+    codes_r = codes.reshape(C, 1, out, ng).to(torch.int16).cpu()
+    # write full-256 arrays: cold updated, hot slots keep encode-part values (unused for hot)
+    full_codes = part["w13_codes"].clone()
+    full_codes[torch.tensor(cold), 0] = codes_r[:, 0]
+    return {
+        "expert_ids": part["expert_ids"],
+        "w13_codes": full_codes,
+        "w13_codebooks": [cb.half().cpu()],
+        "w13_book_entries": part["w13_book_entries"],
+        "w13_scales": part["w13_scales"],
+        "cold_experts": cold,
+        "rel_err_before": err_before, "rel_err_after": err,
+    }
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layer", type=int, required=True)
+    ap.add_argument("--iters", type=int, default=2)
+    ap.add_argument("--save", action="store_true")
+    a = ap.parse_args()
+    os.makedirs(OUT, exist_ok=True)
+    t0 = time.time()
+    r = converge_layer(a.layer, a.iters)
+    print(f"L{a.layer} w13: {r['rel_err_before']:.4f} -> {r['rel_err_after']:.4f}  ({time.time()-t0:.0f}s)")
+    if a.save:
+        torch.save(r, f"{OUT}/layer_{a.layer}_w13.pt"); print("saved.")

code/converge_all.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""Full w13 activation-aware converge (Phase E1) across all MoE layers.
+Shards layers across CUDA_VISIBLE_DEVICES GPUs, 1 spawn-worker/GPU, resumable
+(skips layers already in parts_conv/). 2 iters (gains plateau after iter0).
+Run: CUDA_VISIBLE_DEVICES=4,5,6,7 python converge_all.py
+"""
+import os, time
+import torch.multiprocessing as mp
+import ink_common as ic
+
+import os as _os
+OUT = _os.environ.get("INK_OUT", "/data/inkling-quant/parts_conv")
+ITERS = 2
+
+def worker(gpu, layers):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    import torch, converge as cv
+    for L in layers:
+        out = f"{OUT}/layer_{L}_w13.pt"
+        if os.path.exists(out):
+            print(f"[gpu{gpu}] skip L{L} (exists)", flush=True); continue
+        t0 = time.time()
+        try:
+            r = cv.converge_layer(L, ITERS, device="cuda")
+            torch.save(r, out)
+            print(f"[gpu{gpu}] L{L} w13 {r['rel_err_before']:.4f}->{r['rel_err_after']:.4f} "
+                  f"({time.time()-t0:.0f}s) saved", flush=True)
+        except Exception as ex:
+            print(f"[gpu{gpu}] L{L} FAILED: {ex}", flush=True)
+
+def main():
+    os.makedirs(OUT, exist_ok=True)
+    gpus = [int(x) for x in os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")]
+    layers = ic.moe_layers()
+    todo = [L for L in layers if not os.path.exists(f"{OUT}/layer_{L}_w13.pt")]
+    print(f"{len(layers)} layers, {len(todo)} remaining, {len(gpus)} GPUs {gpus}", flush=True)
+    shards = {g: [] for g in gpus}
+    for i, L in enumerate(todo):
+        shards[gpus[i % len(gpus)]].append(L)
+    mp.set_start_method("spawn", force=True)
+    ps = [mp.Process(target=worker, args=(g, shards[g])) for g in gpus]
+    for p in ps: p.start()
+    for p in ps: p.join()
+    done = sum(1 for L in layers if os.path.exists(f"{OUT}/layer_{L}_w13.pt"))
+    print(f"CONVERGE DONE: {done}/{len(layers)} layers", flush=True)
+
+if __name__ == "__main__":
+    main()

code/converge_w2.py

@@ -0,0 +1,238 @@
+#!/usr/bin/env python3
+"""Activation-aware AQLM convergence (Phase E1) for Inkling cold experts, w2.
+
+Mirrors the validated w13 converger (converge.py) with the TWO w2 differences:
+
+1. HESSIAN INPUT is the per-expert intermediate (down-proj input), not the MoE
+   hidden x. For each cold expert e with routed calib tokens x_e [n_e,6144]:
+       gate_e = teacher_w13_e[0::2]  ([3072,6144])   # row-interleaved gate|up
+       up_e   = teacher_w13_e[1::2]  ([3072,6144])
+       inter  = silu(gate_e @ x_e^T) * (up_e @ x_e^T)      -> [3072, n_e]
+       h_e    = mean_tokens(inter^2)                       -> [3072]  (clamp 1e-8)
+   w2 teacher weight is [6144,3072] (out=6144, in=3072) so groups-of-8 run along
+   the 3072 input dim -> ng=384.
+
+2. TWO BOOKS.  v = W/s (row-normalized teacher) is reconstructed additively as
+   cb0[c0] + cb1[c1]  (book0 @65536 encodes v, book1 @256 encodes the residual
+   v - cb0[c0]).  Coordinate descent per iter:
+       (a) c0 = weighted re-encode of v against cb0
+       (b) resid = v - cb0[c0]
+       (c) c1 = weighted re-encode of resid against cb1
+       (d) update cb0 (weighted mean of v over c0), then cb1 (weighted mean of
+           v - cb0_new[c0] over c1);  weight = h_e[group cols] * s^2 per row.
+
+Objective (same as w13): min sum_{e,o,g} s[e,o]^2 * || (v - recon) * sqrt(h_e[g]) ||^2.
+Warm-starts codes+codebooks from parts/layer_{L}_w2.pt.  Only cold experts are
+converged (per assign_512k.json complement); hot slots keep encode values.
+
+Writes /data/inkling-quant/parts_conv/layer_{L}_w2.pt (same schema as the encode
+part) + rel_err_before/after (unweighted) and werr_before/after (weighted).
+
+NOTE vs w13: wencode / codebook-update / err reductions run in fp32 here (w13
+used fp16). w2's Hessian is E[intermediate^2] which is far larger in magnitude
+than w13's E[x^2]; fp32 in the h-touching ops avoids fp16 overflow->nan. This
+does not change the objective, only the arithmetic precision.
+"""
+from __future__ import annotations
+import argparse, os, time, json
+import torch
+import torch.nn.functional as F
+import ink_common as ic
+
+import os as _os
+PARTS = _os.environ.get("INK_PARTS", "/data/inkling-quant/parts")
+OUT = _os.environ.get("INK_OUT", "/data/inkling-quant/parts_conv")
+ACTS = _os.environ.get("INK_ACTS", "/data/inkling-quant/acts")
+ASSIGN = _os.environ.get("INK_ASSIGN", "/data/inkling-quant/assign_512k.json")
+G = 8; CHUNK = 16384
+OUTD = 6144; IND = 3072; NG = IND // G   # w2: out=6144, in=3072, ng=384
+
+
+def cold_experts(layer, assign):
+    hot = set(assign["hot_experts"].get(str(layer), []))
+    return [e for e in range(ic.N_EXPERTS) if e not in hot]
+
+
+def chunk_weight(st, n, out_dim, ng, hsqrt):
+    """Per-chunk sqrt-weight[n,8] = sqrt(h_e[group cols]); depends only on (e,g)."""
+    idx = torch.arange(st, st + n, device=hsqrt.device)
+    e = idx // (out_dim * ng); g = idx % ng
+    return hsqrt[e, g]                                   # [n,8]
+
+
+def wencode(vecs, cb, out_dim, ng, hsqrt, resid_cb=None, resid_codes=None):
+    """codes[N] int32 = weighted nearest of (vecs - resid) against cb.
+    If resid_cb/resid_codes given, target = vecs - resid_cb[resid_codes] (book1).
+    Weights (sqrt(h)) applied on the fly; fp32 to stay overflow-safe for w2."""
+    cbf = cb.float(); cb2 = (cbf * cbf)
+    out = torch.empty(vecs.shape[0], dtype=torch.int32, device=vecs.device)
+    for s in range(0, vecs.shape[0], CHUNK):
+        v = vecs[s:s + CHUNK].float()
+        if resid_cb is not None:
+            v = v - resid_cb[resid_codes[s:s + CHUNK].long()].float()
+        w = chunk_weight(s, v.shape[0], out_dim, ng, hsqrt).float()   # [n,8]
+        score = 2.0 * ((v * w) @ cbf.t()) - (w @ cb2.t())             # [n,K]
+        out[s:s + CHUNK] = score.argmax(-1).to(torch.int32)
+    return out
+
+
+def _wupdate(vecs, codes, K, out_dim, ng, hsqrt, s2, sub_cb=None, sub_codes=None):
+    """Weighted codebook update: cb_k = sum(r*t)/sum(r), r = h(e,g)*s^2(e,o),
+    t = vecs - sub_cb[sub_codes] (residual target for book1; vecs for book0).
+    Returns new cb [K,8]; entries with no assigned weight left at 0 (caller keeps
+    the old row)."""
+    newcb = torch.zeros(K, G, device=vecs.device)
+    wsum = torch.zeros(K, G, device=vecs.device)
+    for st in range(0, vecs.shape[0], CHUNK * 8):
+        n = min(CHUNK * 8, vecs.shape[0] - st)
+        idx = torch.arange(st, st + n, device=vecs.device)
+        e = idx // (out_dim * ng); o = (idx // ng) % out_dim; g = idx % ng
+        r = (hsqrt[e, g].float() ** 2) * s2[e, o].unsqueeze(1)        # [n,8] = h*s^2
+        t = vecs[st:st + n].float()
+        if sub_cb is not None:
+            t = t - sub_cb[sub_codes[st:st + n].long()].float()
+        c = codes[st:st + n].long()
+        newcb.index_add_(0, c, t * r); wsum.index_add_(0, c, r)
+    return newcb, wsum
+
+
+def _recon_chunk(c0, c1, cb0, cb1):
+    return cb0[c0.long()] + cb1[c1.long()]
+
+
+def rel_err2(vecs, c0, c1, cb0, cb1):
+    """Unweighted weight-space rel err of the 2-book recon (chunked, fp32)."""
+    num = torch.zeros((), device=vecs.device); den = torch.zeros((), device=vecs.device)
+    for s in range(0, vecs.shape[0], CHUNK * 8):
+        n = min(CHUNK * 8, vecs.shape[0] - s)
+        v = vecs[s:s + n].float()
+        r = _recon_chunk(c0[s:s + n], c1[s:s + n], cb0, cb1).float()
+        num += (v - r).pow(2).sum(); den += v.pow(2).sum()
+    return float((num.sqrt() / den.sqrt()))
+
+
+def werr2(vecs, c0, c1, cb0, cb1, out_dim, ng, hsqrt, s2):
+    """Weighted (activation-aware) rel err of the 2-book recon = the objective:
+    w = h(e,g)*s^2(e,o). Scale-invariant to any global rescale of h."""
+    num = torch.zeros((), device=vecs.device); den = torch.zeros((), device=vecs.device)
+    for st in range(0, vecs.shape[0], CHUNK * 8):
+        n = min(CHUNK * 8, vecs.shape[0] - st)
+        idx = torch.arange(st, st + n, device=vecs.device)
+        e = idx // (out_dim * ng); o = (idx // ng) % out_dim; g = idx % ng
+        w = (hsqrt[e, g].float() ** 2) * s2[e, o].unsqueeze(1)        # [n,8]
+        v = vecs[st:st + n].float()
+        r = _recon_chunk(c0[st:st + n], c1[st:st + n], cb0, cb1).float()
+        num += (w * (v - r).pow(2)).sum(); den += (w * v.pow(2)).sum()
+    return float((num.sqrt() / den.sqrt()))
+
+
+def build_w2(layer, cold, device):
+    """Per-cold-expert: normalized vecs [C*out*ng,8] half, per-row scale s [C,out],
+    and diagonal input Hessian hsqrt [C,ng,8] (sqrt of E[intermediate^2]).
+    Built per expert to avoid a full float teacher copy (OOM)."""
+    a = torch.load(f"{ACTS}/acts_layer_{layer}.pt")
+    x = a["x"].to(device).float()                     # [N,6144]
+    ids = a["ids"].to(device)                         # [N,6]
+    tc = ic.ShardReader(ic.BF16_DIR)
+    C = len(cold)
+    s = torch.empty(C, OUTD, device=device)
+    vecs = torch.empty(C * OUTD * NG, G, dtype=torch.float16, device=device)
+    hsqrt = torch.empty(C, NG, G, device=device)
+    for ci, e in enumerate(cold):
+        # ---- weight vectors (teacher w2) ----
+        w2e = ic.teacher_expert(tc, layer, "w2", e, device=device).float()   # [6144,3072]
+        se = w2e.abs().mean(-1).clamp_min(1e-8)                               # [6144]
+        s[ci] = se
+        base = ci * OUTD * NG
+        vecs[base:base + OUTD * NG] = (w2e / se.unsqueeze(-1)).reshape(OUTD * NG, G).half()
+        del w2e
+        # ---- Hessian: intermediate = silu(gate@x^T)*(up@x^T), teacher w13 ----
+        m = (ids == e).any(1)
+        if int(m.sum()) >= 8:
+            xe = x[m]                                                        # [n_e,6144]
+            w13e = ic.teacher_expert(tc, layer, "w13", e, device=device).float()  # [6144,6144]
+            gate = w13e[0::2]; up = w13e[1::2]                               # [3072,6144] each
+            del w13e
+            g_out = gate @ xe.t(); u_out = up @ xe.t()                       # [3072,n_e]
+            inter = F.silu(g_out) * u_out                                    # [3072,n_e]
+            h_e = inter.pow(2).mean(1).clamp_min(1e-8)                       # [3072]
+            del xe, g_out, u_out, inter
+        else:
+            h_e = torch.ones(IND, device=device)
+        hsqrt[ci] = h_e.reshape(NG, G).sqrt()
+    del x, ids
+    torch.cuda.empty_cache()
+    return vecs, s, hsqrt
+
+
+def converge_layer(layer, iters, device="cuda"):
+    assign = json.load(open(ASSIGN))
+    cold = cold_experts(layer, assign)
+    part = torch.load(f"{PARTS}/layer_{layer}_w2.pt")
+    cb0 = part["w2_codebooks"][0].float().to(device)          # [65536,8]
+    cb1 = part["w2_codebooks"][1].float().to(device)          # [256,8]
+    K0 = cb0.shape[0]; K1 = cb1.shape[0]
+    C = len(cold)
+    coldt = torch.tensor(cold)
+
+    vecs, s, hsqrt = build_w2(layer, cold, device)
+    s2 = (s * s)                                              # [C,out]
+
+    # warm-start codes for cold experts (int16 -> positive index; book0 wraps mod 65536)
+    c0 = (part["w2_codes"][:, 0][coldt].reshape(-1).to(device).long() % K0).to(torch.int32)
+    c1 = (part["w2_codes"][:, 1][coldt].reshape(-1).to(device).long() % K1).to(torch.int32)
+
+    err_before = rel_err2(vecs, c0, c1, cb0, cb1)
+    werr_before = werr2(vecs, c0, c1, cb0, cb1, OUTD, NG, hsqrt, s2)
+    print(f"  L{layer} w2 warm-start: unweighted {err_before:.4f} | WEIGHTED {werr_before:.4f}", flush=True)
+
+    err = err_before; we = werr_before
+    for it in range(iters):
+        # (a) re-encode book0 against cb0 (target v)
+        c0 = wencode(vecs, cb0, OUTD, NG, hsqrt)
+        # (b)+(c) re-encode book1 against cb1 (target v - cb0[c0])
+        c1 = wencode(vecs, cb1, OUTD, NG, hsqrt, resid_cb=cb0, resid_codes=c0)
+        # (d) update cb0 (weighted mean of v over c0)
+        newcb0, wsum0 = _wupdate(vecs, c0, K0, OUTD, NG, hsqrt, s2)
+        nz0 = wsum0[:, 0] > 0
+        cb0[nz0] = newcb0[nz0] / wsum0[nz0].clamp_min(1e-12)
+        # then update cb1 (weighted mean of v - cb0_new[c0] over c1)
+        newcb1, wsum1 = _wupdate(vecs, c1, K1, OUTD, NG, hsqrt, s2, sub_cb=cb0, sub_codes=c0)
+        nz1 = wsum1[:, 0] > 0
+        cb1[nz1] = newcb1[nz1] / wsum1[nz1].clamp_min(1e-12)
+
+        err = rel_err2(vecs, c0, c1, cb0, cb1)
+        we = werr2(vecs, c0, c1, cb0, cb1, OUTD, NG, hsqrt, s2)
+        print(f"  L{layer} w2 iter{it}: unweighted {err:.4f} | WEIGHTED {we:.4f}", flush=True)
+
+    # reshape codes back to [C,out,ng] and store as int16 (book0 wraps to signed)
+    c0r = c0.reshape(C, OUTD, NG).to(torch.int16).cpu()
+    c1r = c1.reshape(C, OUTD, NG).to(torch.int16).cpu()
+    full_codes = part["w2_codes"].clone()                    # [256,2,out,ng]
+    full_codes[coldt, 0] = c0r
+    full_codes[coldt, 1] = c1r
+    return {
+        "expert_ids": part["expert_ids"],
+        "w2_codes": full_codes,
+        "w2_codebooks": [cb0.half().cpu(), cb1.half().cpu()],
+        "w2_book_entries": part["w2_book_entries"],
+        "w2_scales": part["w2_scales"],
+        "cold_experts": cold,
+        "rel_err_before": err_before, "rel_err_after": err,
+        "werr_before": werr_before, "werr_after": we,
+    }
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layer", type=int, required=True)
+    ap.add_argument("--iters", type=int, default=2)
+    ap.add_argument("--save", action="store_true")
+    a = ap.parse_args()
+    os.makedirs(OUT, exist_ok=True)
+    t0 = time.time()
+    r = converge_layer(a.layer, a.iters)
+    print(f"L{a.layer} w2: unw {r['rel_err_before']:.4f}->{r['rel_err_after']:.4f} | "
+          f"WEIGHTED {r['werr_before']:.4f}->{r['werr_after']:.4f}  ({time.time()-t0:.0f}s)")
+    if a.save:
+        torch.save(r, f"{OUT}/layer_{a.layer}_w2.pt"); print("saved.")

code/converge_w2_all.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""Full w2 activation-aware converge (Phase E1) across all MoE layers.
+Shards layers across CUDA_VISIBLE_DEVICES GPUs, 1 spawn-worker/GPU, resumable
+(skips layers already in parts_conv/). 2 iters.
+Run: CUDA_VISIBLE_DEVICES=4,5,6,7 python converge_w2_all.py
+"""
+import os, time
+import torch.multiprocessing as mp
+import ink_common as ic
+
+import os as _os
+OUT = _os.environ.get("INK_OUT", "/data/inkling-quant/parts_conv")
+ITERS = 2
+
+
+def worker(gpu, layers):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    import torch, converge_w2 as cv
+    for L in layers:
+        out = f"{OUT}/layer_{L}_w2.pt"
+        if os.path.exists(out):
+            print(f"[gpu{gpu}] skip L{L} (exists)", flush=True); continue
+        t0 = time.time()
+        try:
+            r = cv.converge_layer(L, ITERS, device="cuda")
+            tmp = out + ".tmp"
+            torch.save(r, tmp); os.replace(tmp, out)
+            print(f"[gpu{gpu}] L{L} w2 unw {r['rel_err_before']:.4f}->{r['rel_err_after']:.4f} "
+                  f"| WEIGHTED {r['werr_before']:.4f}->{r['werr_after']:.4f} "
+                  f"({time.time()-t0:.0f}s) saved", flush=True)
+        except Exception as ex:
+            import traceback
+            print(f"[gpu{gpu}] L{L} FAILED: {ex}", flush=True)
+            traceback.print_exc()
+
+
+def main():
+    os.makedirs(OUT, exist_ok=True)
+    gpus = [int(x) for x in os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")]
+    layers = ic.moe_layers()
+    todo = [L for L in layers if not os.path.exists(f"{OUT}/layer_{L}_w2.pt")]
+    print(f"{len(layers)} layers, {len(todo)} remaining, {len(gpus)} GPUs {gpus}", flush=True)
+    shards = {g: [] for g in gpus}
+    for i, L in enumerate(todo):
+        shards[gpus[i % len(gpus)]].append(L)
+    mp.set_start_method("spawn", force=True)
+    ps = [mp.Process(target=worker, args=(g, shards[g])) for g in gpus]
+    for p in ps: p.start()
+    for p in ps: p.join()
+    done = sum(1 for L in layers if os.path.exists(f"{OUT}/layer_{L}_w2.pt"))
+    print(f"CONVERGE_W2 DONE: {done}/{len(layers)} layers", flush=True)
+
+
+if __name__ == "__main__":
+    main()

code/encode_all.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+"""Encode ALL routed experts (AQLM) for every MoE layer, ship scheme
+(w13=[65536], w2=[65536,256]). Shards (layer,proj) jobs across the GPUs given
+in CUDA_VISIBLE_DEVICES, one worker process per GPU. Resumable: skips parts
+already on disk. Codebooks are shared per layer, so we encode all 256 experts;
+the checkpoint build later slices hot vs cold per assign_512k.json.
+Run: CUDA_VISIBLE_DEVICES=4,5,6,7 python encode_all.py
+"""
+import os, sys, time, json
+import torch
+import torch.multiprocessing as mp
+import ink_common as ic
+
+PARTS = "/data/inkling-quant/parts"
+BOOK = {"w13": [65536], "w2": [65536, 256]}
+
+def worker(gpu, jobs, ret):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    import torch
+    import ink_aqlm as aq
+    for (layer, proj) in jobs:
+        out = f"{PARTS}/layer_{layer}_{proj}.pt"
+        if os.path.exists(out):
+            print(f"[gpu{gpu}] skip L{layer} {proj} (exists)", flush=True); continue
+        t0 = time.time()
+        try:
+            r = aq.quantize_layer(layer, proj, BOOK[proj], device="cuda")
+            torch.save(r, out)
+            print(f"[gpu{gpu}] L{layer} {proj} rel_err={r['rel_err']:.4f} ({time.time()-t0:.0f}s) -> saved", flush=True)
+        except Exception as e:
+            print(f"[gpu{gpu}] L{layer} {proj} FAILED: {e}", flush=True)
+
+def main():
+    os.makedirs(PARTS, exist_ok=True)
+    gpus = [int(x) for x in os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")]
+    layers = ic.moe_layers()
+    jobs = [(L, p) for L in layers for p in ("w13", "w2")]
+    todo = [j for j in jobs if not os.path.exists(f"{PARTS}/layer_{j[0]}_{j[1]}.pt")]
+    print(f"{len(jobs)} total jobs, {len(todo)} remaining, {len(gpus)} GPUs {gpus}", flush=True)
+    # round-robin assign
+    shards = {g: [] for g in gpus}
+    for i, j in enumerate(todo):
+        shards[gpus[i % len(gpus)]].append(j)
+    mp.set_start_method("spawn", force=True)
+    procs = []
+    for g in gpus:
+        p = mp.Process(target=worker, args=(g, shards[g], None)); p.start(); procs.append(p)
+    for p in procs: p.join()
+    done = sum(1 for j in jobs if os.path.exists(f"{PARTS}/layer_{j[0]}_{j[1]}.pt"))
+    print(f"DONE: {done}/{len(jobs)} parts present", flush=True)
+
+if __name__ == "__main__":
+    main()

code/gpqa40.py

@@ -0,0 +1,35 @@
+import glob, re, json, time, concurrent.futures as cf, urllib.request
+import pyarrow.parquet as pq
+K="sk-98f2447ba021acf1c9ac0aea25f65bb58cf3066dcda734e5"
+URL="http://localhost:8001/v1/chat/completions"; MODEL="GLM-5.2-NVFP4-AQLM-hybrid"
+N=40; CONC=6; MAXTOK=20000
+p=glob.glob("/home/coder/.cache/huggingface/hub/datasets--fingertap--GPQA-Diamond/snapshots/*/test/*.parquet")[0]
+d=pq.read_table(p).to_pydict(); Q=d["question"]; A=d["answer"]
+rows=list(range(N))
+def ask(i):
+    q=Q[i]
+    prompt=(q.strip()+"\n\nThink step by step, then end with a single line exactly:\nAnswer: <letter>")
+    body=json.dumps({"model":MODEL,"messages":[{"role":"user","content":prompt}],
+                     "max_tokens":MAXTOK,"temperature":0}).encode()
+    req=urllib.request.Request(URL,data=body,headers={"Authorization":f"Bearer {K}","Content-Type":"application/json"})
+    for _ in range(2):
+        try:
+            r=json.load(urllib.request.urlopen(req,timeout=900))
+            txt=r["choices"][0]["message"]["content"]; fin=r["choices"][0]["finish_reason"]
+            m=re.findall(r"Answer:\s*\**([A-Da-d])\b", txt)
+            if not m: m=re.findall(r"\b([A-Da-d])\b\s*$", txt.strip())
+            pred=(m[-1].upper() if m else "?")
+            return i,pred,(pred==A[i].upper()),r["usage"]["completion_tokens"],fin
+        except Exception as e:
+            last=str(e)[:60]; time.sleep(2)
+    return i,"ERR",False,0,last
+t0=time.time(); res={}
+with cf.ThreadPoolExecutor(max_workers=CONC) as ex:
+    for i,pred,ok,tok,fin in ex.map(ask,rows):
+        res[i]=(pred,ok,tok,fin)
+        print(f"q{i:02d} gold={A[i]} pred={pred} {'OK ' if ok else 'x  '} tok={tok} {fin}",flush=True)
+c=sum(1 for i in rows if res[i][1]); trunc=sum(1 for i in rows if res[i][3]=="length")
+print(f"\nGPQA-40 (fingertap MC, current TP4-no-DCP serve): {c}/{N} = {c/N:.3f}")
+print(f"reference fixA full-198 = 0.889 ; e4m3 = 0.884 ; (this is a 40-subset, expect +-0.05 noise)")
+print(f"truncated(length-capped) answers: {trunc}  (these may undercount)  | wall {time.time()-t0:.0f}s")
+json.dump({str(i):res[i] for i in rows}, open("/data/inkling-quant/gpqa40_current.json","w"))

code/ink_aqlm.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""AQLM (additive k-means) 2-bit quantizer for Inkling routed experts.
+
+Per (layer, proj) over a set of experts, minimize ||W - What||^2 (init-grade;
+activation-aware convergence is a later pass). Overfit to Inkling's FUSED
+layout: teacher W is [E, out, in] bf16 read from /data/inkling.
+
+Scheme (mirrors the GLM hybrid, adapted to fused tensors):
+  - per-(expert,out-channel) scale s[e,o] = mean|W[e,o,:]|  (rows normalized by s)
+  - groups of GDIM=8 along the INPUT dim -> vectors in R^8
+  - ENTRIES=65536 shared codebook(s) fit by GPU k-means on a row sample
+    (shared across all experts of the layer); books=2 encodes a 2nd-codebook
+    residual (used for w2), books=1 for w13
+  - least-squares refit of s after encoding
+Output part: /data/inkling-quant/parts/layer_{L}_{proj}.pt with
+  expert_ids, {proj}_codes int16 [E, books, out, in/8], {proj}_codebooks fp16
+  [books, ENTRIES, 8], {proj}_scales fp16 [E, out], err_before/after.
+"""
+from __future__ import annotations
+import argparse, os, time
+import torch
+import ink_common as ic
+
+ENTRIES = 65536
+GDIM = 8
+SAMPLE_VECS = 4_000_000
+KMEANS_ITERS = 12
+CHUNK = 32_768  # scores [CHUNK, ENTRIES] fp16 ~4.3GB
+
+def _argmax_codes(v16, cent16, cn16):
+    d = v16 @ cent16.t()          # [n, ENTRIES]
+    d.mul_(2).sub_(cn16)          # 2 v.c - |c|^2  (argmax == nearest)
+    return d.argmax(-1).to(torch.int32)
+
+def _encode(vecs16, cent16):
+    cn16 = (cent16.float() * cent16.float()).sum(-1).half()
+    out = torch.empty(vecs16.shape[0], dtype=torch.int32, device=vecs16.device)
+    for s in range(0, vecs16.shape[0], CHUNK):
+        out[s:s+CHUNK] = _argmax_codes(vecs16[s:s+CHUNK], cent16, cn16)
+    return out
+
+def _kmeans(vecs16, iters, gen, entries=ENTRIES):
+    n = vecs16.shape[0]; dev = vecs16.device
+    cent = vecs16[torch.randperm(n, generator=gen, device=dev)[:entries]].float()
+    for _ in range(iters):
+        cent16 = cent.half(); cn16 = (cent*cent).sum(-1).half()
+        sums = torch.zeros(entries, GDIM, device=dev); cnts = torch.zeros(entries, device=dev)
+        for s in range(0, n, CHUNK):
+            v = vecs16[s:s+CHUNK]; a = _argmax_codes(v, cent16, cn16).long()
+            sums.index_add_(0, a, v.float()); cnts.index_add_(0, a, torch.ones_like(a, dtype=torch.float32))
+        nz = cnts > 0; cent[nz] = sums[nz] / cnts[nz].unsqueeze(1)
+    return cent.half()
+
+def quantize_layer(layer, proj, book_entries, device="cuda", seed=0):
+    """book_entries: list of codebook sizes per book, e.g. w13=[65536], w2=[65536,256]."""
+    books = len(book_entries)
+    tc = ic.ShardReader(ic.BF16_DIR)
+    gen = torch.Generator(device=device); gen.manual_seed(seed + 1000*layer + (0 if proj=="w13" else 1))
+    # stack teacher experts in bf16 (avoid full float copies -> OOM on w13): [E, out, in]
+    W = torch.stack([ic.teacher_expert(tc, layer, proj, e, device=device) for e in range(ic.N_EXPERTS)])
+    E, out, ind = W.shape
+    assert ind % GDIM == 0
+    s = W.abs().float().mean(-1).clamp_min(1e-8)      # [E, out] per-row scale (fp32)
+    vecs = (W.float() / s.unsqueeze(-1)).reshape(E*out*(ind//GDIM), GDIM).half()  # normalized groups-of-8
+    del W; torch.cuda.empty_cache()
+    codes = torch.empty(books, vecs.shape[0], dtype=torch.int32, device=device)
+    codebooks = []  # variable sizes per book
+    recon = None
+    for b in range(books):
+        target = vecs if b == 0 else (vecs - recon).half()
+        idx = torch.randperm(target.shape[0], generator=gen, device=device)[:SAMPLE_VECS]
+        cbb = _kmeans(target[idx], KMEANS_ITERS, gen, entries=book_entries[b])
+        codes[b] = _encode(target, cbb); codebooks.append(cbb.cpu())
+        contrib = cbb[codes[b].long()]
+        recon = contrib if recon is None else (recon + contrib).half()
+        del target
+    # rel err (weight-space, unweighted) — chunked to avoid full float copies
+    num = torch.zeros((), device=device); den = torch.zeros((), device=device)
+    for st in range(0, vecs.shape[0], CHUNK*8):
+        v = vecs[st:st+CHUNK*8].float(); r = recon[st:st+CHUNK*8].float()
+        num += (v-r).pow(2).sum(); den += v.pow(2).sum()
+    rel = (num.sqrt() / den.sqrt())
+    codes = codes.reshape(books, E, out, ind//GDIM).permute(1,0,2,3).contiguous()  # [E,books,out,in/8]
+    return {
+        "expert_ids": torch.arange(E, dtype=torch.int32),
+        f"{proj}_codes": codes.to(torch.int16).cpu(),
+        f"{proj}_codebooks": codebooks,           # list of [entries_b, 8] fp16
+        f"{proj}_book_entries": list(book_entries),
+        f"{proj}_scales": s.to(torch.float16).cpu(),
+        "rel_err": float(rel),
+    }
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layer", type=int, default=3)
+    ap.add_argument("--proj", default="w2", choices=["w13","w2"])
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--save", action="store_true")
+    a = ap.parse_args()
+    # ship scheme: w13=[65536], w2=[65536,256]
+    be = [65536] if a.proj=="w13" else [65536, 256]
+    t0=time.time()
+    r = quantize_layer(a.layer, a.proj, be, a.device)
+    print(f"L{a.layer} {a.proj} book_entries={be}: weight-space rel_err={r['rel_err']:.4f}  ({time.time()-t0:.1f}s)")
+    if a.save:
+        os.makedirs("/data/inkling-quant/parts", exist_ok=True)
+        torch.save(r, f"/data/inkling-quant/parts/layer_{a.layer}_{a.proj}.pt")
+        print("saved.")

code/ink_common.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+"""Foundational tooling for the Inkling NVFP4+AQLM hybrid.
+
+Handles Inkling's FUSED per-layer expert layout:
+  NVFP4 (/data/inkling-nvfp4): mlp.experts.w13_weight u8 [E, out, in/2]
+     + .scale f8_e4m3 [E, out, in/16] (block-16) + .scale2 f32 [E] (per-expert global)
+  BF16 teacher (/data/inkling): mlp.experts.w13_weight bf16 [E, out, in]
+w13: out=6144 (gate|up=2*3072), in=6144.  w2: out=6144, in=3072.
+"""
+from __future__ import annotations
+import json, os
+import torch
+
+NVFP4_DIR = os.environ.get("INK_NVFP4_DIR", "/data/inkling-nvfp4")
+BF16_DIR  = os.environ.get("INK_BF16_DIR",  "/data/inkling")
+N_EXPERTS = 256
+G = 8  # AQLM group dim
+# NVFP4 e2m1 value table (low nibble = first of each pair), matches GLM/modelopt.
+FP4_LUT = torch.tensor([0.0,0.5,1.0,1.5,2.0,3.0,4.0,6.0,
+                        -0.0,-0.5,-1.0,-1.5,-2.0,-3.0,-4.0,-6.0], dtype=torch.float32)
+
+class ShardReader:
+    """Random access to tensors across a repo's safetensors shards."""
+    def __init__(self, repo_dir: str):
+        idx = json.load(open(os.path.join(repo_dir, "model.safetensors.index.json")))
+        self.weight_map = idx["weight_map"]
+        self.repo_dir = repo_dir
+        self._open = {}
+    def has(self, name: str) -> bool:
+        return name in self.weight_map
+    def get(self, name: str) -> torch.Tensor:
+        from safetensors import safe_open
+        shard = self.weight_map[name]
+        if shard not in self._open:
+            self._open[shard] = safe_open(os.path.join(self.repo_dir, shard), framework="pt")
+        return self._open[shard].get_tensor(name)
+
+def _pfx(layer: int, proj: str) -> str:
+    return f"model.llm.layers.{layer}.mlp.experts.{proj}_weight"
+
+def dequant_nvfp4_expert(reader: ShardReader, layer: int, proj: str, e: int, device="cuda") -> torch.Tensor:
+    """Dequantize one NVFP4 expert (proj in {'w13','w2'}) -> bf16 [out, in]."""
+    p = _pfx(layer, proj)
+    lut = FP4_LUT.to(device)
+    packed = reader.get(p)[e].to(device)                 # u8 [out, in/2]
+    scale  = reader.get(f"{p}.scale")[e].to(device)      # f8_e4m3 [out, in/16]
+    scale2 = reader.get(f"{p}.scale2")[e].to(torch.float32).to(device)  # f32 scalar
+    lo = (packed & 0x0F).long(); hi = (packed >> 4).long()
+    out, half = packed.shape
+    vals = torch.empty(out, half * 2, dtype=torch.float32, device=device)
+    vals[:, 0::2] = lut[lo]; vals[:, 1::2] = lut[hi]
+    del lo, hi
+    blk = scale.to(torch.float32).repeat_interleave(16, dim=1)
+    return (vals * blk * scale2).to(torch.bfloat16)
+
+def teacher_expert(reader: ShardReader, layer: int, proj: str, e: int, device="cuda") -> torch.Tensor:
+    """BF16 teacher expert -> bf16 [out, in]."""
+    return reader.get(_pfx(layer, proj))[e].to(device).to(torch.bfloat16)
+
+def quantize_nvfp4(W: torch.Tensor):
+    """bf16/fp32 [out,in] -> ModelOpt NVFP4 (inverse of dequant_nvfp4_expert):
+    packed u8 [out,in/2], scale f8_e4m3 [out,in/16] (block-16), scale2 f32 scalar.
+    W ~= lut[code] * scale.repeat_interleave(16) * scale2."""
+    dev = W.device; Wf = W.float()
+    out, ind = Wf.shape; assert ind % 16 == 0
+    lut = FP4_LUT.to(dev)
+    amax = Wf.abs().amax().clamp_min(1e-8)
+    scale2 = amax / (6.0 * 448.0)                       # global fp32
+    w = Wf / scale2
+    wb = w.reshape(out, ind // 16, 16)
+    bmax = wb.abs().amax(-1, keepdim=True).clamp_min(1e-8)
+    bscale = (bmax / 6.0)                                # per-block, -> f8
+    bscale_f8 = bscale.to(torch.float8_e4m3fn).float().clamp_min(1e-8)
+    q = wb / bscale_f8                                   # in fp4 dynamic range
+    # nearest fp4 level
+    diff = (q.unsqueeze(-1) - lut.view(1, 1, 1, -1)).abs()
+    codes = diff.argmin(-1).to(torch.uint8)             # [out, nb, 16] indices 0..15
+    codes = codes.reshape(out, ind)
+    lo = codes[:, 0::2]; hi = codes[:, 1::2]
+    packed = (lo | (hi << 4)).to(torch.uint8)           # [out, in/2]
+    scale_f8 = bscale.reshape(out, ind // 16).to(torch.float8_e4m3fn)
+    return packed, scale_f8, scale2.to(torch.float32).reshape(1)
+
+def aqlm_dequant_expert(part, proj, e, device="cuda"):
+    """Reconstruct one AQLM expert -> bf16 [out,in] from a converge/encode part.
+    part has {proj}_codes [E,books,out,ng] int16, {proj}_codebooks list of
+    [entries,8] fp16, {proj}_scales [E,out] fp16."""
+    codes = part[f"{proj}_codes"][e].to(device).long()      # [books,out,ng]
+    cbs = part[f"{proj}_codebooks"]                          # list of [entries,8]
+    scales = part[f"{proj}_scales"][e].to(device).float()   # [out]
+    books, out, ng = codes.shape
+    recon = torch.zeros(out, ng, G, device=device)
+    for b in range(books):
+        cb = cbs[b].to(device).float()                      # [entries,8]; neg idx wraps ok for 65536
+        recon += cb[codes[b]]                               # [out,ng,8]
+    recon = recon.reshape(out, ng * G) * scales.unsqueeze(-1)
+    return recon.to(torch.bfloat16)
+
+def dequant_nvfp4_raw(packed, scale_f8, scale2, device="cuda"):
+    """Dequant raw NVFP4 tensors (for round-trip checks)."""
+    lut = FP4_LUT.to(device)
+    packed = packed.to(device); scale_f8 = scale_f8.to(device); scale2 = scale2.to(device).float()
+    lo = (packed & 0x0F).long(); hi = (packed >> 4).long()
+    out, half = packed.shape
+    vals = torch.empty(out, half * 2, dtype=torch.float32, device=device)
+    vals[:, 0::2] = lut[lo]; vals[:, 1::2] = lut[hi]
+    blk = scale_f8.to(torch.float32).repeat_interleave(16, dim=1)
+    return (vals * blk * scale2).to(torch.bfloat16)
+
+def moe_layers() -> list[int]:
+    """Layers that carry routed experts (skip the 2 dense layers)."""
+    idx = json.load(open(os.path.join(NVFP4_DIR, "model.safetensors.index.json")))["weight_map"]
+    ls = set()
+    for k in idx:
+        if ".mlp.experts.w13_weight" == k[k.find(".mlp"):] if ".mlp.experts.w13_weight" in k else False:
+            pass
+    import re
+    for k in idx:
+        m = re.match(r"model\.llm\.layers\.(\d+)\.mlp\.experts\.w13_weight$", k)
+        if m: ls.add(int(m.group(1)))
+    return sorted(ls)

code/merge_acts.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+"""Merge text + image + audio acts into ONE dir per split, so REAP / converge / PV
+see the full multimodal routing. Concatenates x/ids and sums n per layer, capped.
+
+  train: acts_v3 + acts_v3_img + acts_v3_audio           -> acts_v3_all
+  held : acts_v3_heldout + acts_v3_img_held + acts_v3_audio_held -> acts_v3_all_held
+"""
+import os, glob, torch
+Q = "/data/inkling-quant"
+CAP = 160000   # max token-rows/layer in the merged file (text-heavy but mm represented)
+SETS = {
+    "acts_v3_all":      ["acts_v3", "acts_v3_img", "acts_v3_audio"],
+    "acts_v3_all_held": ["acts_v3_heldout", "acts_v3_img_held", "acts_v3_audio_held"],
+}
+
+def layers(dirs):
+    s = set()
+    for d in dirs:
+        for f in glob.glob(f"{Q}/{d}/acts_layer_*.pt"):
+            s.add(int(f.split("_")[-1].split(".")[0]))
+    return sorted(s)
+
+def main():
+    for out, srcs in SETS.items():
+        od = f"{Q}/{out}"; os.makedirs(od, exist_ok=True)
+        present = [s for s in srcs if os.path.isdir(f"{Q}/{s}")]
+        print(f">>> {out} <- {present}")
+        for L in layers(present):
+            xs, ids, n = [], [], 0
+            for s in present:
+                p = f"{Q}/{s}/acts_layer_{L}.pt"
+                if not os.path.exists(p): continue
+                a = torch.load(p, map_location="cpu")
+                xs.append(a["x"]); ids.append(a["ids"]); n += int(a.get("n", a["x"].shape[0]))
+            if not xs: continue
+            x = torch.cat(xs, 0); i = torch.cat(ids, 0)
+            if x.shape[0] > CAP:
+                sel = torch.randperm(x.shape[0])[:CAP]; x = x[sel]; i = i[sel]
+            torch.save({"x": x, "ids": i, "n": x.shape[0]}, f"{od}/acts_layer_{L}.pt")
+        print(f"    wrote {len(layers(present))} layers -> {od}")
+
+if __name__ == "__main__":
+    main()

code/ppl_eval.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+"""Perplexity + coherence gate for a checkpoint (baseline NVFP4 vs 512k gate).
+Offline vLLM: teacher-forced prompt_logprobs over a fixed held-out text set ->
+mean NLL -> perplexity. Same text for both models => the RATIO is the quality
+signal (absolute may be optimistic since eval text overlaps calib distribution).
+Usage: CUDA_VISIBLE_DEVICES=4,5,6,7 python ppl_eval.py --model DIR --tag NAME
+"""
+import argparse, json, math, os
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--tag", required=True)
+    ap.add_argument("--n", type=int, default=48)
+    # v3: default to the TRULY HELD-OUT set (wikitext-2 test, disjoint from all
+    # calib/tuning) so pre/post ppl is honest. Pass --eval-file for others.
+    ap.add_argument("--eval-file", default="/data/inkling-quant/ppl_v3.jsonl")
+    a = ap.parse_args()
+    os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+    from vllm import LLM, SamplingParams
+    # fixed eval set: first N chunks (identical text for every model -> fair ratio).
+    # v3 ppl_v3.jsonl carries a per-chunk "cat" -> report per-domain ppl too.
+    rows = [json.loads(l) for l in open(a.eval_file)][:a.n]
+    prompts = [r["text"] for r in rows]
+    cats = [r.get("cat", "all") for r in rows]
+    print(f"PPL_EVAL model={a.model} tag={a.tag} eval_file={a.eval_file} n={len(prompts)}")
+    llm = LLM(model=a.model, tensor_parallel_size=4, max_model_len=8192,
+              gpu_memory_utilization=0.9, enforce_eager=True, trust_remote_code=True)
+    outs = llm.generate(prompts, SamplingParams(max_tokens=1, temperature=0, prompt_logprobs=0))
+    tot_nll = 0.0; tot_tok = 0
+    by = {}   # cat -> [nll_sum, tok]
+    for cat, o in zip(cats, outs):
+        c = by.setdefault(cat, [0.0, 0])
+        for d in o.prompt_logprobs:      # list per prompt token; [0] is None
+            if not d:
+                continue
+            lp = next(iter(d.values())).logprob   # teacher-forced -> actual token
+            tot_nll += -lp; tot_tok += 1
+            c[0] += -lp; c[1] += 1
+    ppl = math.exp(tot_nll / max(1, tot_tok))
+    print(f"PPL_RESULT tag={a.tag} tokens={tot_tok} mean_nll={tot_nll/max(1,tot_tok):.4f} ppl={ppl:.4f}")
+    for cat in sorted(by):
+        n, t = by[cat]
+        print(f"  PPL_CAT tag={a.tag} cat={cat:12s} tokens={t} ppl={math.exp(n/max(1,t)):.4f}")
+    # coherence probes
+    for q in ["The capital of France is", "def fibonacci(n):", "Water is made of hydrogen and"]:
+        r = llm.generate([q], SamplingParams(max_tokens=24, temperature=0))[0].outputs[0].text
+        print(f"COHERENCE [{q[:22]!r}] -> {r[:60]!r}")
+
+if __name__ == "__main__":
+    main()

code/pv.py

@@ -0,0 +1,256 @@
+#!/usr/bin/env python3
+"""Output-preserving (PV) tuning (Phase E2) for Inkling 512k v2 cold experts.
+
+For each layer, tune the SHARED AQLM codebook VALUES + per-(cold-expert,out-row)
+SCALES so the dequantized expert's OUTPUT matches the BF16 teacher's output on
+calib activations.  CODES stay FIXED (loaded from the converged parts).
+
+Trainable nn.Parameters (init from parts_conv):
+  cb13     [65536,8]        w13 codebook (1 book)
+  cb2_0    [65536,8]        w2 codebook book0
+  cb2_1    [256,8]          w2 codebook book1
+  s13      [ncold,6144]     w13 per-(cold-expert,out-row) scales
+  s2       [ncold,6144]     w2  per-(cold-expert,out-row) scales
+
+Differentiable dequant (codes fixed):
+  Wq = scales.unsqueeze(-1) * sum_b codebook_b[idx_b] reshaped to [out,in]
+  idx = codes.long() % cb.shape[0]   (int16 negative-wrap for the 65536 book)
+
+Expert forward (w13 ROW-INTERLEAVED: gate=w13[0::2], up=w13[1::2]):
+  h = silu(gate @ x^T) * (up @ x^T)     [3072, n]
+  f = w2 @ h                            [6144, n]
+Teacher forward uses the BF16 teacher weights (fixed, precomputed).
+
+Loss = sum over cold experts of MSE(f_quant(x_e), f_teacher(x_e)) (unweighted).
+Optimizer: Adam, lr 1e-3 (scales) / 1e-4 (codebooks), param groups.
+
+Writes /data/inkling-quant/parts_pv/layer_{L}_{w13,w2}.pt with the SAME schema as
+parts_conv ({proj}_codes [unchanged], {proj}_codebooks [tuned], {proj}_book_entries,
+{proj}_scales [tuned], cold_experts, expert_ids) + output_mse_before/after and the
+weight-space rel_err sanity numbers.
+"""
+from __future__ import annotations
+import argparse, os, time, json
+import torch
+import torch.nn.functional as F
+import ink_common as ic
+
+CONV = "/data/inkling-quant/parts_conv"
+OUT = "/data/inkling-quant/parts_pv"
+ACTS = "/data/inkling-quant/acts"
+ASSIGN = "/data/inkling-quant/assign_512k.json"
+G = 8
+OUTD = 6144
+W13_IN = 6144; W13_NG = W13_IN // G      # 768
+W2_IN = 3072;  W2_NG = W2_IN // G        # 384
+MIN_TOK = 8
+MAX_TOK = 4096
+
+
+def cold_experts(layer, assign):
+    hot = set(assign["hot_experts"].get(str(layer), []))
+    return [e for e in range(ic.N_EXPERTS) if e not in hot]
+
+
+def dequant_w13(ci, codes13, cb13, s13):
+    """codes13[ci,0] -> [6144,768]; returns Wq13 [6144,6144] (differentiable)."""
+    idx = codes13[ci, 0].long() % cb13.shape[0]              # [6144,768]
+    recon = cb13[idx].reshape(OUTD, W13_IN)                  # [6144,6144]
+    return s13[ci].unsqueeze(-1) * recon
+
+
+def dequant_w2(ci, codes2, cb2_0, cb2_1, s2):
+    """codes2[ci] -> books; returns Wq2 [6144,3072] (differentiable)."""
+    idx0 = codes2[ci, 0].long() % cb2_0.shape[0]             # [6144,384]
+    idx1 = codes2[ci, 1].long() % cb2_1.shape[0]             # [6144,384]
+    recon = (cb2_0[idx0] + cb2_1[idx1]).reshape(OUTD, W2_IN)  # [6144,3072]
+    return s2[ci].unsqueeze(-1) * recon
+
+
+def expert_forward(Wq13, Wq2, xe):
+    """xe [n,6144] -> expert output [6144,n]."""
+    gate = Wq13[0::2]                                        # [3072,6144]
+    up = Wq13[1::2]                                          # [3072,6144]
+    xt = xe.t()                                              # [6144,n]
+    h = F.silu(gate @ xt) * (up @ xt)                        # [3072,n]
+    return Wq2 @ h                                           # [6144,n]
+
+
+def build_data(layer, cold, device):
+    """Per cold expert with >=MIN_TOK routed tokens: xe [n,6144] fp32, teacher
+    output ftea [6144,n] fp32 (precomputed from BF16 teacher weights, fixed).
+    Returns list of (ci, e, xe, ftea) and the count of data experts."""
+    a = torch.load(f"{ACTS}/acts_layer_{layer}.pt")
+    x = a["x"].to(device).float()                            # [N,6144]
+    ids = a["ids"].to(device)                                # [N,6]
+    tc = ic.ShardReader(ic.BF16_DIR)
+    data = []
+    gen = torch.Generator(device=device); gen.manual_seed(1234 + layer)
+    for ci, e in enumerate(cold):
+        m = (ids == e).any(1)
+        n = int(m.sum())
+        if n < MIN_TOK:
+            continue
+        xe = x[m]                                            # [n,6144]
+        if n > MAX_TOK:
+            sel = torch.randperm(n, generator=gen, device=device)[:MAX_TOK]
+            xe = xe[sel]
+        w13 = ic.teacher_expert(tc, layer, "w13", e, device=device).float()   # [6144,6144]
+        w2 = ic.teacher_expert(tc, layer, "w2", e, device=device).float()     # [6144,3072]
+        with torch.no_grad():
+            ftea = expert_forward(w13, w2, xe)               # [6144,n]
+        data.append((ci, e, xe.contiguous(), ftea.contiguous()))
+        del w13, w2
+    del x, ids
+    torch.cuda.empty_cache()
+    return data
+
+
+@torch.no_grad()
+def eval_output_mse(data, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2):
+    """Mean over data experts of per-expert MSE(f_quant, f_teacher), plus the
+    mean normalized rel-err (||fq-ft||^2 / ||ft||^2)."""
+    mse_sum = 0.0; rel_sum = 0.0
+    for (ci, e, xe, ftea) in data:
+        Wq13 = dequant_w13(ci, codes13, cb13, s13)
+        Wq2 = dequant_w2(ci, codes2, cb2_0, cb2_1, s2)
+        fq = expert_forward(Wq13, Wq2, xe)
+        diff = (fq - ftea)
+        mse_sum += float(diff.pow(2).mean())
+        rel_sum += float(diff.pow(2).sum() / ftea.pow(2).sum().clamp_min(1e-12))
+    ne = len(data)
+    return mse_sum / ne, rel_sum / ne
+
+
+@torch.no_grad()
+def weight_relerr(layer, data, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2, device):
+    """Aggregate plain weight-space rel-err ||Wq-Wt||/||Wt|| over the data experts,
+    for w13 and w2 separately (reloads BF16 teacher weights)."""
+    tc = ic.ShardReader(ic.BF16_DIR)
+    n13 = d13 = n2 = d2 = 0.0
+    for (ci, e, _, _) in data:
+        Wt13 = ic.teacher_expert(tc, layer, "w13", e, device=device).float()
+        Wq13 = dequant_w13(ci, codes13, cb13, s13)
+        n13 += float((Wq13 - Wt13).pow(2).sum()); d13 += float(Wt13.pow(2).sum())
+        del Wt13, Wq13
+        Wt2 = ic.teacher_expert(tc, layer, "w2", e, device=device).float()
+        Wq2 = dequant_w2(ci, codes2, cb2_0, cb2_1, s2)
+        n2 += float((Wq2 - Wt2).pow(2).sum()); d2 += float(Wt2.pow(2).sum())
+        del Wt2, Wq2
+    return (n13 / d13) ** 0.5, (n2 / d2) ** 0.5
+
+
+def pv_layer(layer, steps=300, lr_cb=1e-4, lr_s=1e-3, device="cuda", verbose=True):
+    assign = json.load(open(ASSIGN))
+    cold = cold_experts(layer, assign)
+    coldt = torch.tensor(cold)
+    p13 = torch.load(f"{CONV}/layer_{layer}_w13.pt", map_location="cpu")
+    p2 = torch.load(f"{CONV}/layer_{layer}_w2.pt", map_location="cpu")
+
+    # fixed codes (cold-indexed) on device
+    codes13 = p13["w13_codes"][coldt].to(device)             # [ncold,1,6144,768] int16
+    codes2 = p2["w2_codes"][coldt].to(device)                # [ncold,2,6144,384] int16
+
+    # trainable params (init from converged parts)
+    cb13 = p13["w13_codebooks"][0].float().to(device).clone().requires_grad_()
+    cb2_0 = p2["w2_codebooks"][0].float().to(device).clone().requires_grad_()
+    cb2_1 = p2["w2_codebooks"][1].float().to(device).clone().requires_grad_()
+    s13 = p13["w13_scales"][coldt].float().to(device).clone().requires_grad_()
+    s2 = p2["w2_scales"][coldt].float().to(device).clone().requires_grad_()
+
+    data = build_data(layer, cold, device)
+    if verbose:
+        print(f"  L{layer}: {len(cold)} cold, {len(data)} with >={MIN_TOK} tokens", flush=True)
+    if len(data) == 0:
+        raise RuntimeError(f"L{layer}: no cold experts with calib tokens")
+
+    # baseline (converged-parts init)
+    mse0, rel0 = eval_output_mse(data, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2)
+    w13rel0, w2rel0 = weight_relerr(layer, data, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2, device)
+    if verbose:
+        print(f"  L{layer} BEFORE: out_mse={mse0:.6e} rel={rel0:.4f} | "
+              f"w13_relerr={w13rel0:.4f} w2_relerr={w2rel0:.4f}", flush=True)
+
+    opt = torch.optim.Adam([
+        {"params": [cb13, cb2_0, cb2_1], "lr": lr_cb},
+        {"params": [s13, s2], "lr": lr_s},
+    ])
+
+    for step in range(steps):
+        opt.zero_grad(set_to_none=True)
+        tot = 0.0
+        for (ci, e, xe, ftea) in data:
+            Wq13 = dequant_w13(ci, codes13, cb13, s13)
+            Wq2 = dequant_w2(ci, codes2, cb2_0, cb2_1, s2)
+            fq = expert_forward(Wq13, Wq2, xe)
+            loss = F.mse_loss(fq, ftea)
+            loss.backward()
+            tot += float(loss.detach())
+        opt.step()
+        if verbose and (step % 50 == 0 or step == steps - 1):
+            print(f"  L{layer} step{step}: sum_mse={tot:.6e} (mean={tot/len(data):.6e})", flush=True)
+
+    mse1, rel1 = eval_output_mse(data, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2)
+    w13rel1, w2rel1 = weight_relerr(layer, data, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2, device)
+    if verbose:
+        print(f"  L{layer} AFTER:  out_mse={mse1:.6e} rel={rel1:.4f} | "
+              f"w13_relerr={w13rel1:.4f} w2_relerr={w2rel1:.4f}", flush=True)
+
+    # assemble full-256 tuned arrays (cold rows tuned; hot rows carried from conv part)
+    with torch.no_grad():
+        full_s13 = p13["w13_scales"].clone()
+        full_s13[coldt] = s13.detach().half().cpu()
+        full_s2 = p2["w2_scales"].clone()
+        full_s2[coldt] = s2.detach().half().cpu()
+
+    common = dict(
+        cold_experts=cold,
+        output_mse_before=mse0, output_mse_after=mse1,
+        output_relerr_before=rel0, output_relerr_after=rel1,
+        conv_rel_err_after_w13=float(p13.get("rel_err_after", float("nan"))),
+        conv_rel_err_after_w2=float(p2.get("rel_err_after", float("nan"))),
+    )
+    out13 = {
+        "expert_ids": p13["expert_ids"],
+        "w13_codes": p13["w13_codes"],                       # UNCHANGED
+        "w13_codebooks": [cb13.detach().half().cpu()],       # tuned
+        "w13_book_entries": p13["w13_book_entries"],
+        "w13_scales": full_s13,                              # tuned
+        "weight_relerr_before": w13rel0, "weight_relerr_after": w13rel1,
+        **common,
+    }
+    out2 = {
+        "expert_ids": p2["expert_ids"],
+        "w2_codes": p2["w2_codes"],                          # UNCHANGED
+        "w2_codebooks": [cb2_0.detach().half().cpu(), cb2_1.detach().half().cpu()],  # tuned
+        "w2_book_entries": p2["w2_book_entries"],
+        "w2_scales": full_s2,                                # tuned
+        "weight_relerr_before": w2rel0, "weight_relerr_after": w2rel1,
+        **common,
+    }
+    return out13, out2, dict(
+        mse0=mse0, mse1=mse1, rel0=rel0, rel1=rel1,
+        w13rel0=w13rel0, w13rel1=w13rel1, w2rel0=w2rel0, w2rel1=w2rel1,
+    )
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layer", type=int, required=True)
+    ap.add_argument("--steps", type=int, default=300)
+    ap.add_argument("--lr_cb", type=float, default=1e-4)
+    ap.add_argument("--lr_s", type=float, default=1e-3)
+    ap.add_argument("--save", action="store_true")
+    a = ap.parse_args()
+    os.makedirs(OUT, exist_ok=True)
+    t0 = time.time()
+    out13, out2, m = pv_layer(a.layer, a.steps, a.lr_cb, a.lr_s)
+    dt = time.time() - t0
+    print(f"L{a.layer} PV: out_mse {m['mse0']:.6e} -> {m['mse1']:.6e} "
+          f"(rel {m['rel0']:.4f}->{m['rel1']:.4f}) | "
+          f"w13_relerr {m['w13rel0']:.4f}->{m['w13rel1']:.4f} "
+          f"w2_relerr {m['w2rel0']:.4f}->{m['w2rel1']:.4f}  ({dt:.0f}s)")
+    if a.save:
+        torch.save(out13, f"{OUT}/layer_{a.layer}_w13.pt")
+        torch.save(out2, f"{OUT}/layer_{a.layer}_w2.pt")
+        print("saved.")

code/pv_all.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""Full output-preserving (PV) tuning (Phase E2) across all MoE layers.
+Shards layers across CUDA_VISIBLE_DEVICES GPUs, 1 spawn-worker/GPU, resumable
+(skips layers whose BOTH parts_pv files already exist).
+Run: CUDA_VISIBLE_DEVICES=4,5,6,7 python pv_all.py [--steps 300]
+"""
+import os, sys, time, argparse
+import torch.multiprocessing as mp
+import ink_common as ic
+
+OUT = "/data/inkling-quant/parts_pv"
+
+
+def done(L):
+    return (os.path.exists(f"{OUT}/layer_{L}_w13.pt")
+            and os.path.exists(f"{OUT}/layer_{L}_w2.pt"))
+
+
+def worker(gpu, layers, steps):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    import torch, pv
+    for L in layers:
+        if done(L):
+            print(f"[gpu{gpu}] skip L{L} (exists)", flush=True); continue
+        t0 = time.time()
+        try:
+            out13, out2, m = pv.pv_layer(L, steps=steps, device="cuda", verbose=False)
+            for proj, obj in (("w13", out13), ("w2", out2)):
+                dst = f"{OUT}/layer_{L}_{proj}.pt"
+                tmp = dst + ".tmp"
+                torch.save(obj, tmp); os.replace(tmp, dst)
+            improved = "OK" if m["mse1"] < m["mse0"] else "WORSE"
+            print(f"[gpu{gpu}] L{L} PV out_mse {m['mse0']:.4e}->{m['mse1']:.4e} "
+                  f"(rel {m['rel0']:.4f}->{m['rel1']:.4f}) {improved} | "
+                  f"w13rel {m['w13rel0']:.4f}->{m['w13rel1']:.4f} "
+                  f"w2rel {m['w2rel0']:.4f}->{m['w2rel1']:.4f} "
+                  f"({time.time()-t0:.0f}s) saved", flush=True)
+        except Exception as ex:
+            import traceback
+            print(f"[gpu{gpu}] L{L} FAILED: {ex}", flush=True)
+            traceback.print_exc()
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--steps", type=int, default=300)
+    a = ap.parse_args()
+    os.makedirs(OUT, exist_ok=True)
+    gpus = [int(x) for x in os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")]
+    layers = ic.moe_layers()
+    todo = [L for L in layers if not done(L)]
+    print(f"{len(layers)} layers, {len(todo)} remaining, {len(gpus)} GPUs {gpus}", flush=True)
+    shards = {g: [] for g in gpus}
+    for i, L in enumerate(todo):
+        shards[gpus[i % len(gpus)]].append(L)
+    mp.set_start_method("spawn", force=True)
+    ps = [mp.Process(target=worker, args=(g, shards[g], a.steps)) for g in gpus]
+    for p in ps: p.start()
+    for p in ps: p.join()
+    d = sum(1 for L in layers if done(L))
+    print(f"PV DONE: {d}/{len(layers)} layers", flush=True)
+
+
+if __name__ == "__main__":
+    main()

code/pv_all_v3.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+"""Full PV-v3 across all MoE layers, with held-out early-stop (pv_v3).
+Shards layers across CUDA_VISIBLE_DEVICES, 1 spawn-worker/GPU, resumable.
+Needs acts_v3/ (train) + acts_v3_heldout/ (val) captured first.
+Run: CUDA_VISIBLE_DEVICES=4,5,6,7 python pv_all_v3.py [--max-steps 2000] [--patience 6]
+"""
+import os, time, argparse
+import torch.multiprocessing as mp
+import ink_common as ic
+
+import os as _os
+OUT = _os.environ.get("INK_PV_OUT", "/data/inkling-quant/parts_pv_v3")
+
+def done(L):
+    return (os.path.exists(f"{OUT}/layer_{L}_w13.pt")
+            and os.path.exists(f"{OUT}/layer_{L}_w2.pt"))
+
+def worker(gpu, layers, max_steps, eval_every, patience, tune_cb):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    import torch, pv_v3
+    for L in layers:
+        if done(L):
+            print(f"[gpu{gpu}] skip L{L} (exists)", flush=True); continue
+        t0 = time.time()
+        try:
+            out13, out2, m = pv_v3.pv_layer_v3(L, max_steps=max_steps, eval_every=eval_every,
+                                               patience=patience, tune_cb=tune_cb,
+                                               device="cuda", verbose=False)
+            for proj, obj in (("w13", out13), ("w2", out2)):
+                dst = f"{OUT}/layer_{L}_{proj}.pt"; tmp = dst + ".tmp"
+                torch.save(obj, tmp); os.replace(tmp, dst)
+            tag = "OK" if (m["val1"] <= m["val0"] or m["val0"] != m["val0"]) else "WORSE"
+            print(f"[gpu{gpu}] L{L} PVv3 val_rel {m['val0']:.4f}->{m['val1']:.4f} "
+                  f"best@{m['best_step']} {tag} ({time.time()-t0:.0f}s) saved", flush=True)
+        except Exception as ex:
+            import traceback
+            print(f"[gpu{gpu}] L{L} FAILED: {ex}", flush=True); traceback.print_exc()
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--max-steps", type=int, default=2000)
+    ap.add_argument("--eval-every", type=int, default=25)
+    ap.add_argument("--patience", type=int, default=6)
+    ap.add_argument("--no-tune-cb", action="store_true")
+    a = ap.parse_args()
+    os.makedirs(OUT, exist_ok=True)
+    gpus = [int(x) for x in os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")]
+    layers = ic.moe_layers()
+    todo = [L for L in layers if not done(L)]
+    print(f"{len(layers)} layers, {len(todo)} remaining, {len(gpus)} GPUs {gpus}", flush=True)
+    shards = {g: [] for g in gpus}
+    for i, L in enumerate(todo):
+        shards[gpus[i % len(gpus)]].append(L)
+    mp.set_start_method("spawn", force=True)
+    ps = [mp.Process(target=worker, args=(g, shards[g], a.max_steps, a.eval_every,
+                                          a.patience, not a.no_tune_cb)) for g in gpus]
+    for p in ps: p.start()
+    for p in ps: p.join()
+    d = sum(1 for L in layers if done(L))
+    print(f"PVv3 DONE: {d}/{len(layers)} layers", flush=True)
+
+if __name__ == "__main__":
+    main()

code/pv_gen_check.py

@@ -0,0 +1,33 @@
+import torch, json, random, statistics as st
+import ink_common as ic
+random.seed(0); dev="cuda"
+assign=json.load(open("/data/inkling-quant/assign_512k_v2.json"))["hot_experts"]
+tc=ic.ShardReader(ic.BF16_DIR)
+def expert_out(W13,W2,x):  # x[n,6144]; W13[6144,6144] interleaved; W2[6144,3072]
+    gate=W13[0::2]; up=W13[1::2]
+    h=torch.nn.functional.silu(x@gate.t())*(x@up.t())
+    return h@W2.t()
+layers=[3,8,15,22,30,38,45,52,58,63]
+conv=[]; pv=[]
+for L in layers:
+    a=torch.load(f"/data/inkling-quant/acts_heldout/acts_layer_{L}.pt",map_location=dev)
+    x=a["x"].float(); ids=a["ids"]
+    hot=set(assign[str(L)]); cold=[e for e in range(256) if e not in hot]
+    pc13=torch.load(f"/data/inkling-quant/parts_conv/layer_{L}_w13.pt",map_location="cpu")
+    pc2 =torch.load(f"/data/inkling-quant/parts_conv/layer_{L}_w2.pt",map_location="cpu")
+    pp13=torch.load(f"/data/inkling-quant/parts_pv/layer_{L}_w13.pt",map_location="cpu")
+    pp2 =torch.load(f"/data/inkling-quant/parts_pv/layer_{L}_w2.pt",map_location="cpu")
+    for e in random.sample(cold,min(12,len(cold))):
+        m=(ids==e).any(1)
+        if int(m.sum())<8: continue
+        xe=x[m][:2048]
+        ft=expert_out(ic.teacher_expert(tc,L,"w13",e,device=dev).float(),ic.teacher_expert(tc,L,"w2",e,device=dev).float(),xe)
+        nt=ft.norm()
+        conv.append(((expert_out(ic.aqlm_dequant_expert(pc13,"w13",e,device=dev).float(),ic.aqlm_dequant_expert(pc2,"w2",e,device=dev).float(),xe)-ft).norm()/nt).item())
+        pv.append(((expert_out(ic.aqlm_dequant_expert(pp13,"w13",e,device=dev).float(),ic.aqlm_dequant_expert(pp2,"w2",e,device=dev).float(),xe)-ft).norm()/nt).item())
+c=st.mean(conv); p=st.mean(pv)
+print(f"HELD-OUT output rel-err over {len(conv)} cold experts:")
+print(f"  converged (parts_conv): {c:.4f}")
+print(f"  PV-tuned  (parts_pv):   {p:.4f}")
+print(f"  PV vs converged (held-out): {100*(1-p/c):+.1f}%")
+print(f"  VERDICT: {'PV GENERALIZES -> use parts_pv' if p < c*0.98 else 'PV OVERFIT / no held-out gain -> use parts_conv'}")

code/pv_generalize.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""PV generalization check (the v2 ship decision).
+
+On HELD-OUT tokens (acts_heldout/, disjoint from calib), for a SAMPLE of cold
+experts across shallow/mid/deep layers, compute the expert OUTPUT rel-err vs the
+BF16 teacher for TWO weight sources:
+  (a) converged parts_conv   (REAP+converge only)
+  (b) PV-tuned  parts_pv      (tuned on calib tokens)
+
+Cold = experts in the v2 complement (assign_512k_v2.json) AND actually PV-tuned
+(in the part's stored cold_experts). Identical held-out tokens used for (a)/(b).
+
+Expert output rel-err = ||f_e(x)-f_teacher(x)||^2 / ||f_teacher(x)||^2, where
+f_e(x)=w2 @ (silu(gate@x^T)*(up@x^T)), w13 ROW-INTERLEAVED gate=w13[0::2]/up=w13[1::2].
+
+VERDICT: if PV held-out rel-err is LOWER than converged by >2% relative -> PV
+GENERALIZES (ship parts_pv). Else PV OVERFIT (ship parts_conv).
+"""
+from __future__ import annotations
+import json, argparse, time
+import torch
+import ink_common as ic
+import pv  # expert_forward
+
+ACTS_HELDOUT = "/data/inkling-quant/acts_heldout"
+CONV = "/data/inkling-quant/parts_conv"
+PV = "/data/inkling-quant/parts_pv"
+ASSIGN_V2 = "/data/inkling-quant/assign_512k_v2.json"
+MIN_TOK = 16
+MAX_TOK = 4096
+DEFAULT_LAYERS = [3, 8, 12, 20, 28, 36, 44, 52, 58, 65]
+N_EXP_PER_LAYER = 16
+
+
+@torch.no_grad()
+def expert_relerr(part13, part2, e, xe, ftea, device):
+    Wq13 = ic.aqlm_dequant_expert(part13, "w13", e, device=device).float()
+    Wq2 = ic.aqlm_dequant_expert(part2, "w2", e, device=device).float()
+    fq = pv.expert_forward(Wq13, Wq2, xe)
+    diff = fq - ftea
+    return float(diff.pow(2).sum() / ftea.pow(2).sum().clamp_min(1e-12))
+
+
+@torch.no_grad()
+def eval_layer(layer, n_exp, device):
+    v2 = json.load(open(ASSIGN_V2))["hot_experts"]
+    v2_cold = set(range(ic.N_EXPERTS)) - set(v2.get(str(layer), []))
+    c13 = torch.load(f"{CONV}/layer_{layer}_w13.pt", map_location="cpu")
+    c2 = torch.load(f"{CONV}/layer_{layer}_w2.pt", map_location="cpu")
+    p13 = torch.load(f"{PV}/layer_{layer}_w13.pt", map_location="cpu")
+    p2 = torch.load(f"{PV}/layer_{layer}_w2.pt", map_location="cpu")
+    pv_tuned = set(c13["cold_experts"])                 # experts PV actually tuned (v1 cold)
+    cand = sorted(v2_cold & pv_tuned)                   # v2-cold AND PV-tuned
+
+    a = torch.load(f"{ACTS_HELDOUT}/acts_layer_{layer}.pt")
+    x = a["x"].to(device).float()                       # [N,6144]
+    ids = a["ids"].to(device)                           # [N,6]
+
+    # eligible = enough held-out tokens; sample a seeded subset
+    counts = {e: int((ids == e).any(1).sum()) for e in cand}
+    elig = [e for e in cand if counts[e] >= MIN_TOK]
+    gen = torch.Generator().manual_seed(2024 + layer)
+    if len(elig) > n_exp:
+        perm = torch.randperm(len(elig), generator=gen).tolist()
+        sample = sorted(elig[i] for i in perm[:n_exp])
+    else:
+        sample = elig
+
+    tc = ic.ShardReader(ic.BF16_DIR)
+    gsel = torch.Generator(device=device).manual_seed(99 + layer)
+    conv_errs, pv_errs, toks = [], [], []
+    for e in sample:
+        m = (ids == e).any(1)
+        xe = x[m]
+        if xe.shape[0] > MAX_TOK:
+            sel = torch.randperm(xe.shape[0], generator=gsel, device=device)[:MAX_TOK]
+            xe = xe[sel]
+        xe = xe.contiguous()
+        wt13 = ic.teacher_expert(tc, layer, "w13", e, device=device).float()
+        wt2 = ic.teacher_expert(tc, layer, "w2", e, device=device).float()
+        ftea = pv.expert_forward(wt13, wt2, xe)
+        del wt13, wt2
+        conv_errs.append(expert_relerr(c13, c2, e, xe, ftea, device))
+        pv_errs.append(expert_relerr(p13, p2, e, xe, ftea, device))
+        toks.append(int(xe.shape[0]))
+        del ftea, xe
+    del x, ids
+    torch.cuda.empty_cache()
+
+    conv_m = sum(conv_errs) / len(conv_errs)
+    pv_m = sum(pv_errs) / len(pv_errs)
+    # calib numbers stored in the PV part (combined-expert output relerr on calib)
+    calib_conv = float(p13["output_relerr_before"])
+    calib_pv = float(p13["output_relerr_after"])
+    return dict(layer=layer, n=len(sample), tot_tok=sum(toks),
+                conv=conv_m, pv=pv_m, calib_conv=calib_conv, calib_pv=calib_pv,
+                experts=sample)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layers", type=int, nargs="+", default=DEFAULT_LAYERS)
+    ap.add_argument("--n_exp", type=int, default=N_EXP_PER_LAYER)
+    a = ap.parse_args()
+    device = "cuda"
+    print(f"HELD-OUT PV-generalization check | layers {a.layers} | <= {a.n_exp} cold experts/layer")
+    print(f"{'layer':>5} {'nexp':>4} {'tok':>6} | {'HELDOUT conv':>13} {'HELDOUT pv':>11} {'d%':>7} || "
+          f"{'CALIB conv':>10} {'CALIB pv':>9} {'d%':>7}")
+    rows = []
+    for L in a.layers:
+        t0 = time.time()
+        r = eval_layer(L, a.n_exp, device)
+        rows.append(r)
+        hd = 100 * (r["pv"] - r["conv"]) / r["conv"]
+        cd = 100 * (r["calib_pv"] - r["calib_conv"]) / r["calib_conv"]
+        print(f"{r['layer']:>5} {r['n']:>4} {r['tot_tok']:>6} | "
+              f"{r['conv']:>13.4f} {r['pv']:>11.4f} {hd:>+6.1f}% || "
+              f"{r['calib_conv']:>10.4f} {r['calib_pv']:>9.4f} {cd:>+6.1f}%  ({time.time()-t0:.0f}s)", flush=True)
+
+    conv_mean = sum(r["conv"] for r in rows) / len(rows)
+    pv_mean = sum(r["pv"] for r in rows) / len(rows)
+    cconv = sum(r["calib_conv"] for r in rows) / len(rows)
+    cpv = sum(r["calib_pv"] for r in rows) / len(rows)
+    hd = 100 * (pv_mean - conv_mean) / conv_mean
+    cd = 100 * (cpv - cconv) / cconv
+    print("-" * 96)
+    print(f"{'MEAN':>5} {sum(r['n'] for r in rows):>4} {sum(r['tot_tok'] for r in rows):>6} | "
+          f"{conv_mean:>13.4f} {pv_mean:>11.4f} {hd:>+6.1f}% || "
+          f"{cconv:>10.4f} {cpv:>9.4f} {cd:>+6.1f}%")
+    print()
+    print(f"HELD-OUT: converged rel-err {conv_mean:.4f}  vs  PV rel-err {pv_mean:.4f}  ({hd:+.1f}%)")
+    print(f"CALIB   : converged rel-err {cconv:.4f}  vs  PV rel-err {cpv:.4f}  ({cd:+.1f}%)")
+    if hd < -2.0:
+        print(f"\nVERDICT: PV GENERALIZES (held-out {hd:+.1f}% <= -2%). SHIP parts_pv for v2 cold experts.")
+    elif hd > 2.0:
+        print(f"\nVERDICT: PV OVERFIT (held-out {hd:+.1f}% WORSE). SHIP parts_conv for v2 cold experts.")
+    else:
+        print(f"\nVERDICT: PV ~= converged on held-out ({hd:+.1f}%, within +-2%) despite winning calib "
+              f"({cd:+.1f}%) => PV did NOT generalize. SHIP parts_conv (simpler, no overfit risk).")
+
+
+if __name__ == "__main__":
+    main()

code/pv_v3.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+"""PV-tune v3 — output-preserving tune WITH a held-out early-stop.
+
+Why v3 differs from pv.py: v2's PV overfit (calib output-err -32% but held-out
++10.5% WORSE; weight-err rose 0.30->0.43) because it ran a FIXED step count on a
+tiny per-expert calib sample. v3 fixes that two ways:
+  1. TRAIN on the larger calib_v3 acts (acts_v3/), so cold experts see more tokens.
+  2. EARLY-STOP on a DISJOINT held-out acts capture (acts_v3_heldout/): every
+     EVAL_EVERY steps, measure aggregate held-out output rel-err; keep the BEST-val
+     params; stop after PATIENCE evals without improvement. We SAVE best-val, not last.
+
+Same schema/mechanics as pv.py (tunes shared codebooks + per-cold-expert scales,
+codes fixed, w13 row-interleaved gate/up). Assignment + converged init = v2's
+(assign_512k_v2.json + parts_conv) since v3 is strictly PV-focused (tiering/MTP unchanged).
+
+GPU phase — needs acts captured FIRST (see run_capture_acts for calib_v3 -> acts_v3/
+and heldout_v3 -> acts_v3_heldout/). Writes parts_pv_v3/layer_{L}_{w13,w2}.pt.
+"""
+from __future__ import annotations
+import argparse, os, time, json
+import torch
+import torch.nn.functional as F
+import ink_common as ic
+from pv import dequant_w13, dequant_w2, expert_forward   # reuse pure dequant/forward
+
+# PV trains on the UNION of text + image + audio acts (multimodal coverage so
+# image/audio-routed experts aren't degraded); each modality captured to its own
+# dir. Only dirs that exist are used.
+import os as _os
+# default = merged multimodal acts + v3 tiering/converge (env-overridable)
+TRAIN_ACT_DIRS = _os.environ.get("INK_ACTS_TRAIN", "/data/inkling-quant/acts_v3_all").split(",")
+VAL_ACT_DIRS   = _os.environ.get("INK_ACTS_VAL", "/data/inkling-quant/acts_v3_all_held").split(",")
+CONV       = _os.environ.get("INK_CONV", "/data/inkling-quant/parts_conv_v3")
+OUT        = _os.environ.get("INK_PV_OUT", "/data/inkling-quant/parts_pv_v3")
+ASSIGN     = _os.environ.get("INK_ASSIGN", "/data/inkling-quant/assign_512k_v3.json")
+OUTD = 6144
+MIN_TOK = 8            # min TRAIN tokens to tune an expert
+MIN_VAL = 4            # min held-out tokens for an expert to count toward the val signal
+MAX_TOK = 2048         # train-token cap/expert (was 8192 -> ~4x faster steps; still ample for a codebook tune)
+MAX_VAL = 1024
+
+
+def cold_experts(layer, assign):
+    hot = set(assign["hot_experts"].get(str(layer), []))
+    return [e for e in range(ic.N_EXPERTS) if e not in hot]
+
+
+def _tokens_multi(dirs, layer, cold, device, cap, seed):
+    """Union tokens across acts dirs (text+image+audio) per cold expert, then cap."""
+    acc = {}   # ci -> list of xe tensors
+    for d in dirs:
+        p = f"{d}/acts_layer_{layer}.pt"
+        if not os.path.exists(p):
+            continue
+        a = torch.load(p, map_location="cpu")
+        x = a["x"].to(device).float(); ids = a["ids"].to(device)
+        for ci, e in enumerate(cold):
+            m = (ids == e).any(1)
+            if int(m.sum()) == 0:
+                continue
+            acc.setdefault(ci, []).append(x[m].contiguous())
+        del x, ids; torch.cuda.empty_cache()
+    gen = torch.Generator(device=device); gen.manual_seed(seed)
+    out = {}
+    for ci, parts in acc.items():
+        xe = torch.cat(parts, 0); n = xe.shape[0]
+        if n > cap:
+            sel = torch.randperm(n, generator=gen, device=device)[:cap]; xe = xe[sel]
+        out[ci] = xe.contiguous()
+    torch.cuda.empty_cache()
+    return out
+
+
+def build_split(layer, cold, device):
+    """Train + val data (multimodal union). For each expert build teacher outputs once.
+    train: [(ci,e,xe,ftea)] with >=MIN_TOK ; val: [(ci,e,xe,ftea)] with >=MIN_VAL."""
+    xt_by = _tokens_multi(TRAIN_ACT_DIRS, layer, cold, device, MAX_TOK, 1234+layer)
+    xv_by = _tokens_multi(VAL_ACT_DIRS, layer, cold, device, MAX_VAL, 9876+layer)
+    tc = ic.ShardReader(ic.BF16_DIR)
+    train, val = [], []
+    for ci, e in enumerate(cold):
+        need = (ci in xt_by and xt_by[ci].shape[0] >= MIN_TOK) or \
+               (ci in xv_by and xv_by[ci].shape[0] >= MIN_VAL)
+        if not need:
+            continue
+        w13 = ic.teacher_expert(tc, layer, "w13", e, device=device).float()
+        w2  = ic.teacher_expert(tc, layer, "w2",  e, device=device).float()
+        if ci in xt_by and xt_by[ci].shape[0] >= MIN_TOK:
+            xe = xt_by[ci]
+            with torch.no_grad(): ft = expert_forward(w13, w2, xe)
+            train.append((ci, e, xe, ft.contiguous()))
+        if ci in xv_by and xv_by[ci].shape[0] >= MIN_VAL:
+            xe = xv_by[ci]
+            with torch.no_grad(): ft = expert_forward(w13, w2, xe)
+            val.append((ci, e, xe, ft.contiguous()))
+        del w13, w2
+    torch.cuda.empty_cache()
+    return train, val
+
+
+@torch.no_grad()
+def agg_relerr(data, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2):
+    if not data: return float("nan")
+    rel = 0.0
+    for (ci, e, xe, ft) in data:
+        Wq13 = dequant_w13(ci, codes13, cb13, s13)
+        Wq2  = dequant_w2(ci, codes2, cb2_0, cb2_1, s2)
+        fq = expert_forward(Wq13, Wq2, xe)
+        rel += float((fq-ft).pow(2).sum() / ft.pow(2).sum().clamp_min(1e-12))
+    return rel / len(data)
+
+
+def pv_layer_v3(layer, max_steps=2000, eval_every=25, patience=6,
+                lr_cb=1e-4, lr_s=1e-3, tune_cb=True, device="cuda", verbose=True):
+    assign = json.load(open(ASSIGN))
+    cold = cold_experts(layer, assign); coldt = torch.tensor(cold)
+    p13 = torch.load(f"{CONV}/layer_{layer}_w13.pt", map_location="cpu")
+    p2  = torch.load(f"{CONV}/layer_{layer}_w2.pt",  map_location="cpu")
+    codes13 = p13["w13_codes"][coldt].to(device)
+    codes2  = p2["w2_codes"][coldt].to(device)
+    cb13  = p13["w13_codebooks"][0].float().to(device).clone().requires_grad_(tune_cb)
+    cb2_0 = p2["w2_codebooks"][0].float().to(device).clone().requires_grad_(tune_cb)
+    cb2_1 = p2["w2_codebooks"][1].float().to(device).clone().requires_grad_(tune_cb)
+    s13 = p13["w13_scales"][coldt].float().to(device).clone().requires_grad_()
+    s2  = p2["w2_scales"][coldt].float().to(device).clone().requires_grad_()
+
+    train, val = build_split(layer, cold, device)
+    if not train:
+        raise RuntimeError(f"L{layer}: no cold experts with >= {MIN_TOK} train tokens")
+    if verbose:
+        print(f"  L{layer}: {len(cold)} cold | train {len(train)} | val {len(val)}"
+              + ("" if val else "  (WARN no held-out acts -> falls back to fixed-step, NO early-stop)"), flush=True)
+
+    val0 = agg_relerr(val, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2) if val else float("nan")
+    params = [{"params": [s13, s2], "lr": lr_s}]
+    if tune_cb: params.append({"params": [cb13, cb2_0, cb2_1], "lr": lr_cb})
+    opt = torch.optim.Adam(params)
+
+    def snapshot():
+        return (cb13.detach().clone(), cb2_0.detach().clone(), cb2_1.detach().clone(),
+                s13.detach().clone(), s2.detach().clone())
+    best_val = val0; best = snapshot(); best_step = 0; bad = 0
+
+    for step in range(1, max_steps+1):
+        opt.zero_grad(set_to_none=True); tot = 0.0
+        for (ci, e, xe, ft) in train:
+            Wq13 = dequant_w13(ci, codes13, cb13, s13)
+            Wq2  = dequant_w2(ci, codes2, cb2_0, cb2_1, s2)
+            loss = F.mse_loss(expert_forward(Wq13, Wq2, xe), ft)
+            loss.backward(); tot += float(loss.detach())
+        opt.step()
+        if val and (step % eval_every == 0 or step == max_steps):
+            vr = agg_relerr(val, codes13, codes2, cb13, cb2_0, cb2_1, s13, s2)
+            improved = vr < best_val - 1e-5
+            if improved:
+                best_val = vr; best = snapshot(); best_step = step; bad = 0
+            else:
+                bad += 1
+            if verbose:
+                print(f"  L{layer} step{step}: train_mse={tot/len(train):.3e} "
+                      f"val_rel={vr:.4f} best={best_val:.4f}@{best_step} bad={bad}", flush=True)
+            if bad >= patience:
+                if verbose: print(f"  L{layer} EARLY-STOP at {step} (best val {best_val:.4f} @ {best_step})", flush=True)
+                break
+
+    # restore best-val params
+    cb13d, cb2_0d, cb2_1d, s13d, s2d = best
+    val1 = agg_relerr(val, codes13, codes2, cb13d, cb2_0d, cb2_1d, s13d, s2d) if val else float("nan")
+    if verbose:
+        print(f"  L{layer} DONE: val_rel {val0:.4f} -> {val1:.4f} (best @ step {best_step})", flush=True)
+
+    with torch.no_grad():
+        full_s13 = p13["w13_scales"].clone(); full_s13[coldt] = s13d.half().cpu()
+        full_s2  = p2["w2_scales"].clone();   full_s2[coldt]  = s2d.half().cpu()
+    common = dict(cold_experts=cold, val_relerr_before=val0, val_relerr_after=val1,
+                  best_step=best_step, max_steps=max_steps, tuned_cb=tune_cb)
+    out13 = {"expert_ids": p13["expert_ids"], "w13_codes": p13["w13_codes"],
+             "w13_codebooks": [cb13d.half().cpu()], "w13_book_entries": p13["w13_book_entries"],
+             "w13_scales": full_s13, **common}
+    out2  = {"expert_ids": p2["expert_ids"], "w2_codes": p2["w2_codes"],
+             "w2_codebooks": [cb2_0d.half().cpu(), cb2_1d.half().cpu()],
+             "w2_book_entries": p2["w2_book_entries"], "w2_scales": full_s2, **common}
+    return out13, out2, dict(val0=val0, val1=val1, best_step=best_step)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layer", type=int, required=True)
+    ap.add_argument("--max-steps", type=int, default=2000)
+    ap.add_argument("--eval-every", type=int, default=25)
+    ap.add_argument("--patience", type=int, default=6)
+    ap.add_argument("--lr_cb", type=float, default=1e-4)
+    ap.add_argument("--lr_s", type=float, default=1e-3)
+    ap.add_argument("--no-tune-cb", action="store_true", help="scales-only (most overfit-safe)")
+    ap.add_argument("--save", action="store_true")
+    a = ap.parse_args()
+    os.makedirs(OUT, exist_ok=True)
+    t0 = time.time()
+    out13, out2, m = pv_layer_v3(a.layer, a.max_steps, a.eval_every, a.patience,
+                                 a.lr_cb, a.lr_s, tune_cb=not a.no_tune_cb)
+    print(f"L{a.layer} PVv3: val_rel {m['val0']:.4f} -> {m['val1']:.4f} "
+          f"(best@{m['best_step']})  ({time.time()-t0:.0f}s)")
+    if a.save:
+        torch.save(out13, f"{OUT}/layer_{a.layer}_w13.pt")
+        torch.save(out2,  f"{OUT}/layer_{a.layer}_w2.pt")
+        print("saved ->", OUT)

code/reap_merge.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Merge per-layer REAP tmp files -> /data/inkling-quant/reap_scores.pt and validate."""
+import torch, glob, os, re
+import ink_common as ic
+
+TMP = os.environ.get("INK_REAP_TMP", "/data/inkling-quant/reap_tmp")
+OUT = os.environ.get("INK_REAP_OUT", "/data/inkling-quant/reap_scores.pt")
+
+def spearman(a, b):
+    a = a.double(); b = b.double()
+    ra = a.argsort().argsort().double()
+    rb = b.argsort().argsort().double()
+    ra = ra - ra.mean(); rb = rb - rb.mean()
+    return float((ra @ rb) / (ra.norm() * rb.norm()).clamp_min(1e-12))
+
+layers = ic.moe_layers()
+scores = {}
+details = {}
+for L in layers:
+    f = f"{TMP}/layer_{L}.pt"
+    if not os.path.exists(f):
+        print(f"MISSING layer {L}"); continue
+    d = torch.load(f, map_location="cpu")
+    scores[L] = d["score"].float()
+    details[L] = d
+
+torch.save(scores, OUT)
+print(f"WROTE {OUT} with {len(scores)} layers (expected 64)")
+
+rs = torch.load("/data/inkling-quant/route_stats.pt", map_location="cpu")
+print("\n== validation: spearman(reap_score, route wmass) for 3 layers ==")
+for L in [layers[0], layers[len(layers)//2], layers[-1]]:
+    wm = rs[L]["wmass"].float()
+    sc = scores[L]
+    print(f"L{L}: spearman={spearman(sc, wm):+.3f}  "
+          f"score mean={sc.mean():.4g} max={sc.max():.4g} min={sc.min():.4g}  "
+          f"relerr mean={details[L]['relerr'].mean():.4f}  "
+          f"experts_with_tokens={int((details[L]['ntok']>0).sum())}/256")
+
+print("\n== per-layer score sanity (all layers) ==")
+allmean = torch.tensor([scores[L].mean() for L in scores])
+allmax  = torch.tensor([scores[L].max() for L in scores])
+zerocnt = sum(int((scores[L]==0).sum()) for L in scores)
+print(f"layers={len(scores)}  score.mean range [{allmean.min():.4g},{allmean.max():.4g}]  "
+      f"score.max range [{allmax.min():.4g},{allmax.max():.4g}]  total zero-experts={zerocnt}")
+print("shapes ok:", all(scores[L].shape==(256,) for L in scores),
+      " finite:", all(torch.isfinite(scores[L]).all() for L in scores))

code/reap_worker.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""Compute TRUE REAP expert-importance scores for a set of MoE layers on ONE gpu.
+score_e = saliency_e * relerr_e
+saliency_e = sum_t g_{t,e} * ||f_e(x_t)||_2   (teacher expert output L2 norm)
+relerr_e   = param-weighted 2-bit AQLM reconstruction error of expert e
+Writes /data/inkling-quant/reap_tmp/layer_{L}.pt = fp32 tensor[256].
+"""
+import os, sys, time, json, traceback
+import torch
+import torch.nn.functional as F
+sys.path.insert(0, "/data/inkling-quant/tools")
+import ink_common as ic
+
+DEV = "cuda"
+R, S, K = 256, 2, 6          # routed, shared(sink), top-k
+ROUTE_SCALE = 8.0
+N13 = 6144 * 6144
+N2  = 6144 * 3072
+TMP = "/data/inkling-quant/reap_tmp"
+os.makedirs(TMP, exist_ok=True)
+
+
+def gate_weights(x_f32, gw, bias, scale):
+    """Replicate InklingGate.select_experts. Returns (topi[T,6] long, rw[T,6] fp32)."""
+    logits = x_f32 @ gw.T                                   # [T,258] fp32
+    sel = torch.sigmoid(logits[:, :R]) + bias.view(1, -1)   # [T,256]
+    _, topi = sel.topk(K, dim=1)                            # [T,6]
+    routed_raw = torch.gather(logits[:, :R], 1, topi)       # [T,6]
+    sink_raw = logits[:, R:R + S]                           # [T,2]
+    active = torch.cat([routed_raw, sink_raw], dim=1)       # [T,8]
+    lp = torch.minimum(active, torch.zeros_like(active)) - torch.log1p(torch.exp(-active.abs()))
+    lp = lp - lp.max(1, keepdim=True).values
+    e = torch.exp(lp)
+    w = e / e.sum(1, keepdim=True) * scale                  # [T,8]
+    return topi, w[:, :K].contiguous()
+
+
+def process_layer(L, reader, treader):
+    t0 = time.time()
+    a = torch.load(f"{__import__('os').environ.get('INK_ACTS','/data/inkling-quant/acts')}/acts_layer_{L}.pt", map_location="cpu")
+    n = int(a["n"])
+    x_f16 = a["x"][:n].to(DEV)                     # [T,6144] fp16
+    x_f32 = x_f16.float()
+
+    gw   = reader.get(f"model.llm.layers.{L}.mlp.gate.weight").to(DEV).float()
+    bias = reader.get(f"model.llm.layers.{L}.mlp.gate.bias").to(DEV).float()
+    gs   = float(reader.get(f"model.llm.layers.{L}.mlp.gate.global_scale")[0])
+    scale = ROUTE_SCALE * gs
+
+    topi, rw = gate_weights(x_f32, gw, bias, scale)     # [T,6],[T,6]
+    del x_f32, gw, bias
+
+    # Full BF16 TEACHER weights for this layer (from BF16_DIR).
+    w13 = treader.get(ic._pfx(L, "w13")).to(DEV).to(torch.bfloat16)   # [256,6144,6144]
+    w2  = treader.get(ic._pfx(L, "w2")).to(DEV).to(torch.bfloat16)    # [256,6144,3072]
+
+    # AQLM parts (w13 CONVERGED, w2).
+    p13 = torch.load(f"{__import__('os').environ.get('INK_CONV_REF','/data/inkling-quant/parts_conv')}/layer_{L}_w13.pt", map_location="cpu")
+    p2  = torch.load(f"/data/inkling-quant/parts/layer_{L}_w2.pt", map_location="cpu")
+
+    saliency = torch.zeros(R, dtype=torch.float32, device=DEV)
+    relerr   = torch.zeros(R, dtype=torch.float32, device=DEV)
+    ntok     = torch.zeros(R, dtype=torch.long, device=DEV)
+
+    for e in range(R):
+        # ---- saliency ----
+        slot_mask = (topi == e)                     # [T,6] bool
+        tok_mask = slot_mask.any(1)                 # [T]
+        idx = tok_mask.nonzero(as_tuple=True)[0]
+        ne = idx.numel()
+        ntok[e] = ne
+        if ne > 0:
+            g_e = (rw * slot_mask).sum(1)[idx]      # [ne] fp32
+            xe = x_f16.index_select(0, idx)         # [ne,6144] fp16
+            we = w13[e]                             # [6144,6144] bf16
+            gate_e = we[0::2].to(torch.bfloat16)    # [3072,6144]
+            up_e   = we[1::2].to(torch.bfloat16)    # [3072,6144]
+            xt = xe.to(torch.bfloat16).T            # [6144,ne]
+            ge = gate_e @ xt                        # [3072,ne]
+            ue = up_e @ xt
+            h = F.silu(ge.float()) * ue.float()     # [3072,ne] fp32
+            fe = w2[e].float() @ h                  # [6144,ne] fp32
+            norms = fe.norm(dim=0)                  # [ne]
+            saliency[e] = (g_e * norms).sum()
+            del xe, xt, ge, ue, h, fe, norms, g_e
+
+        # ---- relerr ----
+        W13 = w13[e].float()
+        dq13 = ic.aqlm_dequant_expert(p13, "w13", e, device=DEV).float()
+        rel13 = torch.linalg.norm(W13 - dq13) / torch.linalg.norm(W13).clamp_min(1e-12)
+        del W13, dq13
+        W2 = w2[e].float()
+        dq2 = ic.aqlm_dequant_expert(p2, "w2", e, device=DEV).float()
+        rel2 = torch.linalg.norm(W2 - dq2) / torch.linalg.norm(W2).clamp_min(1e-12)
+        del W2, dq2
+        relerr[e] = (rel13 * N13 + rel2 * N2) / (N13 + N2)
+
+    score = (saliency * relerr).cpu()
+    out = {
+        "score": score,
+        "saliency": saliency.cpu(),
+        "relerr": relerr.cpu(),
+        "ntok": ntok.cpu(),
+    }
+    torch.save(out, f"{TMP}/layer_{L}.pt")
+    dt = time.time() - t0
+    nz = int((ntok > 0).sum())
+    print(f"[L{L}] done {dt:.1f}s  score mean={score.mean():.4g} max={score.max():.4g} "
+          f"| saliency max={saliency.max():.4g} | relerr mean={relerr.mean():.4f} "
+          f"| experts_with_tokens={nz}/256", flush=True)
+    del w13, w2, x_f16, topi, rw, saliency, relerr, ntok
+    torch.cuda.empty_cache()
+
+
+def main():
+    layers = [int(x) for x in sys.argv[1].split(",")]
+    reader = ic.ShardReader(ic.NVFP4_DIR)      # gate params
+    treader = ic.ShardReader(ic.BF16_DIR)      # bf16 teacher experts
+    print(f"[worker pid={os.getpid()} gpu={os.environ.get('CUDA_VISIBLE_DEVICES')}] layers={layers}", flush=True)
+    for L in layers:
+        try:
+            process_layer(L, reader, treader)
+        except Exception:
+            print(f"[L{L}] ERROR\n{traceback.format_exc()}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

code/run_capture_acts.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+"""Capture per-layer activation samples (routed hidden inputs x + expert ids)
+over the calib corpus, for offline Hessian construction (Phase E converge).
+Dumps /data/inkling-quant/acts/acts_layer_{L}.pt atexit (per TP worker; the
+residual stream is replicated so any worker's dump is complete).
+"""
+import json, os
+
+def main():
+    os.environ.setdefault("INK_ACT_DIR", "/data/inkling-quant/acts")
+    os.environ.setdefault("INK_ACT_CAP", "24000")
+    os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+    from vllm import LLM, SamplingParams
+    prompts = [json.loads(l)["text"] for l in open("/data/inkling-quant/calib.jsonl")]
+    print(f"loaded {len(prompts)} calib prompts; capturing acts (cap {os.environ['INK_ACT_CAP']}/layer)")
+    llm = LLM(model="/data/inkling-nvfp4", tensor_parallel_size=4, max_model_len=8192,
+              gpu_memory_utilization=0.9, enforce_eager=True, trust_remote_code=True)
+    llm.generate(prompts, SamplingParams(max_tokens=1, temperature=0))
+    print("done; acts dump atexit.")
+
+if __name__ == "__main__":
+    main()

code/run_capture_acts_heldout.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+"""Capture per-layer activation samples over the HELD-OUT corpus (heldout.jsonl),
+into /data/inkling-quant/acts_heldout/.  Same hook as run_capture_acts.py but a
+disjoint text set + a cap BELOW the held-out token count (25723) so EVERY layer
+hits the cap mid-run and dumps via the safe incremental path (no atexit truncation).
+GPUs 4-7 only (set via CUDA_VISIBLE_DEVICES before launch)."""
+import json, os
+
+def main():
+    os.environ.setdefault("INK_ACT_DIR", "/data/inkling-quant/acts_heldout")
+    os.environ.setdefault("INK_ACT_CAP", "16000")
+    os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+    from vllm import LLM, SamplingParams
+    prompts = [json.loads(l)["text"] for l in open("/data/inkling-quant/heldout.jsonl")]
+    print(f"loaded {len(prompts)} held-out prompts; capturing acts "
+          f"(cap {os.environ['INK_ACT_CAP']}/layer) -> {os.environ['INK_ACT_DIR']}")
+    llm = LLM(model="/data/inkling-nvfp4", tensor_parallel_size=4, max_model_len=8192,
+              gpu_memory_utilization=0.9, enforce_eager=True, trust_remote_code=True)
+    llm.generate(prompts, SamplingParams(max_tokens=1, temperature=0))
+    print("done; residual layers dump atexit.")
+
+if __name__ == "__main__":
+    main()

code/run_capture_acts_mm.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Multimodal activation capture for Inkling v3 PV.
+
+Inkling is inkling_mm_model (vision hmlp + audio dMel towers -> 6144-dim, feeding
+the SAME MoE backbone as text). A text-only PV pass would leave the image/audio-
+routed experts uncalibrated and could DEGRADE multimodal performance. This driver
+feeds image+text and audio+text prompts through vLLM so the acts hook (INK_ACT_DIR)
+captures MoE activations on multimodal tokens too.
+
+Modality is chosen by INK_MM=image|audio. Writes acts to INK_ACT_DIR (keep image
+and audio in SEPARATE dirs so pv_v3 can merge text+image+audio per expert):
+
+  # image acts (chest X-rays + radiology-style prompt)
+  INK_MM=image INK_ACT_DIR=/data/inkling-quant/acts_v3_img INK_ACT_CAP=40000 \
+    CUDA_VISIBLE_DEVICES=4,5,6,7 python run_capture_acts_mm.py
+  # audio acts (LibriSpeech + transcribe prompt)
+  INK_MM=audio INK_ACT_DIR=/data/inkling-quant/acts_v3_audio INK_ACT_CAP=40000 \
+    CUDA_VISIBLE_DEVICES=4,5,6,7 python run_capture_acts_mm.py
+
+Placeholders (from the model): image=<|content_image|>, audio=<|content_audio_input|>.
+Holds out the last HOLDOUT items of each modality (INK_MM_HOLDOUT dir) for mm ppl.
+"""
+import os, glob, json, wave, random
+import numpy as np
+
+IMG_DIR = "/data/inkling-quant/corpus_v3/raw/radiology/images"
+LS_DIR  = "/data/asr/datasets-asr-normalized/librispeech"
+IMG_PROMPTS = [
+    "<|content_image|>\nDescribe the findings in this chest radiograph.",
+    "<|content_image|>\nReport any abnormalities visible in this chest X-ray.",
+    "<|content_image|>\nWhat pathology, if any, is present in this image?",
+]
+AUD_PROMPT = "<|content_audio_input|>\nTranscribe the spoken audio."
+
+def load_images(n, off=0):
+    """Prefer the multimodal manifest (path,domain,prompt across all image domains);
+    fall back to the chest-X-ray glob if no manifest is given."""
+    from PIL import Image
+    manifest = os.environ.get("INK_MM_MANIFEST")
+    out = []
+    if manifest and os.path.exists(manifest):
+        rows = [json.loads(l) for l in open(manifest)][off:off+n]
+        for r in rows:
+            try:
+                out.append((r["prompt"], {"image": Image.open(r["path"]).convert("RGB")}))
+            except Exception as e:
+                print("  img skip", os.path.basename(r.get("path","?")), str(e)[:50])
+        return out
+    fs = sorted(glob.glob(f"{IMG_DIR}/*.png"))[off:off+n]
+    for i, f in enumerate(fs):
+        try:
+            out.append((IMG_PROMPTS[i % len(IMG_PROMPTS)], {"image": Image.open(f).convert("RGB")}))
+        except Exception as e:
+            print("  img skip", os.path.basename(f), str(e)[:50])
+    return out
+
+def load_audio(n, off=0):
+    fs = sorted(glob.glob(f"{LS_DIR}/dev-clean/*.wav"))[off:off+n]
+    out = []
+    for f in fs:
+        try:
+            w = wave.open(f); sr = w.getframerate()
+            a = np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16).astype(np.float32) / 32768.0
+            out.append((AUD_PROMPT, {"audio": (a, sr)}))
+        except Exception as e:
+            print("  aud skip", os.path.basename(f), str(e)[:50])
+    return out
+
+def main():
+    mm = os.environ.get("INK_MM", "image")
+    os.environ.setdefault("INK_ACT_DIR", f"/data/inkling-quant/acts_v3_{mm}")
+    os.environ.setdefault("INK_ACT_CAP", "40000")
+    os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+    n = int(os.environ.get("INK_MM_N", "300"))
+    off = int(os.environ.get("INK_MM_OFFSET", "0"))   # disjoint held-out slice
+    items = load_images(n, off) if mm == "image" else load_audio(n, off)
+    print(f"INK_MM={mm} items={len(items)} -> {os.environ['INK_ACT_DIR']} (cap {os.environ['INK_ACT_CAP']})", flush=True)
+    if not items:
+        raise SystemExit(f"no {mm} inputs found")
+    from vllm import LLM, SamplingParams
+    llm = LLM(model="/data/inkling-nvfp4", tensor_parallel_size=4, max_model_len=8192,
+              gpu_memory_utilization=0.9, enforce_eager=True, trust_remote_code=True,
+              limit_mm_per_prompt={mm: 1})
+    reqs = [{"prompt": p, "multi_modal_data": d} for p, d in items]
+    llm.generate(reqs, SamplingParams(max_tokens=1, temperature=0))
+    print("done; acts dump atexit.", flush=True)
+
+if __name__ == "__main__":
+    main()

code/run_capture_acts_v3.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+"""v3 activation capture — parameterized corpus / out-dir / per-layer cap.
+
+The v2 PV starvation came from INK_ACT_CAP=24000 tokens/layer (the "~25k" that
+overfit). v3 raises the cap AND uses the bigger disjoint corpora:
+
+  # TRAIN acts (what PV tunes on) — big cap for cold-expert coverage
+  INK_CORPUS=/data/inkling-quant/calib_v3.jsonl INK_ACT_DIR=/data/inkling-quant/acts_v3 \
+    INK_ACT_CAP=80000 CUDA_VISIBLE_DEVICES=4,5,6,7 python run_capture_acts_v3.py
+
+  # VAL acts (PV early-stop signal) — disjoint corpus, smaller cap is fine
+  INK_CORPUS=/data/inkling-quant/heldout_v3.jsonl INK_ACT_DIR=/data/inkling-quant/acts_v3_heldout \
+    INK_ACT_CAP=24000 CUDA_VISIBLE_DEVICES=4,5,6,7 python run_capture_acts_v3.py
+
+Disk: ~ cap * 6144 * 2 bytes * 64 layers  (e.g. cap=80000 -> ~63 GB). Check df first.
+"""
+import json, os
+
+def main():
+    corpus = os.environ.setdefault("INK_CORPUS", "/data/inkling-quant/calib_v3.jsonl")
+    os.environ.setdefault("INK_ACT_DIR", "/data/inkling-quant/acts_v3")
+    os.environ.setdefault("INK_ACT_CAP", "80000")
+    os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+    from vllm import LLM, SamplingParams
+    prompts = [json.loads(l)["text"] for l in open(corpus)]
+    print(f"corpus={corpus} n={len(prompts)} -> {os.environ['INK_ACT_DIR']} "
+          f"(cap {os.environ['INK_ACT_CAP']}/layer)", flush=True)
+    llm = LLM(model="/data/inkling-nvfp4", tensor_parallel_size=4, max_model_len=8192,
+              gpu_memory_utilization=0.9, enforce_eager=True, trust_remote_code=True)
+    llm.generate(prompts, SamplingParams(max_tokens=1, temperature=0))
+    print("done; acts dump atexit.", flush=True)
+
+if __name__ == "__main__":
+    main()

code/run_route_stats.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+"""Run the calib corpus through Inkling-NVFP4 offline (vLLM LLM) with
+INK_ROUTE_STATS set, so moe.py accumulates per-layer per-expert routing
+frequency + weight mass. Prefill dominates -> every calib token exercises
+routing. Stats dumped atexit to $INK_ROUTE_STATS (in each TP worker; the
+gate is replicated so any worker's dump is complete).
+"""
+import json, os
+
+def main():
+    os.environ.setdefault("INK_ROUTE_STATS", "/data/inkling-quant/route_stats.pt")
+    os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+    from vllm import LLM, SamplingParams
+    CALIB = "/data/inkling-quant/calib.jsonl"
+    prompts = [json.loads(l)["text"] for l in open(CALIB)]
+    print(f"loaded {len(prompts)} calib prompts")
+    llm = LLM(
+        model="/data/inkling-nvfp4",
+        tensor_parallel_size=4,
+        max_model_len=8192,
+        gpu_memory_utilization=0.9,
+        enforce_eager=True,
+        trust_remote_code=True,
+    )
+    sp = SamplingParams(max_tokens=1, temperature=0)  # prefill routing only
+    out = llm.generate(prompts, sp)
+    print(f"processed {len(out)} prompts; route stats dump atexit.")
+
+if __name__ == "__main__":
+    main()

code/solve_tier.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+"""Byte-exact expert tiering solver for the Inkling hybrid.
+Ranks experts globally by routing mass; upgrades the highest-mass experts from
+AQLM (cold) to NVFP4 (hot) until the expert VRAM budget is spent, with per-layer
+floor/cap. Reports achievable hot% + avg bpw for a target deployment budget.
+
+Deployment target = 4x96GB RTX6000 = 384GB (we emulate on B200; ship for 384GB).
+"""
+import argparse, json, math
+import torch
+import ink_common as ic
+
+# exact per-expert shapes (in params)
+W13 = (6144, 6144); W2 = (6144, 3072)
+P13 = W13[0]*W13[1]; P2 = W2[0]*W2[1]
+GDIM = 8
+
+def nvfp4_bytes():
+    # 4-bit packed + fp8 block-16 scale + fp32 scale2 (negligible)
+    def lin(o,i): return o*i*0.5 + o*(i/16)*1.0   # u8 packed + f8 scale
+    return lin(*W13) + lin(*W2)
+
+def aqlm_bytes(w13_books, w2_books, book2_entries):
+    code_b = 2  # book0 int16 (65536)
+    b2 = 1 if book2_entries <= 256 else 2
+    def proj(o,i,books):
+        groups = o*(i//GDIM)
+        by = groups*code_b + (groups*b2*(books-1) if books>1 else 0)
+        return by + o*2  # + fp16 per-row scale
+    return proj(*W13, w13_books) + proj(*W2, w2_books)
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--budget-gb", type=float, default=384.0)
+    ap.add_argument("--nonexpert-gb", type=float, default=40.0)
+    ap.add_argument("--kv-gb", type=float, default=11.0)      # ~512k fp8
+    ap.add_argument("--workspace-gb", type=float, default=25.0)
+    ap.add_argument("--w13-books", type=int, default=1)
+    ap.add_argument("--w2-books", type=int, default=2)
+    ap.add_argument("--book2-entries", type=int, default=65536)
+    ap.add_argument("--floor", type=int, default=4)   # min hot per layer
+    ap.add_argument("--cap", type=int, default=220)   # max hot per layer
+    ap.add_argument("--save", default="")
+    ap.add_argument("--route-stats", default="/data/inkling-quant/route_stats.pt")
+    ap.add_argument("--scores", default="", help="optional pt {L:tensor[256]} (e.g. REAP) to rank by instead of routing mass")
+    a = ap.parse_args()
+
+    st = torch.load(a.route_stats)
+    layers = sorted(st.keys()); E = ic.N_EXPERTS
+    mass = {L: st[L]["wmass"].clone() for L in layers}  # router-weight mass per expert
+    if a.scores:
+        sc = torch.load(a.scores)
+        mass = {L: sc[L].clone().float() for L in layers}   # rank by supplied scores (REAP)
+        print(f"ranking by --scores {a.scores} (overrides routing mass)")
+
+    hotB = nvfp4_bytes(); coldB = aqlm_bytes(a.w13_books, a.w2_books, a.book2_entries)
+    n_exp = len(layers)*E
+    expert_budget = (a.budget_gb - a.nonexpert_gb - a.kv_gb - a.workspace_gb) * 1e9
+    all_cold = n_exp * coldB
+    print(f"per-expert: NVFP4(hot)={hotB/1e6:.1f}MB  AQLM(cold w13={a.w13_books}b/w2={a.w2_books}b@{a.book2_entries})={coldB/1e6:.1f}MB")
+    print(f"experts total={n_exp}  all-cold={all_cold/1e9:.1f}GB  expert budget={expert_budget/1e9:.1f}GB")
+    spare = expert_budget - all_cold
+    if spare < 0:
+        print(f"INFEASIBLE: all-cold ({all_cold/1e9:.1f}GB) exceeds budget ({expert_budget/1e9:.1f}GB) by {-spare/1e9:.1f}GB. Need cheaper cold scheme.")
+        return
+    max_upgrades = int(spare // (hotB - coldB))
+    print(f"spare after all-cold={spare/1e9:.1f}GB  -> up to {max_upgrades} hot upgrades ({100*max_upgrades/n_exp:.1f}% hot)")
+
+    # global rank by mass, apply floor first (guarantee floor hot per layer), then fill by mass under cap
+    assign = {L: torch.zeros(E, dtype=torch.bool) for L in layers}
+    used = 0
+    # floor
+    for L in layers:
+        top = torch.argsort(mass[L], descending=True)[:a.floor]
+        assign[L][top] = True; used += a.floor
+    # remaining budget by global mass ranking (respect cap)
+    pool = []
+    for L in layers:
+        order = torch.argsort(mass[L], descending=True).tolist()
+        for rank,e in enumerate(order):
+            if not assign[L][e]:
+                pool.append((float(mass[L][e]), L, e, rank))
+    pool.sort(reverse=True)
+    for m,L,e,rank in pool:
+        if used >= max_upgrades: break
+        if int(assign[L].sum()) >= a.cap: continue
+        assign[L][e] = True; used += 1
+    hot_total = used
+    # per-band report
+    def band(lo,hi): return sum(int(assign[L].sum()) for L in layers if lo<=L<=hi)/max(1,len([L for L in layers if lo<=L<=hi]))
+    print(f"\nHOT experts assigned: {hot_total} ({100*hot_total/n_exp:.1f}%)")
+    print(f"per-layer mean hot: early(2-15)={band(2,15):.0f}  mid(30-45)={band(30,45):.0f}  deep(55-65)={band(55,65):.0f}  (of 256)")
+    avg_bpw = ((hot_total*hotB + (n_exp-hot_total)*coldB)*8) / (n_exp*(P13+P2))
+    print(f"avg expert bpw={avg_bpw:.2f}  total expert VRAM={((hot_total*hotB+(n_exp-hot_total)*coldB))/1e9:.1f}GB")
+    if a.save:
+        out = {str(L): sorted(torch.nonzero(assign[L]).flatten().tolist()) for L in layers}
+        json.dump({"hot_experts": out, "config": vars(a), "hot_total": hot_total,
+                   "hot_bytes": hotB, "cold_bytes": coldB}, open(a.save,"w"))
+        print("saved assignment ->", a.save)
+
+if __name__ == "__main__":
+    main()

code/v1_vs_v2_heldout.py

@@ -0,0 +1,33 @@
+import torch, json, random, statistics as st
+import ink_common as ic
+random.seed(0); dev="cuda"
+# v1 cold = converged w13 (parts_conv) + INIT w2 (parts);  v2 cold = converged w13 + CONVERGED w2 (parts_conv)
+tc=ic.ShardReader(ic.BF16_DIR)
+def eout(W13,W2,x):
+    g=W13[0::2]; u=W13[1::2]
+    h=torch.nn.functional.silu(x@g.t())*(x@u.t()); return h@W2.t()
+layers=[3,8,15,22,30,38,45,52,58,63]
+v1=[]; v2=[]
+for L in layers:
+    a=torch.load(f"/data/inkling-quant/acts_heldout/acts_layer_{L}.pt",map_location=dev)
+    x=a["x"].float(); ids=a["ids"]
+    # cold = experts not hot in v2 assign (same set for both comparisons; tiering handled separately)
+    hot=set(json.load(open("/data/inkling-quant/assign_512k_v2.json"))["hot_experts"][str(L)])
+    cold=[e for e in range(256) if e not in hot]
+    c13=torch.load(f"/data/inkling-quant/parts_conv/layer_{L}_w13.pt",map_location="cpu")  # converged w13 (both)
+    w2_init=torch.load(f"/data/inkling-quant/parts/layer_{L}_w2.pt",map_location="cpu")     # v1 w2 (init)
+    w2_conv=torch.load(f"/data/inkling-quant/parts_conv/layer_{L}_w2.pt",map_location="cpu")# v2 w2 (converged)
+    W13=None
+    for e in random.sample(cold,min(12,len(cold))):
+        m=(ids==e).any(1)
+        if int(m.sum())<8: continue
+        xe=x[m][:2048]
+        ft=eout(ic.teacher_expert(tc,L,"w13",e,device=dev).float(),ic.teacher_expert(tc,L,"w2",e,device=dev).float(),xe); nt=ft.norm()
+        Wc13=ic.aqlm_dequant_expert(c13,"w13",e,device=dev).float()
+        v1.append(((eout(Wc13,ic.aqlm_dequant_expert(w2_init,"w2",e,device=dev).float(),xe)-ft).norm()/nt).item())
+        v2.append(((eout(Wc13,ic.aqlm_dequant_expert(w2_conv,"w2",e,device=dev).float(),xe)-ft).norm()/nt).item())
+a=st.mean(v1); b=st.mean(v2)
+print(f"HELD-OUT per-expert output rel-err over {len(v1)} cold experts:")
+print(f"  v1 cold (conv-w13 + INIT-w2):      {a:.4f}")
+print(f"  v2 cold (conv-w13 + CONVERGED-w2): {b:.4f}")
+print(f"  v2 vs v1 on HELD-OUT: {100*(1-b/a):+.1f}%  -> {'v2 cold experts BETTER (w2-converge generalizes)' if b < a*0.995 else 'no held-out gain'}")

code/validate_gate_ckpt.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Validate /data/inkling-512k-gate:
+ (a) all shards referenced by the output index exist + non-weight files copied;
+ (b) output index present & its referenced shards all exist;
+ (c) cold expert w13 dequant rel_err vs BF16 teacher ~0.30-0.32 (carries AQLM err);
+ (d) hot expert unchanged vs /data/inkling-nvfp4 (rel_err ~0.09, byte-identical).
+"""
+import json, os, torch, filecmp
+import ink_common as ic
+from safetensors import safe_open
+
+OUT = "/data/inkling-512k-gate"
+SRC = ic.NVFP4_DIR
+ASSIGN = "/data/inkling-quant/assign_512k.json"
+
+
+def main():
+    dev = "cuda"
+    ok = True
+    # ---- (a)/(b) file inventory ----
+    idxp = f"{OUT}/model.safetensors.index.json"
+    assert os.path.exists(idxp), "output index missing"
+    wm = json.load(open(idxp))["weight_map"]
+    shards = sorted(set(wm.values()))
+    missing = [s for s in shards if not os.path.exists(f"{OUT}/{s}")]
+    print(f"[a/b] index references {len(shards)} shards; missing: {missing}")
+    ok &= not missing
+    n_st = len([f for f in os.listdir(OUT) if f.endswith('.safetensors')])
+    print(f"[a] .safetensors files in output: {n_st}")
+    SKIP = {".cache"}  # HF download cache dir, intentionally not copied
+    nonweight = [f for f in os.listdir(SRC) if not f.endswith('.safetensors') and f not in SKIP]
+    miss_nw = [f for f in nonweight if not os.path.exists(f"{OUT}/{f}")]
+    print(f"[a] non-weight source files: {len(nonweight)}; missing in output: {miss_nw}")
+    ok &= not miss_nw
+
+    hot = json.load(open(ASSIGN))["hot_experts"]
+    L = 10
+    hot_L = set(hot[str(L)])
+    cold_e = next(e for e in range(ic.N_EXPERTS) if e not in hot_L)
+    hot_e = sorted(hot_L)[0]
+
+    tR = ic.ShardReader(ic.BF16_DIR)
+    oR = ic.ShardReader(OUT)   # the new checkpoint
+    nR = ic.ShardReader(SRC)   # original nvfp4
+
+    # ---- (c) cold expert now carries AQLM error ----
+    proj = "w13"
+    teach = ic.teacher_expert(tR, L, proj, cold_e, dev).float()
+    coldW = ic.dequant_nvfp4_expert(oR, L, proj, cold_e, dev).float()
+    cold_err = ((coldW - teach).norm() / teach.norm()).item()
+    print(f"[c] COLD e{cold_e} {proj} rel_err vs teacher = {cold_err:.4f}  (expect 0.30-0.32)")
+    ok &= 0.29 <= cold_err <= 0.33
+
+    # ---- (d) hot expert unchanged ----
+    teach_h = ic.teacher_expert(tR, L, proj, hot_e, dev).float()
+    hotW = ic.dequant_nvfp4_expert(oR, L, proj, hot_e, dev).float()
+    hot_err = ((hotW - teach_h).norm() / teach_h.norm()).item()
+    print(f"[d] HOT  e{hot_e} {proj} rel_err vs teacher = {hot_err:.4f}  (expect ~0.09)")
+    ok &= 0.07 <= hot_err <= 0.11
+    # byte-identical hot slice vs original nvfp4 (all three tensors)
+    p = f"model.llm.layers.{L}.mlp.experts.{proj}_weight"
+    for suf in ["", ".scale", ".scale2"]:
+        a = nR.get(p + suf)[hot_e]
+        b = oR.get(p + suf)[hot_e]
+        same = torch.equal(a, b)
+        print(f"[d] HOT byte-identical {p+suf}[{hot_e}] vs original nvfp4: {same}")
+        ok &= same
+    # sanity: cold slice DIFFERS from original nvfp4
+    diff = not torch.equal(nR.get(p)[cold_e], oR.get(p)[cold_e])
+    print(f"[c] COLD packed[{cold_e}] differs from original nvfp4: {diff}")
+    ok &= diff
+
+    # ---- (e) layer 2 (full-precision bf16 special layer) ----
+    L2 = 2
+    hot2 = set(json.load(open(ASSIGN))["hot_experts"]["2"])
+    c2 = next(e for e in range(ic.N_EXPERTS) if e not in hot2)
+    h2 = sorted(hot2)[0]
+    p2 = f"model.llm.layers.{L2}.mlp.experts.w13_weight"
+    print(f"\n[e] layer2 w13_weight dtype in output: {oR.get(p2).dtype} (expect bfloat16, unpacked)")
+    t2 = ic.teacher_expert(tR, L2, "w13", c2, dev).float()
+    got2 = oR.get(p2)[c2].to(dev).float()
+    e2 = ((got2 - t2).norm() / t2.norm()).item()
+    print(f"[e] layer2 COLD e{c2} w13 rel_err vs teacher = {e2:.4f}  (expect ~0.30, pure AQLM bf16)")
+    ok &= 0.28 <= e2 <= 0.32
+    same2 = torch.equal(nR.get(p2)[h2], oR.get(p2)[h2])
+    print(f"[e] layer2 HOT e{h2} w13 byte-identical vs original: {same2}")
+    ok &= same2
+
+    print("\nVALIDATION:", "PASS" if ok else "FAIL")
+    return 0 if ok else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/validate_hybrid.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Validate the assembled /data/inkling-512k-hybrid checkpoint before upload."""
+import json, os, glob
+import torch
+from safetensors import safe_open
+import build_hybrid as B
+import ink_common as ic
+
+OUT = B.OUT
+
+def validate():
+    ok = True
+    idx = json.load(open(f"{OUT}/model.safetensors.index.json"))
+    wm = idx["weight_map"]
+    total = idx["metadata"]["total_size"]
+    ref_shards = sorted(set(wm.values()))
+    disk_shards = sorted(os.path.basename(p) for p in glob.glob(f"{OUT}/model-*.safetensors"))
+
+    print(f"[chk] index tensors: {len(wm)}  total_size: {total/1e9:.1f} GB")
+    print(f"[chk] shards referenced: {len(ref_shards)}  on disk: {len(disk_shards)}")
+
+    missing = [s for s in ref_shards if not os.path.exists(f"{OUT}/{s}")]
+    extra = [s for s in disk_shards if s not in set(ref_shards)]
+    if missing:
+        ok = False; print(f"  !! MISSING shards: {missing}")
+    if extra:
+        print(f"  ~~ extra shards on disk (not in index): {extra}")
+    if not missing and not extra:
+        print("  ok: shard set matches index exactly")
+
+    # size sanity: smaller build must fit a 384 GB (4x96GB) machine.
+    gb = total / 1e9
+    # per-book uint8 cold codes (~ -38 GB) + layer-2 compressed (~ -25 GB) => ~345-355 GB
+    if not (330 <= gb <= 375):
+        ok = False; print(f"  !! total size {gb:.1f} GB out of expected band (330-375)")
+    else:
+        print(f"  ok: total size {gb:.1f} GB in range")
+    if gb < 384:
+        print(f"  ok: total size {gb:.1f} GB < 384 GB (fits 4x96GB machine)")
+    else:
+        ok = False; print(f"  !! total size {gb:.1f} GB does NOT fit 384 GB machine")
+
+    # every tensor resolvable in its shard
+    open_cache = {}
+    def keys_of(sh):
+        if sh not in open_cache:
+            open_cache[sh] = set(safe_open(f"{OUT}/{sh}", framework="pt").keys())
+        return open_cache[sh]
+    bad = [n for n, sh in wm.items() if n not in keys_of(sh)]
+    if bad:
+        ok = False; print(f"  !! {len(bad)} tensors not found in their shard, e.g. {bad[:3]}")
+    else:
+        print(f"  ok: all {len(wm)} tensors resolvable in their shards")
+
+    # readback: hot slabs, per-book cold codes, hot_ids from a packed layer
+    packed = sorted(B.packed_layers()); hot = B.hot_map()
+    L = packed[0]
+    nv = ic.ShardReader(OUT)  # reads assembled output via its index
+    checks = {
+        f"model.llm.layers.{L}.mlp.experts.w13_hot_weight": ("uint8", 3),
+        f"model.llm.layers.{L}.mlp.experts.w13_cold_codes.0": ("int16", 3),  # w13 book0 @65536
+        f"model.llm.layers.{L}.mlp.experts.w2_cold_codes.0": ("int16", 3),   # w2  book0 @65536
+        f"model.llm.layers.{L}.mlp.experts.w2_cold_codes.1": ("uint8", 3),   # w2  book1 @256 -> uint8
+        f"model.llm.layers.{L}.mlp.experts.hot_ids": ("int32", 1),
+        f"model.llm.layers.{L}.mlp.experts.cold_ids": ("int32", 1),
+    }
+    for name, (dt, ndim) in checks.items():
+        t = nv.get(name)
+        print(f"  readback {name.split('.mlp.')[-1]}: {t.dtype} {tuple(t.shape)}")
+        if str(t.dtype).split(".")[-1] != dt or t.dim() != ndim:
+            ok = False; print(f"    !! unexpected dtype/ndim (want {dt} ndim {ndim})")
+    # confirm uint8 book1 values are within 0..255 (lossless) and codebook has <=256 rows
+    c1 = nv.get(f"model.llm.layers.{L}.mlp.experts.w2_cold_codes.1")
+    cb1 = nv.get(f"model.llm.layers.{L}.mlp.experts.w2_cold_codebook.1")
+    if int(c1.max()) > cb1.shape[0] - 1:
+        ok = False; print(f"  !! w2 book1 code {int(c1.max())} out of range for {cb1.shape[0]} entries")
+    else:
+        print(f"  ok: w2 book1 uint8 codes in 0..{int(c1.max())} (codebook {cb1.shape[0]} rows)")
+    # hot_ids/cold_ids consistency vs assignment
+    hids = nv.get(f"model.llm.layers.{L}.mlp.experts.hot_ids").tolist()
+    cids = nv.get(f"model.llm.layers.{L}.mlp.experts.cold_ids").tolist()
+    if sorted(hids) != sorted(hot[L]):
+        ok = False; print(f"  !! hot_ids mismatch layer {L}")
+    if sorted(hids + cids) != list(range(ic.N_EXPERTS)):
+        ok = False; print(f"  !! hot_ids+cold_ids do not partition 0..255")
+    else:
+        print(f"  ok: hot_ids({len(hids)})+cold_ids({len(cids)}) partition 0..255 for layer {L}")
+    # hot_weight n matches n_hot
+    hw = nv.get(f"model.llm.layers.{L}.mlp.experts.w13_hot_weight")
+    if hw.shape[0] != len(hot[L]):
+        ok = False; print(f"  !! hot_weight expert-dim {hw.shape[0]} != n_hot {len(hot[L])}")
+
+    # layer 2 is now a hybrid layer: bf16 HOT slabs (no scales), AQLM cold, fused weights dropped
+    for proj in ("w13", "w2"):
+        hb = f"model.llm.layers.2.mlp.experts.{proj}_hot_bf16"
+        if hb not in wm:
+            ok = False; print(f"  !! layer-2 {proj}_hot_bf16 MISSING")
+        else:
+            t = nv.get(hb)
+            print(f"  layer2 {proj}_hot_bf16: {t.dtype} {tuple(t.shape)}")
+            if t.dtype != torch.bfloat16 or t.dim() != 3:
+                ok = False; print(f"    !! layer-2 {proj}_hot_bf16 not bf16/3-D")
+            if t.shape[0] != len(hot[2]):
+                ok = False; print(f"    !! layer-2 {proj}_hot_bf16 expert-dim {t.shape[0]} != n_hot {len(hot[2])}")
+        if f"{hb}.scale" in wm or f"{hb}.scale2" in wm:
+            ok = False; print(f"    !! layer-2 {proj}_hot_bf16 unexpectedly has scale")
+        # cold codes present for layer 2
+        if f"model.llm.layers.2.mlp.experts.{proj}_cold_codes.0" not in wm:
+            ok = False; print(f"  !! layer-2 {proj}_cold_codes.0 MISSING")
+        # fused uncompressed bf16 must be dropped
+        if f"model.llm.layers.2.mlp.experts.{proj}_weight" in wm:
+            ok = False; print(f"  !! layer-2 {proj}_weight fused bf16 NOT dropped")
+    # NO layer (incl. layer 2) may carry fused expert weights in the output
+    built = set(hot)
+    leaked = [n for n in wm if B.EXP_ANY.match(n) and int(B.EXP_ANY.match(n).group(1)) in built]
+    if leaked:
+        ok = False; print(f"  !! fused expert weights leaked into output: {leaked[:3]}")
+    else:
+        print(f"  ok: no fused expert weights in output")
+
+    # side files
+    for fn in ("config.json", "hf_quant_config.json", "tokenizer.json", "README.md",
+               "tokenizer_config.json", "chat_template.jinja", "tiktoken/tokenizer.model"):
+        if not os.path.exists(f"{OUT}/{fn}"):
+            print(f"  ~~ side file missing: {fn}")
+    hq = json.load(open(f"{OUT}/hf_quant_config.json"))
+    if "aqlm_hybrid" not in hq:
+        ok = False; print("  !! hf_quant_config.json missing aqlm_hybrid block")
+    cfg = json.load(open(f"{OUT}/config.json"))
+    if cfg.get("quantization_config", {}).get("quant_method") != "inkling_nvfp4_aqlm_hybrid":
+        ok = False; print("  !! config.json missing quantization_config.quant_method")
+
+    print(f"\n[VALIDATE] {'PASS' if ok else 'FAIL'}")
+    return ok
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(0 if validate() else 1)

config.json

@@ -0,0 +1,140 @@
+{
+  "architectures": [
+    "InklingForConditionalGeneration"
+  ],
+  "model_type": "inkling_mm_model",
+  "eos_token_id": 200006,
+  "text_config": {
+    "model_max_length": 1048576,
+    "torch_dtype": "bfloat16",
+    "hidden_size": 6144,
+    "num_hidden_layers": 66,
+    "vocab_size": 201024,
+    "num_attention_heads": 64,
+    "num_key_value_heads": 8,
+    "head_dim": 128,
+    "d_rel": 16,
+    "rel_extent": 1024,
+    "q_bias": false,
+    "o_bias": false,
+    "log_scaling_n_floor": 128000,
+    "log_scaling_alpha": 0.1,
+    "rms_norm_eps": 1e-06,
+    "use_embed_norm": true,
+    "local_layer_ids": [
+      0,
+      1,
+      2,
+      3,
+      4,
+      6,
+      7,
+      8,
+      9,
+      10,
+      12,
+      13,
+      14,
+      15,
+      16,
+      18,
+      19,
+      20,
+      21,
+      22,
+      24,
+      25,
+      26,
+      27,
+      28,
+      30,
+      31,
+      32,
+      33,
+      34,
+      36,
+      37,
+      38,
+      39,
+      40,
+      42,
+      43,
+      44,
+      45,
+      46,
+      48,
+      49,
+      50,
+      51,
+      52,
+      54,
+      55,
+      56,
+      57,
+      58,
+      60,
+      61,
+      62,
+      63,
+      64
+    ],
+    "dense_mlp_idx": 2,
+    "use_sconv": true,
+    "sconv_kernel_size": 4,
+    "unpadded_vocab_size": 200058,
+    "logits_mup_width_multiplier": 24.0,
+    "final_logit_softcapping": null,
+    "swa_head_dim": 128,
+    "swa_num_attention_heads": 64,
+    "swa_num_key_value_heads": 16,
+    "sliding_window_size": 512,
+    "n_routed_experts": 256,
+    "num_experts_per_tok": 6,
+    "n_shared_experts": 2,
+    "shared_expert_sink": true,
+    "dense_intermediate_size": 24576,
+    "intermediate_size": 3072,
+    "route_scale": 8.0,
+    "use_gate_bias": true,
+    "gate_activation": "sigmoid",
+    "norm_after_topk": true,
+    "use_global_scale": true
+  },
+  "audio_config": {
+    "decoder_dmodel": 6144,
+    "n_mel_bins": 80,
+    "mel_vocab_size": 16,
+    "bias": false,
+    "dmel_min_value": -7.0,
+    "dmel_max_value": 2.0,
+    "use_audio_norm": true,
+    "audio_mode": "dmel"
+  },
+  "vision_config": {
+    "vision_encoder_type": "hmlp",
+    "decoder_dmodel": 6144,
+    "patch_size": 40,
+    "temporal_patch_size": 2,
+    "n_channels": 3,
+    "n_layers": 4,
+    "use_vision_norm": true
+  },
+  "mtp_config": {
+    "num_nextn_predict_layers": 8,
+    "chain_hidden_post_norm": false,
+    "local_layer_ids": [
+      0,
+      2,
+      4,
+      5,
+      6,
+      7
+    ]
+  },
+  "quantization_config": {
+    "quant_method": "inkling_nvfp4_aqlm_hybrid",
+    "base_nvfp4": "thinkingmachines/Inkling-NVFP4",
+    "group_size": 8,
+    "note": "Mixed per-expert quant: hot experts ModelOpt-NVFP4, cold experts AQLM-2bit. Full spec in hf_quant_config.json['aqlm_hybrid'] and README.md. Requires a custom fused-MoE kernel; does NOT load on stock vLLM."
+  }
+}

hf_quant_config.json

@@ -0,0 +1,1049 @@
+{
+  "quantization": {
+    "quant_algo": "NVFP4",
+    "kv_cache_quant_algo": "none",
+    "exclude_modules": [
+      "model.audio.encoder",
+      "model.audio.final_norm",
+      "model.llm.embed",
+      "model.llm.embed_norm",
+      "model.llm.norm",
+      "model.llm.unembed",
+      "model.visual.final_norm",
+      "model.visual.layers.linear_0",
+      "model.visual.layers.linear_1",
+      "model.visual.layers.linear_2",
+      "model.visual.layers.linear_3",
+      "model.visual.layers.norm_0",
+      "model.visual.layers.norm_1",
+      "model.visual.layers.norm_2",
+      "model.llm.layers.0.attn",
+      "model.llm.layers.0.attn_norm",
+      "model.llm.layers.0.attn_sconv",
+      "model.llm.layers.0.mlp.global_scale",
+      "model.llm.layers.0.mlp.w13_dn",
+      "model.llm.layers.0.mlp.w2_md",
+      "model.llm.layers.0.mlp_norm",
+      "model.llm.layers.0.mlp_sconv",
+      "model.llm.layers.1.attn",
+      "model.llm.layers.1.attn_norm",
+      "model.llm.layers.1.attn_sconv",
+      "model.llm.layers.1.mlp.global_scale",
+      "model.llm.layers.1.mlp.w13_dn",
+      "model.llm.layers.1.mlp.w2_md",
+      "model.llm.layers.1.mlp_norm",
+      "model.llm.layers.1.mlp_sconv",
+      "model.llm.layers.10.attn",
+      "model.llm.layers.10.attn_norm",
+      "model.llm.layers.10.attn_sconv",
+      "model.llm.layers.10.mlp.gate",
+      "model.llm.layers.10.mlp.shared_experts",
+      "model.llm.layers.10.mlp_norm",
+      "model.llm.layers.10.mlp_sconv",
+      "model.llm.layers.11.attn",
+      "model.llm.layers.11.attn_norm",
+      "model.llm.layers.11.attn_sconv",
+      "model.llm.layers.11.mlp.gate",
+      "model.llm.layers.11.mlp.shared_experts",
+      "model.llm.layers.11.mlp_norm",
+      "model.llm.layers.11.mlp_sconv",
+      "model.llm.layers.12.attn",
+      "model.llm.layers.12.attn_norm",
+      "model.llm.layers.12.attn_sconv",
+      "model.llm.layers.12.mlp.gate",
+      "model.llm.layers.12.mlp.shared_experts",
+      "model.llm.layers.12.mlp_norm",
+      "model.llm.layers.12.mlp_sconv",
+      "model.llm.layers.13.attn",
+      "model.llm.layers.13.attn_norm",
+      "model.llm.layers.13.attn_sconv",
+      "model.llm.layers.13.mlp.gate",
+      "model.llm.layers.13.mlp.shared_experts",
+      "model.llm.layers.13.mlp_norm",
+      "model.llm.layers.13.mlp_sconv",
+      "model.llm.layers.14.attn",
+      "model.llm.layers.14.attn_norm",
+      "model.llm.layers.14.attn_sconv",
+      "model.llm.layers.14.mlp.gate",
+      "model.llm.layers.14.mlp.shared_experts",
+      "model.llm.layers.14.mlp_norm",
+      "model.llm.layers.14.mlp_sconv",
+      "model.llm.layers.15.attn",
+      "model.llm.layers.15.attn_norm",
+      "model.llm.layers.15.attn_sconv",
+      "model.llm.layers.15.mlp.gate",
+      "model.llm.layers.15.mlp.shared_experts",
+      "model.llm.layers.15.mlp_norm",
+      "model.llm.layers.15.mlp_sconv",
+      "model.llm.layers.16.attn",
+      "model.llm.layers.16.attn_norm",
+      "model.llm.layers.16.attn_sconv",
+      "model.llm.layers.16.mlp.gate",
+      "model.llm.layers.16.mlp.shared_experts",
+      "model.llm.layers.16.mlp_norm",
+      "model.llm.layers.16.mlp_sconv",
+      "model.llm.layers.17.attn",
+      "model.llm.layers.17.attn_norm",
+      "model.llm.layers.17.attn_sconv",
+      "model.llm.layers.17.mlp.gate",
+      "model.llm.layers.17.mlp.shared_experts",
+      "model.llm.layers.17.mlp_norm",
+      "model.llm.layers.17.mlp_sconv",
+      "model.llm.layers.18.attn",
+      "model.llm.layers.18.attn_norm",
+      "model.llm.layers.18.attn_sconv",
+      "model.llm.layers.18.mlp.gate",
+      "model.llm.layers.18.mlp.shared_experts",
+      "model.llm.layers.18.mlp_norm",
+      "model.llm.layers.18.mlp_sconv",
+      "model.llm.layers.19.attn",
+      "model.llm.layers.19.attn_norm",
+      "model.llm.layers.19.attn_sconv",
+      "model.llm.layers.19.mlp.gate",
+      "model.llm.layers.19.mlp.shared_experts",
+      "model.llm.layers.19.mlp_norm",
+      "model.llm.layers.19.mlp_sconv",
+      "model.llm.layers.2.attn",
+      "model.llm.layers.2.attn_norm",
+      "model.llm.layers.2.attn_sconv",
+      "model.llm.layers.2.mlp.experts",
+      "model.llm.layers.2.mlp.gate",
+      "model.llm.layers.2.mlp.shared_experts",
+      "model.llm.layers.2.mlp_norm",
+      "model.llm.layers.2.mlp_sconv",
+      "model.llm.layers.20.attn",
+      "model.llm.layers.20.attn_norm",
+      "model.llm.layers.20.attn_sconv",
+      "model.llm.layers.20.mlp.gate",
+      "model.llm.layers.20.mlp.shared_experts",
+      "model.llm.layers.20.mlp_norm",
+      "model.llm.layers.20.mlp_sconv",
+      "model.llm.layers.21.attn",
+      "model.llm.layers.21.attn_norm",
+      "model.llm.layers.21.attn_sconv",
+      "model.llm.layers.21.mlp.gate",
+      "model.llm.layers.21.mlp.shared_experts",
+      "model.llm.layers.21.mlp_norm",
+      "model.llm.layers.21.mlp_sconv",
+      "model.llm.layers.22.attn",
+      "model.llm.layers.22.attn_norm",
+      "model.llm.layers.22.attn_sconv",
+      "model.llm.layers.22.mlp.gate",
+      "model.llm.layers.22.mlp.shared_experts",
+      "model.llm.layers.22.mlp_norm",
+      "model.llm.layers.22.mlp_sconv",
+      "model.llm.layers.23.attn",
+      "model.llm.layers.23.attn_norm",
+      "model.llm.layers.23.attn_sconv",
+      "model.llm.layers.23.mlp.gate",
+      "model.llm.layers.23.mlp.shared_experts",
+      "model.llm.layers.23.mlp_norm",
+      "model.llm.layers.23.mlp_sconv",
+      "model.llm.layers.24.attn",
+      "model.llm.layers.24.attn_norm",
+      "model.llm.layers.24.attn_sconv",
+      "model.llm.layers.24.mlp.gate",
+      "model.llm.layers.24.mlp.shared_experts",
+      "model.llm.layers.24.mlp_norm",
+      "model.llm.layers.24.mlp_sconv",
+      "model.llm.layers.25.attn",
+      "model.llm.layers.25.attn_norm",
+      "model.llm.layers.25.attn_sconv",
+      "model.llm.layers.25.mlp.gate",
+      "model.llm.layers.25.mlp.shared_experts",
+      "model.llm.layers.25.mlp_norm",
+      "model.llm.layers.25.mlp_sconv",
+      "model.llm.layers.26.attn",
+      "model.llm.layers.26.attn_norm",
+      "model.llm.layers.26.attn_sconv",
+      "model.llm.layers.26.mlp.gate",
+      "model.llm.layers.26.mlp.shared_experts",
+      "model.llm.layers.26.mlp_norm",
+      "model.llm.layers.26.mlp_sconv",
+      "model.llm.layers.27.attn",
+      "model.llm.layers.27.attn_norm",
+      "model.llm.layers.27.attn_sconv",
+      "model.llm.layers.27.mlp.gate",
+      "model.llm.layers.27.mlp.shared_experts",
+      "model.llm.layers.27.mlp_norm",
+      "model.llm.layers.27.mlp_sconv",
+      "model.llm.layers.28.attn",
+      "model.llm.layers.28.attn_norm",
+      "model.llm.layers.28.attn_sconv",
+      "model.llm.layers.28.mlp.gate",
+      "model.llm.layers.28.mlp.shared_experts",
+      "model.llm.layers.28.mlp_norm",
+      "model.llm.layers.28.mlp_sconv",
+      "model.llm.layers.29.attn",
+      "model.llm.layers.29.attn_norm",
+      "model.llm.layers.29.attn_sconv",
+      "model.llm.layers.29.mlp.gate",
+      "model.llm.layers.29.mlp.shared_experts",
+      "model.llm.layers.29.mlp_norm",
+      "model.llm.layers.29.mlp_sconv",
+      "model.llm.layers.3.attn",
+      "model.llm.layers.3.attn_norm",
+      "model.llm.layers.3.attn_sconv",
+      "model.llm.layers.3.mlp.gate",
+      "model.llm.layers.3.mlp.shared_experts",
+      "model.llm.layers.3.mlp_norm",
+      "model.llm.layers.3.mlp_sconv",
+      "model.llm.layers.30.attn",
+      "model.llm.layers.30.attn_norm",
+      "model.llm.layers.30.attn_sconv",
+      "model.llm.layers.30.mlp.gate",
+      "model.llm.layers.30.mlp.shared_experts",
+      "model.llm.layers.30.mlp_norm",
+      "model.llm.layers.30.mlp_sconv",
+      "model.llm.layers.31.attn",
+      "model.llm.layers.31.attn_norm",
+      "model.llm.layers.31.attn_sconv",
+      "model.llm.layers.31.mlp.gate",
+      "model.llm.layers.31.mlp.shared_experts",
+      "model.llm.layers.31.mlp_norm",
+      "model.llm.layers.31.mlp_sconv",
+      "model.llm.layers.32.attn",
+      "model.llm.layers.32.attn_norm",
+      "model.llm.layers.32.attn_sconv",
+      "model.llm.layers.32.mlp.gate",
+      "model.llm.layers.32.mlp.shared_experts",
+      "model.llm.layers.32.mlp_norm",
+      "model.llm.layers.32.mlp_sconv",
+      "model.llm.layers.33.attn",
+      "model.llm.layers.33.attn_norm",
+      "model.llm.layers.33.attn_sconv",
+      "model.llm.layers.33.mlp.gate",
+      "model.llm.layers.33.mlp.shared_experts",
+      "model.llm.layers.33.mlp_norm",
+      "model.llm.layers.33.mlp_sconv",
+      "model.llm.layers.34.attn",
+      "model.llm.layers.34.attn_norm",
+      "model.llm.layers.34.attn_sconv",
+      "model.llm.layers.34.mlp.gate",
+      "model.llm.layers.34.mlp.shared_experts",
+      "model.llm.layers.34.mlp_norm",
+      "model.llm.layers.34.mlp_sconv",
+      "model.llm.layers.35.attn",
+      "model.llm.layers.35.attn_norm",
+      "model.llm.layers.35.attn_sconv",
+      "model.llm.layers.35.mlp.gate",
+      "model.llm.layers.35.mlp.shared_experts",
+      "model.llm.layers.35.mlp_norm",
+      "model.llm.layers.35.mlp_sconv",
+      "model.llm.layers.36.attn",
+      "model.llm.layers.36.attn_norm",
+      "model.llm.layers.36.attn_sconv",
+      "model.llm.layers.36.mlp.gate",
+      "model.llm.layers.36.mlp.shared_experts",
+      "model.llm.layers.36.mlp_norm",
+      "model.llm.layers.36.mlp_sconv",
+      "model.llm.layers.37.attn",
+      "model.llm.layers.37.attn_norm",
+      "model.llm.layers.37.attn_sconv",
+      "model.llm.layers.37.mlp.gate",
+      "model.llm.layers.37.mlp.shared_experts",
+      "model.llm.layers.37.mlp_norm",
+      "model.llm.layers.37.mlp_sconv",
+      "model.llm.layers.38.attn",
+      "model.llm.layers.38.attn_norm",
+      "model.llm.layers.38.attn_sconv",
+      "model.llm.layers.38.mlp.gate",
+      "model.llm.layers.38.mlp.shared_experts",
+      "model.llm.layers.38.mlp_norm",
+      "model.llm.layers.38.mlp_sconv",
+      "model.llm.layers.39.attn",
+      "model.llm.layers.39.attn_norm",
+      "model.llm.layers.39.attn_sconv",
+      "model.llm.layers.39.mlp.gate",
+      "model.llm.layers.39.mlp.shared_experts",
+      "model.llm.layers.39.mlp_norm",
+      "model.llm.layers.39.mlp_sconv",
+      "model.llm.layers.4.attn",
+      "model.llm.layers.4.attn_norm",
+      "model.llm.layers.4.attn_sconv",
+      "model.llm.layers.4.mlp.gate",
+      "model.llm.layers.4.mlp.shared_experts",
+      "model.llm.layers.4.mlp_norm",
+      "model.llm.layers.4.mlp_sconv",
+      "model.llm.layers.40.attn",
+      "model.llm.layers.40.attn_norm",
+      "model.llm.layers.40.attn_sconv",
+      "model.llm.layers.40.mlp.gate",
+      "model.llm.layers.40.mlp.shared_experts",
+      "model.llm.layers.40.mlp_norm",
+      "model.llm.layers.40.mlp_sconv",
+      "model.llm.layers.41.attn",
+      "model.llm.layers.41.attn_norm",
+      "model.llm.layers.41.attn_sconv",
+      "model.llm.layers.41.mlp.gate",
+      "model.llm.layers.41.mlp.shared_experts",
+      "model.llm.layers.41.mlp_norm",
+      "model.llm.layers.41.mlp_sconv",
+      "model.llm.layers.42.attn",
+      "model.llm.layers.42.attn_norm",
+      "model.llm.layers.42.attn_sconv",
+      "model.llm.layers.42.mlp.gate",
+      "model.llm.layers.42.mlp.shared_experts",
+      "model.llm.layers.42.mlp_norm",
+      "model.llm.layers.42.mlp_sconv",
+      "model.llm.layers.43.attn",
+      "model.llm.layers.43.attn_norm",
+      "model.llm.layers.43.attn_sconv",
+      "model.llm.layers.43.mlp.gate",
+      "model.llm.layers.43.mlp.shared_experts",
+      "model.llm.layers.43.mlp_norm",
+      "model.llm.layers.43.mlp_sconv",
+      "model.llm.layers.44.attn",
+      "model.llm.layers.44.attn_norm",
+      "model.llm.layers.44.attn_sconv",
+      "model.llm.layers.44.mlp.gate",
+      "model.llm.layers.44.mlp.shared_experts",
+      "model.llm.layers.44.mlp_norm",
+      "model.llm.layers.44.mlp_sconv",
+      "model.llm.layers.45.attn",
+      "model.llm.layers.45.attn_norm",
+      "model.llm.layers.45.attn_sconv",
+      "model.llm.layers.45.mlp.gate",
+      "model.llm.layers.45.mlp.shared_experts",
+      "model.llm.layers.45.mlp_norm",
+      "model.llm.layers.45.mlp_sconv",
+      "model.llm.layers.46.attn",
+      "model.llm.layers.46.attn_norm",
+      "model.llm.layers.46.attn_sconv",
+      "model.llm.layers.46.mlp.gate",
+      "model.llm.layers.46.mlp.shared_experts",
+      "model.llm.layers.46.mlp_norm",
+      "model.llm.layers.46.mlp_sconv",
+      "model.llm.layers.47.attn",
+      "model.llm.layers.47.attn_norm",
+      "model.llm.layers.47.attn_sconv",
+      "model.llm.layers.47.mlp.gate",
+      "model.llm.layers.47.mlp.shared_experts",
+      "model.llm.layers.47.mlp_norm",
+      "model.llm.layers.47.mlp_sconv",
+      "model.llm.layers.48.attn",
+      "model.llm.layers.48.attn_norm",
+      "model.llm.layers.48.attn_sconv",
+      "model.llm.layers.48.mlp.gate",
+      "model.llm.layers.48.mlp.shared_experts",
+      "model.llm.layers.48.mlp_norm",
+      "model.llm.layers.48.mlp_sconv",
+      "model.llm.layers.49.attn",
+      "model.llm.layers.49.attn_norm",
+      "model.llm.layers.49.attn_sconv",
+      "model.llm.layers.49.mlp.gate",
+      "model.llm.layers.49.mlp.shared_experts",
+      "model.llm.layers.49.mlp_norm",
+      "model.llm.layers.49.mlp_sconv",
+      "model.llm.layers.5.attn",
+      "model.llm.layers.5.attn_norm",
+      "model.llm.layers.5.attn_sconv",
+      "model.llm.layers.5.mlp.gate",
+      "model.llm.layers.5.mlp.shared_experts",
+      "model.llm.layers.5.mlp_norm",
+      "model.llm.layers.5.mlp_sconv",
+      "model.llm.layers.50.attn",
+      "model.llm.layers.50.attn_norm",
+      "model.llm.layers.50.attn_sconv",
+      "model.llm.layers.50.mlp.gate",
+      "model.llm.layers.50.mlp.shared_experts",
+      "model.llm.layers.50.mlp_norm",
+      "model.llm.layers.50.mlp_sconv",
+      "model.llm.layers.51.attn",
+      "model.llm.layers.51.attn_norm",
+      "model.llm.layers.51.attn_sconv",
+      "model.llm.layers.51.mlp.gate",
+      "model.llm.layers.51.mlp.shared_experts",
+      "model.llm.layers.51.mlp_norm",
+      "model.llm.layers.51.mlp_sconv",
+      "model.llm.layers.52.attn",
+      "model.llm.layers.52.attn_norm",
+      "model.llm.layers.52.attn_sconv",
+      "model.llm.layers.52.mlp.gate",
+      "model.llm.layers.52.mlp.shared_experts",
+      "model.llm.layers.52.mlp_norm",
+      "model.llm.layers.52.mlp_sconv",
+      "model.llm.layers.53.attn",
+      "model.llm.layers.53.attn_norm",
+      "model.llm.layers.53.attn_sconv",
+      "model.llm.layers.53.mlp.gate",
+      "model.llm.layers.53.mlp.shared_experts",
+      "model.llm.layers.53.mlp_norm",
+      "model.llm.layers.53.mlp_sconv",
+      "model.llm.layers.54.attn",
+      "model.llm.layers.54.attn_norm",
+      "model.llm.layers.54.attn_sconv",
+      "model.llm.layers.54.mlp.gate",
+      "model.llm.layers.54.mlp.shared_experts",
+      "model.llm.layers.54.mlp_norm",
+      "model.llm.layers.54.mlp_sconv",
+      "model.llm.layers.55.attn",
+      "model.llm.layers.55.attn_norm",
+      "model.llm.layers.55.attn_sconv",
+      "model.llm.layers.55.mlp.gate",
+      "model.llm.layers.55.mlp.shared_experts",
+      "model.llm.layers.55.mlp_norm",
+      "model.llm.layers.55.mlp_sconv",
+      "model.llm.layers.56.attn",
+      "model.llm.layers.56.attn_norm",
+      "model.llm.layers.56.attn_sconv",
+      "model.llm.layers.56.mlp.gate",
+      "model.llm.layers.56.mlp.shared_experts",
+      "model.llm.layers.56.mlp_norm",
+      "model.llm.layers.56.mlp_sconv",
+      "model.llm.layers.57.attn",
+      "model.llm.layers.57.attn_norm",
+      "model.llm.layers.57.attn_sconv",
+      "model.llm.layers.57.mlp.gate",
+      "model.llm.layers.57.mlp.shared_experts",
+      "model.llm.layers.57.mlp_norm",
+      "model.llm.layers.57.mlp_sconv",
+      "model.llm.layers.58.attn",
+      "model.llm.layers.58.attn_norm",
+      "model.llm.layers.58.attn_sconv",
+      "model.llm.layers.58.mlp.gate",
+      "model.llm.layers.58.mlp.shared_experts",
+      "model.llm.layers.58.mlp_norm",
+      "model.llm.layers.58.mlp_sconv",
+      "model.llm.layers.59.attn",
+      "model.llm.layers.59.attn_norm",
+      "model.llm.layers.59.attn_sconv",
+      "model.llm.layers.59.mlp.gate",
+      "model.llm.layers.59.mlp.shared_experts",
+      "model.llm.layers.59.mlp_norm",
+      "model.llm.layers.59.mlp_sconv",
+      "model.llm.layers.6.attn",
+      "model.llm.layers.6.attn_norm",
+      "model.llm.layers.6.attn_sconv",
+      "model.llm.layers.6.mlp.gate",
+      "model.llm.layers.6.mlp.shared_experts",
+      "model.llm.layers.6.mlp_norm",
+      "model.llm.layers.6.mlp_sconv",
+      "model.llm.layers.60.attn",
+      "model.llm.layers.60.attn_norm",
+      "model.llm.layers.60.attn_sconv",
+      "model.llm.layers.60.mlp.gate",
+      "model.llm.layers.60.mlp.shared_experts",
+      "model.llm.layers.60.mlp_norm",
+      "model.llm.layers.60.mlp_sconv",
+      "model.llm.layers.61.attn",
+      "model.llm.layers.61.attn_norm",
+      "model.llm.layers.61.attn_sconv",
+      "model.llm.layers.61.mlp.gate",
+      "model.llm.layers.61.mlp.shared_experts",
+      "model.llm.layers.61.mlp_norm",
+      "model.llm.layers.61.mlp_sconv",
+      "model.llm.layers.62.attn",
+      "model.llm.layers.62.attn_norm",
+      "model.llm.layers.62.attn_sconv",
+      "model.llm.layers.62.mlp.gate",
+      "model.llm.layers.62.mlp.shared_experts",
+      "model.llm.layers.62.mlp_norm",
+      "model.llm.layers.62.mlp_sconv",
+      "model.llm.layers.63.attn",
+      "model.llm.layers.63.attn_norm",
+      "model.llm.layers.63.attn_sconv",
+      "model.llm.layers.63.mlp.gate",
+      "model.llm.layers.63.mlp.shared_experts",
+      "model.llm.layers.63.mlp_norm",
+      "model.llm.layers.63.mlp_sconv",
+      "model.llm.layers.64.attn",
+      "model.llm.layers.64.attn_norm",
+      "model.llm.layers.64.attn_sconv",
+      "model.llm.layers.64.mlp.gate",
+      "model.llm.layers.64.mlp.shared_experts",
+      "model.llm.layers.64.mlp_norm",
+      "model.llm.layers.64.mlp_sconv",
+      "model.llm.layers.65.attn",
+      "model.llm.layers.65.attn_norm",
+      "model.llm.layers.65.attn_sconv",
+      "model.llm.layers.65.mlp.gate",
+      "model.llm.layers.65.mlp.shared_experts",
+      "model.llm.layers.65.mlp_norm",
+      "model.llm.layers.65.mlp_sconv",
+      "model.llm.layers.7.attn",
+      "model.llm.layers.7.attn_norm",
+      "model.llm.layers.7.attn_sconv",
+      "model.llm.layers.7.mlp.gate",
+      "model.llm.layers.7.mlp.shared_experts",
+      "model.llm.layers.7.mlp_norm",
+      "model.llm.layers.7.mlp_sconv",
+      "model.llm.layers.8.attn",
+      "model.llm.layers.8.attn_norm",
+      "model.llm.layers.8.attn_sconv",
+      "model.llm.layers.8.mlp.gate",
+      "model.llm.layers.8.mlp.shared_experts",
+      "model.llm.layers.8.mlp_norm",
+      "model.llm.layers.8.mlp_sconv",
+      "model.llm.layers.9.attn",
+      "model.llm.layers.9.attn_norm",
+      "model.llm.layers.9.attn_sconv",
+      "model.llm.layers.9.mlp.gate",
+      "model.llm.layers.9.mlp.shared_experts",
+      "model.llm.layers.9.mlp_norm",
+      "model.llm.layers.9.mlp_sconv"
+    ],
+    "modelopt_quant_config": {
+      "quant_cfg": {
+        "*weight_quantizer": {
+          "num_bits": [
+            2,
+            1
+          ],
+          "block_sizes": {
+            "-1": 16,
+            "type": "dynamic",
+            "scale_bits": [
+              4,
+              3
+            ]
+          },
+          "axis": null,
+          "enable": true
+        },
+        "*input_quantizer": {
+          "num_bits": [
+            2,
+            1
+          ],
+          "block_sizes": {
+            "-1": 16,
+            "type": "dynamic",
+            "scale_bits": [
+              4,
+              3
+            ]
+          },
+          "axis": null,
+          "enable": true
+        },
+        "nn.BatchNorm1d": {
+          "*": {
+            "enable": false
+          }
+        },
+        "nn.BatchNorm2d": {
+          "*": {
+            "enable": false
+          }
+        },
+        "nn.BatchNorm3d": {
+          "*": {
+            "enable": false
+          }
+        },
+        "nn.LeakyReLU": {
+          "*": {
+            "enable": false
+          }
+        },
+        "*lm_head*": {
+          "enable": false
+        },
+        "*proj_out.*": {
+          "enable": false
+        },
+        "*block_sparse_moe.gate*": {
+          "enable": false
+        },
+        "*router*": {
+          "enable": false
+        },
+        "*mlp.gate.*": {
+          "enable": false
+        },
+        "*mlp.shared_expert_gate.*": {
+          "enable": false
+        },
+        "*linear_attn.conv1d*": {
+          "enable": false
+        },
+        "*mixer.conv1d*": {
+          "enable": false
+        },
+        "*output_layer*": {
+          "enable": false
+        },
+        "output.*": {
+          "enable": false
+        },
+        "default": {
+          "enable": false
+        }
+      },
+      "algorithm": "max"
+    },
+    "group_size": 16
+  },
+  "aqlm_hybrid": {
+    "quant_method": "inkling_nvfp4_aqlm_hybrid",
+    "base_nvfp4": "thinkingmachines/Inkling-NVFP4",
+    "group_size": 8,
+    "w13_book_entries": [
+      65536
+    ],
+    "w2_book_entries": [
+      65536,
+      256
+    ],
+    "w13_code_dtypes": [
+      "int16"
+    ],
+    "w2_code_dtypes": [
+      "int16",
+      "uint8"
+    ],
+    "cold_code_layout": "per-book tensors {proj}_cold_codes.{b}, dtype uint8 if book<=256 entries else int16",
+    "aqlm_layer_books": {
+      "2": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": false,
+        "hot_format": "bf16"
+      },
+      "3": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "4": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "5": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "6": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "7": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "8": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "9": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "10": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "11": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "12": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "13": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "14": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "15": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "16": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "17": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "18": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "19": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "20": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "21": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "22": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "23": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "24": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "25": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "26": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "27": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "28": {
+        "n_hot": 4,
+        "n_nvfp4": 4,
+        "n_cold": 252,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "29": {
+        "n_hot": 8,
+        "n_nvfp4": 8,
+        "n_cold": 248,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "30": {
+        "n_hot": 10,
+        "n_nvfp4": 10,
+        "n_cold": 246,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "31": {
+        "n_hot": 11,
+        "n_nvfp4": 11,
+        "n_cold": 245,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "32": {
+        "n_hot": 18,
+        "n_nvfp4": 18,
+        "n_cold": 238,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "33": {
+        "n_hot": 21,
+        "n_nvfp4": 21,
+        "n_cold": 235,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "34": {
+        "n_hot": 34,
+        "n_nvfp4": 34,
+        "n_cold": 222,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "35": {
+        "n_hot": 23,
+        "n_nvfp4": 23,
+        "n_cold": 233,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "36": {
+        "n_hot": 27,
+        "n_nvfp4": 27,
+        "n_cold": 229,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "37": {
+        "n_hot": 26,
+        "n_nvfp4": 26,
+        "n_cold": 230,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "38": {
+        "n_hot": 30,
+        "n_nvfp4": 30,
+        "n_cold": 226,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "39": {
+        "n_hot": 37,
+        "n_nvfp4": 37,
+        "n_cold": 219,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "40": {
+        "n_hot": 45,
+        "n_nvfp4": 45,
+        "n_cold": 211,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "41": {
+        "n_hot": 39,
+        "n_nvfp4": 39,
+        "n_cold": 217,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "42": {
+        "n_hot": 35,
+        "n_nvfp4": 35,
+        "n_cold": 221,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "43": {
+        "n_hot": 39,
+        "n_nvfp4": 39,
+        "n_cold": 217,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "44": {
+        "n_hot": 32,
+        "n_nvfp4": 32,
+        "n_cold": 224,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "45": {
+        "n_hot": 45,
+        "n_nvfp4": 45,
+        "n_cold": 211,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "46": {
+        "n_hot": 50,
+        "n_nvfp4": 50,
+        "n_cold": 206,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "47": {
+        "n_hot": 45,
+        "n_nvfp4": 45,
+        "n_cold": 211,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "48": {
+        "n_hot": 47,
+        "n_nvfp4": 47,
+        "n_cold": 209,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "49": {
+        "n_hot": 55,
+        "n_nvfp4": 55,
+        "n_cold": 201,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "50": {
+        "n_hot": 48,
+        "n_nvfp4": 48,
+        "n_cold": 208,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "51": {
+        "n_hot": 59,
+        "n_nvfp4": 59,
+        "n_cold": 197,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "52": {
+        "n_hot": 79,
+        "n_nvfp4": 79,
+        "n_cold": 177,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "53": {
+        "n_hot": 78,
+        "n_nvfp4": 78,
+        "n_cold": 178,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "54": {
+        "n_hot": 79,
+        "n_nvfp4": 79,
+        "n_cold": 177,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "55": {
+        "n_hot": 84,
+        "n_nvfp4": 84,
+        "n_cold": 172,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "56": {
+        "n_hot": 92,
+        "n_nvfp4": 92,
+        "n_cold": 164,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "57": {
+        "n_hot": 106,
+        "n_nvfp4": 106,
+        "n_cold": 150,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "58": {
+        "n_hot": 118,
+        "n_nvfp4": 118,
+        "n_cold": 138,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "59": {
+        "n_hot": 107,
+        "n_nvfp4": 107,
+        "n_cold": 149,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "60": {
+        "n_hot": 110,
+        "n_nvfp4": 110,
+        "n_cold": 146,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "61": {
+        "n_hot": 118,
+        "n_nvfp4": 118,
+        "n_cold": 138,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "62": {
+        "n_hot": 110,
+        "n_nvfp4": 110,
+        "n_cold": 146,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "63": {
+        "n_hot": 147,
+        "n_nvfp4": 147,
+        "n_cold": 109,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "64": {
+        "n_hot": 178,
+        "n_nvfp4": 178,
+        "n_cold": 78,
+        "packed": true,
+        "hot_format": "nvfp4"
+      },
+      "65": {
+        "n_hot": 119,
+        "n_nvfp4": 119,
+        "n_cold": 137,
+        "packed": true,
+        "hot_format": "nvfp4"
+      }
+    },
+    "note": "cold experts: AQLM, per-book code tensors {proj}_cold_codes.{b} (w2 book0=int16@65536, book1=uint8@256; w13 book0=int16@65536) + per-book codebooks + per-row scales. hot experts: NVFP4 (hot_format=nvfp4) except layer 2 which has NO NVFP4 base and stores bf16 hot slabs {proj}_hot_bf16 (no scale/scale2, hot_format=bf16). Its cold experts are AQLM like every other layer. hot_ids/cold_ids give the routed-expert index for each slab."
+  }
+}

model-00022-of-00045.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e317fa6fff7fc7a670fb796045e691026362e6bf204989266f2b26d42568eff7
+size 8578233880

model-00023-of-00045.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0a8b03eba80f9d802da22f0010092c598eb241b9d1a562bf719bda3fab7963f
+size 8352080344

model-00027-of-00045.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88a6b29aa4c54147274b122bedc55e2bbbb4d4e3ce0687d4a132d503e60233e6
+size 8490002496

model-00031-of-00045.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5dfe68d85e7fa49b5e679fe89809436d7e63d31fbb7f7c237e1da594672512d
+size 7576180336

model-00032-of-00045.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62d0720eab719e7506b650200d60d76971b778e3fba27bf7bb707c4fc704df28
+size 8364945840

model-00041-of-00045.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28df293809977e0966741702a04939f0b5e8f58b2e8f8505ee4675b7e5b03d67
+size 6977939008

model-00043-of-00045.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a904ed8c26e9621ea4e2071791a128fd7861058657fc5fac3d0346dfa32176e
+size 7737135388

processor_config.json

@@ -0,0 +1,46 @@
+{
+  "audio_token": "<|unused_200053|>",
+  "audio_bos_token": "<|content_audio_input|>",
+  "dmel_max_value": 2.0,
+  "dmel_min_value": -7.0,
+  "feature_extractor": {
+    "audio_token_duration_s": 0.05,
+    "feature_extractor_type": "InklingFeatureExtractor",
+    "feature_size": 80,
+    "hop_length": 800,
+    "n_fft": 1600,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "return_attention_mask": true,
+    "sampling_rate": 16000,
+    "window_size": 1600,
+    "window_size_multiplier": 2.0
+  },
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.48145466,
+      0.4578275,
+      0.40821073
+    ],
+    "image_processor_type": "InklingImageProcessor",
+    "image_std": [
+      0.26862954,
+      0.26130258,
+      0.27577711
+    ],
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+      "height": 40,
+      "width": 40
+    }
+  },
+  "image_token": "<|unused_200054|>",
+  "image_bos_token": "<|content_image|>",
+  "num_dmel_bins": 16,
+  "processor_class": "InklingProcessor"
+}

special_tokens_map.json

@@ -0,0 +1,22 @@
+{
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|message_user|>",
+    "<|message_model|>",
+    "<|message_system|>",
+    "<|message_tool|>",
+    "<|content_text|>",
+    "<|content_image|>",
+    "<|content_model_end_sampling|>",
+    "<|content_thinking|>",
+    "<|end_message|>",
+    "<|content_audio_input|>",
+    "<|content_tool_error|>",
+    "<|audio|>",
+    "<|content_xml|>",
+    "<|begin_of_text|>",
+    "<|audio_end|>",
+    "<|content_invoke_tool_json|>",
+    "<|content_invoke_tool_text|>"
+  ]
+}

tiktoken/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc253fd2b702f7a6da7105eaa8f3463b2f1247e83614f23e5323b921088bed2a
+size 3615874

tokenizer_config.json

@@ -0,0 +1,508 @@
+{
+  "added_tokens_decoder": {
+    "199998": {
+      "content": "<|unused|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "199999": {
+      "content": "<|endoftext|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200000": {
+      "content": "<|message_user|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200001": {
+      "content": "<|message_model|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200002": {
+      "content": "<|message_system|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200003": {
+      "content": "<|message_tool|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200004": {
+      "content": "<|content_text|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200005": {
+      "content": "<|content_image|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200006": {
+      "content": "<|content_model_end_sampling|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200007": {
+      "content": "<|unused_200007|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200008": {
+      "content": "<|content_thinking|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200009": {
+      "content": "<|unused_200009|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200010": {
+      "content": "<|end_message|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200011": {
+      "content": "<|unused_200011|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200012": {
+      "content": "<|unused_200012|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200013": {
+      "content": "<|unused_200013|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200014": {
+      "content": "<|unused_200014|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200015": {
+      "content": "<|unused_200015|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200016": {
+      "content": "<|unused_200016|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200017": {
+      "content": "<|unused_200017|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200018": {
+      "content": "<|unused_200018|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200019": {
+      "content": "<|unused_200019|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200020": {
+      "content": "<|content_audio_input|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200021": {
+      "content": "<|unused_200021|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200022": {
+      "content": "<|content_tool_error|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200023": {
+      "content": "<|audio|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200024": {
+      "content": "<|content_xml|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200025": {
+      "content": "<|unused_200025|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200026": {
+      "content": "<|unused_200026|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200027": {
+      "content": "<|unused_200027|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200028": {
+      "content": "<|begin_of_text|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200029": {
+      "content": "<|unused_200029|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200030": {
+      "content": "<|unused_200030|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200031": {
+      "content": "<|unused_200031|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200032": {
+      "content": "<|unused_200032|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200033": {
+      "content": "<|unused_200033|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200034": {
+      "content": "<|unused_200034|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200035": {
+      "content": "<|unused_200035|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200036": {
+      "content": "<|unused_200036|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200037": {
+      "content": "<|unused_200037|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200038": {
+      "content": "<|unused_200038|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200039": {
+      "content": "<|unused_200039|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200040": {
+      "content": "<|unused_200040|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200041": {
+      "content": "<|unused_200041|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200042": {
+      "content": "<|unused_200042|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200043": {
+      "content": "<|audio_end|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200044": {
+      "content": "<|unused_200044|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200045": {
+      "content": "<|unused_200045|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200046": {
+      "content": "<|unused_200046|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200047": {
+      "content": "<|unused_200047|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200048": {
+      "content": "<|unused_200048|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200049": {
+      "content": "<|content_invoke_tool_json|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200050": {
+      "content": "<|unused_200050|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200051": {
+      "content": "<|unused_200051|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200052": {
+      "content": "<|unused_200052|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200053": {
+      "content": "<|unused_200053|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200054": {
+      "content": "<|unused_200054|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200055": {
+      "content": "<|unused_200055|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200056": {
+      "content": "<|unused_200056|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    "200057": {
+      "content": "<|content_invoke_tool_text|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "fix_mistral_regex": false,
+  "extra_special_tokens": {
+    "endoftext": "<|endoftext|>",
+    "message_user": "<|message_user|>",
+    "message_model": "<|message_model|>",
+    "message_system": "<|message_system|>",
+    "message_tool": "<|message_tool|>",
+    "content_text": "<|content_text|>",
+    "content_image": "<|content_image|>",
+    "content_model_end_sampling": "<|content_model_end_sampling|>",
+    "content_thinking": "<|content_thinking|>",
+    "end_message": "<|end_message|>",
+    "content_audio_input": "<|content_audio_input|>",
+    "content_tool_error": "<|content_tool_error|>",
+    "audio": "<|audio|>",
+    "content_xml": "<|content_xml|>",
+    "begin_of_text": "<|begin_of_text|>",
+    "audio_end": "<|audio_end|>",
+    "content_invoke_tool_json": "<|content_invoke_tool_json|>",
+    "content_invoke_tool_text": "<|content_invoke_tool_text|>"
+  },
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}