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!

[mikekgfb]

Replacing MATH with FLASH_ATTENTION in https://github.com/pytorch/torchchat/blob/main/torchchat/export.py#L70 worked fine for me. (On a Mac, AOTI/DSO exporting to CPU.)

I'm actually not sure that the backend specification in https://github.com/pytorch/torchchat/blob/main/torchchat/export.py#L317 matters for ET export, because ET replaces all attention modules at https://github.com/pytorch/torchchat/blob/main/torchchat/export.py#L314 with the ET optimized CustomSDPAAttention (although I don't know if that in turn might call F.SDPA somewhere deep in the definition such that MATH vs FLASH_ATTENTION would matter?)

[mingfeima]

Oh I see. ET will replace it with CustomSDPAAttention, which goes into sdpa_with_kv_cache

torchchat/torchchat/model.py

Line 1028 in cbc72a4

from executorch.extension.llm.custom_ops import sdpa_with_kv_cache # no-qa

In that case, I think ET will overwrite the algo you defined from torchchat, as you can see algo is not passed down here
https://github.com/pytorch/torchchat/blob/main/torchchat/export.py#L213-L221

[Jack-Khuu]

Merged in #1456

Thanks everyone!! We'll add a matrix for tested use cases