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

move dflash for it own file

Browse Source
This commit is contained in:
SamuelOliveirads
2026-06-02 10:22:13 -03:00
parent 1250f522ed
commit dc43cdf06b
13 changed files with 4533 additions and 4052 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/speculative-impl.h
+1740
View File
File diff suppressed because it is too large Load Diff
common/speculative.cpp
+2 -1739
View File
File diff suppressed because it is too large Load Diff
src/CMakeLists.txt
+2
View File
@@ -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
src/llama-build-context.cpp
+5 -5
View File
@@ -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();
src/llama-context.h
+154 -73
View File
@@ -278,77 +278,162 @@ struct llama_context {
size_t draft_input_hidden_state_n_floats = 0;
std::vector<float> 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<float> dflash_target_features_owned;
std::vector<float> dflash_target_append_features_owned;
std::vector<llama_pos> dflash_target_positions_owned;
std::vector<float> dflash_target_features_padded;
std::vector<float> dflash_feature_view_buffer;
std::vector<llama_pos> dflash_pos_ctx_data;
std::vector<float> dflash_kq_mask_data;
std::vector<float> dflash_kq_mask_swa_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;
std::vector<struct ggml_tensor *> dflash_k_ctx_workspace;
std::vector<struct ggml_tensor *> dflash_v_ctx_workspace;
struct ggml_context * dflash_cache_ctx = nullptr;
std::vector<ggml_backend_buffer_t> dflash_cache_bufs;
int32_t dflash_kv_cache_write_pos = 0;
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<uint8_t> dflash_buf_compute_meta;
std::vector<uint8_t> dflash_workspace_buf_compute_meta;
ggml_backend_sched_t dflash_sched = nullptr;
ggml_backend_sched_t dflash_workspace_sched = nullptr;
ggml_cgraph * dflash_kv_graph = nullptr;
ggml_cgraph * dflash_kv_workspace_graph = nullptr;
int32_t dflash_kv_graph_rows = 0;
int32_t dflash_kv_graph_write_pos = 0;
int32_t dflash_kv_workspace_graph_rows = 0;
int32_t dflash_kv_workspace_graph_write_pos = 0;
struct ggml_tensor * dflash_kv_input_target_features = nullptr;
struct ggml_tensor * dflash_kv_input_pos_ctx = nullptr;
struct ggml_tensor * dflash_kq_mask_tensor = nullptr;
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<float> features_owned;
std::vector<float> append_features_owned;
std::vector<llama_pos> positions_owned;
std::vector<float> features_padded;
std::vector<llama_pos> pos_ctx_data;
std::vector<float> kq_mask_data;
std::vector<float> kq_mask_swa_data;
};
struct dflash_capture_state {
std::vector<int32_t> layer_ids;
std::vector<std::vector<float>> layer_rows;
int32_t row_count = 0;
int32_t row_width = 0;
uint64_t capture_batch_id = 0;
std::vector<uint64_t> 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<struct ggml_tensor *> k_ctx_cache;
std::vector<struct ggml_tensor *> v_ctx_cache;
std::vector<struct ggml_tensor *> k_ctx_workspace;
std::vector<struct ggml_tensor *> v_ctx_workspace;
struct ggml_context * cache_ctx = nullptr;
std::vector<ggml_backend_buffer_t> 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<uint8_t> cache_compute_meta;
std::vector<uint8_t> 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<int32_t> layer_ids;
std::vector<std::vector<float>> layer_rows;
int32_t row_count = 0;
int32_t row_width = 0;
uint64_t capture_batch_id = 0;
std::vector<uint64_t> 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_state> capture;
std::vector<float> feature_view_buffer;
input_state inputs;
int32_t visible_cross_ctx = 0;
llama_dflash_profile_stats profile;
};
std::unique_ptr<dflash_capture_state> 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<float> & dflash_target_features_owned = dflash.target.features_owned;
std::vector<float> & dflash_target_append_features_owned = dflash.target.append_features_owned;
std::vector<llama_pos> & dflash_target_positions_owned = dflash.target.positions_owned;
std::vector<float> & dflash_target_features_padded = dflash.target.features_padded;
std::vector<float> & dflash_feature_view_buffer = dflash.feature_view_buffer;
std::vector<llama_pos> & dflash_pos_ctx_data = dflash.target.pos_ctx_data;
std::vector<float> & dflash_kq_mask_data = dflash.target.kq_mask_data;
std::vector<float> & dflash_kq_mask_swa_data = dflash.target.kq_mask_swa_data;
int32_t & dflash_visible_cross_ctx = dflash.visible_cross_ctx;
std::vector<struct ggml_tensor *> & dflash_k_ctx_cache = dflash.kv.k_ctx_cache;
std::vector<struct ggml_tensor *> & dflash_v_ctx_cache = dflash.kv.v_ctx_cache;
std::vector<struct ggml_tensor *> & dflash_k_ctx_workspace = dflash.kv.k_ctx_workspace;
std::vector<struct ggml_tensor *> & dflash_v_ctx_workspace = dflash.kv.v_ctx_workspace;
struct ggml_context * & dflash_cache_ctx = dflash.kv.cache_ctx;
std::vector<ggml_backend_buffer_t> & 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<uint8_t> & dflash_buf_compute_meta = dflash.kv.cache_compute_meta;
std::vector<uint8_t> & 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_state> & 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);
src/llama-dflash.cpp
+1240
View File
File diff suppressed because it is too large Load Diff
src/llama-dflash.h
+8
View File
@@ -0,0 +1,8 @@
#pragma once
#include <cstdint>
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);
src/llama-quantize.cpp
+3 -2
View File
@@ -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;
}
}
src/llama-spec-features-dflash.cpp
+1097
View File
File diff suppressed because it is too large Load Diff
src/llama-spec-features-dflash.h
+279
View File
@@ -0,0 +1,279 @@
#pragma once
#include "llama.h"
#include <algorithm>
#include <cstdint>
#include <vector>
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<int32_t>(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<int32_t>(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<int32_t> & output_indices,
std::vector<float> & hidden_rows);
src/llama-spec-features.cpp
-1084
View File
File diff suppressed because it is too large Load Diff
src/llama-spec-features.h
+1 -303
View File
@@ -2,7 +2,6 @@
#include "llama.h"
#include <algorithm>
#include <cstdint>
#include <vector>
@@ -25,316 +24,20 @@ struct llama_spec_feature_view {
std::vector<llama_spec_feature_row_view> 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<int32_t>(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<int32_t>(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<int32_t> & output_indices,
std::vector<float> & hidden_rows);
bool llama_spec_copy_dflash_rows_from_output_indices(
struct llama_context * ctx,
const std::vector<int32_t> & output_indices,
std::vector<float> & hidden_rows);
src/llama.cpp
+2 -846
View File
@@ -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<rpc_device> extract_device_from_rpc_device(std::vector<std::string> 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<int32_t>(1, cross_ctx);
const int32_t target_token_capacity = std::max<int32_t>(1, (int32_t) model.hparams.dflash_block_size);
const int32_t target_workspace_n_kv_total = GGML_PAD(target_cross_ctx + target_token_capacity, cparams.flash_attn ? 256 : 32);
const int32_t n_layer = model.hparams.n_layer;
const int64_t n_embd_head_k = model.hparams.n_embd_head_k(0);
const int64_t n_embd_head_v = model.hparams.n_embd_head_v(0);
const int64_t n_head_kv = model.hparams.n_head_kv();
if (dflash_cache_ctx != nullptr && !dflash_k_ctx_cache.empty()) {
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<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;
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<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);
}
if (lctx.dflash_sched == nullptr || lctx.dflash_kv_cache_reserved_rows != cross_ctx) {
std::vector<ggml_backend_buffer_type_t> backend_buft;
backend_buft.reserve(lctx.backends.size());
for (auto * backend : lctx.backends) {
if (ggml_backend_is_cpu(backend)) {
backend_buft.push_back(llama_default_buffer_type_cpu(true));
} else {
backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
}
}
if (lctx.dflash_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<int32_t>(1, cross_ctx)) + 16 * lctx.model.hparams.n_layer;
const size_t meta_size = ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false);
if (lctx.dflash_workspace_buf_compute_meta.size() != meta_size) {
lctx.dflash_workspace_buf_compute_meta.resize(meta_size);
}
ggml_cgraph * gf_workspace = nullptr;
const bool can_reuse_workspace_graph = lctx.dflash_kv_workspace_graph != nullptr &&
lctx.dflash_kv_workspace_graph_rows == lctx.dflash_kv_cache_view_n_filled &&
lctx.dflash_kv_workspace_graph_write_pos == lctx.dflash_kv_cache_view_write_pos;
if (can_reuse_workspace_graph) {
gf_workspace = lctx.dflash_kv_workspace_graph;
} else {
const int64_t t_build_us = ggml_time_us();
gf_workspace = llm_build_context::llama_build_graph_dflash_kv_workspace(lctx);
profile.graph_kv_workspace_build_us += (uint64_t) (ggml_time_us() - t_build_us);
if (gf_workspace == nullptr) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to build DFlash K/V workspace graph\n", __func__);
return false;
}
std::vector<ggml_backend_buffer_type_t> backend_buft;
backend_buft.reserve(lctx.backends.size());
for (auto * backend : lctx.backends) {
if (ggml_backend_is_cpu(backend)) {
backend_buft.push_back(llama_default_buffer_type_cpu(true));
} else {
backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
}
}
if (lctx.dflash_workspace_sched == nullptr) {
lctx.dflash_workspace_sched = ggml_backend_sched_new(lctx.backends.data(), backend_buft.data(), lctx.backends.size(), max_nodes, false);
}
if (lctx.dflash_kv_workspace_reserved_rows != cross_ctx) {
const bool saved_view_valid = lctx.dflash_kv_cache_view_valid;
const int32_t saved_view_rows = lctx.dflash_kv_cache_view_n_filled;
const int32_t saved_view_write_pos = lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_cache_view_valid = true;
lctx.dflash_kv_cache_view_n_filled = cross_ctx;
lctx.dflash_kv_cache_view_write_pos = cross_ctx > 1 ? 1 : 0;
const int64_t t_reserve_build_us = ggml_time_us();
ggml_cgraph * gf_workspace_reserve = llm_build_context::llama_build_graph_dflash_kv_workspace(lctx);
profile.graph_kv_workspace_build_us += (uint64_t) (ggml_time_us() - t_reserve_build_us);
lctx.dflash_kv_cache_view_valid = saved_view_valid;
lctx.dflash_kv_cache_view_n_filled = saved_view_rows;
lctx.dflash_kv_cache_view_write_pos = saved_view_write_pos;
const int64_t t_reserve_us = ggml_time_us();
const bool reserved = lctx.dflash_workspace_sched != nullptr &&
gf_workspace_reserve != nullptr &&
ggml_backend_sched_reserve(lctx.dflash_workspace_sched, gf_workspace_reserve);
profile.graph_kv_workspace_reserve_us += (uint64_t) (ggml_time_us() - t_reserve_us);
if (!reserved) {
profile.graph_shape_failures++;
LLAMA_LOG_ERROR("%s: failed to initialize DFlash K/V workspace scheduler\n", __func__);
return false;
}
lctx.dflash_kv_workspace_reserved_rows = cross_ctx;
}
const int64_t t_reset_us = ggml_time_us();
ggml_backend_sched_reset(lctx.dflash_workspace_sched);
profile.graph_kv_workspace_reset_us += (uint64_t) (ggml_time_us() - t_reset_us);
const int64_t t_alloc_us = ggml_time_us();
ggml_backend_sched_alloc_graph(lctx.dflash_workspace_sched, gf_workspace);
profile.graph_kv_workspace_alloc_us += (uint64_t) (ggml_time_us() - t_alloc_us);
lctx.dflash_kv_workspace_graph = gf_workspace;
lctx.dflash_kv_workspace_graph_rows = lctx.dflash_kv_cache_view_n_filled;
lctx.dflash_kv_workspace_graph_write_pos = lctx.dflash_kv_cache_view_write_pos;
}
const int64_t t_workspace_us = ggml_time_us();
llama_graph_compute_sched(lctx, lctx.dflash_workspace_sched, gf_workspace, lctx.cparams.n_threads);
profile.graph_kv_workspace_compute_us += (uint64_t) (ggml_time_us() - t_workspace_us);
lctx.dflash_kv_workspace_sync_pending = true;
profile.graph_kv_workspace_calls++;
lctx.dflash_kv_workspace_n_filled = lctx.dflash_kv_cache_view_n_filled;
lctx.dflash_kv_workspace_write_pos = lctx.dflash_kv_cache_view_write_pos;
lctx.dflash_kv_workspace_applied_window_version = lctx.dflash_kv_cache_applied_window_version;
lctx.dflash_kv_workspace_valid = true;
}
}
} else {
ggml_backend_tensor_set(lctx.inp_dflash_target_features, lctx.dflash_target_features_padded.data(), 0, ggml_nbytes(lctx.inp_dflash_target_features));
ggml_backend_tensor_set(lctx.inp_dflash_pos_ctx, lctx.dflash_pos_ctx_data.data(), 0, ggml_nbytes(lctx.inp_dflash_pos_ctx));
}
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<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) {
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;