8.4 Causal masking & the KV cache

import json, math

# The real microgpt after 1,000 training steps (same weights the widgets use).
M = json.load(open("microgpt-trained.json"))
W, UCHARS = M["weights"], M["uchars"]
BOS = len(UCHARS)

def linear(x, w):  # microgpt line 94
    return [sum(wi * xi for wi, xi in zip(row, x)) for row in w]

def softmax(z):  # lesson 6.4
    m = max(z)
    e = [math.exp(v - m) for v in z]
    t = sum(e)
    return [v / t for v in e]


def rmsnorm(x):  # given here; you build it in lesson 8.5
    ms = sum(v * v for v in x) / len(x)
    return [v * (ms + 1e-5) ** -0.5 for v in x]


def attend(q, keys, values):  # lesson 8.3
    # assert cond, message: carry on if cond is True, otherwise stop with message
    assert keys, ("the cache is empty, so there is nothing to attend to yet: append this "
                  "token's key and value first (the TODO in read)")
    scores = [sum(qi * ki for qi, ki in zip(q, k)) / math.sqrt(len(q)) for k in keys]
    weights = softmax(scores)
    return [sum(w * v[j] for w, v in zip(weights, values)) for j in range(len(values[0]))], weights


# Read a name one token at a time. Each new token adds its key and value to
# the cache, then attends over everything cached so far - never the future.
def read(name, head=0):
    tokens = [BOS] + [UCHARS.index(c) for c in name]  # .index(c) = where c sits in UCHARS: its token id
    keys, values = [], []                            # the KV cache, empty at the start
    all_weights, outputs = [], []
    s, e = head * 4, head * 4 + 4                    # this head's 4 numbers
    for pos, tok in enumerate(tokens):  # enumerate gives (0, first), (1, second), ...: position and item
        x = [t + p for t, p in zip(W['wte'][tok], W['wpe'][pos])]
        x = rmsnorm(rmsnorm(x))  # exactly as microgpt does before attention (lines 112, 117)
        q = linear(x, W['layer0.attn_wq'])[s:e]
        # TODO: append this token's key and value to the cache (same pattern as q,
        # with W['layer0.attn_wk'] and W['layer0.attn_wv'], and the same [s:e] slice)
        out, weights = attend(q, keys, values)
        all_weights.append(weights)
        outputs.append(out)
    return all_weights, outputs


# Given, for comparison: a reader that CHEATS. It makes the keys and values of
# the WHOLE name first, then lets every position attend over all of them,
# future letters included. This is exactly what the cache prevents.
def read_cheating(name, head=0):
    tokens = [BOS] + [UCHARS.index(c) for c in name]
    s, e = head * 4, head * 4 + 4
    xs = [rmsnorm(rmsnorm([t + p for t, p in zip(W['wte'][tok], W['wpe'][pos])]))
          for pos, tok in enumerate(tokens)]
    keys = [linear(x, W['layer0.attn_wk'])[s:e] for x in xs]
    values = [linear(x, W['layer0.attn_wv'])[s:e] for x in xs]
    all_weights, outputs = [], []
    for x in xs:
        out, weights = attend(linear(x, W['layer0.attn_wq'])[s:e], keys, values)
        all_weights.append(weights)
        outputs.append(out)
    return all_weights, outputs


# No peeking? Read "em", then "emma". The first 3 rows (., e, m) should not change.
def biggest_change(reader):
    short, _ = reader('em')
    longer, _ = reader('emma')
    biggest = 0.0
    for r1, r2 in zip(short, longer):
        for a, b in zip(r1, r2):
            biggest = max(biggest, abs(a - b))  # abs(x) drops the sign: abs(-0.3) is 0.3
    return biggest


weights, _ = read('emma')
for pos, w in enumerate(weights):
    print('.emma'[pos], [round(x, 2) for x in w])
print('biggest change to the first 3 rows when "em" grows to "emma":')
print('  read          ', round(biggest_change(read), 4))
print('  read_cheating ', round(biggest_change(read_cheating), 4), ' <- the future leaked in')