diff --git a/.github/workflows/benchmark_rocm.yml b/.github/workflows/benchmark_rocm.yml
new file mode 100644
index 000000000000..74b51589e753
--- /dev/null
+++ b/.github/workflows/benchmark_rocm.yml
@@ -0,0 +1,198 @@
+# CI - Benchmark ROCm
+#
+# This workflow runs the ROCm benchmarks in ROCm team's GHCR containers.
+# It can be triggered manually via workflow_dispatch or called by other workflows
+# via workflow_call.
+#
+# It consists of the following job:
+# run-benchmarks:
+#    - Runs in ROCm team's container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest)
+#    - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from S3.
+#    - Executes the target benchmark scripts at `targets/<target>/run.sh`.
+name: CI - Benchmark ROCm
+on:
+  workflow_dispatch:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: choice
+        default: "linux-x86-64-8gpu-amd"
+        options:
+          - "linux-x86-64-1gpu-amd"
+          - "linux-x86-64-4gpu-amd"
+          - "linux-x86-64-8gpu-amd"
+      python:
+        description: "Which Python version to use?"
+        type: choice
+        default: "3.12"
+        options:
+          - "3.11"
+          - "3.12"
+      rocm-version:
+        description: "Which ROCm version to benchmark?"
+        type: choice
+        default: "7.2.0"
+        options:
+          - "7.2.0"
+      rocm-tag:
+        description: "ROCm tag for container image (e.g., rocm720)"
+        type: string
+        default: "rocm720"
+      target:
+        description: "Benchmark target"
+        type: choice
+        default: "maxtext"
+        options:
+          - "maxtext"
+      workload:
+        description: "Benchmark workload"
+        type: string
+        default: "llama3_8b"
+      jaxlib-version:
+        description: "Which jaxlib version to use? (head/pypi_latest)"
+        type: choice
+        default: "head"
+        options:
+          - "head"
+          - "pypi_latest"
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)"
+        type: choice
+        default: '0'
+        options:
+          - '0'
+          - '1'
+      gcs_download_uri:
+        description: "GCS location prefix from where the artifacts should be downloaded"
+        type: string
+        default: 'gs://jax-nightly-artifacts/latest'
+      s3_download_uri:
+        description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)"
+        type: string
+        default: 'latest'
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: string
+        default: 'no'
+  workflow_call:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: string
+        default: "linux-x86-64-8gpu-amd"
+      python:
+        description: "Which Python version to use?"
+        type: string
+        default: "3.12"
+      rocm-version:
+        description: "Which ROCm version to benchmark?"
+        type: string
+        default: "7.2.0"
+      rocm-tag:
+        description: "ROCm tag for container image (e.g., rocm720)"
+        type: string
+        default: "rocm720"
+      target:
+        description: "Benchmark target"
+        type: string
+        default: "maxtext"
+      workload:
+        description: "Benchmark workload"
+        type: string
+        default: "llama3_8b"
+      jaxlib-version:
+        description: "Which jaxlib version to use? (head/pypi_latest)"
+        type: string
+        default: "head"
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)"
+        type: string
+        default: '0'
+      gcs_download_uri:
+        description: "GCS location prefix from where the artifacts should be downloaded"
+        type: string
+        default: 'gs://jax-nightly-artifacts/latest'
+      s3_download_uri:
+        description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)"
+        type: string
+        default: 'latest'
+permissions:
+  id-token: write
+  contents: read
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
+jobs:
+  run-benchmarks:
+    defaults:
+      run:
+        # Set the shell to bash as GitHub actions run with /bin/sh by default
+        shell: bash
+    runs-on: ${{ inputs.runner }}
+    continue-on-error: true
+    # Run in ROCm team's GHCR container with GPU access
+    container:
+      image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest  # zizmor: ignore[unpinned-images]
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+      options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings
+    name: "${{ (contains(inputs.runner, '1gpu') && '1gpu') ||
+        (contains(inputs.runner, '4gpu') && '4gpu') ||
+        (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}"
+
+    env:
+      JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
+      JAXCI_PYTHON: "python${{ inputs.python }}"
+      JAXCI_ENABLE_X64: "0"
+      INPUT_ROCM_VERSION: "${{ inputs.rocm-version }}"
+      INPUT_RUNNER: "${{ inputs.runner }}"
+      INPUT_ROCM_TAG: "${{ inputs.rocm-tag }}"
+      
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+        with:
+          persist-credentials: false
+      - name: Download JAX ROCm wheels
+        uses: ./.github/actions/download-jax-rocm-wheels
+        with:
+          python: ${{ inputs.python }}
+          rocm-version: ${{ inputs.rocm-version }}
+          jaxlib-version: ${{ inputs.jaxlib-version }}
+          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 }}
+      - name: Install Python dependencies
+        run: |
+          $JAXCI_PYTHON -m pip install uv~=0.11.2
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Run ROCm benchmarks
+        env:
+          TARGET: ${{ inputs.target }}
+          WORKLOAD: ${{ inputs.workload }}
+        timeout-minutes: 120
+        run: ./ci/benchmark_targets/${TARGET}_rocm/run_${TARGET}_rocm.sh "${WORKLOAD}"
+      - name: Upload GitHub artifacts
+        if: always()
+        continue-on-error: true
+        uses: actions/upload-artifact@v4
+        with:
+          name: benchmark-artifacts-${{ inputs.target }}-${{ inputs.workload }}-${{ inputs.runner }}-py${{ inputs.python }}-rocm${{ inputs.rocm-version }}
+          path: ci/benchmark_targets/${{ inputs.target }}_rocm/run_artifacts/${{ inputs.workload }}/result.json
+          if-no-files-found: warn
+      # - name: Upload CI artifacts to S3
+      #   if: always()
+      #   env:
+      #     S3_BUCKET_NAME: jax-ci-amd
+      #   run: |
+      #     RUN_KEY="$(date -u +%F)_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}"
+      #     PREFIX="jax-ci-bench-logs/${GITHUB_REPOSITORY}/${GITHUB_REF_NAME}/${RUN_KEY}"
+      #     ./ci/upload_rocm_artifacts.sh \
+      #       "ci/benchmark_targets/${{ inputs.target }}_rocm/run_artifacts/${{ inputs.workload }}" \
+      #       "${PREFIX}" \
+      #       success
\ No newline at end of file
diff --git a/.github/workflows/pytest_rocm_abort.yml b/.github/workflows/pytest_rocm_abort.yml
new file mode 100644
index 000000000000..854b6681978b
--- /dev/null
+++ b/.github/workflows/pytest_rocm_abort.yml
@@ -0,0 +1,175 @@
+# CI - Pytest ROCm (Abort Support)
+#
+# This workflow runs the ROCm tests with Pytest in ROCm GHCR containers,
+# using the ROCm `pytest-abort` retry wrapper to detect/retry aborts/crashes.
+#
+# It can be triggered manually via workflow_dispatch or called by other workflows
+# via workflow_call.
+#
+# It consists of the following job:
+# run-tests:
+#    - Runs in ROCm container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest)
+#    - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from latest release.
+#    - Executes the `run_pytest_rocm_abort.sh` script, which installs wheel artifacts and
+#      runs the ROCm tests with Pytest under `pytest-abort-retry`.
+name: CI - Pytest ROCm (Abort Support)
+
+on:
+  workflow_dispatch:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: choice
+        default: "linux-x86-64-4gpu-amd"
+        options:
+          - "linux-x86-64-1gpu-amd"
+          - "linux-x86-64-4gpu-amd"
+          - "linux-x86-64-8gpu-amd"
+      python:
+        description: "Which Python version to use?"
+        type: choice
+        default: "3.11"
+        options:
+          - "3.11"
+          - "3.12"
+      rocm-version:
+        description: "Which ROCm version to test?"
+        type: choice
+        default: "7.2.0"
+        options:
+          - "7.2.0"
+      rocm-tag:
+        description: "ROCm tag for container image (e.g., rocm720)"
+        type: string
+        default: "rocm720"
+      jaxlib-version:
+        description: "Which jaxlib version to use? (head/pypi_latest)"
+        type: choice
+        default: "head"
+        options:
+          - "head"
+          - "pypi_latest"
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)"
+        type: choice
+        default: '0'
+        options:
+          - '0'
+          - '1'
+      gcs_download_uri:
+        description: "GCS location prefix from where the artifacts should be downloaded"
+        type: string
+        default: 'gs://jax-nightly-artifacts/latest'
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: string
+        default: 'no'
+      max-worker-restart:
+        description: "Max xdist worker restarts (passed to pytest --max-worker-restart)"
+        type: string
+        default: '50'
+  workflow_call:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: string
+        default: "linux-x86-64-4gpu-amd"
+      python:
+        description: "Which Python version to use?"
+        type: string
+        default: "3.11"
+      rocm-version:
+        description: "Which ROCm version to test?"
+        type: string
+        default: "7.2.0"
+      rocm-tag:
+        description: "ROCm tag for container image (e.g., rocm720)"
+        type: string
+        default: "rocm720"
+      jaxlib-version:
+        description: "Which jaxlib version to use? (head/pypi_latest)"
+        type: string
+        default: "head"
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)"
+        default: '0'
+        type: string
+      gcs_download_uri:
+        description: "GCS location prefix from where the artifacts should be downloaded"
+        default: 'gs://jax-nightly-artifacts/latest'
+        type: string
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: string
+        default: 'no'
+      max-worker-restart:
+        description: "Max xdist worker restarts (passed to pytest --max-worker-restart)"
+        type: string
+        default: '50'
+
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
+jobs:
+  run-tests:
+    defaults:
+      run:
+        # Set the shell to bash as GitHub actions run with /bin/sh by default
+        shell: bash
+    runs-on: ${{ inputs.runner }}
+    continue-on-error: true
+    # Run in ROCm GHCR container with GPU access
+    container:
+      image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+      options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings
+    name: "${{ (contains(inputs.runner, '1gpu') && '1gpu') ||
+        (contains(inputs.runner, '4gpu') && '4gpu') ||
+        (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}"
+
+    env:
+      JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
+      JAXCI_PYTHON: "python${{ inputs.python }}"
+      JAXCI_ENABLE_X64: "0"
+      MAX_WORKER_RESTART: "${{ inputs['max-worker-restart'] }}"
+
+    steps:
+      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+        with:
+          persist-credentials: false
+      - name: Download JAX ROCm wheels
+        uses: ./.github/actions/download-jax-rocm-wheels
+        with:
+          python: ${{ inputs.python }}
+          rocm-version: ${{ inputs.rocm-version }}
+          jaxlib-version: ${{ inputs.jaxlib-version }}
+          skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
+          gcs_download_uri: ${{ inputs.gcs_download_uri }}
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      - name: Install Python dependencies
+        run: |
+          $JAXCI_PYTHON -m pip install uv~=0.5.30
+          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Run Pytest ROCm tests (abort support)
+        timeout-minutes: 180
+        run: ./ci/run_pytest_rocm_abort.sh
+      - name: Upload pytest results to artifact
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: logs_abort
+          path: |
+            logs_abort/
+          if-no-files-found: warn
+          retention-days: 2
+          overwrite: true
diff --git a/.github/workflows/wheel_benchmarks_nightly_release.yml b/.github/workflows/wheel_benchmarks_nightly_release.yml
new file mode 100644
index 000000000000..76a7b0584552
--- /dev/null
+++ b/.github/workflows/wheel_benchmarks_nightly_release.yml
@@ -0,0 +1,63 @@
+# CI - Wheel Tests (Nightly/Release)
+#
+# This workflow is used to benchmark the JAX wheels that were built by internal CI jobs.
+#
+# 1. run-benchmark-rocm: Calls the `benchmark_rocm.yml` workflow which downloads the JAX wheels and
+#                     ROCm plugins from S3 and runs ROCm benchmarks.
+name: CI - Wheel Benchmarks (Nightly/Release)
+
+on:
+  workflow_dispatch:
+    inputs:
+      gcs_download_uri:
+        description: "GCS location URI from where the artifacts should be downloaded"
+        required: true
+        default: 'gs://jax-nightly-artifacts/latest'
+        type: string
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to download only the jax wheel from GCS (e.g for testing a jax only release)"
+        required: true
+        default: '0'
+        type: string
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection? (yes/no)'
+        required: false
+        default: 'no'
+        type: string
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
+jobs:
+  run-benchmark-rocm:
+    permissions:
+      id-token: write
+      contents: read
+      checks: write
+    uses: ./.github/workflows/benchmark_rocm.yml
+    strategy:
+        fail-fast: false # don't cancel all jobs on failure
+        matrix:
+          runner: ["linux-x86-64-8gpu-amd"]
+          python: ["3.12"]
+          rocm: [
+            {version: "7.2.0", tag: "rocm720"},
+          ]
+          target: ["maxtext"]
+          workload: ["gemma3-4b"]
+    name: "Benchmark ROCm (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
+    with:
+      runner: ${{ matrix.runner }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm.version }}
+      rocm-tag: ${{ matrix.rocm.tag }}
+      target: ${{ matrix.target }}
+      workload: ${{ matrix.workload }}
+      jaxlib-version: "head"
+      skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
+      gcs_download_uri: ${{ inputs.gcs_download_uri }}
diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml
index ef5fb02f7ced..c043a8c80eb3 100644
--- a/.github/workflows/wheel_tests_continuous.yml
+++ b/.github/workflows/wheel_tests_continuous.yml
@@ -34,8 +34,8 @@ permissions:
   actions: read
 
 on:
-  schedule:
-    - cron: "0 */3 * * *" # Run once every 3 hours
+  # schedule:
+  #  - cron: "0 */3 * * *" # Run once every 3 hours
   workflow_dispatch: # allows triggering the workflow run manually
 
 concurrency:
diff --git a/.github/workflows/wheel_tests_nightly_release_abort.yml b/.github/workflows/wheel_tests_nightly_release_abort.yml
new file mode 100644
index 000000000000..518306956eff
--- /dev/null
+++ b/.github/workflows/wheel_tests_nightly_release_abort.yml
@@ -0,0 +1,53 @@
+# CI - Wheel Tests (Nightly/Release) (ROCm abort only)
+#
+# This workflow runs only the ROCm wheel tests using the abort/retry wrapper workflow.
+name: CI - Wheel Tests (Nightly/Release) (ROCm abort only)
+
+on:
+  workflow_dispatch:
+    inputs:
+      gcs_download_uri:
+        description: "GCS location URI from where the artifacts should be downloaded"
+        required: true
+        default: 'gs://jax-nightly-artifacts/latest'
+        type: string
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to download only the jax wheel from GCS (e.g for testing a jax only release)"
+        required: true
+        default: '0'
+        type: string
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection? (yes/no)'
+        required: false
+        default: 'no'
+        type: string
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
+jobs:
+  run-pytest-rocm:
+    uses: ./.github/workflows/pytest_rocm_abort.yml
+    strategy:
+        fail-fast: false # don't cancel all jobs on failure
+        matrix:
+          runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"]
+          python: ["3.11", "3.12", "3.13", "3.14"]
+          rocm: [
+            {version: "7.2.0", tag: "rocm720"},
+          ]
+    name: "Pytest ROCm abort (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
+    with:
+      runner: ${{ matrix.runner }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm.version }}
+      rocm-tag: ${{ matrix.rocm.tag }}
+      jaxlib-version: "head"
+      skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
+      gcs_download_uri: ${{inputs.gcs_download_uri}}
+      halt-for-connection: ${{inputs.halt-for-connection}}
diff --git a/.gitignore b/.gitignore
index d30c019b31e8..e55a9a5de101 100644
--- a/.gitignore
+++ b/.gitignore
@@ -31,3 +31,7 @@ jax.iml
 /include/
 /lib/
 /share/
+
+/compile_commands.json
+/strace.txt
+/external
diff --git a/build/build.py b/build/build.py
index 38a11cedff46..d4368c444caa 100755
--- a/build/build.py
+++ b/build/build.py
@@ -667,6 +667,10 @@ async def main():
       )
 
   if "rocm" in args.wheels:
+    if not args.configure_only:
+      print("ERROR: This repo is not used for building the ROCm JAX plugins. Please use the new plugin repo: https://github.com/ROCm/rocm-jax")
+      exit(1)
+
     wheel_build_command_base.append("--config=rocm_base")
     wheel_build_command_base.append("--config=rocm")
     if clang_local:
diff --git a/build/rocm-test-requirements.txt b/build/rocm-test-requirements.txt
new file mode 100644
index 000000000000..399237175957
--- /dev/null
+++ b/build/rocm-test-requirements.txt
@@ -0,0 +1,24 @@
+absl-py
+build
+cloudpickle
+colorama>=0.4.4
+filelock
+flatbuffers
+hypothesis
+mpmath>=1.3
+pillow>=10.4.0
+# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t
+portpicker; python_version<"3.13"
+pytest-xdist
+pytest-json-report
+pytest-html
+pytest-csv
+pytest-rerunfailures
+pytest-html-merger
+pytest-reportlog
+wheel
+rich
+setuptools
+matplotlib
+opt-einsum
+auditwheel
diff --git a/ci/benchmark_targets/collect_bench_manifest_rocm.py b/ci/benchmark_targets/collect_bench_manifest_rocm.py
new file mode 100644
index 000000000000..f9844333a48c
--- /dev/null
+++ b/ci/benchmark_targets/collect_bench_manifest_rocm.py
@@ -0,0 +1,52 @@
+# Copyright 2026 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import argparse
+import json
+from pathlib import Path
+
+
+def read(path):
+    return Path(path).read_text(errors="replace")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--target", required=True)
+    parser.add_argument("--workload", required=True)
+    parser.add_argument("--run-code", required=True, type=int)
+    parser.add_argument("--run-started-at", required=True)
+    parser.add_argument("--run-completed-at", required=True)
+    parser.add_argument("--raw", action="append", default=[])
+    parser.add_argument("--out", required=True)
+    args = parser.parse_args()
+
+    benchmark = {
+        "target": args.target,
+        "workload": args.workload,
+        "run_code": args.run_code,
+        "run_started_at": args.run_started_at,
+        "run_completed_at": args.run_completed_at,
+    }
+
+    for raw in args.raw:
+        name, path = raw.split("=", 1)
+        benchmark[f"{name}_raw"] = read(path)
+
+    Path(args.out).write_text(json.dumps({"benchmark": benchmark}, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py
new file mode 100644
index 000000000000..ce425a31773a
--- /dev/null
+++ b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py
@@ -0,0 +1,122 @@
+# Copyright 2026 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# Compares workload benchmark metrics against configured baselines.
+#
+# Each workload defines:
+#   - which metrics are evaluated
+#   - how metric values are extracted from logs
+#   - whether lower or higher values are better
+#   - acceptable regression thresholds
+#
+# The output contains one comparison result per metric.
+import argparse
+import json
+import statistics
+from pathlib import Path
+import yaml
+
+
+def read(path):
+    return Path(path).read_text(errors="replace")
+
+
+def load_workload_config(path, workload):
+    config = yaml.safe_load(read(path)) or {}
+    if workload not in config:
+        raise KeyError(f"Missing workload in baseline config: {workload}")
+    return config[workload]
+
+
+def parse_metric_values(log, metric_config):
+    values = []
+    pattern = metric_config["log_pattern"]
+    field_index = int(metric_config["field_index"])
+    warmup_steps = int(metric_config.get("warmup_steps", 0))
+    for line in log.splitlines():
+        if pattern not in line:
+            continue
+        try:
+            value = line.split(",")[field_index]
+            values.append(float(value.split(":", 1)[1].strip()))
+        except (IndexError, ValueError):
+            continue
+
+    return values[warmup_steps:]
+
+
+def regression_percent(value, baseline, direction):
+    if direction == "lower_is_better":
+        regression = ((value - baseline) / baseline) * 100.0
+    elif direction == "higher_is_better":
+        regression((baseline - value) / baseline) * 100.0
+    else:
+        raise ValueError(f"Unknown comparison direction: {direction}")
+
+    return abs(regression)
+
+
+def evaluate_metric(log, name, metric_config):
+    values = parse_metric_values(log, metric_config)
+    value = statistics.median(values) if values else None
+
+    baseline = float(metric_config["baseline"])
+    threshold = float(metric_config["threshold_percent"])
+    direction = metric_config["direction"]
+
+    regression = None
+    cmp_code = 1
+    if value is not None and baseline != 0:
+        regression = regression_percent(value, baseline, direction)
+        cmp_code = int(regression > threshold)
+
+    return {
+        "name": name,
+        "value": round(value, 2) if value is not None else None,
+        "baseline": baseline,
+        "threshold_percent": threshold,
+        "direction": direction,
+        "regression_percent": (
+            round(regression, 4) if regression is not None else None
+        ),
+        "cmp_code": cmp_code,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--log", required=True)
+    parser.add_argument("--baseline", required=True)
+    parser.add_argument("--workload", required=True)
+    parser.add_argument("--out", required=True)
+    args = parser.parse_args()
+
+    log = read(args.log)
+    workload = load_workload_config(args.baseline, args.workload)
+
+    metrics = {
+        evaluate_metric(log, name, config)
+        for name, config in workload["metrics"].items()
+    }
+
+    Path(args.out).write_text(
+        json.dumps({"benchmark": {"metrics": metrics}}, indent=2) + "\n"
+    )
+
+    return max(metric["cmp_code"] for metric in metrics)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml
new file mode 100644
index 000000000000..5e887517a525
--- /dev/null
+++ b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml
@@ -0,0 +1,24 @@
+# Copyright 2026 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# Baseline for JAX benchmarks
+gemma3_4b:
+  benchmark_metric:
+    baseline: 2500
+    threshold: 0.10        # regression allowance (%)
+    direction: higher_is_better
+    log_pattern: completed_step
+    field_index: 3
+    warmup_step: 4
\ No newline at end of file
diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh
new file mode 100755
index 000000000000..dfb65578ec98
--- /dev/null
+++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh
@@ -0,0 +1,129 @@
+#!/bin/bash
+# Copyright 2026 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# Runs a MaxText ROCm benchmark workload.
+#
+# This script:
+#   - prepares the benchmark environment
+#   - executes the workload
+#   - collects benchmark run metadata
+#   - evaluates benchmark metrics
+#   - generates the final ROCm run manifest
+set -euo pipefail
+
+WORKLOAD="${1:-gemma3_4b}"
+
+JAX_DIR="${JAX_DIR:-$PWD}"
+
+PYTHON="${JAXCI_PYTHON:-python3}"
+PYTHON_VERSION="${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}"
+
+TARGET="maxtext_rocm"
+TARGET_DIR="${JAX_DIR}/ci/benchmark_targets/${TARGET}"
+
+MAXTEXT="${TARGET_DIR}/maxtext"
+RUN_DIR="${TARGET_DIR}/run_artifacts/${WORKLOAD}"
+
+mkdir -p "${RUN_DIR}"
+
+source "${JAX_DIR}/ci/envs/default.env"
+source "${JAX_DIR}/ci/utilities/install_wheels_locally.sh"
+
+if [[ ! -d "${MAXTEXT}/.git" ]]; then
+  git clone \
+    --depth 1 \
+    --branch add-rocm-benchmark-configs \
+    https://github.com/ROCm/maxtext.git \
+    "${MAXTEXT}"
+fi
+
+for file in \
+  "${MAXTEXT}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" \
+  "${MAXTEXT}/src/dependencies/requirements/requirements_rocm_benchmark.txt" \
+  "${TARGET_DIR}/baseline.yml"; do
+  [[ -f "${file}" ]] || {
+    echo "missing required file: ${file}" >&2
+    exit 2
+  }
+done
+
+"${PYTHON}" -m pip install -r \
+  "${MAXTEXT}/src/dependencies/requirements/requirements_rocm_benchmark.txt"
+
+export PY_COLORS=1
+export NCCL_DEBUG=WARN
+export TF_CPP_MIN_LOG_LEVEL=0
+
+export JAX_ENABLE_X64="${JAXCI_ENABLE_X64:-0}"
+
+export XLA_PYTHON_CLIENT_ALLOCATOR=platform
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+
+RUN_STARTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
+
+set +e
+
+pushd "${MAXTEXT}/src" >/dev/null
+
+"${PYTHON}" -m maxtext.trainers.pre_train.train \
+  "$(realpath "maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml")" \
+> "${RUN_DIR}/run.log" 2>&1
+
+RUN_CODE=$?
+
+popd >/dev/null
+
+set -e
+
+RUN_COMPLETED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
+
+"${PYTHON}" "${JAX_DIR}/ci/collect_bench_manifest_rocm.py" \
+  --target "${TARGET}" \
+  --workload "${WORKLOAD}" \
+  --run-code "${RUN_CODE}" \
+  --run-started-at "${RUN_STARTED_AT}" \
+  --run-completed-at "${RUN_COMPLETED_AT}" \
+  --raw workload_config="${MAXTEXT}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" \
+  --raw requirements="${MAXTEXT}/src/dependencies/requirements/requirements_rocm_benchmark.txt" \
+  --raw baseline_config="${TARGET_DIR}/baseline.yml" \
+  --out "${RUN_DIR}/benchmark_meta.json"
+
+CMP_CODE=0
+
+"${PYTHON}" "${TARGET_DIR}/cmp_maxtext_rocm.py" \
+  --log "${RUN_DIR}/run.log" \
+  --baseline "${TARGET_DIR}/baseline.yml" \
+  --workload "${WORKLOAD}" \
+  --out "${RUN_DIR}/benchmark_metrics.json" || CMP_CODE=$?
+
+"${PYTHON}" "${JAX_DIR}/ci/collect_run_manifest_rocm.py" \
+  --runner "${INPUT_RUNNER}" \
+  --python-version "${PYTHON_VERSION}" \
+  --python-bin "${PYTHON}" \
+  --rocm-version "${INPUT_ROCM_VERSION}" \
+  --rocm-tag "${INPUT_ROCM_TAG}" \
+  --extra "${RUN_DIR}/benchmark_meta.json" \
+  --extra "${RUN_DIR}/benchmark_metrics.json" \
+  --out "${RUN_DIR}/result.json"
+
+rm -f \
+  "${RUN_DIR}/run.log" \
+  "${RUN_DIR}/benchmark_meta.json" \
+  "${RUN_DIR}/benchmark_metrics.json"
+
+[[ -s "${RUN_DIR}/result.json" ]] && touch "${RUN_DIR}/_SUCCESS"
+
+exit $(( RUN_CODE != 0 || CMP_CODE != 0 ))
diff --git a/ci/collect_run_manifest_rocm.py b/ci/collect_run_manifest_rocm.py
new file mode 100644
index 000000000000..3175e41fd06e
--- /dev/null
+++ b/ci/collect_run_manifest_rocm.py
@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+
+# Collect ROCm CI run metadata and optionally merges additional JSON
+# payloads into the final result manifest.
+#
+# Extra JSON files are merged using dict.update() semantics. Fields in
+# extra payloads may override existing manifest fields if names collide.
+#
+# Benchmark/test-specific payloads should avoid using reserved top-level
+# run metadata keys such as:
+#  - github_*
+#  - run_*
+#  - python_version
+#  - rocm_
+
+#!/usr/bin/env python3
+# Collects ROCm CI run metadata and writes the final result manifest.
+#
+# Optional --extra JSON payloads are merged into the manifest using
+# dict.update() semantics. Extra payloads should avoid reserved run metadata
+# keys such as github_*, run_*, python_version, rocm_*, runner, and gpu_count.
+
+import argparse
+import json
+import os
+import re
+import subprocess
+import urllib.request
+from datetime import datetime, timezone
+from pathlib import Path
+
+
+def env(name, default="unknown"):
+    return os.environ.get(name, default)
+
+
+def run(args):
+    try:
+        return subprocess.check_output(
+            args, text=True, stderr=subprocess.DEVNULL
+        ).strip()
+    except Exception:
+        return ""
+
+
+def sh(command):
+    return run(["bash", "-lc", command])
+
+
+def utc_now():
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def get_json(url, headers=None):
+    request = urllib.request.Request(url, headers=headers or {})
+    with urllib.request.urlopen(request, timeout=30) as response:
+        return json.loads(response.read().decode())
+
+
+def get_response_headers(url, headers=None):
+    request = urllib.request.Request(url, headers=headers or {})
+    with urllib.request.urlopen(request, timeout=30) as response:
+        return response.headers
+
+
+def github_run_started_at():
+    repo = env("GITHUB_REPOSITORY", "")
+    run_id = env("GITHUB_RUN_ID", "")
+    try:
+        data = get_json(
+            f"https://api.github.com/repos/{repo}/actions/runs/{run_id}",
+            headers={"Accept": "application/vnd.github+json"},
+        )
+        return data.get("run_started_at", "")
+    except Exception:
+        return ""
+
+
+def gpu_count(runner):
+    match = re.search(r"([0-9]+)gpu", runner or "")
+    return int(match.group(1)) if match else None
+
+
+def image_digest(rocm_tag):
+    repo = f"rocm/jax-base-ubu24.{rocm_tag}"
+    try:
+        token_data = get_json(
+            f"https://ghcr.io/token?service=ghcr.io&scope=repository:{repo}:pull"
+        )
+        token = token_data.get("token", "")
+        if not token:
+            return ""
+        headers = get_response_headers(
+            f"https://ghcr.io/v2/{repo}/manifests/latest",
+            headers={
+                "Authorization": f"Bearer {token}",
+                "Accept": "application/vnd.docker.distribution.manifest.v2+json",
+            },
+        )
+        return headers.get("Docker-Content-Digest", "")
+    except Exception:
+        return ""
+
+
+def package_snapshot(python_bin):
+    pkgs = run([python_bin, "-m", "pip", "list", "--format=freeze"])
+    return "|".join(
+        line
+        for line in pkgs.splitlines()
+        if re.search(r"^(jax|jaxlib)==|pjrt|plugin", line)
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--runner", required=True)
+    parser.add_argument("--python-version", required=True)
+    parser.add_argument("--python-bin", default="python3")
+    parser.add_argument("--rocm-version", required=True)
+    parser.add_argument("--rocm-tag", required=True)
+    parser.add_argument("--extra", action="append", default=[])
+    parser.add_argument("--out", required=True)
+    args = parser.parse_args()
+    started = github_run_started_at()
+    date = started.split("T", 1)[0] if started else sh("date -u +%F")
+    repo = env("GITHUB_REPOSITORY")
+    run_id = env("GITHUB_RUN_ID")
+    attempt = env("GITHUB_RUN_ATTEMPT")
+    workflow = env("GITHUB_WORKFLOW", "")
+    manifest = {
+        "schema_version": 1,
+        "run_key": f"{date}_{run_id}_{attempt}",
+        "run_started_at": started,
+        "run_completed_at": utc_now(),
+        "github_run_url": f"https://github.com/{repo}/actions/runs/{run_id}",
+        "github_repository": repo,
+        "github_ref_name": env("GITHUB_REF_NAME"),
+        "github_ref": env("GITHUB_REF"),
+        "github_sha": env("GITHUB_SHA"),
+        "github_event_name": env("GITHUB_EVENT_NAME"),
+        "github_run_id": run_id,
+        "github_run_attempt": attempt,
+        "github_run_number": env("GITHUB_RUN_NUMBER"),
+        "github_workflow": workflow,
+        "is_nightly": "nightly" if "nightly" in workflow.lower() else "continuous",
+        "github_job": env("GITHUB_JOB"),
+        "python_version": args.python_version,
+        "rocm_version": args.rocm_version,
+        "rocm_tag": args.rocm_tag,
+        "gpu_count": gpu_count(args.runner),
+        "runner": args.runner,
+        "base_image_name": f"ghcr.io/rocm/jax-base-ubu24.{args.rocm_tag}:latest",
+        "base_image_digest": image_digest(args.rocm_tag),
+        "jax_packages_raw": package_snapshot(args.python_bin),
+        "wheels_sha_raw": sh("sha256sum dist/*.whl 2>/dev/null || true").replace(
+            "\n", "|"
+        ),
+    }
+
+    for extra in args.extra:
+        path = Path(extra)
+        if path.exists():
+            manifest.update(json.loads(path.read_text()))
+
+    Path(args.out).write_text(json.dumps(manifest, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/ci/run_pytest_rocm_abort.sh b/ci/run_pytest_rocm_abort.sh
new file mode 100755
index 000000000000..6c6ab6d18281
--- /dev/null
+++ b/ci/run_pytest_rocm_abort.sh
@@ -0,0 +1,182 @@
+#!/bin/bash
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Runs Pytest ROCm tests (with ROCm pytest-abort retry wrapper).
+# Requires the jaxlib and ROCm plugin wheels to be present inside $JAXCI_OUTPUT_DIR (../dist)
+#
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -x: log all commands
+# -o history: record shell history
+# -o allexport: export all functions and variables to be available to subscripts
+set -exu -o history -o allexport
+
+source ci/envs/default.env
+
+# Install jaxlib and ROCm plugin wheels inside the $JAXCI_OUTPUT_DIR directory
+echo "Installing wheels locally..."
+source ./ci/utilities/install_wheels_locally.sh
+
+# Print all the installed packages
+echo "Installed packages:"
+"$JAXCI_PYTHON" -m uv pip freeze
+
+"$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))"
+
+rocm-smi
+
+# ==============================================================================
+# Set up the generic test environment variables
+# ==============================================================================
+export PY_COLORS=1
+export JAX_SKIP_SLOW_TESTS=true
+export NCCL_DEBUG=WARN
+export TF_CPP_MIN_LOG_LEVEL=0
+export JAX_ENABLE_X64="$JAXCI_ENABLE_X64"
+
+# ==============================================================================
+# Calculate the optimal number of parallel processes for pytest
+# This will be the minimum of: GPU capacity, CPU core count, and a system RAM limit.
+# ==============================================================================
+
+export gpu_count=$(rocminfo | egrep -c "Device Type:\\s+GPU")
+echo "Number of GPUs detected: $gpu_count"
+
+# Query GPU 0 memory using rocm-smi
+export memory_per_gpu_mib=$(rocm-smi -d 0 --showmeminfo vram | grep -i "vram total" | awk '{print int($NF/1024/1024)}' | head -1)
+echo "Reported memory per GPU: $memory_per_gpu_mib MiB"
+
+# Convert effective memory from MiB to GiB.
+export memory_per_gpu_gib=$((memory_per_gpu_mib / 1024))
+echo "Effective memory per GPU: $memory_per_gpu_gib GiB"
+
+# Allow 2 GiB of GPU RAM per test.
+export max_tests_per_gpu=$((memory_per_gpu_gib / 2))
+echo "Max tests per GPU (assuming 2GiB/test): $max_tests_per_gpu"
+
+export num_processes=$((gpu_count * max_tests_per_gpu))
+echo "Initial number of processes based on GPU capacity: $num_processes"
+
+export num_cpu_cores=$(nproc)
+echo "Number of CPU cores available: $num_cpu_cores"
+
+# Reads total memory from /proc/meminfo (in KiB) and converts to GiB.
+export total_ram_gib=$(awk '/MemTotal/ {printf "%.0f", $2/1048576}' /proc/meminfo)
+echo "Total system RAM: $total_ram_gib GiB"
+
+# Set a safety limit for system RAM usage, e.g., 1/6th of total.
+export host_memory_limit=$((total_ram_gib / 6))
+echo "Host memory process limit (1/6th of total RAM): $host_memory_limit"
+
+if [[ $num_cpu_cores -lt $num_processes ]]; then
+  num_processes=$num_cpu_cores
+  echo "Adjusting num_processes to match CPU core count: $num_processes"
+fi
+
+if [[ $host_memory_limit -lt $num_processes ]]; then
+  num_processes=$host_memory_limit
+  echo "Adjusting num_processes to match host memory limit: $num_processes"
+fi
+
+if [[ 16 -lt $num_processes ]]; then
+  num_processes=16
+  echo "Reducing num_processes to $num_processes"
+fi
+
+echo "Final number of processes to run: $num_processes"
+
+export JAX_ENABLE_ROCM_XDIST="$gpu_count"
+export XLA_PYTHON_CLIENT_ALLOCATOR=platform
+export XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1 --xla_gpu_enable_nccl_comm_splitting=false --xla_gpu_enable_command_buffer="
+
+# Disable core dumps just in case
+ulimit -c 0
+
+# Keep deselected tests in one place for the abort wrapper.
+ROCM_PYTEST_DESELECT_ARGS=(
+  --deselect=tests/multi_device_test.py::MultiDeviceTest::test_computation_follows_data
+  --deselect=tests/multiprocess_gpu_test.py::MultiProcessGpuTest::test_distributed_jax_visible_devices
+  --deselect=tests/compilation_cache_test.py::CompilationCacheTest::test_task_using_cache_metric
+)
+
+# --max-runs: retry the entire pytest run up to N times on abort/crash.
+# --max-worker-restart: restart crashed xdist workers up to N times.
+# --maxfail: stop the run after N test failures.
+rocm_test_cmd() {
+  local abort_flag="${1:-0}"
+  shift
+  if [[ "$abort_flag" == "1" ]]; then
+    pytest-abort-retry --max-runs 3 --clear-crash-log -- "$@"
+  else
+    "$@"
+  fi
+}
+
+rocm_log_tail_on_failure() {
+  local logfile="$1"
+  local status="$2"
+  if [[ "$status" -ne 0 ]]; then
+    echo "Pytest failed (exit=$status). Showing last 200 lines of $logfile:"
+    tail -n 200 "$logfile" || true
+  else
+    echo "Pytest output saved to $logfile (uploaded as artifact)."
+  fi
+}
+
+rocm_install_extra_requirements() {
+  if [[ -n "${GITHUB_WORKSPACE:-}" ]]; then
+    cd "$GITHUB_WORKSPACE"
+  fi
+
+  # Install extra requirements.
+  "$JAXCI_PYTHON" -m uv pip install pytest-timeout pytest-html pytest-csv pytest-json-report pytest-abort
+}
+
+rocm_install_extra_requirements
+
+echo "Running ROCm tests (with abort/retry wrapper)..."
+mkdir -p logs_abort
+logfile="logs_abort/jax_ToT_UT_abort.log"
+
+# Allow the workflow to override worker restart limit.
+max_worker_restart="${MAX_WORKER_RESTART:-50}"
+
+# pytest-abort output directories (must be set before running pytest).
+export PYTEST_ABORT_LAST_RUNNING_DIR="logs_abort/last_running"
+export PYTEST_ABORT_CRASHED_TESTS_LOG="logs_abort/crashed_tests.jsonl"
+mkdir -p "$PYTEST_ABORT_LAST_RUNNING_DIR"
+
+set +e
+rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart="$max_worker_restart" --tb=short --timeout=1200 --timeout-method=thread tests \
+  "${ROCM_PYTEST_DESELECT_ARGS[@]}" \
+  --json-report \
+  --json-report-file=logs_abort/tests-report-abort.json \
+  --csv=logs_abort/tests-report-abort.csv \
+  --html=logs_abort/tests-report-abort.html \
+  --self-contained-html \
+  >"$logfile" 2>&1
+pytest_status=$?
+set -e
+rocm_log_tail_on_failure "$logfile" "$pytest_status"
+
+echo "Postprocessing reports with crashed tests..."
+pytest-abort-postprocess \
+  --crash-log "$PYTEST_ABORT_CRASHED_TESTS_LOG" \
+  --json-report logs_abort/tests-report-abort.json \
+  --html-report logs_abort/tests-report-abort.html \
+  --csv-report logs_abort/tests-report-abort.csv \
+  >>"$logfile" 2>&1
+
+exit "$pytest_status"
diff --git a/ci/upload_rocm_artifacts.sh b/ci/upload_rocm_artifacts.sh
new file mode 100644
index 000000000000..bdae41f03703
--- /dev/null
+++ b/ci/upload_rocm_artifacts.sh
@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Uploads a file or directory to S3.
+#
+# Usage:
+#   upload_rocm_artifacts.sh <source_path> <s3_dest_prefix> [success]
+#
+# If the optional third argument is "success", an empty _SUCCESS marker is
+# written after the upload completes.
+
+: "${S3_BUCKET_NAME:?}"
+
+SRC="${1:?missing source path}"
+DEST_PREFIX="${2:?missing S3 destination prefix}"
+WRITE_SUCCESS="${3:-}"
+
+[[ -e "${SRC}" ]] || {
+ echo "missing source path: ${SRC}" >&2
+ exit 2
+}
+
+DEST="s3://${S3_BUCKET_NAME}/${DEST_PREFIX}"
+
+echo "[upload] ${SRC} -> ${DEST}"
+
+if [[ -d "${SRC}" ]]; then
+  aws s3 cp --only-show-errors "${SRC}" "${DEST}" --recursive
+else
+  aws s3 cp --only-show-errors "${SRC}" "${DEST}/$(basename "${SRC}")"
+fi
+
+if [[ "${WRITE_SUCCESS}" == "success" ]]; then
+  printf '' | aws s3 cp --only-show-errors - "${DEST}/_SUCCESS"
+fi
+
+echo "[done]"
diff --git a/conftest.py b/conftest.py
index cd8414369b89..8f905bf53af8 100644
--- a/conftest.py
+++ b/conftest.py
@@ -16,7 +16,6 @@
 import os
 import pytest
 
-
 @pytest.fixture(autouse=True)
 def add_imports(doctest_namespace):
   import jax
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 1c348250dd7d..1433bb540ecf 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -634,9 +634,8 @@ def make_ir_context() -> ir.Context:
   # multi threaded execution aborts the process if we try to register a new
   # dialect after this point. The dialect registry in a context is not thread
   # safe, and a fatal error is much better than a data race.
-  # jax_mlir_ext.enter_multi_threaded_execution(context)
-  # TODO(phawkins): clean up users who add their own dialects to JAX's contexts
-  # and enable this.
+  # if jaxlib_version >= (0, 8):
+  #  jax_mlir_ext.enter_multi_threaded_execution(context)
   return context
 
 
diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py
index 5d11761386bb..ab0cc3849503 100644
--- a/jax/_src/test_util.py
+++ b/jax/_src/test_util.py
@@ -25,6 +25,7 @@
 
 import math
 import os
+from pathlib import Path
 import platform
 import re
 import sys
@@ -389,6 +390,15 @@ def supported_dtypes() -> set[DTypeLike]:
 def is_device_rocm() -> bool:
   return 'rocm' in xla_bridge.get_backend().platform_version
 
+def get_rocm_version():
+  rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
+  version_path = Path(rocm_path) / ".info" / "version"
+  if not version_path.exists():
+    raise FileNotFoundError(f"Expected ROCm version file at {version_path}")
+  version_str = version_path.read_text().strip()
+  major, minor, *_ = version_str.split(".")
+  return int(major), int(minor)
+
 def is_device_cuda() -> bool:
   return 'cuda' in xla_bridge.get_backend().platform_version
 
@@ -610,7 +620,14 @@ def skip_on_devices(*disabled_devices, skip_reason=None):
     skip_reason: Optional custom skip message when test is skipped.
   """
   if skip_reason is None:
-    skip_reason = "Skipped on devices with tags: " + ", ".join(disabled_devices)
+    skip_messages = {
+      ("gpu",): "Skipped on all GPUs.",
+      ("cpu",): "Skipped on CPU.",
+      ("tpu",): "Skipped on TPU.",
+      ("cuda",): "Skipped on CUDA GPUs.",
+      ("rocm",): "Skipped on ROCm GPUs.",
+    }
+    skip_reason = skip_messages.get(disabled_devices)
   return _device_filter(lambda: not test_device_matches(disabled_devices), skip_reason)
 
 def run_on_devices(*enabled_devices, skip_reason=None):
@@ -621,9 +638,14 @@ def run_on_devices(*enabled_devices, skip_reason=None):
     skip_reason: Optional custom skip message when test is skipped.
   """
   if skip_reason is None:
-    skip_reason = (
-      "Skipped unless running on devices with tags: " + ", ".join(enabled_devices)
-    )
+    device_specific_skip_reasons = {
+        ("cpu",): "Skipped: CPU-only test.",
+        ("tpu",): "Skipped: TPU-only test.",
+        ("gpu",): "Skipped: GPU-only test.",
+        ("rocm",): "Skipped: ROCm-only test.",
+        ("cuda",): "Skipped: CUDA-only test.",
+    }
+    skip_reason = device_specific_skip_reasons.get(enabled_devices)
   return _device_filter(lambda: test_device_matches(enabled_devices), skip_reason)
 
 def device_supports_buffer_donation():
diff --git a/jaxlib/BUILD b/jaxlib/BUILD
index 1dd304f5bbff..562cf538a514 100644
--- a/jaxlib/BUILD
+++ b/jaxlib/BUILD
@@ -41,7 +41,7 @@ licenses(["notice"])
 
 package(
     default_applicable_licenses = [],
-    default_visibility = ["//jax:internal"],
+    default_visibility = ["//visibility:public"],
 )
 
 package_group(
diff --git a/jaxlib/gpu/solver_interface.cc b/jaxlib/gpu/solver_interface.cc
index 7ebcf7483652..e530e5bc1f2f 100644
--- a/jaxlib/gpu/solver_interface.cc
+++ b/jaxlib/gpu/solver_interface.cc
@@ -63,8 +63,8 @@ JAX_GPU_DEFINE_GETRF(gpuDoubleComplex, gpusolverDnZgetrf);
 
 JAX_GPU_DEFINE_GETRF_BATCHED(float, gpublasSgetrfBatched);
 JAX_GPU_DEFINE_GETRF_BATCHED(double, gpublasDgetrfBatched);
-JAX_GPU_DEFINE_GETRF_BATCHED(gpublasComplex, gpublasCgetrfBatched);
-JAX_GPU_DEFINE_GETRF_BATCHED(gpublasDoubleComplex, gpublasZgetrfBatched);
+JAX_GPU_DEFINE_GETRF_BATCHED(gpuComplex, gpublasCgetrfBatched);
+JAX_GPU_DEFINE_GETRF_BATCHED(gpuDoubleComplex, gpublasZgetrfBatched);
 #undef JAX_GPU_DEFINE_GETRF_BATCHED
 
 // QR decomposition: geqrf
@@ -102,8 +102,8 @@ JAX_GPU_DEFINE_GEQRF(gpuDoubleComplex, gpusolverDnZgeqrf);
 
 JAX_GPU_DEFINE_GEQRF_BATCHED(float, gpublasSgeqrfBatched);
 JAX_GPU_DEFINE_GEQRF_BATCHED(double, gpublasDgeqrfBatched);
-JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasComplex, gpublasCgeqrfBatched);
-JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasDoubleComplex, gpublasZgeqrfBatched);
+JAX_GPU_DEFINE_GEQRF_BATCHED(gpuComplex, gpublasCgeqrfBatched);
+JAX_GPU_DEFINE_GEQRF_BATCHED(gpuDoubleComplex, gpublasZgeqrfBatched);
 #undef JAX_GPU_DEFINE_GEQRF_BATCHED
 
 // Householder transformations: orgqr
@@ -302,8 +302,8 @@ JAX_GPU_DEFINE_SYEVD(gpuDoubleComplex, gpusolverDnZheevd);
 
 JAX_GPU_DEFINE_SYRK(float, gpublasSsyrk);
 JAX_GPU_DEFINE_SYRK(double, gpublasDsyrk);
-JAX_GPU_DEFINE_SYRK(gpublasComplex, gpublasCsyrk);
-JAX_GPU_DEFINE_SYRK(gpublasDoubleComplex, gpublasZsyrk);
+JAX_GPU_DEFINE_SYRK(gpuComplex, gpublasCsyrk);
+JAX_GPU_DEFINE_SYRK(gpuDoubleComplex, gpublasZsyrk);
 #undef JAX_GPU_DEFINE_SYRK
 
 // Singular Value Decomposition: gesvd
diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h
index 9adb1f11776f..fb9c1f407145 100644
--- a/jaxlib/gpu/vendor.h
+++ b/jaxlib/gpu/vendor.h
@@ -457,6 +457,8 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32;
 
 #elif defined(JAX_GPU_HIP)
 
+#define HIPBLAS_V2 1
+
 // IWYU pragma: begin_exports
 #include "rocm/include/hip/hip_cooperative_groups.h"
 #include "rocm/include/hip/hip_runtime_api.h"
@@ -477,17 +479,11 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32;
 // MIOpen lib. Remove when MIOpen support is complete.
 #define MIOPEN_STATUS_SUCCESS 0
 
-typedef hipFloatComplex gpuComplex;
+typedef hipComplex gpuComplex;
 typedef hipDoubleComplex gpuDoubleComplex;
 
-#if TF_ROCM_VERSION >= 70000
-typedef hipFloatComplex gpublasComplex;
+typedef hipComplex gpublasComplex;
 typedef hipDoubleComplex gpublasDoubleComplex;
-#else
-typedef hipblasComplex gpublasComplex;
-typedef hipblasDoubleComplex gpublasDoubleComplex;
-#endif  // TF_ROCM_VERSION >= 70000
-
 typedef struct hipsolverHandle_* gpusolverDnHandle_t;
 typedef hipblasFillMode_t gpublasFillMode_t;
 typedef hipsolverFillMode_t gpusolverFillMode_t;
@@ -545,6 +541,7 @@ inline hipblasStatus_t gpublasCreate(gpublasHandle_t* handle) {
   return hipblasCreate(reinterpret_cast<hipblasHandle_t*>(handle));
 }
 }  // namespace jax::hip
+
 #define gpublasCreate ::jax::hip::gpublasCreate
 #define gpublasSetStream hipblasSetStream
 #define gpublasSgeqrfBatched hipblasSgeqrfBatched
@@ -601,6 +598,7 @@ inline hipsolverStatus_t gpusolverDnCreate(gpusolverDnHandle_t* handle) {
   return hipsolverCreate(reinterpret_cast<hipsolverHandle_t*>(handle));
 }
 }  // namespace jax::hip
+
 #define gpusolverDnCreate ::jax::hip::gpusolverDnCreate
 #define gpusolverDnSetStream hipsolverSetStream
 #define gpusolverDnCreateSyevjInfo hipsolverCreateSyevjInfo
@@ -780,10 +778,10 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) {
 #define GPUSPARSE_INDEX_32I HIPSPARSE_INDEX_32I
 #define GPUSPARSE_INDEX_64I HIPSPARSE_INDEX_64I
 #define GPUSPARSE_DENSETOSPARSE_ALG_DEFAULT HIPSPARSE_DENSETOSPARSE_ALG_DEFAULT
-#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_MV_ALG_DEFAULT
-#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_MV_ALG_DEFAULT
-#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_ALG_DEFAULT
-#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_ALG_DEFAULT
+#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_COOMV_ALG
+#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_CSRMV_ALG1
+#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_COO_ALG1
+#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_CSR_ALG1
 #define GPUSPARSE_INDEX_BASE_ZERO HIPSPARSE_INDEX_BASE_ZERO
 #define GPUSPARSE_OPERATION_NON_TRANSPOSE HIPSPARSE_OPERATION_NON_TRANSPOSE
 #define GPUSPARSE_OPERATION_TRANSPOSE HIPSPARSE_OPERATION_TRANSPOSE
@@ -796,7 +794,7 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) {
 #define GPU_STREAM_NON_BLOCKING hipStreamNonBlocking
 
 #define gpuMalloc hipMalloc
-#define gpuGetLastError hipGetLastError
+#define gpuGetLastError hipExtGetLastError
 #define gpuGetErrorString hipGetErrorString
 #define gpuMemcpyAsync hipMemcpyAsync
 #define gpuMemcpyDeviceToDevice hipMemcpyDeviceToDevice
diff --git a/jaxlib/jax.bzl b/jaxlib/jax.bzl
index d07287cf6684..fb9ff22ca06c 100644
--- a/jaxlib/jax.bzl
+++ b/jaxlib/jax.bzl
@@ -22,10 +22,6 @@ load("@jax_wheel//:wheel.bzl", "WHEEL_VERSION")
 load("@jax_wheel_version_suffix//:wheel_version_suffix.bzl", "WHEEL_VERSION_SUFFIX")
 load("@local_config_cuda//cuda:build_defs.bzl", _cuda_library = "cuda_library", _if_cuda_is_configured = "if_cuda_is_configured")
 load("@local_config_rocm//rocm:build_defs.bzl", _if_rocm_is_configured = "if_rocm_is_configured", _rocm_library = "rocm_library")
-
-# TODO(Intel-tf): Update `sycl` with `oneapi` when xla changes to use `oneapi`.
-load("@local_config_sycl//sycl:build_defs.bzl", _if_oneapi_is_configured = "if_sycl_is_configured", _oneapi_library = "sycl_library")
-load("@nvidia_wheel_versions//:versions.bzl", "NVIDIA_WHEEL_VERSIONS")
 load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION", "HERMETIC_PYTHON_VERSION_KIND")
 load("@rocm_external_test_deps//:external_deps.bzl", "EXTERNAL_DEPS")
 load("@rocm_prebuilt_test_deps//:external_deps.bzl", PREBUILT_EXTERNAL_DEPS = "EXTERNAL_DEPS")
@@ -494,7 +490,6 @@ def _jax_wheel_impl(ctx):
         if ctx.attr.platform_version == "":
             fail("platform_version must be set to a valid cuda version for cuda wheels")
         args.add("--platform_version", ctx.attr.platform_version)  # required for gpu wheels
-        args.add("--nvidia_wheel_versions_data", NVIDIA_WHEEL_VERSIONS)  # required for gpu wheels
     if ctx.attr.enable_rocm:
         args.add("--enable-rocm", "True")
         if ctx.attr.platform_version == "":
diff --git a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc
index 090a74e95056..243122889f14 100644
--- a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc
+++ b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc
@@ -310,24 +310,23 @@ NB_MODULE(_jax_mlir_ext, m) {
         MlirDialectRegistry c_registry = registry.get();
 #define REGISTER_DIALECT(name)                                           \
   MlirDialectHandle name##_dialect = mlirGetDialectHandle__##name##__(); \
-  mlirDialectHandleInsertDialect(name##_dialect, c_registry)
-        REGISTER_DIALECT(arith);
-        REGISTER_DIALECT(func);
-        REGISTER_DIALECT(math);
-        REGISTER_DIALECT(memref);
-        REGISTER_DIALECT(scf);
-        REGISTER_DIALECT(vector);
-        // TODO(jpienaar): these don't seem to have C API targets known to Bazel
-        unwrap(c_registry)->insert<mlir::shape::ShapeDialect>();
-        unwrap(c_registry)->insert<mlir::tensor::TensorDialect>();
-        unwrap(c_registry)->insert<mlir::vhlo::VhloDialect>();
-
-        // For Mosaic GPU
-        REGISTER_DIALECT(cf);
-        REGISTER_DIALECT(gpu);
-        REGISTER_DIALECT(nvgpu);
-        REGISTER_DIALECT(nvvm);
-        REGISTER_DIALECT(llvm);
+  mlirDialectHandleInsertDialect(name##_dialect, registry)
+    REGISTER_DIALECT(arith);
+    REGISTER_DIALECT(func);
+    REGISTER_DIALECT(math);
+    REGISTER_DIALECT(memref);
+    REGISTER_DIALECT(scf);
+    REGISTER_DIALECT(vector);
+    // TODO(jpienaar): these don't seem to have C API targets known to Bazel
+    unwrap(registry)->insert<mlir::shape::ShapeDialect>();
+    unwrap(registry)->insert<mlir::tensor::TensorDialect>();
+    unwrap(registry)->insert<mlir::vhlo::VhloDialect>();
+    // For Mosaic GPU
+    REGISTER_DIALECT(cf);
+    REGISTER_DIALECT(gpu);
+    REGISTER_DIALECT(nvgpu);
+    REGISTER_DIALECT(nvvm);
+    REGISTER_DIALECT(llvm);
 #undef REGISTER_DIALECT
 
         mlirMosaicGpuRegisterSerdePass();
diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD
index a27c138bd42c..5deb8c568ef1 100644
--- a/jaxlib/rocm/BUILD
+++ b/jaxlib/rocm/BUILD
@@ -28,7 +28,7 @@ licenses(["notice"])
 
 package(
     default_applicable_licenses = [],
-    default_visibility = ["//:__subpackages__"],
+    default_visibility = ["//visibility:public"],
 )
 
 cc_library(
@@ -425,6 +425,10 @@ rocm_nanobind_extension(
         "@xla//xla/ffi/api:ffi",
         "@xla//xla/python:safe_static_init",
     ],
+    linkopts = [
+        "-L/opt/rocm/lib",
+        "-lamdhip64",
+    ],
 )
 
 cc_library(
@@ -558,6 +562,7 @@ rocm_nanobind_extension(
     enable_stub_generation = False,
     module_name = "rocm_plugin_extension",
     deps = [
+        ":hip_gpu_kernel_helpers",
         ":py_client_gpu",
         "//jaxlib:kernel_nanobind_helpers",
         "//jaxlib/gpu:gpu_plugin_extension",
diff --git a/jaxlib/rocm/rocm_plugin_extension.cc b/jaxlib/rocm/rocm_plugin_extension.cc
index 7f3a6dd12b2e..525aea1011b6 100644
--- a/jaxlib/rocm/rocm_plugin_extension.cc
+++ b/jaxlib/rocm/rocm_plugin_extension.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "jaxlib/gpu/gpu_plugin_extension.h"
 #include "jaxlib/gpu/py_client_gpu.h"
 #include "jaxlib/kernel_nanobind_helpers.h"
+#include "jaxlib/gpu/gpu_kernel_helpers.h"
 
 namespace nb = nanobind;
 
@@ -98,6 +99,13 @@ nb::dict FfiHandlers() {
   return dict;
 }
 
+int ROCmDeviceCount() {
+  int device_count = -1;
+  JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipInit(0)));
+  JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipGetDeviceCount(&device_count)));
+  return device_count;
+}
+
 }  // namespace
 
 NB_MODULE(rocm_plugin_extension, m) {
@@ -124,5 +132,6 @@ NB_MODULE(rocm_plugin_extension, m) {
         return device_ordinal;
       },
       nb::arg("data_value"));
+  m.def("get_device_count", &ROCmDeviceCount);
 }
 }  // namespace jax
diff --git a/jaxlib/rocm/rocm_rpath.bzl b/jaxlib/rocm/rocm_rpath.bzl
index 354082c0d200..1da2d640ca42 100644
--- a/jaxlib/rocm/rocm_rpath.bzl
+++ b/jaxlib/rocm/rocm_rpath.bzl
@@ -35,8 +35,11 @@ _ROCM_LINK_ONLY = "@local_config_rocm//rocm:link_only"
 
 _WHEEL_RPATHS = [
     "-Wl,-rpath,$$ORIGIN/../rocm/lib",
+    "-Wl,-rpath,$$ORIGIN/../rocm/lib/rocm_sysdeps/lib",
     "-Wl,-rpath,$$ORIGIN/../../rocm/lib",
+    "-Wl,-rpath,$$ORIGIN/../../rocm/lib/rocm_sysdeps/lib",
     "-Wl,-rpath,/opt/rocm/lib",
+    "-Wl,-rpath,/opt/rocm/lib/rocm_sysdeps/lib",
 ]
 
 def _wheel_features():
diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py
index d033a490f547..b14e9171f952 100644
--- a/tests/array_interoperability_test.py
+++ b/tests/array_interoperability_test.py
@@ -312,7 +312,7 @@ def testCudaArrayInterfaceOnNonCudaFails(self):
     self.assertFalse(hasattr(x, "__cuda_array_interface__"))
     with self.assertRaisesRegex(
         AttributeError,
-        "__cuda_array_interface__ is only defined for .*GPU buffers.",
+        "__cuda_array_interface__ is only defined for GPU buffers.",
     ):
       _ = x.__cuda_array_interface__
 
diff --git a/tests/array_test.py b/tests/array_test.py
index cf08765c5893..890b701dd589 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1550,7 +1550,7 @@ class RngShardingTest(jtu.JaxTestCase):
   # tests that the PRNGs are automatically sharded as expected
 
   @parameterized.named_parameters(("3", 3), ("4", 4), ("5", 5))
-  @jtu.skip_on_devices("gpu")
+  @jtu.skip_on_devices("cuda")
   def test_random_bits_is_pure_map_1d(self, num_devices):
     @jax.jit
     def f(x):
@@ -1584,7 +1584,7 @@ def f(x):
        "mesh_shape": mesh_shape, "pspec": pspec}
       for mesh_shape in [(3, 2), (4, 2), (2, 3)]
       for pspec in [P('x', None), P(None, 'y'), P('x', 'y')])
-  @jtu.skip_on_devices("gpu")
+  @jtu.skip_on_devices("cuda")
   def test_random_bits_is_pure_map_2d(self, mesh_shape, pspec):
     @jax.jit
     def f(x):
diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py
index 81c7f5103f44..fa168d61b03c 100644
--- a/tests/lax_vmap_test.py
+++ b/tests/lax_vmap_test.py
@@ -759,6 +759,8 @@ def testSort(self, shape, dimension, arity, bdims, is_stable):
   # TODO Collapse
   # TODO Scatter
 
+  # b/183233858: variadic reduce-window not implemented on XLA:CUDA
+  @jtu.skip_on_devices("cuda")
   def test_variadic_reduce_window(self):
     # https://github.com/jax-ml/jax/discussions/9818 and
     # https://github.com/jax-ml/jax/issues/9837
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 867168b1c300..b99492724fd5 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -497,6 +497,247 @@ def testEigBatching(self, shape, dtype):
     self.assertTrue(np.all(np.linalg.norm(
         np.matmul(args, vs) - ws[..., None, :] * vs) < 1e-3))
 
+  @jtu.sample_product(
+      n=[0, 4, 5, 50, 512],
+      dtype=float_types + complex_types,
+      lower=[True, False],
+  )
+  def testEigh(self, n, dtype, lower):
+    rng = jtu.rand_default(self.rng())
+    eps = np.finfo(dtype).eps
+    args_maker = lambda: [rng((n, n), dtype)]
+
+    uplo = "L" if lower else "U"
+
+    a, = args_maker()
+    a = (a + np.conj(a.T)) / 2
+    w, v = jnp.linalg.eigh(np.tril(a) if lower else np.triu(a),
+                           UPLO=uplo, symmetrize_input=False)
+    w = w.astype(v.dtype)
+    tol = 2 * n * eps
+    self.assertAllClose(
+        np.eye(n, dtype=v.dtype),
+        np.matmul(np.conj(T(v)), v),
+        atol=tol,
+        rtol=tol,
+    )
+
+    with jax.numpy_rank_promotion('allow'):
+      tol = 100 * eps
+      self.assertLessEqual(
+          np.linalg.norm(np.matmul(a, v) - w * v), tol * np.linalg.norm(a)
+      )
+
+    self._CompileAndCheck(
+        partial(jnp.linalg.eigh, UPLO=uplo), args_maker, rtol=eps
+    )
+
+    # Compare eigenvalues against Numpy using double precision. We do not compare
+    # eigenvectors because they are not uniquely defined, but the two checks above
+    # guarantee that that they satisfy the conditions for being eigenvectors.
+    double_type = dtype
+    if dtype == np.float32:
+      double_type = np.float64
+    if dtype == np.complex64:
+      double_type = np.complex128
+    w_np = np.linalg.eigvalsh(a.astype(double_type))
+    tol = 8 * eps
+    self.assertAllClose(
+        w_np.astype(w.dtype), w, atol=tol * np.linalg.norm(a), rtol=tol
+    )
+
+  @jax._src.config.explicit_x64_dtypes("allow")
+  @jtu.run_on_devices("gpu")
+  @unittest.skip("Needs a large amount of GPU memory, doesn't work in CI")
+  def testEighLargeMatrix(self):
+    # https://github.com/jax-ml/jax/issues/33062
+    n = 16384
+    A = jnp.eye(n, dtype=jnp.float64)
+    jax.block_until_ready(jax.lax.linalg.eigh(A))
+
+  @jtu.sample_product(
+      start=[0, 1, 63, 64, 65, 255],
+      end=[1, 63, 64, 65, 256],
+  )
+  @jtu.run_on_devices("tpu")  # TODO(rmlarsen: enable on other devices)
+  def testEighSubsetByIndex(self, start, end):
+    if start >= end:
+      return
+    dtype = np.float32
+    n = 256
+    rng = jtu.rand_default(self.rng())
+    eps = np.finfo(dtype).eps
+    args_maker = lambda: [rng((n, n), dtype)]
+    subset_by_index = (start, end)
+    k = end - start
+    (a,) = args_maker()
+    a = (a + np.conj(a.T)) / 2
+
+    v, w = lax.linalg.eigh(
+        a, symmetrize_input=False, subset_by_index=subset_by_index
+    )
+    w = w.astype(v.dtype)
+
+    self.assertEqual(v.shape, (n, k))
+    self.assertEqual(w.shape, (k,))
+    with jax.numpy_rank_promotion("allow"):
+      tol = 200 * eps
+      self.assertLessEqual(
+          np.linalg.norm(np.matmul(a, v) - w * v), tol * np.linalg.norm(a)
+      )
+    tol = 3 * n * eps
+    self.assertAllClose(
+        np.eye(k, dtype=v.dtype),
+        np.matmul(np.conj(T(v)), v),
+        atol=tol,
+        rtol=tol,
+    )
+
+    self._CompileAndCheck(partial(jnp.linalg.eigh), args_maker, rtol=eps)
+
+    # Compare eigenvalues against Numpy. We do not compare eigenvectors because
+    # they are not uniquely defined, but the two checks above guarantee that
+    # that they satisfy the conditions for being eigenvectors.
+    double_type = dtype
+    if dtype == np.float32:
+      double_type = np.float64
+    if dtype == np.complex64:
+      double_type = np.complex128
+    w_np = np.linalg.eigvalsh(a.astype(double_type))[
+        subset_by_index[0] : subset_by_index[1]
+    ]
+    tol = 20 * eps
+    self.assertAllClose(
+        w_np.astype(w.dtype), w, atol=tol * np.linalg.norm(a), rtol=tol
+    )
+
+  def testEighZeroDiagonal(self):
+    a = np.array([[0., -1., -1.,  1.],
+                  [-1.,  0.,  1., -1.],
+                  [-1.,  1.,  0., -1.],
+                  [1., -1., -1.,  0.]], dtype=np.float32)
+    w, v = jnp.linalg.eigh(a)
+    w = w.astype(v.dtype)
+    eps = jnp.finfo(a.dtype).eps
+    with jax.numpy_rank_promotion('allow'):
+      self.assertLessEqual(
+          np.linalg.norm(np.matmul(a, v) - w * v), 2.5 * eps * np.linalg.norm(a)
+      )
+
+
+  def testEighTinyNorm(self):
+    if jtu.is_device_rocm():
+      # numerical errors seen as of ROCm 7.2 due to hipSolver issue
+      # TODO: re-enable the test once the hipSolver issue is fixed
+      self.skipTest("testEighNorm not supported on ROCm due to hipSOLVER issue")
+    rng = jtu.rand_default(self.rng())
+    a = rng((300, 300), dtype=np.float32)
+    eps = jnp.finfo(a.dtype).eps
+    a = eps * (a + np.conj(a.T))
+    w, v = jnp.linalg.eigh(a)
+    w = w.astype(v.dtype)
+    with jax.numpy_rank_promotion("allow"):
+      self.assertLessEqual(
+          np.linalg.norm(np.matmul(a, v) - w * v), 80 * eps * np.linalg.norm(a)
+      )
+
+  @jtu.sample_product(
+      rank=[1, 3, 299],
+  )
+  def testEighRankDeficient(self, rank):
+    rng = jtu.rand_default(self.rng())
+    eps = jnp.finfo(np.float32).eps
+    a = rng((300, rank), dtype=np.float32)
+    a = a @ np.conj(a.T)
+    w, v = jnp.linalg.eigh(a)
+    w = w.astype(v.dtype)
+    with jax.numpy_rank_promotion("allow"):
+      self.assertLessEqual(
+          np.linalg.norm(np.matmul(a, v) - w * v),
+          85 * eps * np.linalg.norm(a),
+      )
+
+  @jtu.sample_product(
+    n=[0, 4, 5, 50, 512],
+    dtype=float_types + complex_types,
+    lower=[True, False],
+  )
+  def testEighIdentity(self, n, dtype, lower):
+    tol = np.finfo(dtype).eps
+    uplo = "L" if lower else "U"
+
+    a = jnp.eye(n, dtype=dtype)
+    w, v = jnp.linalg.eigh(a, UPLO=uplo, symmetrize_input=False)
+    w = w.astype(v.dtype)
+    self.assertLessEqual(
+        np.linalg.norm(np.eye(n) - np.matmul(np.conj(T(v)), v)), tol
+    )
+    with jax.numpy_rank_promotion('allow'):
+      self.assertLessEqual(np.linalg.norm(np.matmul(a, v) - w * v),
+                           tol * np.linalg.norm(a))
+
+  @jtu.sample_product(
+    shape=[(4, 4), (5, 5), (50, 50)],
+    dtype=float_types + complex_types,
+  )
+  def testEigvalsh(self, shape, dtype):
+    rng = jtu.rand_default(self.rng())
+    n = shape[-1]
+    def args_maker():
+      a = rng((n, n), dtype)
+      a = (a + np.conj(a.T)) / 2
+      return [a]
+    self._CheckAgainstNumpy(
+        np.linalg.eigvalsh, jnp.linalg.eigvalsh, args_maker, tol=2e-5
+    )
+
+  @jtu.sample_product(
+      shape=[(1, 1), (4, 4), (5, 5), (25, 25), (2, 10, 10)],
+      dtype=float_types + complex_types,
+      lower=[True, False],
+  )
+  def testEighGrad(self, shape, dtype, lower):
+    if platform.system() == "Windows":
+      self.skipTest("Skip on Windows due to tolerance issues.")
+    rng = jtu.rand_default(self.rng())
+    a = rng(shape, dtype)
+    a = (a + np.conj(T(a))) / 2
+    ones = np.ones((a.shape[-1], a.shape[-1]), dtype=dtype)
+    a *= np.tril(ones) if lower else np.triu(ones)
+    # Gradient checks will fail without symmetrization as the eigh jvp rule
+    # is only correct for tangents in the symmetric subspace, whereas the
+    # checker checks against unconstrained (co)tangents.
+    f = partial(_normalizing_eigh, lower=lower, symmetrize_input=True)
+    norm_a = jnp.linalg.norm(a)
+    eps = 2e-5 * norm_a
+    atol = 5e-3 * norm_a
+    rtol = 0.025
+    jtu.check_grads(f, (a,), 2, atol=atol, rtol=rtol, eps=eps)
+
+  def testEighGradPrecision(self):
+    rng = jtu.rand_default(self.rng())
+    a = rng((3, 3), np.float32)
+    jtu.assert_dot_precision(
+        lax.Precision.HIGHEST, partial(jvp, jnp.linalg.eigh), (a,), (a,))
+
+  def testEighGradRankPromotion(self):
+    rng = jtu.rand_default(self.rng())
+    a = rng((10, 3, 3), np.float32)
+    jvp(jnp.linalg.eigh, (a,), (a,))  # doesn't crash
+
+  @jtu.sample_product(
+    shape=[(1, 1), (4, 4), (5, 5), (300, 300)],
+    dtype=float_types + complex_types,
+  )
+  def testEighBatching(self, shape, dtype):
+    rng = jtu.rand_default(self.rng())
+    shape = (10,) + shape
+    args = rng(shape, dtype)
+    args = (args + np.conj(T(args))) / 2
+    ws, vs = vmap(jsp.linalg.eigh)(args)
+    ws = ws.astype(vs.dtype)
+    norm = np.max(np.linalg.norm(np.matmul(args, vs) - ws[..., None, :] * vs))
+    self.assertLess(norm, 1.4e-2)
 
   @jtu.sample_product(
     shape=[(1,), (4,), (5,)],
diff --git a/tests/pallas/gpu_paged_attention_test.py b/tests/pallas/gpu_paged_attention_test.py
index 1b778c787a6d..6a1d8de22a78 100644
--- a/tests/pallas/gpu_paged_attention_test.py
+++ b/tests/pallas/gpu_paged_attention_test.py
@@ -112,6 +112,64 @@ class PagedAttentionKernelTest(PallasBaseTest):
   def setUp(self):
     super().setUp()
 
+  def _estimate_shared_memory_bytes(self, block_h, pages_per_compute_block, 
+                                     page_size, head_dim, dtype):
+    """Estimate shared memory usage for paged attention kernel."""
+    dtype_size = jnp.dtype(dtype).itemsize
+    # Approximate calculation based on kernel's memory usage
+    # Q block: block_h * head_dim
+    # K/V blocks: pages_per_compute_block * page_size * head_dim
+    # Plus accumulators and intermediate values
+    block_k = pages_per_compute_block * page_size
+    estimated = dtype_size * (
+        block_h * head_dim +  # Q
+        2 * block_k * head_dim +  # K and V
+        block_h * block_k +  # logits/attention weights
+        block_h * 8  # accumulators (m, l, etc.) in float32
+    )
+    return estimated
+
+  def _adjust_params_for_shared_memory(self, block_h, pages_per_compute_block, 
+                                       page_size, head_dim, dtype):
+    """Adjust parameters to fit within device shared memory limits.
+    
+    Uses XLA's DeviceDescription.shared_memory_per_block_optin() to query
+    the actual device capability rather than hardcoding values.
+    """
+    try:
+      device = jax.local_devices()[0]
+      # Query XLA DeviceDescription for max shared memory per block
+      # This is exposed from stream_executor::DeviceDescription::shared_memory_per_block_optin()
+      max_smem = device.shared_memory_per_block_optin
+    except (AttributeError, IndexError):
+      # Fallback if XLA doesn't expose shared_memory_per_block_optin (older versions)
+      # or if no devices are available. Use conservative 48KB (safe for most GPUs).
+      max_smem = 48 * 1024
+    
+    estimated = self._estimate_shared_memory_bytes(
+        block_h, pages_per_compute_block, page_size, head_dim, dtype)
+    
+    # If within limits, no adjustment needed
+    if estimated <= max_smem:
+      return block_h, pages_per_compute_block, page_size
+    
+    # Try to reduce parameters to fit
+    while estimated > max_smem:
+      if pages_per_compute_block > 2:
+        pages_per_compute_block = pages_per_compute_block // 2
+      elif page_size > 8:
+        page_size = page_size // 2
+      elif block_h > 8:
+        block_h = block_h // 2
+      else:
+        # Can't reduce further, will need to skip
+        return None, None, None
+      
+      estimated = self._estimate_shared_memory_bytes(
+          block_h, pages_per_compute_block, page_size, head_dim, dtype)
+    
+    return block_h, pages_per_compute_block, page_size
+
   @jtu.sample_product(
       dtype=(jnp.float16,),
       page_size=(8, 16, 32),
@@ -201,6 +259,17 @@ def test_quantized_paged_attention(
     if (quant_dtype == jnp.float8_e4m3fn
         and not jtu.is_cuda_compute_capability_at_least("8.9")):
       self.skipTest("Skipping since float8_e4m3fn is not supported on < sm89")
+    
+    # Check and adjust parameters if needed to fit device limits for ROCm
+    if jtu.is_device_rocm():
+      adjusted = self._adjust_params_for_shared_memory(
+          block_h, pages_per_compute_block, page_size, head_dim, dtype)
+      
+      if adjusted == (None, None, None):
+        self.skipTest("Cannot adjust parameters to fit ROCm device shared memory limits")
+      
+      block_h, pages_per_compute_block, page_size = adjusted
+    
     max_kv_len = 2048
     seq_lens = np.asarray([3, 256, 513, 1023, 2048], dtype=jnp.int32)
     q, k_pages, v_pages, block_tables = _generate_qkv(
@@ -218,7 +287,7 @@ def test_quantized_paged_attention(
 
     k_, k_scales = (_quantize(k_pages, quant_dtype)
                     if quantize_k else (k_pages, None))
-    v_, v_scales = (_quantize(k_pages, quant_dtype)
+    v_, v_scales = (_quantize(v_pages, quant_dtype)
                     if quantize_v else (v_pages, None))
 
     o = paged_attention.paged_attention(
diff --git a/tests/pallas/gpu_pallas_distributed_test.py b/tests/pallas/gpu_pallas_distributed_test.py
index f47e14e7ac21..169e24b2062b 100644
--- a/tests/pallas/gpu_pallas_distributed_test.py
+++ b/tests/pallas/gpu_pallas_distributed_test.py
@@ -133,14 +133,17 @@ def setUp(self):
     if jtu.is_device_rocm():
       self.skipTest("Mosaic GPU is not supported on ROCm.")
 
-    if (not jtu.is_device_cuda() or
+    # Check mosaic support first (before GPU capability check)
+    if not mgpu.supports_cross_device_collectives():
+      if jtu.test_device_matches(["rocm"]):
+        self.skipTest("Mosaic not supported on ROCm currently.")
+      else:
+        self.skipTest("NVSHMEM library unavailable.")
+    if (not jtu.test_device_matches(["cuda"]) or
         not jtu.is_cuda_compute_capability_at_least("9.0")):
       self.skipTest("Only works on GPU with capability >= sm90")
-    if not mgpu.supports_cross_device_collectives():
-      self.skipTest(
-          "Skip test since cross-device collectives are not supported"
-          " (either NVSHMEM is not available in multi-process mode, or mixed"
-          " mode is used).")
+    if jax.process_count() == 1:
+      self.skipTest("Test requires multiple processes.")
     if os.environ.get("XLA_PYTHON_CLIENT_ALLOCATOR", "") == "platform":
       self.skipTest("NVSHMEM doesn't work with the platform allocator.")
 
diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py
index 9fea8249e525..a1af7d6f252a 100644
--- a/tests/pallas/ops_test.py
+++ b/tests/pallas/ops_test.py
@@ -1047,6 +1047,13 @@ def test_is_finite(self, dtype):
     if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0):
       self.skipTest("is_finite not in Triton lowering for jax 0.8.0")
 
+    if jtu.test_device_matches(["cuda"]):
+      self.skipTest("Not tested on CUDA") # set this b/c this how the test was
+      # originally configured. Have no way to test cuda.
+
+    if jtu.is_device_rocm():
+      self.skipTest("is_finite not in Triton lowering for jax 0.8.0")
+
     size = len(self.IS_FINITE_TEST_VALUES)
 
     @functools.partial(
@@ -1099,6 +1106,9 @@ def test_is_finite_scalar(self, dtype):
     if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0):
       self.skipTest("is_finite not in Triton lowering for jax 0.8.0")
 
+    if jtu.is_device_rocm():
+      self.skipTest("is_finite not in Triton lowering for jax 0.8.0")
+
     size = len(self.IS_FINITE_TEST_VALUES)
 
     @functools.partial(
@@ -2151,12 +2161,20 @@ def test_dot(self, lhs_and_rhs_shape, dtype, trans_x, trans_y):
         self.skipTest("bfloat16 type are not supported on GPU")
       # Check shared memory limit: Triton loads lhs + rhs into shared memory
       if jtu.is_device_rocm():
+        # ROCm: use correct formula with dynamic limit from rocminfo
         dtype_size = jnp.dtype(dtype).itemsize
-        if (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size > get_rocm_shared_memory_limit():
+        shared_mem_bytes = (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size
+        shared_mem_limit = get_rocm_shared_memory_limit()
+        if shared_mem_bytes > shared_mem_limit:
           self.skipTest("Shared memory size limit exceeded")
-      elif math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape) > (256 * 256) * 2:
-        self.skipTest("Shared memory size limit exceeded")
-      if (jax.local_devices()[0].shared_memory_per_block_optin == 99 * 1024 and
+      else:
+        # NVIDIA: keep original check
+        if (
+            math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape)
+            > (256 * 256) * 2
+        ):
+          self.skipTest("Shared memory size limit exceeded")
+      if (jax.local_devices()[0].device_kind == "NVIDIA L4" and
           dtype == jnp.float32 and
           lhs_and_rhs_shape in [
             ((128, 16), (128, 256)),
diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py
index 860f09a7f988..f7b68a483103 100644
--- a/tests/random_lax_test.py
+++ b/tests/random_lax_test.py
@@ -1024,7 +1024,7 @@ def testMultivariateNormalCovariance(self):
                         check_dtypes=False)
 
   @jtu.sample_product(method=['cholesky', 'eigh', 'svd'])
-  @jtu.skip_on_devices('cuda', 'tpu')  # Some NaNs on accelerators.
+  @jtu.skip_on_devices('cuda', 'tpu')  # Some NaNs on accelerators. ROCm supported
   def testMultivariateNormalSingularCovariance(self, method):
     # Singular covariance matrix https://github.com/jax-ml/jax/discussions/13293
     mu = jnp.zeros((2,))
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
index 3823ca26d8f9..a7bf9373c34c 100644
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -139,8 +139,6 @@ def test_csr_fromdense_ad(self, shape, dtype):
   @jax.default_matmul_precision("float32")
   def test_csr_matmul_ad(self, shape, dtype, bshape):
     if jtu.is_device_rocm():
-      # hipSPARSE segfault observed as of ROCm 7.2.
-      # TODO(ROCm): Re-enable once hipSPARSE issue is fixed.
       self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue")
     csr_matmul = sparse_csr._csr_matvec if len(bshape) == 1 else sparse_csr._csr_matmat
     tol = {np.float32: 2E-5, np.float64: 1E-12, np.complex64: 1E-5,
@@ -221,8 +219,6 @@ def test_csr_fromdense(self, shape, dtype):
   )
   def test_csr_matvec(self, shape, dtype, transpose):
     if jtu.is_device_rocm():
-      # hipSPARSE segfault observed as of ROCm 7.2.
-      # TODO(ROCm): Re-enable once hipSPARSE issue is fixed.
       self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue")
 
     op = lambda M: M.T if transpose else M
@@ -594,8 +590,6 @@ def test_coo_spmm(self, shape, dtype, transpose):
   @jtu.run_on_devices("gpu")
   def test_csr_spmv(self, shape, dtype, transpose):
     if jtu.is_device_rocm():
-      # hipSPARSE segfault observed as of ROCm 7.2.
-      # TODO(ROCm): Re-enable once hipSPARSE issue is fixed.
       self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue")
     tol = {np.float32: 2E-5, np.float64: 2E-14}
 
@@ -1056,8 +1050,6 @@ def test_transpose(self, shape, dtype, Obj):
   @jax.default_matmul_precision("float32")
   def test_matmul(self, shape, dtype, Obj, bshape):
     if jtu.is_device_rocm():
-      # hipSPARSE segfault observed as of ROCm 7.2.
-      # TODO(ROCm): Re-enable once hipSPARSE issue is fixed.
       self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue")
     rng = sptu.rand_sparse(self.rng(), post=jnp.array)
     rng_b = jtu.rand_default(self.rng())