Use workflow matrix to deduplicate code (#168)

Author: oulgen

.github/workflows/test.yml

@@ -12,11 +12,12 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  test_cuda126_py310_a10g:
-    name: test-cuda12.6-py3.10-a10g
+  test_cuda126_a10g:
+    name: test-cuda12.6-py${{ matrix.python-version }}-a10g
     strategy:
       fail-fast: true
       matrix:
+        python-version: ["3.10", "3.12"]
         include:
           - name: A10G
             runs-on: linux.g5.4xlarge.nvidia.gpu
@@ -31,33 +32,7 @@ jobs:
       gpu-arch-version: ${{ matrix.gpu-arch-version }}
       submodules: recursive
       script: |
-        conda create -n venv python=3.10 -y
-        conda activate venv
-        python -m pip install --upgrade pip
-        pip install ${{ matrix.torch-spec }}
-        ./.github/scripts/install_triton.sh
-        pip install -r requirements.txt
-        python -m unittest discover -s test/ -p "*.py" -v -t .
-  test_cuda126_py312_a10g:
-    name: test-cuda12.6-py3.12-a10g
-    strategy:
-      fail-fast: true
-      matrix:
-        include:
-          - name: A10G
-            runs-on: linux.g5.4xlarge.nvidia.gpu
-            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu126'
-            gpu-arch-type: "cuda"
-            gpu-arch-version: "12.6"
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
-    with:
-      timeout: 60
-      runner: ${{ matrix.runs-on }}
-      gpu-arch-type: ${{ matrix.gpu-arch-type }}
-      gpu-arch-version: ${{ matrix.gpu-arch-version }}
-      submodules: recursive
-      script: |
-        conda create -n venv python=3.12 -y
+        conda create -n venv python=${{ matrix.python-version }} -y
         conda activate venv
         python -m pip install --upgrade pip
         pip install ${{ matrix.torch-spec }}