mirror of
https://github.com/ikawrakow/ik_llama.cpp.git
synced 2026-07-21 10:15:35 +00:00
Fix DSV4 quantized cache accounting
This commit is contained in:
@@ -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) {
|
||||
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user