Decode One Token against the Cached Prefix

Trace one incremental attention step from the current residual record through Q/K/V projection, cache append, causal attention over the prefix, residual updates, and one vocabulary-logit row.

Suppose a layer already caches $t$ positions. Normalize the current token's layer input and project one query, key, and value per head:

q_t,k_t,v_t:(B,h,1,d_h).

Append the new key and value:

K_{0:t}=[K_{0:t-1};k_t], \qquad V_{0:t}=[V_{0:t-1};v_t].

The current query is allowed to read all $t+1$ entries, so its score tensor is

S_t=\frac{q_tK_{0:t}^{\mathsf T}}{\sqrt{d_h}} :(B,h,1,t+1).

No upper-triangular mask is needed for this single final query if the cache contains only its valid prefix. A padding or batched-validity mask may still be needed.

Calculate One Head

Let $d_h=2$ , current query $q=[1,0]$ , keys

k_0=[1,0],\quad k_1=[0,1],\quad k_2=[1,1],

and values

v_0=[1,0],\quad v_1=[0,2],\quad v_2=[1,2].

The scaled scores are approximately $[0.7071,0,0.7071]$ . Their softmax is $[0.4011,0.1978,0.4011]$ . The mixed value is

0.4011v_0+0.1978v_1+0.4011v_2 \approx[0.8022,1.1978].

Join all heads, apply $W_O$ , add the attention residual, apply the current token's norm and MLP, and continue to the next layer. Only the current token travels through these updates during decode.

Q1. Calculate a cached score row

With $d_h=2$ , $q=[1,1]$ , and cached keys $[1,0]$ and $[0,2]$ , calculate the two scaled dot-product scores.

Your answer

Compute it first, then check your number.

Hint

The unscaled dot products are 1 and 2.

Solution

The scores are

[1/\sqrt2,2/\sqrt2]\approx[0.7071,1.4142]

Append Exactly Once

Appending the current key twice, appending along the head axis, or forgetting one layer's append can all leave later operations with plausible ranks. Assert that every layer's cache length grows by exactly one per consumed input token.

Calculate One Head

Q1. Calculate a cached score row

Append Exactly Once

Pause and reflect