LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Add logging for DFlash statistics and clean up workspace handling

Browse Source
This commit is contained in:
SamuelOliveirads
2026-06-13 20:14:08 -03:00
parent 3a1d46c4d1
commit 3b1a0f88d5
5 changed files with 352 additions and 519 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/speculative-impl.h
+3 -7
View File
@@ -266,7 +266,6 @@ struct common_speculative_state_dflash : public common_speculative_state {
return;
}
const bool use_kv_cache = dflash_use_kv_cache_experiment();
const float * target_features = nullptr;
size_t target_feature_floats = 0;
llama_dflash_window_update window_update = {
@@ -277,16 +276,13 @@ struct common_speculative_state_dflash : public common_speculative_state {
target_window_append_features.empty() ? nullptr : target_window_append_features.data(),
target_window_append_features.size(),
};
const llama_dflash_kv_cache_transition cache_plan = use_kv_cache
? llama_plan_dflash_kv_cache_transition_for_ctx(ctx_dft, window_update, target_window_rows)
: llama_dflash_kv_cache_transition{};
const llama_dflash_kv_cache_transition cache_plan =
llama_plan_dflash_kv_cache_transition_for_ctx(ctx_dft, window_update, target_window_rows);
if (!use_kv_cache || cache_plan.rebuild_cache) {
if (cache_plan.rebuild_cache) {
dflash_materialize_target_window_features(*this);
target_features = target_window.data();
target_feature_floats = target_window.size();
}
if (use_kv_cache && cache_plan.rebuild_cache) {
window_update.append_features = target_window.data();
window_update.append_floats = target_window.size();
window_update.append_rows = target_window_rows;
common/speculative.cpp
+3 -3
View File
@@ -309,8 +309,8 @@ static bool dflash_contract_log_enabled() {
std::strcmp(env, "off") != 0;
}
static bool dflash_use_kv_cache_experiment() {
const char * env = std::getenv("IK_DFLASH_KV_CACHE");
static bool dflash_stats_log_enabled() {
const char * env = std::getenv("IK_DFLASH_STATS_LOG");
if (env == nullptr || *env == '\0') {
return false;
}
@@ -1318,7 +1318,7 @@ void common_speculative_print_stats(const common_speculative * spec, double slot
if (impl->type == COMMON_SPECULATIVE_TYPE_DFLASH) {
const auto * dflash_state = dynamic_cast<const common_speculative_state_dflash *>(impl.get());
if (dflash_state != nullptr) {
if (dflash_state != nullptr && dflash_stats_log_enabled()) {
llama_dflash_profile_stats capture_stats;
llama_dflash_profile_stats graph_stats;
const bool have_capture = llama_dflash_profile_get_stats(dflash_state->ctx_tgt, &capture_stats);
src/graphs/build_dflash.cpp
+42 -173
View File
@@ -3,29 +3,6 @@
#include "../llama-model.h"
#include <cmath>
#include <cstdlib>
static bool dflash_use_kv_cache_experiment() {
const char * env = std::getenv("IK_DFLASH_KV_CACHE");
if (env == nullptr || *env == '\0') {
return false;
}
return std::strcmp(env, "0") != 0 &&
std::strcmp(env, "false") != 0 &&
std::strcmp(env, "off") != 0;
}
static bool dflash_use_kv_workspace_experiment() {
const char * env = std::getenv("IK_DFLASH_KV_WORKSPACE");
if (env == nullptr || *env == '\0') {
return false;
}
return std::strcmp(env, "0") != 0 &&
std::strcmp(env, "false") != 0 &&
std::strcmp(env, "off") != 0;
}
ggml_cgraph * llm_build_context::build_dflash_kv_workspace() {
const int64_t n_embd_head_k = hparams.n_embd_head_k(0);
@@ -272,13 +249,10 @@ ggml_cgraph * llm_build_context::build_dflash() {
const int64_t n_embd_head_v = hparams.n_embd_head_v(0);
const int64_t n_target_features = hparams.dflash_n_target_features;
auto & profile = lctx.dflash_profile;
const bool use_kv_cache = dflash_use_kv_cache_experiment();
const bool use_kv_workspace = use_kv_cache && dflash_use_kv_workspace_experiment();
const int64_t ctx_len = lctx.dflash_visible_cross_ctx > 0
? (int64_t) lctx.dflash_visible_cross_ctx
: std::max<int64_t>(1, (int64_t) cparams.n_ctx - (int64_t) hparams.dflash_block_size);
const int32_t cache_rows = use_kv_cache ? std::clamp(lctx.dflash_kv_cache_view_n_filled, 0, (int32_t) ctx_len) : 0;
const int32_t cache_write_pos = use_kv_cache && ctx_len > 0
const int32_t cache_write_pos = ctx_len > 0
? ((lctx.dflash_kv_cache_view_write_pos % (int32_t) ctx_len) + (int32_t) ctx_len) % (int32_t) ctx_len
: 0;
const int64_t n_kv_total = GGML_PAD(ctx_len + n_tokens, flash_attn ? 256 : 32);
@@ -286,8 +260,8 @@ ggml_cgraph * llm_build_context::build_dflash() {
GGML_ASSERT(n_embd_head_k == n_embd_head_v);
GGML_ASSERT(n_target_features > 0);
GGML_ASSERT(!use_kv_cache || lctx.ensure_dflash_kv_cache_tensors((int32_t) ctx_len));
GGML_ASSERT(!use_kv_cache || (cache_write_pos >= 0 && cache_write_pos < ctx_len));
GGML_ASSERT(lctx.ensure_dflash_kv_cache_tensors((int32_t) ctx_len));
GGML_ASSERT(cache_write_pos >= 0 && cache_write_pos < ctx_len);
ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes((int) std::max<int64_t>(n_tokens, ctx_len)) + 32 * n_layer, false);
@@ -316,22 +290,6 @@ ggml_cgraph * llm_build_context::build_dflash() {
dflash_kq_mask_swa = flash_attn ? ggml_cast(ctx0, lctx.inp_dflash_kq_mask_swa, GGML_TYPE_F16) : lctx.inp_dflash_kq_mask_swa;
}
ggml_tensor * fused_target = nullptr;
ggml_tensor * pos_ctx = nullptr;
if (!use_kv_cache) {
lctx.inp_dflash_target_features = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_target_features, ctx_len);
ggml_set_input(lctx.inp_dflash_target_features);
cb(lctx.inp_dflash_target_features, "dflash_target_features", -1);
lctx.inp_dflash_pos_ctx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ctx_len);
ggml_set_input(lctx.inp_dflash_pos_ctx);
cb(lctx.inp_dflash_pos_ctx, "dflash_pos_ctx", -1);
fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, lctx.inp_dflash_target_features);
fused_target = llm_build_norm(ctx0, fused_target, hparams, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, cb, -1);
pos_ctx = lctx.inp_dflash_pos_ctx;
}
ggml_tensor * tok_embd = model.tok_embd;
if (tok_embd == nullptr) {
tok_embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_Q4_0, n_embd, hparams.n_vocab);
@@ -370,147 +328,58 @@ ggml_cgraph * llm_build_context::build_dflash() {
Vcur_noise = ggml_reshape_3d(ctx0, Vcur_noise, n_embd_head_v, n_head_kv, n_tokens);
cb(Vcur_noise, "Vcur_noise", il);
const int64_t t_cache_read_us = use_kv_cache ? ggml_time_us() : 0;
ggml_tensor * Kcur_ctx = nullptr;
ggml_tensor * Vcur_ctx = nullptr;
const bool have_workspace_ctx = use_kv_workspace &&
(size_t) il < lctx.dflash_k_ctx_workspace.size() &&
(size_t) il < lctx.dflash_v_ctx_workspace.size() &&
lctx.dflash_k_ctx_workspace[(size_t) il] != nullptr &&
lctx.dflash_v_ctx_workspace[(size_t) il] != nullptr;
const int64_t t_cache_read_us = ggml_time_us();
GGML_ASSERT((size_t) il < lctx.dflash_k_ctx_workspace.size());
GGML_ASSERT((size_t) il < lctx.dflash_v_ctx_workspace.size());
GGML_ASSERT(lctx.dflash_k_ctx_workspace[(size_t) il] != nullptr);
GGML_ASSERT(lctx.dflash_v_ctx_workspace[(size_t) il] != nullptr);
if (have_workspace_ctx) {
Kcur_ctx = ggml_view_3d(ctx0, lctx.dflash_k_ctx_workspace[(size_t) il],
lctx.dflash_k_ctx_workspace[(size_t) il]->ne[0],
ctx_len,
lctx.dflash_k_ctx_workspace[(size_t) il]->ne[2],
lctx.dflash_k_ctx_workspace[(size_t) il]->nb[1],
lctx.dflash_k_ctx_workspace[(size_t) il]->nb[2],
0);
Vcur_ctx = ggml_view_3d(ctx0, lctx.dflash_v_ctx_workspace[(size_t) il],
lctx.dflash_v_ctx_workspace[(size_t) il]->ne[0],
ctx_len,
lctx.dflash_v_ctx_workspace[(size_t) il]->ne[2],
lctx.dflash_v_ctx_workspace[(size_t) il]->nb[1],
lctx.dflash_v_ctx_workspace[(size_t) il]->nb[2],
0);
cb(Kcur_ctx, "Kcur_ctx_workspace", il);
cb(Vcur_ctx, "Vcur_ctx_workspace", il);
} else if (use_kv_cache) {
auto build_ordered_cache_view = [&](ggml_tensor * cache) -> ggml_tensor * {
if (!lctx.dflash_kv_cache_view_valid || cache_rows <= 0) {
return cache;
}
ggml_tensor * Kcur_ctx = ggml_view_3d(ctx0, lctx.dflash_k_ctx_workspace[(size_t) il],
lctx.dflash_k_ctx_workspace[(size_t) il]->ne[0],
ctx_len,
lctx.dflash_k_ctx_workspace[(size_t) il]->ne[2],
lctx.dflash_k_ctx_workspace[(size_t) il]->nb[1],
lctx.dflash_k_ctx_workspace[(size_t) il]->nb[2],
0);
ggml_tensor * Vcur_ctx = ggml_view_3d(ctx0, lctx.dflash_v_ctx_workspace[(size_t) il],
lctx.dflash_v_ctx_workspace[(size_t) il]->ne[0],
ctx_len,
lctx.dflash_v_ctx_workspace[(size_t) il]->ne[2],
lctx.dflash_v_ctx_workspace[(size_t) il]->nb[1],
lctx.dflash_v_ctx_workspace[(size_t) il]->nb[2],
0);
cb(Kcur_ctx, "Kcur_ctx_workspace", il);
cb(Vcur_ctx, "Vcur_ctx_workspace", il);
if (cache_rows < ctx_len) {
ggml_tensor * zero_pad = ggml_view_3d(ctx0, cache,
cache->ne[0],
cache->ne[1],
ctx_len - cache_rows,
cache->nb[1],
cache->nb[2],
(size_t) cache_rows * cache->nb[2]);
ggml_tensor * valid = ggml_view_3d(ctx0, cache,
cache->ne[0],
cache->ne[1],
cache_rows,
cache->nb[1],
cache->nb[2],
0);
return ggml_concat(ctx0, zero_pad, valid, 2);
}
ggml_tensor * Kcur_draft = ggml_cont(ctx0, ggml_permute(ctx0, Kcur_noise, 0, 2, 1, 3));
ggml_tensor * Vcur_draft = ggml_cont(ctx0, ggml_permute(ctx0, Vcur_noise, 0, 2, 1, 3));
cb(Kcur_draft, "dflash_main_k_perm_cont", il);
cb(Vcur_draft, "dflash_main_v_perm_cont", il);
if (cache_write_pos == 0) {
return cache;
}
ggml_tensor * Kcur = ggml_concat(ctx0, Kcur_ctx, Kcur_draft, 1);
ggml_tensor * Vcur = ggml_concat(ctx0, Vcur_ctx, Vcur_draft, 1);
cb(Kcur, "dflash_main_k_concat", il);
cb(Vcur, "dflash_main_v_concat", il);
ggml_tensor * tail = ggml_view_3d(ctx0, cache,
cache->ne[0],
cache->ne[1],
ctx_len - cache_write_pos,
cache->nb[1],
cache->nb[2],
(size_t) cache_write_pos * cache->nb[2]);
ggml_tensor * head = ggml_view_3d(ctx0, cache,
cache->ne[0],
cache->ne[1],
cache_write_pos,
cache->nb[1],
cache->nb[2],
0);
return ggml_concat(ctx0, tail, head, 2);
};
Kcur_ctx = build_ordered_cache_view(lctx.dflash_k_ctx_cache[(size_t) il]);
Vcur_ctx = build_ordered_cache_view(lctx.dflash_v_ctx_cache[(size_t) il]);
cb(Kcur_ctx, "Kcur_ctx_cache", il);
cb(Vcur_ctx, "Vcur_ctx_cache", il);
} else {
Kcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, fused_target);
Kcur_ctx = ggml_reshape_3d(ctx0, Kcur_ctx, n_embd_head_k, n_head_kv, ctx_len);
Kcur_ctx = llm_build_norm(ctx0, Kcur_ctx, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il);
Kcur_ctx = ggml_rope_ext(ctx0, Kcur_ctx, pos_ctx, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target);
Vcur_ctx = ggml_reshape_3d(ctx0, Vcur_ctx, n_embd_head_v, n_head_kv, ctx_len);
cb(Kcur_ctx, "Kcur_ctx", il);
cb(Vcur_ctx, "Vcur_ctx", il);
if (n_kv_pad > 0) {
Kcur = ggml_pad(ctx0, Kcur, 0, (int) n_kv_pad, 0, 0);
Vcur = ggml_pad(ctx0, Vcur, 0, (int) n_kv_pad, 0, 0);
cb(Kcur, "dflash_main_k_pad", il);
cb(Vcur, "dflash_main_v_pad", il);
}
ggml_tensor * Kcur = nullptr;
ggml_tensor * Vcur = nullptr;
if (have_workspace_ctx) {
ggml_tensor * Kcur_draft = ggml_cont(ctx0, ggml_permute(ctx0, Kcur_noise, 0, 2, 1, 3));
ggml_tensor * Vcur_draft = ggml_cont(ctx0, ggml_permute(ctx0, Vcur_noise, 0, 2, 1, 3));
cb(Kcur_draft, "dflash_main_k_perm_cont", il);
cb(Vcur_draft, "dflash_main_v_perm_cont", il);
Kcur = ggml_concat(ctx0, Kcur_ctx, Kcur_draft, 1);
Vcur = ggml_concat(ctx0, Vcur_ctx, Vcur_draft, 1);
cb(Kcur, "dflash_main_k_concat", il);
cb(Vcur, "dflash_main_v_concat", il);
if (n_kv_pad > 0) {
Kcur = ggml_pad(ctx0, Kcur, 0, (int) n_kv_pad, 0, 0);
Vcur = ggml_pad(ctx0, Vcur, 0, (int) n_kv_pad, 0, 0);
cb(Kcur, "dflash_main_k_pad", il);
cb(Vcur, "dflash_main_v_pad", il);
}
} else {
ggml_tensor * Kcur_concat = ggml_concat(ctx0, Kcur_ctx, Kcur_noise, 2);
ggml_tensor * Vcur_concat = ggml_concat(ctx0, Vcur_ctx, Vcur_noise, 2);
cb(Kcur_concat, "dflash_main_k_concat", il);
cb(Vcur_concat, "dflash_main_v_concat", il);
Kcur = Kcur_concat;
Vcur = Vcur_concat;
if (n_kv_pad > 0) {
Kcur = ggml_pad(ctx0, Kcur, 0, 0, (int) n_kv_pad, 0);
Vcur = ggml_pad(ctx0, Vcur, 0, 0, (int) n_kv_pad, 0);
cb(Kcur, "dflash_main_k_pad", il);
cb(Vcur, "dflash_main_v_pad", il);
}
}
if (use_kv_cache) {
profile.graph_kv_cache_read_concat_pad_us += (uint64_t) (ggml_time_us() - t_cache_read_us);
profile.graph_kv_cache_read_concat_pad_calls++;
profile.graph_kv_cache_cached_bytes += ggml_nbytes(lctx.dflash_k_ctx_cache[(size_t) il]) + ggml_nbytes(lctx.dflash_v_ctx_cache[(size_t) il]);
}
profile.graph_kv_cache_read_concat_pad_us += (uint64_t) (ggml_time_us() - t_cache_read_us);
profile.graph_kv_cache_read_concat_pad_calls++;
profile.graph_kv_cache_cached_bytes += ggml_nbytes(lctx.dflash_k_ctx_cache[(size_t) il]) + ggml_nbytes(lctx.dflash_v_ctx_cache[(size_t) il]);
cb(Qcur, "Qcur", il);
ggml_tensor * q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
ggml_tensor * k = have_workspace_ctx ? Kcur : ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 0, 2, 1, 3));
ggml_tensor * v = have_workspace_ctx ? Vcur : ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 0, 2, 1, 3));
ggml_tensor * k = Kcur;
ggml_tensor * v = Vcur;
ggml_tensor * dflash_kq_mask_l = (hparams.swa_layers[il] && dflash_kq_mask_swa != nullptr)
? dflash_kq_mask_swa
: dflash_kq_mask_full;
cb(q, "q", il);
if (!have_workspace_ctx) {
cb(k, "dflash_main_k_perm_cont", il);
cb(v, "dflash_main_v_perm_cont", il);
}
cur = ggml_flash_attn_ext(ctx0, q, k, v, dflash_kq_mask_l, kq_scale, hparams.f_max_alibi_bias,
hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
src/llama-dflash.cpp
+288 -332
View File
@@ -23,6 +23,10 @@ static bool llama_env_flag_enabled_local(const char * name) {
std::strcmp(env, "off") != 0;
}
static bool llama_dflash_stats_log_enabled() {
return llama_env_flag_enabled_local("IK_DFLASH_STATS_LOG");
}
enum llama_dflash_kv_node_kind {
LLAMA_DFLASH_KV_NODE_NONE = 0,
LLAMA_DFLASH_KV_NODE_FUSED_TARGET,
@@ -359,10 +363,6 @@ static bool llama_dflash_main_node_eval_callback(struct ggml_tensor * tensor, bo
return prev_result || tracked;
}
static bool llama_dflash_use_kv_workspace_experiment() {
return llama_env_flag_enabled_local("IK_DFLASH_KV_WORKSPACE");
}
void llama_sync_dflash_workspace_if_pending(struct llama_context & lctx) {
if (!lctx.dflash_kv_workspace_sync_pending || lctx.dflash_workspace_sched == nullptr) {
return;
@@ -413,7 +413,6 @@ static ggml_backend_t llama_backend_for_tensor(const llama_context & lctx, const
}
bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
const bool use_kv_workspace = llama_env_flag_enabled_local("IK_DFLASH_KV_WORKSPACE");
const int32_t target_cross_ctx = std::max<int32_t>(1, cross_ctx);
const int32_t target_token_capacity = std::max<int32_t>(1, (int32_t) model.hparams.dflash_block_size);
const int32_t target_workspace_n_kv_total = GGML_PAD(target_cross_ctx + target_token_capacity, cparams.flash_attn ? 256 : 32);
@@ -426,11 +425,9 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
const bool cache_matches = (int32_t) dflash_k_ctx_cache.size() == n_layer &&
dflash_k_ctx_cache.front() != nullptr &&
(int32_t) dflash_k_ctx_cache.front()->ne[2] == target_cross_ctx;
const bool workspace_matches = use_kv_workspace
? ((int32_t) dflash_k_ctx_workspace.size() == n_layer &&
dflash_k_ctx_workspace.front() != nullptr &&
(int32_t) dflash_k_ctx_workspace.front()->ne[1] == target_workspace_n_kv_total)
: dflash_k_ctx_workspace.empty() && dflash_v_ctx_workspace.empty();
const bool workspace_matches = (int32_t) dflash_k_ctx_workspace.size() == n_layer &&
dflash_k_ctx_workspace.front() != nullptr &&
(int32_t) dflash_k_ctx_workspace.front()->ne[1] == target_workspace_n_kv_total;
if (cache_matches && workspace_matches) {
return true;
@@ -457,7 +454,7 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
}
ggml_init_params params = {
/*.mem_size =*/ (size_t) ((use_kv_workspace ? 4 : 2) * std::max(1, n_layer)) * ggml_tensor_overhead(),
/*.mem_size =*/ (size_t) (4 * std::max(1, n_layer)) * ggml_tensor_overhead(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
@@ -471,12 +468,10 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
dflash_v_ctx_cache.resize((size_t) n_layer);
dflash_k_ctx_workspace.clear();
dflash_v_ctx_workspace.clear();
if (use_kv_workspace) {
dflash_k_ctx_workspace.resize((size_t) n_layer);
dflash_v_ctx_workspace.resize((size_t) n_layer);
}
dflash_k_ctx_workspace.resize((size_t) n_layer);
dflash_v_ctx_workspace.resize((size_t) n_layer);
dflash_cache_bufs.clear();
dflash_cache_bufs.reserve((size_t) std::max(1, n_layer) * (use_kv_workspace ? 4 : 2));
dflash_cache_bufs.reserve((size_t) std::max(1, n_layer) * 4);
int32_t host_layers = 0;
const char * first_buft_name = nullptr;
const char * last_buft_name = nullptr;
@@ -524,54 +519,54 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
ggml_backend_buffer_clear(v_buf, 0);
dflash_cache_bufs.push_back(v_buf);
if (use_kv_workspace) {
dflash_k_ctx_workspace[(size_t) il] = ggml_new_tensor_3d(dflash_cache_ctx, GGML_TYPE_F32, n_embd_head_k, target_workspace_n_kv_total, n_head_kv);
dflash_v_ctx_workspace[(size_t) il] = ggml_new_tensor_3d(dflash_cache_ctx, GGML_TYPE_F32, n_embd_head_v, target_workspace_n_kv_total, n_head_kv);
if (dflash_k_ctx_workspace[(size_t) il] == nullptr || dflash_v_ctx_workspace[(size_t) il] == nullptr) {
free_dflash_kv_cache_tensors();
return false;
}
ggml_set_input(dflash_k_ctx_workspace[(size_t) il]);
ggml_set_input(dflash_v_ctx_workspace[(size_t) il]);
ggml_format_name(dflash_k_ctx_workspace[(size_t) il], "dflash_k_ctx_workspace_%d", il);
ggml_format_name(dflash_v_ctx_workspace[(size_t) il], "dflash_v_ctx_workspace_%d", il);
const size_t k_workspace_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_k_ctx_workspace[(size_t) il]);
ggml_backend_buffer_t k_workspace_buf = ggml_backend_buft_alloc_buffer(layer_buft, k_workspace_bytes);
if (k_workspace_buf == nullptr) {
free_dflash_kv_cache_tensors();
return false;
}
ggml_backend_buffer_set_usage(k_workspace_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE);
ggml_backend_tensor_alloc(k_workspace_buf, dflash_k_ctx_workspace[(size_t) il], ggml_backend_buffer_get_base(k_workspace_buf));
ggml_backend_buffer_clear(k_workspace_buf, 0);
dflash_cache_bufs.push_back(k_workspace_buf);
const size_t v_workspace_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_v_ctx_workspace[(size_t) il]);
ggml_backend_buffer_t v_workspace_buf = ggml_backend_buft_alloc_buffer(layer_buft, v_workspace_bytes);
if (v_workspace_buf == nullptr) {
free_dflash_kv_cache_tensors();
return false;
}
ggml_backend_buffer_set_usage(v_workspace_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE);
ggml_backend_tensor_alloc(v_workspace_buf, dflash_v_ctx_workspace[(size_t) il], ggml_backend_buffer_get_base(v_workspace_buf));
ggml_backend_buffer_clear(v_workspace_buf, 0);
dflash_cache_bufs.push_back(v_workspace_buf);
dflash_k_ctx_workspace[(size_t) il] = ggml_new_tensor_3d(dflash_cache_ctx, GGML_TYPE_F32, n_embd_head_k, target_workspace_n_kv_total, n_head_kv);
dflash_v_ctx_workspace[(size_t) il] = ggml_new_tensor_3d(dflash_cache_ctx, GGML_TYPE_F32, n_embd_head_v, target_workspace_n_kv_total, n_head_kv);
if (dflash_k_ctx_workspace[(size_t) il] == nullptr || dflash_v_ctx_workspace[(size_t) il] == nullptr) {
free_dflash_kv_cache_tensors();
return false;
}
ggml_set_input(dflash_k_ctx_workspace[(size_t) il]);
ggml_set_input(dflash_v_ctx_workspace[(size_t) il]);
ggml_format_name(dflash_k_ctx_workspace[(size_t) il], "dflash_k_ctx_workspace_%d", il);
ggml_format_name(dflash_v_ctx_workspace[(size_t) il], "dflash_v_ctx_workspace_%d", il);
const size_t k_workspace_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_k_ctx_workspace[(size_t) il]);
ggml_backend_buffer_t k_workspace_buf = ggml_backend_buft_alloc_buffer(layer_buft, k_workspace_bytes);
if (k_workspace_buf == nullptr) {
free_dflash_kv_cache_tensors();
return false;
}
ggml_backend_buffer_set_usage(k_workspace_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE);
ggml_backend_tensor_alloc(k_workspace_buf, dflash_k_ctx_workspace[(size_t) il], ggml_backend_buffer_get_base(k_workspace_buf));
ggml_backend_buffer_clear(k_workspace_buf, 0);
dflash_cache_bufs.push_back(k_workspace_buf);
const size_t v_workspace_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_v_ctx_workspace[(size_t) il]);
ggml_backend_buffer_t v_workspace_buf = ggml_backend_buft_alloc_buffer(layer_buft, v_workspace_bytes);
if (v_workspace_buf == nullptr) {
free_dflash_kv_cache_tensors();
return false;
}
ggml_backend_buffer_set_usage(v_workspace_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE);
ggml_backend_tensor_alloc(v_workspace_buf, dflash_v_ctx_workspace[(size_t) il], ggml_backend_buffer_get_base(v_workspace_buf));
ggml_backend_buffer_clear(v_workspace_buf, 0);
dflash_cache_bufs.push_back(v_workspace_buf);
}
dflash_profile.last_kv_cache_host_layers = host_layers;
dflash_kv_workspace_token_capacity = use_kv_workspace ? target_token_capacity : 0;
dflash_kv_workspace_n_kv_total = use_kv_workspace ? target_workspace_n_kv_total : 0;
dflash_kv_workspace_token_capacity = target_token_capacity;
dflash_kv_workspace_n_kv_total = target_workspace_n_kv_total;
llama_reset_dflash_kv_cache_state(this);
LLAMA_LOG_INFO("%s: DFlash K/V cache placement cross_ctx=%d host_layers=%d/%d first=%s last=%s\n",
__func__,
target_cross_ctx,
host_layers,
n_layer,
first_buft_name != nullptr ? first_buft_name : "(none)",
last_buft_name != nullptr ? last_buft_name : "(none)");
if (llama_dflash_stats_log_enabled()) {
LLAMA_LOG_INFO("%s: DFlash K/V cache placement cross_ctx=%d host_layers=%d/%d first=%s last=%s\n",
__func__,
target_cross_ctx,
host_layers,
n_layer,
first_buft_name != nullptr ? first_buft_name : "(none)",
last_buft_name != nullptr ? last_buft_name : "(none)");
}
return true;
}
@@ -758,8 +753,6 @@ static bool validate_dflash_graph_contract(const llama_context & lctx) {
bool llama_prepare_dflash_graph_inputs(
struct llama_context & lctx,
uint32_t n_tokens) {
const bool use_kv_cache = llama_env_flag_enabled_local("IK_DFLASH_KV_CACHE");
const bool use_kv_workspace = use_kv_cache && llama_dflash_use_kv_workspace_experiment();
const bool kv_node_timing = llama_env_flag_enabled_local("IK_DFLASH_KV_NODE_TIMING");
auto & profile = lctx.dflash_profile;
const int32_t cross_ctx = lctx.dflash_visible_cross_ctx > 0
@@ -778,13 +771,8 @@ bool llama_prepare_dflash_graph_inputs(
return false;
}
if (use_kv_cache) {
if (!lctx.ensure_dflash_kv_cache_tensors(cross_ctx) || lctx.dflash_k_ctx_cache.empty() || lctx.dflash_v_ctx_cache.empty()) {
LLAMA_LOG_ERROR("%s: DFlash K/V cache inputs are not initialized\n", __func__);
return false;
}
} else if (lctx.inp_dflash_target_features == nullptr || lctx.inp_dflash_pos_ctx == nullptr) {
LLAMA_LOG_ERROR("%s: DFlash inline inputs are not initialized\n", __func__);
if (!lctx.ensure_dflash_kv_cache_tensors(cross_ctx) || lctx.dflash_k_ctx_cache.empty() || lctx.dflash_v_ctx_cache.empty()) {
LLAMA_LOG_ERROR("%s: DFlash K/V cache inputs are not initialized\n", __func__);
return false;
}
@@ -797,9 +785,9 @@ bool llama_prepare_dflash_graph_inputs(
const int32_t n_rows = lctx.dflash_target_features_n_rows;
const int32_t append_rows_available = lctx.dflash_target_append_features_n_rows;
const int32_t width = (int32_t) lctx.model.hparams.dflash_n_target_features;
const int32_t graph_cross_ctx = use_kv_cache
? (lctx.dflash_k_ctx_cache.front() != nullptr ? (int32_t) lctx.dflash_k_ctx_cache.front()->ne[2] : 0)
: (lctx.inp_dflash_target_features != nullptr ? (int32_t) lctx.inp_dflash_target_features->ne[1] : 0);
const int32_t graph_cross_ctx = lctx.dflash_k_ctx_cache.front() != nullptr
? (int32_t) lctx.dflash_k_ctx_cache.front()->ne[2]
: 0;
const int32_t n_mask_tokens = (int32_t) kq_mask->ne[1];
const int32_t n_kv_total = (int32_t) kq_mask->ne[0];
const int64_t t_total_us = ggml_time_us();
@@ -811,9 +799,7 @@ bool llama_prepare_dflash_graph_inputs(
profile.last_n_tokens = (int32_t) n_tokens;
profile.last_n_kv_total = n_kv_total;
if (use_kv_workspace) {
llama_sync_dflash_workspace_if_pending(lctx);
}
llama_sync_dflash_workspace_if_pending(lctx);
if (graph_cross_ctx != cross_ctx) {
profile.graph_shape_failures++;
@@ -836,12 +822,6 @@ bool llama_prepare_dflash_graph_inputs(
return false;
}
if (!use_kv_cache && !have_full_src) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: missing contiguous DFlash target features for inline path\n", __func__);
return false;
}
if (n_kv_total < cross_ctx + (int32_t) n_tokens) {
profile.graph_mask_overflow++;
LLAMA_LOG_ERROR("%s: invalid DFlash mask shape (n_kv_total=%d < cross_ctx+n_tokens=%d)\n",
@@ -851,25 +831,6 @@ bool llama_prepare_dflash_graph_inputs(
const int32_t left_pad = cross_ctx - n_rows;
profile.last_left_pad = left_pad;
if (!use_kv_cache) {
const size_t padded_floats = (size_t) cross_ctx * (size_t) width;
const size_t dst_offset = (size_t) left_pad * (size_t) width;
const int64_t t_feature_us = ggml_time_us();
if (lctx.dflash_target_features_padded.size() != padded_floats) {
lctx.dflash_target_features_padded.resize(padded_floats);
}
if (left_pad == 0 && total_floats == padded_floats) {
std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin());
} else {
if (dst_offset > 0) {
std::fill(lctx.dflash_target_features_padded.begin(),
lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset, 0.0f);
}
std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset);
}
profile.graph_feature_copy_us += (uint64_t) (ggml_time_us() - t_feature_us);
profile.graph_feature_bytes += padded_floats * sizeof(float);
}
const int64_t t_pos_us = ggml_time_us();
lctx.dflash_pos_ctx_data.resize((size_t) cross_ctx);
@@ -903,274 +864,269 @@ bool llama_prepare_dflash_graph_inputs(
profile.graph_pos_copy_us += (uint64_t) (ggml_time_us() - t_pos_us);
profile.graph_pos_bytes += lctx.dflash_pos_ctx_data.size() * sizeof(llama_pos);
if (use_kv_cache) {
const llama_dflash_kv_cache_transition cache_plan = llama_plan_dflash_kv_cache_transition(
cross_ctx,
lctx.dflash_kv_cache_n_filled,
lctx.dflash_kv_cache_write_pos,
lctx.dflash_kv_cache_valid,
lctx.dflash_kv_cache_applied_window_version,
lctx.dflash_target_window_version,
lctx.dflash_target_window_keep_rows,
lctx.dflash_target_window_append_rows,
lctx.dflash_target_window_replace,
n_rows);
const llama_dflash_kv_cache_transition cache_plan = llama_plan_dflash_kv_cache_transition(
cross_ctx,
lctx.dflash_kv_cache_n_filled,
lctx.dflash_kv_cache_write_pos,
lctx.dflash_kv_cache_valid,
lctx.dflash_kv_cache_applied_window_version,
lctx.dflash_target_window_version,
lctx.dflash_target_window_keep_rows,
lctx.dflash_target_window_append_rows,
lctx.dflash_target_window_replace,
n_rows);
const bool have_append_src = append_src != nullptr &&
append_rows_available == cache_plan.append_rows &&
append_floats == (size_t) cache_plan.append_rows * (size_t) width;
const bool have_append_src = append_src != nullptr &&
append_rows_available == cache_plan.append_rows &&
append_floats == (size_t) cache_plan.append_rows * (size_t) width;
const int32_t update_rows = cache_plan.cache_up_to_date
? 0
: (cache_plan.rebuild_cache ? n_rows : cache_plan.append_rows);
const size_t max_nodes = lctx.model.max_nodes((int) std::max<int32_t>(1, cross_ctx)) + 24 * lctx.model.hparams.n_layer;
const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false);
if (lctx.dflash_buf_compute_meta.size() != meta_size) {
lctx.dflash_buf_compute_meta.resize(meta_size);
const int32_t update_rows = cache_plan.cache_up_to_date
? 0
: (cache_plan.rebuild_cache ? n_rows : cache_plan.append_rows);
const size_t max_nodes = lctx.model.max_nodes((int) std::max<int32_t>(1, cross_ctx)) + 24 * lctx.model.hparams.n_layer;
const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false);
if (lctx.dflash_buf_compute_meta.size() != meta_size) {
lctx.dflash_buf_compute_meta.resize(meta_size);
}
if (lctx.dflash_sched == nullptr || lctx.dflash_kv_cache_reserved_rows != cross_ctx) {
std::vector<ggml_backend_buffer_type_t> backend_buft;
backend_buft.reserve(lctx.backends.size());
for (auto * backend : lctx.backends) {
if (ggml_backend_is_cpu(backend)) {
backend_buft.push_back(llama_default_buffer_type_cpu(true));
} else {
backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
}
}
if (lctx.dflash_sched == nullptr || lctx.dflash_kv_cache_reserved_rows != cross_ctx) {
std::vector<ggml_backend_buffer_type_t> backend_buft;
backend_buft.reserve(lctx.backends.size());
for (auto * backend : lctx.backends) {
if (ggml_backend_is_cpu(backend)) {
backend_buft.push_back(llama_default_buffer_type_cpu(true));
} else {
backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
}
}
if (lctx.dflash_sched != nullptr) {
ggml_backend_sched_free(lctx.dflash_sched);
lctx.dflash_sched = nullptr;
}
lctx.dflash_kv_graph = nullptr;
lctx.dflash_kv_graph_rows = 0;
lctx.dflash_kv_graph_write_pos = 0;
if (lctx.dflash_sched != nullptr) {
ggml_backend_sched_free(lctx.dflash_sched);
lctx.dflash_sched = nullptr;
}
lctx.dflash_kv_graph = nullptr;
lctx.dflash_kv_graph_rows = 0;
lctx.dflash_kv_graph_write_pos = 0;
const int32_t saved_update_rows = lctx.dflash_kv_cache_update_rows;
lctx.dflash_kv_cache_update_rows = cross_ctx;
const int64_t t_build_us = ggml_time_us();
ggml_cgraph * gf_reserve = llm_build_context::llama_build_graph_dflash_kv_cache(lctx);
profile.graph_kv_cache_build_us += (uint64_t) (ggml_time_us() - t_build_us);
lctx.dflash_kv_cache_update_rows = saved_update_rows;
if (gf_reserve == nullptr) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache reserve graph\n", __func__);
return false;
}
const int32_t saved_update_rows = lctx.dflash_kv_cache_update_rows;
lctx.dflash_kv_cache_update_rows = cross_ctx;
const int64_t t_reserve_us = ggml_time_us();
lctx.dflash_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false);
const bool reserved = lctx.dflash_sched != nullptr && ggml_backend_sched_reserve(lctx.dflash_sched, gf_reserve);
profile.graph_kv_cache_reserve_us += (uint64_t) (ggml_time_us() - t_reserve_us);
if (!reserved) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V scheduler\n", __func__);
return false;
}
lctx.dflash_kv_cache_reserved_rows = cross_ctx;
}
if (update_rows > 0) {
const float * update_src = nullptr;
if (have_append_src && update_rows == cache_plan.append_rows) {
update_src = append_src;
} else if (have_full_src) {
update_src = src + (size_t) (n_rows - update_rows) * (size_t) width;
}
const llama_pos * update_pos = src_pos + (n_rows - update_rows);
if (update_src == nullptr) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: missing DFlash appended target features for cached update (rows=%d append_rows=%d floats=%zu)\n",
__func__, n_rows, update_rows, append_floats);
return false;
}
if (cache_plan.rebuild_cache) {
llama_reset_dflash_kv_cache_state(&lctx);
}
lctx.dflash_kv_cache_update_rows = update_rows;
ggml_cgraph * gf_kv = nullptr;
const bool can_reuse_kv_graph = lctx.dflash_kv_graph != nullptr &&
lctx.dflash_kv_graph_rows == update_rows &&
lctx.dflash_kv_graph_write_pos == lctx.dflash_kv_cache_write_pos;
if (can_reuse_kv_graph) {
gf_kv = lctx.dflash_kv_graph;
} else {
const int64_t t_build_us = ggml_time_us();
ggml_cgraph * gf_reserve = llm_build_context::llama_build_graph_dflash_kv_cache(lctx);
gf_kv = llm_build_context::llama_build_graph_dflash_kv_cache(lctx);
profile.graph_kv_cache_build_us += (uint64_t) (ggml_time_us() - t_build_us);
lctx.dflash_kv_cache_update_rows = saved_update_rows;
if (gf_reserve == nullptr) {
if (gf_kv == nullptr || lctx.dflash_kv_input_target_features == nullptr || lctx.dflash_kv_input_pos_ctx == nullptr) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache reserve graph\n", __func__);
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache graph\n", __func__);
return false;
}
const int64_t t_reserve_us = ggml_time_us();
lctx.dflash_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false);
const bool reserved = lctx.dflash_sched != nullptr && ggml_backend_sched_reserve(lctx.dflash_sched, gf_reserve);
profile.graph_kv_cache_reserve_us += (uint64_t) (ggml_time_us() - t_reserve_us);
if (!reserved) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V scheduler\n", __func__);
return false;
}
lctx.dflash_kv_cache_reserved_rows = cross_ctx;
const int64_t t_reset_us = ggml_time_us();
ggml_backend_sched_reset(lctx.dflash_sched);
profile.graph_kv_cache_reset_us += (uint64_t) (ggml_time_us() - t_reset_us);
const int64_t t_alloc_us = ggml_time_us();
ggml_backend_sched_alloc_graph(lctx.dflash_sched, gf_kv);
profile.graph_kv_cache_alloc_us += (uint64_t) (ggml_time_us() - t_alloc_us);
lctx.dflash_kv_graph = gf_kv;
lctx.dflash_kv_graph_rows = update_rows;
lctx.dflash_kv_graph_write_pos = lctx.dflash_kv_cache_write_pos;
}
if (update_rows > 0) {
const float * update_src = nullptr;
if (have_append_src && update_rows == cache_plan.append_rows) {
update_src = append_src;
} else if (have_full_src) {
update_src = src + (size_t) (n_rows - update_rows) * (size_t) width;
}
const llama_pos * update_pos = src_pos + (n_rows - update_rows);
ggml_backend_t kv_feature_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_target_features);
const int64_t t_feature_upload_us = ggml_time_us();
if (kv_feature_backend != nullptr) {
ggml_backend_tensor_set_async(kv_feature_backend, lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features));
} else {
ggml_backend_tensor_set(lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features));
}
profile.graph_kv_cache_feature_upload_us += (uint64_t) (ggml_time_us() - t_feature_upload_us);
profile.graph_feature_bytes += (size_t) update_rows * (size_t) width * sizeof(float);
if (update_src == nullptr) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: missing DFlash appended target features for cached update (rows=%d append_rows=%d floats=%zu)\n",
__func__, n_rows, update_rows, append_floats);
return false;
ggml_backend_t kv_pos_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_pos_ctx);
const int64_t t_pos_upload_us = ggml_time_us();
if (kv_pos_backend != nullptr) {
ggml_backend_tensor_set_async(kv_pos_backend, lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx));
} else {
ggml_backend_tensor_set(lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx));
}
profile.graph_kv_cache_pos_upload_us += (uint64_t) (ggml_time_us() - t_pos_upload_us);
const int64_t t_kv_cache_us = ggml_time_us();
llama_dflash_kv_node_profiler kv_node_profiler;
if (kv_node_timing) {
kv_node_profiler.profile = &profile;
ggml_backend_sched_set_eval_callback(lctx.dflash_sched, llama_dflash_kv_node_eval_callback, &kv_node_profiler);
}
llama_graph_compute_sched(lctx, lctx.dflash_sched, gf_kv, lctx.cparams.n_threads);
if (kv_node_timing) {
ggml_backend_sched_set_eval_callback(lctx.dflash_sched, nullptr, nullptr);
}
profile.graph_kv_cache_compute_us += (uint64_t) (ggml_time_us() - t_kv_cache_us);
const int64_t t_sync_us = ggml_time_us();
ggml_backend_sched_synchronize(lctx.dflash_sched);
profile.graph_kv_cache_sync_us += (uint64_t) (ggml_time_us() - t_sync_us);
profile.graph_kv_cache_calls++;
lctx.dflash_kv_cache_n_filled = std::min(cross_ctx, lctx.dflash_kv_cache_n_filled + update_rows);
lctx.dflash_kv_cache_write_pos = (lctx.dflash_kv_cache_write_pos + update_rows) % cross_ctx;
lctx.dflash_kv_cache_applied_window_version = lctx.dflash_target_window_version;
lctx.dflash_kv_cache_valid = true;
lctx.dflash_kv_cache_view_n_filled = lctx.dflash_kv_cache_n_filled;
lctx.dflash_kv_cache_view_write_pos = lctx.dflash_kv_cache_write_pos;
lctx.dflash_kv_cache_view_valid = true;
}
if (lctx.dflash_kv_cache_view_valid &&
!lctx.dflash_k_ctx_workspace.empty() && !lctx.dflash_v_ctx_workspace.empty()) {
const bool need_workspace_refresh = !lctx.dflash_kv_workspace_valid ||
lctx.dflash_kv_workspace_n_filled != lctx.dflash_kv_cache_view_n_filled ||
lctx.dflash_kv_workspace_write_pos != lctx.dflash_kv_cache_view_write_pos ||
lctx.dflash_kv_workspace_applied_window_version != lctx.dflash_kv_cache_applied_window_version;
if (need_workspace_refresh) {
const size_t max_nodes = lctx.model.max_nodes((int) std::max<int32_t>(1, cross_ctx)) + 16 * lctx.model.hparams.n_layer;
const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false);
if (lctx.dflash_workspace_buf_compute_meta.size() != meta_size) {
lctx.dflash_workspace_buf_compute_meta.resize(meta_size);
}
if (cache_plan.rebuild_cache) {
llama_reset_dflash_kv_cache_state(&lctx);
}
ggml_cgraph * gf_workspace = nullptr;
const bool can_reuse_workspace_graph = lctx.dflash_kv_workspace_graph != nullptr &&
lctx.dflash_kv_workspace_graph_rows == lctx.dflash_kv_cache_view_n_filled &&
lctx.dflash_kv_workspace_graph_write_pos == lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_cache_update_rows = update_rows;
ggml_cgraph * gf_kv = nullptr;
const bool can_reuse_kv_graph = lctx.dflash_kv_graph != nullptr &&
lctx.dflash_kv_graph_rows == update_rows &&
lctx.dflash_kv_graph_write_pos == lctx.dflash_kv_cache_write_pos;
if (can_reuse_kv_graph) {
gf_kv = lctx.dflash_kv_graph;
if (can_reuse_workspace_graph) {
gf_workspace = lctx.dflash_kv_workspace_graph;
} else {
const int64_t t_build_us = ggml_time_us();
gf_kv = llm_build_context::llama_build_graph_dflash_kv_cache(lctx);
profile.graph_kv_cache_build_us += (uint64_t) (ggml_time_us() - t_build_us);
if (gf_kv == nullptr || lctx.dflash_kv_input_target_features == nullptr || lctx.dflash_kv_input_pos_ctx == nullptr) {
gf_workspace = llm_build_context::llama_build_graph_dflash_kv_workspace(lctx);
profile.graph_kv_workspace_build_us += (uint64_t) (ggml_time_us() - t_build_us);
if (gf_workspace == nullptr) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache graph\n", __func__);
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V workspace graph\n", __func__);
return false;
}
const int64_t t_reset_us = ggml_time_us();
ggml_backend_sched_reset(lctx.dflash_sched);
profile.graph_kv_cache_reset_us += (uint64_t) (ggml_time_us() - t_reset_us);
const int64_t t_alloc_us = ggml_time_us();
ggml_backend_sched_alloc_graph(lctx.dflash_sched, gf_kv);
profile.graph_kv_cache_alloc_us += (uint64_t) (ggml_time_us() - t_alloc_us);
lctx.dflash_kv_graph = gf_kv;
lctx.dflash_kv_graph_rows = update_rows;
lctx.dflash_kv_graph_write_pos = lctx.dflash_kv_cache_write_pos;
}
ggml_backend_t kv_feature_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_target_features);
const int64_t t_feature_upload_us = ggml_time_us();
if (kv_feature_backend != nullptr) {
ggml_backend_tensor_set_async(kv_feature_backend, lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features));
} else {
ggml_backend_tensor_set(lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features));
}
profile.graph_kv_cache_feature_upload_us += (uint64_t) (ggml_time_us() - t_feature_upload_us);
profile.graph_feature_bytes += (size_t) update_rows * (size_t) width * sizeof(float);
ggml_backend_t kv_pos_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_pos_ctx);
const int64_t t_pos_upload_us = ggml_time_us();
if (kv_pos_backend != nullptr) {
ggml_backend_tensor_set_async(kv_pos_backend, lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx));
} else {
ggml_backend_tensor_set(lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx));
}
profile.graph_kv_cache_pos_upload_us += (uint64_t) (ggml_time_us() - t_pos_upload_us);
const int64_t t_kv_cache_us = ggml_time_us();
llama_dflash_kv_node_profiler kv_node_profiler;
if (kv_node_timing) {
kv_node_profiler.profile = &profile;
ggml_backend_sched_set_eval_callback(lctx.dflash_sched, llama_dflash_kv_node_eval_callback, &kv_node_profiler);
}
llama_graph_compute_sched(lctx, lctx.dflash_sched, gf_kv, lctx.cparams.n_threads);
if (kv_node_timing) {
ggml_backend_sched_set_eval_callback(lctx.dflash_sched, nullptr, nullptr);
}
profile.graph_kv_cache_compute_us += (uint64_t) (ggml_time_us() - t_kv_cache_us);
const int64_t t_sync_us = ggml_time_us();
ggml_backend_sched_synchronize(lctx.dflash_sched);
profile.graph_kv_cache_sync_us += (uint64_t) (ggml_time_us() - t_sync_us);
profile.graph_kv_cache_calls++;
lctx.dflash_kv_cache_n_filled = std::min(cross_ctx, lctx.dflash_kv_cache_n_filled + update_rows);
lctx.dflash_kv_cache_write_pos = (lctx.dflash_kv_cache_write_pos + update_rows) % cross_ctx;
lctx.dflash_kv_cache_applied_window_version = lctx.dflash_target_window_version;
lctx.dflash_kv_cache_valid = true;
lctx.dflash_kv_cache_view_n_filled = lctx.dflash_kv_cache_n_filled;
lctx.dflash_kv_cache_view_write_pos = lctx.dflash_kv_cache_write_pos;
lctx.dflash_kv_cache_view_valid = true;
}
if (use_kv_workspace && lctx.dflash_kv_cache_view_valid &&
!lctx.dflash_k_ctx_workspace.empty() && !lctx.dflash_v_ctx_workspace.empty()) {
const bool need_workspace_refresh = !lctx.dflash_kv_workspace_valid ||
lctx.dflash_kv_workspace_n_filled != lctx.dflash_kv_cache_view_n_filled ||
lctx.dflash_kv_workspace_write_pos != lctx.dflash_kv_cache_view_write_pos ||
lctx.dflash_kv_workspace_applied_window_version != lctx.dflash_kv_cache_applied_window_version;
if (need_workspace_refresh) {
const size_t max_nodes = lctx.model.max_nodes((int) std::max<int32_t>(1, cross_ctx)) + 16 * lctx.model.hparams.n_layer;
const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false);
if (lctx.dflash_workspace_buf_compute_meta.size() != meta_size) {
lctx.dflash_workspace_buf_compute_meta.resize(meta_size);
std::vector<ggml_backend_buffer_type_t> backend_buft;
backend_buft.reserve(lctx.backends.size());
for (auto * backend : lctx.backends) {
if (ggml_backend_is_cpu(backend)) {
backend_buft.push_back(llama_default_buffer_type_cpu(true));
} else {
backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
}
}
ggml_cgraph * gf_workspace = nullptr;
const bool can_reuse_workspace_graph = lctx.dflash_kv_workspace_graph != nullptr &&
lctx.dflash_kv_workspace_graph_rows == lctx.dflash_kv_cache_view_n_filled &&
lctx.dflash_kv_workspace_graph_write_pos == lctx.dflash_kv_cache_view_write_pos;
if (lctx.dflash_workspace_sched == nullptr) {
lctx.dflash_workspace_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false);
}
if (can_reuse_workspace_graph) {
gf_workspace = lctx.dflash_kv_workspace_graph;
} else {
const int64_t t_build_us = ggml_time_us();
gf_workspace = llm_build_context::llama_build_graph_dflash_kv_workspace(lctx);
profile.graph_kv_workspace_build_us += (uint64_t) (ggml_time_us() - t_build_us);
if (gf_workspace == nullptr) {
if (lctx.dflash_kv_workspace_reserved_rows != cross_ctx) {
const bool saved_view_valid = lctx.dflash_kv_cache_view_valid;
const int32_t saved_view_rows = lctx.dflash_kv_cache_view_n_filled;
const int32_t saved_view_write_pos = lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_cache_view_valid = true;
lctx.dflash_kv_cache_view_n_filled = cross_ctx;
lctx.dflash_kv_cache_view_write_pos = cross_ctx > 1 ? 1 : 0;
const int64_t t_reserve_build_us = ggml_time_us();
ggml_cgraph * gf_workspace_reserve = llm_build_context::llama_build_graph_dflash_kv_workspace(lctx);
profile.graph_kv_workspace_build_us += (uint64_t) (ggml_time_us() - t_reserve_build_us);
lctx.dflash_kv_cache_view_valid = saved_view_valid;
lctx.dflash_kv_cache_view_n_filled = saved_view_rows;
lctx.dflash_kv_cache_view_write_pos = saved_view_write_pos;
const int64_t t_reserve_us = ggml_time_us();
const bool reserved = lctx.dflash_workspace_sched != nullptr &&
gf_workspace_reserve != nullptr &&
ggml_backend_sched_reserve(lctx.dflash_workspace_sched, gf_workspace_reserve);
profile.graph_kv_workspace_reserve_us += (uint64_t) (ggml_time_us() - t_reserve_us);
if (!reserved) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V workspace graph\n", __func__);
LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V workspace scheduler\n", __func__);
return false;
}
std::vector<ggml_backend_buffer_type_t> backend_buft;
backend_buft.reserve(lctx.backends.size());
for (auto * backend : lctx.backends) {
if (ggml_backend_is_cpu(backend)) {
backend_buft.push_back(llama_default_buffer_type_cpu(true));
} else {
backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
}
}
if (lctx.dflash_workspace_sched == nullptr) {
lctx.dflash_workspace_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false);
}
if (lctx.dflash_kv_workspace_reserved_rows != cross_ctx) {
const bool saved_view_valid = lctx.dflash_kv_cache_view_valid;
const int32_t saved_view_rows = lctx.dflash_kv_cache_view_n_filled;
const int32_t saved_view_write_pos = lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_cache_view_valid = true;
lctx.dflash_kv_cache_view_n_filled = cross_ctx;
lctx.dflash_kv_cache_view_write_pos = cross_ctx > 1 ? 1 : 0;
const int64_t t_reserve_build_us = ggml_time_us();
ggml_cgraph * gf_workspace_reserve = llm_build_context::llama_build_graph_dflash_kv_workspace(lctx);
profile.graph_kv_workspace_build_us += (uint64_t) (ggml_time_us() - t_reserve_build_us);
lctx.dflash_kv_cache_view_valid = saved_view_valid;
lctx.dflash_kv_cache_view_n_filled = saved_view_rows;
lctx.dflash_kv_cache_view_write_pos = saved_view_write_pos;
const int64_t t_reserve_us = ggml_time_us();
const bool reserved = lctx.dflash_workspace_sched != nullptr &&
gf_workspace_reserve != nullptr &&
ggml_backend_sched_reserve(lctx.dflash_workspace_sched, gf_workspace_reserve);
profile.graph_kv_workspace_reserve_us += (uint64_t) (ggml_time_us() - t_reserve_us);
if (!reserved) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V workspace scheduler\n", __func__);
return false;
}
lctx.dflash_kv_workspace_reserved_rows = cross_ctx;
}
const int64_t t_reset_us = ggml_time_us();
ggml_backend_sched_reset(lctx.dflash_workspace_sched);
profile.graph_kv_workspace_reset_us += (uint64_t) (ggml_time_us() - t_reset_us);
const int64_t t_alloc_us = ggml_time_us();
ggml_backend_sched_alloc_graph(lctx.dflash_workspace_sched, gf_workspace);
profile.graph_kv_workspace_alloc_us += (uint64_t) (ggml_time_us() - t_alloc_us);
lctx.dflash_kv_workspace_graph = gf_workspace;
lctx.dflash_kv_workspace_graph_rows = lctx.dflash_kv_cache_view_n_filled;
lctx.dflash_kv_workspace_graph_write_pos = lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_workspace_reserved_rows = cross_ctx;
}
const int64_t t_workspace_us = ggml_time_us();
llama_graph_compute_sched(lctx, lctx.dflash_workspace_sched, gf_workspace, lctx.cparams.n_threads);
profile.graph_kv_workspace_compute_us += (uint64_t) (ggml_time_us() - t_workspace_us);
lctx.dflash_kv_workspace_sync_pending = true;
profile.graph_kv_workspace_calls++;
const int64_t t_reset_us = ggml_time_us();
ggml_backend_sched_reset(lctx.dflash_workspace_sched);
profile.graph_kv_workspace_reset_us += (uint64_t) (ggml_time_us() - t_reset_us);
lctx.dflash_kv_workspace_n_filled = lctx.dflash_kv_cache_view_n_filled;
lctx.dflash_kv_workspace_write_pos = lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_workspace_applied_window_version = lctx.dflash_kv_cache_applied_window_version;
lctx.dflash_kv_workspace_valid = true;
const int64_t t_alloc_us = ggml_time_us();
ggml_backend_sched_alloc_graph(lctx.dflash_workspace_sched, gf_workspace);
profile.graph_kv_workspace_alloc_us += (uint64_t) (ggml_time_us() - t_alloc_us);
lctx.dflash_kv_workspace_graph = gf_workspace;
lctx.dflash_kv_workspace_graph_rows = lctx.dflash_kv_cache_view_n_filled;
lctx.dflash_kv_workspace_graph_write_pos = lctx.dflash_kv_cache_view_write_pos;
}
const int64_t t_workspace_us = ggml_time_us();
llama_graph_compute_sched(lctx, lctx.dflash_workspace_sched, gf_workspace, lctx.cparams.n_threads);
profile.graph_kv_workspace_compute_us += (uint64_t) (ggml_time_us() - t_workspace_us);
lctx.dflash_kv_workspace_sync_pending = true;
profile.graph_kv_workspace_calls++;
lctx.dflash_kv_workspace_n_filled = lctx.dflash_kv_cache_view_n_filled;
lctx.dflash_kv_workspace_write_pos = lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_workspace_applied_window_version = lctx.dflash_kv_cache_applied_window_version;
lctx.dflash_kv_workspace_valid = true;
}
} else {
ggml_backend_tensor_set(lctx.inp_dflash_target_features, lctx.dflash_target_features_padded.data(), 0, ggml_nbytes(lctx.inp_dflash_target_features));
ggml_backend_tensor_set(lctx.inp_dflash_pos_ctx, lctx.dflash_pos_ctx_data.data(), 0, ggml_nbytes(lctx.inp_dflash_pos_ctx));
}
const int64_t t_mask_us = ggml_time_us();
@@ -1221,7 +1177,7 @@ bool llama_prepare_dflash_graph_inputs(
profile.graph_visible_kv_max = std::max<uint64_t>(profile.graph_visible_kv_max, (uint64_t) visible_kv_max);
profile.graph_prepare_total_us += (uint64_t) (ggml_time_us() - t_total_us);
if (profile.graph_prepare_calls == 1) {
if (profile.graph_prepare_calls == 1 && llama_dflash_stats_log_enabled()) {
int32_t n_swa_layers = 0;
for (int32_t il = 0; il < lctx.model.hparams.n_layer; ++il) {
n_swa_layers += lctx.model.hparams.swa_layers[(size_t) il] ? 1 : 0;
src/llama-spec-features-dflash.cpp
+16 -4
View File
@@ -10,6 +10,14 @@
#include "llama-model.h"
#include "llama-context.h"
static bool llama_dflash_stats_log_enabled() {
const char * env = std::getenv("IK_DFLASH_STATS_LOG");
return env != nullptr && *env != '\0' &&
std::strcmp(env, "0") != 0 &&
std::strcmp(env, "false") != 0 &&
std::strcmp(env, "off") != 0;
}
static bool llama_dflash_positions_strictly_increasing(
const llama_pos * positions,
int32_t n_rows,
@@ -295,12 +303,16 @@ bool llama_model_share_dflash_io_tensors(
const struct ggml_tensor * output = llama_model_dflash_output_tensor(draft_model);
if (draft_model->tok_embd != nullptr && output != nullptr) {
LLAMA_LOG_INFO("%s: DFlash IO mode=%s output_head=%s tensor=%s type=%s\n",
LLAMA_LOG_INFO("%s: DFlash ready io=%s output_head=%s\n",
__func__,
llama_dflash_io_mode_name(llama_model_dflash_io_mode(draft_model, target_model)),
llama_dflash_output_head_kind(draft_model, target_model),
output->name[0] != '\0' ? output->name : "(unnamed)",
ggml_type_name(output->type));
llama_dflash_output_head_kind(draft_model, target_model));
if (llama_dflash_stats_log_enabled()) {
LLAMA_LOG_INFO("%s: DFlash IO tensor=%s type=%s\n",
__func__,
output->name[0] != '\0' ? output->name : "(unnamed)",
ggml_type_name(output->type));
}
}
return draft_model->tok_embd != nullptr && output != nullptr;