fix(heretic2-nvfp4): quant as ConditionalGeneration (namespace fix) + modelopt recipe for working MTP

Root-caused the NVFP4 gibberish to a quant-namespace bug: quant_nvfp4.py loaded
via AutoModelForCausalLM -> text-only Qwen3_5ForCausalLM -> flat model.layers.* keys,
but vLLM 0.24 serves only Qwen3_5ForConditionalGeneration (whose weight mapper needs
model.language_model.*). Fixed by loading as AutoModelForImageTextToText; NVFP4 now
serves coherent (validated greedy on ana-ml2 GPU0).

Base NVFP4 (compressed-tensors) measured ~53 tok/s (~= GGUF at batch-1, no single-stream
win) and its MTP is 0% acceptance (vLLM's Qwen3_5MTP drafter loads the bf16 mtp head only
off a modelopt main-model checkpoint). Added quant_modelopt.py (nvidia-modelopt PTQ,
matches AEON's NVFP4 W4A4 g16 + lm_head/linear_attn/visual exclusions) as the path to
working native MTP; graft + splice + serve otherwise unchanged.

services/heretic2-nvfp4-quant/quant_modelopt.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""NVFP4-quantize the MTP-grafted Heretic2 model via NVIDIA nvidia-modelopt (MODELOPT format).
+
+Sibling of quant_nvfp4.py (llm-compressor / compressed-tensors) but produces the **modelopt**
+NVFP4 format instead. Why it exists: the compressed-tensors path (quant_nvfp4.py) serves COHERENT
+but its MTP is 0% acceptance — vLLM's `Qwen3_5MTP` speculative-decode drafter only loads the bf16
+MTP head off a **modelopt** main-model checkpoint (the mtp tensors are byte-identical between the
+two formats; the difference is purely how the main model's quantized weights/scales are stored).
+So working native MTP (the ~2-4x goal) requires this format. (2026-07-14 modelopt pivot.)
+
+Reference = AEON `/tank/aimodels/qwen36-27b-aeon-nvfp4` (served by vllm-aeon-rp). Its
+hf_quant_config.json says: quant_algo NVFP4 (W4A4, group_size 16), targets Linear, exclude
+`lm_head` + `model.visual*` + every `linear_attn*` (the GDN). This script matches that EXACTLY.
+`embed_tokens` is an Embedding (not a Linear target) so it is never quantized — no need to exclude.
+
+Pipeline (unchanged except THIS quant step swaps llm-compressor -> modelopt):
+  graft_mtp.py  ->  quant_modelopt.py (this)  ->  splice_mtp.py (bf16 mtp)  ->  serve
+  serve: vllm --quantization modelopt --speculative-config
+         '{"method":"qwen3_5_mtp","num_speculative_tokens":3}'  (AEON vllm-aeon-rp is the ref)
+
+The MTP head is NOT in the module tree at load (transformers builds no mtp module for either
+Qwen3_5 class), so the 15 bf16 mtp.* tensors are spliced back AFTER export — same as pantheon/AEON.
+
+Load class is AutoModelForImageTextToText (= Qwen3_5ForConditionalGeneration) so weight keys are
+born `model.language_model.*` + `model.visual.*` — the namespace vLLM's ConditionalGeneration
+loader expects. (The compressed-tensors-path gibberish was a text-only-namespace bug; same fix.)
+
+Run in a vLLM container on the freed GPU0 (nvidia-modelopt[hf] per the transformers-compat warning):
+  docker run --gpus '"device=0"' --ipc host -v /tank/aimodels:/tank/aimodels -v /home/lkraven:/lk \
+    --entrypoint bash vllm/vllm-openai:v0.24.0 -c \
+    "pip install -q 'nvidia-modelopt[hf]' tiktoken sentencepiece && python3 /lk/quant_modelopt.py \
+       --model /tank/aimodels/heretic2-nvfp4-work/heretic2-mtp-bf16 \
+       --calib-mode chat --calib /tank/aimodels/heretic2-nvfp4-work/production_calib_512.jsonl \
+       --num-samples 512 --seqlen 8192 \
+       --out /tank/aimodels/heretic2-nvfp4-work/heretic2-modelopt-nvfp4"
+"""
+import argparse
+import copy
+import json
+import sys
+
+
+def load_calib_text(name, tokenizer, n, seqlen):
+    from datasets import load_dataset
+    ds = load_dataset(name, split="train").shuffle(seed=42).select(range(min(n, 100000)))
+    col = "text" if "text" in ds.column_names else ds.column_names[0]
+    return [tokenizer(x[col], truncation=True, max_length=seqlen) for x in ds.select(range(n))]
+
+
+def load_calib_chat(path, tokenizer, seqlen, n):
+    # Identical to quant_nvfp4.py: render each row via the Qwen3.6 chat template with thinking on
+    # so the forward-pass sees the seat's native tool-call XML activations. tool_call arguments may
+    # arrive as OpenAI wire-form JSON strings; the template does .items() on them -> parse to dict.
+    rows = [json.loads(l) for l in open(path) if l.strip()][:n]
+    out = []
+    for r in rows:
+        for m in r["messages"]:
+            for tc in (m.get("tool_calls") or []):
+                a = tc.get("function", {}).get("arguments")
+                if isinstance(a, str):
+                    tc["function"]["arguments"] = json.loads(a)
+        text = tokenizer.apply_chat_template(
+            r["messages"], tools=r.get("tools"),
+            tokenize=False, add_generation_prompt=False, enable_thinking=True,
+        )
+        out.append(tokenizer(text, truncation=True, max_length=seqlen))
+    return out
+
+
+def build_nvfp4_cfg(mtq):
+    """NVFP4 W4A4 (group_size 16) matching AEON's exclusions. NVFP4_DEFAULT_CFG already disables
+    lm_head + linear_attn.conv1d/in_proj_a/in_proj_b + mlp.gate; append the FULL linear_attn (GDN)
+    and the vision tower so only the standard attn/MLP Linears get NVFP4 (later entries override)."""
+    cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG)
+    cfg["quant_cfg"].append({"quantizer_name": "*linear_attn*", "enable": False})
+    cfg["quant_cfg"].append({"quantizer_name": "*visual*", "enable": False})
+    cfg["quant_cfg"].append({"quantizer_name": "*mtp*", "enable": False})  # moot (not in tree); belt+suspenders
+    return cfg
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--calib-mode", choices=["text", "chat"], required=True)
+    ap.add_argument("--calib", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--num-samples", type=int, default=512)
+    ap.add_argument("--seqlen", type=int, default=8192)
+    args = ap.parse_args()
+
+    import torch
+    from transformers import AutoModelForImageTextToText, AutoTokenizer
+    import modelopt.torch.quantization as mtq
+    from modelopt.torch.export import export_hf_checkpoint
+
+    print(f"loading grafted model (multimodal ConditionalGeneration): {args.model}", flush=True)
+    model = AutoModelForImageTextToText.from_pretrained(
+        args.model, torch_dtype="auto", device_map="auto", trust_remote_code=True,
+    )
+    model.eval()
+    tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+
+    print(f"building calibration ({args.calib_mode}, up to {args.num_samples} @ seq {args.seqlen})", flush=True)
+    if args.calib_mode == "text":
+        calib = load_calib_text(args.calib, tok, args.num_samples, args.seqlen)
+    else:
+        calib = load_calib_chat(args.calib, tok, args.seqlen, args.num_samples)
+    print(f"  {len(calib)} calibration rows", flush=True)
+
+    def forward_loop(m):
+        with torch.no_grad():
+            for i, row in enumerate(calib):
+                ids = torch.tensor([row["input_ids"]], device=m.device)
+                m(input_ids=ids)
+                if (i + 1) % 64 == 0:
+                    print(f"  calib {i + 1}/{len(calib)}", flush=True)
+
+    cfg = build_nvfp4_cfg(mtq)
+    print("running modelopt NVFP4 PTQ (W4A4 g16; lm_head/linear_attn/visual kept BF16)", flush=True)
+    mtq.quantize(model, cfg, forward_loop=forward_loop)
+
+    print(f"exporting modelopt HF checkpoint -> {args.out}", flush=True)
+    export_hf_checkpoint(model, export_dir=args.out)
+    tok.save_pretrained(args.out)
+    print("DONE. Next: splice_mtp.py <out> <graft>  then serve "
+          "--quantization modelopt --speculative-config "
+          "'{\"method\":\"qwen3_5_mtp\",\"num_speculative_tokens\":3}' "
+          "--reasoning-parser qwen3 --tool-call-parser qwen3_coder --enable-auto-tool-choice", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())