PIXELBANKv8.2.1
Menu
Sign In

KV-Cache Attention

HardAttention Mechanisms

Implement KV-cache for efficient autoregressive inference.

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

  1. We only have a new single query vector (1, d)
  2. We append the new K and V to the cache
  3. 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:
2
2
1 0
1 0
1 0
0 1
0 1
0 1
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]q = [1, 0], k=[1,0]k = [1, 0], and v=[1,0]v = [1, 0]. We compute attention using these values: attention=q⋅kT∑q⋅kT⋅v=[1,0]⋅[1,0]T[1,0]⋅[1,0]T⋅[1,0]=[1,0]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 kk and vv to the cache, so K=[[1,0],[0,1]]K = [[1, 0], [0, 1]] and V=[[1,0],[0,1]]V = [[1, 0], [0, 1]]. We compute attention using the new q=[0,1]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
Editor

Test Results

0/0
Run code to see test results.