From dc43cdf06b0f1761090a579027f4bab352c01046 Mon Sep 17 00:00:00 2001 From: SamuelOliveirads Date: Tue, 2 Jun 2026 10:22:13 -0300 Subject: [PATCH] move dflash for it own file --- common/speculative-impl.h | 1740 +++++++++++++++++++++++++++ common/speculative.cpp | 1741 +--------------------------- src/CMakeLists.txt | 2 + src/llama-build-context.cpp | 10 +- src/llama-context.h | 227 ++-- src/llama-dflash.cpp | 1240 ++++++++++++++++++++ src/llama-dflash.h | 8 + src/llama-quantize.cpp | 5 +- src/llama-spec-features-dflash.cpp | 1097 ++++++++++++++++++ src/llama-spec-features-dflash.h | 279 +++++ src/llama-spec-features.cpp | 1084 ----------------- src/llama-spec-features.h | 304 +---- src/llama.cpp | 848 +------------- 13 files changed, 4533 insertions(+), 4052 deletions(-) create mode 100644 common/speculative-impl.h create mode 100644 src/llama-dflash.cpp create mode 100644 src/llama-dflash.h create mode 100644 src/llama-spec-features-dflash.cpp create mode 100644 src/llama-spec-features-dflash.h diff --git a/common/speculative-impl.h b/common/speculative-impl.h new file mode 100644 index 00000000..47603461 --- /dev/null +++ b/common/speculative-impl.h @@ -0,0 +1,1740 @@ +// DFlash runtime state and draft path. +struct common_speculative_state_dflash : public common_speculative_state { + llama_context * ctx_tgt; + llama_context * ctx_dft; + + llama_batch batch = {}; + + int32_t block_size = 0; + int32_t mask_token_id = -1; + int32_t n_target_features = 0; + int32_t cross_ctx = 0; + bool ready = false; + + std::vector target_layer_ids; + std::vector target_window; + std::vector target_window_pos; + std::vector target_window_stage; + std::vector target_window_pos_stage; + std::vector target_window_ring; + std::vector target_window_append_features; + int32_t target_window_rows = 0; + int32_t target_window_ring_write_pos = 0; + int32_t target_window_ring_filled = 0; + uint64_t target_window_version = 0; + int32_t target_window_keep_rows = 0; + int32_t target_window_append_rows = 0; + bool target_window_replace = false; + bool target_window_materialized = false; + 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; + + uint64_t t_draft_decode_us = 0; + uint64_t t_draft_sample_us = 0; + uint64_t t_warmup_collect_us = 0; + uint64_t t_warmup_append_us = 0; + uint64_t t_accept_output_copy_us = 0; + uint64_t t_accept_commit_us = 0; + uint64_t t_accept_append_us = 0; + uint64_t t_accept_append_filter_us = 0; + uint64_t t_accept_append_window_alloc_us = 0; + uint64_t t_accept_append_replace_us = 0; + uint64_t t_accept_append_keep_old_us = 0; + uint64_t t_accept_append_new_rows_us = 0; + uint64_t t_accept_append_commit_detail_us = 0; + uint64_t t_accept_append_log_us = 0; + size_t n_warmup_collect_calls = 0; + size_t n_warmup_collect_rows = 0; + size_t n_warmup_append_calls = 0; + size_t n_warmup_append_rows = 0; + size_t n_accept_output_copy_calls = 0; + size_t n_accept_output_copy_rows = 0; + size_t n_accept_commit_calls = 0; + size_t n_accept_commit_rows = 0; + size_t n_accept_append_calls = 0; + size_t n_accept_append_rows = 0; + size_t n_accept_append_replace_calls = 0; + size_t n_accept_append_slide_calls = 0; + + common_speculative_state_dflash( + enum common_speculative_type type, + llama_context * ctx_tgt, + llama_context * ctx_dft, + int32_t cross_ctx) + : common_speculative_state(type) + , ctx_tgt(ctx_tgt) + , ctx_dft(ctx_dft) + , cross_ctx(std::max(1, cross_ctx)) + { + const llama_model * model_tgt = llama_get_model(ctx_tgt); + const llama_model * model_dft = llama_get_model(ctx_dft); + + if (!common_speculative_are_dflash_compatible(model_tgt, model_dft)) { + LOG_ERR("%s: DFlash draft model vocab/tokenizer is incompatible with the target model\n", __func__); + return; + } + + block_size = llama_model_dflash_block_size(model_dft); + mask_token_id = llama_model_dflash_mask_token_id(model_dft); + n_target_features = llama_model_dflash_n_target_features(model_dft); + const int32_t n_target_layers = llama_model_dflash_n_target_layers(model_dft); + + if (block_size <= 0 || mask_token_id < 0 || n_target_features <= 0 || n_target_layers <= 0) { + LOG_ERR("%s: invalid DFlash metadata (block_size=%d, mask_token_id=%d, n_target_features=%d, n_target_layers=%d)\n", + __func__, block_size, mask_token_id, n_target_features, n_target_layers); + return; + } + + target_layer_ids.resize((size_t) n_target_layers); + if (llama_model_dflash_target_layer_ids(model_dft, target_layer_ids.data(), n_target_layers) != n_target_layers) { + LOG_ERR("%s: failed to read DFlash target layer ids\n", __func__); + target_layer_ids.clear(); + return; + } + + const auto * vocab_tgt = llama_model_get_vocab(model_tgt); + const auto * vocab_dft = llama_model_get_vocab(model_dft); + const int32_t target_vocab_size = llama_vocab_n_tokens(vocab_tgt); + const int32_t draft_vocab_size = llama_vocab_n_tokens(vocab_dft); + const int32_t target_hidden_size = llama_model_n_embd(model_tgt); + const int32_t draft_hidden_size = llama_model_n_embd(model_dft); + const int32_t target_mask_token_id = llama_model_dflash_target_mask_token_id(model_tgt); + const int32_t expected_n_target_features = target_hidden_size > 0 ? target_hidden_size * n_target_layers : 0; + + if (target_mask_token_id != (int32_t) LLAMA_TOKEN_NULL && mask_token_id != target_mask_token_id) { + LOG_ERR("%s: DFlash mask token mismatch (draft=%d target=%d)\n", + __func__, mask_token_id, target_mask_token_id); + return; + } + + if (target_hidden_size <= 0 || draft_hidden_size <= 0) { + LOG_ERR("%s: invalid DFlash hidden sizes (draft=%d target=%d)\n", + __func__, draft_hidden_size, target_hidden_size); + return; + } + + if (expected_n_target_features <= 0 || n_target_features != expected_n_target_features) { + LOG_ERR("%s: DFlash target feature width mismatch (metadata=%d expected=%d target_hidden=%d target_layers=%d)\n", + __func__, n_target_features, expected_n_target_features, target_hidden_size, n_target_layers); + return; + } + + std::vector sorted_target_layer_ids = target_layer_ids; + std::sort(sorted_target_layer_ids.begin(), sorted_target_layer_ids.end()); + if (std::adjacent_find(sorted_target_layer_ids.begin(), sorted_target_layer_ids.end()) != sorted_target_layer_ids.end()) { + LOG_ERR("%s: duplicate DFlash target layer ids survived into runtime validation\n", __func__); + target_layer_ids.clear(); + return; + } + + const int32_t n_target_model_layers = llama_n_layer(model_tgt); + for (int32_t layer_id : target_layer_ids) { + if (layer_id < 0 || layer_id >= n_target_model_layers) { + LOG_ERR("%s: invalid DFlash target layer id %d for target model with %d layers\n", + __func__, layer_id, n_target_model_layers); + target_layer_ids.clear(); + return; + } + } + + const int32_t io_mode = llama_model_dflash_io_mode(model_dft, model_tgt); + if (io_mode == LLAMA_DFLASH_IO_MODE_INVALID) { + LOG_ERR("%s: DFlash draft is missing required IO tensors after target sharing\n", __func__); + return; + } + + if (io_mode == LLAMA_DFLASH_IO_MODE_MIXED) { + LOG_ERR("%s: DFlash IO contract must be fully shared or fully self-contained, but resolved to mixed mode\n", __func__); + return; + } + + if (io_mode == LLAMA_DFLASH_IO_MODE_SELF_CONTAINED && !llama_model_dflash_io_tensors_match(model_dft, target_hidden_size, target_vocab_size)) { + LOG_ERR("%s: DFlash self-contained IO tensors do not match the target hidden/vocab contract (target_hidden=%d target_vocab=%d)\n", + __func__, + target_hidden_size, + target_vocab_size); + return; + } + + if (!llama_set_dflash_capture_layers(ctx_tgt, target_layer_ids.data(), (int32_t) target_layer_ids.size())) { + LOG_ERR("%s: failed to configure DFlash target capture callback\n", __func__); + return; + } + + batch = llama_batch_init(std::max(1, block_size), 0, 1); + target_window.reserve((size_t) this->cross_ctx * (size_t) n_target_features); + target_window_stage.reserve((size_t) this->cross_ctx * (size_t) n_target_features); + target_window_ring.resize((size_t) this->cross_ctx * (size_t) n_target_features); + target_window_append_features.reserve((size_t) this->cross_ctx * (size_t) n_target_features); + target_window_pos.reserve((size_t) this->cross_ctx); + target_window_pos_stage.reserve((size_t) this->cross_ctx); + ready = true; + + 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]; + } + + const char * io_mode_name = io_mode == LLAMA_DFLASH_IO_MODE_SHARED ? "shared" : "self-contained"; + 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()); + LOG_INF("%s: DFlash artifact io=%s draft_vocab=%d target_vocab=%d draft_hidden=%d target_hidden=%d mask_token_id=%d target_mask_token_id=%d\n", + __func__, io_mode_name, draft_vocab_size, target_vocab_size, draft_hidden_size, target_hidden_size, mask_token_id, target_mask_token_id); + } + + ~common_speculative_state_dflash() override { + llama_clear_dflash_capture(ctx_tgt); + if (ctx_dft) { + llama_free(ctx_dft); + } + if (batch.token != nullptr) { + llama_batch_free(batch); + } + } + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + llama_kv_cache_clear(ctx_dft); + llama_reset_dflash_kv_cache_state(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; + t_draft_decode_us = 0; + t_draft_sample_us = 0; + t_warmup_collect_us = 0; + t_warmup_append_us = 0; + t_accept_output_copy_us = 0; + t_accept_commit_us = 0; + t_accept_append_us = 0; + t_accept_append_filter_us = 0; + t_accept_append_window_alloc_us = 0; + t_accept_append_replace_us = 0; + t_accept_append_keep_old_us = 0; + t_accept_append_new_rows_us = 0; + t_accept_append_commit_detail_us = 0; + t_accept_append_log_us = 0; + n_warmup_collect_calls = 0; + n_warmup_collect_rows = 0; + n_warmup_append_calls = 0; + n_warmup_append_rows = 0; + n_accept_output_copy_calls = 0; + n_accept_output_copy_rows = 0; + n_accept_commit_calls = 0; + n_accept_commit_rows = 0; + n_accept_append_calls = 0; + n_accept_append_rows = 0; + n_accept_append_replace_calls = 0; + n_accept_append_slide_calls = 0; + llama_dflash_profile_reset(ctx_tgt); + llama_dflash_profile_reset(ctx_dft); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + GGML_UNUSED(prompt_tgt); + + result.clear(); + if (!ready || target_window_rows <= 0) { + n_draft_empty++; + return; + } + + const int32_t n_keep = std::min(params.n_max, block_size - 1); + if (n_keep <= 0) { + return; + } + + const bool use_kv_cache = dflash_use_kv_cache_experiment(); + const float * target_features = nullptr; + size_t target_feature_floats = 0; + llama_dflash_window_update window_update = { + target_window_version, + target_window_keep_rows, + target_window_append_rows, + target_window_replace, + target_window_append_features.empty() ? nullptr : target_window_append_features.data(), + target_window_append_features.size(), + }; + const llama_dflash_kv_cache_transition cache_plan = use_kv_cache + ? llama_plan_dflash_kv_cache_transition_for_ctx(ctx_dft, window_update, target_window_rows) + : llama_dflash_kv_cache_transition{}; + + if (!use_kv_cache || cache_plan.rebuild_cache) { + dflash_materialize_target_window_features(*this); + target_features = target_window.data(); + target_feature_floats = target_window.size(); + } + if (use_kv_cache && cache_plan.rebuild_cache) { + window_update.append_features = target_window.data(); + window_update.append_floats = target_window.size(); + window_update.append_rows = target_window_rows; + } + + if (!llama_set_dflash_target_features_view(ctx_dft, target_features, target_feature_floats, target_window_rows, target_window_pos.data(), &window_update)) { + 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 int32_t batch_len = n_keep + 1; + const llama_pos draft_pos_base = last_target_pos >= 0 ? last_target_pos + 1 : (llama_pos) target_window_rows; + const llama_pos seed_pos = last_target_pos >= 0 ? last_target_pos : draft_pos_base - 1; + last_draft_pos_base = draft_pos_base; + common_batch_add(batch, id_last, seed_pos, { 0 }, false); + for (int32_t i = 1; i < batch_len; ++i) { + common_batch_add(batch, mask_token_id, draft_pos_base + (i - 1), { 0 }, i <= n_keep); + } + + const int64_t t_decode_us = ggml_time_us(); + 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; + } + t_draft_decode_us += (uint64_t) (ggml_time_us() - t_decode_us); + + result.reserve((size_t) n_keep); + const int64_t t_sample_us = ggml_time_us(); + for (int32_t i = 0; i < n_keep; ++i) { + result.push_back(common_sampler_sample_speculative(nullptr, ctx_dft, i + 1, nullptr)); + } + t_draft_sample_us += (uint64_t) (ggml_time_us() - t_sample_us); + + batch.n_tokens = 0; + dflash_contract_log_draft(*this, n_keep, result.size()); + } + + void accept(uint16_t n_accepted) override { + GGML_UNUSED(n_accepted); + } +}; + +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; + const llama_pos seed_pos = state.last_target_pos; + const llama_pos mask_first_pos = state.last_draft_pos_base; + const llama_pos mask_last_pos = state.last_draft_pos_base >= 0 + ? state.last_draft_pos_base + n_keep - 1 + : -1; + + LOG_INF("dflash contract draft[%llu]: window_rows=%d window_pos=%s pos=[%d..%d] gaps=%d nonmono=%d last_target_pos=%d seed_pos=%d mask_pos=[%d..%d] sample_rows=[1..%d] output_rows=[1..%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) seed_pos, + (int) mask_first_pos, + (int) mask_last_pos, + n_keep, + n_keep, + (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; + + common_sampler * smpl; + + llama_batch batch; + llama_tokens prompt_dft; + + bool vocab_cmpt = true; // whether retokenization is needed + std::unordered_map vocab_map; + + common_speculative_state_draft( + enum common_speculative_type type, + llama_context * ctx_tgt, + llama_context * ctx_dft, + const std::vector> & replacements) + : common_speculative_state(type) + , ctx_tgt(ctx_tgt) + , ctx_dft(ctx_dft) + { + batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1); + smpl = nullptr; + { + struct common_params_sampling params; + params.top_k = 10; + params.samplers_sequence = { + llama_sampler_type::TOP_K, + llama_sampler_type::DIST, // needed to get probabilities + }; + smpl = common_sampler_init(llama_get_model(ctx_dft), params); + } + + vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft)); + LOG_DBG("vocab_cmpt = %d\n", vocab_cmpt); + + if (!vocab_cmpt) { + LOG_WRN("the target and draft vocabs are not compatible - tokens will be translated between the two\n"); + + for (const auto & pair : replacements) { + vocab_map[pair.first] = pair.second; + } + } + } + + ~common_speculative_state_draft() override { + llama_free(ctx_dft); + + common_sampler_free(smpl); + + llama_batch_free(batch); + } + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + auto * spec = this; + + auto & batch = spec->batch; + auto & ctx_tgt = spec->ctx_tgt; + auto & ctx_dft = spec->ctx_dft; + auto & smpl = spec->smpl; + auto & prompt_dft = spec->prompt_dft; + + int reuse_i = 0; + int reuse_n = 0; + + const int n_ctx = llama_n_ctx(ctx_dft) - params.n_max; + + llama_tokens prompt_cnv; + if (!spec->vocab_cmpt) { + // convert id_last to draft vocab. llama_detokenize is called directly to avoid an allocation + const auto * model_tgt = llama_get_model(ctx_tgt); + const auto * vocab_tgt = llama_model_get_vocab(model_tgt); + + std::string text; + + text = common_detokenize(ctx_tgt, prompt_tgt, true); + text = replace_to_dft(text); + + LOG_DBG("%s: main->draft detokenized string: '%s'\n", __func__, text.c_str()); + + prompt_cnv = common_tokenize(ctx_dft, text, false, true); + + + + int32_t n_chars = llama_detokenize(vocab_tgt, &id_last, 1, nullptr, 0, false, false); + GGML_ASSERT(n_chars < 0 && "failed to detokenize id_last"); + + text.resize(-n_chars); + llama_detokenize(vocab_tgt, &id_last, 1, text.data(), text.size(), false, false); + text = replace_to_dft(text); + + LOG_DBG("main->draft detokenized id_last(%d): '%s'\n", id_last, text.c_str()); + id_last = common_tokenize(ctx_dft, text, false, true)[0]; + } + + const llama_tokens & prompt_cur = spec->vocab_cmpt ? prompt_tgt : prompt_cnv; + + const int i_start = std::max(0, (int) prompt_cur.size() - n_ctx); + + // reuse as much as possible from the old draft context + // ideally, the draft context should be as big as the target context and we will always reuse the entire prompt + for (int i = 0; i < (int) prompt_dft.size(); ++i) { + int cur = 0; + while (i_start + cur < (int) prompt_cur.size() && + i + cur < (int) prompt_dft.size() && + prompt_cur[i_start + cur] == prompt_dft[i + cur]) { + cur++; + } + + if ((cur >= 256 || n_ctx >= (int) prompt_cur.size()) && cur > reuse_n) { + reuse_i = i; + reuse_n = cur; + } + } + + LOG_DBG("%s: reuse_i = %d, reuse_n = %d, prompt = %d\n", __func__, reuse_i, reuse_n, (int) prompt_dft.size()); + + result.clear(); + result.reserve(params.n_max); + + if (reuse_n == 0) { + llama_kv_cache_clear(ctx_dft); + prompt_dft.clear(); + } else { + // this happens when a previous draft has been discarded (for example, due to being too small), but the + // target model agreed with it. in this case, we simply pass back the previous results to save compute + if (reuse_i + reuse_n < (int) prompt_dft.size() && prompt_dft[reuse_i + reuse_n] == id_last) { + for (int i = reuse_i + reuse_n + 1; i < (int) prompt_dft.size(); ++i) { + result.push_back(prompt_dft[i]); + + if (params.n_max <= (int) result.size()) { + break; + } + } + + return; + } + + if (reuse_i > 0) { + llama_kv_cache_seq_rm (ctx_dft, 0, 0, reuse_i); + llama_kv_cache_seq_add(ctx_dft, 0, reuse_i, -1, -reuse_i); + + prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); + } + + if (reuse_n < (int) prompt_dft.size()) { + llama_kv_cache_seq_rm (ctx_dft, 0, reuse_n, -1); + prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); + } + } + + // prepare a batch to evaluate any new tokens in the prompt + common_batch_clear(batch); + + for (size_t i = i_start + reuse_n; i < prompt_cur.size(); ++i) { + //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_cur[i]); + common_batch_add(batch, prompt_cur[i], i - i_start, { 0 }, false); + + prompt_dft.push_back(prompt_cur[i]); + } + + // we should rarely end-up here during normal decoding + if (batch.n_tokens > 0) { + //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str()); + + llama_decode(ctx_dft, batch); + } + + const llama_pos n_past = prompt_dft.size(); + + LOG_DBG("%s: n_past = %d\n", __func__, n_past); + + common_batch_clear(batch); + common_batch_add (batch, id_last, n_past, { 0 }, true); + + prompt_dft.push_back(id_last); + + //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx_dft, prompt_dft).c_str()); + + llama_decode(ctx_dft, batch); + + common_sampler_reset(smpl); + + // sample n_draft tokens from the draft model + for (int i = 0; i < params.n_max; ++i) { + common_batch_clear(batch); + + common_sampler_sample(smpl, ctx_dft, 0, true); + + const auto * cur_p = common_sampler_get_candidates(smpl, true); + + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + } + + // add drafted token for each sequence + const llama_token id = cur_p->data[0].id; + + common_sampler_accept(smpl, nullptr, id, true); + + // only collect very high-confidence draft tokens + if (cur_p->data[0].p < params.p_min) { + if (i == 0) { + result.push_back(id); + } + break; + } + + result.push_back(id); + + if (params.n_max <= (int) result.size()) { + break; + } + + + common_batch_add(batch, id, n_past + i + 1, { 0 }, true); + + // evaluate the drafted tokens on the draft model + llama_decode(ctx_dft, batch); + + prompt_dft.push_back(id); + } + + if (!spec->vocab_cmpt) { + std::string detokenized = common_detokenize(ctx_dft, result, true); + detokenized = replace_to_tgt(detokenized); + LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); + result = common_tokenize(ctx_tgt, detokenized, false, true); + if (result.size() > (size_t)params.n_max) { + result.resize(params.n_max); + } + } + } + + void accept(uint16_t n_accepted) override { + // noop + GGML_UNUSED(n_accepted); + } + + std::string replace_to_dft(const std::string & input) const { + std::string result = input; + + for (const auto & pair : this->vocab_map) { + size_t pos = result.find(pair.first); + while (pos != std::string::npos) { + result.replace(pos, pair.first.length(), pair.second); + pos = result.find(pair.first, pos + pair.second.length()); + } + } + + return result; + } + + std::string replace_to_tgt(const std::string & input) const { + std::string result = input; + + for (const auto & pair : this->vocab_map) { + size_t pos = result.find(pair.second); + while (pos != std::string::npos) { + result.replace(pos, pair.second.length(), pair.first); + pos = result.find(pair.second, pos + pair.first.length()); + } + } + + return result; + } +}; + +struct common_speculative_state_eagle3 : public common_speculative_state { + common_speculative_state_eagle3(enum common_speculative_type type) : common_speculative_state(type) {} + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & draft_tokens) override { + // TODO: implement + GGML_UNUSED(params); + GGML_UNUSED(prompt_tgt); + GGML_UNUSED(id_last); + GGML_UNUSED(draft_tokens); + } + + void accept(uint16_t n_accepted) override { + // noop + GGML_UNUSED(n_accepted); + } +}; + +// state of self-speculation (simple implementation, not ngram-map) +struct common_speculative_state_ngram_simple : public common_speculative_state { + common_ngram_simple_config config; + + common_speculative_state_ngram_simple( + enum common_speculative_type type, + common_ngram_simple_config config) + : common_speculative_state(type), config(config) {} + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + + result = common_ngram_simple_draft(config, prompt_tgt, id_last); + GGML_UNUSED(params); + } + + void accept(uint16_t n_accepted) override { + // noop + GGML_UNUSED(n_accepted); + } +}; + +struct common_speculative_state_ngram_map_k : public common_speculative_state { + // draft ngram map for speculative decoding without draft model + common_ngram_map map; + + common_speculative_state_ngram_map_k( + enum common_speculative_type type, + common_ngram_map map) + : common_speculative_state(type), map(std::move(map)) {} + + void begin(const llama_tokens & prompt) override { + common_ngram_map_begin(map, prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + common_ngram_map_draft(map, prompt_tgt, id_last, result); + GGML_UNUSED(params); + } + + void accept(uint16_t n_accepted) override { + common_ngram_map_accept(map, n_accepted); + } +}; + +struct common_speculative_state_ngram_mod : public common_speculative_state { + common_ngram_mod & mod; + + // the last position in the prompt that was added to the ngram container + size_t i_last = 0; + + // length of the last drafted n‑gram (number of tokens returned by draft) + size_t n_draft_last = 0; + + // consecutive accept rounds with low acceptance fraction (< 0.5) + int n_low = 0; + + // enable trace logging if LLAMA_TRACE is set + const bool verbose; + + common_speculative_state_ngram_mod(enum common_speculative_type type, common_ngram_mod & mod) + : common_speculative_state(type), mod(mod), verbose(std::getenv("LLAMA_TRACE") != nullptr) { + static_assert(sizeof(llama_token) == sizeof(common_ngram_mod::entry_t)); + } + + void begin(const llama_tokens & prompt) override { + i_last = 0; + + n_draft_last = 0; + n_low = 0; + + const size_t n = mod.get_n(); + + if (prompt.size() < n) { + return; + } + + for (size_t i = 0; i < prompt.size() - n; ++i) { + mod.add(prompt.data() + i); + } + + i_last = prompt.size() - n; + + const double f = (double)mod.get_used() / (double)mod.size(); + LOG_INF("%s: ngram_mod occupancy = %zu/%zu (%.2f)\n", __func__, mod.get_used(), mod.size(), f); + + constexpr double f_thold = 0.25; + if (f > f_thold) { + LOG_WRN("%s: ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", __func__, f, f_thold); + + mod.reset(); + } + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + GGML_UNUSED(params); + + n_draft_last = 0; + + const size_t cur_len = prompt_tgt.size(); + if (cur_len < mod.get_n()) { + return; + } + + const size_t n = mod.get_n(); + + // add new ngrams in chunks + if (i_last + 32 < cur_len) { + for (size_t i = i_last; i < cur_len - n; ++i) { + mod.add(prompt_tgt.data() + i); + } + + i_last = cur_len - n; + } + + result.resize(n + params.n_max); + for (size_t i = 0; i < n - 1; ++i) { + result[i] = prompt_tgt[cur_len - n + 1 + i]; + } + result[n - 1] = id_last; + + for (int i = 0; i < params.n_max; ++i) { + const llama_token token = mod.get(result.data() + i); + if (token == common_ngram_mod::EMPTY) { + if (i < params.n_min) { + result.clear(); + return; + } + + result.resize(n + i); + break; + } + result[n + i] = token; + } + + // only return the m tokens that were drafted + for (size_t i = 0; n + i < result.size(); ++i) { + result[i] = result[n + i]; + } + result.resize(result.size() - n); + + // store length of drafted n‑gram for later acceptance analysis + n_draft_last = result.size(); + } + + void accept(uint16_t n_accepted) override { + if (verbose) { + LOG_INF("%s: accepted %d tokens from %zu drafted tokens\n", __func__, n_accepted, n_draft_last); + } + + // compute acceptance fraction if we have a recorded draft length + if (n_draft_last > 0) { + const double f_acc = (double)n_accepted / (double)n_draft_last; + if (f_acc < 0.5) { + n_low++; + if (n_low >= 3) { + LOG_WRN("%s: low acceptance streak (%d) – resetting ngram_mod\n", __func__, n_low); + + mod.reset(); + n_low = 0; + i_last = 0; + } + } else { + n_low = 0; + } + } + } +}; + +struct common_speculative_state_ngram_cache : public common_speculative_state { + uint16_t n_draft; + bool save_dynamic; + bool save_static; + + common_ngram_cache ngram_cache_context; + common_ngram_cache ngram_cache_dynamic; + common_ngram_cache ngram_cache_static; + + size_t cache_size = 0; // number of tokens in n-gram cache + + common_speculative_state_ngram_cache( + const enum common_speculative_type type, + const std::string & path_static, + const std::string & path_dynamic, + uint16_t n_draft, + bool save_dynamic, + bool save_static) + : common_speculative_state(type) + , n_draft(n_draft) + , save_dynamic(save_dynamic) + , save_static(save_static) + { + if (!path_static.empty()) { + try { + ngram_cache_static = common_ngram_cache_load(path_static); + } catch (...) { + LOG_ERR("failed to open static lookup cache: %s", path_static.c_str()); + GGML_ABORT("Couldn't read static lookup cache"); + } + } + + if (!path_dynamic.empty()) { + try { + ngram_cache_dynamic = common_ngram_cache_load(path_dynamic); + } catch (...) { + LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str()); + GGML_ABORT("Couldn't read dynamic lookup cache"); + } + } + } + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + GGML_UNUSED(params); + + if (cache_size < prompt_tgt.size() + 1) { + llama_tokens tokens_new; + tokens_new.reserve(prompt_tgt.size() + 1 - cache_size); + for (size_t j = cache_size; j < prompt_tgt.size(); ++j) { + tokens_new.push_back(prompt_tgt[j]); + } + tokens_new.push_back(id_last); // add the last token + + // Update context ngram cache with new prompt_tgt: + common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, + tokens_new, tokens_new.size(), false); + cache_size = prompt_tgt.size() + 1; + } + + llama_tokens inp; + inp.reserve(prompt_tgt.size() + 1); + for (size_t j = 0; j < prompt_tgt.size(); ++j) { + inp.push_back(prompt_tgt[j]); + } + inp.push_back(id_last); + + result.push_back(id_last); + + common_ngram_cache_draft(inp, result, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, + ngram_cache_context, + ngram_cache_dynamic, + ngram_cache_static); + + if (result.size() > 0) { + // delete first token in result (which is the id_last token) + result.erase(result.begin()); + } + } + + void accept(uint16_t n_accepted) override { + // TODO: noop + GGML_UNUSED(n_accepted); + } +}; + +struct common_speculative_state_suffix : public common_speculative_state { + common_suffix_tree tree; + common_suffix_tree corpus_tree; + bool has_corpus = false; + size_t cache_size = 0; + + // Acceptance feedback + size_t n_draft_last = 0; + bool had_accept = false; + int n_low = 0; + float base_p_min = 0.1f; + float eff_p_min = 0.1f; + + common_speculative_state_suffix( + enum common_speculative_type type, + int max_depth, + const std::string & corpus_path, + const llama_model * model) + : common_speculative_state(type) + , tree(max_depth) + , corpus_tree(max_depth) + { + if (!corpus_path.empty()) { + std::function(const std::string &)> tokenize_fn; + if (model) { + tokenize_fn = [model](const std::string & text) -> std::vector { + return common_tokenize(model, text, false, true); + }; + } + has_corpus = corpus_tree.load_corpus(corpus_path, tokenize_fn); + } + } + + void begin(const llama_tokens & prompt) override { + cache_size = 0; + n_draft_last = 0; + had_accept = false; + n_low = 0; + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + + base_p_min = params.p_min; + if (n_draft_last > 0 && !had_accept) { + if (++n_low >= 3) { + eff_p_min = std::min(eff_p_min + 0.1f, 0.5f); + n_low = 0; + } + } + had_accept = false; + + if (cache_size < prompt_tgt.size() + 1) { + llama_tokens tokens_new; + tokens_new.reserve(prompt_tgt.size() + 1 - cache_size); + for (size_t j = cache_size; j < prompt_tgt.size(); ++j) { + tokens_new.push_back(prompt_tgt[j]); + } + tokens_new.push_back(id_last); + + tree.extend(tokens_new.data(), (int)tokens_new.size()); + cache_size = prompt_tgt.size() + 1; + } + + const int ctx_len = std::min((int)(prompt_tgt.size() + 1), tree.max_depth()); + llama_tokens context; + context.reserve(ctx_len); + const int ctx_start = (int)prompt_tgt.size() + 1 - ctx_len; + for (int j = ctx_start; j < (int)prompt_tgt.size(); ++j) { + context.push_back(prompt_tgt[j]); + } + context.push_back(id_last); + const int min_match_len = std::max(1, params.suffix_min_match_len); + + result = tree.speculate( + context.data(), (int)context.size(), + params.n_max, + eff_p_min, + 1, + min_match_len); + + if (has_corpus) { + auto corpus_result = corpus_tree.speculate( + context.data(), (int)context.size(), + params.n_max, + eff_p_min, + 1, + min_match_len); + if (corpus_result.size() > result.size()) { + result = std::move(corpus_result); + } + } + + n_draft_last = result.size(); + } + + void accept(uint16_t n_accepted) override { + if (n_draft_last == 0) { + return; + } + had_accept = true; + const double f_acc = (double)n_accepted / (double)n_draft_last; + if (f_acc < 0.5) { + if (++n_low >= 3) { + eff_p_min = std::min(eff_p_min + 0.1f, 0.5f); + n_low = 0; + } + } else { + n_low = 0; + if (eff_p_min > base_p_min) { + eff_p_min = std::max(eff_p_min - 0.05f, base_p_min); + } + } + } +}; + +struct common_speculative { + std::vector configs; // resolved stage config for each implementation + std::vector> impls; // list of implementations to use and their states + common_speculative_state * curr_impl = nullptr; // current implementation in use (for stats) + std::unique_ptr tuner; + int last_n_drafted = 0; + int64_t t_step_start_us = 0; +}; + +static bool common_speculative_stage_chain_matches( + const std::vector & stages, + const std::vector & configs) { + if (stages.size() != configs.size()) { + return false; + } + + for (size_t i = 0; i < stages.size(); ++i) { + if (stages[i].type != configs[i].type) { + return false; + } + } + + return true; +} + +static common_params_speculative common_speculative_get_runtime_params( + const common_speculative_config & config, + const common_params_speculative & params, + const common_speculative_stage_params & stage) { + common_params_speculative result = config.params; + + result.type = config.type; + result.n_max = stage.has_n_max_override() ? stage.n_max : params.n_max; + result.n_min = stage.has_n_min_override() ? stage.n_min : params.n_min; + result.p_min = stage.has_p_min_override() ? stage.p_min : params.p_min; + + if (config.type == COMMON_SPECULATIVE_TYPE_SUFFIX) { + result.suffix_min_match_len = stage.has_suffix_min_match_len_override() + ? stage.suffix_min_match_len + : params.suffix_min_match_len; + } + + result.n_max = std::max(result.n_max, 0); + result.n_min = std::max(0, std::min(result.n_min, result.n_max)); + result.stages.clear(); + + return result; +} + +static common_ngram_map get_common_ngram_map(const common_speculative_config & config) { + uint16_t size_key = config.params.ngram_size_n; + uint16_t size_value = config.params.ngram_size_m; + bool key_only = (config.type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K); + uint16_t min_hits = config.params.ngram_min_hits; + + return common_ngram_map(size_key, size_value, key_only, min_hits); +} + +static common_speculative_state_ngram_cache create_state_ngram_cache( + const std::string & path_static, const std::string & path_dynamic, + const common_speculative_config & config) { + uint16_t n_draft = 8; // TODO get from config? + + // TODO bool param in common/common.h to set save_static/save_dynamic? + bool save_static = false; + bool save_dynamic = false; + + common_speculative_state_ngram_cache state(config.type, path_static, path_dynamic, n_draft, save_static, save_dynamic); + + return state; +} + +std::string common_speculative_type_name_str() { + std::string result; + for (size_t i = 0; i < common_speculative_types.size(); i++) { + if (i > 0) { + result += ", "; + } + result += common_speculative_type_to_str(common_speculative_types[i]); + } + return result; +} + +std::string common_speculative_type_to_str(enum common_speculative_type type) { + switch (type) { + case COMMON_SPECULATIVE_TYPE_NONE: return "none"; + case COMMON_SPECULATIVE_TYPE_DRAFT: return "draft"; + case COMMON_SPECULATIVE_TYPE_DFLASH: return "dflash"; + case COMMON_SPECULATIVE_TYPE_MTP: return "mtp"; + case COMMON_SPECULATIVE_TYPE_EAGLE3: return "eagle3"; + case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram_simple"; + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram_map_k"; + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram_map_k4v"; + case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: return "ngram_mod"; + case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: return "ngram_cache"; + case COMMON_SPECULATIVE_TYPE_SUFFIX: return "suffix"; + default: return "unknown"; + } +} + +enum common_speculative_type common_speculative_type_from_name(const std::string & name) { + std::string normalized = name; + std::replace(normalized.begin(), normalized.end(), '-', '_'); + + const auto it = common_speculative_type_from_name_map.find(normalized); + if (it == common_speculative_type_from_name_map.end()) { + return COMMON_SPECULATIVE_TYPE_COUNT; + } + return it->second; +} + +bool common_speculative_is_compat(llama_context * ctx_tgt) { + bool res = true; + + llama_kv_cache_clear(ctx_tgt); + + // eval 2 tokens to check if the context is compatible + std::vector tmp; + tmp.push_back(0); + tmp.push_back(0); + + int ret = llama_decode(ctx_tgt, llama_batch_get_one(tmp.data(), tmp.size(), 0, 0)); + if (ret != 0) { + LOG_ERR("%s: llama_decode() failed: %d\n", __func__, ret); + res = false; + goto done; + } + + // try to remove the last tokens + if (!llama_kv_cache_seq_rm(ctx_tgt, 0, 1, -1)) { + LOG_WRN("%s: the target context does not support partial sequence removal\n", __func__); + res = false; + goto done; + } + +done: + llama_kv_cache_clear(ctx_tgt); + llama_synchronize(ctx_tgt); + + return res; +} + +// initialization of the speculative decoding system +// +common_speculative * common_speculative_init( + common_params_speculative & params, + llama_context * ctx_tgt) { + std::string chain_error; + if (!common_speculative_validate_chain(params, &chain_error)) { + LOG_ERR("%s: invalid speculative stage chain: %s\n", __func__, chain_error.c_str()); + return nullptr; + } + + const auto stages = params.get_resolved_stages(); + if (params.model_dft && llama_model_is_gemma4_mtp_assistant(params.model_dft)) { + const bool has_draft_stage = std::any_of(stages.begin(), stages.end(), [](const common_speculative_stage_params & stage) { + return stage.type == COMMON_SPECULATIVE_TYPE_DRAFT; + }); + + if (has_draft_stage) { + LOG_ERR("%s: Gemma4 assistant models only support MTP stages; omit -md for self-spec-only runs or use -mtp/--spec-stage mtp for assistant-backed MTP\n", __func__); + return nullptr; + } + } + + const bool has_dflash_stage = std::any_of(stages.begin(), stages.end(), [](const common_speculative_stage_params & stage) { + return stage.type == COMMON_SPECULATIVE_TYPE_DFLASH; + }); + + const bool needs_draft_ctx = std::any_of(stages.begin(), stages.end(), [¶ms](const common_speculative_stage_params & stage) { + return stage.type == COMMON_SPECULATIVE_TYPE_DRAFT || + stage.type == COMMON_SPECULATIVE_TYPE_DFLASH || + (stage.type == COMMON_SPECULATIVE_TYPE_MTP && params.model_dft != nullptr); + }); + + llama_context * ctx_dft = nullptr; + if (needs_draft_ctx) { + if (!params.model_dft) { + LOG_ERR("%s: draft speculative stage requires a loaded draft model\n", __func__); + return nullptr; + } + + llama_context_params cparams_dft = params.cparams_dft; + + if (has_dflash_stage) { + if (!llama_model_share_dflash_io_tensors(params.model_dft, llama_get_model(ctx_tgt))) { + LOG_ERR("%s: failed to share target IO tensors with DFlash draft model\n", __func__); + return nullptr; + } + + int32_t max_cross_ctx = 0; + for (const auto & stage : stages) { + if (stage.type != COMMON_SPECULATIVE_TYPE_DFLASH) { + continue; + } + + max_cross_ctx = std::max(max_cross_ctx, params.with_stage_overrides(stage).dflash_cross_ctx); + } + + const int32_t block_size = llama_model_dflash_block_size(params.model_dft); + if (block_size <= 0) { + LOG_ERR("%s: invalid DFlash draft block size\n", __func__); + return nullptr; + } + + const int64_t required_n_ctx = (int64_t) max_cross_ctx + (int64_t) block_size; + if (required_n_ctx > std::numeric_limits::max()) { + LOG_ERR("%s: invalid DFlash draft context size cross_ctx=%d block_size=%d required_n_ctx=%lld\n", + __func__, max_cross_ctx, block_size, (long long) required_n_ctx); + return nullptr; + } + + cparams_dft.n_ctx = (uint32_t) required_n_ctx; + } + + ctx_dft = llama_init_from_model(params.model_dft, cparams_dft); + if (ctx_dft == nullptr) { + LOG_ERR("%s", "failed to create draft context\n"); + return nullptr; + } + } + + // Compute the implementations to use based on the resolved stage chain. + std::vector configs = {}; + configs.reserve(stages.size()); + + for (const auto & stage : stages) { + common_params_speculative stage_params = params.with_stage_overrides(stage); + + if (stage.type == COMMON_SPECULATIVE_TYPE_NGRAM_MOD && !stage_params.ngram_mod) { + stage_params.ngram_mod = std::make_shared(stage_params.ngram_size_n, 4*1024*1024); + + LOG_INF("%s: initialized ngram_mod with n=%d, size=%zu (%.3f MB)\n", __func__, + stage_params.ngram_size_n, stage_params.ngram_mod->size(), + (float)(stage_params.ngram_mod->size_bytes())/1024/1024); + + if (stage_params.ngram_size_n < 16) { + LOG_WRN("%s: ngram_mod n=%d is too small - poor quality is possible, see: https://github.com/ggml-org/llama.cpp/pull/19164\n", __func__, stage_params.ngram_size_n); + } + } + + configs.push_back(common_speculative_config(stage, stage_params)); + } + + if (!configs.empty() && llama_model_has_recurrent(llama_get_model(ctx_tgt))) { + const int ckpt_tokens = std::max(1, params.get_max_stage_n_max() + 1); + const int actual_mode = llama_spec_ckpt_init(ctx_tgt, params.recurrent_ckpt_mode, ckpt_tokens); + if (actual_mode == LLAMA_SPEC_CKPT_NONE) { + LOG_ERR("%s: failed to prepare recurrent checkpoint mode '%s' during speculative init (max_tokens=%d)\n", + __func__, + params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_PER_STEP ? "per-step" : + params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_GPU_FALLBACK ? "gpu-fallback" : + params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_CPU ? "cpu" : "auto", + ckpt_tokens); + if (ctx_dft != nullptr) { + llama_free(ctx_dft); + } + return nullptr; + } + llama_spec_ckpt_discard(ctx_tgt); + params.recurrent_ckpt_mode = actual_mode; + } + + std::vector> impls = {}; + + for (const common_speculative_config & config : configs) { + LOG_DBG("%s: adding implementation %s\n", __func__, common_speculative_type_to_str(config.type).c_str()); + switch (config.type) { + case COMMON_SPECULATIVE_TYPE_NONE: + break; + case COMMON_SPECULATIVE_TYPE_DRAFT: { + impls.push_back(std::make_unique(config.type, + /* .ctx_tgt = */ ctx_tgt, + /* .ctx_dft = */ ctx_dft, + /* .replacements = */ config.params.replacements + )); + break; + } + case COMMON_SPECULATIVE_TYPE_DFLASH: { + auto state = std::make_unique( + config.type, + ctx_tgt, + ctx_dft, + config.params.dflash_cross_ctx); + if (!state->ready) { + LOG_ERR("%s: failed to initialize DFlash speculative state\n", __func__); + return nullptr; + } + impls.push_back(std::move(state)); + ctx_dft = nullptr; + break; + } + case COMMON_SPECULATIVE_TYPE_MTP: { + llama_context * ctx_mtp = ctx_dft; + if (!ctx_mtp) { + const llama_model * model = llama_get_model(ctx_tgt); + ctx_mtp = llama_init_from_model(const_cast(model), config.params.cparams_dft); + if (!ctx_mtp) { + LOG_ERR("%s: failed to create MTP context\n", __func__); + return nullptr; + } + } + ctx_dft = nullptr; + + const bool use_constant_draft_positions = llama_model_is_gemma4_mtp_assistant(llama_get_model(ctx_mtp)); + impls.push_back(std::make_unique( + config.type, ctx_tgt, ctx_mtp, use_constant_draft_positions)); + break; + } + case COMMON_SPECULATIVE_TYPE_EAGLE3: { + impls.push_back(std::make_unique(config.type)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { + common_ngram_map ngram_map = get_common_ngram_map(config); + + uint16_t ngram_size_key = ngram_map.size_key; + uint16_t mgram_size_value = ngram_map.size_value; + + auto config_simple = common_ngram_simple_config { + /* .size_ngram = */ ngram_size_key, + /* .size_mgram = */ mgram_size_value + }; + auto state = std::make_unique( + /* .type = */ config.type, + /* .state = */ config_simple + ); + impls.push_back(std::move(state)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: { + impls.push_back(std::make_unique( + (config.type), + get_common_ngram_map(config) + )); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: { + GGML_ASSERT(config.params.ngram_mod); + impls.push_back(std::make_unique(config.type, *config.params.ngram_mod)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: { + auto state = create_state_ngram_cache( + config.params.lookup_cache_static, config.params.lookup_cache_dynamic, config); + impls.push_back(std::make_unique(state)); + break; + } + case COMMON_SPECULATIVE_TYPE_SUFFIX: { + int depth = config.params.suffix_max_depth > 0 ? config.params.suffix_max_depth : 64; + const llama_model * model = llama_get_model(ctx_tgt); + impls.push_back(std::make_unique( + config.type, depth, config.params.suffix_corpus, model)); + break; + } + default: + break; + } + } + + if (impls.empty()) { + LOG_WRN("%s", "no implementations specified for speculative decoding\n"); + return nullptr; + } + + auto * result = new common_speculative { + /* .configs = */ std::move(configs), + /* .impls = */ std::move(impls) + }; + + // initialize autotune if requested + if (params.autotune && params.has_composite_stage_chain()) { + LOG_WRN("Autotune disabled — explicit speculative stage chains are not supported yet\n"); + } else if (params.autotune && !result->impls.empty()) { + auto actual_type = result->impls[0]->type; + if (actual_type != COMMON_SPECULATIVE_TYPE_NONE && + actual_type != COMMON_SPECULATIVE_TYPE_EAGLE3) { + result->tuner = std::make_unique(); + result->tuner->init(actual_type, params, llama_get_model(ctx_tgt)); + LOG_DBG("Autotune initialized for %s, tuning %zu parameters\n", + common_speculative_type_to_str(actual_type).c_str(), + result->tuner->coords.size()); + } else { + LOG_WRN("Autotune disabled — speculative type %s is not supported for autotuning\n", + common_speculative_type_to_str(actual_type).c_str()); + } + } + + return result; +} + +void common_speculative_free(common_speculative * spec) { + if (spec == nullptr) { + return; + } + + delete spec; +} + +void common_speculative_begin(common_speculative * spec, const llama_tokens & prompt) { + if (spec == nullptr) { + return; + } + + for (auto & impl : spec->impls) { + common_time_meas tm(impl->t_begin_us, !impl->gen_perf); + impl->begin(prompt); + impl->n_call_begin++; + } +} + +llama_tokens common_speculative_draft( + common_speculative * spec, + common_params_speculative & params, + const llama_tokens & prompt_tgt, // specified in target model vocab + llama_token id_last, + llama_pos draft_base_pos, + llama_seq_id draft_seq_id) { + llama_tokens result; + + spec->t_step_start_us = ggml_time_us(); + + // apply autotune proposal if enabled + if (spec->tuner && spec->tuner->enabled) { + spec->tuner->propose(params); + } + + const auto runtime_stages = params.get_resolved_stages(); + const bool use_runtime_stage_overrides = common_speculative_stage_chain_matches(runtime_stages, spec->configs); + + spec->curr_impl = nullptr; // reset current implementation + + for (size_t i = 0; i < spec->impls.size(); ++i) { + auto & impl = spec->impls[i]; + const auto & runtime_stage = use_runtime_stage_overrides ? runtime_stages[i] : spec->configs[i].stage; + common_params_speculative impl_params = common_speculative_get_runtime_params(spec->configs[i], params, runtime_stage); + result.clear(); + + { + common_time_meas tm(impl->t_draft_us, !impl->gen_perf); + impl->draft(impl_params, prompt_tgt, id_last, draft_base_pos, draft_seq_id, result); + impl->n_call_draft++; + } + + if (result.empty()) { + continue; + } + + if (common_speculative_type_is_self_spec(impl->type) && impl_params.n_min > 0 && (int)result.size() < impl_params.n_min) { + LOG_DBG("%s: impl %s drafted %zu tokens, below fallback threshold %d - trying next implementation\n", + __func__, common_speculative_type_to_str(impl->type).c_str(), result.size(), impl_params.n_min); + result.clear(); + continue; + } + LOG_DBG("%s: called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", __func__, + common_speculative_type_to_str(impl.get()->type).c_str(), prompt_tgt.size(), + impl.get()->n_call_draft, result.size()); + + spec->curr_impl = impl.get(); + impl->n_gen_drafts++; + impl->n_gen_tokens += result.size(); + + break; // We have a draft, so break out of the loop and return it. + } + + // store draft count for tuner feedback + if (spec->tuner && spec->tuner->enabled) { + spec->last_n_drafted = (int)result.size(); + } + + return result; +} + +void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) { + if (spec->tuner && spec->tuner->enabled && spec->t_step_start_us > 0) { + int64_t step_time_us = ggml_time_us() - spec->t_step_start_us; + double step_tps = (step_time_us > 100) + ? (n_accepted + 1.0) * 1e6 / (double)step_time_us + : 0.0; + spec->tuner->accept_feedback(n_accepted, spec->last_n_drafted, step_tps); + spec->t_step_start_us = 0; + } + + common_speculative_state * impl = spec->curr_impl; + + if (!impl) { + return; + } + + { + common_time_meas tm(impl->t_accept_us, !impl->gen_perf); + if (n_accepted > 0) { + impl->n_acc_drafts++; + impl->n_acc_tokens += n_accepted; + } + + impl->accept(n_accepted); + impl->n_call_accept++; + } + + if (impl->type != COMMON_SPECULATIVE_TYPE_MTP) { + if (auto * mtp_state = common_speculative_get_mtp_state(spec); mtp_state != nullptr) { + mtp_invalidate_cached_drafts(*mtp_state); + } + } +} + +static bool common_speculative_has_type(const common_speculative * spec, common_speculative_type type) { + if (spec == nullptr) { + return false; + } + + return std::any_of(spec->configs.begin(), spec->configs.end(), [type](const common_speculative_config & config) { + return config.type == type; + }); +} + +static int common_speculative_ctx_mtp_n_embd(llama_context * ctx) { + return ctx ? (int) llama_mtp_state_n_embd(ctx) : 0; +} + +static bool common_speculative_batch_token_has_seq_id( + const llama_batch & batch, + int token_index, + llama_seq_id seq_id) { + if (batch.n_seq_id == nullptr || batch.seq_id == nullptr || batch.n_seq_id[token_index] <= 0 || batch.seq_id[token_index] == nullptr) { + return false; + } + + for (int i = 0; i < batch.n_seq_id[token_index]; ++i) { + if (batch.seq_id[token_index][i] == seq_id) { + return true; + } + } + + return false; +} + +static bool common_speculative_batch_is_exact_single_seq( + const llama_batch & batch, + llama_seq_id seq_id) { + if (batch.n_tokens <= 0 || batch.n_seq_id == nullptr || batch.seq_id == nullptr) { + return false; + } + + for (int i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] != 1 || batch.seq_id[i] == nullptr || batch.seq_id[i][0] != seq_id) { + return false; + } + } + + return true; +} + +static int common_speculative_copy_seq_batch( + const llama_batch & batch, + llama_seq_id seq_id, + llama_batch & seq_batch) { + if (batch.token == nullptr || batch.pos == nullptr) { + return -1; + } + + if (batch.n_tokens < 1) { + return 0; + } + + std::vector token_indices; + token_indices.reserve(batch.n_tokens); + for (int i = 0; i < batch.n_tokens; ++i) { + if (common_speculative_batch_token_has_seq_id(batch, i, seq_id)) { + token_indices.push_back(i); + } + } + + if (token_indices.empty()) { + return 0; + } + + seq_batch = llama_batch_init((int) token_indices.size(), 0, 1); + for (const int i : token_indices) { + common_batch_add(seq_batch, batch.token[i], batch.pos[i], { seq_id }, batch.logits != nullptr && batch.logits[i]); + } + + return (int) token_indices.size(); +} + +static bool common_speculative_feature_view_copy_batch_rows( + const common_speculative_feature_view & view, + const llama_batch & batch, + llama_seq_id seq_id, + std::vector * hidden_rows) { + if (hidden_rows == nullptr || view.kind != COMMON_SPECULATIVE_FEATURE_HIDDEN_STATE || view.width <= 0 || batch.n_tokens <= 0 || batch.pos == nullptr) { + return false; + } + + std::unordered_map rows_by_pos; + rows_by_pos.reserve(view.rows.size()); + for (const auto & row : view.rows) { + if (row.seq_id == seq_id && row.data != nullptr) { + rows_by_pos[row.pos] = row.data; + } + } + + hidden_rows->clear(); + hidden_rows->reserve((size_t) batch.n_tokens * view.width); + for (int i = 0; i < batch.n_tokens; ++i) { + auto it = rows_by_pos.find(batch.pos[i]); + if (it == rows_by_pos.end()) { + hidden_rows->clear(); + return false; + } + + hidden_rows->insert(hidden_rows->end(), it->second, it->second + view.width); + } + + return hidden_rows->size() == (size_t) batch.n_tokens * view.width; +} diff --git a/common/speculative.cpp b/common/speculative.cpp index 016aeaa3..e7e6559f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -477,1745 +477,7 @@ struct common_speculative_state_mtp : public common_speculative_state { } }; -struct common_speculative_state_dflash : public common_speculative_state { - llama_context * ctx_tgt; - llama_context * ctx_dft; - - llama_batch batch = {}; - - int32_t block_size = 0; - int32_t mask_token_id = -1; - int32_t n_target_features = 0; - int32_t cross_ctx = 0; - bool ready = false; - - std::vector target_layer_ids; - std::vector target_window; - std::vector target_window_pos; - std::vector target_window_stage; - std::vector target_window_pos_stage; - std::vector target_window_ring; - std::vector target_window_append_features; - int32_t target_window_rows = 0; - int32_t target_window_ring_write_pos = 0; - int32_t target_window_ring_filled = 0; - uint64_t target_window_version = 0; - int32_t target_window_keep_rows = 0; - int32_t target_window_append_rows = 0; - bool target_window_replace = false; - bool target_window_materialized = false; - 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; - - uint64_t t_draft_decode_us = 0; - uint64_t t_draft_sample_us = 0; - uint64_t t_warmup_collect_us = 0; - uint64_t t_warmup_append_us = 0; - uint64_t t_accept_output_copy_us = 0; - uint64_t t_accept_commit_us = 0; - uint64_t t_accept_append_us = 0; - uint64_t t_accept_append_filter_us = 0; - uint64_t t_accept_append_window_alloc_us = 0; - uint64_t t_accept_append_replace_us = 0; - uint64_t t_accept_append_keep_old_us = 0; - uint64_t t_accept_append_new_rows_us = 0; - uint64_t t_accept_append_commit_detail_us = 0; - uint64_t t_accept_append_log_us = 0; - size_t n_warmup_collect_calls = 0; - size_t n_warmup_collect_rows = 0; - size_t n_warmup_append_calls = 0; - size_t n_warmup_append_rows = 0; - size_t n_accept_output_copy_calls = 0; - size_t n_accept_output_copy_rows = 0; - size_t n_accept_commit_calls = 0; - size_t n_accept_commit_rows = 0; - size_t n_accept_append_calls = 0; - size_t n_accept_append_rows = 0; - size_t n_accept_append_replace_calls = 0; - size_t n_accept_append_slide_calls = 0; - - common_speculative_state_dflash( - enum common_speculative_type type, - llama_context * ctx_tgt, - llama_context * ctx_dft, - int32_t cross_ctx) - : common_speculative_state(type) - , ctx_tgt(ctx_tgt) - , ctx_dft(ctx_dft) - , cross_ctx(std::max(1, cross_ctx)) - { - const llama_model * model_tgt = llama_get_model(ctx_tgt); - const llama_model * model_dft = llama_get_model(ctx_dft); - - if (!common_speculative_are_dflash_compatible(model_tgt, model_dft)) { - LOG_ERR("%s: DFlash draft model vocab/tokenizer is incompatible with the target model\n", __func__); - return; - } - - block_size = llama_model_dflash_block_size(model_dft); - mask_token_id = llama_model_dflash_mask_token_id(model_dft); - n_target_features = llama_model_dflash_n_target_features(model_dft); - const int32_t n_target_layers = llama_model_dflash_n_target_layers(model_dft); - - if (block_size <= 0 || mask_token_id < 0 || n_target_features <= 0 || n_target_layers <= 0) { - LOG_ERR("%s: invalid DFlash metadata (block_size=%d, mask_token_id=%d, n_target_features=%d, n_target_layers=%d)\n", - __func__, block_size, mask_token_id, n_target_features, n_target_layers); - return; - } - - target_layer_ids.resize((size_t) n_target_layers); - if (llama_model_dflash_target_layer_ids(model_dft, target_layer_ids.data(), n_target_layers) != n_target_layers) { - LOG_ERR("%s: failed to read DFlash target layer ids\n", __func__); - target_layer_ids.clear(); - return; - } - - const auto * vocab_tgt = llama_model_get_vocab(model_tgt); - const auto * vocab_dft = llama_model_get_vocab(model_dft); - const int32_t target_vocab_size = llama_vocab_n_tokens(vocab_tgt); - const int32_t draft_vocab_size = llama_vocab_n_tokens(vocab_dft); - const int32_t target_hidden_size = llama_model_n_embd(model_tgt); - const int32_t draft_hidden_size = llama_model_n_embd(model_dft); - const int32_t target_mask_token_id = llama_model_dflash_target_mask_token_id(model_tgt); - const int32_t expected_n_target_features = target_hidden_size > 0 ? target_hidden_size * n_target_layers : 0; - - if (target_mask_token_id != (int32_t) LLAMA_TOKEN_NULL && mask_token_id != target_mask_token_id) { - LOG_ERR("%s: DFlash mask token mismatch (draft=%d target=%d)\n", - __func__, mask_token_id, target_mask_token_id); - return; - } - - if (target_hidden_size <= 0 || draft_hidden_size <= 0) { - LOG_ERR("%s: invalid DFlash hidden sizes (draft=%d target=%d)\n", - __func__, draft_hidden_size, target_hidden_size); - return; - } - - if (expected_n_target_features <= 0 || n_target_features != expected_n_target_features) { - LOG_ERR("%s: DFlash target feature width mismatch (metadata=%d expected=%d target_hidden=%d target_layers=%d)\n", - __func__, n_target_features, expected_n_target_features, target_hidden_size, n_target_layers); - return; - } - - std::vector sorted_target_layer_ids = target_layer_ids; - std::sort(sorted_target_layer_ids.begin(), sorted_target_layer_ids.end()); - if (std::adjacent_find(sorted_target_layer_ids.begin(), sorted_target_layer_ids.end()) != sorted_target_layer_ids.end()) { - LOG_ERR("%s: duplicate DFlash target layer ids survived into runtime validation\n", __func__); - target_layer_ids.clear(); - return; - } - - const int32_t n_target_model_layers = llama_n_layer(model_tgt); - for (int32_t layer_id : target_layer_ids) { - if (layer_id < 0 || layer_id >= n_target_model_layers) { - LOG_ERR("%s: invalid DFlash target layer id %d for target model with %d layers\n", - __func__, layer_id, n_target_model_layers); - target_layer_ids.clear(); - return; - } - } - - const int32_t io_mode = llama_model_dflash_io_mode(model_dft, model_tgt); - if (io_mode == LLAMA_DFLASH_IO_MODE_INVALID) { - LOG_ERR("%s: DFlash draft is missing required IO tensors after target sharing\n", __func__); - return; - } - - if (io_mode == LLAMA_DFLASH_IO_MODE_MIXED) { - LOG_ERR("%s: DFlash IO contract must be fully shared or fully self-contained, but resolved to mixed mode\n", __func__); - return; - } - - if (io_mode == LLAMA_DFLASH_IO_MODE_SELF_CONTAINED && !llama_model_dflash_io_tensors_match(model_dft, target_hidden_size, target_vocab_size)) { - LOG_ERR("%s: DFlash self-contained IO tensors do not match the target hidden/vocab contract (target_hidden=%d target_vocab=%d)\n", - __func__, - target_hidden_size, - target_vocab_size); - return; - } - - if (!llama_set_dflash_capture_layers(ctx_tgt, target_layer_ids.data(), (int32_t) target_layer_ids.size())) { - LOG_ERR("%s: failed to configure DFlash target capture callback\n", __func__); - return; - } - - batch = llama_batch_init(std::max(1, block_size), 0, 1); - target_window.reserve((size_t) this->cross_ctx * (size_t) n_target_features); - target_window_stage.reserve((size_t) this->cross_ctx * (size_t) n_target_features); - target_window_ring.resize((size_t) this->cross_ctx * (size_t) n_target_features); - target_window_append_features.reserve((size_t) this->cross_ctx * (size_t) n_target_features); - target_window_pos.reserve((size_t) this->cross_ctx); - target_window_pos_stage.reserve((size_t) this->cross_ctx); - ready = true; - - 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]; - } - - const char * io_mode_name = io_mode == LLAMA_DFLASH_IO_MODE_SHARED ? "shared" : "self-contained"; - 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()); - LOG_INF("%s: DFlash artifact io=%s draft_vocab=%d target_vocab=%d draft_hidden=%d target_hidden=%d mask_token_id=%d target_mask_token_id=%d\n", - __func__, io_mode_name, draft_vocab_size, target_vocab_size, draft_hidden_size, target_hidden_size, mask_token_id, target_mask_token_id); - } - - ~common_speculative_state_dflash() override { - llama_clear_dflash_capture(ctx_tgt); - if (ctx_dft) { - llama_free(ctx_dft); - } - if (batch.token != nullptr) { - llama_batch_free(batch); - } - } - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - llama_kv_cache_clear(ctx_dft); - llama_reset_dflash_kv_cache_state(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; - t_draft_decode_us = 0; - t_draft_sample_us = 0; - t_warmup_collect_us = 0; - t_warmup_append_us = 0; - t_accept_output_copy_us = 0; - t_accept_commit_us = 0; - t_accept_append_us = 0; - t_accept_append_filter_us = 0; - t_accept_append_window_alloc_us = 0; - t_accept_append_replace_us = 0; - t_accept_append_keep_old_us = 0; - t_accept_append_new_rows_us = 0; - t_accept_append_commit_detail_us = 0; - t_accept_append_log_us = 0; - n_warmup_collect_calls = 0; - n_warmup_collect_rows = 0; - n_warmup_append_calls = 0; - n_warmup_append_rows = 0; - n_accept_output_copy_calls = 0; - n_accept_output_copy_rows = 0; - n_accept_commit_calls = 0; - n_accept_commit_rows = 0; - n_accept_append_calls = 0; - n_accept_append_rows = 0; - n_accept_append_replace_calls = 0; - n_accept_append_slide_calls = 0; - llama_dflash_profile_reset(ctx_tgt); - llama_dflash_profile_reset(ctx_dft); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - GGML_UNUSED(prompt_tgt); - - result.clear(); - if (!ready || target_window_rows <= 0) { - n_draft_empty++; - return; - } - - const int32_t n_keep = std::min(params.n_max, block_size - 1); - if (n_keep <= 0) { - return; - } - - const bool use_kv_cache = dflash_use_kv_cache_experiment(); - const float * target_features = nullptr; - size_t target_feature_floats = 0; - llama_dflash_window_update window_update = { - target_window_version, - target_window_keep_rows, - target_window_append_rows, - target_window_replace, - target_window_append_features.empty() ? nullptr : target_window_append_features.data(), - target_window_append_features.size(), - }; - const llama_dflash_kv_cache_transition cache_plan = use_kv_cache - ? llama_plan_dflash_kv_cache_transition_for_ctx(ctx_dft, window_update, target_window_rows) - : llama_dflash_kv_cache_transition{}; - - if (!use_kv_cache || cache_plan.rebuild_cache) { - dflash_materialize_target_window_features(*this); - target_features = target_window.data(); - target_feature_floats = target_window.size(); - } - if (use_kv_cache && cache_plan.rebuild_cache) { - window_update.append_features = target_window.data(); - window_update.append_floats = target_window.size(); - window_update.append_rows = target_window_rows; - } - - if (!llama_set_dflash_target_features_view(ctx_dft, target_features, target_feature_floats, target_window_rows, target_window_pos.data(), &window_update)) { - 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 int32_t batch_len = n_keep + 1; - const llama_pos draft_pos_base = last_target_pos >= 0 ? last_target_pos + 1 : (llama_pos) target_window_rows; - const llama_pos seed_pos = last_target_pos >= 0 ? last_target_pos : draft_pos_base - 1; - last_draft_pos_base = draft_pos_base; - common_batch_add(batch, id_last, seed_pos, { 0 }, false); - for (int32_t i = 1; i < batch_len; ++i) { - common_batch_add(batch, mask_token_id, draft_pos_base + (i - 1), { 0 }, i <= n_keep); - } - - const int64_t t_decode_us = ggml_time_us(); - 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; - } - t_draft_decode_us += (uint64_t) (ggml_time_us() - t_decode_us); - - result.reserve((size_t) n_keep); - const int64_t t_sample_us = ggml_time_us(); - for (int32_t i = 0; i < n_keep; ++i) { - result.push_back(common_sampler_sample_speculative(nullptr, ctx_dft, i + 1, nullptr)); - } - t_draft_sample_us += (uint64_t) (ggml_time_us() - t_sample_us); - - batch.n_tokens = 0; - dflash_contract_log_draft(*this, n_keep, result.size()); - } - - void accept(uint16_t n_accepted) override { - GGML_UNUSED(n_accepted); - } -}; - -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; - const llama_pos seed_pos = state.last_target_pos; - const llama_pos mask_first_pos = state.last_draft_pos_base; - const llama_pos mask_last_pos = state.last_draft_pos_base >= 0 - ? state.last_draft_pos_base + n_keep - 1 - : -1; - - LOG_INF("dflash contract draft[%llu]: window_rows=%d window_pos=%s pos=[%d..%d] gaps=%d nonmono=%d last_target_pos=%d seed_pos=%d mask_pos=[%d..%d] sample_rows=[1..%d] output_rows=[1..%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) seed_pos, - (int) mask_first_pos, - (int) mask_last_pos, - n_keep, - n_keep, - (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; - - common_sampler * smpl; - - llama_batch batch; - llama_tokens prompt_dft; - - bool vocab_cmpt = true; // whether retokenization is needed - std::unordered_map vocab_map; - - common_speculative_state_draft( - enum common_speculative_type type, - llama_context * ctx_tgt, - llama_context * ctx_dft, - const std::vector> & replacements) - : common_speculative_state(type) - , ctx_tgt(ctx_tgt) - , ctx_dft(ctx_dft) - { - batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1); - smpl = nullptr; - { - struct common_params_sampling params; - params.top_k = 10; - params.samplers_sequence = { - llama_sampler_type::TOP_K, - llama_sampler_type::DIST, // needed to get probabilities - }; - smpl = common_sampler_init(llama_get_model(ctx_dft), params); - } - - vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft)); - LOG_DBG("vocab_cmpt = %d\n", vocab_cmpt); - - if (!vocab_cmpt) { - LOG_WRN("the target and draft vocabs are not compatible - tokens will be translated between the two\n"); - - for (const auto & pair : replacements) { - vocab_map[pair.first] = pair.second; - } - } - } - - ~common_speculative_state_draft() override { - llama_free(ctx_dft); - - common_sampler_free(smpl); - - llama_batch_free(batch); - } - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - auto * spec = this; - - auto & batch = spec->batch; - auto & ctx_tgt = spec->ctx_tgt; - auto & ctx_dft = spec->ctx_dft; - auto & smpl = spec->smpl; - auto & prompt_dft = spec->prompt_dft; - - int reuse_i = 0; - int reuse_n = 0; - - const int n_ctx = llama_n_ctx(ctx_dft) - params.n_max; - - llama_tokens prompt_cnv; - if (!spec->vocab_cmpt) { - // convert id_last to draft vocab. llama_detokenize is called directly to avoid an allocation - const auto * model_tgt = llama_get_model(ctx_tgt); - const auto * vocab_tgt = llama_model_get_vocab(model_tgt); - - std::string text; - - text = common_detokenize(ctx_tgt, prompt_tgt, true); - text = replace_to_dft(text); - - LOG_DBG("%s: main->draft detokenized string: '%s'\n", __func__, text.c_str()); - - prompt_cnv = common_tokenize(ctx_dft, text, false, true); - - - - int32_t n_chars = llama_detokenize(vocab_tgt, &id_last, 1, nullptr, 0, false, false); - GGML_ASSERT(n_chars < 0 && "failed to detokenize id_last"); - - text.resize(-n_chars); - llama_detokenize(vocab_tgt, &id_last, 1, text.data(), text.size(), false, false); - text = replace_to_dft(text); - - LOG_DBG("main->draft detokenized id_last(%d): '%s'\n", id_last, text.c_str()); - id_last = common_tokenize(ctx_dft, text, false, true)[0]; - } - - const llama_tokens & prompt_cur = spec->vocab_cmpt ? prompt_tgt : prompt_cnv; - - const int i_start = std::max(0, (int) prompt_cur.size() - n_ctx); - - // reuse as much as possible from the old draft context - // ideally, the draft context should be as big as the target context and we will always reuse the entire prompt - for (int i = 0; i < (int) prompt_dft.size(); ++i) { - int cur = 0; - while (i_start + cur < (int) prompt_cur.size() && - i + cur < (int) prompt_dft.size() && - prompt_cur[i_start + cur] == prompt_dft[i + cur]) { - cur++; - } - - if ((cur >= 256 || n_ctx >= (int) prompt_cur.size()) && cur > reuse_n) { - reuse_i = i; - reuse_n = cur; - } - } - - LOG_DBG("%s: reuse_i = %d, reuse_n = %d, prompt = %d\n", __func__, reuse_i, reuse_n, (int) prompt_dft.size()); - - result.clear(); - result.reserve(params.n_max); - - if (reuse_n == 0) { - llama_kv_cache_clear(ctx_dft); - prompt_dft.clear(); - } else { - // this happens when a previous draft has been discarded (for example, due to being too small), but the - // target model agreed with it. in this case, we simply pass back the previous results to save compute - if (reuse_i + reuse_n < (int) prompt_dft.size() && prompt_dft[reuse_i + reuse_n] == id_last) { - for (int i = reuse_i + reuse_n + 1; i < (int) prompt_dft.size(); ++i) { - result.push_back(prompt_dft[i]); - - if (params.n_max <= (int) result.size()) { - break; - } - } - - return; - } - - if (reuse_i > 0) { - llama_kv_cache_seq_rm (ctx_dft, 0, 0, reuse_i); - llama_kv_cache_seq_add(ctx_dft, 0, reuse_i, -1, -reuse_i); - - prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); - } - - if (reuse_n < (int) prompt_dft.size()) { - llama_kv_cache_seq_rm (ctx_dft, 0, reuse_n, -1); - prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); - } - } - - // prepare a batch to evaluate any new tokens in the prompt - common_batch_clear(batch); - - for (size_t i = i_start + reuse_n; i < prompt_cur.size(); ++i) { - //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_cur[i]); - common_batch_add(batch, prompt_cur[i], i - i_start, { 0 }, false); - - prompt_dft.push_back(prompt_cur[i]); - } - - // we should rarely end-up here during normal decoding - if (batch.n_tokens > 0) { - //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str()); - - llama_decode(ctx_dft, batch); - } - - const llama_pos n_past = prompt_dft.size(); - - LOG_DBG("%s: n_past = %d\n", __func__, n_past); - - common_batch_clear(batch); - common_batch_add (batch, id_last, n_past, { 0 }, true); - - prompt_dft.push_back(id_last); - - //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx_dft, prompt_dft).c_str()); - - llama_decode(ctx_dft, batch); - - common_sampler_reset(smpl); - - // sample n_draft tokens from the draft model - for (int i = 0; i < params.n_max; ++i) { - common_batch_clear(batch); - - common_sampler_sample(smpl, ctx_dft, 0, true); - - const auto * cur_p = common_sampler_get_candidates(smpl, true); - - for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", - k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); - } - - // add drafted token for each sequence - const llama_token id = cur_p->data[0].id; - - common_sampler_accept(smpl, nullptr, id, true); - - // only collect very high-confidence draft tokens - if (cur_p->data[0].p < params.p_min) { - if (i == 0) { - result.push_back(id); - } - break; - } - - result.push_back(id); - - if (params.n_max <= (int) result.size()) { - break; - } - - - common_batch_add(batch, id, n_past + i + 1, { 0 }, true); - - // evaluate the drafted tokens on the draft model - llama_decode(ctx_dft, batch); - - prompt_dft.push_back(id); - } - - if (!spec->vocab_cmpt) { - std::string detokenized = common_detokenize(ctx_dft, result, true); - detokenized = replace_to_tgt(detokenized); - LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); - result = common_tokenize(ctx_tgt, detokenized, false, true); - if (result.size() > (size_t)params.n_max) { - result.resize(params.n_max); - } - } - } - - void accept(uint16_t n_accepted) override { - // noop - GGML_UNUSED(n_accepted); - } - - std::string replace_to_dft(const std::string & input) const { - std::string result = input; - - for (const auto & pair : this->vocab_map) { - size_t pos = result.find(pair.first); - while (pos != std::string::npos) { - result.replace(pos, pair.first.length(), pair.second); - pos = result.find(pair.first, pos + pair.second.length()); - } - } - - return result; - } - - std::string replace_to_tgt(const std::string & input) const { - std::string result = input; - - for (const auto & pair : this->vocab_map) { - size_t pos = result.find(pair.second); - while (pos != std::string::npos) { - result.replace(pos, pair.second.length(), pair.first); - pos = result.find(pair.second, pos + pair.first.length()); - } - } - - return result; - } -}; - -struct common_speculative_state_eagle3 : public common_speculative_state { - common_speculative_state_eagle3(enum common_speculative_type type) : common_speculative_state(type) {} - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & draft_tokens) override { - // TODO: implement - GGML_UNUSED(params); - GGML_UNUSED(prompt_tgt); - GGML_UNUSED(id_last); - GGML_UNUSED(draft_tokens); - } - - void accept(uint16_t n_accepted) override { - // noop - GGML_UNUSED(n_accepted); - } -}; - -// state of self-speculation (simple implementation, not ngram-map) -struct common_speculative_state_ngram_simple : public common_speculative_state { - common_ngram_simple_config config; - - common_speculative_state_ngram_simple( - enum common_speculative_type type, - common_ngram_simple_config config) - : common_speculative_state(type), config(config) {} - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - - result = common_ngram_simple_draft(config, prompt_tgt, id_last); - GGML_UNUSED(params); - } - - void accept(uint16_t n_accepted) override { - // noop - GGML_UNUSED(n_accepted); - } -}; - -struct common_speculative_state_ngram_map_k : public common_speculative_state { - // draft ngram map for speculative decoding without draft model - common_ngram_map map; - - common_speculative_state_ngram_map_k( - enum common_speculative_type type, - common_ngram_map map) - : common_speculative_state(type), map(std::move(map)) {} - - void begin(const llama_tokens & prompt) override { - common_ngram_map_begin(map, prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - common_ngram_map_draft(map, prompt_tgt, id_last, result); - GGML_UNUSED(params); - } - - void accept(uint16_t n_accepted) override { - common_ngram_map_accept(map, n_accepted); - } -}; - -struct common_speculative_state_ngram_mod : public common_speculative_state { - common_ngram_mod & mod; - - // the last position in the prompt that was added to the ngram container - size_t i_last = 0; - - // length of the last drafted n‑gram (number of tokens returned by draft) - size_t n_draft_last = 0; - - // consecutive accept rounds with low acceptance fraction (< 0.5) - int n_low = 0; - - // enable trace logging if LLAMA_TRACE is set - const bool verbose; - - common_speculative_state_ngram_mod(enum common_speculative_type type, common_ngram_mod & mod) - : common_speculative_state(type), mod(mod), verbose(std::getenv("LLAMA_TRACE") != nullptr) { - static_assert(sizeof(llama_token) == sizeof(common_ngram_mod::entry_t)); - } - - void begin(const llama_tokens & prompt) override { - i_last = 0; - - n_draft_last = 0; - n_low = 0; - - const size_t n = mod.get_n(); - - if (prompt.size() < n) { - return; - } - - for (size_t i = 0; i < prompt.size() - n; ++i) { - mod.add(prompt.data() + i); - } - - i_last = prompt.size() - n; - - const double f = (double)mod.get_used() / (double)mod.size(); - LOG_INF("%s: ngram_mod occupancy = %zu/%zu (%.2f)\n", __func__, mod.get_used(), mod.size(), f); - - constexpr double f_thold = 0.25; - if (f > f_thold) { - LOG_WRN("%s: ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", __func__, f, f_thold); - - mod.reset(); - } - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - GGML_UNUSED(params); - - n_draft_last = 0; - - const size_t cur_len = prompt_tgt.size(); - if (cur_len < mod.get_n()) { - return; - } - - const size_t n = mod.get_n(); - - // add new ngrams in chunks - if (i_last + 32 < cur_len) { - for (size_t i = i_last; i < cur_len - n; ++i) { - mod.add(prompt_tgt.data() + i); - } - - i_last = cur_len - n; - } - - result.resize(n + params.n_max); - for (size_t i = 0; i < n - 1; ++i) { - result[i] = prompt_tgt[cur_len - n + 1 + i]; - } - result[n - 1] = id_last; - - for (int i = 0; i < params.n_max; ++i) { - const llama_token token = mod.get(result.data() + i); - if (token == common_ngram_mod::EMPTY) { - if (i < params.n_min) { - result.clear(); - return; - } - - result.resize(n + i); - break; - } - result[n + i] = token; - } - - // only return the m tokens that were drafted - for (size_t i = 0; n + i < result.size(); ++i) { - result[i] = result[n + i]; - } - result.resize(result.size() - n); - - // store length of drafted n‑gram for later acceptance analysis - n_draft_last = result.size(); - } - - void accept(uint16_t n_accepted) override { - if (verbose) { - LOG_INF("%s: accepted %d tokens from %zu drafted tokens\n", __func__, n_accepted, n_draft_last); - } - - // compute acceptance fraction if we have a recorded draft length - if (n_draft_last > 0) { - const double f_acc = (double)n_accepted / (double)n_draft_last; - if (f_acc < 0.5) { - n_low++; - if (n_low >= 3) { - LOG_WRN("%s: low acceptance streak (%d) – resetting ngram_mod\n", __func__, n_low); - - mod.reset(); - n_low = 0; - i_last = 0; - } - } else { - n_low = 0; - } - } - } -}; - -struct common_speculative_state_ngram_cache : public common_speculative_state { - uint16_t n_draft; - bool save_dynamic; - bool save_static; - - common_ngram_cache ngram_cache_context; - common_ngram_cache ngram_cache_dynamic; - common_ngram_cache ngram_cache_static; - - size_t cache_size = 0; // number of tokens in n-gram cache - - common_speculative_state_ngram_cache( - const enum common_speculative_type type, - const std::string & path_static, - const std::string & path_dynamic, - uint16_t n_draft, - bool save_dynamic, - bool save_static) - : common_speculative_state(type) - , n_draft(n_draft) - , save_dynamic(save_dynamic) - , save_static(save_static) - { - if (!path_static.empty()) { - try { - ngram_cache_static = common_ngram_cache_load(path_static); - } catch (...) { - LOG_ERR("failed to open static lookup cache: %s", path_static.c_str()); - GGML_ABORT("Couldn't read static lookup cache"); - } - } - - if (!path_dynamic.empty()) { - try { - ngram_cache_dynamic = common_ngram_cache_load(path_dynamic); - } catch (...) { - LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str()); - GGML_ABORT("Couldn't read dynamic lookup cache"); - } - } - } - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - GGML_UNUSED(params); - - if (cache_size < prompt_tgt.size() + 1) { - llama_tokens tokens_new; - tokens_new.reserve(prompt_tgt.size() + 1 - cache_size); - for (size_t j = cache_size; j < prompt_tgt.size(); ++j) { - tokens_new.push_back(prompt_tgt[j]); - } - tokens_new.push_back(id_last); // add the last token - - // Update context ngram cache with new prompt_tgt: - common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, - tokens_new, tokens_new.size(), false); - cache_size = prompt_tgt.size() + 1; - } - - llama_tokens inp; - inp.reserve(prompt_tgt.size() + 1); - for (size_t j = 0; j < prompt_tgt.size(); ++j) { - inp.push_back(prompt_tgt[j]); - } - inp.push_back(id_last); - - result.push_back(id_last); - - common_ngram_cache_draft(inp, result, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, - ngram_cache_context, - ngram_cache_dynamic, - ngram_cache_static); - - if (result.size() > 0) { - // delete first token in result (which is the id_last token) - result.erase(result.begin()); - } - } - - void accept(uint16_t n_accepted) override { - // TODO: noop - GGML_UNUSED(n_accepted); - } -}; - -struct common_speculative_state_suffix : public common_speculative_state { - common_suffix_tree tree; - common_suffix_tree corpus_tree; - bool has_corpus = false; - size_t cache_size = 0; - - // Acceptance feedback - size_t n_draft_last = 0; - bool had_accept = false; - int n_low = 0; - float base_p_min = 0.1f; - float eff_p_min = 0.1f; - - common_speculative_state_suffix( - enum common_speculative_type type, - int max_depth, - const std::string & corpus_path, - const llama_model * model) - : common_speculative_state(type) - , tree(max_depth) - , corpus_tree(max_depth) - { - if (!corpus_path.empty()) { - std::function(const std::string &)> tokenize_fn; - if (model) { - tokenize_fn = [model](const std::string & text) -> std::vector { - return common_tokenize(model, text, false, true); - }; - } - has_corpus = corpus_tree.load_corpus(corpus_path, tokenize_fn); - } - } - - void begin(const llama_tokens & prompt) override { - cache_size = 0; - n_draft_last = 0; - had_accept = false; - n_low = 0; - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - - base_p_min = params.p_min; - if (n_draft_last > 0 && !had_accept) { - if (++n_low >= 3) { - eff_p_min = std::min(eff_p_min + 0.1f, 0.5f); - n_low = 0; - } - } - had_accept = false; - - if (cache_size < prompt_tgt.size() + 1) { - llama_tokens tokens_new; - tokens_new.reserve(prompt_tgt.size() + 1 - cache_size); - for (size_t j = cache_size; j < prompt_tgt.size(); ++j) { - tokens_new.push_back(prompt_tgt[j]); - } - tokens_new.push_back(id_last); - - tree.extend(tokens_new.data(), (int)tokens_new.size()); - cache_size = prompt_tgt.size() + 1; - } - - const int ctx_len = std::min((int)(prompt_tgt.size() + 1), tree.max_depth()); - llama_tokens context; - context.reserve(ctx_len); - const int ctx_start = (int)prompt_tgt.size() + 1 - ctx_len; - for (int j = ctx_start; j < (int)prompt_tgt.size(); ++j) { - context.push_back(prompt_tgt[j]); - } - context.push_back(id_last); - const int min_match_len = std::max(1, params.suffix_min_match_len); - - result = tree.speculate( - context.data(), (int)context.size(), - params.n_max, - eff_p_min, - 1, - min_match_len); - - if (has_corpus) { - auto corpus_result = corpus_tree.speculate( - context.data(), (int)context.size(), - params.n_max, - eff_p_min, - 1, - min_match_len); - if (corpus_result.size() > result.size()) { - result = std::move(corpus_result); - } - } - - n_draft_last = result.size(); - } - - void accept(uint16_t n_accepted) override { - if (n_draft_last == 0) { - return; - } - had_accept = true; - const double f_acc = (double)n_accepted / (double)n_draft_last; - if (f_acc < 0.5) { - if (++n_low >= 3) { - eff_p_min = std::min(eff_p_min + 0.1f, 0.5f); - n_low = 0; - } - } else { - n_low = 0; - if (eff_p_min > base_p_min) { - eff_p_min = std::max(eff_p_min - 0.05f, base_p_min); - } - } - } -}; - -struct common_speculative { - std::vector configs; // resolved stage config for each implementation - std::vector> impls; // list of implementations to use and their states - common_speculative_state * curr_impl = nullptr; // current implementation in use (for stats) - std::unique_ptr tuner; - int last_n_drafted = 0; - int64_t t_step_start_us = 0; -}; - -static bool common_speculative_stage_chain_matches( - const std::vector & stages, - const std::vector & configs) { - if (stages.size() != configs.size()) { - return false; - } - - for (size_t i = 0; i < stages.size(); ++i) { - if (stages[i].type != configs[i].type) { - return false; - } - } - - return true; -} - -static common_params_speculative common_speculative_get_runtime_params( - const common_speculative_config & config, - const common_params_speculative & params, - const common_speculative_stage_params & stage) { - common_params_speculative result = config.params; - - result.type = config.type; - result.n_max = stage.has_n_max_override() ? stage.n_max : params.n_max; - result.n_min = stage.has_n_min_override() ? stage.n_min : params.n_min; - result.p_min = stage.has_p_min_override() ? stage.p_min : params.p_min; - - if (config.type == COMMON_SPECULATIVE_TYPE_SUFFIX) { - result.suffix_min_match_len = stage.has_suffix_min_match_len_override() - ? stage.suffix_min_match_len - : params.suffix_min_match_len; - } - - result.n_max = std::max(result.n_max, 0); - result.n_min = std::max(0, std::min(result.n_min, result.n_max)); - result.stages.clear(); - - return result; -} - -static common_ngram_map get_common_ngram_map(const common_speculative_config & config) { - uint16_t size_key = config.params.ngram_size_n; - uint16_t size_value = config.params.ngram_size_m; - bool key_only = (config.type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K); - uint16_t min_hits = config.params.ngram_min_hits; - - return common_ngram_map(size_key, size_value, key_only, min_hits); -} - -static common_speculative_state_ngram_cache create_state_ngram_cache( - const std::string & path_static, const std::string & path_dynamic, - const common_speculative_config & config) { - uint16_t n_draft = 8; // TODO get from config? - - // TODO bool param in common/common.h to set save_static/save_dynamic? - bool save_static = false; - bool save_dynamic = false; - - common_speculative_state_ngram_cache state(config.type, path_static, path_dynamic, n_draft, save_static, save_dynamic); - - return state; -} - -std::string common_speculative_type_name_str() { - std::string result; - for (size_t i = 0; i < common_speculative_types.size(); i++) { - if (i > 0) { - result += ", "; - } - result += common_speculative_type_to_str(common_speculative_types[i]); - } - return result; -} - -std::string common_speculative_type_to_str(enum common_speculative_type type) { - switch (type) { - case COMMON_SPECULATIVE_TYPE_NONE: return "none"; - case COMMON_SPECULATIVE_TYPE_DRAFT: return "draft"; - case COMMON_SPECULATIVE_TYPE_DFLASH: return "dflash"; - case COMMON_SPECULATIVE_TYPE_MTP: return "mtp"; - case COMMON_SPECULATIVE_TYPE_EAGLE3: return "eagle3"; - case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram_simple"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram_map_k"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram_map_k4v"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: return "ngram_mod"; - case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: return "ngram_cache"; - case COMMON_SPECULATIVE_TYPE_SUFFIX: return "suffix"; - default: return "unknown"; - } -} - -enum common_speculative_type common_speculative_type_from_name(const std::string & name) { - std::string normalized = name; - std::replace(normalized.begin(), normalized.end(), '-', '_'); - - const auto it = common_speculative_type_from_name_map.find(normalized); - if (it == common_speculative_type_from_name_map.end()) { - return COMMON_SPECULATIVE_TYPE_COUNT; - } - return it->second; -} - -bool common_speculative_is_compat(llama_context * ctx_tgt) { - bool res = true; - - llama_kv_cache_clear(ctx_tgt); - - // eval 2 tokens to check if the context is compatible - std::vector tmp; - tmp.push_back(0); - tmp.push_back(0); - - int ret = llama_decode(ctx_tgt, llama_batch_get_one(tmp.data(), tmp.size(), 0, 0)); - if (ret != 0) { - LOG_ERR("%s: llama_decode() failed: %d\n", __func__, ret); - res = false; - goto done; - } - - // try to remove the last tokens - if (!llama_kv_cache_seq_rm(ctx_tgt, 0, 1, -1)) { - LOG_WRN("%s: the target context does not support partial sequence removal\n", __func__); - res = false; - goto done; - } - -done: - llama_kv_cache_clear(ctx_tgt); - llama_synchronize(ctx_tgt); - - return res; -} - -// initialization of the speculative decoding system -// -common_speculative * common_speculative_init( - common_params_speculative & params, - llama_context * ctx_tgt) { - std::string chain_error; - if (!common_speculative_validate_chain(params, &chain_error)) { - LOG_ERR("%s: invalid speculative stage chain: %s\n", __func__, chain_error.c_str()); - return nullptr; - } - - const auto stages = params.get_resolved_stages(); - if (params.model_dft && llama_model_is_gemma4_mtp_assistant(params.model_dft)) { - const bool has_draft_stage = std::any_of(stages.begin(), stages.end(), [](const common_speculative_stage_params & stage) { - return stage.type == COMMON_SPECULATIVE_TYPE_DRAFT; - }); - - if (has_draft_stage) { - LOG_ERR("%s: Gemma4 assistant models only support MTP stages; omit -md for self-spec-only runs or use -mtp/--spec-stage mtp for assistant-backed MTP\n", __func__); - return nullptr; - } - } - - const bool has_dflash_stage = std::any_of(stages.begin(), stages.end(), [](const common_speculative_stage_params & stage) { - return stage.type == COMMON_SPECULATIVE_TYPE_DFLASH; - }); - - const bool needs_draft_ctx = std::any_of(stages.begin(), stages.end(), [¶ms](const common_speculative_stage_params & stage) { - return stage.type == COMMON_SPECULATIVE_TYPE_DRAFT || - stage.type == COMMON_SPECULATIVE_TYPE_DFLASH || - (stage.type == COMMON_SPECULATIVE_TYPE_MTP && params.model_dft != nullptr); - }); - - llama_context * ctx_dft = nullptr; - if (needs_draft_ctx) { - if (!params.model_dft) { - LOG_ERR("%s: draft speculative stage requires a loaded draft model\n", __func__); - return nullptr; - } - - llama_context_params cparams_dft = params.cparams_dft; - - if (has_dflash_stage) { - if (!llama_model_share_dflash_io_tensors(params.model_dft, llama_get_model(ctx_tgt))) { - LOG_ERR("%s: failed to share target IO tensors with DFlash draft model\n", __func__); - return nullptr; - } - - int32_t max_cross_ctx = 0; - for (const auto & stage : stages) { - if (stage.type != COMMON_SPECULATIVE_TYPE_DFLASH) { - continue; - } - - max_cross_ctx = std::max(max_cross_ctx, params.with_stage_overrides(stage).dflash_cross_ctx); - } - - const int32_t block_size = llama_model_dflash_block_size(params.model_dft); - if (block_size <= 0) { - LOG_ERR("%s: invalid DFlash draft block size\n", __func__); - return nullptr; - } - - const int64_t required_n_ctx = (int64_t) max_cross_ctx + (int64_t) block_size; - if (required_n_ctx > std::numeric_limits::max()) { - LOG_ERR("%s: invalid DFlash draft context size cross_ctx=%d block_size=%d required_n_ctx=%lld\n", - __func__, max_cross_ctx, block_size, (long long) required_n_ctx); - return nullptr; - } - - cparams_dft.n_ctx = (uint32_t) required_n_ctx; - } - - ctx_dft = llama_init_from_model(params.model_dft, cparams_dft); - if (ctx_dft == nullptr) { - LOG_ERR("%s", "failed to create draft context\n"); - return nullptr; - } - } - - // Compute the implementations to use based on the resolved stage chain. - std::vector configs = {}; - configs.reserve(stages.size()); - - for (const auto & stage : stages) { - common_params_speculative stage_params = params.with_stage_overrides(stage); - - if (stage.type == COMMON_SPECULATIVE_TYPE_NGRAM_MOD && !stage_params.ngram_mod) { - stage_params.ngram_mod = std::make_shared(stage_params.ngram_size_n, 4*1024*1024); - - LOG_INF("%s: initialized ngram_mod with n=%d, size=%zu (%.3f MB)\n", __func__, - stage_params.ngram_size_n, stage_params.ngram_mod->size(), - (float)(stage_params.ngram_mod->size_bytes())/1024/1024); - - if (stage_params.ngram_size_n < 16) { - LOG_WRN("%s: ngram_mod n=%d is too small - poor quality is possible, see: https://github.com/ggml-org/llama.cpp/pull/19164\n", __func__, stage_params.ngram_size_n); - } - } - - configs.push_back(common_speculative_config(stage, stage_params)); - } - - if (!configs.empty() && llama_model_has_recurrent(llama_get_model(ctx_tgt))) { - const int ckpt_tokens = std::max(1, params.get_max_stage_n_max() + 1); - const int actual_mode = llama_spec_ckpt_init(ctx_tgt, params.recurrent_ckpt_mode, ckpt_tokens); - if (actual_mode == LLAMA_SPEC_CKPT_NONE) { - LOG_ERR("%s: failed to prepare recurrent checkpoint mode '%s' during speculative init (max_tokens=%d)\n", - __func__, - params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_PER_STEP ? "per-step" : - params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_GPU_FALLBACK ? "gpu-fallback" : - params.recurrent_ckpt_mode == LLAMA_SPEC_CKPT_CPU ? "cpu" : "auto", - ckpt_tokens); - if (ctx_dft != nullptr) { - llama_free(ctx_dft); - } - return nullptr; - } - llama_spec_ckpt_discard(ctx_tgt); - params.recurrent_ckpt_mode = actual_mode; - } - - std::vector> impls = {}; - - for (const common_speculative_config & config : configs) { - LOG_DBG("%s: adding implementation %s\n", __func__, common_speculative_type_to_str(config.type).c_str()); - switch (config.type) { - case COMMON_SPECULATIVE_TYPE_NONE: - break; - case COMMON_SPECULATIVE_TYPE_DRAFT: { - impls.push_back(std::make_unique(config.type, - /* .ctx_tgt = */ ctx_tgt, - /* .ctx_dft = */ ctx_dft, - /* .replacements = */ config.params.replacements - )); - break; - } - case COMMON_SPECULATIVE_TYPE_DFLASH: { - auto state = std::make_unique( - config.type, - ctx_tgt, - ctx_dft, - config.params.dflash_cross_ctx); - if (!state->ready) { - LOG_ERR("%s: failed to initialize DFlash speculative state\n", __func__); - return nullptr; - } - impls.push_back(std::move(state)); - ctx_dft = nullptr; - break; - } - case COMMON_SPECULATIVE_TYPE_MTP: { - llama_context * ctx_mtp = ctx_dft; - if (!ctx_mtp) { - const llama_model * model = llama_get_model(ctx_tgt); - ctx_mtp = llama_init_from_model(const_cast(model), config.params.cparams_dft); - if (!ctx_mtp) { - LOG_ERR("%s: failed to create MTP context\n", __func__); - return nullptr; - } - } - ctx_dft = nullptr; - - const bool use_constant_draft_positions = llama_model_is_gemma4_mtp_assistant(llama_get_model(ctx_mtp)); - impls.push_back(std::make_unique( - config.type, ctx_tgt, ctx_mtp, use_constant_draft_positions)); - break; - } - case COMMON_SPECULATIVE_TYPE_EAGLE3: { - impls.push_back(std::make_unique(config.type)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { - common_ngram_map ngram_map = get_common_ngram_map(config); - - uint16_t ngram_size_key = ngram_map.size_key; - uint16_t mgram_size_value = ngram_map.size_value; - - auto config_simple = common_ngram_simple_config { - /* .size_ngram = */ ngram_size_key, - /* .size_mgram = */ mgram_size_value - }; - auto state = std::make_unique( - /* .type = */ config.type, - /* .state = */ config_simple - ); - impls.push_back(std::move(state)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: { - impls.push_back(std::make_unique( - (config.type), - get_common_ngram_map(config) - )); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: { - GGML_ASSERT(config.params.ngram_mod); - impls.push_back(std::make_unique(config.type, *config.params.ngram_mod)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: { - auto state = create_state_ngram_cache( - config.params.lookup_cache_static, config.params.lookup_cache_dynamic, config); - impls.push_back(std::make_unique(state)); - break; - } - case COMMON_SPECULATIVE_TYPE_SUFFIX: { - int depth = config.params.suffix_max_depth > 0 ? config.params.suffix_max_depth : 64; - const llama_model * model = llama_get_model(ctx_tgt); - impls.push_back(std::make_unique( - config.type, depth, config.params.suffix_corpus, model)); - break; - } - default: - break; - } - } - - if (impls.empty()) { - LOG_WRN("%s", "no implementations specified for speculative decoding\n"); - return nullptr; - } - - auto * result = new common_speculative { - /* .configs = */ std::move(configs), - /* .impls = */ std::move(impls) - }; - - // initialize autotune if requested - if (params.autotune && params.has_composite_stage_chain()) { - LOG_WRN("Autotune disabled — explicit speculative stage chains are not supported yet\n"); - } else if (params.autotune && !result->impls.empty()) { - auto actual_type = result->impls[0]->type; - if (actual_type != COMMON_SPECULATIVE_TYPE_NONE && - actual_type != COMMON_SPECULATIVE_TYPE_EAGLE3) { - result->tuner = std::make_unique(); - result->tuner->init(actual_type, params, llama_get_model(ctx_tgt)); - LOG_DBG("Autotune initialized for %s, tuning %zu parameters\n", - common_speculative_type_to_str(actual_type).c_str(), - result->tuner->coords.size()); - } else { - LOG_WRN("Autotune disabled — speculative type %s is not supported for autotuning\n", - common_speculative_type_to_str(actual_type).c_str()); - } - } - - return result; -} - -void common_speculative_free(common_speculative * spec) { - if (spec == nullptr) { - return; - } - - delete spec; -} - -void common_speculative_begin(common_speculative * spec, const llama_tokens & prompt) { - if (spec == nullptr) { - return; - } - - for (auto & impl : spec->impls) { - common_time_meas tm(impl->t_begin_us, !impl->gen_perf); - impl->begin(prompt); - impl->n_call_begin++; - } -} - -llama_tokens common_speculative_draft( - common_speculative * spec, - common_params_speculative & params, - const llama_tokens & prompt_tgt, // specified in target model vocab - llama_token id_last, - llama_pos draft_base_pos, - llama_seq_id draft_seq_id) { - llama_tokens result; - - spec->t_step_start_us = ggml_time_us(); - - // apply autotune proposal if enabled - if (spec->tuner && spec->tuner->enabled) { - spec->tuner->propose(params); - } - - const auto runtime_stages = params.get_resolved_stages(); - const bool use_runtime_stage_overrides = common_speculative_stage_chain_matches(runtime_stages, spec->configs); - - spec->curr_impl = nullptr; // reset current implementation - - for (size_t i = 0; i < spec->impls.size(); ++i) { - auto & impl = spec->impls[i]; - const auto & runtime_stage = use_runtime_stage_overrides ? runtime_stages[i] : spec->configs[i].stage; - common_params_speculative impl_params = common_speculative_get_runtime_params(spec->configs[i], params, runtime_stage); - result.clear(); - - { - common_time_meas tm(impl->t_draft_us, !impl->gen_perf); - impl->draft(impl_params, prompt_tgt, id_last, draft_base_pos, draft_seq_id, result); - impl->n_call_draft++; - } - - if (result.empty()) { - continue; - } - - if (common_speculative_type_is_self_spec(impl->type) && impl_params.n_min > 0 && (int)result.size() < impl_params.n_min) { - LOG_DBG("%s: impl %s drafted %zu tokens, below fallback threshold %d - trying next implementation\n", - __func__, common_speculative_type_to_str(impl->type).c_str(), result.size(), impl_params.n_min); - result.clear(); - continue; - } - LOG_DBG("%s: called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", __func__, - common_speculative_type_to_str(impl.get()->type).c_str(), prompt_tgt.size(), - impl.get()->n_call_draft, result.size()); - - spec->curr_impl = impl.get(); - impl->n_gen_drafts++; - impl->n_gen_tokens += result.size(); - - break; // We have a draft, so break out of the loop and return it. - } - - // store draft count for tuner feedback - if (spec->tuner && spec->tuner->enabled) { - spec->last_n_drafted = (int)result.size(); - } - - return result; -} - -void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) { - if (spec->tuner && spec->tuner->enabled && spec->t_step_start_us > 0) { - int64_t step_time_us = ggml_time_us() - spec->t_step_start_us; - double step_tps = (step_time_us > 100) - ? (n_accepted + 1.0) * 1e6 / (double)step_time_us - : 0.0; - spec->tuner->accept_feedback(n_accepted, spec->last_n_drafted, step_tps); - spec->t_step_start_us = 0; - } - - common_speculative_state * impl = spec->curr_impl; - - if (!impl) { - return; - } - - { - common_time_meas tm(impl->t_accept_us, !impl->gen_perf); - if (n_accepted > 0) { - impl->n_acc_drafts++; - impl->n_acc_tokens += n_accepted; - } - - impl->accept(n_accepted); - impl->n_call_accept++; - } - - if (impl->type != COMMON_SPECULATIVE_TYPE_MTP) { - if (auto * mtp_state = common_speculative_get_mtp_state(spec); mtp_state != nullptr) { - mtp_invalidate_cached_drafts(*mtp_state); - } - } -} - -static bool common_speculative_has_type(const common_speculative * spec, common_speculative_type type) { - if (spec == nullptr) { - return false; - } - - return std::any_of(spec->configs.begin(), spec->configs.end(), [type](const common_speculative_config & config) { - return config.type == type; - }); -} - -static int common_speculative_ctx_mtp_n_embd(llama_context * ctx) { - return ctx ? (int) llama_mtp_state_n_embd(ctx) : 0; -} - -static bool common_speculative_batch_token_has_seq_id( - const llama_batch & batch, - int token_index, - llama_seq_id seq_id) { - if (batch.n_seq_id == nullptr || batch.seq_id == nullptr || batch.n_seq_id[token_index] <= 0 || batch.seq_id[token_index] == nullptr) { - return false; - } - - for (int i = 0; i < batch.n_seq_id[token_index]; ++i) { - if (batch.seq_id[token_index][i] == seq_id) { - return true; - } - } - - return false; -} - -static bool common_speculative_batch_is_exact_single_seq( - const llama_batch & batch, - llama_seq_id seq_id) { - if (batch.n_tokens <= 0 || batch.n_seq_id == nullptr || batch.seq_id == nullptr) { - return false; - } - - for (int i = 0; i < batch.n_tokens; ++i) { - if (batch.n_seq_id[i] != 1 || batch.seq_id[i] == nullptr || batch.seq_id[i][0] != seq_id) { - return false; - } - } - - return true; -} - -static int common_speculative_copy_seq_batch( - const llama_batch & batch, - llama_seq_id seq_id, - llama_batch & seq_batch) { - if (batch.token == nullptr || batch.pos == nullptr) { - return -1; - } - - if (batch.n_tokens < 1) { - return 0; - } - - std::vector token_indices; - token_indices.reserve(batch.n_tokens); - for (int i = 0; i < batch.n_tokens; ++i) { - if (common_speculative_batch_token_has_seq_id(batch, i, seq_id)) { - token_indices.push_back(i); - } - } - - if (token_indices.empty()) { - return 0; - } - - seq_batch = llama_batch_init((int) token_indices.size(), 0, 1); - for (const int i : token_indices) { - common_batch_add(seq_batch, batch.token[i], batch.pos[i], { seq_id }, batch.logits != nullptr && batch.logits[i]); - } - - return (int) token_indices.size(); -} - -static bool common_speculative_feature_view_copy_batch_rows( - const common_speculative_feature_view & view, - const llama_batch & batch, - llama_seq_id seq_id, - std::vector * hidden_rows) { - if (hidden_rows == nullptr || view.kind != COMMON_SPECULATIVE_FEATURE_HIDDEN_STATE || view.width <= 0 || batch.n_tokens <= 0 || batch.pos == nullptr) { - return false; - } - - std::unordered_map rows_by_pos; - rows_by_pos.reserve(view.rows.size()); - for (const auto & row : view.rows) { - if (row.seq_id == seq_id && row.data != nullptr) { - rows_by_pos[row.pos] = row.data; - } - } - - hidden_rows->clear(); - hidden_rows->reserve((size_t) batch.n_tokens * view.width); - for (int i = 0; i < batch.n_tokens; ++i) { - auto it = rows_by_pos.find(batch.pos[i]); - if (it == rows_by_pos.end()) { - hidden_rows->clear(); - return false; - } - - hidden_rows->insert(hidden_rows->end(), it->second, it->second + view.width); - } - - return hidden_rows->size() == (size_t) batch.n_tokens * view.width; -} +#include "speculative-impl.h" static bool common_speculative_capture_target_features( common_speculative * spec, @@ -3009,6 +1271,7 @@ static void mtp_clear_target_hidden(common_speculative_state_mtp & state, llama_ state.draft_cache_by_seq.erase(seq_id); } +// DFlash target-window replay and maintenance helpers. struct dflash_append_breakdown { uint64_t filter_us = 0; uint64_t window_alloc_us = 0; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 035dd8e6..87f375bd 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -41,6 +41,8 @@ add_library(llama ../include/llama.h llama.cpp llama-spec-features.cpp + llama-spec-features-dflash.cpp + llama-dflash.cpp llama-vocab.cpp llama-grammar.cpp llama-sampling.cpp diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index eff7d675..7f3e4f33 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -116,9 +116,9 @@ void llm_build_context::init() { 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; + lctx.dflash.inputs.target_features = nullptr; + lctx.dflash.inputs.pos_ctx = nullptr; + lctx.dflash.inputs.kq_mask = nullptr; } } @@ -2195,7 +2195,7 @@ struct ggml_cgraph * llm_build_context::llama_build_graph_dflash_kv_cache(llama_ } }; - struct llm_build_context llm(lctx, dummy, cb, false, false, 0, false, &lctx.dflash_buf_compute_meta); + struct llm_build_context llm(lctx, dummy, cb, false, false, 0, false, &lctx.dflash.kv.cache_compute_meta); llm.init(); @@ -2232,7 +2232,7 @@ struct ggml_cgraph * llm_build_context::llama_build_graph_dflash_kv_workspace(ll } }; - struct llm_build_context llm(lctx, dummy, cb, false, false, 0, false, &lctx.dflash_workspace_buf_compute_meta); + struct llm_build_context llm(lctx, dummy, cb, false, false, 0, false, &lctx.dflash.kv.workspace_compute_meta); llm.init(); diff --git a/src/llama-context.h b/src/llama-context.h index 8ad9d74b..ebd4ded3 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -278,77 +278,162 @@ struct llama_context { size_t draft_input_hidden_state_n_floats = 0; std::vector draft_input_hidden_state_owned; - const float * dflash_target_features = nullptr; - size_t dflash_target_features_n_floats = 0; - int32_t dflash_target_features_n_rows = 0; - const float * dflash_target_append_features = nullptr; - size_t dflash_target_append_features_n_floats = 0; - int32_t dflash_target_append_features_n_rows = 0; - const llama_pos * dflash_target_positions = nullptr; - size_t dflash_target_positions_n = 0; - uint64_t dflash_target_window_version = 0; - int32_t dflash_target_window_keep_rows = 0; - int32_t dflash_target_window_append_rows = 0; - bool dflash_target_window_replace = false; - std::vector dflash_target_features_owned; - std::vector dflash_target_append_features_owned; - std::vector dflash_target_positions_owned; - std::vector dflash_target_features_padded; - std::vector dflash_feature_view_buffer; - std::vector dflash_pos_ctx_data; - std::vector dflash_kq_mask_data; - std::vector dflash_kq_mask_swa_data; - 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; - int32_t dflash_kv_cache_n_filled = 0; - int32_t dflash_kv_cache_update_rows = 0; - int32_t dflash_kv_cache_reserved_rows = 0; - int32_t dflash_kv_cache_view_write_pos = 0; - int32_t dflash_kv_cache_view_n_filled = 0; - 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; - struct ggml_tensor * dflash_kq_mask_swa_tensor = nullptr; + struct dflash_runtime { + struct target_window_state { + const float * features = nullptr; + size_t features_n_floats = 0; + int32_t features_n_rows = 0; + const float * append_features = nullptr; + size_t append_features_n_floats = 0; + int32_t append_features_n_rows = 0; + const llama_pos * positions = nullptr; + size_t positions_n = 0; + uint64_t version = 0; + int32_t keep_rows = 0; + int32_t append_rows = 0; + bool replace = false; + std::vector features_owned; + std::vector append_features_owned; + std::vector positions_owned; + std::vector features_padded; + std::vector pos_ctx_data; + std::vector kq_mask_data; + std::vector kq_mask_swa_data; + }; - struct dflash_capture_state { - std::vector layer_ids; - std::vector> layer_rows; - int32_t row_count = 0; - int32_t row_width = 0; - uint64_t capture_batch_id = 0; - std::vector layer_seen_batch_id; - ggml_backend_sched_eval_callback prev_cb_eval = nullptr; - void * prev_cb_eval_user_data = nullptr; + struct kv_runtime_state { + std::vector k_ctx_cache; + std::vector v_ctx_cache; + std::vector k_ctx_workspace; + std::vector v_ctx_workspace; + struct ggml_context * cache_ctx = nullptr; + std::vector cache_bufs; + int32_t cache_write_pos = 0; + int32_t cache_n_filled = 0; + int32_t cache_update_rows = 0; + int32_t cache_reserved_rows = 0; + int32_t cache_view_write_pos = 0; + int32_t cache_view_n_filled = 0; + uint64_t cache_applied_window_version = 0; + bool cache_valid = false; + bool cache_view_valid = false; + int32_t workspace_write_pos = 0; + int32_t workspace_n_filled = 0; + int32_t workspace_reserved_rows = 0; + int32_t workspace_token_capacity = 0; + int32_t workspace_n_kv_total = 0; + uint64_t workspace_applied_window_version = 0; + bool workspace_valid = false; + bool workspace_sync_pending = false; + std::vector cache_compute_meta; + std::vector workspace_compute_meta; + ggml_backend_sched_t cache_sched = nullptr; + ggml_backend_sched_t workspace_sched = nullptr; + ggml_cgraph * cache_graph = nullptr; + ggml_cgraph * workspace_graph = nullptr; + int32_t cache_graph_rows = 0; + int32_t cache_graph_write_pos = 0; + int32_t workspace_graph_rows = 0; + int32_t workspace_graph_write_pos = 0; + struct ggml_tensor * cache_input_target_features = nullptr; + struct ggml_tensor * cache_input_pos_ctx = nullptr; + struct ggml_tensor * kq_mask_tensor = nullptr; + struct ggml_tensor * kq_mask_swa_tensor = nullptr; + }; + + struct capture_state { + std::vector layer_ids; + std::vector> layer_rows; + int32_t row_count = 0; + int32_t row_width = 0; + uint64_t capture_batch_id = 0; + std::vector layer_seen_batch_id; + ggml_backend_sched_eval_callback prev_cb_eval = nullptr; + void * prev_cb_eval_user_data = nullptr; + }; + + struct input_state { + struct ggml_tensor * target_features = nullptr; // F32 [n_target_features, cross_ctx] + struct ggml_tensor * pos_ctx = nullptr; // I32 [cross_ctx] + struct ggml_tensor * kq_mask = nullptr; // F32 [cross_ctx + n_batch, GGML_PAD(n_batch)] + struct ggml_tensor * kq_mask_swa = nullptr; // F32 [cross_ctx + n_batch, GGML_PAD(n_batch)] + }; + + target_window_state target; + kv_runtime_state kv; + std::unique_ptr capture; + std::vector feature_view_buffer; + input_state inputs; + int32_t visible_cross_ctx = 0; + llama_dflash_profile_stats profile; }; - std::unique_ptr dflash_capture; - llama_dflash_profile_stats dflash_profile; + dflash_runtime dflash; + using dflash_capture_state = dflash_runtime::capture_state; + + const float * & dflash_target_features = dflash.target.features; + size_t & dflash_target_features_n_floats = dflash.target.features_n_floats; + int32_t & dflash_target_features_n_rows = dflash.target.features_n_rows; + const float * & dflash_target_append_features = dflash.target.append_features; + size_t & dflash_target_append_features_n_floats = dflash.target.append_features_n_floats; + int32_t & dflash_target_append_features_n_rows = dflash.target.append_features_n_rows; + const llama_pos * & dflash_target_positions = dflash.target.positions; + size_t & dflash_target_positions_n = dflash.target.positions_n; + uint64_t & dflash_target_window_version = dflash.target.version; + int32_t & dflash_target_window_keep_rows = dflash.target.keep_rows; + int32_t & dflash_target_window_append_rows = dflash.target.append_rows; + bool & dflash_target_window_replace = dflash.target.replace; + std::vector & dflash_target_features_owned = dflash.target.features_owned; + std::vector & dflash_target_append_features_owned = dflash.target.append_features_owned; + std::vector & dflash_target_positions_owned = dflash.target.positions_owned; + std::vector & dflash_target_features_padded = dflash.target.features_padded; + std::vector & dflash_feature_view_buffer = dflash.feature_view_buffer; + std::vector & dflash_pos_ctx_data = dflash.target.pos_ctx_data; + std::vector & dflash_kq_mask_data = dflash.target.kq_mask_data; + std::vector & dflash_kq_mask_swa_data = dflash.target.kq_mask_swa_data; + int32_t & dflash_visible_cross_ctx = dflash.visible_cross_ctx; + std::vector & dflash_k_ctx_cache = dflash.kv.k_ctx_cache; + std::vector & dflash_v_ctx_cache = dflash.kv.v_ctx_cache; + std::vector & dflash_k_ctx_workspace = dflash.kv.k_ctx_workspace; + std::vector & dflash_v_ctx_workspace = dflash.kv.v_ctx_workspace; + struct ggml_context * & dflash_cache_ctx = dflash.kv.cache_ctx; + std::vector & dflash_cache_bufs = dflash.kv.cache_bufs; + int32_t & dflash_kv_cache_write_pos = dflash.kv.cache_write_pos; + int32_t & dflash_kv_cache_n_filled = dflash.kv.cache_n_filled; + int32_t & dflash_kv_cache_update_rows = dflash.kv.cache_update_rows; + int32_t & dflash_kv_cache_reserved_rows = dflash.kv.cache_reserved_rows; + int32_t & dflash_kv_cache_view_write_pos = dflash.kv.cache_view_write_pos; + int32_t & dflash_kv_cache_view_n_filled = dflash.kv.cache_view_n_filled; + uint64_t & dflash_kv_cache_applied_window_version = dflash.kv.cache_applied_window_version; + bool & dflash_kv_cache_valid = dflash.kv.cache_valid; + bool & dflash_kv_cache_view_valid = dflash.kv.cache_view_valid; + int32_t & dflash_kv_workspace_write_pos = dflash.kv.workspace_write_pos; + int32_t & dflash_kv_workspace_n_filled = dflash.kv.workspace_n_filled; + int32_t & dflash_kv_workspace_reserved_rows = dflash.kv.workspace_reserved_rows; + int32_t & dflash_kv_workspace_token_capacity = dflash.kv.workspace_token_capacity; + int32_t & dflash_kv_workspace_n_kv_total = dflash.kv.workspace_n_kv_total; + uint64_t & dflash_kv_workspace_applied_window_version = dflash.kv.workspace_applied_window_version; + bool & dflash_kv_workspace_valid = dflash.kv.workspace_valid; + bool & dflash_kv_workspace_sync_pending = dflash.kv.workspace_sync_pending; + std::vector & dflash_buf_compute_meta = dflash.kv.cache_compute_meta; + std::vector & dflash_workspace_buf_compute_meta = dflash.kv.workspace_compute_meta; + ggml_backend_sched_t & dflash_sched = dflash.kv.cache_sched; + ggml_backend_sched_t & dflash_workspace_sched = dflash.kv.workspace_sched; + ggml_cgraph * & dflash_kv_graph = dflash.kv.cache_graph; + ggml_cgraph * & dflash_kv_workspace_graph = dflash.kv.workspace_graph; + int32_t & dflash_kv_graph_rows = dflash.kv.cache_graph_rows; + int32_t & dflash_kv_graph_write_pos = dflash.kv.cache_graph_write_pos; + int32_t & dflash_kv_workspace_graph_rows = dflash.kv.workspace_graph_rows; + int32_t & dflash_kv_workspace_graph_write_pos = dflash.kv.workspace_graph_write_pos; + struct ggml_tensor * & dflash_kv_input_target_features = dflash.kv.cache_input_target_features; + struct ggml_tensor * & dflash_kv_input_pos_ctx = dflash.kv.cache_input_pos_ctx; + struct ggml_tensor * & dflash_kq_mask_tensor = dflash.kv.kq_mask_tensor; + struct ggml_tensor * & dflash_kq_mask_swa_tensor = dflash.kv.kq_mask_swa_tensor; + std::unique_ptr & dflash_capture = dflash.capture; + llama_dflash_profile_stats & dflash_profile = dflash.profile; + struct ggml_tensor * & inp_dflash_target_features = dflash.inputs.target_features; + struct ggml_tensor * & inp_dflash_pos_ctx = dflash.inputs.pos_ctx; + struct ggml_tensor * & inp_dflash_kq_mask = dflash.inputs.kq_mask; + struct ggml_tensor * & inp_dflash_kq_mask_swa = dflash.inputs.kq_mask_swa; // input tensors struct ggml_tensor * inp_tokens; // I32 [n_batch] @@ -369,10 +454,6 @@ struct llama_context { struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch] struct ggml_tensor * inp_scale = nullptr; // F32 [n_tokens] struct ggml_tensor * inp_mtp_states = nullptr; - struct ggml_tensor * inp_dflash_target_features = nullptr; // F32 [n_target_features, cross_ctx] - struct ggml_tensor * inp_dflash_pos_ctx = nullptr; // I32 [cross_ctx] - struct ggml_tensor * inp_dflash_kq_mask = nullptr; // F32 [cross_ctx + n_batch, GGML_PAD(n_batch)] - struct ggml_tensor * inp_dflash_kq_mask_swa = nullptr; // F32 [cross_ctx + n_batch, GGML_PAD(n_batch)] ggml_backend_t ggml_backend_by_name(const char * name); diff --git a/src/llama-dflash.cpp b/src/llama-dflash.cpp new file mode 100644 index 00000000..aed84a25 --- /dev/null +++ b/src/llama-dflash.cpp @@ -0,0 +1,1240 @@ +#include "llama-dflash.h" + +#include "llama-impl.h" +#include "llama-build-context.h" +#include "llama-context.h" +#include "llama-model.h" +#include "llama-spec-features.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include + +static bool llama_env_flag_enabled_local(const char * name) { + const char * env = std::getenv(name); + return env != nullptr && *env != '\0' && + std::strcmp(env, "0") != 0 && + std::strcmp(env, "false") != 0 && + std::strcmp(env, "off") != 0; +} + +enum llama_dflash_kv_node_kind { + LLAMA_DFLASH_KV_NODE_NONE = 0, + LLAMA_DFLASH_KV_NODE_FUSED_TARGET, + LLAMA_DFLASH_KV_NODE_K_PROJ, + LLAMA_DFLASH_KV_NODE_K_NORM, + LLAMA_DFLASH_KV_NODE_K_ROPE, + LLAMA_DFLASH_KV_NODE_V_PROJ, + LLAMA_DFLASH_KV_NODE_K_STORE, + 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; + } + + 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; + } + if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_proj")) { + return LLAMA_DFLASH_KV_NODE_K_PROJ; + } + if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_norm")) { + return LLAMA_DFLASH_KV_NODE_K_NORM; + } + if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_rope")) { + return LLAMA_DFLASH_KV_NODE_K_ROPE; + } + if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_v_proj")) { + return LLAMA_DFLASH_KV_NODE_V_PROJ; + } + if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_store")) { + return LLAMA_DFLASH_KV_NODE_K_STORE; + } + if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_v_store")) { + return LLAMA_DFLASH_KV_NODE_V_STORE; + } + + return LLAMA_DFLASH_KV_NODE_NONE; +} + +static void llama_dflash_kv_node_profile_add( + llama_dflash_profile_stats & profile, + llama_dflash_kv_node_kind kind, + uint64_t elapsed_us) { + switch (kind) { + case LLAMA_DFLASH_KV_NODE_FUSED_TARGET: + profile.graph_kv_node_fused_target_calls++; + profile.graph_kv_node_fused_target_us += elapsed_us; + break; + case LLAMA_DFLASH_KV_NODE_K_PROJ: + profile.graph_kv_node_k_proj_calls++; + profile.graph_kv_node_k_proj_us += elapsed_us; + break; + case LLAMA_DFLASH_KV_NODE_K_NORM: + profile.graph_kv_node_k_norm_calls++; + profile.graph_kv_node_k_norm_us += elapsed_us; + break; + case LLAMA_DFLASH_KV_NODE_K_ROPE: + profile.graph_kv_node_k_rope_calls++; + profile.graph_kv_node_k_rope_us += elapsed_us; + break; + case LLAMA_DFLASH_KV_NODE_V_PROJ: + profile.graph_kv_node_v_proj_calls++; + profile.graph_kv_node_v_proj_us += elapsed_us; + break; + case LLAMA_DFLASH_KV_NODE_K_STORE: + profile.graph_kv_node_k_store_calls++; + profile.graph_kv_node_k_store_us += elapsed_us; + break; + case LLAMA_DFLASH_KV_NODE_V_STORE: + profile.graph_kv_node_v_store_calls++; + profile.graph_kv_node_v_store_us += elapsed_us; + break; + case LLAMA_DFLASH_KV_NODE_NONE: + break; + } +} + +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) { + return false; + } + + const llama_dflash_kv_node_kind kind = llama_dflash_kv_node_kind_from_tensor(tensor); + if (ask) { + if (kind == LLAMA_DFLASH_KV_NODE_NONE) { + return false; + } + + profiler->active_kind = kind; + profiler->t_start_us = ggml_time_us(); + return true; + } + + if (kind != LLAMA_DFLASH_KV_NODE_NONE && profiler->active_kind == kind && profiler->t_start_us > 0) { + llama_dflash_kv_node_profile_add(*profiler->profile, kind, (uint64_t) (ggml_time_us() - profiler->t_start_us)); + } + + profiler->active_kind = LLAMA_DFLASH_KV_NODE_NONE; + profiler->t_start_us = 0; + 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_local("IK_DFLASH_KV_WORKSPACE"); +} + +void llama_sync_dflash_workspace_if_pending(struct llama_context & lctx) { + if (!lctx.dflash_kv_workspace_sync_pending || lctx.dflash_workspace_sched == nullptr) { + return; + } + + 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; +} + +static ggml_backend_buffer_type_t llama_dflash_kv_cache_layer_buft(const llama_context & lctx, int32_t il) { + if (il >= 0 && (size_t) il < lctx.model.buft_layer.size() && lctx.model.buft_layer[(size_t) il].buft != nullptr) { + return lctx.model.buft_layer[(size_t) il].buft; + } + + if (il >= 0 && (size_t) il < lctx.model.layers.size()) { + const ggml_tensor * wk = lctx.model.layers[(size_t) il].wk; + if (wk != nullptr && wk->buffer != nullptr) { + return ggml_backend_buffer_get_type(wk->buffer); + } + } + + return llama_default_buffer_type_cpu(true); +} + +static ggml_backend_t llama_backend_for_tensor(const llama_context & lctx, const ggml_tensor * tensor) { + if (tensor == nullptr) { + return nullptr; + } + + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (buf == nullptr) { + return nullptr; + } + + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf); + for (ggml_backend_t backend : lctx.backends) { + ggml_backend_buffer_type_t backend_buft = ggml_backend_is_cpu(backend) + ? llama_default_buffer_type_cpu(true) + : ggml_backend_get_default_buffer_type(backend); + if (backend_buft == buft) { + return backend; + } + } + + return nullptr; +} + +bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) { + const bool use_kv_workspace = llama_env_flag_enabled_local("IK_DFLASH_KV_WORKSPACE"); + const int32_t target_cross_ctx = std::max(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()) { + 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; + } + + free_dflash_kv_cache_tensors(); + if (dflash_sched != nullptr) { + 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) ((use_kv_workspace ? 4 : 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); + 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) * (use_kv_workspace ? 4 : 2)); + int32_t host_layers = 0; + const char * first_buft_name = nullptr; + const char * last_buft_name = nullptr; + for (int32_t il = 0; il < n_layer; ++il) { + ggml_backend_buffer_type_t layer_buft = llama_dflash_kv_cache_layer_buft(*this, il); + if (ggml_backend_buft_is_host(layer_buft)) { + host_layers++; + } + if (first_buft_name == nullptr) { + first_buft_name = ggml_backend_buft_name(layer_buft); + } + last_buft_name = ggml_backend_buft_name(layer_buft); + + 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); + + const size_t k_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_k_ctx_cache[(size_t) il]); + ggml_backend_buffer_t k_buf = ggml_backend_buft_alloc_buffer(layer_buft, k_bytes); + if (k_buf == nullptr) { + free_dflash_kv_cache_tensors(); + return false; + } + ggml_backend_buffer_set_usage(k_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + ggml_backend_tensor_alloc(k_buf, dflash_k_ctx_cache[(size_t) il], ggml_backend_buffer_get_base(k_buf)); + ggml_backend_buffer_clear(k_buf, 0); + dflash_cache_bufs.push_back(k_buf); + + const size_t v_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_v_ctx_cache[(size_t) il]); + ggml_backend_buffer_t v_buf = ggml_backend_buft_alloc_buffer(layer_buft, v_bytes); + if (v_buf == nullptr) { + free_dflash_kv_cache_tensors(); + return false; + } + ggml_backend_buffer_set_usage(v_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + 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__, + target_cross_ctx, + host_layers, + n_layer, + first_buft_name != nullptr ? first_buft_name : "(none)", + last_buft_name != nullptr ? last_buft_name : "(none)"); + + return true; +} + +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; + dflash_kv_cache_reserved_rows = 0; + dflash_kv_cache_view_write_pos = 0; + dflash_kv_cache_view_n_filled = 0; + 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); + } + } + dflash_cache_bufs.clear(); + if (dflash_cache_ctx != nullptr) { + ggml_free(dflash_cache_ctx); + dflash_cache_ctx = nullptr; + } +} + +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 dflash_layer_has_attention_bias(const llama_layer & layer) { + return layer.bq != nullptr || + layer.bk != nullptr || + layer.bv != nullptr || + layer.bo != nullptr || + layer.bqkv != nullptr || + layer.bqk != nullptr || + layer.bkv != nullptr; +} + +static bool validate_dflash_graph_contract(const llama_context & lctx) { + const auto & model = lctx.model; + const auto & hparams = model.hparams; + + auto rope_dim_for_layer = [&hparams](int32_t il) -> uint32_t { + if (hparams.rope_dim_per_layer[(size_t) il] != 0) { + return hparams.rope_dim_per_layer[(size_t) il]; + } + + return hparams.swa_layers[(size_t) il] ? hparams.n_rot_swa : hparams.n_rot; + }; + + auto rope_base_for_layer = [&hparams](int32_t il) -> float { + if (hparams.has_rope_freq_base_per_layer) { + return hparams.rope_freq_base_per_layer[(size_t) il]; + } + + return hparams.swa_layers[(size_t) il] ? hparams.rope_freq_base_train_swa : hparams.rope_freq_base_train; + }; + + auto rope_scale_for_layer = [&hparams](int32_t il) -> float { + return hparams.swa_layers[(size_t) il] ? hparams.rope_freq_scale_train_swa : hparams.rope_freq_scale_train; + }; + + const uint32_t ref_n_head = hparams.n_head(0); + const uint32_t ref_n_head_kv = hparams.n_head_kv(0); + const uint32_t ref_n_embd_head_k = hparams.n_embd_head_k(0); + const uint32_t ref_n_embd_head_v = hparams.n_embd_head_v(0); + const uint32_t ref_rope_dim = rope_dim_for_layer(0); + const float ref_rope_base = rope_base_for_layer(0); + const float ref_rope_scale = rope_scale_for_layer(0); + + for (int32_t il = 0; il < (int32_t) hparams.n_layer; ++il) { + if (hparams.n_head((uint32_t) il) != ref_n_head || + hparams.n_head_kv((uint32_t) il) != ref_n_head_kv || + hparams.n_embd_head_k(il) != ref_n_embd_head_k || + hparams.n_embd_head_v(il) != ref_n_embd_head_v) { + LLAMA_LOG_ERROR("%s: DFlash graph assumes layer-invariant head config, but layer %d differs (n_head=%u/%u n_head_kv=%u/%u head_k=%u/%u head_v=%u/%u)\n", + __func__, + il, + hparams.n_head((uint32_t) il), ref_n_head, + hparams.n_head_kv((uint32_t) il), ref_n_head_kv, + hparams.n_embd_head_k(il), ref_n_embd_head_k, + hparams.n_embd_head_v(il), ref_n_embd_head_v); + return false; + } + + const uint32_t rope_dim = rope_dim_for_layer(il); + const float rope_base = rope_base_for_layer(il); + const float rope_scale = rope_scale_for_layer(il); + if (rope_dim != ref_rope_dim || std::fabs(rope_base - ref_rope_base) > 1e-6f || std::fabs(rope_scale - ref_rope_scale) > 1e-6f) { + LLAMA_LOG_ERROR("%s: DFlash graph assumes layer-invariant RoPE config, but layer %d differs (dim=%u/%u base=%g/%g scale=%g/%g)\n", + __func__, + il, + rope_dim, ref_rope_dim, + (double) rope_base, (double) ref_rope_base, + (double) rope_scale, (double) ref_rope_scale); + return false; + } + + if (model.layers[(size_t) il].attn_norm == nullptr || + model.layers[(size_t) il].attn_q_norm == nullptr || + model.layers[(size_t) il].attn_k_norm == nullptr) { + LLAMA_LOG_ERROR("%s: DFlash graph requires attn_norm, attn_q_norm, and attn_k_norm weights, but layer %d is missing one or more of them\n", + __func__, il); + return false; + } + + const bool has_q_norm = model.layers[(size_t) il].attn_q_norm != nullptr; + const bool has_k_norm = model.layers[(size_t) il].attn_k_norm != nullptr; + if (has_q_norm != has_k_norm) { + LLAMA_LOG_ERROR("%s: DFlash graph requires symmetric Q/K norm presence, but layer %d has q_norm=%d k_norm=%d\n", + __func__, il, (int) has_q_norm, (int) has_k_norm); + return false; + } + + if (model.layers[(size_t) il].attn_norm_b != nullptr || + model.layers[(size_t) il].attn_q_norm_b != nullptr || + model.layers[(size_t) il].attn_k_norm_b != nullptr) { + LLAMA_LOG_ERROR("%s: DFlash graph does not implement norm-bias tensors, but layer %d requires attn_norm_b/q_norm_b/k_norm_b\n", + __func__, il); + return false; + } + + if (dflash_layer_has_attention_bias(model.layers[(size_t) il])) { + LLAMA_LOG_ERROR("%s: DFlash graph does not implement attention bias tensors, but layer %d requires them\n", + __func__, il); + return false; + } + } + + return true; +} + +bool llama_prepare_dflash_graph_inputs( + struct llama_context & lctx, + uint32_t n_tokens) { + const bool use_kv_cache = llama_env_flag_enabled_local("IK_DFLASH_KV_CACHE"); + const bool use_kv_workspace = use_kv_cache && llama_dflash_use_kv_workspace_experiment(); + const bool kv_node_timing = llama_env_flag_enabled_local("IK_DFLASH_KV_NODE_TIMING"); + auto & profile = lctx.dflash_profile; + const int32_t cross_ctx = lctx.dflash_visible_cross_ctx > 0 + ? 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; + ggml_tensor * kq_mask_swa = lctx.dflash_kq_mask_swa_tensor; + + if (kq_mask == nullptr) { + LLAMA_LOG_ERROR("%s: DFlash graph inputs are not initialized\n", __func__); + return false; + } + + if (!validate_dflash_graph_contract(lctx)) { + profile.graph_shape_failures++; + 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 float * append_src = lctx.dflash_target_append_features; + const llama_pos * src_pos = lctx.dflash_target_positions; + const size_t total_floats = lctx.dflash_target_features_n_floats; + const size_t append_floats = lctx.dflash_target_append_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 append_rows_available = lctx.dflash_target_append_features_n_rows; + const int32_t width = (int32_t) lctx.model.hparams.dflash_n_target_features; + const int32_t graph_cross_ctx = use_kv_cache + ? (lctx.dflash_k_ctx_cache.front() != nullptr ? (int32_t) lctx.dflash_k_ctx_cache.front()->ne[2] : 0) + : (lctx.inp_dflash_target_features != nullptr ? (int32_t) lctx.inp_dflash_target_features->ne[1] : 0); + const int32_t n_mask_tokens = (int32_t) kq_mask->ne[1]; + const int32_t n_kv_total = (int32_t) kq_mask->ne[0]; + const int64_t t_total_us = ggml_time_us(); + + 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 (use_kv_workspace) { + llama_sync_dflash_workspace_if_pending(lctx); + } + + 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 (n_rows <= 0) { + profile.graph_shape_failures++; + LLAMA_LOG_ERROR("%s: missing DFlash target feature rows\n", __func__); + return false; + } + + const bool have_full_src = src != nullptr && total_floats == (size_t) n_rows * (size_t) width; + if (n_rows > cross_ctx || (src != nullptr && !have_full_src)) { + 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; + } + + if (!use_kv_cache && !have_full_src) { + profile.graph_shape_failures++; + LLAMA_LOG_ERROR("%s: missing contiguous DFlash target features for inline path\n", __func__); + return false; + } + + if (n_kv_total < cross_ctx + (int32_t) n_tokens) { + profile.graph_mask_overflow++; + LLAMA_LOG_ERROR("%s: invalid DFlash mask shape (n_kv_total=%d < cross_ctx+n_tokens=%d)\n", + __func__, n_kv_total, cross_ctx + (int32_t) n_tokens); + return false; + } + + const int32_t left_pad = cross_ctx - n_rows; + profile.last_left_pad = left_pad; + if (!use_kv_cache) { + const size_t padded_floats = (size_t) cross_ctx * (size_t) width; + const size_t dst_offset = (size_t) left_pad * (size_t) width; + const int64_t t_feature_us = ggml_time_us(); + if (lctx.dflash_target_features_padded.size() != padded_floats) { + lctx.dflash_target_features_padded.resize(padded_floats); + } + if (left_pad == 0 && total_floats == padded_floats) { + std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin()); + } else { + if (dst_offset > 0) { + std::fill(lctx.dflash_target_features_padded.begin(), + lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset, 0.0f); + } + std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset); + } + profile.graph_feature_copy_us += (uint64_t) (ggml_time_us() - t_feature_us); + profile.graph_feature_bytes += padded_floats * sizeof(float); + } + + const int64_t t_pos_us = ggml_time_us(); + lctx.dflash_pos_ctx_data.resize((size_t) cross_ctx); + 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) { + profile.graph_pos_fallbacks++; + profile.graph_shape_failures++; + profile.last_pos_first = -1; + profile.last_pos_last = -1; + if (profile.graph_pos_fallbacks <= 3) { + LLAMA_LOG_ERROR("%s: missing DFlash target positions (rows=%d positions=%zu cross_ctx=%d)\n", + __func__, n_rows, total_positions, cross_ctx); + } + return false; + } + + profile.last_pos_first = src_pos[0]; + profile.last_pos_last = src_pos[n_rows - 1]; + for (int32_t i = 1; i < n_rows; ++i) { + if (src_pos[i] <= src_pos[i - 1]) { + profile.graph_pos_non_monotonic++; + profile.graph_shape_failures++; + if (profile.graph_pos_non_monotonic <= 3) { + LLAMA_LOG_ERROR("%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]); + } + return false; + } + } + std::copy(src_pos, src_pos + n_rows, lctx.dflash_pos_ctx_data.begin() + (ptrdiff_t) left_pad); + profile.graph_pos_copy_us += (uint64_t) (ggml_time_us() - t_pos_us); + profile.graph_pos_bytes += lctx.dflash_pos_ctx_data.size() * sizeof(llama_pos); + + if (use_kv_cache) { + const llama_dflash_kv_cache_transition cache_plan = llama_plan_dflash_kv_cache_transition( + cross_ctx, + lctx.dflash_kv_cache_n_filled, + lctx.dflash_kv_cache_write_pos, + lctx.dflash_kv_cache_valid, + lctx.dflash_kv_cache_applied_window_version, + lctx.dflash_target_window_version, + lctx.dflash_target_window_keep_rows, + lctx.dflash_target_window_append_rows, + lctx.dflash_target_window_replace, + n_rows); + + const bool have_append_src = append_src != nullptr && + append_rows_available == cache_plan.append_rows && + append_floats == (size_t) cache_plan.append_rows * (size_t) width; + + const int32_t update_rows = cache_plan.cache_up_to_date + ? 0 + : (cache_plan.rebuild_cache ? n_rows : cache_plan.append_rows); + const size_t max_nodes = lctx.model.max_nodes((int) std::max(1, cross_ctx)) + 24 * lctx.model.hparams.n_layer; + const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false); + if (lctx.dflash_buf_compute_meta.size() != meta_size) { + lctx.dflash_buf_compute_meta.resize(meta_size); + } + + if (lctx.dflash_sched == nullptr || lctx.dflash_kv_cache_reserved_rows != cross_ctx) { + std::vector backend_buft; + backend_buft.reserve(lctx.backends.size()); + for (auto * backend : lctx.backends) { + if (ggml_backend_is_cpu(backend)) { + backend_buft.push_back(llama_default_buffer_type_cpu(true)); + } else { + backend_buft.push_back(ggml_backend_get_default_buffer_type(backend)); + } + } + + if (lctx.dflash_sched != nullptr) { + ggml_backend_sched_free(lctx.dflash_sched); + lctx.dflash_sched = nullptr; + } + lctx.dflash_kv_graph = nullptr; + lctx.dflash_kv_graph_rows = 0; + lctx.dflash_kv_graph_write_pos = 0; + + const int32_t saved_update_rows = lctx.dflash_kv_cache_update_rows; + lctx.dflash_kv_cache_update_rows = cross_ctx; + const int64_t t_build_us = ggml_time_us(); + ggml_cgraph * gf_reserve = llm_build_context::llama_build_graph_dflash_kv_cache(lctx); + profile.graph_kv_cache_build_us += (uint64_t) (ggml_time_us() - t_build_us); + lctx.dflash_kv_cache_update_rows = saved_update_rows; + if (gf_reserve == nullptr) { + profile.graph_shape_failures++; + LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache reserve graph\n", __func__); + return false; + } + + const int64_t t_reserve_us = ggml_time_us(); + lctx.dflash_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false); + const bool reserved = lctx.dflash_sched != nullptr && ggml_backend_sched_reserve(lctx.dflash_sched, gf_reserve); + profile.graph_kv_cache_reserve_us += (uint64_t) (ggml_time_us() - t_reserve_us); + if (!reserved) { + profile.graph_shape_failures++; + LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V scheduler\n", __func__); + return false; + } + lctx.dflash_kv_cache_reserved_rows = cross_ctx; + } + + if (update_rows > 0) { + const float * update_src = nullptr; + if (have_append_src && update_rows == cache_plan.append_rows) { + update_src = append_src; + } else if (have_full_src) { + update_src = src + (size_t) (n_rows - update_rows) * (size_t) width; + } + const llama_pos * update_pos = src_pos + (n_rows - update_rows); + + if (update_src == nullptr) { + profile.graph_shape_failures++; + LLAMA_LOG_ERROR("%s: missing DFlash appended target features for cached update (rows=%d append_rows=%d floats=%zu)\n", + __func__, n_rows, update_rows, append_floats); + return false; + } + + if (cache_plan.rebuild_cache) { + llama_reset_dflash_kv_cache_state(&lctx); + } + + lctx.dflash_kv_cache_update_rows = update_rows; + ggml_cgraph * gf_kv = nullptr; + const bool can_reuse_kv_graph = lctx.dflash_kv_graph != nullptr && + lctx.dflash_kv_graph_rows == update_rows && + lctx.dflash_kv_graph_write_pos == lctx.dflash_kv_cache_write_pos; + if (can_reuse_kv_graph) { + gf_kv = lctx.dflash_kv_graph; + } else { + const int64_t t_build_us = ggml_time_us(); + 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; + } + + ggml_backend_t kv_feature_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_target_features); + const int64_t t_feature_upload_us = ggml_time_us(); + if (kv_feature_backend != nullptr) { + ggml_backend_tensor_set_async(kv_feature_backend, lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features)); + } else { + ggml_backend_tensor_set(lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features)); + } + profile.graph_kv_cache_feature_upload_us += (uint64_t) (ggml_time_us() - t_feature_upload_us); + profile.graph_feature_bytes += (size_t) update_rows * (size_t) width * sizeof(float); + + ggml_backend_t kv_pos_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_pos_ctx); + const int64_t t_pos_upload_us = ggml_time_us(); + if (kv_pos_backend != nullptr) { + ggml_backend_tensor_set_async(kv_pos_backend, lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx)); + } else { + ggml_backend_tensor_set(lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx)); + } + profile.graph_kv_cache_pos_upload_us += (uint64_t) (ggml_time_us() - t_pos_upload_us); + + const int64_t t_kv_cache_us = ggml_time_us(); + llama_dflash_kv_node_profiler kv_node_profiler; + if (kv_node_timing) { + kv_node_profiler.profile = &profile; + ggml_backend_sched_set_eval_callback(lctx.dflash_sched, llama_dflash_kv_node_eval_callback, &kv_node_profiler); + } + llama_graph_compute_sched(lctx, lctx.dflash_sched, gf_kv, lctx.cparams.n_threads); + if (kv_node_timing) { + ggml_backend_sched_set_eval_callback(lctx.dflash_sched, nullptr, nullptr); + } + profile.graph_kv_cache_compute_us += (uint64_t) (ggml_time_us() - t_kv_cache_us); + + const int64_t t_sync_us = ggml_time_us(); + ggml_backend_sched_synchronize(lctx.dflash_sched); + profile.graph_kv_cache_sync_us += (uint64_t) (ggml_time_us() - t_sync_us); + profile.graph_kv_cache_calls++; + + lctx.dflash_kv_cache_n_filled = std::min(cross_ctx, lctx.dflash_kv_cache_n_filled + update_rows); + lctx.dflash_kv_cache_write_pos = (lctx.dflash_kv_cache_write_pos + update_rows) % cross_ctx; + lctx.dflash_kv_cache_applied_window_version = lctx.dflash_target_window_version; + lctx.dflash_kv_cache_valid = true; + lctx.dflash_kv_cache_view_n_filled = lctx.dflash_kv_cache_n_filled; + lctx.dflash_kv_cache_view_write_pos = lctx.dflash_kv_cache_write_pos; + lctx.dflash_kv_cache_view_valid = true; + } + + if (use_kv_workspace && lctx.dflash_kv_cache_view_valid && + !lctx.dflash_k_ctx_workspace.empty() && !lctx.dflash_v_ctx_workspace.empty()) { + const bool need_workspace_refresh = !lctx.dflash_kv_workspace_valid || + lctx.dflash_kv_workspace_n_filled != lctx.dflash_kv_cache_view_n_filled || + lctx.dflash_kv_workspace_write_pos != lctx.dflash_kv_cache_view_write_pos || + lctx.dflash_kv_workspace_applied_window_version != lctx.dflash_kv_cache_applied_window_version; + + if (need_workspace_refresh) { + const size_t max_nodes = lctx.model.max_nodes((int) std::max(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)); + } + + const int64_t t_mask_us = ggml_time_us(); + const int32_t full_visible_first = left_pad; + const int32_t full_visible_last = cross_ctx + (int32_t) n_tokens - 1; + 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) n_tokens; + visible_kv_max = std::max(visible_kv_max, visible_kv); + profile.graph_visible_kv_sum += (uint64_t) visible_kv; + for (int32_t i = full_visible_first; i <= full_visible_last; ++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); + + if (kq_mask_swa != nullptr) { + lctx.dflash_kq_mask_swa_data.assign((size_t) n_kv_total * (size_t) n_mask_tokens, -INFINITY); + const int32_t swa_window = (int32_t) lctx.model.hparams.n_swa; + const int32_t draft_pos_base = (int32_t) profile.last_pos_last; + for (uint32_t j = 0; j < n_tokens; ++j) { + float * row = lctx.dflash_kq_mask_swa_data.data() + (size_t) j * (size_t) n_kv_total; + const int32_t q_pos = draft_pos_base + (int32_t) j; + + for (int32_t k = left_pad; k < cross_ctx; ++k) { + const int32_t k_pos = (int32_t) lctx.dflash_pos_ctx_data[(size_t) k]; + if (q_pos - k_pos < swa_window) { + row[k] = 0.0f; + } + } + + for (int32_t k = cross_ctx; k < cross_ctx + (int32_t) n_tokens; ++k) { + const int32_t block_k = k - cross_ctx; + if (block_k <= (int32_t) j) { + row[k] = 0.0f; + } + } + } + + ggml_backend_tensor_set(kq_mask_swa, lctx.dflash_kq_mask_swa_data.data(), 0, ggml_nbytes(kq_mask_swa)); + profile.graph_mask_bytes += ggml_nbytes(kq_mask_swa); + } + + 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) { + int32_t n_swa_layers = 0; + for (int32_t il = 0; il < lctx.model.hparams.n_layer; ++il) { + n_swa_layers += lctx.model.hparams.swa_layers[(size_t) il] ? 1 : 0; + } + + 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] full_mask=[%d..%d] swa_window=%u swa_layers=%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, + full_visible_first, full_visible_last, + lctx.model.hparams.n_swa, + n_swa_layers); + } + + return true; +} diff --git a/src/llama-dflash.h b/src/llama-dflash.h new file mode 100644 index 00000000..8280c6ca --- /dev/null +++ b/src/llama-dflash.h @@ -0,0 +1,8 @@ +#pragma once + +#include + +struct llama_context; + +bool llama_prepare_dflash_graph_inputs(llama_context & lctx, uint32_t n_tokens); +void llama_sync_dflash_workspace_if_pending(llama_context & lctx); diff --git a/src/llama-quantize.cpp b/src/llama-quantize.cpp index 1f538882..367b7225 100644 --- a/src/llama-quantize.cpp +++ b/src/llama-quantize.cpp @@ -616,7 +616,9 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n if (qs.model.hparams.n_vocab >= 127999 && (qs.model.type == MODEL_8B || qs.model.type == MODEL_70B)) new_type = GGML_TYPE_IQ6_K; } - else if (qs.model.hparams.n_gqa() >= 4) { + else if (qs.model.hparams.n_gqa() >= 4 && + !(arch == LLM_ARCH_DFLASH_DRAFT && + (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q5_K_M))) { if (new_type == GGML_TYPE_Q2_K || new_type == GGML_TYPE_IQ3_XXS) new_type = GGML_TYPE_IQ3_S; else if (new_type == GGML_TYPE_Q2_K_R4 || new_type == GGML_TYPE_IQ3_XXS_R4) new_type = GGML_TYPE_IQ3_K_R4; else if (new_type == GGML_TYPE_Q3_K || new_type == GGML_TYPE_IQ3_S) new_type = GGML_TYPE_Q4_K; @@ -1778,4 +1780,3 @@ uint32_t llama_model_quantize( return 1; } } - diff --git a/src/llama-spec-features-dflash.cpp b/src/llama-spec-features-dflash.cpp new file mode 100644 index 00000000..088f6b2d --- /dev/null +++ b/src/llama-spec-features-dflash.cpp @@ -0,0 +1,1097 @@ +#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; +} + +void llama_dflash_profile_reset(struct llama_context * ctx) { + if (ctx == nullptr) { + return; + } + + ctx->dflash.profile = {}; +} + +void llama_reset_dflash_kv_cache_state(struct llama_context * ctx) { + if (ctx == nullptr) { + return; + } + + ctx->dflash.kv.cache_write_pos = 0; + ctx->dflash.kv.cache_n_filled = 0; + ctx->dflash.kv.cache_update_rows = 0; + ctx->dflash.kv.cache_view_write_pos = 0; + ctx->dflash.kv.cache_view_n_filled = 0; + 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.kv.cache_bufs) { + if (buf != nullptr) { + ggml_backend_buffer_clear(buf, 0); + } + } +} + +llama_dflash_kv_cache_transition llama_plan_dflash_kv_cache_transition_for_ctx( + const struct llama_context * ctx, + const llama_dflash_window_update & window_update, + int32_t n_rows) { + if (ctx == nullptr) { + llama_dflash_kv_cache_transition plan; + plan.rebuild_cache = true; + plan.append_rows = std::clamp(window_update.append_rows, 0, n_rows); + plan.next_n_filled = n_rows; + return plan; + } + + const int32_t cross_ctx = ctx->dflash.visible_cross_ctx > 0 + ? ctx->dflash.visible_cross_ctx + : std::max(1, (int32_t) ctx->cparams.n_ctx - (int32_t) ctx->model.hparams.dflash_block_size); + + return llama_plan_dflash_kv_cache_transition( + cross_ctx, + ctx->dflash.kv.cache_n_filled, + ctx->dflash.kv.cache_write_pos, + ctx->dflash.kv.cache_valid, + ctx->dflash.kv.cache_applied_window_version, + window_update.version, + window_update.keep_rows, + window_update.append_rows, + window_update.replace, + n_rows); +} + +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; +} + +int32_t llama_model_dflash_mask_token_id(const struct llama_model * model) { + return model ? (int32_t) model->hparams.dflash_mask_token_id : -1; +} + +int32_t llama_model_dflash_n_target_layers(const struct llama_model * model) { + return model ? (int32_t) model->hparams.dflash_n_target_layers : 0; +} + +int32_t llama_model_dflash_n_target_features(const struct llama_model * model) { + return model ? (int32_t) model->hparams.dflash_n_target_features : 0; +} + +int32_t llama_model_dflash_target_layer_ids( + const struct llama_model * model, + int32_t * layer_ids, + int32_t capacity) { + if (model == nullptr || layer_ids == nullptr || capacity <= 0) { + return 0; + } + + const int32_t n_layers = std::min((int32_t) model->hparams.dflash_n_target_layers, capacity); + for (int32_t i = 0; i < n_layers; ++i) { + layer_ids[i] = (int32_t) model->hparams.dflash_target_layer_ids[i]; + } + + return n_layers; +} + +int32_t llama_model_dflash_target_mask_token_id(const struct llama_model * model) { + if (model == nullptr) { + return (int32_t) LLAMA_TOKEN_NULL; + } + + 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) { + if (draft_model == nullptr || target_model == nullptr || draft_model->arch != LLM_ARCH_DFLASH_DRAFT) { + return LLAMA_DFLASH_IO_MODE_INVALID; + } + + 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_output == target_output; + if (shared_tok && shared_output) { + return LLAMA_DFLASH_IO_MODE_SHARED; + } + + if (!shared_tok && !shared_output) { + return LLAMA_DFLASH_IO_MODE_SELF_CONTAINED; + } + + return LLAMA_DFLASH_IO_MODE_MIXED; +} + +bool llama_model_dflash_io_tensors_match( + const struct llama_model * draft_model, + int32_t n_embd, + int32_t n_vocab) { + 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) output->ne[0] == n_embd && + (int32_t) output->ne[1] == n_vocab; +} + +bool llama_model_share_dflash_io_tensors( + struct llama_model * draft_model, + const struct llama_model * target_model) { + if (draft_model == nullptr || target_model == nullptr) { + return false; + } + + if (draft_model->arch != LLM_ARCH_DFLASH_DRAFT) { + return true; + } + + if (draft_model->tok_embd == nullptr) { + draft_model->tok_embd = target_model->tok_embd; + } + + if (draft_model->output == nullptr) { + draft_model->output = target_model->output ? target_model->output : target_model->tok_embd; + if (draft_model->output == nullptr) { + draft_model->output = draft_model->tok_embd; + } + } + + 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; +} + +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, + bool copy_data, + const llama_dflash_window_update * window_update) { + const bool have_full_features = target_features != nullptr && n_floats > 0; + const bool have_append_features = window_update != nullptr && + window_update->append_features != nullptr && + window_update->append_floats > 0 && + window_update->append_rows > 0; + + if (ctx == nullptr || n_rows <= 0 || (!have_full_features && !have_append_features)) { + return false; + } + + auto & profile = ctx->dflash.profile; + const int64_t t_start_us = ggml_time_us(); + const int32_t row_width = have_full_features + ? (n_rows > 0 ? (int32_t) (n_floats / (size_t) n_rows) : 0) + : (window_update->append_rows > 0 ? (int32_t) (window_update->append_floats / (size_t) window_update->append_rows) : 0); + llama_pos first_pos = -1; + llama_pos last_pos = -1; + + if (have_full_features && 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 if (have_full_features) { + ctx->dflash.target.features_owned.clear(); + ctx->dflash.target.features = target_features; + } else { + ctx->dflash.target.features_owned.clear(); + ctx->dflash.target.features = nullptr; + } + ctx->dflash.target.features_n_floats = have_full_features ? n_floats : 0; + ctx->dflash.target.features_n_rows = n_rows; + if (have_append_features && copy_data) { + ctx->dflash.target.append_features_owned.assign( + window_update->append_features, + window_update->append_features + window_update->append_floats); + ctx->dflash.target.append_features = ctx->dflash.target.append_features_owned.data(); + } else if (have_append_features) { + ctx->dflash.target.append_features_owned.clear(); + ctx->dflash.target.append_features = window_update->append_features; + } else { + ctx->dflash.target.append_features_owned.clear(); + ctx->dflash.target.append_features = nullptr; + } + ctx->dflash.target.append_features_n_floats = have_append_features ? window_update->append_floats : 0; + ctx->dflash.target.append_features_n_rows = have_append_features ? window_update->append_rows : 0; + ctx->dflash.target.version = window_update != nullptr && window_update->version > 0 + ? window_update->version + : ctx->dflash.target.version + 1; + ctx->dflash.target.keep_rows = window_update != nullptr + ? std::max(0, std::min(n_rows, window_update->keep_rows)) + : 0; + ctx->dflash.target.append_rows = window_update != nullptr + ? std::max(0, std::min(n_rows, window_update->append_rows)) + : n_rows; + ctx->dflash.target.replace = window_update != nullptr + ? window_update->replace + : true; + if (ctx->dflash.target.keep_rows + ctx->dflash.target.append_rows > n_rows) { + ctx->dflash.target.keep_rows = std::max(0, n_rows - ctx->dflash.target.append_rows); + } + + const int32_t cross_ctx = ctx->dflash.visible_cross_ctx > 0 + ? ctx->dflash.visible_cross_ctx + : std::max(1, (int32_t) ctx->cparams.n_ctx - (int32_t) ctx->model.hparams.dflash_block_size); + const llama_dflash_window_update cache_window_update = { + ctx->dflash.target.version, + ctx->dflash.target.keep_rows, + ctx->dflash.target.append_rows, + ctx->dflash.target.replace, + ctx->dflash.target.append_features, + ctx->dflash.target.append_features_n_floats, + }; + const llama_dflash_kv_cache_transition cache_plan = llama_plan_dflash_kv_cache_transition_for_ctx(ctx, cache_window_update, n_rows); + + if (cache_plan.cache_up_to_date) { + ctx->dflash.kv.cache_view_n_filled = ctx->dflash.kv.cache_n_filled; + ctx->dflash.kv.cache_view_write_pos = ctx->dflash.kv.cache_write_pos; + ctx->dflash.kv.cache_view_valid = ctx->dflash.kv.cache_valid; + } else if (cross_ctx > 0) { + ctx->dflash.kv.cache_view_n_filled = cache_plan.next_n_filled; + ctx->dflash.kv.cache_view_write_pos = cache_plan.next_write_pos; + ctx->dflash.kv.cache_view_valid = cache_plan.next_n_filled > 0; + } + + if (target_positions != nullptr) { + 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 += + (have_full_features ? n_floats : 0) * sizeof(float) + + (have_append_features ? window_update->append_floats : 0) * 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, + const llama_dflash_window_update * window_update) { + return llama_set_dflash_target_features_impl(ctx, target_features, n_floats, n_rows, target_positions, true, window_update); +} + +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, + const llama_dflash_window_update * window_update) { + return llama_set_dflash_target_features_impl(ctx, target_features, n_floats, n_rows, target_positions, false, window_update); +} + +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; + } + + static constexpr const char * prefix = "l_out-"; + if (std::strncmp(tensor->name, prefix, std::strlen(prefix)) != 0) { + return false; + } + + char * end = nullptr; + const long raw = std::strtol(tensor->name + std::strlen(prefix), &end, 10); + if (end == tensor->name + std::strlen(prefix) || *end != '\0') { + return false; + } + + layer_id = (int32_t) raw; + if (layer_id >= 1000) { + layer_id %= 1000; + } + + return layer_id >= 0; +} + +static int32_t llama_dflash_find_layer_index(const struct llama_context * ctx, int32_t layer_id) { + if (ctx == nullptr || !ctx->dflash.capture) { + return -1; + } + + const auto & layer_ids = ctx->dflash.capture->layer_ids; + const auto it = std::find(layer_ids.begin(), layer_ids.end(), layer_id); + return it == layer_ids.end() ? -1 : (int32_t) std::distance(layer_ids.begin(), it); +} + +static bool llama_dflash_capture_eval_callback(struct ggml_tensor * tensor, bool ask, void * user_data) { + auto * ctx = static_cast(user_data); + if (ctx == nullptr || !ctx->dflash.capture) { + return false; + } + + int32_t layer_id = -1; + if (!llama_dflash_parse_layer_id(tensor, layer_id)) { + return false; + } + + const int32_t layer_idx = llama_dflash_find_layer_index(ctx, layer_id); + if (layer_idx < 0) { + return false; + } + + if (ask) { + return true; + } + + const int32_t row_width = (int32_t) tensor->ne[0]; + const int32_t row_count = row_width > 0 ? (int32_t) (ggml_nelements(tensor) / (int64_t) row_width) : 0; + if (row_width <= 0 || row_count <= 0) { + return false; + } + + auto & capture = *ctx->dflash.capture; + if (capture.capture_batch_id == 0) { + capture.capture_batch_id = 1; + } + if (capture.layer_seen_batch_id.size() != capture.layer_ids.size()) { + capture.layer_seen_batch_id.assign(capture.layer_ids.size(), 0); + } + + auto & rows = capture.layer_rows[(size_t) layer_idx]; + rows.resize((size_t) row_count * (size_t) row_width); + ggml_backend_tensor_get(tensor, rows.data(), 0, ggml_nbytes(tensor)); + capture.row_width = row_width; + capture.row_count = row_count; + capture.layer_seen_batch_id[(size_t) layer_idx] = capture.capture_batch_id; + return true; +} + +bool llama_set_dflash_capture_layers( + struct llama_context * ctx, + const int32_t * layer_ids, + int32_t n_layers) { + if (ctx == nullptr || layer_ids == nullptr || n_layers <= 0) { + return false; + } + + auto capture = std::make_unique(); + capture->layer_ids.assign(layer_ids, layer_ids + n_layers); + capture->layer_rows.resize((size_t) n_layers); + capture->layer_seen_batch_id.assign((size_t) n_layers, 0); + capture->prev_cb_eval = ctx->cparams.cb_eval; + capture->prev_cb_eval_user_data = ctx->cparams.cb_eval_user_data; + ctx->dflash.capture = std::move(capture); + ctx->dflash.feature_view_buffer.clear(); + + ctx->cparams.cb_eval = llama_dflash_capture_eval_callback; + ctx->cparams.cb_eval_user_data = ctx; + if (ctx->sched != nullptr) { + ggml_backend_sched_set_eval_callback(ctx->sched, ctx->cparams.cb_eval, ctx->cparams.cb_eval_user_data); + } + + return true; +} + +void llama_clear_dflash_capture(struct llama_context * ctx) { + if (ctx == nullptr) { + return; + } + + ggml_backend_sched_eval_callback prev_cb_eval = nullptr; + void * prev_cb_eval_user_data = nullptr; + if (ctx->dflash.capture) { + prev_cb_eval = ctx->dflash.capture->prev_cb_eval; + prev_cb_eval_user_data = ctx->dflash.capture->prev_cb_eval_user_data; + } + + ctx->dflash.capture.reset(); + ctx->dflash.feature_view_buffer.clear(); + + if (ctx->cparams.cb_eval == llama_dflash_capture_eval_callback && ctx->cparams.cb_eval_user_data == ctx) { + ctx->cparams.cb_eval = prev_cb_eval; + ctx->cparams.cb_eval_user_data = prev_cb_eval_user_data; + if (ctx->sched != nullptr) { + ggml_backend_sched_set_eval_callback(ctx->sched, prev_cb_eval, prev_cb_eval_user_data); + } + } +} + +void llama_begin_dflash_capture_batch(struct llama_context * ctx) { + if (ctx == nullptr || !ctx->dflash.capture) { + return; + } + + auto & capture = *ctx->dflash.capture; + capture.capture_batch_id++; + capture.row_count = 0; + capture.row_width = 0; + std::fill(capture.layer_seen_batch_id.begin(), capture.layer_seen_batch_id.end(), 0); +} + +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, + int32_t & row_width, + int32_t & n_layers) { + if (ctx == nullptr || !ctx->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; + } + + if (capture.capture_batch_id == 0 || capture.layer_seen_batch_id.size() != (size_t) n_layers) { + profile.capture_prepare_failures++; + profile.capture_layer_batch_mismatch++; + if (profile.capture_layer_batch_mismatch <= 3) { + LLAMA_LOG_WARN("%s: DFlash capture batch markers are not initialized (batch_id=%llu layers=%zu expected=%d)\n", + __func__, + (unsigned long long) capture.capture_batch_id, + capture.layer_seen_batch_id.size(), + n_layers); + } + return false; + } + + for (int32_t layer_idx = 0; layer_idx < n_layers; ++layer_idx) { + if (capture.layer_seen_batch_id[(size_t) layer_idx] != capture.capture_batch_id) { + profile.capture_prepare_failures++; + profile.capture_layer_batch_mismatch++; + if (profile.capture_layer_batch_mismatch <= 3) { + LLAMA_LOG_WARN("%s: DFlash capture is stale for layer %d (seen_batch=%llu current_batch=%llu rows=%d width=%d)\n", + __func__, + capture.layer_ids[(size_t) layer_idx], + (unsigned long long) capture.layer_seen_batch_id[(size_t) layer_idx], + (unsigned long long) capture.capture_batch_id, + row_count, + row_width); + } + return false; + } + + 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; + } + } + + 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_prepared( + struct llama_context * ctx, + int32_t row_count, + int32_t row_width, + int32_t n_layers, + const std::vector & row_indices, + std::vector & rows_out, + int32_t & combined_width); + +static bool llama_spec_materialize_dflash_rows( + struct llama_context * ctx, + const std::vector & row_indices, + std::vector & rows_out, + int32_t & combined_width) { + 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)) { + if (ctx != nullptr) { + ctx->dflash.profile.capture_materialize_failures++; + } + return false; + } + + return llama_spec_materialize_dflash_rows_prepared(ctx, row_count, row_width, n_layers, row_indices, rows_out, combined_width); +} + +static bool llama_spec_materialize_dflash_rows_prepared( + struct llama_context * ctx, + int32_t row_count, + int32_t row_width, + int32_t n_layers, + const std::vector & row_indices, + std::vector & rows_out, + int32_t & combined_width) { + rows_out.clear(); + combined_width = 0; + if (ctx == nullptr || row_indices.empty()) { + return false; + } + + auto & profile = ctx->dflash.profile; + profile.capture_materialize_calls++; + const int64_t t_start_us = ggml_time_us(); + + if (row_count <= 0 || row_width <= 0 || n_layers <= 0 || ctx->dflash.capture == nullptr) { + profile.capture_materialize_failures++; + return false; + } + + combined_width = row_width * n_layers; + rows_out.resize((size_t) row_indices.size() * (size_t) combined_width); + + const auto & layer_rows = ctx->dflash.capture->layer_rows; + for (size_t out_row = 0; out_row < row_indices.size(); ++out_row) { + int32_t row_index = row_indices[out_row]; + if (row_index < 0) { + row_index += row_count; + } + if (row_index < 0 || row_index >= row_count) { + rows_out.clear(); + combined_width = 0; + profile.capture_materialize_failures++; + return false; + } + + float * dst = rows_out.data() + out_row * (size_t) combined_width; + for (int32_t layer_idx = 0; layer_idx < n_layers; ++layer_idx) { + const float * src = layer_rows[(size_t) layer_idx].data() + (size_t) row_index * (size_t) row_width; + std::memcpy(dst + (size_t) layer_idx * (size_t) row_width, src, (size_t) row_width * sizeof(float)); + } + } + + 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; +} + + +bool llama_spec_get_dflash_feature_view( + struct llama_context * ctx, + const llama_batch & batch, + llama_spec_feature_view & view) { + if (ctx == nullptr || batch.n_tokens <= 0 || batch.pos == nullptr || batch.n_seq_id == nullptr || batch.seq_id == nullptr) { + return false; + } + + 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)) { + return false; + } + + const int32_t batch_row_offset = std::max(0, batch.n_tokens - row_count); + std::vector row_indices; + std::vector batch_indices; + row_indices.reserve((size_t) (batch.n_tokens - batch_row_offset)); + batch_indices.reserve((size_t) (batch.n_tokens - batch_row_offset)); + for (int32_t i = batch_row_offset; i < batch.n_tokens; ++i) { + row_indices.push_back(i - batch_row_offset); + batch_indices.push_back(i); + } + + if (row_indices.empty()) { + return false; + } + + view = {}; + view.kind = LLAMA_SPEC_FEATURE_HIDDEN_STATE; + if (!llama_spec_materialize_dflash_rows_prepared(ctx, row_count, row_width, n_layers, row_indices, ctx->dflash.feature_view_buffer, view.width)) { + return false; + } + + view.rows.reserve(batch_indices.size()); + for (int32_t batch_index : batch_indices) { + if (batch.n_seq_id[batch_index] <= 0 || batch.seq_id[batch_index] == nullptr) { + view.rows.clear(); + return false; + } + + view.rows.push_back({ + /* .seq_id = */ batch.seq_id[batch_index][0], + /* .pos = */ batch.pos[batch_index], + /* .data = */ ctx->dflash.feature_view_buffer.data() + view.rows.size() * (size_t) view.width, + }); + } + + 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; +} + +bool llama_spec_get_dflash_feature_view_for_seq( + struct llama_context * ctx, + const llama_batch & batch, + llama_seq_id seq_id, + llama_spec_feature_view & view) { + if (ctx == nullptr || batch.n_tokens <= 0 || batch.pos == nullptr || batch.n_seq_id == nullptr || batch.seq_id == nullptr) { + return false; + } + + 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)) { + return false; + } + + const int32_t batch_row_offset = std::max(0, batch.n_tokens - row_count); + std::vector row_indices; + row_indices.reserve((size_t) batch.n_tokens); + std::vector batch_indices; + batch_indices.reserve((size_t) batch.n_tokens); + for (int32_t i = batch_row_offset; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] <= 0 || batch.seq_id[i] == nullptr) { + return false; + } + + for (int32_t j = 0; j < batch.n_seq_id[i]; ++j) { + if (batch.seq_id[i][j] == seq_id) { + row_indices.push_back(i - batch_row_offset); + batch_indices.push_back(i); + break; + } + } + } + + if (row_indices.empty()) { + return false; + } + + view = {}; + view.kind = LLAMA_SPEC_FEATURE_HIDDEN_STATE; + if (!llama_spec_materialize_dflash_rows_prepared(ctx, row_count, row_width, n_layers, row_indices, ctx->dflash.feature_view_buffer, view.width)) { + return false; + } + + view.rows.reserve(row_indices.size()); + for (size_t i = 0; i < batch_indices.size(); ++i) { + const int32_t batch_index = batch_indices[i]; + view.rows.push_back({ + /* .seq_id = */ seq_id, + /* .pos = */ batch.pos[batch_index], + /* .data = */ ctx->dflash.feature_view_buffer.data() + i * (size_t) view.width, + }); + } + + 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; +} + +bool llama_spec_copy_dflash_rows_from_output_indices( + struct llama_context * ctx, + const std::vector & output_indices, + std::vector & hidden_rows) { + int32_t combined_width = 0; + if (!llama_spec_materialize_dflash_rows(ctx, output_indices, hidden_rows, combined_width)) { + hidden_rows.clear(); + 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-dflash.h b/src/llama-spec-features-dflash.h new file mode 100644 index 00000000..02e709d1 --- /dev/null +++ b/src/llama-spec-features-dflash.h @@ -0,0 +1,279 @@ +#pragma once + +#include "llama.h" + +#include +#include +#include + +struct llama_context; +struct llama_model; +struct ggml_tensor; +struct llama_spec_feature_view; + +struct llama_dflash_profile_stats { + uint64_t decode_internal_chunks = 0; + uint64_t decode_graph_rebuilds = 0; + uint64_t decode_sync_profile_points = 0; + uint64_t decode_prelude_us = 0; + uint64_t decode_sched_reset_us = 0; + uint64_t decode_build_graph_us = 0; + uint64_t decode_sched_alloc_graph_us = 0; + uint64_t decode_set_inputs_us = 0; + uint64_t decode_graph_compute_us = 0; + uint64_t decode_result_us = 0; + uint64_t decode_embedding_us = 0; + uint64_t decode_final_sched_reset_us = 0; + + uint64_t decode_output_reserve_calls = 0; + uint64_t decode_output_reserve_us = 0; + uint64_t decode_output_reserve_reallocs = 0; + uint64_t decode_output_reserve_realloc_bytes = 0; + uint64_t decode_prepare_calls = 0; + uint64_t decode_prepare_us = 0; + uint64_t decode_prepare_failures = 0; + + 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_layer_batch_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_build_us = 0; + uint64_t graph_kv_cache_reserve_us = 0; + uint64_t graph_kv_cache_reset_us = 0; + uint64_t graph_kv_cache_alloc_us = 0; + uint64_t graph_kv_cache_feature_upload_us = 0; + uint64_t graph_kv_cache_pos_upload_us = 0; + uint64_t graph_kv_cache_compute_us = 0; + uint64_t graph_kv_cache_sync_us = 0; + uint64_t graph_kv_cache_read_concat_pad_us = 0; + 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; + uint64_t graph_kv_node_k_proj_us = 0; + uint64_t graph_kv_node_k_norm_calls = 0; + uint64_t graph_kv_node_k_norm_us = 0; + uint64_t graph_kv_node_k_rope_calls = 0; + uint64_t graph_kv_node_k_rope_us = 0; + uint64_t graph_kv_node_v_proj_calls = 0; + uint64_t graph_kv_node_v_proj_us = 0; + uint64_t graph_kv_node_k_store_calls = 0; + 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; + 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 last_kv_cache_host_layers = 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; +}; + +struct llama_dflash_window_update { + uint64_t version = 0; + int32_t keep_rows = 0; + int32_t append_rows = 0; + bool replace = false; + const float * append_features = nullptr; + size_t append_floats = 0; +}; + +struct llama_dflash_kv_cache_transition { + bool cache_up_to_date = false; + bool rebuild_cache = false; + int32_t append_rows = 0; + int32_t next_n_filled = 0; + int32_t next_write_pos = 0; +}; + +static inline llama_dflash_kv_cache_transition llama_plan_dflash_kv_cache_transition( + int32_t cross_ctx, + int32_t current_n_filled, + int32_t current_write_pos, + bool cache_valid, + uint64_t applied_window_version, + uint64_t target_window_version, + int32_t keep_rows, + int32_t append_rows, + bool replace, + int32_t n_rows) { + llama_dflash_kv_cache_transition plan; + + const int32_t safe_cross_ctx = std::max(1, cross_ctx); + const int32_t bounded_n_filled = std::clamp(current_n_filled, 0, safe_cross_ctx); + const int32_t bounded_append_rows = std::clamp(append_rows, 0, n_rows); + const int32_t bounded_keep_rows = std::clamp(keep_rows, 0, n_rows); + const int32_t expected_keep_rows = std::min(bounded_n_filled, std::max(0, safe_cross_ctx - bounded_append_rows)); + + plan.cache_up_to_date = cache_valid && applied_window_version == target_window_version; + plan.rebuild_cache = !cache_valid || replace || bounded_append_rows <= 0 || bounded_append_rows > n_rows; + if (!plan.rebuild_cache && bounded_keep_rows != expected_keep_rows) { + plan.rebuild_cache = true; + } + + plan.append_rows = bounded_append_rows; + if (plan.cache_up_to_date) { + plan.next_n_filled = bounded_n_filled; + plan.next_write_pos = safe_cross_ctx > 0 + ? ((current_write_pos % safe_cross_ctx) + safe_cross_ctx) % safe_cross_ctx + : 0; + } else if (plan.rebuild_cache) { + plan.next_n_filled = std::min(safe_cross_ctx, n_rows); + plan.next_write_pos = plan.next_n_filled % safe_cross_ctx; + } else { + plan.next_n_filled = std::min(safe_cross_ctx, bounded_n_filled + bounded_append_rows); + plan.next_write_pos = (current_write_pos + bounded_append_rows) % safe_cross_ctx; + } + + return plan; +} + +llama_dflash_kv_cache_transition llama_plan_dflash_kv_cache_transition_for_ctx( + const struct llama_context * ctx, + const llama_dflash_window_update & window_update, + int32_t n_rows); + +void llama_dflash_profile_reset(struct llama_context * ctx); +void llama_reset_dflash_kv_cache_state(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); +int32_t llama_model_dflash_n_target_layers(const struct llama_model * model); +int32_t llama_model_dflash_n_target_features(const struct llama_model * model); +int32_t llama_model_dflash_target_layer_ids(const struct llama_model * model, int32_t * layer_ids, int32_t capacity); +int32_t llama_model_dflash_target_mask_token_id(const struct llama_model * model); +const struct ggml_tensor * llama_model_dflash_output_tensor(const struct llama_model * model); + +enum llama_dflash_io_mode { + LLAMA_DFLASH_IO_MODE_INVALID = 0, + LLAMA_DFLASH_IO_MODE_SHARED, + LLAMA_DFLASH_IO_MODE_SELF_CONTAINED, + LLAMA_DFLASH_IO_MODE_MIXED, +}; + +int32_t llama_model_dflash_io_mode(const struct llama_model * draft_model, const struct llama_model * target_model); +bool llama_model_dflash_io_tensors_match(const struct llama_model * draft_model, int32_t n_embd, int32_t n_vocab); +bool llama_model_share_dflash_io_tensors(struct llama_model * draft_model, const struct llama_model * target_model); + +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, + const llama_dflash_window_update * window_update = nullptr); + +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, + const llama_dflash_window_update * window_update = nullptr); + +bool llama_set_dflash_capture_layers(struct llama_context * ctx, const int32_t * layer_ids, int32_t n_layers); +void llama_clear_dflash_capture(struct llama_context * ctx); +void llama_begin_dflash_capture_batch(struct llama_context * ctx); +void llama_finish_dflash_capture_batch(struct llama_context * ctx, bool is_prompt_warmup); + +bool llama_spec_get_dflash_feature_view( + struct llama_context * ctx, + const llama_batch & batch, + llama_spec_feature_view & view); + +bool llama_spec_get_dflash_feature_view_for_seq( + struct llama_context * ctx, + const llama_batch & batch, + llama_seq_id seq_id, + llama_spec_feature_view & view); + +bool llama_spec_copy_dflash_rows_from_output_indices( + struct llama_context * ctx, + const std::vector & output_indices, + std::vector & hidden_rows); diff --git a/src/llama-spec-features.cpp b/src/llama-spec-features.cpp index 00c4b6e2..933f3d15 100644 --- a/src/llama-spec-features.cpp +++ b/src/llama-spec-features.cpp @@ -10,30 +10,6 @@ #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; @@ -47,278 +23,6 @@ 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_reset_dflash_kv_cache_state(struct llama_context * ctx) { - if (ctx == nullptr) { - return; - } - - ctx->dflash_kv_cache_write_pos = 0; - ctx->dflash_kv_cache_n_filled = 0; - ctx->dflash_kv_cache_update_rows = 0; - ctx->dflash_kv_cache_view_write_pos = 0; - ctx->dflash_kv_cache_view_n_filled = 0; - 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) { - ggml_backend_buffer_clear(buf, 0); - } - } -} - -llama_dflash_kv_cache_transition llama_plan_dflash_kv_cache_transition_for_ctx( - const struct llama_context * ctx, - const llama_dflash_window_update & window_update, - int32_t n_rows) { - if (ctx == nullptr) { - llama_dflash_kv_cache_transition plan; - plan.rebuild_cache = true; - plan.append_rows = std::clamp(window_update.append_rows, 0, n_rows); - plan.next_n_filled = n_rows; - return plan; - } - - const int32_t cross_ctx = ctx->dflash_visible_cross_ctx > 0 - ? ctx->dflash_visible_cross_ctx - : std::max(1, (int32_t) ctx->cparams.n_ctx - (int32_t) ctx->model.hparams.dflash_block_size); - - return llama_plan_dflash_kv_cache_transition( - cross_ctx, - ctx->dflash_kv_cache_n_filled, - ctx->dflash_kv_cache_write_pos, - ctx->dflash_kv_cache_valid, - ctx->dflash_kv_cache_applied_window_version, - window_update.version, - window_update.keep_rows, - window_update.append_rows, - window_update.replace, - n_rows); -} - -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; -} - -int32_t llama_model_dflash_mask_token_id(const struct llama_model * model) { - return model ? (int32_t) model->hparams.dflash_mask_token_id : -1; -} - -int32_t llama_model_dflash_n_target_layers(const struct llama_model * model) { - return model ? (int32_t) model->hparams.dflash_n_target_layers : 0; -} - -int32_t llama_model_dflash_n_target_features(const struct llama_model * model) { - return model ? (int32_t) model->hparams.dflash_n_target_features : 0; -} - -int32_t llama_model_dflash_target_layer_ids( - const struct llama_model * model, - int32_t * layer_ids, - int32_t capacity) { - if (model == nullptr || layer_ids == nullptr || capacity <= 0) { - return 0; - } - - const int32_t n_layers = std::min((int32_t) model->hparams.dflash_n_target_layers, capacity); - for (int32_t i = 0; i < n_layers; ++i) { - layer_ids[i] = (int32_t) model->hparams.dflash_target_layer_ids[i]; - } - - return n_layers; -} - -int32_t llama_model_dflash_target_mask_token_id(const struct llama_model * model) { - if (model == nullptr) { - return (int32_t) LLAMA_TOKEN_NULL; - } - - 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) { - if (draft_model == nullptr || target_model == nullptr || draft_model->arch != LLM_ARCH_DFLASH_DRAFT) { - return LLAMA_DFLASH_IO_MODE_INVALID; - } - - 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_output == target_output; - if (shared_tok && shared_output) { - return LLAMA_DFLASH_IO_MODE_SHARED; - } - - if (!shared_tok && !shared_output) { - return LLAMA_DFLASH_IO_MODE_SELF_CONTAINED; - } - - return LLAMA_DFLASH_IO_MODE_MIXED; -} - -bool llama_model_dflash_io_tensors_match( - const struct llama_model * draft_model, - int32_t n_embd, - int32_t n_vocab) { - 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) output->ne[0] == n_embd && - (int32_t) output->ne[1] == n_vocab; -} - -bool llama_model_share_dflash_io_tensors( - struct llama_model * draft_model, - const struct llama_model * target_model) { - if (draft_model == nullptr || target_model == nullptr) { - return false; - } - - if (draft_model->arch != LLM_ARCH_DFLASH_DRAFT) { - return true; - } - - if (draft_model->tok_embd == nullptr) { - draft_model->tok_embd = target_model->tok_embd; - } - - if (draft_model->output == nullptr) { - draft_model->output = target_model->output ? target_model->output : target_model->tok_embd; - if (draft_model->output == nullptr) { - draft_model->output = draft_model->tok_embd; - } - } - - 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( struct llama_context * ctx, const float * hidden_state, @@ -333,648 +37,6 @@ bool llama_set_draft_input_hidden_state_copy( return true; } -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, - bool copy_data, - const llama_dflash_window_update * window_update) { - const bool have_full_features = target_features != nullptr && n_floats > 0; - const bool have_append_features = window_update != nullptr && - window_update->append_features != nullptr && - window_update->append_floats > 0 && - window_update->append_rows > 0; - - if (ctx == nullptr || n_rows <= 0 || (!have_full_features && !have_append_features)) { - return false; - } - - auto & profile = ctx->dflash_profile; - const int64_t t_start_us = ggml_time_us(); - const int32_t row_width = have_full_features - ? (n_rows > 0 ? (int32_t) (n_floats / (size_t) n_rows) : 0) - : (window_update->append_rows > 0 ? (int32_t) (window_update->append_floats / (size_t) window_update->append_rows) : 0); - llama_pos first_pos = -1; - llama_pos last_pos = -1; - - if (have_full_features && 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 if (have_full_features) { - ctx->dflash_target_features_owned.clear(); - ctx->dflash_target_features = target_features; - } else { - ctx->dflash_target_features_owned.clear(); - ctx->dflash_target_features = nullptr; - } - ctx->dflash_target_features_n_floats = have_full_features ? n_floats : 0; - ctx->dflash_target_features_n_rows = n_rows; - if (have_append_features && copy_data) { - ctx->dflash_target_append_features_owned.assign( - window_update->append_features, - window_update->append_features + window_update->append_floats); - ctx->dflash_target_append_features = ctx->dflash_target_append_features_owned.data(); - } else if (have_append_features) { - ctx->dflash_target_append_features_owned.clear(); - ctx->dflash_target_append_features = window_update->append_features; - } else { - ctx->dflash_target_append_features_owned.clear(); - ctx->dflash_target_append_features = nullptr; - } - ctx->dflash_target_append_features_n_floats = have_append_features ? window_update->append_floats : 0; - ctx->dflash_target_append_features_n_rows = have_append_features ? window_update->append_rows : 0; - ctx->dflash_target_window_version = window_update != nullptr && window_update->version > 0 - ? window_update->version - : ctx->dflash_target_window_version + 1; - ctx->dflash_target_window_keep_rows = window_update != nullptr - ? std::max(0, std::min(n_rows, window_update->keep_rows)) - : 0; - ctx->dflash_target_window_append_rows = window_update != nullptr - ? std::max(0, std::min(n_rows, window_update->append_rows)) - : n_rows; - ctx->dflash_target_window_replace = window_update != nullptr - ? window_update->replace - : true; - if (ctx->dflash_target_window_keep_rows + ctx->dflash_target_window_append_rows > n_rows) { - ctx->dflash_target_window_keep_rows = std::max(0, n_rows - ctx->dflash_target_window_append_rows); - } - - const int32_t cross_ctx = ctx->dflash_visible_cross_ctx > 0 - ? ctx->dflash_visible_cross_ctx - : std::max(1, (int32_t) ctx->cparams.n_ctx - (int32_t) ctx->model.hparams.dflash_block_size); - const llama_dflash_window_update cache_window_update = { - ctx->dflash_target_window_version, - ctx->dflash_target_window_keep_rows, - ctx->dflash_target_window_append_rows, - ctx->dflash_target_window_replace, - ctx->dflash_target_append_features, - ctx->dflash_target_append_features_n_floats, - }; - const llama_dflash_kv_cache_transition cache_plan = llama_plan_dflash_kv_cache_transition_for_ctx(ctx, cache_window_update, n_rows); - - if (cache_plan.cache_up_to_date) { - ctx->dflash_kv_cache_view_n_filled = ctx->dflash_kv_cache_n_filled; - ctx->dflash_kv_cache_view_write_pos = ctx->dflash_kv_cache_write_pos; - ctx->dflash_kv_cache_view_valid = ctx->dflash_kv_cache_valid; - } else if (cross_ctx > 0) { - ctx->dflash_kv_cache_view_n_filled = cache_plan.next_n_filled; - ctx->dflash_kv_cache_view_write_pos = cache_plan.next_write_pos; - ctx->dflash_kv_cache_view_valid = cache_plan.next_n_filled > 0; - } - - if (target_positions != nullptr) { - 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 += - (have_full_features ? n_floats : 0) * sizeof(float) + - (have_append_features ? window_update->append_floats : 0) * 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, - const llama_dflash_window_update * window_update) { - return llama_set_dflash_target_features_impl(ctx, target_features, n_floats, n_rows, target_positions, true, window_update); -} - -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, - const llama_dflash_window_update * window_update) { - return llama_set_dflash_target_features_impl(ctx, target_features, n_floats, n_rows, target_positions, false, window_update); -} - -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; - } - - static constexpr const char * prefix = "l_out-"; - if (std::strncmp(tensor->name, prefix, std::strlen(prefix)) != 0) { - return false; - } - - char * end = nullptr; - const long raw = std::strtol(tensor->name + std::strlen(prefix), &end, 10); - if (end == tensor->name + std::strlen(prefix) || *end != '\0') { - return false; - } - - layer_id = (int32_t) raw; - if (layer_id >= 1000) { - layer_id %= 1000; - } - - return layer_id >= 0; -} - -static int32_t llama_dflash_find_layer_index(const struct llama_context * ctx, int32_t layer_id) { - if (ctx == nullptr || !ctx->dflash_capture) { - return -1; - } - - const auto & layer_ids = ctx->dflash_capture->layer_ids; - const auto it = std::find(layer_ids.begin(), layer_ids.end(), layer_id); - return it == layer_ids.end() ? -1 : (int32_t) std::distance(layer_ids.begin(), it); -} - -static bool llama_dflash_capture_eval_callback(struct ggml_tensor * tensor, bool ask, void * user_data) { - auto * ctx = static_cast(user_data); - if (ctx == nullptr || !ctx->dflash_capture) { - return false; - } - - int32_t layer_id = -1; - if (!llama_dflash_parse_layer_id(tensor, layer_id)) { - return false; - } - - const int32_t layer_idx = llama_dflash_find_layer_index(ctx, layer_id); - if (layer_idx < 0) { - return false; - } - - if (ask) { - return true; - } - - const int32_t row_width = (int32_t) tensor->ne[0]; - const int32_t row_count = row_width > 0 ? (int32_t) (ggml_nelements(tensor) / (int64_t) row_width) : 0; - if (row_width <= 0 || row_count <= 0) { - return false; - } - - auto & capture = *ctx->dflash_capture; - if (capture.capture_batch_id == 0) { - capture.capture_batch_id = 1; - } - if (capture.layer_seen_batch_id.size() != capture.layer_ids.size()) { - capture.layer_seen_batch_id.assign(capture.layer_ids.size(), 0); - } - - auto & rows = capture.layer_rows[(size_t) layer_idx]; - rows.resize((size_t) row_count * (size_t) row_width); - ggml_backend_tensor_get(tensor, rows.data(), 0, ggml_nbytes(tensor)); - capture.row_width = row_width; - capture.row_count = row_count; - capture.layer_seen_batch_id[(size_t) layer_idx] = capture.capture_batch_id; - return true; -} - -bool llama_set_dflash_capture_layers( - struct llama_context * ctx, - const int32_t * layer_ids, - int32_t n_layers) { - if (ctx == nullptr || layer_ids == nullptr || n_layers <= 0) { - return false; - } - - auto capture = std::make_unique(); - capture->layer_ids.assign(layer_ids, layer_ids + n_layers); - capture->layer_rows.resize((size_t) n_layers); - capture->layer_seen_batch_id.assign((size_t) n_layers, 0); - capture->prev_cb_eval = ctx->cparams.cb_eval; - capture->prev_cb_eval_user_data = ctx->cparams.cb_eval_user_data; - ctx->dflash_capture = std::move(capture); - ctx->dflash_feature_view_buffer.clear(); - - ctx->cparams.cb_eval = llama_dflash_capture_eval_callback; - ctx->cparams.cb_eval_user_data = ctx; - if (ctx->sched != nullptr) { - ggml_backend_sched_set_eval_callback(ctx->sched, ctx->cparams.cb_eval, ctx->cparams.cb_eval_user_data); - } - - return true; -} - -void llama_clear_dflash_capture(struct llama_context * ctx) { - if (ctx == nullptr) { - return; - } - - ggml_backend_sched_eval_callback prev_cb_eval = nullptr; - void * prev_cb_eval_user_data = nullptr; - if (ctx->dflash_capture) { - prev_cb_eval = ctx->dflash_capture->prev_cb_eval; - prev_cb_eval_user_data = ctx->dflash_capture->prev_cb_eval_user_data; - } - - ctx->dflash_capture.reset(); - ctx->dflash_feature_view_buffer.clear(); - - if (ctx->cparams.cb_eval == llama_dflash_capture_eval_callback && ctx->cparams.cb_eval_user_data == ctx) { - ctx->cparams.cb_eval = prev_cb_eval; - ctx->cparams.cb_eval_user_data = prev_cb_eval_user_data; - if (ctx->sched != nullptr) { - ggml_backend_sched_set_eval_callback(ctx->sched, prev_cb_eval, prev_cb_eval_user_data); - } - } -} - -void llama_begin_dflash_capture_batch(struct llama_context * ctx) { - if (ctx == nullptr || !ctx->dflash_capture) { - return; - } - - auto & capture = *ctx->dflash_capture; - capture.capture_batch_id++; - capture.row_count = 0; - capture.row_width = 0; - std::fill(capture.layer_seen_batch_id.begin(), capture.layer_seen_batch_id.end(), 0); -} - -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, - int32_t & row_width, - int32_t & n_layers) { - if (ctx == nullptr || !ctx->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; - } - - if (capture.capture_batch_id == 0 || capture.layer_seen_batch_id.size() != (size_t) n_layers) { - profile.capture_prepare_failures++; - profile.capture_layer_batch_mismatch++; - if (profile.capture_layer_batch_mismatch <= 3) { - LLAMA_LOG_WARN("%s: DFlash capture batch markers are not initialized (batch_id=%llu layers=%zu expected=%d)\n", - __func__, - (unsigned long long) capture.capture_batch_id, - capture.layer_seen_batch_id.size(), - n_layers); - } - return false; - } - - for (int32_t layer_idx = 0; layer_idx < n_layers; ++layer_idx) { - if (capture.layer_seen_batch_id[(size_t) layer_idx] != capture.capture_batch_id) { - profile.capture_prepare_failures++; - profile.capture_layer_batch_mismatch++; - if (profile.capture_layer_batch_mismatch <= 3) { - LLAMA_LOG_WARN("%s: DFlash capture is stale for layer %d (seen_batch=%llu current_batch=%llu rows=%d width=%d)\n", - __func__, - capture.layer_ids[(size_t) layer_idx], - (unsigned long long) capture.layer_seen_batch_id[(size_t) layer_idx], - (unsigned long long) capture.capture_batch_id, - row_count, - row_width); - } - return false; - } - - 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; - } - } - - 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_prepared( - struct llama_context * ctx, - int32_t row_count, - int32_t row_width, - int32_t n_layers, - const std::vector & row_indices, - std::vector & rows_out, - int32_t & combined_width); - -static bool llama_spec_materialize_dflash_rows( - struct llama_context * ctx, - const std::vector & row_indices, - std::vector & rows_out, - int32_t & combined_width) { - 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)) { - if (ctx != nullptr) { - ctx->dflash_profile.capture_materialize_failures++; - } - return false; - } - - return llama_spec_materialize_dflash_rows_prepared(ctx, row_count, row_width, n_layers, row_indices, rows_out, combined_width); -} - -static bool llama_spec_materialize_dflash_rows_prepared( - struct llama_context * ctx, - int32_t row_count, - int32_t row_width, - int32_t n_layers, - const std::vector & row_indices, - std::vector & rows_out, - int32_t & combined_width) { - rows_out.clear(); - combined_width = 0; - if (ctx == nullptr || row_indices.empty()) { - return false; - } - - auto & profile = ctx->dflash_profile; - profile.capture_materialize_calls++; - const int64_t t_start_us = ggml_time_us(); - - if (row_count <= 0 || row_width <= 0 || n_layers <= 0 || ctx->dflash_capture == nullptr) { - profile.capture_materialize_failures++; - return false; - } - - combined_width = row_width * n_layers; - rows_out.resize((size_t) row_indices.size() * (size_t) combined_width); - - const auto & layer_rows = ctx->dflash_capture->layer_rows; - for (size_t out_row = 0; out_row < row_indices.size(); ++out_row) { - int32_t row_index = row_indices[out_row]; - if (row_index < 0) { - row_index += row_count; - } - if (row_index < 0 || row_index >= row_count) { - rows_out.clear(); - combined_width = 0; - profile.capture_materialize_failures++; - return false; - } - - float * dst = rows_out.data() + out_row * (size_t) combined_width; - for (int32_t layer_idx = 0; layer_idx < n_layers; ++layer_idx) { - const float * src = layer_rows[(size_t) layer_idx].data() + (size_t) row_index * (size_t) row_width; - std::memcpy(dst + (size_t) layer_idx * (size_t) row_width, src, (size_t) row_width * sizeof(float)); - } - } - - 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; -} - static bool llama_spec_prepare_hidden_feature_view( struct llama_context * ctx, int32_t n_rows, @@ -1031,137 +93,6 @@ bool llama_spec_get_hidden_feature_view( return true; } -bool llama_spec_get_dflash_feature_view( - struct llama_context * ctx, - const llama_batch & batch, - llama_spec_feature_view & view) { - if (ctx == nullptr || batch.n_tokens <= 0 || batch.pos == nullptr || batch.n_seq_id == nullptr || batch.seq_id == nullptr) { - return false; - } - - 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)) { - return false; - } - - const int32_t batch_row_offset = std::max(0, batch.n_tokens - row_count); - std::vector row_indices; - std::vector batch_indices; - row_indices.reserve((size_t) (batch.n_tokens - batch_row_offset)); - batch_indices.reserve((size_t) (batch.n_tokens - batch_row_offset)); - for (int32_t i = batch_row_offset; i < batch.n_tokens; ++i) { - row_indices.push_back(i - batch_row_offset); - batch_indices.push_back(i); - } - - if (row_indices.empty()) { - return false; - } - - view = {}; - view.kind = LLAMA_SPEC_FEATURE_HIDDEN_STATE; - if (!llama_spec_materialize_dflash_rows_prepared(ctx, row_count, row_width, n_layers, row_indices, ctx->dflash_feature_view_buffer, view.width)) { - return false; - } - - view.rows.reserve(batch_indices.size()); - for (int32_t batch_index : batch_indices) { - if (batch.n_seq_id[batch_index] <= 0 || batch.seq_id[batch_index] == nullptr) { - view.rows.clear(); - return false; - } - - view.rows.push_back({ - /* .seq_id = */ batch.seq_id[batch_index][0], - /* .pos = */ batch.pos[batch_index], - /* .data = */ ctx->dflash_feature_view_buffer.data() + view.rows.size() * (size_t) view.width, - }); - } - - 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; -} - -bool llama_spec_get_dflash_feature_view_for_seq( - struct llama_context * ctx, - const llama_batch & batch, - llama_seq_id seq_id, - llama_spec_feature_view & view) { - if (ctx == nullptr || batch.n_tokens <= 0 || batch.pos == nullptr || batch.n_seq_id == nullptr || batch.seq_id == nullptr) { - return false; - } - - 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)) { - return false; - } - - const int32_t batch_row_offset = std::max(0, batch.n_tokens - row_count); - std::vector row_indices; - row_indices.reserve((size_t) batch.n_tokens); - std::vector batch_indices; - batch_indices.reserve((size_t) batch.n_tokens); - for (int32_t i = batch_row_offset; i < batch.n_tokens; ++i) { - if (batch.n_seq_id[i] <= 0 || batch.seq_id[i] == nullptr) { - return false; - } - - for (int32_t j = 0; j < batch.n_seq_id[i]; ++j) { - if (batch.seq_id[i][j] == seq_id) { - row_indices.push_back(i - batch_row_offset); - batch_indices.push_back(i); - break; - } - } - } - - if (row_indices.empty()) { - return false; - } - - view = {}; - view.kind = LLAMA_SPEC_FEATURE_HIDDEN_STATE; - if (!llama_spec_materialize_dflash_rows_prepared(ctx, row_count, row_width, n_layers, row_indices, ctx->dflash_feature_view_buffer, view.width)) { - return false; - } - - view.rows.reserve(row_indices.size()); - for (size_t i = 0; i < batch_indices.size(); ++i) { - const int32_t batch_index = batch_indices[i]; - view.rows.push_back({ - /* .seq_id = */ seq_id, - /* .pos = */ batch.pos[batch_index], - /* .data = */ ctx->dflash_feature_view_buffer.data() + i * (size_t) view.width, - }); - } - - 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; -} bool llama_spec_get_hidden_feature_view_for_seq( struct llama_context * ctx, @@ -1255,18 +186,3 @@ bool llama_spec_copy_hidden_rows_from_output_indices( return hidden_rows.size() == (size_t) output_indices.size() * view.width; } - -bool llama_spec_copy_dflash_rows_from_output_indices( - struct llama_context * ctx, - const std::vector & output_indices, - std::vector & hidden_rows) { - int32_t combined_width = 0; - if (!llama_spec_materialize_dflash_rows(ctx, output_indices, hidden_rows, combined_width)) { - hidden_rows.clear(); - 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 1c327049..b1342fed 100644 --- a/src/llama-spec-features.h +++ b/src/llama-spec-features.h @@ -2,7 +2,6 @@ #include "llama.h" -#include #include #include @@ -25,316 +24,20 @@ struct llama_spec_feature_view { std::vector rows; }; -struct llama_dflash_profile_stats { - uint64_t decode_internal_chunks = 0; - uint64_t decode_graph_rebuilds = 0; - uint64_t decode_sync_profile_points = 0; - uint64_t decode_prelude_us = 0; - uint64_t decode_sched_reset_us = 0; - uint64_t decode_build_graph_us = 0; - uint64_t decode_sched_alloc_graph_us = 0; - uint64_t decode_set_inputs_us = 0; - uint64_t decode_graph_compute_us = 0; - uint64_t decode_result_us = 0; - uint64_t decode_embedding_us = 0; - uint64_t decode_final_sched_reset_us = 0; - - uint64_t decode_output_reserve_calls = 0; - uint64_t decode_output_reserve_us = 0; - uint64_t decode_output_reserve_reallocs = 0; - uint64_t decode_output_reserve_realloc_bytes = 0; - uint64_t decode_prepare_calls = 0; - uint64_t decode_prepare_us = 0; - uint64_t decode_prepare_failures = 0; - - 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_layer_batch_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_build_us = 0; - uint64_t graph_kv_cache_reserve_us = 0; - uint64_t graph_kv_cache_reset_us = 0; - uint64_t graph_kv_cache_alloc_us = 0; - uint64_t graph_kv_cache_feature_upload_us = 0; - uint64_t graph_kv_cache_pos_upload_us = 0; - uint64_t graph_kv_cache_compute_us = 0; - uint64_t graph_kv_cache_sync_us = 0; - uint64_t graph_kv_cache_read_concat_pad_us = 0; - 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; - uint64_t graph_kv_node_k_proj_us = 0; - uint64_t graph_kv_node_k_norm_calls = 0; - uint64_t graph_kv_node_k_norm_us = 0; - uint64_t graph_kv_node_k_rope_calls = 0; - uint64_t graph_kv_node_k_rope_us = 0; - uint64_t graph_kv_node_v_proj_calls = 0; - uint64_t graph_kv_node_v_proj_us = 0; - uint64_t graph_kv_node_k_store_calls = 0; - 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; - 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 last_kv_cache_host_layers = 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; -}; - -struct llama_dflash_window_update { - uint64_t version = 0; - int32_t keep_rows = 0; - int32_t append_rows = 0; - bool replace = false; - const float * append_features = nullptr; - size_t append_floats = 0; -}; - -struct llama_dflash_kv_cache_transition { - bool cache_up_to_date = false; - bool rebuild_cache = false; - int32_t append_rows = 0; - int32_t next_n_filled = 0; - int32_t next_write_pos = 0; -}; - -static inline llama_dflash_kv_cache_transition llama_plan_dflash_kv_cache_transition( - int32_t cross_ctx, - int32_t current_n_filled, - int32_t current_write_pos, - bool cache_valid, - uint64_t applied_window_version, - uint64_t target_window_version, - int32_t keep_rows, - int32_t append_rows, - bool replace, - int32_t n_rows) { - llama_dflash_kv_cache_transition plan; - - const int32_t safe_cross_ctx = std::max(1, cross_ctx); - const int32_t bounded_n_filled = std::clamp(current_n_filled, 0, safe_cross_ctx); - const int32_t bounded_append_rows = std::clamp(append_rows, 0, n_rows); - const int32_t bounded_keep_rows = std::clamp(keep_rows, 0, n_rows); - const int32_t expected_keep_rows = std::min(bounded_n_filled, std::max(0, safe_cross_ctx - bounded_append_rows)); - - plan.cache_up_to_date = cache_valid && applied_window_version == target_window_version; - plan.rebuild_cache = !cache_valid || replace || bounded_append_rows <= 0 || bounded_append_rows > n_rows; - if (!plan.rebuild_cache && bounded_keep_rows != expected_keep_rows) { - plan.rebuild_cache = true; - } - - plan.append_rows = bounded_append_rows; - if (plan.cache_up_to_date) { - plan.next_n_filled = bounded_n_filled; - plan.next_write_pos = safe_cross_ctx > 0 - ? ((current_write_pos % safe_cross_ctx) + safe_cross_ctx) % safe_cross_ctx - : 0; - } else if (plan.rebuild_cache) { - plan.next_n_filled = std::min(safe_cross_ctx, n_rows); - plan.next_write_pos = plan.next_n_filled % safe_cross_ctx; - } else { - plan.next_n_filled = std::min(safe_cross_ctx, bounded_n_filled + bounded_append_rows); - plan.next_write_pos = (current_write_pos + bounded_append_rows) % safe_cross_ctx; - } - - return plan; -} - -llama_dflash_kv_cache_transition llama_plan_dflash_kv_cache_transition_for_ctx( - const struct llama_context * ctx, - const llama_dflash_window_update & window_update, - int32_t n_rows); +#include "llama-spec-features-dflash.h" uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx); -void llama_dflash_profile_reset(struct llama_context * ctx); - -void llama_reset_dflash_kv_cache_state(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); - -int32_t llama_model_dflash_n_target_layers(const struct llama_model * model); - -int32_t llama_model_dflash_n_target_features(const struct llama_model * model); - -int32_t llama_model_dflash_target_layer_ids( - const struct llama_model * model, - int32_t * layer_ids, - int32_t capacity); - -enum llama_dflash_io_mode { - LLAMA_DFLASH_IO_MODE_INVALID = 0, - LLAMA_DFLASH_IO_MODE_SHARED, - LLAMA_DFLASH_IO_MODE_SELF_CONTAINED, - LLAMA_DFLASH_IO_MODE_MIXED, -}; - -int32_t llama_model_dflash_target_mask_token_id(const struct llama_model * model); - -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, - int32_t n_vocab); - -bool llama_model_share_dflash_io_tensors( - struct llama_model * draft_model, - const struct llama_model * target_model); - bool llama_set_draft_input_hidden_state_copy( struct llama_context * ctx, const float * hidden_state, size_t n_floats); -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, - const llama_dflash_window_update * window_update = nullptr); - -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, - const llama_dflash_window_update * window_update = nullptr); - -bool llama_set_dflash_capture_layers( - struct llama_context * ctx, - const int32_t * layer_ids, - int32_t n_layers); - -void llama_clear_dflash_capture(struct llama_context * ctx); - -void llama_begin_dflash_capture_batch(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, llama_spec_feature_view & view); -bool llama_spec_get_dflash_feature_view( - struct llama_context * ctx, - const llama_batch & batch, - llama_spec_feature_view & view); - -bool llama_spec_get_dflash_feature_view_for_seq( - struct llama_context * ctx, - const llama_batch & batch, - llama_seq_id seq_id, - llama_spec_feature_view & view); - bool llama_spec_get_hidden_feature_view_for_seq( struct llama_context * ctx, const llama_batch & batch, @@ -352,8 +55,3 @@ bool llama_spec_copy_hidden_rows_from_output_indices( struct llama_context * ctx, const std::vector & output_indices, std::vector & hidden_rows); - -bool llama_spec_copy_dflash_rows_from_output_indices( - struct llama_context * ctx, - const std::vector & output_indices, - std::vector & hidden_rows); diff --git a/src/llama.cpp b/src/llama.cpp index a1b63a73..75482d80 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -18,6 +18,7 @@ #include "llama-hparams.h" #include "llama-context.h" #include "llama-spec-features.h" +#include "llama-dflash.h" #include "llama-quantize.h" #include "unicode.h" @@ -515,20 +516,6 @@ static bool llama_dflash_main_node_eval_callback(struct ggml_tensor * tensor, bo return prev_result || tracked; } -static bool llama_dflash_use_kv_workspace_experiment() { - return llama_env_flag_enabled("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) { @@ -924,259 +911,6 @@ void llama_context::reset_scheduler() { prev_mtp.reset(); } -static ggml_backend_buffer_type_t llama_dflash_kv_cache_layer_buft(const llama_context & lctx, int32_t il) { - if (il >= 0 && (size_t) il < lctx.model.buft_layer.size() && lctx.model.buft_layer[(size_t) il].buft != nullptr) { - return lctx.model.buft_layer[(size_t) il].buft; - } - - if (il >= 0 && (size_t) il < lctx.model.layers.size()) { - const ggml_tensor * wk = lctx.model.layers[(size_t) il].wk; - if (wk != nullptr && wk->buffer != nullptr) { - return ggml_backend_buffer_get_type(wk->buffer); - } - } - - return llama_default_buffer_type_cpu(true); -} - -static ggml_backend_t llama_backend_for_tensor(const llama_context & lctx, const ggml_tensor * tensor) { - if (tensor == nullptr) { - return nullptr; - } - - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (buf == nullptr) { - return nullptr; - } - - ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf); - for (ggml_backend_t backend : lctx.backends) { - ggml_backend_buffer_type_t backend_buft = ggml_backend_is_cpu(backend) - ? llama_default_buffer_type_cpu(true) - : ggml_backend_get_default_buffer_type(backend); - if (backend_buft == buft) { - return backend; - } - } - - return nullptr; -} - -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()) { - 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; - } - - free_dflash_kv_cache_tensors(); - if (dflash_sched != nullptr) { - 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) ((use_kv_workspace ? 4 : 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); - 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) * (use_kv_workspace ? 4 : 2)); - int32_t host_layers = 0; - const char * first_buft_name = nullptr; - const char * last_buft_name = nullptr; - for (int32_t il = 0; il < n_layer; ++il) { - ggml_backend_buffer_type_t layer_buft = llama_dflash_kv_cache_layer_buft(*this, il); - if (ggml_backend_buft_is_host(layer_buft)) { - host_layers++; - } - if (first_buft_name == nullptr) { - first_buft_name = ggml_backend_buft_name(layer_buft); - } - last_buft_name = ggml_backend_buft_name(layer_buft); - - 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); - - const size_t k_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_k_ctx_cache[(size_t) il]); - ggml_backend_buffer_t k_buf = ggml_backend_buft_alloc_buffer(layer_buft, k_bytes); - if (k_buf == nullptr) { - free_dflash_kv_cache_tensors(); - return false; - } - ggml_backend_buffer_set_usage(k_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE); - ggml_backend_tensor_alloc(k_buf, dflash_k_ctx_cache[(size_t) il], ggml_backend_buffer_get_base(k_buf)); - ggml_backend_buffer_clear(k_buf, 0); - dflash_cache_bufs.push_back(k_buf); - - const size_t v_bytes = ggml_backend_buft_get_alloc_size(layer_buft, dflash_v_ctx_cache[(size_t) il]); - ggml_backend_buffer_t v_buf = ggml_backend_buft_alloc_buffer(layer_buft, v_bytes); - if (v_buf == nullptr) { - free_dflash_kv_cache_tensors(); - return false; - } - ggml_backend_buffer_set_usage(v_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE); - 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__, - target_cross_ctx, - host_layers, - n_layer, - first_buft_name != nullptr ? first_buft_name : "(none)", - last_buft_name != nullptr ? last_buft_name : "(none)"); - - return true; -} - -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; - dflash_kv_cache_reserved_rows = 0; - dflash_kv_cache_view_write_pos = 0; - dflash_kv_cache_view_n_filled = 0; - 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); - } - } - dflash_cache_bufs.clear(); - 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; @@ -5631,584 +5365,6 @@ static bool dflash_layer_has_attention_bias(const llama_layer & layer) { layer.bkv != nullptr; } -static bool validate_dflash_graph_contract(const llama_context & lctx) { - const auto & model = lctx.model; - const auto & hparams = model.hparams; - - auto rope_dim_for_layer = [&hparams](int32_t il) -> uint32_t { - if (hparams.rope_dim_per_layer[(size_t) il] != 0) { - return hparams.rope_dim_per_layer[(size_t) il]; - } - - return hparams.swa_layers[(size_t) il] ? hparams.n_rot_swa : hparams.n_rot; - }; - - auto rope_base_for_layer = [&hparams](int32_t il) -> float { - if (hparams.has_rope_freq_base_per_layer) { - return hparams.rope_freq_base_per_layer[(size_t) il]; - } - - return hparams.swa_layers[(size_t) il] ? hparams.rope_freq_base_train_swa : hparams.rope_freq_base_train; - }; - - auto rope_scale_for_layer = [&hparams](int32_t il) -> float { - return hparams.swa_layers[(size_t) il] ? hparams.rope_freq_scale_train_swa : hparams.rope_freq_scale_train; - }; - - const uint32_t ref_n_head = hparams.n_head(0); - const uint32_t ref_n_head_kv = hparams.n_head_kv(0); - const uint32_t ref_n_embd_head_k = hparams.n_embd_head_k(0); - const uint32_t ref_n_embd_head_v = hparams.n_embd_head_v(0); - const uint32_t ref_rope_dim = rope_dim_for_layer(0); - const float ref_rope_base = rope_base_for_layer(0); - const float ref_rope_scale = rope_scale_for_layer(0); - - for (int32_t il = 0; il < (int32_t) hparams.n_layer; ++il) { - if (hparams.n_head((uint32_t) il) != ref_n_head || - hparams.n_head_kv((uint32_t) il) != ref_n_head_kv || - hparams.n_embd_head_k(il) != ref_n_embd_head_k || - hparams.n_embd_head_v(il) != ref_n_embd_head_v) { - LLAMA_LOG_ERROR("%s: DFlash graph assumes layer-invariant head config, but layer %d differs (n_head=%u/%u n_head_kv=%u/%u head_k=%u/%u head_v=%u/%u)\n", - __func__, - il, - hparams.n_head((uint32_t) il), ref_n_head, - hparams.n_head_kv((uint32_t) il), ref_n_head_kv, - hparams.n_embd_head_k(il), ref_n_embd_head_k, - hparams.n_embd_head_v(il), ref_n_embd_head_v); - return false; - } - - const uint32_t rope_dim = rope_dim_for_layer(il); - const float rope_base = rope_base_for_layer(il); - const float rope_scale = rope_scale_for_layer(il); - if (rope_dim != ref_rope_dim || std::fabs(rope_base - ref_rope_base) > 1e-6f || std::fabs(rope_scale - ref_rope_scale) > 1e-6f) { - LLAMA_LOG_ERROR("%s: DFlash graph assumes layer-invariant RoPE config, but layer %d differs (dim=%u/%u base=%g/%g scale=%g/%g)\n", - __func__, - il, - rope_dim, ref_rope_dim, - (double) rope_base, (double) ref_rope_base, - (double) rope_scale, (double) ref_rope_scale); - return false; - } - - if (model.layers[(size_t) il].attn_norm == nullptr || - model.layers[(size_t) il].attn_q_norm == nullptr || - model.layers[(size_t) il].attn_k_norm == nullptr) { - LLAMA_LOG_ERROR("%s: DFlash graph requires attn_norm, attn_q_norm, and attn_k_norm weights, but layer %d is missing one or more of them\n", - __func__, il); - return false; - } - - const bool has_q_norm = model.layers[(size_t) il].attn_q_norm != nullptr; - const bool has_k_norm = model.layers[(size_t) il].attn_k_norm != nullptr; - if (has_q_norm != has_k_norm) { - LLAMA_LOG_ERROR("%s: DFlash graph requires symmetric Q/K norm presence, but layer %d has q_norm=%d k_norm=%d\n", - __func__, il, (int) has_q_norm, (int) has_k_norm); - return false; - } - - if (model.layers[(size_t) il].attn_norm_b != nullptr || - model.layers[(size_t) il].attn_q_norm_b != nullptr || - model.layers[(size_t) il].attn_k_norm_b != nullptr) { - LLAMA_LOG_ERROR("%s: DFlash graph does not implement norm-bias tensors, but layer %d requires attn_norm_b/q_norm_b/k_norm_b\n", - __func__, il); - return false; - } - - if (dflash_layer_has_attention_bias(model.layers[(size_t) il])) { - LLAMA_LOG_ERROR("%s: DFlash graph does not implement attention bias tensors, but layer %d requires them\n", - __func__, il); - return false; - } - } - - return true; -} - -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 - ? 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; - ggml_tensor * kq_mask_swa = lctx.dflash_kq_mask_swa_tensor; - - if (kq_mask == nullptr) { - LLAMA_LOG_ERROR("%s: DFlash graph inputs are not initialized\n", __func__); - return false; - } - - if (!validate_dflash_graph_contract(lctx)) { - profile.graph_shape_failures++; - 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 float * append_src = lctx.dflash_target_append_features; - const llama_pos * src_pos = lctx.dflash_target_positions; - const size_t total_floats = lctx.dflash_target_features_n_floats; - const size_t append_floats = lctx.dflash_target_append_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 append_rows_available = lctx.dflash_target_append_features_n_rows; - const int32_t width = (int32_t) lctx.model.hparams.dflash_n_target_features; - const int32_t graph_cross_ctx = use_kv_cache - ? (lctx.dflash_k_ctx_cache.front() != nullptr ? (int32_t) lctx.dflash_k_ctx_cache.front()->ne[2] : 0) - : (lctx.inp_dflash_target_features != nullptr ? (int32_t) lctx.inp_dflash_target_features->ne[1] : 0); - const int32_t n_mask_tokens = (int32_t) kq_mask->ne[1]; - const int32_t n_kv_total = (int32_t) kq_mask->ne[0]; - const int64_t t_total_us = ggml_time_us(); - - 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 (use_kv_workspace) { - llama_sync_dflash_workspace_if_pending(lctx); - } - - 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 (n_rows <= 0) { - profile.graph_shape_failures++; - LLAMA_LOG_ERROR("%s: missing DFlash target feature rows\n", __func__); - return false; - } - - const bool have_full_src = src != nullptr && total_floats == (size_t) n_rows * (size_t) width; - if (n_rows > cross_ctx || (src != nullptr && !have_full_src)) { - 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; - } - - if (!use_kv_cache && !have_full_src) { - profile.graph_shape_failures++; - LLAMA_LOG_ERROR("%s: missing contiguous DFlash target features for inline path\n", __func__); - return false; - } - - if (n_kv_total < cross_ctx + (int32_t) n_tokens) { - profile.graph_mask_overflow++; - LLAMA_LOG_ERROR("%s: invalid DFlash mask shape (n_kv_total=%d < cross_ctx+n_tokens=%d)\n", - __func__, n_kv_total, cross_ctx + (int32_t) n_tokens); - return false; - } - - const int32_t left_pad = cross_ctx - n_rows; - profile.last_left_pad = left_pad; - if (!use_kv_cache) { - const size_t padded_floats = (size_t) cross_ctx * (size_t) width; - const size_t dst_offset = (size_t) left_pad * (size_t) width; - const int64_t t_feature_us = ggml_time_us(); - if (lctx.dflash_target_features_padded.size() != padded_floats) { - lctx.dflash_target_features_padded.resize(padded_floats); - } - if (left_pad == 0 && total_floats == padded_floats) { - std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin()); - } else { - if (dst_offset > 0) { - std::fill(lctx.dflash_target_features_padded.begin(), - lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset, 0.0f); - } - std::copy(src, src + total_floats, lctx.dflash_target_features_padded.begin() + (ptrdiff_t) dst_offset); - } - profile.graph_feature_copy_us += (uint64_t) (ggml_time_us() - t_feature_us); - profile.graph_feature_bytes += padded_floats * sizeof(float); - } - - const int64_t t_pos_us = ggml_time_us(); - lctx.dflash_pos_ctx_data.resize((size_t) cross_ctx); - 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) { - profile.graph_pos_fallbacks++; - profile.graph_shape_failures++; - profile.last_pos_first = -1; - profile.last_pos_last = -1; - if (profile.graph_pos_fallbacks <= 3) { - LLAMA_LOG_ERROR("%s: missing DFlash target positions (rows=%d positions=%zu cross_ctx=%d)\n", - __func__, n_rows, total_positions, cross_ctx); - } - return false; - } - - profile.last_pos_first = src_pos[0]; - profile.last_pos_last = src_pos[n_rows - 1]; - for (int32_t i = 1; i < n_rows; ++i) { - if (src_pos[i] <= src_pos[i - 1]) { - profile.graph_pos_non_monotonic++; - profile.graph_shape_failures++; - if (profile.graph_pos_non_monotonic <= 3) { - LLAMA_LOG_ERROR("%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]); - } - return false; - } - } - std::copy(src_pos, src_pos + n_rows, lctx.dflash_pos_ctx_data.begin() + (ptrdiff_t) left_pad); - profile.graph_pos_copy_us += (uint64_t) (ggml_time_us() - t_pos_us); - profile.graph_pos_bytes += lctx.dflash_pos_ctx_data.size() * sizeof(llama_pos); - - if (use_kv_cache) { - const llama_dflash_kv_cache_transition cache_plan = llama_plan_dflash_kv_cache_transition( - cross_ctx, - lctx.dflash_kv_cache_n_filled, - lctx.dflash_kv_cache_write_pos, - lctx.dflash_kv_cache_valid, - lctx.dflash_kv_cache_applied_window_version, - lctx.dflash_target_window_version, - lctx.dflash_target_window_keep_rows, - lctx.dflash_target_window_append_rows, - lctx.dflash_target_window_replace, - n_rows); - - const bool have_append_src = append_src != nullptr && - append_rows_available == cache_plan.append_rows && - append_floats == (size_t) cache_plan.append_rows * (size_t) width; - - const int32_t update_rows = cache_plan.cache_up_to_date - ? 0 - : (cache_plan.rebuild_cache ? n_rows : cache_plan.append_rows); - const size_t max_nodes = lctx.model.max_nodes((int) std::max(1, cross_ctx)) + 24 * lctx.model.hparams.n_layer; - const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false); - if (lctx.dflash_buf_compute_meta.size() != meta_size) { - lctx.dflash_buf_compute_meta.resize(meta_size); - } - - if (lctx.dflash_sched == nullptr || lctx.dflash_kv_cache_reserved_rows != cross_ctx) { - std::vector backend_buft; - backend_buft.reserve(lctx.backends.size()); - for (auto * backend : lctx.backends) { - if (ggml_backend_is_cpu(backend)) { - backend_buft.push_back(llama_default_buffer_type_cpu(true)); - } else { - backend_buft.push_back(ggml_backend_get_default_buffer_type(backend)); - } - } - - if (lctx.dflash_sched != nullptr) { - ggml_backend_sched_free(lctx.dflash_sched); - lctx.dflash_sched = nullptr; - } - lctx.dflash_kv_graph = nullptr; - lctx.dflash_kv_graph_rows = 0; - lctx.dflash_kv_graph_write_pos = 0; - - const int32_t saved_update_rows = lctx.dflash_kv_cache_update_rows; - lctx.dflash_kv_cache_update_rows = cross_ctx; - const int64_t t_build_us = ggml_time_us(); - ggml_cgraph * gf_reserve = llm_build_context::llama_build_graph_dflash_kv_cache(lctx); - profile.graph_kv_cache_build_us += (uint64_t) (ggml_time_us() - t_build_us); - lctx.dflash_kv_cache_update_rows = saved_update_rows; - if (gf_reserve == nullptr) { - profile.graph_shape_failures++; - LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache reserve graph\n", __func__); - return false; - } - - const int64_t t_reserve_us = ggml_time_us(); - lctx.dflash_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false); - const bool reserved = lctx.dflash_sched != nullptr && ggml_backend_sched_reserve(lctx.dflash_sched, gf_reserve); - profile.graph_kv_cache_reserve_us += (uint64_t) (ggml_time_us() - t_reserve_us); - if (!reserved) { - profile.graph_shape_failures++; - LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V scheduler\n", __func__); - return false; - } - lctx.dflash_kv_cache_reserved_rows = cross_ctx; - } - - if (update_rows > 0) { - const float * update_src = nullptr; - if (have_append_src && update_rows == cache_plan.append_rows) { - update_src = append_src; - } else if (have_full_src) { - update_src = src + (size_t) (n_rows - update_rows) * (size_t) width; - } - const llama_pos * update_pos = src_pos + (n_rows - update_rows); - - if (update_src == nullptr) { - profile.graph_shape_failures++; - LLAMA_LOG_ERROR("%s: missing DFlash appended target features for cached update (rows=%d append_rows=%d floats=%zu)\n", - __func__, n_rows, update_rows, append_floats); - return false; - } - - if (cache_plan.rebuild_cache) { - llama_reset_dflash_kv_cache_state(&lctx); - } - - lctx.dflash_kv_cache_update_rows = update_rows; - ggml_cgraph * gf_kv = nullptr; - const bool can_reuse_kv_graph = lctx.dflash_kv_graph != nullptr && - lctx.dflash_kv_graph_rows == update_rows && - lctx.dflash_kv_graph_write_pos == lctx.dflash_kv_cache_write_pos; - if (can_reuse_kv_graph) { - gf_kv = lctx.dflash_kv_graph; - } else { - const int64_t t_build_us = ggml_time_us(); - 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; - } - - ggml_backend_t kv_feature_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_target_features); - const int64_t t_feature_upload_us = ggml_time_us(); - if (kv_feature_backend != nullptr) { - ggml_backend_tensor_set_async(kv_feature_backend, lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features)); - } else { - ggml_backend_tensor_set(lctx.dflash_kv_input_target_features, update_src, 0, ggml_nbytes(lctx.dflash_kv_input_target_features)); - } - profile.graph_kv_cache_feature_upload_us += (uint64_t) (ggml_time_us() - t_feature_upload_us); - profile.graph_feature_bytes += (size_t) update_rows * (size_t) width * sizeof(float); - - ggml_backend_t kv_pos_backend = llama_backend_for_tensor(lctx, lctx.dflash_kv_input_pos_ctx); - const int64_t t_pos_upload_us = ggml_time_us(); - if (kv_pos_backend != nullptr) { - ggml_backend_tensor_set_async(kv_pos_backend, lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx)); - } else { - ggml_backend_tensor_set(lctx.dflash_kv_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash_kv_input_pos_ctx)); - } - profile.graph_kv_cache_pos_upload_us += (uint64_t) (ggml_time_us() - t_pos_upload_us); - - const int64_t t_kv_cache_us = ggml_time_us(); - llama_dflash_kv_node_profiler kv_node_profiler; - if (kv_node_timing) { - kv_node_profiler.profile = &profile; - ggml_backend_sched_set_eval_callback(lctx.dflash_sched, llama_dflash_kv_node_eval_callback, &kv_node_profiler); - } - llama_graph_compute_sched(lctx, lctx.dflash_sched, gf_kv, lctx.cparams.n_threads); - if (kv_node_timing) { - ggml_backend_sched_set_eval_callback(lctx.dflash_sched, nullptr, nullptr); - } - profile.graph_kv_cache_compute_us += (uint64_t) (ggml_time_us() - t_kv_cache_us); - - const int64_t t_sync_us = ggml_time_us(); - ggml_backend_sched_synchronize(lctx.dflash_sched); - profile.graph_kv_cache_sync_us += (uint64_t) (ggml_time_us() - t_sync_us); - profile.graph_kv_cache_calls++; - - lctx.dflash_kv_cache_n_filled = std::min(cross_ctx, lctx.dflash_kv_cache_n_filled + update_rows); - lctx.dflash_kv_cache_write_pos = (lctx.dflash_kv_cache_write_pos + update_rows) % cross_ctx; - lctx.dflash_kv_cache_applied_window_version = lctx.dflash_target_window_version; - lctx.dflash_kv_cache_valid = true; - lctx.dflash_kv_cache_view_n_filled = lctx.dflash_kv_cache_n_filled; - lctx.dflash_kv_cache_view_write_pos = lctx.dflash_kv_cache_write_pos; - lctx.dflash_kv_cache_view_valid = true; - } - - if (use_kv_workspace && lctx.dflash_kv_cache_view_valid && - !lctx.dflash_k_ctx_workspace.empty() && !lctx.dflash_v_ctx_workspace.empty()) { - const bool need_workspace_refresh = !lctx.dflash_kv_workspace_valid || - lctx.dflash_kv_workspace_n_filled != lctx.dflash_kv_cache_view_n_filled || - lctx.dflash_kv_workspace_write_pos != lctx.dflash_kv_cache_view_write_pos || - lctx.dflash_kv_workspace_applied_window_version != lctx.dflash_kv_cache_applied_window_version; - - if (need_workspace_refresh) { - const size_t max_nodes = lctx.model.max_nodes((int) std::max(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)); - } - - const int64_t t_mask_us = ggml_time_us(); - const int32_t full_visible_first = left_pad; - const int32_t full_visible_last = cross_ctx + (int32_t) n_tokens - 1; - 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) n_tokens; - visible_kv_max = std::max(visible_kv_max, visible_kv); - profile.graph_visible_kv_sum += (uint64_t) visible_kv; - for (int32_t i = full_visible_first; i <= full_visible_last; ++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); - - if (kq_mask_swa != nullptr) { - lctx.dflash_kq_mask_swa_data.assign((size_t) n_kv_total * (size_t) n_mask_tokens, -INFINITY); - const int32_t swa_window = (int32_t) lctx.model.hparams.n_swa; - const int32_t draft_pos_base = (int32_t) profile.last_pos_last; - for (uint32_t j = 0; j < n_tokens; ++j) { - float * row = lctx.dflash_kq_mask_swa_data.data() + (size_t) j * (size_t) n_kv_total; - const int32_t q_pos = draft_pos_base + (int32_t) j; - - for (int32_t k = left_pad; k < cross_ctx; ++k) { - const int32_t k_pos = (int32_t) lctx.dflash_pos_ctx_data[(size_t) k]; - if (q_pos - k_pos < swa_window) { - row[k] = 0.0f; - } - } - - for (int32_t k = cross_ctx; k < cross_ctx + (int32_t) n_tokens; ++k) { - const int32_t block_k = k - cross_ctx; - if (block_k <= (int32_t) j) { - row[k] = 0.0f; - } - } - } - - ggml_backend_tensor_set(kq_mask_swa, lctx.dflash_kq_mask_swa_data.data(), 0, ggml_nbytes(kq_mask_swa)); - profile.graph_mask_bytes += ggml_nbytes(kq_mask_swa); - } - - 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) { - int32_t n_swa_layers = 0; - for (int32_t il = 0; il < lctx.model.hparams.n_layer; ++il) { - n_swa_layers += lctx.model.hparams.swa_layers[(size_t) il] ? 1 : 0; - } - - 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] full_mask=[%d..%d] swa_window=%u swa_layers=%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, - full_visible_first, full_visible_last, - lctx.model.hparams.n_swa, - n_swa_layers); - } - - return true; -} - // decode a batch of tokens by evaluating the transformer // // - lctx: llama context @@ -6548,7 +5704,7 @@ static int llama_decode_internal( if (dflash_profile != nullptr) { dflash_profile->decode_prepare_calls++; const int64_t t_prepare_dflash_us = ggml_time_us(); - if (!prepare_dflash_graph_inputs(lctx, n_tokens)) { + if (!llama_prepare_dflash_graph_inputs(lctx, n_tokens)) { dflash_profile->decode_prepare_failures++; dflash_profile->decode_prepare_us += (uint64_t) (ggml_time_us() - t_prepare_dflash_us); return GGML_STATUS_FAILED;