Xinya Zhang 5bddbed399 Initial Flash Attention support on ROCM (#114309) ... This pull requests add initial Flash Attention support for AMD/ROCM platform. It added a specialized Triton repository/branch as a compile-time dependency for Flash Attention math library on AMD/ROCM. This triton submodule is not used at runtime and will not be shipped to the final pytorch package. We have the plan to release this specialized Triton as a separate project. Know limitations: - [ ] Only supports MI200 series GPU (i.e., `gcnArchName == gfx90a:sramecc+:xnack-`. - [ ] Only supports power of two sequence lengths. - [ ] No support for varlen APIs. - [ ] Only support head dimension 16,32,64,128. - [ ] Performance is still being optimized. Fixes https://github.com/pytorch/pytorch/issues/112997 Pull Request resolved: https://github.com/pytorch/pytorch/pull/114309 Approved by: https://github.com/jeffdaily, https://github.com/malfet --------- Co-authored-by: Joseph Groenenboom <joseph.groenenboom@amd.com>		2023-12-14 08:52:57 -08:00
..
_internal	Initial Flash Attention support on ROCM (#114309)	2023-12-14 08:52:57 -08:00
__init__.py
_comparison.py	[BE]: Update lintrunner mypy to 1.6.0 (#111375)	2023-10-17 01:22:06 +00:00
_creation.py	Add meta func for scaled mm (#112609)	2023-11-03 03:44:22 +00:00