[Feature] Support for Triton on Windows CUDA, Linux ROCm, and Linux XPU

[iwr-redmond]

The Triton Windows fork is now well advanced, and has releases for 3.1 and 3.2. The main outstanding issue from an ease-of-use perspective is documented in the Readme file as:

TODO: Make a minimal bundle of MSVC, Windows SDK, and CUDA SDK in the wheels (help wanted)

While there is no Pip index file (yet) for this fork, such a file could easily be created using the CI script originally written by @AgainstEntropy for the Nexa SDK. As with your PCI database updates, this could be triggered automatically whenever a new release is published upstream, with the package itself referencing the index file published here on Github.

[iwr-redmond]

Interestingly, Triton is installed by default on Linux (CUDA) already.

torchruntime install
> ['/nfs/homes/iwr/Applications/Diffusion-Inference/venv/bin/python', '-m', 'pip', 'install', 'torch', 'torchaudio', 'torchvision', '--index-url', 'https://download.pytorch.org/whl/cu124']

Installing collected packages: triton, nvidia-cusparselt-cu12, mpmath, typing-extensions, sympy, pillow, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, numpy, networkx, MarkupSafe, fsspec, filelock, nvidia-cusparse-cu12, nvidia-cudnn-cu12, jinja2, nvidia-cusolver-cu12, torch, torchvision, torchaudio

[iwr-redmond]

Good news: triton-windows is now being published to PyPI, and should be much easier to integrate into torchruntime.

[iwr-redmond]

More good news: as of the latest release, the triton-windows package is now entirely self-contained, fixing the outstanding TODO noted above.

See also this comment by @woct0rdho.

[cmdr2]

Sorry for my ignorance, I haven't followed triton too closely. Does it speed up torch? Does torch have the required triton programs to speed up common operations?

If you have any quick links that I could read for why this improves performance, that would be really helpful! Thanks :)

[woct0rdho]

torch.compile requires Triton to compile Python code to GPUs. Also, packages like SageAttention and Unsloth directly use Triton to write high-performance kernels, which have accelerated recent LLMs by a lot.

[iwr-redmond]

Here are two repos that use torch.compile extensively:

• diffusers-torchao by @sayakpaul
• Comfy-WaveSpeed by @chengzeyi

[iwr-redmond]

While I was plowing the south forty over in FooocusPlus#102, I noticed that there are PyTorch builds of Triton 3.2, the current release, for both ROCm and Intel XPU on Linux.