nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

apply workspace support for KV cache

This commit is contained in:
SamuelOliveirads 2026-06-01 09:55:34 -03:00
parent ed403dca27
commit 3d73312d9d
9 changed files with 969 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

99
common/speculative.cpp
View File

@ -2596,6 +2596,32 @@ void common_speculative_print_stats(const common_speculative * spec, double slot
const double kv_upload_total_ms = kv_upload_feature_ms + kv_upload_pos_ms;
const double kv_compute_ms = (double) graph_stats.graph_kv_cache_compute_us / 1000.0;
const double kv_sync_ms = (double) graph_stats.graph_kv_cache_sync_us / 1000.0;
const double kv_workspace_total_ms = (double) (
graph_stats.graph_kv_workspace_build_us +
graph_stats.graph_kv_workspace_reserve_us +
graph_stats.graph_kv_workspace_reset_us +
graph_stats.graph_kv_workspace_alloc_us +
graph_stats.graph_kv_workspace_compute_us +
graph_stats.graph_kv_workspace_sync_us) / 1000.0;
const double draft_kv_traffic_ms = (double) (
graph_stats.graph_main_node_k_ctx_view_us +
graph_stats.graph_main_node_v_ctx_view_us +
graph_stats.graph_main_node_k_concat_us +
graph_stats.graph_main_node_v_concat_us +
graph_stats.graph_main_node_k_pad_us +
graph_stats.graph_main_node_v_pad_us +
graph_stats.graph_main_node_k_perm_cont_us +
graph_stats.graph_main_node_v_perm_cont_us) / 1000.0;
const double draft_main_profiled_ms = (double) (
graph_stats.graph_main_node_qcur_us +
graph_stats.graph_main_node_k_draft_us +
graph_stats.graph_main_node_v_draft_us +
graph_stats.graph_main_node_flash_attn_us +
graph_stats.graph_main_node_attn_out_us +
graph_stats.graph_main_node_ffn_us +
graph_stats.graph_main_node_result_rows_us +
graph_stats.graph_main_node_result_norm_us +
graph_stats.graph_main_node_result_us) / 1000.0;
const double replay_append_ms = (double) dflash_state->t_accept_append_us / 1000.0;
const double feature_path_ms = (double) (
capture_stats.capture_prepare_sync_us +
@ -2704,6 +2730,18 @@ void common_speculative_print_stats(const common_speculative * spec, double slot
(unsigned long long) graph_stats.graph_kv_cache_cached_bytes,
graph_stats.last_kv_cache_host_layers);
if (graph_stats.graph_kv_workspace_calls > 0) {
LOG_INF("statistics dflash kv workspace: total=%.3f ms build/reserve/reset/alloc/compute/sync=%.3f/%.3f/%.3f/%.3f/%.3f/%.3f ms calls=%llu\n",
kv_workspace_total_ms,
(double) graph_stats.graph_kv_workspace_build_us / 1000.0,
(double) graph_stats.graph_kv_workspace_reserve_us / 1000.0,
(double) graph_stats.graph_kv_workspace_reset_us / 1000.0,
(double) graph_stats.graph_kv_workspace_alloc_us / 1000.0,
(double) graph_stats.graph_kv_workspace_compute_us / 1000.0,
(double) graph_stats.graph_kv_workspace_sync_us / 1000.0,
(unsigned long long) graph_stats.graph_kv_workspace_calls);
}
if (graph_stats.decode_internal_chunks > 0) {
LOG_INF("statistics dflash decode: llama_decode(total)=%.3f ms calls=%zu chunks=%llu rebuilds=%llu sync_points=%llu internal(total/prelude/sched_reset/build/alloc/prepare/set_inputs/compute/get_result/get_embedding/final_reset)=%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f ms\n",
(double) dflash_state->t_draft_decode_us / 1000.0,
@ -2748,6 +2786,67 @@ void common_speculative_print_stats(const common_speculative * spec, double slot
(unsigned long long) graph_stats.graph_kv_node_v_store_calls);
}
if (graph_stats.graph_main_node_qcur_calls > 0 ||
graph_stats.graph_main_node_k_draft_calls > 0 ||
graph_stats.graph_main_node_v_draft_calls > 0 ||
graph_stats.graph_main_node_flash_attn_calls > 0 ||
graph_stats.graph_main_node_attn_out_calls > 0 ||
graph_stats.graph_main_node_ffn_calls > 0 ||
graph_stats.graph_main_node_result_rows_calls > 0 ||
graph_stats.graph_main_node_result_norm_calls > 0 ||
graph_stats.graph_main_node_result_calls > 0) {
LOG_INF("statistics dflash draft nodes: profiled=%.3f ms graph_compute=%.3f ms qcur/k_draft/v_draft/flash_attn/attn_out/ffn/result_rows/result_norm/result=%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f ms calls=%llu/%llu/%llu/%llu/%llu/%llu/%llu/%llu/%llu\n",
draft_main_profiled_ms,
(double) graph_stats.decode_graph_compute_us / 1000.0,
(double) graph_stats.graph_main_node_qcur_us / 1000.0,
(double) graph_stats.graph_main_node_k_draft_us / 1000.0,
(double) graph_stats.graph_main_node_v_draft_us / 1000.0,
(double) graph_stats.graph_main_node_flash_attn_us / 1000.0,
(double) graph_stats.graph_main_node_attn_out_us / 1000.0,
(double) graph_stats.graph_main_node_ffn_us / 1000.0,
(double) graph_stats.graph_main_node_result_rows_us / 1000.0,
(double) graph_stats.graph_main_node_result_norm_us / 1000.0,
(double) graph_stats.graph_main_node_result_us / 1000.0,
(unsigned long long) graph_stats.graph_main_node_qcur_calls,
(unsigned long long) graph_stats.graph_main_node_k_draft_calls,
(unsigned long long) graph_stats.graph_main_node_v_draft_calls,
(unsigned long long) graph_stats.graph_main_node_flash_attn_calls,
(unsigned long long) graph_stats.graph_main_node_attn_out_calls,
(unsigned long long) graph_stats.graph_main_node_ffn_calls,
(unsigned long long) graph_stats.graph_main_node_result_rows_calls,
(unsigned long long) graph_stats.graph_main_node_result_norm_calls,
(unsigned long long) graph_stats.graph_main_node_result_calls);
}
if (graph_stats.graph_main_node_k_ctx_view_calls > 0 ||
graph_stats.graph_main_node_v_ctx_view_calls > 0 ||
graph_stats.graph_main_node_k_concat_calls > 0 ||
graph_stats.graph_main_node_v_concat_calls > 0 ||
graph_stats.graph_main_node_k_pad_calls > 0 ||
graph_stats.graph_main_node_v_pad_calls > 0 ||
graph_stats.graph_main_node_k_perm_cont_calls > 0 ||
graph_stats.graph_main_node_v_perm_cont_calls > 0) {
LOG_INF("statistics dflash draft kv traffic: total=%.3f ms graph_compute=%.3f ms k_ctx_view/v_ctx_view/k_concat/v_concat/k_pad/v_pad/k_perm_cont/v_perm_cont=%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f/%.3f ms calls=%llu/%llu/%llu/%llu/%llu/%llu/%llu/%llu\n",
draft_kv_traffic_ms,
(double) graph_stats.decode_graph_compute_us / 1000.0,
(double) graph_stats.graph_main_node_k_ctx_view_us / 1000.0,
(double) graph_stats.graph_main_node_v_ctx_view_us / 1000.0,
(double) graph_stats.graph_main_node_k_concat_us / 1000.0,
(double) graph_stats.graph_main_node_v_concat_us / 1000.0,
(double) graph_stats.graph_main_node_k_pad_us / 1000.0,
(double) graph_stats.graph_main_node_v_pad_us / 1000.0,
(double) graph_stats.graph_main_node_k_perm_cont_us / 1000.0,
(double) graph_stats.graph_main_node_v_perm_cont_us / 1000.0,
(unsigned long long) graph_stats.graph_main_node_k_ctx_view_calls,
(unsigned long long) graph_stats.graph_main_node_v_ctx_view_calls,
(unsigned long long) graph_stats.graph_main_node_k_concat_calls,
(unsigned long long) graph_stats.graph_main_node_v_concat_calls,
(unsigned long long) graph_stats.graph_main_node_k_pad_calls,
(unsigned long long) graph_stats.graph_main_node_v_pad_calls,
(unsigned long long) graph_stats.graph_main_node_k_perm_cont_calls,
(unsigned long long) graph_stats.graph_main_node_v_perm_cont_calls);
}
LOG_INF("statistics dflash hot: kv(upload_f/upload_p/upload/compute/sync)=%.3f/%.3f/%.3f/%.3f/%.3f ms calls=%llu replay(accepted_prefix_append)=%.3f ms calls=%zu rows=%zu\n",
kv_upload_feature_ms,
kv_upload_pos_ms,

218
src/graphs/build_dflash.cpp
View File

@ -16,6 +16,119 @@ static bool dflash_use_kv_cache_experiment() {
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);
const int64_t n_embd_head_v = hparams.n_embd_head_v(0);
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 = std::clamp(lctx.dflash_kv_cache_view_n_filled, 0, (int32_t) ctx_len);
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;
GGML_ASSERT(n_embd_head_k == n_embd_head_v);
GGML_ASSERT(lctx.ensure_dflash_kv_cache_tensors((int32_t) ctx_len));
GGML_ASSERT((int32_t) lctx.dflash_k_ctx_workspace.size() == n_layer);
GGML_ASSERT((int32_t) lctx.dflash_v_ctx_workspace.size() == n_layer);
ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes((int) std::max<int64_t>(1, ctx_len)) + 16 * n_layer, false);
auto build_ordered_cache_view = [&](ggml_tensor * cache) -> ggml_tensor * {
if (!lctx.dflash_kv_cache_view_valid || cache_rows <= 0) {
return cache;
}
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);
}
if (cache_write_pos == 0) {
return cache;
}
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);
};
for (int il = 0; il < n_layer; ++il) {
GGML_ASSERT((size_t) il < lctx.dflash_k_ctx_cache.size());
GGML_ASSERT((size_t) il < lctx.dflash_v_ctx_cache.size());
ggml_tensor * Kordered = build_ordered_cache_view(lctx.dflash_k_ctx_cache[(size_t) il]);
ggml_tensor * Vordered = build_ordered_cache_view(lctx.dflash_v_ctx_cache[(size_t) il]);
cb(Kordered, "dflash_workspace_k_ctx_view", il);
cb(Vordered, "dflash_workspace_v_ctx_view", il);
ggml_tensor * Kworkspace = ggml_cont(ctx0, ggml_permute(ctx0, Kordered, 0, 2, 1, 3));
ggml_tensor * Vworkspace = ggml_cont(ctx0, ggml_permute(ctx0, Vordered, 0, 2, 1, 3));
cb(Kworkspace, "dflash_workspace_k_perm_cont", il);
cb(Vworkspace, "dflash_workspace_v_perm_cont", il);
ggml_tensor * Kdst = 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 * Vdst = 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);
ggml_tensor * Kstore = ggml_cpy(ctx0, Kworkspace, Kdst);
ggml_tensor * Vstore = ggml_cpy(ctx0, Vworkspace, Vdst);
cb(Kstore, "dflash_workspace_k_store", il);
cb(Vstore, "dflash_workspace_v_store", il);
ggml_build_forward_expand(gf, Kstore);
ggml_build_forward_expand(gf, Vstore);
}
return gf;
}
ggml_cgraph * llm_build_context::build_dflash_kv_cache() {
const int64_t n_embd_head_k = hparams.n_embd_head_k(0);
const int64_t n_embd_head_v = hparams.n_embd_head_v(0);
@ -160,6 +273,7 @@ ggml_cgraph * llm_build_context::build_dflash() {
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);
@ -226,6 +340,7 @@ ggml_cgraph * llm_build_context::build_dflash() {
ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, tok_embd, cb);
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = (n_tokens > 1 && n_outputs < n_tokens) ? build_inp_out_ids() : nullptr;
bool result_rows_selected = false;
const float kq_scale = 1.0f / std::sqrt((float) n_embd_head_k);
@ -258,7 +373,30 @@ ggml_cgraph * llm_build_context::build_dflash() {
const int64_t t_cache_read_us = use_kv_cache ? ggml_time_us() : 0;
ggml_tensor * Kcur_ctx = nullptr;
ggml_tensor * Vcur_ctx = nullptr;
if (use_kv_cache) {
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;
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;
@ -267,19 +405,19 @@ ggml_cgraph * llm_build_context::build_dflash() {
if (cache_rows < ctx_len) {
ggml_tensor * zero_pad = ggml_view_3d(ctx0, cache,
cache->ne[0],
cache->ne[1],
cache->ne[1],
ctx_len - cache_rows,
cache->nb[1],
cache->nb[2],
(size_t) cache_rows * 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->ne[1],
cache_rows,
cache->nb[1],
cache->nb[2],
0);
return ggml_concat(ctx0, zero_pad, valid, 2);
return ggml_concat(ctx0, zero_pad, valid, 2);
}
if (cache_write_pos == 0) {
@ -288,19 +426,19 @@ ggml_cgraph * llm_build_context::build_dflash() {
ggml_tensor * tail = ggml_view_3d(ctx0, cache,
cache->ne[0],
cache->ne[1],
cache->ne[1],
ctx_len - cache_write_pos,
cache->nb[1],
cache->nb[2],
(size_t) cache_write_pos * 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->ne[1],
cache_write_pos,
cache->nb[1],
cache->nb[2],
0);
return ggml_concat(ctx0, tail, head, 2);
return ggml_concat(ctx0, tail, head, 2);
};
Kcur_ctx = build_ordered_cache_view(lctx.dflash_k_ctx_cache[(size_t) il]);
@ -321,32 +459,58 @@ ggml_cgraph * llm_build_context::build_dflash() {
cb(Vcur_ctx, "Vcur_ctx", il);
}
ggml_tensor * Kcur = ggml_concat(ctx0, Kcur_ctx, Kcur_noise, 2);
ggml_tensor * Vcur = ggml_concat(ctx0, Vcur_ctx, Vcur_noise, 2);
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);
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]);
}
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur_f16", il);
cb(Vcur, "Vcur_f16", il);
ggml_tensor * q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
ggml_tensor * k = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 0, 2, 1, 3));
ggml_tensor * v = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 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 * 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);
cb(k, "k", il);
cb(v, "v", 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);
@ -362,6 +526,12 @@ ggml_cgraph * llm_build_context::build_dflash() {
cur = ggml_add(ctx0, cur, inpSA);
cb(cur, "attn_residual", il);
if (inp_out_ids != nullptr && il == n_layer - 1) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cb(cur, "result_output_rows", -1);
result_rows_selected = true;
}
ggml_tensor * ffn_residual = cur;
cur = llm_build_norm(ctx0, cur, hparams, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, cb, il);
cb(cur, "attn_post_norm", il);
@ -380,13 +550,13 @@ ggml_cgraph * llm_build_context::build_dflash() {
inpL = cur;
}
ggml_tensor * output = model.output;
ggml_tensor * output = const_cast<ggml_tensor *>(llama_model_dflash_output_tensor(&model));
if (output == nullptr) {
output = ggml_new_tensor_2d(ctx0, GGML_TYPE_Q4_0, n_embd, hparams.n_vocab);
}
ggml_tensor * result_input = inpL;
if (inp_out_ids) {
if (inp_out_ids && !result_rows_selected) {
result_input = ggml_get_rows(ctx0, result_input, inp_out_ids);
cb(result_input, "result_output_rows", -1);
}

37
src/llama-build-context.cpp
View File

@ -2206,6 +2206,43 @@ struct ggml_cgraph * llm_build_context::llama_build_graph_dflash_kv_cache(llama_
return result;
}
struct ggml_cgraph * llm_build_context::llama_build_graph_dflash_kv_workspace(llama_context & lctx) {
llama_batch dummy;
dummy.n_tokens = 0;
llm_build_cb cb = [&](struct ggml_tensor * cur, const char * name, int il) {
if (il >= 0) {
int j = 0;
for (; j < GGML_MAX_NAME - 1; ++j) {
cur->name[j] = name[j];
if (!name[j]) {
break;
}
}
if (j < GGML_MAX_NAME - 3) {
cur->name[j++] = '-';
auto sil = std::to_string(il);
for (int k = 0; k < (int) sil.size() && j < GGML_MAX_NAME - 1; ++k) {
cur->name[j++] = sil[k];
}
}
cur->name[j] = 0;
} else {
ggml_set_name(cur, name);
}
};
struct llm_build_context llm(lctx, dummy, cb, false, false, 0, false, &lctx.dflash_workspace_buf_compute_meta);
llm.init();
struct ggml_cgraph * result = llm.build_dflash_kv_workspace();
llm.free();
return result;
}
ggml_cgraph * llm_build_context::llama_build_graph(
llama_context & lctx,
const llama_batch & batch,

4
src/llama-build-context.h
View File

@ -249,6 +249,8 @@ struct llm_build_context {
ggml_cgraph * build_dflash_kv_cache();
ggml_cgraph * build_dflash_kv_workspace();
ggml_cgraph * build_starcoder2();
ggml_cgraph * build_mamba();
@ -466,6 +468,8 @@ llm_expert_gating_func_type gating_op,
static ggml_cgraph * llama_build_graph_dflash_kv_cache(llama_context & lctx);
static ggml_cgraph * llama_build_graph_dflash_kv_workspace(llama_context & lctx);
static ggml_cgraph * llama_build_graph(llama_context & lctx, const llama_batch & batch, bool worst_case, int n_outputs = 0);
ggml_tensor * build_std_attention(ggml_cgraph * gf, ggml_tensor * attn_norm, ggml_tensor * cur,

18
src/llama-context.h
View File

@ -301,6 +301,8 @@ struct llama_context {
int32_t dflash_visible_cross_ctx = 0;
std::vector<struct ggml_tensor *> dflash_k_ctx_cache;
std::vector<struct ggml_tensor *> dflash_v_ctx_cache;
std::vector<struct ggml_tensor *> dflash_k_ctx_workspace;
std::vector<struct ggml_tensor *> dflash_v_ctx_workspace;
struct ggml_context * dflash_cache_ctx = nullptr;
std::vector<ggml_backend_buffer_t> dflash_cache_bufs;
int32_t dflash_kv_cache_write_pos = 0;
@ -312,8 +314,24 @@ struct llama_context {
uint64_t dflash_kv_cache_applied_window_version = 0;
bool dflash_kv_cache_valid = false;
bool dflash_kv_cache_view_valid = false;
int32_t dflash_kv_workspace_write_pos = 0;
int32_t dflash_kv_workspace_n_filled = 0;
int32_t dflash_kv_workspace_reserved_rows = 0;
int32_t dflash_kv_workspace_token_capacity = 0;
int32_t dflash_kv_workspace_n_kv_total = 0;
uint64_t dflash_kv_workspace_applied_window_version = 0;
bool dflash_kv_workspace_valid = false;
bool dflash_kv_workspace_sync_pending = false;
std::vector<uint8_t> dflash_buf_compute_meta;
std::vector<uint8_t> dflash_workspace_buf_compute_meta;
ggml_backend_sched_t dflash_sched = nullptr;
ggml_backend_sched_t dflash_workspace_sched = nullptr;
ggml_cgraph * dflash_kv_graph = nullptr;
ggml_cgraph * dflash_kv_workspace_graph = nullptr;
int32_t dflash_kv_graph_rows = 0;
int32_t dflash_kv_graph_write_pos = 0;
int32_t dflash_kv_workspace_graph_rows = 0;
int32_t dflash_kv_workspace_graph_write_pos = 0;
struct ggml_tensor * dflash_kv_input_target_features = nullptr;
struct ggml_tensor * dflash_kv_input_pos_ctx = nullptr;
struct ggml_tensor * dflash_kq_mask_tensor = nullptr;

1
src/llama-load-tensors.cpp
View File

@ -2202,6 +2202,7 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) {
model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
model.output_mtp = create_tensor(ctx_output, "output_extra.weight", {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
if (model.output == nullptr && model.tok_embd != nullptr) {
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED);
}

98
src/llama-spec-features.cpp
View File

@ -68,6 +68,11 @@ void llama_reset_dflash_kv_cache_state(struct llama_context * ctx) {
ctx->dflash_kv_cache_applied_window_version = 0;
ctx->dflash_kv_cache_valid = false;
ctx->dflash_kv_cache_view_valid = false;
ctx->dflash_kv_workspace_write_pos = 0;
ctx->dflash_kv_workspace_n_filled = 0;
ctx->dflash_kv_workspace_applied_window_version = 0;
ctx->dflash_kv_workspace_valid = false;
ctx->dflash_kv_workspace_sync_pending = false;
for (ggml_backend_buffer_t buf : ctx->dflash_cache_bufs) {
if (buf != nullptr) {
@ -171,6 +176,65 @@ int32_t llama_model_dflash_target_mask_token_id(const struct llama_model * model
return (int32_t) model->vocab.token_mask();
}
const struct ggml_tensor * llama_model_dflash_output_tensor(
const struct llama_model * model) {
if (model == nullptr) {
return nullptr;
}
if (model->output_mtp != nullptr) {
return model->output_mtp;
}
if (model->output != nullptr) {
return model->output;
}
return model->tok_embd;
}
static const char * llama_dflash_io_mode_name(int32_t io_mode) {
switch (io_mode) {
case LLAMA_DFLASH_IO_MODE_SHARED:
return "shared";
case LLAMA_DFLASH_IO_MODE_SELF_CONTAINED:
return "self-contained";
case LLAMA_DFLASH_IO_MODE_MIXED:
return "mixed";
default:
return "invalid";
}
}
static const char * llama_dflash_output_head_kind(
const struct llama_model * draft_model,
const struct llama_model * target_model) {
const struct ggml_tensor * output = llama_model_dflash_output_tensor(draft_model);
if (output == nullptr) {
return "missing";
}
if (output == draft_model->tok_embd) {
return draft_model->tok_embd == (target_model ? target_model->tok_embd : nullptr)
? "shared_token_embedding"
: "token_embedding";
}
if (draft_model->output_mtp != nullptr && output == draft_model->output_mtp) {
if (target_model != nullptr && target_model->output_mtp != nullptr && output == target_model->output_mtp) {
return "output_mtp";
}
if (std::strcmp(output->name, "output_extra.weight") == 0) {
return "output_extra";
}
return "output_mtp";
}
return "output";
}
int32_t llama_model_dflash_io_mode(
const struct llama_model * draft_model,
const struct llama_model * target_model) {
@ -178,13 +242,14 @@ int32_t llama_model_dflash_io_mode(
return LLAMA_DFLASH_IO_MODE_INVALID;
}
const ggml_tensor * target_output = target_model->output != nullptr ? target_model->output : target_model->tok_embd;
if (draft_model->tok_embd == nullptr || draft_model->output == nullptr || target_model->tok_embd == nullptr || target_output == nullptr) {
const ggml_tensor * draft_output = llama_model_dflash_output_tensor(draft_model);
const ggml_tensor * target_output = llama_model_dflash_output_tensor(target_model);
if (draft_model->tok_embd == nullptr || draft_output == nullptr || target_model->tok_embd == nullptr || target_output == nullptr) {
return LLAMA_DFLASH_IO_MODE_INVALID;
}
const bool shared_tok = draft_model->tok_embd == target_model->tok_embd;
const bool shared_output = draft_model->output == target_output;
const bool shared_output = draft_output == target_output;
if (shared_tok && shared_output) {
return LLAMA_DFLASH_IO_MODE_SHARED;
}
@ -200,14 +265,15 @@ bool llama_model_dflash_io_tensors_match(
const struct llama_model * draft_model,
int32_t n_embd,
int32_t n_vocab) {
if (draft_model == nullptr || draft_model->tok_embd == nullptr || draft_model->output == nullptr || n_embd <= 0 || n_vocab <= 0) {
const ggml_tensor * output = llama_model_dflash_output_tensor(draft_model);
if (draft_model == nullptr || draft_model->tok_embd == nullptr || output == nullptr || n_embd <= 0 || n_vocab <= 0) {
return false;
}
return (int32_t) draft_model->tok_embd->ne[0] == n_embd &&
(int32_t) draft_model->tok_embd->ne[1] == n_vocab &&
(int32_t) draft_model->output->ne[0] == n_embd &&
(int32_t) draft_model->output->ne[1] == n_vocab;
(int32_t) output->ne[0] == n_embd &&
(int32_t) output->ne[1] == n_vocab;
}
bool llama_model_share_dflash_io_tensors(
@ -232,7 +298,25 @@ bool llama_model_share_dflash_io_tensors(
}
}
return draft_model->tok_embd != nullptr && draft_model->output != nullptr;
const bool uses_shared_tok = draft_model->tok_embd == target_model->tok_embd;
const bool uses_shared_output = draft_model->output == target_model->output ||
draft_model->output == target_model->tok_embd;
if (draft_model->output_mtp == nullptr && target_model->output_mtp != nullptr && uses_shared_tok && uses_shared_output) {
draft_model->output_mtp = target_model->output_mtp;
}
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",
__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));
}
return draft_model->tok_embd != nullptr && output != nullptr;
}
bool llama_set_draft_input_hidden_state_copy(

44
src/llama-spec-features.h
View File

@ -86,6 +86,13 @@ struct llama_dflash_profile_stats {
uint64_t graph_kv_cache_read_concat_pad_calls = 0;
uint64_t graph_kv_cache_cached_bytes = 0;
uint64_t graph_kv_cache_calls = 0;
uint64_t graph_kv_workspace_build_us = 0;
uint64_t graph_kv_workspace_reserve_us = 0;
uint64_t graph_kv_workspace_reset_us = 0;
uint64_t graph_kv_workspace_alloc_us = 0;
uint64_t graph_kv_workspace_compute_us = 0;
uint64_t graph_kv_workspace_sync_us = 0;
uint64_t graph_kv_workspace_calls = 0;
uint64_t graph_kv_node_fused_target_calls = 0;
uint64_t graph_kv_node_fused_target_us = 0;
uint64_t graph_kv_node_k_proj_calls = 0;
@ -100,6 +107,40 @@ struct llama_dflash_profile_stats {
uint64_t graph_kv_node_k_store_us = 0;
uint64_t graph_kv_node_v_store_calls = 0;
uint64_t graph_kv_node_v_store_us = 0;
uint64_t graph_main_node_qcur_calls = 0;
uint64_t graph_main_node_qcur_us = 0;
uint64_t graph_main_node_k_draft_calls = 0;
uint64_t graph_main_node_k_draft_us = 0;
uint64_t graph_main_node_v_draft_calls = 0;
uint64_t graph_main_node_v_draft_us = 0;
uint64_t graph_main_node_k_ctx_view_calls = 0;
uint64_t graph_main_node_k_ctx_view_us = 0;
uint64_t graph_main_node_v_ctx_view_calls = 0;
uint64_t graph_main_node_v_ctx_view_us = 0;
uint64_t graph_main_node_k_concat_calls = 0;
uint64_t graph_main_node_k_concat_us = 0;
uint64_t graph_main_node_v_concat_calls = 0;
uint64_t graph_main_node_v_concat_us = 0;
uint64_t graph_main_node_k_pad_calls = 0;
uint64_t graph_main_node_k_pad_us = 0;
uint64_t graph_main_node_v_pad_calls = 0;
uint64_t graph_main_node_v_pad_us = 0;
uint64_t graph_main_node_k_perm_cont_calls = 0;
uint64_t graph_main_node_k_perm_cont_us = 0;
uint64_t graph_main_node_v_perm_cont_calls = 0;
uint64_t graph_main_node_v_perm_cont_us = 0;
uint64_t graph_main_node_flash_attn_calls = 0;
uint64_t graph_main_node_flash_attn_us = 0;
uint64_t graph_main_node_attn_out_calls = 0;
uint64_t graph_main_node_attn_out_us = 0;
uint64_t graph_main_node_ffn_calls = 0;
uint64_t graph_main_node_ffn_us = 0;
uint64_t graph_main_node_result_rows_calls = 0;
uint64_t graph_main_node_result_rows_us = 0;
uint64_t graph_main_node_result_norm_calls = 0;
uint64_t graph_main_node_result_norm_us = 0;
uint64_t graph_main_node_result_calls = 0;
uint64_t graph_main_node_result_us = 0;
uint64_t graph_feature_bytes = 0;
uint64_t graph_pos_bytes = 0;
uint64_t graph_mask_bytes = 0;
@ -232,6 +273,9 @@ int32_t llama_model_dflash_io_mode(
const struct llama_model * draft_model,
const struct llama_model * target_model);
const struct ggml_tensor * llama_model_dflash_output_tensor(
const struct llama_model * model);
bool llama_model_dflash_io_tensors_match(
const struct llama_model * draft_model,
int32_t n_embd,

502
src/llama.cpp
View File

@ -190,12 +190,42 @@ enum llama_dflash_kv_node_kind {
LLAMA_DFLASH_KV_NODE_V_STORE,
};
enum llama_dflash_main_node_kind {
LLAMA_DFLASH_MAIN_NODE_NONE = 0,
LLAMA_DFLASH_MAIN_NODE_QCUR,
LLAMA_DFLASH_MAIN_NODE_K_DRAFT,
LLAMA_DFLASH_MAIN_NODE_V_DRAFT,
LLAMA_DFLASH_MAIN_NODE_K_CTX_VIEW,
LLAMA_DFLASH_MAIN_NODE_V_CTX_VIEW,
LLAMA_DFLASH_MAIN_NODE_K_CONCAT,
LLAMA_DFLASH_MAIN_NODE_V_CONCAT,
LLAMA_DFLASH_MAIN_NODE_K_PAD,
LLAMA_DFLASH_MAIN_NODE_V_PAD,
LLAMA_DFLASH_MAIN_NODE_K_PERM_CONT,
LLAMA_DFLASH_MAIN_NODE_V_PERM_CONT,
LLAMA_DFLASH_MAIN_NODE_FLASH_ATTN,
LLAMA_DFLASH_MAIN_NODE_ATTN_OUT,
LLAMA_DFLASH_MAIN_NODE_FFN,
LLAMA_DFLASH_MAIN_NODE_RESULT_ROWS,
LLAMA_DFLASH_MAIN_NODE_RESULT_NORM,
LLAMA_DFLASH_MAIN_NODE_RESULT,
};
struct llama_dflash_kv_node_profiler {
llama_dflash_profile_stats * profile = nullptr;
int64_t t_start_us = 0;
llama_dflash_kv_node_kind active_kind = LLAMA_DFLASH_KV_NODE_NONE;
};
struct llama_dflash_main_node_profiler {
llama_dflash_profile_stats * profile = nullptr;
ggml_backend_sched_eval_callback prev_callback = nullptr;
void * prev_user_data = nullptr;
bool prev_active = false;
int64_t t_start_us = 0;
llama_dflash_main_node_kind active_kind = LLAMA_DFLASH_MAIN_NODE_NONE;
};
static bool llama_dflash_tensor_name_has_prefix(const struct ggml_tensor * tensor, const char * prefix) {
if (tensor == nullptr || prefix == nullptr || prefix[0] == '\0') {
return false;
@ -204,6 +234,16 @@ static bool llama_dflash_tensor_name_has_prefix(const struct ggml_tensor * tenso
return std::strncmp(tensor->name, prefix, std::strlen(prefix)) == 0;
}
static bool llama_dflash_tensor_name_matches_label(const struct ggml_tensor * tensor, const char * label) {
if (!llama_dflash_tensor_name_has_prefix(tensor, label)) {
return false;
}
const size_t label_len = std::strlen(label);
const char next = tensor->name[label_len];
return next == '\0' || next == '-';
}
static llama_dflash_kv_node_kind llama_dflash_kv_node_kind_from_tensor(const struct ggml_tensor * tensor) {
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_fused_target")) {
return LLAMA_DFLASH_KV_NODE_FUSED_TARGET;
@ -268,6 +308,146 @@ static void llama_dflash_kv_node_profile_add(
}
}
static llama_dflash_main_node_kind llama_dflash_main_node_kind_from_tensor(const struct ggml_tensor * tensor) {
if (llama_dflash_tensor_name_has_prefix(tensor, "Qcur")) {
return LLAMA_DFLASH_MAIN_NODE_QCUR;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Kcur_noise")) {
return LLAMA_DFLASH_MAIN_NODE_K_DRAFT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Vcur_noise")) {
return LLAMA_DFLASH_MAIN_NODE_V_DRAFT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Kcur_ctx_cache")) {
return LLAMA_DFLASH_MAIN_NODE_K_CTX_VIEW;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Vcur_ctx_cache")) {
return LLAMA_DFLASH_MAIN_NODE_V_CTX_VIEW;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_k_concat")) {
return LLAMA_DFLASH_MAIN_NODE_K_CONCAT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_v_concat")) {
return LLAMA_DFLASH_MAIN_NODE_V_CONCAT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_k_pad")) {
return LLAMA_DFLASH_MAIN_NODE_K_PAD;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_v_pad")) {
return LLAMA_DFLASH_MAIN_NODE_V_PAD;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_k_perm_cont")) {
return LLAMA_DFLASH_MAIN_NODE_K_PERM_CONT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_v_perm_cont")) {
return LLAMA_DFLASH_MAIN_NODE_V_PERM_CONT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "flash_attn_reshaped")) {
return LLAMA_DFLASH_MAIN_NODE_NONE;
}
if (llama_dflash_tensor_name_matches_label(tensor, "flash_attn")) {
return LLAMA_DFLASH_MAIN_NODE_FLASH_ATTN;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "kqv_out")) {
return LLAMA_DFLASH_MAIN_NODE_ATTN_OUT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "ffn_out")) {
return LLAMA_DFLASH_MAIN_NODE_FFN;
}
if (llama_dflash_tensor_name_matches_label(tensor, "result_output_rows")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT_ROWS;
}
if (llama_dflash_tensor_name_matches_label(tensor, "result_norm")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT_NORM;
}
if (llama_dflash_tensor_name_matches_label(tensor, "output")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT;
}
if (llama_dflash_tensor_name_matches_label(tensor, "result_output")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT;
}
return LLAMA_DFLASH_MAIN_NODE_NONE;
}
static void llama_dflash_main_node_profile_add(
llama_dflash_profile_stats & profile,
llama_dflash_main_node_kind kind,
uint64_t elapsed_us) {
switch (kind) {
case LLAMA_DFLASH_MAIN_NODE_QCUR:
profile.graph_main_node_qcur_calls++;
profile.graph_main_node_qcur_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_DRAFT:
profile.graph_main_node_k_draft_calls++;
profile.graph_main_node_k_draft_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_DRAFT:
profile.graph_main_node_v_draft_calls++;
profile.graph_main_node_v_draft_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_CTX_VIEW:
profile.graph_main_node_k_ctx_view_calls++;
profile.graph_main_node_k_ctx_view_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_CTX_VIEW:
profile.graph_main_node_v_ctx_view_calls++;
profile.graph_main_node_v_ctx_view_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_CONCAT:
profile.graph_main_node_k_concat_calls++;
profile.graph_main_node_k_concat_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_CONCAT:
profile.graph_main_node_v_concat_calls++;
profile.graph_main_node_v_concat_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_PAD:
profile.graph_main_node_k_pad_calls++;
profile.graph_main_node_k_pad_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_PAD:
profile.graph_main_node_v_pad_calls++;
profile.graph_main_node_v_pad_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_PERM_CONT:
profile.graph_main_node_k_perm_cont_calls++;
profile.graph_main_node_k_perm_cont_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_PERM_CONT:
profile.graph_main_node_v_perm_cont_calls++;
profile.graph_main_node_v_perm_cont_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_FLASH_ATTN:
profile.graph_main_node_flash_attn_calls++;
profile.graph_main_node_flash_attn_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_ATTN_OUT:
profile.graph_main_node_attn_out_calls++;
profile.graph_main_node_attn_out_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_FFN:
profile.graph_main_node_ffn_calls++;
profile.graph_main_node_ffn_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_RESULT_ROWS:
profile.graph_main_node_result_rows_calls++;
profile.graph_main_node_result_rows_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_RESULT_NORM:
profile.graph_main_node_result_norm_calls++;
profile.graph_main_node_result_norm_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_RESULT:
profile.graph_main_node_result_calls++;
profile.graph_main_node_result_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_NONE:
break;
}
}
static bool llama_dflash_kv_node_eval_callback(struct ggml_tensor * tensor, bool ask, void * user_data) {
auto * profiler = static_cast<llama_dflash_kv_node_profiler *>(user_data);
if (profiler == nullptr || profiler->profile == nullptr) {
@ -294,6 +474,62 @@ static bool llama_dflash_kv_node_eval_callback(struct ggml_tensor * tensor, bool
return true;
}
static bool llama_dflash_main_node_eval_callback(struct ggml_tensor * tensor, bool ask, void * user_data) {
auto * profiler = static_cast<llama_dflash_main_node_profiler *>(user_data);
if (profiler == nullptr || profiler->profile == nullptr) {
return false;
}
const llama_dflash_main_node_kind kind = llama_dflash_main_node_kind_from_tensor(tensor);
if (ask) {
profiler->prev_active = profiler->prev_callback != nullptr
? profiler->prev_callback(tensor, ask, profiler->prev_user_data)
: false;
if (kind == LLAMA_DFLASH_MAIN_NODE_NONE) {
profiler->active_kind = LLAMA_DFLASH_MAIN_NODE_NONE;
profiler->t_start_us = 0;
return profiler->prev_active;
}
profiler->active_kind = kind;
profiler->t_start_us = ggml_time_us();
return true;
}
bool prev_result = false;
if (profiler->prev_active && profiler->prev_callback != nullptr) {
prev_result = profiler->prev_callback(tensor, ask, profiler->prev_user_data);
}
const bool tracked = kind != LLAMA_DFLASH_MAIN_NODE_NONE &&
profiler->active_kind == kind &&
profiler->t_start_us > 0;
if (tracked) {
llama_dflash_main_node_profile_add(*profiler->profile, kind, (uint64_t) (ggml_time_us() - profiler->t_start_us));
}
profiler->prev_active = false;
profiler->active_kind = LLAMA_DFLASH_MAIN_NODE_NONE;
profiler->t_start_us = 0;
return prev_result || tracked;
}
static bool llama_dflash_use_kv_workspace_experiment() {
return llama_env_flag_enabled("IK_DFLASH_KV_WORKSPACE");
}
static void llama_sync_dflash_workspace_if_pending(struct llama_context & lctx) {
if (!lctx.dflash_kv_workspace_sync_pending || lctx.dflash_workspace_sched == nullptr) {
return;
}
const int64_t t_workspace_sync_us = ggml_time_us();
ggml_backend_sched_synchronize(lctx.dflash_workspace_sched);
lctx.dflash_profile.graph_kv_workspace_sync_us += (uint64_t) (ggml_time_us() - t_workspace_sync_us);
lctx.dflash_kv_workspace_sync_pending = false;
}
// extract ip and port from RPC[ip:port] for rpc and keep other device names
static std::vector<rpc_device> extract_device_from_rpc_device(std::vector<std::string> devices) {
std::vector<rpc_device> rpc_servers;
@ -727,16 +963,26 @@ 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("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);
const int32_t n_layer = model.hparams.n_layer;
const int64_t n_embd_head_k = model.hparams.n_embd_head_k(0);
const int64_t n_embd_head_v = model.hparams.n_embd_head_v(0);
const int64_t n_head_kv = model.hparams.n_head_kv();
if (dflash_cache_ctx != nullptr && !dflash_k_ctx_cache.empty()) {
if ((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 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();
if (cache_matches && workspace_matches) {
return true;
}
@ -745,11 +991,23 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
ggml_backend_sched_free(dflash_sched);
dflash_sched = nullptr;
}
if (dflash_workspace_sched != nullptr) {
ggml_backend_sched_free(dflash_workspace_sched);
dflash_workspace_sched = nullptr;
}
dflash_kv_graph = nullptr;
dflash_kv_workspace_graph = nullptr;
dflash_kv_graph_rows = 0;
dflash_kv_graph_write_pos = 0;
dflash_kv_workspace_graph_rows = 0;
dflash_kv_workspace_graph_write_pos = 0;
dflash_kv_workspace_reserved_rows = 0;
dflash_buf_compute_meta.clear();
dflash_workspace_buf_compute_meta.clear();
}
ggml_init_params params = {
/*.mem_size =*/ (size_t) (2 * std::max(1, n_layer)) * ggml_tensor_overhead(),
/*.mem_size =*/ (size_t) ((use_kv_workspace ? 4 : 2) * std::max(1, n_layer)) * ggml_tensor_overhead(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
@ -761,8 +1019,14 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
dflash_k_ctx_cache.resize((size_t) n_layer);
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_cache_bufs.clear();
dflash_cache_bufs.reserve((size_t) std::max(1, n_layer) * 2);
dflash_cache_bufs.reserve((size_t) std::max(1, n_layer) * (use_kv_workspace ? 4 : 2));
int32_t host_layers = 0;
const char * first_buft_name = nullptr;
const char * last_buft_name = nullptr;
@ -809,9 +1073,47 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
ggml_backend_tensor_alloc(v_buf, dflash_v_ctx_cache[(size_t) il], ggml_backend_buffer_get_base(v_buf));
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_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;
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__,
@ -827,6 +1129,8 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
void llama_context::free_dflash_kv_cache_tensors() {
dflash_k_ctx_cache.clear();
dflash_v_ctx_cache.clear();
dflash_k_ctx_workspace.clear();
dflash_v_ctx_workspace.clear();
dflash_kv_cache_write_pos = 0;
dflash_kv_cache_n_filled = 0;
dflash_kv_cache_update_rows = 0;
@ -836,11 +1140,31 @@ void llama_context::free_dflash_kv_cache_tensors() {
dflash_kv_cache_applied_window_version = 0;
dflash_kv_cache_valid = false;
dflash_kv_cache_view_valid = false;
dflash_kv_workspace_write_pos = 0;
dflash_kv_workspace_n_filled = 0;
dflash_kv_workspace_reserved_rows = 0;
dflash_kv_workspace_token_capacity = 0;
dflash_kv_workspace_n_kv_total = 0;
dflash_kv_workspace_applied_window_version = 0;
dflash_kv_workspace_valid = false;
dflash_kv_workspace_sync_pending = false;
dflash_kv_graph = nullptr;
dflash_kv_workspace_graph = nullptr;
dflash_kv_graph_rows = 0;
dflash_kv_graph_write_pos = 0;
dflash_kv_workspace_graph_rows = 0;
dflash_kv_workspace_graph_write_pos = 0;
dflash_kv_input_target_features = nullptr;
dflash_kv_input_pos_ctx = nullptr;
dflash_kq_mask_tensor = nullptr;
dflash_kq_mask_swa_tensor = nullptr;
if (dflash_workspace_sched != nullptr) {
ggml_backend_sched_synchronize(dflash_workspace_sched);
ggml_backend_sched_free(dflash_workspace_sched);
dflash_workspace_sched = nullptr;
}
for (ggml_backend_buffer_t buf : dflash_cache_bufs) {
if (buf != nullptr) {
ggml_backend_buffer_free(buf);
@ -4229,7 +4553,7 @@ static bool llm_load_tensors(
if (model.arch == LLM_ARCH_GEMMA4) {
llm_scale_gate_inp_s(model, use_mmap_buffer);
}
if ((model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) && extra_output_type != GGML_TYPE_COUNT) {
if ((model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DFLASH_DRAFT) && extra_output_type != GGML_TYPE_COUNT) {
llm_requantize_output_tensor(model, extra_output_type);
}
@ -5405,6 +5729,7 @@ static bool prepare_dflash_graph_inputs(
struct llama_context & lctx,
uint32_t n_tokens) {
const bool use_kv_cache = llama_env_flag_enabled("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("IK_DFLASH_KV_NODE_TIMING");
auto & profile = lctx.dflash_profile;
const int32_t cross_ctx = lctx.dflash_visible_cross_ctx > 0
@ -5456,6 +5781,10 @@ static bool 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);
}
if (graph_cross_ctx != cross_ctx) {
profile.graph_shape_failures++;
@ -5585,6 +5914,9 @@ static bool prepare_dflash_graph_inputs(
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;
@ -5631,23 +5963,35 @@ static bool prepare_dflash_graph_inputs(
}
lctx.dflash_kv_cache_update_rows = update_rows;
const int64_t t_build_us = ggml_time_us();
ggml_cgraph * 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) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache graph\n", __func__);
return false;
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();
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) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache 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;
}
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);
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) {
@ -5692,6 +6036,108 @@ static bool prepare_dflash_graph_inputs(
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);
}
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 (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) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V workspace graph\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;
}
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));
@ -5806,6 +6252,7 @@ static int llama_decode_internal(
auto * dflash_profile = lctx.model.arch == LLM_ARCH_DFLASH_DRAFT ? &lctx.dflash_profile : nullptr;
const bool dflash_decode_timing = dflash_profile != nullptr && llama_env_flag_enabled("IK_DFLASH_DECODE_TIMING");
const bool dflash_draft_node_timing = dflash_profile != nullptr && llama_env_flag_enabled("IK_DFLASH_DRAFT_NODE_TIMING");
auto & kv_self = lctx.kv_self;
@ -6165,8 +6612,21 @@ static int llama_decode_internal(
#if IK_PRINT_TIMING
tim1 = ggml_time_us();
#endif
if (lctx.dflash_kv_workspace_sync_pending) {
llama_sync_dflash_workspace_if_pending(lctx);
}
const int64_t t_dflash_graph_compute_us = dflash_decode_timing ? ggml_time_us() : 0;
llama_dflash_main_node_profiler draft_node_profiler;
if (dflash_draft_node_timing) {
draft_node_profiler.profile = dflash_profile;
draft_node_profiler.prev_callback = lctx.cparams.cb_eval;
draft_node_profiler.prev_user_data = lctx.cparams.cb_eval_user_data;
ggml_backend_sched_set_eval_callback(lctx.sched, llama_dflash_main_node_eval_callback, &draft_node_profiler);
}
llama_graph_compute(lctx, gf, n_threads);
if (dflash_draft_node_timing) {
ggml_backend_sched_set_eval_callback(lctx.sched, lctx.cparams.cb_eval, lctx.cparams.cb_eval_user_data);
}
if (dflash_decode_timing) {
llama_synchronize(&lctx);
dflash_profile->decode_sync_profile_points++;