KV-Cache Attention

hardFREEAttention MechanismsLLM 1: Foundations

Implement KV-cache for efficient autoregressive inference.

During autoregressive generation, we cache previous K and V values. At each new step:

We only have a new single query vector (1, d)
We append the new K and V to the cache
We compute scaled dot-product attention with the full cached K, V: $\text{softmax}(qK^T/\sqrt{d})\,V$

Input: kv_cache_attention(d, steps)

d: the dimension
steps: a list of (q, k, v) tuples, one per generation step, each vector holding d floats

The example input lists the same data as lines: d, then num_steps, then the q, k and v vectors of each step.

Output: A list with one attention output vector (d floats) per step. Values are printed rounded to 4 decimal places.

Use causal attention (each step sees all previous + current).

Example

Input

Output

[1.0000 0.0000]
[0.3302 0.6698]

Reasoning

Inputs: $d = 2$ and two steps. Step 1 brings $q = k = v = (1, 0)$ ; step 2 brings $q = k = v = (0, 1)$ .
Step 1: the cache holds a single key, $K = V = \begin{pmatrix}1 & 0\end{pmatrix}$ . Softmax over one score is always $1$ , so the output is exactly that value: $(1, 0)$ .
Step 2: the cache now has two rows, $K = V = \begin{pmatrix}1 & 0 \\ 0 & 1\end{pmatrix}$ . Scores: $q \cdot k_1 / \sqrt{2} = 0$ and $q \cdot k_2 / \sqrt{2} = 1/\sqrt{2} \approx 0.7071$ .
Softmax: $e^{0} = 1$ , $e^{0.7071} \approx 2.0281$ , total $3.0281$ , giving weights $1/3.0281 \approx 0.3302$ and $2.0281/3.0281 \approx 0.6698$ .
Output: $0.3302\,(1, 0) + 0.6698\,(0, 1) = (0.3302, 0.6698)$ , printed as [0.3302 0.6698].
Notice: unscaled scores $0$ and $1$ would give weights $0.2689$ and $0.7311$ . Dividing by $\sqrt{2}$ narrows the gap to $0.7071$ , pushing the weights closer to even.

Constraints

1 <= d <= 8, 1 <= num_steps <= 5
Each step produces a 1xd output
Cache grows by one K,V pair per step
Round to 4 decimal places

Learn this in

Deployment & OptimizationLLM Study Plan Transformer ArchitectureLLM Study Plan NLP FundamentalsML Study Plan Advanced NLP TopicsNLP Study Plan

More from LLM 1: Foundations

Transformer Block Forward Passhard Feed-Forward Networkeasy Layer Normalizationeasy Multi-Head Attentionmedium Causal Attention Maskmedium

Concept primer

Background Knowledge

The problem revolves around Attention Mechanisms, a key component in many modern neural network architectures, especially in Natural Language Processing (NLP) and sequence-to-sequence models. Attention allows the model to focus on different parts of the input sequence when generating each output element, rather than using a fixed context. In the context of autoregressive generation, where the model predicts one output element at a time based on the previously generated elements, causal attention is particularly relevant. Causal attention ensures that at each step, the model only attends to the previous and current elements in the sequence, not to future elements.

The KV-Cache Attention mechanism is an optimization technique for efficient autoregressive inference. It involves caching the previously computed Key (K) and Value (V) vectors, which are essential components in computing attention weights. By caching these vectors, the model avoids redundant computations at each step, significantly improving inference speed. The Query (Q) vector, which represents the context in which the attention is being applied, is used to compute attention weights with respect to the cached K and V vectors.

Understanding the basics of vector operations, such as dot product, softmax function, and matrix multiplication, is crucial for implementing the KV-Cache Attention mechanism. Additionally, familiarity with autoregressive models and how they generate sequences one step at a time is necessary. The problem also involves dynamic caching, where the cache is updated at each generation step by appending new K and V vectors.

Algorithm/Approach

The general approach to solving this problem involves implementing a loop that iterates over each generation step. At each step, the algorithm should:

Append the new K and V vectors to the cache.
Compute the attention weights using the new Q vector and the cached K vectors.
Calculate the attention output by taking the dot product of the attention weights and the cached V vectors.

This process leverages the efficiency of caching to minimize redundant computations, making it suitable for autoregressive generation tasks.

Step-by-Step Strategy

Initialize the cache for K and V vectors.
For each generation step:

Read the new Q, K, and V vectors.
Append the new K and V vectors to their respective caches.
Compute the attention weights using the new Q vector and the cached K vectors.
Calculate the attention output by combining the attention weights with the cached V vectors.
Output the attention output vector, rounded to 4 decimal places.

Common Pitfalls

Incorrectly updating the cache: Ensure that the cache is updated with the new K and V vectors at each step.
Miscalculating attention weights: Verify that the attention weights are computed correctly using the Q vector and the cached K vectors.
Improperly handling vector operations: Pay attention to the dimensions of vectors during operations like dot product and matrix multiplication.

Time & Space Complexity

Time Complexity: The time complexity is expected to be $O(n \cdot d^2)$ , where $n$ is the number of generation steps and $d$ is the dimension of the vectors. This is because at each step, the model performs operations involving the Q, K, and V vectors, which have a dimension of $d$ .
Space Complexity: The space complexity is $O(n \cdot d)$ , as the model needs to store the cached K and V vectors for all previous steps.

solution.py

Python 3.13

Test results