Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

RMSNorm Blocked Implementation #638

Open
wants to merge 25 commits into
base: main_perf
Choose a base branch
from
Open

Conversation

rahulbatra85
Copy link

Switch to using blocks for columns. This gives much better performance specially for larger n_cols sizes

micmelesse and others added 19 commits July 17, 2024 05:04
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
* 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
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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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}")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)])
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants