diff --git a/common/speculative.cpp b/common/speculative.cpp index e7ce71f9..016aeaa3 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -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, diff --git a/src/graphs/build_dflash.cpp b/src/graphs/build_dflash.cpp index a5b9a815..80c45c1e 100644 --- a/src/graphs/build_dflash.cpp +++ b/src/graphs/build_dflash.cpp @@ -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(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(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(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(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); } diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index fc03353c..eff7d675 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -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, diff --git a/src/llama-build-context.h b/src/llama-build-context.h index ec7cbbb9..a33fdf39 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -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, diff --git a/src/llama-context.h b/src/llama-context.h index 1a7a9d80..8ad9d74b 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -301,6 +301,8 @@ struct llama_context { int32_t dflash_visible_cross_ctx = 0; std::vector dflash_k_ctx_cache; std::vector dflash_v_ctx_cache; + std::vector dflash_k_ctx_workspace; + std::vector dflash_v_ctx_workspace; struct ggml_context * dflash_cache_ctx = nullptr; std::vector 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 dflash_buf_compute_meta; + std::vector 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; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index bcd08ff4..ab84302d 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -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); } diff --git a/src/llama-spec-features.cpp b/src/llama-spec-features.cpp index ab8efddb..00c4b6e2 100644 --- a/src/llama-spec-features.cpp +++ b/src/llama-spec-features.cpp @@ -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( diff --git a/src/llama-spec-features.h b/src/llama-spec-features.h index d976c89a..1c327049 100644 --- a/src/llama-spec-features.h +++ b/src/llama-spec-features.h @@ -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, diff --git a/src/llama.cpp b/src/llama.cpp index e53940a2..a1b63a73 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -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(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(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 extract_device_from_rpc_device(std::vector devices) { std::vector 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(1, cross_ctx); + const int32_t target_token_capacity = std::max(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(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 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++;