SIGN IN SIGN UP

🚨[wav2vec2] Support attn_implementation=sdpa dispatch (#46196)

* [SeamlessM4T/v2] Support attn_implementation=sdpa dispatch

* simplify code

* Init

* config

* prepare relative bias function

* TP friendly

* func comments

* parameter scale

* Fix multi-device error

* Code quality

* Refactor

* Apply a refactor to wav2vec2

* Fix copy consistency

* Doc

* Skip flash sdpa UT

* Update

* rename func

* scale->scaling

* Refine

* copy and code quality

* copy and ruff

* Fix CI

* Fix CI

* Refine data2vec_audio

* Refine

* refine

* Fix CI

* Fix `is_causal` SymBool leak in SDPA under torch.export

Gate the query-length check behind the always-concrete conditions so a
non-causal module or a provided mask short-circuits to a plain `False`
before touching `query.shape[2]`, which can be a data-dependent `SymBool`
under torch.export (e.g. sub-sampled audio in the Conformer encoder) that
SDPA rejects. When causality is still possible, `query.shape[2] > 1` stays
the first `and` operand so backed symints are still specialized to a bool.

* Fix SymBool error and conflicts

* comments

* refine

* ruff

* refine

* Add TODO

* refine TODO comments

* Skip dynamic export UT for SeamlessM4T/v2 speech models

---------

Co-authored-by: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com>
Y
YangKai0616 committed
2b296305d7bcf040831661b7b689e13a2e2e8d67
Parent: 9ad4b85
Committed by GitHub <noreply@github.com> on 8/20/2026, 10:40:30 AM