Each row of the post-softmax attention weight matrix corresponds to which side of the QK product?
Each row corresponds to one query token; entries along the row are that token's weights over all key tokens; row sums to 1 because softmax normalizes per row.
Picture a table where each row belongs to one person who is asking a question. That row lists how much attention this person pays to every other person in the room. The numbers in the row are slices of a pie that has to add up to a whole pie, because the person is dividing 100 percent of their attention across the room. Different rows belong to different askers, each splitting their own pie independently.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
The row/column convention of the attention matrix is one of those small details that everything else in attention rests on. Get it right and masking, visualization, sparse patterns, and FlashAttention all make sense in a single consistent picture. Get it backward and every subsequent operation seems to do the opposite of what it should.
The convention is: rows index queries (who is asking), columns index keys (who is being asked about), softmax runs along the row axis (the column dimension), and each row independently sums to exactly 1. This comes directly from the matrix arithmetic of Q K^T and the choice of softmax axis.
This deep dive walks through the matrix shape, the row-sum invariant, how masking preserves it, the degenerate case during autoregressive decoding, and the operational consequences for everything from attention visualization to FlashAttention's tiling strategy.
Where the row/column convention comes from
The attention score matrix is computed as scores = Q K^T / sqrt(d_k). The shapes determine the convention.
Shape arithmetic
Q has shape (seq_len_q, d_k). K transposed has shape (d_k, seq_len_k). Their product has shape (seq_len_q, seq_len_k). Standard linear algebra: the output rows correspond to the rows of the left operand (Q), and the output columns correspond to the columns of the right operand (K^T, which is the original rows of K).
So entry (i, j) in the score matrix is q_i dot k_j / sqrt(d_k), the alignment between query i and key j.
Softmax along the row axis
Softmax is applied along the last dimension, which is the key dimension. For each row i, softmax converts the seq_len_k scores into a probability distribution over keys. The row now represents 'how much attention does query i pay to each key j?'
Why this choice
The alternative (softmax along the column axis) would mean each key gets normalized across queries, which doesn't match the semantics of attention. Attention is a soft retrieval: each query asks 'where should I look?' and the answer is a distribution over keys. The row-axis softmax matches this semantic exactly.
The row-axis convention is a choice, but it's the only choice that makes attention semantically a per-query distribution. Every diagnostic, visualization, and kernel takes this convention as given.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- Attention visualization tools like BertViz and the HuggingFace inspect APIs all use the row=query, column=key convention.
- FlashAttention 2 and 3 tile over rows of Q in the outer loop, processing one query block at a time against streamed K blocks.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhat changes about the matrix during autoregressive decoding when T_q equals 1?
The matrix becomes a single row of shape (1, T_k_total), where T_k_total is the full context length stored in the KV cache. Softmax over that one row is still standard, and the row sums to 1 over all cached positions. This is why decode is memory-bandwidth bound: one query streams over the entire cache.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Swapping rows and columns. Rows index queries (who is asking) and columns index keys (who is being asked about). Flipping this breaks every subsequent operation.
60 second bullets to scan on the way to the call.
Shape of the attention score matrix as a function of Q and K shapes
Why softmax is applied along the key axis (the last dimension)
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.