Paged Attention & Continuous Batching
Relevant source files
The following files were used as context for generating this wiki page:
Paged Attention is a memory management system in oxidize-core that implements vLLM-style KV cache optimization. It partitions the Key-Value (KV) cache into fixed-size blocks to eliminate memory fragmentation and allow flexible sequence scheduling. This system enables continuous batching, which processes new requests and running sequences simultaneously to maximize server throughput.
Sources: oxidize-core/src/paged_attention/AGENTS.md, PRD.md:164-165
Memory Management via BlockPool
The BlockPool manages physical memory for the KV cache by dividing it into discrete blocks. Each block typically stores a fixed number of tokens, defined as BLOCK_SIZE. This architecture supports Copy-on-Write (CoW) sharing, where blocks with identical KV contents share physical memory through reference counting and block hashing.
Sources: oxidize-core/src/paged_attention/AGENTS.md, oxidize-core/src/paged_attention/block_pool.rs
Core Components
- BlockId: A unique identifier for a physical memory block.
- PhysicalTokenBlock: The underlying storage unit for a fixed number of tokens.
- BlockTable: A mapping structure that links logical sequences to their respective physical blocks.
- Reference Counting: Tracks how many sequences utilize a specific block to enable efficient memory reuse and deallocation.
The diagram shows how logical sequences map to physical blocks through block tables, allowing block sharing.
Sources: oxidize-core/src/paged_attention/block_pool.rs, oxidize-core/src/paged_attention/mod.rs
Request Scheduler
The Scheduler manages the sequence lifecycle and token budget. It dispatches requests based on a First-Come-First-Served (FCFS) strategy, ensuring that older requests are prioritized during resource allocation. The scheduler tracks sequence states and prevents memory exhaustion by managing a strict token budget.
Sources: oxidize-core/src/paged_attention/scheduler.rs, oxidize-core/src/paged_attention/AGENTS.md
Scheduling Logic
- Waiting Queue: New requests enter a waiting queue based on their arrival order.
- Resource Check: The scheduler verifies if sufficient blocks are available in the
BlockPoolfor the next request. - Prioritization: Prefill sequences (initial prompt processing) generally take priority, followed by running decode sequences (token generation).
- Continuous Batching: The system integrates new sequences into the active processing batch without waiting for currently running sequences to finish.
Sources: oxidize-core/src/paged_attention/scheduler.rs, oxidize-server/src/AGENTS.md
Sequence State Machine
The scheduler manages sequences through a strictly forward-moving state machine.
| State | Description |
|---|---|
| Waiting | The sequence is in the queue, awaiting resource allocation. |
| Running | The sequence has allocated blocks and is currently being processed. |
| Finished | The sequence has reached an EOS token or limit and its blocks are freed. |
Sources: oxidize-core/src/paged_attention/scheduler.rs, oxidize-core/src/paged_attention/AGENTS.md
The state machine depicts the lifecycle of an inference request within the paged attention system.
Sources: oxidize-core/src/paged_attention/scheduler.rs
Continuous Batching & Server Integration
Continuous batching is implemented at the server level using a wait queue and the paged attention scheduler. It allows the oxidize-server to handle concurrent requests by overlapping the prefill phase of new requests with the decoding phase of existing requests.
Sources: oxidize-server/src/AGENTS.md, PRD.md:164-165
Configuration & Limits
The system uses specific parameters to maintain stability and performance.
| Parameter | Role |
|---|---|
BLOCK_SIZE | Fixed number of tokens per block (Default: 16). |
max_num_seqs | Maximum number of concurrent sequences in a batch. |
max_num_batched_tokens | Total token budget across all active sequences. |
Sources: oxidize-core/src/paged_attention/scheduler.rs, oxidize-core/src/paged_attention/mod.rs
Implementation Details
The implementation utilizes tokio::sync::Mutex (or StdMutex in some candidates) to manage state in asynchronous contexts. The scheduler.rs module contains the core loop that interacts with the BlockPool to allocate memory for every new token generated during the Running state. If allocation fails, the scheduler must preempt sequences to free resources.
Sources: oxidize-core/src/paged_attention/scheduler.rs, oxidize-core/src/paged_attention/AGENTS.md, oxidize-server/src/AGENTS.md
This sequence shows the interaction between the Scheduler, BlockPool, and Model during a single token generation step.
Sources: oxidize-core/src/paged_attention/scheduler.rs, oxidize-core/src/paged_attention/block_pool.rs
Paged Attention and Continuous Batching provide the infrastructure necessary for high-throughput, memory-efficient LLM inference within the oxidize project, specifically supporting server-side deployments and complex scheduling requirements.