Hard End-to-End

Causal RoPE and SwiGLU Block

Build a causal attention block with rotary position embeddings, followed by a gated feed-forward network. Each branch receives a normalized input, but its output is added to the unnormalized residual stream.

The components are often used in language models. This exercise specifies a simplified block with ordinary multi-head attention and no learned RMSNorm scale. Follow the operations below rather than copying a complete model's architecture or parameter layout.

Inputs and output

Implement llama_block(x, w_q, w_k, w_v, w_o, w_gate, w_up, w_down, num_heads, freqs_cos, freqs_sin). Keep the supplied function name.

x has shape (N, T, d_model).
w_q, w_k, w_v, w_o each have shape (d_model, d_model).
w_gate and w_up have shape (d_model, d_ff).
w_down has shape (d_ff, d_model).
num_heads divides d_model; d_head = d_model // num_heads is even.
freqs_cos and freqs_sin each have shape (T, d_head // 2).

All arrays are floating-point and all dimensions are positive. Return shape (N, T, d_model). Use the supplied matrix orientations and cosine/sine tables directly. There are no biases, dropout, padding masks or learned normalization parameters. Keep batch items independent.

Normalize without subtracting the mean

Use this RMSNorm separately for each token, reducing over the final axis:

rmsnorm(t) = t / sqrt(mean(t*t, axis=-1, keepdims=True) + 1e-6)

Unlike layer normalization, this does not subtract the mean. For example, a nonzero constant vector stays nonzero after RMSNorm. Add epsilon inside the square root and do not apply a learned scale.

Attention branch

Compute norm1 = rmsnorm(x). Project Q, K and V from norm1 using their respective weight matrices. Keep the original x for the residual.
Reshape each projection to (N, T, num_heads, d_head), then swap the head and sequence axes to get (N, num_heads, T, d_head).
Rotate adjacent feature pairs in Q and K only. For pair (a, b) at position t and pair index k, take c = freqs_cos[t, k] and s = freqs_sin[t, k]. Replace the pair with (a*c - b*s, a*s + b*c). Use both original coordinates, broadcast tables over batches and heads, and interleave the results. Leave V unchanged.
Compute scores = Q_rot @ transpose_last_two_axes(K_rot) / sqrt(d_head). The score shape is (N, num_heads, T, T). Allow key j for query i only when j <= i. Replace future scores with negative infinity, then apply a stable softmax over keys and multiply by V.
Swap the head and sequence axes back and join features in their original order. Compute y = x + joined @ w_o.

Create the causal mask on the input device in PyTorch. Every row includes its diagonal, so none is fully masked. Assume allowed scores are finite. A finite sentinel such as -1e9 is not exact exclusion when valid scores can be even smaller.

Gated feed-forward branch

Normalize the updated residual y, then compute:

norm2 = rmsnorm(y)
gate = silu(norm2 @ w_gate)
up = norm2 @ w_up
ffn_out = (gate * up) @ w_down
output = y + ffn_out

Here silu(t) = t * sigmoid(t). The gate and up arrays both have shape (N, T, d_ff). Their multiplication is elementwise; multiplication by w_down is a matrix product that restores width d_model. Use a numerically stable sigmoid or the framework's SiLU operation.

Check the residual order

There is no final normalization after either residual addition. The first residual adds to x; the second adds to y, which already includes attention. If both w_o and w_down are zero, the complete block returns x unchanged. Neither branch can use future positions: attention is causally masked, and RMSNorm and the feed-forward network work separately at each position.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Give attention information about token order

Self-attention based only on token content has no absolute position input. Position information can be added to token vectors, added to attention scores, or introduced by transforming queries and keys.

Learned absolute embeddings provide a vector for each position. Sinusoidal encodings use sine and cosine pairs at different frequencies. When adding either to token embeddings, the feature widths must match. In a sinusoidal encoding, position determines the angle while the feature-pair index determines the frequency.

Describe the distance between query and key

Relative-position methods represent the relationship between two positions. A learned score table, for example, can use their signed distance as an index. For query position 3 and key position 1, key - query is -2, while query - key is 2. Use the task's sign convention.

Clipping distances to a fixed range makes all more-distant pairs share an entry at the corresponding end of that range.

Rotate vectors or add a score bias

RoPE rotates pairs of query and key coordinates according to position. Rotations preserve vector lengths. Applying the same frequency scheme to queries and keys makes their dot product depend on the relative rotation between their positions. RoPE transforms the vectors rather than adding an embedding to them.

Some implementations pair adjacent features; others pair corresponding features from the two halves of a vector. Use the specified pairing, because it changes the result.

ALiBi adds a linear distance bias to attention scores, with a slope for each head. It affects preferences among keys without creating a position-embedding vector. Causal masking is still needed separately: a finite negative score bias does not give a future key exactly zero attention.

Check position zero, equal positions, and a small offset. With a cache during generation, use the new token's actual sequence position; restarting at zero on every step gives the wrong positional relationship to cached keys.

Combine attention with a feed-forward network

A transformer block uses attention to combine information across token positions. Its feed-forward network transforms features at each position independently, usually using shared weights across positions. Residual connections add a sublayer's output to its input, and normalization controls the values seen by that sublayer or passed to the next one.

Split attention into heads

Multi-head attention projects inputs into queries, keys, and values for several heads. Each head computes its own attention result. The block concatenates those results and applies an output projection.

For shape (batch, time, heads * head_width), first reshape to (batch, time, heads, head_width), then move the head axis if the attention operation expects (batch, heads, time, head_width). A reshape alone cannot exchange time and head axes while preserving their meanings.

Match masks to the source of the keys

In self-attention, queries, keys, and values come from the same sequence. In encoder-decoder cross-attention, decoder states supply queries and encoder outputs supply keys and values. Output length follows the queries, even when the source sequence has a different length.

A causal mask blocks future keys in decoder self-attention. Cross-attention to an available encoder sequence usually needs a source-padding mask instead. Use the mask appropriate to the information the model is allowed to read.

Keep normalization in the specified place

Pre-normalization computes x + F(Norm(x)). Post-normalization computes Norm(x + F(x)). Moving the normalization changes the function.

The same applies to replacing LayerNorm with RMSNorm or changing the feed-forward activation to SwiGLU. Follow the exercise's equations rather than filling in details from another model with a similar name. Write each intermediate shape, and check residual additions, packed weight order, epsilon, and activations.

Detecting runtime…