A causal mask with a KV cache cannot always be built by drawing a lower triangle from query row zero. Queries may contain only the new suffix while keys contain the cached prefix and the suffix. The two axes have different origins. Compare their logical positions: a key is visible only if it is not later than the query, subject to other validity constraints.

A two-by-five example below turns a silent alignment error into distinct numerical outputs, with a runnable scalar reference. 阅读中文版.

Label the actual positions on both axes

Let P tokens already be cached and let the current input contain L new tokens, using contiguous zero-based logical positions. Query row i has position P+i; concatenated key column j has position j; the key length is S=P+L. In this standard self-attention example, visibility is j≤P+i.

With no cached prefix, P=0 and the condition reduces to j≤i. An upper-left lower triangle then works. Caching does not change the causal rule; it changes where the query axis begins. Confusing a local row index with a logical position drops that offset.

The Transformers caching explanation requires masks in iterative generation to cover past and current keys. The site’s KV cache introduction explains stored state; this article isolates which stored keys the current query may access.

For this example, with contiguous keys and queries corresponding exactly to the final L positions, the offset is also S−L. A truncated window, unused slots, or another position mapping can invalidate reconstruction from dimensions alone. Position comparison is the general check; geometric alignment is shorthand under explicit assumptions.

What a two-by-five mask should allow

For P=3 and L=2, the first query is at position 3. It may attend to keys 0 through 3 but not future key 4. The second query is at position 4 and may attend to all five keys. The allowed counts are four and five.

An ordinary upper-left lower triangle on a two-by-five matrix instead keeps only key 0 in the first row and keys 0 and 1 in the second. The shape remains valid and matrix multiplication may succeed, while required historical context is silently excluded.

A single query can be restricted incorrectly too

If L=1, S=P+1, and keys consist only of valid past and current tokens, the only query is at position P. Every key in this particular tensor satisfies the causal condition. Causality alone therefore excludes none of those columns.

This does not mean single-token decoding can always ignore masks. Unwritten static-cache slots, padding, sliding windows, and model-specific restrictions may still exclude keys. Establish what the key tensor actually contains before deciding whether additional conditions are necessary.

Choose values that expose the wrong mask

Set all five attention logits to zero and use values 10,20,30,40,50. Softmax gives equal weight to allowed positions. The correct rows produce 25, the mean of the first four values, and 30, the mean of all five. The incorrect triangle produces 10 and 15.

These inputs are deliberately constructed for debugging. Distinct values make omitted context observable. If every value were identical, correct and incorrect masks could accidentally produce the same output, allowing a weak test to pass.

# Pure-Python reference: True means allowed. No model or GPU is used.
import math

def causal_mask(query_positions, key_positions, key_valid=None):
    valid = [True]*len(key_positions) if key_valid is None else key_valid
    if len(valid) != len(key_positions):
        raise ValueError("key validity length mismatch")
    return [[bool(ok) and kp <= qp for kp,ok in zip(key_positions,valid)]
            for qp in query_positions]

def attend(logits, values, allowed):
    if not (len(logits)==len(values)==len(allowed)):
        raise ValueError("row length mismatch")
    kept = [(x,v) for x,v,ok in zip(logits,values,allowed) if ok]
    if not kept:
        raise ValueError("no valid key for this query")
    peak = max(x for x,_ in kept)
    weights = [math.exp(x-peak) for x,_ in kept]
    return sum(w*v for w,(_,v) in zip(weights,kept))/sum(weights)

keys = list(range(5))
correct = causal_mask([3,4],keys)
wrong = [[j <= i for j in keys] for i in range(2)]
assert correct == [[True,True,True,True,False], [True]*5]
values, logits = [10,20,30,40,50], [0.0]*5
assert [attend(logits,values,m) for m in correct] == [25.,30.]
assert [attend(logits,values,m) for m in wrong] == [10.,15.]
assert causal_mask([4],keys) == [[True]*5]
# Padding is independent of causality, even when its position is old.
assert causal_mask([4],keys,[False,True,True,True,True])[0][0] is False
try:
    attend([0.],[10.],[False])
except ValueError:
    pass
else:
    raise AssertionError("all-masked row was accepted")
print("offset, single-token, padding, and empty-row checks passed")

The code was executed locally for chunk offsets, single-query visibility, padding intersection, and an all-masked row. It is pure Python scalar logic, without PyTorch calls, model inference, or GPU measurements. The outputs are results of the specified teaching example, not a benchmark.

An API causal flag may use a different alignment

The PyTorch 2.14 SDPA reference specifies upper-left causal bias for a non-square is_causal mask. A parameter name does not establish that an operator knows a query came from the cache’s suffix. Confirm how the required position relation is represented for this example.

For this interface, an explicit boolean attn_mask uses True to permit attention. The documentation contrasts that with MultiheadAttention key_padding_mask, where the boolean meaning is reversed. Equal shapes and boolean types do not make the masks interchangeable. Print allowed positions in a small matrix before interpreting generated text.

No unexecuted library-call snippet is presented here. Integration requires checking the installed version, mask broadcasting, backend, and restrictions on combining an explicit mask with is_causal. A logical reference helps define acceptance but does not validate a particular kernel.

Keep padding validity separate from causality

Causality answers whether a key is in the future. Padding validity answers whether a column represents a real token. An old logical position does not make padding valid. The reference intersects position comparison with key_valid, preserving the distinction between the two conditions.

Batch members may have different prefix lengths. A batch-wide maximum cannot stand in for each sample’s true positions. Left padding, evicted window entries, and paged storage can make physical indices misleading. List the logical query position for each row and key position for each column when checking access.

If a query row is itself padding, its output or downstream loss also needs the model’s prescribed handling. Inserting an arbitrary visible key merely to avoid an error does not make the query valid. Error handling must not silently redefine the attention problem.

Setting a masked logit to zero does not exclude it

A zero logit still has positive softmax weight. Suppose two logits are zero, their values are 10 and 1000, and the second key is forbidden. The correct output is 10. Multiplying the second logit by zero leaves it at zero, assigning half the weight to each key and producing 505. This counterexample was checked separately with the same scalar function.

The reference excludes forbidden entries from normalization. Tensor implementations often express exclusion through an appropriate additive bias, subject to the interface’s numerical contract. A boolean visibility relation is not interchangeable with arbitrary elementwise multiplication, and a valid output shape does not establish that masking took effect.

Test an all-masked row explicitly

With no valid key, softmax does not describe an ordinary normalized distribution over allowed entries. Different numerical paths may produce NaN, zero, or a defined special-case result. None should be assumed universal. The teaching function raises an error to expose incomplete conditions.

A production system needs behavior defined by its model and interface contract. If such rows are allowed, specify and test the intended result. Replacing negative infinity with an arbitrary large negative constant does not automatically create the correct probability semantics for an all-masked row.

Compare full and chunked forward paths

Compute the same inputs in one full pass and in cached chunks, then compare the corresponding query outputs. Hold weights, positions, precision, and other conditions fixed. If dropout remains active in evaluation, establish its disabling rule first; randomness can obscure or imitate a mismatch.

Begin with P=0, then P>0 with L=1, then P>0 with L>1. Add padding, static unused slots, and any window restriction actually required by the model one at a time. A failure then identifies which additional condition changed the result.

Today’s RoPE verification article checks which positions transform the vectors; this article checks which keys participate. Both must agree, and either can fail independently. If these examples save you time debugging cached decoding, voluntary support below is welcome.

Sources and verification scope

The PyTorch 2.14 SDPA reference and Transformers caching explanation. Diagrams, scalar functions, and numerical examples are original teaching content. The checks validate mask reference logic, not a GPU kernel, model quality, or inference speed.