Implementation:Ollama Ollama Llama KV Cache Header

Knowledge Sources	Ollama
Domains	LLM Inference, Memory Management
Last Updated	2025-02-15 00:00 GMT

Overview

Header declaring the llama_kv_cache class and its supporting types for managing key-value attention state during transformer inference.

Description

Declares llama_kv_cache inheriting from llama_memory_i, with nested types: slot_info (maps tokens to cache cell indices with stream support), stream_copy_info (for copying between streams), kv_layer (per-layer key/value tensor pairs), and cell_ranges_t (ranges of cell indices). The llama_kv_cache_context manages per-batch cache context including slot allocation, attention mask generation, and graph input preparation.

Usage

Include this header when working with the KV cache, building attention compute graphs, or implementing new attention-based architectures.

Code Reference

Source Location

• Repository: Ollama
• File: llama/llama.cpp/src/llama-kv-cache.h
• Lines: 1-390

Signature

class llama_kv_cache : public llama_memory_i {
public:
    struct slot_info {
        uint32_t s0, s1;
        std::vector<llama_seq_id> strm;
        std::vector<idx_vec_t>    idxs;
        uint32_t head() const;
        bool is_contiguous() const;
    };

    using slot_info_vec_t = std::vector<slot_info>;

    // llama_memory_i interface
    llama_memory_context_ptr init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) override;
    slot_info_vec_t prepare(const std::vector<llama_ubatch> & ubatches);
    slot_info find_slot(const llama_ubatch & ubatch, bool cont) const;

    // graph_build API
    ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
    ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
    ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
    ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
};

Import

#include "llama-kv-cache.h"

I/O Contract

Inputs

Name	Type	Required	Description
model	const llama_model &	Yes	Model for layer/tensor configuration
type_k	ggml_type	Yes	Key tensor data type
type_v	ggml_type	Yes	Value tensor data type
kv_size	uint32_t	Yes	Number of KV cache cells

Outputs

Name	Type	Description
slot_info	struct	Token-to-cell index mapping for ubatches
ggml_tensor*	ggml_tensor*	Key and value tensor views for compute graph

Usage Examples

#include "llama-kv-cache.h"

// Access KV cache through the context:
const llama_kv_cache_context * kv_ctx = ...;
ggml_tensor * k = kv_ctx->get_k(ctx, il);
ggml_tensor * v = kv_ctx->get_v(ctx, il);

Related Pages

• Principle:Ollama_Ollama_LLM_Memory_Architecture