diff --git a/common/arg.cpp b/common/arg.cpp index 72750a3cba0a..4297a65036b4 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1305,6 +1305,38 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str())); } }).set_env("LLAMA_ARG_FLASH_ATTN")); + // KV cache sizing - allocate space for fewer tokens when full context isn't needed + // Note: This reduces upfront allocation, NOT per-token memory cost + add_opt(common_arg( + {"--kv-cache-tokens"}, "N", + "limit KV cache to N tokens (reduces memory when full context not needed, 0 = use context size)", + [](common_params & params, int value) { + // Use block size of 64 by default for efficient memory alignment + params.paged_attn_block_size = 64; + params.paged_attn_max_blocks = (value + 63) / 64; // Round up to nearest block + } + ).set_env("LLAMA_KV_CACHE_TOKENS")); + add_opt(common_arg( + {"--kv-block-size"}, "N", + string_format("KV cache block size in tokens for block tracking (default: %d, 0 = disabled)", params.paged_attn_block_size), + [](common_params & params, int value) { + params.paged_attn_block_size = value; + } + ).set_env("LLAMA_KV_BLOCK_SIZE")); + add_opt(common_arg( + {"--kv-max-blocks"}, "N", + string_format("max KV cache blocks to allocate (default: %d, 0 = unlimited)", params.paged_attn_max_blocks), + [](common_params & params, int value) { + params.paged_attn_max_blocks = value; + } + ).set_env("LLAMA_KV_MAX_BLOCKS")); + add_opt(common_arg( + {"--kv-cache-demand-paged"}, + "use demand-paged KV cache (Linux: physical memory grows with usage, not allocated upfront)", + [](common_params & params) { + params.kv_cache_demand_paged = true; + } + ).set_env("LLAMA_KV_DEMAND_PAGED")); add_opt(common_arg( {"-p", "--prompt"}, "PROMPT", "prompt to start generation with; for system message, use -sys", diff --git a/common/common.cpp b/common/common.cpp index 744f0b4eeb49..a8c6e1012ac7 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -1425,6 +1426,18 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.type_k = params.cache_type_k; cparams.type_v = params.cache_type_v; + // Set KV cache sizing environment variables if CLI flags are used + // This reduces upfront KV cache allocation (NOT per-token memory cost) + if (params.paged_attn_block_size > 0) { + setenv("LLAMA_PAGED_ATTN", std::to_string(params.paged_attn_block_size).c_str(), 1); + } + if (params.paged_attn_max_blocks > 0) { + setenv("LLAMA_PAGED_ATTN_MAX_BLOCKS", std::to_string(params.paged_attn_max_blocks).c_str(), 1); + } + if (params.kv_cache_demand_paged) { + setenv("LLAMA_KV_DEMAND_PAGED", "1", 1); + } + return cparams; } diff --git a/common/common.h b/common/common.h index 7794c0268bd4..de3189e6af48 100644 --- a/common/common.h +++ b/common/common.h @@ -356,6 +356,11 @@ struct common_params { enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention + // KV cache sizing parameters (reduces upfront allocation, NOT per-token memory cost) + uint32_t paged_attn_block_size = 0; // block size in tokens for tracking (0 = disabled) + uint32_t paged_attn_max_blocks = 0; // max blocks to allocate (0 = use full context) + bool kv_cache_demand_paged = false; // use demand-paged memory (Linux: true on-demand, memory grows with usage) + struct common_params_sampling sampling; struct common_params_speculative speculative; struct common_params_vocoder vocoder; diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index a9d1778641e8..50b6eb3da550 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -368,6 +368,9 @@ extern "C" { GGML_API ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size); GGML_API ggml_backend_buffer_type_t ggml_backend_cpu_buffer_type(void); + // Demand-paged CPU buffer (Linux/macOS: true on-demand via mmap, Windows: fallback to regular) + GGML_API ggml_backend_buffer_type_t ggml_backend_cpu_buffer_mmap_type(void); + #ifdef __cplusplus } #endif diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index b69583dd3fde..21a1e62d51f4 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2342,6 +2342,21 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); + // Paged flash attention with indirect KV access via block table + // block_table: I32 tensor [max_blocks, n_seqs] mapping logical → physical blocks + // op_params[4] encodes block_size + GGML_API struct ggml_tensor * ggml_flash_attn_ext_paged( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + struct ggml_tensor * block_table, // I32 [max_blocks, n_seqs] logical→physical mapping + float scale, + float max_bias, + float logit_softcap, + int32_t block_size); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 1b59924b8cb6..3dd91a53c40e 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -25,6 +25,11 @@ #ifdef __APPLE__ #include #include +#include +#endif + +#if defined(__linux__) +#include #endif @@ -2265,3 +2270,86 @@ ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size) GGML_ASSERT((uintptr_t)ptr % TENSOR_ALIGNMENT == 0 && "buffer pointer must be aligned"); return ggml_backend_buffer_init(ggml_backend_cpu_buffer_from_ptr_type(), ggml_backend_cpu_buffer_from_ptr_i, ptr, size); } + +// Demand-paged CPU buffer using mmap (Linux/macOS) +// Physical memory is allocated by the kernel on first access, not upfront +// This provides TRUE memory-proportional allocation for KV caches +// +// Linux: Uses MAP_NORESERVE for guaranteed lazy allocation +// macOS: Uses mmap without MAP_NORESERVE (macOS does lazy allocation by default) +// Windows: Falls back to regular allocation (would need VirtualAlloc implementation) + +#if defined(__linux__) || defined(__APPLE__) + +static void ggml_backend_cpu_buffer_mmap_free_buffer(ggml_backend_buffer_t buffer) { + if (buffer->context) { + munmap(buffer->context, buffer->size); + } +} + +static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_mmap_i = { + /* .free_buffer = */ ggml_backend_cpu_buffer_mmap_free_buffer, + /* .get_base = */ ggml_backend_cpu_buffer_get_base, + /* .init_tensor = */ NULL, + /* .memset_tensor = */ ggml_backend_cpu_buffer_memset_tensor, + /* .set_tensor = */ ggml_backend_cpu_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_cpu_buffer_get_tensor, + /* .cpy_tensor = */ ggml_backend_cpu_buffer_cpy_tensor, + /* .clear = */ ggml_backend_cpu_buffer_clear, + /* .reset = */ NULL, +}; + +static const char * ggml_backend_cpu_buffer_mmap_type_get_name(ggml_backend_buffer_type_t buft) { + return "CPU_DemandPaged"; + GGML_UNUSED(buft); +} + +static ggml_backend_buffer_t ggml_backend_cpu_buffer_mmap_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + // Round up to page size for mmap alignment + const size_t page_size = 4096; + size_t aligned_size = ((size + page_size - 1) / page_size) * page_size; + + // Use mmap for demand paging - physical pages allocated on first touch + int flags = MAP_PRIVATE | MAP_ANON; +#if defined(__linux__) + // MAP_NORESERVE: don't reserve swap space, allow overcommit (Linux-specific) + flags |= MAP_NORESERVE; +#endif + // macOS does lazy allocation by default with MAP_PRIVATE | MAP_ANON + + void * data = mmap(NULL, aligned_size, PROT_READ | PROT_WRITE, flags, -1, 0); + + if (data == MAP_FAILED) { + GGML_LOG_ERROR("%s: mmap failed for size %zu\n", __func__, aligned_size); + return NULL; + } + + return ggml_backend_buffer_init(buft, ggml_backend_cpu_buffer_mmap_i, data, aligned_size); +} + +static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_mmap_type_instance = { + /* .iface = */ { + /* .get_name = */ ggml_backend_cpu_buffer_mmap_type_get_name, + /* .alloc_buffer = */ ggml_backend_cpu_buffer_mmap_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_buffer_type_get_alignment, + /* .get_max_size = */ NULL, + /* .get_alloc_size = */ NULL, + /* .is_host = */ ggml_backend_cpu_buffer_type_is_host, + }, + /* .device = */ NULL, + /* .context = */ NULL, +}; + +ggml_backend_buffer_type_t ggml_backend_cpu_buffer_mmap_type(void) { + return &ggml_backend_cpu_buffer_mmap_type_instance; +} + +#else + +// Fallback for Windows/other: just use regular CPU buffer +// Windows would need VirtualAlloc with MEM_RESERVE for similar behavior +ggml_backend_buffer_type_t ggml_backend_cpu_buffer_mmap_type(void) { + return ggml_backend_cpu_buffer_type(); +} + +#endif // __linux__ || __APPLE__ diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 3032783971dc..b542cb5884c1 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -8022,11 +8022,21 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( const ggml_compute_params * params, ggml_tensor * dst, int ir0, int ir1) { - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; - const ggml_tensor * mask = dst->src[3]; - const ggml_tensor * sinks = dst->src[4]; + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * src4 = dst->src[4]; // sinks (F32) or block_table (I32) + + // Detect paging mode: src[4] is block_table if I32, sinks if F32 + const bool use_paging = (src4 && src4->type == GGML_TYPE_I32); + const ggml_tensor * sinks = use_paging ? NULL : src4; + const ggml_tensor * block_table = use_paging ? src4 : NULL; + + // Block table setup for paged attention + const int32_t * bt_data = use_paging ? (const int32_t *) block_table->data : NULL; + const int64_t max_blocks = use_paging ? block_table->ne[0] : 0; + const int32_t block_size = use_paging ? ggml_get_op_params_i32(dst, 4) : 0; GGML_TENSOR_LOCALS(int64_t, neq, q, ne) GGML_TENSOR_LOCALS(size_t, nbq, q, nb) @@ -8144,9 +8154,24 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( continue; } + // Compute physical position for paged attention + int64_t kv_pos = ic; // default: identity mapping + if (use_paging) { + const int64_t logical_block = ic / block_size; + const int64_t offset_in_block = ic % block_size; + if (logical_block >= max_blocks) { + continue; // beyond allocated blocks + } + const int32_t physical_block = bt_data[iq3 * max_blocks + logical_block]; + if (physical_block < 0) { + continue; // invalid block + } + kv_pos = physical_block * block_size + offset_in_block; + } + float s; // KQ value - const char * k_data = (const char *) k->data + ( ic*nbk1 + ik2*nbk2 + ik3*nbk3); + const char * k_data = (const char *) k->data + (kv_pos*nbk1 + ik2*nbk2 + ik3*nbk3); kq_vec_dot(DK, &s, 0, k_data, 0, Q_q, 0, 1); s = s*scale; // scale KQ value @@ -8162,7 +8187,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value float vs = 1.0f; // post-softmax KQ value, expf(s - M) - const char * v_data = ((const char *) v->data + (ic*nbv1 + iv2*nbv2 + iv3*nbv3)); + const char * v_data = ((const char *) v->data + (kv_pos*nbv1 + iv2*nbv2 + iv3*nbv3)); if (v->type == GGML_TYPE_F16) { if (s > M) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 09b8eb466d38..99d795c8b923 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5323,6 +5323,60 @@ void ggml_flash_attn_ext_add_sinks( a->src[4] = sinks; } +// ggml_flash_attn_ext_paged + +struct ggml_tensor * ggml_flash_attn_ext_paged( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + struct ggml_tensor * block_table, + float scale, + float max_bias, + float logit_softcap, + int32_t block_size) { + GGML_ASSERT(ggml_can_mul_mat(k, q)); + // TODO: check if vT can be multiplied by (k*qT) + + GGML_ASSERT(q->ne[3] == k->ne[3]); + GGML_ASSERT(q->ne[3] == v->ne[3]); + + if (mask) { + GGML_ASSERT(ggml_is_contiguous(mask)); + + GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); + GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); + } + + if (max_bias > 0.0f) { + GGML_ASSERT(mask); + } + + // block_table must be I32 tensor [max_blocks, n_seqs] + GGML_ASSERT(block_table != NULL); + GGML_ASSERT(block_table->type == GGML_TYPE_I32); + GGML_ASSERT(block_size > 0); + + // permute(0, 2, 1, 3) + int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + float params[] = { scale, max_bias, logit_softcap }; + ggml_set_op_params(result, params, sizeof(params)); + ggml_set_op_params_i32(result, 3, GGML_PREC_DEFAULT); // precision + ggml_set_op_params_i32(result, 4, block_size); // block size for paging + + result->op = GGML_OP_FLASH_ATTN_EXT; // Uses same op; kernel detects paging via src[4] type + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = mask; + result->src[4] = block_table; + + return result; +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 374ff1ebf3a2..a71beb7fee0c 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -391,6 +391,10 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { mctx->set_input_v_idxs(self_v_idxs, ubatch); mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + + if (self_block_table != nullptr) { + mctx->set_input_block_table(self_block_table); + } } bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) { @@ -1126,7 +1130,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * weights = ggml_get_rows(ctx0, probs, selected_experts); // [1, n_expert_used, n_tokens] cb(weights, "ffn_moe_weights", il); - if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT) { weights = ggml_reshape_2d(ctx0, weights, n_expert_used, n_tokens); weights = ggml_soft_max(ctx0, weights); // [n_expert_used, n_tokens] @@ -1248,26 +1251,27 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * cur_experts[LLAMA_MAX_EXPERTS] = { nullptr }; - assert(n_expert_used > 0); + // Use hparams.n_expert_used for aggregation to avoid dynamic allocation issues during warmup + // ref: https://github.com/ggml-org/llama.cpp/pull/14753 + const uint32_t n_expert_agg = hparams.n_expert_used; + + assert(n_expert_agg > 0); // order the views before the adds - for (uint32_t i = 0; i < hparams.n_expert_used; ++i) { + for (uint32_t i = 0; i < n_expert_agg; ++i) { cur_experts[i] = ggml_view_2d(ctx0, experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1]); ggml_build_forward_expand(gf, cur_experts[i]); } // aggregate experts - // note: here we explicitly use hparams.n_expert_used instead of n_expert_used - // to avoid potentially a large number of add nodes during warmup - // ref: https://github.com/ggml-org/llama.cpp/pull/14753 ggml_tensor * moe_out = cur_experts[0]; - for (uint32_t i = 1; i < hparams.n_expert_used; ++i) { + for (uint32_t i = 1; i < n_expert_agg; ++i) { moe_out = ggml_add(ctx0, moe_out, cur_experts[i]); } - if (hparams.n_expert_used == 1) { + if (n_expert_agg == 1) { // avoid returning a non-contiguous tensor moe_out = ggml_cont(ctx0, moe_out); } @@ -1485,7 +1489,9 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, - int il) const { + int il, + ggml_tensor * block_table, + int32_t block_size) const { const bool v_trans = v->nb[1] > v->nb[2]; // split the batch into streams if needed @@ -1515,12 +1521,21 @@ ggml_tensor * llm_graph_context::build_attn_mha( v = ggml_cast(ctx0, v, GGML_TYPE_F16); } - cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, - hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); - cb(cur, LLAMA_TENSOR_NAME_FATTN, il); - - ggml_flash_attn_ext_add_sinks(cur, sinks); - ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); + if (block_table != nullptr && block_size > 0) { + // Use paged attention when block table is provided + cur = ggml_flash_attn_ext_paged(ctx0, q, k, v, kq_mask, block_table, kq_scale, + hparams.f_max_alibi_bias, + hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f, + block_size); + cb(cur, "fattn_paged", il); + // Note: paged attention uses src[4] for block_table, so sinks are not supported + } else { + cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, + hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + cb(cur, LLAMA_TENSOR_NAME_FATTN, il); + ggml_flash_attn_ext_add_sinks(cur, sinks); + } + ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32); if (v_mla) { #if 0 @@ -1706,6 +1721,14 @@ static std::unique_ptr build_attn_inp_kv_impl( ggml_set_input(inp->self_kq_mask); inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask; + + // Create block table tensor for paged attention if enabled + if (mctx_cur->has_block_tracking() && cparams.flash_attn) { + const uint32_t block_size = mctx_cur->get_block_size(); + const uint32_t max_blocks_per_seq = (n_kv + block_size - 1) / block_size; + const uint32_t n_seqs = ubatch.n_seqs_unq; + inp->self_block_table = mctx_cur->build_block_table_tensor(ctx0, n_seqs, max_blocks_per_seq); + } } return inp; @@ -1755,7 +1778,10 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * block_table = inp->get_block_table(); + int32_t block_size = block_table ? mctx_cur->get_block_size() : 0; + + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il, block_table, block_size); cb(cur, "kqv_out", il); if (wo) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 503ffd695aa7..fe6e9b704c82 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -299,12 +299,16 @@ class llm_graph_input_attn_kv : public llm_graph_input_i { ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; } + ggml_tensor * get_block_table() const { return self_block_table; } + ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa] ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream] ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_block_table = nullptr; // I32 [max_blocks, n_seqs] (for paged attention) + // note: these have to be copies because in order to be able to reuse a graph, its inputs // need to carry these parameters with them. otherwise, they can point to freed // llm_graph_params from a previous batch, causing stack-use-after-return @@ -773,7 +777,9 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] float kq_scale, - int il) const; + int il, + ggml_tensor * block_table = nullptr, // [max_blocks, n_seqs] for paged attention + int32_t block_size = 0) const; llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; diff --git a/src/llama-kv-block.h b/src/llama-kv-block.h new file mode 100644 index 000000000000..357f5deb397a --- /dev/null +++ b/src/llama-kv-block.h @@ -0,0 +1,263 @@ +#pragma once + +#include "llama.h" + +#include +#include +#include +#include + +// Block configuration for paged KV cache +struct llama_kv_block_config { + uint32_t tokens_per_block = 64; + uint32_t min_free_blocks = 4; + size_t block_size_bytes = 0; // computed during init +}; + +// Statistics for block pool usage +struct llama_kv_block_stats { + uint32_t n_blocks_total = 0; + uint32_t n_blocks_used = 0; + uint32_t n_tokens_total = 0; + uint32_t n_tokens_used = 0; + uint32_t n_tokens_wasted = 0; + float fragmentation = 0.0f; + float utilization = 0.0f; + size_t memory_allocated = 0; + size_t memory_used = 0; + size_t memory_wasted = 0; +}; + +// Simple block metadata +struct llama_kv_block { + uint32_t n_tokens = 0; // valid tokens in block + bool allocated = false; + + void reset() { + n_tokens = 0; + allocated = false; + } +}; + +// Stack-based block pool with O(1) allocation/deallocation. +// NOT thread-safe: designed for single-threaded llama_context access. +class llama_kv_block_pool { +public: + llama_kv_block_pool() = default; + + void init(uint32_t n_blocks) { + blocks.resize(n_blocks); + free_stack.reserve(n_blocks); + for (uint32_t i = n_blocks; i > 0; --i) { + blocks[i - 1].reset(); + free_stack.push_back(i - 1); + } + } + + // Returns block index, or -1 if unavailable + int32_t allocate() { + if (free_stack.empty()) { + return -1; + } + const uint32_t idx = free_stack.back(); + free_stack.pop_back(); + assert(!blocks[idx].allocated); + blocks[idx].allocated = true; + return static_cast(idx); + } + + std::vector allocate_batch(uint32_t n) { + if (n > free_stack.size()) { + return {}; + } + std::vector result; + result.reserve(n); + for (uint32_t i = 0; i < n; ++i) { + int32_t idx = allocate(); + if (idx < 0) { + for (int32_t allocated_idx : result) { + deallocate(allocated_idx); + } + return {}; + } + result.push_back(idx); + } + return result; + } + + void deallocate(int32_t idx) { + assert(idx >= 0 && static_cast(idx) < blocks.size()); + assert(blocks[idx].allocated); + blocks[idx].reset(); + free_stack.push_back(static_cast(idx)); + } + + llama_kv_block & get(int32_t idx) { + assert(idx >= 0 && static_cast(idx) < blocks.size()); + return blocks[idx]; + } + + const llama_kv_block & get(int32_t idx) const { + assert(idx >= 0 && static_cast(idx) < blocks.size()); + return blocks[idx]; + } + + uint32_t n_free() const { return static_cast(free_stack.size()); } + uint32_t n_total() const { return static_cast(blocks.size()); } + uint32_t n_used() const { return n_total() - n_free(); } + + llama_kv_block_stats compute_stats(uint32_t tokens_per_block, size_t bytes_per_token = 0) const { + llama_kv_block_stats stats; + stats.n_blocks_total = n_total(); + stats.n_blocks_used = n_used(); + + uint32_t total_tokens_in_blocks = 0; + uint32_t actual_tokens_stored = 0; + + for (const auto & block : blocks) { + if (block.allocated) { + total_tokens_in_blocks += tokens_per_block; + actual_tokens_stored += block.n_tokens; + } + } + + stats.n_tokens_total = stats.n_blocks_total * tokens_per_block; + stats.n_tokens_used = actual_tokens_stored; + stats.n_tokens_wasted = total_tokens_in_blocks - actual_tokens_stored; + + if (total_tokens_in_blocks > 0) { + stats.fragmentation = static_cast(stats.n_tokens_wasted) / + static_cast(total_tokens_in_blocks); + } + if (stats.n_tokens_total > 0) { + stats.utilization = static_cast(stats.n_tokens_used) / + static_cast(stats.n_tokens_total); + } + if (bytes_per_token > 0) { + stats.memory_allocated = stats.n_blocks_used * tokens_per_block * bytes_per_token; + stats.memory_used = stats.n_tokens_used * bytes_per_token; + stats.memory_wasted = stats.n_tokens_wasted * bytes_per_token; + } + return stats; + } + + void clear() { + for (auto & block : blocks) { + block.reset(); + } + free_stack.clear(); + free_stack.reserve(blocks.size()); + for (uint32_t i = static_cast(blocks.size()); i > 0; --i) { + free_stack.push_back(i - 1); + } + } + +private: + std::vector blocks; + std::vector free_stack; +}; + +// Maps logical blocks to physical blocks per sequence. +// NOT thread-safe: same design as llama_kv_block_pool. +class llama_kv_block_table { +public: + llama_kv_block_table() = default; + + int32_t get_physical(llama_seq_id seq_id, uint32_t logical_idx) const { + auto it = table.find(seq_id); + if (it == table.end() || logical_idx >= it->second.size()) { + return -1; + } + return it->second[logical_idx]; + } + + void set_mapping(llama_seq_id seq_id, uint32_t logical_idx, int32_t physical_idx) { + auto & seq_blocks = table[seq_id]; + if (logical_idx >= seq_blocks.size()) { + seq_blocks.resize(logical_idx + 1, -1); + } + seq_blocks[logical_idx] = physical_idx; + } + + uint32_t append_block(llama_seq_id seq_id, int32_t physical_idx) { + auto & seq_blocks = table[seq_id]; + uint32_t logical_idx = static_cast(seq_blocks.size()); + seq_blocks.push_back(physical_idx); + return logical_idx; + } + + const std::vector * get_sequence_blocks(llama_seq_id seq_id) const { + auto it = table.find(seq_id); + return it != table.end() ? &it->second : nullptr; + } + + uint32_t get_sequence_n_blocks(llama_seq_id seq_id) const { + auto it = table.find(seq_id); + return it != table.end() ? static_cast(it->second.size()) : 0; + } + + void clear_sequence(llama_seq_id seq_id) { + table.erase(seq_id); + } + + std::vector truncate_sequence(llama_seq_id seq_id, uint32_t n_blocks_to_remove) { + std::vector removed; + auto it = table.find(seq_id); + if (it == table.end()) { + return removed; + } + auto & seq_blocks = it->second; + uint32_t n_to_remove = std::min(n_blocks_to_remove, static_cast(seq_blocks.size())); + removed.reserve(n_to_remove); + for (uint32_t i = 0; i < n_to_remove; ++i) { + removed.push_back(seq_blocks.back()); + seq_blocks.pop_back(); + } + if (seq_blocks.empty()) { + table.erase(it); + } + return removed; + } + + void copy_sequence(llama_seq_id seq_src, llama_seq_id seq_dst) { + auto it = table.find(seq_src); + if (it != table.end()) { + table[seq_dst] = it->second; + } + } + + bool has_sequence(llama_seq_id seq_id) const { + return table.find(seq_id) != table.end(); + } + + std::vector get_sequences() const { + std::vector seqs; + seqs.reserve(table.size()); + for (const auto & [seq_id, _] : table) { + seqs.push_back(seq_id); + } + return seqs; + } + + void clear() { table.clear(); } + + size_t total_mappings() const { + size_t total = 0; + for (const auto & [_, blocks] : table) { + total += blocks.size(); + } + return total; + } + + uint32_t n_sequences() const { + return static_cast(table.size()); + } + +private: + std::unordered_map> table; +}; + +// Utility functions +inline uint32_t tokens_to_blocks(uint32_t n_tokens, uint32_t tokens_per_block) { + return (n_tokens + tokens_per_block - 1) / tokens_per_block; +} diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3186242d60fd..a53ec134d24d 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include // @@ -34,6 +35,45 @@ llama_kv_cache::llama_kv_cache( model(model), hparams(model.hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type) { + // KV cache sizing: allocate for fewer tokens when full context isn't needed + // Note: This reduces upfront allocation, NOT per-token memory cost + // Use --kv-cache-tokens N or environment variables to configure + const char * env_block_size = getenv("LLAMA_PAGED_ATTN"); + const char * env_max_blocks = getenv("LLAMA_PAGED_ATTN_MAX_BLOCKS"); + const char * env_demand_paged = getenv("LLAMA_KV_DEMAND_PAGED"); + + uint32_t kv_block_size = 0; + uint32_t kv_max_blocks = 0; + bool use_demand_paged = (env_demand_paged && atoi(env_demand_paged) != 0); + + if (use_demand_paged) { + LLAMA_LOG_INFO("%s: demand-paged KV cache enabled (physical memory grows with usage)\n", __func__); + } + + if (env_block_size) { + kv_block_size = (uint32_t) atoi(env_block_size); + } + + if (env_max_blocks) { + kv_max_blocks = (uint32_t) atoi(env_max_blocks); + } + + // If configured, limit KV cache to specified number of tokens + if (kv_block_size > 0 && kv_max_blocks > 0) { + uint32_t kv_size_limited = kv_max_blocks * kv_block_size; + + // Ensure alignment with n_pad + if (n_pad > 1 && kv_size_limited % n_pad != 0) { + kv_size_limited = ((kv_size_limited + n_pad - 1) / n_pad) * n_pad; + } + + if (kv_size_limited < kv_size) { + LLAMA_LOG_INFO("%s: KV cache limited to %u tokens (model supports %u)\n", + __func__, kv_size_limited, kv_size); + kv_size = kv_size_limited; + } + } + GGML_ASSERT(kv_size % n_pad == 0); const uint32_t n_layer_kv = hparams.n_layer_kv(); @@ -114,13 +154,18 @@ llama_kv_cache::llama_kv_cache( const char * dev_name = "CPU"; - ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_buffer_type_t buft; if (offload) { auto * dev = model.dev_layer(il); buft = ggml_backend_dev_buffer_type(dev); - dev_name = ggml_backend_dev_name(dev); + } else if (use_demand_paged) { + // Use demand-paged buffer for CPU: physical memory allocated on first access + buft = ggml_backend_cpu_buffer_mmap_type(); + dev_name = "CPU_DemandPaged"; + } else { + buft = ggml_backend_cpu_buffer_type(); } LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); @@ -206,6 +251,17 @@ llama_kv_cache::llama_kv_cache( const char * LLAMA_KV_CACHE_DEBUG = getenv("LLAMA_KV_CACHE_DEBUG"); debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0; + + // Enable block tracking for KV cache management + // Usage: --kv-block-size N or LLAMA_PAGED_ATTN=N (block size in tokens) + const char * env_kv_block_size = getenv("LLAMA_PAGED_ATTN"); + if (env_kv_block_size) { + const uint32_t block_size_env = (uint32_t) atoi(env_kv_block_size); + if (block_size_env > 0) { + enable_blocks(block_size_env); + LLAMA_LOG_INFO("%s: KV cache block tracking enabled (block_size = %u tokens)\n", __func__, block_size_env); + } + } } void llama_kv_cache::clear(bool data) { @@ -281,6 +337,58 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { } } + // Update block tracking after removal + if (enable_block_tracking && block_size > 0) { + if (seq_id >= 0) { + // Check if we removed the entire sequence (all positions) + const uint32_t stream_idx = seq_to_stream[seq_id]; + if (stream_idx < v_block_tables.size() && stream_idx < v_block_pools.size()) { + auto & table = v_block_tables[stream_idx]; + auto & pool = v_block_pools[stream_idx]; + auto & cells = v_cells[stream_idx]; + + // Check if this sequence still has any cells + bool seq_has_cells = false; + for (uint32_t i = 0; i < cells.size() && !seq_has_cells; ++i) { + if (cells.seq_has(i, seq_id)) { + seq_has_cells = true; + } + } + + if (!seq_has_cells) { + // Sequence is completely removed, deallocate all its blocks + const auto * blocks = table.get_sequence_blocks(seq_id); + if (blocks) { + for (int32_t physical_block : *blocks) { + if (physical_block >= 0) { + pool.deallocate(physical_block); + LLAMA_LOG_DEBUG("%s: deallocated block %d for seq %d\n", + __func__, physical_block, seq_id); + } + } + } + table.clear_sequence(seq_id); + } + } + } else { + // Removed all sequences - clear all block tracking + for (uint32_t s = 0; s < n_stream; ++s) { + if (s < v_block_pools.size()) { + v_block_pools[s].clear(); + } + if (s < v_block_tables.size()) { + v_block_tables[s].clear(); + } + } + LLAMA_LOG_DEBUG("%s: cleared all block tracking\n", __func__); + } + + // Print stats in debug mode + if (debug > 0) { + print_block_stats(); + } + } + return true; } @@ -815,6 +923,14 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, // for non-continuous slots, we test the tokens one by one const uint32_t n_test = cont ? n_tokens : 1; + // For SWA caches: compute min position in this batch for sequence s + // This ensures all tokens in the batch have their full attention window + // (the token at min position has the most demanding context requirement) + llama_pos pos_batch_min = 0; + if (n_swa > 0) { + pos_batch_min = ubatch.pos[s * n_tokens]; + } + while (true) { if (head_cur + n_test > cells.size()) { n_tested += cells.size() - head_cur; @@ -835,7 +951,7 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, // - the cell is empty // - the cell is occupied only by one sequence: // - (disabled) mask causally, if the sequence is the same as the one we are inserting - // - mask SWA, using current max pos for that sequence in the cache + // - mask SWA, using the batch's min pos (ensures all tokens have context) // always insert in the cell with minimum pos bool can_use = cells.is_empty(idx); @@ -848,11 +964,11 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, // can_use = pos_cell >= pos; //} - if (!can_use) { - const llama_seq_id seq_id_cell = cells.seq_get(idx); - - // SWA mask - if (is_masked_swa(pos_cell, cells.seq_pos_max(seq_id_cell) + 1)) { + // SWA mask - check if cell position is outside attention window + // Use batch's min position to ensure all tokens have their full context + // (min position token has the most demanding attention window) + if (!can_use && n_swa > 0) { + if (is_masked_swa(pos_cell, pos_batch_min + 1)) { can_use = true; } } @@ -963,6 +1079,9 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & head = sinfo.idxs[s].back() + 1; } + + // Update block tracking metadata + update_block_tokens(sinfo); } bool llama_kv_cache::get_can_shift() const { @@ -1974,6 +2093,206 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return true; } +// +// llama_kv_cache block tracking for KV cache management +// + +void llama_kv_cache::enable_blocks(uint32_t tokens_per_block) { + if (tokens_per_block == 0) { + return; + } + + enable_block_tracking = true; + block_size = tokens_per_block; + + // Initialize block pools and tables for each stream + const uint32_t kv_size = get_size(); + const uint32_t n_blocks = (kv_size + block_size - 1) / block_size; + + v_block_pools.resize(n_stream); + v_block_tables.resize(n_stream); + + for (uint32_t s = 0; s < n_stream; ++s) { + v_block_pools[s].init(n_blocks); + } + + LLAMA_LOG_INFO("%s: enabled block tracking with block_size=%u, n_blocks=%u\n", + __func__, block_size, n_blocks); +} + +bool llama_kv_cache::has_block_tracking() const { + return enable_block_tracking; +} + +uint32_t llama_kv_cache::get_block_size() const { + return block_size; +} + +ggml_tensor * llama_kv_cache::build_block_table_tensor( + ggml_context * ctx, + uint32_t n_seqs, + uint32_t max_blocks_per_seq) const { + // Create I32 tensor [max_blocks_per_seq, n_seqs] + ggml_tensor * block_table = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, max_blocks_per_seq, n_seqs); + ggml_set_input(block_table); + ggml_set_name(block_table, "block_table"); + return block_table; +} + +void llama_kv_cache::set_input_block_table(ggml_tensor * dst, const slot_info & sinfo) const { + int32_t * data = (int32_t *) dst->data; + const int64_t max_blocks = dst->ne[0]; + const int64_t n_seqs = dst->ne[1]; + + // Initialize all entries to identity mapping (physical = logical) + // This ensures correct behavior even when actual block mappings aren't available + for (int64_t s = 0; s < n_seqs; ++s) { + for (int64_t b = 0; b < max_blocks; ++b) { + data[s * max_blocks + b] = (int32_t) b; + } + } + + // If block tracking is enabled and we have block tables, use actual mappings + if (enable_block_tracking && !v_block_tables.empty()) { + for (uint32_t s = 0; s < sinfo.n_stream() && s < (uint32_t) n_seqs; ++s) { + const uint32_t stream_idx = sinfo.strm[s]; + if (stream_idx >= v_block_tables.size()) { + continue; + } + + const auto & block_table = v_block_tables[stream_idx]; + const auto sequences = block_table.get_sequences(); + + for (const auto & seq_id : sequences) { + const auto * blocks = block_table.get_sequence_blocks(seq_id); + if (!blocks || blocks->empty()) { + continue; + } + + // Map this sequence's blocks to the table row + const int64_t table_row = s; + if (table_row >= n_seqs) { + continue; + } + + for (size_t b = 0; b < blocks->size() && (int64_t) b < max_blocks; ++b) { + data[table_row * max_blocks + b] = (*blocks)[b]; + } + break; // Only process first sequence per stream for now + } + } + } +} + +void llama_kv_cache::print_block_stats() const { + if (!enable_block_tracking || block_size == 0) { + LLAMA_LOG_INFO("%s: block tracking not enabled\n", __func__); + return; + } + + // Compute bytes per token (K + V per layer) + size_t bytes_per_token = 0; + for (const auto & layer : layers) { + const uint32_t il = layer.il; + const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il); + const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); + bytes_per_token += n_embd_k_gqa * ggml_type_size(layer.k->type); + bytes_per_token += n_embd_v_gqa * ggml_type_size(layer.v->type); + } + + LLAMA_LOG_INFO("%s: === Block Pool Statistics ===\n", __func__); + LLAMA_LOG_INFO("%s: block_size = %u tokens, bytes_per_token = %zu\n", + __func__, block_size, bytes_per_token); + + for (uint32_t s = 0; s < n_stream && s < v_block_pools.size(); ++s) { + const auto & pool = v_block_pools[s]; + const auto stats = pool.compute_stats(block_size, bytes_per_token); + + LLAMA_LOG_INFO("%s: stream %u: blocks %u/%u (%.1f%% used), tokens %u/%u (%.1f%% util)\n", + __func__, s, + stats.n_blocks_used, stats.n_blocks_total, + stats.n_blocks_total > 0 ? 100.0f * stats.n_blocks_used / stats.n_blocks_total : 0.0f, + stats.n_tokens_used, stats.n_tokens_total, + 100.0f * stats.utilization); + + if (stats.n_tokens_wasted > 0) { + LLAMA_LOG_INFO("%s: wasted: %u tokens (%.1f%% fragmentation)\n", + __func__, stats.n_tokens_wasted, 100.0f * stats.fragmentation); + } + + if (bytes_per_token > 0 && stats.memory_allocated > 0) { + LLAMA_LOG_INFO("%s: memory: %.2f MiB allocated, %.2f MiB used, %.2f MiB wasted\n", + __func__, + stats.memory_allocated / (1024.0f * 1024.0f), + stats.memory_used / (1024.0f * 1024.0f), + stats.memory_wasted / (1024.0f * 1024.0f)); + } + } + LLAMA_LOG_INFO("%s: =============================\n", __func__); +} + +void llama_kv_cache::update_block_tokens(const slot_info & sinfo) { + if (!enable_block_tracking || block_size == 0) { + return; + } + + // For each stream in the slot info, update block allocations + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { + const uint32_t stream_idx = sinfo.strm[s]; + if (stream_idx >= v_block_pools.size() || stream_idx >= v_block_tables.size()) { + continue; + } + + auto & pool = v_block_pools[stream_idx]; + auto & table = v_block_tables[stream_idx]; + const auto & cells = v_cells[stream_idx]; + + // Track which logical blocks we've seen for each sequence + std::unordered_map> seq_logical_blocks; + + // Process each cell index in this stream + for (uint32_t ii = 0; ii < sinfo.idxs[s].size(); ++ii) { + const uint32_t cell_idx = sinfo.idxs[s][ii]; + const uint32_t logical_block = cell_idx / block_size; + + // Get sequence IDs for this cell + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq_max; ++seq_id) { + if (!cells.seq_has(cell_idx, seq_id)) { + continue; + } + + // Check if this logical block already has a physical block allocated + int32_t physical_block = table.get_physical(seq_id, logical_block); + + if (physical_block < 0) { + // Need to allocate a new physical block + physical_block = pool.allocate(); + + if (physical_block < 0) { + // Pool exhausted - this shouldn't happen with proper sizing + LLAMA_LOG_WARN("%s: block pool exhausted for stream %u\n", __func__, stream_idx); + continue; + } + + // Initialize block and update block table mapping + pool.get(physical_block).n_tokens = 0; + table.set_mapping(seq_id, logical_block, physical_block); + + LLAMA_LOG_DEBUG("%s: allocated block %d for seq %d, logical %u\n", + __func__, physical_block, seq_id, logical_block); + } + + // Increment token count for this block + auto & block = pool.get(physical_block); + if (seq_logical_blocks[seq_id].insert(logical_block).second) { + // First time seeing this logical block for this sequence in this batch + block.n_tokens++; + } + } + } + } +} + // // llama_kv_cache_context // @@ -2098,3 +2417,27 @@ void llama_kv_cache_context::set_input_kq_mask(ggml_tensor * dst, const llama_ub void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const { kv->set_input_pos_bucket(dst, ubatch); } + +// +// llama_kv_cache_context block tracking forwarding methods +// + +bool llama_kv_cache_context::has_block_tracking() const { + return kv->has_block_tracking(); +} + +uint32_t llama_kv_cache_context::get_block_size() const { + return kv->get_block_size(); +} + +ggml_tensor * llama_kv_cache_context::build_block_table_tensor( + ggml_context * ctx, + uint32_t n_seqs, + uint32_t max_blocks_per_seq) const { + return kv->build_block_table_tensor(ctx, n_seqs, max_blocks_per_seq); +} + +void llama_kv_cache_context::set_input_block_table(ggml_tensor * dst) const { + GGML_ASSERT(!sinfos.empty() && "set_input_block_table called without slot infos"); + kv->set_input_block_table(dst, sinfos[i_cur]); +} diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 0c4ed6484564..9f05b7696024 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -3,6 +3,7 @@ #include "llama-batch.h" #include "llama-graph.h" #include "llama-kv-cells.h" +#include "llama-kv-block.h" #include "llama-memory.h" #include @@ -199,6 +200,29 @@ class llama_kv_cache : public llama_memory_i { void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + // + // block tracking API (for paged attention) + // + + // Enable block-based tracking with the given block size (tokens per block) + void enable_blocks(uint32_t tokens_per_block); + + // Check if block tracking is enabled + bool has_block_tracking() const; + + // Get tokens per block (0 if not enabled) + uint32_t get_block_size() const; + + // Build a block table tensor for graph construction + // Returns I32 tensor [max_blocks_per_seq, n_seqs] + ggml_tensor * build_block_table_tensor(ggml_context * ctx, uint32_t n_seqs, uint32_t max_blocks_per_seq) const; + + // Populate the block table tensor with current block mappings + void set_input_block_table(ggml_tensor * dst, const slot_info & sinfo) const; + + // Print block pool statistics (for debugging) + void print_block_stats() const; + private: const llama_model & model; const llama_hparams & hparams; @@ -252,6 +276,20 @@ class llama_kv_cache : public llama_memory_i { // model layer id -> KV cache layer id std::unordered_map map_layer_ids; + // + // block tracking for paged attention + // + + bool enable_block_tracking = false; + uint32_t block_size = 0; // tokens per block + + // per-stream block pool and tables + std::vector v_block_pools; // [n_stream] + std::vector v_block_tables; // [n_stream] + + // Update block token counts after apply_ubatch + void update_block_tokens(const slot_info & sinfo); + size_t total_size() const; size_t size_k_bytes() const; @@ -355,6 +393,15 @@ class llama_kv_cache_context : public llama_memory_context_i { void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + // + // block tracking API (for paged attention) - forwarding methods + // + + bool has_block_tracking() const; + uint32_t get_block_size() const; + ggml_tensor * build_block_table_tensor(ggml_context * ctx, uint32_t n_seqs, uint32_t max_blocks_per_seq) const; + void set_input_block_table(ggml_tensor * dst) const; + private: llama_memory_status status; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6245cd967afe..3202180cc590 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -188,6 +188,7 @@ llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) llama_build_and_test(test-chat-template.cpp) llama_build_and_test(test-json-partial.cpp) llama_build_and_test(test-log.cpp) +llama_build_and_test(test-kv-block.cpp) llama_build_and_test( test-peg-parser.cpp peg-parser/simple-tokenize.cpp diff --git a/tests/test-kv-block.cpp b/tests/test-kv-block.cpp new file mode 100644 index 000000000000..bf50e1556c9f --- /dev/null +++ b/tests/test-kv-block.cpp @@ -0,0 +1,444 @@ +// Unit tests for llama_kv_block_pool and llama_kv_block_table +// Tests: allocation, deallocation, sequence management + +#include "../src/llama-kv-block.h" + +#include +#include +#include +#include + +#define TEST_ASSERT(cond) do { \ + if (!(cond)) { \ + fprintf(stderr, "FAIL: %s at %s:%d\n", #cond, __FILE__, __LINE__); \ + return false; \ + } \ +} while(0) + +// +// Block Pool Tests +// + +static bool test_pool_init() { + printf(" test_pool_init... "); + + llama_kv_block_pool pool; + pool.init(100); + + TEST_ASSERT(pool.n_total() == 100); + TEST_ASSERT(pool.n_free() == 100); + TEST_ASSERT(pool.n_used() == 0); + + printf("OK\n"); + return true; +} + +static bool test_pool_allocate_single() { + printf(" test_pool_allocate_single... "); + + llama_kv_block_pool pool; + pool.init(10); + + int32_t idx = pool.allocate(); + TEST_ASSERT(idx >= 0 && idx < 10); + TEST_ASSERT(pool.n_free() == 9); + TEST_ASSERT(pool.n_used() == 1); + TEST_ASSERT(pool.get(idx).allocated); + + printf("OK\n"); + return true; +} + +static bool test_pool_allocate_all() { + printf(" test_pool_allocate_all... "); + + llama_kv_block_pool pool; + pool.init(5); + + std::vector allocated; + for (int i = 0; i < 5; i++) { + int32_t idx = pool.allocate(); + TEST_ASSERT(idx >= 0); + allocated.push_back(idx); + } + + TEST_ASSERT(pool.n_free() == 0); + TEST_ASSERT(pool.n_used() == 5); + + // Next allocation should fail + int32_t idx = pool.allocate(); + TEST_ASSERT(idx == -1); + + // Verify all indices are unique + std::set unique(allocated.begin(), allocated.end()); + TEST_ASSERT(unique.size() == 5); + + printf("OK\n"); + return true; +} + +static bool test_pool_deallocate() { + printf(" test_pool_deallocate... "); + + llama_kv_block_pool pool; + pool.init(10); + + int32_t idx1 = pool.allocate(); + int32_t idx2 = pool.allocate(); + TEST_ASSERT(pool.n_used() == 2); + + pool.deallocate(idx1); + TEST_ASSERT(pool.n_used() == 1); + TEST_ASSERT(!pool.get(idx1).allocated); + TEST_ASSERT(pool.get(idx2).allocated); + + // Reallocate - should get same block back (LIFO) + int32_t idx3 = pool.allocate(); + TEST_ASSERT(idx3 == idx1); + + printf("OK\n"); + return true; +} + +static bool test_pool_batch_allocate() { + printf(" test_pool_batch_allocate... "); + + llama_kv_block_pool pool; + pool.init(10); + + auto blocks = pool.allocate_batch(5); + TEST_ASSERT(blocks.size() == 5); + TEST_ASSERT(pool.n_used() == 5); + + // Allocate more than available should fail + auto blocks2 = pool.allocate_batch(6); + TEST_ASSERT(blocks2.empty()); + TEST_ASSERT(pool.n_used() == 5); // No change + + printf("OK\n"); + return true; +} + +static bool test_pool_stats() { + printf(" test_pool_stats... "); + + llama_kv_block_pool pool; + pool.init(100); + + const uint32_t tokens_per_block = 64; + + // Allocate some blocks and set token counts + for (int i = 0; i < 10; i++) { + int32_t idx = pool.allocate(); + pool.get(idx).n_tokens = (i == 9) ? 32 : 64; // Last block half full + } + + auto stats = pool.compute_stats(tokens_per_block); + + TEST_ASSERT(stats.n_blocks_total == 100); + TEST_ASSERT(stats.n_blocks_used == 10); + TEST_ASSERT(stats.n_tokens_total == 6400); // 100 * 64 + TEST_ASSERT(stats.n_tokens_used == 9 * 64 + 32); // 608 + TEST_ASSERT(stats.n_tokens_wasted == 32); // 10 * 64 - 608 + + printf("OK\n"); + return true; +} + +static bool test_pool_clear() { + printf(" test_pool_clear... "); + + llama_kv_block_pool pool; + pool.init(10); + + pool.allocate(); + pool.allocate(); + pool.allocate(); + TEST_ASSERT(pool.n_used() == 3); + + pool.clear(); + TEST_ASSERT(pool.n_used() == 0); + TEST_ASSERT(pool.n_free() == 10); + + printf("OK\n"); + return true; +} + +// +// Block Table Tests +// + +static bool test_table_mapping() { + printf(" test_table_mapping... "); + + llama_kv_block_table table; + + // Set mapping for seq 0 + table.set_mapping(0, 0, 5); + table.set_mapping(0, 1, 10); + table.set_mapping(0, 2, 15); + + TEST_ASSERT(table.get_physical(0, 0) == 5); + TEST_ASSERT(table.get_physical(0, 1) == 10); + TEST_ASSERT(table.get_physical(0, 2) == 15); + TEST_ASSERT(table.get_physical(0, 3) == -1); // Not set + + // Non-existent sequence + TEST_ASSERT(table.get_physical(1, 0) == -1); + + printf("OK\n"); + return true; +} + +static bool test_table_append() { + printf(" test_table_append... "); + + llama_kv_block_table table; + + uint32_t idx0 = table.append_block(0, 100); + uint32_t idx1 = table.append_block(0, 200); + uint32_t idx2 = table.append_block(0, 300); + + TEST_ASSERT(idx0 == 0); + TEST_ASSERT(idx1 == 1); + TEST_ASSERT(idx2 == 2); + + TEST_ASSERT(table.get_physical(0, 0) == 100); + TEST_ASSERT(table.get_physical(0, 1) == 200); + TEST_ASSERT(table.get_physical(0, 2) == 300); + + printf("OK\n"); + return true; +} + +static bool test_table_sequence_blocks() { + printf(" test_table_sequence_blocks... "); + + llama_kv_block_table table; + + table.append_block(0, 10); + table.append_block(0, 20); + table.append_block(1, 30); + + TEST_ASSERT(table.get_sequence_n_blocks(0) == 2); + TEST_ASSERT(table.get_sequence_n_blocks(1) == 1); + TEST_ASSERT(table.get_sequence_n_blocks(2) == 0); + + const auto * blocks = table.get_sequence_blocks(0); + TEST_ASSERT(blocks != nullptr); + TEST_ASSERT(blocks->size() == 2); + TEST_ASSERT((*blocks)[0] == 10); + TEST_ASSERT((*blocks)[1] == 20); + + printf("OK\n"); + return true; +} + +static bool test_table_clear_sequence() { + printf(" test_table_clear_sequence... "); + + llama_kv_block_table table; + + table.append_block(0, 10); + table.append_block(0, 20); + table.append_block(1, 30); + + TEST_ASSERT(table.has_sequence(0)); + TEST_ASSERT(table.has_sequence(1)); + + table.clear_sequence(0); + + TEST_ASSERT(!table.has_sequence(0)); + TEST_ASSERT(table.has_sequence(1)); + TEST_ASSERT(table.get_physical(0, 0) == -1); + + printf("OK\n"); + return true; +} + +static bool test_table_truncate() { + printf(" test_table_truncate... "); + + llama_kv_block_table table; + + table.append_block(0, 10); + table.append_block(0, 20); + table.append_block(0, 30); + table.append_block(0, 40); + + auto removed = table.truncate_sequence(0, 2); + + TEST_ASSERT(removed.size() == 2); + TEST_ASSERT(removed[0] == 40); // LIFO order + TEST_ASSERT(removed[1] == 30); + + TEST_ASSERT(table.get_sequence_n_blocks(0) == 2); + TEST_ASSERT(table.get_physical(0, 0) == 10); + TEST_ASSERT(table.get_physical(0, 1) == 20); + + printf("OK\n"); + return true; +} + +static bool test_table_copy_sequence() { + printf(" test_table_copy_sequence... "); + + llama_kv_block_table table; + + table.append_block(0, 10); + table.append_block(0, 20); + table.append_block(0, 30); + + table.copy_sequence(0, 1); + + TEST_ASSERT(table.has_sequence(1)); + TEST_ASSERT(table.get_sequence_n_blocks(1) == 3); + TEST_ASSERT(table.get_physical(1, 0) == 10); + TEST_ASSERT(table.get_physical(1, 1) == 20); + TEST_ASSERT(table.get_physical(1, 2) == 30); + + printf("OK\n"); + return true; +} + +static bool test_table_get_sequences() { + printf(" test_table_get_sequences... "); + + llama_kv_block_table table; + + table.append_block(0, 10); + table.append_block(5, 20); + table.append_block(10, 30); + + auto seqs = table.get_sequences(); + TEST_ASSERT(seqs.size() == 3); + + std::set seq_set(seqs.begin(), seqs.end()); + TEST_ASSERT(seq_set.count(0) == 1); + TEST_ASSERT(seq_set.count(5) == 1); + TEST_ASSERT(seq_set.count(10) == 1); + + printf("OK\n"); + return true; +} + +static bool test_table_total_mappings() { + printf(" test_table_total_mappings... "); + + llama_kv_block_table table; + + table.append_block(0, 10); + table.append_block(0, 20); + table.append_block(1, 30); + table.append_block(1, 40); + table.append_block(1, 50); + + TEST_ASSERT(table.total_mappings() == 5); + TEST_ASSERT(table.n_sequences() == 2); + + printf("OK\n"); + return true; +} + +// +// Helper Function Tests +// + +static bool test_tokens_to_blocks() { + printf(" test_tokens_to_blocks... "); + + const uint32_t tpb = 64; // tokens per block + + TEST_ASSERT(tokens_to_blocks(0, tpb) == 0); + TEST_ASSERT(tokens_to_blocks(1, tpb) == 1); + TEST_ASSERT(tokens_to_blocks(64, tpb) == 1); + TEST_ASSERT(tokens_to_blocks(65, tpb) == 2); + TEST_ASSERT(tokens_to_blocks(128, tpb) == 2); + TEST_ASSERT(tokens_to_blocks(129, tpb) == 3); + + printf("OK\n"); + return true; +} + +// +// Integration Tests +// + +static bool test_pool_table_integration() { + printf(" test_pool_table_integration... "); + + llama_kv_block_pool pool; + llama_kv_block_table table; + + pool.init(100); + + // Simulate allocating blocks for a sequence + const llama_seq_id seq_id = 0; + const uint32_t tokens_per_block = 64; + + // Allocate 3 blocks for sequence 0 + for (int i = 0; i < 3; i++) { + int32_t phys = pool.allocate(); + TEST_ASSERT(phys >= 0); + + table.append_block(seq_id, phys); + pool.get(phys).n_tokens = tokens_per_block; + } + + TEST_ASSERT(pool.n_used() == 3); + TEST_ASSERT(table.get_sequence_n_blocks(seq_id) == 3); + + // Deallocate sequence + const auto * blocks = table.get_sequence_blocks(seq_id); + for (int32_t phys : *blocks) { + pool.deallocate(phys); + } + table.clear_sequence(seq_id); + + TEST_ASSERT(pool.n_used() == 0); + TEST_ASSERT(!table.has_sequence(seq_id)); + + printf("OK\n"); + return true; +} + +// +// Main +// + +int main() { + printf("Running KV block tests...\n\n"); + + int n_pass = 0; + int n_fail = 0; + + printf("Block Pool Tests:\n"); + if (test_pool_init()) n_pass++; else n_fail++; + if (test_pool_allocate_single()) n_pass++; else n_fail++; + if (test_pool_allocate_all()) n_pass++; else n_fail++; + if (test_pool_deallocate()) n_pass++; else n_fail++; + if (test_pool_batch_allocate()) n_pass++; else n_fail++; + if (test_pool_stats()) n_pass++; else n_fail++; + if (test_pool_clear()) n_pass++; else n_fail++; + + printf("\nBlock Table Tests:\n"); + if (test_table_mapping()) n_pass++; else n_fail++; + if (test_table_append()) n_pass++; else n_fail++; + if (test_table_sequence_blocks()) n_pass++; else n_fail++; + if (test_table_clear_sequence()) n_pass++; else n_fail++; + if (test_table_truncate()) n_pass++; else n_fail++; + if (test_table_copy_sequence()) n_pass++; else n_fail++; + if (test_table_get_sequences()) n_pass++; else n_fail++; + if (test_table_total_mappings()) n_pass++; else n_fail++; + + printf("\nHelper Function Tests:\n"); + if (test_tokens_to_blocks()) n_pass++; else n_fail++; + + printf("\nIntegration Tests:\n"); + if (test_pool_table_integration()) n_pass++; else n_fail++; + + printf("\n========================================\n"); + printf("Results: %d passed, %d failed\n", n_pass, n_fail); + + return n_fail > 0 ? 1 : 0; +}