SDPA (#1540)
Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/1540
Added,
* Support for basic partitioner - just graph matching based
* Added basic AoT support to recompose sdpa and lower
* Serialization support
* Runtime logic to map to xnnpack node
* A few possible e2e test somewhat real
* Llama2 e2e test (just the core model)
TODOs:
* ConvertToSDPAPass asserts but it shouldn't
* Dynamic shape support
* {3,4}D mask tensor (see tests)
* Doesn't support SDPA without mask (see tests)
* MHA, MQA tests
Reviewed By: kimishpatel, mcr229, larryliu0820
Differential Revision: D52531038
fbshipit-source-id: 75290e3633dcba6e954037200e16e16043b840fa D
Digant Desai committed
fa50ded6707848ca41388828463ade5528b55272
Parent: 62c6a43
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 1/19/2024, 5:00:30 AM