A fast attention function is useful only if it lets each query read the right keys. A call can run without an error, return the expected shape, and still hide almost the entire prompt. Cached decoding is where this mistake is easy to miss.
Scaled dot-product attention, or SDPA, is PyTorch's function for the calculation from Day 14: score keys against queries, scale those scores, apply visibility rules, use softmax, and mix values. A backend is the implementation that performs that calculation on the available hardware. We choose the meaning of the calculation. PyTorch can select a compatible implementation.
We will use one small sequence with four known positions, then extend it by one position. Track the key columns each query may read throughout the lesson.
The visibility rule is part of the computation, not a speed setting. A faster backend cannot correct the wrong rule.
Import the Tensor and Attention Tools
A tensor is an array with named roles for its axes, such as batch, heads, positions, and features. We import PyTorch to create those arrays and its functional module to call attention.
import torch is needed even though the attention function lives under F. Importing the functional module alone does not define the name torch.
import torch
import torch.nn.functional as FThe examples use small CPU tensors and require PyTorch with SDPA support. They do not require a downloaded model or GPU. The values are teaching inputs, not trained language features.
SDPA does not create the learned query, key, and value projections for us. The surrounding attention layer must already have produced those tensors. Nor does this function include multi-head output projection, a residual addition, normalization, or the feed-forward network. Those are separate parts of the block from Day 19.
Q, K, and V enter after projection. SDPA returns the weighted values before the attention output projection.
Read the Shapes Before Writing the Call
For ordinary multi-head attention, a common layout is:
B counts sequences, H counts heads, and T_q and T_k count query and key positions. The last axes hold the features within a head. Queries and keys need the same feature width because their dot product multiplies matching entries. Keys and values need the same number of positions because a weight for key position 2 must select value position 2.
The value width, , may differ from the query/key width, . Therefore the output is , not always the query's exact shape. The function supports some broadcasting and grouped-head arrangements too; we keep these axes equal here so the first example has one clear interpretation.
Four query rows and four key columns create a four-by-four score grid for each head.
We create one sequence, two heads, four positions, and eight features per head. These small sizes make the axes easy to inspect. The seed makes the example inputs repeatable on the same setup; it does not promise identical floating-point results across all hardware.
torch.manual_seed(7)
query = torch.randn(1, 2, 4, 8)
key = torch.randn(1, 2, 4, 8)
value = torch.randn(1, 2, 4, 8)All three tensors share a shape because we deliberately chose equal query and value widths for this example. Their contents differ. A model would obtain them through three learned projections, not three random draws.
Full-Sequence Causal Attention
Number the four positions 0, 1, 2, and 3. During causal self-attention, query 0 may read key 0. Query 1 may read keys 0 and 1. Query 3 may read all four. The diagonal is allowed because the known token at position 3 is used to predict the token at position 4.
The following call creates that rule. We explicitly set dropout to zero because we are checking inference behavior.
full_output = F.scaled_dot_product_attention(
query, key, value, dropout_p=0.0, is_causal=True
)
assert full_output.shape == (1, 2, 4, 8)The output contains a vector for every query, for each head. It contains neither vocabulary probabilities nor generated token IDs. The rest of the Transformer must still combine heads and eventually apply its vocabulary output layer.
Each output row mixes only the values at its allowed key columns.
A shape assertion catches an axis mistake. It does not prove that the allowed columns were correct. For that, we need a test where changing a forbidden value cannot change the earlier output.
Why One Cached Query Is a Different Case
Now process position 4 after the four-position prefix. This forward pass makes one new Q, K, and V. After adding the new K and V to the prefix cache, the query length is 1 while the key length is 5.
The new query is at absolute position 4, meaning its position in the complete sequence. It is row 0 only within this small one-row tensor. These are different coordinates.
For a non-square attention grid, SDPA's is_causal=True uses an upper-left causal alignment. With one query row and five key columns, that mask allows only column 0. It does not automatically infer that this query belongs after a four-token cache. This behavior is specified in the PyTorch SDPA reference.
The newest query needs five allowed columns, not the single column allowed by a one-row upper-left triangle.
There are no future keys in a correctly sliced, fully filled single-step cache. In that specific case, every supplied key is valid, so we use is_causal=False. This does not remove causality from the model. The supplied keys already contain only the allowed past and current position.
We first construct the next position and append only its K and V. The -2 axis counts positions in our layout.
next_query = torch.randn(1, 2, 1, 8)
next_key = torch.randn(1, 2, 1, 8)
next_value = torch.randn(1, 2, 1, 8)
cached_key = torch.cat([key, next_key], dim=-2)
cached_value = torch.cat([value, next_value], dim=-2)The query remains one row. Concatenating queries would ask SDPA to recompute old output rows that this next-token step does not need.
next_output = F.scaled_dot_product_attention(
next_query, cached_key, cached_value,
dropout_p=0.0, is_causal=False
)
assert next_output.shape == (1, 2, 1, 8)Do not apply this shortcut to a cache with unused reserved slots, padded keys, or several new query positions. Those cases need a visibility mask or correctly limited inputs.
Several New Queries Need Absolute Positions
Suppose a forward pass processes positions 4 and 5 together after our four-token prefix. Query 4 may read key columns 0 through 4, but not column 5. Query 5 may read all six columns.
We build the rule from actual positions. [:, None] makes a column of query positions; [None, :] makes a row of key positions. Their comparison produces a two-by-six Boolean grid.
query_positions = torch.tensor([4, 5])
key_positions = torch.arange(6)
allowed = key_positions[None, :] <= query_positions[:, None]Its first row is [True, True, True, True, True, False]. Its second row is all True. For SDPA, a Boolean True means allowed, not blocked. Some other PyTorch attention interfaces use the opposite convention, so do not copy a mask between APIs without checking its meaning.
Compare key positions with query positions in the full sequence. Local tensor row numbers are not enough.
When passing this mask as attn_mask, leave is_causal=False. The custom mask already expresses the causal rule. The documented SDPA contract does not allow setting both an attention mask and is_causal=True.
A padding rule can be combined with allowed using a Boolean AND with a correctly shaped valid-key mask. That lets a key contribute only if it is both real and visible. Keep at least one valid key per real query; the mathematical softmax of an entirely blocked row is undefined.
Dropout Does Not Read Your Module's Mode
Dropout randomly removes contributions during training. It can help a model avoid relying too heavily on particular connections. During an inference check, it should be absent so we compare the intended attention calculation.
SDPA applies the dropout probability passed to the function. Calling a surrounding module's eval() does not override a nonzero dropout_p argument in that call. A real module's forward method must pass its training-dependent probability explicitly.
Visibility rules say what may be read. Dropout is a separate training operation on the resulting attention weights.
For these examples, dropout_p=0.0 is the correct choice. Do not add a model.training expression to a standalone snippet without first defining a model. That introduces a new object without helping us understand attention.
Backend Choice Comes After Correctness
A kernel is a low-level routine that performs an operation on a processor. SDPA may use a fused implementation that combines operations to reduce memory traffic, or a general math implementation. The available choices depend on the PyTorch build, hardware, dtype, shapes, masks, and other input properties.
This CPU example makes no claim that a FlashAttention kernel ran. A call succeeding proves that some compatible path worked, not that it was the fastest one. Different paths can also produce small rounding differences because floating-point addition depends on operation order.
Fix the inputs and mask first. Then inspect the actual execution path before making a speed claim.
A useful performance comparison keeps the model, device, dtype, prompt lengths, masks, and dropout unchanged. GPU timing also needs warm-up and synchronization because GPU operations may finish after the CPU has moved on. Otherwise a timing result can measure task submission instead of completed work.
Check Cached Output Against a Full Pass
We can test the cached step without training a model. Append the new query to the previous four queries, run full causal attention over all five positions, then compare its final row with our cached result.
The earlier four K/V rows are identical in both calculations. The last query sees all five keys in both. Their last-row outputs should therefore agree within a small numerical tolerance.
all_queries = torch.cat([query, next_query], dim=-2)
reference = F.scaled_dot_product_attention(
all_queries, cached_key, cached_value,
dropout_p=0.0, is_causal=True
)
torch.testing.assert_close(
next_output, reference[:, :, -1:, :],
rtol=1e-5, atol=1e-6
)The tolerances allow small float32 rounding differences. They are suitable for this small CPU example, not a universal threshold for every dtype. Lower-precision checks need tolerances chosen for their expected error.
Both paths use the same last query, five keys, five values, and five allowed positions.
Try the cached call again with is_causal=True. It normally disagrees with the reference because it reads only the first value. That is an information-flow bug, even though the returned tensor still has the expected shape.
Check Your Understanding
Why does a query of length 1 not mean absolute position 0? Because tensor length describes the current slice, not where that slice begins in the complete sequence.
Why can a value head have a different feature width? Scores come from Q and K. Each score then weights a V row, whose width determines the output width.
Why is is_causal=False correct for our single-step cache but wrong for a full training sequence? The cache supplied no future keys. The full training sequence does.
The two calls preserve the same causal dependency with different tensor shapes and mask settings.
Recap
An SDPA call needs compatible axes, the intended key visibility, and an explicit dropout setting. Full-sequence causal attention can use the built-in triangle. Cached attention must account for the query's position in the complete sequence.
On Day 26, we keep this calculation but change how many query heads share each key/value head. That reduces cache storage without reducing the number of query heads.
