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, Debug and Performance Tools, the JAX Scaling Book for AMD GPUs, and serves as a central index for all ROCm JAX/XLA release branches.
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
- Scaling Book
- Workflows
For supported ROCm versions and GFX targets per release, see the AMD ROCm compatibility matrix.
Use the same ROCm-less build container as CI for reproducible local builds and CI parity:
docker pull ghcr.io/rocm/jax-build-ubu24.rocmless@sha256:68854ed30b800d3e8c390f6e651721e2bead58830da1c59e1b828354eb8f077fThis is the digest pinned in the rocm-config job of the main CI workflows; the image's latest tag is a newer build that CI does not use. The image is Ubuntu 24.04 with Python 3.12, Bazel (via Bazelisk), and the build dependencies needed by JAX and XLA, but it has no ROCm install (no /opt/rocm). Instead, Bazel downloads a TheRock ROCm tarball and builds against it with the hermetic Clang toolchain, so neither a system ROCm nor a system Clang is required. Mount your local jax and xla source checkouts into the container (e.g. as /tf/jax and /tf/xla), and pass --device=/dev/kfd --device=/dev/dri --group-add video to docker run so tests can see the GPU, before running the build commands below.
Bazel picks the TheRock tarball from these flags, which the main CI jobs all pass:
--repo_env=ROCM_PATH=
--repo_env=ROCM_DISTRO_URL=https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz
--repo_env=ROCM_DISTRO_HASH=1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6-
Keep
ROCM_PATHempty, and don't pass--rocm_pathtobuild/build.py. A non-emptyROCM_PATHtakes priority overROCM_DISTRO_URL, so Bazel would use a local ROCm install instead of the tarball. -
TheRock 10.0.0 is also XLA's default hermetic ROCm (
rocm_latest_multiarchinthird_party/gpus/rocm/rocm_redist.bzl); the twoROCM_DISTRO_*flags pin it explicitly. -
To build against a TheRock nightly, set
ROCM_DISTRO_HASH=no_hash(which skips the checksum) and pointROCM_DISTRO_URLat that nightly's tarball. The nightly URL differs between release lines and is subject to change, so check thetherock_nightlyentries inrocm_xla_ci_nightly.ymlandrocm_xla_ci_jax.ymlfor the current ones. At the time of writing:- 10.1:
https://rocm.nightlies.amd.com/tarball-multi-arch/therock-dist-linux-multiarch-<version>.tar.gz, with the version index athttps://rocm.nightlies.amd.com/whl-multi-arch/. - 10.2:
https://nightly.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-<version>.tar.gz, with the version index athttps://nightly.repo.amd.com/rocm/whl-next/.
To print the latest
<version>, pass the line's version index and number, for example for 10.2:GPU_TARGET=device-all THEROCK_INDEX_URL=https://nightly.repo.amd.com/rocm/whl-next/ ROCM_SDK_VERSION=10.2 \ bash build_tools/rocm/detect-therock-version.sh
- 10.1:
-
CI runs the XLA unit tests with
build_tools/rocm/run_xla_ci_build.shfromopenxla/xla, and builds and tests JAX withci/run_bazel_test_rocm_rbe.shfromjax-ml/jax, both with these flags on top of the ROCm RBE configs. Seerocm_xla_ci_nightly.ymlandrocm_xla_ci_jax.ymlfor the exact commands.
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-rocm10-pjrt jax-rocm10-plugin jax-rocm7-pjrt jax-rocm7-plugin -y
rm -rf dist/*
python build/build.py build \
--wheels=jax,jaxlib,jax-rocm-plugin,jax-rocm-pjrt \
--rocm_amdgpu_targets=gfx950 \
--rocm_version=10 \
--bazel_options=--override_module=xla=/tf/xla \
--bazel_options=--config=rocm_clang_hermetic \
--bazel_options=--repo_env=ROCM_PATH= \
--bazel_options=--repo_env=ROCM_DISTRO_URL=https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz \
--bazel_options=--repo_env=ROCM_DISTRO_HASH=1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6 \
--bazel_options=--repo_env=TF_ROCM_AMDGPU_TARGETS=gfx950 \
--bazel_startup_options="--bazelrc=build/rocm/rocm.bazelrc"
pip install dist/*.whl
# The rocmless image has no ROCm runtime, so point the loader at the TheRock
# libraries Bazel extracted during the build.
ROCM_DIST="$(bazel --bazelrc=build/rocm/rocm.bazelrc info output_base)/external/xla++rocm_configure_ext+local_config_rocm/rocm/rocm_dist"
export LD_LIBRARY_PATH="${ROCM_DIST}/lib:${ROCM_DIST}/lib/rocm_sysdeps/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
# Run outside /tf/jax so Python imports the installed wheels, not the source tree.
cd /tmp && python3 -c "import jax; print(jax.devices())"Notes:
- Use
--override_module=xla=/tf/xla, NOT--override_repository=xla=.... Current JAX builds run under bzlmod (.bazelrcsets--enable_bzlmod --noenable_workspace), where XLA is a Bazel module whose canonical repo name isxla+(note the trailing+). The legacy--override_repository=xla=...targets a repo literally namedxla, which does not exist under bzlmod, so it is silently ignored — Bazel then builds XLA from the bzlmod-pinned commit and your local checkout is never compiled. - The example builds for
gfx950, the architecture CI builds and tests on. Adjust--rocm_amdgpu_targetsand--repo_env=TF_ROCM_AMDGPU_TARGETSto match your hardware (e.g.gfx90a,gfx942,gfx950,gfx1100), keeping both the same. Multiple targets can be passed comma-separated.--rocm_amdgpu_targetsalone does not reach Bazel's ROCm setup, becausebuild/rocm/rocm.bazelrcsetsTF_ROCM_AMDGPU_TARGETSto 27 architectures with--repo_env, which makes the build much slower and the wheels much larger. - The
--wheels=jax,jaxlib,jax-rocm-plugin,jax-rocm-pjrtlist produces all four required wheels indist/;pip install dist/*.whlinstalls the matching set. Built against TheRock 10.0.0, the ROCm wheels are namedjax_rocm10_pluginandjax_rocm10_pjrt. - The last command should print
[RocmDevice(id=0), ...]. Without theLD_LIBRARY_PATHstep, JAX prints a plugin error (librocprofiler-sdk.so.1: cannot open shared object file) and falls back to[CpuDevice(id=0)]. SetLD_LIBRARY_PATHagain in every new shell, and rebuild after abazel clean --expunge, which deletes the extracted libraries.
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 main and the 0.11.x+ branches, use the build/build.py invocation shown in the previous section instead. The 0.10.x branches' XLA predates TheRock tarball support (ROCM_DISTRO_URL), so build those in the ROCm 7.2.4 image (ghcr.io/rocm/jax-ubu22.rocm7.2.4) with the same command, replacing the three --repo_env=ROCM_* options with --rocm_path=/opt/rocm and using --rocm_version=7.
Track pending and in-progress ROCm upstream contributions to OpenXLA on the OpenXLA:GPU - AMD/ROCm project board.
Every openxla/xla PR runs the upstream rocm_ci.yml workflow on a gfx950 runner, using the ROCm-less image and TheRock 10.0.0 tarball described in Recommended Docker Image && Build Instructions. It has two jobs:
| Job | Entry point | Purpose |
|---|---|---|
xla |
build_tools/rocm/run_xla_ci_build.sh (in openxla/xla) |
XLA test sweep. Runs the //xla/... single-GPU and CPU test suites. |
jax |
ci/run_bazel_test_rocm_rbe.sh (in jax-ml/jax) |
JAX test sweep against XLA-from-source. Builds jaxlib and JAX with Bazel against the PR's XLA and runs the JAX GPU tests. Runs only for PRs targeting main. |
Both jobs use the EngFlow internal.actinolite.cluster.engflow.com RBE cluster through the rocm_rbe / rocm_rbe_dynamic configs in XLA's build_tools/rocm/rocm_xla.bazelrc and JAX's build/rocm/rocm.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. |
The scaling book is a companion to Google's JAX Scaling Book for AMD GPU users, covering performance features for JAX training at the implementation layer (ROCm libraries, MaxText, XLA internals, and JAX APIs).
The scaling_book/ directory contains the Jekyll site served on GitHub Pages. Page content can be found at scaling_book/pages/ with images at scaling_book/pages/img/.
Reference reproductions for the case studies can be found at the following locations:
| Folder | Description |
|---|---|
scaling_book/llama7b-jax-fundamentals |
Rematerialisation, attention backends, and profiling examples |
scaling_book/llama70b-mixed-precision-training |
Mixed quantization experiments in BF16, FP16, FP8, MXFP8, and MXFP4 |
scaling_book/mixtral8-22b-distributed-strategies |
Sharding and grouped GEMM experiments |
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-5.5 powered code review on every pull request |
rocm_xla_ci_nightly.yml |
Daily at 3:47 AM UTC | Runs the latest XLA (openxla/xla) unit test suite, the TSAN and ASAN sanitizer builds (rocm_xla_ci_sanitizers.yml) and the latest JAX (jax-ml/jax) unit tests (rocm_xla_ci_jax.yml), each against the TheRock 10.0.0 release tarball and the latest TheRock nightly tarball (the XLA tests track the 10.2 nightlies, the sanitizer and JAX jobs the 10.1 nightlies) |
rocm_infra_monitor.yml |
Every 1h | Analyzes xla upstream PRs to detect infrastructure problems. If infra problem is detected then new issue is created. |
scaling_book_pages.yml |
Push to rocm-dev-infra when scaling_book/ changes |
Builds the scaling book from scaling_book/ and deploys to GitHub Pages at rocm.github.io/xla/scaling-book |
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). |