This repo implements all of OpenAI Whisper's forward pass in under 150 lines of Numpy using Einsum / Einops.
Transcription (tiny model): <|startoftranscript|><|notimestamps|> The little tales they tell are false. The door was barred, locked and bolted as well. Right pears are fit for a queen's table. A big wet stain was on the round carpet. The kite dipped and swayed but stayed aloft. The pleasant hours fly by much too soon. The room was crowded with a
Compare to main.py, the key changes in main_kv.py are
+ kv_cache = {}
+ if name not in kv_cache:
kv_cache[name] = np.array([kv_x @ W_k.T, kv_x @ W_v.T + B_v]) # prefill
elif is_casual:
kv_cache[name], _ = pack([kv_cache[name], np.array([kv_x @ W_k.T, kv_x @ W_v.T + B_v])], 'm b * c') # decode
is_casual = False # casual attention reduces to cross attention
Implements KV cache for cross attention (prefill-only), and decode in masked attention by viewing it as cross attention with one query
And of course
tokens_input, _ = pack([tokens_input, x[:, -1:]], 'b *') --> tokens_input = x[:, -1:]
Is what actually buys us the reduction in complexity, by only doing the "new" work incurred for each new token