#!/usr/bin/env python3
"""Splice the BF16 MTP head tensors into the NVFP4-quantized output.

The Qwen3_5ForConditionalGeneration model class does NOT expose the MTP head as a
loadable module (it's a spec-decode drafting head vLLM reads separately from the
checkpoint), so AutoModelForCausalLM.from_pretrained DROPS the top-level `mtp.*`
state-dict keys on load — the quant output ends up with 0 mtp tensors. The working
reference (pantheon-27b-mtp-nvfp4) keeps them by SPLICING the BF16 mtp.* tensors
into the safetensors after quantization. This does the same: rewrites the single
output shard to include the 15 BF16 mtp.* tensors (from the graft's model-mtp
shard), matching pantheon's structure. (R36 fast-seat spike.)

Usage: python3 splice_mtp.py <nvfp4_out_dir> <grafted_dir_with_model-mtp.safetensors>
"""
import sys

from safetensors import safe_open
from safetensors.torch import save_file

out_dir, graft_dir = sys.argv[1], sys.argv[2]
out_st = f"{out_dir}/model.safetensors"
mtp_st = f"{graft_dir}/model-mtp.safetensors"

# load the full quantized shard (round-trips packed/scale tensors exactly)
tensors = {}
with safe_open(out_st, framework="pt") as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)
n_main = len(tensors)

# add the BF16 mtp.* tensors
with safe_open(mtp_st, framework="pt") as f:
    mtp_keys = list(f.keys())
    for k in mtp_keys:
        tensors[k] = f.get_tensor(k)

assert not any("mtp" in k.lower() for k in list(tensors)[:n_main]), "output already had mtp?"
print(f"main tensors: {n_main} | splicing {len(mtp_keys)} BF16 mtp tensors")
print("mtp dtypes:", {str(tensors[k].dtype) for k in mtp_keys})

save_file(tensors, out_st, metadata={"format": "pt"})
print(f"DONE — rewrote {out_st} with {len(tensors)} tensors ({n_main} quantized + {len(mtp_keys)} BF16 mtp)")