Implementation:Ollama Ollama KVCache Causal

Knowledge Sources	Ollama
Domains	Inference Runtime, KV Cache Management
Last Updated	2025-02-15 00:00 GMT

Overview

Implements the causal (autoregressive) KV cache that stores key and value tensors indexed by sequence position, supporting sliding window attention, chunked attention, prefix caching, and cache defragmentation.

Description

The Causal struct implements the Cache interface for decoder-style autoregressive models. It manages a cell-based storage system where each cell tracks a sequence ID, position, and whether it is in use. The cache supports sliding window attention (with separate window and memory sizes), chunked attention (fixed-size context windows), and prefix caching (reusing previously computed KV pairs for shared prompt prefixes). The StartForward method finds available cache slots, builds attention masks, and handles cache eviction. Get and Put retrieve and store K/V tensors per layer. The cache includes a defragmentation mechanism that compacts fragmented entries and a Remove method for evicting sequences.

Usage

Used as the primary KV cache for all causal (decoder-only) language models during inference. Created via NewCausal or NewSWA (for sliding window) and initialized with backend, dtype, and capacity parameters.

Code Reference

Source Location

• Repository: Ollama
• File: kvcache/causal.go
• Lines: 1-666

Signature

type Causal struct {
    DType         ml.DType
    swaWindowSize int32
    swaMemorySize int32
    chunkSize     int32
    opts          CausalOptions
    // ...
}

func NewCausal(opts ...CausalOption) *Causal
func NewSWA(windowSize, memorySize int32, opts ...CausalOption) *Causal
func (c *Causal) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int)
func (c *Causal) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error
func (c *Causal) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor)
func (c *Causal) Put(ctx ml.Context, key, value ml.Tensor)
func (c *Causal) SetLayer(layer int)
func (c *Causal) Remove(seqId int, beginIndex, endIndex int32) error
func (c *Causal) Close()

Import

import "github.com/ollama/ollama/kvcache"

I/O Contract

Inputs

Name	Type	Required	Description
backend	ml.Backend	Yes	ML backend used for tensor allocation and graph operations
dtype	ml.DType	Yes	Data type for cache storage (e.g., f16, f32)
maxSequences	int	Yes	Maximum number of concurrent sequences
capacity	int	Yes	Number of cache entries per sequence
batch	input.Batch	Yes	Current batch with positions, sequences, and outputs

Outputs

Name	Type	Description
key	ml.Tensor	Historical key tensor for the current layer
value	ml.Tensor	Historical value tensor for the current layer
mask	ml.Tensor	Attention mask tensor (history_size x batch_size)

Usage Examples

// Create a standard causal cache
cache := kvcache.NewCausal()
cache.Init(backend, ml.DTypeF16, 4, 2048, 512)

// Create a sliding window cache
swaCache := kvcache.NewSWA(4096, 8192)
swaCache.Init(backend, ml.DTypeF16, 4, 2048, 512)

// During inference
cache.StartForward(ctx, batch, false)
cache.SetLayer(0)
key, value, mask := cache.Get(ctx)
// ... compute attention ...
cache.Put(ctx, newKey, newValue)

Related Pages

• Principle:Ollama_Ollama_KVCache_Causal_Attention