Add Triton TileIR kernels (rms_norm, rope, layer_norm_legacy, dropout) + TileIR CI pass#157
Open
hannahli-nv wants to merge 2 commits into
Open
Add Triton TileIR kernels (rms_norm, rope, layer_norm_legacy, dropout) + TileIR CI pass#157hannahli-nv wants to merge 2 commits into
hannahli-nv wants to merge 2 commits into
Conversation
…) + TileIR CI pass Add Triton implementations of rms_norm, rope, layer_norm_legacy and dropout (wired into the dispatch + covered by test_op), and exercise them under the Triton TileIR backend in CI: - modeling/transformers/Dockerfile: build the TileIR Triton backend (github.com/triton-lang/Triton-to-tile-IR) from source side-by-side at /opt/nvtriton. - .github/workflows/tilegym-ci.yml: add an ops test pass that runs the Triton backend tests with PYTHONPATH=/opt/nvtriton and ENABLE_TILE=1.
Collaborator
Author
|
/ok to test ad2b382 |
hannahli-nv
force-pushed
the
triton-kernels-nvt-ci
branch
3 times, most recently
from
c2f76da to
4524c74
Compare
…uilt wheel on CUDA 13.3 / Ubuntu 24.04 - ruff --select I, ruff format, and end-of-file-fixer on the new triton kernels and the selector helper - Base image -> nvcr cuda 13.3.0-devel-ubuntu24.04 (GCC 13, Python 3.12) - Install the prebuilt nvtriton (TileIR) cp312 wheel at /opt/nvtriton instead of building from source, which was timing out the CI build job - Pin uv sync to Python 3.12 so the venv matches the wheel - test-ops: raise job timeout to 90m, default ops step to 45m, and run the TileIR ops step with -n 12 (the added triton params grew the matrix and the job cap was too low for two test steps)
hannahli-nv
force-pushed
the
triton-kernels-nvt-ci
branch
from
4524c74 to
b053eab
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add Triton implementations of
rms_norm,rope,layer_norm_legacy, anddropout(wired into the dispatch and covered bytest_op), and validate them under the Triton TileIR backend in CI.Changes
src/tilegym/ops/triton/{rms_norm,rope,layer_norm_legacy,dropout}.pyplus dispatch wiring (ops/ops.py,ops/__init__.py,ops/triton/__init__.py) and backend detection (backend/selector.py,backend/__init__.py).modeling/transformers/Dockerfile): build the TileIR Triton backend (github.com/triton-lang/Triton-to-tile-IR) from source, installed side-by-side at/opt/nvtriton..github/workflows/tilegym-ci.yml): add an ops test pass that runs the Triton-backend tests withPYTHONPATH=/opt/nvtritonandENABLE_TILE=1.Notes