Fix mask creation not being skipped under torch.compile

[jiqing-feng]

Fixes #48924
Fixes #48925

Problem

create_causal_mask and create_bidirectional_mask return None when SDPA can express the mask
through its own is_causal argument. This lets SDPA dispatch to flash attention. Passing a dense
attn_mask instead selects a slower path, which materializes a [batch, 1, q_len, kv_len] tensor on
every forward.

Two independent bugs prevented the skip:

1. OPTDecoder.forward overwrote its own attention_mask (OPT training with torch.compile decomposes SDPA into Triton kernels due to attention_mask #48924). When the caller passed None,
   OPT replaced it with a dense all-ones tensor and forwarded that into create_causal_mask. The
   tensor is only needed to derive position_ids.
2. masking_utils bailed out of every skip condition under is_tracing(padding_mask) (4D mask created under torch.compile when padding_mask is None #48925).
   That guard also covers torch.compile, so a compiled model never skipped mask creation. Yet the
   unpadded case reads no tensor values: the presence of a padding mask, q_length, kv_length and
   q_offset are static properties that dynamo already guards on. Only fast_all(padding_mask) is
   data-dependent.

Fix

1. OPTDecoder keeps the dense tensor in a local position_attention_mask. The caller's
   attention_mask now reaches create_causal_mask unchanged.
2. _ignore_causal_mask_sdpa, _ignore_bidirectional_mask_sdpa and _can_skip_bidirectional_mask_xpu
   are reordered to: export guard, static conditions, tracing guard, value-dependent conditions.
   torch.export still bails out first, since export hard-codes is_causal into the exported program
   (pytorch#108108).

Both fixes are required. Spying on create_causal_mask inside OPT without a padding mask:

	eager	torch.compile
main	skipped	materialized
fix 1 only	skipped	materialized
fix 2 only	skipped	materialized
both	skipped	skipped

Fix 2 takes effect on torch >= 2.14. Dynamo hard-coded torch.compiler.is_exporting() to True
until pytorch#176499, so the export guard also
triggered under torch.compile. On older versions the mask is still materialized, which is the
current behavior. The new compile test is therefore gated with @require_torch_greater_or_equal("2.14").

Result

import time

import torch
from transformers import AutoConfig, OPTForCausalLM

config = AutoConfig.from_pretrained("facebook/opt-125m", attn_implementation="sdpa")
model = OPTForCausalLM(config).cuda()
compiled = torch.compile(model, dynamic=False)
input_ids = torch.randint(0, config.vocab_size, (8, 1024), device="cuda")
padded = torch.ones_like(input_ids)
padded[0, :17] = 0

# Skipping the mask must not change the outputs, with or without padding.
model.eval()
for mask in (None, padded):
    with torch.no_grad():
        expected = model(input_ids, attention_mask=mask).logits
        actual = compiled(input_ids, attention_mask=mask).logits
    torch.testing.assert_close(actual, expected)
print("outputs match eager")

model.to(torch.bfloat16).train()

def step():
    compiled(input_ids, labels=input_ids).loss.backward()
    model.zero_grad(set_to_none=True)

for _ in range(5):
    step()
torch.cuda.synchronize()

timings = []
for _ in range(3):
    start = time.perf_counter()
    for _ in range(20):
        step()
    torch.cuda.synchronize()
    timings.append((time.perf_counter() - start) / 20 * 1000)
print(f"opt-125m fwd+bwd: {min(timings):.2f} ms")

On one A100 80GB with torch 2.14.0: 63.78 ms -> 49.17 ms, a 1.30x speedup. The assertion passes
on both revisions, so the skip does not change the results. Padded masks are still materialized.

Dynamo reports the same two graph breaks before and after, both from unrelated data-dependent
branching.

Tests

test_attention_mask_is_not_overwritten_for_causal_mask (OPT) and
test_mask_skip_without_padding_mask_under_compile (masking utils) are new. Both fail on main.

Full runs on A100 with torch 2.14: tests/models/opt 150 passed, tests/models/{llama,bert,gemma2}
593 passed, tests/utils/test_masking_utils.py all passed except the pre-existing
test_packed_sequence_mask_flex_attention failure.

The bidirectional helpers now return a Python bool via bool(fast_all(padding_mask)) instead of a
0-dim tensor, which matches their -> bool annotation. No replacement guard was added for the
torch.jit.is_tracing() coverage that is_tracing() provided, as torch.jit is deprecated and
is_jit_tracing has no other call site.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: opt

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 35566453195:2
Result: failure | Jobs: 16 | Tests: 190,606 | Failures: 0 | Duration: 15h 11m