[dynamo] [torch.compile] sdpa become slower after torch._dynamo.run in dynamo benchmark when attention_mask is None
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.6k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
`python benchmarks/dynamo/huggingface.py --performance --bfloat16 -dcuda -n3 --no-skip --dashboard --batch-size 2 --training --total-partitions 1 --partition-id 0 --only AlbertForQuestionAnswering --backend=eager --timeout=7200 --disable-cudagraphs --export-profiler-trace` I run this on 4080s for `backend=eager` or `indcutor`.
In the exported trace, I can see it calls `aten::_flash_attention_forward` in eager run but calls `aten::_efficient_attention_forward` in the dynamo run. `attention_mask` changes from `None` to a tensor of all `True`
This will hurt performance, i.e. The total device time of sdpa bwd changes from 3.6ms to 6.4ms, fwd changes from 1.4ms to 2.0ms. Why we set attention_mask even it's None?
sdpa in trace marked as expected:
sdpa in trace marked as actual:
### Versions
torch 2.12.0.dev20260317+cu128
cc @jerryzh168 @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @amjames @Lucaskabela @jataylo @drisspg @liangel-02 @howardzhang-cv
Contributor guide
Assessment
This issue has not been assessed yet.