A bounded cache backed by fixed pre-allocated buffers written in place.
Parameters
max_cache_lenintAttributes
Notes
get_seq_length is tracked by a per-layer counter (not inferred from the
buffer shape, which is always max_cache_len). update stores into the
full buffer; when the consumer opts in (passes a read_len in
cache_kwargs), it returns a filled-prefix view so attention attends
only over real keys — O(filled), not O(max_cache_len). The stored buffer
stays full, so the fixed buffer shape is preserved.
Used by 2
Constructors
1Properties
1Instance methods
8Repeat every batch row repeats times along the batch dimension.
Keep only the batch rows named by indices along the batch dim.
Logically truncate the cache to max_length tokens.
The buffer is fixed-size, so cropping only rewinds the length counters
(cropped positions become 'future' and are masked by the causal mask).
A negative max_length drops that many tokens off the end.
Return the fixed buffer capacity max_cache_len (always bounded).
Number of tokens written so far for layer_idx, read from the
per-layer counter (the buffer shape is always max_cache_len, so the
count cannot be inferred from it).
Reindex every layer's buffers along the batch dimension with
beam_idx (in place) to keep beams aligned after a beam permutation.
Zero every buffer and reset the length counters in place.
update(key_states: Tensor, value_states: Tensor, layer_idx: int, cache_kwargs: dict[str, object] | None = None)Write key_states / value_states into the fixed buffer at
cache_position and return a filled-prefix view of it.
The write goes through the traceable lucid.index_copy (not
in-place __setitem__), so the decode forward stays compile-visible
and the stored buffer keeps a constant max_cache_len shape every step.
The returned keys/values are narrowed (dim 2) to the filled-prefix
width (the read_len in cache_kwargs) so attention attends over the
real keys, not the zero-padded tail — recovering O(filled) attention cost.
Parameters
key_statesTensor(B, num_heads, T_new, head_dim).value_statesTensorlayer_idxintcache_kwargsdict or None= None"cache_position" (a fixed-shape int64 Tensor of
absolute write indices; absent → positions from the running counter)
and "read_len" (an int — narrow the returned view to this width;
the decoder-only attention passes past_len + T here to attend over
only the filled prefix). When no read_len is given the full buffer
is returned (back-compatible).Returns
Dunder methods
3Return the full (keys, values) buffers for layer_idx.
Iterate the per-layer (keys, values) buffers.
Number of layers currently allocated.