Brainwaves MTP Head: Recipe-Reconstruction Approach

TL;DR

We benchmarked four variants of Brainwaves' MTP (Multi-Token Prediction) head to compare the original restored head against a nuslerp-recipe-reconstructed head. The recipe-merged head consistently shows ~3 percentage points higher speculative decoding acceptance.

Key findings:

• Recipe heads: 64.35–64.59% acceptance vs 61.40–62.75% for as-shipped heads
• Imatrix quantization provides marginally higher throughput (+2–4 t/s decode) [This was done with the standard mradermacher imatrix dataset]
• Recipe i1 Q4 had the highest acceptance: 64.59%

What we did: Applied the same 3-level nuslerp merge recipe used for the main model to the MTP head tensors from the four parent models, then benchmarked against the original restored head.

---

Benchmark Results

All models benchmarked on SPEED-Bench coding category (32 samples, production sampling: temp=1.0, top_p=0.95, top_k=20). Server config: --spec-type draft-mtp,ngram-mod,ngram-map-k4v --spec-draft-n-max 3.

Nightmedia's models (as-shipped head):

Our recipe-merged head models:

Per-Position Acceptance (Recipe i1 Q4)

---

Background

Brainwaves is a nuslerp merge of four parent models. The published release originally had the mtp.fc.weight tensor (which maps to nextn.eh_proj in llama.cpp) missing from the shards. It was subsequently restored by the author, but we wanted to verify whether applying the same nuslerp merge recipe to the MTP head tensors would produce better speculative decoding performance.

The MTP head contains 15 tensors (8 weight matrices, 7 norm tensors) that enable llama.cpp's speculative decoding. Higher acceptance rates mean fewer rejected draft tokens and faster generation.

---

Approach: Recipe-Reconstruction

We applied the same 3-level hierarchical nuslerp merge recipe to the MTP head tensors from all four parent models:

1. nbeerbower/Wichtel-Qwen3.6-27B
2. trohrbaugh/Qwen3.8-27B-heretic-ara
3. DavidAU/Qwen3.8-27B-Cold-Fusion-GAIN-V1.1
4. DavidAU/Qwen3.6-27B-V1.1-FF711-Darker-Hero-GAIN-H2.0

All four parents had complete MTP heads with the same 15-tensor layout. The merge recipe was applied identically to the main model tensors and the MTP head tensors.

Verification

The reconstructed MTP head was compared against the restored head from the Brainwaves repository:

• 6 norm tensors: bit-identical (rel_l2 = 0.0)
• 8 weight matrices: cosine ≥ 0.99995, rel_l2 ≤ 1.6e-2, max_abs ≤ 1.6e-2
• Overall: The recipe merge faithfully reproduces the published head to within bf16 numerical noise

This confirms our nuslerp implementation is correct and the recipe was applied correctly. The ~3pp acceptance advantage suggests the recipe merge is capturing complementary strengths from the parent models' MTP heads.

---

Conversion Process

Step 1: BF16 GGUF Trunk

convert_hf_to_gguf.py <model-dir> --outfile brainwaves-bf16-trunk.gguf --outtype bf16 --no-mtp

Step 2: Static Q4_K_M Quantization

llama-quantize brainwaves-bf16-trunk.gguf brainwaves-q4_k_m.gguf Q4_K_M

Step 3: Imatrix Q4_K_M Quantization

Download imatrix from mradermacher/Qwen3.8-27B-i1-GGUF and apply:

llama-quantize --imatrix Qwen3.8-27B.imatrix.gguf brainwaves-bf16-trunk.gguf brainwaves-i1-q4_k_m.gguf Q4_K_M

Step 4: MTP Head Preparation

Build MTP-only index for each parent and apply the nuslerp recipe:

python build_mtp_index.py <parent-dir> --out <parent-dir>/model.safetensors.index.json
python merge_mtp_hf_sources.py --recipe brainwaves-mtp-recipe.json --output brainwaves-mtp-merged

Step 5: MTP Grafting

python graft_mtp_gguf.py brainwaves-q4_k_m.gguf brainwaves-mtp-merged brainwaves-mtp-q4_k_m.gguf
python graft_mtp_gguf.py brainwaves-i1-q4_k_m.gguf brainwaves-mtp-merged brainwaves-mtp-i1-q4_k_m.gguf

---

Merge Recipe

The Brainwaves model uses a 3-level hierarchical nuslerp merge:

Level 1: Wichtel-Heretic-B

models:
  - model: nbeerbower/Wichtel-Qwen3.6-27B
    parameters:
      weight: 1.6
  - model: trohrbaugh/Qwen3.8-27B-heretic-ara
    parameters:
      weight: 0.4
merge_method: nuslerp
dtype: bfloat16
name: Qwen3.8-27B-Wichtel-Heretic-B

Level 2: Cold-Fusion-FF711-Darker-Hero-GAIN-B

models:
  - model: DavidAU/Qwen3.8-27B-Cold-Fusion-GAIN-V1.1
    parameters:
      weight: 1.6
  - model: DavidAU/Qwen3.6-27B-V1.1-FF711-Darker-Hero-GAIN-H2.0
    parameters:
      weight: 0.4
merge_method: nuslerp
dtype: bfloat16
name: Qwen3.8-27B-Cold-Fusion-FF711-Darker-Hero-GAIN-B

Level 3: Final Brainwaves

models:
  - model: Qwen3.8-27B-Wichtel-Heretic-B
    parameters:
      weight: 1.6
  - model: Qwen3.8-27B-Cold-Fusion-FF711-Darker-Hero-GAIN-B
    parameters:
      weight: 0.4
merge_method: nuslerp
dtype: bfloat16
name: Qwen3.8-27B-Architect-Wichtel-B-Cold-Fusion-FF711-Darker-Hero-GAIN-B

---

Notes for Nightmedia

If You Want to Verify/Reproduce

1. Download the four parent models' MTP head shards
2. Apply the same nuslerp recipe as the main model
3. Verify against the restored head using the tensor comparison metrics above

If You Want to Improve the MTP Head

• The recipe merge shows that the nuslerp approach generalizes well to MTP heads
• Consider fine-tuning the MTP head separately after the merge to optimize speculative decoding performance
• The ~3pp acceptance advantage suggests the recipe merge may be capturing complementary strengths from the parent models

Technical Details

• MTP head contains 15 tensors (8 weight matrices, 7 norm tensors)
• All tensors are bf16 in the source model
• Norm tensors should be kept at F32 precision in GGUF for stability
• Weight matrices quantized to Q8_0 for the MTP block in llama.cpp

---

Scripts Used

All the scripts are custom in a private git repo, however I have made the key ones available here:

• build_mtp_index.py - Build MTP-only index from parent shards
• merge_mtp_hf_sources.py - Apply nuslerp recipe to MTP tensors
• assemble_mtp_hf_source.py - Overlay MTP tensors onto a base model
• graft_mtp_gguf.py - Graft MTP head onto quantized trunk

These scripts contain no sensitive data and can be made available for verification/reproduction.

Benchmark Script

Standardized MTP acceptance measurement is done from a custom script that parses data from llama-server's endpoints, and is not included here as reference. It's a bit longer. However if you'd want it, i can post it as well. It:

• Auto-detects server spec-decode configuration
• Runs SPEED-Bench coding category
• Captures both bench results and /metrics counters
• Computes per-position acceptance rates and MTP vs fallback split

---

Appendix: Key Scripts

build_mtp_index.py

Builds a minimal index for MTP tensors in a directory.

"""Build a minimal model.safetensors.index.json for the mtp.* tensors present in a directory.

Scans every *.safetensors header in the directory and writes an index whose weight_map
covers only mtp.* / model.mtp.* tensors. Used to prepare partial downloads (only the
shards that carry MTP tensors) for merge_mtp_hf_sources.py and convert_hf_to_gguf.py --mtp.
"""

from __future__ import annotations

import argparse
import json
import struct
from pathlib import Path


def safetensors_header(path: Path) -> dict:
    with open(path, "rb") as f:
        (header_len,) = struct.unpack("<Q", f.read(8))
        return json.loads(f.read(header_len))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("directory", type=Path)
    parser.add_argument("--out", default="", help="Index path (default: <directory>/model.safetensors.index.json)")
    args = parser.parse_args()

    weight_map: dict[str, str] = {}
    total_size = 0
    for shard in sorted(args.directory.glob("*.safetensors")):
        header = safetensors_header(shard)
        for name, info in header.items():
            if name == "__metadata__":
                continue
            if not (name.startswith("mtp.") or name.startswith("model.mtp.")):
                continue
            shape = info.get("shape", [])
            numel = 1
            for dim in shape:
                numel *= dim
            dtype_bytes = {"F32": 4, "F16": 2, "BF16": 2, "F64": 8}.get(info.get("dtype", "F32"), 4)
            weight_map[name] = shard.name
            total_size += numel * dtype_bytes

    if not weight_map:
        raise SystemExit(f"No mtp.* tensors found in {args.directory}")
    out_path = Path(args.out) if args.out else args.directory / "model.safetensors.index.json"
    index = {"metadata": {"total_size": total_size}, "weight_map": weight_map}
    out_path.write_text(json.dumps(index, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {out_path}: {len(weight_map)} tensors, {total_size / 1e6:.1f} MB")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

merge_mtp_hf_sources.py

Applies a hierarchical nuslerp recipe to parent MTP tensors.

"""Reconstruct an MTP block by applying a hierarchical nuslerp recipe to parent HF sources.

Reads mtp.* tensors from each parent's safetensors (via its weight index), evaluates the
merge tree with mergekit's nuslerp formulation, and writes a minimal HF directory
(config + index + one safetensors) compatible with convert_hf_to_gguf.py --mtp.

The recipe JSON mirrors mergekit's nested model structure:

{
  "name": "final",
  "models": [
    {"dir": "D:/AI/CONVERSION/owner__model-a", "weight": 1.6},
    {"name": "mid-b", "weight": 0.4, "models": [
       {"dir": "D:/AI/CONVERSION/owner__model-b", "weight": 1.6},
       {"dir": "D:/AI/CONVERSION/owner__model-c", "weight": 0.4}
    ]}
  ]
}

Leaves carry a local HF directory; internal nodes merge their children. Each models
entry requires a weight. Internal nodes may carry a name for diagnostics.
"""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file

_EPS = 1e-7
_SLEXP_EPS = 1e-8


def _normalize(x: torch.Tensor) -> torch.Tensor:
    return x / torch.norm(x, dim=-1, keepdim=True).clamp(min=_EPS)


def nuslerp(t: float, v0: torch.Tensor, v1: torch.Tensor) -> torch.Tensor:
    out_shape = v0.shape
    v0 = v0.view(-1)
    v1 = v1.view(-1)
    v0_u = _normalize(v0)
    v1_u = _normalize(v1)
    cos_theta = torch.sum(v0_u * v1_u, dim=-1, keepdim=True)
    theta = torch.acos(cos_theta.clamp(-1, 1))
    sin_theta = torch.sin(theta)
    colinear = sin_theta.abs() < _SLEXP_EPS
    res = (torch.sin((1 - t) * theta) * v0 + torch.sin(t * theta) * v1) / sin_theta
    res = torch.where(colinear, (1 - t) * v0 + t * v1, res)
    return res.view(out_shape)


def canonical_mtp_name(name: str) -> str:
    return name.removeprefix("model.")


def mtp_weight_map(source_dir: Path) -> dict[str, str]:
    index_path = source_dir / "model.safetensors.index.json"
    index = json.loads(index_path.read_text(encoding="utf-8"))
    return {
        canonical_mtp_name(name): shard
        for name, shard in index["weight_map"].items()
        if name.startswith("mtp.") or name.startswith("model.mtp.")
    }


class TensorStore:
    def __init__(self) -> None:
        self._maps: dict[Path, dict[str, str]] = {}
        self._shards: dict[tuple[Path, str], dict[str, torch.Tensor]] = {}

    def names(self, source_dir: Path) -> set[str]:
        source_dir = source_dir.resolve()
        if source_dir not in self._maps:
            self._maps[source_dir] = mtp_weight_map(source_dir)
        return set(self._maps[source_dir])

    def load(self, source_dir: Path, name: str) -> torch.Tensor:
        source_dir = source_dir.resolve()
        weight_map = self._maps.get(source_dir) or mtp_weight_map(source_dir)
        if weight_map:
            self._maps[source_dir] = weight_map
        if name not in weight_map:
            raise KeyError(f"{source_dir} has no tensor {name!r}")
        shard = weight_map[name]
        key = (source_dir, shard)
        if key not in self._shards:
            self._shards[key] = load_file(str(source_dir / shard))
        return self._shards[key][name]


def collect_leaf_dirs(node: dict, found: list[Path]) -> None:
    for child in node["models"]:
        if "dir" in child:
            found.append(Path(child["dir"]))
        else:
            collect_leaf_dirs(child, found)


def evaluate(node: dict, name: str, store: TensorStore, compute_dtype: torch.dtype) -> torch.Tensor:
    models = node["models"]
    if len(models) != 2:
        raise ValueError(f"nuslerp node '{node.get('name', '<unnamed>')}' has {len(models)} children; expected 2")

    results = []
    for child in models:
        if "dir" in child:
            tensor = store.load(Path(child["dir"]), name)
        else:
            tensor = evaluate(child, name, store, compute_dtype)
        results.append(tensor.to(compute_dtype))

    w0 = float(models[0]["weight"])
    w1 = float(models[1]["weight"])
    total = w0 + w1
    t = 0.5 if abs(total) < 1e-6 else w1 / total
    label = node.get("name", "<unnamed>")
    merged = nuslerp(t, results[0], results[1])
    print(f"  [{label}] {name}: t={t:.6f}")
    return merged


def compare(report: list[dict], built_in: torch.Tensor, candidate: torch.Tensor) -> None:
    a = built_in.detach().float().view(-1)
    b = candidate.detach().float().view(-1)
    if a.numel() != b.numel():
        report.append({"error": "size-mismatch", "built_in": a.numel(), "candidate": b.numel()})
        return
    dot = torch.sum(a * b)
    cos = (dot / (torch.norm(a) * torch.norm(b))).item()
    diff = (a - b).abs()
    rel_l2 = ((a - b).norm() / a.norm()).item()
    report.append({
        "cosine_similarity": round(cos, 10),
        "max_abs_diff": round(diff.max().item(), 8),
        "mean_abs_diff": round(diff.mean().item(), 8),
        "relative_l2": round(rel_l2, 10),
    })


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--recipe", type=Path, required=True, help="Nested nuslerp recipe JSON")
    parser.add_argument("--output", type=Path, required=True, help="Output minimal HF directory")
    parser.add_argument("--reference-dir", type=Path, required=True,
                        help="HF directory supplying config.json (and tokenizer files) for the output")
    parser.add_argument("--compare-dir", default="",
                        help="Optional HF directory whose built-in mtp.* tensors are compared against the merge")
    parser.add_argument("--compute-dtype", choices=("bfloat16", "float32"), default="bfloat16",
                        help="Dtype used for the slerp math (default bfloat16, matching the merge)")
    parser.add_argument("--report", default="", help="Optional path to write the comparison JSON")
    args = parser.parse_args()
    compare_dir = Path(args.compare_dir) if args.compare_dir else None
    report_path = Path(args.report) if args.report else None

    recipe = json.loads(args.recipe.read_text(encoding="utf-8"))
    compute_dtype = torch.bfloat16 if args.compute_dtype == "bfloat16" else torch.float32

    leaf_dirs: list[Path] = []
    collect_leaf_dirs(recipe, leaf_dirs)
    store = TensorStore()
    name_sets = [store.names(d) for d in leaf_dirs]
    common = set.intersection(*name_sets)
    if not common:
        raise SystemExit("No mtp.* tensors common to all leaves")
    missing = {str(d): sorted(ns - common) for d, ns in zip(leaf_dirs, name_sets) if ns != common}
    if missing:
        print(f"Warning: tensor sets differ across leaves: {missing}")
    names = sorted(common)
    print(f"Merging {len(names)} MTP tensors from {len(leaf_dirs)} leaves (dtype={args.compute_dtype})")

    merged = {}
    for name in names:
        merged[name] = evaluate(recipe, name, store, compute_dtype).to(torch.bfloat16)
        print(f"  merged {name}: {tuple(merged[name].shape)}")

    args.output.mkdir(parents=True, exist_ok=True)
    # The stock converter only enumerates shards named model*.safetensors.
    shard_name = "model-mtp.safetensors"
    save_file({k: v.contiguous() for k, v in merged.items()}, args.output / shard_name)
    total_size = sum(t.numel() * t.element_size() for t in merged.values())
    index = {"metadata": {"total_size": total_size}, "weight_map": {n: shard_name for n in names}}
    (args.output / "model.safetensors.index.json").write_text(json.dumps(index, indent=2) + "\n", encoding="utf-8")
    for file_name in ("config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt", "generation_config.json", "chat_template.jinja"):
        source_file = args.reference_dir / file_name
        if source_file.exists():
            shutil.copy2(source_file, args.output / file_name)
    print(f"Wrote minimal HF directory: {args.output} ({total_size / 1e6:.1f} MB BF16)")

    if compare_dir:
        built_in_map = mtp_weight_map(compare_dir)
        report: list[dict] = []
        print(f"\nComparing against built-in MTP in {compare_dir}:")
        for name in names:
            if name not in built_in_map:
                print(f"  {name}: MISSING from built-in index")
                report.append({"tensor": name, "status": "missing_from_index"})
                continue
            try:
                built_in = store.load(compare_dir, name)
            except KeyError:
                print(f"  {name}: referenced by built-in index but ABSENT from shard (incomplete release)")
                report.append({"tensor": name, "status": "absent_from_shard"})
                continue
            per_tensor: list[dict] = []
            compare(per_tensor, built_in, merged[name])
            entry = per_tensor[0] if per_tensor else {}
            report.append({"tensor": name, **entry})
            print(f"  {name}: cosine={entry.get('cosine_similarity')} rel_l2={entry.get('relative_l2')} max_abs={entry.get('max_abs_diff')}")
        if report:
            cosines = [r["cosine_similarity"] for r in report if "cosine_similarity" in r]
            rel_l2s = [r["relative_l2"] for r in report if "relative_l2" in r]
            if cosines:
                print(f"\n  min cosine={min(cosines):.10f}  mean cosine={sum(cosines) / len(cosines):.10f}")
                print(f"  max rel_l2={max(rel_l2s):.10f}  mean rel_l2={sum(rel_l2s) / len(rel_l2s):.10f}")
        if report_path:
            report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
            print(f"  report written: {report_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

assemble_mtp_hf_source.py

Overlays one MTP tensor set onto another (e.g., published release + reconstructed fill).

"""Assemble a minimal HF MTP source dir by overlaying one MTP tensor set onto another.

Base dir supplies its mtp.* tensors verbatim (byte-identical copies). Fill dir (optional)
supplies tensors missing from base; omit it when base is complete and the output simply
equals base's tensors. A supplied fill dir must have a model.safetensors.index.json
containing at least one mtp.* tensor, and the script exits with an error when base is
missing tensors but no fill dir was supplied.

The output shard is named model-mtp.safetensors so the stock converter's shard enumeration
(filename startswith 'model') picks it up. Config/tokenizer sidecar files are copied from
the base dir by whitelist (.json/.txt/.jinja/.model/.yml; *.safetensors is never copied),
so the output dir is directly usable as a -MtpSource for the converter or Invoke-MtpGraft.ps1.

With --verify, re-reads the written shard and confirms every tensor is byte-identical
to its source tensor before reporting success.

Typical use: base dir = the published release (which may be missing an MTP tensor the
merge dropped), fill dir = the complete nuslerp reconstruction built by
merge_mtp_hf_sources.py, --expected-missing = the dropped tensor name (e.g. mtp.fc.weight).
"""

import argparse
import hashlib
import json
import shutil
from pathlib import Path

import torch

from safetensors import safe_open
from safetensors.torch import save_file


def mtp_tensors_from_index(index_path):
    index = json.loads(index_path.read_text(encoding='utf-8'))
    weight_map = index.get('weight_map', {})
    return {name: shard for name, shard in weight_map.items() if name.startswith('mtp.')}


def read_tensor(shard_path, name):
    with safe_open(shard_path, framework='pt') as handle:
        tensor = handle.get_tensor(name)
    return tensor


def tensor_bytes(tensor):
    # BFloat16 has no direct numpy conversion; cast through float32 (deterministic on both sides).
    return tensor.detach().to(torch.float32).contiguous().cpu().numpy().tobytes()


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--base-dir', required=True,
                        help='HF dir whose mtp.* tensors are copied verbatim')
    parser.add_argument('--fill-dir', default=None,
                        help='HF dir supplying tensors missing from base (omit if base is complete)')
    parser.add_argument('--output', required=True, help='Output minimal HF dir')
    parser.add_argument('--expected-missing', default='',
                        help='Comma-separated tensor names expected to come from fill (empty = accept any)')
    parser.add_argument('--verify', action='store_true',
                        help='Re-read output shard and verify byte-identity with sources')
    args = parser.parse_args()

    base_dir = Path(args.base_dir)
    fill_dir = Path(args.fill_dir) if args.fill_dir else None
    out_dir = Path(args.output)

    base_tensors = mtp_tensors_from_index(base_dir / 'model.safetensors.index.json')
    fill_tensors = {}
    if fill_dir is not None:
        fill_index = fill_dir / 'model.safetensors.index.json'
        if not fill_index.is_file():
            raise SystemExit(f'Fill dir has no model.safetensors.index.json: {fill_dir}')
        fill_tensors = mtp_tensors_from_index(fill_index)
        if not fill_tensors:
            raise SystemExit(f'Fill dir has no mtp.* tensors in its index: {fill_dir}')

    if not base_tensors:
        raise SystemExit(f'Base dir has no mtp.* tensors in its index: {base_dir}')

    all_names = sorted(set(base_tensors) | set(fill_tensors))
    missing = [name for name in all_names if name not in base_tensors]
    expected = [item.strip() for item in args.expected_missing.split(',') if item.strip()]
    if expected and missing != expected:
        raise SystemExit(f'Expected fill tensors {expected}, got {missing}.')
    if not missing:
        if fill_dir is None:
            print('No fill dir supplied; base is complete, output equals base tensors.')
        else:
            print('WARNING: fill dir contributed nothing; output equals base tensors.')
    elif fill_dir is None:
        raise SystemExit(f'Base dir is missing tensors {missing} but no --fill-dir was supplied.')

    out_dir.mkdir(parents=True, exist_ok=True)
    shard_name = 'model-mtp.safetensors'
    shard_path = out_dir / shard_name

    tensors = {}
    sources = {}
    for name in all_names:
        if name in base_tensors:
            shard = base_tensors[name]
            tensors[name] = read_tensor(base_dir / shard, name)
            sources[name] = (base_dir / shard, name)
            print(f'{name}  {tuple(tensors[name].shape)}  {tensors[name].dtype}  <- base:{shard}')
        else:
            shard = fill_tensors[name]
            tensors[name] = read_tensor(fill_dir / shard, name)
            sources[name] = (fill_dir / shard, name)
            print(f'{name}  {tuple(tensors[name].shape)}  {tensors[name].dtype}  <- fill:{shard}')

    save_file(tensors, str(shard_path), metadata={'format': 'pt'})
    total_size = sum(t.numel() * t.element_size() for t in tensors.values())
    index = {
        'metadata': {'total_size': int(total_size)},
        'weight_map': {name: shard_name for name in all_names},
    }
    (out_dir / 'model.safetensors.index.json').write_text(json.dumps(index, indent=2), encoding='utf-8')

    skip = {'model.safetensors.index.json', 'mergekit_config.yml'}
    sidecar_suffixes = {'.json', '.txt', '.jinja', '.model', '.yml'}
    copied = []
    for filename in sorted(base_dir.iterdir()):
        if (filename.is_file() and filename.name not in skip
                and filename.suffix in sidecar_suffixes and '.safetensors' not in filename.name):
            shutil.copy2(filename, out_dir / filename.name)
            copied.append(filename.name)
    print(f'Copied sidecar files from base dir: {", ".join(copied) if copied else "(none)"}')

    if args.verify:
        mismatches = []
        for name in all_names:
            source_shard, _ = sources[name]
            with safe_open(source_shard, framework='pt') as source:
                source_bytes = tensor_bytes(source.get_tensor(name))
            with safe_open(shard_path, framework='pt') as out:
                out_bytes = tensor_bytes(out.get_tensor(name))
            if sha256(source_bytes) != sha256(out_bytes):
                mismatches.append(name)
        if mismatches:
            raise SystemExit(f'Verification FAILED for: {", ".join(mismatches)}')
        print(f'Verification OK: all {len(all_names)} tensors byte-identical to their sources.')

    print(f'Wrote {len(all_names)} tensors ({total_size / (1024 * 1024):.1f} MiB) -> {out_dir}')


if __name__ == '__main__':
    main()

graft_mtp_gguf.py

Grafts an MTP block from one GGUF into another.

"""Graft a compatible Qwen-style MTP block into a target GGUF."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "gguf-py"))

from gguf import GGMLQuantizationType, GGUFReader, GGUFWriter, GGUFValueType


def field_value(reader: GGUFReader, name: str, default=None):
    field = reader.fields.get(name)
    return default if field is None else field.contents()


def arch_key(reader: GGUFReader, suffix: str) -> str:
    architecture = field_value(reader, "general.architecture")
    if not architecture:
        raise ValueError("GGUF has no general.architecture metadata")
    return f"{architecture}.{suffix}"


def copy_metadata(reader: GGUFReader, writer: GGUFWriter, replacements: set[str]):
    for field in reader.fields.values():
        if field.name in ("GGUF.architecture", "general.architecture") or field.name.startswith("GGUF."):
            continue
        if field.name in replacements:
            continue
        value_type = field.types[0]
        subtype = field.types[-1] if value_type == GGUFValueType.ARRAY else None
        value = field.contents()
        if value is not None:
            writer.add_key_value(field.name, value, value_type, sub_type=subtype)


def tensor_bytes(reader: GGUFReader, tensor) -> bytes:
    data = np.ascontiguousarray(np.asarray(tensor.data))
    return data.view(np.uint8).tobytes(order="C")


def writer_shape(tensor):
    raw_types = {
        GGMLQuantizationType.F32: 4,
        GGMLQuantizationType.F16: 2,
    }
    byte_size = raw_types.get(tensor.tensor_type)
    if byte_size is None:
        return tensor.data.shape
    shape = list(tensor.data.shape)
    shape[-1] *= byte_size
    return tuple(shape)


def find_mtp_tensors(reader: GGUFReader, target_block_count: int) -> list[tuple[str, object]]:
    tensors = []
    expected_prefix = f"blk.{target_block_count}."
    for tensor in reader.tensors:
        name = tensor.name
        if name.startswith(expected_prefix):
            tensors.append((name, tensor))
        elif name.startswith("mtp.") or name.startswith("model.mtp."):
            raise ValueError(
                "Source contains HF-style mtp.* tensors. Convert it with the stock "
                "converter --mtp first; direct GGUF grafting expects blk.N.* names."
            )
    if not tensors:
        raise ValueError(f"Source contains no MTP tensors at expected block {target_block_count}")
    return tensors


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("target", type=Path)
    parser.add_argument("mtp_source", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--allow-family-mismatch", action="store_true")
    args = parser.parse_args()

    for path in (args.target, args.mtp_source):
        if not path.exists():
            raise SystemExit(f"GGUF not found: {path}")
    if args.output.exists():
        raise SystemExit(f"Refusing to overwrite existing output: {args.output}")

    target = GGUFReader(args.target)
    source = GGUFReader(args.mtp_source)
    target_arch = field_value(target, "general.architecture")
    source_arch = field_value(source, "general.architecture")
    target_blocks_metadata = int(field_value(target, arch_key(target, "block_count"), 0))
    target_nextn = int(field_value(target, arch_key(target, "nextn_predict_layers"), 0))
    target_blocks = target_blocks_metadata - target_nextn if target_nextn > 0 else target_blocks_metadata
    source_blocks = int(field_value(source, arch_key(source, "block_count"), 0))
    target_hidden = field_value(target, arch_key(target, "embedding_length"))
    source_hidden = field_value(source, arch_key(source, "embedding_length"))

    if not args.allow_family_mismatch:
        if target_arch != source_arch:
            raise ValueError(f"Architecture mismatch: target={target_arch}, source={source_arch}")
        if target_hidden != source_hidden:
            raise ValueError(f"Hidden-size mismatch: target={target_hidden}, source={source_hidden}")
    if target_blocks <= 0:
        raise ValueError("Target has no usable block_count metadata")
    if source_blocks not in (target_blocks, target_blocks + 1):
        raise ValueError(
            f"Unexpected source block_count={source_blocks}; expected {target_blocks} or {target_blocks + 1}"
        )

    mtp_tensors = find_mtp_tensors(source, target_blocks)
    target_names = {tensor.name for tensor in target.tensors}
    collisions = sorted(name for name, _ in mtp_tensors if name in target_names)
    if collisions:
        raise ValueError(f"Target already contains MTP tensors: {collisions[:5]}")

    output_arch_block = arch_key(target, "block_count")
    output_arch_nextn = arch_key(target, "nextn_predict_layers")
    replacements = {output_arch_block, output_arch_nextn}
    writer = GGUFWriter(args.output, arch=target_arch, endianess=target.endianess)
    copy_metadata(target, writer, replacements)
    writer.add_key_value(output_arch_block, target_blocks + 1, GGUFValueType.UINT32)
    writer.add_key_value(output_arch_nextn, len(mtp_tensors) > 0 and 1 or 0, GGUFValueType.UINT32)

    all_specs = [(tensor.name, tensor) for tensor in target.tensors] + mtp_tensors
    for name, tensor in all_specs:
        raw = tensor_bytes(target if name in target_names else source, tensor)
        writer.add_tensor_info(name, writer_shape(tensor), np.dtype(np.uint8), len(raw), raw_dtype=tensor.tensor_type)

    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_ti_data_to_file()
    writer.write_padding(writer.fout[0], writer.fout[0].tell())
    for name, tensor in all_specs:
        reader = target if name in target_names else source
        raw = tensor_bytes(reader, tensor)
        writer.fout[0].write(raw)
        writer.write_padding(writer.fout[0], len(raw))
    writer.close()

    print(f"Grafted {len(mtp_tensors)} MTP tensors into {args.output}")
    print(f"Architecture: {target_arch}; block_count: {target_blocks} -> {target_blocks + 1}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())