Fill the blank: for a single head, QK^T has shape (T_query, ___).

Fill in blank·Easy·4.0 · 0·~1 min·Asked atAi4bharat Autodesk Inflection Ai·Relevant atMicrosoft

Attempt it

For a single attention head with Q of shape (T_query, d_head) and K of shape (T_key, d_head), the product QK^T has shape (T_query, ).

TL;DR

QK^T has shape (T_query, T_key). Each entry is the dot-product similarity between one query token and one key token.

Memory aid

Easy to grasp

Imagine a spreadsheet where rows are people asking questions and columns are books that might answer them. Every cell holds a score for how well that book matches that question. If five people ask questions and there are twenty books, you get a five by twenty grid. That grid is QK^T. The row count comes from how many askers there are (queries), and the column count comes from how many books exist (keys). In self-attention the askers and the books are the same set, so the grid is square. In cross-attention they are different sets, so the grid is rectangular.

Key concepts

Tensor shapes Cross-attention

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 shape of QK^T is one of those mechanical facts every transformer engineer should be able to state without thinking, because almost every interesting consequence of attention's design follows from it. The (T_query, T_key) shape is the source of the quadratic memory cost, the canvas for causal and bidirectional masks, the structure that cross-attention uses to bridge encoder and decoder, and the tensor that FlashAttention painstakingly avoids ever materializing in HBM.

This deep dive walks the matmul shape rule, distinguishes self-attention from cross-attention in shape conventions, lifts the canonical attention formula, derives the memory footprint that motivates FlashAttention, and lays out where sparse attention variants diverge from the dense (T_query, T_key) baseline.

The goal is not just to remember 'the shape is (T_q, T_k)' but to understand why every downstream architectural choice in modern attention engineering traces back to that shape.

The matmul shape rule, applied

Matrix multiplication contracts (sums over) the shared inner dimension. The rule: (a, b) @ (b, c) -> (a, c).

Applying to QK^T

Q has shape (T_query, d_head).
K has shape (T_key, d_head).
K^T has shape (d_head, T_key).
Q @ K^T has shape (T_query, T_key).

The d_head dimension gets summed over by the matmul. It does NOT appear in the output. What remains is the two sequence-length axes: one for queries (rows), one for keys (columns).

Interpretation

Entry (i, j) of QK^T is q_i . k_j (the dot product of the i-th query vector and the j-th key vector). This is the unnormalized similarity score, what we then scale and softmax.

Canonical formula

\text{Attn}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

The softmax acts row-wise (across the T_key axis), producing a (T_query, T_key) matrix where each row is a probability distribution over keys. The final matmul with V of shape (T_key, d_head) contracts T_key, giving output (T_query, d_head).

Self-attention versus cross-attention

Where heads and batches sit

Why FlashAttention exists

Sparse attention: when (T_q, T_k) is the wrong baseline

Situations where this technique stops working.

2–4 min · Everything important, quickly.

Real products, models, and research that use this idea.

Llama 4 with H_q=64 heads at 128k context has per-layer score tensor of shape (B, 64, 131072, 131072), which is why FlashAttention is mandatory.
Encoder-decoder translation in NLLB-200 has cross-attention score matrices of shape (target_seq_len, source_seq_len) per layer.

What an interviewer would ask next. Try answering before peeking at the approach.

QIf the score matrix is (T_query, T_key), what is the shape of the attention output, and how do you get there?

Softmax along the T_key axis preserves the (T_query, T_key) shape. Multiplying by V of shape (T_key, d_head) contracts T_key, producing (T_query, d_head) output per head. Then concat across heads to (T_query, H * d_head) = (T_query, d_model) and apply W_O.

1 more follow-up an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Writing d_head as the second dimension. The d_head axis gets contracted by the matmul, what remains is the two sequence lengths.

60 second bullets to scan on the way to the call.

Matmul shape rule and which dimension gets contracted
Score matrix shape for self-attention versus cross-attention

Primary sources. Browse if you want the original framing.