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)

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).
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.
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.
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.
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).
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.
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.