Use lower-right causal mask alignment consistently

[Anri-Lombard]

Summary

• Document that MLX's mask="causal" uses lower-right alignment
• Clarify the difference from PyTorch's default is_causal=True (upper-left)

When T_q != T_kv, this distinction matters:

• MLX (lower-right): Last query aligns with last key
• PyTorch default (upper-left): First query aligns with first key

References:

• PyTorch causal_lower_right
• PyTorch causal_upper_left

Relates to #2835

[zcbenz]

I don't think PyTorch has a causal_lower_right option for SDPA and the description is not really right.

[Anri-Lombard]

Hey @zcbenz, it does have causal_lower_right since 2.3 and can be used with SDPA via the attn_mask parameter. I ran a script with:

from torch.nn.attention.bias import causal_lower_right
bias = causal_lower_right(T_q, T_kv)
F.scaled_dot_product_attention(q, k, v, attn_mask=bias)

to verify.

Here is the tutorial that documents this explicitly: https://docs.pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html.

I also verified masks are mathematically identical. For example with T_q=2, T_kv=4:

  MLX's mask (using q_off = max(0, kL - qL)):
        k0  k1  k2  k3
  q0 [  1   1   1   0  ]
  q1 [  1   1   1   1  ]

  PyTorch's causal_lower_right(2, 4):
        k0  k1  k2  k3
  q0 [  1   1   1   0  ]
  q1 [  1   1   1   1  ]

  PyTorch's is_causal=True (upper_left):
        k0  k1  k2  k3
  q0 [  1   0   0   0  ]
  q1 [  1   1   0   0  ]
  

The first two are identical; the third is different. This is also consistent with MLX's CUDA backend which uses cuDNN's set_causal_mask_bottom_right.

Is there something specific about the description you think is incorrect? if your concern is that causal_lower_right isn't a direct SDPA parameter (like is_causal=True) but rather a separate utility class, I could clarify the wording to use the full module path torch.nn.attention.bias.causal_lower_right.

[zcbenz]

Thanks for linking the docs, this is a new learn for me. On the behavior, it actually depends on whether T_q is larger or smaller than T_kv:

mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp

Lines 204 to 208 in 9052f67

	if (q.shape(2) > k.shape(2)) {
	options.set_causal_mask(do_causal);
	} else {
	options.set_causal_mask_bottom_right(do_causal);
	}

[Anri-Lombard]

Thanks! Fixed to describe the conditional alignment behavior 🙏

[zcbenz]

Looks good to me. /cc @awni for a second look.

[awni]

The comment definitely makes sense. But I also find it a bit strange that we switch from lower right to upper left depending on if query is longer or shorter than the keys. It's quite rare for the query to be longer than the keys which is why we never really looked at it carefully.

I'm wondering if we should change the behavior in that case rather than documenting something that is a bit unusual? Or maybe it's a good idea to keep it this way?

[zcbenz]

I agree current behavior is unusual, and using lower right for all should be a better choice.

[awni]

@Anri-Lombard what do you think about changing the behavior to always be lower right even when QL > KL? Do you want to send a patch to this PR / send a new one instead of this?

[Anri-Lombard]

Hey @awni, always lower-right makes sense. The change is minimal (unless I'm missing somethign) - just two cuDNN locations (forward/backward) and the CPU fallback offset calculation. I'll update this PR to make the behavior change instead of just documenting it 👍

[awni]

Yes the change should be pretty straight-forward. We may also need to update the mask index calculation in the Metal kernels. If you add a test for this case as well (qL > kL) that would be great. I can help with the metal kernels if needed.

[Anri-Lombard]

@awni and @zcbenz updated and took a stab at the Metal kernels as well - feel free to push changes directly or point out where I deviated if you don't mind the extra time so I can learn the convention preferences more 🙏

For qL > kL, early queries have no keys to attend. Softmax of all-masked values gives uniform weights (exp(finite_min - finite_min) = 1), not zeros. Following PyTorch's pytorch/pytorch#108108 convention, we explicitly zero these rows... I think this is the only "big" change.

[Anri-Lombard]

@awni you mentioned the tests, the existing test shapes (127, 65, ...) with mask="causal" cover the qL > kL case. Would you prefer an explicit test that verifies early queries output zeros? 🙏

[awni]

Nope if it's already tested that is fine!

[awni]

I don't think we should ensure 0s in the qL > kL case. It's a problem I've looked at in the past is what to do if every key position for a given query is masked. And right now it's not consistent. For now let's leave it as undefined behavior and then look into a more principled fix if necessary. I also would rather not reduce performance overall to handle an edge case we don't really care much about).

[Anri-Lombard]

@awni done - removed the zero-row handling. The qL > kL + causal case is now undefined behavior as suggested. Tests skip that case.