[CI] Pin py3.10 and one py3.12 on pytorch2.9 (#858)

Author: oulgen

.github/matrix.json

@@ -7,6 +7,7 @@
       "image": "nvidia/cuda:12.8.1-devel-ubuntu24.04",
       "runtime-version": "cu128",
       "container-options": "--gpus all",
+      "pytorch-version": "pytorch-2.9",
       "alias": "a10g"
     },
     {
@@ -16,6 +17,7 @@
       "image": "nvidia/cuda:12.8.1-devel-ubuntu24.04",
       "runtime-version": "cu128",
       "container-options": "--gpus all",
+      "pytorch-version": "pytorch-nightly",
       "alias": "a10g"
     },
     {
@@ -25,6 +27,7 @@
       "image": "nvidia/cuda:12.8.1-devel-ubuntu24.04",
       "runtime-version": "cu128",
       "container-options": "--gpus all",
+      "pytorch-version": "pytorch-2.9",
       "alias": "a10g-dtype-asserts",
       "dtype-asserts": true,
       "expecttest-accept": true
@@ -36,6 +39,7 @@
       "image": "nvidia/cuda:12.8.1-devel-ubuntu24.04",
       "runtime-version": "cu128",
       "container-options": "--gpus all",
+      "pytorch-version": "pytorch-nightly",
       "alias": "a10g-ref-eager"
     },
     {
@@ -45,6 +49,7 @@
       "image": "nvidia/cuda:12.8.1-devel-ubuntu24.04",
       "runtime-version": "cu128",
       "container-options": "--gpus all",
+      "pytorch-version": "pytorch-nightly",
       "alias": "h100"
     },
     {
@@ -54,6 +59,7 @@
       "image": "nvidia/cuda:13.0.1-devel-ubuntu24.04",
       "runtime-version": "cu130",
       "container-options": "--gpus all",
+      "pytorch-version": "pytorch-nightly",
       "alias": "b200"
     },
     {
@@ -63,6 +69,7 @@
       "image": "rocm/dev-ubuntu-24.04:7.0-complete",
       "runtime-version": "rocm7.0",
       "container-options": "--device=/dev/kfd --device=/dev/dri",
+      "pytorch-version": "pytorch-nightly",
       "alias": "mi325x"
     }
   ]

.github/workflows/test.yml

@@ -33,7 +33,7 @@ jobs:
       fail-fast: false
       matrix: ${{ fromJSON(needs.load-matrix.outputs.matrix) }}
 
-    name: test-${{ matrix.runtime-version }}-py${{ matrix.python-version }}-${{ matrix.alias }}
+    name: test-${{ matrix.runtime-version }}-py${{ matrix.python-version }}-${{ matrix.pytorch-version }}-${{ matrix.alias }}
 
     container:
       image: ${{ matrix.image }}
@@ -83,15 +83,21 @@ jobs:
           path: |
             ~/.cache/uv
             ~/.venv
-          key: ${{ matrix.python-version }}-${{ matrix.runtime-version }}-${{ hashFiles('.github/workflows/test.yml', 'requirements.txt') }}-${{ steps.date.outputs.month }}
+          key: ${{ matrix.python-version }}-${{ matrix.runtime-version }}-${{ matrix.pytorch-version }}-${{ hashFiles('.github/workflows/test.yml', 'requirements.txt') }}-${{ steps.date.outputs.month }}
 
       - name: Install PyTorch
         run: |
           source .venv/bin/activate
-          uv pip install -U --pre torch --index-url https://download.pytorch.org/whl/nightly/${{ matrix.runtime-version }}
+          if [[ "${{ matrix.pytorch-version }}" == "pytorch-2.9" ]]; then
+            # Install stable 2.9 from test channel
+            uv pip install -U "torch==2.9.*" --index-url https://download.pytorch.org/whl/test/${{ matrix.runtime-version }}
+          else
+            # Default to nightly
+            uv pip install -U --pre torch --index-url https://download.pytorch.org/whl/nightly/${{ matrix.runtime-version }}
+          fi
 
       - name: Install Triton
-        if: steps.cache.outputs.cache-hit != 'true'
+        if: steps.cache.outputs.cache-hit != 'true' && matrix.pytorch-version != 'pytorch-2.9'
         run: |
           set -x
           source .venv/bin/activate