[Draft] feat: implement paged KV cache and attention - #22569
Conversation
github-actions
Bot
added
model
|
Hi @matiaslin, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
|
Multiple backend changes in one PR concern: AI-generated content: Large PR: |
Add an opt-in paged KV cache backed by a fixed pool of fixed-size blocks, with a continuous-batching scheduler that manages admission, eviction, and swap-in/out between GPU and CPU block pools. Enabled via --kv-paged (disabled by default). The existing KV paths are left untouched. This is phase 1 of paged KV cache and attention intended for design review. Phase 1 restrictions (enforced at context creation): - single CUDA device only - full offload only - n_batch == n_ubatch Phase 2 (not yet implemented): - CoW / prefix caching - seq_cp / seq_keep / seq_div / seq_add - state_write / state_read Core components: - llama_block_manager: physical block lifecycle with watermark and refcounts. - llama_kv_cache_paged: block-based KV cache with GPU/CPU block pools, watermark-based allocation, and swap-in/out under memory pressure. - llama_paged_scheduler: FCFS scheduler with token-budget admission control, deadlock and livelock detection, exposed via a new C API. - ggml paged_attn operation: CPU and CUDA backends - fit_params: auto-sizes n_gpu_blocks from available VRAM when the user does not pass -ngpub explicitly. - examples/paged: end-to-end demo of the continuous-batching loop - tests/test-paged-kv: unit tests for block manager, paged KV cache, and scheduler. - tests/test-paged-kv-e2e: end-to-end test of greedy-equivalence against the non-paged path. Verified via test-paged-kv-e2e: llama, qwen2, qwen3 (separate QKV projections) and falcon, starcoder2 (fused QKV projections). SWA architectures (gemma3, llama4, etc.) not supported in this phase. Signed-off-by: matiaslin <matiasenoclin@gmail.com>
matiaslin
changed the title
|
If the only reason to do this to enable more than 25 parallel requests, it does not make sense to do this as most users of llama.cpp are not going to run that many parallel requests. But it could be useful if it brings some improvements to sub 20 batched range |
|
your PR and your initial proposal is pretty much slop. AFAIK paged attn is only beneficial when the incoming requests have the same prompt prefix (usually system prompt). we can obviously implement the same scheduling for cross-request KV reuse at server level this PR will likely be rejected, and I personally don't wish to proceed with contributors who don't discuss complex changes in advanced as specified in the guideline. |
|
Thank you @am17an and @ngxson for the comments. I responded to you guys' questions in the original discussion post: #21961. We can continue the discussion over there if/when you guys get the chance. Thank you again. (I would urge further commentors to continue the discussion in the original post, thank you!). |
|
Related to #17579 |
Cherry-pick new files only (no conflict with existing turbo/triattention code): - ggml/src/ggml-cuda/pagedattn.cu/.cuh -- CUDA paged attention kernel - src/llama-block-manager.cpp/.h -- GPU/CPU block pool allocator - src/llama-kv-cache-paged.cpp/.h -- paged KV cache implementation - src/llama-paged-scheduler*.cpp/.h -- FCFS scheduler with deadlock detection - include/llama.h -- public API declarations - src/CMakeLists.txt -- register new sources - tests/test-paged-kv*.cpp -- unit and e2e tests Not yet wired to context init or CLI flags. Engine present, UI pending. Based on arXiv:2309.06180. Draft PR: ggml-org/llama.cpp#22569. Assisted-by: Claude Sonnet
PR ggml-org#22569 (ggml-org/llama.cpp) ist DRAFT, dirty, 2+ Monate inaktiv, CUDA-only. Der 2.5x-Durchsatz entsteht bei 247 concurrent sequences (vLLM Cloud-Serving-Design). Unser Edge/LAN-Setup fährt -np 1/-np 2 — bei 1-2 sequences ist Paged ~3% langsamer (reiner Overhead). TurboQuant (turbo3/turbo4, 3-4 bit KV-Kompression) ist die richtige KV-Optimierung für unseren Use-Case: komprimiert KV-Daten direkt statt nur Allokation zu managen. Mars (Vulkan) und Venus (Vulkan) würden von CUDA-only PagedAttention nicht profitieren. M5 (Coopmat2 + Multi-GPU) → ✅ abgeschlossen: AtomicBot-ai#12✅, AtomicBot-ai#20✅, AtomicBot-ai#21❌.
matiaslin's paged KV cache + continuous-batching scheduler (draft PR ggml-org#22569) was written against upstream f5c643f on 2026-04-08. Mainline has moved 294 commits since. This commit carries the fixes needed to land it on current master. Conflicts resolved (17 blocks across 11 files). Most were additive, but four were not and would have compiled into silently wrong behaviour: 1. ggml-cuda.cu dispatch switch: keeping both cases naively let GGML_OP_LIGHTNING_INDEXER FALL THROUGH into GGML_OP_PAGED_ATTN, because the two sides shared the trailing 'break;'. Added an explicit break. 2. llama-model.cpp: the branch wraps cache creation in if(kv_paged){...}else{...}, but its else-branch called the April llama_kv_cache ctor. Mainline's has since gained 'hparams' and a 'filter' argument. Took their structure with HEAD's call. 3. hparams.n_layer became a METHOD in mainline (it returns n_layer_all minus n_layer_nextn). The branch used it as a field, in llama-context.cpp and llama-model.cpp. 4. Three conflicts split a function or class body, with the closing brace living outside the conflict region, so both sides relied on it: ops.cpp (fwht/paged_attn), llama-graph.h (llm_graph_input_dsv4 / attn_kv_paged), llama-graph.cpp (build_attn_inp_k_dsa / build_attn_inp_kv_paged). Each needed a brace inserted between the two sides. common.cpp additionally carried the stale April common_init_result signature alongside mainline's model_only one. create_memory keeps the extended signature (backend_gpu, backend_cpu): the paged block manager allocates its GPU and CPU pools up front and needs the handles. Verified: builds on macOS with Metal, and test-paged-kv passes end to end (block manager, allocator, scheduler, deadlock detection, oversize-prompt rejection). NOT yet verified on CUDA, and no throughput number claimed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
4 participants
Overview
Add an opt-in paged KV cache backed by a fixed pool of fixed-size blocks, with a continuous-batching scheduler that manages admission, eviction, and swap-in/out between GPU and CPU blocks. Enabled via
--kv-paged(disabled by default). The existing KV paths are left untouched.Paged KV cache addresses the memory over-allocation that the unified KV cache introduces by splitting the KV buffer into fixed-size blocks (default 16 tokens each). Because these blocks are allocated on demand, a given sequence only holds the blocks that it truly needs. In turn, this allows the admittance of several more sequences than the unified layout can accommodate.
Single
NVIDIA A10G,Llama-3(8B f16), same KV budget for both modes,n_batch = n_ubatch = 1024:At the same concurrency level,
pagedandunifiedare within ~3% (agg. tok/s), but we observe thatunifiedOOMs at 26 concurrent sequences, whilepagedcan reach up to 247 sequences. This higher concurrency helps us see that at a peak-vs-peak comparisonpagedreaches a 2.5x aggregate throughput.Additional information
A while back, I started a design discussion in
Ideas, see: #21961. That post provides more details on motivation backed by benchmarks as well as, design decisions of what was included and excluded for this first implementation phase.I recognize that maintainers are busy, so I hope this draft PR will make the changes I propose more tangible and a bit easier to review.
The feedback I would appreciate are:
llama_paged_scheduler_*C API the right approach?llama_paged_batch_inforemain public or internal with getters?llama.cpp. I'm open to pull those changes in for phase 1, if we decide it's necessary).ggml_paged_attnop as one PR, then scheduler + tests)?(There are more questions in the discussion post, but above are the most relevant to this PR).
Requirements