#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_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_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);
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);
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);