Mixed KV cache (#1599)

This commit is contained in:
Kawrakow
2026-04-09 09:33:17 +02:00
committed by GitHub
parent 5950d0259e
commit 9db5d9907e
4 changed files with 148 additions and 7 deletions
+70
View File
@@ -1260,6 +1260,50 @@ 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 == "-ctk-first" || arg == "--cache-type-k-first") {
CHECK_ARG
auto p = string_split(argv[i], ",");
if (p.size() != 2) {
invalid_param = true;
} else {
params.type_k_first = p[0];
params.n_k_first = std::stoi(p[1].c_str());
}
return true;
}
if (arg == "-ctk-last" || arg == "--cache-type-k-last") {
CHECK_ARG
auto p = string_split(argv[i], ",");
if (p.size() != 2) {
invalid_param = true;
} else {
params.type_k_last = p[0];
params.n_k_last = std::stoi(p[1].c_str());
}
return true;
}
if (arg == "-ctv-first" || arg == "--cache-type-v-first") {
CHECK_ARG
auto p = string_split(argv[i], ",");
if (p.size() != 2) {
invalid_param = true;
} else {
params.type_v_first = p[0];
params.n_v_first = std::stoi(p[1].c_str());
}
return true;
}
if (arg == "-ctv-last" || arg == "--cache-type-v-last") {
CHECK_ARG
auto p = string_split(argv[i], ",");
if (p.size() != 2) {
invalid_param = true;
} else {
params.type_v_last = p[0];
params.n_v_last = std::stoi(p[1].c_str());
}
return true;
}
if (arg == "-ctkd" || arg == "--cache-type-k-draft") {
params.speculative.cache_type_k = argv[++i];
return true;
@@ -2470,6 +2514,10 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
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({ "*", "-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() });
options.push_back({ "*", "-ctv-first, --cache-type-v-first TYPE,N", "KV cache data type for the first N layers of V (default: %s,-1)", params.type_v_first.c_str() });
options.push_back({ "*", "-ctk-last, --cache-type-v-last TYPE,N", "KV cache data type for the last N layers of V (default: %s,-1)", params.type_v_last.c_str() });
options.push_back({ "*", "-ctkd, --cache-type-k-draft TYPE", "KV cache data type for K for the draft model" });
options.push_back({ "*", "-ctvd, --cache-type-v-draft TYPE", "KV cache data type for V for the draft model" });
@@ -3357,6 +3405,14 @@ 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.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);
mparams.type_v_last = kv_cache_type_from_str(params.type_v_last );
mparams.n_k_first = params.n_k_first;
mparams.n_k_last = params.n_k_last;
mparams.n_v_first = params.n_v_first;
mparams.n_v_last = params.n_v_last;
mparams.max_ctx_size = params.n_ctx;
mparams.n_seq_max = params.n_parallel;
mparams.n_ubatch = get_batch_ubatch(params).second;
@@ -3464,6 +3520,20 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa
if (!cparams.flash_attn && ggml_is_quantized(cparams.type_v)) {
throw std::runtime_error("Quantized V cache cannot be used without flash attention");
}
cparams.type_k_first = kv_cache_type_from_str(params.type_k_first);
cparams.type_k_last = kv_cache_type_from_str(params.type_k_last );
cparams.type_v_first = kv_cache_type_from_str(params.type_v_first);
cparams.type_v_last = kv_cache_type_from_str(params.type_v_last );
cparams.n_k_first = params.n_k_first;
cparams.n_k_last = params.n_k_last;
cparams.n_v_first = params.n_v_first;
cparams.n_v_last = params.n_v_last;
if (!cparams.flash_attn && ggml_is_quantized(cparams.type_v_first) && cparams.n_v_first > 0) {
throw std::runtime_error("Quantized V cache cannot be used without flash attention");
}
if (!cparams.flash_attn && ggml_is_quantized(cparams.type_v_last) && cparams.n_v_last > 0) {
throw std::runtime_error("Quantized V cache cannot be used without flash attention");
}
if (!params.offload_policy.empty()) cparams.offload_policy = (void *)&params.offload_policy;
if (!params.cuda_params.empty()) cparams.cuda_params = (void *)params.cuda_params.data();