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

KVCacheAssignments

KVCacheAssignments

class max.nn.kv_cache.KVCacheAssignments(cache_lengths_by_device, staged_by_device, max_prompt_length, max_cache_length, batch_characteristics)

source

Bases: object

Assignments of request blocks to KV cache pages for a replica.

batch_characteristics carries the effective (batch_size, max_prompt_length, max_cache_valid_length) used to build this assignment (after any graph-capture upper-bound override) so that KVCacheParamInterface.build_runtime_inputs() can resolve the decode attention dispatch keys from the same values.

Parameters:

• cache_lengths_by_device (list[Buffer])
• staged_by_device (list[dict[str, Buffer]])
• max_prompt_length (Buffer)
• max_cache_length (Buffer)
• batch_characteristics (BatchCharacteristics)

batch_characteristics

batch_characteristics: BatchCharacteristics

source

cache_lengths_by_device

cache_lengths_by_device: list[Buffer]

source

max_cache_length

max_cache_length: Buffer

source

max_prompt_length

max_prompt_length: Buffer

source

staged_by_device

staged_by_device: list[dict[str, Buffer]]

source

Each leaf’s own inputs, per device.

A leaf’s own key holds the table its ops index the pool by. A leaf wanting more than one keys them under itself, the way quantized attention keys <leaf>/scales.