# Fused NextN draft graph (task #590 / bug-858, bug-867)

**Status: BUILT + SHIPPED (patch `0093-mtp-dualctx-fused-nextn.patch`) —
flag-gated, default OFF, experimental.** The fused builder + driver
(`llama_decode_mtp_fused_nextn`, `src/llama-context.cpp`; dispatch in
`common/speculative.cpp` behind env **`OPENCOTI_MTP_FUSED_NEXTN=1`**, NextN
non-shared-mem only, `n_max ≥ 2`) landed via tasks S0–S6 (#593–#599).
**Verdict (2026-07-14 S5 re-validation): fused-on is measurably SLOWER than
the shipped AR draft loop** — the fused graph is not shape-invariant
(`can_reuse` fails), so it rebuilds O(n_layers×n_steps) every cycle; the
single-launch design never pays. The bug-858 Qwen gap was closed instead by
patch `0128` (spec-clone-cheap + mmvf verify lift), which took the AR path to
b9859 parity+. The fused path is kept as infrastructure (a shape-invariant
padded-graph rework is the only way it could win); do NOT enable it in
production. The blueprint below and the "Measured" table PRE-DATE 0128 —
their −12–15% gap no longer exists on the current binary.

**Problem.** Qwen NextN self-speculative MTP (`--spec-type draft-mtp`) on 35B-A3B
decodes ~16% slower than upstream llama.cpp `b9859` on the *same* GGUF/HW, with
**acceptance and cycle count at parity** (ours 0.79–0.80 accept, 307 draft_n, 244
accepted @ 244 t/s; upstream 0.815, 303, 247 @ 283 t/s). Base single-decode is at
parity. The gap is therefore **per-cycle host-sync/launch overhead**, not drafter
quality and not verify criterion — see `docs/evaluations/three-way-tps.md` §"#590
diagnostic" and buglog bug-864/866/867. The overhead lives in the **un-fused
per-step draft loop**: `common_speculative_impl_draft_mtp::draft`
(`common/speculative.cpp:702-803`) does, *per drafted token*, a
`llama_decode(ctx_dft)` + a `llama_get_embeddings_pre_norm_ith` **host readback**
(`:751-752, :796`) — one launch + one D2H sync per step.

**Fix.** Give NextN a **fused draft graph**: the NextN analogue of the
Gemma-assistant `decode_mtp_fused` (`src/llama-context.cpp:2900`). Build the N draft
steps as *one* graph with an on-device argmax→embed→hidden chain between steps, so
the whole block costs **one launch + one readback** (`get_argmax()` I32[N] +
`get_embd()`), eliminating the per-step host round-trip. Payoff scales with n_max
(modest at the n=2 cap, expected to fix the n=3 collapse).

`decode_mtp_fused` is **not reusable as-is** — it is hard-wired to Gemma
(`model.mtp_assistant` at `:2908`; `dynamic_cast<llama_kv_cache_iswa>` at `:2916`).
The NextN drafter is a *separate* draft context (`ctx_dft`) with its own model, its
own **unified** (non-iSWA) KV cache, and its own NextN block. So this is a
from-scratch build that *reuses the scheduler/readback machinery* but adds a new
Qwen fused builder + a new driver.

---

## The two MTP lanes (why reuse is partial)

`src/llama-graph.h:37-39` graph types:

| type | drafter | context | KV cache | fused today? |
|---|---|---|---|---|
| `LLM_GRAPH_TYPE_MTP` | Gemma-4 **assistant** (sub-model in target ctx) | target | iSWA (shared) | **yes** (`decode_mtp_fused`) |
| `LLM_GRAPH_TYPE_DECODER_MTP` | Qwen **NextN** (own block in draft model) | separate `ctx_dft` | unified | **no** (per-step loop) |

The Gemma assistant shares the *target's* iSWA KV via cross-attention
(`build_attn_mtp`, `src/llama-graph.cpp:2977`); Qwen NextN runs a full draft context
with its own unified KV (`src/llama-model.cpp:1979-1995`) and reuses the trunk
layer's weights for its single NextN block (`src/models/qwen35.cpp:510-656`,
`qwen35moe.cpp:575`).

---

## Reuse as-is (model-agnostic machinery)

From `src/llama-context.cpp`, all shared once the iSWA/assistant specifics are
generalized out:

- `process_ubatch_mtp` (`:1538-1588`) — dedicated `sched_mtp` + reusable
  `gf_res_prev_mtp`, reuse gate + `set_inputs`.
- `graph_compute_mtp` (`:1591`).
- `ensure_sched_mtp` N×-nodes reserve pattern (`:1419-1445`:
  `max_nodes = graph_max_nodes * n_steps`).
- The readback contract: `res->get_argmax()` → I32[n_steps] +
  `res->get_embd()` (`:3000-3013`); accessors `src/llama-graph.h:679-719`
  (`t_argmax` is set by **direct field assignment**, no `set_argmax` method —
  reset at `src/llama-graph.cpp:898`).

The **only** things inside these that must be generalized: the
`dynamic_cast<llama_kv_cache_iswa>` guards (`:1451`, `:2916`) and the
`mtp_assistant`-derived arch/width (`:2681-2685`, `:2930`).

---

## The 7 deltas (what must be built)

1. **N-step chaining loop in the NextN builder.** `qwen35.cpp:graph_mtp`
   (`:510-656`) is strictly single-block/single-pass. Add a `for k in n_steps`
   unroll modeled on `gemma4-assistant.cpp:341-357`, re-invoking
   eh_proj→attn→ffn→head N times in one graph. Factor the current body into a
   `graph_mtp_build_one_step` primitive (mirror
   `gemma4_mtp_build_one_step`, `gemma4-assistant.cpp:45-258`).

2. **`n_mtp_steps` plumbing to the Qwen path.** `n_mtp_steps`
   (`llama-graph.h:595-599`) is today only read by the Gemma builder and only set
   via `mtp_fused_steps`→`graph_params_mtp` (`llama-context.cpp:2696`). The Qwen
   draft ctx builds through the *normal* `process_ubatch`/`build_graph`, so there
   is **no wire** carrying a step count into `graph_mtp`. Add a Qwen equivalent of
   `graph_params_mtp` (arch = draft model's own arch, `gtype =
   LLM_GRAPH_TYPE_DECODER_MTP`, `n_mtp_steps = N`) used by the new driver.

3. **In-graph argmax on the NextN head.** `graph_mtp` sets only `t_logits`
   (`:654`) + `t_h_pre_norm` (`:637`), never `t_argmax`. Add per-step
   `ggml_argmax` on the LM-head output + concat into I32[N] (mirror
   `gemma4-assistant.cpp:252, 362-368`), publish `res->t_argmax` (direct assign).
   **Consequence:** the fused path is **greedy-only** and drops the host `p_min`
   early-stop confidence gate (`speculative.cpp:767`) — see Risks.

4. **On-device token→embed + hidden chaining.** `graph_mtp` already does
   `ggml_get_rows(embed_tokens, tokens)` (`:540-543`) but on a *host-supplied*
   `tokens`. Fused: step k+1's token = step k's on-device argmax feeding
   `ggml_get_rows` (mirror `:355` + `:72`); step k+1's `h_input` = step k's
   `t_h_pre_norm` (mirror `:356`) — entirely on-device.

5. **NextN per-step position input — RESOLVED: reuse `llm_graph_input_mtp`.**
   Its setter (`llama-graph.cpp:116-131`) is **width-agnostic** — `inp_h_prev` is
   written `inp_h_prev->ne[0]` floats — so Qwen reuses it verbatim by creating
   `inp_h_prev` at `n_embd` width (not Gemma's `n_embd_out_impl`). The only change
   needed: extend `llm_graph_input_mtp::can_reuse` (`:134-145`) to also accept
   `gtype==LLM_GRAPH_TYPE_DECODER_MTP` (today it hard-returns false unless
   `LLM_GRAPH_TYPE_MTP` at `:135`), else the fused Qwen graph rebuilds every draft
   and the perf win is lost (correct but slow).

6. **KV for the N fused steps — RESOLVED (Option A, 2026-07-02): read-only
   cross-attend the frozen prefix; NO new cache surgery.** Reading `build_attn_mtp`
   (`llama-graph.cpp:2977`, esp. `:3011` `get_k`/`:3012` `get_v`) showed the
   proven Gemma fused path is a **read-only cross-attention into the frozen prefix
   KV** — it never writes draft KV; the recurrence is carried entirely by the
   hidden-state chain (`h_post`→next `h`). That is why one slot suffices:
   `llama_kv_cache::mtp_slot_info` (`llama-kv-cache.cpp:1985`) returns a **single**
   idx (the `pmax` cell), and it already exists on the **unified** cache (iSWA's
   `init_mtp` just calls it on `kv_base`/`kv_swa`). So the unified path needs **no**
   new `init_mtp`/N-cell/causal-mask work — build one `llama_kv_cache_context` from
   `mtp_slot_info(seq_id)` (mirror `llama-kv-cache-iswa.cpp:238`).

   **Fork decided: Option A "Gemma-mirror" (prototype-and-measure).** Qwen's
   per-step loop uses *self-attention that writes each draft token's KV* (so step
   k sees drafts 0..k-1); Option A cross-attends only the frozen prefix (step k
   sees the prefix, not 0..k-1) → drafts differ (`real_frac≠0`) and acceptance may
   dip, but correctness is guaranteed by verify. Chosen because it's the proven
   pattern with no KV-cache surgery. Measure acceptance in S5; only build the
   self-attn N-cell exact-match variant (Option B) if A regresses materially vs
   the 0.80 baseline. **Impl:** the fused steps need a Qwen read-only cross-attn
   into the unified cache (an analog of `build_attn_mtp` — get_k/get_v from the
   unified `llama_kv_cache_context` + the `build_attn_inp_kv` mask exposing the
   prefix + `build_attn_mha`, no KV write), branched inside `build_one_step` on a
   `fused` flag so the single-step path stays byte-identical.

7. **New driver = `decode_mtp_fused` twin without the Gemma guards.** Drop:
   `model.mtp_assistant` checks (`:2908, 1428`) → replace with "draft ctx is
   `LLAMA_CONTEXT_TYPE_MTP`"; the iSWA cast (`:2916-2920, 1451-1457`) → unified
   cache + its `init_mtp`; seed width `n_bb = n_embd_out_impl` (`:2930`) → plain
   `n_embd` (Qwen's `t_h_pre_norm` is `[n_embd, n_outputs]`, read via
   `get_embeddings_pre_norm_ith`, `:1061`); `graph_params_mtp`'s
   `arch/gtype` (`:2681-2685`) → draft model's own arch + `DECODER_MTP`.

Then **swap the driver into the loop**: `common_speculative_impl_draft_mtp::draft`
(`speculative.cpp:702-803`) calls the new fused driver once instead of the per-step
`llama_decode`+readback loop, and reads back I32[N] drafts + last hidden.

---

## Progress (2026-07-02)

- **S0 — DONE.** vendor-backup `20260702-143936-pre-fused-nextn-590` (9338
  entries, 37 patches, verified).
- **S1 — DONE + gated byte-identical.** `build_one_step` lambda extracted in
  `qwen35.cpp` + `qwen35moe.cpp` (`this`-capturing generic lambda, minimal-diff,
  no re-indent). Incremental host build OK (2 TUs, host-only — no CUDA DSO
  touched). bs2 35B NextN n2 chash `f80af37c9a` == pre-refactor baseline; decode
  243–245, accept 0.79–0.80. **Build cadence confirmed cheap:** edit → incremental
  `bun run build:llamafile:make` (2 TUs + relink) → rsync binary to bs2 (reuses
  cached sm120 DSO, correct since host-only) → gate ~5 min.
- **S2 — DONE + gated byte-identical.** `build_one_step` now takes the token index
  (embeds internally via `ggml_get_rows`) and returns `{h_pre_norm, logits, arg}`
  where `arg = ggml_argmax(logits)`. Added the **dormant** fused N-step branch
  (gated `n_mtp_steps>1`, inert at default 1) in both files: reuses
  `llm_graph_input_mtp` (N per-step I32[1] positions), chains `tok_k=arg_k` /
  `h_k=h_pre_norm_k` on-device, concats argmaxes → `res->t_argmax`. Build OK; bs2
  35B NextN n2 chash `f80af37c9a` (single-step path byte-identical, argmax node
  pruned when unreferenced). **All graph-builder work is now in place, dormant.**

## Staged plan (de-risk ordering)

- **S0 — vendor-backup WHOLE tree** (`scripts/vendor-backup.sh backup
  pre-fused-nextn`). Mandatory before source edits + build cycles. ✅
- **S1 — refactor, no behaviour change.** Extract `build_one_step` from
  `qwen35.cpp:graph_mtp` (and the MoE twin `qwen35moe.cpp:575`); the single-step
  path calls it once. Build + prove **byte-identical** decode vs current
  (chash) — pure refactor gate. ✅ (chash `f80af37c9a`).
- **S2 — fused builder (greedy N-step) + input + argmax.** Deltas 1,3,4,5. Publish
  `t_argmax`/`t_embd`. Still driven single-step (N=1) first to prove the argmax
  path matches host `common_sampler_sample` (real_frac=0 vs the per-step loop).
- **S3 — KV slots + driver.** Deltas 2,6,7. New `decode_mtp_fused_nextn` +
  unified-cache `init_mtp`. Wire into `speculative.cpp` behind an env gate
  (`OPENCOTI_MTP_FUSED_NEXTN`, default **off** until proven, then flip on like the
  Gemma default per the standing "if it wins, default it" rule).
- **S4 — build + restamp.** Host `rm -rf o` + make on solidPC (consistent tree),
  CUDA DSO if any device code touched (argmax/get_rows are existing ops → likely
  **no** new kernel, host-only), restamp both DSO paths, `nm -D` verify.
- **S5 — gates.** (a) **Correctness:** logit-equiv / `real_frac=0` between fused
  and per-step drafts on the same prompt (drafts must be identical token IDs since
  both are greedy-argmax over the same NextN head); niah retrieval unchanged;
  byte-identical final decode with fused **off**. (b) **Perf:** ours-fused vs
  upstream b9859 on 35B NextN n=2 **and** n=3 — target: close the +16% at n2 and
  kill the n3 collapse. Deploy binary to bs2 (GPU0), reuse `m592-accept.sh`
  harness shape.
- **S6 — ship.** Additive patch(es) into `vendors/patches/llamafile/` (backup +
  snapshot-diff, byte-identical re-apply), `opencoti-hook:` markers +
  UPSTREAM_SYNC registry, `docs/evaluations/mtp.md` + this doc's status, README
  row, `.wolf` (anatomy/cerebrum/memory), pgvector. Commit on dev only when asked.

---

## Risks / open questions

- **Loses `p_min` adaptive early-stop — negligible at the n=2 cap, and
  recoverable.** The fused path is greedy-argmax with a fixed N (like Gemma), so
  the host confidence gate (`speculative.cpp:767`, `cur_p->data[0].p < p_min` →
  stop drafting; llama.cpp default ~0.75) is gone and it always drafts exactly N.
  - **Correctness impact: none.** The **target verify pass rejects bad drafts**
    (`llama-context.cpp:2896-2899`); output is exactly the target's regardless of
    draft depth or greedy-vs-sampled. Fused greedy draft was measured
    **byte-identical** to the sequential path in #454.
  - **Throughput impact: near-neutral at n=2.** `p_min` exists to avoid spending a
    *sequential* per-step `llama_decode`+readback on a token that will likely be
    rejected. The fused path computes all N steps in **one** launch, so that
    per-step cost the gate protected against no longer exists. Always-drafting-2
    costs only the fused graph's marginal step-2 compute + one extra position in
    the verify batch (3 vs 2 — negligible), and can *gain* acceptance when a
    sub-threshold second token is actually correct. At large n_max the waste
    grows, but that's where fused's cheap launch wins most, and we cap n anyway.
  - **Recovery if measurement ever shows it matters (no per-step sync):** emit a
    second small in-graph tensor — the per-step softmax-max (confidence) alongside
    the argmax — read it back in the *same* single D2H, and trim the proposed
    draft block host-side by `p_min` before submitting to verify. That restores
    adaptive depth for the price of one extra tiny readback, not N syncs. Keep as a
    follow-up lever; do not build unless S5's acceptance-delta check regresses.
- **Unified-cache `init_mtp` is new surface.** iSWA has `init_mtp`; the unified
  cache does not. This is the highest-uncertainty delta (S3) — the N draft
  positions must be placeable and each step's KV visible to the next step's
  attention. De-risk by reading `llama_kv_cache_iswa::init_mtp` and mirroring on
  the unified type.
- **Likely host-only (no CUDA rebuild).** The chain uses existing ops
  (`ggml_argmax`, `ggml_get_rows`, attention). If true, S4 skips the CUDA DSO
  rebuild — confirm no new device kernel is introduced before assuming so.
- **Effort:** multi-session. S1 (refactor) + S2 (builder) are the bulk; S3
  (KV/driver) is the risk. Payoff is ~16% on the 35B-class NextN spec path only —
  base and the KV/DCA/turbo moat are unaffected either way.

---

## Measured — Qwen3.6-35B-A3B NextN, ours vs b9859 (bs2, 2026-07-14, #654 short4)

Coherent prompt, `n_ctx=4096`, greedy, `cache_prompt=false`, `n_predict=300`, `--spec-type
draft-mtp --spec-draft-n-max {1,2,3}`. ours = `llamafile.dualctx 7abd684039a3`; b9859 =
upstream `llama-server`.

| config | ours decode t/s | ours accept | b9859 decode t/s | b9859 accept |
|--------|-----------------|-------------|------------------|--------------|
| base (no spec) | 208.3 | — | 212.0 | — |
| n1 | 226.4 | 0.892 | 258.3 | 0.869 |
| n2 | 230.4 | 0.739 | 269.9 | 0.770 |
| n3 | 240.1 | 0.638 | 277.3 | 0.735 |

**Read:** base is at **parity** (208 vs 212). ours' NextN decode still trails b9859 by **~12–15 %**
(n1 −12 %, n2 −15 %, n3 −13 %) — the unchanged **bug-858 Qwen verify/CUDA-graph gap** (accept is at
or above upstream at n1; the tps deficit is spec-orchestration/verify host overhead, not draft
quality). This supersedes the pre-#614 "~16 % slower (244 vs 283)" note with current-binary numbers.
The earlier #654 pass BOOT-FAILED on ours (harness omitted `--server`; llamafile came up in CLI chat
mode — buglog bug-2179); re-run with the conditional `--server` fix produced the table above.

---

## Anchors (quick index)

- Un-fused loop: `common/speculative.cpp:702-803` (readback `:752`).
- Qwen NextN builder: `src/models/qwen35.cpp:510-656` (dispatch `:133-138`);
  MoE `src/models/qwen35moe.cpp:156-158, 575`.
- Gemma fused model: `src/models/gemma4-assistant.cpp:45-258` (one-step),
  `:260-374` (N-step chain), argmax concat `:362-368`.
- Gemma fused driver: `src/llama-context.cpp:2900-3017`;
  runner `:1538-1588`; params `:2673-2698`; reserve `:1419-1470`.
- Graph types: `src/llama-graph.h:37-39`; `n_mtp_steps` `:595-599`;
  `llm_graph_input_mtp` `:128-145`; result accessors `:679-719`.