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