Skip to content

[Draft] feat: implement paged KV cache and attention #22569

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
matiaslin wants to merge 1 commit into
base: master
Choose a base branch
from
+4,029 −70
Draft

[Draft] feat: implement paged KV cache and attention #22569

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 45 additions & 2 deletions common/arg.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1328,7 +1328,50 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
[](common_params & params, bool value) {
params.kv_unified = value;
}
).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL}));
).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL, LLAMA_EXAMPLE_PAGED}));
add_opt(common_arg(
{"-kvp", "--kv-paged"},
{"-no-kvp", "--no-kv-paged"},
"use paged KV buffer shared across all sequences (default: disabled)",
[](common_params & params, bool value) {
params.kv_paged = value;
}
).set_env("LLAMA_ARG_KV_PAGED").set_examples({LLAMA_EXAMPLE_PAGED}));
add_opt(common_arg(
{"-ncpub", "--n-cpu-blocks"}, "N",
"number of physical CPU blocks for paged KV cache (default: 1)",
[](common_params & params, int value) {
params.n_cpu_blocks = value;
}
).set_examples({LLAMA_EXAMPLE_PAGED}));
add_opt(common_arg(
{"-ngpub", "--n-gpu-blocks"}, "N",
"number of physical GPU blocks for paged KV cache (default: 1)",
[](common_params & params, int value) {
params.n_gpu_blocks = value;
}
).set_examples({LLAMA_EXAMPLE_PAGED}));
add_opt(common_arg(
{"-kvbls", "--kv-block-size"}, "N",
"fixed number of tokens for a given paged block (default: 16)",
[](common_params & params, int value) {
if (value <= 0 || (value & (value - 1)) != 0) {
throw std::invalid_argument("--kv-block-size must be a positive power of 2");
}
params.block_size = value;
}
).set_examples({LLAMA_EXAMPLE_PAGED}));
add_opt(common_arg(
{"--kv-paged-watermark"}, "N",
"fraction of blocks reserved before processing new requests (default: 0.05, range [0.0, 1.0))",
[](common_params & params, const std::string & value) {
float potential_watermark = std::stof(value);
if (potential_watermark < 0.0f || potential_watermark >= 1.0f) {
throw std::invalid_argument("--kv-paged-watermark must be in range [0.0, 1.0)");
}
params.kv_paged_watermark = potential_watermark;
}
).set_examples({LLAMA_EXAMPLE_PAGED}));
add_opt(common_arg(
{"--cache-idle-slots"},
{"--no-cache-idle-slots"},
Expand Down Expand Up @@ -2152,7 +2195,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
[](common_params & params, int value) {
params.n_sequences = value;
}
).set_examples({LLAMA_EXAMPLE_PARALLEL}));
).set_examples({LLAMA_EXAMPLE_PARALLEL, LLAMA_EXAMPLE_PAGED}));
add_opt(common_arg(
{"-cb", "--cont-batching"},
{"-nocb", "--no-cont-batching"},
Expand Down
113 changes: 87 additions & 26 deletions common/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1141,6 +1141,56 @@ struct common_init_result::impl {
std::vector<llama_sampler_seq_config> samplers_seq_config;
};

static void common_fit_paged_kv_blocks(common_params& params, const llama_model * model) {
GGML_ASSERT(model && "model must be loaded before fitting paged KV blocks.");
ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU);
if (!dev) {
LOG_WRN("%s: no GPU device found, cannot fit paged KV blocks.\n", __func__);
return;
}

size_t free_vram = 0;
size_t total_vram = 0;
ggml_backend_dev_memory(dev, &free_vram, &total_vram);

const uint32_t n_heads_kv = llama_model_n_head_kv(model);
const uint32_t n_layers = llama_model_n_layer(model);
const uint32_t head_dim = llama_model_n_embd(model) / llama_model_n_head(model);
const uint32_t block_size = params.block_size;

const size_t bytes_per_block = (size_t)2 * head_dim * n_heads_kv * block_size * n_layers * ggml_type_size(GGML_TYPE_F16);

const size_t margin = params.fit_params_target.empty()
? (size_t)(total_vram * 0.05f)
: (size_t)params.fit_params_target[0];

if (free_vram <= margin) {
LOG_ERR("%s: not enough free VRAM for paged KV blocks. "
"free_vram=%.1f MiB <= margin=%.1f MiB. "
"Try reducing --margin or offloading fewer layers to GPU.\n",
__func__, free_vram / 1024.0f / 1024.0f, margin / 1024.0f / 1024.0f);
return; // leave params.n_gpu_blocks at its existing value
}

const size_t available = (free_vram > margin) ? free_vram - margin : 0;

if (bytes_per_block == 0 || available < bytes_per_block) {
LOG_ERR("%s: available VRAM (%.1f MiB) is less than one block (%.1f MiB). "
"Try increasing n_gpu_blocks manually or reducing block_size.\n",
__func__, available / 1024.0f / 1024.0f, bytes_per_block / 1024.0f / 1024.0f);
return;
}

const uint32_t n_gpu_blocks = (uint32_t)(available / bytes_per_block);
const uint32_t n_cpu_blocks = (uint32_t)(n_gpu_blocks * params.cpu_to_gpu_blocks_ratio);

LOG_INF("%s: free_vram=%0.1f MiB, bytes_per_block=%ld, n_gpu_blocks=%d, n_cpu_blocks=%d\n",
__func__, free_vram / 1024.0f / 1024.0f, bytes_per_block, n_gpu_blocks, n_cpu_blocks);

params.n_gpu_blocks = n_gpu_blocks;
params.n_cpu_blocks = n_cpu_blocks;
}

common_init_result::common_init_result(common_params & params) :
pimpl(new impl{}) {
auto mparams = common_model_params_to_llama(params);
Expand All @@ -1163,6 +1213,12 @@ common_init_result::common_init_result(common_params & params) :

pimpl->model.reset(model);

if (params.fit_params && params.kv_paged) {
LOG_INF("%s: fitting KV paged params to device memory\n", __func__);
common_fit_paged_kv_blocks(params, pimpl->model.get());
cparams = common_context_params_to_llama(params); // re-derive this params to reflect changes
}

const llama_vocab * vocab = llama_model_get_vocab(model);

// load and optionally apply lora adapters
Expand Down Expand Up @@ -1488,32 +1544,37 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
struct llama_context_params common_context_params_to_llama(const common_params & params) {
auto cparams = llama_context_default_params();

cparams.n_ctx = params.n_ctx;
cparams.n_seq_max = params.n_parallel;
cparams.n_batch = params.n_batch;
cparams.n_ubatch = params.n_ubatch;
cparams.n_threads = params.cpuparams.n_threads;
cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ?
params.cpuparams.n_threads : params.cpuparams_batch.n_threads;
cparams.embeddings = params.embedding;
cparams.rope_scaling_type = params.rope_scaling_type;
cparams.rope_freq_base = params.rope_freq_base;
cparams.rope_freq_scale = params.rope_freq_scale;
cparams.yarn_ext_factor = params.yarn_ext_factor;
cparams.yarn_attn_factor = params.yarn_attn_factor;
cparams.yarn_beta_fast = params.yarn_beta_fast;
cparams.yarn_beta_slow = params.yarn_beta_slow;
cparams.yarn_orig_ctx = params.yarn_orig_ctx;
cparams.pooling_type = params.pooling_type;
cparams.attention_type = params.attention_type;
cparams.flash_attn_type = params.flash_attn_type;
cparams.cb_eval = params.cb_eval;
cparams.cb_eval_user_data = params.cb_eval_user_data;
cparams.offload_kqv = !params.no_kv_offload;
cparams.no_perf = params.no_perf;
cparams.op_offload = !params.no_op_offload;
cparams.swa_full = params.swa_full;
cparams.kv_unified = params.kv_unified;
cparams.n_ctx = params.n_ctx;
cparams.n_seq_max = params.n_parallel;
cparams.n_batch = params.n_batch;
cparams.n_ubatch = params.n_ubatch;
cparams.n_threads = params.cpuparams.n_threads;
cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ?
params.cpuparams.n_threads : params.cpuparams_batch.n_threads;
cparams.embeddings = params.embedding;
cparams.rope_scaling_type = params.rope_scaling_type;
cparams.rope_freq_base = params.rope_freq_base;
cparams.rope_freq_scale = params.rope_freq_scale;
cparams.yarn_ext_factor = params.yarn_ext_factor;
cparams.yarn_attn_factor = params.yarn_attn_factor;
cparams.yarn_beta_fast = params.yarn_beta_fast;
cparams.yarn_beta_slow = params.yarn_beta_slow;
cparams.yarn_orig_ctx = params.yarn_orig_ctx;
cparams.pooling_type = params.pooling_type;
cparams.attention_type = params.attention_type;
cparams.flash_attn_type = params.flash_attn_type;
cparams.cb_eval = params.cb_eval;
cparams.cb_eval_user_data = params.cb_eval_user_data;
cparams.offload_kqv = !params.no_kv_offload;
cparams.no_perf = params.no_perf;
cparams.op_offload = !params.no_op_offload;
cparams.swa_full = params.swa_full;
cparams.kv_unified = params.kv_unified;
cparams.kv_paged = params.kv_paged;
cparams.block_size = params.block_size;
cparams.n_gpu_blocks = params.n_gpu_blocks;
cparams.n_cpu_blocks = params.n_cpu_blocks;
cparams.kv_paged_watermark = params.kv_paged_watermark;

cparams.type_k = params.cache_type_k;
cparams.type_v = params.cache_type_v;
Expand Down
11 changes: 10 additions & 1 deletion common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ enum llama_example {
LLAMA_EXAMPLE_FIT_PARAMS,
LLAMA_EXAMPLE_RESULTS,
LLAMA_EXAMPLE_EXPORT_GRAPH_OPS,

LLAMA_EXAMPLE_PAGED,
LLAMA_EXAMPLE_COUNT,
};

Expand Down Expand Up @@ -539,6 +539,15 @@ struct common_params {
bool ctx_shift = false; // context shift on infinite text generation
bool swa_full = false; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
bool kv_unified = false; // enable unified KV cache
bool kv_paged = false; // enable paged KV cache

int32_t block_size = 16;

uint32_t n_gpu_blocks = 1;
uint32_t n_cpu_blocks = 1;

float cpu_to_gpu_blocks_ratio = 0.25;
float kv_paged_watermark = 0.05; // percentage

bool input_prefix_bos = false; // prefix BOS to user inputs, preceding input_prefix
bool use_mmap = true; // enable mmap to use filesystem cache
Expand Down
2 changes: 2 additions & 0 deletions examples/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ else()
add_subdirectory(gen-docs)
add_subdirectory(training)
add_subdirectory(diffusion)
add_subdirectory(paged)
add_subdirectory(continuous-batch)
if (NOT GGML_BACKEND_DL)
add_subdirectory(convert-llama2c-to-ggml)
# these examples use the backends directly and cannot be built with dynamic loading
Expand Down
5 changes: 5 additions & 0 deletions examples/continuous-batch/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
set(TARGET llama-continuous-batch)
add_executable(${TARGET} continuous-batch.cpp)
install(TARGETS ${TARGET} RUNTIME)
target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT})
target_compile_features(${TARGET} PRIVATE cxx_std_17)
Loading