mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
[ROCm] Update to AOTriton 0.8b (#140172)
Notable new features for SDPA operators on AMD systems from AOTriton 0.8b: 1. Nestedtensor support; 2. MQA/GQA support; 3. Restore Efficient attention support for causal=True and seqlen_q != seqlen_k cases; + The kernel should use top-left alignment, bottom right alignment will be added later 4. Move gfx1100 (RX7900/W7800/W7900) out of experimental support status. However, users are strongly recommended to update to ROCM 6.2.4, notably for its firmware updates. Related unit tests are enabled as well. Notable related changes from AOTriton 0.8b: 1. AOTriton 0.8b moves the GPU kernel out of libaotriton.so to a separate directory `aotriton.images`; 2. LZMA replaces ZSTD as GPU kernel compression algorithm for better compression ratio: aotriton0.8b (.so + aotriton.images take 350MB) compared to aotriton0.7b .so: 800MB 3. The compression cannot be disabled now, and `liblzma` is hard run-time dependency. + Should not be a problem, since `lzma` is part of Python Standard Library Pull Request resolved: https://github.com/pytorch/pytorch/pull/140172 Approved by: https://github.com/jithunnair-amd, https://github.com/jeffdaily Co-authored-by: Jithun Nair <37884920+jithunnair-amd@users.noreply.github.com>
This commit is contained in:
committed by
PyTorch MergeBot
parent
0a619a212f
commit
424156c26c
7
setup.py
7
setup.py
@ -1373,6 +1373,13 @@ def main():
|
||||
"lib/*.lib",
|
||||
]
|
||||
)
|
||||
aotriton_image_path = os.path.join(lib_path, "aotriton.images")
|
||||
aks2_files = []
|
||||
for root, dirs, files in os.walk(aotriton_image_path):
|
||||
subpath = os.path.relpath(root, start=aotriton_image_path)
|
||||
for fn in files:
|
||||
aks2_files.append(os.path.join("lib/aotriton.images", subpath, fn))
|
||||
torch_package_data += aks2_files
|
||||
if get_cmake_cache_vars()["USE_TENSORPIPE"]:
|
||||
torch_package_data.extend(
|
||||
[
|
||||
|
Reference in New Issue
Block a user