furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 02:05:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

GLM-DSA: add ability to use quantized indexer cache (#2075)

Browse Source 
* GLM-DSA: improve TG performance even more

* Fix crash when using MTP

* GLM-DSA: much better PP performance (CPU-only)

* GLM-DSA: allow for quantized indexer cache
This commit is contained in:
Kawrakow
2026-07-06 12:04:15 +02:00
committed by GitHub
parent eacd450d2b
commit a0859ce5ed
4 changed files with 17 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+7
View File
@@ -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)) {
common/common.h
+1
View File
@@ -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";
include/llama.h
+2
View File
@@ -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;
src/llama.cpp
+7 -3
View File
@@ -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));
}
}