Qwen: prepare MTP prefix caches and preserve predictor state

Qwen prompt sync previously ran only the trunk. The first predictor step
then wrote at the decode position while attention read the entire nextn
prefix, including unwritten rows or data left by a previous session.
Target verification still gated the generated tokens, but the drafter's
context was incomplete and could reduce proposal acceptance.

Build the predictor K/V, raw indexer keys and pooled block keys while each
prefill chunk's trunk hidden states are available. Share the existing MTP
input projection and attention preparation, using at most three rows of
draft scratch. Skip the predictor attention output, MoE and vocabulary
head, avoiding routed-expert SSD reads for this cache preparation.

Retain the last committed trunk row until its actual successor is known:
- Repair the last predictor row before prompt append, ordinary decode or
  verification, including a sampled token that differs from the argmax.
- Produce the first draft from this retained row rather than requiring an
  initial ordinary decode cycle to seed the predictor.
- Preserve the row in all three verifier snapshots and restore it during
  rejection, exact-sampling replacement and caller-requested rewind.
- Track a contiguous predictor prefix and invalidate provisional chain
  rows and staged inputs when restoring or preparing only the cache.
- Initialize future RoPE positions for recursive drafts without advancing
  the trunk's multimodal position state.

Qwen payload v3 stores only the canonical predictor prefix plus its pending
trunk row. Reject old payload tags that cannot prove a complete prefix;
old Qwen disk checkpoints must be rebuilt. Account for the retained row
and verifier copies in memory estimates: 40 KiB per row on this model,
up to 160 KiB including all snapshots. Clarify that timing statistics count
first-draft acceptance, and document MTP usage with Metal SSD streaming.

Add a model-backed Metal regression comparing cache-only preparation with
the complete MTP layer on identical trunk batches. Cover poisoned caches,
one-token prompts, chunk lookahead, non-argmax append, reused sessions,
MTP on/off payload interoperability, legacy-tag rejection and depth-2,
depth-3 and exact-sampling continuations. Restore FORCE_ACCEPT after the
rewind test and remove assumptions that the first MTP call commits only
one token. Allow the small-buffer CLI regression to use SSD streaming.

Validation on M1 Max 32 GiB, Qwen3.8-Flash-Next-Q2.gguf:
- Standard make and Qwen memory-estimate test passed.
- MTP prefix oracle passed with bitwise cache, residual and prefill logits
  agreement; continuation logits stayed within the established tolerance
  (maximum absolute error below 0.001 across 248320 vocabulary entries).
- Prefill checkpoints, reused restore, rewind, exact-sampling resample,
  n-gram read-failure recovery and small/unfused buffer tests passed.
- Two runs per arm in ABBA order, ctx4096/chunk128, SSD streaming, temp0:
  short-prompt decode 8.13 -> 8.09 t/s; 575-token-prompt decode
  8.87 -> 9.88 t/s (+11.4%), using harmonic means. Output was identical.
  Long-prompt first-draft acceptance rose from 45.5% to 56.2%; prefill
  throughput fell 2.5% on the long prompt and 6.4% on the short prompt.
  Mean total process time improved only 1.1% on the long prompt and
  increased 3.4% on the short prompt, including startup and prefill.
  These are local observations with uncontrolled OS file-cache/background
  variation, not a general 10% speed guarantee. CUDA/ROCm and vision were
  not exercised; the affected Qwen execution path is Metal.

.gitignore

@@ -72,6 +72,7 @@
 /tests/test_deepseek41_gguf
 /tests/test_qwen4_ngrams
 /tests/test_qwen4_ngram_state
+/tests/test_qwen4_mtp_prefill
 /tests/test_deepseek41_graph
 /tests/test_deepseek41_prefill
 /tests/test_deepseek41_cli

Makefile

@@ -739,6 +739,12 @@ tests/test_qwen4_ngram_state.o: tests/test_qwen4_ngram_state.c ds4.c ds4.h
 tests/test_qwen4_ngram_state: tests/test_qwen4_ngram_state.o $(filter-out ds4.o,$(CORE_OBJS))
 	$(CC) $(filter-out -ffast-math,$(CFLAGS)) -o $@ $^ $(METAL_LDLIBS)
 
+tests/test_qwen4_mtp_prefill.o: tests/test_qwen4_mtp_prefill.c ds4.c ds4.h
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) -Wno-unused-function -I. -c -o $@ $<
+
+tests/test_qwen4_mtp_prefill: tests/test_qwen4_mtp_prefill.o $(filter-out ds4.o,$(CORE_OBJS))
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) -o $@ $^ $(METAL_LDLIBS)
+
 ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_gpu_tp.h ds4_gpu_mgpu.h ds4_linux_memory.h ds4_deepseek41_gpu.h ds4_deepseek41_cuda.cuh ds4_glm53_vision_gpu.cuh ds4_deepseek4_vision_gpu.cuh ds4_image.h ds4_iq2_tables_cuda.inc cuda/mmq/ds4_mmq.h
 	$(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu
 
@@ -1044,6 +1050,7 @@ clean:
 	rm -f tests/test_qwen4_ngrams
 	rm -f tests/test_qwen4_memory tests/test_qwen4_ssd_experts
 	rm -f tests/test_qwen4_ngram_state
+	rm -f tests/test_qwen4_mtp_prefill
 	rm -f tests/test_web_recovery
 	rm -f tests/test_metal_ssd_experts
 	rm -f tests/test_metal_command_memory

docs/QWEN38_FLASH_NEXT.md

@@ -204,6 +204,19 @@ For a smaller memory footprint, run
 256-token context, an eight-token prefill chunk and a 1024-expert cache request
 while keeping both comparison sessions and MTP enabled.
 
+To compare the predictor's cache-only prefill with a complete MTP pass:
+
+```sh
+make tests/test_qwen4_mtp_prefill
+./tests/test_qwen4_mtp_prefill gguf/Qwen3.8-Flash-Next-Q2.gguf --ssd-streaming
+```
+
+This checks one-token prompts, chunk boundaries, non-argmax continuation,
+stale cache contents and payload restore, followed by greedy, depth-three
+and exact-sampling continuation checks. It uses a 128-token context and a
+1024-expert cache request. The small-buffer MTP CLI regression also accepts
+`--ssd-streaming`.
+
 [Checkpoint-fix benchmark charts and measurements](../speed-bench/qwen38-checkpoints/README.md)
 compare prefill, ordinary decode, and MTP decode against the preceding PR head.
 

docs/SPECULATIVE_DECODING.md

@@ -70,6 +70,39 @@ Ordinary decode uses the same file with `--mtp` omitted. For non-zero
 temperature, add `--mtp-exact-sampling` to preserve the target sampling
 distribution. See [Qwen setup](QWEN38_FLASH_NEXT.md) for the Metal runtime.
 
+For the Q2 model with Metal SSD streaming:
+
+```sh
+./ds4 --metal -m gguf/Qwen3.8-Flash-Next-Q2.gguf \
+  --ssd-streaming --ctx 4096 --prefill-chunk 128 \
+  --mtp-timing --temp 0 --nothink -n 100 \
+  -p "narrami la storia di roma"
+```
+
+`--mtp-timing` enables MTP and reports verification cycles and first-draft
+acceptance. Measure generation throughput against the same command without
+that option: speculative decoding can lose time when too few proposals are
+accepted, especially when loading the predictor's experts from SSD.
+
+Qwen's [official model configuration](https://huggingface.co/Qwen/Qwen3.8-Flash-Next/blob/main/config.json)
+contains one MTP layer with no dedicated token embeddings. The same layer
+can be called repeatedly to propose multiple tokens; the number of MTP
+layers is not the draft length. The
+[official blog](https://www.alibabacloud.com/blog/qwen3-8-flash-next-a-new-architecture-towardsultimate-cost-efficiency_603501)
+describes training this module over multiple steps and using QSA in it.
+These weights are already included in the GGUF; no separate draft model is
+needed.
+
+Prompt prefill prepares the predictor's K/V and indexer cache from the
+trunk hidden states and the next prompt token. It skips the predictor's
+attention output, MoE and vocabulary head, so this preparation does not
+stream predictor experts from SSD. The final hidden row is retained until
+the actual next token is known. Snapshot restore and speculative rollback
+retain that row too; this costs 40 KiB per row for this model, at most
+160 KiB including the three verifier snapshots. Qwen payload version 3
+stores the prepared prefix and retained row; older Qwen checkpoints must
+be rebuilt.
+
 The cycle drafts one token ahead by default and engages a **second, chained
 draft** (one extra nextn-layer step conditioned on the predictor's own
 stream, verified in a 3-row pass) while recent first-draft acceptance is

ds4.c

@@ -1073,7 +1073,7 @@ static ds4_context_memory qwen4_graph_memory_estimate(uint32_t ctx, uint32_t pre
      * model allocations, not graph workspace. */
     uint64_t fixed = (mtp ? 4u : 1u) * state + attn_part + DS4_N_EXPERT +
         (uint64_t)(mtp ? 12u : 4u) * DS4_N_VOCAB;
-    if (mtp) fixed += 12u * (hc + 1u) * E + 1u +
+    if (mtp) fixed += 12u * (hc + 1u) * E + 4u * hc_dim + 1u +
         2u * (((uint64_t)DS4_N_VOCAB + 4095u) / 4096u);
     m.scratch_bytes = (T * ((gpu_row > old_row ? gpu_row : old_row) + hc_dim + 4u) + fixed) * 4u;
     m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes;
@@ -57542,6 +57542,12 @@ typedef struct {
     /* MTP: staged predictor input, its residual, and the recurrent-state
      * snapshot that verify rollback restores */
     ds4_gpu_tensor *mtp_e, *mtp_cat, *mtp_proj, *mtp_R, *mtp_argmax, *mtp_argmax_tmp;
+    /* Last committed trunk stream: the next token supplies its predictor
+     * embedding. Keep one row across append, sampling and payload restore. */
+    ds4_gpu_tensor *mtp_tail_R, *snap_tail_R, *snap2_tail_R, *snap0_tail_R;
+    uint32_t mtp_tail_pos;
+    int mtp_tail_next_token;
+    bool mtp_tail_valid, mtp_tail_cached;
     /* MTP draft head over a ranked vocabulary subset: gathered Q8 rows, their
      * token ids, and whether the lazy load was attempted */
     ds4_gpu_tensor *draft_head;
@@ -57688,8 +57694,9 @@ static void qwen4_graph_free(ds4_qwen4_gpu_graph *g) {
         &g->score, &g->tile_max, &g->sel_blocks, &g->sel_tokens, &g->n_sel, &g->attn_part,
         &g->router, &g->selected, &g->weights, &g->mid, &g->part, &g->sh_gate_logit, &g->logits,
         &g->moe_lists, &g->moe_counts, &g->sh_gate, &g->sh_up, &g->sh_mid, &g->sh_out, &g->hc_u, &g->hc_lo_act,
-        &g->inj_alt, &g->mtp_e, &g->mtp_cat, &g->mtp_proj, &g->mtp_R, &g->mtp_argmax, &g->mtp_argmax_tmp,
-        &g->snap_ple_hist, &g->snap2_ple_hist, &g->snap0_ple_hist, &g->pos3,
+        &g->inj_alt, &g->mtp_e, &g->mtp_cat, &g->mtp_proj, &g->mtp_R, &g->mtp_argmax, &g->mtp_argmax_tmp, &g->mtp_tail_R,
+        &g->snap_ple_hist, &g->snap2_ple_hist, &g->snap0_ple_hist,
+        &g->snap_tail_R, &g->snap2_tail_R, &g->snap0_tail_R, &g->pos3,
         &g->draft_head, &g->steer_dirs,
     };
     free(g->host_pos3);
@@ -57799,6 +57806,7 @@ static bool qwen4_graph_alloc(ds4_qwen4_gpu_graph *g, const ds4_weights *w, uint
     QWEN4_ALLOC(hc_u, T * hc_dim);
     QWEN4_ALLOC(hc_lo_act, T * DS4_N_HC_LOWRANK);
     QWEN4_ALLOC(logits, (uint64_t)g->n_logit_rows * DS4_N_VOCAB);
+    if (DS4_N_NEXTN_PREDICT) { QWEN4_ALLOC(mtp_tail_R, hc_dim); }
     if (mtp) {
         QWEN4_ALLOC(mtp_e, 3u * E);
         QWEN4_ALLOC(mtp_cat, 3u * (hc + 1u) * 2u * E);
@@ -57807,6 +57815,7 @@ static bool qwen4_graph_alloc(ds4_qwen4_gpu_graph *g, const ds4_weights *w, uint
         QWEN4_ALLOC(mtp_argmax, 1u);
         QWEN4_ALLOC(mtp_argmax_tmp, ((DS4_N_VOCAB + 4095u) / 4096u) * 2u);
         QWEN4_ALLOC(snap_ple_hist, (uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * hc_dim);
+        QWEN4_ALLOC(snap_tail_R, hc_dim);
     }
 #undef QWEN4_ALLOC
     for (uint32_t il = 0; il < DS4_N_LAYER && ok; il++) {
@@ -57941,6 +57950,7 @@ static void qwen4_graph_reset(ds4_qwen4_gpu_graph *g) {
     g->pos = 0;
     g->mtp_pos = 0;
     g->mtp_last_rows = 0;
+    g->mtp_tail_valid = g->mtp_tail_cached = false;
     g->mrope_delta = 0;
     g->snap_valid = false;
     g->snap2_valid = false;
@@ -58223,11 +58233,10 @@ static bool qwen4_graph_attention_core(ds4_qwen4_gpu_graph *g, uint32_t il,
     return ok;
 }
 
-static bool qwen4_graph_attention(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l,
+static bool qwen4_graph_attention_cache(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l,
                                   uint32_t il, uint32_t pos0, uint32_t T) {
     const uint32_t ratio = 4u;
     const uint32_t last = pos0 + T - 1u;
-    const uint32_t q_dim = DS4_N_HEAD * DS4_N_HEAD_DIM;
     {
         bool ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T);
         if (ok && qwen4_graph_fused(g, T)) {
@@ -58264,6 +58273,15 @@ static bool qwen4_graph_attention(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
                                             DS4_N_ROT, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS)) {
         return false;
     }
+    return true;
+}
+
+static bool qwen4_graph_attention(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l,
+                                  uint32_t il, uint32_t pos0, uint32_t T) {
+    const uint32_t ratio = 4u;
+    const uint32_t n_blocks_after = (pos0 + T) / ratio;
+    const uint32_t q_dim = DS4_N_HEAD * DS4_N_HEAD_DIM;
+    if (!qwen4_graph_attention_cache(g, m, l, il, pos0, T)) return false;
     if (T == 3u && g->verify_rows_exact) {
         /* 3-row speculative verify: run the attention core as 2/1-row
          * sub-batches so every dispatch keeps the exact T <= 2 kernel paths
@@ -58467,6 +58485,26 @@ static bool qwen4_graph_stage_inputs(ds4_qwen4_gpu_graph *g, const ds4_model *m,
     return ds4_gpu_tensor_write(g->ple_emb, 0, row, (uint64_t)T * E * sizeof(float)) != 0;
 }
 
+static bool qwen4_graph_mtp_cache_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m,
+                                        const ds4_weights *w, uint32_t row,
+                                        const int *next_tokens, uint32_t T, uint32_t idx);
+static bool qwen4_graph_mtp_boundary(ds4_qwen4_gpu_graph *g, const ds4_model *m,
+                                     const ds4_weights *w, int next_token) {
+    if (!g->mtp_R || !g->pos) return true;
+    if (!g->mtp_tail_valid || g->mtp_tail_pos + 1u != g->pos || g->mtp_pos + 1u < g->pos)
+        return false;
+    if (!g->mtp_tail_cached || g->mtp_tail_next_token != next_token || g->mtp_pos < g->pos) {
+        ds4_gpu_tensor *R = g->R;
+        g->R = g->mtp_tail_R;
+        const bool ok = qwen4_graph_mtp_cache_steps(g, m, w, 0, &next_token, 1u, g->pos - 1u);
+        g->R = R;
+        if (!ok) return false;
+    }
+    /* Rows produced by a recursive draft beyond this boundary are provisional. */
+    g->mtp_pos = g->pos;
+    return true;
+}
+
 /* Forward T tokens at g->pos..; logits (optional) receive the last token's
  * row, or all T rows with all_rows (T <= n_logit_rows).  T <= cap_tokens;
  * everything is causal by construction because the recurrent kernels walk
@@ -58502,7 +58540,8 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
         }
     }
     const double t0 = timing ? now_sec() : 0.0;
-    if (!qwen4_graph_stage_inputs(g, m, w, tokens, T)) return false;
+    if (!qwen4_graph_mtp_boundary(g, m, w, tokens[0]) ||
+        !qwen4_graph_stage_inputs(g, m, w, tokens, T)) return false;
     const double t1 = timing ? now_sec() : 0.0;
     if (!glm_graph_begin_commands_if_needed()) return false;
     bool ok = true;
@@ -58600,6 +58639,14 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
         }
         if (ok) ok = qwen4_gemv(g->logits, m, w->output, g->mixed, 1);
     }
+    if (ok && g->mtp_R) {
+        const uint64_t row_bytes = (uint64_t)hc_dim * sizeof(float);
+        ok = ds4_gpu_tensor_copy(g->mtp_tail_R, 0, g->R, (uint64_t)(T - 1u) * row_bytes, row_bytes) != 0;
+        if (ok && g->snap_after_first)
+            ok = ds4_gpu_tensor_copy(g->snap_tail_R, 0, g->R, 0, row_bytes) != 0;
+        if (ok && g->snap_after_second)
+            ok = ds4_gpu_tensor_copy(g->snap2_tail_R, 0, g->R, row_bytes, row_bytes) != 0;
+    }
     const double t2 = timing ? now_sec() : 0.0;
     if (!ds4_gpu_end_commands()) ok = false;
     const double t3 = timing ? now_sec() : 0.0;
@@ -58617,7 +58664,14 @@ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *
             memset(acc, 0, sizeof(acc));
         }
     }
-    if (ok) g->pos += T;
+    if (ok) {
+        g->pos += T;
+        if (g->mtp_R) {
+            g->mtp_tail_pos = g->pos - 1u;
+            g->mtp_tail_valid = true;
+            g->mtp_tail_cached = false;
+        }
+    }
     if (g->snap_after_first) {
         g->snap_valid = ok;
         g->snap_after_first = false;
@@ -58639,9 +58693,9 @@ static bool qwen4_graph_forward_token(ds4_qwen4_gpu_graph *g, const ds4_model *m
  * saving to it (save) or restoring from it. */
 static bool qwen4_graph_state_copy_set(ds4_qwen4_gpu_graph *g, bool save,
                                        ds4_gpu_tensor **snap_state, ds4_gpu_tensor **snap_hist,
-                                       ds4_gpu_tensor **snap_ple, int *snap_prev,
+                                       ds4_gpu_tensor **snap_ple, ds4_gpu_tensor *snap_tail, int *snap_prev,
                                        uint32_t *snap_pos, int32_t *snap_mrope) {
-    if (!*snap_ple) return false;
+    if (!*snap_ple || !snap_tail) return false;
     const uint64_t conv_dim = DS4_N_LIN_CONV_DIM;
     const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM;
     const uint64_t state_bytes = v_dim * DS4_N_LIN_HEAD_DIM * sizeof(float);
@@ -58660,6 +58714,11 @@ static bool qwen4_graph_state_copy_set(ds4_qwen4_gpu_graph *g, bool save,
         ok = ds4_gpu_tensor_copy(save ? *snap_ple : g->ple_hist, 0,
                                  save ? g->ple_hist : *snap_ple, 0, ple_bytes) != 0;
     }
+    if (ok) {
+        const uint64_t bytes = (uint64_t)DS4_N_EMBD * DS4_N_HC * sizeof(float);
+        ok = ds4_gpu_tensor_copy(save ? snap_tail : g->mtp_tail_R, 0,
+                                 save ? g->mtp_tail_R : snap_tail, 0, bytes) != 0;
+    }
     if (!ds4_gpu_end_commands()) ok = false;
     if (save) {
         memcpy(snap_prev, g->ple_prev, sizeof(g->ple_prev));
@@ -58669,25 +58728,30 @@ static bool qwen4_graph_state_copy_set(ds4_qwen4_gpu_graph *g, bool save,
         memcpy(g->ple_prev, snap_prev, sizeof(g->ple_prev));
         g->pos = *snap_pos;
         g->mrope_delta = *snap_mrope;
+        g->mtp_tail_valid = ok && g->pos != 0;
+        g->mtp_tail_cached = false;
+        g->mtp_tail_pos = g->pos ? g->pos - 1u : 0u;
+        g->mtp_last_rows = 0;
+        if (g->mtp_pos >= g->pos) g->mtp_pos = g->pos ? g->pos - 1u : 0u;
     }
     return ok;
 }
 
 static bool qwen4_graph_state_copy(ds4_qwen4_gpu_graph *g, bool save) {
     return qwen4_graph_state_copy_set(g, save, g->snap_lin_state, g->snap_lin_hist,
-                                      &g->snap_ple_hist, g->snap_ple_prev,
+                                      &g->snap_ple_hist, g->snap_tail_R, g->snap_ple_prev,
                                       &g->snap_pos, &g->snap_mrope_delta);
 }
 
 static bool qwen4_graph_state_copy2(ds4_qwen4_gpu_graph *g, bool save) {
     return qwen4_graph_state_copy_set(g, save, g->snap2_lin_state, g->snap2_lin_hist,
-                                      &g->snap2_ple_hist, g->snap2_ple_prev,
+                                      &g->snap2_ple_hist, g->snap2_tail_R, g->snap2_ple_prev,
                                       &g->snap2_pos, &g->snap2_mrope_delta);
 }
 
 static bool qwen4_graph_state_copy0(ds4_qwen4_gpu_graph *g, bool save) {
     return qwen4_graph_state_copy_set(g, save, g->snap0_lin_state, g->snap0_lin_hist,
-                                      &g->snap0_ple_hist, g->snap0_ple_prev,
+                                      &g->snap0_ple_hist, g->snap0_tail_R, g->snap0_ple_prev,
                                       &g->snap0_pos, &g->snap0_mrope_delta);
 }
 
@@ -58716,6 +58780,12 @@ static bool qwen4_graph_state_swap(ds4_qwen4_gpu_graph *g) {
     memcpy(g->ple_prev, g->snap_ple_prev, sizeof(g->ple_prev));
     g->pos = g->snap_pos;
     g->mrope_delta = g->snap_mrope_delta;
+    t = g->mtp_tail_R; g->mtp_tail_R = g->snap_tail_R; g->snap_tail_R = t;
+    g->mtp_tail_valid = g->pos != 0;
+    g->mtp_tail_cached = false;
+    g->mtp_tail_pos = g->pos ? g->pos - 1u : 0u;
+    g->mtp_last_rows = 0;
+    if (g->mtp_pos >= g->pos) g->mtp_pos = g->pos ? g->pos - 1u : 0u;
     g->snap_valid = false;
     return true;
 }
@@ -58724,7 +58794,7 @@ static bool qwen4_graph_state_swap(ds4_qwen4_gpu_graph *g) {
  * become a supposedly valid snapshot on the next attempt. */
 static bool qwen4_graph_ensure_snapshot(ds4_qwen4_gpu_graph *g,
                                         ds4_gpu_tensor **state, ds4_gpu_tensor **hist,
-                                        ds4_gpu_tensor **ple) {
+                                        ds4_gpu_tensor **ple, ds4_gpu_tensor **tail) {
     if (*ple) return true;
     if (!g->snap_ple_hist) return false;
     const uint64_t conv_dim = DS4_N_LIN_CONV_DIM;
@@ -58734,6 +58804,8 @@ static bool qwen4_graph_ensure_snapshot(ds4_qwen4_gpu_graph *g,
     *ple = qwen4_graph_alloc_f32((uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM *
                                               (uint64_t)DS4_N_EMBD * DS4_N_HC);
     if (!*ple) return false;
+    *tail = qwen4_graph_alloc_f32((uint64_t)DS4_N_EMBD * DS4_N_HC);
+    if (!*tail) { ds4_gpu_tensor_free(*ple); *ple = NULL; return false; }
     for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
         if (!g->snap_lin_state[il]) continue;
         state[il] = qwen4_graph_alloc_f32(state_n);
@@ -58745,7 +58817,8 @@ static bool qwen4_graph_ensure_snapshot(ds4_qwen4_gpu_graph *g,
                 state[j] = hist[j] = NULL;
             }
             ds4_gpu_tensor_free(*ple);
-            *ple = NULL;
+            ds4_gpu_tensor_free(*tail);
+            *ple = *tail = NULL;
             return false;
         }
     }
@@ -58754,7 +58827,7 @@ static bool qwen4_graph_ensure_snapshot(ds4_qwen4_gpu_graph *g,
 
 /* Depth-2 sessions never pay for the second verifier snapshot. */
 static bool qwen4_graph_ensure_snap2(ds4_qwen4_gpu_graph *g) {
-    return qwen4_graph_ensure_snapshot(g, g->snap2_lin_state, g->snap2_lin_hist, &g->snap2_ple_hist);
+    return qwen4_graph_ensure_snapshot(g, g->snap2_lin_state, g->snap2_lin_hist, &g->snap2_ple_hist, &g->snap2_tail_R);
 }
 
 /* Allocate the pre-verify snapshot set on the first exact-sampling cycle.
@@ -58762,7 +58835,7 @@ static bool qwen4_graph_ensure_snap2(ds4_qwen4_gpu_graph *g) {
  * failed allocation only means boundary resamples fall back to the old
  * reset-and-replay rewind. */
 static bool qwen4_graph_ensure_snap0(ds4_qwen4_gpu_graph *g) {
-    return qwen4_graph_ensure_snapshot(g, g->snap0_lin_state, g->snap0_lin_hist, &g->snap0_ple_hist);
+    return qwen4_graph_ensure_snapshot(g, g->snap0_lin_state, g->snap0_lin_hist, &g->snap0_ple_hist, &g->snap0_tail_R);
 }
 
 /* Same swap-restore against the second snapshot (state after row 1) for a
@@ -58784,6 +58857,12 @@ static bool qwen4_graph_state_swap2(ds4_qwen4_gpu_graph *g) {
     memcpy(g->ple_prev, g->snap2_ple_prev, sizeof(g->ple_prev));
     g->pos = g->snap2_pos;
     g->mrope_delta = g->snap2_mrope_delta;
+    t = g->mtp_tail_R; g->mtp_tail_R = g->snap2_tail_R; g->snap2_tail_R = t;
+    g->mtp_tail_valid = g->pos != 0;
+    g->mtp_tail_cached = false;
+    g->mtp_tail_pos = g->pos ? g->pos - 1u : 0u;
+    g->mtp_last_rows = 0;
+    if (g->mtp_pos >= g->pos) g->mtp_pos = g->pos ? g->pos - 1u : 0u;
     g->snap2_valid = false;
     return true;
 }
@@ -58795,14 +58874,10 @@ static bool qwen4_graph_state_swap2(ds4_qwen4_gpu_graph *g) {
  * The nextn layer runs on its own residual; want_logits returns the last
  * row only. Batching catches up accepted drafts and predicts again with
  * one command submission and one set of weight reads. */
-static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
-                                 uint32_t row, const int *next_tokens, uint32_t T, uint32_t idx, bool want_logits,
-                                 float *logits_out, int *draft_out) {
-    if (!g->mtp_R || T == 0u || T > 3u || idx > g->ctx_cap || T > g->ctx_cap - idx ||
-        row > g->cap_tokens || T > g->cap_tokens - row) return false;
+static bool qwen4_graph_mtp_input(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
+                                   uint32_t row, const int *next_tokens, uint32_t T) {
     const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC;
-    const uint32_t il = DS4_N_LAYER - 1u;
-    const ds4_layer_weights *l = &w->layer[il];
+    const ds4_layer_weights *l = &w->layer[DS4_N_LAYER - 1u];
     for (uint32_t t = 0; t < T; t++) {
         if (next_tokens[t] < 0 || next_tokens[t] >= (int)DS4_N_VOCAB) return false;
         qwen4_ref_row(m, w->token_embd, (uint64_t)next_tokens[t], g->host_row + (uint64_t)t * E);
@@ -58836,6 +58911,48 @@ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
     }
     g->R = g->mtp_R;
     if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, T);
+    g->R = R_save;
+    return ok;
+}
+
+/* Build only the predictor's attention cache. Historical KV depends on
+ * its input projection, not its attention output or routed FFN. Three rows
+ * reuse the existing draft scratch and never stream predictor experts. */
+static bool qwen4_graph_mtp_cache_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m,
+                                        const ds4_weights *w, uint32_t row,
+                                        const int *next_tokens, uint32_t T, uint32_t idx) {
+    if (!g->mtp_R || !T || T > 3u || idx > g->mtp_pos || idx > g->ctx_cap ||
+        T > g->ctx_cap - idx || row > g->cap_tokens || T > g->cap_tokens - row) return false;
+    ds4_gpu_tensor *R = g->R;
+    bool ok = qwen4_graph_mtp_input(g, m, w, row, next_tokens, T);
+    g->R = g->mtp_R;
+    if (ok) ok = qwen4_graph_attention_cache(g, m, &w->layer[DS4_N_LAYER - 1u],
+                                             DS4_N_LAYER - 1u, idx, T);
+    g->R = R;
+    if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false;
+    g->mtp_last_rows = 0;  /* staged inputs are not recursive predictor outputs */
+    if (ok) {
+        g->mtp_pos = idx + T;
+        if (g->mtp_tail_valid && idx + T - 1u == g->mtp_tail_pos) {
+            g->mtp_tail_cached = true;
+            g->mtp_tail_next_token = next_tokens[T - 1u];
+        }
+    }
+    return ok;
+}
+
+static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
+                                 uint32_t row, const int *next_tokens, uint32_t T, uint32_t idx, bool want_logits,
+                                 float *logits_out, int *draft_out) {
+    if (!g->mtp_R || T == 0u || T > 3u || idx > g->mtp_pos || idx > g->ctx_cap || T > g->ctx_cap - idx ||
+        row > g->cap_tokens || T > g->cap_tokens - row) return false;
+    const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC;
+    const uint32_t il = DS4_N_LAYER - 1u;
+    const ds4_layer_weights *l = &w->layer[il];
+    ds4_gpu_tensor *R_save = g->R;
+    const uint64_t emb_bytes = (uint64_t)E * sizeof(float);
+    bool ok = qwen4_graph_mtp_input(g, m, w, row, next_tokens, T);
+    g->R = g->mtp_R;
     if (ok) ok = qwen4_graph_attention(g, m, l, il, idx, T);
     if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, T, E, hc) != 0;
     if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, T);
@@ -58877,6 +58994,10 @@ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
     if (ok) {
         g->mtp_pos = idx + T;
         g->mtp_last_rows = T;
+        if (g->mtp_tail_valid && idx + T - 1u == g->mtp_tail_pos) {
+            g->mtp_tail_cached = true;
+            g->mtp_tail_next_token = next_tokens[T - 1u];
+        }
     }
     return ok;
 }
@@ -58890,13 +59011,17 @@ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, co
 static bool qwen4_graph_mtp_chain_step(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w,
                                        int next_token, uint32_t idx, int *draft_out) {
     if (!g->mtp_R || g->mtp_last_rows == 0u || g->mtp_last_rows > 3u || !draft_out ||
-        idx >= g->ctx_cap || next_token < 0 || next_token >= (int)DS4_N_VOCAB) return false;
+        idx > g->mtp_pos || idx >= g->ctx_cap || next_token < 0 || next_token >= (int)DS4_N_VOCAB) return false;
     const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC;
     const uint32_t il = DS4_N_LAYER - 1u;
     const ds4_layer_weights *l = &w->layer[il];
     const uint64_t emb_bytes = (uint64_t)E * sizeof(float);
     const uint64_t cat_bytes = (hc + 1u) * 2u * emb_bytes;
     const uint64_t proj_bytes = (hc + 1u) * emb_bytes;
+    uint32_t next_pos[4] = {0};
+    int32_t delta = g->mrope_delta;
+    qwen4_mrope_pos(NULL, 0, idx, &delta, next_pos);
+    if (!ds4_gpu_tensor_write(g->pos3, (uint64_t)idx * sizeof(next_pos), next_pos, sizeof(next_pos))) return false;
     qwen4_ref_row(m, w->token_embd, (uint64_t)next_token, g->host_row);
     if (!ds4_gpu_tensor_write(g->mtp_e, 0, g->host_row, emb_bytes) ||
         !glm_graph_begin_commands_if_needed()) return false;
@@ -62216,7 +62341,7 @@ static int ds41_load_payload(ds4_session *s, FILE *fp, const uint32_t *h,
 }
 #endif
 #ifdef DS4_HAS_QWEN4_METAL
-static uint64_t qwen4_payload_tensor_bytes(uint32_t rows, uint32_t mtp_rows);
+static uint64_t qwen4_payload_tensor_bytes(uint32_t rows, uint32_t mtp_rows, bool tail);
 #endif
 
 uint64_t ds4_session_payload_bytes(ds4_session *s) {
@@ -62230,8 +62355,11 @@ uint64_t ds4_session_payload_bytes(ds4_session *s) {
         bytes += (uint64_t)s->checkpoint.len * sizeof(uint32_t);
         bytes += (uint64_t)DS4_N_VOCAB * sizeof(float);
         const uint32_t rows = (uint32_t)s->checkpoint.len;
-        const uint32_t mtp_rows = s->qwen4_graph.mtp_pos < rows ? s->qwen4_graph.mtp_pos : rows;
-        bytes += qwen4_payload_tensor_bytes(rows, mtp_rows);
+        const ds4_qwen4_gpu_graph *g = &s->qwen4_graph;
+        const bool tail = g->mtp_R && rows && g->mtp_tail_valid &&
+            g->mtp_tail_pos == rows - 1u && g->mtp_pos >= rows - 1u;
+        const uint32_t mtp_rows = tail ? rows - 1u : 0u;
+        bytes += qwen4_payload_tensor_bytes(rows, mtp_rows, tail);
         return bytes;
 #endif
     }
@@ -62375,7 +62503,8 @@ int ds4_session_stage_payload(ds4_session *s, ds4_session_payload_file *out,
  * pooled block keys (attention layers, incl. the MTP block), then the PLE
  * conv history and n-gram context.  The header's raw_live field carries the
  * family tag so a DeepSeek payload is never read as a Qwen one. */
-#define DS4_QWEN4_PAYLOAD_TAG 0x51573802u
+/* v3 stores a contiguous canonical MTP prefix and its pending trunk row. */
+#define DS4_QWEN4_PAYLOAD_TAG 0x51573803u
 
 static uint64_t qwen4_payload_lin_state_bytes(void) {
     return (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM * DS4_N_LIN_HEAD_DIM * sizeof(float);
@@ -62397,8 +62526,9 @@ static uint64_t qwen4_payload_ple_hist_bytes(void) {
     return (uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * DS4_N_EMBD * DS4_N_HC * sizeof(float);
 }
 
-static uint64_t qwen4_payload_tensor_bytes(uint32_t rows, uint32_t mtp_rows) {
-    uint64_t bytes = sizeof(uint32_t);   /* mtp_rows */
+static uint64_t qwen4_payload_tensor_bytes(uint32_t rows, uint32_t mtp_rows, bool tail) {
+    uint64_t bytes = 2u * sizeof(uint32_t);   /* canonical mtp_rows, tail present */
+    if (tail) bytes += (uint64_t)DS4_N_EMBD * DS4_N_HC * sizeof(float);
     for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
         if (ds4_qwen4_layer_is_linear(il)) {
             bytes += qwen4_payload_lin_state_bytes() + qwen4_payload_lin_hist_bytes();
@@ -62453,8 +62583,18 @@ static int qwen4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_
     if (payload_write_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), err, errlen) != 0) return 1;
     uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK);
     int rc = 0;
-    const uint32_t mtp_rows = g->mtp_pos < rows ? g->mtp_pos : rows;
+    const bool tail = g->mtp_R && rows && g->mtp_tail_valid &&
+        g->mtp_tail_pos == rows - 1u && g->mtp_pos >= rows - 1u;
+    const uint32_t mtp_rows = tail ? rows - 1u : 0u;
+    if (g->mtp_R && rows && !tail) {
+        free(buf);
+        payload_set_err(err, errlen, "Qwen3.8 predictor prefix requires replay before snapshot");
+        return 1;
+    }
     if (rc == 0) rc = payload_write_u32(fp, mtp_rows, err, errlen);
+    if (rc == 0) rc = payload_write_u32(fp, tail ? 1u : 0u, err, errlen);
+    if (rc == 0 && tail) rc = payload_write_tensor_span(fp, g->mtp_tail_R, 0,
+        (uint64_t)DS4_N_EMBD * DS4_N_HC * sizeof(float), buf, DS4_SESSION_IO_CHUNK, err, errlen);
     for (uint32_t il = 0; rc == 0 && il < DS4_N_LAYER; il++) {
         if (ds4_qwen4_layer_is_linear(il)) {
             rc = payload_write_tensor_span(fp, g->layer_lin_state[il], 0, qwen4_payload_lin_state_bytes(),
@@ -62529,6 +62669,8 @@ static int qwen4_session_load_payload(ds4_session *s, FILE *fp, const uint32_t *
     s->glm_mtp_have = 0;
     s->glm_mtp_have2 = false;
     g->snap_valid = g->snap2_valid = g->snap0_valid = false;
+    g->mtp_tail_valid = g->mtp_tail_cached = false;
+    g->mtp_last_rows = 0;
     if (payload_read_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), remaining, err, errlen) != 0) {
         token_vec_free(&new_checkpoint);
         return 1;
@@ -62538,13 +62680,16 @@ static int qwen4_session_load_payload(ds4_session *s, FILE *fp, const uint32_t *
         payload_set_err(err, errlen, "failed to synchronize accelerator before Qwen3.8 restore");
         return 1;
     }
-    uint32_t mtp_rows = 0;
+    uint32_t mtp_rows = 0, tail = 0;
     int rc = payload_read_u32(fp, &mtp_rows, remaining, err, errlen);
-    if (rc == 0 && mtp_rows > rows) {
-        payload_set_err(err, errlen, "KV checkpoint MTP rows exceed the token count");
+    if (!rc) rc = payload_read_u32(fp, &tail, remaining, err, errlen);
+    if (!rc && (tail > 1u || (tail ? (!rows || mtp_rows != rows - 1u || !g->mtp_tail_R) : mtp_rows != 0u))) {
+        payload_set_err(err, errlen, "KV checkpoint predictor prefix or pending row is invalid");
         rc = 1;
     }
     uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK);
+    if (!rc && tail) rc = payload_read_tensor_span(fp, g->mtp_tail_R, 0,
+        (uint64_t)DS4_N_EMBD * DS4_N_HC * sizeof(float), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen);
     for (uint32_t il = 0; rc == 0 && il < DS4_N_LAYER; il++) {
         if (ds4_qwen4_layer_is_linear(il)) {
             rc = payload_read_tensor_span(fp, g->layer_lin_state[il], 0, qwen4_payload_lin_state_bytes(),
@@ -62596,6 +62741,9 @@ static int qwen4_session_load_payload(ds4_session *s, FILE *fp, const uint32_t *
     }
     g->pos = rows;
     g->mtp_pos = mtp_rows;
+    g->mtp_tail_valid = tail != 0;
+    g->mtp_tail_pos = rows ? rows - 1u : 0u;
+    g->mtp_tail_cached = false;
     token_vec_free(&s->checkpoint);
     s->checkpoint = new_checkpoint;
     s->checkpoint_valid = true;
@@ -73057,7 +73205,7 @@ void ds4_session_free(ds4_session *s) {
         if (ds4_session_is_qwen4(s)) {
             s->engine->qwen4_session_bytes -= s->qwen4_reserved_graph_bytes;
             if (s->engine && s->engine->glm_mtp_timing && s->qwen4_spec_cycles) {
-                fprintf(stderr, "ds4: Qwen3.8 mtp: %" PRIu64 " verify cycles, %" PRIu64 " drafts accepted (%.1f%%)\n",
+                fprintf(stderr, "ds4: Qwen3.8 mtp: %" PRIu64 " verify cycles, %" PRIu64 " first drafts accepted (%.1f%%)\n",
                         s->qwen4_spec_cycles, s->qwen4_spec_accepted,
                         100.0 * (double)s->qwen4_spec_accepted / (double)s->qwen4_spec_cycles);
             }
@@ -73753,7 +73901,10 @@ static int ds4_session_glm_spec_cycle(ds4_session *s, int first_token,
 /* after a rewind past the verify snapshot the graph was reset: re-prefill the
  * kept transcript so the state matches the checkpoint again */
 static int qwen4_session_replay_if_stale(ds4_session *s, char *err, size_t errlen) {
-    if (!s->qwen4_graph_ready || s->qwen4_graph.pos == (uint32_t)s->checkpoint.len) return 0;
+    if (!s->qwen4_graph_ready) return 0;
+    if (s->qwen4_graph.pos == (uint32_t)s->checkpoint.len &&
+        (!s->qwen4_graph.mtp_R || !s->checkpoint.len ||
+         (s->qwen4_graph.mtp_tail_valid && s->qwen4_graph.mtp_pos + 1u >= (uint32_t)s->checkpoint.len))) return 0;
     if (s->checkpoint_image_count) {
         if (errlen) snprintf(err, errlen, "Qwen3.8 image rewind requires multimodal sync");
         s->checkpoint_valid = false;
@@ -73845,7 +73996,23 @@ static int ds4_session_qwen4_spec_cycle(ds4_session *s, int first_token, float t
     const uint32_t V = DS4_N_VOCAB;
     const int depth = (exact_sampling && temperature > 0.0f) || getenv("DS4_QWEN4_NO_MTP_BATCH") != NULL
         ? 2 : qwen4_spec_depth(s);
-    if (s->glm_mtp_have && first_token != s->glm_mtp_parent) s->glm_mtp_have = 0;
+    if (s->glm_mtp_have && first_token != s->glm_mtp_parent) {
+        s->glm_mtp_have = 0;
+        s->glm_mtp_have2 = false;
+    }
+    if (!s->glm_mtp_have && accepted_cap >= 2 && pos && pos + 2u <= g->ctx_cap &&
+        g->cap_tokens >= 2u && g->mtp_R && g->mtp_tail_valid && qwen4_graph_fused(g, 2u)) {
+        ds4_gpu_tensor *R = g->R;
+        g->R = g->mtp_tail_R;
+        int draft = -1;
+        const bool drafted = qwen4_graph_mtp_step(g, m, w, 0, first_token, pos - 1u, true, NULL, &draft);
+        g->R = R;
+        if (drafted) {
+            s->glm_mtp_parent = first_token;
+            s->glm_mtp_draft = draft;
+            s->glm_mtp_have = 1;
+        }
+    }
     if (depth == 3 && s->glm_mtp_have && s->glm_mtp_have2 && !g->snap2_ple_hist &&
         accepted_cap >= 3 && pos + 3u <= g->ctx_cap && g->cap_tokens >= 3u) {
         (void)qwen4_graph_ensure_snap2(g);
@@ -75022,6 +75189,8 @@ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, c
         s->glm_mtp_have2 = false;
         if (s->checkpoint_valid &&
             s->qwen4_graph.pos == (uint32_t)s->checkpoint.len &&
+            (!s->qwen4_graph.mtp_R || !s->checkpoint.len ||
+             (s->qwen4_graph.mtp_tail_valid && s->qwen4_graph.mtp_pos + 1u >= (uint32_t)s->checkpoint.len)) &&
             (!s->qwen4_rewound || prompt->len > s->checkpoint.len) &&
             prompt->len >= s->checkpoint.len &&
             ds4_tokens_starts_with(prompt, &s->checkpoint)) {
@@ -75063,6 +75232,23 @@ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, c
                 prefill_rc = 1;
                 break;
             }
+            if (s->qwen4_graph.mtp_R) {
+                /* Prompt lookahead covers chunk boundaries. The final row
+                 * waits in mtp_tail_R for the actual next sampled token. */
+                uint32_t known = chunk;
+                if (i + (int)chunk == prompt->len) known--;
+                for (uint32_t row = 0; row < known;) {
+                    const uint32_t n = known - row > 3u ? 3u : known - row;
+                    if (!qwen4_graph_mtp_cache_steps(&s->qwen4_graph, &e->model, &e->weights,
+                            row, prompt->v + i + row + 1u, n, (uint32_t)i + row)) {
+                        snprintf(err, errlen, "Qwen3.8 predictor cache prefill failed at token %u", (uint32_t)i + row);
+                        prefill_rc = 1;
+                        break;
+                    }
+                    row += n;
+                }
+                if (prefill_rc) break;
+            }
             for (uint32_t j = 0; j < chunk; j++) token_vec_push(&s->checkpoint, prompt->v[i + (int)j]);
             i += (int)chunk;
             s->checkpoint_valid = true;

tests/ds4_test.c

@@ -163,6 +163,7 @@ static void test_session_rewind_replay(void) {
     char err[192] = {0};
     ds4_token_score got[8], want[8];
     const bool mtp = test_env_bool("DS4_TEST_GLM_MTP");
+    char *saved_force = test_save_env("DS4_QWEN4_SPEC_FORCE_ACCEPT");
     if (mtp) setenv("DS4_QWEN4_SPEC_FORCE_ACCEPT", "1", 1);
 
     ds4_chat_begin(engine, &prompt);
@@ -174,8 +175,8 @@ static void test_session_rewind_replay(void) {
     TEST_ASSERT(ds4_session_sync(live, &prompt, err, sizeof(err)) == 0);
     for (int i = 0; i < prompt.len; i++) ds4_tokens_push(&replay, prompt.v[i]);
 
-    /* three greedy steps; under MTP the first cycle seeds a draft and the
-     * forced-accept cycles after it commit token pairs */
+    /* Three greedy steps; prepared MTP prefixes can verify immediately,
+     * and forced acceptance ensures the final cycle commits a token pair. */
     int last_n = 1;
     for (int step = 0; step < 3; step++) {
         const int first = ds4_session_argmax(live);
@@ -250,6 +251,7 @@ cleanup:
     ds4_tokens_free(&prompt);
     ds4_session_free(fresh);
     ds4_session_free(live);
+    test_restore_env("DS4_QWEN4_SPEC_FORCE_ACCEPT", saved_force);
 }
 
 /* Issue #8 regression: the server rewinds a verified block to its start and
@@ -462,7 +464,7 @@ static void test_qwen_restore_reused_session(void) {
         int accepted[2];
         int n = ds4_session_eval_speculative_argmax(live, ds4_session_argmax(live),
                                                    2, -1, accepted, 2, err, sizeof(err));
-        TEST_ASSERT(n == (step == 0 ? 1 : 2));
+        TEST_ASSERT(step == 0 ? n >= 1 && n <= 2 : n == 2);
         if (n < 1) goto cleanup;
     }
     const ds4_tokens *tokens = ds4_session_tokens(live);

tests/test_qwen4_mtp_limits.py

@@ -11,6 +11,7 @@ def main():
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--model", type=Path, required=True)
     parser.add_argument("--output", type=Path)
+    parser.add_argument("--ssd-streaming", action="store_true")
     args = parser.parse_args()
     root = Path(__file__).resolve().parents[1]
     out = args.output or Path(tempfile.mkdtemp(prefix="qwen-mtp-limits-"))
@@ -18,6 +19,8 @@ def main():
     base = [str(root / "ds4"), "-m", str(args.model.resolve()),
             "--ctx", "256", "--temp", "0",
             "--nothink", "-p", "Count from one to ten.", "-n", "24"]
+    if args.ssd_streaming:
+        base += ["--ssd-streaming", "--ssd-streaming-cache-experts", "1024"]
     for name, chunk, unfused in [("one-row", 1, False),
                                   ("two-rows-depth-three", 2, False),
                                   ("unfused", 128, True)]:

tests/test_qwen4_mtp_prefill.c

@@ -0,0 +1,341 @@
+#include "../ds4.c"
+#include <assert.h>
+
+#ifdef DS4_HAS_QWEN4_METAL
+/* Compare cache-only prefix preparation with the existing complete MTP
+ * layer on exactly the same trunk rows. This uses real Metal/model weights;
+ * it never runs the full CPU model or substitutes a mock dequantizer. */
+enum { TEST_CTX = 128, TEST_CHUNK = 8, TEST_PREFIX = 13, TEST_APPEND = 16, MAX_FRAMES = 4 };
+static char error[256];
+
+static void need(bool ok, const char *what) {
+    if (!ok) {
+        fprintf(stderr, "Qwen MTP prefill: %s: %s\n", what, error);
+        exit(1);
+    }
+}
+
+static void *read_bytes(const ds4_gpu_tensor *t, uint64_t bytes) {
+    void *p = malloc(bytes ? (size_t)bytes : 1u);
+    need(p != NULL, "host comparison allocation");
+    need(!bytes || ds4_gpu_tensor_read(t, 0, p, bytes), "read comparison tensor");
+    return p;
+}
+
+static void exact(const char *what, const void *a, const void *b, size_t bytes) {
+    if (!memcmp(a, b, bytes)) return;
+    const unsigned char *x = a, *y = b;
+    size_t first = 0;
+    while (first < bytes && x[first] == y[first]) first++;
+    fprintf(stderr, "Qwen MTP prefill: %s differs at byte %zu/%zu\n", what, first, bytes);
+    exit(1);
+}
+
+typedef struct {
+    uint32_t rows, mtp_rows, residual_rows;
+    uint64_t kv_bytes, ik_bytes, block_bytes, residual_bytes;
+    void *k, *v, *ik, *block, *tail, *residual;
+    float *logits;
+} frontier;
+
+static frontier capture(ds4_session *s, uint32_t residual_rows) {
+    ds4_qwen4_gpu_graph *g = &s->qwen4_graph;
+    const uint32_t il = DS4_N_LAYER - 1u;
+    const uint64_t row_bytes = (uint64_t)DS4_N_EMBD * DS4_N_HC * sizeof(float);
+    frontier f = {.rows = g->pos, .mtp_rows = g->mtp_pos, .residual_rows = residual_rows};
+    f.kv_bytes = qwen4_payload_kv_bytes(f.mtp_rows);
+    f.ik_bytes = qwen4_payload_ik_bytes(f.mtp_rows);
+    f.block_bytes = qwen4_payload_block_key_bytes(f.mtp_rows);
+    f.residual_bytes = residual_rows * row_bytes;
+    f.k = read_bytes(g->layer_k_cache[il], f.kv_bytes);
+    f.v = read_bytes(g->layer_v_cache[il], f.kv_bytes);
+    f.ik = read_bytes(g->layer_ik_cache[il], f.ik_bytes);
+    f.block = read_bytes(g->layer_block_key[il], f.block_bytes);
+    f.tail = read_bytes(g->mtp_tail_R, row_bytes);
+    f.residual = read_bytes(g->R, f.residual_bytes);
+    f.logits = malloc((size_t)DS4_N_VOCAB * sizeof(float));
+    need(f.logits != NULL, "logits comparison allocation");
+    memcpy(f.logits, s->logits, (size_t)DS4_N_VOCAB * sizeof(float));
+    return f;
+}
+
+static void free_frontier(frontier *f) {
+    free(f->k); free(f->v); free(f->ik); free(f->block);
+    free(f->tail); free(f->residual); free(f->logits);
+    memset(f, 0, sizeof(*f));
+}
+
+static void compare_frontier(ds4_session *s, const frontier *want, bool residual) {
+    ds4_qwen4_gpu_graph *g = &s->qwen4_graph;
+    need(s->checkpoint_valid && g->pos == want->rows &&
+         ds4_session_pos(s) == (int)want->rows, "target frontier position");
+    need(g->mtp_pos == want->mtp_rows, "contiguous MTP prefix position");
+    need(g->mtp_tail_valid && g->mtp_tail_pos + 1u == g->pos, "last trunk row retained");
+    frontier got = capture(s, residual ? want->residual_rows : 0);
+    exact("nextn K", want->k, got.k, want->kv_bytes);
+    exact("nextn V", want->v, got.v, want->kv_bytes);
+    exact("nextn indexer keys", want->ik, got.ik, want->ik_bytes);
+    exact("nextn pooled keys", want->block, got.block, want->block_bytes);
+    exact("retained trunk row", want->tail, got.tail,
+          (size_t)DS4_N_EMBD * DS4_N_HC * sizeof(float));
+    if (residual) exact("unmodified trunk residual", want->residual, got.residual, want->residual_bytes);
+    exact("unmodified target logits", want->logits, got.logits, (size_t)DS4_N_VOCAB * sizeof(float));
+    free_frontier(&got);
+}
+
+static void poison_nextn(ds4_session *s, unsigned char value) {
+    ds4_qwen4_gpu_graph *g = &s->qwen4_graph;
+    const uint32_t il = DS4_N_LAYER - 1u;
+    ds4_gpu_tensor *t[] = {g->layer_k_cache[il], g->layer_v_cache[il],
+                          g->layer_ik_cache[il], g->layer_block_key[il]};
+    for (unsigned i = 0; i < 4; i++) {
+        const uint64_t bytes = ds4_gpu_tensor_bytes(t[i]);
+        void *p = malloc((size_t)bytes);
+        need(p != NULL, "poison allocation"); memset(p, value, (size_t)bytes);
+        need(ds4_gpu_tensor_write(t[i], 0, p, bytes), "poison old nextn cache");
+        free(p);
+    }
+}
+
+/* Full MTP is the oracle. With a known lookahead it also fills the last row
+ * of an intermediate chunk, so graph_forward_tokens never needs to repair
+ * that boundary through the new cache-only helper. */
+static unsigned reference_sync(ds4_session *s, const ds4_tokens *tokens,
+                               uint32_t start, frontier out[MAX_FRAMES]) {
+    ds4_qwen4_gpu_graph *g = &s->qwen4_graph;
+    ds4_engine *e = s->engine;
+    unsigned n = 0;
+    if (!start) {
+        ds4_session_invalidate(s); qwen4_graph_reset(g); poison_nextn(s, 0x55);
+    } else if (g->mtp_pos < start) {
+        need(g->mtp_pos + 1u == start && g->mtp_tail_valid, "reference append boundary");
+        ds4_gpu_tensor *saved = g->R;
+        g->R = g->mtp_tail_R;
+        const bool ok = qwen4_graph_mtp_step(g, &e->model, &e->weights,
+            0, tokens->v[start], start - 1u, false, NULL, NULL);
+        g->R = saved;
+        need(ok, "full MTP append boundary oracle");
+    }
+    for (uint32_t pos = start; pos < (uint32_t)tokens->len;) {
+        uint32_t T = (uint32_t)tokens->len - pos;
+        if (T > TEST_CHUNK) T = TEST_CHUNK;
+        need(qwen4_graph_forward_tokens(g, &e->model, &e->weights,
+            tokens->v + pos, T, s->logits, false), "reference trunk forward");
+        uint32_t rows = T;
+        if (pos + T == (uint32_t)tokens->len) rows--;
+        for (uint32_t row = 0; row < rows;) {
+            const uint32_t count = rows - row < 3u ? rows - row : 3u;
+            need(qwen4_graph_mtp_steps(g, &e->model, &e->weights, row,
+                tokens->v + pos + row + 1u, count, pos + row, false, NULL, NULL),
+                "full MTP prefix oracle");
+            row += count;
+        }
+        for (uint32_t row = 0; row < T; row++) token_vec_push(&s->checkpoint, tokens->v[pos + row]);
+        s->checkpoint_valid = true; s->qwen4_rewound = false;
+        need(n < MAX_FRAMES, "bounded expected frontiers");
+        out[n++] = capture(s, T);
+        pos += T;
+    }
+    return n;
+}
+
+typedef struct { ds4_session *s; const frontier *frames; unsigned n, seen; } progress_check;
+static void check_progress(void *opaque, const char *phase, int completed, int total) {
+    progress_check *p = opaque;
+    if (strcmp(phase, "prefill_chunk")) return;
+    need(p->seen < p->n, "unexpected prefill callback");
+    need(completed == (int)p->frames[p->seen].rows && completed <= total, "chunk callback frontier");
+    compare_frontier(p->s, p->frames + p->seen++, true);
+}
+
+static void checked_sync(ds4_session *s, const ds4_tokens *tokens, frontier *f, unsigned n) {
+    progress_check p = {.s = s, .frames = f, .n = n};
+    ds4_session_set_progress(s, check_progress, &p);
+    need(ds4_session_sync(s, tokens, error, sizeof(error)) == 0, "candidate prefix sync");
+    ds4_session_set_progress(s, NULL, NULL);
+    need(p.seen == n, "every completed chunk checked");
+    need(s->qwen4_graph.mtp_last_rows == 0, "cache-only prefill does not publish predictor output");
+    compare_frontier(s, f + n - 1u, true);
+}
+
+static void close_logits(const float *got, const float *want) {
+    /* The existing 2/3-row verifier and one-row target have distinct
+     * reductions; use the established MTP batch tolerance, not text alone. */
+    float worst = 0.0f;
+    for (uint32_t i = 0; i < DS4_N_VOCAB; i++) {
+        const float d = fabsf(got[i] - want[i]);
+        if (!isfinite(got[i]) || !isfinite(want[i]) || d > 2e-3f + 1e-4f * fabsf(want[i])) {
+            fprintf(stderr, "Qwen MTP target mismatch id=%u got=%g want=%g error=%g\n",
+                    i, (double)got[i], (double)want[i], (double)d);
+            exit(1);
+        }
+        if (d > worst) worst = d;
+    }
+    printf("  all %u target logits: max error %.6g\n", DS4_N_VOCAB, (double)worst);
+}
+
+static void check_cycles(ds4_session *live, ds4_session *target,
+                         const ds4_session_snapshot *snapshot, unsigned mode) {
+    need(ds4_session_load_snapshot(live, snapshot, error, sizeof(error)) == 0 &&
+         ds4_session_load_snapshot(target, snapshot, error, sizeof(error)) == 0, "restore cycle frontier");
+    const bool forced = mode == 1u, stochastic = mode == 2u;
+    need(setenv("DS4_QWEN4_MTP_DEPTH", mode ? "3" : "2", 1) == 0, "set depth case");
+    need((forced ? setenv("DS4_QWEN4_SPEC_FORCE_ACCEPT", "1", 1) :
+                   unsetenv("DS4_QWEN4_SPEC_FORCE_ACCEPT")) == 0, "set acceptance case");
+    live->engine->dspark_exact_sampling = stochastic;
+    uint64_t rng = 0x3456789abcdefull;
+    bool saw_verify = false, saw_three = false;
+    for (unsigned step = 0; step < 3u; step++) {
+        const int first = ds4_session_sample(target, stochastic ? 0.8f : 0.0f,
+                                             0, 0.95f, 0.0f, &rng);
+        int accepted[3] = {-1, -1, -1};
+        const int before = ds4_session_pos(live);
+        const int n = ds4_session_eval_speculative(live, first, 3, -1,
+            stochastic ? 0.8f : 0.0f, 0, 0.95f, 0.0f, &rng, accepted, 3, error, sizeof(error));
+        need(n >= 1 && n <= (stochastic ? 2 : 3), "bounded accepted count");
+        need(accepted[0] == first && ds4_session_pos(live) == before + n, "committed prefix count");
+        saw_verify |= n > 1; saw_three |= n == 3;
+        for (int i = 0; i < n; i++) {
+            need(accepted[i] >= 0 && accepted[i] < (int)DS4_N_VOCAB, "accepted token range");
+            if (!forced && !stochastic)
+                need(accepted[i] == ds4_session_argmax(target), "greedy proposal agrees with target");
+            need(ds4_session_eval(target, accepted[i], error, sizeof(error)) == 0,
+                 "independent one-row target continuation");
+        }
+        close_logits(live->logits, target->logits);
+    }
+    if (forced) need(saw_verify && saw_three, "depth-three verifier actually exercised");
+    printf("PASS MTP continuation %s, independent target logits\n",
+           forced ? "forced depth3 state fixture" : stochastic ? "exact sampling depth2" : "greedy depth2");
+}
+
+static void check_payload_modes(ds4_session *live, ds4_session *reference,
+                                const ds4_tokens *prefix, const ds4_session_snapshot *snapshot) {
+    ds4_engine *e = live->engine;
+    /* Share one model mapping, but allocate the third graph exactly as a
+     * session created with MTP disabled. Ordinary eval uses that graph only. */
+    ds4_session *off = NULL;
+    const bool enabled = e->glm_mtp;
+    e->glm_mtp = false;
+    const int create_rc = ds4_session_create(&off, e, TEST_CTX);
+    e->glm_mtp = enabled;
+    need(create_rc == 0 && off && !off->qwen4_graph.mtp_R, "allocate MTP-disabled graph");
+    need(ds4_session_load_snapshot(off, snapshot, error, sizeof(error)) == 0 &&
+         ds4_session_load_snapshot(reference, snapshot, error, sizeof(error)) == 0,
+         "MTP-on payload loads with MTP disabled");
+    exact("MTP-disabled restored target logits", off->logits, reference->logits,
+          (size_t)DS4_N_VOCAB * sizeof(float));
+    const int next = (ds4_session_argmax(reference) + 1) % (int)DS4_N_VOCAB;
+    need(ds4_session_eval(off, next, error, sizeof(error)) == 0 &&
+         ds4_session_eval(reference, next, error, sizeof(error)) == 0,
+         "MTP-disabled continuation from MTP-on payload");
+    exact("MTP-disabled target continuation", off->logits, reference->logits,
+          (size_t)DS4_N_VOCAB * sizeof(float));
+
+    need(ds4_session_load_snapshot(off, snapshot, error, sizeof(error)) == 0,
+         "restore MTP-disabled snapshot source");
+    ds4_session_snapshot plain = {0};
+    need(ds4_session_save_snapshot(off, &plain, error, sizeof(error)) == 0,
+         "save actual MTP-disabled payload");
+    ds4_session_free(off);
+    poison_nextn(live, 0x77);
+    need(ds4_session_load_snapshot(live, &plain, error, sizeof(error)) == 0,
+         "MTP-off payload loads with MTP enabled");
+    need(live->qwen4_graph.mtp_pos == 0 && !live->qwen4_graph.mtp_tail_valid,
+         "MTP-off payload does not claim a predictor prefix");
+    frontier frames[MAX_FRAMES] = {{0}};
+    const unsigned n = reference_sync(reference, prefix, 0, frames);
+    checked_sync(live, prefix, frames, n);
+    for (unsigned i = 0; i < n; i++) free_frontier(frames + i);
+    ds4_session_snapshot_free(&plain);
+
+    /* The old tag cannot describe the deferred final row. Reject it before
+     * any tensor writes instead of interpreting its high-water mark as KV. */
+    ds4_session_snapshot old = {.len = snapshot->len, .cap = snapshot->len};
+    old.ptr = malloc((size_t)old.len);
+    need(old.ptr != NULL, "old payload allocation");
+    memcpy(old.ptr, snapshot->ptr, (size_t)old.len);
+    uint32_t tag = 0;
+    memcpy(&tag, old.ptr + 12u * sizeof(uint32_t), sizeof(tag));
+    need(tag == DS4_QWEN4_PAYLOAD_TAG, "payload family tag location");
+    tag = 0x51573802u;
+    memcpy(old.ptr + 12u * sizeof(uint32_t), &tag, sizeof(tag));
+    need(ds4_session_load_snapshot(live, &old, error, sizeof(error)) != 0,
+         "reject legacy nextn payload without a retained row");
+    error[0] = '\0';
+    ds4_session_snapshot_free(&old);
+    puts("PASS MTP on/off payload compatibility, prefix replay and legacy-tag rejection");
+}
+
+int main(int argc, char **argv) {
+    const bool streaming = argc == 3 && !strcmp(argv[2], "--ssd-streaming");
+    if (argc != 2 && !streaming) {
+        fprintf(stderr, "usage: %s QWEN_GGUF [--ssd-streaming]\n", argv[0]); return 2;
+    }
+    const char *names[] = {"DS4_QWEN4_SPEC_FORCE_ACCEPT", "DS4_QWEN4_MTP_DEPTH"};
+    char *saved[2] = {NULL, NULL};
+    for (unsigned i = 0; i < 2; i++) {
+        const char *v = getenv(names[i]); saved[i] = v ? strdup(v) : NULL;
+        need(!v || saved[i], "save test environment"); need(unsetenv(names[i]) == 0, "clear test override");
+    }
+    ds4_engine_options opt = {.model_path = argv[1], .backend = DS4_BACKEND_METAL,
+        .context_size = TEST_CTX, .prefill_chunk = TEST_CHUNK, .glm_mtp = true,
+        .ssd_streaming = streaming, .ssd_streaming_cache_experts = streaming ? 1024u : 0u};
+    ds4_engine *engine = NULL;
+    need(ds4_engine_open(&engine, &opt) == 0 && ds4_engine_is_qwen4(engine), "open Qwen model");
+    ds4_session *live = NULL, *reference = NULL;
+    need(ds4_session_create(&live, engine, TEST_CTX) == 0 &&
+         ds4_session_create(&reference, engine, TEST_CTX) == 0, "allocate two bounded sessions");
+    need(live->qwen4_graph.cap_tokens == TEST_CHUNK, "requested chunk size");
+    ds4_tokens prompt = {0};
+    ds4_encode_chat_prompt(engine, NULL, "Count from one to ten and explain the pattern.", DS4_THINK_NONE, &prompt);
+    need(prompt.len >= TEST_APPEND, "enough real prompt token IDs");
+    ds4_tokens prefix = prompt; prefix.len = 1;
+    frontier frames[MAX_FRAMES] = {{0}};
+    unsigned n = reference_sync(reference, &prefix, 0, frames);
+    poison_nextn(live, 0xaa);
+    checked_sync(live, &prefix, frames, n);
+    need(live->qwen4_graph.mtp_pos == 0 && !live->qwen4_graph.mtp_tail_cached,
+         "one-token prefix has only a deferred row");
+    for (unsigned i = 0; i < n; i++) free_frontier(frames + i);
+    ds4_session_invalidate(live);
+    prefix.len = TEST_PREFIX;
+    n = reference_sync(reference, &prefix, 0, frames);
+    poison_nextn(live, 0xaa);
+    checked_sync(live, &prefix, frames, n);
+    need(live->qwen4_graph.mtp_pos + 1u == TEST_PREFIX, "last unknown token is not guessed");
+    for (unsigned i = 0; i < n; i++) free_frontier(frames + i);
+    puts("PASS cache-only prefix against full MTP: chunk lookahead and final retained row");
+
+    /* Force the deferred row to use a real next input, not a guessed argmax. */
+    prompt.v[TEST_PREFIX] = (ds4_session_argmax(reference) + 1) % (int)DS4_N_VOCAB;
+    prefix.len = TEST_APPEND;
+    n = reference_sync(reference, &prefix, TEST_PREFIX, frames);
+    checked_sync(live, &prefix, frames, n);
+    ds4_session_snapshot snapshot = {0};
+    need(ds4_session_save_snapshot(live, &snapshot, error, sizeof(error)) == 0, "save prepared prefix");
+    /* Load into a graph containing unrelated nextn rows and a different tail. */
+    ds4_tokens other = {0};
+    for (int i = 0; i < TEST_PREFIX; i++) ds4_tokens_push(&other, prompt.v[(i + 2) % TEST_APPEND]);
+    need(ds4_session_sync(live, &other, error, sizeof(error)) == 0, "reuse with unrelated prefix");
+    poison_nextn(live, 0x33);
+    need(ds4_session_load_snapshot(live, &snapshot, error, sizeof(error)) == 0, "restore over stale cache");
+    compare_frontier(live, frames + n - 1u, false);
+    for (unsigned i = 0; i < n; i++) free_frontier(frames + i);
+    puts("PASS append boundary and checkpoint restore over reused state");
+
+    check_payload_modes(live, reference, &prefix, &snapshot);
+    for (unsigned mode = 0; mode < 3; mode++) check_cycles(live, reference, &snapshot, mode);
+    ds4_session_snapshot_free(&snapshot);
+    ds4_tokens_free(&other); ds4_tokens_free(&prompt);
+    ds4_session_free(reference); ds4_session_free(live); ds4_engine_close(engine);
+    for (unsigned i = 0; i < 2; i++) {
+        need((saved[i] ? setenv(names[i], saved[i], 1) : unsetenv(names[i])) == 0, "restore test environment");
+        free(saved[i]);
+    }
+    puts("Qwen MTP prefill: all checks passed");
+    return 0;
+}
+#else
+int main(void) { fputs("Qwen MTP prefill test requires Metal\n", stderr); return 2; }
+#endif

tests/test_qwen4_ngram_state.c

@@ -33,8 +33,16 @@ int main(int argc, char **argv) {
         assert(ds4_session_sync(live, &prompt, error, sizeof(error)) == 0);
         int accepted[3];
         if (mode == 2) {
-            assert(ds4_session_eval_speculative_argmax(live, ds4_session_argmax(live),
-                3, -1, accepted, 3, error, sizeof(error)) == 1);
+            /* Prefix preparation may already seed a proposal. Otherwise run
+             * a cycle, checking its committed frontier rather than assuming
+             * the first call always consumes exactly one token. */
+            for (int attempt = 0; !live->glm_mtp_have && attempt < 3; attempt++) {
+                const int before = ds4_session_pos(live);
+                const int n = ds4_session_eval_speculative_argmax(live, ds4_session_argmax(live),
+                    3, -1, accepted, 3, error, sizeof(error));
+                assert(n >= 1 && n <= 3);
+                assert(ds4_session_pos(live) == before + n);
+            }
             assert(live->glm_mtp_have);
         }
         const ds4_tokens *current = ds4_session_tokens(live);