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
PagedCacheValues
PagedCacheValues
class max.nn.common_layers.kv_cache.PagedCacheValues(kv_blocks, cache_lengths, lookup_table, max_prompt_length, max_cache_length, page_stride, kv_scales=None, scales_page_stride=None, scales_lookup_table=None, attention_dispatch_metadata=None, draft_attention_dispatch_metadata=None, mla_num_partitions=None, draft_mla_num_partitions=None, split_head_dims=None, mixed_step_split=None)
Bases: KVCacheInputsPerDevice[Tensor, Tensor]
Distributed view of KV cache inputs.
-
Parameters:
-
- kv_blocks (_Buffer)
- cache_lengths (_Tensor)
- lookup_table (_Tensor)
- max_prompt_length (_Tensor)
- max_cache_length (_Tensor)
- page_stride (_Tensor)
- kv_scales (_Buffer | None)
- scales_page_stride (_Tensor | None)
- scales_lookup_table (_Tensor | None)
- attention_dispatch_metadata (_Tensor | None)
- draft_attention_dispatch_metadata (_Tensor | None)
- mla_num_partitions (_Tensor | None)
- draft_mla_num_partitions (_Tensor | None)
- split_head_dims (SplitKVHeadDims | None)
- mixed_step_split (_Tensor | None)
from_upstream()
classmethod from_upstream(per_device, mapping)
Constructs from per-device caches, as graph values or tensors.
Each device’s fields are its shards: graph values, single-device tensors in the graph being built, or, for a decoder called eagerly, real tensors.
-
Parameters:
-
- per_device (Sequence[KVCacheInputsPerDevice[TensorValue, BufferValue]] | Sequence[KVCacheInputsPerDevice[Tensor, Tensor]])
- mapping (DeviceMapping)
-
Return type:
n_devices
property n_devices: int
Returns the number of devices the paged KV cache is located on.
to_graph_values()
to_graph_values()
Returns this single-device cache as graph values.
The blocks and scales become buffers, since attention kernels write them in place.
-
Return type: