wmantly/llama.cpp-xrip
1
0 Fork 0
Code Issues Pull Requests Actions 36 Packages Projects Releases Wiki Activity

kv-cache: fix restoring mismatched KV cache rotation by saving exact rotation metadata (#28498)

Browse Source 
* kv-cache: save exact KV rotation metadata, reject restoring mismatched rotation

* tests : move the state rotation test to test-save-load-state

the test is now part of the save/load test matrix and runs against
every model under test, like the rest of the suite

it probes the KV cache type combinations supported by the model and
treats models that do not use attention rotation as passing vacuously

Assisted-by: pi:llama.cpp/Qwen3.8-27B

* cont : skip unsupported KV caches

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
Evan HuusandGeorgi Gerganov authored and GitHub committed 2026-10-05 11:27:30 +03:00
1 parent e5983d6704
commit 210791069b
4 files changed
+140 -34
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/llama-kv-cache.cpp
+45 -28
View File
@@ -309,8 +309,8 @@ llama_kv_cache::llama_kv_cache(
n_embd_head_k_all = other->n_embd_head_k_all;
n_embd_head_v_all = other->n_embd_head_v_all;
attn_rot_k = other->attn_rot_k;
attn_rot_v = other->attn_rot_v;
n_rot_k = other->n_rot_k;
n_rot_v = other->n_rot_v;
} else {
const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");
const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;
@@ -318,7 +318,7 @@ llama_kv_cache::llama_kv_cache(
LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);
}
attn_rot_k =
bool attn_rot_k =
!attn_rot_disable &&
n_embd_head_k_all > 0 &&
ggml_is_quantized(type_k) &&
@@ -331,19 +331,37 @@ llama_kv_cache::llama_kv_cache(
attn_rot_k = true;
}
attn_rot_v =
const bool attn_rot_v =
!attn_rot_disable &&
n_embd_head_v_all > 0 &&
ggml_is_quantized(type_v) &&
hparams.n_embd_head_v() % 64 == 0;
if (attn_rot_k) {
n_rot_k = 64;
// TODO: investigate if using the smallest rotation matrix is beneficial also for K (similar as for V)
// ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4141323088
while (n_embd_head_k_all % (2*n_rot_k) == 0) {
n_rot_k *= 2;
}
}
if (attn_rot_v) {
n_rot_v = 64;
// using smaller rotation matrices for V seems beneficial
// ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4146397570
//while (hparams.n_embd_head_v() % (2*n_rot_v) == 0) {
// n_rot_v *= 2;
//}
}
}
LLAMA_LOG_INFO("%s: attn_rot_k = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_k, n_embd_head_k_all);
LLAMA_LOG_INFO("%s: attn_rot_v = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_v, n_embd_head_v_all);
LLAMA_LOG_INFO("%s: n_rot_k = %u, n_embd_head_k_all = %d\n", __func__, n_rot_k, n_embd_head_k_all);
LLAMA_LOG_INFO("%s: n_rot_v = %u, n_embd_head_k_all = %d\n", __func__, n_rot_v, n_embd_head_v_all);
// pre-compute the haramard matrices and keep them in host memory
// TODO: in the future, we can make copies in the backend buffers to avoid host -> device transfers
if (attn_rot_k || attn_rot_v) {
if (n_rot_k || n_rot_v) {
for (int64_t n = 64; n <= std::max(n_embd_head_k_all, n_embd_head_v_all); n *= 2) {
attn_rot_hadamard[n] = std::vector<float>(n*n);
@@ -1449,17 +1467,8 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama
ggml_tensor * llama_kv_cache::build_input_k_rot(ggml_context * ctx) const {
ggml_tensor * res = nullptr;
if (attn_rot_k) {
int nrot = 64;
// TODO: investigate if using the smallest rotation matrix is beneficial also for K (similar as for V)
// ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4141323088
do {
nrot *= 2;
} while (n_embd_head_k_all % nrot == 0);
nrot /= 2;
res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot);
if (n_rot_k) {
res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_rot_k, n_rot_k);
ggml_set_input(res);
ggml_set_name(res, "attn_inp_k_rot");
}
@@ -1470,16 +1479,8 @@ ggml_tensor * llama_kv_cache::build_input_k_rot(ggml_context * ctx) const {
ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const {
ggml_tensor * res = nullptr;
if (attn_rot_v) {
int nrot = 64;
// using smaller rotation matrices for V seems beneficial
// ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4146397570
//do {
// nrot *= 2;
//} while (hparams.n_embd_head_v() % nrot == 0);
//nrot /= 2;
res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot);
if (n_rot_v) {
res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_rot_v, n_rot_v);
ggml_set_input(res);
ggml_set_name(res, "attn_inp_v_rot");
}
@@ -2254,6 +2255,8 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
io.write(&v_trans, sizeof(v_trans));
io.write(&n_layer, sizeof(n_layer));
io.write(&n_rot_k, sizeof(n_rot_k));
io.write(&n_rot_v, sizeof(n_rot_v));
// Iterate and write all the keys first, each row is a cell
// Get whole range at a time
@@ -2546,9 +2549,13 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
uint32_t v_trans;
uint32_t n_layer;
uint32_t n_rot_k_ref;
uint32_t n_rot_v_ref;
io.read(&v_trans, sizeof(v_trans));
io.read(&n_layer, sizeof(n_layer));
io.read(&n_rot_k_ref, sizeof(n_rot_k_ref));
io.read(&n_rot_v_ref, sizeof(n_rot_v_ref));
if (n_layer != layers.size()) {
LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, (uint32_t) layers.size());
@@ -2565,6 +2572,16 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
return false;
}
if (n_rot_k_ref != n_rot_k) {
LLAMA_LOG_ERROR("%s: incompatible key rotation (%u instead of %u)\n", __func__, n_rot_k_ref, n_rot_k);
return false;
}
if (n_rot_v_ref != n_rot_v) {
LLAMA_LOG_ERROR("%s: incompatible value rotation (%u instead of %u)\n", __func__, n_rot_v_ref, n_rot_v);
return false;
}
// For each layer, read the keys for each cell, one row is one cell, read as one contiguous block
for (const auto & layer : layers) {
const uint32_t il = layer.il;