From 842e4c951af17ba0bf2785cc66594ea03d106332 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:10:48 -0700 Subject: [PATCH 1/3] memory : keep a pending rollback across seq_cp With n_rs_seq != 0, a partial seq_rm records the rollback in rs_idx and the next decode applies it. seq_cp shared the cell but not rs_idx, so a sequence forked right after a rollback read the state from before the rollback. Copy rs_idx from the source in llama_memory_recurrent::seq_cp. The DSV4 compressor state had the same problem: seq_cp copied the current state plane of the source and reset rs_idx of the destination. Copy the snapshot plane that the pending rollback points to into the current plane of the destination instead. Add a test case that forks a sequence with a pending rollback and compares its state with a sequence that rolled back directly. --- src/llama-kv-cache-dsv4.cpp | 18 ++++--- src/llama-kv-cache-dsv4.h | 2 +- src/llama-memory-recurrent.cpp | 4 ++ tests/test-recurrent-state-rollback.cpp | 72 ++++++++++++++++++++++++- 4 files changed, 87 insertions(+), 9 deletions(-) diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 948d08146fbf..31b25dc712a7 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/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(2u*(1 + 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 kv_stream; std::vector 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,9 +1022,10 @@ 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; @@ -1032,7 +1033,8 @@ void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ 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; diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index ce39867c0342..ee706d1dd86c 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/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; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7e..57f6b46e729e 100644 --- a/src/llama-memory-recurrent.cpp +++ b/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]; + } } } diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index c6f599e584ca..7d1a5355dd4a 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -9,7 +9,7 @@ #include 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 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; } From f8d62ec32ed796bc5f1bdb570b466818e8d3d7bb Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:10:48 -0700 Subject: [PATCH 2/3] common : default MTP draft depth to 2 on CUDA On CUDA the cost of a verify batch grows faster with the draft depth than on Metal. A depth of 3 drafts tokens that do not pay back their verify cost, and 1 or 2 is faster. On Metal, 3 is still the fastest of 1, 2 and 3. When draft-mtp is enabled, the main model offloads to a CUDA device and --spec-draft-n-max is not given, use a depth of 2. An explicit --spec-draft-n-max always wins. --- common/arg.cpp | 5 +++-- common/common.cpp | 28 ++++++++++++++++++++++++++++ common/common.h | 4 ++++ 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 94dead3ff3ba..f95e32ecf649 100644 --- a/common/arg.cpp +++ b/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( diff --git a/common/common.cpp b/common/common.cpp index 4045903ae1a3..a24c68a1cd95 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1431,7 +1431,35 @@ std::vector & 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 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(); diff --git a/common/common.h b/common/common.h index 0ab2a3ffb3d1..5ad033298f92 100644 --- a/common/common.h +++ b/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 From e5f7a1e844699fe187255e8bbb9d17d0c7104a5e Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:45:33 -0700 Subject: [PATCH 3/3] dsv4 : compute the compressor state context size in size_t --- src/llama-kv-cache-dsv4.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 31b25dc712a7..c6538120f79c 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/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*(1 + n_rs_seq))*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, };