diff --git a/common/common.cpp b/common/common.cpp index fc6779fb3..b1fa0f12a 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1794,6 +1794,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa params.cache_type_v = argv[++i]; return true; } + if (arg == "-ictk" || arg == "--indexer-cache-type-k") { + params.indexer_cache_type_k = argv[++i]; + return true; + } if (arg == "-ctk-first" || arg == "--cache-type-k-first") { CHECK_ARG auto p = string_split(argv[i], ","); @@ -3184,6 +3188,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", "-dkvc, --dump-kv-cache", "verbose print of the KV cache" }); options.push_back({ "*", "-nkvo, --no-kv-offload", "disable KV offload" }); options.push_back({ "*", "-ctk, --cache-type-k TYPE", "KV cache data type for K (default: %s)", params.cache_type_k.c_str() }); + options.push_back({ "*", "-ictk, --indexer-cache-type-k TYPE", "indexer K-cache data type (default: %s)", params.indexer_cache_type_k.c_str() }); options.push_back({ "*", "-ctv, --cache-type-v TYPE", "KV cache data type for V (default: %s)", params.cache_type_v.c_str() }); options.push_back({ "*", "-ctk-first, --cache-type-k-first TYPE,N", "KV cache data type for the first N layers of K (default: %s,-1)", params.type_k_first.c_str() }); options.push_back({ "*", "-ctv-last, --cache-type-k-last TYPE,N", "KV cache data type for the last N layers of K (default: %s,-1)", params.type_k_last.c_str() }); @@ -4161,6 +4166,7 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params mparams.worst_graph_tokens = params.worst_graph_tokens; mparams.type_k = kv_cache_type_from_str(params.cache_type_k); mparams.type_v = kv_cache_type_from_str(params.cache_type_v); + mparams.idx_type_k = kv_cache_type_from_str(params.indexer_cache_type_k); mparams.type_k_first = kv_cache_type_from_str(params.type_k_first); mparams.type_k_last = kv_cache_type_from_str(params.type_k_last ); mparams.type_v_first = kv_cache_type_from_str(params.type_v_first); @@ -4283,6 +4289,7 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa cparams.type_k = kv_cache_type_from_str(params.cache_type_k); cparams.type_v = kv_cache_type_from_str(params.cache_type_v); + cparams.idx_type_k = kv_cache_type_from_str(params.indexer_cache_type_k); cparams.type_reduce = ggml_type_from_str(params.reduce_type); cparams.type_graph_attn = ggml_type_from_str(params.graph_attn_precision); if (!cparams.flash_attn && ggml_is_quantized(cparams.type_v)) { diff --git a/common/common.h b/common/common.h index 21ef179b3..24a041377 100644 --- a/common/common.h +++ b/common/common.h @@ -452,6 +452,7 @@ struct gpt_params { std::string cache_type_k = "f16"; // KV cache data type for the K std::string cache_type_v = "f16"; // KV cache data type for the V + std::string indexer_cache_type_k = "f16"; // indexer K-cache data type std::string reduce_type = "f16"; std::string graph_attn_precision = "f16"; diff --git a/include/llama.h b/include/llama.h index ff6915f1d..fc4e5d470 100644 --- a/include/llama.h +++ b/include/llama.h @@ -378,6 +378,7 @@ extern "C" { enum ggml_type type_k; enum ggml_type type_v; + enum ggml_type idx_type_k; uint32_t max_ctx_size; int32_t n_seq_max; int32_t n_ubatch; @@ -465,6 +466,7 @@ extern "C" { enum ggml_type type_k; // data type for K cache [EXPERIMENTAL] enum ggml_type type_v; // data type for V cache [EXPERIMENTAL] + enum ggml_type idx_type_k; // data type for indexer K cache [EXPERIMENTAL] enum ggml_type type_reduce; // data type for reduce operations enum ggml_type type_graph_attn; // flash-attn precision under -sm graph enum ggml_type type_k_first; diff --git a/src/llama.cpp b/src/llama.cpp index de6a5f2ea..63702d3cf 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -840,6 +840,7 @@ static bool llama_kv_cache_init( const llama_context * ctx, ggml_type type_k, ggml_type type_v, + ggml_type idx_type_k, uint32_t kv_size, bool offload, ggml_type type_k_first, @@ -1030,7 +1031,7 @@ static bool llama_kv_cache_init( // DSA lightning-indexer key cache (MQA, single head). Store the Hadamard-rotated // indexer keys in F16 so a decoded token can score against ALL past keys. if (has_dsa_indexer && model.layers[i].indexer_attn_k && !is_mtp_tail_layer) { - ggml_tensor * kr = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hparams.indexer_head_size, kv_size); + ggml_tensor * kr = ggml_new_tensor_2d(ctx, idx_type_k, hparams.indexer_head_size, kv_size); ggml_format_name(kr, "cache_kr_l%d", i); cache.kr_l[i] = kr; } @@ -6543,6 +6544,7 @@ struct llama_model_params llama_model_default_params() { /*.ncmoe =*/ 0, /*.type_k =*/ GGML_TYPE_F16, /*.type_v =*/ GGML_TYPE_F16, + /*.idx_type_k =*/ GGML_TYPE_F16, /*.max_ctx_size =*/ 0, /*.n_seq_max =*/ 1, /*.n_ubatch =*/ 512, @@ -6615,6 +6617,7 @@ struct llama_context_params llama_context_default_params() { /*.cb_eval_user_data =*/ nullptr, /*.type_k =*/ GGML_TYPE_F16, /*.type_v =*/ GGML_TYPE_F16, + /*.idx_type_k =*/ GGML_TYPE_F16, /*.type_reduce =*/ GGML_TYPE_F16, /*.type_graph_attn =*/ GGML_TYPE_F16, /*.type_first_k =*/ GGML_TYPE_F16, @@ -7411,7 +7414,7 @@ struct llama_context * llama_init_from_model( } ctx->backends.push_back(ctx->backend_cpu); - if (!llama_kv_cache_init(ctx->kv_self, ctx, type_k, type_v, kv_size, cparams.offload_kqv, + if (!llama_kv_cache_init(ctx->kv_self, ctx, type_k, type_v, params.idx_type_k, kv_size, cparams.offload_kqv, params.type_k_first, params.type_k_last, params.type_v_first, params.type_v_last, params.n_k_first, params.n_k_last, params.n_v_first, params.n_v_last)) { LLAMA_LOG_ERROR("%s: llama_kv_cache_init() failed for self-attention cache\n", __func__); @@ -7459,7 +7462,8 @@ struct llama_context * llama_init_from_model( } } if (memory_size_k_indexer > 0) { - LLAMA_LOG_INFO("%s: KV self indexer size = %7.2f MiB\n", __func__, memory_size_k_indexer/1024./1024.); + LLAMA_LOG_INFO("%s: KV self indexer size = %7.2f MiB (%s)\n", __func__, memory_size_k_indexer/1024./1024., + ggml_type_name(params.idx_type_k)); } }