Commit 202c1a229 for llama.cpp
commit 202c1a229f5a83d1c9d758ed0f2d43760fa4bab0
Author: Oliver Simons <osimons@nvidia.com>
Date: Tue Nov 10 17:08:32 2026 +0100
llama : Add FP8 KV-cache support (#28392)
* Add FP8 KV-cache support
* Add tests for scaled, FP8 KV-cache to test-backend-ops
* Fix get_*_scale mapping and dequnatize FP8 also before rope shift
* Normalizes cache-scale names independently of projection layout.
While fused KQV is currently impractival for non-shared scales, this is
an improvement for future use regardless
* Enable Gemma 4 Assistant with FP8 by storing and fetching scale ptrs
* Update comment wording
* Move E2E FP8 KV cache tests to test-llama-archs
They belong there as they are not just a test for ggml's backend ops
* Increase support surface, remove seed
* make scaled kv generic
* Fix CI failings for test-kv-cache
https://github.com/ggml-org/llama.cpp/actions/runs/38718763447/job/116367842578?pr=28392
* CI: execute kv-cache-scale tests
* exclude fp8-kv tests from OpenVINO
* Fix linting errors of test_kv_cache_scales
* Disable kv-ffload=false + multi-GPU tests for FP8
See https://github.com/ggml-org/llama.cpp/actions/runs/38952722084/job/117148353699?pr=28413
* Adjust test label for test-kv-cache.cpp
* Fix conflict-resolution mistakes in test-kv-cache.cpp
https://github.com/ggml-org/llama.cpp/actions/runs/40554674463/job/122367426461?pr=28392
* Move KV tests to test-fp8.cpp
https://github.com/ggml-org/llama.cpp/pull/27512#discussion_r4410518158
* Remove tmp file for test-kv-cache-scales.py on Windows as well
* Fix fp8 kv cache test by population kv_source_impl
* Exclude models that consume KV cache directly
* Remove obsolete comment
* Allow BF16 scales as well for KV cache
* Rename *_cache_scale to *_cache_s
https://github.com/ggml-org/llama.cpp/pull/28392#discussion_r4517194999
* Move `is_type_scaled()` to `llama-kv-cache.cpp`
diff --git a/.github/workflows/build-openvino.yml b/.github/workflows/build-openvino.yml
index f4c13b611..009bde4e4 100644
--- a/.github/workflows/build-openvino.yml
+++ b/.github/workflows/build-openvino.yml
@@ -37,7 +37,7 @@ env:
LLAMA_ARG_LOG_PREFIX: 1
LLAMA_ARG_LOG_TIMESTAMPS: 1
# TODO: fix failing tests on OpenVINO backend
- CTEST_EXCLUDE: "test-llama-archs|test-llama-context"
+ CTEST_EXCLUDE: "test-llama-archs|test-llama-context|test-fp8-kv"
jobs:
ubuntu-24-openvino:
diff --git a/ci/run.sh b/ci/run.sh
index 346b116c7..1beb54b07 100755
--- a/ci/run.sh
+++ b/ci/run.sh
@@ -180,7 +180,7 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"
# TODO: fix failing tests on OpenVINO backend
- CTEST_EXTRA="-E test-llama-archs|test-llama-context"
+ CTEST_EXTRA="-E test-llama-archs|test-llama-context|test-fp8-kv"
fi
## helpers
diff --git a/common/arg.cpp b/common/arg.cpp
index 3a49d86bc..4b716398e 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -282,6 +282,7 @@ const std::vector<ggml_type> kv_cache_types = {
GGML_TYPE_F32,
GGML_TYPE_F16,
GGML_TYPE_BF16,
+ GGML_TYPE_F8_E4M3,
GGML_TYPE_Q8_0,
GGML_TYPE_Q4_0,
GGML_TYPE_Q4_1,
diff --git a/conversion/base.py b/conversion/base.py
index 2cfdba30b..6dee2e878 100644
--- a/conversion/base.py
+++ b/conversion/base.py
@@ -7,6 +7,7 @@ import ast
import logging
import contextlib
import json
+import math
import os
import re
import sys
@@ -343,6 +344,36 @@ class ModelBase:
logger.info(f" + {scale_name} (per-expert scale, shape [{len(scales)}])")
self.gguf_writer.add_tensor(scale_name, scale_vals)
+ def _prepare_kv_cache_scales(self):
+ for name in list(self.model_tensors.keys()):
+ if not name.endswith((".k_scale", ".v_scale")):
+ continue
+
+ mapped = self.tensor_map.get_type_and_name(key=name, try_suffixes=(".k_scale", ".v_scale"))
+ if mapped is None:
+ continue
+
+ tensor_type, new_name = mapped
+ if tensor_type not in (gguf.MODEL_TENSOR.ATTN_K, gguf.MODEL_TENSOR.ATTN_V, gguf.MODEL_TENSOR.ATTN_QKV):
+ continue
+
+ bid = int(new_name.split(".")[1])
+ scale_type = gguf.MODEL_TENSOR.ATTN_K if name.endswith(".k_scale") else gguf.MODEL_TENSOR.ATTN_V
+ suffix = ".k_scale" if name.endswith(".k_scale") else ".v_scale"
+ new_name = gguf.TENSOR_NAMES[scale_type].format(bid=bid) + suffix
+
+ scale = LazyTorchTensor.to_eager(self.model_tensors.pop(name)())
+ if scale.dtype not in (torch.bfloat16, torch.float32) or scale.numel() != 1:
+ raise ValueError(f"KV cache scale {name!r} must be a scalar BF16 or FP32 tensor")
+
+ scale = scale.float()
+ value = float(scale.item())
+ if not math.isfinite(value) or value == 0.0:
+ raise ValueError(f"KV cache scale {name!r} must be finite and nonzero")
+
+ logger.info(f" + {new_name} (KV cache scale, shape [1])")
+ self.gguf_writer.add_tensor(new_name, scale.flatten().numpy())
+
def dequant_model(self):
# If all quantized tensors were already handled (e.g. pure NVFP4), skip
if self._is_nvfp4 and not any(k.endswith((".weight_scale", ".weight_scale_inv", ".input_scale", ".activation_scale", "_activation_scale", ".k_scale", ".v_scale")) for k in self.model_tensors):
@@ -1455,6 +1486,8 @@ class ModelBase:
unmatched = layer_algo.keys() - matched
assert not unmatched, f"quantized_layers entries not in the model tensors: {sorted(unmatched)}"
+ self._prepare_kv_cache_scales()
+
# NVFP4 weights are repacked and written directly to gguf_writer.
# This must run before dequant_model so NVFP4 tensors are removed
# from model_tensors, leaving only non-NVFP4 (e.g. FP8) for dequant.
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 5bb886454..670adcc02 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -4531,6 +4531,45 @@ llama_context * llama_init_from_model(
return nullptr;
}
+ const bool fp8_cache = params.type_k == GGML_TYPE_F8_E4M3 || params.type_v == GGML_TYPE_F8_E4M3;
+ // Direct cache consumers do not use the shared attention scaling helpers.
+ if (fp8_cache && (model->hparams.is_mla() ||
+ model->arch == LLM_ARCH_DEEPSEEK32 ||
+ model->arch == LLM_ARCH_DEEPSEEK4 ||
+ model->arch == LLM_ARCH_DEEPSEEK41 ||
+ model->arch == LLM_ARCH_GLM_DSA ||
+ model->arch == LLM_ARCH_GLM5_NEXT ||
+ model->arch == LLM_ARCH_DOTS3NOTE ||
+ model->arch == LLM_ARCH_MISTRAL_LARGE4 ||
+ model->arch == LLM_ARCH_HY_V4 ||
+ model->arch == LLM_ARCH_MINIMAX_M3 ||
+ model->arch == LLM_ARCH_QWEN4EXP ||
+ model->arch == LLM_ARCH_DFLASH)) {
+ LLAMA_LOG_ERROR("%s: FP8 cache is not supported for architecture '%s'\n", __func__, llm_arch_name(model->arch));
+ return nullptr;
+ }
+
+ if (fp8_cache) {
+ for (uint32_t il = 0; il < model->hparams.n_layer_all; ++il) {
+ if (!model->hparams.has_kv(il)) {
+ continue;
+ }
+
+ for (ggml_tensor * scale : { model->layers[il].k_cache_s, model->layers[il].v_cache_s }) {
+ if (!scale || !scale->data) {
+ continue;
+ }
+
+ float value;
+ ggml_backend_tensor_get(scale, &value, 0, sizeof(value));
+ if (!std::isfinite(value) || value == 0.0f) {
+ LLAMA_LOG_ERROR("%s: FP8 cache scale '%s' must be finite and nonzero\n", __func__, ggml_get_name(scale));
+ return nullptr;
+ }
+ }
+ }
+ }
+
if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__);
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 6ba773464..847a3e1ec 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -3315,6 +3315,43 @@ 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_kv_inputs(
+ ggml_context * ctx,
+ const llama_kv_cache_context * mctx,
+ ggml_tensor * & q,
+ ggml_tensor * & k,
+ ggml_tensor * & v,
+ int32_t il) {
+ if (ggml_tensor * scale = mctx->get_k_scale(il)) {
+ // Store K/s and move its dequant scale to Q: (Q*s) dot (K/s) = Q dot K.
+ q = ggml_mul(ctx, q, scale);
+ if (k) {
+ k = ggml_div(ctx, k, scale);
+ }
+ }
+
+ if (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_kv_output(
+ ggml_context * ctx,
+ const llama_kv_cache_context * mctx,
+ ggml_tensor * cur,
+ int32_t il) {
+ if (ggml_tensor * scale = mctx->get_v_scale(il)) {
+ // P*(V/s)*s = P*V.
+ cur = ggml_mul(ctx, cur, scale);
+ }
+
+ return cur;
+}
+
ggml_tensor * llm_graph_context::build_attn(
llm_graph_input_attn_kv * inp,
ggml_tensor * wo,
@@ -3341,6 +3378,10 @@ ggml_tensor * llm_graph_context::build_attn(
v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot[il]);
}
+ const auto * mctx_cur = inp->mctx;
+
+ build_attn_scale_kv_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
@@ -3348,8 +3389,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;
-
ggml_tensor * q = q_cur;
ggml_tensor * k;
ggml_tensor * v;
@@ -3376,6 +3415,7 @@ ggml_tensor * llm_graph_context::build_attn(
const auto & kv_rows = inp->self_attn_mask.rows;
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
+ cur = build_attn_scale_kv_output(ctx0, mctx_cur, cur, il);
cb(cur, "kqv_out", il);
if (inp->self_v_rot[il]) {
@@ -3635,6 +3675,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_kv_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);
@@ -3647,10 +3692,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();
-
// whole seq fits into batch in training mode
const bool use_kv_cur = cparams.training && k_cur && v_cur;
if (use_kv_cur) {
@@ -3680,6 +3721,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * v = use_kv_cur ? v_cur : mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
+ cur = build_attn_scale_kv_output(ctx0, mctx_cur, cur, il);
cb(cur, "kqv_out", il);
if (v_rot) {
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index a1ec67e26..4ec7fb89c 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -246,7 +246,7 @@ llama_kv_cache::llama_kv_cache(
map_layer_ids[il] = layers.size();
- layers.push_back({ il, k, v, k_stream, v_stream, });
+ layers.push_back({ il, k, v, model.layers[il].k_cache_s, model.layers[il].v_cache_s, k_stream, v_stream, });
}
if (reuse) {
@@ -1470,6 +1470,30 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k
ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
}
+static bool is_type_scaled(ggml_type type) {
+ return type == GGML_TYPE_F8_E4M3;
+}
+
+ggml_tensor * llama_kv_cache::get_k_scale(int32_t il) const {
+ if (!is_type_scaled(type_k())) {
+ return nullptr;
+ }
+
+ const int32_t ikv = map_layer_ids.at(il);
+
+ return layers[ikv].k_scale;
+}
+
+ggml_tensor * llama_kv_cache::get_v_scale(int32_t il) const {
+ if (!is_type_scaled(type_v())) {
+ return nullptr;
+ }
+
+ const int32_t ikv = map_layer_ids.at(il);
+
+ return layers[ikv].v_scale;
+}
+
ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(sinfo);
@@ -2109,7 +2133,7 @@ ggml_tensor * llama_kv_cache::build_rope_shift(
const int64_t n_embd_head_k = hparams.n_embd_head_k(il);
const int64_t n_cells = ggml_nelements(k)/n_embd_k_gqa;
- if (!ggml_is_quantized(k->type) && !rot && !rp.k_norm) {
+ if (!ggml_is_quantized(k->type) && k->type != GGML_TYPE_F8_E4M3 && !rot && !rp.k_norm) {
// we rotate only the n_rot dimensions
ggml_tensor * k_rope =
ggml_view_3d(ctx, k,
@@ -3049,6 +3073,14 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons
return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
}
+ggml_tensor * llama_kv_cache_context::get_k_scale(int32_t il) const {
+ return kv->get_k_scale(il);
+}
+
+ggml_tensor * llama_kv_cache_context::get_v_scale(int32_t il) const {
+ return kv->get_v_scale(il);
+}
+
ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
}
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index 5aafbebb3..a81cd8498 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -207,6 +207,9 @@ public:
ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
+ ggml_tensor * get_k_scale(int32_t il) const;
+ ggml_tensor * get_v_scale(int32_t il) const;
+
// store k_cur and v_cur in the cache based on the provided head location
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
@@ -270,6 +273,9 @@ private:
ggml_tensor * k;
ggml_tensor * v;
+ ggml_tensor * k_scale;
+ ggml_tensor * v_scale;
+
std::vector<ggml_tensor *> k_stream;
std::vector<ggml_tensor *> v_stream;
@@ -420,6 +426,9 @@ public:
ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il) const;
+ ggml_tensor * get_k_scale(int32_t il) const;
+ ggml_tensor * get_v_scale(int32_t il) const;
+
// store k_cur and v_cur in the cache based on the provided head location
// note: the heads in k_cur and v_cur should be laid out contiguously in memory
// - k_cur [n_embd_head_k, n_head_k, n_tokens]
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index dbc91a137..cd49ccdd0 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -1813,6 +1813,16 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
for (int i = 0; i < n_layer_all; ++i) {
auto & layer = layers[i];
+ if (hparams.has_kv(i)) {
+ layer.k_cache_s = create_tensor(tn(LLM_TENSOR_ATTN_K, "k_scale", i), {1}, TENSOR_NOT_REQUIRED);
+ layer.v_cache_s = create_tensor(tn(LLM_TENSOR_ATTN_V, "v_scale", i), {1}, TENSOR_NOT_REQUIRED);
+
+ if ((layer.k_cache_s && layer.k_cache_s->type != GGML_TYPE_F32) ||
+ (layer.v_cache_s && layer.v_cache_s->type != GGML_TYPE_F32)) {
+ throw std::runtime_error(format("KV cache scales for layer %d must be F32", i));
+ }
+ }
+
// attention weight scales (per-tensor, shape {1})
if (!layer.wq_s && layer.wq) {
layer.wq_s = create_tensor(tn(LLM_TENSOR_ATTN_Q, "scale", i), {1}, TENSOR_NOT_REQUIRED);
diff --git a/src/llama-model.h b/src/llama-model.h
index 9c0a51f3d..1b7ea1ce7 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -290,6 +290,8 @@ struct llama_layer {
struct ggml_tensor * wq = nullptr;
struct ggml_tensor * wk = nullptr;
struct ggml_tensor * wv = nullptr;
+ struct ggml_tensor * k_cache_s = nullptr;
+ struct ggml_tensor * v_cache_s = nullptr;
struct ggml_tensor * wo = nullptr;
struct ggml_tensor * wqkv = nullptr;
struct ggml_tensor * wg = nullptr;
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 4ba52d3cf..d0eb01daf 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -190,6 +190,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
llama_build_and_test(test-grammar-integration.cpp LABEL behavior)
llama_build_and_test(test-llama-grammar.cpp LABEL behavior LINK llama)
llama_build_and_test(test-batch-alloc.cpp LABEL behavior LINK llama)
+ llama_build_and_test(test-kv-cache.cpp LABEL behavior)
+ target_include_directories(test-kv-cache PRIVATE ${PROJECT_SOURCE_DIR}/src)
llama_build_and_test(test-chat.cpp LABEL behavior WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server)
target_link_libraries(test-chat PRIVATE server-context)
@@ -237,7 +239,9 @@ if (GGML_CPU AND (NOT WIN32 OR NOT BUILD_SHARED_LIBS))
set_tests_properties(test-generate-fp8-models PROPERTIES FIXTURES_SETUP fp8-models)
llama_test(test-fp8 NAME test-fp8-requantize LABEL "main behavior backend"
ARGS --models "${FP8_MODEL_DIR}")
- set_tests_properties(test-fp8-requantize PROPERTIES FIXTURES_REQUIRED fp8-models)
+ llama_test(test-fp8 NAME test-fp8-kv LABEL "main behavior backend"
+ ARGS --models "${FP8_MODEL_DIR}" --kv)
+ set_tests_properties(test-fp8-requantize test-fp8-kv PROPERTIES FIXTURES_REQUIRED fp8-models)
endif()
set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/")
@@ -296,6 +300,12 @@ llama_test_cmd(python3
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
ARGS -m unittest discover -s tests -p test_conversion_fp8.py -v
)
+llama_test_cmd(python3
+ NAME test-kv-cache-scales
+ LABEL "python behavior"
+ WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
+ ARGS -m unittest discover -s tests -p test_kv_cache_scales.py -v
+)
llama_build_and_test(test-json-schema.cpp LABEL behavior)
llama_build_and_test(test-chat-auto-parser.cpp LABEL behavior WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
llama_build_and_test(test-chat-template.cpp LABEL behavior)
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 676aa9ed2..e7314fee4 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12917,6 +12917,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_Q4_0));
test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q2_0));
test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16));
+ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F8_E4M3, GGML_TYPE_F8_E4M3));
+ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F8_E4M3, GGML_TYPE_F16));
+ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F8_E4M3));
// q8_0 KV cases: decode and prompt batches, KV pad, permuted KV, feature flags, and long context
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0));
diff --git a/tests/test-fp8.cpp b/tests/test-fp8.cpp
index a92787836..d66c2140c 100644
--- a/tests/test-fp8.cpp
+++ b/tests/test-fp8.cpp
@@ -7,8 +7,12 @@
#include "gguf.h"
#include "llama.h"
#include "llama-cpp.h"
+#include "../src/llama-context.h"
+#include "../src/llama-kv-cache.h"
+#include "../src/llama-model.h"
#include <algorithm>
+#include <cmath>
#include <cstdio>
#include <cstring>
#include <filesystem>
@@ -286,6 +290,299 @@ static void check_fp8_requantize(
}
}
+static void check_fp8_attention(llama_model * model, bool offload_kqv = true) {
+ // TODO: enable CPU KV cache tests with tensor split after the meta backend handles the attention transfers.
+ if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR && !offload_kqv) {
+ LOG("FP8 attention: tensor split with CPU KV cache is not supported, skipping\n");
+ return;
+ }
+
+ ggml_context_ptr scale_ctx(ggml_init({ 4096, nullptr, true }));
+ ggml_tensor * k_scale = ggml_new_tensor_1d(scale_ctx.get(), GGML_TYPE_F32, 1);
+ ggml_tensor * v_scale = ggml_new_tensor_1d(scale_ctx.get(), GGML_TYPE_F32, 1);
+ ggml_set_name(k_scale, "blk.0.attn_k.k_scale");
+ ggml_set_name(v_scale, "blk.0.attn_v.v_scale");
+ ggml_backend_buffer_ptr scale_buf(ggml_backend_alloc_ctx_tensors_from_buft(scale_ctx.get(), ggml_backend_dev_buffer_type(model->dev_layer(0))));
+ ggml_tensor * old_k_scale = model->layers[0].k_cache_s;
+ ggml_tensor * old_v_scale = model->layers[0].v_cache_s;
+
+ const int64_t head_size = model->hparams.n_embd_head_k();
+ const int64_t n_head = model->hparams.n_head();
+ const int64_t n_tokens = 4;
+ GGML_ASSERT(n_head == model->hparams.n_head_kv());
+ GGML_ASSERT(head_size == model->hparams.n_embd_head_v());
+
+ for (llama_flash_attn_type flash : { LLAMA_FLASH_ATTN_TYPE_DISABLED, LLAMA_FLASH_ATTN_TYPE_ENABLED }) {
+ if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR && flash == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
+ LOG("FP8 attention: tensor split requires flash attention, skipping disabled cases\n");
+ continue;
+ }
+ for (int scaled : { 1, 2, 3 }) {
+ const float sk = scaled & 1 ? 0.03125f : 1.0f;
+ const float sv = scaled & 2 ? 0.125f : 1.0f;
+ ggml_backend_tensor_set(k_scale, &sk, 0, sizeof(sk));
+ ggml_backend_tensor_set(v_scale, &sv, 0, sizeof(sv));
+ model->layers[0].k_cache_s = scaled & 1 ? k_scale : nullptr;
+ model->layers[0].v_cache_s = scaled & 2 ? v_scale : nullptr;
+
+ std::vector<float> tensors[4];
+ auto params = llama_context_default_params();
+ params.n_ctx = 256;
+ params.n_batch = n_tokens;
+ params.n_ubatch = n_tokens;
+ params.n_threads = 1;
+ params.n_threads_batch = 1;
+ params.type_k = scaled & 1 ? GGML_TYPE_F8_E4M3 : GGML_TYPE_F16;
+ params.type_v = scaled & 2 ? GGML_TYPE_F8_E4M3 : GGML_TYPE_F16;
+ params.flash_attn_type = flash;
+ params.offload_kqv = offload_kqv;
+ if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR) {
+ ggml_context_ptr ctx(ggml_init({ 8*ggml_tensor_overhead(), nullptr, true }));
+ ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, head_size, n_tokens, n_head, 1);
+ ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), params.type_k, head_size, params.n_ctx, n_head, 1);
+ ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), params.type_v, head_size, params.n_ctx, n_head, 1);
+ ggml_tensor * mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, params.n_ctx, params.n_ctx, 1, 1);
+ ggml_tensor * attn = ggml_flash_attn_ext(ctx.get(), q, k, v, mask, 1.0f/sqrtf(head_size), 0.0f, 0.0f);
+ auto * dev = model->dev_layer(0);
+ if (!ggml_backend_dev_supports_op(dev, attn)) {
+ LOG("FP8 attention: scaled=%d offload=%d: tensor-split flash attention not supported on %s, skipping\n",
+ scaled, offload_kqv, ggml_backend_dev_name(dev));
+ continue;
+ }
+ }
+ if (offload_kqv) {
+ // Cache writes cannot fall back from a preallocated device buffer.
+ llama_kv_cache probe(*model, false, model->hparams, params.type_k, params.type_v,
+ flash == LLAMA_FLASH_ATTN_TYPE_DISABLED, false, true, params.n_ctx, 1,
+ 0, LLAMA_SWA_TYPE_NONE, nullptr, nullptr, nullptr, nullptr);
+ ggml_context_ptr ctx(ggml_init({ 64*ggml_tensor_overhead(), nullptr, true }));
+ llama_kv_cache::slot_info slots;
+ bool supported = true;
+ for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
+ ggml_tensor * k = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, head_size, n_head, n_tokens);
+ ggml_tensor * v = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, head_size, n_head, n_tokens);
+ ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I64, n_tokens);
+ ggml_tensor * v_idxs = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I64,
+ flash == LLAMA_FLASH_ATTN_TYPE_DISABLED ? n_tokens*head_size*n_head : n_tokens);
+ auto * dev = model->dev_layer(il);
+ if (!ggml_backend_dev_supports_op(dev, probe.cpy_k(ctx.get(), k, k_idxs, il, slots)) ||
+ !ggml_backend_dev_supports_op(dev, probe.cpy_v(ctx.get(), v, v_idxs, il, slots))) {
+ LOG("FP8 attention: flash=%s scaled=%d offload=1: cache writes not supported on %s, skipping\n",
+ llama_flash_attn_type_name(flash), scaled, ggml_backend_dev_name(dev));
+ supported = false;
+ break;
+ }
+ }
+ if (!supported) {
+ continue;
+ }
+ }
+ params.cb_eval_user_data = tensors;
+ params.cb_eval = [](ggml_tensor * t, bool ask, void * user_data) {
+ const char * names[] = { "Qcur-0", "Kcur-0", "Vcur-0", "kqv_out-0" };
+ for (int i = 0; i < 4; ++i) {
+ if (strcmp(t->name, names[i]) != 0) {
+ continue;
+ }
+ if (!ask) {
+ GGML_ASSERT(t->type == GGML_TYPE_F32 && ggml_is_contiguous(t));
+ auto & data = static_cast<std::vector<float> *>(user_data)[i];
+ data.resize(ggml_nelements(t));
+ ggml_backend_tensor_get(t, data.data(), 0, data.size()*sizeof(float));
+ }
+ return true;
+ }
+ return false;
+ };
+ llama_context_ptr lctx(llama_init_from_model(model, params));
+ GGML_ASSERT(lctx);
+ GGML_ASSERT(lctx->get_cparams().flash_attn == (flash == LLAMA_FLASH_ATTN_TYPE_ENABLED));
+ get_logits(model, lctx.get(), { 1, 2, 3, 4 });
+ llama_synchronize(lctx.get());
+ for (const auto & data : tensors) {
+ GGML_ASSERT(data.size() == (size_t) (head_size*n_head*n_tokens));
+ }
+
+ // Round the original K/V independently of the cache write graph.
+ for (int i : { 1, 2 }) {
+ const float scale = i == 1 ? sk : sv;
+ const ggml_type type = i == 1 ? params.type_k : params.type_v;
+ auto & data = tensors[i];
+ if (type == GGML_TYPE_F8_E4M3) {
+ for (float & value : data) {
+ value /= scale;
+ }
+ std::vector<uint8_t> quantized(data.size());
+ ggml_get_type_traits(type)->from_float_ref(data.data(), quantized.data(), data.size());
+ ggml_get_type_traits(type)->to_float(quantized.data(), data.data(), data.size());
+ for (float & value : data) {
+ value *= scale;
+ }
+ } else {
+ for (float & value : data) {
+ value = ggml_fp16_to_fp32(ggml_fp32_to_fp16(value));
+ }
+ }
+ }
+
+ std::vector<float> expected(tensors[3].size(), 0.0f);
+ for (int64_t token = 0; token < n_tokens; ++token) {
+ for (int64_t head = 0; head < n_head; ++head) {
+ const int64_t q_offset = (token*n_head + head)*head_size;
+ std::vector<double> scores(token + 1);
+ double max_score = -INFINITY;
+ for (int64_t past = 0; past <= token; ++past) {
+ const int64_t kv_offset = (past*n_head + head)*head_size;
+ double dot = 0.0;
+ for (int64_t d = 0; d < head_size; ++d) {
+ dot += (double) tensors[0][q_offset + d]*tensors[1][kv_offset + d];
+ }
+ scores[past] = dot/sqrt((double) head_size);
+ max_score = std::max(max_score, scores[past]);
+ }
+ double sum = 0.0;
+ for (double & score : scores) {
+ score = exp(score - max_score);
+ sum += score;
+ }
+ for (int64_t past = 0; past <= token; ++past) {
+ const int64_t kv_offset = (past*n_head + head)*head_size;
+ for (int64_t d = 0; d < head_size; ++d) {
+ expected[q_offset + d] += scores[past]/sum*tensors[2][kv_offset + d];
+ }
+ }
+ }
+ }
+ const double error = nmse(expected, tensors[3]);
+ if (!std::isfinite(error) || error >= 1e-5) {
+ LOG_ERR("FP8 attention: flash=%s scaled=%d offload=%d nmse=%.2e\n", llama_flash_attn_type_name(flash), scaled, offload_kqv, error);
+ }
+ GGML_ASSERT(std::isfinite(error) && error < 1e-5);
+ }
+ }
+ model->layers[0].k_cache_s = old_k_scale;
+ model->layers[0].v_cache_s = old_v_scale;
+}
+
+static void check_fp8_context_shift(llama_model * model) {
+ ggml_context_ptr scale_ctx(ggml_init({ 4096, nullptr, true }));
+ ggml_tensor * scale = ggml_new_tensor_1d(scale_ctx.get(), GGML_TYPE_F32, 1);
+ ggml_backend_buffer_ptr scale_buf(ggml_backend_alloc_ctx_tensors_from_buft(scale_ctx.get(), ggml_backend_cpu_buffer_type()));
+ ggml_tensor * old_scale = model->layers[0].k_cache_s;
+
+ for (float scale_value : { 1.0f, 0.5f }) {
+ ggml_backend_tensor_set(scale, &scale_value, 0, sizeof(scale_value));
+ model->layers[0].k_cache_s = scale;
+
+ auto params = llama_context_default_params();
+ params.n_ctx = 256;
+ params.n_batch = 4;
+ params.n_ubatch = 4;
+ params.n_threads = 1;
+ params.n_threads_batch = 1;
+ params.type_k = GGML_TYPE_F8_E4M3;
+ params.type_v = GGML_TYPE_F8_E4M3;
+ params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
+ params.offload_kqv = false;
+ llama_context_ptr lctx(llama_init_from_model(model, params));
+ GGML_ASSERT(lctx);
+ get_logits(model, lctx.get(), { 1, 2, 3, 4 });
+ llama_synchronize(lctx.get());
+
+ auto * kv = dynamic_cast<llama_kv_cache *>(llama_get_memory(lctx.get()));
+ GGML_ASSERT(kv && kv->get_can_shift());
+ ggml_tensor * k = kv->get_k_storage(0);
+ std::vector<uint8_t> before(ggml_nbytes(k));
+ ggml_backend_tensor_get(k, before.data(), 0, before.size());
+
+ ggml_context_ptr ctx(ggml_init({ 2*1024*1024, nullptr, false }));
+ const auto & hp = model->hparams;
+ const auto & cp = lctx->get_cparams();
+ ggml_tensor * keys = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, hp.n_embd_head_k(), hp.n_head_kv(), kv->get_size());
+ ggml_get_type_traits(GGML_TYPE_F8_E4M3)->to_float(before.data(), (float *) keys->data, ggml_nelements(keys));
+ ggml_tensor * shifts = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, kv->get_size());
+ ggml_set_i32(shifts, 0);
+ for (int i = 1; i < 4; ++i) {
+ ((int32_t *) shifts->data)[i] = -1;
+ }
+ ggml_tensor * expected_f32 = ggml_rope_ext(ctx.get(), keys, shifts, model->get_rope_factors(cp, 0), hp.n_rot(), hp.rope_type,
+ cp.n_ctx_orig_yarn, model->get_rope_freq_base(cp, 0), model->get_rope_freq_scale(cp, 0),
+ cp.yarn_ext_factor, cp.yarn_attn_factor, cp.yarn_beta_fast, cp.yarn_beta_slow);
+ ggml_cgraph * gf = ggml_new_graph(ctx.get());
+ ggml_build_forward_expand(gf, expected_f32);
+ GGML_ASSERT(ggml_graph_compute_with_ctx(ctx.get(), gf, 1) == GGML_STATUS_SUCCESS);
+ std::vector<uint8_t> expected(before.size());
+ ggml_quantize_chunk(GGML_TYPE_F8_E4M3, (float *) expected_f32->data, expected.data(), 0,
+ ggml_nelements(expected_f32)/keys->ne[0], keys->ne[0], nullptr);
+
+ llama_memory_seq_add(kv, 0, 1, 4, -1);
+ GGML_ASSERT(kv->get_has_shift());
+ GGML_ASSERT(lctx->memory_update(false));
+ llama_synchronize(lctx.get());
+ GGML_ASSERT(!kv->get_has_shift());
+ std::vector<uint8_t> after(before.size());
+ ggml_backend_tensor_get(k, after.data(), 0, after.size());
+ GGML_ASSERT(after != before);
+ GGML_ASSERT(after == expected);
+ lctx.reset();
+ }
+
+ model->layers[0].k_cache_s = old_scale;
+}
+
+enum class test_status {
+ PASS,
+ FAIL,
+ SKIP,
+};
+
+static test_status test_fp8_kv(const std::string & path, const char * target_backend) {
+ ggml_backend_dev_t cpu_devices[] = {nullptr};
+ auto params = llama_model_default_params();
+ params.devices = cpu_devices;
+ params.progress_callback = silent_model_load_progress;
+ {
+ llama_model_ptr model = load_model(path, params);
+ if (model->arch != LLM_ARCH_LLAMA || model->hparams.n_expert > 0) {
+ LOG("FP8 KV: only dense Llama fixtures are supported, skipping %s\n", path.c_str());
+ return test_status::SKIP;
+ }
+ LOG("FP8 KV: CPU attention and context shift\n");
+ check_fp8_attention(model.get());
+ check_fp8_context_shift(model.get());
+ }
+
+ std::vector<ggml_backend_dev_t> devices_meta;
+ for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
+ auto * dev = ggml_backend_dev_get(i);
+ if (target_backend && strcmp(target_backend, ggml_backend_dev_name(dev)) != 0) {
+ continue;
+ }
+ if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) {
+ devices_meta.push_back(dev);
+ }
+ if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
+ continue;
+ }
+ ggml_backend_dev_t devices[] = {dev, nullptr};
+ params.devices = devices;
+ LOG("FP8 KV: device=%s\n", ggml_backend_dev_name(dev));
+ llama_model_ptr model = load_model(path, params);
+ check_fp8_attention(model.get());
+ check_fp8_attention(model.get(), false);
+ }
+ if (!target_backend && !devices_meta.empty()) {
+ devices_meta.push_back(nullptr);
+ params.devices = devices_meta.data();
+ params.split_mode = LLAMA_SPLIT_MODE_TENSOR;
+ LOG("FP8 KV: tensor split\n");
+ llama_model_ptr model = load_model(path, params);
+ check_fp8_attention(model.get());
+ check_fp8_attention(model.get(), false);
+ }
+ return test_status::PASS;
+}
+
static void test_fp8(const std::string & model_path) {
const auto dir = std::filesystem::temp_directory_path() / ("llama-fp8-" + std::to_string(std::random_device{}()));
if (!std::filesystem::create_directory(dir)) {
@@ -324,51 +621,93 @@ static void test_fp8(const std::string & model_path) {
check_fp8_requantize(input, reference, output, reference_output, n_embd, model_params);
}
+static const char * test_status_str(test_status status) {
+ switch (status) {
+ case test_status::PASS: return "PASS";
+ case test_status::FAIL: return "FAIL";
+ case test_status::SKIP: return "SKIP";
+ }
+ GGML_ABORT("invalid test status");
+}
+
int main(int argc, char ** argv) {
if (argc == 2 && (strcmp(argv[1], "-h") == 0 || strcmp(argv[1], "--help") == 0)) {
- printf("Usage: %s -m model.gguf | --models DIR\n", argv[0]);
+ printf("Usage: %s (-m model.gguf | --models DIR) [--kv [-b backend]]\n", argv[0]);
return 0;
}
- if (argc != 3 || (strcmp(argv[1], "-m") != 0 && strcmp(argv[1], "--model") != 0 && strcmp(argv[1], "--models") != 0)) {
- fprintf(stderr, "Usage: %s -m model.gguf | --models DIR\n", argv[0]);
+ std::string model_file;
+ std::string models_dir;
+ const char * target_backend = nullptr;
+ bool kv = false;
+ bool valid = true;
+ for (int i = 1; i < argc; ++i) {
+ if ((strcmp(argv[i], "-m") == 0 || strcmp(argv[i], "--model") == 0 || strcmp(argv[i], "--models") == 0) && i + 1 < argc) {
+ if (!model_file.empty() || !models_dir.empty()) {
+ valid = false;
+ break;
+ }
+ const bool directory = strcmp(argv[i], "--models") == 0;
+ if (directory) {
+ models_dir = argv[++i];
+ } else {
+ model_file = argv[++i];
+ }
+ } else if (strcmp(argv[i], "--kv") == 0) {
+ kv = true;
+ } else if ((strcmp(argv[i], "-b") == 0 || strcmp(argv[i], "--backend") == 0) && i + 1 < argc) {
+ target_backend = argv[++i];
+ } else {
+ valid = false;
+ break;
+ }
+ }
+ if (!valid || (model_file.empty() && models_dir.empty()) || (target_backend && !kv)) {
+ fprintf(stderr, "Usage: %s (-m model.gguf | --models DIR) [--kv [-b backend]]\n", argv[0]);
return 1;
}
common_init();
llama_backend_init();
int result = 0;
try {
+ if (target_backend && !ggml_backend_dev_by_name(target_backend)) {
+ throw std::runtime_error("unknown backend device: " + std::string(target_backend));
+ }
std::vector<std::string> models;
- if (strcmp(argv[1], "--models") == 0) {
- if (!std::filesystem::is_directory(argv[2])) {
- throw std::runtime_error("models directory does not exist: " + std::string(argv[2]));
+ if (!models_dir.empty()) {
+ if (!std::filesystem::is_directory(models_dir)) {
+ throw std::runtime_error("models directory does not exist: " + models_dir);
}
- for (const auto & entry : std::filesystem::directory_iterator(argv[2])) {
+ for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
models.push_back(entry.path().string());
}
}
std::sort(models.begin(), models.end());
if (models.empty()) {
- throw std::runtime_error("no .gguf models found in: " + std::string(argv[2]));
+ throw std::runtime_error("no .gguf models found in: " + models_dir);
}
} else {
- models.push_back(argv[2]);
+ models.push_back(model_file);
}
- std::vector<bool> passed;
+ std::vector<test_status> results;
for (const auto & model : models) {
try {
- test_fp8(model);
- passed.push_back(true);
+ if (kv) {
+ results.push_back(test_fp8_kv(model, target_backend));
+ } else {
+ test_fp8(model);
+ results.push_back(test_status::PASS);
+ }
} catch (const std::exception & err) {
fprintf(stderr, "FP8 test failed for '%s': %s\n", model.c_str(), err.what());
- passed.push_back(false);
+ results.push_back(test_status::FAIL);
result = 1;
}
}
printf("\n| Model | Result |\n|-------|--------|\n");
for (size_t i = 0; i < models.size(); ++i) {
- printf("| %s | %s |\n", std::filesystem::path(models[i]).filename().string().c_str(), passed[i] ? "PASS" : "FAIL");
+ printf("| %s | %s |\n", std::filesystem::path(models[i]).filename().string().c_str(), test_status_str(results[i]));
}
} catch (const std::exception & err) {
fprintf(stderr, "FP8 test failed: %s\n", err.what());
diff --git a/tests/test-kv-cache.cpp b/tests/test-kv-cache.cpp
new file mode 100644
index 000000000..5cf74801b
--- /dev/null
+++ b/tests/test-kv-cache.cpp
@@ -0,0 +1,105 @@
+#include "llama-kv-cache.h"
+#include "llama-model.h"
+
+#include <numeric>
+#include <string>
+
+static void init_model(llama_model & model, uint32_t n_layer) {
+ model.hparams = {};
+ model.hparams.n_layer_all = n_layer;
+ model.hparams.n_embd_head_k_full = 4;
+ model.hparams.n_embd_head_v_full = 4;
+ model.hparams.n_head_arr.fill(1);
+ model.hparams.n_head_kv_arr.fill(1);
+ std::iota(model.hparams.kv_source_impl.begin(), model.hparams.kv_source_impl.end(), 0);
+ model.layers.resize(n_layer);
+}
+
+int main() {
+ llama_backend_init();
+
+ {
+ std::string errors;
+ llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
+ if (level == GGML_LOG_LEVEL_ERROR) {
+ *static_cast<std::string *>(user_data) += text;
+ }
+ }, &errors);
+
+ for (llm_arch arch : { LLM_ARCH_DEEPSEEK2, LLM_ARCH_DEEPSEEK32, LLM_ARCH_DEEPSEEK4,
+ LLM_ARCH_DEEPSEEK41, LLM_ARCH_GLM_DSA, LLM_ARCH_GLM5_NEXT, LLM_ARCH_DOTS3NOTE,
+ LLM_ARCH_MISTRAL_LARGE4, LLM_ARCH_HY_V4, LLM_ARCH_MINIMAX_M3, LLM_ARCH_QWEN4EXP,
+ LLM_ARCH_DFLASH }) {
+ std::unique_ptr<llama_model> model(llama_model_create(arch, llama_model_default_params()));
+ init_model(*model, 1);
+ if (arch == LLM_ARCH_DEEPSEEK2) {
+ model->hparams.n_embd_head_k_mla_impl = 4;
+ model->hparams.n_embd_head_v_mla_impl = 4;
+ }
+
+ for (ggml_type type_k : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+ for (ggml_type type_v : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+ if (type_k == GGML_TYPE_F16 && type_v == GGML_TYPE_F16) {
+ continue;
+ }
+ auto params = llama_context_default_params();
+ params.type_k = type_k;
+ params.type_v = type_v;
+ errors.clear();
+ llama_context * ctx = llama_init_from_model(model.get(), params);
+ GGML_ASSERT(ctx == nullptr);
+ if (type_k == type_v || (!model->hparams.is_mla() && arch != LLM_ARCH_DEEPSEEK4 && arch != LLM_ARCH_DEEPSEEK41)) {
+ GGML_ASSERT(errors.find("FP8 cache is not supported for architecture") != std::string::npos);
+ } else {
+ GGML_ASSERT(errors.find("model does not support different K") != std::string::npos);
+ }
+ }
+ }
+ }
+
+ llama_log_set(nullptr, nullptr);
+ }
+
+ {
+ std::unique_ptr<llama_model> target(llama_model_create(LLM_ARCH_GEMMA4, llama_model_default_params()));
+ std::unique_ptr<llama_model> assistant(llama_model_create(LLM_ARCH_GEMMA4_ASSISTANT, llama_model_default_params()));
+ init_model(*target, 3);
+ init_model(*assistant, 2);
+
+ ggml_context_ptr ctx(ggml_init({ 4096, nullptr, false }));
+ target->layers[2].k_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+ target->layers[2].v_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+ assistant->layers[0].k_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+ assistant->layers[0].v_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+ assistant->layers[1].k_cache_s = assistant->layers[0].k_cache_s;
+ assistant->layers[1].v_cache_s = assistant->layers[0].v_cache_s;
+ ggml_set_f32(target->layers[2].k_cache_s, 0.5f);
+ ggml_set_f32(target->layers[2].v_cache_s, 0.25f);
+ ggml_set_f32(assistant->layers[0].k_cache_s, 2.0f);
+ ggml_set_f32(assistant->layers[0].v_cache_s, 4.0f);
+
+ for (ggml_type type_k : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+ for (ggml_type type_v : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+ llama_kv_cache source(*target, false, target->hparams, type_k, type_v,
+ false, false, true, 32, 1, 0, LLAMA_SWA_TYPE_NONE, nullptr, nullptr,
+ [](uint32_t il) { return il == 0 ? 2 : -1; }, nullptr);
+ llama_kv_cache shared(*assistant, false, assistant->hparams, type_k, type_v,
+ false, false, true, 32, 1, 0, LLAMA_SWA_TYPE_NONE, &source, nullptr, nullptr,
+ [](int32_t il) { return il == 0 ? 2 : 1; });
+
+ GGML_ASSERT(source.get_k_scale(2) == (type_k == GGML_TYPE_F8_E4M3 ? target->layers[2].k_cache_s : nullptr));
+ GGML_ASSERT(source.get_v_scale(2) == (type_v == GGML_TYPE_F8_E4M3 ? target->layers[2].v_cache_s : nullptr));
+ GGML_ASSERT(source.get_k_scale(0) == source.get_k_scale(2));
+ GGML_ASSERT(source.get_v_scale(0) == source.get_v_scale(2));
+ GGML_ASSERT(shared.get_k_storage(0) == source.get_k_storage(2));
+ GGML_ASSERT(shared.get_k_scale(0) == source.get_k_scale(2));
+ GGML_ASSERT(shared.get_v_scale(0) == source.get_v_scale(2));
+ GGML_ASSERT(shared.get_k_scale(1) == nullptr);
+ GGML_ASSERT(shared.get_v_scale(1) == nullptr);
+ }
+ }
+ }
+
+ llama_backend_free();
+ return 0;
+}
diff --git a/tests/test_kv_cache_scales.py b/tests/test_kv_cache_scales.py
new file mode 100644
index 000000000..7dedf4b31
--- /dev/null
+++ b/tests/test_kv_cache_scales.py
@@ -0,0 +1,57 @@
+"""Test KV cache scales with synthetic tensors; no GPU or model downloads required.
+
+Install dependencies and run from the repository root:
+ python -m pip install -r requirements/requirements-convert_hf_to_gguf.txt
+ python -m pip install -e ./gguf-py
+ python -m unittest discover -s tests -p 'test_kv_cache_scales.py' -v
+
+Use unittest discovery so the repository root is available for imports.
+"""
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+
+from conversion.base import ModelBase
+import gguf
+
+
+class TestKVCacheScales(unittest.TestCase):
+ def test_projection_names(self):
+ cases = (
+ (gguf.MODEL_ARCH.PHI3, "model.layers.2.self_attn.qkv_proj", "model.layers.2.self_attn.qkv_proj"),
+ (gguf.MODEL_ARCH.GPT2, "transformer.h.2.attn.c_attn", "transformer.h.2.attn.c_attn"),
+ (gguf.MODEL_ARCH.LLAMA, "model.layers.2.self_attn.k_proj", "model.layers.2.self_attn.v_proj"),
+ )
+ for arch, k_name, v_name in cases:
+ for dtype in (torch.float32, torch.bfloat16):
+ with self.subTest(arch=arch, dtype=dtype), tempfile.TemporaryDirectory() as tmp:
+ path = Path(tmp) / "scales.gguf"
+ model = object.__new__(ModelBase)
+ model.tensor_map = gguf.get_tensor_name_map(arch, 3)
+ model.gguf_writer = gguf.GGUFWriter(path, gguf.MODEL_ARCH_NAMES[arch])
+ model.model_tensors = {
+ k_name + ".k_scale": lambda: torch.tensor(0.5, dtype=dtype),
+ v_name + ".v_scale": lambda: torch.tensor(1.0, dtype=dtype),
+ }
+ model._prepare_kv_cache_scales()
+ self.assertFalse(model.model_tensors)
+ model.gguf_writer.write_header_to_file()
+ model.gguf_writer.write_kv_data_to_file()
+ model.gguf_writer.write_tensors_to_file()
+ model.gguf_writer.close()
+
+ # Release the reader's memory map before Windows removes the temporary directory.
+ self._check_scales(path)
+
+ def _check_scales(self, path):
+ reader = gguf.GGUFReader(path)
+ self.assertEqual([t.name for t in reader.tensors], ["blk.2.attn_k.k_scale", "blk.2.attn_v.v_scale"])
+ self.assertEqual([t.tensor_type for t in reader.tensors], [gguf.GGMLQuantizationType.F32] * 2)
+ self.assertEqual([float(t.data[0]) for t in reader.tensors], [0.5, 1.0])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp
index bd87644c3..5cbdda14c 100644
--- a/tools/llama-bench/llama-bench.cpp
+++ b/tools/llama-bench/llama-bench.cpp
@@ -500,6 +500,9 @@ static ggml_type ggml_type_from_name(const std::string & s) {
if (s == "bf16") {
return GGML_TYPE_BF16;
}
+ if (s == "f8_e4m3") {
+ return GGML_TYPE_F8_E4M3;
+ }
if (s == "q8_0") {
return GGML_TYPE_Q8_0;
}