Module 5. Modern Attention Mechanisms

Solution. Each retained token needs both a key and a value on every layer:

$2H_{kv}d_h\cdot2\ \text{bytes}=2\cdot8\cdot64\cdot2=2048\ \text{B/token}.$

Thus all-full cache is $36\cdot131072\cdot2048=9$ GiB. The 18 full layers use $18\cdot131072\cdot2048=4.5$ GiB, while 18 local layers retain only their window and use $18\cdot128\cdot2048=4.5$ MiB. Total: 4.5043945 GiB. Query-head count does not enter because K/V are stored per KV head.

Solution. The first tile has maximum $m_1=2$:

$\ell_1=e^{2-2}+e^{1-2}=1+e^{-1}=1.367879,$ $n_1=1e^0+2e^{-1}=1.735759.$

For the second tile, $m_2=3$, so $\ell_2=1+e^{-2.5}=1.082085$ and $n_2=3+4e^{-2.5}=3.328340$. The global maximum is $M=3$; move tile 1 into this scale with $e^{m_1-M}=e^{-1}$:

$\ell=e^{-1}\ell_1+\ell_2=1.585300,\qquad n=e^{-1}n_1+n_2=3.966890.$

Therefore $O=n/\ell=\mathbf{2.502297}$. This rescaling is what lets tiled attention combine blocks without retaining all logits.

Solution. Append a logit $s$ to the ordinary scores $z_i=q^\top k_i/\sqrt{d_h}$ and pair it with $v_s=0$. The output becomes

The extra column can absorb probability mass without adding content. A fixed key instead produces $q^\top k_s/\sqrt{d_h}$, which changes with the query. Setting $k_s=0$ gives only the constant 0; an arbitrary learned constant $s$ for every query must be represented as a separate scalar/bias, not as an ordinary key.

Solution. The group size is $g=H_q/H_{kv}=4$, so query head $h$ uses KV head $\lfloor h/4\rfloor$: heads 0–3 map to KV0 and 4–7 to KV1. Label the KV heads [10, 20]. Correct expansion is

repeat_interleave(4): [10, 10, 10, 10, 20, 20, 20, 20]

which matches the grouped formula. tile(4) produces [10, 20, 10, 20, 10, 20, 10, 20]; the shape is valid but the head mapping is wrong. A shape check alone therefore cannot detect the bug.

Demo for this project exercise: English project-demo notebook for Module 5, sections “Demo 1” and “Demo 2” as starting points for this module’s projects. It shows a compact reproducible example on small inputs, not a completed solution to the full-scale project.

Solution. Causality compares absolute positions: a key is visible when $k\le5$. The full mask row is therefore [1,1,1,1,1,1]. A window of width 3 also requires $5-3+1\le k$, retaining keys 3, 4, and 5: [0,0,0,1,1,1].

The only row of tril(ones(1,6)) has local index 0, so it yields [1,0,0,0,0,0]. The operation has no knowledge of the five preceding positions. Construct the mask from absolute indices or supply the appropriate offset. Extend the notebook experiment to a chunk of new queries, assigning each row its own absolute position.

Solution. Stack the queries as rows of $Q\in\mathbb R^{3\times6}$ and the latents as rows of $C\in\mathbb R^{7\times4}$. The direct path constructs $K=CA^\top$ of shape $7\times6$ and $V=CB^\top$ of shape $7\times5$. Its content scores satisfy

$QK^\top=QAC^\top.$

The rearranged path first computes $QA$, a $3\times4$ array, then the $3\times7$ score matrix. Add the same positional scores and use identical scaling, masks, and softmax on both paths to obtain the same weights $P\in\mathbb R^{3\times7}$. The value computation satisfies

$PV=P(CB^\top)=(PC)B^\top.$

The direct path reconstructs seven width-5 values before mixing them. The rearranged path mixes four-dimensional latents first, producing $PC$ of shape $3\times4$, then expands the result to $3\times5$. With an output matrix $O\in\mathbb R^{5\times d_{out}}$, the product $B^\top O$ can be formed in advance.

Rotating reconstructed keys changes the effective query to $A^\top R_{n-m}^\top q$, which depends on key position $n$. It cannot be computed once and reused for every key. A separate positional branch avoids this particular dependency. Notebook §8 checks scores, values, and output projection at tolerance $10^{-12}$; a changed-key-position counterexample tests the limit of the rearrangement.

For the GPU project, establish accuracy on a short sequence with one head and an FP64 reference first. Benchmark larger shapes with matched inputs, masks and precision, warm-up, no_grad, CUDA synchronization around timing, and reset peak counters. State whether projections, mask construction and cache allocation are included; record the selected kernel and distinguish allocated from reserved memory. Reduce shapes on memory exhaustion; do not construct a large FP64 reference merely to measure speed.