memory : keep a pending rollback across seq_cp, and default MTP draft depth to 2 on CUDA

common/arg.cpp

@@ -4103,12 +4103,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
 
     add_opt(common_arg(
         {"--spec-draft-n-max"}, "N",
-        string_format("number of tokens to draft for speculative decoding (default: %d)", params.speculative.draft.n_max),
+        string_format("number of tokens to draft for speculative decoding (default: %d, %d for draft-mtp on CUDA)", params.speculative.draft.n_max, COMMON_SPEC_MTP_N_MAX_CUDA),
         [](common_params & params, int value) {
             if (value < 0) {
                 throw std::invalid_argument("invalid value");
             }
-            params.speculative.draft.n_max = value;
+            params.speculative.draft.n_max     = value;
+            params.speculative.draft.n_max_set = true;
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX"));
     add_opt(common_arg(

common/common.cpp

@@ -1431,7 +1431,35 @@ std::vector<llama_adapter_lora_ptr> & common_init_result::lora() {
     return pimpl->lora;
 }
 
+// true if the main model offloads to a CUDA device
+static bool common_params_offload_cuda(const common_params & params) {
+    if (params.n_gpu_layers == 0) {
+        return false;
+    }
+
+    // no --device list: the model goes to the first GPU
+    std::vector<ggml_backend_dev_t> devs = params.devices;
+    if (devs.empty()) {
+        devs.push_back(ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU));
+    }
+
+    for (auto * dev : devs) {
+        if (dev != nullptr && strcmp(ggml_backend_reg_name(ggml_backend_dev_backend_reg(dev)), "CUDA") == 0) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 common_init_result_ptr common_init_from_params(common_params & params, bool model_only) {
+    auto & spec = params.speculative;
+    if (!spec.draft.n_max_set && common_params_offload_cuda(params) &&
+        std::find(spec.types.begin(), spec.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != spec.types.end()) {
+        spec.draft.n_max = COMMON_SPEC_MTP_N_MAX_CUDA;
+        COM_INF("using MTP draft depth %d on CUDA (set --spec-draft-n-max to override)\n", spec.draft.n_max);
+    }
+
     common_init_result_ptr res(new common_init_result(params, model_only));
 
     llama_model * model = res->model();

common/common.h

@@ -321,9 +321,13 @@ struct common_params_model {
     }
 };
 
+// default MTP draft depth when the target runs on CUDA: there the verify cost grows faster with depth than on Metal
+#define COMMON_SPEC_MTP_N_MAX_CUDA 2
+
 // draft-model-based speculative decoding parameters
 struct common_params_speculative_draft {
     int32_t n_max = 3; // maximum number of tokens to draft during speculative decoding
+    bool    n_max_set = false; // set by --spec-draft-n-max; if false, MTP on CUDA uses COMMON_SPEC_MTP_N_MAX_CUDA
     int32_t n_min = 0; // minimum number of draft tokens to use for speculative decoding
 
     float p_split = 0.1f; // speculative decoding split probability

src/llama-kv-cache-dsv4.cpp

@@ -920,7 +920,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
         auto it = ctx_map.find(buft);
         if (it == ctx_map.end()) {
             ggml_init_params params = {
-                /*.mem_size   =*/ size_t(2u*(1 + n_stream)*hparams.n_layer()*ggml_tensor_overhead()),
+                /*.mem_size   =*/ size_t(2)*(1 + size_t(n_stream)*(1 + n_rs_seq))*hparams.n_layer()*ggml_tensor_overhead(),
                 /*.mem_buffer =*/ NULL,
                 /*.no_alloc   =*/ true,
             };
@@ -971,7 +971,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
         std::vector<ggml_tensor *> kv_stream;
         std::vector<ggml_tensor *> score_stream;
 
-        for (uint32_t s = 0; s < n_stream; ++s) {
+        for (uint32_t s = 0; s < n_planes; ++s) {
             kv_stream.push_back(ggml_view_2d(ctx, kv, n_embd_state, state_size, kv->nb[1], s*kv->nb[2]));
             score_stream.push_back(ggml_view_2d(ctx, score, n_embd_state, state_size, score->nb[1], s*score->nb[2]));
         }
@@ -1022,17 +1022,19 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
     }
 }
 
-void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) {
+void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, uint32_t rs_src) {
     GGML_ASSERT(seq_id_src >= 0 && (uint32_t) seq_id_src < n_stream);
     GGML_ASSERT(seq_id_dst >= 0 && (uint32_t) seq_id_dst < n_stream);
+    GGML_ASSERT(rs_src <= n_rs_seq);
 
     if (seq_id_src == seq_id_dst) {
         return;
     }
 
     clear(seq_id_dst, true);
 
-    sc_info.ssrc.push_back((uint32_t) seq_id_src);
+    // copy the snapshot plane of a pending rollback, so the destination starts from the rolled-back state
+    sc_info.ssrc.push_back(rs_src*n_stream + (uint32_t) seq_id_src);
     sc_info.sdst.push_back((uint32_t) seq_id_dst);
 }
 
@@ -1517,9 +1519,11 @@ void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ds
     kv_hca->seq_cp(seq_id_src, seq_id_dst, -1, -1);
     kv_lid->seq_cp(seq_id_src, seq_id_dst, -1, -1);
 
-    csa_state->seq_cp(seq_id_src, seq_id_dst);
-    hca_state->seq_cp(seq_id_src, seq_id_dst);
-    lid_state->seq_cp(seq_id_src, seq_id_dst);
+    const uint32_t rs_src = seq_id_src != seq_id_dst ? rs_idx[seq_id_src] : 0;
+
+    csa_state->seq_cp(seq_id_src, seq_id_dst, rs_src);
+    hca_state->seq_cp(seq_id_src, seq_id_dst, rs_src);
+    lid_state->seq_cp(seq_id_src, seq_id_dst, rs_src);
 
     if (seq_id_src != seq_id_dst) {
         rs_idx[seq_id_dst] = 0;

src/llama-kv-cache-dsv4.h

@@ -27,7 +27,7 @@ class llama_dsv4_comp_state {
         const llama_memory_i::layer_filter_cb & filter);
 
     void clear(llama_seq_id seq_id, bool data);
-    void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst);
+    void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, uint32_t rs_src);
     void apply_copies(const stream_copy_info & sc_info) const;
 
     uint32_t get_ratio()      const;

src/llama-memory-recurrent.cpp

@@ -269,6 +269,10 @@ void llama_memory_recurrent::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id
             cell_src.seq_id.insert(seq_id_dst);
             tail_dst.tail = tail_src.tail;
         }
+        // a pending partial rollback is applied at the next decode, so the copy must use the same snapshot row
+        if (n_rs_seq != 0 && (uint32_t) seq_id_src < rs_idx.size() && (uint32_t) seq_id_dst < rs_idx.size()) {
+            rs_idx[seq_id_dst] = rs_idx[seq_id_src];
+        }
     }
 }
 

tests/test-recurrent-state-rollback.cpp

@@ -9,7 +9,7 @@
 #include <vector>
 
 static llama_context * make_ctx(const common_params & params, llama_model * model) {
-    auto cparams = common_context_params_to_llama(params);
+    auto cparams       = common_context_params_to_llama(params);
     cparams.n_seq_max = 1;
     cparams.n_rs_seq  = 8;
     cparams.n_batch   = std::max(cparams.n_batch,  (uint32_t) (cparams.n_rs_seq + 1));
@@ -207,6 +207,72 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode
     return true;
 }
 
+// Fork a sequence with seq_cp while its partial rollback is pending, then decode one token on the fork.
+// The reference rolls back seq 1 directly. The fork must match it bitwise.
+// The tiny test weights make the recurrent part too small to show in the logits, so compare the state bytes.
+static bool test_seq_cp_pending_rollback(const common_params & params, llama_model * model, const int n_vocab) {
+    constexpr uint32_t  n_prompt   = 12;
+    constexpr uint32_t  n_rollback = 3;
+    constexpr llama_pos p0         = n_prompt - n_rollback;
+
+    auto cparams       = common_context_params_to_llama(params);
+    cparams.n_seq_max  = 2;
+    cparams.n_rs_seq   = 8;
+    cparams.n_ctx      = 256;
+    cparams.n_batch    = 256;
+    cparams.n_ubatch   = 256;
+    cparams.kv_unified = false;
+
+    // fork: seq 0 -> rollback -> seq_cp to seq 1, ref: seq 1 -> rollback
+    std::vector<uint8_t> state[2];
+    for (int fork : { 1, 0 }) {
+        llama_context * ctx = llama_init_from_model(model, cparams);
+        if (ctx == nullptr) {
+            fprintf(stderr, "%s : failed to init context\n", __func__);
+            return false;
+        }
+
+        llama_memory_t     mem = llama_get_memory(ctx);
+        const llama_seq_id src = fork ? 0 : 1;
+
+        llama_batch batch = llama_batch_init(n_prompt, 0, 1);
+        for (llama_pos pos = 0; pos < (llama_pos) n_prompt; ++pos) {
+            common_batch_add(batch, (llama_token) ((5 * pos + 3) % n_vocab), pos, { src }, false);
+        }
+        bool ok = llama_decode(ctx, batch) == 0;
+
+        ok = ok && llama_memory_seq_rm(mem, src, p0, -1);
+        if (fork) {
+            llama_memory_seq_cp(mem, 0, 1, -1, -1);
+        }
+
+        common_batch_clear(batch);
+        common_batch_add(batch, (llama_token) (p0 % n_vocab), p0, { 1 }, true);
+        ok = ok && llama_decode(ctx, batch) == 0;
+        llama_batch_free(batch);
+
+        if (ok) {
+            state[fork].resize(llama_state_seq_get_size_ext(ctx, 1, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
+            llama_state_seq_get_data_ext(ctx, state[fork].data(), state[fork].size(), 1, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+        }
+        llama_free(ctx);
+
+        if (!ok) {
+            fprintf(stderr, "%s : decode/rollback failed (fork = %d)\n", __func__, fork);
+            return false;
+        }
+    }
+
+    if (state[0].empty() || state[0] != state[1]) {
+        fprintf(stderr, "%s : forked seq did not continue from the rolled-back state (%zu vs %zu bytes)\n",
+                __func__, state[1].size(), state[0].size());
+        return false;
+    }
+
+    fprintf(stderr, "%s : forked seq matched the rolled-back state (%zu bytes)\n", __func__, state[0].size());
+    return true;
+}
+
 int main(int argc, char ** argv) {
     std::setlocale(LC_NUMERIC, "C");
 
@@ -397,5 +463,9 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    if (!test_seq_cp_pending_rollback(params, model, n_vocab)) {
+        return 1;
+    }
+
     return 0;
 }