Use configure_pytorch_release_matrix.py to drive the CI job matrix

.github/workflows/multi_arch_ci.yml

@@ -46,6 +46,10 @@ on:
         type: string
         default: ""
         description: "Comma-separated list of changed projects to filter tests (empty = run all tests)"
+      build_pytorch:
+        type: boolean
+        default: true
+        description: "Build PyTorch wheels"
   pull_request:
     types:
       - labeled
@@ -78,6 +82,7 @@ jobs:
       windows_test_labels: ${{ inputs.windows_test_labels || '' }}
       prebuilt_stages: ${{ inputs.prebuilt_stages || '' }}
       baseline_run_id: ${{ inputs.baseline_run_id || '' }}
+      build_pytorch: ${{ github.event_name != 'workflow_dispatch' || inputs.build_pytorch }}
 
   linux_build_and_test:
     name: Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }}

.github/workflows/multi_arch_ci_asan.yml

@@ -58,6 +58,7 @@ jobs:
       linux_test_labels: ${{ inputs.linux_test_labels || '' }}
       prebuilt_stages: ${{ inputs.prebuilt_stages || '' }}
       baseline_run_id: ${{ inputs.baseline_run_id || '' }}
+      build_pytorch: false
 
   linux_build_and_test:
     name: Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }}

.github/workflows/multi_arch_ci_linux.yml

@@ -14,6 +14,7 @@ on:
           build_variant_label, build_variant_cmake_preset,
           build_variant_suffix, build_pytorch,
           build_native_linux, test_python_packages_matrix,
+          pytorch_build_matrix,
           prebuilt_stages, baseline_run_id.
       test_labels:
         type: string
@@ -289,20 +290,26 @@ jobs:
   # inputs.kpack_split.
   build_pytorch_wheel_fat:
     needs: [build_python_packages]
-    name: Build PyTorch (fat + split) | ${{ matrix.pytorch_git_ref }}
-    if: ${{ !failure() && !cancelled() && fromJSON(inputs.build_config).build_pytorch == true }}
+    name: Build PyTorch (fat + split)
+    if: >-
+      ${{
+        !failure() &&
+        !cancelled() &&
+        fromJSON(inputs.build_config).build_pytorch == true &&
+        toJSON(fromJSON(inputs.build_config).pytorch_build_matrix) != '[]'
+      }}
     strategy:
       fail-fast: false
       matrix:
-        pytorch_git_ref: ["release/2.10", "release/2.11", "release/2.12"]
+        include: ${{ fromJSON(inputs.build_config).pytorch_build_matrix }}
     uses: ./.github/workflows/multi_arch_build_portable_linux_pytorch_wheels_ci.yml
     with:
       artifact_group: ${{ fromJSON(inputs.build_config).artifact_group }}
-      python_version: "3.12"
+      python_version: ${{ matrix.python_version }}
       pytorch_git_ref: ${{ matrix.pytorch_git_ref }}
       # The reusable workflow expands this to gfx targets via
       # cmake/therock_amdgpu_targets.cmake.
-      amdgpu_families: ${{ fromJSON(inputs.build_config).dist_amdgpu_families }}
+      amdgpu_families: ${{ matrix.amdgpu_families }}
       kpack_split: ${{ needs.build_python_packages.outputs.kpack_split }}
       rocm_package_find_links_url: ${{ needs.build_python_packages.outputs.package_find_links_url }}
       rocm_version: ${{ inputs.rocm_package_version }}

.github/workflows/multi_arch_ci_windows.yml

@@ -14,6 +14,7 @@ on:
           build_variant_label, build_variant_cmake_preset,
           build_variant_suffix, build_pytorch,
           build_native_linux, test_python_packages_matrix,
+          pytorch_build_matrix,
           prebuilt_stages, baseline_run_id.
       test_labels:
         type: string
@@ -206,15 +207,25 @@ jobs:
   build_pytorch_wheel_fat:
     needs: [build_python_packages]
     name: Build PyTorch (fat + split)
-    if: ${{ !failure() && !cancelled() && fromJSON(inputs.build_config).build_pytorch == true }}
+    if: >-
+      ${{
+        !failure() &&
+        !cancelled() &&
+        fromJSON(inputs.build_config).build_pytorch == true &&
+        toJSON(fromJSON(inputs.build_config).pytorch_build_matrix) != '[]'
+      }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include: ${{ fromJSON(inputs.build_config).pytorch_build_matrix }}
     uses: ./.github/workflows/multi_arch_build_windows_pytorch_wheels_ci.yml
     with:
       artifact_group: ${{ fromJSON(inputs.build_config).artifact_group }}
-      python_version: "3.12"
-      pytorch_git_ref: "release/2.10"
+      python_version: ${{ matrix.python_version }}
+      pytorch_git_ref: ${{ matrix.pytorch_git_ref }}
       # The reusable workflow expands this to gfx targets via
       # cmake/therock_amdgpu_targets.cmake.
-      amdgpu_families: ${{ fromJSON(inputs.build_config).dist_amdgpu_families }}
+      amdgpu_families: ${{ matrix.amdgpu_families }}
       kpack_split: ${{ needs.build_python_packages.outputs.kpack_split }}
       rocm_package_find_links_url: ${{ needs.build_python_packages.outputs.package_find_links_url }}
       rocm_version: ${{ inputs.rocm_package_version }}

.github/workflows/multi_arch_release.yml

@@ -105,6 +105,8 @@ jobs:
       build_variant: "release"
       release_type: ${{ inputs.release_type }}
       prerelease_version: ${{ inputs.prerelease_version }}
+      build_pytorch: ${{ inputs.build_pytorch }}
+      python_version: ${{ inputs.python_version }}
       linux_amdgpu_families: ${{ inputs.linux_amdgpu_families }}
       windows_amdgpu_families: ${{ inputs.windows_amdgpu_families }}
       repository: ${{ inputs.repository || github.repository }}

.github/workflows/multi_arch_release_asan.yml

@@ -86,6 +86,8 @@ jobs:
       prerelease_version: ${{ inputs.prerelease_version }}
       linux_amdgpu_families: ${{ inputs.linux_amdgpu_families || 'all' }}
       windows_amdgpu_families: "none"
+      build_pytorch: false
+      python_version: ${{ inputs.python_version }}
       repository: ${{ inputs.repository || github.repository }}
       ref: ${{ inputs.ref }}
 

.github/workflows/multi_arch_release_linux.yml

@@ -231,7 +231,7 @@ jobs:
   trigger_release_pytorch_wheels:
     needs: [build_artifacts, build_python_packages]
     name: Trigger Release PyTorch Wheels
-    if: ${{ inputs.build_pytorch != false }}
+    if: ${{ inputs.build_pytorch != false && fromJSON(inputs.build_config).build_pytorch == true }}
     runs-on: ubuntu-24.04
     permissions:
       actions: write

.github/workflows/multi_arch_release_linux_pytorch_wheels.yml

@@ -131,11 +131,13 @@ jobs:
           python ./build_tools/github_actions/configure_pytorch_release_matrix.py \
             --python-versions="${{ inputs.python_version }}" \
             --platform=linux \
+            --release-type="${{ inputs.release_type }}" \
             --amdgpu-families="${{ inputs.amdgpu_families }}"
 
   build_pytorch_wheels:
     needs: setup_matrix
     name: Build | py ${{ matrix.python_version }} | torch ${{ matrix.pytorch_git_ref }}
+    if: ${{ needs.setup_matrix.outputs.pytorch_matrix != '[]' }}
     strategy:
       fail-fast: false
       matrix:

.github/workflows/multi_arch_release_windows.yml

@@ -156,7 +156,7 @@ jobs:
   trigger_release_pytorch_wheels:
     needs: [build_artifacts, build_python_packages]
     name: Trigger Release PyTorch Wheels
-    if: ${{ inputs.build_pytorch != false }}
+    if: ${{ inputs.build_pytorch != false && fromJSON(inputs.build_config).build_pytorch == true }}
     runs-on: ubuntu-24.04
     permissions:
       actions: write

.github/workflows/multi_arch_release_windows_pytorch_wheels.yml

@@ -123,11 +123,13 @@ jobs:
           python ./build_tools/github_actions/configure_pytorch_release_matrix.py \
             --python-versions="${{ inputs.python_version }}" \
             --platform=windows \
+            --release-type="${{ inputs.release_type }}" \
             --amdgpu-families="${{ inputs.amdgpu_families }}"
 
   build_pytorch_wheels:
     needs: setup_matrix
     name: Build | py ${{ matrix.python_version }} | torch ${{ matrix.pytorch_git_ref }}
+    if: ${{ needs.setup_matrix.outputs.pytorch_matrix != '[]' }}
     strategy:
       fail-fast: false
       matrix:

.github/workflows/setup_multi_arch.yml

@@ -43,6 +43,14 @@ on:
         description: "Workflow run ID to copy prebuilt stage artifacts from."
         type: string
         default: ""
+      build_pytorch:
+        description: "Build PyTorch wheels."
+        type: boolean
+        default: true
+      python_version:
+        description: "Single Python version for wheel builds; empty uses the default matrix."
+        type: string
+        default: ""
       repository:
         description: "Repository to checkout. Defaults to github.repository."
         type: string
@@ -138,6 +146,8 @@ jobs:
         env:
           BUILD_VARIANT: ${{ inputs.build_variant }}
           RELEASE_TYPE: ${{ inputs.release_type }}
+          BUILD_PYTORCH: ${{ inputs.build_pytorch }}
+          PYTHON_VERSION: ${{ inputs.python_version }}
           LINUX_AMDGPU_FAMILIES: ${{ inputs.linux_amdgpu_families }}
           WINDOWS_AMDGPU_FAMILIES: ${{ inputs.windows_amdgpu_families }}
           LINUX_TEST_LABELS: ${{ inputs.linux_test_labels }}

build_tools/github_actions/configure_multi_arch_ci.py

@@ -65,6 +65,7 @@
     get_git_submodule_paths,
     is_ci_run_required,
 )
+from configure_pytorch_release_matrix import generate_pytorch_matrix_for_release_type
 from configure_rocm_python_test_matrix import build_rocm_python_test_matrix
 from github_actions_api import (
     gha_append_step_summary,
@@ -130,6 +131,8 @@ class CIInputs:
     base_ref: str  # Git ref for the workflow run (PR base or HEAD^1, used for diffing)
     build_variant: str  # Build variant label, e.g. "release", "asan", "tsan"
     release_type: str = "ci"  # "ci", or "dev", "nightly", "prerelease" for releases
+    build_pytorch: bool = True
+    python_versions: list[str] = field(default_factory=list)
 
     # PR labels (from event payload for pull_request events)
     pr_labels: list[str] = field(default_factory=list)
@@ -181,6 +184,8 @@ def from_environ() -> "CIInputs":
         # push before-commit) comes from the event payload.
         build_variant = os.environ.get("BUILD_VARIANT", "release")
         release_type = os.environ.get("RELEASE_TYPE", "ci")
+        build_pytorch = os.environ.get("BUILD_PYTORCH", "true").lower() != "false"
+        python_version = os.environ.get("PYTHON_VERSION", "").strip()
 
         pr_labels: list[str] = []
         base_ref = "HEAD^1"
@@ -215,6 +220,8 @@ def from_environ() -> "CIInputs":
             base_ref=base_ref,
             build_variant=build_variant,
             release_type=release_type,
+            build_pytorch=build_pytorch,
+            python_versions=[python_version] if python_version else [],
             pr_labels=pr_labels,
             linux_amdgpu_families=_parse_comma_list(
                 os.environ.get("LINUX_AMDGPU_FAMILIES", "")
@@ -429,6 +436,7 @@ class BuildConfig:
     build_native_linux: bool
     build_pytorch: bool
     test_python_packages_matrix: list[dict[str, str]] = field(default_factory=list)
+    pytorch_build_matrix: list[dict[str, str]] = field(default_factory=list)
     # Build runner label for this platform/variant combination
     build_runs_on: str = ""
     # Prebuilt stage configuration — set by configure() from JobDecisions.
@@ -663,15 +671,17 @@ def decide_jobs(
                 test_type_reason="ASAN tests skipped due to non-nightly trigger",
             )
 
+    build_pytorch_action = JobAction.RUN if ci_inputs.build_pytorch else JobAction.SKIP
+
     # Other jobs run unconditionally with no configuration.
     # TODO: job pruning: skip pytorch if only JAX has been edited, etc.
 
     return JobDecisions(
         build_rocm=build_rocm,
         test_rocm=test_rocm,
         build_rocm_python=JobGroupDecision(action=JobAction.RUN),
-        build_pytorch=JobGroupDecision(action=JobAction.RUN),
-        test_pytorch=JobGroupDecision(action=JobAction.RUN),
+        build_pytorch=JobGroupDecision(action=build_pytorch_action),
+        test_pytorch=JobGroupDecision(action=build_pytorch_action),
     )
 
 
@@ -952,11 +962,24 @@ def _expand_build_config_for_platform(
         return None
 
     family_names = [f["amdgpu_family"] for f in per_family_info]
+    dist_amdgpu_families = ";".join(family_names)
     suffix = variant_config.get("build_variant_suffix", "")
 
     # Select build runner using weighted distribution
     build_runs_on = select_build_runner(platform, build_variant)
 
+    pytorch_build_matrix: list[dict[str, str]] = []
+    build_pytorch = jobs.build_pytorch.action == JobAction.RUN
+    if build_pytorch:
+        pytorch_build_matrix = generate_pytorch_matrix_for_release_type(
+            release_type=ci_inputs.release_type,
+            amdgpu_families=dist_amdgpu_families,
+            platform=platform,
+            python_versions=ci_inputs.python_versions or None,
+        )
+        # Flip back to False if the generated matrix is empty.
+        build_pytorch = bool(pytorch_build_matrix)
+
     test_python_packages_matrix = build_rocm_python_test_matrix(
         per_family_info=per_family_info,
         platform=platform,
@@ -967,13 +990,14 @@ def _expand_build_config_for_platform(
 
     return BuildConfig(
         per_family_info=per_family_info,
-        dist_amdgpu_families=";".join(family_names),
+        dist_amdgpu_families=dist_amdgpu_families,
         artifact_group=f"multi-arch-{suffix or 'release'}",
         build_variant_label=variant_config["build_variant_label"],
         build_variant_suffix=suffix,
         build_variant_cmake_preset=variant_config["build_variant_cmake_preset"],
         build_native_linux=(suffix != "asan"),
-        build_pytorch=(suffix != "asan"),
+        build_pytorch=build_pytorch,
+        pytorch_build_matrix=pytorch_build_matrix,
         build_runs_on=build_runs_on,
         test_python_packages_matrix=test_python_packages_matrix,
         prebuilt_stages=jobs.build_rocm.prebuilt_stages,

build_tools/github_actions/configure_multi_arch_ci_summary.py

@@ -68,6 +68,10 @@ def format_summary(
     lines.append("")
     _append_test_rocm(lines, outputs)
 
+    lines.append("### build-pytorch")
+    lines.append("")
+    _append_build_pytorch(lines, outputs)
+
     return "\n".join(lines)
 
 
@@ -199,6 +203,32 @@ def _append_build_rocm(
         )
 
 
+def _append_build_pytorch(lines: list[str], outputs: CIOutputs) -> None:
+    lines.append("| Platform | Python | PyTorch ref | Families |")
+    lines.append("|----------|--------|-------------|----------|")
+
+    rows = 0
+    for platform, config in [
+        ("Linux", outputs.builds.linux),
+        ("Windows", outputs.builds.windows),
+    ]:
+        if config is None:
+            continue
+        for row in config.pytorch_build_matrix:
+            families = ", ".join(
+                f"`{family}`" for family in row["amdgpu_families"].split(";")
+            )
+            lines.append(
+                f"| {platform} | `{row['python_version']}` | "
+                f"`{row['pytorch_git_ref']}` | {families} |"
+            )
+            rows += 1
+
+    if rows == 0:
+        lines.append("| — | — | — | — |")
+    lines.append("")
+
+
 def _append_test_rocm(lines: list[str], outputs: CIOutputs) -> None:
     # Note: this assumes that the test_rocm job is never skipped.
     # We may decide to skip it under certain conditions in the future