From Scratch · C++ · Transformers · Building Blocks

SingleHeadAttention

April 2026 · C++ / LibTorch · Transformers

Operations Pipeline

This class takes in token embeddings and returns attended representations. Here is every operation and its shape, top to bottom:

x (batch, S, n_embd) input embeddings | |-- Q(x) linear projection --> (batch, S, H) query vectors |-- K(x) linear projection --> (batch, S, H) key vectors |-- V(x) linear projection --> (batch, S, H) value vectors | |-- matmul(Q, K^T) --> (batch, S, S) raw attention scores |-- / sqrt(H) --> (batch, S, S) scaled scores |-- masked_fill (causal) --> (batch, S, S) future tokens set to -1e9 |-- softmax (row-wise) --> (batch, S, S) attention weights (rows sum to 1) |-- matmul(weights, V) --> (batch, S, H) weighted sum of value vectors | output (batch, S, H) attended representations

Where S = sequence length, H = head_size. Input dimension n_embd gets projected down to H. In a full Transformer, multiple heads run in parallel with smaller H, and their outputs are concatenated back to n_embd.

The rest of this page walks through each operation in the C++ code, focusing on what happens to the tensor shapes at every step.

The Code

Here is the complete class. Read it once to get the full picture, then we will break down each part below.

class SingleHeadAttention : public torch::nn::Module{ int n_embd, head_size; torch::nn::Linear K{nullptr}, Q{nullptr}, V{nullptr}; public: SingleHeadAttention(int n_embd, int head_size) : n_embd(n_embd), head_size(head_size) { K = register_module("K", torch::nn::Linear(torch::nn::LinearOptions(n_embd, head_size).bias(false))); Q = register_module("Q", torch::nn::Linear(torch::nn::LinearOptions(n_embd, head_size).bias(false))); V = register_module("V", torch::nn::Linear(torch::nn::LinearOptions(n_embd, head_size).bias(false))); } torch::Tensor forward(torch::Tensor x){ int seq_len = x.size(1); torch::Tensor output, k_vec, q_vec, v_vec, raw_attn, mask; q_vec = Q(x); k_vec = K(x); v_vec = V(x); raw_attn = torch::matmul(q_vec, k_vec.transpose(-2, -1)) / ((float) sqrt(head_size)); mask = torch::tril(torch::ones({seq_len, seq_len}, x.options())); raw_attn.masked_fill_(mask == 0, -1e9f); output = torch::softmax(raw_attn, -1); output = torch::matmul(output, v_vec); return output; } };

Constructor: Setting Up Layers

Three linear layers (Q, K, V), each mapping n_embd inputs to head_size outputs, no bias.

Layer	Weight shape	Parameters
`Q`	`(head_size, n_embd)`	head_size * n_embd
`K`	`(head_size, n_embd)`	head_size * n_embd
`V`	`(head_size, n_embd)`	head_size * n_embd

torch::nn::Linear K{nullptr} creates an empty placeholder. We cannot construct the layer in the class body because the dimensions (n_embd, head_size) come from constructor arguments. So we initialize as null first, then assign real layers inside the constructor body.

register_module("K", ...) does three things: (1) adds the layer's parameters to the module's parameter list so the optimizer can find them, (2) gives it a named key for checkpoint saving/loading, and (3) ensures it moves to GPU when you call model->to(torch::kCUDA).

LinearOptions(n_embd, head_size).bias(false) creates a pure matrix multiply with no bias term. This follows the original Transformer paper. The LinearOptions pattern is LibTorch's C++ equivalent of PyTorch's keyword arguments.

Under the hood, torch::nn::Linear(LinearOptions(n_embd, head_size).bias(false)) creates a weight tensor of shape (head_size, n_embd) with Kaiming uniform initialization. When called as Q(x), it computes x @ weight.T.

Q, K, V Projections

int seq_len = x.size(1); // extract S from dim 1 of (batch, S, n_embd) torch::Tensor output, k_vec, q_vec, v_vec, raw_attn, mask; q_vec = Q(x); // (batch, S, n_embd) @ (n_embd, H) = (batch, S, H) k_vec = K(x); // (batch, S, n_embd) @ (n_embd, H) = (batch, S, H) v_vec = V(x); // (batch, S, n_embd) @ (n_embd, H) = (batch, S, H)

Each linear layer multiplies the last dimension of x by its weight matrix. The operation is applied independently to every token in every sequence. It is a per-token projection from n_embd dimensions down to head_size dimensions.

(batch, S, n_embd) @ (n_embd, head_size) = (batch, S, head_size) Token 0 in batch 0 gets the same weight matrix as token 5 in batch 3. No cross-token interaction happens here, just a per-token linear transform.

Why transpose(-2, -1)

This is the single most important mechanical detail in the attention computation. If you understand why we transpose, you understand how attention scores are computed. Let us build the reasoning from first principles.

What We Are Trying to Compute

We want to compute a score between every pair of tokens. Specifically, for token i and token j, the attention score is the dot product of token i's query vector and token j's key vector:

score[i][j] = dot(Q_i, K_j) = sum over d of (Q_i[d] * K_j[d])

If we have S tokens, we need to compute S x S scores: every query dotted with every key. This is exactly the definition of matrix multiplication. If we stack all query vectors as rows of a matrix Q and all key vectors as columns of a matrix K^T, the matrix product Q @ K^T gives us all S x S scores in one operation.

Q = [ Q_0 ] K^T = [ K_0 | K_1 | K_2 ] [ Q_1 ] [ Q_2 ] Q @ K^T = [ dot(Q_0,K_0) dot(Q_0,K_1) dot(Q_0,K_2) ] [ dot(Q_1,K_0) dot(Q_1,K_1) dot(Q_1,K_2) ] [ dot(Q_2,K_0) dot(Q_2,K_1) dot(Q_2,K_2) ]

Row i, column j of the result is exactly the attention score from token i to token j. This is what we want. But there is a problem with shapes.

The Shape Problem

After the projections, q_vec and k_vec both have shape (batch, S, H), where S is the sequence length and H is head_size. Let us try to multiply them directly:

q_vec: (batch, S, H) k_vec: (batch, S, H) matmul(q_vec, k_vec): Last two dims: (S, H) @ (S, H) Inner dimensions: H vs S H != S -> DIMENSION MISMATCH. Cannot multiply.

Recall the fundamental rule of matrix multiplication: for 2D matrices, (m x n) @ (n x k) = (m x k). The inner dimensions must match: the last dimension of the first matrix must equal the first dimension of the second matrix.

Standard 2D matmul rule: A: (m x n) B: (n x k) Result: (m x k) ^ ^ these must match (the "inner dimensions")

Applying this to our shapes: for (S, H) @ (?, ?), the second matrix must have its first dimension equal to H. But k_vec's last two dimensions are (S, H), so the first of those is S, not H. The multiplication is impossible as-is.

We need the second matrix to have shape (H, S) so that the inner dimensions align: (S, H) @ (H, S). The inner H dimensions match, and the result is (S, S), which is exactly the S x S attention score matrix we want.

What transpose(-2, -1) Does

k_vec.transpose(-2, -1) swaps the last two dimensions of the tensor. Before and after:

Before: k_vec shape = (batch, S, H) After: k_vec^T shape = (batch, H, S)

Now the matmul works:

matmul(q_vec, k_vec.transpose(-2, -1)): q_vec: (batch, S, H) k_vec^T: (batch, H, S) Last two dims: (S, H) @ (H, S) Inner dimension: H == H -> MATCH Result: (batch, S, S)

The result is a (batch, S, S) tensor. For each sequence in the batch, we get an S x S matrix where entry [i][j] is the dot product of query i with key j. That is the raw attention score.

An important detail: .transpose(-2, -1) does not copy or rearrange any data in memory. It changes the stride metadata of the tensor. LibTorch tensors store data as a flat 1D array plus a shape and stride descriptor. Transposing simply swaps the stride values for the two dimensions, so the same memory is read in a different order. This makes the transpose operation nearly free, just a metadata update, not a data copy.

Why -2 and -1, Not 0 and 1?

Negative indices count from the end of the dimension list:

For a 3D tensor with dims [0, 1, 2]: -1 = dimension 2 (the last dim, H) -2 = dimension 1 (second-to-last, S) -3 = dimension 0 (third-to-last, batch) transpose(-2, -1) swaps dim 1 and dim 2 -> swaps S and H

The reason for using negative indices instead of .transpose(1, 2) is generality. If the tensor has additional batch dimensions (for example, in multi-head attention the shape might be (batch, n_heads, S, H), a 4D tensor), .transpose(-2, -1) still swaps the correct pair (S and H). Using .transpose(1, 2) on a 4D tensor would swap n_heads and S, which is wrong.

3D case: (batch, S, H) .transpose(-2, -1) -> swaps dim 1, dim 2 -> (batch, H, S) CORRECT .transpose(1, 2) -> swaps dim 1, dim 2 -> (batch, H, S) CORRECT 4D case: (batch, n_heads, S, H) .transpose(-2, -1) -> swaps dim 2, dim 3 -> (batch, n_heads, H, S) CORRECT .transpose(1, 2) -> swaps dim 1, dim 2 -> (batch, S, n_heads, H) WRONG

By always using -2 and -1, the code says "swap the last two dimensions, no matter how many dimensions exist." This is robust and future-proof.

Batched matmul

When torch::matmul receives tensors with more than 2 dimensions, it does not follow the simple "first dim of B matches last dim of A" rule from 2D linear algebra. Instead, it does batched matrix multiplication: all dimensions except the last two are treated as batch dimensions, and the actual multiplication only happens on the last two dimensions.

The matching rule for 3D+ tensors: PyTorch matches the last dimension of A with the second-to-last dimension of B. The batch dimensions are carried through unchanged.

2D matmul (standard linear algebra): A: (m, n) @ B: (n, k) = (m, k) ^ ^ last dim of A matches first dim of B 3D+ matmul (PyTorch batched rule): A: (batch, m, n) @ B: (batch, n, k) = (batch, m, k) ^ ^ last dim of A matches second-to-last dim of B batch dimensions carry through untouched

Applying this to attention:

q_vec: (batch, S, H) k_vec.transpose: (batch, H, S) PyTorch looks at the last two dims only: (S, H) @ (H, S) = (S, S) ^ ^ last dim of q_vec (H) matches second-to-last dim of k_vec^T (H) Batch dim just carries through: (batch, S, S)

Concretely, for each index in the batch dimension:

For batch index 0: q_vec[0] is (S, H), k_vec^T[0] is (H, S) -> result[0] = (S, S) For batch index 1: q_vec[1] is (S, H), k_vec^T[1] is (H, S) -> result[1] = (S, S) ... Final result: (batch, S, S)

Each sequence in the batch gets its own independent matrix multiplication. The batch dimension is never mixed: tokens in batch 0 never interact with tokens in batch 1. This is correct, because each sequence is an independent input and should be processed independently.

Efficiency note: Batched matmul is not a loop. It dispatches all the multiplications to GPU in a single kernel call, running them in parallel. Processing a batch of 32 sequences costs almost the same time as processing one, because the GPU has thousands of cores that execute the independent matrix multiplications simultaneously.

Interactive Animation: Transpose and Matmul

Step through the transpose and matmul operation visually. Watch how the dimensions align and produce the attention score matrix.

Transpose(-2, -1) and Matmul Walkthrough

Step 0 / 5

Step 0: q_vec has shape (2, 3, 4). Batch=2 sequences, S=3 tokens each, H=4 features per token.

To summarize the transpose: k_vec.transpose(-2, -1) swaps the sequence length and head_size dimensions. This turns each key row vector into a column, so that matrix multiplication produces one dot product per (query, key) pair. The batch dimension passes through untouched. The result is an (S x S) attention score matrix per sequence.

Scaling by sqrt(head_size)

raw_attn = torch::matmul(q_vec, k_vec.transpose(-2, -1)) / ((float) sqrt(head_size));

After the raw dot products, we divide by sqrt(head_size). Without this, the dot products grow with the dimension: if head_size = 64, typical scores range from -16 to +16. Large inputs push softmax into saturation (nearly one-hot), which kills gradient flow.

head_size = 64 => dot product std dev = sqrt(64) = 8 => scores in [-16, +16] head_size = 512 => dot product std dev = sqrt(512) = 22 => scores in [-45, +45] After dividing by sqrt(head_size): std dev = 1, scores in moderate range.

head_size is int. C++ sqrt on an int returns double. Dividing a float tensor by a double can cause type promotion issues. The (float) cast ensures the division stays in float32.

Causal Mask

mask = torch::tril(torch::ones({seq_len, seq_len}, x.options())); raw_attn.masked_fill_(mask == 0, -1e9f);

Prevents tokens from attending to future positions. torch::tril zeros out everything above the diagonal. masked_fill_ sets those positions to -1e9, which softmax converts to ~0 weight.

mask (4x4): After masked_fill_: [ 1 0 0 0 ] [ 0.5 -1e9 -1e9 -1e9 ] [ 1 1 0 0 ] [ 0.2 0.7 -1e9 -1e9 ] [ 1 1 1 0 ] [ 0.9 0.1 0.5 -1e9 ] [ 1 1 1 1 ] [ 0.4 0.8 0.2 0.7 ] Token i can attend to tokens 0..i only.

x.options() creates the mask on the same device/dtype as the input (avoids CPU vs GPU mismatch). The trailing underscore in masked_fill_ means in-place modification, avoiding a tensor copy. We use -1e9 instead of -INFINITY to avoid potential NaN issues in edge cases.

The mask is (S, S) but raw_attn is (batch, S, S). Broadcasting treats the mask as (1, S, S) and applies it identically across the batch.

Softmax and Weighted Sum

output = torch::softmax(raw_attn, -1); // (batch, S, S) - each row sums to 1 output = torch::matmul(output, v_vec); // (batch, S, S) @ (batch, S, H) = (batch, S, H)

softmax(..., -1) normalizes each row of the score matrix into a probability distribution. Masked positions (-1e9) become ~0 weight. Then the final matmul computes a weighted sum of value vectors for each token:

For token 2 with weights [0.47, 0.21, 0.32, 0.00]: output[2] = 0.47 * V_0 + 0.21 * V_1 + 0.32 * V_2 + 0.00 * V_3 Result shape: (batch, S, head_size)

The variable output is reused for both softmax and the final matmul. After the second assignment, the intermediate (batch, S, S) softmax tensor is released, which matters for long sequences where that matrix is large.

Complete Shape Trace

Here is every variable in the forward pass with its shape, traced from input to output. Use S for sequence length and H for head_size.

Line	Variable	Shape
`x` (input)	`x`	`(batch, S, n_embd)`
`q_vec = Q(x)`	`q_vec`	`(batch, S, head_size)`
`k_vec = K(x)`	`k_vec`	`(batch, S, head_size)`
`v_vec = V(x)`	`v_vec`	`(batch, S, head_size)`
`k_vec.transpose(-2, -1)`	(intermediate)	`(batch, head_size, S)`
`matmul(q_vec, k_vec^T)`	`raw_attn`	`(batch, S, S)`
`/ sqrt(head_size)`	`raw_attn`	`(batch, S, S)`
`torch::tril(torch::ones(...))`	`mask`	`(S, S)`
`raw_attn.masked_fill_(...)`	`raw_attn`	`(batch, S, S)`
`torch::softmax(raw_attn, -1)`	`output`	`(batch, S, S)`
`matmul(output, v_vec)`	`output`	`(batch, S, head_size)`

Input to output: The tensor enters as (batch, S, n_embd) and leaves as (batch, S, head_size). The sequence length S is preserved (every token still has its own representation). The embedding dimension changes from n_embd to head_size. This dimensionality reduction is intentional: in multi-head attention, each head works in a smaller subspace, and the outputs are concatenated to recover the full dimension.

Full Code

The complete SingleHeadAttention class, ready to compile with LibTorch:

class SingleHeadAttention : public torch::nn::Module{ /* The input embedding vector comes in (size n_embd) and the output vector is of size head_size */ int n_embd, head_size; torch::nn::Linear K{nullptr}, Q{nullptr}, V{nullptr}; public: SingleHeadAttention(int n_embd, int head_size) : n_embd(n_embd), head_size(head_size) { K = register_module( "K", torch::nn::Linear( torch::nn::LinearOptions(n_embd, head_size).bias(false) ) ); Q = register_module( "Q", torch::nn::Linear( torch::nn::LinearOptions(n_embd, head_size).bias(false) ) ); V = register_module( "V", torch::nn::Linear( torch::nn::LinearOptions(n_embd, head_size).bias(false) ) ); } torch::Tensor forward(torch::Tensor x){ int seq_len = x.size(1); torch::Tensor output, k_vec, q_vec, v_vec, raw_attn, mask; q_vec = Q(x); k_vec = K(x); v_vec = V(x); raw_attn = torch::matmul(q_vec, k_vec.transpose(-2, -1)) / ((float) sqrt(head_size)); mask = torch::tril(torch::ones({seq_len, seq_len}, x.options())); raw_attn.masked_fill_(mask == 0, -1e9f); output = torch::softmax(raw_attn, -1); output = torch::matmul(output, v_vec); return output; } };

Usage example

// Create a single-head attention with n_embd=64, head_size=16 auto head = std::make_shared<SingleHeadAttention>(64, 16); // Create a random input: batch=2, seq_len=8, n_embd=64 torch::Tensor x = torch::randn({2, 8, 64}); // Forward pass torch::Tensor out = head->forward(x); // out shape: (2, 8, 16) std::cout << out.sizes() << std::endl; // [2, 8, 16]

With n_embd=64 and head_size=16, each linear layer has a weight matrix of shape (16, 64) = 1024 parameters. Three layers (Q, K, V) give 3 * 1024 = 3072 total trainable parameters. No bias parameters since .bias(false) is set.

What Comes Next

A single attention head captures one type of relationship between tokens. But language has many simultaneous relationships: syntactic structure, coreference, semantic similarity, positional proximity, and more. A single head with a small head_size cannot capture all of these at once.

MultiHeadAttention solves this by running multiple SingleHeadAttention modules in parallel, each with its own Q, K, V weight matrices. Each head learns to focus on a different type of relationship. Their outputs are concatenated along the last dimension and projected back to the original embedding size:

n_heads = 4, head_size = 16, n_embd = 64 Head 0 output: (batch, S, 16) -- maybe learns syntactic dependencies Head 1 output: (batch, S, 16) -- maybe learns coreference Head 2 output: (batch, S, 16) -- maybe learns semantic similarity Head 3 output: (batch, S, 16) -- maybe learns positional patterns Concatenated: (batch, S, 64) -- all relationship types combined Final linear: (batch, S, 64) -- projected and mixed

The concatenation recovers the full embedding dimension: n_heads * head_size = 4 * 16 = 64 = n_embd. The final linear projection lets the model learn how to combine information from all the heads.

Key insight: The single-head attention you just built is the fundamental unit. Multi-head attention is just multiple instances of it running side by side. Understanding this one module means you understand the core of the Transformer.