Streamk v0.2 #646

xiaohuguo2023 · 2024-09-25T16:17:49Z

streamk v0.2:

new streamk tuning script to reduce compiling and profiling time
use load/store cache modifier to reimplement spinning lock
add CI test for streamk-kernel
able to use streampipelineV2

Add Perf Kernels This is a combination of 2 commits. Add Perf Kernels Add Perf Kernels This is a combination of 6 commits. add perf-kernels fix formating issues fix unused variables and other bugs fix other issues remove scripts save check changes format save save try pre-commit check save

Change all block pointers to tensor pointers Block pointers are for nvidia TMAs. They are useful for regular loads as well but not well supported. Also cleaned up some code I came across along the way and updated comment at the top.

Add support for layouts commonly used by users. Add option for varlen / thd layout to specify equal context lengths for all batches. Also often used by users.

* remove on push for Integration Tests * rename * add post merge test * save * dtype params * skip bad config * fix more stuff

Increase CI timeout

Couple of FA optimizations Set SM scale multiplication to a constexpr. Minor asm improvement. Changed acc scaling to adjust for softmax division to multiplication with reciprocal. ~10% perf improvement. --------- Co-authored-by: Michael Melesse <[email protected]>

* streamk v0.1 * remove unused variable * fix format issues * add README * fix format issue * change num_sms to num_cus

* Add explicit multiply-reduce GEMM kernel * Remove `SPLIT_K` argument from kernel * Remove `GROUP_SIZE_M` argument from kernel * Remove conditional call to `tl.dot` from kernel * Remove table with performance data from README

* Copy *tune_gemm* from `triton-mlir` branch to `main_perf` branch The source commit in `triton-mlir` branch is the following one: ``` commit cf44637 Author: Lixun Zhang <[email protected]> Date: Tue Jul 23 14:22:01 2024 -0500 [tuning] gemm tuning script v3.3 (#606) ``` *tune_gemm* was copied from the source branch directory `scripts/amd/gemm` to the destination branch directory `python/perf-kernels/tune_gemm`. The SHA-256 hashes of *tune_gemm* files are the following ones: ``` 423aef1deb6c60f6578a1ecfc94d2473f8746b00d0368c553d31641fcfa5e354 README.md 46ab93978fee33f75df23332f12546dae7910478c391f08b7b1ebd415d8266b7 icache_flush.py f18711544641b810a652e6a6629bfa2b613f6ade87399e88fdf05b81d4af58a4 matmul.py 84a1c80ede36d3154e51188276eda2d2d0f52ed4f496ff69349c390d83b8ec10 matmul_kernel.py 2812b40183637bc8d7e47d283c7d66b1792134a43de76f3eacf7b9b3e1c2431a one_config.py 0ac09c33b0173cea06ddabbf9f4e3afa1816781dea4fdcce5894a7e7d6a80e19 rocprof_gemm.py 00eff41cf1c0bfc41d623e42b51706af67639fec76146741e2067d2a93e0148a utils/file_generator.py cb7afb773ccee835b00396cccf87e0d44fe513131161f031fae42453725b3c82 utils/utils.py 59f23811b660e49e566927853926a21f02a7014bb19c8ea67e6b382db6c59900 tune_gemm.py e787f35d750b869f113b3c01692f64243a9cb8a71a18ade2f0465f614f7284e4 tune_gemm.sh ``` The files were kept as-is despite `pre-commit` intentions to change them. After that, *tune_gemm* directory in code and documentation was fixed to reflect it's new location.

* Reformat *tune_gemm* files with Triton's pre-commit The following command was executed to reformat the files: ``` $ pre-commit run --files \ python/perf-kernels/tune_gemm/* \ python/perf-kernels/tune_gemm/utils/* ``` * Fix *tune_gemm* issue with (1, 1) bias tensors * Fix `ruff` F405 errors Fix the following linter error: F405 `identifier` may be undefined, or defined from star imports * Fix `ruff` F841 errors Fix the following linter error: F841 Local variable `identifier` is assigned to but never used * Fix minor issues in README file * Add `--` to `num_threads` argument. * Replace `--icahe` argument (non-existent argument) with `--icache_flush` (existent argument). * Remove old files from *tune_gemm* V1 * Add dependency graph to README file * Selectively disable `yapf` for parts of `one_config.py`

…ofv1 (#630) * Change to rocprofv1 * improve post processing of rocprof results - set --iters=200 as default. This is enough since the time is stable after the first few runs. - Filter out kernel time that is too large. We use the first kernel time as the threshold. There must be something wrong with the kernel if its elapsedTime is larger than the first run. We need to investigate the reason. For now, just filter them out. * Add xcd-based pid remapping * Enable EVEN_K=false for large gemms * Update readme

Softmax kernel

* Move utility tools from triton-mlir to main_perf branch - Plot layout script - occ.sh - amdgcn-cfg * yapf format * More formats * remove executablility of plot_layout.py * Address ruff complains * Move tune_gemm to tools

Add rmsnorm kernel

Online softmax implementation

Additionally changes locks to use uint8 instead of int32 for smaller space footprint.

Add Layernorm kernel

neoblizz

I put in some comments.

neoblizz · 2024-09-26T21:18:40Z

python/perf-kernels/streamk/streamk_kernel.py

Note, for gfx90a the load/stores with cache_modifiers do not work. Documented here: https://github.com/ROCm/triton-internal/issues/311

neoblizz · 2024-09-26T21:19:07Z

python/perf-kernels/streamk/streamk_kernel.py

+            rm1 = tl.max_contiguous(tl.multiple_of(rm1, BLOCK_SIZE_M), BLOCK_SIZE_M)
+            rn1 = tl.max_contiguous(tl.multiple_of(rn1, BLOCK_SIZE_N), BLOCK_SIZE_N)
+            P_ = P + pid * BLOCK_SIZE_M * BLOCK_SIZE_N + rm1[:, None] * BLOCK_SIZE_N + rn1[None, :]
+            tl.store(P_, acc, cache_modifier=".wt")


Note, for gfx90a the load/stores with cache_modifiers do not work. Documented here: https://github.com/ROCm/triton-internal/issues/311

neoblizz · 2024-09-26T21:19:27Z

python/perf-kernels/streamk/streamk_kernel.py

-                # todo: try use tl.load once cache modifier landed upstream
-                while tl.atomic_cas(locks + next_pid, 1, 1) != 1:
+            while (end < tile_iter_end and next_pid < NUM_SMS):
+                while tl.load(locks + next_pid, cache_modifier=".cv", volatile=True) != 1:


This also does not work in gfx90a: https://github.com/ROCm/triton-internal/issues/311

neoblizz · 2024-09-26T21:20:02Z

python/perf-kernels/streamk/streamk_kernel.py

    EVEN_K: tl.constexpr,
 ):
    pid = tl.program_id(0)
-    pid = get_new_pid(pid, num_cus)
+    pid = (pid % 8) * (NUM_SMS // 8) + (pid // 8)


This is not needed for anything but gfx942, so we will actually remove this if the arch was gfx90a.

neoblizz · 2024-09-26T21:20:37Z

python/perf-kernels/streamk/tune_streamk.py

-    P = torch.zeros((num_cus, block_m * block_n), device="cuda", dtype=torch.float32)
-    triton_output = matmul(a, b, c, P, locks, num_cus, block_m, block_n, block_k, group_m, num_warps, num_stages,
-                           waves_per_eu, mfmaInstrSize, kpack, EVEN_K)
+    locks = torch.zeros((num_sms, ), device="cuda", dtype=torch.int32)


locks can be less than int32 type, we only need 1 byte: uint8 should work.

* Added regression tests to tune_gemm * Add regression tests to pipelines * Add missing imports * Use warnings to signal that no performance comparison is found * Split regression tests into separate file * Disable github pipeline in favour of jenkins * Improve output and skip tests if no performance reference can be found * Add testcase for overall mean regression * Extend parameters which can be adjusted for perf regression tests * Switch to geo mean for overall result * Always recompile kernels in perf regression tests in case the user does no specify otherwise * Report default values in exported result to support changing them in the future

xiaohuguo2023 · 2024-10-29T22:18:38Z

let's close this PR as there are too many difference with new main_perf, it's not safe to merge anymore. I have created a new ]PR]( #652) for v0.2

micmelesse and others added 30 commits July 17, 2024 05:04

skip backward (#586)

17575ea

Change all block pointers to tensor pointers (#585)

a3d784a

Change all block pointers to tensor pointers Block pointers are for nvidia TMAs. They are useful for regular loads as well but not well supported. Also cleaned up some code I came across along the way and updated comment at the top.

Add support for bshd layout (#587)

aa6685a

Add support for layouts commonly used by users. Add option for varlen / thd layout to specify equal context lengths for all batches. Also often used by users.

Post-Merge CI (#612)

dbe1173

* remove on push for Integration Tests * rename * add post merge test * save * dtype params * skip bad config * fix more stuff

Increase CI timeout (#615)

23ba546

Increase CI timeout

streamk v0.1 (#619)

52a908f

* streamk v0.1 * remove unused variable * fix format issues * add README * fix format issue * change num_sms to num_cus

Add explicit multiply-reduce GEMM kernel (#621)

1d2e066

* Add explicit multiply-reduce GEMM kernel * Remove `SPLIT_K` argument from kernel * Remove `GROUP_SIZE_M` argument from kernel * Remove conditional call to `tl.dot` from kernel * Remove table with performance data from README

add barrier to fix racing for spinning locks (#632)

177d0bd

add Lixun's occ.sh to here, will move to proper location later

83f298a

the v0.2 streamk kernel and tuning script

b927376

update readme

50fc4f6

Softmax kernel

e42690d

Merge pull request #634 from ROCm/main_perf-softmax

6d283a2

Softmax kernel

Move utility tools from triton-mlir to main_perf branch (#635)

3704738

* Move utility tools from triton-mlir to main_perf branch - Plot layout script - occ.sh - amdgcn-cfg * yapf format * More formats * remove executablility of plot_layout.py * Address ruff complains * Move tune_gemm to tools

Add rmsnorm kernel

f80aed7

Merge branch 'main_perf' into main_perf-rmsnorm

9da4278

Merge pull request #633 from ROCm/main_perf-rmsnorm

c4bd738

Add rmsnorm kernel

fix manual k-loop peeling bug and experiment with cache modifier

5711d6e

Online softmax implementation

a782caf

add missing dependencies

556aa1a

update tune_streamk script

ec58261

Merge pull request #639 from ROCm/softmax_updates

96b3d37

Online softmax implementation

add wrapper for peopel to test

64535f7

Add Layernorm kernel

042aa91

tl.load() now doesn't result in a race cond.

d686de8

Additionally changes locks to use uint8 instead of int32 for smaller space footprint.

xiaohuguo2023 and others added 12 commits September 20, 2024 06:43

reduce register usage using tl.multiple_of before tl.load

b29688f

Add use mask

ccb3538

Merge pull request #641 from ROCm/main_perf-layernorm

e13fc4c

Add Layernorm kernel

change num_sms to compiletime constant

71c5793

add num_sms to tuning space

72fc9f3

Merge branch 'main_perf' into streamk_v0.2

1838c8f

add unit test for streamk kernel

cca9c8e

add CI tests for streamk kernel

b7a99c0

fix format issues

7272ed6

remove unused wrapper and occ.sh

1edd889

fix the format issues

e98f789

update README

3275cdd

xiaohuguo2023 requested review from zhanglx13, scxiao and jayfurmanek September 25, 2024 16:18

xiaohuguo2023 added 2 commits September 25, 2024 13:01

fix git issue

2755b72

more change to make git work

4ff9faa

neoblizz self-requested a review September 26, 2024 21:17

neoblizz reviewed Sep 26, 2024

View reviewed changes

micmelesse and others added 7 commits October 1, 2024 17:25

add RDNA Config (#640)

b36e072

need cache_modifer when load P_

8add102

add back matmul wrapper for now

bf8cfe8

Add test and benchmark for explicit dot GEMM (#637)

b6633f3

fix wrapper

c129a0a

Merge branch 'main_perf' into streamk_v0.2

2ac73f6

micmelesse force-pushed the main_perf branch from 16b0bbf to 628e09b Compare October 28, 2024 15:11

xiaohuguo2023 closed this Oct 29, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Streamk v0.2 #646

Streamk v0.2 #646

xiaohuguo2023 commented Sep 25, 2024

neoblizz left a comment

neoblizz Sep 26, 2024

neoblizz Sep 26, 2024

neoblizz Sep 26, 2024

neoblizz Sep 26, 2024

neoblizz Sep 26, 2024

xiaohuguo2023 commented Oct 29, 2024

Streamk v0.2 #646

Streamk v0.2 #646

Conversation

xiaohuguo2023 commented Sep 25, 2024

neoblizz left a comment

Choose a reason for hiding this comment

neoblizz Sep 26, 2024

Choose a reason for hiding this comment

neoblizz Sep 26, 2024

Choose a reason for hiding this comment

neoblizz Sep 26, 2024

Choose a reason for hiding this comment

neoblizz Sep 26, 2024

Choose a reason for hiding this comment

neoblizz Sep 26, 2024

Choose a reason for hiding this comment

xiaohuguo2023 commented Oct 29, 2024