📘

KV-Cache Attention

HardAttention Mechanisms 📚 LLM 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 attention with the full cached K, V

Input:

Line 1: d (dimension)
Line 2: num_steps (number of generation steps)
For each step:
- A line with the new q vector (d floats)
- A line with the new k vector (d floats)
- A line with the new v vector (d floats)

Output: For each step, the attention output vector (d floats), rounded to 4 decimal places.

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

Example:

Input:

Output:

[1.0000 0.0000]
[0.2689 0.7311]

Reasoning:

We start with an empty cache and at the first step, we have $q = [1, 0]$ , $k = [1, 0]$ , and $v = [1, 0]$ . We compute attention using these values: $attention = \frac{q \cdot k^T}{\sum q \cdot k^T} \cdot v = \frac{[1, 0] \cdot [1, 0]^T}{[1, 0] \cdot [1, 0]^T} \cdot [1, 0] = [1, 0]$ .
At the second step, we append the new $k$ and $v$ to the cache, so $K = [[1, 0], [0, 1]]$ and $V = [[1, 0], [0, 1]]$ . We compute attention using the new $q = [0, 1]$ : $attention = \frac{[0, 1] \cdot [[1, 0], [0, 1]]^T}{[0, 1] \cdot [[1, 0], [0, 1]]^T} \cdot [[1, 0], [0, 1]] = \frac{[0, 1]}{1} \cdot [[1, 0], [0, 1]] = [0, 1] \cdot \frac{[[1, 0], [0, 1]]}{1} = \frac{[0, 1]}{\sqrt{2}} \cdot \frac{[[1, 0], [0, 1]]}{\sqrt{2}} = \frac{[0, 1]}{2} \cdot [[1, 0], [0, 1]] = \frac{1}{2} \cdot [0, 1] = \frac{1}{2} \cdot \begin{bmatrix} 0 \ 1 \end{bmatrix} + \frac{1}{2} \cdot \begin{bmatrix} 0 \ 1 \end{bmatrix} = \begin{bmatrix} 0 \ \frac{1}{2} \end{bmatrix} + \begin{bmatrix} 0 \ \frac{1}{2} \end{b

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

More from LLM 1: Foundations

Transformer Block Forward PassHard Feed-Forward NetworkEasy Layer NormalizationEasy Multi-Head AttentionMedium Causal Attention MaskMedium

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

0/0

Run code to see test results.