From f665be9a4ce8a7a887bead4c7eb4d5a68482a0a2 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Wed, 15 Jul 2026 13:32:34 -0500 Subject: [PATCH 1/9] Enable additional single-GPU bazel tests in ROCm CI (#816) Drop 43 of the 45 test-target exclusions in build/rocm/ci_test_targets.txt, leaving only two: - buffer_callback_test_gpu, whose py_import deps are incompatible with pre-built plugin wheels (build_jaxlib=false). - pallas:ops_test_mgpu_gpu, a multi-accelerator test that is excluded from single_gpu runs anyway. Most of the dropped exclusions traced back to a complex-GEMM gap that no longer reproduces: C64 and C128 matmul return correct results on ROCm 7.14. The three teardown-hang exclusions (cholesky_update_test, pallas:pallas_vmap_test and pallas:gpu_paged_attention_test) each passed five consecutive runs with no hang. This brings the suite to 114 single-GPU GPU targets. Validated on MI300X (gfx942); the bazel CI pool is gfx950, so CI is the real gate. --- build/rocm/ci_test_targets.txt | 46 +++------------------------------- 1 file changed, 3 insertions(+), 43 deletions(-) diff --git a/build/rocm/ci_test_targets.txt b/build/rocm/ci_test_targets.txt index 55b28e7fc425..9174ea492bd7 100644 --- a/build/rocm/ci_test_targets.txt +++ b/build/rocm/ci_test_targets.txt @@ -7,50 +7,10 @@ //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test --//tests/pallas:pallas_test_gpu --//tests/pallas:ops_test_gpu + +# Multi-accelerator test (not single-GPU; excluded from single_gpu runs anyway). -//tests/pallas:ops_test_mgpu_gpu --//tests/pallas:pallas_shape_poly_test_gpu --//tests/pallas:pallas_vmap_test_gpu --//tests/pallas:triton_pallas_test_gpu --//tests/pallas:gpu_paged_attention_test_gpu --//tests:export_harnesses_multi_platform_test_gpu --//tests:jet_test_gpu --//tests:lax_autodiff_test_gpu --//tests:lax_numpy_setops_test_gpu --//tests:lax_numpy_test_gpu --//tests:lax_test_gpu --//tests:linalg_test_gpu --//tests:logging_test_gpu --//tests:random_lax_test_gpu --//tests:scipy_signal_test_gpu --//tests:stax_test_gpu --//tests:ode_test_gpu --//tests:lobpcg_test_gpu --//tests:scipy_stats_test_gpu --//tests:nn_test_gpu --//tests:lax_scipy_sparse_test_gpu --//tests:lax_scipy_spectral_dac_test_gpu --//tests:lax_scipy_special_functions_test_gpu --//tests:cholesky_update_test_gpu --//tests:api_test_gpu --//tests:ann_test_gpu --//tests:experimental_rnn_test_gpu --//tests:lax_vmap_test_gpu --//tests:qdwh_test_gpu --//tests:scaled_dot_test_gpu --//tests:scipy_spatial_test_gpu --//tests:shape_poly_test_gpu --//tests:sparsify_test_gpu --//tests:lax_numpy_reducers_test_gpu --//tests:scipy_optimize_test_gpu --//tests:lax_numpy_einsum_test_gpu --//tests:lax_scipy_test_gpu --//tests:sparse_test_gpu --//tests:sparse_bcoo_bcsr_test_gpu --//tests:svd_test_gpu --//tests:eigh_test_gpu --//tests:image_test_gpu + # buffer_callback_test_gpu uses py_import deps that are incompatible with # pre-built plugin wheels (build_jaxlib=false); exclude it unconditionally. -//tests:buffer_callback_test_gpu From 9e9f1a9af5d77946a0055119751595b19bccf733 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Tue, 4 Aug 2026 18:10:41 +0000 Subject: [PATCH 2/9] [ROCm] nightly: build jax/jaxlib on fork runners Upstream CI has stopped publishing jax/jaxlib nightly wheels to GCS (gs://jax-nightly-artifacts/latest), leaving the ROCm nightly test legs running against missing/stale base wheels. On the ROCm fork, build the base jax and jaxlib wheels on the same amd-do-linux.jax.cpu runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, so the ROCm test legs consume freshly built wheels instead of the broken GCS bucket. Upstream (jax-ml) behavior is unchanged. - ci/build_rocm_artifacts.sh: allow "jax"/"jaxlib" artifacts; skip auditwheel for the pure-python jax wheel. - wheel_tests_nightly_release.yml: new fork-only build-rocm-jax-jaxlib job; wire it into the ROCm test legs and select S3 as the jax/jaxlib source on the fork. - pytest_rocm.yml / bazel_rocm.yml: new optional jax-jaxlib-source input (default gcs), forwarded to the download action. - download-jax-rocm-wheels: when source is s3, fetch jax/jaxlib from the CloudFront listing alongside the ROCm wheels instead of GCS. (cherry picked from commit 8654db624806ff686405864af739aae8d888a431) --- .../download-jax-rocm-wheels/action.yml | 39 +++++++++++++-- .github/workflows/bazel_rocm.yml | 5 ++ .github/workflows/pytest_rocm.yml | 9 ++++ .../workflows/wheel_tests_nightly_release.yml | 50 ++++++++++++++++++- ci/build_rocm_artifacts.sh | 15 ++++-- 5 files changed, 110 insertions(+), 8 deletions(-) diff --git a/.github/actions/download-jax-rocm-wheels/action.yml b/.github/actions/download-jax-rocm-wheels/action.yml index 1b9c18c70302..b1f91eceed5f 100644 --- a/.github/actions/download-jax-rocm-wheels/action.yml +++ b/.github/actions/download-jax-rocm-wheels/action.yml @@ -27,6 +27,9 @@ inputs: s3_download_uri: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' + jax-jaxlib-source: + description: "Where to fetch the base jax/jaxlib wheels from: 'gcs' (upstream nightly bucket) or 's3' (built alongside the ROCm wheels on the fork's runners)" + default: 'gcs' permissions: {} runs: using: "composite" @@ -74,7 +77,11 @@ runs: mkdir -p $(pwd)/dist # Disable Authentication on the ROCm runners. gcloud config set auth/disable_credentials True - if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then + # jax is fetched from S3/CloudFront in the jaxlib "head" block below, + # alongside jaxlib and the ROCm plugin/pjrt wheels. + echo "jax wheel will be downloaded from S3/CloudFront alongside the ROCm wheels." + elif [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}"/jax*py3*none*any.whl $(pwd)/dist/ else echo "JAX wheel won't be downloaded, only jaxlib pre-built wheel is tested." @@ -86,7 +93,9 @@ runs: echo "JAX only release. Only downloading the jax wheel from the release bucket." else if [[ ${INPUTS_JAXLIB_VERSION} == "head" ]]; then - gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" != "s3" ]]; then + gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ + fi # Resolve ROCm plugin/PJRT wheels path via CloudFront. ROCM_WHEELS_BASE_URL="https://d22q5eopkfeftw.cloudfront.net" @@ -112,8 +121,31 @@ runs: exit 1 fi + DOWNLOAD_URLS=("${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}") + + # On the fork, the base jax/jaxlib wheels are built alongside the + # ROCm wheels and uploaded to the same S3 prefix, so fetch them from + # the same CloudFront listing instead of the (broken) GCS bucket. + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then + JAXLIB_FILE=$(echo "$FILES" | grep "jaxlib-" | grep "${PYTHON_MAJOR_MINOR}" | head -n1 || true) + if [[ -z "${JAXLIB_FILE}" ]]; then + echo "Could not find jaxlib wheel in ${WHEELS_URL}" + exit 1 + fi + DOWNLOAD_URLS+=("${WHEELS_URL}/${JAXLIB_FILE}") + + if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then + JAX_FILE=$(echo "$FILES" | grep -E "jax-[0-9].*-py3-none-any\.whl" | head -n1 || true) + if [[ -z "${JAX_FILE}" ]]; then + echo "Could not find jax wheel in ${WHEELS_URL}" + exit 1 + fi + DOWNLOAD_URLS+=("${WHEELS_URL}/${JAX_FILE}") + fi + fi + PYTHON_BIN="python${INPUTS_PYTHON}" - "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}" + "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${DOWNLOAD_URLS[@]}" elif [[ ${INPUTS_JAXLIB_VERSION} == "pypi_latest" ]]; then PYTHON=python${INPUTS_PYTHON} $PYTHON -m pip download jaxlib jax-rocm${JAXCI_ROCM_VERSION}-pjrt jax-rocm${JAXCI_ROCM_VERSION}-plugin --dest $(pwd)/dist/ --no-deps @@ -134,6 +166,7 @@ runs: INPUTS_PYTHON: ${{ inputs.python }} INPUTS_S3_DOWNLOAD_URI: ${{ inputs.s3_download_uri }} INPUTS_ROCM_FULL_VERSION: ${{ inputs.rocm-release-version || inputs.rocm-version }} + INPUTS_JAX_JAXLIB_SOURCE: ${{ inputs.jax-jaxlib-source }} - name: Skip the test run if the wheel artifacts were not downloaded successfully shell: bash if: steps.download-wheel-artifacts.outcome == 'failure' diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index 15d5bab2b5a0..13d57551ab46 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -87,6 +87,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' type: string + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + default: 'gcs' + type: string gate-targets-file: description: "Path to a file listing Bazel targets to run (one per line). If empty, runs the full default test suite." type: string @@ -181,6 +185,7 @@ jobs: jaxlib-version: ${{ inputs.jaxlib-version }} gcs_download_uri: ${{ inputs.gcs_download_uri }} s3_download_uri: ${{ inputs.s3_download_uri }} + jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }} # Halt for testing - name: Wait For Connection uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index 36efc218112e..2d68747bb714 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -66,6 +66,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" type: string default: 'latest' + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + type: string + default: 'gcs' halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -108,6 +112,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' type: string + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + default: 'gcs' + type: string halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -168,6 +176,7 @@ jobs: skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }} gcs_download_uri: ${{ inputs.gcs_download_uri }} s3_download_uri: ${{ inputs.s3_download_uri }} + jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }} - name: Install Python dependencies run: | $JAXCI_PYTHON -m pip install uv~=0.12.3 diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index cd0aa2a90422..e49c6e91abdb 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -221,9 +221,49 @@ jobs: s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' git_commit: ${{ inputs.git_commit }} + # Upstream CI frequently stops publishing jax/jaxlib nightly wheels to GCS. + # On the ROCm fork, build the base jax and jaxlib wheels ourselves on the same + # runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, + # so the ROCm test legs can consume them instead of the (stale) GCS bucket. + build-rocm-jax-jaxlib: + if: github.repository_owner == 'ROCm' + uses: ./.github/workflows/build_rocm_artifacts.yml + permissions: + id-token: write + contents: read + actions: read + strategy: + fail-fast: false + matrix: + runner: ["amd-do-linux.jax.cpu"] + artifact: ["jax", "jaxlib"] + python: ["3.12", "3.13", "3.14"] + rocm: [ + {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, + ] + exclude: + # jax is a pure-python (py3-none-any) wheel; build it once per ROCm + # variant rather than once per Python version. + - artifact: "jax" + python: "3.13" + - artifact: "jax" + python: "3.14" + name: "Build ROCm jax/jaxlib (${{ matrix.rocm.label }})" + with: + runner: ${{ matrix.runner }} + artifact: ${{ matrix.artifact }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.wheel-version }} + rocm-release-version: ${{ matrix.rocm.release-version }} + manylinux-image: ${{ matrix.rocm.manylinux-image }} + clone_main_xla: 0 + upload_artifacts_to_s3: true + s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + run-pytest-rocm: if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }} - needs: [build-rocm-artifacts] + needs: [build-rocm-artifacts, build-rocm-jax-jaxlib] permissions: id-token: write contents: read @@ -250,10 +290,13 @@ jobs: gcs_download_uri: ${{inputs.gcs_download_uri}} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' git_commit: ${{ inputs.git_commit }} + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} run-bazel-test-rocm: if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }} - needs: [build-rocm-artifacts] + needs: [build-rocm-artifacts, build-rocm-jax-jaxlib] permissions: id-token: write contents: read @@ -279,6 +322,9 @@ jobs: gcs_download_uri: ${{inputs.gcs_download_uri}} halt-for-connection: ${{inputs.halt-for-connection}} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} build_jaxlib: ${{ 'false' }} build_jax: ${{ 'false' }} jaxlib-version: "head" diff --git a/ci/build_rocm_artifacts.sh b/ci/build_rocm_artifacts.sh index 4bfacf04616a..77097036ce9e 100755 --- a/ci/build_rocm_artifacts.sh +++ b/ci/build_rocm_artifacts.sh @@ -37,7 +37,13 @@ fi # Set up the build environment. source "ci/utilities/setup_build_environment.sh" -allowed_artifacts=("jax-rocm-plugin" "jax-rocm-pjrt") +# The base "jax" and "jaxlib" wheels are also built here on the ROCm fork's +# runners so the ROCm test legs no longer depend on the (frequently broken) +# upstream nightly wheels published to GCS. jax is a pure-python wheel and +# jaxlib is the backend-neutral CPU base (ROCm support lives in the separate +# plugin/pjrt wheels), so both build correctly with the ROCm bazel config used +# below. +allowed_artifacts=("jax-rocm-plugin" "jax-rocm-pjrt" "jax" "jaxlib") if [[ ! " ${allowed_artifacts[*]} " =~ " ${artifact} " ]]; then echo "Error: Invalid artifact: $artifact. Allowed values are: ${allowed_artifacts[*]}" @@ -101,5 +107,8 @@ python build/build.py build --wheels="$artifact" \ $wheel_version_suffix_flag \ "${rocm_path_flags[@]}" -# Verify manylinux compliance. -./ci/utilities/run_auditwheel.sh +# Verify manylinux compliance. The pure-python "jax" wheel has no platform +# wheels to inspect (run_auditwheel.sh errors when it finds none), so skip it. +if [[ "$artifact" != "jax" ]]; then + ./ci/utilities/run_auditwheel.sh +fi From 864e4c9bbfb71879fc701ad029d1449f27edaae3 Mon Sep 17 00:00:00 2001 From: magaonka Date: Sun, 16 Aug 2026 22:17:56 -0500 Subject: [PATCH 3/9] [ROCm] continuous: build jax/jaxlib on fork runners The nightly flow already builds the base jax/jaxlib wheels on the fork's own runners (see "[ROCm] nightly: build jax/jaxlib on fork runners"), because upstream stopped reliably publishing them to gs://jax-nightly-artifacts/latest. The continuous flow was left behind: build-jax-artifact and build-jaxlib-artifact are gated on the jax-ml owner, so on the fork the ROCm legs silently fell back to that same stale GCS bucket. Apply the nightly pattern to the continuous flow so both consume freshly built wheels. Upstream (jax-ml) behavior is unchanged. - wheel_tests_continuous.yml: new fork-only build-rocm-jax-jaxlib job uploading to the same wheel-tests-continuous S3 prefix as the ROCm plugin/pjrt wheels; wired into the needs of both ROCm test legs, which now pass jax-jaxlib-source=s3 on the fork. - select_continuous_wheel_test_jobs.py: register build_rocm_jax_jaxlib as a build job required by run_pytest_rocm and run_bazel_test_rocm, and exclude it from per-job connection halts like the other ROCm jobs. Only Python 3.12 is built, matching the Python versions the continuous ROCm test legs actually consume; unlike plugin/pjrt these wheels are not published, so building 3.13/3.14 every three hours would be wasted work. --- .../select_continuous_wheel_test_jobs.py | 7 +++ .github/workflows/wheel_tests_continuous.yml | 59 +++++++++++++++++-- 2 files changed, 62 insertions(+), 4 deletions(-) diff --git a/.github/workflows/scripts/select_continuous_wheel_test_jobs.py b/.github/workflows/scripts/select_continuous_wheel_test_jobs.py index 41da70cc47db..99348f826108 100644 --- a/.github/workflows/scripts/select_continuous_wheel_test_jobs.py +++ b/.github/workflows/scripts/select_continuous_wheel_test_jobs.py @@ -39,6 +39,10 @@ class JobSpec: 'build_jaxlib_artifact', 'build_cuda_artifacts', 'build_rocm_artifacts', + # Fork-only. Upstream builds jax/jaxlib in build_jax_artifact and + # build_jaxlib_artifact, which are gated on the jax-ml owner; on the ROCm + # fork this job builds them instead. See wheel_tests_continuous.yml. + 'build_rocm_jax_jaxlib', ) # requires_builds lists artifact jobs whose uploaded wheels are consumed by @@ -118,6 +122,7 @@ class JobSpec: 'build_jax_artifact', 'build_jaxlib_artifact', 'build_rocm_artifacts', + 'build_rocm_jax_jaxlib', ), ), # LINT.ThenChange(.github/workflows/wheel_tests_continuous.yml:run_pytest_rocm) @@ -130,6 +135,7 @@ class JobSpec: 'build_jax_artifact', 'build_jaxlib_artifact', 'build_rocm_artifacts', + 'build_rocm_jax_jaxlib', ), ), # LINT.ThenChange(.github/workflows/wheel_tests_continuous.yml:run_bazel_test_rocm) @@ -145,6 +151,7 @@ class JobSpec: # callers keep halt-for-connection pinned to 'no'. ROCM_JOB_IDS = frozenset(( 'build_rocm_artifacts', + 'build_rocm_jax_jaxlib', 'run_pytest_rocm', 'run_bazel_test_rocm', )) diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 43b777f35adf..f5e6c4845cd0 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -24,9 +24,12 @@ # runs Bazel TPU tests with py_import. # 10. build-rocm-artifacts: Calls the `build_rocm_artifacts.yml` workflow to build ROCm plugin/pjrt # wheels and uploads them to an S3 bucket. -# 11. run-pytest-rocm: Calls the `pytest_rocm.yml` workflow which downloads the jaxlib and +# 11. build-rocm-jax-jaxlib: Fork-only. Calls the `build_rocm_artifacts.yml` workflow to build the +# base jax/jaxlib wheels into the same S3 prefix, since the upstream +# jax/jaxlib build jobs above only run under the jax-ml owner. +# 12. run-pytest-rocm: Calls the `pytest_rocm.yml` workflow which downloads the jaxlib and # ROCm artifacts and runs the ROCm tests. -# 12. run-bazel-test-rocm: Calls the `bazel_rocm.yml` workflow which runs the ROCm Bazel tests. +# 13. run-bazel-test-rocm: Calls the `bazel_rocm.yml` workflow which runs the ROCm Bazel tests. name: CI - Wheel Tests (Continuous) permissions: @@ -524,9 +527,51 @@ jobs: upload_artifacts_to_s3: true s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}' + # build-jax-artifact and build-jaxlib-artifact above are gated on the jax-ml + # owner, so on the fork the ROCm test legs would fall back to the upstream + # nightly GCS bucket, which is frequently stale or missing. Build the base + # jax and jaxlib wheels here instead, on the same runners and into the same + # S3 prefix as the ROCm plugin/pjrt wheels, mirroring what + # wheel_tests_nightly_release.yml already does. + build-rocm-jax-jaxlib: + needs: [select-jobs] + if: >- + ${{ + fromJSON(needs.select-jobs.outputs.required_builds || '{}').build_rocm_jax_jaxlib == true + && github.repository_owner == 'ROCm' + }} + uses: ./.github/workflows/build_rocm_artifacts.yml + permissions: + id-token: write + contents: read + actions: read + strategy: + fail-fast: false + matrix: + runner: ["amd-do-linux.jax.cpu"] + artifact: ["jax", "jaxlib"] + # Only the Python versions consumed by the ROCm test legs below are + # built here; unlike the plugin/pjrt wheels these are not published. + python: ["3.12"] + rocm: [ + {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + ] + name: "Build ROCm jax/jaxlib ${{ matrix.rocm.label }}" + with: + runner: ${{ matrix.runner }} + artifact: ${{ matrix.artifact }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.wheel-version }} + rocm-release-version: ${{ matrix.rocm.release-version }} + manylinux-image: ${{ matrix.rocm.manylinux-image }} + clone_main_xla: 0 + halt-for-connection: 'no' + upload_artifacts_to_s3: true + s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}' + # LINT.IfChange(run_pytest_rocm) run-pytest-rocm: - needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts] + needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts, build-rocm-jax-jaxlib] # LINT.ThenChange(.github/workflows/scripts/select_continuous_wheel_test_jobs.py:run_pytest_rocm) if: >- ${{ @@ -557,11 +602,14 @@ jobs: jaxlib-version: "head" gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri || 'gs://jax-nightly-artifacts/latest' }} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}' + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} halt-for-connection: 'no' # LINT.IfChange(run_bazel_test_rocm) run-bazel-test-rocm: - needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts] + needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts, build-rocm-jax-jaxlib] # LINT.ThenChange(.github/workflows/scripts/select_continuous_wheel_test_jobs.py:run_bazel_test_rocm) if: >- ${{ @@ -593,6 +641,9 @@ jobs: enable-x64: ${{ matrix.enable-x64 }} gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri || 'gs://jax-nightly-artifacts/latest' }} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}' + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} build_jaxlib: ${{ 'false' }} build_jax: ${{ 'false' }} jaxlib-version: "head" From 342ac1a8692724140cf46d75b6d6fe21e4605571 Mon Sep 17 00:00:00 2001 From: magaonka Date: Sun, 16 Aug 2026 22:21:47 -0500 Subject: [PATCH 4/9] [ROCm] Enable the 35 device-agnostic bazel tests in ROCm CI build/rocm/ci_test_targets.txt has always listed //tests:backend_independent_tests and //tests/pallas:backend_independent_tests, but none of those targets ever ran. jaxlib/jax.bzl defines those suites as exactly the targets carrying no jax_test_ tag, while run_bazel_test_rocm_rbe.sh passed --build_tag_filters/--test_tag_filters=jax_test_gpu,...,gpu for single_gpu. Bazel requires a target to match at least one positive tag, so the two suites were silently filtered down to nothing. Drop the positive tags from the single_gpu filter and keep only the exclusions. The target pattern file already scopes which suites run, so it is what should decide the set. This is also what the CUDA CI does -- ci/run_bazel_test_cuda_rbe.sh passes --test_tag_filters=-multiaccelerator with no positive tag -- which is why CUDA has been running these tests all along. The multi_gpu filter string is unchanged. Verified with bazel on gfx942 in ghcr.io/rocm/jax-base-ubu24.therock-7.14, using --nobuild to count selected targets: single_gpu, old filter: 112 test targets single_gpu, new filter: 147 test targets (+35) ci_blocking_test_targets.txt (PR gate): 1 target, unchanged The 35 newly enabled targets are all genuinely device-agnostic, e.g. config_test, tree_util_test, util_test, typing_test, xla_bridge_test, mesh_utils_test and pallas:pallas_cost_estimate_test. Nothing else is added: //jaxlib/tools:check_gpu_wheel_sources_test is tagged manual and stays excluded by -manual as before. --- ci/run_bazel_test_rocm_rbe.sh | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/ci/run_bazel_test_rocm_rbe.sh b/ci/run_bazel_test_rocm_rbe.sh index 1595d638145e..f2957905cee8 100755 --- a/ci/run_bazel_test_rocm_rbe.sh +++ b/ci/run_bazel_test_rocm_rbe.sh @@ -43,7 +43,10 @@ fi # Run Bazel GPU tests with RBE (single accelerator tests with one GPU apiece). echo "Running RBE GPU tests..." -TAG_FILTERS="jax_test_gpu,-config-cuda-only,-manual" +# Exclusions that apply to every config. Note that Bazel ORs *positive* tag +# filters and requires a target to match at least one of them, so adding a +# positive tag here narrows the run rather than widening it. +TAG_FILTERS="-config-cuda-only,-manual" # JAXCI_GATE_TARGETS_FILE selects which Bazel target pattern file to use. # Defaults to the full CI suite; set to build/rocm/ci_blocking_test_targets.txt @@ -52,10 +55,14 @@ TARGETS_FILE="${JAXCI_GATE_TARGETS_FILE:-build/rocm/ci_test_targets.txt}" for arg in "$@"; do if [[ "$arg" == "--config=multi_gpu" ]]; then - TAG_FILTERS="${TAG_FILTERS},multiaccelerator" + TAG_FILTERS="jax_test_gpu,${TAG_FILTERS},multiaccelerator" fi if [[ "$arg" == "--config=single_gpu" ]]; then - TAG_FILTERS="${TAG_FILTERS},gpu,-multiaccelerator" + # Deliberately no positive tag filter. The backend-independent suites + # in TARGETS_FILE are defined in jaxlib/jax.bzl as exactly the targets + # that carry no jax_test_ tag, so a positive jax_test_gpu or + # gpu filter drops all of them. TARGETS_FILE is what scopes the run. + TAG_FILTERS="${TAG_FILTERS},-multiaccelerator" fi done From 21783329e3dba6979e4c3113d140d5465bdbf153 Mon Sep 17 00:00:00 2001 From: magaonka Date: Sun, 16 Aug 2026 23:18:07 -0500 Subject: [PATCH 5/9] [ROCm] Enable buffer_callback_test_gpu in ROCm CI The exclusion claimed buffer_callback_test_gpu "uses py_import deps that are incompatible with pre-built plugin wheels (build_jaxlib=false)". That does not reproduce. The target was originally swept into the 45-target blanket TESTS_TO_IGNORE list added by ce621cf53a ("[ROCm] Include prebuilt rocm plugin pkgs as deps to tests"), and it was one of only two entries left behind when db2242f82f dropped the other 43 as stale; the stated reason was never verified against a real run. Verified on gfx942 in ghcr.io/rocm/jax-base-ubu24.therock-7.14, using the exact CI configuration -- --config=rocm_clang_hermetic --config=single_gpu with --//jax:build_jaxlib=false --//jax:build_jax=false and the prebuilt jax_rocm7_plugin/jax_rocm7_pjrt wheels staged in dist/ (the same wheels the CI download action stages, which is where the rocm_wheels repo rule picks them up): 51 cases: 23 pass, 28 skipped, 0 failures 5/5 consecutive runs pass, avg 3.5s, no teardown hang The 28 skips are @jtu.run_on_devices("cpu") cases that are CPU-only by design and skip identically on CUDA. The 23 that run are the 22 test_cuda_array_interface cases (11 dtypes x command_buffer_compatible) plus test_side_effect, so this adds real coverage of the __cuda_array_interface__ import path on ROCm. Brings the single-GPU suite to 148 targets. The blocking PR gate is unchanged at 1 target. Also corrects the comment on the remaining exclusion. ops_test_mgpu_gpu is a Mosaic GPU test, not a multi-accelerator one, and it carries no multiaccelerator tag -- verified by query -- so the claim that it is "excluded from single_gpu runs anyway" was wrong: the explicit exclusion is the only thing keeping it out. --- build/rocm/ci_test_targets.txt | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/build/rocm/ci_test_targets.txt b/build/rocm/ci_test_targets.txt index 9174ea492bd7..26b06ad599fc 100644 --- a/build/rocm/ci_test_targets.txt +++ b/build/rocm/ci_test_targets.txt @@ -8,9 +8,7 @@ //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test -# Multi-accelerator test (not single-GPU; excluded from single_gpu runs anyway). +# Mosaic GPU (mgpu) ops test. Mosaic GPU is an NVIDIA-only codegen path, so +# this never passes on ROCm. Note it carries no multiaccelerator tag, so this +# exclusion is what keeps it out of single_gpu runs -- not the tag filter. -//tests/pallas:ops_test_mgpu_gpu - -# buffer_callback_test_gpu uses py_import deps that are incompatible with -# pre-built plugin wheels (build_jaxlib=false); exclude it unconditionally. --//tests:buffer_callback_test_gpu From 9cb1d3da285b9796968825a45a472f88160215ec Mon Sep 17 00:00:00 2001 From: Jian Li Date: Wed, 26 Aug 2026 10:15:36 +0000 Subject: [PATCH 6/9] [ROCm] Enable previously skipped pallas tests on gfx950 Several fused attention tests were failing due to WAR hazard between VALU and GDS instructions which is fixed by https://github.com/llvm/llvm-project/pull/203770. So reenable these tests on ROCm platform. (cherry picked from commit fd31e390f222c6f75cb2c75291f333502cfbf587) --- tests/pallas/gpu_ops_test.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/pallas/gpu_ops_test.py b/tests/pallas/gpu_ops_test.py index 17681fe5b725..dae2b61f1594 100644 --- a/tests/pallas/gpu_ops_test.py +++ b/tests/pallas/gpu_ops_test.py @@ -115,19 +115,20 @@ def test_fused_attention_fwd( use_fwd, use_segment_ids, ): - # Skip specific failing tests due to autotuner issue on ROCm. Issue #34711 - # TODO(GulsumGudukbay): Unskip once fixed. + # Skip specific failing tests due to resource limits under sharding on ROCm. if jtu.is_device_rocm(): + device = jax.local_devices(backend="gpu")[0] + cc = getattr(device, "compute_capability", "").split(":", 1)[0] key = (batch_size, seq_len, num_heads, head_dim, tuple(block_sizes), causal, use_fwd, use_segment_ids) - if (key in { + if (cc == "gfx942" and key in { (1, 384, 1, 72, (("block_q", 64), ("block_k", 64)), False, True, True), # fwd5 (1, 384, 1, 72, (("block_q", 64), ("block_k", 128)), False, False, True), # fwd7 (2, 384, 1, 64, (("block_q", 64), ("block_k", 64)), True, False, True), # fwd8 (1, 384, 2, 72, (("block_q", 128), ("block_k", 128)), False, True, True), # fwd0 }): - self.skipTest("Skipped on ROCm due to autotuner issue.") + self.skipTest("Skipped on gfx942 because the test does not fit.") k1, k2, k3 = random.split(random.key(0), 3) q = random.normal( k1, (batch_size, seq_len, num_heads, head_dim), dtype=jnp.float16 @@ -227,13 +228,14 @@ def test_fused_attention_bwd( causal, use_segment_ids, ): - # Skip specific failing tests on ROCm due to autotuner issue. Issue #34711 - # TODO(GulsumGudukbay): Unskip once fixed. + # Skip specific failing tests due to resource limits under sharding on ROCm. if jtu.is_device_rocm(): + device = jax.local_devices(backend="gpu")[0] + cc = getattr(device, "compute_capability", "").split(":", 1)[0] key = (batch_size, seq_len, num_heads, head_dim, tuple(block_sizes), causal, use_segment_ids) - if key in { + if cc == "gfx942" and key in { (1, 128, 2, 64, (("block_q", 64), ("block_k", 128), ("block_q_dkv", 64), ("block_kv_dkv", 32), ("block_q_dq", 32), ("block_kv_dq", 64)), True, False), # bwd0 @@ -253,7 +255,7 @@ def test_fused_attention_bwd( ("block_kv_dkv", 64), ("block_q_dq", 64), ("block_kv_dq", 64)), True, False), # bwd9 }: - self.skipTest("Skipped on ROCm due to autotuner issue.") + self.skipTest("Skipped on gfx942 because the test does not fit.") if jtu.is_cuda_compute_capability_at_least("8.0"): # TODO(b/416306534) From 4c293031a7f7a8c90242dfee25becc05bb07181d Mon Sep 17 00:00:00 2001 From: magaonka-amd Date: Thu, 27 Aug 2026 14:18:31 -0500 Subject: [PATCH 7/9] [ROCm] Build jax/jaxlib against the TheRock 10.0 images The fork-only build-rocm-jax-jaxlib job was still pinned to TheRock 7.14 while every leg that consumes its output had moved to 10.0. Since release-version is interpolated into the shared S3 prefix, the pytest and bazel legs looked for jax/jaxlib under the 10.0.0 prefix while the build uploaded to the 7.14.0 one, and failed with "Could not find jaxlib wheel". This job block has no upstream counterpart, so it was not covered by the upstream switch of the pinned ROCm legs to 10.0. --- .github/workflows/wheel_tests_continuous.yml | 2 +- .github/workflows/wheel_tests_nightly_release.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index f5e6c4845cd0..89eb0d17b25e 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -554,7 +554,7 @@ jobs: # built here; unlike the plugin/pjrt wheels these are not published. python: ["3.12"] rocm: [ - {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "TheRock 10.0.0", wheel-version: "10", release-version: "10.0.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-10.0:latest" }, ] name: "Build ROCm jax/jaxlib ${{ matrix.rocm.label }}" with: diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index e49c6e91abdb..baa8c9c53e1f 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -239,7 +239,7 @@ jobs: artifact: ["jax", "jaxlib"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "TheRock 10.0.0", wheel-version: "10", release-version: "10.0.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-10.0:latest" }, {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, ] exclude: From e3b1a958108a5d15520cb0a94c716a0215f7929e Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Tue, 1 Sep 2026 19:49:50 +0000 Subject: [PATCH 8/9] Test PR CI trigger --- build/rocm/rocm.bazelrc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/build/rocm/rocm.bazelrc b/build/rocm/rocm.bazelrc index c493cb030986..20949f6ea156 100644 --- a/build/rocm/rocm.bazelrc +++ b/build/rocm/rocm.bazelrc @@ -45,6 +45,8 @@ test:rocm --flaky_test_attempts=3 test:rocm --test_verbose_timeout_warnings test:rocm --run_under=@xla//build_tools/rocm:parallel_gpu_execute +# Add this to test PR execution + # Configs for ROCm RBE tests build:single_gpu --repo_env=TF_ROCM_RBE_POOL="linux_x64_gpu_do_gfx950" build:single_gpu --repo_env=TF_ROCM_RBE_SINGLE_GPU_POOL="linux_x64_gpu_do_gfx950" From de4f966ee13df49d24969a5bf35e06272962bdcc Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Tue, 1 Sep 2026 19:55:50 +0000 Subject: [PATCH 9/9] Remove PR gate --- .github/workflows/bazel_rocm_presubmit.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index 6161df94c6e0..572b59ac92a1 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -28,7 +28,6 @@ concurrency: permissions: {} jobs: run-tests: - if: github.event_name == 'workflow_dispatch' || (github.event.repository.fork == false && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm')) uses: ./.github/workflows/bazel_rocm.yml # Begin Presubmit Naming Check - name modification requires internal check to be updated strategy: @@ -59,7 +58,6 @@ jobs: clone_main_xla: 0 blocking-gate: - if: github.event_name == 'workflow_dispatch' || (github.event.repository.fork == false && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm')) uses: ./.github/workflows/bazel_rocm.yml strategy: fail-fast: false