the next step is tracing what actually happens during a forward pass. Every token flows through a series of tensor transformations until it becomes the next-token prediction. Most of the code is simply reshaping tensors and making GPU computation efficient.
Instead of computing Query, Key, and Value separately, GPT performs one large linear transformation:
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
Recall that a linear layer computes
y = xWᵀ + b
Here the output dimension is three times larger:
768 → 2304
because
2304 = 768 × 3
The resulting tensor is simply split into three equal parts:
QKV
This is computationally cheaper than performing three independent matrix multiplications.
Splitting into Multiple Attention Heads
Each Q, K, and V still has shape
(B, T, C)
GPT then reshapes them:
k = k.view(B, T, n_head, C // n_head).transpose(1, 2)
Suppose
C = 768n_head = 12
Each head receives
768 / 12 = 64 dimensions
After reshaping:
(B, T, C)↓(B, T, n_head, head_size)↓(B, n_head, T, head_size)
The transpose places the head dimension before the sequence dimension, making later matrix multiplication much more efficient.
The attention tensor now has shape
(B, nh, T, hs)
where:
- nh = number of heads
- hs = head size
Each head learns different relationships between tokens in parallel.
Flash Attention: The GPU Does Everything in One Pass
Modern PyTorch can invoke Flash Attention directly:
torch.nn.functional.scaled_dot_product_attention( q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
Although this looks like one function call, internally it launches highly optimized CUDA kernels.
Instead of performing separate operations:
QKᵀ↓scale↓mask↓softmax↓dropout↓× V
Flash Attention fuses them into one GPU kernel.
The GPU avoids materializing the enormous attention matrix in memory, dramatically reducing memory traffic while producing exactly the same mathematical result.
This is why Flash Attention is so much faster—not because the algorithm changes, but because the implementation is engineered around modern GPU architecture.
What Happens Without Flash Attention?
The original implementation makes each step explicit.
First compute attention scores:
att = (q @ k.transpose(-2, -1))att *= 1.0 / math.sqrt(k.size(-1))
This is simply
Attention Scores = QKᵀ / √d
where d is the dimension of each attention head.
The scaling factor prevents the dot products from becoming too large before the softmax.
Applying the Causal Mask
GPT must never look into the future.
The mask is a lower-triangular matrix stored inside the model:
self.bias
It looks like
1 0 0 01 1 0 01 1 1 01 1 1 1
The masking step is
att = att.masked_fill( self.bias[:, :, :T, :T] == 0, float("-inf"))
The lower half remains unchanged.
The upper half becomes
−∞
So the attention matrix becomes something like
0.8 -∞ -∞ -∞0.2 0.6 -∞ -∞0.3 0.4 0.7 -∞...
Why Fill with Negative Infinity?
Softmax computes
eˣ / Σ eˣ
Since
e^(−∞) = 0
those masked positions receive exactly zero probability.
Meanwhile
e⁰ = 1
and finite attention scores contribute normally.
The result is that future tokens are mathematically ignored without requiring any special logic during the softmax.
Dropout on Attention
After softmax, GPT may randomly remove some attention connections:
attn_dropout(att)
Dropout randomly sets some attention weights to zero during training.
This prevents the model from relying too heavily on any single token-to-token relationship, acting as a form of regularization that improves generalization.
Notice that dropout is only applied during training. During inference, it is disabled so predictions remain deterministic.
Putting All the Heads Back Together
Each attention head produces an output of shape
(B, nh, T, hs)
Eventually we want
(B, T, C)
again.
First:
y.transpose(1, 2)
changes
(B, nh, T, hs)↓(B, T, nh, hs)
Now the heads sit beside each token.
Why .contiguous() Is Needed
The transpose only changes how PyTorch views the tensor. It does not rearrange the underlying memory.
After transposing, the desired logical order no longer matches the physical layout in memory.
Since
.view(...)
requires contiguous memory, calling it immediately would fail.
So GPT first does
.contiguous()
which copies the tensor into a new block of memory with the correct layout.
Only then can
.view(B, T, C)
flatten
(B, T, nh, hs)↓(B, T, C)
into the original embedding dimension.
An alternative is
.reshape(...)
which performs this check automatically and creates a copy if necessary.
One Final Linear Projection
After concatenating all attention heads, GPT applies one more linear layer:
y = self.resid_dropout(self.c_proj(y))
This does two things:
c_projmixes information across all attention heads with a final linear projection.resid_dropoutapplies dropout for additional regularization before the result is added back through the residual connection.
At this point, the self-attention block is complete.
The Feed-Forward Network (MLP)
Every Transformer block contains not only attention, but also a small neural network:
class MLP(nn.Module):
Its first layer is
self.c_fc = nn.Linear(...)
Here, fc stands for fully connected.
Unlike attention, which lets tokens exchange information, the MLP processes each token independently, adding nonlinear computation to each token’s representation.
Why Expand to 4× the Embedding Size?
The first linear layer expands
768 → 3072
Why not simply make the embedding dimension 3072 everywhere?
Because that would increase the cost of every operation:
- token embeddings
- Q, K, V projections
- attention computation
- output projections
- residual connections
Everything would become roughly four times larger and much more expensive.
Instead, the Transformer briefly expands each token into a higher-dimensional “workspace”:
768↓3072↓activation↓768
This temporary wide space gives the network more capacity to learn complex nonlinear transformations while keeping the rest of the model efficient.
Token Embeddings and Position Embeddings
The input to GPT starts as token IDs:
idx
with shape
(B, T)
These are simply integers representing vocabulary entries.
The word token embedding (wte) is a lookup table that maps each token ID to a learned vector of length n_embd.
Since Transformers have no inherent notion of sequence order, GPT also learns a position embedding (wpe), another lookup table that assigns each position in the sequence its own vector.
The actual input to the Transformer is the sum of the two:
input = token_embedding + position_embedding
Each token therefore knows both what it is and where it appears in the sequence.
LayerNorm: Normalizing Each Token
Before attention and the MLP, GPT applies Layer Normalization.
Unlike Batch Normalization, LayerNorm works within each token. For every token vector, it computes the mean and variance across the feature (embedding) dimension, then rescales the values to have a stable distribution.
This normalization makes training more stable and allows very deep Transformer stacks to converge effectively.
The Bigger Picture
By this point, a single Transformer block is no longer mysterious. Each input token goes through a predictable pipeline:
Token IDs ↓Word Embedding + Position Embedding ↓LayerNorm ↓Q, K, V Projection ↓Multi-Head Self-Attention ↓Concatenate Heads ↓Linear Projection ↓Residual Connection ↓LayerNorm ↓MLP (Expand → Nonlinearity → Compress) ↓Residual Connection
Stack dozens of these blocks together, and the model gradually transforms simple token embeddings into rich contextual representations capable of predicting the next token with remarkable accuracy.