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(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> 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 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 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 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; }