Mutable execution cache#

Date:

2026-08

discussion

Objective#

A KV-cache may occupy most of the available memory. Updating it must not allocate an output cache of the same size. The runtime needs a guaranteed in-place update, not the current best-effort buffer-reuse heuristic.

Value model#

The cache is a mutable graph input owned by the caller. The graph keeps SSA names, but two names may identify the same storage:

updated_cache = CacheUpdate(cache, values, position)

updated_cache must alias cache. The operator changes the valid cache content but does not allocate another full buffer.

The cache remains an ordinary tensor or sequence. Its capacity is part of its shape, while its valid length and update position are explicit scalar inputs. Mutability and storage aliasing are execution contracts, not new value types.

ShapesContext#

No SymCache class is needed. ShapesContext handles the cache with its existing SymTensor or sequence representation and infers that updated_cache has the same type and shape as cache. When the relevant dimensions are known, it also checks position + update_length <= capacity.

ValueInfoProto.access records whether the graph input is mutable. NodeProto.output_alias records the relation between operator ports, while the output’s ValueInfoProto.alias_of records the same relation by value name. The execution plan and runtime track the concrete storage identity and verify pointer equality. Subgraphs and local functions propagate the same type and alias annotations.

Shape inference maintains alias equivalence classes over value names. For each alias it:

  1. resolves the node port indices to input and output names;

  2. verifies that ValueInfoProto.alias_of names the same input;

  3. merges the input and output type and shape constraints;

  4. rejects incompatible element types, ranks, dimensions, or alias cycles;

  5. records the common alias root for memory planning.

Type and shape information learned for either name is therefore visible through the other. Control-flow branches may merge an aliased output only when they resolve it to the same alias root. ShapesContext stores names and constraints only; concrete pointer identity remains a runtime check.

Required alias#

The execution plan needs a mandatory alias distinct from opportunistic inplace_reuse:

updated_cache.alias_of = "cache"

alias_of means:

  • no output allocation is planned;

  • the kernel receives a writable view of the input storage;

  • pointer identity is checked by the runtime;

  • execution fails if the caller supplied immutable or shared storage;

  • there is no silent copy fallback.

The existing last-use rule does not apply: the cache is persistent state and remains live after the node. Safety comes from explicit mutable ownership, not from the input becoming dead.

Proto additions#

The port-level and value-level contracts are serialized as follows:

message ValueAliasProto {
    enum Kind {
        UNDEFINED = 0;
        MUST_ALIAS = 1;
    }
    int32 output_index = 1;
    int32 input_index = 2;
    Kind kind = 3;
}

message NodeProto {
    ...
    repeated ValueAliasProto output_alias = <N>;
}

message ValueInfoProto {
    enum Access {
        READ_ONLY = 0;
        READ_WRITE = 1;
    }
    ...
    Access access = <N>;
    string alias_of = <N+1>;  // empty when the value owns distinct storage
}

output_alias identifies the input and output ports of the producing node. alias_of names the value whose storage must be shared and exposes the relation at graph or function boundaries. When both annotations describe the same output, a checker requires them to identify the same input. READ_WRITE is valid for a graph input and tells the caller that execution may modify its storage.

No CacheProto is needed. Dense caches remain TensorProto values and quantized pages are QuantizedTensorProto values. Capacity and valid length are described by shapes or explicit scalar inputs.

Alternative alias designs#

Design

Advantage

Limitation

ValueInfoProto.alias_of

Relates named values directly and works at graph boundaries.

Requires ValueInfoProto for every aliased output.

NodeProto.output_alias

Describes the operator contract using stable port indices.

Does not alone expose the alias at a graph boundary.

Operator-schema alias rule

Declared once for CacheUpdate.

Requires the runtime schema; less suitable for dynamic rules.

FunctionProto parameter alias

Preserves aliases through local functions.

Does not describe primitive custom nodes.

Mutable reference type

Naturally represents persistent state.

Introduces reference semantics into ONNX types.

Graph state slot

Hides cache threading between graph calls.

Adds state, concurrency, and checkpoint semantics.

Metadata property

Requires no proto change and is easy to prototype.

Weak contract that generic tools may ignore.

The recommended design keeps NodeProto.output_alias for the local operator contract and ValueInfoProto.alias_of for the named graph contract. Metadata is suitable only for experimentation.

Model example#

The model exposes the cache as a mutable input and returns a different SSA name backed by the same storage:

graph input:
    cache: FLOAT16[2, batch, heads, max_length, head_size]
        access: READ_WRITE
    values: FLOAT16[2, batch, heads, new_length, head_size]
    position: INT64[]

node:
    op_type: "CacheUpdate"
    input: ["cache", "values", "position"]
    output: ["updated_cache"]
    output_alias: {
        output_index: 0
        input_index: 0
        kind: MUST_ALIAS
    }

graph output:
    updated_cache: FLOAT16[2, batch, heads, max_length, head_size]
        alias_of: "cache"

Before execution, the caller allocates cache for max_length and binds it as writable storage. The runtime binds updated_cache to the same address, and CacheUpdate writes only the range starting at position. On the next invocation, that same storage is bound again as cache.

position + new_length is checked against max_length. The execution plan allocates only ordinary outputs and kernel workspace; it never allocates a second cache.

Heterogeneous paged cache#

A paged cache is a sequence or map of QuantizedTensorProto pages. The container accepts several quantization types, while every page selects one exact type:

cache: Sequence<QuantizedTensor {
    allowed_quantized_type: []  // any registered quantization
    elem_type: FLOAT16
    shape: [2, heads, page_size, head_size]
}>
    access: READ_WRITE

pages: [
    QuantizedTensorProto {
        quantized_type: INT4_PAGE
        dims: [2, heads, page_size, head_size]
        raw_data: ...
    },
    QuantizedTensorProto {
        quantized_type: FP8_PAGE
        dims: [2, heads, page_size, head_size]
        raw_data: ...
    },
    QuantizedTensorProto {
        quantized_type: INT2_PAGE
        dims: [2, heads, page_size, head_size]
        raw_data: ...
    }
]

All page types must decode to the same logical element type and page shape, but their physical layouts and byte sizes may differ.

The update node aliases the page table:

updated_cache = PagedCacheUpdate(cache, page_id, values, position)
updated_cache.alias_of = "cache"

ShapesContext keeps only the common logical page type and shape. The runtime owns the page storage identities. The kernel dispatches from quantized_type and quantizes the new values directly into the selected page.

Updating an existing page allocates nothing. Creating a page allocates only that page. Changing a page to a quantization type requiring a different byte size replaces only that page and updates the aliased page table; it never copies the complete cache.

ModelProto.quantizations contains the referenced quantization declarations. Different physical page formats remain valid because they all decode to the container’s common logical element type and shape.

Runtime contract#

The caller opts into mutation when binding the cache. One writable cache binding cannot be used concurrently by multiple sessions. Read-only model initializers are never accepted as mutable caches.

The kernel writes only inside the declared capacity. An overflow is an error; the runtime must not grow the cache by allocating and copying a full replacement.

For a paged cache, only a missing page may be allocated. Existing pages are updated in place, and page-table growth must not duplicate page payloads. Each quantized page keeps its own QuantizedTensorProto.quantized_type.

Existing implementations#

ONNX Runtime GenAI#

ONNX Runtime GenAI exposes the cache as past_key_values.* inputs and present.* outputs. In the basic mode, each present becomes the next iteration’s past.

Buffer sharing is requested in genai_config.json:

{
  "search": {
    "past_present_share_buffer": true,
    "num_beams": 1,
    "max_length": 4096
  }
}

GenAI then allocates one capacity-sized OrtValue per K/V layer and binds the same object to both names:

past_key_values.0.key ──┐
                        ├── OrtValue B
present.0.key ──────────┘

The compatible attention operator writes new values at the current position. DefaultKeyValueCache::Update returns immediately because neither the pointer nor the allocation changes. The ONNX model itself contains no alias annotation; GenAI creates the alias while binding session inputs and outputs.

The option is enabled only when requested and when num_beams == 1, except for the specialized Whisper case. Graph capture requires shared buffers.

It also supports windowed caches, model-managed state, and paged-attention metadata such as block_table and past_sequence_lengths. Quantized caches use a cache-wide or layer-wide representation rather than a different type for every page.

Its PagedKeyValueCache is logically, not physically, paged. For each layer it preallocates two contiguous tensors:

key_cache:   [num_blocks, block_size, num_kv_heads, head_size]
value_cache: [num_blocks, block_size, num_kv_heads, head_size]

The block_table maps each request’s logical pages to slices on the num_blocks axis. Blocks can be assigned and released independently, but the backing tensors never grow after allocation. Once the free block pool is empty, no additional cache capacity is available without creating a larger cache and copying the old content.

All blocks of one cache tensor have the same element type, layout, and physical size. Quantization parameters may differ by layer, but not by page. Independent page allocations and heterogeneous page quantization are not supported. With one active request, continuous batching provides no batching benefit; paging still avoids contiguous cache growth, but it does not reduce the preallocated physical pool.

See kv_cache.cpp, kv_cache.h, paged_key_value_cache.h, the activation rule, and the model builder.

llama.cpp#

llama.cpp keeps the KV-cache in runtime-owned llama_kv_cache state rather than model inputs and outputs. It allocates K/V storage per layer, gives the compute graph views over the destination ranges, and writes new values directly through cpy_k and cpy_v. Sequence operations primarily update cell metadata; defragmentation moves only required ranges.

The capacity kv_size is fixed when the context is created. Each layer owns K/V tensors covering that capacity, while llama_kv_cells maps tokens and sequences to slots. A sequence may therefore occupy non-contiguous cells, but the underlying tensors do not grow. find_slot reuses free cells; context shifting, eviction, and defragmentation recover space. If no suitable slot is available, cache preparation fails rather than extending the allocation.

type_k and type_v are selected for the cache as a whole. They may differ from each other, but not from cell to cell. Specialized block caches exist for some architectures, but there is no generic page table with independently allocated or heterogeneously quantized pages.

See llama-kv-cache.h and llama-kv-cache.cpp.

Design consequence#

output_alias and alias_of serialize the shared-buffer optimization used by ONNX Runtime GenAI instead of leaving it to runtime configuration. A mutable state handle would be closer to llama.cpp but would move cache semantics outside the graph. Sequence<QuantizedTensorProto> additionally permits each page to select a different quantized_type.

Memory planning#

Peak-memory analysis counts an in-place cache update as zero additional cache bytes. It includes only temporary kernel workspace and newly allocated pages. The execution plan keeps the cache alive across graph invocations and excludes it from ordinary release and reuse candidates.

Implementation order#

  1. Add NodeProto.output_alias, ValueInfoProto.alias_of, and execution-plan alias metadata.

  2. Add mutable caller bindings and writable kernel views.

  3. Extend shape inference with alias propagation and cache-capacity checks.

  4. Support dense and heterogeneous quantized paged-cache updates.

  5. Test pointer identity, capacity errors, concurrency rejection, and peak memory without a second cache allocation.