Fix DSV4 quantized cache accounting

This commit is contained in:
samuel
2026-07-14 20:01:50 -03:00
parent 74fcc717df
commit e2b0340742
5 changed files with 122 additions and 13 deletions
+12 -4
View File
@@ -248,6 +248,10 @@ static ggml_tensor * dsv4_require_f32_rows(
return ggml_cast(ctx, t, GGML_TYPE_F32);
}
static ggml_tensor * dsv4_cache_read_f32(ggml_context * ctx, ggml_tensor * t) {
return t != nullptr && ggml_is_quantized(t->type) ? ggml_cast(ctx, t, GGML_TYPE_F32) : t;
}
static ggml_tensor * dsv4_cache_stream_view_3d(
ggml_context * ctx,
ggml_tensor * cache,
@@ -320,7 +324,9 @@ static ggml_tensor * dsv4_raw_get_k(
ggml_tensor * idxs = raw_k_read_idxs->type == GGML_TYPE_I32
? raw_k_read_idxs : ggml_cast(ctx, raw_k_read_idxs, GGML_TYPE_I32);
ggml_tensor * rows = ggml_get_rows(ctx, cache_2d, idxs);
if (rows->type != cache->type && !ggml_is_quantized(cache->type)) {
if (ggml_is_quantized(cache->type) && rows->type != GGML_TYPE_F32) {
rows = ggml_cast(ctx, rows, GGML_TYPE_F32);
} else if (rows->type != cache->type && !ggml_is_quantized(cache->type)) {
rows = ggml_cast(ctx, rows, cache->type);
}
@@ -400,10 +406,10 @@ static ggml_tensor * dsv4_comp_get_k(
}
if (comp.sinfo.n_stream() == 0) {
return ggml_reshape_4d(ctx, dsv4_cache_view_3d(ctx, cache, n_embd_head, n_kv), n_embd_head, 1, n_kv, 1);
return dsv4_cache_read_f32(ctx, ggml_reshape_4d(ctx, dsv4_cache_view_3d(ctx, cache, n_embd_head, n_kv), n_embd_head, 1, n_kv, 1));
}
return dsv4_cache_stream_view_4d(ctx, cache, n_embd_head, n_kv, kv_size, comp.sinfo.s0, (int64_t) comp.sinfo.n_stream());
return dsv4_cache_read_f32(ctx, dsv4_cache_stream_view_4d(ctx, cache, n_embd_head, n_kv, kv_size, comp.sinfo.s0, (int64_t) comp.sinfo.n_stream()));
}
static ggml_tensor * dsv4_comp_cpy_k(
@@ -1202,7 +1208,9 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
if (raw_k_write == nullptr) {
llm_build_kv_store(lctx, ctx0, hparams, cparams, kv_self, gf, kv, nullptr, n_tokens, kv_head, cb, il);
}
llm_build_kv_store(lctx, ctx0, hparams, cparams, kv_self, gf, nullptr, kv, n_tokens, kv_head, cb, il);
if (il < (int64_t) kv_self.v_l.size() && kv_self.v_l[(size_t) il] != nullptr) {
llm_build_kv_store(lctx, ctx0, hparams, cparams, kv_self, gf, nullptr, kv, n_tokens, kv_head, cb, il);
}
ggml_tensor * raw_k = nullptr;
if (hparams.n_head_kv(il) == 1 && lctx.dsv4.inputs.raw_k_read_idxs != nullptr) {
+1
View File
@@ -56,6 +56,7 @@ struct llama_cparams {
enum ggml_type reduce_type;
enum ggml_type graph_attn_precision;
enum ggml_type idx_type_k = GGML_TYPE_F16;
enum llama_pooling_type pooling_type;
enum llama_mtp_op_type mtp_op_type;
+49 -1
View File
@@ -14,6 +14,23 @@
#include <type_traits>
#include <unordered_set>
static bool dsv4_cache_type_supported(ggml_type type) {
return type == GGML_TYPE_F16 || type == GGML_TYPE_BF16 || type == GGML_TYPE_Q8_0;
}
static bool dsv4_validate_cache_type(ggml_type type, int64_t width, const char * name) {
if (!dsv4_cache_type_supported(type)) {
LLAMA_LOG_ERROR("%s: unsupported DSV4 %s cache type %s\n", __func__, name, ggml_type_name(type));
return false;
}
if (ggml_is_quantized(type) && width % ggml_blck_size(type) != 0) {
LLAMA_LOG_ERROR("%s: DSV4 %s cache width %lld is not aligned to %d elements for %s\n",
__func__, name, (long long) width, ggml_blck_size(type), ggml_type_name(type));
return false;
}
return true;
}
static ggml_backend_buffer_type_t llama_dsv4_layer_buft(const llama_context & lctx, int32_t il) {
if (il >= 0 && il < (int32_t) lctx.model.buft_layer.size() && lctx.model.buft_layer[il].buft != nullptr) {
return lctx.model.buft_layer[il].buft;
@@ -794,6 +811,11 @@ bool llama_context::ensure_dsv4_cache_tensors() {
const uint32_t csa_kv = GGML_PAD(dsv4_comp_size(cparams.n_ctx, dsv4_runtime::CSA_RATIO), 256u);
const uint32_t hca_kv = GGML_PAD(dsv4_comp_size(cparams.n_ctx, dsv4_runtime::HCA_RATIO), 256u);
if (!dsv4_validate_cache_type(kv_self.type_k, n_embd_head, "raw/CSA/HCA") ||
!dsv4_validate_cache_type(cparams.idx_type_k, n_indexer_head, "LID")) {
return false;
}
if (dsv4.cache.cache_ctx != nullptr &&
(int32_t) dsv4.cache.csa_k.size() == n_layer &&
dsv4.cache.n_stream == n_stream) {
@@ -845,7 +867,7 @@ bool llama_context::ensure_dsv4_cache_tensors() {
if (ratio == dsv4_runtime::CSA_RATIO) {
cache.csa_k[(size_t) il] = ggml_new_tensor_3d(cache.cache_ctx, kv_self.type_k, n_embd_head, csa_kv*n_stream, 1);
cache.lid_k[(size_t) il] = ggml_new_tensor_3d(cache.cache_ctx, kv_self.type_k, n_indexer_head, csa_kv*n_stream, 1);
cache.lid_k[(size_t) il] = ggml_new_tensor_3d(cache.cache_ctx, cparams.idx_type_k, n_indexer_head, csa_kv*n_stream, 1);
cache.csa_state_kv[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, 2*n_embd_head, 2*dsv4_runtime::CSA_RATIO*n_stream);
cache.csa_state_score[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, 2*n_embd_head, 2*dsv4_runtime::CSA_RATIO*n_stream);
cache.lid_state_kv[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, 2*n_indexer_head, 2*dsv4_runtime::CSA_RATIO*n_stream);
@@ -876,6 +898,32 @@ bool llama_context::ensure_dsv4_cache_tensors() {
}
}
auto bytes = [](const auto & tensors) {
size_t total = 0;
for (const ggml_tensor * tensor : tensors) {
if (tensor != nullptr) {
total += ggml_nbytes(tensor);
}
}
return total;
};
const size_t csa_k_bytes = bytes(cache.csa_k);
const size_t hca_k_bytes = bytes(cache.hca_k);
const size_t lid_k_bytes = bytes(cache.lid_k);
const size_t csa_state_bytes = bytes(cache.csa_state_kv) + bytes(cache.csa_state_score);
const size_t hca_state_bytes = bytes(cache.hca_state_kv) + bytes(cache.hca_state_score);
const size_t lid_state_bytes = bytes(cache.lid_state_kv) + bytes(cache.lid_state_score);
LLAMA_LOG_INFO("%s: DSV4 cache: CSA K=%7.2f MiB (%s), HCA K=%7.2f MiB (%s), LID K=%7.2f MiB (%s), states=%7.2f MiB, total=%7.2f MiB, streams=%u\n",
__func__,
(float) csa_k_bytes / (1024.0f * 1024.0f), ggml_type_name(kv_self.type_k),
(float) hca_k_bytes / (1024.0f * 1024.0f), ggml_type_name(kv_self.type_k),
(float) lid_k_bytes / (1024.0f * 1024.0f), ggml_type_name(cparams.idx_type_k),
(float) (csa_state_bytes + hca_state_bytes + lid_state_bytes) / (1024.0f * 1024.0f),
(float) (csa_k_bytes + hca_k_bytes + lid_k_bytes + csa_state_bytes + hca_state_bytes + lid_state_bytes) / (1024.0f * 1024.0f),
n_stream);
return true;
}
+25
View File
@@ -2280,6 +2280,31 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_
}
return size;
}
if (arch == LLM_ARCH_DEEPSEEK4) {
constexpr uint32_t csa_ratio = 4;
constexpr uint32_t hca_ratio = 128;
constexpr uint32_t cache_pad = 256;
const uint32_t n_stream = std::max<uint32_t>(1, n_seq_max);
const uint32_t csa_kv = GGML_PAD(std::max<uint32_t>(1, (kv_size + csa_ratio - 1)/csa_ratio), cache_pad);
const uint32_t hca_kv = GGML_PAD(std::max<uint32_t>(1, (kv_size + hca_ratio - 1)/hca_ratio), cache_pad);
const uint32_t ratio = hparams.dsv4_compress_ratios[(size_t) il];
const int64_t n_embd_head = hparams.n_embd_head_k(il);
const int64_t n_indexer_head = hparams.indexer_head_size;
size_t size = ggml_row_size(type_k, n_embd_head) * hparams.n_head_kv(il) * kv_size;
if (ratio == csa_ratio) {
size += ggml_row_size(type_k, n_embd_head) * csa_kv * n_stream;
size += ggml_row_size(idx_type_k, n_indexer_head) * csa_kv * n_stream;
size += (size_t) 2 * n_embd_head * 2 * csa_ratio * n_stream * sizeof(float) * 2;
size += (size_t) 2 * n_indexer_head * 2 * csa_ratio * n_stream * sizeof(float) * 2;
} else if (ratio == hca_ratio) {
size += ggml_row_size(type_k, n_embd_head) * hca_kv * n_stream;
size += (size_t) n_embd_head * hca_ratio * n_stream * sizeof(float) * 2;
}
return size;
}
auto n_head_kv = hparams.n_head_kv(il);
auto k_size = ggml_row_size(type_k, hparams.n_embd_head_k(il)) * n_head_kv*kv_size;
auto v_size = ggml_row_size(type_v, hparams.n_embd_v_gqa(il)) * kv_size;
+35 -8
View File
@@ -1092,11 +1092,12 @@ static bool llama_kv_cache_init(
}
}
const bool is_dsv4_k_only = model.arch == LLM_ARCH_DEEPSEEK4;
bool is_mla_attn = model.is_mla_model();
bool split_cache = false;
bool replicate_mla = false;
if ((model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN) && !is_mla_attn && offload) {
if ((model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN) && !is_mla_attn && offload && !is_dsv4_k_only) {
cache.split_k_l.reserve(n_layer);
cache.split_v_l.reserve(n_layer);
if (llama_model_has_recurrent(&model)) {
@@ -1180,7 +1181,7 @@ static bool llama_kv_cache_init(
if (is_mla_attn && cparams.mla_attn) {
needs_v_cache = cparams.mla_attn == 1 && !cparams.flash_attn;
}
if (needs_v_cache) cache.v_l.reserve(n_layer);
if (needs_v_cache && !is_dsv4_k_only) cache.v_l.reserve(n_layer);
cache.s_l.resize(n_layer, nullptr);
// DSA indexer-key cache: one [indexer_head_size, kv_size] tensor per indexer layer.
@@ -1207,7 +1208,7 @@ static bool llama_kv_cache_init(
// For MTP-only context, skip KV allocation for non-MTP layers
if (cparams.mtp_op_type != MTP_OP_NONE && i < n_mtp_first_layer) {
cache.k_l.push_back(nullptr);
if (model.arch != LLM_ARCH_OPENPANGU &&
if (!is_dsv4_k_only && model.arch != LLM_ARCH_OPENPANGU &&
(!is_mla_attn || !cparams.mla_attn || (cparams.mla_attn == 1 && !cparams.flash_attn))) {
cache.v_l.push_back(nullptr);
}
@@ -1278,7 +1279,9 @@ static bool llama_kv_cache_init(
const bool is_mtp_layer = (cparams.mtp_op_type != MTP_OP_NONE && i >= (int)n_mtp_first_layer);
if (!hparams.has_kv(i) && !is_mtp_layer) {
cache.k_l.push_back(nullptr);
cache.v_l.push_back(nullptr);
if (!is_dsv4_k_only) {
cache.v_l.push_back(nullptr);
}
continue;
}
if (qnext_recurrent) {
@@ -1353,6 +1356,8 @@ static bool llama_kv_cache_init(
// The value-side latent is rederived from k_l per graph; no persistent V store.
const int64_t n_lat = (int64_t) hparams.n_lora_kv + hparams.n_rot; // 576
k = ggml_new_tensor_2d(ctx, this_type_k, n_lat, kv_size);
} else if (is_dsv4_k_only) {
k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*kv_size);
} else {
k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*kv_size);
v = ggml_new_tensor_1d(ctx, this_type_v, v_ne);
@@ -1428,7 +1433,7 @@ static bool llama_kv_cache_init(
v->extra = (void *)&split_v_l.ggml;
}
cache.k_l.push_back(k);
if (model.arch != LLM_ARCH_OPENPANGU) {
if (!is_dsv4_k_only && model.arch != LLM_ARCH_OPENPANGU) {
cache.v_l.push_back(v);
}
}
@@ -7566,19 +7571,36 @@ struct llama_context * llama_init_from_model(
// params.flash_attn = false;
//}
if (model->arch == LLM_ARCH_DEEPSEEK4 && params.type_v != GGML_TYPE_F16) {
LLAMA_LOG_WARN("%s: DeepSeek4 has no independent V-cache; ignoring requested V-cache type %s\n",
__func__, ggml_type_name(params.type_v));
params.type_v = GGML_TYPE_F16;
}
if (model->arch != LLM_ARCH_OPENPANGU &&
params.type_v != GGML_TYPE_F16 && params.type_v != GGML_TYPE_BF16 && !params.flash_attn) {
LLAMA_LOG_ERROR("%s: V cache quantization requires flash_attn\n", __func__);
return nullptr;
}
if (model->arch == LLM_ARCH_DEEPSEEK4 &&
(ggml_is_quantized(params.type_k) || params.k_cache_hadamard)) {
LLAMA_LOG_ERROR("%s: DeepSeek4 compressed attention currently supports only F16/BF16 K-cache; quantized K-cache and K-cache Hadamard are not implemented\n",
if (model->arch == LLM_ARCH_DEEPSEEK4 && params.k_cache_hadamard) {
LLAMA_LOG_ERROR("%s: DeepSeek4 K-cache Hadamard is not supported; use an untransformed K-cache\n",
__func__);
return nullptr;
}
if (model->arch == LLM_ARCH_DEEPSEEK4 &&
params.type_k != GGML_TYPE_F16 && params.type_k != GGML_TYPE_BF16 && params.type_k != GGML_TYPE_Q8_0) {
LLAMA_LOG_ERROR("%s: DeepSeek4 K-cache supports only F16, BF16, and Q8_0 (requested %s)\n",
__func__, ggml_type_name(params.type_k));
return nullptr;
}
if (model->arch == LLM_ARCH_DEEPSEEK4 && params.v_cache_hadamard) {
LLAMA_LOG_WARN("%s: DeepSeek4 has no independent V-cache; ignoring -vhad\n", __func__);
params.v_cache_hadamard = false;
}
if (params.k_cache_hadamard && !ggml_is_quantized(params.type_k)) {
LLAMA_LOG_WARN("%s: there is no point in Hadamard transforms with not quantized K-cache. Turning K-cache Hadamard off\n", __func__);
params.k_cache_hadamard = false;
@@ -7676,6 +7698,7 @@ struct llama_context * llama_init_from_model(
cparams.reduce_type = params.type_reduce;
cparams.graph_attn_precision = params.type_graph_attn;
cparams.idx_type_k = params.idx_type_k;
if (cparams.graph_attn_precision != GGML_TYPE_F16 && cparams.graph_attn_precision != GGML_TYPE_F32) {
throw std::runtime_error(format("--graph-attn-precision must be f16 or f32, got %s",
ggml_type_name(cparams.graph_attn_precision)));
@@ -8043,6 +8066,10 @@ struct llama_context * llama_init_from_model(
LLAMA_LOG_INFO("%s: KV self size = %7.2f MiB, c^KV (%s): %7.2f MiB, kv^T: not used\n", __func__,
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f),
ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f));
} else if (model->arch == LLM_ARCH_DEEPSEEK4) {
LLAMA_LOG_INFO("%s: KV self size = %7.2f MiB, K-only (%s): %7.2f MiB; independent V-cache: not used\n", __func__,
(float) memory_size_k / (1024.0f * 1024.0f),
ggml_type_name(type_k), (float) memory_size_k / (1024.0f * 1024.0f));
} else {
LLAMA_LOG_INFO("%s: KV self size = %7.2f MiB, K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__,
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f),