V-cache Hadamard transform (#1527)

This commit is contained in:
Kawrakow 2026-03-28 11:28:50 +01:00 committed by GitHub
parent 798af8676a
commit a95981013a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 41 additions and 1 deletions

View File

@ -499,6 +499,7 @@ void gpt_params_parse_from_env(gpt_params & params) {
get_env("LLAMA_ARG_CACHE_TYPE_V", params.cache_type_v);
get_env("LLAMA_ARG_MLOCK", params.use_mlock);
get_env("LLAMA_ARG_K_CACHE_HADAMARD", params.k_cache_hadamard);
get_env("LLAMA_ARG_V_CACHE_HADAMARD", params.v_cache_hadamard);
}
@ -1545,6 +1546,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
params.k_cache_hadamard = true;
return true;
}
if (arg == "-vhad" || arg == "--v-cache-hadamard") {
params.v_cache_hadamard = true;
return true;
}
if (arg == "-smgs" || arg == "--split-mode-graph-scheduling") {
params.split_mode_graph_scheduling = true;
return true;
@ -2313,6 +2318,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
options.push_back({ "*", "-mqkv, --merge-qkv,", "merge Q,K,V (default: %d)", params.merge_qkv});
options.push_back({ "*", "-muge, --merge-up-gate-experts,","merge ffn_up/gate_exps (default: %d)", params.merge_up_gate_exps});
options.push_back({ "*", "-khad, --k-cache-hadamard,", "Use Hadamard transform for K-cache (default: %d)", params.k_cache_hadamard});
options.push_back({ "*", "-vhad, --v-cache-hadamard,", "Use Hadamard transform for V-cache (default: %d)", params.v_cache_hadamard});
options.push_back({ "*", "-smf16, --split-mode-f16,", "Use f16 for data exchange between GPUs (default: %d)", true});
options.push_back({ "*", "-smf32, --split-mode-f32,", "Use f32 for data exchange between GPUs (default: %d)", false});
options.push_back({ "*", "-grt, --graph-reduce-type", "Type for data exchange between GPUs (default: %s)", "f32"});
@ -3422,6 +3428,7 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa
cparams.rope_cache = params.rope_cache;
cparams.graph_reuse = params.graph_reuse;
cparams.k_cache_hadamard = params.k_cache_hadamard;
cparams.v_cache_hadamard = params.v_cache_hadamard;
cparams.split_mode_graph_scheduling = params.split_mode_graph_scheduling;
//cparams.split_mode_f16 = params.split_mode_f16;
cparams.scheduler_async = params.scheduler_async;
@ -4437,6 +4444,7 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l
fprintf(stream, "rope_cache: %s # default: false\n", params.rope_cache ? "true" : "false");
fprintf(stream, "graph_reuse: %s # default: false\n", params.graph_reuse ? "true" : "false");
fprintf(stream, "k_cache_hadamard: %s # default: false\n", params.k_cache_hadamard ? "true" : "false");
fprintf(stream, "v_cache_hadamard: %s # default: false\n", params.v_cache_hadamard ? "true" : "false");
fprintf(stream, "split_mode_graph_scheduling: %s # default: false\n", params.split_mode_graph_scheduling ? "true" : "false");
//fprintf(stream, "split_mode_f16: %s # default: true\n", params.split_mode_f16 ? "true" : "false");
fprintf(stream, "reduce_type: %s # default f16\n", params.reduce_type.c_str());

View File

@ -358,6 +358,7 @@ struct gpt_params {
bool merge_qkv = false; // if true, merge separate Q, K, V tensors into a single, contiguous tensor
bool merge_up_gate_exps= false; // if true, merge ffn_up_exps and ffn_gate_exps into a single, contiguous tensor
bool k_cache_hadamard = false; // if true, use Hadamard transform for the K-cache (only makes sense with quantized cache)
bool v_cache_hadamard = false; // if true, use Hadamard transform for the V-cache (only makes sense with quantized cache, which requires FA)
bool split_mode_graph_scheduling = false; // if true, force split mode graph scheduling
//bool split_mode_f16 = true; // if true, intermediate results will be cast to f16 before copying to other GPUs to perform reduce ops
bool scheduler_async = false; // if true, in split mode graph the scheduler will use multiple threads to evaluate the graph

View File

@ -465,6 +465,7 @@ extern "C" {
float thresh_experts;
bool only_active_experts;
bool k_cache_hadamard; // if true, apply Hadamard transfrom to K-cache
bool v_cache_hadamard; // if true, apply Hadamard transfrom to V-cache (needs FA)
bool split_mode_graph_scheduling; // if true, force split mode graph scheduling
//bool split_mode_f16; // if true, cast intermediate results to f16 before copying to other GPUs
bool scheduler_async; // if true, with split mode "graph" graph evaluation will be done using multiple threads

View File

@ -1581,6 +1581,10 @@ static ggml_tensor * llm_build_kqv(
}
//ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32);
if (cparams.v_cache_hadamard) {
cur = ggml_hadamard(ctx, cur, n_embd_head_v);
cb(cur, "fa_h", il);
}
cur = ggml_reshape_2d(ctx, cur, n_embd_head_v*n_head, n_tokens);
} else {
@ -1744,6 +1748,9 @@ ggml_tensor * llm_build_context::llm_build_kv(
cb(q_cur, "Qcur_hadamard", il);
cb(k_cur, "Kcur_hadamard", il);
}
if (cparams.v_cache_hadamard) {
v_cur = ggml_hadamard(ctx, v_cur, hparams.n_embd_head_v);
}
// these nodes are added to the graph together so that they are not reordered
// by doing so, the number of splits in the graph is reduced
@ -9322,6 +9329,11 @@ ggml_cgraph* llm_build_context::build_minimaxm2() {
}
ggml_build_forward_expand(gf, Qcur);
ggml_build_forward_expand(gf, Kcur);
if (cparams.v_cache_hadamard) {
Vcur = ggml_hadamard(ctx0, Vcur, hparams.n_embd_head_v);
cb(Vcur, "Vcur_hadamard", il_id);
ggml_build_forward_expand(gf, Vcur);
}
// Store K, V in KV cache
auto idx = 2*wq->n_device*il + 2*id;
@ -10093,6 +10105,10 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
cb(Qcur, "Qcur_hadamard", il_cb);
cb(Kcur, "Kcur_hadamard", il_cb);
}
if (cparams.v_cache_hadamard) {
Vcur = ggml_hadamard(ctx0, Vcur, hparams.n_embd_head_v);
cb(Vcur, "Vcur_hadamard", il_cb);
}
ggml_build_forward_expand(gf, Qcur);
ggml_build_forward_expand(gf, Kcur);
ggml_build_forward_expand(gf, Vcur);
@ -10166,6 +10182,11 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32);
}
if (cparams.v_cache_hadamard) {
cur = ggml_hadamard(ctx0, cur, n_embd_head_v);
cb(cur, "flash_attn_h", il_cb);
}
if (model.layers[il].wqkv_gate) {
auto wqkv_gate = (ggml_split_tensor_t *)model.layers[il].wqkv_gate->extra;
GGML_ASSERT(wqkv_gate && wqkv_gate->splits[id]);

View File

@ -40,6 +40,7 @@ struct llama_cparams {
bool rope_cache;
bool graph_reuse;
bool k_cache_hadamard;
bool v_cache_hadamard;
bool split_mode_graph_scheduling;
//bool split_mode_f16;
bool scheduler_async;

View File

@ -4964,6 +4964,7 @@ struct llama_context_params llama_context_default_params() {
/*.thtesh_experts =*/ 0.0f,
/*.only_active_experts =*/ false,
/*.k_cache_hadamard =*/ false,
/*.v_cache_hadamard =*/ false,
/*.split_mode_graph_scheduling =*/ false,
// /*.split_mode_f16 =*/ true,
/*.scheduler_async =*/ false,
@ -5300,10 +5301,15 @@ struct llama_context * llama_init_from_model(
}
if (params.k_cache_hadamard && !ggml_is_quantized(params.type_k)) {
LLAMA_LOG_WARN("%s: there is no point in Hadamard transforms with not quantized K-cache. Turning Hadamard off\n", __func__);
LLAMA_LOG_WARN("%s: there is no point in Hadamard transforms with not quantized K-cache. Turning K-cache Hadamard off\n", __func__);
params.k_cache_hadamard = false;
}
if (params.v_cache_hadamard && !ggml_is_quantized(params.type_v)) {
LLAMA_LOG_WARN("%s: there is no point in Hadamard transforms with not quantized V-cache. Turning V-cache Hadamard off\n", __func__);
params.v_cache_hadamard = false;
}
llama_context * ctx = new llama_context(*model);
// add devices to ctx->cparams from model
@ -5338,6 +5344,7 @@ struct llama_context * llama_init_from_model(
cparams.rope_cache = params.rope_cache;
cparams.graph_reuse = params.graph_reuse;
cparams.k_cache_hadamard = params.k_cache_hadamard;
cparams.v_cache_hadamard = params.v_cache_hadamard;
cparams.split_mode_graph_scheduling = params.split_mode_graph_scheduling;
//cparams.split_mode_f16 = params.split_mode_f16;
cparams.scheduler_async = params.scheduler_async;
@ -5443,6 +5450,7 @@ struct llama_context * llama_init_from_model(
LLAMA_LOG_INFO("%s: rope_cache = %d\n", __func__, cparams.rope_cache);
LLAMA_LOG_INFO("%s: graph_reuse = %d\n", __func__, cparams.graph_reuse);
LLAMA_LOG_INFO("%s: k_cache_hadam = %d\n", __func__, cparams.k_cache_hadamard);
LLAMA_LOG_INFO("%s: v_cache_hadam = %d\n", __func__, cparams.v_cache_hadamard);
LLAMA_LOG_INFO("%s: split_mode_graph_scheduling = %d\n", __func__, cparams.split_mode_graph_scheduling);
//LLAMA_LOG_INFO("%s: split_mode_f16= %d\n", __func__, cparams.split_mode_f16);
LLAMA_LOG_INFO("%s: reduce_type = %s\n", __func__, ggml_type_name(cparams.reduce_type));