This is the rocm-dev-infra branch — the default branch of the ROCm/xla fork of openxla/xla.
It hosts GitHub Actions workflows, CI/CD infrastructure, and serves as a central index for all ROCm JAX/XLA release branches and their hardware support matrix.
The main branch is an exact mirror of openxla/xla:main and should never
contain fork-specific commits. This separation ensures the merge-upstream API
(fast-forward only) continues to work for automatic syncing.
- XLA Branches for JAX
- Recommended Docker Image && Build Instructions
- OpenXLA Upstream PR Status
- OpenXLA Upstream CI Checks
- Debug Tools
- Performance Tools
- Workflows
For supported ROCm versions and GFX targets per release, see the AMD ROCm compatibility matrix.
Use the prebuilt ROCm JAX build container for reproducible local builds and CI parity:
docker pull ghcr.io/rocm/jax-ubu22.rocm7.2.4This image ships ROCm 7.2.4 at /opt/rocm, multiple Python versions, Bazel, and all build dependencies needed by JAX and XLA. The build uses the hermetic Clang-18 toolchain, so no system Clang is required. Mount your local jax and xla source checkouts into the container (e.g. as /tf/jax and /tf/xla) before running the build commands below.
Inside the recommended container, with the JAX repo mounted at /tf/jax and the XLA repo mounted at /tf/xla:
cd /tf/jax
# Remove any pre-existing JAX/jaxlib installs first, in case a mixed-up
# jaxlib is present that would otherwise shadow the freshly built wheels.
python3 -m pip uninstall jax jaxlib jax-rocm-pjrt jax-rocm-plugin -y
rm -rf dist/*
python build/build.py build \
--wheels=jax,jaxlib,jax-rocm-plugin,jax-rocm-pjrt \
--rocm_path=/opt/rocm \
--rocm_amdgpu_targets=gfx942 \
--rocm_version=7 \
--bazel_options=--override_repository=xla=/tf/xla \
--bazel_options=--config=rocm_clang_hermetic \
--bazel_startup_options="--bazelrc=build/rocm/rocm.bazelrc"
pip install dist/*.whlNotes:
--bazel_options=--override_repository=xla=/tf/xlaredirects the@xlarepository to your local checkout, replacing the upstream-pinned commit in JAX'sWORKSPACE.- Adjust
--rocm_amdgpu_targetsto match your hardware (e.g.gfx90a,gfx942,gfx950,gfx1100). Multiple targets can be passed comma-separated. - The
--wheels=jax,jaxlib,jax-rocm-plugin,jax-rocm-pjrtlist produces all four required wheels indist/;pip install dist/*.whlinstalls the matching set.
For the 0.8.x and 0.9.x JAX releases (rows in the XLA Branches for JAX table above), the wheel build is driven by the ROCm/rocm-jax repository on its rocm-jax-infra branch — not by build/build.py in ROCm/jax directly. This is the legacy stack.py / build/ci_build path that pairs the matching rocm-jaxlib-v0.8.x / rocm-jaxlib-v0.9.x branches of ROCm/xla and ROCm/jax.
| JAX Version | Repo used for building | Builds jaxlib wheel? |
|---|---|---|
| 0.8.x | ROCm/rocm-jax |
Yes |
| 0.9.x | ROCm/rocm-jax |
Yes |
| 0.10.x and newer | ROCm/jax (artifact workflow / build/build.py) |
No (only when jaxlib changes) |
Starting with JAX 0.10.0, wheel builds move to the ROCm/jax fork's artifact workflow (ci/build_rocm_artifacts.sh → build/build.py); the stack.py entrypoint in rocm-jax is deprecated and the rocm-jax-infra branch is retained only for Dockerfiles and image-build infrastructure.
To build a 0.8.x / 0.9.x wheel set, follow the build and CI scripts in the rocm-jax-infra branch of ROCm/rocm-jax, pointing them at the corresponding rocm-jaxlib-v0.8.x / rocm-jaxlib-v0.9.x branches of this repository (ROCm/xla) and ROCm/jax. For 0.10.x+ branches and main, use the build/build.py invocation shown in the previous section instead.
Track pending and in-progress ROCm upstream contributions to OpenXLA on the OpenXLA:GPU - AMD/ROCm project board.
PRs that touch upstream-tracking branches (e.g. main, which mirrors openxla/xla:main) run two CI entry-points on the ROCm RBE pool to validate the change against the latest upstream XLA and JAX:
| Script | Purpose |
|---|---|
build_tools/rocm/execute_ci_build_upstream.sh |
XLA test sweep. Shim script that runs the //xla/... test suite on ROCm RBE |
build_tools/rocm/execute_ci_build_upstream_jax.sh |
JAX test sweep against XLA-from-source. Shim script that runs the JAX test suite against XLA built from source on ROCm RBE. |
Both scripts run on the EngFlow wardite.cluster.engflow.com RBE cluster and consume the rocm_rbe / rocm_rbe_dynamic configs from build/rocm/rocm.bazelrc (in the JAX repo) and XLA's own RBE bazelrc.
Debugging utilities for triaging issues while running MaxText-TE / JAX / TF workloads on AMD GPUs (ROCm). See debug_tools/ for full details. Individual tools live in their own subdirectories with self-contained usage notes.
| Tool | Purpose |
|---|---|
debug_tools/queue_dump |
Hang investigation. Drives gdb / rocgdb / umr / the ROCR Debug Agent / RDL to inspect a hung process and its GPU state (queues, kernels, wavefronts, fences). |
debug_tools/nan_check_tool |
NaN / numerical-correctness investigation. Scans intermediate buffers for NaN / Inf and localises the first offending HLO instruction, buffer index, and rank. |
Performance and regression micro-benchmarks for AMD GPUs (ROCm). See perf_tools/ for full details.
| Tool | Purpose |
|---|---|
perf_tools/rocm_latency_check_tool |
HIP-runtime-only micro-benchmark measuring end-to-end per-request dispatch latency (H2D → compute → D2D → compute → D2H) and comparing it across ROCm versions. Reports a latency percentile histogram plus QPS and can append version-stamped rows to a CSV for regression tracking. |
perf_tools/hlo_eval_tools |
Per-module HLO micro-benchmarks built on multihost_hlo_runner. Replays the pre-optimization HLO dumps (training and inference) from real JAX models one module at a time to measure each submodule's cost on XLA, and appends the timings to a CSV for regression tracking across ROCm/XLA versions. Covers models across four categories: LLM, vision, multimodal, and science. |
GitHub Actions workflows and PR labels that drive automation in this fork.
| Workflow | Schedule | Description |
|---|---|---|
sync_upstream.yml |
Every 3h (weekdays), daily (weekends) | Syncs main with openxla/xla:main via the GitHub merge-upstream API |
claude_auto_review.yml |
Every PR opened against a branch that has pr_event_dispatch.yml |
Automatically runs Claude Opus 4.6-powered code review on every pull request |
PRs opened against branches in the ROCm/xla fork can opt-in to additional automated checks by attaching labels:
| Label | Effect |
|---|---|
claude-review |
Triggers an automated Claude code review on the PR. |
sanitizers |
Runs sanitizer-instrumented builds and tests: ThreadSanitizer (TSAN) and AddressSanitizer (ASAN). |