opencoti-llamafile / patches /0001-lazy-slot-context-defer.patch
ManniX-ITA's picture
Upload folder using huggingface_hub
b69d9d8 verified
Raw
History Blame
10.9 kB
From: opencoti
Subject: [PATCH 1/N] lazy slot-context allocation β€” Phase 1 (deferred zero-fill)
Phase 1 of the F4 M3 lazy slot-context rollout. The eager
`ggml_backend_buffer_clear(buf, 0)` memset that commits every KV
buffer page to RSS at construction time is split out into a new
`ensure_cleared()` method invoked from inference entry points
(init_batch, prepare, update, state_read).
The buffer's malloc happens eagerly at construction time so the
backend scheduler's `sched_reserve` (which runs at llama_context
construction, before any batch) can resolve `tensor->buffer` for
every KV tensor. On Linux, malloc is essentially free for unused
ranges β€” it reserves address space without committing pages.
Pages only commit to RSS when first written. The eager memset(0)
was what wrote every page, so removing it from the constructor
moves the RSS commit to the first real inference call.
Effect for the production embedder (qwen3-0.6b,
`--parallel 2 --ctx-size 16384`, CPU mode): the KV cache's
~1,792 MiB no longer commits at process startup. Idle RSS drops
from ~3,478 MiB to roughly model + compute scratch only. First
request pays the memset cost (~1 s for 1.8 GiB), then warm-state
RSS matches the un-patched baseline. Phase 1 alone is a
cold-start win; Phases 2 (grow-on-demand) and 3 (shrink-on-idle)
target warm-state RSS.
Smoke validation (Qwen2-Math-1.5B Q4_K_M, --parallel 2 -c 4096
CPU mode): startup emits
llama_kv_cache: CPU KV buffer size = 112.00 MiB
llama_kv_cache: KV cache zero-fill deferred until first batch
First /v1/chat/completions request emits
ensure_cleared: CPU KV buffer zero-fill = 112.00 MiB (deferred)
and produces the expected response.
Coverage of ensure_cleared() call sites:
- init_batch (inference entry)
- prepare (iswa wrapper calls prepare() directly on each inner
kv_cache without going through init_batch; see
llama-kv-cache-iswa.cpp:126-159)
- update (cross-stream seq_cp resolution + RoPE shift)
- state_read (saved-state restore writes tensor data)
state_write is left const + un-cleared-safe: empty v_cells write
zero-cell-count records that state_read's
`if (cell_count == 0) continue` already handles.
clear(data=true) on an un-cleared cache skips the buffer memset
deliberately β€” v_cells.reset() above already represents "no
data," and committing RSS just to zero it would defeat Phase 1's
purpose.
The no_alloc model path is unchanged: zero-byte dummy buffers
are assigned per tensor at construction and trivially "cleared"
(zero-byte memset). `cleared` is set to true in that case.
Bench: docs/decisions/0001-lazy-slot-context.md Β§6.2 Phase 1 row.
Milestone: F4 M3 β€” see docs/features/memory_embedder.md
Design: docs/decisions/0001-lazy-slot-context.md
Upstreaming: deferred; propose to llamafile + llama.cpp
maintainers after F4 M3's full phase-set ships and bench
numbers are public. This is a generally-useful patch.
---
diff --git a/llama.cpp/src/llama-kv-cache.cpp b/llama.cpp/src/llama-kv-cache.cpp
--- a/llama.cpp/src/llama-kv-cache.cpp
+++ b/llama.cpp/src/llama-kv-cache.cpp
@@ -250,7 +250,15 @@ llama_kv_cache::llama_kv_cache(
}
}
- // allocate tensors and initialize the buffers to avoid NaNs in the padding
+ // opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+ // Eager malloc, deferred zero-fill. The address-space reservation done
+ // by ggml_backend_alloc_ctx_tensors_from_buft is essentially free on
+ // Linux β€” pages are not committed to RSS until they are first written.
+ // The `ggml_backend_buffer_clear(buf, 0)` memset below is what commits
+ // every page. By skipping it for the non-no_alloc path we keep the
+ // scheduler-visible tensor->buffer wiring intact (sched_reserve runs
+ // at construction time and asserts buffer_id >= 0) while moving the
+ // RSS commit to the first ensure_cleared() call.
for (auto & [buft, ctx] : ctx_map) {
ggml_backend_buffer_t buf;
if (model.hparams.no_alloc) {
@@ -259,7 +267,7 @@ llama_kv_cache::llama_kv_cache(
t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it
}
} else {
- buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
+ buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer (malloc'd but not zeroed yet)
}
if (!buf) {
throw std::runtime_error("failed to allocate buffer for kv cache");
@@ -267,10 +275,23 @@ llama_kv_cache::llama_kv_cache(
LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0);
- ggml_backend_buffer_clear(buf, 0);
+ if (model.hparams.no_alloc) {
+ // Dummy buffer is size 0; clearing is a no-op but keeps the
+ // upstream invariant that constructed buffers are zeroed.
+ ggml_backend_buffer_clear(buf, 0);
+ }
+ // else: leave the buffer un-cleared. ensure_cleared() will zero it
+ // on the first inference entry point that touches tensor data.
+
ctxs_bufs.emplace_back(std::move(ctx), buf);
}
+ if (model.hparams.no_alloc) {
+ cleared = true;
+ } else {
+ LLAMA_LOG_INFO("%s: KV cache zero-fill deferred until first batch\n", __func__);
+ }
+
{
const size_t memory_size_k = size_k_bytes();
const size_t memory_size_v = size_v_bytes();
@@ -327,6 +348,22 @@ llama_kv_cache::llama_kv_cache(
debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0;
}
+// opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+void llama_kv_cache::ensure_cleared() {
+ if (cleared) {
+ return;
+ }
+
+ for (auto & [_, buf] : ctxs_bufs) {
+ LLAMA_LOG_INFO("%s: %10s KV buffer zero-fill = %8.2f MiB (deferred)\n", __func__,
+ ggml_backend_buffer_name(buf.get()),
+ ggml_backend_buffer_get_size(buf.get())/1024.0/1024.0);
+ ggml_backend_buffer_clear(buf.get(), 0);
+ }
+
+ cleared = true;
+}
+
void llama_kv_cache::clear(bool data) {
for (uint32_t s = 0; s < n_stream; ++s) {
v_cells[s].reset();
@@ -334,8 +371,14 @@ void llama_kv_cache::clear(bool data) {
}
if (data) {
- for (auto & [_, buf] : ctxs_bufs) {
- ggml_backend_buffer_clear(buf.get(), 0);
+ // opencoti F4 M3 Phase 1: if the deferred zero-fill has not happened
+ // yet there is nothing to clear β€” v_cells.reset() above already
+ // satisfies the logical "data cleared" state. Don't commit RSS just
+ // to zero pages that have never been written.
+ if (cleared) {
+ for (auto & [_, buf] : ctxs_bufs) {
+ ggml_backend_buffer_clear(buf.get(), 0);
+ }
}
}
}
@@ -630,6 +673,9 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
bool embd_all) {
GGML_UNUSED(embd_all);
+ // opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+ ensure_cleared();
+
do {
balloc.split_reset();
@@ -674,6 +720,12 @@ llama_memory_context_ptr llama_kv_cache::init_update(llama_context * lctx, bool
}
llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_ubatch> & ubatches) {
+ // opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+ // The iswa wrapper calls prepare() directly on each inner kv_cache without
+ // going through init_batch(), so the ensure_cleared() in init_batch alone
+ // is not enough. Calling here covers both paths and is idempotent.
+ ensure_cleared();
+
llama_kv_cache::slot_info_vec_t res;
struct state_t {
@@ -744,6 +796,15 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co
auto * sched = lctx->get_sched();
+ // opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+ // update() can be reached without init_batch first (e.g. update-only paths
+ // for cross-stream copy or RoPE shift). When that happens with the buffer
+ // still deferred we have nothing to copy or shift β€” make the dependency
+ // explicit so the tensor->data dereferences below are safe.
+ if (!sc_info.empty() || do_shift) {
+ ensure_cleared();
+ }
+
if (!sc_info.empty()) {
assert(n_stream > 1 && "stream copy should never happen with a single stream");
@@ -1900,6 +1961,11 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
+ // opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+ // state_read writes the saved KV blob into our tensors, which requires
+ // backing storage. Force the deferred allocation before we read.
+ ensure_cleared();
+
uint32_t n_stream_cur;
io.read(&n_stream_cur, sizeof(n_stream_cur));
if (n_stream_cur != n_stream) {
diff --git a/llama.cpp/src/llama-kv-cache.h b/llama.cpp/src/llama-kv-cache.h
--- a/llama.cpp/src/llama-kv-cache.h
+++ b/llama.cpp/src/llama-kv-cache.h
@@ -152,6 +152,16 @@ public:
bool get_has_shift() const;
+ // opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+ // Idempotent: zeroes the KV-cache backend buffers on first call.
+ // The buffers are malloc'd eagerly at construction (so the scheduler
+ // can resolve tensor->buffer at reserve time) but NOT cleared β€” on
+ // Linux, malloc reserves address space without committing pages, so
+ // RSS stays low until something writes. This method does the memset
+ // that commits the pages. Called from inference entry points
+ // (init_batch, prepare, update, state_read).
+ void ensure_cleared();
+
ggml_type type_k() const;
ggml_type type_v() const;
@@ -256,6 +266,13 @@ private:
// ggml contexts for the KV cache along with the allocated backend buffers:
std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
+ // opencoti F4 M3 Phase 1 β€” see docs/decisions/0001-lazy-slot-context.md
+ // True once the backend buffers in ctxs_bufs have been zero-filled.
+ // The constructor sets this true for no_alloc models (whose dummy
+ // buffers are size 0 β€” clearing is a no-op) and false otherwise. The
+ // first ensure_cleared() call zeroes the buffers and flips this true.
+ bool cleared = false;
+
// the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot())
// note: this is not part of the KV state and it's only used to speed-up the find_slot() method
std::vector<uint32_t> v_heads;