From 532499836efb11c723dc152633239e6697e2489a Mon Sep 17 00:00:00 2001 From: SamuelOliveirads Date: Sat, 30 May 2026 21:36:10 -0300 Subject: [PATCH] improve DFlash caching and profiling capabilities --- common/speculative.cpp | 276 ++++++++++++++++++++- examples/server/server-context.cpp | 111 ++++++++- src/graphs/build_dflash.cpp | 130 ++++++++-- src/llama-build-context.cpp | 64 +++-- src/llama-build-context.h | 9 +- src/llama-context.h | 14 ++ src/llama-spec-features.cpp | 379 ++++++++++++++++++++++++++++- src/llama-spec-features.h | 77 ++++++ src/llama.cpp | 274 ++++++++++++++++++++- 9 files changed, 1261 insertions(+), 73 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index d740ded9..bd854a7f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -11,9 +11,12 @@ #include "suffix-tree.h" #include +#include +#include #include #include #include +#include #include #define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 @@ -210,6 +213,15 @@ struct common_speculative_state { struct common_speculative_state_mtp; struct common_speculative_state_dflash; +static void dflash_contract_log_append( + const common_speculative_state_dflash & state, + llama_seq_id seq_id, + const std::vector & new_positions); +static void dflash_contract_log_draft( + const common_speculative_state_dflash & state, + int32_t n_keep, + size_t result_size); + static common_speculative_state_mtp * common_speculative_get_mtp_state(common_speculative * spec); static const common_speculative_state_mtp * common_speculative_get_mtp_state(const common_speculative * spec); static common_speculative_state_dflash * common_speculative_get_dflash_state(common_speculative * spec); @@ -235,6 +247,81 @@ static std::vector mtp_speculative_gen_draft( static int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch & batch, bool is_prompt_warmup); +static bool dflash_contract_log_enabled() { + const char * env = std::getenv("IK_DFLASH_CONTRACT_LOG"); + if (env == nullptr || *env == '\0') { + return false; + } + + return std::strcmp(env, "0") != 0 && + std::strcmp(env, "false") != 0 && + std::strcmp(env, "off") != 0; +} + +template +static std::string dflash_contract_format_values( + const std::vector & values, + size_t edge_count = 4) { + std::ostringstream oss; + oss << '['; + if (values.empty()) { + oss << ']'; + return oss.str(); + } + + const size_t head = std::min(edge_count, values.size()); + for (size_t i = 0; i < head; ++i) { + if (i > 0) { + oss << ','; + } + oss << values[i]; + } + + if (values.size() > edge_count * 2) { + oss << ",...,"; + for (size_t i = values.size() - edge_count; i < values.size(); ++i) { + if (i > values.size() - edge_count) { + oss << ','; + } + oss << values[i]; + } + } else { + for (size_t i = head; i < values.size(); ++i) { + oss << ',' << values[i]; + } + } + + oss << ']'; + return oss.str(); +} + +struct dflash_contract_pos_summary { + llama_pos first = -1; + llama_pos last = -1; + int32_t gap_count = 0; + int32_t nonmono_count = 0; +}; + +static dflash_contract_pos_summary dflash_contract_summarize_positions( + const std::vector & positions) { + dflash_contract_pos_summary summary; + if (positions.empty()) { + return summary; + } + + summary.first = positions.front(); + summary.last = positions.back(); + for (size_t i = 1; i < positions.size(); ++i) { + if (positions[i] <= positions[i - 1]) { + summary.nonmono_count++; + } else if (positions[i] != positions[i - 1] + 1) { + summary.gap_count++; + } + } + + return summary; +} + struct mtp_last_embd { std::vector embd; float prob = 0.0f; @@ -368,6 +455,14 @@ struct common_speculative_state_dflash : public common_speculative_state { std::vector target_window_pos; int32_t target_window_rows = 0; llama_pos last_target_pos = -1; + size_t n_window_updates = 0; + size_t n_rows_seen = 0; + size_t n_rows_dropped = 0; + size_t n_context_shifts = 0; + size_t n_draft_empty = 0; + size_t n_set_target_fail = 0; + size_t n_decode_fail = 0; + llama_pos last_draft_pos_base = -1; common_speculative_state_dflash( enum common_speculative_type type, @@ -412,8 +507,20 @@ struct common_speculative_state_dflash : public common_speculative_state { batch = llama_batch_init(std::max(1, block_size), 0, 1); ready = true; - LOG_INF("%s: DFlash context ready (n_ctx=%d, block_size=%d, cross_ctx=%d, n_target_features=%d)\n", - __func__, llama_n_ctx(ctx_dft), block_size, this->cross_ctx, n_target_features); + llama_set_dflash_visible_cross_ctx(ctx_dft, this->cross_ctx); + llama_dflash_profile_reset(ctx_tgt); + llama_dflash_profile_reset(ctx_dft); + + std::ostringstream layers_oss; + for (size_t i = 0; i < target_layer_ids.size(); ++i) { + if (i > 0) { + layers_oss << ","; + } + layers_oss << target_layer_ids[i]; + } + + LOG_INF("%s: DFlash context ready (n_ctx=%d, block_size=%d, cross_ctx=%d, n_target_features=%d, target_layer_ids=[%s])\n", + __func__, llama_n_ctx(ctx_dft), block_size, this->cross_ctx, n_target_features, layers_oss.str().c_str()); } ~common_speculative_state_dflash() override { @@ -429,6 +536,16 @@ struct common_speculative_state_dflash : public common_speculative_state { void begin(const llama_tokens & prompt) override { GGML_UNUSED(prompt); llama_kv_cache_clear(ctx_dft); + n_window_updates = 0; + n_rows_seen = 0; + n_rows_dropped = 0; + n_context_shifts = 0; + n_draft_empty = 0; + n_set_target_fail = 0; + n_decode_fail = 0; + last_draft_pos_base = -1; + llama_dflash_profile_reset(ctx_tgt); + llama_dflash_profile_reset(ctx_dft); } void draft( @@ -441,6 +558,7 @@ struct common_speculative_state_dflash : public common_speculative_state { result.clear(); if (!ready || target_window_rows <= 0) { + n_draft_empty++; return; } @@ -449,20 +567,23 @@ struct common_speculative_state_dflash : public common_speculative_state { return; } - if (!llama_set_dflash_target_features_copy(ctx_dft, target_window.data(), target_window.size(), target_window_rows, target_window_pos.data())) { + if (!llama_set_dflash_target_features_view(ctx_dft, target_window.data(), target_window.size(), target_window_rows, target_window_pos.data())) { LOG_ERR("%s: failed to set DFlash target features\n", __func__); + n_set_target_fail++; return; } llama_kv_cache_clear(ctx_dft); batch.n_tokens = 0; const llama_pos draft_pos_base = last_target_pos >= 0 ? last_target_pos + 1 : (llama_pos) target_window_rows; + last_draft_pos_base = draft_pos_base; for (int32_t i = 0; i < block_size; ++i) { common_batch_add(batch, mask_token_id, draft_pos_base + i, { 0 }, i < n_keep); } if (llama_decode(ctx_dft, batch) != 0) { LOG_ERR("%s: llama_decode() failed for DFlash draft batch\n", __func__); + n_decode_fail++; batch.n_tokens = 0; return; } @@ -473,6 +594,7 @@ struct common_speculative_state_dflash : public common_speculative_state { } batch.n_tokens = 0; + dflash_contract_log_draft(*this, n_keep, result.size()); } void accept(uint16_t n_accepted) override { @@ -480,6 +602,83 @@ struct common_speculative_state_dflash : public common_speculative_state { } }; +static void dflash_contract_log_append( + const common_speculative_state_dflash & state, + llama_seq_id seq_id, + const std::vector & new_positions) { + if (!dflash_contract_log_enabled()) { + return; + } + + static std::atomic counter = 0; + const uint64_t ordinal = counter.fetch_add(1, std::memory_order_relaxed); + if (ordinal >= 8) { + return; + } + + const dflash_contract_pos_summary incoming = dflash_contract_summarize_positions(new_positions); + const dflash_contract_pos_summary window = dflash_contract_summarize_positions(state.target_window_pos); + + LOG_INF("dflash contract append[%llu]: seq=%d incoming_rows=%zu incoming_pos=%s pos=[%d..%d] gaps=%d nonmono=%d window_rows=%d window_pos=%s pos=[%d..%d] gaps=%d nonmono=%d last_target_pos=%d\n", + (unsigned long long) (ordinal + 1), + (int) seq_id, + new_positions.size(), + dflash_contract_format_values(new_positions).c_str(), + (int) incoming.first, + (int) incoming.last, + incoming.gap_count, + incoming.nonmono_count, + state.target_window_rows, + dflash_contract_format_values(state.target_window_pos).c_str(), + (int) window.first, + (int) window.last, + window.gap_count, + window.nonmono_count, + (int) state.last_target_pos); +} + +static void dflash_contract_log_draft( + const common_speculative_state_dflash & state, + int32_t n_keep, + size_t result_size) { + if (!dflash_contract_log_enabled()) { + return; + } + + static std::atomic counter = 0; + const uint64_t ordinal = counter.fetch_add(1, std::memory_order_relaxed); + if (ordinal >= 8) { + return; + } + + const dflash_contract_pos_summary window = dflash_contract_summarize_positions(state.target_window_pos); + llama_dflash_profile_stats graph_stats = {}; + llama_dflash_profile_get_stats(state.ctx_dft, &graph_stats); + const int draft_delta = (state.last_target_pos >= 0 && state.last_draft_pos_base >= 0) + ? (int) (state.last_draft_pos_base - state.last_target_pos) + : -1; + + LOG_INF("dflash contract draft[%llu]: window_rows=%d window_pos=%s pos=[%d..%d] gaps=%d nonmono=%d last_target_pos=%d draft_pos_base=%d delta=%d n_keep=%d result=%zu set_target(missing/nonmono)=%llu/%llu graph(fallback/nonmono)=%llu/%llu graph_pos=[%d..%d]\n", + (unsigned long long) (ordinal + 1), + state.target_window_rows, + dflash_contract_format_values(state.target_window_pos).c_str(), + (int) window.first, + (int) window.last, + window.gap_count, + window.nonmono_count, + (int) state.last_target_pos, + (int) state.last_draft_pos_base, + draft_delta, + n_keep, + result_size, + (unsigned long long) graph_stats.set_target_missing_positions, + (unsigned long long) graph_stats.set_target_non_monotonic_positions, + (unsigned long long) graph_stats.graph_pos_fallbacks, + (unsigned long long) graph_stats.graph_pos_non_monotonic, + (int) graph_stats.last_pos_first, + (int) graph_stats.last_pos_last); +} + struct common_speculative_state_draft : public common_speculative_state { llama_context * ctx_tgt; // only used for retokenizing from ctx_dft llama_context * ctx_dft; @@ -2101,6 +2300,70 @@ void common_speculative_print_stats(const common_speculative * spec, double slot impl->n_gen_tokens, impl->n_acc_tokens, str_perf.c_str()); + + if (impl->type == COMMON_SPECULATIVE_TYPE_DFLASH) { + const auto * dflash_state = dynamic_cast(impl.get()); + if (dflash_state != nullptr) { + llama_dflash_profile_stats capture_stats; + llama_dflash_profile_stats graph_stats; + const bool have_capture = llama_dflash_profile_get_stats(dflash_state->ctx_tgt, &capture_stats); + const bool have_graph = llama_dflash_profile_get_stats(dflash_state->ctx_dft, &graph_stats); + + LOG_INF("statistics dflash detail: cross_ctx=%d, window_rows=%d, pos=[%d..%d], window_updates=%zu, rows_seen=%zu, rows_dropped=%zu, shifts=%zu, draft_fail(empty/set/decode)=%zu/%zu/%zu, next_draft_pos=%d\n", + dflash_state->cross_ctx, + dflash_state->target_window_rows, + dflash_state->target_window_pos.empty() ? -1 : (int) dflash_state->target_window_pos.front(), + dflash_state->target_window_pos.empty() ? -1 : (int) dflash_state->target_window_pos.back(), + dflash_state->n_window_updates, + dflash_state->n_rows_seen, + dflash_state->n_rows_dropped, + dflash_state->n_context_shifts, + dflash_state->n_draft_empty, + dflash_state->n_set_target_fail, + dflash_state->n_decode_fail, + (int) dflash_state->last_draft_pos_base); + + if (have_capture || have_graph) { + LOG_INF("statistics dflash profile: capture(sync/materialize)=%.3f/%.3f ms calls=%llu/%llu bytes=%llu phase(prompt/verify batches changes)=%llu/%llu %llu/%llu, set_target=%.3f ms rows=%llu bytes=%llu, prep(total/features/pos/mask)=%.3f/%.3f/%.3f/%.3f ms kv_cache=%.3f ms calls=%llu/%llu bytes=%llu/%llu/%llu, fallback_pos(copy/graph)=%llu/%llu, nonmono(copy/graph)=%llu/%llu, capture_fail=%llu/%llu, visible_kv_max=%llu, last(rows=%d width=%d left_pad=%d n_tokens=%d n_kv=%d pos=[%d..%d])\n", + (double) capture_stats.capture_prepare_sync_us / 1000.0, + (double) capture_stats.capture_materialize_us / 1000.0, + (unsigned long long) capture_stats.capture_prepare_calls, + (unsigned long long) capture_stats.capture_materialize_calls, + (unsigned long long) capture_stats.capture_materialize_bytes, + (unsigned long long) capture_stats.capture_prompt_batches, + (unsigned long long) capture_stats.capture_prompt_shape_changes, + (unsigned long long) capture_stats.capture_verify_batches, + (unsigned long long) capture_stats.capture_verify_shape_changes, + (double) graph_stats.set_target_copy_us / 1000.0, + (unsigned long long) graph_stats.set_target_rows, + (unsigned long long) graph_stats.set_target_copy_bytes, + (double) graph_stats.graph_prepare_total_us / 1000.0, + (double) graph_stats.graph_feature_copy_us / 1000.0, + (double) graph_stats.graph_pos_copy_us / 1000.0, + (double) graph_stats.graph_mask_build_us / 1000.0, + (double) graph_stats.graph_kv_cache_compute_us / 1000.0, + (unsigned long long) graph_stats.graph_prepare_calls, + (unsigned long long) graph_stats.graph_kv_cache_calls, + (unsigned long long) graph_stats.graph_feature_bytes, + (unsigned long long) graph_stats.graph_pos_bytes, + (unsigned long long) graph_stats.graph_mask_bytes, + (unsigned long long) graph_stats.set_target_missing_positions, + (unsigned long long) graph_stats.graph_pos_fallbacks, + (unsigned long long) graph_stats.set_target_non_monotonic_positions, + (unsigned long long) graph_stats.graph_pos_non_monotonic, + (unsigned long long) capture_stats.capture_prepare_failures, + (unsigned long long) capture_stats.capture_materialize_failures, + (unsigned long long) graph_stats.graph_visible_kv_max, + graph_stats.last_n_rows, + graph_stats.last_width, + graph_stats.last_left_pad, + graph_stats.last_n_tokens, + graph_stats.last_n_kv_total, + (int) graph_stats.last_pos_first, + (int) graph_stats.last_pos_last); + } + } + } } if (spec->tuner && spec->tuner->enabled && slot_tps > 0.0 && n_decoded > 0) { @@ -2251,7 +2514,10 @@ static bool dflash_append_target_features( } const int32_t n_rows = (int32_t) new_positions.size(); + state.n_window_updates++; + state.n_rows_seen += (size_t) n_rows; if (n_rows >= state.cross_ctx) { + state.n_rows_dropped += (size_t) state.target_window_rows + (size_t) (n_rows - state.cross_ctx); const int32_t keep_from = n_rows - state.cross_ctx; state.target_window.assign( new_rows.begin() + (ptrdiff_t) keep_from * (ptrdiff_t) row_width, @@ -2259,10 +2525,12 @@ static bool dflash_append_target_features( state.target_window_pos.assign(new_positions.begin() + keep_from, new_positions.end()); state.target_window_rows = state.cross_ctx; state.last_target_pos = state.target_window_pos.empty() ? -1 : state.target_window_pos.back(); + dflash_contract_log_append(state, seq_id, new_positions); return true; } const int32_t keep_old_rows = std::min(state.target_window_rows, state.cross_ctx - n_rows); + state.n_rows_dropped += (size_t) std::max(0, state.target_window_rows - keep_old_rows); std::vector next_window((size_t) (keep_old_rows + n_rows) * row_width); std::vector next_window_pos((size_t) (keep_old_rows + n_rows)); @@ -2282,6 +2550,7 @@ static bool dflash_append_target_features( state.target_window_pos = std::move(next_window_pos); state.target_window_rows = keep_old_rows + n_rows; state.last_target_pos = state.target_window_pos.empty() ? -1 : state.target_window_pos.back(); + dflash_contract_log_append(state, seq_id, new_positions); return true; } @@ -2329,6 +2598,7 @@ static void dflash_context_shift( state.target_window_pos = std::move(shifted_positions); state.target_window_rows = (int32_t) state.target_window_pos.size(); state.last_target_pos = state.target_window_pos.empty() ? -1 : state.target_window_pos.back(); + state.n_context_shifts++; } static bool common_speculative_capture_target_features(common_speculative * spec, const common_speculative_feature_view & features) { diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index 2a4e8c95..c3171d76 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -13,8 +13,12 @@ #include "mtmd-helper.h" #include +#include +#include +#include #include #include +#include #include static void server_prompt_checkpoint_update(server_prompt_checkpoint & ckpt, llama_context * ctx, int id, int64_t n_tokens, llama_pos pos_min = -1, llama_pos pos_max = -1, int32_t offset = 0) { @@ -45,6 +49,83 @@ static void log_text(const gpt_params & params_base, const std::string & text) { } } +static bool server_dflash_contract_log_enabled() { + const char * env = std::getenv("IK_DFLASH_CONTRACT_LOG"); + if (env == nullptr || *env == '\0') { + return false; + } + + return std::strcmp(env, "0") != 0 && + std::strcmp(env, "false") != 0 && + std::strcmp(env, "off") != 0; +} + +static std::string server_dflash_contract_format_indices( + const std::vector & values, + size_t edge_count = 4) { + std::ostringstream oss; + oss << '['; + if (values.empty()) { + oss << ']'; + return oss.str(); + } + + const size_t head = std::min(edge_count, values.size()); + for (size_t i = 0; i < head; ++i) { + if (i > 0) { + oss << ','; + } + oss << values[i]; + } + + if (values.size() > edge_count * 2) { + oss << ",...,"; + for (size_t i = values.size() - edge_count; i < values.size(); ++i) { + if (i > values.size() - edge_count) { + oss << ','; + } + oss << values[i]; + } + } else { + for (size_t i = head; i < values.size(); ++i) { + oss << ',' << values[i]; + } + } + + oss << ']'; + return oss.str(); +} + +static void server_dflash_contract_log_accept( + const server_slot & slot, + common_speculative_type spec_type_used, + const char * path, + bool any_rejected, + size_t n_draft, + const std::vector & ids, + llama_pos pos_base, + const std::vector & output_indices) { + if (!server_dflash_contract_log_enabled() || spec_type_used != COMMON_SPECULATIVE_TYPE_DFLASH) { + return; + } + + static std::atomic counter = 0; + const uint64_t ordinal = counter.fetch_add(1, std::memory_order_relaxed); + if (ordinal >= 8) { + return; + } + + LLAMA_LOG_INFO("dflash contract accept[%llu]: slot=%d path=%s rejected=%s drafted=%zu accepted=%zu pos_base=%d output_indices=%s\n", + (unsigned long long) (ordinal + 1), + slot.id, + path, + any_rejected ? "true" : "false", + n_draft, + ids.size(), + (int) pos_base, + server_dflash_contract_format_indices(output_indices).c_str()); +} + static bool params_use_gemma4_external_mtp(const gpt_params & params_base) { return params_base.has_mtp && llama_model_is_gemma4_mtp_assistant(params_base.speculative.model_dft); @@ -4226,6 +4307,16 @@ static void restore_speculative_checkpoint( redecoded_indices[j] = j; } + server_dflash_contract_log_accept( + slot, + spec_type_used, + "restore", + true, + n_draft, + ids, + slot.spec_ckpt.n_past, + redecoded_indices); + if (!common_speculative_commit_accepted_output( slot.spec, ctx, @@ -4338,6 +4429,16 @@ void server_context::speculative_decoding_accept() { restore_speculative_checkpoint(slot, ctx, model, spec_type_used, sampled_before, ids, n_draft, spec_feature_rows_pre, spec_n_past_base); } else { if (server_speculative_has_target_features(slot.params.speculative) && !accepted_output_indices.empty()) { + server_dflash_contract_log_accept( + slot, + spec_type_used, + "direct", + false, + n_draft, + ids, + spec_n_past_base, + accepted_output_indices); + if (!common_speculative_commit_accepted_output( slot.spec, ctx, @@ -4729,6 +4830,7 @@ void server_context::process_batch_tokens(int32_t & n_batch) { } if (server_speculative_has_target_features(params_base.speculative)) { + bool finished_prompt_warmup_batch = false; for (auto & slot : slots) { if (!slot.spec || !server_speculative_has_target_features(slot.params.speculative)) { continue; @@ -4738,8 +4840,7 @@ void server_context::process_batch_tokens(int32_t & n_batch) { continue; } - if ((slot.state != SLOT_STATE_PROCESSING || slot.n_decoded != 0) && - (slot.state != SLOT_STATE_IDLE || slot.command != SLOT_COMMAND_LOAD_PROMPT)) { + if (slot.command != SLOT_COMMAND_LOAD_PROMPT) { continue; } @@ -4747,8 +4848,14 @@ void server_context::process_batch_tokens(int32_t & n_batch) { common_speculative_clear_sequence_hidden(slot.spec, slot.id); slot.spec_prompt_warmup_failed = true; LOG_ERROR("failed to warm up speculative target-feature state from prompt batch for slot %d\n", slot.id); + } else { + finished_prompt_warmup_batch = true; } } + + if (finished_prompt_warmup_batch) { + llama_finish_dflash_capture_batch(ctx, true); + } } for (auto& slot : slots) { diff --git a/src/graphs/build_dflash.cpp b/src/graphs/build_dflash.cpp index 542821ad..ef50f868 100644 --- a/src/graphs/build_dflash.cpp +++ b/src/graphs/build_dflash.cpp @@ -3,33 +3,103 @@ #include "../llama-model.h" #include +#include + +static bool dflash_use_kv_cache_experiment() { + const char * env = std::getenv("IK_DFLASH_KV_CACHE"); + if (env == nullptr || *env == '\0') { + return false; + } + + return std::strcmp(env, "0") != 0 && + std::strcmp(env, "false") != 0 && + std::strcmp(env, "off") != 0; +} + +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); + const int64_t n_target_features = hparams.dflash_n_target_features; + 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); + + GGML_ASSERT(n_embd_head_k == n_embd_head_v); + GGML_ASSERT(n_target_features > 0); + GGML_ASSERT(lctx.ensure_dflash_kv_cache_tensors((int32_t) ctx_len)); + + ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes((int) std::max(1, ctx_len)) + 24 * n_layer, false); + + lctx.dflash_kv_input_target_features = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_target_features, ctx_len); + ggml_set_input(lctx.dflash_kv_input_target_features); + + lctx.dflash_kv_input_pos_ctx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ctx_len); + ggml_set_input(lctx.dflash_kv_input_pos_ctx); + + ggml_tensor * fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, lctx.dflash_kv_input_target_features); + fused_target = llm_build_norm(ctx0, fused_target, hparams, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, cb, -1); + + 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 * Kcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, fused_target); + Kcur_ctx = ggml_reshape_3d(ctx0, Kcur_ctx, n_embd_head_k, n_head_kv, ctx_len); + Kcur_ctx = llm_build_norm(ctx0, Kcur_ctx, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il); + Kcur_ctx = ggml_rope_ext(ctx0, Kcur_ctx, lctx.dflash_kv_input_pos_ctx, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target); + Vcur_ctx = ggml_reshape_3d(ctx0, Vcur_ctx, n_embd_head_v, n_head_kv, ctx_len); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur_ctx, lctx.dflash_k_ctx_cache[(size_t) il])); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur_ctx, lctx.dflash_v_ctx_cache[(size_t) il])); + } + + return gf; +} ggml_cgraph * llm_build_context::build_dflash() { 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 n_target_features = hparams.dflash_n_target_features; - const int64_t ctx_len = std::max(1, (int64_t) cparams.n_ctx - (int64_t) hparams.dflash_block_size); - const int64_t n_kv_total = ctx_len + n_tokens; + const bool use_kv_cache = dflash_use_kv_cache_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); + const int64_t n_kv_total = GGML_PAD(ctx_len + n_tokens, flash_attn ? 256 : 32); + const int64_t n_kv_pad = n_kv_total - (ctx_len + n_tokens); GGML_ASSERT(n_embd_head_k == n_embd_head_v); GGML_ASSERT(n_target_features > 0); + GGML_ASSERT(!use_kv_cache || lctx.ensure_dflash_kv_cache_tensors((int32_t) ctx_len)); ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes((int) std::max(n_tokens, ctx_len)) + 32 * n_layer, false); - lctx.inp_dflash_target_features = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_target_features, ctx_len); - ggml_set_input(lctx.inp_dflash_target_features); - cb(lctx.inp_dflash_target_features, "dflash_target_features", -1); - - lctx.inp_dflash_pos_ctx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ctx_len); - ggml_set_input(lctx.inp_dflash_pos_ctx); - cb(lctx.inp_dflash_pos_ctx, "dflash_pos_ctx", -1); - lctx.inp_dflash_kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv_total, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)); + lctx.dflash_kq_mask_tensor = lctx.inp_dflash_kq_mask; ggml_set_input(lctx.inp_dflash_kq_mask); cb(lctx.inp_dflash_kq_mask, "dflash_kq_mask", -1); ggml_tensor * dflash_kq_mask = flash_attn ? ggml_cast(ctx0, lctx.inp_dflash_kq_mask, GGML_TYPE_F16) : lctx.inp_dflash_kq_mask; + ggml_tensor * fused_target = nullptr; + ggml_tensor * pos_ctx = nullptr; + if (!use_kv_cache) { + lctx.inp_dflash_target_features = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_target_features, ctx_len); + ggml_set_input(lctx.inp_dflash_target_features); + cb(lctx.inp_dflash_target_features, "dflash_target_features", -1); + + lctx.inp_dflash_pos_ctx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ctx_len); + ggml_set_input(lctx.inp_dflash_pos_ctx); + cb(lctx.inp_dflash_pos_ctx, "dflash_pos_ctx", -1); + + fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, lctx.inp_dflash_target_features); + fused_target = llm_build_norm(ctx0, fused_target, hparams, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, cb, -1); + pos_ctx = lctx.inp_dflash_pos_ctx; + } + ggml_tensor * tok_embd = model.tok_embd; if (tok_embd == nullptr) { tok_embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_Q4_0, n_embd, hparams.n_vocab); @@ -39,10 +109,6 @@ ggml_cgraph * llm_build_context::build_dflash() { ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = (n_tokens > 1 && n_outputs < n_tokens) ? build_inp_out_ids() : nullptr; - ggml_tensor * fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, lctx.inp_dflash_target_features); - fused_target = llm_build_norm(ctx0, fused_target, hparams, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, cb, -1); - cb(fused_target, "dflash_target_fused", -1); - const float kq_scale = 1.0f / std::sqrt((float) n_embd_head_k); for (int il = 0; il < n_layer; ++il) { @@ -71,25 +137,35 @@ ggml_cgraph * llm_build_context::build_dflash() { Vcur_noise = ggml_reshape_3d(ctx0, Vcur_noise, n_embd_head_v, n_head_kv, n_tokens); cb(Vcur_noise, "Vcur_noise", il); - ggml_tensor * Kcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, fused_target); - Kcur_ctx = ggml_reshape_3d(ctx0, Kcur_ctx, n_embd_head_k, n_head_kv, ctx_len); - Kcur_ctx = llm_build_norm(ctx0, Kcur_ctx, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il); - Kcur_ctx = ggml_rope_ext(ctx0, Kcur_ctx, lctx.inp_dflash_pos_ctx, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - cb(Kcur_ctx, "Kcur_ctx", il); + ggml_tensor * Kcur_ctx = nullptr; + ggml_tensor * Vcur_ctx = nullptr; + if (use_kv_cache) { + Kcur_ctx = lctx.dflash_k_ctx_cache[(size_t) il]; + Vcur_ctx = lctx.dflash_v_ctx_cache[(size_t) il]; + cb(Kcur_ctx, "Kcur_ctx_cache", il); + cb(Vcur_ctx, "Vcur_ctx_cache", il); + } else { + Kcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, fused_target); + Kcur_ctx = ggml_reshape_3d(ctx0, Kcur_ctx, n_embd_head_k, n_head_kv, ctx_len); + Kcur_ctx = llm_build_norm(ctx0, Kcur_ctx, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il); + Kcur_ctx = ggml_rope_ext(ctx0, Kcur_ctx, pos_ctx, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); - ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target); - Vcur_ctx = ggml_reshape_3d(ctx0, Vcur_ctx, n_embd_head_v, n_head_kv, ctx_len); - cb(Vcur_ctx, "Vcur_ctx", il); + Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target); + Vcur_ctx = ggml_reshape_3d(ctx0, Vcur_ctx, n_embd_head_v, n_head_kv, ctx_len); + cb(Kcur_ctx, "Kcur_ctx", il); + cb(Vcur_ctx, "Vcur_ctx", il); + } 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); + } cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); - - Kcur = ggml_cast(ctx0, Kcur, GGML_TYPE_F16); - Vcur = ggml_cast(ctx0, Vcur, GGML_TYPE_F16); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur_f16", il); cb(Vcur, "Vcur_f16", il); diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index ad89e6e2..081215dd 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -35,7 +35,9 @@ llm_build_context::llm_build_context( const llm_build_cb & cb, bool worst_case, bool warmup, - int n_outputs_) : + int n_outputs_, + bool clear_lctx_inputs, + std::vector * buf_compute_meta_override) : model (lctx.model), lctx (lctx), hparams (model.hparams), @@ -82,8 +84,9 @@ llm_build_context::llm_build_context( thresh_experts (cparams.thresh_experts), pooling_type (cparams.pooling_type), rope_type (hparams.rope_type), + clear_lctx_inputs(clear_lctx_inputs), cb (cb), - buf_compute_meta (lctx.buf_compute_meta) { + buf_compute_meta (buf_compute_meta_override ? *buf_compute_meta_override : lctx.buf_compute_meta) { // all initializations should be done in init() } @@ -96,25 +99,27 @@ void llm_build_context::init() { ctx0 = ggml_init(params); - lctx.inp_tokens = nullptr; - lctx.inp_embd = nullptr; - lctx.inp_pos = nullptr; - lctx.inp_out_ids = nullptr; - lctx.inp_KQ_mask = nullptr; - lctx.inp_KQ_mask_swa = nullptr; - lctx.inp_K_shift = nullptr; - lctx.inp_mean = nullptr; - lctx.inp_cls = nullptr; - lctx.inp_s_copy = nullptr; - lctx.inp_s_mask = nullptr; - lctx.inp_s_seq = nullptr; - lctx.inp_s_seq_qnext = nullptr; - lctx.inp_pos_bucket = nullptr; - lctx.inp_embd_enc = nullptr; - lctx.inp_KQ_mask_cross = nullptr; - lctx.inp_dflash_target_features = nullptr; - lctx.inp_dflash_pos_ctx = nullptr; - lctx.inp_dflash_kq_mask = nullptr; + if (clear_lctx_inputs) { + lctx.inp_tokens = nullptr; + lctx.inp_embd = nullptr; + lctx.inp_pos = nullptr; + lctx.inp_out_ids = nullptr; + lctx.inp_KQ_mask = nullptr; + lctx.inp_KQ_mask_swa = nullptr; + lctx.inp_K_shift = nullptr; + lctx.inp_mean = nullptr; + lctx.inp_cls = nullptr; + lctx.inp_s_copy = nullptr; + lctx.inp_s_mask = nullptr; + lctx.inp_s_seq = nullptr; + lctx.inp_s_seq_qnext = nullptr; + lctx.inp_pos_bucket = nullptr; + lctx.inp_embd_enc = nullptr; + lctx.inp_KQ_mask_cross = nullptr; + lctx.inp_dflash_target_features = nullptr; + lctx.inp_dflash_pos_ctx = nullptr; + lctx.inp_dflash_kq_mask = nullptr; + } } void llm_build_context::free() { @@ -2164,6 +2169,23 @@ struct ggml_cgraph * llm_build_context::llama_build_graph_s_copy(llama_context & return result; } +struct ggml_cgraph * llm_build_context::llama_build_graph_dflash_kv_cache(llama_context & lctx) { + llama_batch dummy; + dummy.n_tokens = 0; + + llm_build_cb cb = [&](struct ggml_tensor * , const char * , int ) { }; + + struct llm_build_context llm(lctx, dummy, cb, false, false, 0, false, &lctx.dflash_buf_compute_meta); + + llm.init(); + + struct ggml_cgraph * result = llm.build_dflash_kv_cache(); + + 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 7542aff6..ec7cbbb9 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -89,6 +89,7 @@ struct llm_build_context { const enum llama_pooling_type pooling_type; const enum llama_rope_type rope_type; + const bool clear_lctx_inputs; const llm_build_cb & cb; @@ -103,7 +104,9 @@ struct llm_build_context { const llm_build_cb & cb, bool worst_case, bool warmup, - int n_outputs = 0); + int n_outputs = 0, + bool clear_lctx_inputs = true, + std::vector * buf_compute_meta_override = nullptr); void init(); @@ -244,6 +247,8 @@ struct llm_build_context { ggml_cgraph * build_dflash(); + ggml_cgraph * build_dflash_kv_cache(); + ggml_cgraph * build_starcoder2(); ggml_cgraph * build_mamba(); @@ -459,6 +464,8 @@ llm_expert_gating_func_type gating_op, static ggml_cgraph * llama_build_graph_s_copy(llama_context & lctx); + static ggml_cgraph * llama_build_graph_dflash_kv_cache(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 c4f62ac1..9e902003 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -289,6 +289,16 @@ struct llama_context { std::vector dflash_feature_view_buffer; std::vector dflash_pos_ctx_data; std::vector dflash_kq_mask_data; + int32_t dflash_visible_cross_ctx = 0; + std::vector dflash_k_ctx_cache; + std::vector dflash_v_ctx_cache; + struct ggml_context * dflash_cache_ctx = nullptr; + ggml_backend_buffer_t dflash_cache_buf = nullptr; + std::vector dflash_buf_compute_meta; + ggml_backend_sched_t dflash_sched = nullptr; + 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; struct dflash_capture_state { std::vector layer_ids; @@ -299,6 +309,7 @@ struct llama_context { void * prev_cb_eval_user_data = nullptr; }; std::unique_ptr dflash_capture; + llama_dflash_profile_stats dflash_profile; // input tensors struct ggml_tensor * inp_tokens; // I32 [n_batch] @@ -340,6 +351,9 @@ struct llama_context { bool update_cache_copies(); + bool ensure_dflash_kv_cache_tensors(int32_t cross_ctx); + void free_dflash_kv_cache_tensors(); + bool prepare_mtp_graph_inputs( struct llama_context & lctx); void set_mtp_op_type(llama_mtp_op_type value); diff --git a/src/llama-spec-features.cpp b/src/llama-spec-features.cpp index ccc6fb5d..91b5d41d 100644 --- a/src/llama-spec-features.cpp +++ b/src/llama-spec-features.cpp @@ -1,13 +1,39 @@ #include "llama-spec-features.h" #include +#include #include #include #include +#include #include "llama-model.h" #include "llama-context.h" +static bool llama_dflash_positions_strictly_increasing( + const llama_pos * positions, + int32_t n_rows, + llama_pos & first_pos, + llama_pos & last_pos) { + first_pos = -1; + last_pos = -1; + + if (positions == nullptr || n_rows <= 0) { + return false; + } + + first_pos = positions[0]; + last_pos = positions[n_rows - 1]; + + for (int32_t i = 1; i < n_rows; ++i) { + if (positions[i] <= positions[i - 1]) { + return false; + } + } + + return true; +} + uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx) { if (ctx == nullptr) { return 0; @@ -21,6 +47,40 @@ uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx) { return hparams.n_embd; } +void llama_dflash_profile_reset(struct llama_context * ctx) { + if (ctx == nullptr) { + return; + } + + ctx->dflash_profile = {}; +} + +void llama_set_dflash_visible_cross_ctx( + struct llama_context * ctx, + int32_t cross_ctx) { + if (ctx == nullptr) { + return; + } + + ctx->dflash_visible_cross_ctx = std::max(0, cross_ctx); +} + +int32_t llama_get_dflash_visible_cross_ctx( + const struct llama_context * ctx) { + return ctx != nullptr ? ctx->dflash_visible_cross_ctx : 0; +} + +bool llama_dflash_profile_get_stats( + const struct llama_context * ctx, + llama_dflash_profile_stats * stats) { + if (ctx == nullptr || stats == nullptr) { + return false; + } + + *stats = ctx->dflash_profile; + return true; +} + int32_t llama_model_dflash_block_size(const struct llama_model * model) { return model ? (int32_t) model->hparams.dflash_block_size : 0; } @@ -92,32 +152,116 @@ bool llama_set_draft_input_hidden_state_copy( return true; } -bool llama_set_dflash_target_features_copy( +static bool llama_set_dflash_target_features_impl( struct llama_context * ctx, const float * target_features, size_t n_floats, int32_t n_rows, - const llama_pos * target_positions) { + const llama_pos * target_positions, + bool copy_data) { if (ctx == nullptr || target_features == nullptr || n_floats == 0 || n_rows <= 0) { return false; } - ctx->dflash_target_features_owned.assign(target_features, target_features + n_floats); - ctx->dflash_target_features = ctx->dflash_target_features_owned.data(); + auto & profile = ctx->dflash_profile; + const int64_t t_start_us = ggml_time_us(); + const int32_t row_width = n_rows > 0 ? (int32_t) (n_floats / (size_t) n_rows) : 0; + llama_pos first_pos = -1; + llama_pos last_pos = -1; + + if (copy_data) { + ctx->dflash_target_features_owned.assign(target_features, target_features + n_floats); + ctx->dflash_target_features = ctx->dflash_target_features_owned.data(); + } else { + ctx->dflash_target_features_owned.clear(); + ctx->dflash_target_features = target_features; + } ctx->dflash_target_features_n_floats = n_floats; ctx->dflash_target_features_n_rows = n_rows; if (target_positions != nullptr) { - ctx->dflash_target_positions_owned.assign(target_positions, target_positions + n_rows); - ctx->dflash_target_positions = ctx->dflash_target_positions_owned.data(); + if (copy_data) { + ctx->dflash_target_positions_owned.assign(target_positions, target_positions + n_rows); + ctx->dflash_target_positions = ctx->dflash_target_positions_owned.data(); + } else { + ctx->dflash_target_positions_owned.clear(); + ctx->dflash_target_positions = target_positions; + } ctx->dflash_target_positions_n = (size_t) n_rows; } else { ctx->dflash_target_positions_owned.clear(); ctx->dflash_target_positions = nullptr; ctx->dflash_target_positions_n = 0; } + + profile.set_target_copy_calls++; + profile.set_target_copy_us += (uint64_t) (ggml_time_us() - t_start_us); + profile.set_target_rows += (uint64_t) n_rows; + profile.set_target_copy_bytes += n_floats * sizeof(float) + (target_positions ? (size_t) n_rows * sizeof(llama_pos) : 0); + profile.last_n_rows = n_rows; + profile.last_width = row_width; + + if (target_positions == nullptr) { + profile.set_target_missing_positions++; + profile.last_pos_first = -1; + profile.last_pos_last = -1; + } else { + if (!llama_dflash_positions_strictly_increasing(target_positions, n_rows, first_pos, last_pos)) { + profile.set_target_non_monotonic_positions++; + } + profile.last_pos_first = first_pos; + profile.last_pos_last = last_pos; + } + return true; } +bool llama_set_dflash_target_features_copy( + struct llama_context * ctx, + const float * target_features, + size_t n_floats, + int32_t n_rows, + const llama_pos * target_positions) { + return llama_set_dflash_target_features_impl(ctx, target_features, n_floats, n_rows, target_positions, true); +} + +bool llama_set_dflash_target_features_view( + struct llama_context * ctx, + const float * target_features, + size_t n_floats, + int32_t n_rows, + const llama_pos * target_positions) { + return llama_set_dflash_target_features_impl(ctx, target_features, n_floats, n_rows, target_positions, false); +} + +static void llama_record_dflash_capture_phase( + struct llama_context * ctx, + bool is_prompt_warmup, + int32_t row_count, + int32_t row_width) { + if (ctx == nullptr || row_count <= 0 || row_width <= 0) { + return; + } + + auto & profile = ctx->dflash_profile; + if (is_prompt_warmup) { + profile.capture_prompt_batches++; + if (profile.capture_prompt_last_rows > 0 && profile.capture_prompt_last_width > 0 && + (profile.capture_prompt_last_rows != row_count || profile.capture_prompt_last_width != row_width)) { + profile.capture_prompt_shape_changes++; + } + profile.capture_prompt_last_rows = row_count; + profile.capture_prompt_last_width = row_width; + } else { + profile.capture_verify_batches++; + if (profile.capture_verify_last_rows > 0 && profile.capture_verify_last_width > 0 && + (profile.capture_verify_last_rows != row_count || profile.capture_verify_last_width != row_width)) { + profile.capture_verify_shape_changes++; + } + profile.capture_verify_last_rows = row_count; + profile.capture_verify_last_width = row_width; + } +} + static bool llama_dflash_parse_layer_id(const struct ggml_tensor * tensor, int32_t & layer_id) { if (tensor == nullptr) { return false; @@ -236,6 +380,22 @@ void llama_clear_dflash_capture(struct llama_context * ctx) { } } +void llama_finish_dflash_capture_batch( + struct llama_context * ctx, + bool is_prompt_warmup) { + if (ctx == nullptr || !ctx->dflash_capture) { + return; + } + + auto & capture = *ctx->dflash_capture; + llama_record_dflash_capture_phase(ctx, is_prompt_warmup, capture.row_count, capture.row_width); + + // Reset the batch-local reference shape so the next decode only compares layers within + // the same batch, not against the previous prompt/verify batch. + capture.row_count = 0; + capture.row_width = 0; +} + static bool llama_spec_prepare_dflash_capture( struct llama_context * ctx, int32_t & row_count, @@ -245,18 +405,31 @@ static bool llama_spec_prepare_dflash_capture( return false; } + auto & profile = ctx->dflash_profile; + profile.capture_prepare_calls++; + const int64_t t_sync_us = ggml_time_us(); llama_synchronize(ctx); + profile.capture_prepare_sync_us += (uint64_t) (ggml_time_us() - t_sync_us); auto & capture = *ctx->dflash_capture; row_count = capture.row_count; row_width = capture.row_width; n_layers = (int32_t) capture.layer_ids.size(); if (row_count <= 0 || row_width <= 0 || n_layers <= 0 || capture.layer_rows.size() != (size_t) n_layers) { + profile.capture_prepare_failures++; return false; } - for (const auto & rows : capture.layer_rows) { + for (int32_t layer_idx = 0; layer_idx < n_layers; ++layer_idx) { + const auto & rows = capture.layer_rows[(size_t) layer_idx]; if (rows.size() != (size_t) row_count * (size_t) row_width) { + profile.capture_prepare_failures++; + profile.capture_layer_shape_mismatch++; + if (profile.capture_layer_shape_mismatch <= 3) { + LLAMA_LOG_WARN("%s: DFlash capture rows mismatch for layer %d: got=%zu expected=%zu (rows=%d width=%d)\n", + __func__, capture.layer_ids[(size_t) layer_idx], rows.size(), + (size_t) row_count * (size_t) row_width, row_count, row_width); + } return false; } } @@ -264,6 +437,164 @@ static bool llama_spec_prepare_dflash_capture( return true; } +static bool llama_dflash_contract_log_enabled() { + const char * env = std::getenv("IK_DFLASH_CONTRACT_LOG"); + if (env == nullptr || *env == '\0') { + return false; + } + + return std::strcmp(env, "0") != 0 && + std::strcmp(env, "false") != 0 && + std::strcmp(env, "off") != 0; +} + +template +static std::string llama_dflash_contract_format_values( + const std::vector & values, + size_t edge_count = 4) { + std::ostringstream oss; + oss << '['; + if (values.empty()) { + oss << ']'; + return oss.str(); + } + + const size_t head = std::min(edge_count, values.size()); + for (size_t i = 0; i < head; ++i) { + if (i > 0) { + oss << ','; + } + oss << values[i]; + } + + if (values.size() > edge_count * 2) { + oss << ",...,"; + for (size_t i = values.size() - edge_count; i < values.size(); ++i) { + if (i > values.size() - edge_count) { + oss << ','; + } + oss << values[i]; + } + } else { + for (size_t i = head; i < values.size(); ++i) { + oss << ',' << values[i]; + } + } + + oss << ']'; + return oss.str(); +} + +static std::vector llama_dflash_contract_collect_batch_positions( + const llama_batch & batch, + const std::vector & batch_indices) { + std::vector positions; + positions.reserve(batch_indices.size()); + for (int32_t batch_index : batch_indices) { + positions.push_back(batch.pos[batch_index]); + } + return positions; +} + +static void llama_dflash_contract_summarize_positions( + const std::vector & positions, + llama_pos & first_pos, + llama_pos & last_pos, + int32_t & gap_count, + int32_t & nonmono_count) { + first_pos = -1; + last_pos = -1; + gap_count = 0; + nonmono_count = 0; + if (positions.empty()) { + return; + } + + first_pos = positions.front(); + last_pos = positions.back(); + for (size_t i = 1; i < positions.size(); ++i) { + if (positions[i] <= positions[i - 1]) { + nonmono_count++; + } else if (positions[i] != positions[i - 1] + 1) { + gap_count++; + } + } +} + +static void llama_dflash_contract_log_feature_view( + const char * kind, + llama_seq_id seq_id, + const llama_batch & batch, + int32_t row_count, + int32_t row_width, + int32_t n_layers, + int32_t batch_row_offset, + const std::vector & row_indices, + const std::vector & batch_indices) { + if (!llama_dflash_contract_log_enabled()) { + return; + } + + static std::atomic counter = 0; + const uint64_t ordinal = counter.fetch_add(1, std::memory_order_relaxed); + if (ordinal >= 8) { + return; + } + + const std::vector positions = llama_dflash_contract_collect_batch_positions(batch, batch_indices); + llama_pos first_pos = -1; + llama_pos last_pos = -1; + int32_t gap_count = 0; + int32_t nonmono_count = 0; + llama_dflash_contract_summarize_positions(positions, first_pos, last_pos, gap_count, nonmono_count); + + LLAMA_LOG_INFO("%s[%llu]: kind=%s seq=%d batch_tokens=%d capture_rows=%d row_width=%d layers=%d batch_row_offset=%d row_indices=%s batch_indices=%s batch_pos=%s pos=[%d..%d] gaps=%d nonmono=%d\n", + __func__, + (unsigned long long) (ordinal + 1), + kind, + (int) seq_id, + batch.n_tokens, + row_count, + row_width, + n_layers, + batch_row_offset, + llama_dflash_contract_format_values(row_indices).c_str(), + llama_dflash_contract_format_values(batch_indices).c_str(), + llama_dflash_contract_format_values(positions).c_str(), + (int) first_pos, + (int) last_pos, + gap_count, + nonmono_count); +} + +static void llama_dflash_contract_log_output_indices( + struct llama_context * ctx, + const std::vector & output_indices) { + if (!llama_dflash_contract_log_enabled()) { + return; + } + + static std::atomic counter = 0; + const uint64_t ordinal = counter.fetch_add(1, std::memory_order_relaxed); + if (ordinal >= 8) { + return; + } + + int32_t row_count = 0; + int32_t row_width = 0; + int32_t n_layers = 0; + const bool have_capture = llama_spec_prepare_dflash_capture(ctx, row_count, row_width, n_layers); + + LLAMA_LOG_INFO("%s[%llu]: output_indices=%s capture_rows=%d row_width=%d layers=%d have_capture=%s\n", + __func__, + (unsigned long long) (ordinal + 1), + llama_dflash_contract_format_values(output_indices).c_str(), + row_count, + row_width, + n_layers, + have_capture ? "true" : "false"); +} + static bool llama_spec_materialize_dflash_rows( struct llama_context * ctx, const std::vector & row_indices, @@ -275,10 +606,15 @@ static bool llama_spec_materialize_dflash_rows( return false; } + auto & profile = ctx->dflash_profile; + profile.capture_materialize_calls++; + const int64_t t_start_us = ggml_time_us(); + int32_t row_count = 0; int32_t row_width = 0; int32_t n_layers = 0; if (!llama_spec_prepare_dflash_capture(ctx, row_count, row_width, n_layers)) { + profile.capture_materialize_failures++; return false; } @@ -294,6 +630,7 @@ static bool llama_spec_materialize_dflash_rows( if (row_index < 0 || row_index >= row_count) { rows_out.clear(); combined_width = 0; + profile.capture_materialize_failures++; return false; } @@ -304,6 +641,10 @@ static bool llama_spec_materialize_dflash_rows( } } + profile.capture_materialize_us += (uint64_t) (ggml_time_us() - t_start_us); + profile.capture_materialize_rows += (uint64_t) row_indices.size(); + profile.capture_materialize_bytes += rows_out.size() * sizeof(float); + return true; } @@ -412,6 +753,17 @@ bool llama_spec_get_dflash_feature_view( }); } + llama_dflash_contract_log_feature_view( + "batch", + view.rows.empty() ? -1 : view.rows.front().seq_id, + batch, + row_count, + row_width, + n_layers, + batch_row_offset, + row_indices, + batch_indices); + return true; } @@ -470,6 +822,17 @@ bool llama_spec_get_dflash_feature_view_for_seq( }); } + llama_dflash_contract_log_feature_view( + "seq", + seq_id, + batch, + row_count, + row_width, + n_layers, + batch_row_offset, + row_indices, + batch_indices); + return true; } @@ -576,5 +939,7 @@ bool llama_spec_copy_dflash_rows_from_output_indices( return false; } + llama_dflash_contract_log_output_indices(ctx, output_indices); + return hidden_rows.size() == (size_t) output_indices.size() * (size_t) combined_width; } diff --git a/src/llama-spec-features.h b/src/llama-spec-features.h index 130d3895..20f0ff51 100644 --- a/src/llama-spec-features.h +++ b/src/llama-spec-features.h @@ -23,8 +23,74 @@ struct llama_spec_feature_view { std::vector rows; }; +struct llama_dflash_profile_stats { + uint64_t set_target_copy_calls = 0; + uint64_t set_target_copy_us = 0; + uint64_t set_target_rows = 0; + uint64_t set_target_copy_bytes = 0; + uint64_t set_target_missing_positions = 0; + uint64_t set_target_non_monotonic_positions = 0; + + uint64_t capture_prepare_calls = 0; + uint64_t capture_prepare_sync_us = 0; + uint64_t capture_prepare_failures = 0; + uint64_t capture_layer_shape_mismatch = 0; + uint64_t capture_prompt_batches = 0; + uint64_t capture_prompt_shape_changes = 0; + uint64_t capture_verify_batches = 0; + uint64_t capture_verify_shape_changes = 0; + uint64_t capture_materialize_calls = 0; + uint64_t capture_materialize_rows = 0; + uint64_t capture_materialize_bytes = 0; + uint64_t capture_materialize_us = 0; + uint64_t capture_materialize_failures = 0; + + uint64_t graph_prepare_calls = 0; + uint64_t graph_prepare_total_us = 0; + uint64_t graph_feature_copy_us = 0; + uint64_t graph_pos_copy_us = 0; + uint64_t graph_mask_build_us = 0; + uint64_t graph_kv_cache_compute_us = 0; + uint64_t graph_kv_cache_calls = 0; + uint64_t graph_feature_bytes = 0; + uint64_t graph_pos_bytes = 0; + uint64_t graph_mask_bytes = 0; + uint64_t graph_visible_kv_sum = 0; + uint64_t graph_visible_kv_max = 0; + uint64_t graph_pos_fallbacks = 0; + uint64_t graph_pos_non_monotonic = 0; + uint64_t graph_shape_failures = 0; + uint64_t graph_mask_overflow = 0; + + int32_t last_n_rows = 0; + int32_t last_width = 0; + int32_t last_cross_ctx = 0; + int32_t last_left_pad = 0; + int32_t last_n_tokens = 0; + int32_t last_n_kv_total = 0; + int32_t capture_prompt_last_rows = 0; + int32_t capture_prompt_last_width = 0; + int32_t capture_verify_last_rows = 0; + int32_t capture_verify_last_width = 0; + llama_pos last_pos_first = -1; + llama_pos last_pos_last = -1; +}; + uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx); +void llama_dflash_profile_reset(struct llama_context * ctx); + +void llama_set_dflash_visible_cross_ctx( + struct llama_context * ctx, + int32_t cross_ctx); + +int32_t llama_get_dflash_visible_cross_ctx( + const struct llama_context * ctx); + +bool llama_dflash_profile_get_stats( + const struct llama_context * ctx, + llama_dflash_profile_stats * stats); + int32_t llama_model_dflash_block_size(const struct llama_model * model); int32_t llama_model_dflash_mask_token_id(const struct llama_model * model); @@ -54,6 +120,13 @@ bool llama_set_dflash_target_features_copy( int32_t n_rows, const llama_pos * target_positions); +bool llama_set_dflash_target_features_view( + struct llama_context * ctx, + const float * target_features, + size_t n_floats, + int32_t n_rows, + const llama_pos * target_positions); + bool llama_set_dflash_capture_layers( struct llama_context * ctx, const int32_t * layer_ids, @@ -61,6 +134,10 @@ bool llama_set_dflash_capture_layers( void llama_clear_dflash_capture(struct llama_context * ctx); +void llama_finish_dflash_capture_batch( + struct llama_context * ctx, + bool is_prompt_warmup); + bool llama_spec_get_hidden_feature_view( struct llama_context * ctx, const llama_batch & batch, diff --git a/src/llama.cpp b/src/llama.cpp index a9b443c7..af37211e 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -565,6 +565,84 @@ void llama_context::reset_scheduler() { prev_mtp.reset(); } +bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) { + const int32_t target_cross_ctx = std::max(1, cross_ctx); + 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) { + return true; + } + + free_dflash_kv_cache_tensors(); + if (dflash_sched != nullptr) { + ggml_backend_sched_free(dflash_sched); + dflash_sched = nullptr; + } + dflash_buf_compute_meta.clear(); + } + + ggml_backend_buffer_type_t buft = llama_default_buffer_type_cpu(true); + + ggml_init_params params = { + /*.mem_size =*/ (size_t) (2 * std::max(1, n_layer)) * ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + dflash_cache_ctx = ggml_init(params); + if (dflash_cache_ctx == nullptr) { + return false; + } + + dflash_k_ctx_cache.resize((size_t) n_layer); + dflash_v_ctx_cache.resize((size_t) n_layer); + for (int32_t il = 0; il < n_layer; ++il) { + dflash_k_ctx_cache[(size_t) il] = ggml_new_tensor_3d(dflash_cache_ctx, GGML_TYPE_F32, n_embd_head_k, n_head_kv, target_cross_ctx); + dflash_v_ctx_cache[(size_t) il] = ggml_new_tensor_3d(dflash_cache_ctx, GGML_TYPE_F32, n_embd_head_v, n_head_kv, target_cross_ctx); + if (dflash_k_ctx_cache[(size_t) il] == nullptr || dflash_v_ctx_cache[(size_t) il] == nullptr) { + free_dflash_kv_cache_tensors(); + return false; + } + + ggml_set_input(dflash_k_ctx_cache[(size_t) il]); + ggml_set_input(dflash_v_ctx_cache[(size_t) il]); + ggml_format_name(dflash_k_ctx_cache[(size_t) il], "dflash_k_ctx_cache_%d", il); + ggml_format_name(dflash_v_ctx_cache[(size_t) il], "dflash_v_ctx_cache_%d", il); + } + + dflash_cache_buf = ggml_backend_alloc_ctx_tensors_from_buft(dflash_cache_ctx, buft); + if (dflash_cache_buf == nullptr) { + free_dflash_kv_cache_tensors(); + return false; + } + + ggml_backend_buffer_clear(dflash_cache_buf, 0); + return true; +} + +void llama_context::free_dflash_kv_cache_tensors() { + dflash_k_ctx_cache.clear(); + dflash_v_ctx_cache.clear(); + dflash_kv_input_target_features = nullptr; + dflash_kv_input_pos_ctx = nullptr; + dflash_kq_mask_tensor = nullptr; + + if (dflash_cache_buf != nullptr) { + ggml_backend_buffer_free(dflash_cache_buf); + dflash_cache_buf = nullptr; + } + if (dflash_cache_ctx != nullptr) { + ggml_free(dflash_cache_ctx); + dflash_cache_ctx = nullptr; + } +} + bool llama_context::can_reuse_graph(const llama_batch & u_batch) { if (!cparams.graph_reuse) return false; //if (kv_self.save_per_step_ssm) return false; @@ -687,6 +765,10 @@ void llama_context::set_mtp_op_type(llama_mtp_op_type value) { } llama_context::~llama_context() { + if (dflash_sched != nullptr) { + ggml_backend_sched_free(dflash_sched); + } + free_dflash_kv_cache_tensors(); ggml_backend_sched_free(sched); for (ggml_backend_t backend : backends) { @@ -4934,6 +5016,30 @@ static void llama_graph_compute( // fprintf(stderr, "splits: %d\n", ggml_backend_sched_get_n_splits(lctx.sched)); } +static void llama_graph_compute_sched( + llama_context & lctx, + ggml_backend_sched_t sched, + ggml_cgraph * gf, + int n_threads) { +#ifdef GGML_USE_METAL + if (ggml_backend_is_metal(lctx.backend_metal)) { + ggml_backend_metal_set_n_cb(lctx.backend_metal, n_threads); + } +#endif + + if (lctx.backend_cpu != nullptr) { + ggml_backend_cpu_set_n_threads(lctx.backend_cpu, n_threads); + ggml_backend_cpu_set_abort_callback(lctx.backend_cpu, lctx.abort_callback, lctx.abort_callback_data); + } +#ifdef GGML_USE_BLAS + if (lctx.backend_blas != nullptr) { + ggml_backend_blas_set_n_threads(lctx.backend_blas, n_threads); + } +#endif + + ggml_backend_sched_graph_compute_async(sched, gf); +} + static bool prepare_mtp_graph_inputs( struct llama_context & lctx, uint32_t cur_token, @@ -4984,62 +5090,206 @@ static bool prepare_mtp_graph_inputs( static bool prepare_dflash_graph_inputs( struct llama_context & lctx, uint32_t n_tokens) { - ggml_tensor * target_hidden = lctx.inp_dflash_target_features; - ggml_tensor * pos_ctx = lctx.inp_dflash_pos_ctx; - ggml_tensor * kq_mask = lctx.inp_dflash_kq_mask; + const char * dflash_kv_cache_env = std::getenv("IK_DFLASH_KV_CACHE"); + const bool use_kv_cache = dflash_kv_cache_env != nullptr && *dflash_kv_cache_env != '\0' && + std::strcmp(dflash_kv_cache_env, "0") != 0 && + std::strcmp(dflash_kv_cache_env, "false") != 0 && + std::strcmp(dflash_kv_cache_env, "off") != 0; + const int32_t cross_ctx = lctx.dflash_visible_cross_ctx > 0 + ? lctx.dflash_visible_cross_ctx + : std::max(1, (int32_t) lctx.cparams.n_ctx - (int32_t) lctx.model.hparams.dflash_block_size); + ggml_tensor * kq_mask = lctx.dflash_kq_mask_tensor; - if (target_hidden == nullptr || pos_ctx == nullptr || kq_mask == nullptr) { + if (kq_mask == nullptr) { LLAMA_LOG_ERROR("%s: DFlash graph inputs are not initialized\n", __func__); return false; } + if (use_kv_cache) { + if (!lctx.ensure_dflash_kv_cache_tensors(cross_ctx) || lctx.dflash_k_ctx_cache.empty() || lctx.dflash_v_ctx_cache.empty()) { + LLAMA_LOG_ERROR("%s: DFlash K/V cache inputs are not initialized\n", __func__); + return false; + } + } else if (lctx.inp_dflash_target_features == nullptr || lctx.inp_dflash_pos_ctx == nullptr) { + LLAMA_LOG_ERROR("%s: DFlash inline inputs are not initialized\n", __func__); + return false; + } + const float * src = lctx.dflash_target_features; const llama_pos * src_pos = lctx.dflash_target_positions; const size_t total_floats = lctx.dflash_target_features_n_floats; const size_t total_positions = lctx.dflash_target_positions_n; const int32_t n_rows = lctx.dflash_target_features_n_rows; - const int32_t width = (int32_t) target_hidden->ne[0]; - const int32_t cross_ctx = (int32_t) target_hidden->ne[1]; + const int32_t width = (int32_t) lctx.model.hparams.dflash_n_target_features; + const int32_t graph_cross_ctx = use_kv_cache + ? (lctx.dflash_k_ctx_cache.front() != nullptr ? (int32_t) lctx.dflash_k_ctx_cache.front()->ne[2] : 0) + : (lctx.inp_dflash_target_features != nullptr ? (int32_t) lctx.inp_dflash_target_features->ne[1] : 0); const int32_t n_mask_tokens = (int32_t) kq_mask->ne[1]; const int32_t n_kv_total = (int32_t) kq_mask->ne[0]; + auto & profile = lctx.dflash_profile; + const int64_t t_total_us = ggml_time_us(); + profile.graph_prepare_calls++; + profile.last_n_rows = n_rows; + profile.last_width = width; + profile.last_cross_ctx = cross_ctx; + profile.last_n_tokens = (int32_t) n_tokens; + profile.last_n_kv_total = n_kv_total; + + if (graph_cross_ctx != cross_ctx) { + profile.graph_shape_failures++; + + LLAMA_LOG_ERROR("%s: DFlash graph cross_ctx drift (graph=%d configured=%d)\n", + __func__, graph_cross_ctx, cross_ctx); + return false; + } if (src == nullptr || total_floats == 0 || n_rows <= 0) { + profile.graph_shape_failures++; LLAMA_LOG_ERROR("%s: missing DFlash target features\n", __func__); return false; } if (n_rows > cross_ctx || total_floats != (size_t) n_rows * (size_t) width) { + profile.graph_shape_failures++; LLAMA_LOG_ERROR("%s: invalid DFlash target feature shape (rows=%d width=%d floats=%zu cross_ctx=%d)\n", __func__, n_rows, width, total_floats, cross_ctx); return false; } - lctx.dflash_target_features_padded.assign((size_t) cross_ctx * (size_t) width, 0.0f); - const size_t dst_offset = (size_t) (cross_ctx - n_rows) * (size_t) width; - const int32_t left_pad = cross_ctx - n_rows; - std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset); - ggml_backend_tensor_set(target_hidden, lctx.dflash_target_features_padded.data(), 0, ggml_nbytes(target_hidden)); + if (n_kv_total < cross_ctx + (int32_t) n_tokens) { + profile.graph_mask_overflow++; + LLAMA_LOG_ERROR("%s: invalid DFlash mask shape (n_kv_total=%d < cross_ctx+n_tokens=%d)\n", + __func__, n_kv_total, cross_ctx + (int32_t) n_tokens); + return false; + } + const int32_t left_pad = cross_ctx - n_rows; + const size_t padded_floats = (size_t) cross_ctx * (size_t) width; + const size_t dst_offset = (size_t) left_pad * (size_t) width; + const int64_t t_feature_us = ggml_time_us(); + profile.last_left_pad = left_pad; + if (lctx.dflash_target_features_padded.size() != padded_floats) { + lctx.dflash_target_features_padded.resize(padded_floats); + } + if (left_pad == 0 && total_floats == padded_floats) { + std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin()); + } else { + if (dst_offset > 0) { + std::fill(lctx.dflash_target_features_padded.begin(), + lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset, 0.0f); + } + std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset); + } + profile.graph_feature_copy_us += (uint64_t) (ggml_time_us() - t_feature_us); + profile.graph_feature_bytes += padded_floats * sizeof(float); + + const int64_t t_pos_us = ggml_time_us(); lctx.dflash_pos_ctx_data.resize((size_t) cross_ctx); std::fill(lctx.dflash_pos_ctx_data.begin(), lctx.dflash_pos_ctx_data.end(), 0); if (src_pos != nullptr && total_positions == (size_t) n_rows) { + bool monotonic = true; + for (int32_t i = 1; i < n_rows; ++i) { + if (src_pos[i] <= src_pos[i - 1]) { + monotonic = false; + break; + } + } + if (!monotonic) { + profile.graph_pos_non_monotonic++; + if (profile.graph_pos_non_monotonic <= 3) { + LLAMA_LOG_WARN("%s: DFlash target positions are not strictly increasing (rows=%d first=%d last=%d)\n", + __func__, n_rows, (int) src_pos[0], (int) src_pos[n_rows - 1]); + } + } + profile.last_pos_first = src_pos[0]; + profile.last_pos_last = src_pos[n_rows - 1]; std::copy(src_pos, src_pos + n_rows, lctx.dflash_pos_ctx_data.begin() + (ptrdiff_t) left_pad); } else { + profile.graph_pos_fallbacks++; + profile.last_pos_first = -1; + profile.last_pos_last = -1; + if (profile.graph_pos_fallbacks <= 3) { + LLAMA_LOG_WARN("%s: using synthetic DFlash positions (rows=%d positions=%zu cross_ctx=%d)\n", + __func__, n_rows, total_positions, cross_ctx); + } for (int32_t i = 0; i < n_rows; ++i) { lctx.dflash_pos_ctx_data[(size_t) left_pad + (size_t) i] = i; } } - ggml_backend_tensor_set(pos_ctx, lctx.dflash_pos_ctx_data.data(), 0, ggml_nbytes(pos_ctx)); + profile.graph_pos_copy_us += (uint64_t) (ggml_time_us() - t_pos_us); + profile.graph_pos_bytes += lctx.dflash_pos_ctx_data.size() * sizeof(llama_pos); + if (use_kv_cache) { + const size_t max_nodes = lctx.model.max_nodes((int) std::max(1, cross_ctx)) + 24 * lctx.model.hparams.n_layer; + const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false); + if (lctx.dflash_buf_compute_meta.size() != meta_size) { + lctx.dflash_buf_compute_meta.resize(meta_size); + } + + ggml_cgraph * gf_kv = llm_build_context::llama_build_graph_dflash_kv_cache(lctx); + 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; + } + + if (lctx.dflash_sched == nullptr) { + 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)); + } + } + + lctx.dflash_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false); + if (lctx.dflash_sched == nullptr || !ggml_backend_sched_reserve(lctx.dflash_sched, gf_kv)) { + profile.graph_shape_failures++; + LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V scheduler\n", __func__); + return false; + } + } + + ggml_backend_sched_reset(lctx.dflash_sched); + ggml_backend_sched_alloc_graph(lctx.dflash_sched, gf_kv); + ggml_backend_tensor_set(lctx.dflash_kv_input_target_features, lctx.dflash_target_features_padded.data(), 0, ggml_nbytes(lctx.dflash_kv_input_target_features)); + ggml_backend_tensor_set(lctx.dflash_kv_input_pos_ctx, lctx.dflash_pos_ctx_data.data(), 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx)); + const int64_t t_kv_cache_us = ggml_time_us(); + llama_graph_compute_sched(lctx, lctx.dflash_sched, gf_kv, lctx.cparams.n_threads); + llama_synchronize(&lctx); + profile.graph_kv_cache_compute_us += (uint64_t) (ggml_time_us() - t_kv_cache_us); + profile.graph_kv_cache_calls++; + } else { + ggml_backend_tensor_set(lctx.inp_dflash_target_features, lctx.dflash_target_features_padded.data(), 0, ggml_nbytes(lctx.inp_dflash_target_features)); + ggml_backend_tensor_set(lctx.inp_dflash_pos_ctx, lctx.dflash_pos_ctx_data.data(), 0, ggml_nbytes(lctx.inp_dflash_pos_ctx)); + } + + const int64_t t_mask_us = ggml_time_us(); lctx.dflash_kq_mask_data.assign((size_t) n_kv_total * (size_t) n_mask_tokens, -INFINITY); + int32_t visible_kv_max = 0; for (uint32_t j = 0; j < n_tokens; ++j) { float * row = lctx.dflash_kq_mask_data.data() + (size_t) j * (size_t) n_kv_total; const int32_t visible_kv = cross_ctx + (int32_t) j + 1; + visible_kv_max = std::max(visible_kv_max, visible_kv); + profile.graph_visible_kv_sum += (uint64_t) visible_kv; for (int32_t i = left_pad; i < visible_kv; ++i) { row[i] = 0.0f; } } ggml_backend_tensor_set(kq_mask, lctx.dflash_kq_mask_data.data(), 0, ggml_nbytes(kq_mask)); + profile.graph_mask_build_us += (uint64_t) (ggml_time_us() - t_mask_us); + profile.graph_mask_bytes += ggml_nbytes(kq_mask); + profile.graph_visible_kv_max = std::max(profile.graph_visible_kv_max, (uint64_t) visible_kv_max); + profile.graph_prepare_total_us += (uint64_t) (ggml_time_us() - t_total_us); + + if (profile.graph_prepare_calls == 1) { + LLAMA_LOG_INFO("%s: DFlash graph contract rows=%d width=%d cross_ctx=%d n_tokens=%u left_pad=%d n_kv_total=%d draft_n_ctx=%u pos=%s [%d..%d]\n", + __func__, n_rows, width, cross_ctx, n_tokens, left_pad, n_kv_total, lctx.cparams.n_ctx, + (src_pos != nullptr && total_positions == (size_t) n_rows) ? "target" : "synthetic", + (int) profile.last_pos_first, (int) profile.last_pos_last); + } return true; }