The JAX/ROCm Stack

Chapter 2 of How to Scale Your Model on AMD (Chapter 1: MI355X architecture and system topology | Chapter 3: Profiling a Training Step)

How JAX traces a program into JAXPR and StableHLO, how XLA lowers that onto ROCm backends, and how MaxText flags change the graph before the compiler runs.

Authors

Affiliation

Clarke Chong

ROCm JAX/XLA Team

GPT-5.6-SoL

Published

Sept. 16, 2026

JAX, tracing and JAXPR

JAX adopts a largely functional programming model. Computation is expressed as the composition of pure transformations: functions receive inputs and produce outputs without mutating hidden state. While variables may appear to be reassigned at the Python level, JAX transformations operate on immutable values. The absence of mutable state makes program behaviour easier to reason about because the output of a function depends only on its inputs. More broadly, immutability, referential transparency, and function composition are well-established advantages of the functional programming paradigm; interested readers are referred to Why Functional Programming Matters (Hughes, 1989).

These properties are also valuable from a compiler perspective. Consider a compiler attempting to reorder operations, fuse kernels, eliminate redundant computations, or distribute work across multiple devices. If arbitrary operations could modify shared state, the compiler would need to conservatively preserve execution order. Pure functions instead give the compiler greater freedom to analyze and transform a program while preserving its semantics. In practice, this enables the aggressive whole-program optimizations that form the foundation of modern accelerator software stacks.

JAX builds on this functional model through program transformations. Transformations such as automatic differentiation (grad), vectorization (vmap), and JIT compilation are applied to user-defined functions rather than individual operations. The JAX authors describe JAX as a system for expressing and composing transformations of numerical programs, an idea that underpins much of the framework’s design. Readers interested in this may find the official JAX documentation useful: JAX 101: Expressing Computations.

A key enabling mechanism behind these transformations is tracing. Rather than immediately executing a function, transformations such as jax.jit() first trace the operations performed by the function and construct an IR of the computation. During tracing, concrete values are replaced by tracer objects that record the sequence of primitive operations and the flow of data between them. This representation can then be analyzed, transformed, and compiled for efficient execution on accelerator hardware. The official JAX documentation provides a detailed walkthrough of this process in Tracing.

The first compiler-level representation exposed to users is JAXPR (JAX Expression Representation). JAXPR can be viewed as JAX’s frontend IR, capturing a computation as a sequence of primitive operations and explicit data dependencies. It is intentionally simple, hardware-independent, and well-suited to the transformations that JAX performs, including automatic differentiation, vectorization, and JIT compilation. Readers interested in the broader transformation model may find the JAX documentation on Transformations and Tracing useful supplementary reading.

Below JAXPR lies StableHLO, a portable compiler IR shared across the OpenXLA ecosystem. Although many JAX primitives map naturally to StableHLO operations, the relationship is not necessarily one-to-one.

@jax.jit
def f(x, y):
    return x * y + 1

The traced function is first represented as JAXPR:

{ lambda ; a:f32[] b:f32[].
  let
    c = mul a b
    d = add c 1.0
  in (d,) }

Which is subsequently lowered into StableHLO:

module {
  func.func @main(%arg0: tensor<f32>,
                  %arg1: tensor<f32>)
      -> tensor<f32> {

    %0 = stablehlo.multiply %arg0, %arg1
    %1 = stablehlo.constant dense<1.0>
    %2 = stablehlo.add %0, %1

    return %2
  }
}

StableHLO forms the entry point into the XLA compilation pipeline.

XLA and the GPU pipeline

Take note that going from JAXPR to StableHLO, we transition from JAX-specific representations to the framework-agnostic XLA stack. In other words, HLO is not just a JAX compiler IR, but a shared representation used across multiple frontends including JAX, TensorFlow, and PyTorch integrations.

OpenXLA maintains documentation of XLA internals:

The main points from the ROCm perspective are:

  1. The GPU backend runs platform-specific passes. ROCm and CUDA builds can therefore produce different shardings, fusions, layouts, and library calls from the same StableHLO program.

  1. XLA lowers an HLO operation through an XLA GPU emitter, Triton, or an integrated runtime library.

  1. XLA FFI lets a JAX package or user register an external ROCm implementation. An FFI call appears in HLO as a custom_call, then executes through the registered handler at runtime.

A map of ROCm kernel backends

The three lowering paths, XLA code generation, Triton code generation, and library selection through custom_call, describe how an HLO operation reaches executable device code. In practice, however, it is often more useful to think in terms of what operation is being performed.

A transformer architecture generally consists of a set of operation classes: matrix multiplication, attention, communication, quantization, normalization, activation functions, routing, and numerical primitives. Multiple backend implementations may exist for each of these operations, with the final lowering path determined by shape, datatype, layout, sharding, compiler configuration, and available backend integrations.

The sections below summarize the primary operation classes encountered in modern transformer workloads, together with the ROCm libraries that may implement them.

For example, hipBLAS provides a portable BLAS interface while rocBLAS supplies the underlying ROCm implementation. hipBLASLt extends this model with support for advanced GEMM features such as custom layouts, grouped GEMMs, low-precision datatypes, and fused epilogues. Likewise, Transformer Engine and JAX-AITER expose JAX-facing APIs while relying on lower-level kernel libraries to perform the actual computation.

Dense Matrix Multiplication

GEMMs implement attention projections, MLP projections, output projections, and many optimizer operations.

JAX dot_general and StableHLO dot_general ultimately become HLO dot operations. Depending on the problem being solved, XLA may lower these operations through:

The selected implementation depends on the matrix dimensions, batching structure, layouts, datatypes, epilogues, workspace requirements, and compiler configuration.

Attention

The case studies presented later in this book focus on three attention implementations:

Pallas/Triton-based attention kernels are another possible implementation route.

Although these implementations compute the same attention function, they differ significantly in their tiling strategies, memory movement, supported datatypes, workspace requirements, and backward-pass implementations.

Transformer Engine itself acts as a dispatch layer, selecting an appropriate backend such as Composable Kernel (CK), AITER, or AOTriton depending on the attention configuration and hardware platform.

Low-Precision Training

Modern training increasingly relies on formats such as FP8, MXFP8, and MXFP4 to increase arithmetic throughput and reduce memory consumption.

These formats involve more than a storage datatype. Practical implementations typically define quantization schemes, scaling strategies, accumulation rules, custom kernels, and gradient behavior.

XLA can perform mixed-precision lowering and insert datatype conversions where appropriate. Frameworks such as Transformer Engine and JAX-AITER build additional runtime infrastructure on top of these compiler capabilities to support end-to-end low-precision training.

Collectives and Mixture-of-Experts

Sharding introduces first-class HLO collectives such as all-reduce, all-gather, reduce-scatter, and all-to-all. On ROCm, these operations are implemented through RCCL.

MoE training introduces additional types of work:

  1. Routing, where tokens are assigned to experts.
  2. Dispatch and combine, where tokens move between expert owners.
  3. Expert execution, which performs the expert GEMMs.

Routing is typically expressed using ordinary JAX operations. Dispatch and combine may involve RCCL, MORI, DeepEP, or similar communication libraries. Expert execution typically relies on dense GEMMs, GroupedGEMMs, or specialized MoE kernels.

MORI focuses on expert communication and token movement, while expert computation continues to rely on existing GEMM backends.

General sparse linear algebra is a separate category. Although MoE models are sparse at the routing level, individual experts still tend to execute dense matrix multiplications.

Activations and normalization

RMSNorm, LayerNorm, softmax, SiLU, gating, and related operations are built from pointwise work and reductions. They may become:

Unlike GEMMs or collectives, these operations are often fused into surrounding computation in order to reduce memory traffic and kernel launch overhead.

Other XLA ROCm integrations

XLA’s ROCm runtime also wraps libraries used outside the main transformer path:

Extending the stack

XLA FFI is the extension point used by packages such as JAX-AITER. A package defines a JAX-facing operation, lowers it to an HLO custom call, and registers a ROCm handler that receives the GPU stream and device buffers. The handler can launch its own kernel or call another ROCm library.

This route gives the extension control over its operation contract and runtime implementation. It also gives the extension responsibility for abstract evaluation, batching, autodiff, sharding, layouts, aliases, workspace, and error handling. A forward handler alone is insufficient for training.

Introducing MaxText

MaxText is an open-source large language model training framework developed by Google and built on top of JAX. It serves as both a production-oriented training framework and a reference implementation for modern transformer architectures, distributed training techniques, and accelerator optimizations. In practice, MaxText demonstrates how large-scale foundation models can be implemented within the JAX and OpenXLA ecosystem.

MaxText operates one layer above JAX. Rather than directly modifying the XLA compilation pipeline, MaxText influences the JAX program that enters the compiler stack.

At a high level, MaxText code follows one of two paths:

Standard MaxText layers lower through ordinary JAX primitives. Specialized integrations such as Transformer Engine and JAX-AITER instead lower through FFI-backed custom JAX primitives, which appear in XLA as custom_call operations.

MaxText exposes many of these choices through configuration flags. While such flags may appear as runtime configuration, they frequently alter the generated JAX program itself. Enabling a feature may change which modules are instantiated, which operations are emitted, or whether specialized implementations are selected. Consequently, the effect of a MaxText flag is typically visible before compilation even reaches XLA. The flag modifies the JAX source graph, which in turn changes the generated JAXPR, StableHLO, and subsequent lowering behaviour.

Worked example: Following a MaxText flag though the stack

An instructive example is changing MaxText’s attention implementation from standard JAX attention to Transformer Engine fused attention:

# Standard JAX attention
attention: dot_product

# Transformer Engine fused attention
attention: cudnn_flash_te

The cudnn prefix is a historical compatibility name. On ROCm, this option enters Transformer Engine’s JAX attention path rather than loading cuDNN.

MaxText parses attention as part of its typed configuration, passes it through the model configuration, and uses it to select the attention implementation. These links refer to ROCm/maxtext at b437942a:

  1. types.py defines the attention field.
  2. llama2.py passes the resolved value into the attention module.
  3. attention_op.py contains the implementation dispatch logic.
self.self_attention = Attention(
    config=config,
    attention_kernel=config.attention,
)

This configuration determines which Python code path JAX traces. The dispatch occurs inside AttentionOp.apply_attention:

if self.attention_kernel == "dot_product":
    return self.apply_attention_dot(...)

elif self.attention_kernel == "cudnn_flash_te":
    return self.cudnn_flash_attention(...)

With attention=dot_product, MaxText remains within standard JAX operations. The implementation constructs the attention algorithm by computing QK scores, applying masks, calculating the softmax terms, and multiplying by V:

attn_weights = self.qk_product(query, key, ...)
attn_weights = apply_mask_to_logits(attn_weights, attn_mask)

local_max = jnp.max(logits, axis=-1, keepdims=True)
local_exps_combined = jnp.exp(logits - local_max)
local_sum = jnp.sum(local_exps_combined, axis=-1, keepdims=True)

local_exps = local_exps_combined[..., :s]
local_out = self.wv_product(local_exps, value, ...)

Each operation is independently visible to JAX tracing and therefore appears in the generated HLO at this stage.

With attention=cudnn_flash_te, MaxText instead constructs a Transformer Engine attention module:

from transformer_engine.jax.flax.transformer import DotProductAttention

dpa_layer = DotProductAttention(
    head_dim=head_dim,
    num_attention_heads=self.num_query_heads,
    num_gqa_groups=self.num_kv_heads,
    attn_mask_type=mask_type,
    qkv_layout=qkv_layout,
    # ...
)

return dpa_layer(
    query,
    key,
    value,
    sequence_descriptor=attn_mask,
)

JAX now traces a Transformer Engine operation, which lowers through an FFI-backed custom call rather than exposing the internal attention algorithm. MaxText also enters Transformer Engine’s mesh context, passing the data, tensor, FSDP, and context-parallel axes so the attention implementation observes the same device mesh as the rest of the model.

Changes at the HLO level

The examples below were captured from matched BF16 attention fixtures with Q, K, and V tensors shaped [1, 128, 8, 64]. One fixture uses standard JAX attention, and the other uses the same Transformer Engine path selected by MaxText.

With standard JAX attention, the HLO exposes the attention algorithm directly. The first dot computes QK scores, the reductions and exponential implement the softmax, and the final dot multiplies the resulting probabilities by V:

XLA-rendered HLO subgraph for standard JAX attention

Representative subgraph pruned from XLA’s literal before_optimizations DOT graph.

With Transformer Engine, the same pipeline stage contains the external call. Q, K, V and attention metadata enter custom_call_target="te_fused_attn_forward_ffi", and a get-tuple-element extracts the forward result.

XLA-rendered HLO subgraph for Transformer Engine fused attention

Representative subgraph pruned from XLA’s literal before_optimizations DOT graph. Q, K, V, and metadata converge on the fused-attention custom call.

This is the central pattern behind many MaxText configuration options. The flag does not modify XLA directly, rather it changes the JAX program being traced, producing a different JAXPR, different StableHLO, and a different XLA lowering path.

Next: profiling a training step.

Citation

For attribution in academic contexts, please cite this work as:

    Chong et al., "How to Scale Your Model on AMD", online, 2026.

or as a BibTeX entry:

    @article{scale-your-amd,
      title = {How to Scale Your Model on AMD},
      author = {Chong, Clarke and ROCm JAX/XLA Team and GPT-5.6-SoL},
      howpublished = {Online},
      note = {Retrieved from https://rocm.github.io/xla/scaling-book/},
      year = {2026}
    }