pytorch / pytorch/pytorch

[dynamo] [torch.compile] sdpa become slower after torch._dynamo.run in dynamo benchmark when attention_mask is None

Open
#177,712 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: dynamo module: performance module: sdpa oncall: pt2 triaged
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:
Image

sdpa in trace marked as actual:
Image

### 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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.