[Docker] Install Trition deps

[malfet]

Triton needs a working gcc, so install one from apt
Also, copy ptxas and cuda.h from conda to /usr/local/cuda
Add torchaudio to the matrix
Fix typo in workflow file

Fixes #90377

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 b06e58c:
💚 Looks good so far! There are no failures yet. 💚

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

[ngimel]

should we instead set triton's env var to point to conda's PTXAS, and similarly set envvar to add conda's include to common includes?

[malfet]

Let me try to see if it will work...

[malfet]

Hmm, ok, exporting CUDA_HOME=/opt/conda/ works, though it feels a bit weird, is it does not contain nvcc and other tools...

[ngimel]

For triton the var is TRITON_PTXAS_PATH

[ngimel]

Yeah but /usr/local/cuda/bin that contains only ptxas is no less weird :-)

[malfet]

TRITION_PTXAS_PATH only works for PTXAS, but cuda.h is only search in CUDA_HOME folder

[malfet]

Ok, come to think of it:

• Short term solution should be to create /usr/local/trition-min-cuda-X.Y and symlink to /usr/local/cuda/ (so that if user attempts to install cuda toolkit it will simply override the symlink)
• Mid-term solution is to package those dependencies into triton and make it search there if its not available in other places

[malfet]

@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