Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. Parameters: key_states (`torch

Given the sequence length of the new inputs, returns the usable length of the cache.

Expand KV to query-head granularity before compression/caching, unless the opt-in kv_head cache layout is enabled (see docs/gqa_cache_layout.md).

Expand a kv-head-granular cache for eager/sdpa attention compute. No-op when the cache already holds query-head-granular tensors.

Returns the sequence length of the cached states. A layer index can be optionally passed.

Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token first unpad the input, then compu

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, he

Reduce query-head-granularity scores to kv-head granularity. `scores` has shape (bsz, num_query_heads, ...) with heads laid out the way repea

Detect the query-heads-per-kv-head group count from the tensors themselves. Returns 1 when query and key head counts match (MHA, or GQA tensors t

Build the Hugging Face `generate(..., cache_config=...)` dictionary. KIVI quantizes keys per-channel and values per-token. With the HQQ cache

Return True when ``past_key_values`` holds no cached tokens yet. Handles the three shapes seen across transformers versions: ``None``, a ``Ca

Converts the `DynamicCache` instance into the its equivalent in the legacy cache format. Used for backward compatibility.

Run update_kv twice on cloned identical inputs; return both results.

Returns the maximum sequence length of the cached states, if there is any.

Small synthetic per-QUERY-head capacity table, as run_longbench builds (integer, shape (num_layers, num_heads)); values stay well below q_len

Quantizes a key/value using a defined quantization method.

Top-k select past tokens per head and keep the observation window. Works at whatever head granularity key/value/attn_cache share; in kv_head

Converts a cache in the legacy cache format into an equivalent `DynamicCache`. Used for backward compatibility.

Unrepeated GQA tensors: query at Q_HEADS, key/value at KV_HEADS.

Collect every stop token id (e.g. Llama-3 dual terminators, issue #46).

Return top Quest page indices for the non-recent prefix. Recent tokens are protected outside the page score path and should be appended by `s

Return the per-layer key cache tensors from a Cache or legacy tuple.

Thread the per-layer knobs onto self_attn.config the way run_longbench.py does.

Dequantizes back the tensor that was quantized by `self._quantize()`

SnapKV-style observation-window scores at kv-head granularity. Keys are repeated TRANSIENTLY (for the score matmul only); the returned cache

Mirror the prepare_inputs_for_generation reset guard used by the patches. Returns True when a reset happened so tests can assert the branch.

Repeat the cache `repeats` times in the batch dimension. Used in contrastive search.

Only keep the `indices` in the batch dimension of the cache. Used in contrastive search.

Split the current instance into a list of `DynamicCache` by the batch size. This will be used by `_split_model_inputs()` in `generation.utils`

Build per-page key minima and maxima. Args: key_states: Tensor with shape `[batch, heads, seq_len, head_dim]`. page_size: Number

Moves the previous layer cache to the CPU

Returns the sequence length of the cached states. A layer index can be optionally passed.

Merge adjacent-layer states with MiniCache SLERP direction sharing. Args: current_states: KV states from layer `l`, shape `[batch, heads,

Independent reference for the SnapKV kv_head selection: transient repeat -> window scores with the SnapKV causal mask -> fp32 softmax -> row-sum -

Lower text and remove punctuation, articles and extra whitespace.

Lower text and remove punctuation, extra whitespace.

Install local cache overrides required by selected quantization methods.

Helper function to plot a heatmap for a tensor of shape (seq_len, seq_len).

Reduce attention scores column-wise using NACL proxy tokens. Args: attn_scores: Either `[batch, heads, query, key]` raw attention scores

Reduce recent attention history into per-token importance scores. `attn_scores` may be raw attention with shape `[batch, heads, query, key]`

Reorders the cache for beam search, given the selected beam indices.

Estimate query-aware criticality for each page. The score is the upper bound of `query @ key` over the page's min/max box: positive query dim

Select high-angular-distance token positions to keep unmerged. `angular_distance` is expected to be normalized by pi and shaped `[batch, head

Choose proxy token positions for NACL scoring. `suffix` matches the common long-context QA layout where the user question is placed at the en

Select retained KV token indices with NACL proxy/random eviction. Returned indices have shape `[batch, heads, retained]` and are sorted in as

Return Quest-selected token indices with an exact token budget. The prefix selection is page-based, then expanded to ascending token indices.

Select persistent pivotal tokens under a fixed KV budget. Returned indices have shape `[batch, heads, retained]` and are sorted in chronologi

Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. Parameters: key_states (`torch