Merge branch 'feat/blockwise_fp8_quant_triton_gemm_ker' of github.com:Degnel/ao into feat/blockwise_fp8_quant_triton_gemm_ker

Author: Degnel

.github/workflows/build-wheels_m1.yml

@@ -30,8 +30,8 @@ jobs:
       with-cuda: enable
       with-rocm: enable
       with-xpu: enable
-      # please note: excluding 3.13t for aarch64 builds for now
-      python-versions: '["3.9", "3.10", "3.11", "3.12", "3.13"]'
+      # Note: if free-threaded python is required add py3.13t here
+      python-versions: '["3.9"]'
 
   build:
     needs: generate-matrix

.github/workflows/build_wheels_aarch64_linux.yml

@@ -7,14 +7,12 @@ on:
       - 'gh/**'
     paths:
       - 'torchao/prototype/float8nocompile/**'
-      - '!torchao/prototype/float8nocompile/**'
   pull_request:
     branches:
       - main
       - 'gh/**'
     paths:
       - 'torchao/prototype/float8nocompile/**'
-      - '!torchao/prototype/float8nocompile/**'
 
 concurrency:
   group: floatnocompile_test-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}

.github/workflows/build_wheels_linux.yml

@@ -11,7 +11,7 @@ on:
       - 'gh/**'
 
 jobs:
-  test:
+  test-cpu-ops:
     strategy:
       matrix:
         runner: [macos-14]
@@ -33,9 +33,13 @@ jobs:
       - name: Install requirements
         run: |
           conda activate venv
-          pip install --extra-index-url "https://download.pytorch.org/whl/nightly/cpu" torch=="2.6.0.dev20250104"
+          # Install executorch first because it installs its own version
+          # of torch and torchao, which we do not want to use
+          pip install executorch
+          pip install torch --index-url "https://download.pytorch.org/whl/nightly/cpu" --force-reinstall
           pip install numpy
           pip install pytest
+          pip install parameterized
           USE_CPP=1 pip install .
       - name: Run python tests
         run: |
@@ -56,3 +60,61 @@ jobs:
           sh build_and_run_tests.sh
           rm -rf /tmp/cmake-out
           popd
+      - name: ET ops build
+        run: |
+          conda activate venv
+          pushd torchao/experimental
+          sh build_torchao_ops.sh executorch
+          popd
+
+  test-mps-ops:
+    strategy:
+      matrix:
+        runner: [macos-m1-stable]
+    runs-on: ${{matrix.runner}}
+    steps:
+      - name: Print machine info
+        run: |
+          uname -a
+          if [ $(uname -s) == Darwin ]; then
+            sysctl machdep.cpu.brand_string
+            sysctl machdep.cpu.core_count
+          fi
+      - name: Checkout repo
+        uses: actions/checkout@v3
+        with:
+          submodules: true
+      - name: Create conda env
+        run: |
+          conda create -yn test-mps-ops-env python=3.11
+      - name: Activate conda env
+        run: |
+          source activate base
+          conda activate test-mps-ops-env
+      - name: Install torch
+        run: |
+          pip install torch --index-url "https://download.pytorch.org/whl/nightly/cpu"
+      - name: Print torch version
+        run: |
+          python -c "import torch; print(torch.__version__)"
+      - name: Install requirements
+        run: |
+          pip install cmake
+          pip install parameterized
+          pip install pyyaml
+          pip install numpy
+      - name: Print pip freeze
+        run: |
+          pip freeze
+      - name: Print current directory
+        run: |
+          python -c "import os; print(os.getcwd())"
+      - name: Build ao with experimental mps ops
+        run: |
+          USE_CPP=1 TORCHAO_BUILD_EXPERIMENTAL_MPS=1 pip install .
+      - name: Run mps tests
+        run: |
+          pushd torchao/experimental/ops/mps/test
+          python test_lowbit.py
+          python test_quantizer.py
+          popd

.github/workflows/build_wheels_windows.yml

@@ -375,3 +375,4 @@ checkpoints/
 
 # Experimental
 torchao/experimental/cmake-out
+torchao/experimental/deps