Kernel Storage V2 #50

xinyazhang · 2024-10-22T20:14:31Z

The Kernel Storage V2 achieves the following goals

Major changes:

Move GPU kernel image to separate files
- Now organized as aotriton.images/<vendor>-<arch>/<kernel_family>/<kernel_name>/FONLY__<functionals>___<GPU>.aks2'
  - <GPU> can be a family of GPUs, e.g., MI300X／MI300A／MI325X
- This enables per-architecture delivery
- No more linking errors when binaries are bloated.
Introduce AKS2 file format to compress GPU kernels with LZMA. Reduce the total package down to 200MB (MI200/300/Navi31)
- AKS2 means "Aotriton Kernel Storage version 2"
- LZMA is picked over Zstandard for much better compression ratio (~7% vs ~12%) and acceptable performance (<0.1s)
Look up kernel image relative to the .so file (achieved through dladdr)
Can only build the C++ part by setting cmake option AOTRITON_NOIMAGE_MODE to ON

Minor changes:

Release GIL in pybind11
aks2.py tool is added to create AKS2 file.
Update Triton compiler to Oct/23/2024
- This is to fix bug Assertion `(!LeaveBefore || Idx <= LeaveBefore) && "Interference"' failed
- See [AMDGPU] Assertion `(!LeaveBefore || Idx <= LeaveBefore) && "Interference"' failed llvm/llvm-project#109294 for more details

- ROCm/aotriton#50 - upstream repository versions used without modifications: aotriton: https://github.com/ROCm/aotriton.git 93299bd xinyazhang/storagev2 aotriton submodules thirdparty/triton https://github.com/ROCm/triton.git c69fa304e0ba07944376ff5e7da175f32c276d85 origin/aotriton/0.8 thirdparty/pypind11 https://github.com/pybind/pybind11.git tag v2.11.1 8a099e44b3d5f85b20f05828d919d2332a8de841 thirdparty/incbin https://github.com/graphitemaster/incbin.git 6e576cae5ab5810f25e2631f2e0b80cbe7dc8cb original files removed ====================== rm -rf aotriton/.gitmodules aotriton/.git aotriton/third_party/incbin/.git aotriton/third_party/triton/.git aotriton/third_party/pybind11/.git commands used to add all of the required files to rocmsdk aotrition =================================================================== git add -A git add -f third_party/pybind11/docs/Makefile git add -f third_party/triton/.github/workflows/llvm-build/almalinux.Dockerfile git add -f third_party/triton/.github/workflows/llvm-build/centos.Dockerfile git add -f third_party/triton/docs/getting-started/tutorials/grouped_vs_row_major_ordering.png git add -f third_party/triton/docs/getting-started/tutorials/parallel_reduction.png git add -f third_party/triton/docs/getting-started/tutorials/random_bits.png git add -f third_party/triton/python/triton/backends/__init__.py git add -f third_party/triton/python/triton/backends/compiler.py git add -f third_party/triton/python/triton/backends/driver.py Signed-off-by: Mika Laitio <lamikr@gmail.com>

…ing scheme

generate_compile writes necessary information of Kernel Storage V2 to Bare.cluster with this option.

1. Replace ZSTD with LZMA 2. Assume compression and remove code path without compression 3. Call v2python.archive_for_kernel_storage_v2 for the packing kernels.

Unable to test these code right now. Need to implement Tool v2python.archive_for_kernel_storage_v2 first.

This allows building C++ part without re-building GPU kernels.

liblzma is needed instead of libzstd right now. (Notably the header file, the library itself is shipped with Python)

…age V2.

jithunnair-amd · 2024-11-11T22:01:03Z

@ethanwee1 to help with testing of building AOTriton using this PR with and without AOTRITON_NOIMAGE_MODE=ON

New Triton needs this apparently.

jithunnair-amd · 2024-11-13T19:24:44Z

Changes for this PR tested using PR 51: #51 (comment)

ethanwee1 · 2024-11-13T23:25:24Z

Tested with these commands to build
aotriton-e278d4a853170c7a9063cfe847419414cb7b62b6-manylinux_2_28_x86_64-rocm6.2-shared.tar.gz

Commands:

git clone https://github.com/ROCm/aotriton.git
cd aotriton/
git checkout xinyazhang/manylinux_2_28-dockerfile
cd dockerfile/
export AMDGPU_INSTALLER=https://repo.radeon.com/amdgpu-install/6.2.4/el/8.10/amdgpu-install-6.2.60204-1.el8.noarch.rpm
mkdir -p output
cd input
vi install_aotriton.sh
added -DAOTRITON_NOIMAGE_MODE=ON

TRITON_LLVM_HASH="b5cc222d"  bash build.sh input tmpfs output e278d4a853170c7a9063cfe847419414cb7b62b6 "MI300X;MI200" 2>&1 | tee buildlog2.log
tar tvf output/*.tar*

Output:
Size: 1.6MB
aotriton-e278d4a853170c7a9063cfe847419414cb7b62b6-manylinux_2_28_x86_64-rocm6.2-shared.tar.gz
buildlogNOIMAGE.log

jithunnair-amd · 2024-11-14T17:46:55Z

Output: Size: 107MB aotriton-e278d4a853170c7a9063cfe847419414cb7b62b6-manylinux_2_28_x86_64-rocm6.2-shared_WITHAOTRITON_NOIMAGE_MODE_ON.txt buildlog2.log

@ethanwee1 This doesn't look right. The tarball size and txt file indicate that kernel images are still being built. Either your experiement was faulty or the AOTRITON_NOIMAGE_MODE logic is faulty. @xinyazhang, thoughts?

xinyazhang · 2024-11-14T18:00:57Z

Output: Size: 107MB aotriton-e278d4a853170c7a9063cfe847419414cb7b62b6-manylinux_2_28_x86_64-rocm6.2-shared_WITHAOTRITON_NOIMAGE_MODE_ON.txt buildlog2.log

@ethanwee1 This doesn't look right. The tarball size and txt file indicate that kernel images are still being built. Either your experiement was faulty or the AOTRITON_NOIMAGE_MODE logic is faulty. @xinyazhang, thoughts?

I think you need to download the log for review since it's rather large and likely get truncated if viewing with a browser. The trailing -- Install configuration: "Release" lines have indicated this project installed to the target directory.

The size is right. I only specified two architectures and 100MB is about the right size for 2 arches under Kernel Storage V2.

(However non-square causal masks added some more kernels and two arches ~= 130MB when releasing 0.8b)

xinyazhang · 2024-11-14T18:11:11Z

AOTRITON_NOIMAGE_MODE will be tested in another dockerfile related PRs

jithunnair-amd · 2024-11-14T18:22:18Z

Output: Size: 107MB aotriton-e278d4a853170c7a9063cfe847419414cb7b62b6-manylinux_2_28_x86_64-rocm6.2-shared_WITHAOTRITON_NOIMAGE_MODE_ON.txt buildlog2.log

@ethanwee1 This doesn't look right. The tarball size and txt file indicate that kernel images are still being built. Either your experiement was faulty or the AOTRITON_NOIMAGE_MODE logic is faulty. @xinyazhang, thoughts?

I think you need to download the log for review since it's rather large and likely get truncated if viewing with a browser. The trailing -- Install configuration: "Release" lines have indicated this project installed to the target directory.

The size is right. I only specified two architectures and 100MB is about the right size for 2 arches under Kernel Storage V2.

(However non-square causal masks added some more kernels and two arches ~= 130MB when releasing 0.8b)

Just to resolve the confusion here, I was referring to Ethan's initial numbers for the AOTRITON_NOIMAGE_MODE=ON build. He corrected his build steps thereafter and updated his comment to show that the new size is only 1.6MB, which is what I expected. So, in summary, looks like AOTRITON_NOIMAGE_MODE=ON is working as expected.

xinyazhang requested review from wangye805 and jithunnair-amd October 22, 2024 20:14

xinyazhang changed the title ~~[Queued PR] Kernel Storage V2~~ [Draft] Kernel Storage V2 Oct 22, 2024

xinyazhang changed the title ~~[Draft] Kernel Storage V2~~ [Queued PR] Kernel Storage V2 Oct 22, 2024

xinyazhang requested review from pruthvistony and jeffdaily October 22, 2024 20:45

xinyazhang force-pushed the xinyazhang/gqa branch from 80e621b to e16f0f0 Compare October 30, 2024 18:46

xinyazhang force-pushed the xinyazhang/storagev2 branch from 43f5f33 to 1ff178e Compare October 30, 2024 18:47

xinyazhang force-pushed the xinyazhang/gqa branch from e16f0f0 to 12373b7 Compare October 30, 2024 23:13

xinyazhang force-pushed the xinyazhang/storagev2 branch from 1ff178e to 801842e Compare October 30, 2024 23:13

xinyazhang force-pushed the xinyazhang/gqa branch from 12373b7 to 594af28 Compare October 30, 2024 23:27

xinyazhang force-pushed the xinyazhang/storagev2 branch from 801842e to 93299bd Compare October 30, 2024 23:27

xinyazhang added 16 commits November 7, 2024 18:08

generate_compile: Add --test_clustering to figure out optimal cluster…

eeaa826

…ing scheme

generate_compile: add option --generate_cluster_info

c0d29cb

generate_compile writes necessary information of Kernel Storage V2 to Bare.cluster with this option.

build: support kernel storage v2

41abdc5

1. Replace ZSTD with LZMA 2. Assume compression and remove code path without compression 3. Call v2python.archive_for_kernel_storage_v2 for the packing kernels.

Add packed kernel support and defines AKS2 format (Untested).

2a30ad4

Unable to test these code right now. Need to implement Tool v2python.archive_for_kernel_storage_v2 first.

Add the archiving tool aks2.py, also adjust the binary format

d3ff51b

Fix all compiling errors.

0bcd77e

Add some missing checks and cleans in packed_kernel.cc

81ae205

Install aotriton.images

0c9a53e

Fix in progress. python test/test_backward.py (Note NOT pytest) passed.

1d0eceb

Suppress debug output in Release build.

50ef132

Add cmake option AOTRITON_NOIMAGE_MODE

d07e9b3

This allows building C++ part without re-building GPU kernels.

Clang-format C++ Source/Header files

2f17ff6

clean up v2src/CMakeLists.txt

f10ed75

Clean up v2src/triton_kernel.cc

1b0a0c2

Copyright notices of aks2.py

406c64b

Fix the default of no image mode.

bfbdea9

xinyazhang added 3 commits November 7, 2024 18:08

Update Triton to Oct/23/2024

bac1a06

Make _common_test compatible with both PyTorch < 2.5 and >= 2.5

e992093

README: Update the Prerequisites

0dc5705

liblzma is needed instead of libzstd right now. (Notably the header file, the library itself is shipped with Python)

xinyazhang force-pushed the xinyazhang/storagev2 branch from 93299bd to 0dc5705 Compare November 7, 2024 18:08

xinyazhang changed the base branch from xinyazhang/gqa to main November 7, 2024 18:08

xinyazhang marked this pull request as ready for review November 7, 2024 18:10

xinyazhang changed the title ~~[Queued PR] Kernel Storage V2~~ Kernel Storage V2 Nov 7, 2024

xinyazhang mentioned this pull request Nov 7, 2024

Add docker based package builder and switch to system compiler #51

Merged

Remove AOTRITON_NO_SHARED option. It is incompatible with Kernel Stor…

4ef62e4

…age V2.

jithunnair-amd approved these changes Nov 11, 2024

View reviewed changes

xinyazhang added 2 commits November 12, 2024 21:16

Add pybind11 to requirements.txt

536aee4

New Triton needs this apparently.

add missing header <mutex> to packed_kernel.cc

b7c33cd

xinyazhang merged commit 0eb03a6 into main Nov 13, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Kernel Storage V2 #50

Kernel Storage V2 #50

xinyazhang commented Oct 22, 2024 •

edited

Loading

jithunnair-amd commented Nov 11, 2024

jithunnair-amd commented Nov 13, 2024

ethanwee1 commented Nov 13, 2024 •

edited

Loading

jithunnair-amd commented Nov 14, 2024

xinyazhang commented Nov 14, 2024 •

edited

Loading

xinyazhang commented Nov 14, 2024

jithunnair-amd commented Nov 14, 2024

Kernel Storage V2 #50

Kernel Storage V2 #50

Conversation

xinyazhang commented Oct 22, 2024 • edited Loading

jithunnair-amd commented Nov 11, 2024

jithunnair-amd commented Nov 13, 2024

ethanwee1 commented Nov 13, 2024 • edited Loading

jithunnair-amd commented Nov 14, 2024

xinyazhang commented Nov 14, 2024 • edited Loading

xinyazhang commented Nov 14, 2024

jithunnair-amd commented Nov 14, 2024

xinyazhang commented Oct 22, 2024 •

edited

Loading

ethanwee1 commented Nov 13, 2024 •

edited

Loading

xinyazhang commented Nov 14, 2024 •

edited

Loading