Keep token order across cache pages

Separate a token’s position from its cache address

During generation, each decoder layer saves a key and a value for every token it processes. Later queries read those cached rows in the sequence’s logical order. That order gives a token its position; it need not describe where the row lives in a buffer.

The existing contiguous cache keeps a lane’s successive token rows adjacent, so a row’s address follows directly from its logical position. To place neighboring groups of tokens in separate physical regions, the cache needs an extra lookup. A page table supplies it: split the logical sequence into fixed-size pages, then record where each logical page is stored. Attention must continue to read logical positions in order even when the physical pages are out of order.

Here we reserve every page when constructing the cache. The new capability is explicit placement and its accounting; this fixed pool does not automatically use less storage than the contiguous cache.

Follow the same seven positions across page boundaries

The existing tokenizer frames abcde as the token-ID sequence [0,99,100,101,102,103,1]. Its seven logical positions are $0$ through $6$. The model has $4$ query heads and $2$ shared K/V heads, each with $2$ features, and this example has one batch lane. Each of its two decoder layers has logical capacity $12$, with $3$ token positions per page.

Use page table $T=(2,0,3,1)$. Entry $0$ assigns logical page $0$ to physical page $2$; entry $1$ assigns logical page $1$ to physical page $0$, and so on. Within a page, token offsets remain $0,1,2$.

• Logical positions $0,1,2$ occupy physical page $2$ at offsets $0,1,2$.
• Positions $3,4,5$ occupy physical page $0$ at the same three offsets.
• Position $6$ starts logical page $2$, stored in physical page $3$ at offset $0$.
• Logical page $3$ maps to physical page $1$. Its positions $9,10,11$ are reserved but unfilled.

For batch lane $0$, shared K/V head $0$ and its first feature, Rust reports scalar indices $24,26,28,0,2,4,36$ for positions $0$ through $6$ in either array. Thus crossing from position $2$ to $3$ moves from scalar index $28$ back to $0$. Position $6$ begins at index $36$. These are indices into a key array and, separately, a value array; the two arrays do not share their values. The logical sequence still advances from one token to the next.

Locate a row without changing its logical meaning

Let $C>0$ be a layer cache’s logical token capacity and $P>0$ its token positions per page. Reserve $N$ physical pages. For a valid logical position $0\le p<C$, let $q$ be its logical page, $r$ its within-page token offset and $s$ the physical page selected by immutable table $T$:

Ceiling rounds the capacity up to enough complete pages. Integer division selects a logical page; the remainder selects a token within it. The table is a permutation of physical indices $0$ through $N-1$: each appears exactly once. Different logical pages therefore cannot overwrite each other. Different positions within a page have different offsets.

    pub fn address(&self, logical_position: usize) -> Option<KvPageAddress> {
        if logical_position >= self.capacity {
            return None;
        }
        let logical_page = logical_position / self.page_tokens;
        Some(KvPageAddress {
            logical_position,
            logical_page,
            physical_page: self.page_table[logical_page],
            offset_in_page: logical_position % self.page_tokens,
        })
    }

    /// The caller has checked lane, feature width and total buffer size.
    pub(crate) fn row_start(
        &self,
        lane: usize,
        lanes: usize,
        width: usize,
        position: usize,
    ) -> usize {
        let address = self.address(position).expect("checked logical KV position");
        ((address.physical_page * lanes + lane) * self.page_tokens + address.offset_in_page) * width
    }

The constructor rejects zero capacity or page size, the wrong number of table entries, repeated or out-of-range physical pages, and overflowing rounded capacity. Those checks establish the non-aliasing map before a cache can use it.

    pub fn new(
        capacity: usize,
        page_tokens: usize,
        page_table: Vec<usize>,
    ) -> Result<Self, PagedKvError> {
        if capacity == 0 {
            return Err(PagedKvError::ZeroCapacity);
        }
        if page_tokens == 0 {
            return Err(PagedKvError::ZeroPageSize);
        }
        let pages = capacity.div_ceil(page_tokens);
        let physical_tokens = pages
            .checked_mul(page_tokens)
            .ok_or(PagedKvError::Overflow)?;
        if page_table.len() != pages {
            return Err(PagedKvError::PageCount {
                expected: pages,
                actual: page_table.len(),
            });
        }
        let mut used = Vec::new();
        used.try_reserve_exact(pages)
            .map_err(|_| PagedKvError::AllocationFailed { entries: pages })?;
        used.resize(pages, false);
        for (logical_page, &physical_page) in page_table.iter().enumerate() {
            if physical_page >= pages {
                return Err(PagedKvError::PhysicalPage {
                    logical_page,
                    physical_page,
                    pages,
                });
            }
            if used[physical_page] {
                return Err(PagedKvError::AliasedPage { physical_page });
            }
            used[physical_page] = true;
        }
        Ok(Self {
            capacity,
            page_tokens,
            physical_tokens,
            page_table,
        })
    }

Include the batch, shared head and feature

Each physical array has shape $[N,B,H_{kv},P,d]$: physical page, batch lane, shared K/V head, token offset within the page and feature. Here $B$ is the positive batch count, $H_{kv}$ the positive shared-head count and $d$ the positive feature width. Indices $b,h,f$ select a batch lane, shared head and feature, respectively, with $0\le b<B$, $0\le h<H_{kv}$ and $0\le f<d$. After finding $s$ and $r$, the flat scalar index $i$ in either array is

This is the expanded row-major index for those five axes. In the default example each physical page occupies $1\times2\times3\times2=12$ scalar slots in each array. Physical page $2$ therefore starts at index $24$; each token row in head $0$ advances by $2$ scalars. Head $1$ begins another $3\times2=6$ scalars into that page.

A scalar index counts f64 elements in one array. Multiplying it by $8$ gives a byte offset within that array, not a host or device pointer. Keys and values use the same index calculation in separate buffers. A logical snapshot still has shape $[B,H_{kv},n,d]$, where $n$ is the current live prefix length; it gathers only filled rows in logical order.

    /// First scalar of a reserved row in each physical buffer, even if not yet filled.
    pub fn physical_row_start(&self, batch: usize, head: usize, position: usize) -> Option<usize> {
        if batch >= self.batch_size || head >= self.heads || position >= self.capacity {
            return None;
        }
        Some(self.row_start(batch * self.heads + head, position))
    }

    fn row_start(&self, lane: usize, position: usize) -> usize {
        self.paged.as_ref().map_or_else(
            || (lane * self.capacity + position) * self.head_width,
            |layout| {
                layout.row_start(
                    lane,
                    self.batch_size * self.heads,
                    self.head_width,
                    position,
                )
            },
        )
    }

    pub fn paged_layout(&self) -> Option<&PagedKvLayout> {
        self.paged.as_ref()
    }

The attention loop still visits eligible logical positions in increasing order. It performs the same dot products and weighted sums with the same operands and reduction order in this scalar CPU path. RoPE rotates keys and queries using logical positions, not physical page indices. A page size is a storage choice, not an attention-window size: a configured window still determines which logical positions may contribute. Paging neither evicts older rows nor expands their eligible context.

Count live rows and two kinds of unused payload

The number of query heads, $H_q$, does not multiply stored K/V rows. Each logical token position needs one key and one value for every batch lane, shared K/V head and feature. Define $U=2BH_{kv}d$, the combined K/V scalar slots per token position; the factor $2$ counts the two arrays.

For current length $0\le n\le C$, let $A=\lceil n/P\rceil$ be the number of logical pages containing at least one live row. At $n=0$, $A=0$. Write $S_{\mathrm{alloc}}$ for allocated scalar slots, $S_{\mathrm{live}}$ for filled slots, $S_{\mathrm{tail}}$ for unfilled slots in the last active page, and $S_{\mathrm{whole}}$ for slots in wholly inactive pages. Then

$S_{\mathrm{unused}}$ counts every allocated slot that does not currently hold a live logical row. The tail formula gives zero when the live prefix ends exactly at a page boundary or is empty. Every wholly inactive page remains allocated in this eager pool.

For the main example, $U=2\times1\times2\times2=8$, $N=4$, $n=7$ and $A=3$. The observed counts combine both arrays for one layer:

• Allocated: $96$ scalar slots, or $768$ bytes.
• Live: $56$ slots, or $448$ bytes.
• Unused: $40$ slots, or $320$ bytes.
• Of those unused slots, the partial-page tail has $16$ slots, or $128$ bytes, for unfilled positions $7,8$ in physical page $3$.
• The inactive whole page has $24$ slots, or $192$ bytes, for positions $9,10,11$ in physical page $1$.

Each f64 slot occupies $8$ bytes. The two layers together allocate $1536$ bytes in these K/V arrays. These are buffer-length payload counts. They exclude page-table and object metadata, allocator overhead or excess capacity, model weights, activations and temporary tensors. They measure neither process RAM nor GPU memory, traffic or speed.

    /// Counts K and V f64 buffer elements, excluding metadata and other allocations.
    pub fn payload(
        &self,
        length: usize,
        batch: usize,
        kv_heads: usize,
        width: usize,
    ) -> Result<PagedKvPayload, PagedKvError> {
        if length > self.capacity {
            return Err(PagedKvError::Length {
                length,
                capacity: self.capacity,
            });
        }
        let per_position = 2_usize
            .checked_mul(batch)
            .and_then(|n| n.checked_mul(kv_heads))
            .and_then(|n| n.checked_mul(width))
            .ok_or(PagedKvError::Overflow)?;
        let active_pages = length.div_ceil(self.page_tokens);
        let active_positions = active_pages
            .checked_mul(self.page_tokens)
            .ok_or(PagedKvError::Overflow)?;
        let count = |positions: usize| {
            per_position
                .checked_mul(positions)
                .ok_or(PagedKvError::Overflow)
        };
        let allocated_scalar_slots = count(self.physical_tokens)?;
        let live_scalar_slots = count(length)?;
        let active_page_tail_scalar_slots = count(active_positions - length)?;
        let inactive_page_scalar_slots = count(self.physical_tokens - active_positions)?;
        let bytes = |slots: usize| {
            slots
                .checked_mul(std::mem::size_of::<f64>())
                .ok_or(PagedKvError::Overflow)
        };
        Ok(PagedKvPayload {
            reserved_pages: self.page_table.len(),
            active_pages,
            allocated_scalar_slots,
            live_scalar_slots,
            unused_scalar_slots: allocated_scalar_slots - live_scalar_slots,
            active_page_tail_scalar_slots,
            inactive_page_scalar_slots,
            allocated_payload_bytes: bytes(allocated_scalar_slots)?,
            live_payload_bytes: bytes(live_scalar_slots)?,
            unused_payload_bytes: bytes(allocated_scalar_slots - live_scalar_slots)?,
            active_page_tail_payload_bytes: bytes(active_page_tail_scalar_slots)?,
            inactive_page_payload_bytes: bytes(inactive_page_scalar_slots)?,
        })
    }

A rounded page can reserve more than contiguous storage

The report also constructs a separate, empty layer cache with $C=10$ and $P=3$, keeping the same batch count, shared heads and feature width. Four pages reserve room for $12$ token positions. The paged arrays contain $96$ scalar slots, while contiguous arrays with capacity $10$ contain $80$. The extra $16$ slots occupy $128$ bytes. Logical positions $10,11$ remain inaccessible even though their rounded storage exists.

This comparison has length $n=0$, so all of its allocated slots are unused. It does not have the main example’s seven live rows. At the main capacity $12$, which divides evenly into pages, paged and contiguous storage instead reserve the same $96$ slots. An unused reserved page is not an unallocated page; paging by itself is no guarantee of a smaller payload.

From maximum-length reservation to block lookup

Kwon and colleagues’ PagedAttention paper, version 1 dated 12 September 2023, describes the earlier inference design in section 3.1. Earlier LLM inference systems reserved contiguous KV storage for a sequence’s maximum length, leaving reserved capacity unused when its live prefix was shorter. That reservation simplified addressing but tied physical storage to a chosen maximum rather than the currently filled prefix.

Sections 4.1–4.3 of the same version explain the block lookup. PagedAttention separates fixed-token logical blocks from physical blocks and uses a block table to locate keys and values during attention. Its complete system also supplies allocation, sharing and scheduling behavior; the paper’s system gains cannot be inferred from an address calculation alone.

Cache addressing can change without changing token order, rotary positions or eligible context. Retain the existing contiguous cache as the runnable reference and add page lookup to the same layer cache and decoder generation path. This chapter reserves a fixed host pool eagerly. It demonstrates addressing and numerical-payload accounting, without dynamic allocation, sharing, scheduling or measured system gains.

Compare layouts inside the same restored model

LayerKvCache::new_paged reserves the selected layout’s key and value buffers up front. DecoderKvCache::new_paged uses the same checked table in each existing layer cache. It requires the layout’s logical capacity to match the model’s position capacity. Existing parameter and RoPE bindings still prevent using cached rows with an incompatible model.

    /// Select page placement without changing the model or cached decoder computation.
    pub fn new_paged(
        model: &DecoderModel,
        layout: &PagedKvLayout,
    ) -> Result<Self, DecoderKvCacheError> {
        Self::new_with_layout(model, Some(layout))
    }

    fn new_with_layout(
        model: &DecoderModel,
        layout: Option<&PagedKvLayout>,
    ) -> Result<Self, DecoderKvCacheError> {
        let config = model.config();
        validate_paged_capacity(config.max_positions(), layout)?;
        let mut parameter_bindings = Vec::new();
        parameter_bindings
            .try_reserve_exact(model.parameters().len())
            .map_err(|_| DecoderKvCacheError::ParameterAllocationFailed {
                parameters: model.parameters().len(),
            })?;
        parameter_bindings.extend(
            model
                .parameters()
                .iter()
                .map(|parameter| TensorValueBinding::capture(parameter.tensor())),
        );

        let mut layers = Vec::new();
        layers
            .try_reserve_exact(model.blocks().len())
            .map_err(|_| DecoderKvCacheError::LayerAllocationFailed {
                layers: model.blocks().len(),
            })?;
        for (layer, block) in model.blocks().iter().enumerate() {
            layers.push(
                match layout {
                    Some(layout) => LayerKvCache::new_paged(block.attention(), 1, layout.clone()),
                    None => LayerKvCache::new(block.attention(), 1, config.max_positions()),
                }
                .map_err(|source| DecoderKvCacheError::LayerCache { layer, source })?,
            );
        }
        Ok(Self {
            config,
            parameter_bindings,
            layers,
            len: 0,
            prefill_complete: false,
            work: DecoderKvCacheWork::default(),
        })
    }

The cache still checks shapes, capacity and finite candidate values before appending. Size and byte-count overflow and allocation failures return errors. A decoder step prepares its layer rows and completes the output before committing them. A failure therefore cannot leave only some layers advanced across a page boundary. Reset removes the logical prefix while retaining the same buffers, their bytes and the page table; old bytes do not become live again until valid new rows are written.

The default demo constructs the existing text model with width $8$, four query heads, two K/V heads, feed-forward width $12$, two layers, position capacity $12$ and initialization seed $20$. It has no local attention window. The demo saves and restores that model through PortableModel, then gives the same restored weights to the contiguous and paged caches.

For each of the seven framed prefix lengths, it compares every paged vocabulary logit with two references: the contiguous cache and the current last row of a full-prefix forward pass. There are $258$ vocabulary coordinates in each row, so the report counts $7\times258=1806$ comparisons against each reference. Earlier full-prefix rows are not included in that count.

/// Both cache layouts read the very same restored model, in logical token order.
pub fn demonstrate(config: &Config) -> ArchiveResult<Value> {
    let (model, tokenizer) =
        practical_ch18_grouped_query_attention_context::configured_model(config)?;
    let saved = PortableModel::from_text(&model, &tokenizer)?;
    let loaded = PortableModel::from_bytes(&saved.to_bytes()?)?;
    let (restored, tokenizer) = loaded.restore_text()?;
    let cfg = restored.config();
    let layout = selected_layout(config, cfg.max_positions())?;
    let document = config
        .get("text.document")
        .ok_or("missing document")?
        .as_string()?;
    let tokens = tokenizer.encode_utf8_document(document);
    let mut contiguous = DecoderKvCache::new(&restored)?;
    let mut paged = DecoderKvCache::new_paged(&restored, &layout)?;
    let mut compared = 0;
    let mut same_cached_bits = true;
    let mut maximum_cached_error = 0.0_f64;
    let mut maximum_full_error = 0.0_f64;
    let mut graph_free = true;
    {
        let mut a = contiguous.bind(&restored)?;
        let mut b = paged.bind(&restored)?;
        for count in 1..=tokens.len() {
            let ca = if count == 1 {
                a.prefill(&tokens[..1])
            } else {
                a.decode(tokens[count - 1])
            }?;
            let pa = if count == 1 {
                b.prefill(&tokens[..1])
            } else {
                b.decode(tokens[count - 1])
            }?;
            let full = restored.forward_with_context(
                AutogradContext::no_grad(),
                &tokens[..count],
                &[1, count],
            )?;
            let c = ca.logits().value();
            let p = pa.logits().value();
            let f = full.logits().value();
            let last =
                &f.as_slice()[(count - 1) * cfg.vocabulary_size()..count * cfg.vocabulary_size()];
            if c.len() != p.len() || p.len() != last.len() {
                return Err("logit comparison shape mismatch".into());
            }
            for ((&contiguous, &paged), &full) in c.as_slice().iter().zip(p.as_slice()).zip(last) {
                if !contiguous.is_finite() || !paged.is_finite() || !full.is_finite() {
                    return Err("nonfinite compared logit".into());
                }
                compared += 1;
                same_cached_bits &= contiguous.to_bits() == paged.to_bits();
                maximum_cached_error = maximum_cached_error.max((contiguous - paged).abs());
                maximum_full_error = maximum_full_error.max((full - paged).abs());
            }
            graph_free &= !ca.logits().tracks_gradient()
                && !pa.logits().tracks_gradient()
                && !full.logits().tracks_gradient();
        }
    }
    let cache = paged
        .layer_cache(0)
        .ok_or("worked layout requires at least one model layer")?;
    let layer_layout = cache.paged_layout().ok_or("missing selected page layout")?;

Both observed maximum absolute errors are $0$ in this CPU f64 fixture, within the declared absolute tolerance $10^{-10}$. Contiguous and paged cached logits also match bit for bit. All compared outputs are graph-free, and the embedding and text head still share one tied table. The model ID remains 97d068dc931cc123096433f5f5dd0bf45bce815fe5163127155d251f32d5e5bc. Page placement is execution/storage state, not a change to saved logical model weights or identity. This model has undergone no training updates. These comparisons establish neither language quality nor universal or cross-device bitwise equality.

Reuse the existing generation loop

generate_cached and generate_cached_with_layout enter the same cached generation loop; the latter selects an optional page layout. Both use the existing sampler. In the default greedy fixture, prompt IDs [0,99] produce generated IDs [70,56,216,6]. The layouts give the same tokens, per-step records, stop reason and work counts, and their next random-generator values agree.

    let prompt_count = integer(config, "generation.prompt_tokens")?;
    if prompt_count == 0 || prompt_count > tokens.len() {
        return Err("generation prompt count must select a nonempty fixture prefix".into());
    }
    let generation = GenerationConfig::new(
        SamplingMode::Greedy,
        None,
        integer(config, "generation.max_new_tokens")?,
    );
    let mut a = SplitMix64::from_seed(u64::try_from(integer(config, "generation.seed")?)?);
    let mut b = a.clone();
    let contiguous_generation =
        generate_cached(&restored, &tokens[..prompt_count], generation, &mut a)?;
    let paged_generation = generate_cached_with_layout(
        &restored,
        &tokens[..prompt_count],
        generation,
        &mut b,
        Some(&layout),
    )?;
    let generation = json!({"mode":"greedy","prompt_token_ids":tokens[..prompt_count],"generated_token_ids":paged_generation.generated(),
        "same_tokens_steps_stop_work":contiguous_generation==paged_generation,"same_next_rng_value":a.next_u64()==b.next_u64(),
        "final_cache_len":paged_generation.final_cache_len(),"prefill_tokens":paged_generation.work().prefill_tokens(),
        "decode_tokens":paged_generation.work().decode_tokens()});

The prompt contributes $2$ prefill inputs. Only the first $3$ selected tokens need a decode input: after selecting the fourth, generation has reached its token limit and stops. The generation cache therefore ends at length $2+3=5$, not at the prompt length plus every selected token. This is a separate generation run from the seven-position forced-prefix comparison. The untrained token IDs carry no text-quality claim.

The demo uses the existing config-lib and clap paths. In the included native configuration, page_table = 2 0 3 1 supplies the four page indices. An alternate table changes physical placement, while the tested model and generation results remain the same. The text wrapper retains its tokenizer; the shared backbone continues to consume encoded vectors and general layout without requiring text-tokenization logic.

Follow logical order through the physical pool

The logical-page groups below retain token reading order. The physical-page groups show where those same rows reside, including pages and tails that were reserved but have no live rows. Every scalar start belongs to batch $0$, shared head $0$ and the first feature in each separate K/V array.

Reproduce the report and inspect the invariants

Run from the repository root with Docker and the repository’s existing course Docker setup available. The configuration, tokenizer and text fixtures are included; no external corpus or model is needed. This CPU command prints the mapping, per-layer payloads, model comparisons and generation result as JSON:

COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch20-paged-kv-layout

The report includes mappings, physical_pages, per_layer_payloads, model_comparison, generation and rounding_comparison_empty_layer. Both comparison errors are zero, same_cached_bits and within_tolerance are true, and the generated IDs are [70,56,216,6].

With the same working directory and Docker setup, run the demo checks:

COURSE_CORPUS=false ./course run cargo test --offline --locked -p practical-ch20-paged-kv-layout --lib

Exactly $2$ tests pass: the default report checks and an alternate native-config page order. To run the cache-layout checks, use

COURSE_CORPUS=false ./course run cargo test --offline --locked -p llm-from-scratch-practical attention::paged::tests --lib

Exactly $10$ tests pass. They cover page-table validation and overflow, payload conservation and rounding, batched MHA/GQA head sharing with full and windowed context, logical snapshots and reset, model logits, existing greedy and temperature/top-k generation, binding and capacity errors, and failed-step rollback. They include a failure after all layer rows have been prepared, which must leave every physical cache unchanged.

Optional inspection and explanation practice

1. Use mappings from the default command above to trace positions $2$, $3$ and $6$. Identify each logical page, physical page, page offset and first-feature scalar index, then explain which order attention must read.
2. Inspect per_layer_payloads and physical_pages in that report. Reconcile allocated, live and unused K/V slots, then split the unused portion into a partial-page tail and an inactive whole page. State their byte counts.
3. Inspect model_comparison and generation. Explain the logit count per reference, what the zero errors establish, and why selecting four tokens from a two-token prompt leaves a generation cache of length five.
4. Inspect rounding_comparison_empty_layer and run the two test commands above. Explain the extra paged reservation, identify its live length, and relate the reset and failed-step checks to the difference between reserved storage and valid cached state.

Keep storage placement separate from token choice

The same text encoder, source-independent backbone, tied head and cached generator now support explicit page placement without changing logical context or saved model identity. The cache translates addresses while RoPE, eligible keys, shared heads and the sampler retain their existing roles.

Planned Chapter 21, Sampling distributions, will extend how the existing sampler turns vocabulary scores into token choices. Cache placement and sampling policy are separate decisions: this chapter changes where a row is stored, while the next lesson studies how a token is selected from the model’s scores.