[Update] update dev/perf branch to main branch commit 9c2df5 and cherry-pick mla_fix, mmd/fix/fwd_v3

.github/scripts/build_aiter_triton.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+set -ex
+
+echo
+echo "==== ROCm Packages Installed ===="
+dpkg -l | grep rocm || echo "No ROCm packages found."
+
+echo
+echo "==== Install dependencies and aiter ===="
+pip install --upgrade pandas zmq einops numpy==1.26.2
+pip uninstall -y aiter || true
+pip install --upgrade "pybind11>=3.0.1"
+pip install --upgrade "ninja>=1.11.1"
+python3 setup.py develop
+
+# Read BUILD_TRITON env var, default to 1. If 1, install Triton; if 0, skip installation.
+BUILD_TRITON=${BUILD_TRITON:-1}
+
+if [[ "$BUILD_TRITON" == "1" ]]; then
+    echo
+    echo "==== Install triton ===="
+    pip uninstall -y triton || true
+    git clone --depth=1 https://github.com/triton-lang/triton || true
+    cd triton
+    pip install -r python/requirements.txt
+    pip install filecheck
+    MAX_JOBS=64 pip --retries=5 install .
+    cd ..
+else
+    echo
+    echo "[SKIP] Triton installation skipped because BUILD_TRITON=$BUILD_TRITON"
+fi
+
+echo
+echo "==== Show installed packages ===="
+pip list

.github/scripts/check_signal.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# This script attempts to download a pre-checks artifact from a GitHub workflow up to 5 times.
+# If the artifact is found and the signal indicates success, the workflow continues.
+# If the signal indicates failure, the workflow is skipped with details printed.
+# If the artifact cannot be downloaded after all retries, the workflow exits with an error.
+
+set -e
+
+ARTIFACT_NAME="checks-signal-${GITHUB_SHA:-${1}}"
+MAX_RETRIES=5
+
+for i in $(seq 1 $MAX_RETRIES); do
+  echo "Attempt $i: Downloading artifact..."
+  if gh run download --name "$ARTIFACT_NAME"; then
+    if [ -f checks_signal.txt ]; then
+      echo "Artifact $ARTIFACT_NAME downloaded successfully."
+      SIGNAL=$(head -n 1 checks_signal.txt)
+      if [ "$SIGNAL" = "success" ]; then
+        echo "Pre-checks passed, continuing workflow."
+        exit 0
+      else
+        echo "Pre-checks failed, skipping workflow. Details:"
+        tail -n +2 checks_signal.txt
+        exit 78  # 78 = neutral/skip
+      fi
+    fi
+  fi
+  echo "Artifact not found, retrying in 30s..."
+  sleep 30
+done
+
+echo "Failed to download pre-checks artifact after $MAX_RETRIES attempts. Exiting workflow."
+exit 1

.github/scripts/op_tune.sh

@@ -18,17 +18,31 @@ testFailedFiles=()
 declare -a tune_jobs=(
   "ck_batched_gemm_a8w8:csrc/ck_batched_gemm_a8w8:op_tests/test_batched_gemm_a8w8.py:python3 csrc/ck_batched_gemm_a8w8/batched_gemm_a8w8_tune.py -i aiter/configs/a8w8_untuned_batched_gemm.csv -o aiter/configs/a8w8_tuned_batched_gemm.csv"
   "ck_batched_gemm_bf16:csrc/ck_batched_gemm_bf16:op_tests/test_batched_gemm_bf16.py:python3 csrc/ck_batched_gemm_bf16/batched_gemm_bf16_tune.py -i aiter/configs/bf16_untuned_batched_gemm.csv -o aiter/configs/bf16_tuned_batched_gemm.csv"
-#  "csrc/ck_gemm_a4w4_blockscale:op_tests/test_gemm_a4w4_blockscale.py:python3 csrc/ck_gemm_a4w4_blockscale/gemm_a4w4_blockscale_tune.py -i aiter/configs/a4w4_blockscale_untuned_gemm.csv -o aiter/configs/a4w4_blockscale_tuned_gemm.csv"
   "ck_gemm_a8w8:csrc/ck_gemm_a8w8:op_tests/test_gemm_a8w8.py:python3 csrc/ck_gemm_a8w8/gemm_a8w8_tune.py -i aiter/configs/a8w8_untuned_gemm.csv -o aiter/configs/a8w8_tuned_gemm.csv"
   "ck_gemm_a8w8_blockscale:csrc/ck_gemm_a8w8_blockscale:op_tests/test_gemm_a8w8_blockscale.py:python3 csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py -i aiter/configs/a8w8_blockscale_untuned_gemm.csv -o aiter/configs/a8w8_blockscale_tuned_gemm.csv"
   "ck_gemm_a8w8_blockscale_bpreshuffle:csrc/ck_gemm_a8w8_blockscale_bpreshuffle:op_tests/test_gemm_a8w8_blockscale.py:python3 csrc/ck_gemm_a8w8_blockscale_bpreshuffle/gemm_a8w8_blockscale_bpreshuffle_tune.py -i aiter/configs/a8w8_blockscale_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_blockscale_bpreshuffle_tuned_gemm.csv"
   "ck_gemm_a8w8_bpreshuffle:csrc/ck_gemm_a8w8_bpreshuffle:op_tests/test_gemm_a8w8.py:python3 csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py -i aiter/configs/a8w8_bpreshuffle_untuned_gemm.csv -o aiter/configs/a8w8_bpreshuffle_tuned_gemm.csv"
+  #"ck_gemm_a4w4_blockscale:csrc/ck_gemm_a4w4_blockscale:op_tests/test_gemm_a4w4_blockscale.py:python3 csrc/ck_gemm_a4w4_blockscale/gemm_a4w4_blockscale_tune.py -i aiter/configs/a4w4_blockscale_untuned_gemm.csv -o aiter/configs/a4w4_blockscale_tuned_gemm.csv"
 )
 
 for job in "${tune_jobs[@]}"; do
     IFS=':' read -r shape dir test_path tune_cmd <<< "$job"
-    if [ -n "$shape_filter" ] && [ "$shape" != "$shape_filter" ]; then
-        continue
+    # If shape_filter is not empty, check if the current shape exists in the filter list.
+    # shape_filter is a comma-separated list, e.g. "ck_gemm_a8w8,ck_batched_gemm_a8w8"
+    if [ -n "$shape_filter" ]; then
+        # Remove all whitespace from the shape_filter string
+        shape_filter_no_space="${shape_filter//[[:space:]]/}"
+        IFS=',' read -ra filter_shapes <<< "$shape_filter_no_space"
+        found_match=false
+        for filter_shape in "${filter_shapes[@]}"; do
+            if [[ "$shape" == "$filter_shape" ]]; then
+                found_match=true
+                break
+            fi
+        done
+        if [ "$found_match" = false ]; then
+            continue
+        fi
     fi
     echo "============================================================"
     echo "🧪 Processing shape: $shape under directory: $dir"

.github/workflows/aiter-test.yaml

@@ -15,8 +15,21 @@ env:
   DOCKER_IMAGE: "rocm/pytorch:latest"
 
 jobs:
+  check-signal:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Download and check signal artifact
+        run: ./.github/scripts/check_signal.sh
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GITHUB_SHA: ${{ github.sha }}
+
   define-runners:
     runs-on: ubuntu-latest
+    needs: [check-signal]
     outputs:
       standard_runners: ${{ steps.machines.outputs.standard_runners }}
       multigpu_runners: ${{ steps.machines.outputs.multigpu_runners }}
@@ -28,15 +41,15 @@ jobs:
           set -euo pipefail
           pr_title="${{ github.event.pull_request.title }}"
           if [[ "${{ github.ref }}" == "refs/heads/main" ]]; then
-            echo "It's main branch, running tests on MI300 and MI35X..."
+            echo "It's main branch, running tests on MI325 and MI35X..."
             echo 'standard_runners=["aiter-mi355-1gpu"]' >> "$GITHUB_OUTPUT"
             echo 'multigpu_runners=["aiter-mi355-8gpu"]' >> "$GITHUB_OUTPUT"
           elif echo "$pr_title" | grep -qi "mi35x"; then
-            echo "PR title contains 'MI35X', running tests on MI300 and MI35X..."
+            echo "PR title contains 'MI35X', running tests on MI325 and MI35X..."
             echo 'standard_runners=["aiter-mi355-1gpu"]' >> "$GITHUB_OUTPUT"
             echo 'multigpu_runners=["aiter-mi355-8gpu"]' >> "$GITHUB_OUTPUT"
           else
-            echo "Not main branch and PR title does not contain mi35x, only running on MI300..."
+            echo "Not main branch and PR title does not contain mi35x, only running on MI325..."
             echo 'standard_runners=["aiter-mi355-1gpu"]' >> "$GITHUB_OUTPUT"
             echo 'multigpu_runners=["aiter-mi355-8gpu"]' >> "$GITHUB_OUTPUT"
           fi
@@ -91,14 +104,14 @@ jobs:
           --name aiter_test \
           ${{ env.DOCKER_IMAGE }}
 
-      - name: Setup-Triton
+      - name: Setup Aiter
         run: |
           set -ex
-          echo "Setting up Triton..."
+          echo "Setting up Aiter..."
           docker exec \
           -w /workspace \
           aiter_test \
-          ./.github/scripts/build_triton.sh
+          bash -c "BUILD_TRITON=0 ./.github/scripts/build_aiter_triton.sh"
 
       - name: Tests
         run: |
@@ -177,7 +190,7 @@ jobs:
           docker exec \
           -w /workspace \
           aiter_test \
-          ./.github/scripts/build_triton.sh
+          bash -c "BUILD_TRITON=0 ./.github/scripts/build_aiter_triton.sh"
 
       - name: Tests
         run: |

.github/workflows/operators-tuning.yaml

@@ -1,13 +1,10 @@
 name: Operators Tuning
 
 on:
-  pull_request:
-    paths:
-      - 'aiter/configs/*untuned*.csv'
   workflow_dispatch:
     inputs:
       shapes:
-        description: 'Comma separated shape names to run (leave empty for all)'
+        description: 'Comma separated shape names to run, e.g. ck_batched_gemm_a8w8, ck_gemm_a8w8, ck_gemm_a8w8_blockscale, ck_gemm_a8w8_blockscale_bpreshuffle, ck_gemm_a8w8_bpreshuffle etc. (leave empty for all)'
         required: false
         default: ''
       arguments:
@@ -57,14 +54,14 @@ jobs:
           --name operators_tuning_test \
           rocm/pytorch:latest
 
-      - name: Setup-Triton
+      - name: Setup Aiter and Triton
         run: |
           set -ex
-          echo "Setting up Triton..."
+          echo "Setting up Aiter and Triton..."
           docker exec \
           -w /workspace \
           operators_tuning_test \
-          ./.github/scripts/build_triton.sh
+          bash -c "BUILD_TRITON=0 ./.github/scripts/build_aiter_triton.sh"
 
       - name: Show Computing Units
         run: |

.github/workflows/pre-checks.yaml

@@ -0,0 +1,88 @@
+name: Checks
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  check-ck:
+    name: Check Repository Dependency
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          submodules: 'recursive'
+
+      - name: Verify 3rdparty commits
+        run: ./.github/scripts/check_deps.sh
+
+  black:
+    name: Check Code Style with Black
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Run Black
+        uses: psf/black@stable
+
+  ruff:
+    name: Check Code Style with Ruff
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Set up Python environment
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.12"
+      - name: Install dependencies
+        run: pip3 install ruff
+      - name: Install reviewdog
+        uses: reviewdog/action-setup@e04ffabe3898a0af8d0fb1af00c188831c4b5893
+      - name: Run ruff with reviewdog
+        env:
+          REVIEWDOG_GITHUB_API_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          ruff check . -e | reviewdog -efm="%f:%l:%c: %m" -diff="git diff FETCH_HEAD" -reporter=github-pr-check -tee
+
+  upload-success-artifact:
+    name: Upload Success Signal
+    runs-on: ubuntu-latest
+    needs: [check-ck, black, ruff]
+    steps:
+      - name: Create success signal file
+        run: echo "success" > checks_signal.txt
+
+      - name: Upload success artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: checks-signal-${{ github.sha }}
+          path: checks_signal.txt
+
+  upload-failure-artifact:
+    name: Upload Failure Signal
+    runs-on: ubuntu-latest
+    needs: [check-ck, black, ruff]
+    if: ${{ always() && (needs.check-ck.result != 'success' || needs.black.result != 'success' || needs.ruff.result != 'success') }}
+    steps:
+      - name: Create failure signal file with failed jobs
+        run: |
+          echo "failure" > checks_signal.txt
+          if [ "${{ needs.check-ck.result }}" != "success" ]; then
+            echo "FAILED: check-ck (${{ needs.check-ck.result }})" >> checks_signal.txt
+          fi
+          if [ "${{ needs.black.result }}" != "success" ]; then
+            echo "FAILED: black (${{ needs.black.result }})" >> checks_signal.txt
+          fi
+          if [ "${{ needs.ruff.result }}" != "success" ]; then
+            echo "FAILED: ruff (${{ needs.ruff.result }})" >> checks_signal.txt
+          fi
+
+      - name: Upload failure artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: checks-signal-${{ github.sha }}
+          path: checks_signal.txt