Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

DeepseekV4 MTP + DSpark (#25784)

Browse Source
This commit is contained in:
Aman Gupta
2026-08-02 20:55:34 +08:00
committed by GitHub
parent f5919bf458
commit 596a5795bd
14 changed files with 1534 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test-recurrent-state-rollback.cpp
+79 -45
View File
@@ -83,27 +83,33 @@ int main(int argc, char ** argv) {
if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) {
tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 };
} else {
tokens = common_tokenize(ctx_src, "The quick brown fox jumps", true);
tokens = common_tokenize(ctx_src, "The quick brown fox jumps over the lazy dog", true);
}
const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src);
if (tokens.size() > n_rs_seq + 1) {
tokens.resize(n_rs_seq + 1);
constexpr uint32_t n_rollback = 3;
if (n_rs_seq < n_rollback) {
fprintf(stderr, "%s : skipping because n_rs_seq is too small\n", __func__);
llama_free(ctx_src);
llama_free(ctx_dst);
return 0;
}
if (tokens.size() < 2) {
if (tokens.empty()) {
fprintf(stderr, "%s : not enough prompt tokens\n", __func__);
return 1;
}
const uint32_t n_tokens = tokens.size();
const llama_token last_tok = tokens.back();
const llama_pos last_pos = (llama_pos) n_tokens - 2;
tokens.resize(n_rs_seq + 1, tokens.back());
// Decode the full prompt on the source, then roll back the last position.
const uint32_t n_tokens = tokens.size();
const llama_pos rollback_pos = (llama_pos) n_tokens - n_rollback;
// Decode the full prompt on the source, then roll back three positions.
// Replaying them crosses DSV4's ratio-4 compressor boundary.
// Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0).
if (!decode_tokens(ctx_src, tokens, n_tokens)) {
fprintf(stderr, "%s : failed to decode prompt\n", __func__);
return 1;
}
if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, last_pos, -1)) {
if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1)) {
fprintf(stderr, "%s : rollback failed\n", __func__);
return 1;
}
@@ -113,31 +119,56 @@ int main(int argc, char ** argv) {
ckpt.update_tgt(ctx_src, 0, 0);
ckpt.load_tgt(ctx_dst, 0, 0);
// Replay the rolled-back token on both contexts and compare logits.
if (!decode_one(ctx_src, last_tok, last_pos) ||
!decode_one(ctx_dst, last_tok, last_pos)) {
fprintf(stderr, "%s : replay failed\n", __func__);
return 1;
}
const float * logits_src = llama_get_logits_ith(ctx_src, 0);
const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);
if (logits_src == nullptr || logits_dst == nullptr) {
fprintf(stderr, "%s : missing logits\n", __func__);
return 1;
}
constexpr float eps = 1e-5f;
for (int i = 0; i < n_vocab; ++i) {
if (std::fabs(logits_src[i] - logits_dst[i]) > eps) {
fprintf(stderr, "%s : logits mismatch at token %d (%g != %g)\n",
__func__, i, (double) logits_src[i], (double) logits_dst[i]);
return 1;
std::vector<std::vector<float>> logits_src_replay(n_rollback);
const auto replay_and_compare = [&](const char * mode) {
for (uint32_t i = 0; i < n_rollback; ++i) {
const llama_pos pos = rollback_pos + i;
if (!decode_one(ctx_src, tokens[pos], pos) ||
!decode_one(ctx_dst, tokens[pos], pos)) {
fprintf(stderr, "%s : %s replay failed at position %d\n", __func__, mode, pos);
return false;
}
const float * logits_src = llama_get_logits_ith(ctx_src, 0);
const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);
if (logits_src == nullptr || logits_dst == nullptr) {
fprintf(stderr, "%s : missing %s logits at position %d\n", __func__, mode, pos);
return false;
}
logits_src_replay[i].assign(logits_src, logits_src + n_vocab);
for (int token = 0; token < n_vocab; ++token) {
if (std::fabs(logits_src[token] - logits_dst[token]) > eps) {
fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n",
__func__, mode, pos, token, (double) logits_src[token], (double) logits_dst[token]);
return false;
}
}
}
return true;
};
if (!replay_and_compare("full")) {
return 1;
}
if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) ||
!llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) {
fprintf(stderr, "%s : partial rollback failed\n", __func__);
return 1;
}
constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY;
common_prompt_checkpoint ckpt_partial;
ckpt_partial.update_tgt(ctx_src, 0, partial_flags);
ckpt_partial.load_tgt(ctx_dst, 0, partial_flags);
if (!replay_and_compare("partial")) {
return 1;
}
// Repeat the load into a context that already has its own rollback state:
// groups 1..n_rs_seq hold a *different* prompt's history, and rs_idx[0] is
// groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is
// non-zero at load time. The restore must wipe that state and still match.
llama_context * ctx_dirty = make_ctx(params, model);
if (ctx_dirty == nullptr) {
@@ -156,30 +187,33 @@ int main(int argc, char ** argv) {
fprintf(stderr, "%s : dirty prompt decode failed\n", __func__);
return 1;
}
if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, last_pos, -1)) {
if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, rollback_pos, -1)) {
fprintf(stderr, "%s : dirty rollback failed\n", __func__);
return 1;
}
ckpt.load_tgt(ctx_dirty, 0, 0);
if (!decode_one(ctx_dirty, last_tok, last_pos)) {
fprintf(stderr, "%s : dirty replay failed\n", __func__);
return 1;
}
const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);
if (logits_dirty == nullptr) {
fprintf(stderr, "%s : missing dirty logits\n", __func__);
return 1;
}
for (int i = 0; i < n_vocab; ++i) {
if (std::fabs(logits_src[i] - logits_dirty[i]) > eps) {
fprintf(stderr, "%s : dirty-ctx logits mismatch at token %d (%g != %g)\n",
__func__, i, (double) logits_src[i], (double) logits_dirty[i]);
for (uint32_t i = 0; i < n_rollback; ++i) {
const llama_pos pos = rollback_pos + i;
if (!decode_one(ctx_dirty, tokens[pos], pos)) {
fprintf(stderr, "%s : dirty replay failed at position %d\n", __func__, pos);
return 1;
}
const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);
if (logits_dirty == nullptr) {
fprintf(stderr, "%s : missing dirty logits at position %d\n", __func__, pos);
return 1;
}
for (int token = 0; token < n_vocab; ++token) {
if (std::fabs(logits_src_replay[i][token] - logits_dirty[token]) > eps) {
fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, token %d (%g != %g)\n",
__func__, pos, token, (double) logits_src_replay[i][token], (double) logits_dirty[token]);
return 1;
}
}
}
fprintf(stderr, "%s : recurrent rollback checkpoint restored successfully\n", __func__);