LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 14:24:03 +04:00
Code Issues Packages Projects Releases Wiki Activity

DeepSeek V4 spec checkpoints (#2205)

Browse Source 
* add DSV4 speculative checkpoints

* Fix DSV4 checkpoint cleanup indentation
This commit is contained in:
Samuel Oliveira Alves
2026-07-30 13:34:33 +03:00
committed by GitHub
parent 74cccfd71d
commit dd837ff21a
13 changed files with 947 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+14 -10
View File
@@ -1553,21 +1553,24 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
throw common_speculative_legacy_option_error(arg,
"the value inside the relevant repeated --spec-type entry using the canonical key p_min, e.g. --spec-type mtp:p_min=" + std::string(argv[i]));
}
if (arg == "--recurrent-ckpt-mode") {
if (arg == "--spec-ckpt-mode" || arg == "--recurrent-ckpt-mode") {
CHECK_ARG
const std::string val = argv[i];
if (val == "auto" || val == "AUTO") {
params.speculative.recurrent_ckpt_mode = LLAMA_SPEC_CKPT_AUTO;
params.speculative.spec_ckpt_mode = LLAMA_SPEC_CKPT_AUTO;
} else if (val == "per-step" || val == "PER_STEP") {
params.speculative.recurrent_ckpt_mode = LLAMA_SPEC_CKPT_PER_STEP;
params.speculative.spec_ckpt_mode = LLAMA_SPEC_CKPT_PER_STEP;
} else if (val == "gpu-fallback" || val == "GPU_FALLBACK") {
params.speculative.recurrent_ckpt_mode = LLAMA_SPEC_CKPT_GPU_FALLBACK;
params.speculative.spec_ckpt_mode = LLAMA_SPEC_CKPT_GPU_FALLBACK;
} else if (val == "cpu" || val == "CPU") {
params.speculative.recurrent_ckpt_mode = LLAMA_SPEC_CKPT_CPU;
params.speculative.spec_ckpt_mode = LLAMA_SPEC_CKPT_CPU;
} else {
throw std::invalid_argument("unknown --recurrent-ckpt-mode value: " + val +
throw std::invalid_argument("unknown " + std::string(arg) + " value: " + val +
"; expected auto, per-step, gpu-fallback, or cpu");
}
if (arg == "--recurrent-ckpt-mode") {
fprintf(stderr, "warning: --recurrent-ckpt-mode is deprecated; use --spec-ckpt-mode\n");
}
return true;
}
if (arg == "--spec-autotune") {
@@ -3337,11 +3340,12 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
options.push_back({ "*", "-hfr, --hf-repo REPO", "Hugging Face model repository (default: unused)" });
options.push_back({ "*", "-hff, --hf-file FILE", "Hugging Face model file (default: unused)" });
options.push_back({ "*", "-hft, --hf-token TOKEN", "Hugging Face access token (default: value from HF_TOKEN environment variable)" });
options.push_back({ "*", "--recurrent-ckpt-mode MODE", "checkpoint strategy for recurrent/hybrid speculative decoding\n"
options.push_back({ "*", "--spec-ckpt-mode MODE", "checkpoint strategy for speculative decoding\n"
" auto auto-select: per-step if CUDA full-GPU, gpu-fallback otherwise (default)\n"
" per-step save SSM state per draft step in VRAM; no re-decode on rejection\n"
" gpu-fallback copy state to GPU buffer; re-decode on rejection\n"
" cpu serialise state via llama_state_seq; re-decode on rejection" });
" per-step save architecture state per draft step; no re-decode on rejection\n"
" gpu-fallback copy architecture state to a device buffer; re-decode on rejection\n"
" cpu serialise architecture state via host storage; re-decode on rejection\n"
" --recurrent-ckpt-mode remains as a deprecated alias" });
options.push_back({ "*", "--spec-type SPEC[:k=v,...]", "canonical speculative stage entry; repeat for a supported two-stage chain.\n"
"types: none, draft, dflash, mtp, ngram-cache, ngram-simple, ngram-map-k, ngram-map-k4v, ngram-mod, suffix\n"
"canonical keys: n_max,n_min,p_min,heads,cross_ctx,ngram_size_n,ngram_size_m,ngram_min_hits,suffix_min_match_len,suffix_max_depth,suffix_corpus\n"
common/common.h
+2 -2
View File
@@ -209,8 +209,8 @@ struct common_ngram_mod;
struct common_params_speculative {
common_speculative_type type = COMMON_SPECULATIVE_TYPE_NONE; // type of speculative decoding
// Recurrent-model checkpoint strategy for speculative decoding.
int recurrent_ckpt_mode = LLAMA_SPEC_CKPT_AUTO;
// Generic speculative checkpoint mode, recurrent spelling is a compatibility alias.
int spec_ckpt_mode = LLAMA_SPEC_CKPT_AUTO;
std::string devices;
std::string params;
common/speculative.cpp
+40 -24
View File
@@ -17,7 +17,6 @@
#include <iomanip>
#include <limits>
#include <map>
#include <sstream>
#include <unordered_map>
#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128
@@ -55,9 +54,16 @@ const std::map<std::string, enum common_speculative_type> common_speculative_typ
{"suffix", COMMON_SPECULATIVE_TYPE_SUFFIX}
};
bool common_speculative_needs_checkpoint(const llama_model * model) {
return model != nullptr &&
(llama_model_has_recurrent(model) ||
llama_model_is_openpangu(model) ||
llama_model_is_deepseek4(model));
}
void common_speculative_checkpoint::clear() {
valid = false;
per_step_enabled = false;
mode = LLAMA_SPEC_CKPT_NONE;
n_past = 0;
sampled = LLAMA_TOKEN_NULL;
@@ -1338,16 +1344,16 @@ common_speculative * common_speculative_init(
configs.push_back(common_speculative_config(stage, stage_params));
}
if (!configs.empty() && (llama_model_has_recurrent(llama_get_model(ctx_tgt)) ||
llama_model_is_openpangu(llama_get_model(ctx_tgt)))) {
const llama_model * target_model = llama_get_model(ctx_tgt);
if (!configs.empty() && common_speculative_needs_checkpoint(target_model)) {
const int ckpt_tokens = std::max(1, params.get_max_stage_n_max() + 1);
const int actual_mode = llama_spec_ckpt_init(ctx_tgt, params.recurrent_ckpt_mode, ckpt_tokens);
const int actual_mode = llama_spec_ckpt_init(ctx_tgt, params.spec_ckpt_mode, ckpt_tokens);
if (actual_mode == LLAMA_SPEC_CKPT_NONE) {
LOG_ERR("%s: failed to prepare recurrent checkpoint mode '%s' during speculative init (max_tokens=%d)\n",
LOG_ERR("%s: failed to prepare speculative checkpoint mode '%s' during speculative init (max_tokens=%d)\n",
__func__,
params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_PER_STEP ? "per-step" :
params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_GPU_FALLBACK ? "gpu-fallback" :
params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_CPU ? "cpu" : "auto",
params.spec_ckpt_mode == LLAMA_SPEC_CKPT_PER_STEP ? "per-step" :
params.spec_ckpt_mode == LLAMA_SPEC_CKPT_GPU_FALLBACK ? "gpu-fallback" :
params.spec_ckpt_mode == LLAMA_SPEC_CKPT_CPU ? "cpu" : "auto",
ckpt_tokens);
if (ctx_dft != nullptr) {
llama_free(ctx_dft);
@@ -1355,7 +1361,7 @@ common_speculative * common_speculative_init(
return nullptr;
}
llama_spec_ckpt_discard(ctx_tgt);
params.recurrent_ckpt_mode = actual_mode;
params.spec_ckpt_mode = actual_mode;
}
std::vector<std::unique_ptr<common_speculative_state>> impls = {};
@@ -2267,7 +2273,7 @@ static bool common_speculative_checkpoint_save(
if (actual_mode == LLAMA_SPEC_CKPT_NONE) {
return false;
}
ckpt.per_step_enabled = (actual_mode == LLAMA_SPEC_CKPT_PER_STEP);
ckpt.mode = actual_mode;
ckpt.valid = llama_spec_ckpt_save(ctx, seq_id);
if (!ckpt.valid) {
@@ -2299,7 +2305,7 @@ void common_speculative_checkpoint_discard(
llama_spec_ckpt_discard(ctx);
}
void common_speculative_checkpoint_restore(
bool common_speculative_checkpoint_restore(
common_speculative_checkpoint & ckpt,
common_speculative * spec,
llama_context * ctx,
@@ -2312,13 +2318,19 @@ void common_speculative_checkpoint_restore(
const std::vector<float> & mtp_hidden_state_pre,
int32_t mtp_n_past_base) {
if (!ckpt.valid) {
return;
return true;
}
if (ckpt.per_step_enabled) {
const int step = (int) ids.size() - 1;
llama_spec_ckpt_restore(ctx, seq_id, ckpt.n_past, step);
const int step = (int) ids.size() - 1;
const enum llama_spec_ckpt_restore_result restore_result = llama_spec_ckpt_restore_ex(
ctx, seq_id, ckpt.n_past, ckpt.mode == LLAMA_SPEC_CKPT_PER_STEP ? step : 0);
if (restore_result == LLAMA_SPEC_CKPT_RESTORE_FAILED) {
LOG_ERR("%s: seq_id=%d speculative checkpoint restore failed\n", __func__, (int) seq_id);
common_speculative_checkpoint_discard(ckpt, ctx);
return false;
}
if (restore_result == LLAMA_SPEC_CKPT_RESTORE_DIRECT) {
if (ckpt.sampler != nullptr && sampler_dst != nullptr) {
common_sampler_clone(ckpt.sampler, sampler_dst);
}
@@ -2347,8 +2359,6 @@ void common_speculative_checkpoint_restore(
LOG_DBG("%s: seq_id=%d per-step restore: step=%d (rejected %d drafts)\n",
__func__, (int) seq_id, step, (int) (n_draft - (ids.size() - 1)));
} else {
llama_spec_ckpt_restore(ctx, seq_id, ckpt.n_past, 0);
if (ckpt.sampler != nullptr && sampler_dst != nullptr) {
common_sampler_clone(ckpt.sampler, sampler_dst);
}
@@ -2372,6 +2382,9 @@ void common_speculative_checkpoint_restore(
if (ret != 0) {
LOG_ERR("%s: seq_id=%d failed to re-decode accepted tokens after checkpoint restore: %d\n",
__func__, (int) seq_id, ret);
llama_batch_free(re_batch);
common_speculative_checkpoint_discard(ckpt, ctx);
return false;
}
if (common_speculative_has_target_features(spec)) {
@@ -2392,6 +2405,7 @@ void common_speculative_checkpoint_restore(
common_speculative_clear_sequence_hidden(spec, seq_id);
}
}
llama_batch_free(re_batch);
if (sampler_dst != nullptr) {
for (llama_token id : ids) {
@@ -2399,16 +2413,16 @@ void common_speculative_checkpoint_restore(
}
}
llama_batch_free(re_batch);
LOG_DBG("%s: seq_id=%d spec checkpoint restored: re-decoded %d tokens (rejected %d drafts)\n",
__func__, (int) seq_id, n_re, (int) (n_draft - (ids.size() - 1)));
}
}
common_speculative_checkpoint_discard(ckpt, ctx);
return true;
}
void common_speculative_commit(
bool common_speculative_commit(
common_speculative * spec,
llama_context * ctx,
common_sampler * sampler_dst,
@@ -2426,10 +2440,11 @@ void common_speculative_commit(
? spec->curr_impl->type
: COMMON_SPECULATIVE_TYPE_NONE;
const bool any_rejected = (int) ids.size() - 1 < n_draft;
const int n_accepted = (int) ids.size() - 1;
const bool any_rejected = n_accepted < n_draft;
std::vector<float> mtp_hidden_state_pre;
common_speculative_accept(spec, ids.size() - 1);
common_speculative_accept(spec, n_accepted);
if (common_speculative_has_target_features(spec) &&
any_rejected &&
@@ -2441,7 +2456,7 @@ void common_speculative_commit(
}
if (any_rejected && ckpt.valid) {
common_speculative_checkpoint_restore(
const bool restored = common_speculative_checkpoint_restore(
ckpt,
spec,
ctx,
@@ -2453,7 +2468,7 @@ void common_speculative_commit(
n_draft,
mtp_hidden_state_pre,
pos_base);
return;
return restored;
}
if (common_speculative_has_target_features(spec) && !accepted_output_indices.empty()) {
@@ -2475,6 +2490,7 @@ void common_speculative_commit(
llama_kv_cache_seq_rm(ctx, seq_id, pos_base + (llama_pos) (ids.size() - 1), -1);
common_speculative_checkpoint_discard(ckpt, ctx);
return true;
}
void common_speculative_print_stats(const common_speculative * spec, double slot_tps, int n_decoded, int n_past, common_params_speculative * active_params) {
common/speculative.h
+14 -12
View File
@@ -7,6 +7,8 @@
struct common_speculative;
bool common_speculative_needs_checkpoint(const llama_model * model);
enum common_speculative_init_status {
COMMON_SPECULATIVE_INIT_SKIPPED,
COMMON_SPECULATIVE_INIT_READY,
@@ -24,7 +26,7 @@ static constexpr common_speculative_feature_kind COMMON_SPECULATIVE_FEATURE_HIDD
struct common_speculative_checkpoint {
bool valid = false;
bool per_step_enabled = false;
int mode = LLAMA_SPEC_CKPT_NONE;
llama_pos n_past = 0;
llama_token sampled = LLAMA_TOKEN_NULL;
common_sampler * sampler = nullptr;
@@ -170,7 +172,7 @@ void common_speculative_checkpoint_discard(
common_speculative_checkpoint & ckpt,
llama_context * ctx);
void common_speculative_checkpoint_restore(
bool common_speculative_checkpoint_restore(
common_speculative_checkpoint & ckpt,
common_speculative * spec,
llama_context * ctx,
@@ -183,16 +185,16 @@ void common_speculative_checkpoint_restore(
const std::vector<float> & mtp_hidden_state_pre,
int32_t mtp_n_past_base);
void common_speculative_commit(
common_speculative * spec,
llama_context * ctx,
common_sampler * sampler_dst,
llama_seq_id seq_id,
llama_token sampled_before,
const std::vector<llama_token> & ids,
int n_draft,
llama_pos pos_base,
const std::vector<int32_t> & accepted_output_indices);
bool common_speculative_commit(
common_speculative * spec,
llama_context * ctx,
common_sampler * sampler_dst,
llama_seq_id seq_id,
llama_token sampled_before,
const std::vector<llama_token> & ids,
int n_draft,
llama_pos pos_base,
const std::vector<int32_t> & accepted_output_indices);
bool common_speculative_has_sequence_hidden(const common_speculative * spec, llama_seq_id seq_id);
docs/parameters.md
+1 -1
View File
@@ -134,7 +134,7 @@ Check the details [here](./speculative.md).
| `-draft, --draft-params` | Comma-separated list of draft model parameters | - | |
| `--spec-type SPEC[:k=v,...]` | Canonical speculative stage entry; repeat to configure the supported two-stage chain | - | Types: `none`, `draft`, `dflash`, `mtp`, `ngram-cache`, `ngram-simple`, `ngram-map-k`, `ngram-map-k4v`, `ngram-mod`, `suffix`. Canonical keys include `n_max`, `n_min`, `p_min`, `heads`, `cross_ctx`, `ngram_size_n`, `ngram_size_m`, `ngram_min_hits`, `suffix_min_match_len`, `suffix_max_depth`, `suffix_corpus`. For MTP, `heads=1` is the default; values above `1` and `heads=0` (all model heads) are experimental. String values may escape commas as `\,` or quote the value inside the stage payload. Examples: `--spec-type ngram-mod:n_max=64,n_min=2,ngram_size_n=8 --spec-type mtp:n_max=1,p_min=0.0`, `--model-draft draft.gguf --spec-type dflash:n_max=4,cross_ctx=512` |
| `--spec-autotune` | Automatically tune speculative params to maximize tokens/sec | - | Automatically determines the near-optimal arguments for the type of speculation being performed [PR 1595](https://github.com/ikawrakow/ik_llama.cpp/pull/1595) |
| `--recurrent-ckpt-mode MODE` | Checkpoint strategy for recurrent/hybrid speculative decoding | auto | One of: - `auto` auto-select: per-step if CUDA full-GPU, gpu-fallback otherwise - `per-step` save SSM state per draft step in VRAM; no re-decode on rejection - `gpu-fallback` copy state to GPU buffer; re-decode on rejection - `cpu` serialise state via llama_state_seq; re-decode on rejection [PR 1669](https://github.com/ikawrakow/ik_llama.cpp/pull/1669) [PR 1774](https://github.com/ikawrakow/ik_llama.cpp/pull/1774) |
| `--spec-ckpt-mode MODE` (deprecated alias: `--recurrent-ckpt-mode MODE`) | Checkpoint strategy for speculative decoding | auto | One of: - `auto` choose direct per-step, then device fallback, then host fallback - `per-step` save architecture state per draft step; no re-decode on rejection - `gpu-fallback` copy architecture state to a device buffer; re-decode on rejection - `cpu` serialize architecture state in host storage; re-decode on rejection [PR 1669](https://github.com/ikawrakow/ik_llama.cpp/pull/1669) [PR 1774](https://github.com/ikawrakow/ik_llama.cpp/pull/1774) |
Notes:
examples/main/main.cpp
+8 -4
View File
@@ -984,7 +984,7 @@ int main(int argc, char ** argv) {
const int min_usable_draft = params.speculative.get_min_usable_stage_n_min();
if ((int) draft.size() >= min_usable_draft && (!draft.empty() || n_predict_budget > 1)) {
if (llama_model_has_recurrent(model) || llama_model_is_openpangu(model)) {
if (common_speculative_needs_checkpoint(model)) {
if (!common_speculative_before_draft(
spec,
model,
@@ -995,7 +995,7 @@ int main(int argc, char ** argv) {
n_past,
sampled_before,
(int) draft.size() + 1,
params.speculative.recurrent_ckpt_mode)) {
params.speculative.spec_ckpt_mode)) {
LOG_TEE("%s: speculative checkpoint setup failed, falling back to one-token decode\n", __func__);
draft.clear();
}
@@ -1033,7 +1033,7 @@ int main(int argc, char ** argv) {
accepted_output_indices.assign(verify_indices.begin(), verify_indices.begin() + ids.size());
}
common_speculative_commit(
if (!common_speculative_commit(
spec,
ctx,
ctx_sampling,
@@ -1042,7 +1042,11 @@ int main(int argc, char ** argv) {
ids,
(int) draft.size(),
n_past + 1,
accepted_output_indices);
accepted_output_indices)) {
llama_batch_free(verify_batch);
LOG_TEE("%s: speculative checkpoint restore/commit failed\n", __func__);
return 1;
}
llama_batch_free(verify_batch);
examples/server/server-context.cpp
+66 -7
View File
@@ -281,7 +281,7 @@ bool server_context::load_model(const gpt_params& params_) {
void server_context::init() {
const int32_t n_ctx_slot = n_ctx / params_base.n_parallel;
if (!system_prompt.empty() && std::strcmp(llama_model_arch_string(model), "deepseek4") == 0) {
if (!system_prompt.empty() && llama_model_is_deepseek4(model)) {
throw std::runtime_error("DeepSeek4 server system prompts are unsupported because seq_cp does not copy private cache state");
}
@@ -2081,7 +2081,7 @@ void server_context::system_prompt_update() {
}
bool server_context::system_prompt_set(const std::string& sys_prompt) {
if (!sys_prompt.empty() && model != nullptr && std::strcmp(llama_model_arch_string(model), "deepseek4") == 0) {
if (!sys_prompt.empty() && llama_model_is_deepseek4(model)) {
LOG_ERROR("DeepSeek4 server system prompts are unsupported because seq_cp does not copy private cache state", {});
return false;
}
@@ -4238,7 +4238,7 @@ void server_context::speculative_decoding_accept() {
slot.sampled = ids.back(); // last accepted token
slot.n_past = slot.cache_tokens.n_tokens();
common_speculative_commit(
if (!common_speculative_commit(
slot.spec,
ctx,
slot.ctx_sampling,
@@ -4247,7 +4247,18 @@ void server_context::speculative_decoding_accept() {
ids,
n_draft,
spec_pos_base,
accepted_output_indices);
accepted_output_indices)) {
LOG_ERROR("speculative checkpoint restore/commit failed, releasing slot", {
{"id_slot", slot.id},
{"id_task", slot.id_task},
});
send_error(slot, "speculative checkpoint restore failed", ERROR_TYPE_SERVER);
slot.release();
slot.i_batch = -1;
slot.i_batch_dft.clear();
slot.drafted.clear();
continue;
}
slot.spec_target_only = false;
for (size_t i = 0; i < ids.size(); ++i) {
@@ -4830,8 +4841,54 @@ void server_context::update_slots() {
// make sure we're in the right embedding mode
llama_set_embeddings(ctx, batch_type == 1);
if (llama_model_has_recurrent(model) || llama_model_is_openpangu(model)) {
const int ckpt_mode = params_base.speculative.recurrent_ckpt_mode;
if (common_speculative_needs_checkpoint(model)) {
const int ckpt_mode = params_base.speculative.spec_ckpt_mode;
// Remove draft rows if checkpoint setup fails, otherwise rejection is unsafe.
auto make_root_only = [&](server_slot & slot) {
if (slot.i_batch_dft.empty()) {
return;
}
const int32_t root_index = slot.i_batch_dft.front();
const int32_t old_n_tokens = batch.n_tokens;
std::vector<uint8_t> remove(old_n_tokens, 0);
for (size_t i = 1; i < slot.i_batch_dft.size(); ++i) {
const int32_t index = slot.i_batch_dft[i];
if (index >= 0 && index < old_n_tokens) {
remove[index] = 1;
}
}
std::vector<int32_t> remap(old_n_tokens, -1);
int32_t write = 0;
for (int32_t read = 0; read < old_n_tokens; ++read) {
if (remove[read]) {
continue;
}
if (write != read) {
batch.token[write] = batch.token[read];
batch.pos[write] = batch.pos[read];
batch.n_seq_id[write] = batch.n_seq_id[read];
for (size_t seq = 0; seq < batch.n_seq_id[read]; ++seq) {
batch.seq_id[write][seq] = batch.seq_id[read][seq];
}
batch.logits[write] = batch.logits[read];
}
remap[read] = write++;
}
batch.n_tokens = write;
if (root_index >= 0 && root_index < old_n_tokens) {
slot.i_batch = remap[root_index];
}
slot.cache_tokens.keep_first(slot.cache_tokens.n_tokens() - (int32_t) slot.drafted.size());
slot.drafted.clear();
slot.i_batch_dft.clear();
slot.n_past = slot.cache_tokens.n_tokens();
slot.spec_target_only = false;
SLT_WRN(slot, "%s", "spec checkpoint unavailable; removed draft rows and continuing root-only\n");
};
for (auto & slot : slots) {
if (slot.state != SLOT_STATE_PROCESSING || slot.i_batch_dft.empty()) {
@@ -4853,11 +4910,13 @@ void server_context::update_slots() {
ckpt_mode)) {
const common_speculative_checkpoint * ckpt = common_speculative_get_checkpoint(slot.spec);
GGML_ASSERT(ckpt != nullptr);
const char * mode_name = ckpt->per_step_enabled ? "per-step" : "shadow/cpu";
const char * mode_name = ckpt->mode == LLAMA_SPEC_CKPT_PER_STEP ? "per-step" :
ckpt->mode == LLAMA_SPEC_CKPT_GPU_FALLBACK ? "gpu-fallback" : "cpu";
SLT_DBG(slot, "spec checkpoint saved (mode=%s), n_past_pre_spec=%d\n",
mode_name, ckpt->n_past);
} else {
SLT_WRN(slot, "%s", "failed to save spec checkpoint\n");
make_root_only(slot);
}
}
}
include/llama.h
+13
View File
@@ -700,6 +700,9 @@ extern "C" {
LLAMA_API bool llama_model_has_recurrent(const struct llama_model * model);
// Returns whether the model uses the DeepSeek-V4 architecture.
LLAMA_API bool llama_model_is_deepseek4(const struct llama_model * model);
// Returns true if the model is openPangu (conv-only recurrent state that rides the spec-rollback checkpoint)
LLAMA_API bool llama_model_is_openpangu(const struct llama_model * model);
@@ -845,6 +848,12 @@ extern "C" {
LLAMA_SPEC_CKPT_CPU = 3,
};
enum llama_spec_ckpt_restore_result {
LLAMA_SPEC_CKPT_RESTORE_FAILED = 0,
LLAMA_SPEC_CKPT_RESTORE_DIRECT = 1,
LLAMA_SPEC_CKPT_RESTORE_BASE_REPLAY_REQUIRED = 2,
};
// Initialise the checkpoint system for the upcoming speculation window.
LLAMA_API int llama_spec_ckpt_init(struct llama_context * ctx, int mode, int max_tokens);
@@ -855,6 +864,10 @@ extern "C" {
LLAMA_API bool llama_spec_ckpt_restore(struct llama_context * ctx, llama_seq_id seq_id,
llama_pos n_past, int accepted_step);
LLAMA_API enum llama_spec_ckpt_restore_result llama_spec_ckpt_restore_ex(
struct llama_context * ctx, llama_seq_id seq_id,
llama_pos n_past, int accepted_step);
// Discard the saved checkpoint and reset internal mode state.
LLAMA_API void llama_spec_ckpt_discard(struct llama_context * ctx);
src/llama-context.h
+51 -1
View File
@@ -141,6 +141,24 @@ struct llama_kv_cache {
int64_t per_step_conv_dim = 0;
int32_t per_step_d_conv = 0;
// DSV4 per-step compressor-state base and per-row deltas.
std::vector<ggml_tensor *> dsv4_per_step_state;
std::vector<ggml_tensor *> dsv4_per_step_state_shadow;
std::vector<ggml_tensor *> dsv4_per_step_delta;
std::vector<ggml_context *> dsv4_per_step_shadow_ctxs;
std::vector<ggml_backend_buffer_t> dsv4_per_step_shadow_bufs;
std::vector<int32_t> dsv4_per_step_csa_dst;
std::vector<int32_t> dsv4_per_step_hca_dst;
std::vector<int32_t> dsv4_per_step_lid_dst;
std::vector<int32_t> dsv4_per_step_csa_src;
std::vector<int32_t> dsv4_per_step_hca_src;
std::vector<int32_t> dsv4_per_step_lid_src;
bool dsv4_per_step_allocated = false;
bool dsv4_per_step_saved = false;
int32_t dsv4_per_step_max_tokens = 0;
size_t dsv4_per_step_base_bytes = 0;
size_t dsv4_per_step_delta_bytes = 0;
int selected_spec_mode = -1;
int fixed_spec_mode = LLAMA_SPEC_CKPT_NONE;
int32_t fixed_max_tokens = 0;
@@ -148,6 +166,14 @@ struct llama_kv_cache {
// Serialised sequence state for CPU mode
std::vector<uint8_t> cpu_state_data;
// Private DSV4 state snapshot used by GPU/CPU fallback modes.
std::vector<std::vector<uint8_t>> dsv4_state_data;
std::vector<ggml_tensor *> dsv4_state_shadow;
std::vector<struct ggml_context *> dsv4_shadow_ctxs;
std::vector<ggml_backend_buffer_t> dsv4_shadow_bufs;
bool dsv4_shadow_allocated = false;
bool dsv4_shadow_saved = false;
// Separate storage for per-step allocations
std::vector<struct ggml_context *> per_step_ctxs;
std::vector<ggml_backend_buffer_t> per_step_bufs;
@@ -159,19 +185,41 @@ struct llama_kv_cache {
bool shadow_conv_only = false;
bool saved = false;
~gpu_checkpoint() {
void release_dsv4_per_step();
void release_dsv4_snapshot();
void release() {
release_dsv4_per_step();
release_dsv4_snapshot();
for (struct ggml_context * ctx : shadow_ctxs) {
ggml_free(ctx);
}
shadow_ctxs.clear();
for (ggml_backend_buffer_t buf : shadow_bufs) {
ggml_backend_buffer_free(buf);
}
shadow_bufs.clear();
s_l_shadow.clear();
split_s_l_shadow.clear();
allocated = false;
saved = false;
for (struct ggml_context * ctx : per_step_ctxs) {
ggml_free(ctx);
}
per_step_ctxs.clear();
for (ggml_backend_buffer_t buf : per_step_bufs) {
ggml_backend_buffer_free(buf);
}
per_step_bufs.clear();
per_step_ssm.clear();
per_step_conv.clear();
per_step_max_allocated = 0;
}
~gpu_checkpoint() {
release();
}
};
@@ -450,6 +498,8 @@ struct llama_context {
struct comp_plan {
std::vector<int32_t> state_pos;
std::vector<int32_t> state_delta_src_idxs;
std::vector<int32_t> state_delta_dst_idxs;
std::vector<int32_t> state_persist_src_idxs;
std::vector<int32_t> state_persist_dst_idxs;
std::vector<int32_t> state_read_idxs;
src/llama-dsv4.cpp
+648 -5
View File
@@ -1,10 +1,14 @@
#include "llama-dsv4.h"
#include <random>
#include "llama.h"
#include "llama-context.h"
#include "llama-model.h"
#include "llama-impl.h"
#include "ggml.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include <algorithm>
@@ -18,6 +22,10 @@ static bool dsv4_cache_type_supported(ggml_type type) {
return type == GGML_TYPE_F16 || type == GGML_TYPE_BF16 || type == GGML_TYPE_Q8_0;
}
// Per-step capture is limited to the eight-row CSA/LID ring.
// TODO: Expand to a larger number
static constexpr int DSV4_PER_STEP_MAX_STATE_ROWS = 8;
static bool dsv4_validate_cache_type(ggml_type type, int64_t width, const char * name) {
if (!dsv4_cache_type_supported(type)) {
LLAMA_LOG_ERROR("%s: unsupported DSV4 %s cache type %s\n", __func__, name, ggml_type_name(type));
@@ -488,6 +496,8 @@ static llama_context::dsv4_runtime::comp_plan dsv4_build_reserve_comp_plan(
const size_t n_persist = (size_t) std::min<uint64_t>((uint64_t) batch.n_tokens, state_rows);
plan.state_pos.resize((size_t) batch.n_tokens);
plan.state_delta_src_idxs.resize((size_t) batch.n_tokens);
plan.state_delta_dst_idxs.resize((size_t) batch.n_tokens);
plan.state_persist_src_idxs.resize(n_persist);
plan.state_persist_dst_idxs.resize(n_persist);
plan.state_read_idxs.resize((overlap ? 2u : 1u)*ratio*n_blocks);
@@ -540,6 +550,14 @@ static bool dsv4_validate_comp_plan(
return false;
}
if (plan.state_delta_src_idxs.size() != plan.state_pos.size() ||
plan.state_delta_dst_idxs.size() != plan.state_pos.size()) {
LLAMA_LOG_ERROR("%s: DSV4 %s delta row metadata mismatch: state=%zu src=%zu dst=%zu\n",
__func__, tag, plan.state_pos.size(), plan.state_delta_src_idxs.size(),
plan.state_delta_dst_idxs.size());
return false;
}
if (plan.state_persist_src_idxs.size() != plan.state_persist_dst_idxs.size()) {
LLAMA_LOG_ERROR("%s: DSV4 %s persist idx size mismatch: src=%zu dst=%zu\n",
__func__, tag, plan.state_persist_src_idxs.size(), plan.state_persist_dst_idxs.size());
@@ -568,6 +586,14 @@ static bool dsv4_validate_comp_plan(
__func__, tag, i, (long long) pos, ratio);
return false;
}
const int64_t src = plan.state_delta_src_idxs[i];
const int64_t dst = plan.state_delta_dst_idxs[i];
if (src < 0 || src >= batch.n_tokens || dst < 0 || (uint32_t) dst >= state_size*n_stream) {
LLAMA_LOG_ERROR("%s: DSV4 %s delta row[%zu] src=%lld dst=%lld is outside the batch/state ring\n",
__func__, tag, i, (long long) src, (long long) dst);
return false;
}
}
for (size_t i = 0; i < plan.state_persist_src_idxs.size(); ++i) {
@@ -642,11 +668,17 @@ static llama_context::dsv4_runtime::comp_plan dsv4_build_comp_plan(
std::map<std::pair<llama_seq_id, llama_pos>, int32_t> curr_token_idx_map;
for (int32_t i = 0; i < batch.n_tokens; ++i) {
const llama_seq_id seq_id =
batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.n_seq_id[i] > 0 && batch.seq_id[i] != nullptr
? batch.seq_id[i][0]
: 0;
curr_token_idx_map[std::make_pair(seq_id, batch.pos[i])] = i;
const int32_t n_token_seqs =
batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.seq_id[i] != nullptr
? batch.n_seq_id[i]
: 1;
for (int32_t s = 0; s < n_token_seqs; ++s) {
const llama_seq_id seq_id =
batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.seq_id[i] != nullptr
? batch.seq_id[i][s]
: 0;
curr_token_idx_map[std::make_pair(seq_id, batch.pos[i])] = i;
}
}
const auto state_source_idx = [&](llama_seq_id seq_id, llama_pos pos) -> int32_t {
@@ -672,6 +704,14 @@ static llama_context::dsv4_runtime::comp_plan dsv4_build_comp_plan(
plan.state_pos.push_back((int32_t) (pos%ratio));
const llama_seq_id delta_seq_id =
batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.seq_id[i] != nullptr && batch.n_seq_id[i] > 0
? batch.seq_id[i][0]
: 0;
plan.state_delta_src_idxs.push_back(i);
plan.state_delta_dst_idxs.push_back((int32_t) (
dsv4_stream_offset(n_stream, delta_seq_id, state_size) + pos%state_size));
const int64_t n_visible = (int64_t) (pos + 1)/ratio;
plan.n_visible[(size_t) i] = (int32_t) n_visible;
plan.n_kv = std::max(plan.n_kv, n_visible);
@@ -1013,6 +1053,608 @@ void llama_reset_dsv4_state(llama_context * ctx, int32_t seq_id) {
for (ggml_tensor * tensor : ctx->dsv4.cache.lid_state_score) clear_tensor(tensor);
}
static std::vector<ggml_tensor *> dsv4_state_tensors(const llama_context & ctx) {
std::vector<ggml_tensor *> tensors;
const auto append = [&tensors](const std::vector<ggml_tensor *> & group) {
for (ggml_tensor * tensor : group) {
if (tensor != nullptr) {
tensors.push_back(tensor);
}
}
};
append(ctx.dsv4.cache.csa_state_kv);
append(ctx.dsv4.cache.csa_state_score);
append(ctx.dsv4.cache.hca_state_kv);
append(ctx.dsv4.cache.hca_state_score);
append(ctx.dsv4.cache.lid_state_kv);
append(ctx.dsv4.cache.lid_state_score);
return tensors;
}
void llama_kv_cache::gpu_checkpoint::release_dsv4_per_step() {
for (ggml_context * shadow_ctx : dsv4_per_step_shadow_ctxs) {
ggml_free(shadow_ctx);
}
for (ggml_backend_buffer_t buffer : dsv4_per_step_shadow_bufs) {
ggml_backend_buffer_free(buffer);
}
dsv4_per_step_shadow_ctxs.clear();
dsv4_per_step_shadow_bufs.clear();
dsv4_per_step_state.clear();
dsv4_per_step_state_shadow.clear();
dsv4_per_step_delta.clear();
dsv4_per_step_csa_src.clear();
dsv4_per_step_csa_dst.clear();
dsv4_per_step_hca_src.clear();
dsv4_per_step_hca_dst.clear();
dsv4_per_step_lid_src.clear();
dsv4_per_step_lid_dst.clear();
dsv4_per_step_allocated = false;
dsv4_per_step_saved = false;
dsv4_per_step_max_tokens = 0;
dsv4_per_step_base_bytes = 0;
dsv4_per_step_delta_bytes = 0;
}
void llama_kv_cache::gpu_checkpoint::release_dsv4_snapshot() {
for (ggml_context * shadow_ctx : dsv4_shadow_ctxs) {
ggml_free(shadow_ctx);
}
for (ggml_backend_buffer_t buffer : dsv4_shadow_bufs) {
ggml_backend_buffer_free(buffer);
}
dsv4_shadow_ctxs.clear();
dsv4_shadow_bufs.clear();
dsv4_state_data.clear();
dsv4_state_shadow.clear();
dsv4_shadow_allocated = false;
dsv4_shadow_saved = false;
}
static bool dsv4_per_step_alloc(llama_context & ctx, int max_tokens) {
auto & ckpt = ctx.kv_self.ckpt;
const auto states = dsv4_state_tensors(ctx);
if (states.empty() || max_tokens <= 0 || max_tokens > DSV4_PER_STEP_MAX_STATE_ROWS) {
if (max_tokens > DSV4_PER_STEP_MAX_STATE_ROWS) {
LLAMA_LOG_WARN("%s: DSV4 per-step supports at most %d verification rows; requested %d\n",
__func__, DSV4_PER_STEP_MAX_STATE_ROWS, max_tokens);
}
return false;
}
if (ckpt.dsv4_per_step_allocated && ckpt.dsv4_per_step_max_tokens >= max_tokens &&
ckpt.dsv4_per_step_state.size() == states.size()) {
return true;
}
ctx.kv_self.ckpt.release_dsv4_per_step();
ckpt.dsv4_per_step_state = states;
ckpt.dsv4_per_step_state_shadow.assign(states.size(), nullptr);
ckpt.dsv4_per_step_delta.assign(states.size(), nullptr);
struct entry {
size_t index;
ggml_tensor * source;
};
std::map<ggml_backend_buffer_type_t, std::vector<entry>> entries_by_buft;
for (size_t i = 0; i < states.size(); ++i) {
ggml_tensor * source = states[i];
if (source == nullptr || source->buffer == nullptr) {
ctx.kv_self.ckpt.release_dsv4_per_step();
return false;
}
entries_by_buft[ggml_backend_buffer_get_type(source->buffer)].push_back({ i, source });
ckpt.dsv4_per_step_base_bytes += ggml_nbytes(source);
ckpt.dsv4_per_step_delta_bytes += ggml_row_size(source->type, source->ne[0]) * (size_t) max_tokens;
}
for (auto & [buft, entries] : entries_by_buft) {
ggml_init_params params = {
/*.mem_size =*/ entries.size() * 3 * ggml_tensor_overhead(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
ggml_context * graph_ctx = ggml_init(params);
if (graph_ctx == nullptr) {
ctx.kv_self.ckpt.release_dsv4_per_step();
return false;
}
for (const entry & item : entries) {
ggml_tensor * shadow = ggml_dup_tensor(graph_ctx, item.source);
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
shadow->nb[d] = item.source->nb[d];
}
ggml_format_name(shadow, "dsv4_per_step_base_%zu", item.index);
ggml_tensor * delta = ggml_new_tensor_2d(graph_ctx, item.source->type,
item.source->ne[0], max_tokens);
ggml_format_name(delta, "dsv4_per_step_delta_%zu", item.index);
ckpt.dsv4_per_step_state_shadow[item.index] = shadow;
ckpt.dsv4_per_step_delta[item.index] = delta;
}
ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors_from_buft(graph_ctx, buft);
if (buffer == nullptr) {
ggml_free(graph_ctx);
ctx.kv_self.ckpt.release_dsv4_per_step();
return false;
}
ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE);
ggml_backend_buffer_clear(buffer, 0);
ckpt.dsv4_per_step_shadow_ctxs.push_back(graph_ctx);
ckpt.dsv4_per_step_shadow_bufs.push_back(buffer);
}
ckpt.dsv4_per_step_max_tokens = max_tokens;
ckpt.dsv4_per_step_allocated = true;
LLAMA_LOG_INFO("%s: DSV4 per-step base=%8.2f MiB delta=%8.2f MiB max_tokens=%d\n",
__func__, ckpt.dsv4_per_step_base_bytes / (1024.0 * 1024.0),
ckpt.dsv4_per_step_delta_bytes / (1024.0 * 1024.0), max_tokens);
return true;
}
static bool dsv4_per_step_copy_base(llama_context & ctx, bool restore) {
auto & ckpt = ctx.kv_self.ckpt;
if (!ckpt.dsv4_per_step_allocated || ckpt.dsv4_per_step_state.size() != ckpt.dsv4_per_step_state_shadow.size()) {
return false;
}
std::vector<ggml_backend_t> backends;
for (size_t i = 0; i < ckpt.dsv4_per_step_state.size(); ++i) {
ggml_tensor * state = ckpt.dsv4_per_step_state[i];
ggml_tensor * shadow = ckpt.dsv4_per_step_state_shadow[i];
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(ctx.sched, state);
if (state == nullptr || shadow == nullptr || backend == nullptr) {
return false;
}
if (restore) {
ggml_backend_tensor_copy_async(backend, backend, shadow, state);
} else {
ggml_backend_tensor_copy_async(backend, backend, state, shadow);
}
if (std::find(backends.begin(), backends.end(), backend) == backends.end()) {
backends.push_back(backend);
}
}
for (ggml_backend_t backend : backends) {
ggml_backend_synchronize(backend);
}
return true;
}
static bool dsv4_per_step_capture_group(
llama_context & ctx,
const std::vector<ggml_tensor *> & states,
const llama_context::dsv4_runtime::comp_plan & plan) {
auto & ckpt = ctx.kv_self.ckpt;
if (plan.state_delta_src_idxs.size() != plan.state_delta_dst_idxs.size() ||
plan.state_delta_src_idxs.size() > (size_t) ckpt.dsv4_per_step_max_tokens) {
return false;
}
for (ggml_tensor * state : states) {
if (state == nullptr) {
continue;
}
ggml_tensor * delta = llama_dsv4_spec_ckpt_delta(&ctx, state);
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(ctx.sched, state);
if (delta == nullptr || backend == nullptr || delta->ne[0] != state->ne[0]) {
return false;
}
for (size_t row = 0; row < plan.state_delta_src_idxs.size(); ++row) {
const int32_t src_idx = plan.state_delta_src_idxs[row];
const int32_t dst_idx = plan.state_delta_dst_idxs[row];
if (src_idx < 0 || (uint64_t) src_idx >= (uint64_t) delta->ne[1] ||
dst_idx < 0 || (uint64_t) dst_idx >= (uint64_t) state->ne[1]) {
return false;
}
ggml_tensor src_view = *state;
ggml_tensor dst_view = *delta;
src_view.ne[1] = src_view.ne[2] = src_view.ne[3] = 1;
dst_view.ne[1] = dst_view.ne[2] = dst_view.ne[3] = 1;
src_view.nb[2] = src_view.nb[3] = src_view.nb[1];
dst_view.nb[2] = dst_view.nb[3] = dst_view.nb[1];
src_view.data = (char *) state->data + (size_t) dst_idx * state->nb[1];
dst_view.data = (char *) delta->data + (size_t) src_idx * delta->nb[1];
src_view.view_src = nullptr;
dst_view.view_src = nullptr;
src_view.view_offs = 0;
dst_view.view_offs = 0;
ggml_backend_tensor_copy_async(backend, backend, &src_view, &dst_view);
}
}
return true;
}
bool llama_dsv4_spec_ckpt_capture_rows(llama_context * ctx) {
if (ctx == nullptr || ctx->model.arch != LLM_ARCH_DEEPSEEK4) {
return true;
}
const auto & ckpt = ctx->kv_self.ckpt;
if (ckpt.selected_spec_mode != LLAMA_SPEC_CKPT_PER_STEP ||
!ckpt.dsv4_per_step_allocated || !ckpt.dsv4_per_step_saved) {
return true;
}
const bool ok =
dsv4_per_step_capture_group(*ctx, ctx->dsv4.cache.csa_state_kv, ctx->dsv4.csa_plan) &&
dsv4_per_step_capture_group(*ctx, ctx->dsv4.cache.csa_state_score, ctx->dsv4.csa_plan) &&
dsv4_per_step_capture_group(*ctx, ctx->dsv4.cache.hca_state_kv, ctx->dsv4.hca_plan) &&
dsv4_per_step_capture_group(*ctx, ctx->dsv4.cache.hca_state_score, ctx->dsv4.hca_plan) &&
dsv4_per_step_capture_group(*ctx, ctx->dsv4.cache.lid_state_kv, ctx->dsv4.lid_plan) &&
dsv4_per_step_capture_group(*ctx, ctx->dsv4.cache.lid_state_score, ctx->dsv4.lid_plan);
if (!ok) {
LLAMA_LOG_ERROR("%s: failed to queue DSV4 per-step compressor-state row capture\n", __func__);
}
return ok;
}
static bool dsv4_spec_ckpt_alloc_gpu(
llama_context & ctx,
const std::vector<ggml_tensor *> & tensors) {
auto & ckpt = ctx.kv_self.ckpt;
if (ckpt.dsv4_shadow_allocated) {
return ckpt.dsv4_state_shadow.size() == tensors.size();
}
struct tensor_entry {
size_t index;
ggml_tensor * source;
};
std::map<ggml_backend_buffer_type_t, std::vector<tensor_entry>> entries_by_buft;
const auto release_partial = [&]() {
ckpt.release_dsv4_snapshot();
};
for (size_t i = 0; i < tensors.size(); ++i) {
ggml_tensor * tensor = tensors[i];
if (tensor == nullptr) {
continue;
}
if (tensor->buffer == nullptr) {
return false;
}
entries_by_buft[ggml_backend_buffer_get_type(tensor->buffer)].push_back({ i, tensor });
}
ckpt.dsv4_state_shadow.assign(tensors.size(), nullptr);
for (auto & [buft, entries] : entries_by_buft) {
ggml_init_params params = {
/*.mem_size =*/ entries.size() * ggml_tensor_overhead(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
ggml_context * shadow_ctx = ggml_init(params);
if (shadow_ctx == nullptr) {
release_partial();
return false;
}
for (const auto & entry : entries) {
ggml_tensor * shadow = ggml_dup_tensor(shadow_ctx, entry.source);
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
shadow->nb[d] = entry.source->nb[d];
}
ggml_format_name(shadow, "dsv4_spec_shadow_%zu", entry.index);
ckpt.dsv4_state_shadow[entry.index] = shadow;
}
ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors_from_buft(shadow_ctx, buft);
if (buffer == nullptr) {
ggml_free(shadow_ctx);
release_partial();
return false;
}
ggml_backend_buffer_clear(buffer, 0);
LLAMA_LOG_INFO("%s: %10s DSV4 speculative shadow buffer = %8.2f MiB\n",
__func__, ggml_backend_buffer_name(buffer),
ggml_backend_buffer_get_size(buffer) / 1024.0 / 1024.0);
ckpt.dsv4_shadow_ctxs.push_back(shadow_ctx);
ckpt.dsv4_shadow_bufs.push_back(buffer);
}
ckpt.dsv4_shadow_allocated = true;
return true;
}
static bool dsv4_spec_ckpt_copy_gpu(
llama_context & ctx,
const std::vector<ggml_tensor *> & tensors,
bool restore) {
auto & ckpt = ctx.kv_self.ckpt;
if (!ckpt.dsv4_shadow_allocated || ckpt.dsv4_state_shadow.size() != tensors.size()) {
return false;
}
for (size_t i = 0; i < tensors.size(); ++i) {
ggml_tensor * tensor = tensors[i];
ggml_tensor * shadow = ckpt.dsv4_state_shadow[i];
if (tensor == nullptr || shadow == nullptr) {
continue;
}
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(ctx.sched, tensor);
if (backend == nullptr) {
return false;
}
if (restore) {
ggml_backend_tensor_copy_async(backend, backend, shadow, tensor);
} else {
ggml_backend_tensor_copy_async(backend, backend, tensor, shadow);
}
}
return true;
}
bool llama_dsv4_spec_ckpt_prepare(llama_context * ctx, int mode, int max_tokens) {
if (ctx == nullptr || ctx->model.arch != LLM_ARCH_DEEPSEEK4) {
return true;
}
if (mode == LLAMA_SPEC_CKPT_PER_STEP) {
return dsv4_per_step_alloc(*ctx, max_tokens);
}
if (mode == LLAMA_SPEC_CKPT_GPU_FALLBACK) {
return dsv4_spec_ckpt_alloc_gpu(*ctx, dsv4_state_tensors(*ctx));
}
return true;
}
bool llama_dsv4_spec_ckpt_save(llama_context * ctx, bool use_gpu) {
if (ctx == nullptr || ctx->model.arch != LLM_ARCH_DEEPSEEK4) {
return true;
}
if (ctx->kv_self.ckpt.selected_spec_mode == LLAMA_SPEC_CKPT_PER_STEP) {
auto & ckpt = ctx->kv_self.ckpt;
ckpt.dsv4_per_step_saved = false;
ckpt.dsv4_per_step_csa_src.clear();
ckpt.dsv4_per_step_csa_dst.clear();
ckpt.dsv4_per_step_hca_src.clear();
ckpt.dsv4_per_step_hca_dst.clear();
ckpt.dsv4_per_step_lid_src.clear();
ckpt.dsv4_per_step_lid_dst.clear();
if (!use_gpu || !dsv4_per_step_copy_base(*ctx, false)) {
LLAMA_LOG_ERROR("%s: failed to save DSV4 per-step compressor-state base\n", __func__);
return false;
}
ckpt.dsv4_per_step_saved = true;
return true;
}
const auto tensors = dsv4_state_tensors(*ctx);
ctx->kv_self.ckpt.dsv4_shadow_saved = false;
if (use_gpu) {
if (!dsv4_spec_ckpt_alloc_gpu(*ctx, tensors) || !dsv4_spec_ckpt_copy_gpu(*ctx, tensors, false)) {
LLAMA_LOG_ERROR("%s: failed to save DSV4 gpu-fallback checkpoint; explicit GPU mode will not downgrade to CPU\n", __func__);
return false;
}
ctx->kv_self.ckpt.dsv4_state_data.clear();
ctx->kv_self.ckpt.dsv4_shadow_saved = true;
return true;
}
auto & saved = ctx->kv_self.ckpt.dsv4_state_data;
saved.clear();
for (ggml_tensor * tensor : tensors) {
if (tensor == nullptr) {
saved.emplace_back();
continue;
}
const size_t nbytes = ggml_nbytes(tensor);
saved.emplace_back(nbytes);
ggml_backend_tensor_get(tensor, saved.back().data(), 0, nbytes);
}
return true;
}
static enum llama_spec_ckpt_restore_result dsv4_per_step_restore_rows(
llama_context & ctx,
const std::vector<ggml_tensor *> & states,
size_t delta_offset,
const std::vector<ggml_tensor *> & deltas,
const std::vector<int32_t> & src_idxs,
const std::vector<int32_t> & dst_idxs,
int accepted_step,
std::vector<ggml_backend_t> & backends) {
auto & ckpt = ctx.kv_self.ckpt;
if (src_idxs.size() != dst_idxs.size() || src_idxs.size() > (size_t) ckpt.dsv4_per_step_max_tokens ||
delta_offset > deltas.size() || states.size() > deltas.size() - delta_offset) {
LLAMA_LOG_ERROR("%s: invalid DSV4 per-step row restore: states=%zu delta_offset=%zu deltas=%zu src=%zu dst=%zu max=%d\n",
__func__, states.size(), delta_offset, deltas.size(), src_idxs.size(), dst_idxs.size(),
ckpt.dsv4_per_step_max_tokens);
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
for (size_t i = 0; i < states.size(); ++i) {
ggml_tensor * state = states[i];
ggml_tensor * delta = deltas[delta_offset + i];
if (state == nullptr || delta == nullptr) {
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(ctx.sched, state);
if (backend == nullptr) {
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
for (size_t row = 0; row < src_idxs.size(); ++row) {
if (src_idxs[row] > accepted_step) {
continue;
}
// Reject invalid mappings instead of leaving stale compressor state.
if (src_idxs[row] < 0 || dst_idxs[row] < 0 ||
(uint64_t) dst_idxs[row] >= (uint64_t) state->ne[1]) {
LLAMA_LOG_ERROR("%s: invalid visible DSV4 state row src=%d dst=%d accepted_step=%d state_rows=%lld\n",
__func__, src_idxs[row], dst_idxs[row], accepted_step, (long long) state->ne[1]);
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
ggml_tensor src_view = *delta;
ggml_tensor dst_view = *state;
src_view.ne[1] = src_view.ne[2] = src_view.ne[3] = 1;
dst_view.ne[1] = dst_view.ne[2] = dst_view.ne[3] = 1;
src_view.nb[2] = src_view.nb[3] = src_view.nb[1];
dst_view.nb[2] = dst_view.nb[3] = dst_view.nb[1];
src_view.data = (char *) delta->data + (size_t) src_idxs[row] * delta->nb[1];
dst_view.data = (char *) state->data + (size_t) dst_idxs[row] * state->nb[1];
src_view.view_src = nullptr;
dst_view.view_src = nullptr;
src_view.view_offs = 0;
dst_view.view_offs = 0;
ggml_backend_tensor_copy_async(backend, backend, &src_view, &dst_view);
}
if (std::find(backends.begin(), backends.end(), backend) == backends.end()) {
backends.push_back(backend);
}
}
return LLAMA_SPEC_CKPT_RESTORE_DIRECT;
}
enum llama_spec_ckpt_restore_result llama_dsv4_spec_ckpt_restore(llama_context * ctx, bool use_gpu, int accepted_step) {
if (ctx == nullptr || ctx->model.arch != LLM_ARCH_DEEPSEEK4) {
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
auto & ckpt = ctx->kv_self.ckpt;
if (ckpt.selected_spec_mode == LLAMA_SPEC_CKPT_PER_STEP) {
if (!ckpt.dsv4_per_step_saved || !dsv4_per_step_copy_base(*ctx, true)) {
LLAMA_LOG_ERROR("%s: failed to restore DSV4 per-step compressor-state base\n", __func__);
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
const auto compact = [](const std::vector<ggml_tensor *> & source) {
std::vector<ggml_tensor *> result;
for (ggml_tensor * tensor : source) {
if (tensor != nullptr) {
result.push_back(tensor);
}
}
return result;
};
const auto csa_kv = compact(ctx->dsv4.cache.csa_state_kv);
const auto csa_score = compact(ctx->dsv4.cache.csa_state_score);
const auto hca_kv = compact(ctx->dsv4.cache.hca_state_kv);
const auto hca_score = compact(ctx->dsv4.cache.hca_state_score);
const auto lid_kv = compact(ctx->dsv4.cache.lid_state_kv);
const auto lid_score = compact(ctx->dsv4.cache.lid_state_score);
const size_t csa_kv_off = 0;
const size_t csa_score_off = csa_kv_off + csa_kv.size();
const size_t hca_kv_off = csa_score_off + csa_score.size();
const size_t hca_score_off = hca_kv_off + hca_kv.size();
const size_t lid_kv_off = hca_score_off + hca_score.size();
const size_t lid_score_off = lid_kv_off + lid_kv.size();
if (ckpt.dsv4_per_step_delta.size() != lid_score_off + lid_score.size()) {
LLAMA_LOG_ERROR("%s: DSV4 per-step delta tensor layout mismatch\n", __func__);
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
std::vector<ggml_backend_t> backends;
const auto restore_group = [&](const std::vector<ggml_tensor *> & states, size_t offset) {
return dsv4_per_step_restore_rows(*ctx, states, offset, ckpt.dsv4_per_step_delta,
offset == csa_kv_off || offset == csa_score_off ? ckpt.dsv4_per_step_csa_src :
offset == hca_kv_off || offset == hca_score_off ? ckpt.dsv4_per_step_hca_src : ckpt.dsv4_per_step_lid_src,
offset == csa_kv_off || offset == csa_score_off ? ckpt.dsv4_per_step_csa_dst :
offset == hca_kv_off || offset == hca_score_off ? ckpt.dsv4_per_step_hca_dst : ckpt.dsv4_per_step_lid_dst,
accepted_step, backends);
};
if (restore_group(csa_kv, csa_kv_off) == LLAMA_SPEC_CKPT_RESTORE_FAILED ||
restore_group(csa_score, csa_score_off) == LLAMA_SPEC_CKPT_RESTORE_FAILED ||
restore_group(hca_kv, hca_kv_off) == LLAMA_SPEC_CKPT_RESTORE_FAILED ||
restore_group(hca_score, hca_score_off) == LLAMA_SPEC_CKPT_RESTORE_FAILED ||
restore_group(lid_kv, lid_kv_off) == LLAMA_SPEC_CKPT_RESTORE_FAILED ||
restore_group(lid_score, lid_score_off) == LLAMA_SPEC_CKPT_RESTORE_FAILED) {
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
for (ggml_backend_t backend : backends) {
ggml_backend_synchronize(backend);
}
return LLAMA_SPEC_CKPT_RESTORE_DIRECT;
}
const auto tensors = dsv4_state_tensors(*ctx);
if (use_gpu && ctx->kv_self.ckpt.dsv4_shadow_saved) {
return dsv4_spec_ckpt_copy_gpu(*ctx, tensors, true)
? LLAMA_SPEC_CKPT_RESTORE_BASE_REPLAY_REQUIRED
: LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
const auto & saved = ctx->kv_self.ckpt.dsv4_state_data;
if (saved.size() != tensors.size()) {
LLAMA_LOG_ERROR("%s: DSV4 checkpoint tensor count mismatch: saved=%zu current=%zu\n",
__func__, saved.size(), tensors.size());
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
for (size_t i = 0; i < tensors.size(); ++i) {
ggml_tensor * tensor = tensors[i];
if (tensor == nullptr) {
if (!saved[i].empty()) {
LLAMA_LOG_ERROR("%s: DSV4 checkpoint null tensor %zu has saved data\n", __func__, i);
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
continue;
}
if (saved[i].size() != ggml_nbytes(tensor)) {
LLAMA_LOG_ERROR("%s: DSV4 checkpoint tensor %zu size mismatch\n", __func__, i);
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
if (!saved[i].empty()) {
ggml_backend_tensor_set(tensor, saved[i].data(), 0, saved[i].size());
}
}
return LLAMA_SPEC_CKPT_RESTORE_BASE_REPLAY_REQUIRED;
}
ggml_tensor * llama_dsv4_spec_ckpt_delta(llama_context * ctx, ggml_tensor * state_tensor) {
if (ctx == nullptr || state_tensor == nullptr ||
ctx->kv_self.ckpt.selected_spec_mode != LLAMA_SPEC_CKPT_PER_STEP ||
!ctx->kv_self.ckpt.dsv4_per_step_allocated) {
return nullptr;
}
auto & ckpt = ctx->kv_self.ckpt;
for (size_t i = 0; i < ckpt.dsv4_per_step_state.size(); ++i) {
if (ckpt.dsv4_per_step_state[i] == state_tensor) {
return ckpt.dsv4_per_step_delta[i];
}
}
return nullptr;
}
void llama_dsv4_spec_ckpt_record_plan(llama_context * ctx) {
if (ctx == nullptr || ctx->kv_self.ckpt.selected_spec_mode != LLAMA_SPEC_CKPT_PER_STEP) {
return;
}
auto & ckpt = ctx->kv_self.ckpt;
ckpt.dsv4_per_step_csa_src = ctx->dsv4.csa_plan.state_delta_src_idxs;
ckpt.dsv4_per_step_csa_dst = ctx->dsv4.csa_plan.state_delta_dst_idxs;
ckpt.dsv4_per_step_hca_src = ctx->dsv4.hca_plan.state_delta_src_idxs;
ckpt.dsv4_per_step_hca_dst = ctx->dsv4.hca_plan.state_delta_dst_idxs;
ckpt.dsv4_per_step_lid_src = ctx->dsv4.lid_plan.state_delta_src_idxs;
ckpt.dsv4_per_step_lid_dst = ctx->dsv4.lid_plan.state_delta_dst_idxs;
}
void llama_dsv4_spec_ckpt_discard(llama_context * ctx) {
if (ctx != nullptr) {
ctx->kv_self.ckpt.dsv4_state_data.clear();
ctx->kv_self.ckpt.dsv4_shadow_saved = false;
ctx->kv_self.ckpt.dsv4_per_step_saved = false;
ctx->kv_self.ckpt.dsv4_per_step_csa_src.clear();
ctx->kv_self.ckpt.dsv4_per_step_csa_dst.clear();
ctx->kv_self.ckpt.dsv4_per_step_hca_src.clear();
ctx->kv_self.ckpt.dsv4_per_step_hca_dst.clear();
ctx->kv_self.ckpt.dsv4_per_step_lid_src.clear();
ctx->kv_self.ckpt.dsv4_per_step_lid_dst.clear();
}
}
bool llama_prepare_dsv4_graph_inputs(llama_context & lctx, const llama_batch & batch, bool set_tensors, bool reserve_plan) {
if (lctx.model.arch != LLM_ARCH_DEEPSEEK4) {
return true;
@@ -1087,6 +1729,7 @@ bool llama_prepare_dsv4_graph_inputs(llama_context & lctx, const llama_batch & b
set_comp(lctx.dsv4.inputs.csa, lctx.dsv4.csa_plan, true);
set_comp(lctx.dsv4.inputs.hca, lctx.dsv4.hca_plan, true);
set_comp(lctx.dsv4.inputs.lid, lctx.dsv4.lid_plan, false);
llama_dsv4_spec_ckpt_record_plan(&lctx);
//tim2 = ggml_time_us();
//fprintf(stderr, "%s: setting input tensors took %ld us\n", __func__, tim2 - tim1);
src/llama-dsv4.h
+10
View File
@@ -1,8 +1,18 @@
#pragma once
#include "llama.h"
#include <cstdint>
struct llama_batch;
struct llama_context;
struct ggml_tensor;
bool llama_prepare_dsv4_graph_inputs(llama_context & lctx, const llama_batch & batch, bool set_tensors, bool reserve_plan);
void llama_reset_dsv4_state(llama_context * ctx, int32_t seq_id = -1);
bool llama_dsv4_spec_ckpt_prepare(llama_context * ctx, int mode, int max_tokens);
bool llama_dsv4_spec_ckpt_save(llama_context * ctx, bool use_gpu);
bool llama_dsv4_spec_ckpt_capture_rows(llama_context * ctx);
enum llama_spec_ckpt_restore_result llama_dsv4_spec_ckpt_restore(llama_context * ctx, bool use_gpu, int accepted_step);
void llama_dsv4_spec_ckpt_discard(llama_context * ctx);
ggml_tensor * llama_dsv4_spec_ckpt_delta(llama_context * ctx, ggml_tensor * state_tensor);
void llama_dsv4_spec_ckpt_record_plan(llama_context * ctx);
src/llama-model.cpp
+4
View File
@@ -2178,6 +2178,10 @@ bool llama_model_has_recurrent(const llama_model * model) {
return llm_arch_is_hybrid(model->arch) || llm_arch_is_recurrent(model->arch);
}
bool llama_model_is_deepseek4(const llama_model * model) {
return model && model->arch == LLM_ARCH_DEEPSEEK4;
}
bool llama_model_is_openpangu(const llama_model * model) {
return model && model->arch == LLM_ARCH_OPENPANGU;
}
src/llama.cpp
+76 -18
View File
@@ -774,6 +774,7 @@ llama_context::~llama_context() {
if (dflash.kv.cache_sched != nullptr) {
ggml_backend_sched_free(dflash.kv.cache_sched);
}
kv_self.ckpt.release();
free_dflash_kv_cache_tensors();
free_dsv4_cache_tensors();
ggml_backend_sched_free(sched);
@@ -6248,6 +6249,12 @@ static int llama_decode_internal(
//fprintf(stderr, "%s: invoking llama_graph_compute\n", __func__);
llama_graph_compute(lctx, gf, n_threads);
if (lctx.model.arch == LLM_ARCH_DEEPSEEK4 &&
lctx.kv_self.ckpt.selected_spec_mode == LLAMA_SPEC_CKPT_PER_STEP &&
!llama_dsv4_spec_ckpt_capture_rows(&lctx)) {
return GGML_STATUS_FAILED;
}
#if IK_PRINT_TIMING
llama_synchronize(&lctx);
tim2 = ggml_time_us();
@@ -8973,11 +8980,12 @@ static const char * llama_spec_ckpt_mode_name(int mode) {
int llama_spec_ckpt_init(struct llama_context * ctx, int mode, int max_tokens) {
auto & kv = ctx->kv_self;
const bool is_dsv4 = ctx->model.arch == LLM_ARCH_DEEPSEEK4;
kv.save_per_step_ssm = false;
kv.ckpt.selected_spec_mode = LLAMA_SPEC_CKPT_NONE;
if (!kv.checkpoint_supported()) {
if (!kv.checkpoint_supported() && !is_dsv4) {
return (int)LLAMA_SPEC_CKPT_NONE;
}
@@ -8992,19 +9000,37 @@ int llama_spec_ckpt_init(struct llama_context * ctx, int mode, int max_tokens) {
return kv.ckpt.selected_spec_mode;
}
if (is_dsv4 && mode != LLAMA_SPEC_CKPT_AUTO &&
mode != LLAMA_SPEC_CKPT_PER_STEP && mode != LLAMA_SPEC_CKPT_GPU_FALLBACK &&
mode != LLAMA_SPEC_CKPT_CPU) {
LLAMA_LOG_ERROR("%s: unsupported DSV4 checkpoint mode %d\n", __func__, mode);
return (int)LLAMA_SPEC_CKPT_NONE;
}
int requested = mode;
int resolved = LLAMA_SPEC_CKPT_NONE;
const auto prepare_per_step = [&]() {
return is_dsv4
? llama_dsv4_spec_ckpt_prepare(ctx, LLAMA_SPEC_CKPT_PER_STEP, max_tokens)
: spec_ckpt_try_per_step(kv, ctx->model, max_tokens);
};
const auto prepare_gpu_fallback = [&]() {
return is_dsv4
? llama_dsv4_spec_ckpt_prepare(ctx, LLAMA_SPEC_CKPT_GPU_FALLBACK, max_tokens)
: kv.checkpoint_alloc_shadows();
};
// prefer PER_STEP → GPU_FALLBACK → CPU
if (requested == LLAMA_SPEC_CKPT_AUTO) {
requested = LLAMA_SPEC_CKPT_PER_STEP;
}
if (requested == LLAMA_SPEC_CKPT_PER_STEP) {
if (spec_ckpt_try_per_step(kv, ctx->model, max_tokens)) {
if (prepare_per_step()) {
resolved = LLAMA_SPEC_CKPT_PER_STEP;
} else if (mode == LLAMA_SPEC_CKPT_PER_STEP) {
LLAMA_LOG_ERROR("%s: failed to preallocate per-step checkpoint buffers for max_tokens=%d; --recurrent-ckpt-mode=%s requires startup allocation\n",
LLAMA_LOG_ERROR("%s: failed to preallocate per-step checkpoint buffers for max_tokens=%d; --spec-ckpt-mode=%s requires startup allocation\n",
__func__, max_tokens, llama_spec_ckpt_mode_name(mode));
return (int)LLAMA_SPEC_CKPT_NONE;
} else {
@@ -9015,10 +9041,10 @@ int llama_spec_ckpt_init(struct llama_context * ctx, int mode, int max_tokens) {
}
if (resolved == LLAMA_SPEC_CKPT_NONE && requested == LLAMA_SPEC_CKPT_GPU_FALLBACK) {
if (kv.checkpoint_alloc_shadows()) {
if (prepare_gpu_fallback()) {
resolved = LLAMA_SPEC_CKPT_GPU_FALLBACK;
} else if (mode == LLAMA_SPEC_CKPT_GPU_FALLBACK) {
LLAMA_LOG_ERROR("%s: failed to preallocate gpu-fallback checkpoint shadows at startup; --recurrent-ckpt-mode=%s requires startup allocation\n",
LLAMA_LOG_ERROR("%s: failed to preallocate gpu-fallback checkpoint shadows at startup; --spec-ckpt-mode=%s requires startup allocation\n",
__func__, llama_spec_ckpt_mode_name(mode));
return (int)LLAMA_SPEC_CKPT_NONE;
} else {
@@ -9032,7 +9058,7 @@ int llama_spec_ckpt_init(struct llama_context * ctx, int mode, int max_tokens) {
resolved = LLAMA_SPEC_CKPT_CPU;
}
if (resolved == LLAMA_SPEC_CKPT_CPU) {
if (resolved == LLAMA_SPEC_CKPT_CPU && !is_dsv4) {
const size_t cpu_reserve = llama_spec_ckpt_cpu_state_reserve(ctx, 0);
kv.ckpt.cpu_state_data.clear();
kv.ckpt.cpu_state_data.reserve(cpu_reserve);
@@ -9044,8 +9070,8 @@ int llama_spec_ckpt_init(struct llama_context * ctx, int mode, int max_tokens) {
kv.ckpt.fixed_max_tokens = resolved == LLAMA_SPEC_CKPT_PER_STEP ? max_tokens : 0;
kv.ckpt.selected_spec_mode = resolved;
LLAMA_LOG_INFO("%s: fixed recurrent checkpoint mode = %s%s\n",
__func__, llama_spec_ckpt_mode_name(resolved),
LLAMA_LOG_INFO("%s: fixed %s checkpoint mode = %s%s\n",
__func__, is_dsv4 ? "DSV4" : "recurrent", llama_spec_ckpt_mode_name(resolved),
resolved == LLAMA_SPEC_CKPT_PER_STEP ? (std::string(" (max_tokens=") + std::to_string(max_tokens) + ")").c_str() : "");
return resolved;
@@ -9056,13 +9082,22 @@ bool llama_spec_ckpt_save(struct llama_context * ctx, llama_seq_id seq_id) {
switch (kv.ckpt.selected_spec_mode) {
case LLAMA_SPEC_CKPT_PER_STEP:
if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) {
return llama_dsv4_spec_ckpt_save(ctx, true);
}
kv.save_per_step_ssm = true;
return true;
case LLAMA_SPEC_CKPT_GPU_FALLBACK:
if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) {
return llama_dsv4_spec_ckpt_save(ctx, true);
}
return kv.checkpoint_save(ctx->sched);
case LLAMA_SPEC_CKPT_CPU: {
if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) {
return llama_dsv4_spec_ckpt_save(ctx, false);
}
const size_t need = llama_state_seq_get_size(ctx, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
kv.ckpt.cpu_state_data.resize(need);
const size_t written = llama_state_seq_get_data(
@@ -9076,54 +9111,77 @@ bool llama_spec_ckpt_save(struct llama_context * ctx, llama_seq_id seq_id) {
}
}
bool llama_spec_ckpt_restore(struct llama_context * ctx, llama_seq_id seq_id,
llama_pos n_past, int accepted_step) {
enum llama_spec_ckpt_restore_result llama_spec_ckpt_restore_ex(
struct llama_context * ctx, llama_seq_id seq_id,
llama_pos n_past, int accepted_step) {
auto & kv = ctx->kv_self;
switch (kv.ckpt.selected_spec_mode) {
case LLAMA_SPEC_CKPT_PER_STEP: {
if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) {
const llama_pos accepted_pos = n_past + accepted_step;
llama_kv_cache_seq_rm(kv, seq_id, accepted_pos + 1, -1);
return llama_dsv4_spec_ckpt_restore(ctx, true, accepted_step);
}
if (!kv.per_step_restore(ctx->model, ctx->sched, accepted_step)) {
return false;
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
const llama_pos accepted_pos = n_past + accepted_step;
if (seq_id >= 0 && (uint32_t)seq_id < kv.size) {
kv.cells[seq_id].pos = accepted_pos;
}
llama_kv_cache_seq_rm(kv, seq_id, accepted_pos + 1, -1);
return true;
return LLAMA_SPEC_CKPT_RESTORE_DIRECT;
}
case LLAMA_SPEC_CKPT_GPU_FALLBACK:
kv.checkpoint_restore(ctx->sched);
if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) {
llama_kv_cache_seq_rm(kv, seq_id, n_past, -1);
return llama_dsv4_spec_ckpt_restore(ctx, true, 0);
}
if (!kv.checkpoint_restore(ctx->sched)) {
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
llama_kv_cache_seq_rm(kv, seq_id, n_past, -1);
return false;
return LLAMA_SPEC_CKPT_RESTORE_BASE_REPLAY_REQUIRED;
case LLAMA_SPEC_CKPT_CPU:
if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) {
llama_kv_cache_seq_rm(kv, seq_id, n_past, -1);
return llama_dsv4_spec_ckpt_restore(ctx, false, 0);
}
if (!kv.ckpt.cpu_state_data.empty()) {
llama_state_seq_set_data(ctx, kv.ckpt.cpu_state_data.data(),
kv.ckpt.cpu_state_data.size(), seq_id,
LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
llama_kv_cache_seq_rm(kv, seq_id, n_past, -1);
return false;
return LLAMA_SPEC_CKPT_RESTORE_BASE_REPLAY_REQUIRED;
default:
return false;
return LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
}
bool llama_spec_ckpt_restore(struct llama_context * ctx, llama_seq_id seq_id,
llama_pos n_past, int accepted_step) {
return llama_spec_ckpt_restore_ex(ctx, seq_id, n_past, accepted_step) != LLAMA_SPEC_CKPT_RESTORE_FAILED;
}
void llama_spec_ckpt_discard(struct llama_context * ctx) {
auto & kv = ctx->kv_self;
if (kv.ckpt.selected_spec_mode == LLAMA_SPEC_CKPT_PER_STEP) {
kv.save_per_step_ssm = false;
kv.checkpoint_delete();
} else if (kv.ckpt.selected_spec_mode == LLAMA_SPEC_CKPT_GPU_FALLBACK) {
} else if (kv.ckpt.selected_spec_mode == LLAMA_SPEC_CKPT_GPU_FALLBACK &&
ctx->model.arch != LLM_ARCH_DEEPSEEK4) {
kv.checkpoint_delete();
}
kv.ckpt.selected_spec_mode = LLAMA_SPEC_CKPT_NONE;
kv.ckpt.cpu_state_data.clear();
llama_dsv4_spec_ckpt_discard(ctx);
}
bool llama_kv_cache_seq_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
@@ -10317,7 +10375,7 @@ struct llama_data_read_file : llama_data_read {
}
};
// Refuse state I/O when private per-position state is not part of the format.
// Public state I/O excludes private DSV4 state, speculation uses an internal checkpoint.
static bool llama_state_io_supported(const struct llama_context * ctx, const char * func) {
if (ctx->model.arch == LLM_ARCH_OPENPANGU || ctx->model.arch == LLM_ARCH_DEEPSEEK4) {
const char * arch = ctx->model.arch == LLM_ARCH_OPENPANGU ? "openPangu" : "DeepSeek4";