Transformer Attention and SwiGLU Forward Pass

Implement a Transformer attention and SwiGLU forward pass for one sequence of token vectors.

For this exercise, assume the pass consists exactly of bias-free multi-head self-attention followed by a bias-free SwiGLU feed-forward layer. The formulas below define the complete judged computation.

Inputs and shapes

• x has shape n by d.
• wq, wk, wv and wo each have shape d by d.
• wg and wu each have shape d by f; wd has shape f by d.
• heads divides d. Each head uses width = d / heads consecutive feature coordinates.

Multi-head self-attention

1. Project the tokens: Q = x * wq, K = x * wk and V = x * wv.
2. For head h, use coordinates h * width through (h + 1) * width - 1. For query token i and key token j, its score is the dot product of their head slices divided by sqrt(width).
3. Apply softmax across all n key tokens independently for each query and head: p[j] = exp(score[j]) / sum(exp(score)). Form the weighted sum of that head's value slices.
4. Concatenate heads in increasing head order to obtain C of shape n by d, then compute A = C * wo.

Parallel SwiGLU branches

Compute G = A * wg and U = A * wu from the same A. For every coordinate, define SiLU(z) = z / (1 + exp(-z)). Multiply element by element to obtain B = SiLU(G) elementwise-multiplied-by U. Return B * wd, an n by d matrix.

For this exercise, assume there is no batch dimension, causal or padding mask, bias, residual addition, normalization, positional encoding, dropout, training or backward pass. Use ordinary matrix loops and standard mathematical functions without a machine-learning or numerical-array library. A stable softmax may subtract the row maximum before exponentiation.