For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

CacheLeafParamInterface

CacheLeafParamInterface

class max.nn.kv_cache.CacheLeafParamInterface(*args, **kwargs)

source

Bases: Protocol

What every child of a cache tree contributes: leaves, inputs, cost.

allocate_buffers()

allocate_buffers(total_num_pages, _prefix='')

source

Allocates the buffers for the cache, one per replica.

Empty for a cache with no buffer an op indexes.

Parameters:

• total_num_pages (int) – Pages the pool holds, including the null block.
• _prefix (str) – Names the buffers’ leaves the way leaves() does.

Return type:

Sequence[KVCacheBufferInterface]

build_runtime_inputs()

build_runtime_inputs(assignments, buffers, _prefix='', *, staging=None)

source

Builds the runtime cache inputs spanning all replicas.

assignments and buffers are indexed by data-parallel replica. Returns the KVCacheInputs pytree (a tuple of per-device leaves, or a dict of named subtrees for multi-cache models) whose leaves each hold one (replica, TP shard) device’s inputs.

Parameters:

• assignments (Sequence[KVCacheAssignments])
• buffers (Sequence[KVCacheBufferInterface])
• _prefix (str)
• staging (GraphInputStagingInterface | None)

Return type:

KVCacheInputs[Buffer, Buffer]

bytes_per_block

property bytes_per_block: int

source

Number of bytes per cache block.

Zero for a cache whose entry is not a span of tokens.

data_parallel_degree

data_parallel_degree: int

source

Degree of data parallelism.

devices

devices: Sequence[DeviceRef]

source

Devices to use for the cache.

enable_dp_cross_replica_prefix_copy

property enable_dp_cross_replica_prefix_copy: bool

source

enable_prefix_caching

property enable_prefix_caching: bool

source

flattened_kv_inputs()

flattened_kv_inputs()

source

Flattens the symbolic inputs for this cache.

Return type:

list[TensorType | BufferType]

get_symbolic_inputs()

get_symbolic_inputs(namespace='')

source

Returns the symbolic inputs for this cache.

Parameters:

namespace (str) – Prefix that disambiguates this cache’s page-pool symbolic dim from sibling caches in a multi-group tree. Empty for a single-group cache, leaving its names unchanged.

Return type:

KVCacheInputs[TensorType, BufferType]

kv_connector_config

kv_connector_config: KVConnectorConfigInterface

source

kv_hash_algo

property kv_hash_algo: Literal['ahash64', 'sha256', 'sha256_64']

source

kv_hash_seed

property kv_hash_seed: bytes | None

source

leaf_kind

property leaf_kind: CacheLeafKind

source

What this child holds.

leaves()

leaves(_prefix='')

source

Returns the leaves this cache contributes to the pool.

Parameters:

_prefix (str)

Return type:

Mapping[str, KVLeafRegion]

n_devices

property n_devices: int

source

Returns the total number of devices.

num_draft_tokens

num_draft_tokens: int

source

page_size

page_size: int

source

Tokens a block covers, or zero where this child declares no pool.

replicates_kv_across_tp

property replicates_kv_across_tp: bool

source

slab_to_bound_views()

slab_to_bound_views(slabs)

source

Returns the views this cache binds once and never restages.

Keyed the way its leaves read them back. Empty for a cache whose pages the graph reaches through a per-forward table.

Parameters:

slabs (Sequence[Buffer]) – One replica’s slab per device.

Return type:

Mapping[str, list[Buffer]]

slab_to_buffer_views()

slab_to_buffer_views(buffers, padded_page_bytes=None, _prefix='')

source

Converts one replica’s slabs into the pages its leaves occupy.

Parameters:

• buffers (Sequence[Buffer]) – One replica’s slab per device.
• padded_page_bytes (Mapping[str, int] | None) – Each padded leaf’s page stride, or None.
• _prefix (str) – Names the views’ leaves the way leaves() does.

Return type:

KVCacheBufferInterface

speculative_method

speculative_method: Literal['eagle', 'mtp', 'dflash', 'dflash2'] | None

source

staged_dispatch_metadata()

staged_dispatch_metadata(_prefix='')

source

The dispatch metadata this tree stages, keyed by name in a replica.

Empty for a kernel that reads it on the host: nothing is transferred, so there is nothing to stage. _prefix keeps the leaves of a tree apart, the same way leaves() does.

Parameters:

_prefix (str)

Return type:

Mapping[str, DispatchMetadataSpec]

unflatten_kv_inputs()

unflatten_kv_inputs(it)

source

Unflattens the symbolic inputs for this cache.

Parameters:

it (Iterator[Any])

Return type:

KVCacheInputs[TensorValue, BufferValue]