squash xblock for persistent inner reduction

[ngimel]

Currently layer norm kernel performance is pretty bad due to triton perf bug https://gist.github.com/ngimel/c1e7f70f8268f038e710e835b0065f63, but since XBLOCK for persistent reduction is 1 we can just drop this dimension and operate on 1d tensors (and then perf of ln kernels improves a lot)
Perf results http://hud.pytorch.org/benchmark/compilers?startTime=Mon%2C%2022%20May%202023%2001%3A27%3A25%20GMT&stopTime=Mon%2C%2029%20May%202023%2001%3A27%3A25%20GMT&suite=torchbench&mode=training&dtype=amp&lBranch=ngimel/persistent_1d&lCommit=1d5175f5e682f37aae15fd217bc3767e1788bacf&rBranch=main&rCommit=c9f4f01981fd73fcc7c27676cc50230cd1b5bc22, approx 4% on hf

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/102444

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 235b88e:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ngimel]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team

Raised by workflow job

[ngimel]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / linux-focal-rocm5.4.2-py3.8 / test (default, 1, 3, linux.rocm.gpu)

Details for Dev Infra team

Raised by workflow job

[ngimel]

cc @jataylo, I had to exclude rocm from this optimization, because rocm is on the old triton version that doesn't have tl.reduce and workaround doesn't work for 1d tensors (with some fixes it worked for me locally when I forced the workaround path in triton_helpers.py but it was still failing on rocm CI d97b63e). Can you guys update your triton to a later pin so that tl.reduce workaround is not longer needd?

[ngimel]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[jataylo]

cc @jataylo, I had to exclude rocm from this optimization, because rocm is on the old triton version that doesn't have tl.reduce and workaround doesn't work for 1d tensors (with some fixes it worked for me locally when I forced the workaround path in triton_helpers.py but it was still failing on rocm CI d97b63e). Can you guys update your triton to a later pin so that tl.reduce workaround is not longer needd?

Thanks for the heads up @ngimel, we hit a blocker that stopped us updating triton for awhile to bring in the tl.reduce change (ROCm/triton#208), we've fixed that associated issue and we are hoping to be able to move our triton commit forward shortly after merging the changes into our pyt2.0 branch of triton.

I'll remove the conditionalisation of this commit in the PR that bumps our triton commit and add you as a reviewer.

cc: @dllehr-amd