upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-28 17:07:31 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add FP8 KV-cache support

Browse Source
This commit is contained in:
Oliver Simons
2026-09-04 20:51:13 +02:00
parent cef6c528f8
commit b6ccff9b0e
10 changed files with 146 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-graph.cpp
+54 -6
View File
@@ -2788,6 +2788,49 @@ llm_graph_input_attn_kv * llm_graph_context::build_attn_inp_kv() const {
return (llm_graph_input_attn_kv *) res->add_input(std::move(inp));
}
static void build_attn_scale_fp8_inputs(
ggml_context * ctx,
const llama_kv_cache_context * mctx,
ggml_tensor * & q,
ggml_tensor * & k,
ggml_tensor * & v,
int32_t il) {
if (mctx->type_k() == GGML_TYPE_F8_E4M3) {
// Store K/s and move its dequant scale to Q: (Q*s) dot (K/s) = Q dot K.
ggml_tensor * scale = mctx->get_k_scale(il);
if (scale) {
q = ggml_mul(ctx, q, scale);
if (k) {
k = ggml_div(ctx, k, scale);
}
}
}
if (mctx->type_v() == GGML_TYPE_F8_E4M3 && v) {
// Store V/s here and restore its dequant scale after the weighted sum.
ggml_tensor * scale = mctx->get_v_scale(il);
if (scale) {
v = ggml_div(ctx, v, scale);
}
}
}
static ggml_tensor * build_attn_scale_fp8_output(
ggml_context * ctx,
const llama_kv_cache_context * mctx,
ggml_tensor * cur,
int32_t il) {
if (mctx->type_v() == GGML_TYPE_F8_E4M3) {
// P*(V/s)*s = P*V.
ggml_tensor * scale = mctx->get_v_scale(il);
if (scale) {
cur = ggml_mul(ctx, cur, scale);
}
}
return cur;
}
ggml_tensor * llm_graph_context::build_attn(
llm_graph_input_attn_kv * inp,
ggml_tensor * wo,
@@ -2812,6 +2855,10 @@ ggml_tensor * llm_graph_context::build_attn(
v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);
}
const auto * mctx_cur = inp->mctx;
build_attn_scale_fp8_inputs(ctx0, mctx_cur, q_cur, k_cur, v_cur, il);
// 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
// expand k later to enable rope fusion which directly writes into k-v cache
@@ -2819,8 +2866,6 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, v_cur);
ggml_build_forward_expand(gf, k_cur);
const auto * mctx_cur = inp->mctx;
// store to KV cache
{
const auto & k_idxs = inp->get_k_idxs();
@@ -2837,6 +2882,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cur = build_attn_scale_fp8_output(ctx0, mctx_cur, cur, il);
cb(cur, "kqv_out", il);
if (inp->self_v_rot) {
@@ -3056,6 +3102,11 @@ ggml_tensor * llm_graph_context::build_attn(
}
}
const auto * mctx_iswa = inp->mctx;
const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
build_attn_scale_fp8_inputs(ctx0, mctx_cur, q_cur, k_cur, v_cur, il);
// 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
ggml_build_forward_expand(gf, q_cur);
@@ -3068,10 +3119,6 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, v_cur);
}
const auto * mctx_iswa = inp->mctx;
const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
// optionally store to KV cache
if (k_cur) {
const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();
@@ -3092,6 +3139,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cur = build_attn_scale_fp8_output(ctx0, mctx_cur, cur, il);
cb(cur, "kqv_out", il);
if (v_rot) {