Skip to content

fix(opt): preserve attention_mask=None in create_causal_mask for SDPA torch.compile - #48939

Open
ArjunPakhan wants to merge 8 commits into
huggingface:mainfrom
ArjunPakhan:fix/opt-sdpa-torch-compile
Open

ArjunPakhan wants to merge 8 commits into
huggingface:mainfrom
ArjunPakhan:fix/opt-sdpa-torch-compile

Conversation

@ArjunPakhan

@ArjunPakhan ArjunPakhan commented Sep 18, 2026

Copy link
Copy Markdown

CPU CI GPU run-slow

Fixes #48924

Summary

When attention_mask is None, OPT was instantiating an all-ones attention_mask tensor for positional embeddings and overwriting the variable before calling create_causal_mask. This caused create_causal_mask to receive a concrete 2D tensor mask instead of None, preventing SDPA from relying on is_causal=True and triggering Triton kernel fallbacks under torch.compile.

Solution

Decoupled positional embedding mask calculation into pos_attention_mask. This ensures attention_mask remains None when passed into create_causal_mask, preserving the fast-path SDPA tracing logic.

Testing

Ran local OPT unit test suite:
pytest tests/models/opt/test_modeling_opt.py -> 139 passed.

@github-actions

Copy link
Copy Markdown
Contributor

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

run-slow: qwen3_omni_moe

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 35385952536:2
Result: success | Jobs: 16 | Tests: 190,609 | Failures: 0 | Duration: 17h 8m

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

OPT training with torch.compile decomposes SDPA into Triton kernels due to attention_mask

1 participant