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 function

per_device_caches

per_device_caches()

max.nn.common_layers.kv_cache.per_device_caches(kv_params, kv)

source

Returns a decoder’s KV cache inputs as one record per device.

A traced decoder gets records of graph values, as kv_params.unflatten_kv_inputs builds them. A decoder called eagerly gets real tensors, which have no graph values, so each record is a PagedCacheValues of them, whose blocks attention ops can write.

Parameters:

• kv_params (KVCacheParamInterface) – The parameters of the cache kv holds.
• kv (tree.Tree[Tensor]) – The decoder’s KV cache inputs, as its DecoderInputs holds them.

Returns:

One record per device, in device order.

Return type:

list[KVCacheInputsPerDevice[Any, Any]]