ggml, llama : add KV cache size limiting and block tracking infrastructure

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",

common/common.cpp

@@ -17,6 +17,7 @@
 #include <codecvt>
 #include <chrono>
 #include <cstdarg>
+#include <cstdlib>
 #include <cstring>
 #include <ctime>
 #include <filesystem>
@@ -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;
 }
 

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;

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

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,

ggml/src/ggml-backend.cpp

@@ -25,6 +25,11 @@
 #ifdef __APPLE__
 #include <sys/types.h>
 #include <sys/sysctl.h>
+#include <sys/mman.h>
+#endif
+
+#if defined(__linux__)
+#include <sys/mman.h>
 #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__

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) {

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(