Write the causal mask for a decode step with a KV cache, where the query is shorter than the keys. Which corner does it align to?

TL;DR: The T new queries are the last T of S positions, so query row i sits at absolute position S - T + i and may attend to key j exactly when j <= i + (S - T). That is a lower-triangular mask aligned to the bottom-right corner. A top-left-aligned mask lets a single decode query see only key 0 out of 4,096, and PyTorch's is_causal=True is top-left aligned when T < S.

How to approach it

Start from positions, not from the shape of the matrix. Ask how many query tokens arrive per step: one for plain decode, a draft length plus one for speculative verification, a chunk size for chunked prefill. All three are the same problem with different T. Then state the rule in absolute positions and convert it into the row and column indices the kernel actually uses.

A strong answer

During decode the cache already holds keys and values for positions 0 to S - T - 1, and the step appends T new tokens. The score matrix is T by S. Causality says a token may attend to itself and everything before it, so the rule is a comparison between absolute positions:

The offset S - T is the whole answer. With T = S (prefill from empty) it is zero and the mask is the familiar tril. With T = 1 it is S - 1, so the single row is all ones: plain decode needs no mask at all, which is why decode kernels skip masking unless T > 1.

import numpy as np

def decode_causal_mask(T, S):
    """Bool [T, S], True = may attend. The T queries are the LAST T of S positions."""
    q_pos = np.arange(S - T, S)[:, None]   # absolute positions of the new tokens
    k_pos = np.arange(S)[None, :]
    return k_pos <= q_pos                  # bottom-right aligned

The test that matters: run full causal attention over all S tokens, then run the T new queries against the cache with this mask, and compare the last T rows. With S = 37 and T = 5 (deliberately not powers of two) the difference is exactly 0.0. Swapping in np.tril(np.ones((T, S))), the top-left alignment, gives an error of 2.37 on the same inputs, and nothing crashes. That silence is why the bug ships.

speculative verify, T = 4 queries over S = 4,096 keys:
  allowed(i, j) = j <= 4,092 + i
  masked entries = 3 + 2 + 1 + 0 = 6 of 4 x 4,096 = 16,384
top-left mask for the same step:
  row 0 keeps keys 0 only     -> 1 of 4,096 visible

The two conventions have names. FlashAttention aligns causal masks to the bottom-right when the lengths differ (changed in its 2.1 release). PyTorch's scaled_dot_product_attention(..., is_causal=True) aligns top-left: with one query over five keys it puts all weight on key 0 (checked on torch 2.14). PyTorch ships torch.nn.attention.bias.causal_lower_right for the decode case. An engine that swaps attention backends without checking this changes outputs for every speculative or chunked step and leaves T = 1 decode untouched, so a smoke test of plain generation passes.

Inside a tiled kernel the same rule decides three cases per (query tile, key tile) pair. If the key tile's first position exceeds the query tile's last absolute position, skip it before loading anything. If the key tile's last position is at or below the query tile's first, it is fully visible and needs no mask. Only tiles straddling the diagonal apply the comparison element by element. With the offset S - T, the diagonal moves to the right edge of the matrix, so for decode almost every tile is fully visible.

The reversal: if the cache uses a sliding window, the mask gains a lower bound as well, j > (S - T) + i - W, and the cache may be a ring buffer whose physical index is not the absolute position. Then the comparison must use stored positions, not slot indices.

What interviewers probe next

"Now it is tree-based speculative decoding with 8 candidate tokens." Causality is no longer a triangle among the new tokens. Each draft token attends to the full cache plus its own ancestors in the tree, so the T by T corner becomes an explicit ancestor mask built from parent pointers.

"Chunked prefill with the chunk in the middle of a 10K prompt?" Identical rule with T = chunk size and S = tokens cached so far plus the chunk. The offset handles it without special cases.

"Where does the mask live: a tensor, or the kernel?" In the kernel, as index arithmetic. A materialized T by S boolean tensor at S = 128K costs memory traffic for no information the kernel cannot compute from two integers.

How it's scored

Weak: writes tril(ones(T, S)), or says decode needs a mask without noticing that the single row is all ones.

Hire: derives the S - T offset from absolute positions and tests incremental output against full-sequence output.

Strong Hire: names the bottom-right versus top-left convention and the library default, explains why T = 1 hides the bug, and maps the rule to tile skipping in the kernel.

Common mistakes

• Testing only with T = 1, where both alignments produce an all-ones row and the bug cannot show.
• Building the mask from slot indices in a ring-buffer or paged cache instead of from the token positions stored with it.
• Materializing the mask as a dense tensor and passing it to a fused kernel, which can disable the fast path and costs T by S bytes per head.
• Masking with negative infinity and never handling a row that is masked everywhere (padding, or a window that excludes a whole tile), where the row max is negative infinity and the rescale computes NaN.