Add Layernorm kernel #641
Add Layernorm kernel #641
Conversation
rahulbatra85
force-pushed
the
main_perf-layernorm
branch
from
b772444
to
674d526
Compare
rahulbatra85
requested review from
brunomazzottiamd,
micmelesse and
jtang10
There was a problem hiding this comment.
@rahulbatra85, I can't see anything wrong with your PR. I've just have some questions and minor code cleanup suggestions. Feel free to ignore them if you judge appropriate.
Please drop a short line about this new kernel to python/perf-kernels/README.md file.
| @@ -128,8 +128,10 @@ jobs: | |||
| pytest -vvv ./python/perf-kernels/flash-attention.py | |||
| pytest -vvvv ./python/perf-kernels/softmax.py | |||
| pytest -vvv ./python/perf-kernels/rmsnorm.py | |||
| pytest -vvv ./python/perf-kernels/layernorm.py | |||
There was a problem hiding this comment.
What do you think about running all tests with just one pytest invocation? According to https://docs.pytest.org/en/stable/how-to/usage.html, it's possible to do something like pytest -vvvv ./python/perf-kernels. By this way, we'll be editing .github/workflows/amd_perf_kernel_Integration_tests.yml less often and new tests are going to run by default. Do you see any drawback?
Maybe it's worth asking @micmelesse's opinion on this.
There was a problem hiding this comment.
yeah that's a Michael question
There was a problem hiding this comment.
Let's wait for Michael's opinion!
There was a problem hiding this comment.
that is fine. I think some of the tests are broken but maybe worth it to see the state of things
python/perf-kernels/layernorm.py
Outdated
| y = x_hat * w + b | ||
| # Write output | ||
| tl.store(Y + cols, y, mask=mask) | ||
|
|
There was a problem hiding this comment.
Just an idea:
We have three for loops that do masked loads. Do you foresee any benefit of peeling the last iteration of each loop so all iterations except the last one do unmasked loads? I think Shucai and Xiaohu got some performance improvements doing this with GEMMs. I'm not sure if the idea could be beneficial for layer norm.
There was a problem hiding this comment.
ok, yeah, I didn't think of that. Will try this out
There was a problem hiding this comment.
Please let me know if this helped at all.
rahulbatra85
force-pushed
the
main_perf-layernorm
branch
from
d88abbc
to
13c01c4
Compare
brunomazzottiamd
assigned brunomazzottiamd and rahulbatra85 and unassigned brunomazzottiamd
This comment was marked as resolved.
This comment was marked as resolved.
rahulbatra85
force-pushed
the
main_perf-layernorm
branch
from
13c01c4
to
042aa91
Compare
No description provided.