RMSNorm Blocked Implementation #638
Open
+9,591
−0
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 <micmelesse@gmail.com>
* 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 <Lixun.Zhang@amd.com> 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
rahulbatra85
force-pushed
the
rmsnorm_updates
branch
from
ce372d0
to
207f73a
Compare
rahulbatra85
force-pushed
the
rmsnorm_updates
branch
from
207f73a
to
44e9360
Compare
brunomazzottiamd
approved these changes
Oct 1, 2024
| mask = col_offsets < n_cols | ||
| for row_idx in tl.range(row_start, n_rows, row_step): | ||
| row_start_ptr = input_ptr + row_idx * input_row_stride | ||
| row_idx = row_start |
There was a problem hiding this comment.
We're doing row_start = tl.program_id(0) followed by row_idx = row_start. row_start is used only in row_idx = row_start statement. I think it's better to do row_idx = tl.program_id(0) to avoid indirection if someone else is reading the code.
| for b in tl.range(0, n_cols, BLOCK_SIZE): | ||
| col_offsets = b + tl.arange(0, BLOCK_SIZE) | ||
| input_ptrs = row_start_ptr + col_offsets | ||
| mask = col_offsets < n_cols |
There was a problem hiding this comment.
I can't remember if we discussed it in our daily meetings, but have you tried to peel off the last loop iteration so the loop performs unmasked loads and only the last peeled iteration performs masked loads?
| rms_norm = rms_norm * g #element wise multiplication with g | ||
|
|
||
| output_row_start_ptr = output_ptr + row_idx * output_row_stride | ||
| mask = col_offsets < n_cols |
There was a problem hiding this comment.
The peeling of last loop iteration may also be applicable here.
| @@ -110,6 +116,7 @@ def test_rmsnorm(M, N): | |||
| rms_norm = torch.nn.RMSNorm(N, device='cuda') | |||
| y_torch = rms_norm(x) | |||
|
|
|||
| print(f"y_triton={y_triton}") | |||
There was a problem hiding this comment.
Is it a debugging statement that must be removed before merging?
| (1, 8192), | ||
| (873, 1245), | ||
| ]) | ||
| @pytest.mark.parametrize('M, N', [(1, 4), (2, 10), (8192, 4096), (4096, 8192), (1, 8192), (873, 1245), (1, 98304)]) |
There was a problem hiding this comment.
Was this formatting change intentional? Or it was done my yapf? I really think the previous formatting is much more readable. You can do something like this to disable the auto formatter:
# yapf: disable
@pytest.mark.parametrize('M, N', [
(1, 4),
(2, 10),
(8192, 4096),
(4096, 8192),
(1, 8192),
(873, 1245),
])
# yapf: enable
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Switch to using blocks for columns. This gives much better performance specially for larger n_cols sizes