diff --git a/common/arg.cpp b/common/arg.cpp index e6d6450f4a..d02cd6520b 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 05a3172b06..3c822e41e4 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 @@ -328,6 +329,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): @@ -1420,6 +1441,8 @@ class ModelBase: assert titem[0] in modules, f"quantized_layers entry {layer_name!r} is not in the model tensors" self._nvfp4_layer_algo[titem[0]] = entry.get("quant_algo") + 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 455f5198e4..fce4b15fcb 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4079,6 +4079,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 77b0175560..5e14f7f087 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3228,6 +3228,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, @@ -3252,6 +3295,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 @@ -3259,8 +3306,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; @@ -3286,6 +3331,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * kq_mask = inp->get_kq_mask(); 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) { @@ -3527,6 +3573,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); @@ -3539,10 +3590,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) { @@ -3571,6 +3618,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, 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 e108f0423b..25736ba3cd 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -307,6 +307,7 @@ llama_kv_cache::llama_kv_cache( model.attn_rot ? "LLAMA_ATTN_ROT_DISABLE" : "llama_model_set_attn_rot"); } + // Do not rotate scalar FP8 caches. Their static scales are calibrated on unrotated K and V. bool attn_rot_k = !attn_rot_disable && n_embd_head_k_all > 0 && @@ -1320,6 +1321,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); @@ -2878,6 +2887,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 60d2e183c6..f990abf014 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -196,6 +196,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; @@ -404,6 +407,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 0ae58d15bf..77f3f9d1f3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1732,6 +1732,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 222238efed..07c53dc5f9 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -289,6 +289,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 037b1938ca..f7535c91c5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -12229,6 +12229,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 96ae4870ff..fc22d4a211 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; }