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
KVCacheGroupId
KVCacheGroupId
class max.nn.kv_cache.KVCacheGroupId(type, window_size=-1)
Bases: object
Identifies the caches a model reuses and evicts together.
Caches behind the same attention pattern share a prefix-cache hit and an external tier namespace, so this doubles as the key for both.
recurrent names a cache whose entry is a state rather than a span of
tokens. It shares the page pool, the prefix cache and the eviction order
with the attention caches.
scratch names a cache that is never published: one block per request,
drawn when the request first grows and freed when it is released. It
draws from the same page pool as the others, so its bytes stay fungible
with theirs and admission counts them, but it carries no hash and so
takes no part in prefix hits, commits, eviction order or an external
tier.
blocks_in_window()
blocks_in_window(page_size)
full()
classmethod full()
-
Return type:
is_full()
is_full()
-
Return type:
is_recurrent()
is_recurrent()
-
Return type:
is_scratch()
is_scratch()
-
Return type:
is_sliding_window()
is_sliding_window()
-
Return type:
longest_hit()
longest_hit(page_size, resident)
How much of resident a leaf of this shape can reuse.
Residency in, prefix length out. The answer is counted from the start
even for a shape that only reads the tail of it, so the caller can
compare shapes against one another, which
max._kv_core.longest_joint_prefix_hit() does to settle what a
whole tree serves at once.
len(resident) bounds the candidate, so narrowing means passing a
shorter view, not recomputing residency.
Computing eligible Prefix Cache hits for sliding window differs greatly from full attn. Recall that the window size includes the query token. Say the query token is idx=42 and the window size is 10. This means that the query token will attend to tokens from idx=32 to idx=41.
For a concrete example:
[X]: Token is in Prefix Cache
. : Token is not in Prefix Cache
^ : Eligible Prefix Cache hit
Tokens [A] [B] . [D] [E] [F] . . [I] [J] [K] [L] [M]
w_size=2 ^ ^ ^ ^ ^ ^ ^ ^ ^ ^
w_size=3 ^ ^ ^ ^ ^ ^ ^ ^
w_size=4 ^ ^ ^ ^ ^ ^
w_size=5 ^ ^ ^ ^
w_size=6 ^ ^ ^
w_size=7 ^ ^Notice that as window_size increases, the number of indices eligible for a cache hit decreases. Additionally, we can count consecutive runs of window_size-1 tokens to determine eligibility. For example, [DEF] is a run of 3 tokens so token F is a valid cache hit for w_size=4 and below.
Additionally, partial window cache hits is possible if the run starts from the start of sequence. For example, [A] and [AB] are valid cache hits for any window size.
window_size=1 is not a case this has to serve: __post_init__
rejects it, since a query token attending to no historical tokens
would make every block a hit attention never reads back.
The rule itself is
max._kv_core.LeafShape.longest_hit(), applied to
prefix_hit_shape(), so MAX and Mach share one copy.
prefix_hit_shape()
prefix_hit_shape(page_size)
Returns the prefix-hit rule a leaf of this shape answers with.
-
Parameters:
-
page_size (int) – Tokens per block, which turns
window_sizeinto a block count. -
Returns:
-
The shape
max._kv_core’s prefix-hit functions take. -
Return type:
-
LeafShape
recurrent()
classmethod recurrent()
-
Return type:
scratch()
classmethod scratch()
-
Return type:
type
type: Literal['full', 'sliding_window', 'recurrent', 'scratch']
window_size
window_size: int = -1