Homework · Runs in your browser

Softmax, causal attention weights and KV-cache size

Build three small pieces of the lesson in plain Python: softmax, the causal attention weights for a short sequence, and the KV-cache size formula. Edit starter.py, then run the tests. In the browser runner, press Run tests. In Colab or on your own machine, put starter.py and test_starter.py in one folder and run pytest. The same files work everywhere, and nothing needs installing except pytest.

Hints:

Softmax: subtract the largest score before exponentiating. The answer does not change and large numbers do not overflow.
Scores: dot(query i, key j) divided by the square root of d_k, where d_k is the vector length.
Causal mask: query i may only use keys 0 to i. Run softmax over just those scores and put 0.0 in the rest of the row.
KV-cache: 2 x layers x kv_heads x head_dim x seq_len x batch x bytes per value. The lesson’s worked example (32, 8, 128, 2 bytes) is 131,072 bytes per token.

Running the tests needs JavaScript. You can still read the starter code below and run it on your own machine with pytest.

Your code starter.py

import math


def softmax(scores):
    """Turn a list of scores into probabilities that sum to 1."""
    # Wrong on purpose: returns the scores unchanged.
    return list(scores)


def causal_attention_weights(Q, K):
    """Attention weights for a sequence of n tokens.

    Q and K are lists of n vectors (lists of floats) of the same length d_k.
    Return an n x n list of lists. Row i holds the weights token i gives to
    every token. Token i may only look at tokens 0..i, so the weights after
    position i are 0.0. Scores are dot(Q[i], K[j]) / sqrt(d_k).
    """
    raise NotImplementedError


def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch, bytes_per_value):
    """Bytes of KV-cache: a Key and a Value for every layer, head, token and request."""
    raise NotImplementedError

Tab inserts 4 spaces. To move on with the keyboard, press Esc, then Tab. Shift+Tab always moves back.

Show a solution

Try it yourself first. This site is static, so the solution is sent to your browser with the page; it is hidden here, not secret.

import math


def softmax(scores):
    """Turn a list of scores into probabilities that sum to 1."""
    top = max(scores)
    exps = [math.exp(s - top) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]


def causal_attention_weights(Q, K):
    """Attention weights for a sequence of n tokens."""
    n = len(Q)
    d_k = len(Q[0])
    rows = []
    for i in range(n):
        visible = [
            sum(q * k for q, k in zip(Q[i], K[j])) / math.sqrt(d_k)
            for j in range(i + 1)
        ]
        weights = softmax(visible)
        rows.append(weights + [0.0] * (n - i - 1))
    return rows


def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch, bytes_per_value):
    """Bytes of KV-cache: a Key and a Value for every layer, head, token and request."""
    return 2 * layers * kv_heads * head_dim * seq_len * batch * bytes_per_value

Show the tests (test_starter.py)

Python

import math

from starter import causal_attention_weights, kv_cache_bytes, softmax


def test_softmax_sums_to_one_and_keeps_order():
    p = softmax([1.0, 2.0, 3.0])
    assert math.isclose(sum(p), 1.0), "softmax output must sum to 1"
    assert p[0] < p[1] < p[2], "a higher score must get a higher probability"


def test_softmax_does_not_overflow_on_large_scores():
    p = softmax([1000.0, 1000.0])
    assert math.isclose(p[0], 0.5) and math.isclose(p[1], 0.5), (
        "equal scores give equal probability; subtract the max before exp"
    )


def test_future_tokens_get_zero_weight():
    Q = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
    w = causal_attention_weights(Q, Q)
    assert w[0] == [1.0, 0.0, 0.0], "token 0 can only see itself"
    assert w[1][2] == 0.0, "token 1 must not see token 2 (causal mask)"


def test_each_row_sums_to_one():
    Q = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
    for row in causal_attention_weights(Q, Q):
        assert math.isclose(sum(row), 1.0), "each row of weights must sum to 1"


def test_scores_are_scaled_by_sqrt_d_k():
    # Token 1 has raw scores [0, 2]; d_k = 4, so scaled they are [0, 1].
    Q = [[0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]
    K = [[0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]
    w = causal_attention_weights(Q, K)
    expected = 1.0 / (1.0 + math.e)
    assert math.isclose(w[1][0], expected), "divide the dot products by sqrt(d_k)"


def test_kv_cache_matches_lesson_example():
    assert kv_cache_bytes(32, 8, 128, 1, 1, 2) == 131072, (
        "lesson example is 131,072 bytes per token"
    )
    assert kv_cache_bytes(32, 8, 128, 8192, 10, 2) == 10 * 8192 * 131072, (
        "the cache grows with sequence length and with concurrent requests"
    )