Medium

Decode One Token With a KV Cache

~20 mincode completion

Implement kv_cache_decode_step(K_cache, V_cache, q, k, v) and return the output vector as a list of $d$ floats. NumPy only.

An empty cache ( $t = 0$ ) arrives as an empty list. Reshape the caches to $(t, d)$ with .reshape(-1, d) before appending.
With an empty cache the only position is the new token itself, so the output is exactly v.
Scores can be large in either direction. Your softmax must not overflow or return .

Examples

Empty cache: the only position is the new token, so the output is v

Input: kv_cache_decode_step([], [], [1, 0, -1], [0.5, 0.5, 0.5], [2, -1, 3])
Output: [2, -1, 3]

Two cached tokens (the worked example)

Input: kv_cache_decode_step([[1, 0], [0, 1]], [[1, 0], [0, 1]], [1, 0], [1, 1], [2, 2])
Output: [1.203336, 1]

Large scores: exp overflows unless you subtract the max first

Input: kv_cache_decode_step([[3, 1], [2, 2], [1, 3.5]], [[1, -1], [0.5, 0.5], [-2, 4]], [400, 400], [4, 0.2], [3, 1])
Output: [-2, 4]

Hints

Hint 1

Use a matrix product rather than nested loops, and check which operand transposes.

Hint 2

Do not forget to append new key and value. That step is easy to skip.

Requirements

K_cache: Cached keys, shape (t, d). May be empty (t = 0).
V_cache: Cached values, shape (t, d). May be empty (t = 0).
Return Attention output for the new token as a list of d floats.

Constraints

Allowed library: NumPy only
Time limit: 200 ms, Memory: 64 MB

Try similar problems(4)

Medium

Causal Attention Mask

Pro

Attention & Transformers · ~6 min

Medium

Sinusoidal Positional Encoding

Pro

Attention & Transformers · ~12 min

Medium

Softmax Attention Weights

Pro

Attention & Transformers · ~12 min

Easy

Query-Key Attention Scores

Pro

Attention & Transformers · ~12 min

Where this shows up

~20 min

8 employers weight this skill

4 frontier labs, 3 AI product companies, 1 enterprise vendor. Top match scores 92.

Python

import numpy as np

def kv_cache_decode_step(K_cache, V_cache, q, k, v) -> list:
    """
    One decode step of single-head attention with a KV cache.

    Args:
        K_cache: Cached keys, shape (t, d). May be empty (t = 0).
        V_cache: Cached values, shape (t, d). May be empty (t = 0).
        q, k, v: Query, key and value of the new token, each shape (d,)

    Returns:
        Attention output for the new token as a list of d floats.
    """
    q = np.asarray(q, dtype=float)
    d = q.shape[0]
    K_cache = np.asarray(K_cache, dtype=float).reshape(-1, d)
    V_cache = np.asarray(V_cache, dtype=float).reshape(-1, d)
    # YOUR CODE HERE: append k and v, score, stable softmax, mix the values
    pass

Loading docs…

The AI Mentor needs an account

It reads your code and the failing tests and nudges you toward the fix without handing you the answer. Free accounts get it on every problem you're working on today.