furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 10:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

improve DFlash caching and profiling capabilities

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