server/llama: honor per-request speculative.n_max, per-sequence snapshot budget, honest rollback refusal (multi-slot decode with MTP)

common/speculative.cpp

@@ -3148,7 +3148,14 @@ common_speculative_init_result::common_speculative_init_result(
     // the draft context holds as many tokens per sequence as the target context
     cparams.n_ctx = llama_n_ctx(ctx_tgt);
 
-    // n_rs_seq stays as common_context_params_to_llama set it: the draft context needs the same rollback window as the target, with n_rs_seq == 0 its seq_rm fails silently on partial acceptance and keeps stale positions
+    // no rollback snapshot planes in the draft context: every drafter decodes its draft one token per
+    // step (or drives a plain KV cache), and a recurrent memory snapshots only the tokens of a
+    // sequence's last micro-batch (see llama_memory_seq_rm), so the partial removal of a draft can
+    // never be honoured from a snapshot -- with planes it used to read a stale one; without them
+    // common_context_can_seq_rm classifies a recurrent drafter as FULL and the server /
+    // speculative-simple take the checkpoint path, the only correct one. A KV-cache drafter is
+    // classified PART as before; the target keeps its window (common_context_params_to_llama)
+    cparams.n_rs_seq  = 0;
     cparams.ctx_other = ctx_tgt;
 
     std::string model_path;

examples/rs-rollback/rs-rollback-multi.cpp

@@ -6,11 +6,20 @@
 // that class of bug at all.
 //
 //   pass A (reference): run S sequences concurrently, no rollback.
-//   pass B (rollback):  identical, but sequence 0 periodically over-generates R
-//                       tokens which are then seq_rm()'d away and re-generated.
+//   pass B (rollback):  identical, but sequence 0 periodically decodes its next
+//                       token together with R draft tokens as ONE ubatch (the
+//                       server's verify shape) and then seq_rm()s the R draft
+//                       positions away.
 //
 // PASS requires every sequence -- including the untouched ones -- to emit
 // exactly the reference token stream.
+//
+// The rollback stays inside the verify ubatch on purpose: a ubatch snapshots the
+// state after each of its last min(T, K) - 1 tokens and never the state it
+// started from, so a rollback across single-token steps is refused. At a
+// rollback step sequence 0's logits come from an (R + 1)-token ubatch and the
+// other sequences decode without it, so a mismatch exactly there can in rare
+// near-tie cases be matmul rounding rather than a snapshot bug.
 #include "llama.h"
 
 #include <cstdio>
@@ -153,8 +162,62 @@ int main(int argc, char ** argv) {
             out[s].push_back(greedy_ith(ctx, vocab, batch.n_tokens - 1));
         }
 
+        const int n_vocab_run = llama_vocab_n_tokens(vocab);
+
+        // invariant at the top of every iteration: out[s].back() is sampled-but-not-yet-decoded
         int since = 0;
         while ((int) out[0].size() < n_predict) {
+            if (with_rollback && since >= stride && (int) out[0].size() + rewind + 2 < n_predict) {
+                n_rollbacks++;
+                // seq 0 alone takes a verify-style step: its pending token plus `rewind`
+                // draft tokens as ONE ubatch (logits for the pending token's row only),
+                // then the drafts are dropped. A ubatch snapshots the state after each
+                // of its last min(T, K) - 1 tokens and never the state it started from,
+                // so with T = rewind + 1 the rollback of `rewind` positions is the
+                // largest this ubatch supports and lands on the state after the pending
+                // token -- needs n_rs_seq >= rewind. The drafts' content is irrelevant.
+                const int npos0 = npos[0];
+                batch.n_tokens  = 0;
+                for (int r = 0; r <= rewind; r++) {
+                    const llama_token t = r == 0 ? out[0].back() : (llama_token) ((out[0].back() + 7*r + 1) % n_vocab_run);
+                    batch.token[r]     = t;
+                    batch.pos[r]       = npos0 + r;
+                    batch.n_seq_id[r]  = 1;
+                    batch.seq_id[r][0] = 0;
+                    batch.logits[r]    = r == 0;
+                    batch.n_tokens++;
+                }
+                if (llama_decode(ctx, batch)) {
+                    fprintf(stderr, "decode(verify) failed\n");
+                    return false;
+                }
+
+                const llama_pos keep = npos0 + 1;
+                if (!llama_memory_seq_rm(mem, 0, keep, -1)) {
+                    fprintf(stderr, "ROLLBACK REFUSED: keep=%d npos0=%d rewind=%d\n", keep, npos0 + rewind + 1, rewind);
+                    return false;
+                }
+                npos[0] = keep;
+                out[0].push_back(greedy_ith(ctx, vocab, 0));
+
+                // the other sequences take their ordinary single-token step
+                std::vector<std::pair<int, llama_token>> others;
+                for (int s = 1; s < n_seqs; s++) {
+                    others.push_back({ s, out[s].back() });
+                }
+                if (!others.empty()) {
+                    if (!submit(others)) {
+                        fprintf(stderr, "decode(step) failed\n");
+                        return false;
+                    }
+                    for (int s = 1; s < n_seqs; s++) {
+                        out[s].push_back(greedy_ith(ctx, vocab, s - 1));
+                    }
+                }
+                since = 0;
+                continue;
+            }
+
             // one token per sequence, all in one batch
             std::vector<std::pair<int, llama_token>> items;
             for (int s = 0; s < n_seqs; s++) {
@@ -168,36 +231,6 @@ int main(int argc, char ** argv) {
                 out[s].push_back(greedy_ith(ctx, vocab, s));
             }
             since++;
-
-            if (with_rollback && since >= stride && (int) out[0].size() + rewind + 2 < n_predict) {
-                n_rollbacks++;
-                // seq 0 alone over-generates, then discards
-                for (int r = 0; r < rewind; r++) {
-                    std::vector<std::pair<int, llama_token>> one = {
-                        { 0, out[0].back() }
-                    };
-                    if (!submit(one)) {
-                        fprintf(stderr, "decode(draft) failed\n");
-                        return false;
-                    }
-                    out[0].push_back(greedy_ith(ctx, vocab, 0));
-                }
-                for (int r = 0; r < rewind; r++) {
-                    out[0].pop_back();
-                }
-
-                // invariant: out[s].back() is sampled-but-not-yet-decoded. The
-                // drafts occupied positions [npos0 - rewind, npos0), and that
-                // range also covers the decode of the still-pending token, so
-                // discarding it restores the exact pre-draft state -- no replay.
-                const llama_pos keep = npos[0] - rewind;
-                if (!llama_memory_seq_rm(mem, 0, keep, -1)) {
-                    fprintf(stderr, "ROLLBACK REFUSED: keep=%d npos0=%d rewind=%d\n", keep, npos[0], rewind);
-                    return false;
-                }
-                npos[0] = keep;
-                since   = 0;
-            }
         }
 
         llama_batch_free(batch);

examples/rs-rollback/rs-rollback.cpp

@@ -1,16 +1,25 @@
-// Correctness harness for the recurrent-state snapshot ring (n_rs_seq > 0).
+// Correctness harness for the recurrent-state snapshot planes (n_rs_seq > 0).
 //
-// A forward-only generation never reads the snapshot slots back, so it cannot
+// A forward-only generation never reads the snapshot planes back, so it cannot
 // catch a bad snapshot write. This does:
 //
 //   pass A (reference): greedily generate N tokens straight through.
-//   pass B (rollback):  generate the same N tokens, but every S tokens
-//                       over-generate by R, then llama_memory_seq_rm() those R
-//                       positions away and re-generate them.
+//   pass B (rollback):  generate the same N tokens, but every S accepted tokens
+//                       decode the next accepted token together with R draft
+//                       tokens as ONE ubatch (the server's verify shape), then
+//                       llama_memory_seq_rm() the R draft positions away.
 //
 // Pass B only lands on the same tokens if the state restored from the snapshot
-// ring is exactly the state that produced them the first time. Any error in the
+// plane is exactly the state after the accepted token. Any error in the
 // conv-window or SSM snapshot write shows up as a token mismatch.
+//
+// The rollback stays inside the verify ubatch on purpose: a ubatch snapshots the
+// state after each of its last min(T, K) - 1 tokens and never the state it
+// started from, so a rollback across single-token steps is refused. Right after
+// a rollback step the accepted token's logits come from an (R + 1)-token ubatch
+// instead of a single-token one; a mismatch exactly there can in rare near-tie
+// cases be matmul rounding rather than a snapshot bug -- the first-mismatch
+// index tells.
 #include "llama.h"
 
 #include <cstdio>
@@ -39,7 +48,7 @@ int main(int argc, char ** argv) {
     int         n_predict = 128;
     int         stride    = 16;  // roll back every `stride` accepted tokens
     int         rewind    = 4;   // how many tokens to discard and re-generate
-    int         n_rs_seq  = 8;   // snapshot ring depth; must be >= rewind + 1
+    int         n_rs_seq  = 8;   // snapshot planes per seq; must be >= rewind
 
     for (int i = 1; i < argc; i++) {
         if (!strcmp(argv[i], "-m") && i + 1 < argc) {
@@ -98,49 +107,54 @@ int main(int argc, char ** argv) {
         }
         int n_pos = n_prompt;
 
+        const int n_vocab = llama_vocab_n_tokens(vocab);
+
         int since = 0;
         while ((int) out.size() < n_predict) {
             llama_token tok = greedy(ctx, vocab);
             out.push_back(tok);
-
-            batch = llama_batch_get_one(&tok, 1);
-            if (llama_decode(ctx, batch)) {
-                fprintf(stderr, "decode failed\n");
-                return false;
-            }
-            n_pos++;
             since++;
 
             if (with_rollback && since >= stride && (int) out.size() + rewind < n_predict) {
                 n_rollbacks++;
-                // over-generate `rewind` tokens, then throw them away
-                for (int r = 0; r < rewind; r++) {
-                    llama_token t = greedy(ctx, vocab);
-                    batch         = llama_batch_get_one(&t, 1);
-                    if (llama_decode(ctx, batch)) {
-                        fprintf(stderr, "decode(draft) failed\n");
-                        return false;
-                    }
-                    n_pos++;
+                // Verify-style step: decode the accepted token together with `rewind`
+                // draft tokens as ONE ubatch (logits for the accepted token's row only),
+                // then drop the drafts. A ubatch snapshots the state after each of its
+                // last min(T, K) - 1 tokens and never the state it started from, so with
+                // T = rewind + 1 a rollback of `rewind` positions is the largest this
+                // ubatch supports and lands on the state after the accepted token --
+                // needs n_rs_seq >= rewind. The drafts' content is irrelevant.
+                llama_batch vb = llama_batch_init(rewind + 1, 0, 1);
+                for (int r = 0; r <= rewind; r++) {
+                    vb.token[r]     = r == 0 ? tok : (llama_token) ((tok + 7*r + 1) % n_vocab);
+                    vb.pos[r]       = n_pos + r;
+                    vb.n_seq_id[r]  = 1;
+                    vb.seq_id[r][0] = 0;
+                    vb.logits[r]    = r == 0;
+                }
+                vb.n_tokens = rewind + 1;
+                const bool ok_verify = llama_decode(ctx, vb) == 0;
+                llama_batch_free(vb);
+                if (!ok_verify) {
+                    fprintf(stderr, "decode(verify) failed\n");
+                    return false;
                 }
-                // Drop the drafts AND the last accepted token in one go, then
-                // re-decode the accepted token so the logits are valid again.
-                // This forces the state to be restored from a snapshot taken
-                // (rewind + 1) steps back -- needs n_rs_seq >= rewind + 1.
-                const llama_pos keep = n_pos - rewind - 1;
+
+                const llama_pos keep = n_pos + 1;
                 if (!llama_memory_seq_rm(mem, 0, keep, -1)) {
-                    fprintf(stderr, "ROLLBACK REFUSED: keep=%d n_pos=%d rewind=%d\n", keep, n_pos, rewind);
+                    fprintf(stderr, "ROLLBACK REFUSED: keep=%d n_pos=%d rewind=%d\n", keep, n_pos + rewind + 1, rewind);
                     return false;
                 }
                 n_pos = keep;
-
-                batch = llama_batch_get_one(&out.back(), 1);
+                since = 0;
+                // the next greedy() reads the accepted token's logits, the batch's only output row
+            } else {
+                batch = llama_batch_get_one(&tok, 1);
                 if (llama_decode(ctx, batch)) {
-                    fprintf(stderr, "decode(replay) failed\n");
+                    fprintf(stderr, "decode failed\n");
                     return false;
                 }
                 n_pos++;
-                since = 0;
             }
         }
         llama_free(ctx);

ggml/include/ggml.h

@@ -2620,6 +2620,31 @@ extern "C" {
             struct ggml_tensor  * rows,
             int                   n_snap_slots);
 
+    // cache-writing form (CPU backend only): like ggml_gated_delta_net_rows, but
+    // the op MUTATES `states` instead of returning state planes -- the result is
+    // the attention scores only, [S_v*H_v, n_tokens*n_seqs]. For sequence i the
+    // kernel reads cache row rows_in[i] (I32, n_seqs entries), updates the live
+    // state in place in row rows_out[0*n_seqs + i] (I64, min(n_tokens, K)*n_seqs
+    // entries; no copy at all when rows_in[i] == rows_out[i]) and writes snapshot
+    // slot s (state s tokens back) for s in 1..min(n_tokens, K)-1 straight into
+    // row rows_out[s*n_seqs + i]. The caller guarantees that rows_in[i] and every
+    // rows_out[s*n_seqs + i] are planes of ONE cell (row % mem_size equal, checked
+    // by the kernel) and that the cells of different sequences are pairwise
+    // distinct, so no work unit ever reads a row another unit writes. Removes the
+    // gather, the dst state planes and the write-back copy from recurrent decode.
+    GGML_API struct ggml_tensor * ggml_gated_delta_net_cache(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * g,
+            struct ggml_tensor  * beta,
+            struct ggml_tensor  * states,   // 2D cache view [S_v*S_v*H_v, n_rows], mutated
+            struct ggml_tensor  * rows_in,  // I32 [n_seqs]
+            struct ggml_tensor  * rows_out, // I64 [min(n_tokens, K)*n_seqs]
+            int                   n_snap_slots,
+            int                   mem_size);
+
     // fold the per-head gate activations into a gated_delta_net op (scalar gate only):
     //   beta -> sigmoid(beta),  g -> a[h] * softplus(g + dt_bias[h])
     // dt_bias and a are F32 with H_v elements; removes four elementwise ops per layer

ggml/src/ggml-cpu/ggml-cpu.c

@@ -3011,7 +3011,9 @@ struct ggml_cplan ggml_graph_plan(
                     {
                         const int64_t S_v = node->src[2]->ne[0];
                         const int64_t K   = ggml_get_op_params_i32(node, 0);
-                        const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0);
+                        // cache-write mode (op_params[2]) computes on the cache row: no state scratch
+                        const bool cache_write = ggml_get_op_params_i32(node, 2) != 0;
+                        const int64_t per_thread = S_v + ((K > 1 && !cache_write) ? S_v * S_v : 0);
                         cur = per_thread * sizeof(float) * n_tasks;
                     } break;
                 case GGML_OP_COUNT: