Update request: python3Packages.openai-triton 2.0.0 → 2.1.1

[ConnorBaker]

• Package name: python3Packages.openai-triton
• Latest released version: 2.1.1 (*)
• Current version on the unstable channel: 2.0.0
• Current version on the stable/release channel: 2.0.0
• Checked the nixpkgs pull requests

PyTorch's torch.compile is/will be broken by #259068 (which is fine, because python3Packages.openai-triton is a separate package and should be updated separately).

Triton lives doesn't seem to be cutting releases. PyTorch has pinned commits they use:

• CUDA: e6216047b8b0aef1fe8da6ca8667a3ad0a016411
  https://github.com/pytorch/pytorch/blob/4c55dc50355d5e923642c59ad2a23d6ad54711e7/.ci/docker/ci_commit_pins/triton.txt#L1
• ROCm: 34f8189eae57a23cc15b4b4f032fe25757e0db8e
  https://github.com/pytorch/pytorch/blob/4c55dc50355d5e923642c59ad2a23d6ad54711e7/.ci/docker/ci_commit_pins/triton-rocm.txt#L1

Their CI uses these commits to fetch Triton from the corresponding repo (upstream for CUDA, and the AMD's fork for ROCm):

https://github.com/pytorch/pytorch/blob/4c55dc50355d5e923642c59ad2a23d6ad54711e7/.ci/docker/common/install_triton.sh#L15-L21

if [ -n "${ROCM_VERSION}" ]; then
  TRITON_REPO="https://github.com/ROCmSoftwarePlatform/triton"
  TRITON_TEXT_FILE="triton-rocm"
else
  TRITON_REPO="https://github.com/openai/triton"
  TRITON_TEXT_FILE="triton"
fi

This presents some complications for us, at least for the release of Triton which corresponds to PyTorch 2.1.1 -- we have different sources for CUDA and ROCm. (Hopefully this will be resolved as upstream continue to merge fixes for ROCm.)

Relatedly, I had apparently tried to do this bump some time ago. That may be useful: ConnorBaker@3071319

Notify maintainers

@NixOS/cuda-maintainers @NixOS/rocm-maintainers @SomeoneSerge @Madouura

---

Note for maintainers: Please tag this issue in your PR.

Priorities

Add a 👍 reaction to issues you find important.

[Madouura]

We might be able to try https://github.com/openai/triton/tree/release/2.1.x and see if it works for both before resorting to separate commits.
Edit: This won't work as it does not have the ROCm PR's changes.
We may want to change openai-triton to an unstable variant for the time being.

[Madouura]

Also ROCm support should already be somewhat upstreamed?
triton-lang/triton#1983

[Madouura]

This presents some complications for us, at least for the release of Triton which corresponds to PyTorch 2.1.1 -- we have different sources for CUDA and ROCm. (Hopefully this will be resolved as upstream continue to merge fixes for ROCm.)

Shouldn't be a problem if we use a (reasonably stable) commit on upstream. See https://github.com/openai/triton/tree/main/third_party

We'll likely need to set BACKEND and split out triton into openai-triton-cuda and openai-triton-rocm, with openai-triton defaulting to one of the two. Of course, config.(cuda|rocm)Support can handle this as well, but we still want to split for caching purposes.