From 48b02ccaa172630d3f0e501988cddcb85b332246 Mon Sep 17 00:00:00 2001 From: Iwan Kawrakow Date: Mon, 3 Nov 2025 14:37:03 +0200 Subject: [PATCH] Merge Q and K into a single tensor --- src/llama-build-context.cpp | 51 ++++++++++++++++++++++++++----------- src/llama-build-context.h | 1 + src/llama-load-tensors.cpp | 34 +++++++++++++++++++++++++ src/llama-model.h | 4 +++ 4 files changed, 75 insertions(+), 15 deletions(-) diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 24150b11..7474179b 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -1257,6 +1257,7 @@ std::tuple llm_build_context::llm_buil std::tuple llm_build_context::llm_build_mul_mat_qkv(ggml_cgraph * gf, ggml_tensor * cur, ggml_tensor * wqkv, ggml_tensor * bqkv, + ggml_tensor * wqk, ggml_tensor * bqk, ggml_tensor * wq, ggml_tensor * bq, ggml_tensor * wk, ggml_tensor * bk, ggml_tensor * wv, ggml_tensor * bv, @@ -1294,6 +1295,34 @@ std::tuple llm_build_context::llm_buil //ggml_build_forward_expand(gf, Vcur); } + if (wqk) { + auto qk = llm_build_lora_mm(lctx, ctx0, wqk, cur); + cb(qk, "qkv", il); + if (bqk) { + qk = ggml_add(ctx0, qk, bqk); + cb(qk, "qkv_b", il); + } + auto Qcur = ggml_view_3d(ctx0, qk, n_embd_head, n_head, n_tokens, n_embd_head*sizeof(float), qk->nb[1], 0*sizeof(float)*(n_embd)); + auto Kcur = ggml_view_3d(ctx0, qk, n_embd_head, n_head_kv, n_tokens, n_embd_head*sizeof(float), qk->nb[1], 1*sizeof(float)*Qcur->ne[0]*Qcur->ne[1]); + auto Vcur = llm_build_lora_mm(lctx, ctx0, wv, cur); + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + if (q_norm) { + Qcur = llm_build_norm(ctx0, Qcur, hparams, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, cb, il); + cb(Qcur, "Qcur_normed", il); + ggml_build_forward_expand(gf, Qcur); + } + if (k_norm) { + Kcur = llm_build_norm(ctx0, Kcur, hparams, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, cb, il); + cb(Kcur, "Kcur_normed", il); + ggml_build_forward_expand(gf, Kcur); + } + + return {Qcur, Kcur, Vcur}; + + } + auto [Q, K, V] = llm_build_mul_mat_qkv(gf, cur, wq, bq, wk, bk, wv, bv, attention_scale, il); auto Qcur = ggml_reshape_3d(ctx0, Q, n_embd_head, n_head, n_tokens); if (q_norm) { @@ -1361,6 +1390,7 @@ ggml_cgraph * llm_build_context::build_llama() { auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, model.layers[il].bqkv, + model.layers[il].wqk, model.layers[il].bqk, model.layers[il].wq, model.layers[il].bq, model.layers[il].wk, model.layers[il].bk, model.layers[il].wv, model.layers[il].bv, @@ -3383,6 +3413,7 @@ ggml_cgraph * llm_build_context::build_qwen3() { { auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, nullptr, + model.layers[il].wqk, nullptr, model.layers[il].wq, nullptr, model.layers[il].wk, nullptr, model.layers[il].wv, nullptr, @@ -3480,6 +3511,7 @@ ggml_cgraph * llm_build_context::build_qwen3moe() { { auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, nullptr, + model.layers[il].wqk, nullptr, model.layers[il].wq, nullptr, model.layers[il].wk, nullptr, model.layers[il].wv, nullptr, model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, 0, il); @@ -6097,6 +6129,7 @@ ggml_cgraph * llm_build_context::build_glm4_moe() { { auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, model.layers[il].bqkv, + model.layers[il].wqk, model.layers[il].bqk, model.layers[il].wq, model.layers[il].bq, model.layers[il].wk, model.layers[il].bk, model.layers[il].wv, model.layers[il].bv, @@ -6505,6 +6538,7 @@ ggml_cgraph * llm_build_context::build_cohere2() { auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, model.layers[il].bqkv, + model.layers[il].wqk, model.layers[il].bqk, model.layers[il].wq, model.layers[il].bq, model.layers[il].wk, model.layers[il].bk, model.layers[il].wv, model.layers[il].bv, nullptr, nullptr, 0.f, il); @@ -7800,6 +7834,7 @@ ggml_cgraph * llm_build_context::build_openai_moe() { { auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, model.layers[il].bqkv, + model.layers[il].wqk, model.layers[il].bqk, model.layers[il].wq, model.layers[il].bq, model.layers[il].wk, model.layers[il].bk, model.layers[il].wv, model.layers[il].bv, @@ -7815,20 +7850,6 @@ ggml_cgraph * llm_build_context::build_openai_moe() { attn_factor, beta_fast, beta_slow); cb(Kcur, "Kcur", il); - //auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wq, model.layers[il].bq, - // model.layers[il].wk, model.layers[il].bk, - // model.layers[il].wv, model.layers[il].bv, 0.f, il); - - //Qcur = ggml_rope_ext(ctx0, ggml_reshape_3d(ctx0, Qcur, n_rot, n_head, n_tokens), inp_pos, nullptr, - // n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, - // beta_fast, beta_slow); - //cb(Qcur, "Qcur", il); - - //Kcur = ggml_rope_ext(ctx0, ggml_reshape_3d(ctx0, Kcur, n_rot, n_head_kv, n_tokens), inp_pos, nullptr, - // n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, - // attn_factor, beta_fast, beta_slow); - //cb(Kcur, "Kcur", il); - cur = llm_build_kv(ctx0, lctx, kv_self, gf, model.layers[il].wo, model.layers[il].bo, Kcur, Vcur, Qcur, KQ_mask_l, n_tokens, kv_head, n_kv, kq_scale, cb, il, model.layers[il].attn_sinks, is_sliding ? hparams.n_swa : 0); @@ -7926,7 +7947,7 @@ ggml_cgraph * llm_build_context::build_bailingmoe2() { // self_attention { auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, model.layers[il].bqkv, - nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, 0.0f, il); //cur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wqkv, cur); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 150f3591..b79256b6 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -151,6 +151,7 @@ struct llm_build_context { std::tuple llm_build_mul_mat_qkv(ggml_cgraph * gf, ggml_tensor * cur, ggml_tensor * wqkv, ggml_tensor * bqkv, + ggml_tensor * wqk, ggml_tensor * bqk, ggml_tensor * wq, ggml_tensor * bq, ggml_tensor * wk, ggml_tensor * bk, ggml_tensor * wv, ggml_tensor * bv, diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index d1265616..44e948df 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -2476,6 +2476,40 @@ bool create_tensors_helper::merge_qkv(const LLM_TN & tn, int i, int bias) { } } } + if (!fused_qkv && ml.merge_qkv && wq->type == wk->type && hparams.f_attention_scale == 0.0f) { + GGML_ASSERT(wq->ne[0] == n_embd && wq->ne[1] == n_head * n_embd_head_k); + GGML_ASSERT(wk->ne[0] == n_embd && wk->ne[1] == n_embd_gqa); + layer.wqk = ggml_new_tensor_2d(ctx_split, wq->type, n_embd, n_embd_head_k * (n_head + n_head_kv)); + snprintf(layer.wqk->name, GGML_MAX_NAME, "blk.%d.attn_qk.weight", i); + layer.wq = ml.create_tensor_as_view(ctx_split, layer.wqk, wq_name.c_str(), { wq->ne[0], wq->ne[1] }, 0); + layer.wk = ml.create_tensor_as_view(ctx_split, layer.wqk, wk_name.c_str(), { wk->ne[0], wk->ne[1] }, wq->ne[1]*wq->nb[1]); + layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_gqa}); + printf("====================== Merged only Q and K in layer %d because V is of different type\n", i); + fused_qkv = true; + if (bias) { + auto bq_name = tn(LLM_TENSOR_ATTN_Q, "bias", i); + auto bk_name = tn(LLM_TENSOR_ATTN_K, "bias", i); + auto bv_name = tn(LLM_TENSOR_ATTN_V, "bias", i); + auto bq = ml.get_tensor_meta(bq_name.c_str()); + auto bk = ml.get_tensor_meta(bk_name.c_str()); + auto bv = ml.get_tensor_meta(bv_name.c_str()); + if (bias == 2) { + GGML_ASSERT(bq && bk && bv); + } else { + GGML_ASSERT(!bq && !bk && !bv); + } + if (bq && bk && bv) { + GGML_ASSERT(bq->type == GGML_TYPE_F32 && bk->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nrows(bq) == 1 && bq->ne[0] == wq->ne[1]); + GGML_ASSERT(ggml_nrows(bk) == 1 && bk->ne[0] == wk->ne[1]); + layer.bqk = ggml_new_tensor_1d(ctx_layer, bq->type, n_embd_head_k * (n_head + n_head_kv)); + snprintf(layer.bqk->name, GGML_MAX_NAME, "blk.%d.attn_qk.bias", i); + layer.bq = ml.create_tensor_as_view(ctx_layer, layer.bqk, bq_name.c_str(), { bq->ne[0] }, 0); + layer.bk = ml.create_tensor_as_view(ctx_layer, layer.bqk, bk_name.c_str(), { bk->ne[0] }, bq->ne[0]*bq->nb[0]); + layer.bv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_V, "bias", i), {layer.wv->ne[1]}); + } + } + } if (!fused_qkv) { if (ml.merge_qkv) { diff --git a/src/llama-model.h b/src/llama-model.h index 743b0d07..2d4b5610 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -124,6 +124,8 @@ struct llama_layer { struct ggml_tensor * wv = nullptr; struct ggml_tensor * wo = nullptr; struct ggml_tensor * wqkv = nullptr; + struct ggml_tensor * wqk = nullptr; + struct ggml_tensor * wkv = nullptr; struct ggml_tensor * wq_a = nullptr; struct ggml_tensor * wq_b = nullptr; struct ggml_tensor * wkv_a_mqa = nullptr; @@ -146,6 +148,8 @@ struct llama_layer { struct ggml_tensor * bv = nullptr; struct ggml_tensor * bo = nullptr; struct ggml_tensor * bqkv = nullptr; + struct ggml_tensor * bqk = nullptr; + struct ggml_tensor * bkv = nullptr; // relative position bias struct ggml_tensor * attn_rel_b = nullptr;