Why Torchchat uses MATH as SDPA backend?

[yanbing-j]

🐛 Describe the bug

Hi maintainers,

I find that, Torchchat uses MATH as SDPA backend in https://github.com/pytorch/torchchat/blob/main/torchchat/generate.py#L542. However, for other libs like vllm, they all accept flash attention as default backend.

So why Torchchat uses MATH as a default backend? Is this required for accuracy? If not, I can help to add an argument to let user set the backend. Thanks!

Versions

[lucylq]

I can help to add an argument to let user set the backend.

This seems like a good idea! cc @Jack-Khuu if there's any history behind using MATH as default backend?

[lucylq]

Hmn, actually, it seems like there's an issue exporting when the default backend is not MATH.

See issue: pytorch/pytorch#129418

It seems like there's a requirement that decompositions during export must not introduce any mutation ops. SDPBackend.MATH is known to work well with export; with other backends, we may run into issues. cc @angelayi

[yanbing-j]

cc @mingfeima

[mingfeima]

Hmn, actually, it seems like there's an issue exporting when the default backend is not MATH.

See issue: pytorch/pytorch#129418

It seems like there's a requirement that decompositions during export must not introduce any mutation ops. SDPBackend.MATH is known to work well with export; with other backends, we may run into issues. cc @angelayi

I think we may add an argument to let user to decide which backend to use, say "--attention-backend"; so when users are not depending on export can have better performance with flash attention.

[yanbing-j]

I draft a PR #1456 to add an argument attention_backend. The default value of this argument is math. Please take a look, thanks!