Causal LM Forward Pass

Implement causal_lm_forward: convert token IDs into embeddings, apply the supplied causal transformer blocks, and return vocabulary logits at every position. The weights are provided; this function performs a forward pass.

A causal model lets position i use its own token and earlier tokens, but never later tokens. This exercise uses post-normalization blocks: each residual addition is followed by layer normalization. Follow the order below.

Inputs and output

• input_ids: (N, T), integer-valued token IDs in [0, V), delivered as floats.
• w_emb: (V, D), the token embedding table.
• pos_embed: (T, D), one position embedding per sequence position.
• blocks_weights: (L, 6, D, D), containing L transformer blocks.
• w_head: (D, V), the vocabulary projection.
• num_heads: a positive integer dividing D.

N, T, V and D are positive; L may be zero. Shapes are consistent, and all floating-point parameters are finite and share a dtype and device. Token IDs are on the same device; cast them to integers for lookup. Leave inputs unchanged.

Return raw logits of shape (N, T, V), preserving the parameter dtype and device. Use the supplied w_head independently of w_emb. Do not choose tokens, apply a final vocabulary softmax, compute a loss or update weights.

Embeddings

For example n and position t, start with:

x[n,t,:] = w_emb[input_ids[n,t],:] + pos_embed[t,:]

Position embeddings are shared across batch examples. No padding mask is supplied; treat every position as a real token, including token ID zero.

Each transformer block

Apply blocks from 0 through L-1, passing each result into the next. The packed matrices for block i are [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2]. All six have shape (D, D), so the feed-forward width is D in this exercise.

q, k, v = x @ w_q, x @ w_k, x @ w_v
attn_out = causal_multi_head_attention(q, k, v) @ w_o
y = LN(x + attn_out)
ffn_out = GELU(y @ w_mlp1) @ w_mlp2
x = LN(y + ffn_out)

For attention, split the projected feature dimension into consecutive heads of width D/num_heads. Within each head, divide query-key dot products by sqrt(D/num_heads), then normalize over key positions. Concatenate head outputs in their original order before multiplying by w_o. Keep batch examples separate.

For query position i, only keys 0 through i are allowed. Include the diagonal, so the first token attends to itself. Mask future scores with negative infinity before softmax, or compute softmax only over allowed keys. A fixed finite mask value can fail when allowed scores are even more negative.

Compute stable softmax by subtracting each masked score row's maximum before exponentiating and dividing by its sum. Each row has at least its own position available, so it is never fully masked.

LN normalizes each token over its D features, using population variance and (token - mean) / sqrt(variance + 1e-5). It has no learned scale or bias. Use the tanh GELU approximation elementwise:

GELU(u) = 0.5 * u * (1 + tanh(sqrt(2/pi) * (u + 0.044715 * u**3)))

Build attention, normalization, softmax and GELU from basic tensor operations in either framework, without packaged attention or layer-normalization helpers. There are no linear-layer biases, dropout or an additional final normalization.

Vocabulary projection

After all blocks, return x @ w_head. With L=0, project the token-plus-position embeddings directly. The output at position i can be used to predict the next token, but shifting training targets and sampling a token are outside this function. Generation would use the final position's logits for the next choice.

The causal panel below shows the allowed query-key connections for this task.