feat(gguf): load mtp drafts for the qwen3.5 family from nextn blocks and mtp-*.gguf sidecars

docs/cli.md

@@ -102,11 +102,11 @@ See [models.md](models.md#moe-strategies) for what each strategy does, and [perf
 
 | Flag | Default | Meaning |
 |---|---|---|
-| `--speculative-algorithm`, `--spec-algo` | off | Draft method. `mtp` drafts with the checkpoint's multi-token-prediction layer: Qwen3.5 / Qwen3.6 HF checkpoints, and Qwen3.8-Flash-Next HF or GGUF |
+| `--speculative-algorithm`, `--spec-algo` | off | Draft method. `mtp` drafts with the checkpoint's multi-token-prediction layer: Qwen3.5 / Qwen3.6 HF or GGUF checkpoints, and Qwen3.8-Flash-Next HF or GGUF |
 | `--speculative-num-steps`, `--spec-num-steps` | the method's default (3 for `mtp`) | Draft tokens per decode step, 1 to 8 for `mtp` |
-| `--speculative-draft-model-path`, `--spec-draft-model` | the target checkpoint | Where the draft weights load from. A Qwen3.8-Flash-Next target also takes a Strata MTP pack (a GGUF that keeps the checkpoint's `mtp.*` tensors under their HF names), and a GGUF target needs one; the `nextn` MTP blocks llama.cpp writes are not read; a separate checkpoint folder is read with the target's quantization config |
+| `--speculative-draft-model-path`, `--spec-draft-model` | the target checkpoint | Where the draft weights load from. A Qwen3.8-Flash-Next target also takes a Strata MTP pack (a GGUF that keeps the checkpoint's `mtp.*` tensors under their HF names), and a Qwen3.8-Flash-Next GGUF target needs one; a Qwen3.5 / Qwen3.6 GGUF target drafts from its own `nextn` MTP block (the block llama.cpp keeps after the last decoder layer), or from an `mtp-*.gguf` sidecar of the same model (unsloth's `MTP/` releases) given here; a separate checkpoint folder is an HF-target option, read with the target's quantization config, and a GGUF target refuses a folder draft |
 
-Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets with sliding-window or DSV4 window-tier KV are refused at startup. `mtp` also refuses Qwen3.5 / Qwen3.6 GGUF targets and an FTW checkpoint as the draft source, since FTW drops the `mtp.*` tensors.
+Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets with sliding-window or DSV4 window-tier KV are refused at startup. `mtp` refuses an FTW checkpoint as the draft source, since FTW drops the `mtp.*` tensors, and a Qwen3.5 / Qwen3.6 GGUF converted without its MTP layer (llama.cpp `--no-mtp`) when no `mtp-*.gguf` sidecar is given: such a trunk has no draft of its own to read, while a `--no-mtp` trunk plus a sidecar works (the trunk config's shapes, the sidecar's weights).
 
 ### API behaviour
 

python/freetoken/models/qwen3_5_moe/gguf.py

@@ -22,6 +22,7 @@
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
+    from freetoken.moe.expert_banks import ExpertBanks
     from freetoken.models.gguf.config import GgufConfigShim
 
 _LAYER_LEAVES: dict[str, tuple[str, ...]] = {
@@ -51,19 +52,37 @@
     "lm_head": ("output",),
 }
 _LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$")
+# the MTP draft's layer reads the block after the trunk: the same leaves under mtp.layers.0
+_MTP_LAYER_RE = re.compile(r"^mtp\.layers\.0\.(.+)$")
+
+
+def _block_tensors(blk: str, leaf: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """The GGUF tensors (without ``.weight``) of one block's module leaf; stacked experts fuse gate|up when the file stores them fused."""
+    if leaf in ("mlp.experts.0.gate_proj", "mlp.experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors:
+        return (f"{blk}.ffn_gate_up_exps",)
+    return tuple(f"{blk}.{t}" for t in _LAYER_LEAVES.get(leaf, ()))
+
+
+def _nextn_blk(tensors: dict[str, int]) -> str | None:
+    """The block whose ``nextn`` tensors the MTP draft reads (``blk.64``), or None in a file converted without them."""
+    return next((name.split(".nextn.")[0] for name in tensors if ".nextn.eh_proj.weight" in name), None)
 
 
 def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
     """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
     if name in _TOP_LEAVES:
         return _TOP_LEAVES[name]
+    if name == "mtp.fc":
+        nextn = _nextn_blk(tensors)
+        return () if nextn is None else (f"{nextn}.nextn.eh_proj",)
+    m = _MTP_LAYER_RE.match(name)
+    if m is not None:
+        nextn = _nextn_blk(tensors)
+        return () if nextn is None else _block_tensors(nextn, m.group(1), tensors)
     m = _LAYER_RE.match(name)
     if m is None:
         return ()
-    blk, leaf = f"blk.{m.group(1)}", m.group(2)
-    if leaf in ("mlp.experts.0.gate_proj", "mlp.experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors:
-        return (f"{blk}.ffn_gate_up_exps",)
-    return tuple(f"{blk}.{t}" for t in _LAYER_LEAVES.get(leaf, ()))
+    return _block_tensors(f"blk.{m.group(1)}", m.group(2), tensors)
 
 
 def _check_supported(shim: GgufConfigShim) -> None:
@@ -144,6 +163,63 @@ def _untile(t: torch.Tensor, num_k_heads: int, head_dim: int, dim: int = 0) -> t
     return t.index_select(dim, idx)
 
 
+def _iter_block(
+    t,
+    norm,
+    blk: str,
+    pre: str,
+    config: ModelConfig,
+    layer_id: int,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """One decoder block (attention, MLP) as ``{pre}.*`` state-dict entries; the trunk loop and the nextn draft reader share it."""
+    hidden = config.hidden_size
+    yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm")
+    yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.post_attention_norm")
+    if config.is_linear_layer(layer_id):
+        g = config.linear_attention_group()
+        la = f"{pre}.linear_attn"
+        nk, dk = g.num_key_heads, g.key_head_dim
+        qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate")
+        qkv_rows = matrix(qkv)
+        v_start = 2 * nk * dk
+        qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)])
+        z_rows = _untile(matrix(gate), nk, g.value_head_dim)
+        beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha")
+        beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info))
+        if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type):
+            yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden)
+            yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden)
+        else:
+            yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden)
+        conv = float_tensor(t(f"{blk}.ssm_conv1d")).float()
+        conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)])
+        yield f"{la}.conv1d.weight", conv.unsqueeze(1)
+        yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1)
+        yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1)
+        yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm"))
+        yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim)
+    else:
+        sa = f"{pre}.self_attn"
+        parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+        yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden)
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm")
+        yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm")
+
+    mlp = f"{pre}.mlp"
+    if config.moe_enabled:
+        yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp"))
+        sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden)
+        yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size)
+        yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1)
+    else:
+        fg, fu, fd = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.gate_up_proj", [fg, fu], [matrix(fg), matrix(fu)], hidden)
+        yield from fused(f"{mlp}.down_proj", [fd], [matrix(fd)], config.intermediate_size)
+
+
 def iter_gguf_weights(
     model_path: str,
     device: torch.device,
@@ -161,7 +237,6 @@ def iter_gguf_weights(
     vision = include_vision and config.is_multimodal
     log_mmproj(model_path, vision)
     index = gguf_tensor_index(model_path)
-    g = config.linear_attention_group()
     hidden = config.hidden_size
 
     def t(name: str) -> GgufTensorInfo:
@@ -176,54 +251,77 @@ def norm(name: str) -> torch.Tensor:
     yield "model.norm.weight", norm("output_norm")
 
     for layer in range(config.num_layers):
-        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
-        yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm")
-        yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.post_attention_norm")
-        if config.is_linear_layer(layer):
-            la = f"{pre}.linear_attn"
-            nk, dk = g.num_key_heads, g.key_head_dim
-            qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate")
-            qkv_rows = matrix(qkv)
-            v_start = 2 * nk * dk
-            qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)])
-            z_rows = _untile(matrix(gate), nk, g.value_head_dim)
-            beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha")
-            beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info))
-            if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type):
-                yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden)
-                yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden)
-            else:
-                yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden)
-            conv = float_tensor(t(f"{blk}.ssm_conv1d")).float()
-            conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)])
-            yield f"{la}.conv1d.weight", conv.unsqueeze(1)
-            yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1)
-            yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1)
-            yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm"))
-            yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim)
-        else:
-            sa = f"{pre}.self_attn"
-            parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
-            yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden)
-            out = t(f"{blk}.attn_output")
-            yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
-            yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm")
-            yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm")
-
-        mlp = f"{pre}.mlp"
-        if config.moe_enabled:
-            yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp"))
-            sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down"))
-            yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden)
-            yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size)
-            yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1)
-        else:
-            fg, fu, fd = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
-            yield from fused(f"{mlp}.gate_up_proj", [fg, fu], [matrix(fg), matrix(fu)], hidden)
-            yield from fused(f"{mlp}.down_proj", [fd], [matrix(fd)], config.intermediate_size)
+        yield from _iter_block(t, norm, f"blk.{layer}", f"model.layers.{layer}", config, layer)
 
     if vision:
         yield from iter_mmproj_weights(find_mmproj(model_path))
 
 
-__all__ = ["gguf_mm_config", "iter_gguf_weights", "parse_gguf_config", "to_gguf"]
+# ----------------------------------------------------------------------------------------------
+# The trunk's nextn (MTP) block: llama.cpp keeps it after the last decoder layer
+# ----------------------------------------------------------------------------------------------
+
+def nextn_block_index(model_path: str) -> int:
+    """The trunk block that carries the nextn tensors; parse_gguf_config subtracts the same count from block_count."""
+    from freetoken.models.gguf.reader import gguf_architecture, load_gguf_metadata
+
+    md = load_gguf_metadata(model_path)
+    arch = gguf_architecture(model_path)
+    nextn = int(md.get(f"{arch}.nextn_predict_layers", 0))
+    if nextn == 0:
+        raise ValueError(
+            f"{model_path}: no MTP (nextn) block to draft from: the trunk was converted without one "
+            "(llama.cpp --no-mtp), or the file is a Strata-style MTP pack, which is not a draft "
+            "source for this family"
+        )
+    if nextn > 1:
+        raise ValueError(f"{model_path}: {nextn} nextn layers, but the MTP draft reads exactly one")
+    return int(md[f"{arch}.block_count"]) - nextn
+
+
+def iter_nextn_weights(model_path: str, config: ModelConfig) -> Iterator[tuple[str, torch.Tensor]]:
+    """The nextn block as Qwen3_5MTPDraft names it: the block's tensors under the draft's state-dict names, its routed experts excluded (they load as resident banks)."""
+    blk = f"blk.{nextn_block_index(model_path)}"
+    index = gguf_tensor_index(model_path)
+    if f"{blk}.nextn.eh_proj.weight" not in index:
+        raise ValueError(f"{model_path}: {blk} has no nextn tensors")
+
+    def t(name: str) -> GgufTensorInfo:
+        key = name if name in index else f"{name}.weight"
+        if key not in index:
+            raise ValueError(f"{model_path}: the MTP (nextn) block lacks {key}; reconvert with a current llama.cpp")
+        return index[key]
+
+    def norm(name: str) -> torch.Tensor:
+        return float_tensor(t(name)).float() - 1.0
+
+    eh = t(f"{blk}.nextn.eh_proj")
+    yield from fused("fc", [eh], [matrix(eh)], 2 * config.hidden_size)
+    yield "pre_fc_norm_embedding.weight", norm(f"{blk}.nextn.enorm")
+    yield "pre_fc_norm_hidden.weight", norm(f"{blk}.nextn.hnorm")
+    yield "norm.weight", norm(f"{blk}.nextn.shared_head_norm")
+    # the draft's one layer is full attention, whatever the trunk mixes around it
+    yield from _iter_block(t, norm, blk, "layers.0", config, 0)
+
+
+def nextn_expert_banks(model_path: str, device: torch.device) -> "ExpertBanks":
+    """The nextn block's stacked experts as the draft's one resident ggml bank set."""
+    from freetoken.models.gguf.experts import load_gguf_expert_banks
+
+    return load_gguf_expert_banks(model_path, 1, first_layer=nextn_block_index(model_path), device=device)
+
+
+def nextn_quant(model_path: str) -> "GGUFConfig":
+    """The draft's quant over its own file's tensor table: the trunk's when the draft is the
+    trunk's nextn block, the sidecar's when it is an mtp-*.gguf at another ggml mix."""
+    from freetoken.models.gguf.config import GGUF_ARCH_TO_REGISTRY
+    from freetoken.models.gguf.reader import gguf_architecture
+    from freetoken.models.register import checkpoint_quant_config, get_model_spec
+
+    # the HF and GGUF rows share one model class, so the target's class name resolves to the
+    # HF row and its from_hf branch; the file's own architecture names the GGUF row instead
+    spec = get_model_spec(GGUF_ARCH_TO_REGISTRY[gguf_architecture(model_path)])
+    return checkpoint_quant_config(model_path, SimpleNamespace(model_path=model_path), spec)
+
+
+__all__ = ["gguf_mm_config", "iter_gguf_weights", "iter_nextn_weights", "nextn_block_index", "nextn_expert_banks", "parse_gguf_config", "to_gguf"]