[davrens101]

Fixes #15195.

What was wrong

_native_flex_attention selects the mask first and the causality second:

if attn_mask is None or isinstance(attn_mask, flex_attention.BlockMask):
    block_mask = attn_mask
elif is_causal:
    block_mask = flex_attention.create_block_mask(...)

When attn_mask is None the first branch matches, so block_mask stays None and is_causal=True is silently dropped. The backend then computes full (bidirectional) attention instead of causal attention; no error is raised because flex_attention is called without is_causal. Every other backend honours is_causal, and _check_attn_mask_or_causal explicitly allows attn_mask=None, is_causal=True.

Reproduced on a CPU torch 2.15 run:

dispatch_attention_fn(q, k, v, attn_mask=None, is_causal=True, backend="flex")
max abs diff vs SDPA(is_causal=True): 2.57   # completely wrong

Fix

Build the causal block mask when no mask is passed, before falling back to the BlockMask/tensor branches. attn_mask + is_causal is already rejected by the backend constraint, so the separate elif is_causal branch became dead code and is removed.

After the fix the same call matches SDPA to ~3e-7.

Tests

Added TestFlexAttentionCausal in tests/models/test_attention_dispatch.py: causal output vs F.scaled_dot_product_attention(..., is_causal=True) for fp32/bf16, and vs an explicit lower-triangular mask.

python -m pytest tests/models/test_attention_dispatch.py -k Flex -> 3 passed
python -m ruff check/format on both files -> clean

AI assistance was used for this change; I reviewed and ran the tests above.

Coordination status

The linked issue includes a fresh upstream CPU reproduction. Maintainer acknowledgment of scope is still pending; no acknowledgment or compiled/GPU verification is claimed.

[github-actions]

Hi @davrens101, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.