Add tutorial for user defined triton kernels

[oulgen]

No description provided.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 8ee52f7 with merge base 5fbef68 ():
💚 Looks good so far! There are no failures yet. 💚

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

[malfet]

This is not very readable. Why not

if not has_trition:
   print("Skipping because triton is not supported on this device.")
   sys.exit(1)

[oulgen]

Other tutorials are also doing this: https://pytorch.org/tutorials/recipes/torch_logs.html

[malfet]

This makes me sad, let me try to propose a bit more elegant solution

[malfet]

It would be nice to add larger preamble that references back documentation, explains what the triton is, how it will only work with GPUs, what auto-tuning is for, again, link back to some document about dynamic shapes, etc

[malfet]

This makes me sad, let me try to propose a bit more elegant solution

[svekars]

Some editorial suggestions. Let me know if you have questions! Please add a card and an entry to recipes_index.html

[svekars]

Suggested change

	Using User Defined Triton Kernels with ``torch.compile``
	Using User-Defined Triton Kernels with ``torch.compile``

[svekars]

Need to add a better introduction here. Maybe something like this:

Suggested change

	# This tutorial explains how to use user defined Triton kernels with ``torch.compile``.
	# User-defined Triton kernels can be used to optimize specific parts of your
	# model's computation. These kernels are written in Triton's language, which is designed
	# to make it easier to achieve peak hardware performance. By using user-defined Triton
	# kernels with ``torch.compile``, you can integrate these optimized computations into
	# your PyTorch model, potentially achieving significant performance improvements.
	#
	# This recipes demonstrates how you can use user-defined Triton kernels with ``torch.compile``.

[svekars]

Suggested change

	# Reference: https://triton-lang.org/main/getting-started/tutorials/01-vector-add.html
	# For reference, see `Triton documentation <https://triton-lang.org/main/getting-started/tutorials/01-vector-add.html>`__.

[svekars]

Needs to be as long or longer than the title

Suggested change

	# ------------
	# ----------------------

[svekars]

Need to expand the intro a bit. Maybe something like this:

Suggested change

	# It is also possible to triton.autotune with ``torch.compile``.
	# Triton's autotune feature is a powerful tool that automatically optimizes the configuration
	# parameters of your Triton kernels. It explores a range of possible configurations and
	# selects the one that delivers the best performance for your specific use case.
	#
	# When used with ``torch.compile``, ``triton.autotune`` can help ensure that your PyTorch
	# model is running as efficiently as possible. Here is an example of using ``torch.compile``
	# and ``triton.autotune``.

[svekars]

Suggested change

	# As for PyTorch 2.3, the user defined triton kernel support in ``torch.compile``
	# composes with dynamic shapes, ``torch.autograd.Function``, JIT inductor and
	# AOT inductor.
	#
	# The support for tensor subclasses and other advanced features currently do
	# not exist.
	# Support for ``triton.heuristics`` exists when it is used by itself or before
	# ``triton.autotune``; however, support for using ``triton.heuristic`` after
	# ``triton.autotune`` is not yet supported.
	# As of PyTorch 2.3, the support for user-defined Triton kernels in ``torch.compile``
	# includes dynamic shapes, ``torch.autograd.Function``, JIT inductor, and AOT inductor.
	# You can use these features together to build complex, high-performance models.
	#
	# However, there are certain limitations to be aware of:
	#
	# * **Tensor Subclasses:** Currently, there is no support for
	# tensor subclasses and other advanced features.
	# * **Triton Features:** While ``triton.heuristics`` can be used either standalone or
	# before ``triton.autotune``, it cannot be used after ```triton.autotune``. This
	# implies that if ``triton.heuristics`` and ``triton.autotune`` are to be used
	# together, ``triton.heuristics`` must be used first.

[svekars]

Also, need to add a Conclusion:

Suggested change

	# ``triton.autotune`` is not yet supported.
	# ``triton.autotune`` is not yet supported.
	#
	# Conclusion
	# -----------
	# In this recipe, we explored how to utilize user-defined Triton kernels
	# with ``torch.compile``. We delved into the basic usage of a simple
	# vector addition kernel and advanced usage involving Triton's autotune
	# feature. We also discussed the composability of user-defined Triton
	# kernels with other PyTorch features and highlighted some current limitations.

Can you also add what else the user should read?

[svekars]

We have this in the Prerequisites section

Suggested change

	# .. note::
	# This tutorial requires PyTorch 2.3 or later and a GPU that supports Triton.

[svekars]

LGTM from the publishing perspective, please get a technical reviewer to approve. We should not merge until 2.3 binaries are available on the main branch.

[zou3519]

looks great!

[svekars]

Minor editorial fixes