KV Cache¶

The KV cache is the single most impactful optimisation in autoregressive transformer inference. By storing and reusing the Key and Value projections from previous tokens, it converts per-token attention from a quadratic operation to a linear one -- enabling practical generation speeds for sequences of hundreds or thousands of tokens.

1. The Problem¶

During autoregressive generation, we generate tokens one at a time. At step \( t \), the naive approach recomputes the K and V projections for all \( t \) tokens, even though the projections for tokens \( 1, \ldots, t-1 \) have not changed. Over \( T \) generation steps, the total cost is:

For a model with \( L \) layers, \( H \) heads, and head dimension \( d_h \):

2. Key Insight¶

At generation step \( t \), we only need to:

1. Compute \( \mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t \) for the new token.
2. Append \( \mathbf{k}_t, \mathbf{v}_t \) to the cache.
3. Compute attention between \( \mathbf{q}_t \) and all cached \( \mathbf{K}_{1:t}, \mathbf{V}_{1:t} \).

The per-step cost drops from \( O(t \cdot d_h) \) (recomputing all K, V) to \( O(d_h) \) for the projection plus \( O(t \cdot d_h) \) for the attention itself. Critically, the K/V projection is no longer repeated for past tokens.

3. Cache Architecture¶

ZigLlama organises the KV cache as a four-level hierarchy:

graph TD
    MSC["MultiSequenceKVCache<br/>(manages multiple sequences)"]
    MSC --> S1["Sequence 1<br/>[]KVCacheEntry"]
    MSC --> S2["Sequence 2<br/>[]KVCacheEntry"]
    MSC --> SN["Sequence N<br/>..."]

    S1 --> L0["Layer 0<br/>KVCacheEntry"]
    S1 --> L1["Layer 1<br/>KVCacheEntry"]
    S1 --> LM["Layer L<br/>..."]

    L0 --> K0["keys: Tensor(f32)<br/>[seq_len, n_heads, head_dim]"]
    L0 --> V0["values: Tensor(f32)<br/>[seq_len, n_heads, head_dim]"]

3.1 KVCacheEntry¶

The fundamental unit, storing keys and values for a single layer of a single sequence:

pub const KVCacheEntry = struct {
    keys: Tensor(f32),       // [max_length, n_heads, head_dim]
    values: Tensor(f32),     // [max_length, n_heads, head_dim]
    sequence_length: usize,  // Current number of cached positions
    max_length: usize,       // Pre-allocated capacity
    layer_id: usize,         // Which transformer layer this belongs to

    pub fn init(allocator: Allocator, max_length: usize,
                n_heads: usize, head_dim: usize, layer_id: usize) !KVCacheEntry { ... }
    pub fn append(self: *KVCacheEntry, keys: *const Tensor(f32),
                  values: *const Tensor(f32)) !void { ... }
    pub fn get(self: *const KVCacheEntry, start_pos: usize,
               length: usize) !struct { keys: Tensor(f32), values: Tensor(f32) } { ... }
    pub fn clear(self: *KVCacheEntry) void { ... }
    pub fn deinit(self: *KVCacheEntry) void { ... }
};

3.2 MultiSequenceKVCache (ModelKVCache)¶

Manages caches for multiple concurrent sequences (e.g., in a batch or server context), keyed by a u64 sequence ID:

pub const MultiSequenceKVCache = struct {
    sequences: std.AutoHashMap(u64, []KVCacheEntry),
    n_layers: usize,
    n_heads: usize,
    head_dim: usize,
    max_seq_length: usize,
    allocator: Allocator,

    pub fn getOrCreateSequence(self: *MultiSequenceKVCache, sequence_id: u64) ![]KVCacheEntry { ... }
    pub fn appendToSequence(self: *MultiSequenceKVCache, sequence_id: u64,
                            layer_id: usize, keys: *const Tensor(f32),
                            values: *const Tensor(f32)) !void { ... }
    pub fn getFromSequence(self: *MultiSequenceKVCache, sequence_id: u64,
                           layer_id: usize, start_pos: usize, length: usize) !?... { ... }
    pub fn getMemoryUsage(self: *MultiSequenceKVCache) u64 { ... }
};

4. Update Algorithm¶

5. Cache Strategies¶

ZigLlama supports multiple caching strategies to trade memory against computation:

6. Sliding Window Cache¶

For very long sequences, the full cache can exceed available memory. The SlidingWindowKVCache implements a fixed-size window that evicts the oldest tokens when the cache is full:

pub const SlidingWindowKVCache = struct {
    cache: MultiSequenceKVCache,
    window_size: usize,

    pub fn appendWithSliding(self: *SlidingWindowKVCache, sequence_id: u64,
                             layer_id: usize, keys: *const Tensor(f32),
                             values: *const Tensor(f32)) !void { ... }

    pub fn getWithWindow(self: *SlidingWindowKVCache, sequence_id: u64,
                         layer_id: usize, query_pos: usize) !?... { ... }
};

flowchart LR
    subgraph "Before (window=8, length=8)"
        B["t0 t1 t2 t3 t4 t5 t6 t7"]
    end
    subgraph "After appending t8 (evict t0,t1)"
        A["t2 t3 t4 t5 t6 t7 t8 __"]
    end
    B --> A

7. Memory Budget¶

7.1 Concrete Examples¶

ZigLlama provides a utility function for cache size estimation:

pub fn estimateMemoryUsage(
    n_sequences: u64, max_seq_length: u64,
    n_layers: u64, n_heads: u64, head_dim: u64,
) u64 {
    const per_token_memory = n_layers * n_heads * head_dim * 2 * @sizeOf(f32);
    return n_sequences * max_seq_length * per_token_memory;
}

8. Performance Impact¶

For a sequence of length \( T = 1000 \), the KV cache provides approximately 1000x reduction in redundant computation. In practice the speedup is ~100x because other operations (FFN, embedding lookup, etc.) are not affected by caching.

xychart-beta
    title "Per-Token Latency vs Sequence Length"
    x-axis "Sequence Length" [64, 128, 256, 512, 1024, 2048]
    y-axis "Relative Latency"
    line "Without Cache" [1, 2, 4, 8, 16, 32]
    line "With Cache" [1, 1, 1, 1, 1, 1]

9. Cache Growth During Generation¶

The following diagram illustrates how the KV cache grows as tokens are generated, and how the sliding window variant manages this growth:

flowchart TD
    subgraph "Standard Cache Growth"
        direction LR
        S1["Step 1<br/>[P P P]<br/>3 cached"] --> S2["Step 2<br/>[P P P G1]<br/>4 cached"]
        S2 --> S3["Step 3<br/>[P P P G1 G2]<br/>5 cached"]
        S3 --> S4["...<br/>[P P P G1 ... Gn]<br/>n+3 cached"]
    end

    subgraph "Sliding Window (size=4)"
        direction LR
        W1["Step 1<br/>[P P P]<br/>3/4"] --> W2["Step 2<br/>[P P P G1]<br/>4/4 full"]
        W2 --> W3["Step 3<br/>[P P G1 G2]<br/>4/4 evicted P"]
        W3 --> W4["Step 4<br/>[P G1 G2 G3]<br/>4/4 evicted P"]
    end

Legend: P = prompt token, G = generated token.

10. Cache Statistics¶

The KVCacheStats struct provides monitoring data for cache utilisation:

pub const KVCacheStats = struct {
    total_sequences: u32,
    total_memory_bytes: u64,
    average_sequence_length: f32,
    cache_hit_rate: f32,
    cache_efficiency: f32,  // used / allocated ratio
};

References¶