Merge branch 'main' into whitneywhtsang/prefetch_mask

Author: whitneywhtsang

.github/workflows/build-test-reusable.yml

@@ -74,6 +74,12 @@ env:
 jobs:
   integration-tests:
     name: Integration tests
+    strategy:
+      fail-fast: false
+      matrix:
+        suite:
+          - core
+          - rest
     timeout-minutes: 720
     runs-on: ${{ fromJson(inputs.runner_label && format('["linux", "{0}"]', inputs.runner_label) || format('["linux", "{0}", "{1}", "{2}"]', inputs.device, inputs.driver_version, inputs.runner_version)) }}
     defaults:
@@ -111,31 +117,23 @@ jobs:
         with:
           python-version: ${{ inputs.python_version }}
 
-      - name: Identify Python version
-        run: |
-          PYTHON_VERSION="$(python -c 'import sys; print(f"{sys.version_info[0]}.{ sys.version_info[1]}")')"
-          echo "PYTHON_VERSION=$PYTHON_VERSION" | tee -a $GITHUB_ENV
-
       - name: Setup PyTorch
         uses: ./.github/actions/setup-pytorch
         with:
           ref: ${{ inputs.pytorch_ref }}
           mode: ${{ inputs.pytorch_mode }}
 
-      - name: Install pass_rate dependencies
-        run: |
-          pip install defusedxml setproctitle==1.3.5
-
       - name: Setup Triton
         uses: ./.github/actions/setup-triton
         with:
           build_llvm: ${{ inputs.build_llvm }}
           use_spirv_backend: ${{ inputs.use_spirv_backend }}
 
       - name: Report environment details
+        if: matrix.suite == 'core'
         run: |
-          source ./scripts/capture-hw-details.sh --quiet
-          cat <<EOF | tee .env
+          mkdir -p reports
+          cat <<EOF | tee reports/.env
           TIMESTAMP=$(date '+%Y%m%d%H%M%S')
           GITHUB_RUN_ID=$GITHUB_RUN_ID
           GITHUB_RUN_NUMBER=$GITHUB_RUN_NUMBER
@@ -145,11 +143,8 @@ jobs:
           PYTORCH_COMMIT_ID=$PYTORCH_COMMIT_ID
           PYTORCH_VERSION=$PYTORCH_VERSION
           TRITON_REPO=$GITHUB_REPOSITORY
-          LIBIGC1_VERSION=$LIBIGC1_VERSION
-          LEVEL_ZERO_VERSION=$LEVEL_ZERO_VERSION
-          GPU_DEVICE=$GPU_DEVICE
-          AGAMA_VERSION=$AGAMA_VERSION
           EOF
+          ./scripts/capture-hw-details.sh | tee -a reports/.env
 
       - name: Create test-triton command line
         run: |
@@ -169,43 +164,54 @@ jobs:
 
           {
             echo SKIPLIST="$skiplist"
-            echo TRITON_TEST_CMD="bash -x scripts/test-triton.sh --warning-reports --skip-pytorch-install --reports-dir $GITHUB_WORKSPACE/reports ${{ inputs.ignore_errors && '--ignore-errors' || '' }} $skiplist"
+            echo TRITON_TEST_CMD="bash -x scripts/test-triton.sh --skip-pip-install --warning-reports --skip-pytorch-install --reports-dir $GITHUB_WORKSPACE/reports ${{ inputs.ignore_errors && '--ignore-errors' || '' }} $skiplist"
           } | tee -a $GITHUB_ENV
 
+      - name: Install test dependencies
+        run: |
+          pip install -r scripts/requirements-test.txt git+https://github.com/kwasd/pytest-capturewarnings-ng@v1.2.0
+
       - name: Run Proton tests
-        if: ${{ inputs.driver_version == 'rolling' && inputs.device == 'max1100' }}
+        if: matrix.suite == 'rest' && inputs.driver_version == 'rolling' && inputs.device == 'max1100'
         run: |
           cd third_party/proton/test
           pytest test_api.py test_lib.py test_profile.py test_viewer.py test_record.py -s -v
           cd ..
 
       - name: Run unit tests
+        if: matrix.suite == 'rest'
         run: |
           ${{ env.TRITON_TEST_CMD }} --unit
 
       - name: Run core tests
+        if: matrix.suite == 'core'
         run: |
-          ${{ env.TRITON_TEST_CMD }} --core --skip-pip-install
+          ${{ env.TRITON_TEST_CMD }} --core
 
       - name: Run interpreter tests
+        if: matrix.suite == 'rest'
         run: |
-          ${{ env.TRITON_TEST_CMD }} --interpreter --skip-pip-install
+          ${{ env.TRITON_TEST_CMD }} --interpreter
 
       - name: Run Tutorials
+        if: matrix.suite == 'rest'
         run: |
-          ${{ env.TRITON_TEST_CMD }} --tutorial --skip-pip-install
+          ${{ env.TRITON_TEST_CMD }} --tutorial
 
       - name: Run instrumentation tests
+        if: matrix.suite == 'rest'
         run: |
-          ${{ env.TRITON_TEST_CMD }} --instrumentation --skip-pip-install
+          ${{ env.TRITON_TEST_CMD }} --instrumentation
 
       - name: Get transformers version
+        if: matrix.suite == 'rest'
         run: |
           cd pytorch
           TRANSFORMERS_VERSION="$(<.ci/docker/ci_commit_pins/huggingface.txt)"
           echo "TRANSFORMERS_VERSION=$TRANSFORMERS_VERSION" | tee -a $GITHUB_ENV
 
       - name: Install transformers
+        if: matrix.suite == 'rest'
         uses: ./.github/actions/install-dependency
         with:
           package: transformers
@@ -214,8 +220,9 @@ jobs:
           try-tag-prefix: v
 
       - name: Run E2E test
+        if: matrix.suite == 'rest'
         run: |
-          ${{ env.TRITON_TEST_CMD }} --inductor --skip-pip-install
+          ${{ env.TRITON_TEST_CMD }} --inductor
 
       - name: Save pip cache
         if: ${{ steps.pip-cache.outputs.status == 'miss' }}
@@ -224,9 +231,52 @@ jobs:
           path: ${{ steps.pip-cache.outputs.path }}
           dest: ${{ steps.pip-cache.outputs.dest }}
 
+      - name: Upload test reports
+        uses: actions/upload-artifact@v4
+        with:
+          name: test-reports-${{ matrix.suite }}-${{ inputs.python_version }}-${{ inputs.runner_label || inputs.driver_version }}
+          include-hidden-files: true
+          path: reports
+
+  reports:
+    name: Reports
+    runs-on:
+      - linux
+    needs: integration-tests
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Load pip cache
+        id: pip-cache
+        uses: ./.github/actions/load
+        env:
+          # Increase this value to reset cache
+          CACHE_NUMBER: 1
+        with:
+          path: $HOME/.cache/pip
+          key: pip-${{ inputs.python_version }}-${{ hashFiles('python/pyproject.toml', 'python/setup.py') }}-${{ env.CACHE_NUMBER }}
+
+      - name: Download test reports
+        uses: actions/download-artifact@v4
+        with:
+          pattern: test-reports-*-${{ inputs.python_version }}-${{ inputs.runner_label || inputs.driver_version }}
+          path: reports
+          merge-multiple: true
+
+      - name: Install Python (using actions/setup-python) ${{ inputs.python_version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ inputs.python_version }}
+
+      - name: Install pass_rate dependencies
+        run: |
+          pip install defusedxml
+
       - name: Pass rate
         run: |
-          source ./scripts/capture-hw-details.sh
+          ls -la reports
+          source ./reports/.env
           python scripts/pass_rate.py --reports reports
           python scripts/pass_rate.py --reports reports --json > pass_rate.json
           python scripts/pass_rate.py --reports reports --suite tutorials --json > pass_rate_tutorials.json
@@ -247,11 +297,4 @@ jobs:
           include-hidden-files: true
           path: |
             reports/*/*.csv
-            .env
-
-      - name: Upload test reports
-        if: inputs.upload_test_reports
-        uses: actions/upload-artifact@v4
-        with:
-          name: test-reports-${{ inputs.python_version }}-${{ inputs.runner_label || inputs.driver_version }}
-          path: reports
+            reports/.env

.github/workflows/build-test.yml

@@ -100,6 +100,7 @@ jobs:
     name: Integration tests matrix
 
     strategy:
+      fail-fast: false
       matrix:
         python:
           - "3.9"

.github/workflows/integration-tests.yml

@@ -114,7 +114,7 @@ jobs:
         if: env.enable_integration == 'true'
         run: |
           if [ x"${{ github.repository }}" == x"triton-lang/triton" ]; then
-            echo '::set-output name=matrix-CUDA::[["a100-runner-set"], ["h100-runner-set"], ["gb200-runner-set"]]'
+            echo '::set-output name=matrix-CUDA::[["nvidia-a100"], ["nvidia-h100"], ["nvidia-gb200"]]'
             echo '::set-output name=matrix-HIP::[["self-hosted", "gfx90a"], ["amd-gfx942"]]'
             echo '::set-output name=matrix-MACOS::[["macos-latest"]]'
           else
@@ -232,7 +232,7 @@ jobs:
         env:
           CUDA_HOME: "/usr/local/cuda"
         run: |
-          if [ "${{ matrix.runner[0] }}" == "gb200-runner-set" ]; then
+          if [ "${{ matrix.runner[0] }}" == "nvidia-gb200" ]; then
             source /venv/bin/activate
           fi
           echo "PATH is '$PATH'"
@@ -244,23 +244,23 @@ jobs:
         run: make test-lit
       - name: Run python tests on CUDA
         run: |
-          if [ "${{ matrix.runner[0] }}" == "gb200-runner-set" ]; then
+          if [ "${{ matrix.runner[0] }}" == "nvidia-gb200" ]; then
             source /venv/bin/activate
           fi
           make test-unit
       - name: Run interpreter tests
-        if: ${{ matrix.runner[0] == 'h100-runner-set' }}
+        if: ${{ matrix.runner[0] == 'nvidia-h100' }}
         run: make test-interpret
       - name: Run regression tests
         run: |
-          if [ "${{ matrix.runner[0] }}" == "gb200-runner-set" ]; then
+          if [ "${{ matrix.runner[0] }}" == "nvidia-gb200" ]; then
             source /venv/bin/activate
           fi
           make test-regression
       - name: Run C++ unittests
         run: make test-cpp
       - name: Run Proton tests
-        if: ${{ matrix.runner[0] != 'gb200-runner-set' }}
+        if: ${{ matrix.runner[0] != 'nvidia-gb200' }}
         run: make test-proton
       - name: Inspect cache directories
         run: |
@@ -409,7 +409,7 @@ jobs:
           cd python/test/regression
           python3 -m pytest -s -n 8 ./test_cast_matmul.py
       - name: Run Proton tests
-        if: ${{ matrix.runner[0] != 'gb200-runner-set' }}
+        if: ${{ matrix.runner[0] != 'nvidia-gb200' }}
         run: make test-proton
       - name: Run C++ unittests
         run: make test-cpp

.github/workflows/integration-tests.yml.in

@@ -123,7 +123,7 @@ jobs:
         if: env.enable_integration == 'true'
         run: |
           if [ x"${{ github.repository }}" == x"triton-lang/triton" ]; then
-            echo '::set-output name=matrix-CUDA::[["a100-runner-set"], ["h100-runner-set"], ["gb200-runner-set"]]'
+            echo '::set-output name=matrix-CUDA::[["nvidia-a100"], ["nvidia-h100"], ["nvidia-gb200"]]'
             echo '::set-output name=matrix-HIP::[["self-hosted", "gfx90a"], ["amd-gfx942"]]'
             echo '::set-output name=matrix-MACOS::[["macos-latest"]]'
           else
@@ -264,7 +264,7 @@ jobs:
         env:
           CUDA_HOME: "/usr/local/cuda"
         run: |
-          if [ "${{ matrix.runner[0] }}" == "gb200-runner-set" ]; then
+          if [ "${{ matrix.runner[0] }}" == "nvidia-gb200" ]; then
             source /venv/bin/activate
           fi
           echo "PATH is '$PATH'"
@@ -281,18 +281,18 @@ jobs:
 
       - name: Run python tests on CUDA
         run:  |
-          if [ "${{ matrix.runner[0] }}" == "gb200-runner-set" ]; then
+          if [ "${{ matrix.runner[0] }}" == "nvidia-gb200" ]; then
             source /venv/bin/activate
           fi
           make test-unit
 
       - name: Run interpreter tests
-        if: ${{ matrix.runner[0] == 'h100-runner-set' }}
+        if: ${{ matrix.runner[0] == 'nvidia-h100' }}
         run: make test-interpret
 
       - name: Run regression tests
         run:  |
-          if [ "${{ matrix.runner[0] }}" == "gb200-runner-set" ]; then
+          if [ "${{ matrix.runner[0] }}" == "nvidia-gb200" ]; then
             source /venv/bin/activate
           fi
           make test-regression
@@ -303,7 +303,7 @@ jobs:
 
       - &run-proton-tests-step
         name: Run Proton tests
-        if: ${{ matrix.runner[0] != 'gb200-runner-set' }}
+        if: ${{ matrix.runner[0] != 'nvidia-gb200' }}
         run: make test-proton
 
       - *inspect-cache-directories-step

bench/bench/bench_mlp.py

@@ -3,11 +3,11 @@
 import triton.profiler as proton
 import torch
 import triton_bench.swiglu
-from triton_bench.mxfp import downcast_to_mxfp
+from triton_bench.numerics_details.mxfp import downcast_to_mxfp
 from triton_bench.matmul_ogs import MicroscalingCtx, matmul_ogs, PrecisionConfig, FlexCtx
 from triton_bench.numerics import InFlexData
 from triton_bench.routing import routing
-from triton_bench.meta import cuda_capability_geq, is_hip, get_cdna_version
+from triton_bench.target_info import is_hip, get_cdna_version
 
 if torch.cuda.is_available() and not is_hip():
     from triton._C.libtriton import nvidia
@@ -152,5 +152,5 @@ def bench_mlp(batch, dim1, dim2, n_expts_tot, n_expts_act, x_dtype, w_dtype,
     qxdtype = "fp8" if has_native_mx4 else "bf16"
     print(bench_mlp(8192, 8192, 8192, 1, 1, "fp8", "fp8", TP=1, EP=1, name="dense"))
     print(bench_mlp(8192, 8192, 8192, 1, 1, qxdtype, "mx4", TP=1, EP=1, name="dense"))
-    print(bench_mlp(2048, 5120, 8192, 128, 4, "fp8", "fp8", TP=4, EP=2, name="llama4"))
-    print(bench_mlp(2048, 5120, 8192, 128, 4, qxdtype, "mx4", TP=4, EP=2, name="llama4"))
+    print(bench_mlp(2048, 5120, 8192, 128, 4, "fp8", "fp8", TP=4, EP=1, name="llama4"))
+    print(bench_mlp(2048, 5120, 8192, 128, 4, qxdtype, "mx4", TP=4, EP=1, name="llama4"))

bench/tests/test_compact.py bench/tests/test_compaction.pybench/tests/test_compact.py renamed to bench/tests/test_compaction.py

@@ -1,6 +1,6 @@
 import pytest
 import torch
-from triton_bench.compact import masked_compact, masked_compact_torch
+from triton_bench.compaction import compaction, compaction_torch
 
 
 @pytest.mark.parametrize("n_tokens, n_cols, k, p", [
@@ -9,7 +9,7 @@
     (131, 128, 16, 0.6),
     (496, 128, 16, 0.),
 ])
-def test_masked_compact(n_tokens, n_cols, k, p):
+def test_compaction(n_tokens, n_cols, k, p):
     device = "cuda"
     yi = torch.rand((n_tokens, n_cols), device=device).argsort(dim=-1)
     yi = yi[:, :k].to(torch.int32)
@@ -23,7 +23,7 @@ def test_masked_compact(n_tokens, n_cols, k, p):
     chunks = mask.view(*mask.shape[:-1], -1, 32)
     weights = (1 << torch.arange(32, dtype=torch.int32, device=device))
     bitmask = (chunks.int() * weights).sum(dim=-1)
-    yv_ref, yi_ref = masked_compact_torch(yv, yi, bitmask)
-    yv_tri, yi_tri = masked_compact(yv, yi, bitmask)
+    yv_ref, yi_ref = compaction_torch(yv, yi, bitmask)
+    yv_tri, yi_tri = compaction(yv, yi, bitmask)
     assert torch.all(yi_ref == yi_tri)
     assert torch.all(yv_ref == yv_tri)

bench/tests/test_matmul.py

@@ -11,11 +11,11 @@
 from triton_bench.matmul_ogs import matmul_ogs, matmul_ogs_torch
 # numerics utilities
 from triton_bench.numerics import InFlexData, OutFlexData
-from triton_bench.mxfp import downcast_to_mxfp, upcast_from_mxfp
+from triton_bench.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
 # testing utilities
 from triton_bench.testing import assert_close, compute_actual_scale
 # target-specific utilities
-from triton_bench.meta import is_hip
+from triton_bench.target_info import is_hip
 
 # ---------------
 # initialize data