Hugging Face's logo

ManniX-ITA
/
opencoti-llamafile

llamafile
GGUF
llama.cpp
inference-engine
cuda
opencoti
Model card Files
xet
 Community
opencoti-llamafile / docs /features /rolling_kv_compute_reserve.md
ManniX-ITA's picture
ManniX-ITA
Upload folder using huggingface_hub
9cee049 verified
preview code
|
Raw
History Blame
7.31 kB

Rolling-KV: two-pass compute-buffer reserve (fix the resident-allocation OOM)

Status: SHIPPED 2026-07-01 as 0089-rolling-kv-compute-reserve.patch (9 files, byte-identical-verified). Host-only change (llama-kv-cache*, llama-context.cpp, llama-model.cpp, llama-cparams.h, llama-memory-hybrid{,-iswa}.cpp) → build:llamafile:make rebuild ONLY (no CUDA DSO, no restamp). Vendor backup: .opencoti/vendor-backups/20260701-123044-two-pass-compute-reserve.

Gate result (bs2, RTX PRO 6000 96 GB, 2026-07-01): the exact repro config below (-c 524288 --kv-residency-mode window --vram-target 88000 -ctk f16 -ctv f16, KV 96 GB) that OOM'd the pre-0089 binary at graph_reserve now BOOTS and prefills with the streaming window engaged — prefill tps falls 4000→450 tok/s as the host tail grows (the expected cliff, per docs/features/rolling_kv.md). Two-pass gate conditions (offload && flash_attn && !devices.empty() && kv_residency_mode != head) are all met on this path, so boot-success ⇒ the measured reserve replaced the too-small 1536 MiB fixed default. Streaming-window FA math is unchanged (size-only patch), correctness carried by the 0070 M7 logit-equiv/RULER gates.

The bug (bug-1342)

The POSITION_WINDOW resident-window sizer (opencoti_compute_resident_window_cells, llama-kv-cache.cpp:227) reserves a fixed OPENCOTI_HEADINFER_AUTO_RESERVE_BYTES = 1536 MiB (line 120) for the compute scratch that graph_reserve allocates AFTER the KV cache. But the compute buffer scales with context — measured 3.63 GiB (3 800 045 568 B) at 262 144–524 288 ctx on the RTX PRO 6000. So the sizer keeps ~2.1 GiB too much KV resident and graph_reserve OOMs:

ggml_gallocr_reserve_n_impl: failed to allocate CUDA0 buffer of size 3800045568
graph_reserve: failed to allocate compute buffers
llama_init_from_model: failed to initialize the context: failed to allocate compute pp buffers

Repro: RTX6000, Qwen2.5-14B-1M Q8_0, -c 524288 --kv-residency-mode window --vram-target 88000 -ctk f16 -ctv f16. 128k boots (KV ~25.7 GiB fits); 512k (KV ~96 GiB) does not.

The author KNEW (comment at llama-kv-cache.cpp:114-119): "this reserve only covers the not-yet-allocated compute buffers. --vram-target gives the operator a tighter knob when this default is wrong." The fix removes the manual-guess dependency.

Fix: two-pass pre-measure (user-chosen 2026-07-01)

Model-agnostic: MEASURE the real compute buffer, don't estimate it.

  • Pass 1 (measure): create_memory with a measurement flag → the sizer returns a MINIMAL window (walign cells resident, rest host) → tiny device KV → boots for ANY compute size. Run sched_reserve() (creates the sched at llama-context.cpp:525, runs the pp worst-case reserve at :697). Then read the actual GPU compute buffer: sum ggml_backend_sched_get_buffer_size(sched, b) over backends whose ggml_backend_dev_type(...)==GPU.
  • Pass 2 (real): re-create_memory with `kv_compute_reserve_mib = measured
    • safety(256 MiB). Sizer now sizes the window against budget − model − measured_compute. Set sched_need_reserve=true; sched_reserve()` again.
  • Gate the two-pass so the common resident case stays SINGLE-pass / byte-identical: only engage when offload && kv_residency_mode != head(1) && full-KV-demand + a generous compute margin > free_vram (overflow possible). Reuse the kv_demand vs free_vram compare already in the sizer — expose a cheap bool opencoti_kv_may_overflow(...) helper. If it can't overflow → today's single pass, no behavior change.

Compute buffer size is INDEPENDENT of the window/tail split (it depends on n_ubatch × n_kv × n_embd, and n_kv = full ctx regardless of resident split), so the Pass-1 minimal-window measurement equals the Pass-2 compute — one measurement suffices.

Exact edit surface (threading map, verified 2026-07-01)

  1. llama-cparams.h (struct at :9): add uint32_t kv_compute_reserve_mib; (0 = use 1536 default) and bool kv_window_measure_pass; (default false).
  2. llama-context.cpp:
    • Init defaults near :104-108 (beside cparams.vram_target_mib / cparams.kv_residency_mode): kv_compute_reserve_mib=0, kv_window_measure_pass=false.
    • Two-pass orchestration wrapping the memory.reset(model.create_memory(...)) at :385 + the sched_reserve() at :451. Backend-init block (:388-449) is memory-independent — keep it before the first sched_reserve. After the Pass-1 reserve, sum GPU compute buffer, set cparams.kv_compute_reserve_mib, clear kv_window_measure_pass, memory.reset(create_memory(...)) again, sched_need_reserve=true, sched_reserve(). Guard whole block behind opencoti_kv_may_overflow.
  3. llama-kv-cache.h (ctor decl :96-145): add uint32_t compute_reserve_mib and bool window_measure_pass args (beside vram_target_mib :123 / kv_residency_mode :133).
  4. llama-kv-cache.cpp:
    • opencoti_compute_resident_window_cells (:227): take compute_reserve_mib
      • window_measure_pass. Replace the fixed OPENCOTI_HEADINFER_AUTO_RESERVE_BYTES at :271-273 with compute_reserve_mib ? (size_t)compute_reserve_mib<<20 : OPENCOTI_HEADINFER_AUTO_RESERVE_BYTES. If window_measure_pass → early-return the minimal window (walign, using the same align math at :296-298) so Pass-1 device KV is ~0.5 MiB.
    • ctor (:316): accept + forward the 2 args to the sizer call at :396.
    • (optional) expose opencoti_kv_may_overflow (reuse the auto-frac budget math) for the context-level gate.
  5. llama-kv-cache-iswa.{h,cpp}: thread the same 2 args through the wrapper to both inner llama_kv_cache ctors.
  6. llama-model.cpp create_memory (:2082 iswa, :2113 plain): pass cparams.kv_compute_reserve_mib, cparams.kv_window_measure_pass alongside cparams.vram_target_mib (:2098/:2128) and cparams.kv_residency_mode (:2102/:2132).

Gates

  • Boot repro (the fix): the exact 512k config above now BOOTS with a host tail (was OOM). Read the logged position window = C / kv_size cells resident
    • host tail; decode tps > 0, needle held.
  • Resident unchanged (no regression): a config that fits resident (e.g. 128k or --vram-target high) still logs FULLY RESIDENT, single-pass, and decodes at the same ~38 tps — byte-identical path (measure gate not tripped).
  • Curve enabled: with boot fixed, the tps-vs-host-tail curve (#583/#584 on
    1. is runnable on RTX6000: fix a ctx, sweep --vram-target down to grow the tail, read (host-tail MiB, decode tps). NOTE the binary must LOG the device/host split — the position window = ... INFO line at llama-kv-cache.cpp:310 is the measurement source (the bench's grep must match THAT line, not a non-existent CUDA_Host KV buffer line).

Constraints (standing)

Host rebuild only (bun run build:llamafile:make from repo, or the make path) — NOT a CUDA DSO rebuild, NO restamp (llama-kv-cache/llama-context are host TUs). Additive soft-fork + opencoti-hook: markers (two-pass reserve) + UPSTREAM_SYNC registry. Correctness via boot + niah, never greedy byte-equality. Commit only when asked (dev; trailers). Vendor-backup already taken (above).