Preserve sequence identity and global positions across tiles

For packed self-attention, attach a sequence identifier and an in-sequence position to every query and key. An allowed causal pair has matching sequence IDs and a key position no greater than the query position. A cumulative-length array can locate each sequence’s segment in storage, but the array’s global offset must not become the causal position by accident. Tiles should obtain their logical coordinates before comparing positions, including partial boundary tiles.

Padding is another exclusion rule rather than an ordinary token whose score happens to be small. A correct padded reference combines real-token validity with the causal relation. Packing should reproduce the same outputs on valid tokens after unpacking, subject to numerical tolerance. If a tile extends beyond a sequence’s end, its out-of-range lanes must not enter the softmax maximum, denominator, or gradient reductions. Checking only the final output shape cannot detect these internal leaks.

During cached decoding, a short query segment can represent the final positions of a longer key sequence. A bottom-right alignment then allows key index j up to i+Nk−Nq for zero-based query index i. That differs from a top-left triangular rule j≤i. Both rules are coherent for different coordinate interpretations, so an API flag named causal is insufficient to establish which logical positions were intended.

FlashAttention 2.8.3 documents bottom-right alignment for unequal query and key lengths, with zero output for all-masked rows. PyTorch 2.8 scaled_dot_product_attention documents upper-left alignment for its non-square is_causal flag. Construct an explicit mask or choose the correct documented bias when translating between interfaces; do not assume matching flags imply matching operators. Test both Nq<Nk and Nq>Nk, since the latter can create completely empty rows under bottom-right alignment.

Use short ragged sequences with deliberately different value ranges so cross-sequence leakage becomes obvious. Compare a packed kernel with separate per-sequence references, then compare gradients after gathering them back to the original storage order. Include lengths just below and above tile boundaries, empty batches if the API permits them, and noncontiguous layouts where supported. These inputs target indexing mechanisms more directly than a large random square matrix with no padding.

Backward recomputation must rebuild the same logical mask, including position offsets from a cache and any positional biases. A forward-only test can miss a gradient that writes into a neighboring packed sequence. Likewise, dropout replay needs stable logical indexing if it is required to match across layouts. Preserve cumulative lengths, position metadata, and mask conventions with the test case so a numerical disagreement can be distinguished from an intentional change of operator.