kanshan/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-07 19:50:21 +00:00
Code Issues 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));
src/llama-impl.h
+20
View File
@@ -54,6 +54,26 @@ static inline dst_t llama_cast(src_t v) {
}
}
static inline ggml_tensor * llama_mul_mat_hadamard(
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;
}
struct time_meas {
time_meas(int64_t & t_acc, bool disable = false);
~time_meas();
src/llama-kv-cache.cpp
+2 -18
View File
@@ -57,22 +57,6 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) {
}
}
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;
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;
}
//
// llama_kv_cache
//
@@ -1875,14 +1859,14 @@ ggml_tensor * llama_kv_cache::build_rope_shift(
tmp = ggml_cast(ctx, cur, GGML_TYPE_F32);
// rotate back
tmp = ggml_mul_mat_aux(ctx, tmp, rot);
tmp = llama_mul_mat_hadamard(ctx, tmp, rot);
tmp = ggml_rope_ext(ctx, tmp,
shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow);
// rotate fwd
tmp = ggml_mul_mat_aux(ctx, tmp, rot);
tmp = llama_mul_mat_hadamard(ctx, tmp, rot);
tmp = ggml_cpy(ctx, tmp, cur);
} else {
src/models/deepseek4.cpp
+37 -10
View File
@@ -557,7 +557,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k(
cb(indexer_q_pe, "lid_q_pe", il);
indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0);
indexer_q = ggml_mul_mat(ctx0, inp_lid.k_rot, indexer_q);
indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot);
cb(indexer_q, "lid_q_rot", il);
ggml_tensor * indexer_weights = build_lora_mm(layer.indexer_proj, cur);
@@ -652,10 +652,15 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention(
int il) const {
const auto & inp_csa = inp_dsv4->get_csa();
GGML_ASSERT(inp_csa.kq_mask);
GGML_ASSERT(inp_attn->self_k_rot == nullptr);
ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il);
ggml_tensor * k_rot = inp_attn->self_k_rot;
if (k_rot) {
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
}
ggml_build_forward_expand(gf, q);
ggml_build_forward_expand(gf, kv);
@@ -696,6 +701,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention(
ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]);
ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
cb(out, "attn_csa_lid", il);
return out;
@@ -711,7 +719,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention(
int il) const {
const auto & inp_hca = inp_dsv4->get_hca();
GGML_ASSERT(inp_hca.kq_mask);
GGML_ASSERT(inp_attn->self_k_rot == nullptr);
ggml_tensor * k_rot = inp_attn->self_k_rot;
if (k_rot) {
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
}
ggml_build_forward_expand(gf, q);
ggml_build_forward_expand(gf, kv);
@@ -753,6 +766,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention(
ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]);
ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
cb(out, "attn_hca", il);
return out;
@@ -770,8 +786,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention(
ggml_tensor * k_rot = inp_attn->self_k_rot;
if (k_rot) {
q = ggml_mul_mat(ctx0, k_rot, q);
kv = ggml_mul_mat(ctx0, k_rot, kv);
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
}
ggml_build_forward_expand(gf, q);
@@ -788,6 +804,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention(
ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]);
ggml_tensor * out = build_attn_mha(q, k, k, kq_b, kq_mask, sinks, nullptr, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
cb(out, "attn_raw", il);
return out;
@@ -917,6 +936,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
"csa_state_compress",
il);
if (inp_dsv4->get_csa().k_rot) {
kv_comp_csa_state = llama_mul_mat_hadamard(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().k_rot);
cb(kv_comp_csa_state, "csa_state_compress_rot", il);
}
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0,
kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il));
@@ -965,7 +989,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
il);
if (inp_dsv4->get_lid().k_rot) {
kv_comp_lid_state = ggml_mul_mat(ctx0, inp_dsv4->get_lid().k_rot, kv_comp_lid_state);
kv_comp_lid_state = llama_mul_mat_hadamard(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().k_rot);
cb(kv_comp_lid_state, "lid_state_compress_rot", il);
}
@@ -1007,6 +1031,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
"hca_state_compress",
il);
if (inp_dsv4->get_hca().k_rot) {
kv_comp_hca = llama_mul_mat_hadamard(ctx0, kv_comp_hca, inp_dsv4->get_hca().k_rot);
cb(kv_comp_hca, "hca_state_compress_rot", il);
}
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0,
kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il));
hca_state_dep = kv_comp_hca;
@@ -1035,13 +1064,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
if (ratio == DSV4_CSA_RATIO &&
inp_dsv4->get_csa().kq_mask &&
inp_dsv4->get_lid().kq_mask &&
inp_dsv4->get_lid().k_rot &&
inp_attn->self_k_rot == nullptr) {
inp_dsv4->get_lid().k_rot) {
out = build_csa_lid_attention(model, inp_dsv4, inp_attn, q, kv, qr, cur, inp_pos, layer.attn_sinks,
1.0f/sqrtf(float(n_embd_head)), il);
} else if (ratio == DSV4_HCA_RATIO &&
inp_dsv4->get_hca().kq_mask &&
inp_attn->self_k_rot == nullptr) {
inp_dsv4->get_hca().kq_mask) {
out = build_hca_attention(inp_dsv4, inp_attn, q, kv, layer.attn_sinks,
1.0f/sqrtf(float(n_embd_head)), il);
} else {