Commit 210791069 for llama.cpp
commit 210791069bf5460bef3c09b982ee6f0aa9967404
Author: Evan Huus <eapache@gmail.com>
Date: Mon Oct 5 04:27:30 2026 -0400
kv-cache: fix restoring mismatched KV cache rotation by saving exact rotation metadata (#28498)
* 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>
diff --git a/include/llama.h b/include/llama.h
index dea4d4606..9bca8ee10 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -43,10 +43,10 @@
#define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq'
#define LLAMA_SESSION_MAGIC LLAMA_FILE_MAGIC_GGSN
-#define LLAMA_SESSION_VERSION 10
+#define LLAMA_SESSION_VERSION 11
#define LLAMA_STATE_SEQ_MAGIC LLAMA_FILE_MAGIC_GGSQ
-#define LLAMA_STATE_SEQ_VERSION 3
+#define LLAMA_STATE_SEQ_VERSION 4
#ifdef __cplusplus
extern "C" {
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 8891526cd..463510caa 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -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;
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index 52cc8cd32..9bfb8f231 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -278,8 +278,8 @@ private:
const uint32_t n_swa = 0;
// env: LLAMA_ATTN_ROT_DISABLE
- bool attn_rot_k = false;
- bool attn_rot_v = false;
+ uint32_t n_rot_k = 0;
+ uint32_t n_rot_v = 0;
// if all layers participating in the cache have constant head size, the value is stored here
// otherwise the value is -1
diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp
index 02984d185..b9d2fd8e6 100644
--- a/tests/test-save-load-state.cpp
+++ b/tests/test-save-load-state.cpp
@@ -744,6 +744,92 @@ static bool test_state_restore_failure(struct llama_model * model, const struct
}
+// Test 10: state rotation
+// a KV state saved with attention rotation enabled must restore only into a context with the same setting;
+// note: rotation is only active for quantized KV caches with a head size that is a multiple of 64,
+// for other models the restore into the rotation-disabled context is valid and the test passes vacuously
+static bool test_state_rotation(struct llama_model * model, const struct common_params & params) {
+ LOGV(LOG_LEVEL_INFO, "\n=== Test 10: state rotation ===\n");
+
+ const std::string attn_rot_disable = common_get_env("LLAMA_ATTN_ROT_DISABLE");
+ const auto make_context = [&](ggml_type type_k, ggml_type type_v, bool disable_rotation) {
+ common_set_env("LLAMA_ATTN_ROT_DISABLE", disable_rotation ? "1" : "0");
+ auto params_ctx = common_context_params_to_llama(params);
+ params_ctx.n_ctx = 32;
+ params_ctx.n_batch = 1;
+ params_ctx.n_ubatch = 1;
+ params_ctx.type_k = type_k;
+ params_ctx.type_v = type_v;
+ params_ctx.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
+ return llama_context_ptr(llama_init_from_model(model, params_ctx));
+ };
+
+ std::vector<std::pair<ggml_type, ggml_type>> type_pairs;
+ for (const auto & types : { std::pair{GGML_TYPE_Q8_0, GGML_TYPE_Q8_0} }) {
+ if (make_context(types.first, types.second, false)) {
+ type_pairs.push_back(types);
+ }
+ }
+ if (type_pairs.empty()) {
+ LOG_WRN("%s: no supported quantized KV cache type combination - skipping\n", __func__);
+ return true;
+ }
+
+ bool success = true;
+ for (const auto & types : type_pairs) {
+ auto src = make_context(types.first, types.second, false);
+ if (!src) {
+ LOG_ERR("%s: failed to create source context\n", __func__);
+ success = false;
+ break;
+ }
+
+ llama_token token = 0;
+ if (llama_decode(src.get(), llama_batch_get_one(&token, 1))) {
+ LOG_ERR("%s: failed to decode token\n", __func__);
+ success = false;
+ break;
+ }
+
+ const size_t state_size = llama_state_seq_get_size(src.get(), 0);
+ if (state_size == 0) {
+ continue; // no KV state to test
+ }
+
+ std::vector<uint8_t> state(state_size);
+ if (llama_state_seq_get_data(src.get(), state.data(), state.size(), 0) != state.size()) {
+ LOG_ERR("%s: failed to save sequence state\n", __func__);
+ success = false;
+ break;
+ }
+
+ auto matching = make_context(types.first, types.second, false);
+ if (!matching || llama_state_seq_set_data(matching.get(), state.data(), state.size(), 0) != state.size()) {
+ LOG_ERR("%s: failed to restore matching rotation\n", __func__);
+ success = false;
+ break;
+ }
+
+ auto mismatched = make_context(types.first, types.second, true);
+ if (!mismatched) {
+ LOG_ERR("%s: failed to create mismatched rotation context\n", __func__);
+ success = false;
+ break;
+ }
+ if (llama_state_seq_set_data(mismatched.get(), state.data(), state.size(), 0) != 0) {
+ LOG_TRC("%s: state restored into rotation-disabled context, model does not use attention rotation\n", __func__);
+ }
+ }
+ common_set_env("LLAMA_ATTN_ROT_DISABLE", attn_rot_disable);
+
+ if (!success) {
+ return false;
+ }
+
+ LOGV(LOG_LEVEL_INFO, "\nPASS\n");
+ return true;
+}
+
struct test_suite {
std::vector<test_status> results;
@@ -754,10 +840,10 @@ struct test_suite {
// column headers for the --models table, one per test, in the order they are run
static const std::vector<const char *> test_names = {
- "baseline", "seq_rm", "state_load", "cp_h", "cp_d", "cp_h_s", "cp_d_s", "rt", "rf",
+ "baseline", "seq_rm", "state_load", "cp_h", "cp_d", "cp_h_s", "cp_d_s", "rt", "rf", "rot",
};
-// Run the full save/load test suite (tests 1-9) for a single model.
+// Run the full save/load test suite (tests 1-10) for a single model.
// Returns the per-test results.
static test_suite run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
test_suite suite;
@@ -836,6 +922,9 @@ static test_suite run_save_load_tests_for_model(const std::string & model_path,
// Test 9: state restore failure
suite.results.push_back(test_state_restore_failure(model, params, tokens) ? test_status::PASS : test_status::FAIL);
+ // Test 10: state rotation
+ suite.results.push_back(test_state_rotation(model, params) ? test_status::PASS : test_status::FAIL);
+
return suite;
}