PyTorch MPS causal attention leaks future tokens in half precision

PyTorch’s MPS backend can silently break causal masking in scaled dot-product attention when tensors are float16 or bfloat16, letting some query positions attend to future keys.

Dachuan Song of George Mason University reported that torch.nn.functional.scaled_dot_product_attention(..., is_causal=True) on MPS applies the causal mask only once per group of four query positions. Within each block, three positions may see keys past their own index. The op still returns finite, plausible-looking tensors with no warning, error, or NaN.

That matters for any autoregressive model that relies on causal attention on Apple Silicon: training or inference can incorporate information that should be invisible, changing outputs without an obvious failure signal. An explicit lower-triangular boolean attn_mask on the same device and dtypes does not show the same leakage and matches expected causal behavior within normal half-precision noise.

Song’s checks at sequence length 64 found large gaps versus a CPU float64 reference for the is_causal=True path, while the explicit-mask path stayed close. Affected cases aligned with the block-of-four pattern, including heavy future-token weight even at the first query position.

Users who need strict causality on MPS in half precision should prefer an explicit causal mask until the backend path is corrected, or run the attention in a higher precision or on another device.