speculative : capped recurrent planes for DFlash (--spec-draft-rs-planes)

common/arg.cpp

@@ -4208,12 +4208,19 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN"));
     add_opt(common_arg(
-        {"--spec-mtp-rs-planes"}, "N",
-        "total target recurrent-state planes for draft-mtp, including the current state (default: 0, allocate spec-draft-n-max + 1)",
+        {"--spec-draft-rs-planes", "--spec-mtp-rs-planes"}, "N",
+        "total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, allocate spec-draft-n-max + 1)",
         [](common_params & params, int value) {
-            params.speculative.mtp_rs_planes = value;
+            params.speculative.rs_planes = value;
         }
-    ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_MTP_RS_PLANES"));
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_RS_PLANES"));
+    add_opt(common_arg(
+        {"--spec-draft-rs-planes-lead"}, "N",
+        "capped recurrent planes that keep the states after the first draft tokens; the others keep the newest states (default: -1, auto: spec-draft-rs-planes - 2 for draft-dflash, 0 for draft-mtp)",
+        [](common_params & params, int value) {
+            params.speculative.rs_planes_lead = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_RS_PLANES_LEAD"));
     add_opt(common_arg(
         {"--spec-draft-ubatch-size", "--ubatch-size-draft", "-ubd"}, "N",
         "physical maximum batch size for the draft context (default: 0, inherit target ubatch); "

common/common.h

@@ -385,14 +385,27 @@ struct common_params_speculative {
 
     common_params_speculative_ngram_cache ngram_cache;
 
-    int32_t mtp_rs_planes = 0; // total target recurrent planes (0 = draft.n_max + 1)
+    int32_t rs_planes = 0; // total target recurrent planes for draft-mtp or draft-dflash (0 = draft.n_max + 1)
+    int32_t rs_planes_lead = -1; // capped planes that keep the states after the first draft tokens (-1 = auto)
 
     bool has_dft() const {
         return !draft.mparams.empty();
     }
 
-    bool is_mtp_rs_capped() const {
-        return mtp_rs_planes > 0 && int64_t(mtp_rs_planes) < int64_t(draft.n_max) + 1;
+    bool is_rs_capped() const {
+        return rs_planes > 0 && int64_t(rs_planes) < int64_t(draft.n_max) + 1;
+    }
+
+    int32_t get_rs_planes_lead() const {
+        if (!is_rs_capped()) {
+            return 0;
+        }
+        if (rs_planes_lead >= 0) {
+            return rs_planes_lead;
+        }
+        // DFlash rejects most drafts early, so keep one plane for full acceptance and the rest for the first tokens
+        const bool has_dflash = std::find(types.begin(), types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != types.end();
+        return has_dflash ? rs_planes - 2 : 0;
     }
 
     bool has_synth() const {
@@ -404,9 +417,11 @@ struct common_params_speculative {
             return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
         });
 
-        const bool has_mtp = std::find(types.begin(), types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != types.end();
-        if (has_mtp && mtp_rs_planes > 0) {
-            return uint32_t(mtp_rs_planes - 1);
+        const bool can_cap = std::any_of(types.begin(), types.end(), [](auto t) {
+            return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH;
+        });
+        if (can_cap && rs_planes > 0) {
+            return uint32_t(rs_planes - 1);
         }
 
         return needs_rs_seq ? uint32_t(std::max(0, draft.n_max)) : 0u;

common/speculative.cpp

@@ -175,8 +175,8 @@ struct common_speculative_impl {
     virtual bool get_state(llama_seq_id /*seq_id*/, std::vector<uint8_t> & /*data*/) const { return false; }
     virtual void set_state(llama_seq_id /*seq_id*/, const std::vector<uint8_t> & /*data*/) {}
 
-    virtual bool get_mtp_replay_state(llama_seq_id /*seq_id*/, std::vector<uint8_t> & /*data*/) const { return false; }
-    virtual bool set_mtp_replay_state(llama_seq_id /*seq_id*/, const std::vector<uint8_t> & /*data*/) { return false; }
+    virtual bool get_replay_state(llama_seq_id /*seq_id*/, std::vector<uint8_t> & /*data*/) const { return false; }
+    virtual bool set_replay_state(llama_seq_id /*seq_id*/, const std::vector<uint8_t> & /*data*/) { return false; }
 };
 
 struct common_speculative_impl_draft_simple : public common_speculative_impl {
@@ -1328,6 +1328,19 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
     void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
         // noop
     }
+
+    // the draft state is the ctx_dft KV cache, and a capped replay rewinds it with seq_rm
+    bool get_replay_state(llama_seq_id seq_id, std::vector<uint8_t> & data) const override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+            return false;
+        }
+        data.clear();
+        return true;
+    }
+
+    bool set_replay_state(llama_seq_id seq_id, const std::vector<uint8_t> & data) override {
+        return seq_id >= 0 && seq_id < (llama_seq_id) n_seq && data.empty();
+    }
 };
 
 struct common_speculative_impl_draft_mtp : public common_speculative_impl {
@@ -1782,7 +1795,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes);
     }
 
-    bool get_mtp_replay_state(llama_seq_id seq_id, std::vector<uint8_t> & data) const override {
+    bool get_replay_state(llama_seq_id seq_id, std::vector<uint8_t> & data) const override {
         if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq ||
                 pending_h[seq_id].size() != (size_t) n_embd) {
             return false;
@@ -1797,7 +1810,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         return true;
     }
 
-    bool set_mtp_replay_state(llama_seq_id seq_id, const std::vector<uint8_t> & data) override {
+    bool set_replay_state(llama_seq_id seq_id, const std::vector<uint8_t> & data) override {
         if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
             return false;
         }
@@ -2335,40 +2348,58 @@ void common_validate_speculative_params(
                 params.draft.n_ubatch, target_ubatch_raw));
     }
 
-    if (params.mtp_rs_planes == 0) {
+    if (params.rs_planes_lead < -1) {
+        throw std::invalid_argument("spec-draft-rs-planes-lead must be -1 (auto) or >= 0");
+    }
+
+    if (params.rs_planes_lead > 0 && !params.is_rs_capped()) {
+        throw std::invalid_argument("spec-draft-rs-planes-lead requires spec-draft-rs-planes below spec-draft-n-max + 1");
+    }
+
+    if (params.rs_planes == 0) {
         return;
     }
 
-    if (!has_mtp) {
-        throw std::invalid_argument("spec-mtp-rs-planes requires --spec-type draft-mtp");
+    const bool has_dflash = std::find(
+            params.types.begin(), params.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != params.types.end();
+    if (!has_mtp && !has_dflash) {
+        throw std::invalid_argument("spec-draft-rs-planes requires --spec-type draft-mtp or draft-dflash");
     }
 
     if (params.draft.n_max < 1) {
-        throw std::invalid_argument("spec-mtp-rs-planes requires spec-draft-n-max >= 1");
+        throw std::invalid_argument("spec-draft-rs-planes requires spec-draft-n-max >= 1");
     }
 
     const int64_t max_planes = int64_t(params.draft.n_max) + 1;
-    if (params.mtp_rs_planes < 2 || int64_t(params.mtp_rs_planes) > max_planes) {
+    if (params.rs_planes < 2 || int64_t(params.rs_planes) > max_planes) {
         throw std::invalid_argument(string_format(
-                "spec-mtp-rs-planes must be 0 or in [2, %" PRId64 "] for spec-draft-n-max=%d",
+                "spec-draft-rs-planes must be 0 or in [2, %" PRId64 "] for spec-draft-n-max=%d",
                 max_planes, params.draft.n_max));
     }
 
-    const bool has_other_recurrent_mode = std::any_of(
+    // the planes are shared by every mode that rolls back the recurrent state, so only one may use them
+    const auto n_recurrent_modes = std::count_if(
             params.types.begin(), params.types.end(), [](common_speculative_type type) {
-                return type == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 ||
+                return type == COMMON_SPECULATIVE_TYPE_DRAFT_MTP ||
+                       type == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 ||
                        type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH ||
                        type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
             });
-    if (params.is_mtp_rs_capped() && has_other_recurrent_mode) {
+    if (params.is_rs_capped() && params.rs_planes_lead > params.rs_planes - 2) {
+        throw std::invalid_argument(string_format(
+                "spec-draft-rs-planes-lead must be in [0, %d] for spec-draft-rs-planes=%d, one plane keeps the newest state",
+                params.rs_planes - 2, params.rs_planes));
+    }
+
+    if (params.is_rs_capped() && n_recurrent_modes > 1) {
         throw std::invalid_argument(
-                "spec-mtp-rs-planes cannot be combined with another speculative mode that requires recurrent rollback");
+                "spec-draft-rs-planes cannot be combined with another speculative mode that requires recurrent rollback");
     }
 
-    if (params.is_mtp_rs_capped() && target_ubatch_effective > 0 &&
+    if (params.is_rs_capped() && target_ubatch_effective > 0 &&
             int64_t(params.draft.n_max) + 1 > target_ubatch_effective) {
         throw std::invalid_argument(string_format(
-                "capped spec-mtp-rs-planes requires ubatch-size >= spec-draft-n-max + 1 (%" PRId64 ")",
+                "capped spec-draft-rs-planes requires ubatch-size >= spec-draft-n-max + 1 (%" PRId64 ")",
                 int64_t(params.draft.n_max) + 1));
     }
 }
@@ -3081,29 +3112,29 @@ void common_speculative_set_state(common_speculative * spec, llama_seq_id seq_id
     }
 }
 
-bool common_speculative_get_mtp_state(common_speculative * spec, llama_seq_id seq_id, std::vector<uint8_t> & data) {
+bool common_speculative_get_replay_state(common_speculative * spec, llama_seq_id seq_id, std::vector<uint8_t> & data) {
     data.clear();
     if (spec == nullptr) {
         return false;
     }
 
     for (auto & impl : spec->impls) {
-        if (impl->get_mtp_replay_state(seq_id, data)) {
+        if (impl->get_replay_state(seq_id, data)) {
             return true;
         }
     }
 
     return false;
 }
 
-bool common_speculative_set_mtp_state(
+bool common_speculative_set_replay_state(
         common_speculative * spec, llama_seq_id seq_id, const std::vector<uint8_t> & data) {
     if (spec == nullptr) {
         return data.empty();
     }
 
     for (auto & impl : spec->impls) {
-        if (impl->set_mtp_replay_state(seq_id, data)) {
+        if (impl->set_replay_state(seq_id, data)) {
             return true;
         }
     }

common/speculative.h

@@ -94,8 +94,8 @@ void common_speculative_accept(common_speculative * spec, llama_seq_id, uint16_t
 bool common_speculative_get_state(common_speculative * spec, llama_seq_id seq_id, std::vector<uint8_t> & data);
 void common_speculative_set_state(common_speculative * spec, llama_seq_id seq_id, const std::vector<uint8_t> & data);
 
-bool common_speculative_get_mtp_state(common_speculative * spec, llama_seq_id seq_id, std::vector<uint8_t> & data);
-bool common_speculative_set_mtp_state(common_speculative * spec, llama_seq_id seq_id, const std::vector<uint8_t> & data);
+bool common_speculative_get_replay_state(common_speculative * spec, llama_seq_id seq_id, std::vector<uint8_t> & data);
+bool common_speculative_set_replay_state(common_speculative * spec, llama_seq_id seq_id, const std::vector<uint8_t> & data);
 
 // print statistics about the speculative decoding
 void common_speculative_print_stats(const common_speculative * spec);

docs/speculative.md

@@ -247,9 +247,12 @@ Use exactly one of these options:
 --spec-draft-n-min                      N
                                         minimum number of draft tokens to use for speculative decoding (default: 0)
                                         (env: LLAMA_ARG_SPEC_DRAFT_N_MIN)
---spec-mtp-rs-planes                    N
-                                        total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1)
-                                        (env: LLAMA_ARG_SPEC_MTP_RS_PLANES)
+--spec-draft-rs-planes, --spec-mtp-rs-planes N
+                                        total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, use spec-draft-n-max + 1)
+                                        (env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES)
+--spec-draft-rs-planes-lead             N
+                                        capped recurrent planes that keep the states after the first draft tokens; the others keep the newest states (default: -1, auto: spec-draft-rs-planes - 2 for draft-dflash, 0 for draft-mtp)
+                                        (env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES_LEAD)
 --spec-draft-p-split, --draft-p-split   P
                                         speculative decoding split probability (default: 0.10)
                                         (env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT)
@@ -264,15 +267,19 @@ Use exactly one of these options:
                                         (use --list-devices to see available devices)
 ```
 
-#### Capped MTP recurrent planes
+#### Capped recurrent planes
 
-`--spec-mtp-rs-planes` applies only to `draft-mtp`. The default value `0` allocates `--spec-draft-n-max + 1` target recurrent-state planes. An explicit value must be in `[2, --spec-draft-n-max + 1]`; a smaller value enables capped replay. Capped replay cannot be combined with Eagle3, DFlash, or DSpark because those modes also control recurrent rollback.
+`--spec-draft-rs-planes` (formerly `--spec-mtp-rs-planes`, still accepted) applies to `draft-mtp` and `draft-dflash`. The default value `0` allocates `--spec-draft-n-max + 1` target recurrent-state planes. An explicit value must be in `[2, --spec-draft-n-max + 1]`; a smaller value enables capped replay. The planes belong to one speculative mode, so capped replay cannot be combined with a second mode that rolls back the recurrent state (MTP, Eagle3, DFlash or DSpark), and DSpark and Eagle3 cannot use it.
 
-The effective target ubatch is `min(batch-size, ubatch-size)`, or `batch-size` when `ubatch-size` is zero. For capped replay it must be at least `--spec-draft-n-max + 1`. A nonzero draft ubatch override must equal the target ubatch.
+Each plane holds one copy of the recurrent state for every sequence, so the saving is `n_parallel * (n_max + 1 - planes)` copies. DFlash keeps its own state in the draft KV cache, which a replay rewinds together with the target, so no extra speculative state is saved for it.
+
+The effective target ubatch is `min(batch-size, ubatch-size)`, or `batch-size` when `ubatch-size` is zero. For capped replay it must be at least `--spec-draft-n-max + 1`. For `draft-mtp`, a nonzero draft ubatch override must equal the target ubatch.
 
 Every device selected for the recurrent graph operations must support sparse snapshots. The server rejects an unsupported device layout at startup; it does not move the operations to another device or fall back to CPU.
 
-Plane 0 contains the newest trailing state. The following planes contain older trailing states, and plane `K - 1` preserves the state from before verification. If a verification batch has fewer than `K - 1` tokens, the gap between its trailing states and the reserved input plane remains untouched, including the recorded positions. Rollbacks inside the retained trailing range select a plane directly. A deeper partial acceptance restores the reserved pre-verification plane and replays the full verification batch while writing only the selected boundary.
+Plane `K - 1` preserves the state from before verification. The other `K - 1` planes are split into `L = --spec-draft-rs-planes-lead` leading planes and `K - 1 - L` trailing planes. Plane 0 contains the newest state and the following trailing planes contain older states. The leading planes contain the states after the first `L` tokens of the verification batch, which are the states needed when the first drafts are rejected. At least one trailing plane is required, because a fully accepted draft continues from the newest state. If a verification batch has at most `K - 1` tokens, every state is kept in the trailing planes and the gap up to the reserved input plane remains untouched, including the recorded positions.
+
+A rollback to a kept state selects its plane directly. Any other partial acceptance restores the reserved pre-verification plane and replays the full verification batch while writing only the selected boundary. The replay costs a second target decode, so the layout should match where the drafts are usually rejected. MTP usually accepts most of the draft and keeps the default `L = 0`. DFlash usually rejects early, so by default it keeps one trailing plane and uses all other planes for leading states. For example, with `--spec-draft-n-max 7 --spec-draft-rs-planes 4`, DFlash keeps the states after 0, 1 and 7 accepted drafts.
 
 A sparse post-decode failure invalidates every slot that participated in the batch because their state and position planes were committed together. Per-slot recovery would require a separate pre-decode copy of the same recurrent state.
 

ggml/include/ggml.h

@@ -2651,7 +2651,8 @@ extern "C" {
     // By default, slot 0 is the final state and up to min(n_tokens, K) trailing states are written most-recent first.
     // For K > 1, trailing-only mode accepts (trailing_snapshots, selected_token, reserve_input) == ([0, K], -1, false) and writes up to min(n_tokens, trailing_snapshots) states most-recent first.
     // Selected-token mode requires (trailing_snapshots, selected_token, reserve_input) == (0, [0, n_tokens), false) and writes only the state after selected_token to slot 0.
-    // Reserved-input mode requires (trailing_snapshots, selected_token, reserve_input) == (min(n_tokens, K - 1), -1, true) and writes the input state to slot K - 1 and trailing states most-recent first from slot 0.
+    // Reserved-input mode requires (trailing_snapshots + leading_snapshots, selected_token, reserve_input) == (min(n_tokens, K - 1), -1, true) and writes the input state to slot K - 1, trailing states most-recent first from slot 0, and the state after token t < leading_snapshots to slot K - 2 - t.
+    // Other modes require leading_snapshots == 0.
     // State slots not selected by these modes stay untouched.
     // K == 1 requires (trailing_snapshots, selected_token, reserve_input) == (1, -1, false) and writes the final state to slot 0.
     GGML_API struct ggml_tensor * ggml_gated_delta_net(
@@ -2668,7 +2669,8 @@ extern "C" {
             struct ggml_tensor * tensor,
             int32_t              trailing_snapshots,
             int32_t              selected_token,
-            bool                 reserve_input);
+            bool                 reserve_input,
+            int32_t              leading_snapshots);
 
     GGML_API bool ggml_gated_delta_net_has_default_snapshot_params(
             const struct ggml_tensor * tensor);