{"slug": "openxla-and-jax-rocm-support-and-the-state-of-ci", "title": "OpenXLA and JAX - ROCm Support and the State of CI", "summary": "The OpenXLA compiler stack and JAX now run upstream on AMD ROCm, with XLA gating every pull request on real AMD Instinct silicon through GitHub Actions and JAX running hardware tests on every ROCm PR. This development enables AMD Instinct customers to use JAX for pretraining and large fine-tunes without rewriting model code, leveraging XLA's whole-program compilation and JAX's composable transforms for SPMD parallelism.", "body_md": "# OpenXLA and JAX - ROCm Support and the State of CI[#](#openxla-and-jax-rocm-support-and-the-state-of-ci)\n\nThe OpenXLA compiler stack — XLA at the foundation, JAX as the front end — now runs upstream on AMD ROCm. XLA gates every pull request on real AMD Instinct silicon through its GitHub Actions workflow, side by side with the CUDA path; JAX runs the same hardware on every ROCm PR through its own workflows, with the merge gate rolling out next. `pip install \"jax[rocm7-local]\"`\n\nis a first-class entry point. This post documents how that backend is structured, what landed in the last twelve months, and how the CI pipeline that keeps it healthy is wired together. ** Part 1** covers OpenXLA on AMD — the XLA backend, what landed this year, and CI.\n\n**covers JAX on AMD — the plugin architecture, JAX-side changes, and the four-workflow test matrix.**\n\n[Part 2](#part-2-jax-on-amd)\n\nWhat’s in this post\n\nThe OpenXLA stack on ROCm in one diagram.\n\nSupported AMD GPU targets across compiler, default wheels, and CI gating — and why those three lists differ.\n\nThe year’s work on XLA and JAX for ROCm: Triton on AMDGPU, hipBLASLt group-GEMM, FP8, hermetic builds, manylinux wheels.\n\nEnd-to-end CI: how each PR runs on real Instinct hardware and how XLA and JAX cross-pin against each other.\n\nA reproducible three-command quick start.\n\nWhere to file issues, send PRs, and dump HLO when a workload misbehaves.\n\n## Why JAX/OpenXLA on AMD?[#](#why-jax-openxla-on-amd)\n\nJAX on AMD is a deliberate architectural choice. The case for picking it over an eager framework — or over a hand-tuned kernel library — rests on two properties: a whole-program compiler that sees the entire High Level Operations (HLO) graph at lowering time, and a programming model where collectives are inferred from sharding annotations rather than spelled out at the call site.\n\n**The technical case.** XLA is a hybrid ahead-of-time and just-in-time compiler that consumes a whole HLO graph — often hundreds of ops — and lowers it through fusion and autotuning, dispatching to Triton with LLVM-AMDGPU codegen or to ROCm math libraries (hipBLASLt, MIOpen, and others) according to which provides the best performance. For transformer-shaped pretraining and inference — dense matmuls, attention, layer norms, large collectives — that compile-time view produces fused kernels that would otherwise demand hand-tuning per shape and per generation.\n\nJAX provides the programming model on top: pure functions, composable transforms (`jit`\n\n, `grad`\n\n, `vmap`\n\n, `scan`\n\n), and SPMD parallelism through `pjit`\n\n/ `shard_map`\n\nover GSPMD. Collectives are inserted by the compiler from sharding annotations, so the same source runs on CPUs, TPUs, and GPUs — including our AMD Instinct product line.\n\n**Who should keep reading.** This post is most directly useful if you are:\n\nAn\n\n**AMD Instinct customer** running pretraining or large fine-tunes, evaluating which compiler stack to standardise on for MI300 / MI350 capacity.A\n\n**JAX user** adding AMD as a second hardware target without rewriting model code.A\n\n**foundation-model lab** doing SPMD / GSPMD pretraining and weighing Instinct + RCCL against NVIDIA + NCCL.A\n\n**compiler or ML-systems engineer** contributing to OpenXLA or JAX on AMD — the CI sections answer “what will happen to my PR before I open it” for both repositories.A\n\n**maintainer of MaxText / MaxDiffusion**-style reference workloads, or a researcher running scientific simulation on HPC systems (** NumPyro, JAX-MD, Brax**, AlphaFold-shaped models) who wants to target AMD GPUs on leadership-class systems.\n\n## Architecture[#](#architecture)\n\nFigure 1 shows the full path from a JAX program to a GPU kernel — five layers, colour-coded by ownership.\n\n*Figure 1. The OpenXLA stack on AMD ROCm, from your Python code down to the kernel that HIP launches on the GPU.*\n\nBlue is user code; yellow is JAX in Python; green is XLA and its PJRT plugin; red is the ROCm runtime (HIP, math libraries, kernel driver); grey is silicon. The `{N}`\n\nin `jax-rocm{N}-plugin`\n\nand `jax-rocm{N}-pjrt`\n\nis the major ROCm version — today `7`\n\n(for ROCm 7.x), resolved at import time. The two AMD-shaped boxes are the only divergence from the CUDA path; everything above and below is shared. That clean factoring is what makes upstream CI tractable, and what the rest of this post is built around.\n\n## Quick Start[#](#quick-start)\n\nThe shortest path from “I have an Instinct GPU” to “I just compiled a JAX program through XLA on it”:\n\n```\ndocker run -it --rm \\\n  --device=/dev/kfd --device=/dev/dri --group-add video \\\n  --shm-size=64G --ipc=host \\\n  rocm/jax:latest \\\n  python -c \"import jax, jax.numpy as jnp; \\\nprint(jax.devices()); \\\nprint(jax.jit(lambda x: jnp.tanh(x @ x.T))(jnp.ones((1024,1024), jnp.bfloat16)).shape)\"\n```\n\nA `RocmDevice`\n\nin `jax.devices()`\n\nand `(1024, 1024)`\n\non stdout means that a `tanh(x @ x.T)`\n\nHLO graph was compiled through Triton plus the AMDGPU LLVM backend, dispatched through HIP, and returned a result. Full install paths — pip and source — are in [Try It Out and Get Involved](#try-it-out-and-get-involved).\n\n## Hardware Support[#](#hardware-support)\n\nThree lists matter, in order of narrowing scope: what the **compiler recognises**, what the **default wheels build for**, and what **upstream CI exercises on every PR**.\n\n### Compiler-Supported Architectures[#](#compiler-supported-architectures)\n\nThe authoritative list is `kSupportedGfxVersions[]`\n\nin\n[xla/stream_executor/rocm/rocm_compute_capability.h](https://github.com/openxla/xla/blob/main/xla/stream_executor/rocm/rocm_compute_capability.h).\nXLA’s ROCm backend recognises and emits code for:\n\nProduct |\nArchitecture |\nCompiler target |\n|---|---|---|\nInstinct MI200 series (MI210X / MI250X) |\nCDNA 2 |\n|\nInstinct MI300 series (MI300X / MI325X) |\nCDNA 3 |\n|\nInstinct MI350 series (MI350X / MI355X) |\nCDNA 4 |\n|\nRadeon RX 6800 / 6900 |\nRDNA 2 |\n|\nRadeon RX 7900 |\nRDNA 3 |\n|\nRadeon RX 7700 / 7800 |\nRDNA 3 |\n|\nPhoenix |\nRDNA 3 (APU) |\n|\nStrix Point |\nRDNA 3.5 (APU) |\n|\nStrix Halo |\nRDNA 3.5 (APU) |\n|\nRadeon RX 9000 series |\nRDNA 4 |\n|\n\n### Default Wheel Build Targets[#](#default-wheel-build-targets)\n\nJAX’s published ROCm wheels are compiled for the actively supported subset of the compiler list above. From\n[jax/build/rocm/rocm.bazelrc](https://github.com/jax-ml/jax/blob/main/build/rocm/rocm.bazelrc):\n\n```\ngfx908, gfx90a, gfx942, gfx950, gfx1030, gfx1100, gfx1101, gfx1200, gfx1201\n```\n\nThe compiler list above stays broader than the wheel list on purpose: if your target isn’t in the default wheels, `python build/build.py --rocm_amdgpu_targets=…`\n\nlets you build wheels for it locally — including older targets like `gfx900`\n\n/ `gfx906`\n\n.\n\n### CI-Gated Targets[#](#ci-gated-targets)\n\nWe gate each PR before merging by ensuring it runs on our top-of-the-line MI Instinct and Radeon hardware:\n\n**XLA upstream**—`rocm_ci.yml`\n\n`gfx950`\n\n(MI350) on the single-GPU pool, covering the`ci_single_gpu`\n\nconfiguration.**JAX upstream**— AMD-hosted`bazel_rocm.yml`\n\n/`pytest_rocm.yml`\n\n`linux-x86-64-{1,4,8}gpu-amd`\n\npools spanning MI300 in CPX mode (`gfx942`\n\n) for the Bazel/RBE jobs and MI350 (`gfx950`\n\n) for the PyTest jobs.\n\nWhat this means in practice.MI200, MI300, and MI350 Instinct, and RX 6800-class or newer Radeon: the published wheels should work out of the box. Vega (`gfx900`\n\n/`gfx906`\n\n) or an RDNA 3 APU: the compiler still supports you, but expect to build wheels with an explicit`--rocm_amdgpu_targets`\n\n. Upstream CI density today is concentrated on`gfx950`\n\n(MI350); other Instinct and Radeon targets are exercised in AMD’s downstream CI.\n\n## Part 1 — OpenXLA on AMD[#](#part-1-openxla-on-amd)\n\n### Where the AMD Code Lives[#](#where-the-amd-code-lives)\n\n[OpenXLA](https://openxla.org) compiles HLO into fused, hardware-specific\nkernels. The GPU pipeline shares a large surface — HLO optimization,\nfusion, autotuning, command-buffer scheduling — and forks at codegen time\ninto vendor-specific backends.\n\nFor AMD, that backend lives across two main areas of the\n[openxla/xla](https://github.com/openxla/xla) tree:\n\n[xla/stream_executor/rocm/](https://github.com/openxla/xla/tree/main/xla/stream_executor/rocm)— the runtime layer: device discovery, streams, events, command buffers, memory, and wrappers around hipBLASLt, hipFFT, hipSOLVER, hipSPARSE, rocPRIM, MIOpen, and RCCL.[xla/service/gpu/](https://github.com/openxla/xla/tree/main/xla/service/gpu)— the AMD compiler entry points, most notably[amdgpu_compiler.cc](https://github.com/openxla/xla/blob/main/xla/service/gpu/amdgpu_compiler.cc),[custom_kernel_emitter_rocm.cc](https://github.com/openxla/xla/blob/main/xla/service/gpu/custom_kernel_emitter_rocm.cc), and the LLVM AMDGPU backend at[xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc](https://github.com/openxla/xla/blob/main/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc).\n\nThe build glue that makes ROCm a hermetic, reproducible target lives under\n[third_party/gpus/](https://github.com/openxla/xla/tree/main/third_party/gpus), with auto-detection in\n[third_party/gpus/rocm_configure.bzl](https://github.com/openxla/xla/blob/main/third_party/gpus/rocm_configure.bzl)\nand [third_party/gpus/find_rocm_config.py](https://github.com/openxla/xla/blob/main/third_party/gpus/find_rocm_config.py).\n\n### What Landed for ROCm in the Past Year[#](#what-landed-for-rocm-in-the-past-year)\n\nOver the last twelve months, **more than 300 commits** on `main`\n\ntouched ROCm, HIP, AMDGPU codegen, or `gfx9*`\n\ntargets. Four themes account for most of the work.\n\n**Triton on AMDGPU.** Triton is the primary code generator for fused matmul-plus-epilogue patterns in XLA:GPU, and bringing it to production parity on AMD GPUs took a sustained multi-PR effort. An [AMD-specific shared-memory allocation pass](https://github.com/openxla/xla/pull/41407) in the Triton pipeline and a CDNA-aware [ waves_per_eu knob](https://github.com/openxla/xla/pull/40499) in the GEMM autotuner closed the gap on per-shape kernel quality;\n\n[scaled-dot lowering](https://github.com/openxla/xla/pull/40557)brought microscaling MX-format GEMMs to AMD parity; and the early refactors in the Triton AllReduce series (\n\n[1](https://github.com/openxla/xla/pull/40462),\n\n[2](https://github.com/openxla/xla/pull/40460)) laid the groundwork for Triton-fused collectives on ROCm.\n\n**GEMM and FP8 on Instinct.** The headline landing was the [five-part hipBLASLt group-GEMM enablement](https://github.com/openxla/xla/pull/38737) — production group-GEMM through\n\n`hipBLASLt`\n\nis now the default path on Instinct. FP8 is now declared a [first-class ROCm 7 capability](https://github.com/openxla/xla/pull/40702)with fast accumulation, and the compiler-side support check was\n\n[relaxed accordingly](https://github.com/openxla/xla/pull/41176). End-to-end test coverage was extended to cover\n\n[both OCP and NANOO FP8 collective ops](https://github.com/openxla/xla/pull/40490), plus the\n\n[for group-GEMM.](https://github.com/openxla/xla/pull/40502)\n\n`gfx950`\n\nHIP backend requirements**Collectives and rocPRIM.** Hand-rolled fallbacks were replaced with tuned ROCm library primitives where they existed — most visibly [ rocprim::segmented_inclusive_scan](https://github.com/openxla/xla/pull/41229) in the batched row-scan path. Native ROCm collectives also landed in\n\n[: a full](https://github.com/openxla/xla/tree/main/xla/stream_executor/rocm)\n\n`xla/stream_executor/rocm/`\n\n`all_reduce_kernel_rocm.cc`\n\n, multi-GPU barrier, and ragged all-to-all kernels that bring AMD off the CUDA-shim path for these operations.**Build, hermeticity, and runtime hygiene.** The [hermetic LLVM toolchain](https://github.com/openxla/xla/pull/39703) is the largest gain — XLA’s ROCm build no longer depends on the host’s `clang`\n\nversion, which was the single biggest reproducibility hazard for downstream packagers. Other changes in the same vein [streamlined the Bazel targets](https://github.com/openxla/xla/pull/40385) for ROCm libraries, made `LoadKernel`\n\nuse a [ref-counted module path](https://github.com/openxla/xla/pull/40847) so cleanup is correct, propagated [proper error status through the ROCm profiler](https://github.com/openxla/xla/pull/38777), and fixed a subtle [leading-comma bug](https://github.com/openxla/xla/pull/41483) in the AMDGPU feature string passed to LLVM.\n\nNet effect.The JAX and XLA ROCm plugins are now at feature parity with the rest of the backends, and deliver strong performance on AMD Instinct GPUs for bf16 and FP8 transformer training and inference, large-scale collectives, and Triton-fused GEMM epilogues.\n\n### How ROCm Gets Tested in `openxla/xla`\n\n[#](#how-rocm-gets-tested-in-openxla-xla)\n\nA backend without CI is a backend that suffers from bit rot. The defining ROCm investment in `openxla/xla`\n\nover the past year has been the **unification of ROCm CI into a single upstream GitHub Actions workflow** ([PR #36893](https://github.com/openxla/xla/pull/36893)), driven from [.github/workflows/rocm_ci.yml](https://github.com/openxla/xla/blob/main/.github/workflows/rocm_ci.yml). Every PR against `main`\n\nnow runs through it on real AMD silicon before it can be merged.\n\n#### The Workflow at a Glance[#](#the-workflow-at-a-glance)\n\nJob |\nRunner Label |\nAMD Product |\nCoverage |\n|---|---|---|---|\n|\n|\nMI350 ( |\nJAX unit tests built against the PR’s XLA, single-GPU |\n|\n|\nMI350 ( |\nXLA’s own test suite under the |\n\nBoth jobs run inside the `rocm/tensorflow-build:latest-jammy-pythonall-rocm7.2.1-ci_official`\n\ncontainer, pinned by SHA digest for supply-chain hygiene. `/dev/kfd`\n\nand `/dev/dri`\n\nare mapped through, an 80 GiB tmpfs Bazel cache is mounted, and the `video`\n\ngroup is added so HIP can reach the kernel driver. `rocminfo`\n\nis invoked early in the run so a bad host fails the first step rather than burying the error in test logs.\n\n#### Build-System Plumbing[#](#build-system-plumbing)\n\nThe CI is driven by Bazel `--config`\n\nflags defined in [build_tools/rocm/rocm_xla.bazelrc](https://github.com/openxla/xla/blob/main/build_tools/rocm/rocm_xla.bazelrc):\n\n`--config=rocm_rbe`\n\n— Remote Build Execution, parallelising build and test actions across many remote workers.`--config=rocm_rbe_dynamic`\n\n— hybrid mode that builds locally but lets test actions schedule across local and remote, so a single PR can saturate both the on-prem GPU pool and the build farm.`--config=ci_single_gpu`\n\n— wraps tests in[build_tools/rocm/parallel_gpu_execute.sh](https://github.com/openxla/xla/blob/main/build_tools/rocm/parallel_gpu_execute.sh)so multiple test shards can share the GPU safely, plus three flaky-test retries.\n\n#### Test Selection[#](#test-selection)\n\nNot every XLA test is meaningful on AMD GPUs — some are specific to other hardware platforms. The ROCm CI filters in two layers:\n\nexcludes roughly fifteen vendor-specific Bazel tags (`rocm_tag_filters.sh`\n\n`cuda-only`\n\n,`requires-gpu-sm`\n\n, Intel-GPU, and similar) so test discovery stays tractable.The\n\n`test:xla_sgpu`\n\nlist in`rocm_xla.bazelrc`\n\nenumerates the exact targets the single-GPU pool runs, via explicit excludes.\n\nThe XLA job pulls [ execute_ci_build_upstream.sh](https://raw.githubusercontent.com/ROCm/xla/refs/heads/rocm-dev-infra/build_tools/rocm/execute_ci_build_upstream.sh) from AMD’s\n\n`ROCm/xla`\n\nfork at workflow time. That gives the AMD CI team a fast iteration path on the runner-side script (test selection, failure triage, log post-processing) without round-tripping through `openxla/xla`\n\nfor every change. The workflow file, the Bazel configs, and the test target lists remain upstream and reviewable.## Part 2 — JAX on AMD[#](#part-2-jax-on-amd)\n\nJAX uses XLA as its compiler, but the ROCm story is not just “inherit XLA’s backend”. JAX ships a separate plugin, separate wheels, and runs its own four-workflow CI.\n\n### How JAX Loads the ROCm Plugin[#](#how-jax-loads-the-rocm-plugin)\n\nFigure 2 traces the loader path JAX walks on `import jax`\n\n, ending at a registered `RocmDevice`\n\n:\n\n*Figure 2. The JAX ROCm plugin loader path — from import jax down to a registered RocmDevice.*\n\nYellow boxes run in the Python interpreter; the green box is the native shared library compiled into the PJRT wheel; the dashed grey box is the bundled fallback used only if neither dedicated plugin is installed. The loader probes `jax_rocm7_plugin`\n\non import, picking up the ROCm 7 plugin automatically when present.\n\nThe relevant code lives under:\n\n[jax_plugins/rocm/](https://github.com/jax-ml/jax/tree/main/jax_plugins/rocm)— the Python plugin entry point that registers ROCm with`xla_bridge`\n\n.[jaxlib/rocm/](https://github.com/jax-ml/jax/tree/main/jaxlib/rocm)— the native plugin extension (`rocm_plugin_extension.cc`\n\n) that exposes ROCm-specific FFI types and custom-call handlers across the C ABI.[rocm/rocm-jax](https://github.com/rocm/rocm-jax)— AMD’s infrastructure repo, with the Dockerfiles and tooling used to build and ship the`rocm/jax`\n\nimages for each ROCm version.\n\nAt install time, ROCm support ships as **two separate wheels**:\n\nWheel |\nContents |\n|---|---|\n|\nThe native PJRT C-API plugin ( |\n|\nThe Python wrapper that JAX’s |\n\n`{N}`\n\nis the major ROCm version (today `7`\n\n). The user-facing install instructions live in\n[docs/installation.md](https://github.com/jax-ml/jax/blob/main/docs/installation.md);\nthe Dockerfile-based path lives in\n[rocm/rocm-jax](https://github.com/rocm/rocm-jax);\nand a prebuilt image is published as `rocm/jax:latest`\n\n.\n\nWhy two wheels?The split lets AMD ship post-release fixes (`.postN`\n\nbumps) on the PJRT wheel without forcing a JAX version bump, and lets you co-install multiple ROCm-major-version plugins on the same host without conflicts.\n\n### What Landed in ROCm for the Past Year[#](#what-landed-in-rocm-for-the-past-year)\n\nThe JAX-side work has been similarly active over the past year.\n\n**Correctness.** AMD contributors landed a [Pallas inter-block write race fix](https://github.com/jax-ml/jax/pull/37183) for non-range while-loops — a real kernel synchronization bug on ROCm — and added two targeted skips where hipSolver’s semantics diverge from cuSolver: [complex paths in testEighIdentity](https://github.com/jax-ml/jax/pull/36909) and the\n\n[.](https://github.com/jax-ml/jax/pull/36984)\n\n`tridiagonal_solve_perturbed`\n\npath inside `eigh`\n\n**Test infrastructure.** ROCm pytest was split TPU-style into [single- and multi-accelerator passes](https://github.com/jax-ml/jax/pull/36851) (with follow-up parallelization in commit `663efe75a`\n\n); each pytest-xdist worker now gets its own GPU through a [per-worker HIP_VISIBLE_DEVICES override](https://github.com/jax-ml/jax/pull/37054) gated by\n\n`JAX_ENABLE_ROCM_XDIST`\n\n; and the ROCm build wired up [so JAX’s ROCm CI can pin against XLA HEAD instead of JAX’s own XLA pin.](https://github.com/jax-ml/jax/pull/36355)\n\n`clone_main_xla`\n\nplumbing**Wheels and packaging.** ROCm wheels [moved off direct S3 to a CloudFront-backed CDN](https://github.com/jax-ml/jax/pull/36684), `auditwheel`\n\nwas [taught to accept manylinux_2_28](https://github.com/jax-ml/jax/pull/36621) — opening the door to install on a much wider set of Linux distributions out of the box — and\n\n[to track the ROCm-side updates.](https://github.com/jax-ml/jax/pull/37072)\n\n`rules_ml_toolchain`\n\nwas bumped**Workflow hygiene.** The ROCm jobs in [ bazel_rocm.yml](https://github.com/jax-ml/jax/blob/main/.github/workflows/bazel_rocm.yml) and the wheel-download composite action carry explicit\n\n`zizmor`\n\noverrides where the linter’s defaults conflicted with what the ROCm pipeline actually needs to do.The pattern is consistent: correctness fixes, production-grade packaging, and CI plumbing that lets ROCm-side and XLA-side changes ride the same trains as everything else. ROCm is being maintained as a first-class target, not a side branch.\n\n### How ROCm Gets Tested in `jax-ml/jax`\n\n[#](#how-rocm-gets-tested-in-jax-ml-jax)\n\nJAX runs **four ROCm GitHub Actions workflows**:\n\nWorkflow |\nPurpose |\nHardware |\n|---|---|---|\nFull Bazel test sweep on RBE |\n1- and 4-GPU AMD pools |\n|\nLightweight presubmit gate |\nSingle AMD GPU |\n|\nPython-level pytest with multi-accelerator separation |\n1- / 4- / 8-GPU AMD |\n|\nBuilds the |\nmanylinux_2_28 builder |\n\nThe runner-side scripts live in\n[jax/ci/](https://github.com/jax-ml/jax/tree/main/ci):\n\n[run_bazel_test_rocm_rbe.sh](https://github.com/jax-ml/jax/blob/main/ci/run_bazel_test_rocm_rbe.sh)— the Bazel-RBE entry point. Honors`JAXCI_CLONE_MAIN_XLA=1`\n\nto swap in an XLA-HEAD checkout via`--override_repository=xla=…`\n\n, which is how OpenXLA PRs pre-flight against JAX before merge.[run_pytest_rocm.sh](https://github.com/jax-ml/jax/blob/main/ci/run_pytest_rocm.sh)— the pytest entry point. Single-accelerator tests run under pytest-xdist with`JAX_ENABLE_ROCM_XDIST`\n\nset to the GPU count; multi-accelerator tests (`-m \"multiaccelerator\"`\n\n) run serially with the full GPU set.[build_rocm_artifacts.sh](https://github.com/jax-ml/jax/blob/main/ci/build_rocm_artifacts.sh)— drives`python build/build.py --wheels=jax-rocm-plugin,jax-rocm-pjrt`\n\nand runs`auditwheel`\n\nfor manylinux compliance.[upload_rocm_logs.sh](https://github.com/jax-ml/jax/blob/main/ci/upload_rocm_logs.sh)— ships test logs to S3/CloudFront for triage.\n\nContainers used:\n\n`ghcr.io/rocm/jax-dev-ubu24.rocm720:latest`\n\nfor Bazel test workflows.`ghcr.io/rocm/jax-base-ubu24.rocm720:latest`\n\nfor pytest workflows (runtime-trimmed image).`ghcr.io/rocm/jax-manylinux_2_28-rocm-7.2.0:latest`\n\nfor wheel building.\n\nDefault coverage in upstream CI today: ROCm 7.2.x; Python 3.11 through 3.14; MI350 (`gfx950`\n\n) for `pytest_rocm.yml`\n\nand MI300 in CPX mode (`gfx942`\n\n) for `bazel_rocm.yml`\n\n. Other Instinct generations (MI200 `gfx90a`\n\n) and RDNA Radeon targets are exercised in AMD’s downstream CI; upstream coverage expands as runner capacity comes online.\n\nThe xdist isolation pattern.The[hook pins each xdist worker to a single physical GPU by setting both]`conftest.py`\n\n`ROCR_VISIBLE_DEVICES`\n\n(so ROCr enumerates only that GPU) and`HIP_VISIBLE_DEVICES=0`\n\n(so HIP doesn’t re-enumerate hidden agents). Without that pairing, multi-process pytest either crashes on contention or silently colocates workers on device 0. Worth borrowing for any multi-process ROCm test harness.\n\n## The Integrated CI Pipeline[#](#the-integrated-ci-pipeline)\n\nXLA and JAX CI are not independent systems. XLA pre-flights every PR through JAX; JAX can pin against XLA HEAD via Bazel’s `--override_repository`\n\n. Two workflow cadences ride on top of that coupling:\n\n**Nightly**— JAX HEAD built and tested against the XLA commit pinned in JAX’s`WORKSPACE`\n\n. This is the day-to-day regression signal for JAX itself.**Continuous**— JAX HEAD built and tested against XLA HEAD (`JAXCI_CLONE_MAIN_XLA=1`\n\noverrides the pin). This is what catches XLA regressions in the window between XLA-pin bumps.\n\nA **ToT ROCm** axis is being rolled into the same matrix shortly, adding the ROCm release in the container as a third moving piece exercised against JAX HEAD.\n\nThe result is one cross-repo pipeline (Figure 3).\n\n*Figure 3. How a ROCm PR flows through XLA CI, JAX CI, and shared infrastructure to land on AMD Instinct runners. Dotted arrows are the cross-repo integration edges (XLA pre-flighting JAX, and JAX pinning XLA HEAD).*\n\nBlue is a PR trigger; green is XLA CI; yellow is JAX CI; red is shared build and test infrastructure; grey is the physical AMD Instinct runner pool. Solid arrows are intra-workflow control flow. The two **dotted** arrows are what make this one pipeline instead of two:\n\n**XLA → JAX pre-flight.** A PR against`openxla/xla`\n\ntriggers the`jax`\n\njob in, which checks out`.github/workflows/rocm_ci.yml`\n\n`jax-ml/jax`\n\nand runs JAX’s ownwith`run_bazel_test_rocm_rbe.sh`\n\n`--override_repository=xla=$GITHUB_WORKSPACE`\n\n. An XLA change that would silently break JAX gets a red check before merge.**JAX → XLA HEAD pin.** Setting`JAXCI_CLONE_MAIN_XLA=1`\n\nmakes the same script clone the latest XLA`main`\n\nand override the repo, so nightly JAX runs catch XLA regressions in the window between XLA-pin bumps in JAX’s`WORKSPACE`\n\n.\n\nBoth directions terminate at the same backing systems: a third-party RBE cluster for build and test scheduling, and the **AMD Instinct runner pools** for actual GPU execution.\n\n## Try It Out and Get Involved[#](#try-it-out-and-get-involved)\n\nThe fastest path from this post to a JIT-compiled JAX program on Instinct hardware. Pick the entry point that matches your environment.\n\n### Path 1 — Docker (Lowest Friction)[#](#path-1-docker-lowest-friction)\n\nWith an AMD Instinct GPU and a host running ROCm-capable kernel modules, the prebuilt JAX-on-ROCm image is the shortest path:\n\n```\ndocker pull rocm/jax:latest\n\ndocker run -it --rm \\\n  --device=/dev/kfd --device=/dev/dri \\\n  --group-add video --shm-size=64G \\\n  --ipc=host --cap-add=SYS_PTRACE \\\n  --security-opt seccomp=unconfined \\\n  rocm/jax:latest\n```\n\nInside the container:\n\n``` python\nimport jax\nimport jax.numpy as jnp\n\nprint(jax.devices())   # should list ROCm devices\n\n@jax.jit\ndef f(x):\n    return jnp.tanh(x @ x.T)\n\nx = jnp.ones((1024, 1024), dtype=jnp.bfloat16)\nprint(f(x).block_until_ready().shape)\n```\n\nA `RocmDevice`\n\nin `jax.devices()`\n\nconfirms that the plugin loaded and that XLA is compiling through the AMDGPU LLVM backend.\n\n### Path 2 — pip on a Host with ROCm Installed[#](#path-2-pip-on-a-host-with-rocm-installed)\n\nFor an existing ROCm 7 install on the host:\n\n```\npip install --upgrade \"jax[rocm7-local]\"\n```\n\nThis pulls `jax`\n\n, `jax-rocm7-plugin`\n\n, and the matching `jax-rocm7-pjrt`\n\nwheel from PyPI. JAX does **not** install the ROCm toolkit itself — install the runtime first via the [ROCm installation guide](https://rocm.docs.amd.com/). Post-release fixes ship as `jax-rocm7-plugin==X.Y.Z.postN`\n\nand can be upgraded independently of the JAX version.\n\n### Path 3 — Building XLA from Source[#](#path-3-building-xla-from-source)\n\nFor compiler-side work rather than running workloads:\n\n```\ngit clone https://github.com/openxla/xla.git\ncd xla\n./configure.py --backend=ROCM --rocm_path=/opt/rocm\nbazel test --config=rocm //xla/...\n```\n\nThe same `--config=rocm_rbe`\n\nand `--config=ci_single_gpu`\n\noptions that upstream CI uses are available locally; see [build_tools/rocm/rocm_xla.bazelrc](https://github.com/openxla/xla/blob/main/build_tools/rocm/rocm_xla.bazelrc).\n\n### Where to Go from Here[#](#where-to-go-from-here)\n\nIf you want to… |\nStart here |\n|---|---|\nRead installation specifics |\n|\nUnderstand the XLA build |\n|\nLook up ROCm itself |\n|\nWatch CI status |\n|\nFile an XLA bug |\n|\nFile a JAX bug |\n|\nSend a PR |\n|\nSee AMD’s staging branches |\n\nThe highest-leverage contributions from outside AMD, in our experience:\n\n**Performance reports with HLO dumps.** The dump flags and tooling are documented inand`docs/hlo_dumps.md`\n\n. A reproducible HLO module turns “this is slow” into a tractable issue.`docs/tools.md`\n\n**Numerical-divergence reports.** A workload that runs cleanly on CUDA but produces different numerics on ROCm is exactly the kind of signal AMD reviewers prioritise — open an issue with a small reproducer.If you run a target outside the default CI matrix (RDNA in particular), reports of what works and what doesn’t directly inform the next CI expansion.`gfx`\n\ncoverage on the long tail.\n\nEvery ROCm-touching PR against either repo runs through the workflows above and returns real-hardware results within a couple of hours. That feedback loop is the entire point.\n\n## Summary[#](#summary)\n\nIn the past twelve months, AMD contributors and the broader OpenXLA / JAX community landed:\n\n**Triton on AMDGPU** at feature parity for matmul, scaled-dot, and the AllReduce groundwork.**hipBLASLt group-GEMM**,** FP8 fast accumulation**on ROCm 7, and** rocPRIM**integration for batched scans in XLA.** Hermetic LLVM**in the XLA build, clean`hipcc`\n\ntoolchain ordering, and**manylinux_2_28** wheels for JAX.**One unified upstream ROCm CI workflow** for XLA, and a**four-workflow matrix** for JAX with single- and multi-accelerator separation and proper xdist isolation.**JAX ↔ XLA cross-pin**, so changes on either side pre-flight against the other on real Instinct GPUs before merge.\n\nThe roadmap ahead:\n\n**Triton AllReduce.** Complete the four-part series and turn on Triton-fused collectives in production XLA pipelines.Broaden CI coverage past the current`gfx942`\n\n,`gfx950`\n\n, and beyond.`gfx950`\n\nsingle-GPU pool to next-generation Instinct parts as they come online, on both XLA and JAX runners.**Autotuning corpora for AMD.** Extend the persisted-autotuning and collective-perf-table machinery (additions in[PR #40653](https://github.com/openxla/xla/pull/40653)) with AMD-specific tuning data shipped alongside ROCm releases.**Multi-host CI.** Extend coverage from intra-node 4-GPU collectives to multi-host runs so distributed JAX and XLA workflows are validated end to end.**WSL2 graduation.** Move JAX-on-WSL2 from experimental () to a tested CI lane.`docs/installation.md`\n\nIn this blog we covered the full ROCm story for OpenXLA and JAX: the backend architecture, a year of upstream contributions, and the CI infrastructure that gates every PR on real Instinct hardware. ROCm is now a first-class OpenXLA target — upstream, gated on real hardware, and visible in every PR. The quick start at the top of this post is the shortest path from here to a JIT-compiled JAX program on an AMD Instinct GPU, and to filing the next bug or PR that moves the stack forward.\n\n## Disclaimers[#](#disclaimers)\n\nThird-party content is licensed to you directly by the third party that owns the content and is not licensed to you by AMD. ALL LINKED THIRD-PARTY CONTENT IS PROVIDED “AS IS” WITHOUT A WARRANTY OF ANY KIND. USE OF SUCH THIRD-PARTY CONTENT IS DONE AT YOUR SOLE DISCRETION AND UNDER NO CIRCUMSTANCES WILL AMD BE LIABLE TO YOU FOR ANY THIRD-PARTY CONTENT. YOU ASSUME ALL RISK AND ARE SOLELY RESPONSIBLE FOR ANY DAMAGES THAT MAY ARISE FROM YOUR USE OF THIRD-PARTY CONTENT.", "url": "https://wpnews.pro/news/openxla-and-jax-rocm-support-and-the-state-of-ci", "canonical_source": "https://rocm.blogs.amd.com/software-tools-optimization/openxla-jax-rocm/README.html", "published_at": "2026-06-29 00:00:00+00:00", "updated_at": "2026-06-30 20:00:00.738518+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-infrastructure", "ai-tools", "developer-tools"], "entities": ["OpenXLA", "JAX", "AMD", "ROCm", "XLA", "Instinct", "GitHub Actions", "Triton"], "alternates": {"html": "https://wpnews.pro/news/openxla-and-jax-rocm-support-and-the-state-of-ci", "markdown": "https://wpnews.pro/news/openxla-and-jax-rocm-support-and-the-state-of-ci.md", "text": "https://wpnews.pro/news/openxla-and-jax-rocm-support-and-the-state-of-ci.txt", "jsonld": "https://wpnews.pro/news/openxla-and-jax-rocm-support-and-the-state-of-ci.jsonld"}}