📘

KV-Cache Attention

Implement KV-cache for efficient autoregressive inference.

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

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).

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

Editor

Python 3.13.1

0/0

Run code to see test results.

📘