Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

llama: fix quantized kv-cache for dsv4 (#25202)

Browse Source
This commit is contained in:
Aman Gupta
2026-07-07 17:46:57 +08:00
committed by GitHub
parent 108f186d17
commit 024c46ae4e
4 changed files with 77 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-graph.cpp
+18 -28
View File
@@ -63,26 +63,6 @@ static bool can_reuse_kq_mask(
// impl
static ggml_tensor * ggml_mul_mat_aux(
ggml_context * ctx,
ggml_tensor * cur,
ggml_tensor * rot) {
const auto n = rot->ne[0];
ggml_tensor * res;
if (!ggml_is_contiguous(cur)) {
res = ggml_cont_2d (ctx, cur, n, ggml_nelements(cur)/n);
} else {
res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n);
}
res = ggml_mul_mat (ctx, rot, res);
ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD);
res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]);
return res;
}
void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {
if (ubatch->token) {
const int64_t n_tokens = ubatch->n_tokens;
@@ -881,6 +861,14 @@ void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) {
dsv4_set_comp_inputs(inp_hca, plan_hca, "hca", debug > 0, ubatch->n_tokens, n_stream);
dsv4_set_comp_inputs(inp_lid, plan_lid, "lid", debug > 0, ubatch->n_tokens, n_stream);
if (inp_csa.k_rot && inp_csa.k_rot->buffer) {
mctx->get_csa()->set_input_k_rot(inp_csa.k_rot);
}
if (inp_hca.k_rot && inp_hca.k_rot->buffer) {
mctx->get_hca()->set_input_k_rot(inp_hca.k_rot);
}
if (inp_lid.k_rot && inp_lid.k_rot->buffer) {
mctx->get_lid()->set_input_k_rot(inp_lid.k_rot);
}
@@ -2633,12 +2621,12 @@ ggml_tensor * llm_graph_context::build_attn(
GGML_ASSERT(v_mla == nullptr);
if (inp->self_k_rot) {
q_cur = ggml_mul_mat_aux(ctx0, q_cur, inp->self_k_rot);
k_cur = ggml_mul_mat_aux(ctx0, k_cur, inp->self_k_rot);
q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot);
k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot);
}
if (inp->self_v_rot) {
v_cur = ggml_mul_mat_aux(ctx0, v_cur, inp->self_v_rot);
v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);
}
// these nodes are added to the graph together so that they are not reordered
@@ -2669,7 +2657,7 @@ ggml_tensor * llm_graph_context::build_attn(
cb(cur, "kqv_out", il);
if (inp->self_v_rot) {
cur = ggml_mul_mat_aux(ctx0, cur, inp->self_v_rot);
cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
}
if (wo) {
@@ -2874,14 +2862,14 @@ ggml_tensor * llm_graph_context::build_attn(
auto * v_rot = is_swa ? inp->self_v_rot_swa : inp->self_v_rot;
if (k_rot) {
q_cur = ggml_mul_mat_aux(ctx0, q_cur, k_rot);
q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot);
if (k_cur) {
k_cur = ggml_mul_mat_aux(ctx0, k_cur, k_rot);
k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot);
}
}
if (v_rot) {
if (v_cur) {
v_cur = ggml_mul_mat_aux(ctx0, v_cur, v_rot);
v_cur = llama_mul_mat_hadamard(ctx0, v_cur, v_rot);
}
}
@@ -2924,7 +2912,7 @@ ggml_tensor * llm_graph_context::build_attn(
cb(cur, "kqv_out", il);
if (v_rot) {
cur = ggml_mul_mat_aux(ctx0, cur, v_rot);
cur = llama_mul_mat_hadamard(ctx0, cur, v_rot);
}
if (wo) {
@@ -3084,6 +3072,8 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", n_stream);
dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", n_stream);
dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", n_stream);
inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0);
inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0);
inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0);
return (llm_graph_input_dsv4 *) res->add_input(std::move(inp));