UPSTREAM PR #18747: ggml, llama : add CPU paged attention for memory-efficient KV cache

[loci-dev]

Mirrored from ggml-org/llama.cpp#18747

Summary

Add vLLM-inspired paged attention to llama.cpp, enabling block-based KV cache access for improved memory locality on large models and significant memory savings via dynamic allocation.

Motivation

Large language models (70B+) are memory-bound during attention computation. By organizing KV cache access into blocks with prefetching, we achieve better cache utilization and reduced memory stalls. Additionally, paged attention enables memory savings by limiting KV cache allocation to only the blocks actually needed.

This implementation is inspired by PagedAttention (Kwon et al., 2023) used in vLLM.

Related Work

This PR provides a CPU-only implementation, complementing PR #17579 which implements CUDA-only paged attention. The two approaches differ in design:

Aspect	This PR (CPU)	PR #17579 (CUDA)
Backend	CPU with __builtin_prefetch()	CUDA kernels (V1/V2)
GGML Op	GGML_OP_FLASH_ATTN_EXT_PAGED	GGML_OP_PAGED_ATTENTION
Integration	Extends existing flash attention	New paged KV cache classes
Flag	--paged-attn N (block size)	--pagedattention (boolean)

These implementations are independent and can coexist. Per CONTRIBUTING.md guidance to focus on CPU first, this PR establishes the CPU path while PR #17579 addresses CUDA. Future work could unify the APIs if both are merged.

vLLM Design Adaptation

We studied vLLM's PagedAttention implementation to understand the core concepts:

• Block table indirection: Mapping logical token positions to physical memory blocks
• Block prefetching: Pre-loading next block's data while processing current block
• Dynamic allocation: Allocating blocks only as tokens are generated

Key adaptations for llama.cpp:

• Integrated with existing flash attention infrastructure rather than replacing it
• Block table passed as I32 tensor through GGML graph
• Block tracking structures designed for single-threaded llama_context access pattern
• Environment variables for easy experimentation before committing to CLI flags

---

Implementation

1. GGML Kernel (ggml/src/ggml-cpu/ops.cpp)

New op GGML_OP_FLASH_ATTN_EXT_PAGED (~343 lines) with:

• Block table indirection: Maps logical positions → physical blocks
• Block prefetching: __builtin_prefetch() on block boundaries
• Invalid block handling: Skips blocks marked -1

Key difference from standard flash attention:

// Standard: direct linear access
const char * k_data = k->data + ic * nbk1;

// Paged: block table indirection
logical_block = ic / block_size;
physical_block = block_table[seq_idx * max_blocks + logical_block];
physical_pos = physical_block * block_size + (ic % block_size);
const char * k_data = k->data + physical_pos * nbk1;

2. Block Tracking (src/llama-kv-block.h)

Data structures for block-based memory management:

• llama_kv_block_pool: Stack-based O(1) allocator with free list
• llama_kv_block_table: Sequence → physical block mapping
• llama_kv_block_stats: Utilization and fragmentation tracking

3. KV Cache Integration (src/llama-kv-cache.cpp)

• enable_blocks(N): Enable block tracking with N tokens/block
• build_block_table_tensor(): Create I32 tensor for graph
• update_block_tokens(): Dynamic block allocation in apply_ubatch()
• seq_rm(): Block deallocation when sequences are removed

4. Graph Integration (src/llama-graph.cpp)

Routes to paged kernel when block table is available:

if (block_table != nullptr && block_size > 0) {
    cur = ggml_flash_attn_ext_paged(ctx0, q, k, v, mask, block_table, ...);
} else {
    cur = ggml_flash_attn_ext(ctx0, q, k, v, mask, ...);
}

---

Performance

Benchmark Environment

• CPU: AMD EPYC (96 threads)
• Flash attention: enabled
• NUMA: numactl --interleave=all for 70B model

Small Model: Qwen3-1.7B-Q8_0 (16 threads)

Configuration	Prompt (pp256)	Generation (tg32)	pp Change	tg Change
Baseline	660.21 t/s	62.68 t/s	-	-
Paged-64	674.92 t/s	61.27 t/s	+2.2%	-2.2%
Paged-128	673.81 t/s	62.67 t/s	+2.1%	±0%
Paged-256	685.56 t/s	62.02 t/s	+3.8%	-1.0%
Paged-512	684.15 t/s	62.40 t/s	+3.6%	-0.4%

Conclusion: Negligible impact on small compute-bound models.

Medium Model: DeepSeek-R1-Distill-Qwen-32B-Q4_K_M (96 threads)

Configuration	Prompt (pp256)	Generation (tg32)	pp Change	tg Change
Baseline	110.71 t/s	9.46 t/s	-	-
Paged-64	114.00 t/s	8.38 t/s	+3.0%	-11.4%
Paged-128	109.46 t/s	7.22 t/s	-1.1%	-23.7%
Paged-256	112.30 t/s	8.29 t/s	+1.4%	-12.4%
Paged-512	111.45 t/s	5.43 t/s	+0.7%	-42.6%

Conclusion: Generation overhead on compute-bound medium models. Prompt processing slightly improved.

Large Model: Meta-Llama-3.1-70B-Instruct-Q4_K_M (96 threads, numactl)

Configuration	Prompt (pp1024)	Generation (tg64)	pp Change	tg Change
Baseline	50.52 t/s	2.14 t/s	-	-
Paged-64	53.94 t/s	3.77 t/s	+6.8%	+76.2%
Paged-128	54.18 t/s	2.66 t/s	+7.2%	+24.3%
Paged-256	53.62 t/s	2.31 t/s	+6.1%	+7.9%
Paged-512	51.64 t/s	2.54 t/s	+2.2%	+18.7%

Conclusion: Significant speedup on memory-bound large models, especially with smaller block sizes.

Memory Savings

Small Model: Qwen3-1.7B (40960 token default context)

Configuration	Effective Context	KV Buffer	Memory Savings
Baseline	40960 tokens	4480 MiB	-
MAX_BLOCKS=50×64	3200 tokens	350 MiB	92.2%
MAX_BLOCKS=100×64	6400 tokens	700 MiB	84.4%
MAX_BLOCKS=200×64	12800 tokens	1400 MiB	68.8%

Medium Model: DeepSeek-R1-Distill-Qwen-32B (131072 token default context)

Configuration	Effective Context	KV Buffer	Memory Savings
Baseline	131072 tokens	32768 MiB	-
MAX_BLOCKS=100×256	25600 tokens	6400 MiB	80.5%
MAX_BLOCKS=200×256	51200 tokens	12800 MiB	60.9%
MAX_BLOCKS=400×256	102400 tokens	25600 MiB	21.9%

Large Model: Meta-Llama-3.1-70B-Instruct (131072 token default context)

Configuration	Effective Context	KV Buffer	Memory Savings
Baseline	131072 tokens	40960 MiB	-
MAX_BLOCKS=100×256	25600 tokens	8000 MiB	80.5%
MAX_BLOCKS=200×256	51200 tokens	16000 MiB	60.9%
MAX_BLOCKS=400×256	102400 tokens	32000 MiB	21.9%

Key Insight: Memory savings scale with context reduction. For typical interactive use cases (8K-25K tokens), paged attention can reduce KV cache memory by 60-80%.

---

Correctness

• Outputs identical to non-paged at same seed (verified across 1.7B, 32B, 70B models)
• Unit tests: 19 tests in test-kv-block.cpp covering block pool and table operations

---

Usage

Important: Paged attention is a variant of flash attention, not a replacement. It requires --flash-attn on because it operates within the flash attention code path.

# Enable paged attention (block size = 256 tokens)
./llama-cli -m model.gguf --flash-attn on --paged-attn 256 ...

# With memory reduction (100 blocks max = 25600 token context limit)
./llama-cli -m model.gguf --flash-attn on --paged-attn 256 --paged-attn-max-blocks 100 ...

# Environment variables (useful for testing)
LLAMA_PAGED_ATTN=256 ./llama-cli -m model.gguf --flash-attn on ...
LLAMA_PAGED_ATTN=64 LLAMA_PAGED_ATTN_MAX_BLOCKS=100 ./llama-cli -m model.gguf --flash-attn on ...

---

Design Decisions

Why a New GGML Op vs Extending Existing?

We considered adding a conditional path inside ggml_compute_forward_flash_attn_ext but rejected this:

1. Hot path impact: Adding conditionals would degrade performance for the common non-paged case
2. Memory access pattern divergence: Linear vs block-table indirection are fundamentally different
3. Prefetching logic: Block boundary detection has no analog in standard flash attention
4. Precedent: llama.cpp already has multiple attention variants

Why CPU-Only?

Per llama.cpp CONTRIBUTING.md: "When adding support for a new model or feature, focus on CPU support only in the initial PR."

Additionally, PR #17579 already provides a CUDA implementation with a different architecture. Rather than duplicating that work, this PR focuses on CPU where:

• Block prefetching (__builtin_prefetch()) provides cache-line optimization
• Memory-bound 70B+ models show significant speedups (+76% generation)
• The implementation integrates cleanly with existing flash attention infrastructure

The two PRs can coexist independently or be unified later if maintainers prefer a single API.

Memory Overhead of Block Tracking

Component	Size	For 70B (8K context, 128 blocks)
Block metadata	12 bytes/block	1.5 KB
Block table	4 bytes/block/seq	~0.6 KB/seq
Free stack	4 bytes/block	0.5 KB

Total: ~3 KB vs KV cache size of 8-32 GB. Negligible (<0.00001%).

Thread Safety Model

Block pool and block table are explicitly NOT thread-safe. This matches llama_context design:

• Each context owns its own KV cache with independent block pools
• Block operations occur during graph construction, not parallel compute
• Thread safety documented in llama-kv-block.h

Why Single PR?

We considered splitting into multiple PRs (kernel, infrastructure, integration) but chose single PR because:

1. Reviewer feedback: ngxson requested consolidating related changes for easier testing
2. Precedent: PR #12801 (DeepSeek MLA) was ~450 lines as single PR
3. Testability: Partial implementation cannot be meaningfully tested

---

Files Modified

ggml/include/ggml.h          |  16 ++
ggml/src/ggml-cpu/ggml-cpu.c |   6 +
ggml/src/ggml-cpu/ops.cpp    | 343 +++++++++++++++++++++++++++++++
ggml/src/ggml-cpu/ops.h      |   1 +
ggml/src/ggml.c              |  64 ++++-
src/llama-graph.cpp          |  42 ++++-
src/llama-graph.h            |   8 +-
src/llama-kv-cache.cpp       | 137 +++++++++++++
src/llama-kv-cache.h         |  44 +++++
src/llama-kv-block.h         | 403 +++++++++++++++++++++++++++++++++++ (new)
common/arg.cpp               |  20 ++
common/common.cpp            |   8 +
common/common.h              |   4 +
tests/test-kv-block.cpp      | 280 ++++++++++++++++++++++++ (new)

---

Future Work

• GPU backends: PR #17579 addresses CUDA; Metal/Vulkan could follow similar patterns
• API unification: If both PRs merge, consider unifying --paged-attn and --pagedattention flags
• Prefix caching: Block sharing via copy-on-write (CoW infrastructure in place)
• Auto-tuning: Automatic block size selection based on model size and hardware

---

Test Plan

• test-kv-block passes (19 tests)
• test-backend-ops -o FLASH_ATTN_EXT_PAGED passes
• Output determinism verified (same seed → same output)
• Memory savings verified (up to 92% reduction)
• Performance benchmarks across 3 model sizes

---

Acknowledgments

Inspired by vLLM's PagedAttention implementation and the paper "Efficient Memory Management for Large Language Model Serving with PagedAttention" (Kwon et al., 2023).
The authors used claude to construct schematics of vLLM's paged attention infrastructure, review our plan for integrating with existing flash attention rather than replacing it outright in llama.cpp, build unit tests, and draft PR submission.

[loci-review]

Explore the complete analysis inside the Version Insights

I've successfully generated a performance summary report for your project. The report shows:

Key Highlights:

✅ Major Performance Improvements:

• Several functions show dramatic improvements (95% to 307% throughput increases)
• STL tree operations (std::_Rb_tree::end()) improved by over 300%
• Unicode processing function improved by 95%

⚠️ Performance Regressions:

• Some STL vector and iterator operations show 72-75% throughput decreases
• These may indicate changes in usage patterns rather than actual performance issues

The analysis compares version 0089a8f1-ee88-11f0-a055-c529586b3e1a against base version 5eb15641-ee50-11f0-a055-c529586b3e1a for the llama.cpp repository (PR #883).

Would you like me to provide more details about any specific aspect of the report?