diff --git a/common/arg.cpp b/common/arg.cpp index 2669cacd6c..d4682891fb 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -305,6 +305,7 @@ const std::vector 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 a88e8cc8de..e464d6abfb 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 @@ -311,6 +312,26 @@ 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 + + new_name = self.tensor_map.get_name(key=name, try_suffixes=(".k_scale", ".v_scale")) + if new_name is None: + continue + + scale = LazyTorchTensor.to_eager(self.model_tensors.pop(name)()) + if scale.dtype != torch.float32 or scale.numel() != 1: + raise ValueError(f"KV cache scale {name!r} must be a scalar FP32 tensor") + + 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): @@ -968,6 +989,8 @@ class ModelBase: self._is_nvfp4 = quant_algo in ("NVFP4", "W4A16_NVFP4") self._is_mxfp4 = quant_method == "mxfp4" + 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 c1ef12f56b..0feb9f9066 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3696,6 +3696,33 @@ 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; + if (fp8_cache && (model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4)) { + LLAMA_LOG_ERROR("%s: FP8 cache is not supported for MLA models\n", __func__); + 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_scale, model->layers[il].v_cache_scale }) { + 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 a75f632b8c..0f93bee401 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -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) { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index a342ee1191..16e16dd1e4 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -318,6 +318,7 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__); } + // Do not rotate scalar FP8 caches. Their static scales are calibrated on unrotated K and V. attn_rot_k = !attn_rot_disable && n_embd_head_k_all > 0 && @@ -1315,6 +1316,14 @@ 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); } +ggml_tensor * llama_kv_cache::get_k_scale(int32_t il) const { + return model.layers[il].k_cache_scale; +} + +ggml_tensor * llama_kv_cache::get_v_scale(int32_t il) const { + return model.layers[il].v_cache_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); @@ -2758,6 +2767,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 c4d8699def..994a394240 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -189,6 +189,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; @@ -398,6 +401,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 14078fde52..76eb473113 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1524,6 +1524,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_scale = create_tensor(tn(LLM_TENSOR_ATTN_K, "k_scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.v_cache_scale = create_tensor(tn(LLM_TENSOR_ATTN_V, "v_scale", i), {1}, TENSOR_NOT_REQUIRED); + + if ((layer.k_cache_scale && layer.k_cache_scale->type != GGML_TYPE_F32) || + (layer.v_cache_scale && layer.v_cache_scale->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 4c4a30e018..efa70a170a 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -276,6 +276,8 @@ struct llama_layer { struct ggml_tensor * wq = nullptr; struct ggml_tensor * wk = nullptr; struct ggml_tensor * wv = nullptr; + struct ggml_tensor * k_cache_scale = nullptr; + struct ggml_tensor * v_cache_scale = nullptr; struct ggml_tensor * wo = nullptr; struct ggml_tensor * wqkv = nullptr; struct ggml_tensor * wg = nullptr; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index be6ee77161..e1b0c909b7 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10346,6 +10346,9 @@ static std::vector> 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/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 1fff21f701..782ad5e690 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -498,6 +498,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; }