SIGN IN SIGN UP

Fix sdpa flash attention op for et llama deployment (#4322)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/4322

We retropfitted flash attention cpu from aten. The retrofit we did was to
make it work to cacluate attention for a) batched prefill and b) decode with
different start_pos. For b, there was a bug when kv cache's seqlen dim is
split.
As a result attention calculation is not right. There is a detail in the code
to explain the issue.

bypass-github-export-checks
ghstack-source-id: 234634902

Reviewed By: larryliu0820

Differential Revision: D60011925

fbshipit-source-id: 50921846b329e449a4a767cf28c7a55d507217bd
K
Kimish Patel committed
6dbb4dcfa3476088d7b9b24a8ebe12ce5f7a0142
Parent: 9d85965
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 7/22/2024, 5:51:58 PM