Skip to content

OpenXLA Backend Source Status

mlxcel uses MLX as its production compute backend on Apple Silicon and CUDA, and as its experimental backend on ROCm. Alongside it the repository carries a StableHLO/IREE implementation under src/lib/mlxcel-xla/, reached through backend seams guarded by cfg(feature = "xla-backend") and cfg(feature = "xla-iree"). The root Cargo.toml declares those features and the mlxcel-xla dependency, and every one of them is default-off.

Build Configurations

  • A stock cargo build, Metal build, or CUDA build compiles none of the OpenXLA path and contains only the MLX backend.
  • Setting MLXCEL_BACKEND=xla on a stock binary does nothing, because the selector itself is behind the feature gate.
  • cargo build --features xla-backend compiles the seam and the StableHLO graph emitter. It needs no native toolchain, so it builds anywhere the stock binary does, but it has no execution engine.
  • cargo build --features xla-iree adds the executing runtime. Its build script compiles a C shim against a prebuilt IREE distribution, so one of IREE_DIST, IREE_CUDA_HOME (with IREE_CUDA_COMPILE), IREE_MACOS_HOME (with IREE_MACOS_COMPILE), or IREE_ROCM_HOME (with IREE_ROCM_COMPILE) must be set at build time. scripts/iree/setup-cuda.sh, scripts/iree/setup-macos.sh, and scripts/iree/setup-rocm.sh (also make iree-rocm) produce the source-built trees for the last three.
  • xla-diagnostics, xla-micro-oracle and xla-diagnostics-cpu layer instrumentation on top of xla-iree. xla-diagnostics pairs the CPU diagnostics with a CUDA build, and xla-diagnostics-rocm pairs them with a ROCm build (rocm plus xla-diagnostics-cpu) for intermediate-tensor oracle checks on an AMD GPU.
  • On ROCm, the IREE HIP target runs with MLXCEL_XLA_DEVICE=hip and a single gfx target in MLXCEL_XLA_ROCM_TARGET. It is qualified on gfx1151, where 64-token greedy decode of Llama-3.2-1B-Instruct-4bit and Qwen3-0.6B-4bit matches the CPU f32 oracle token for token at the default f16 precision (docs/benchmark_results/xla_iree_rocm_qualification.md).

The XLA source directory and its README stay development material. A build that enables these features is a local, opt-in build, not a CI or release configuration.

Model Coverage

The OpenXLA path runs a fixed set of text architectures through its StableHLO emitter, plus the vision path of a few VLM families on the shared CLI/server prepared-prefill path:

  • LLaVA: the vision tower and projector run on the resident IREE module, with the MLX host path as fallback (MLXCEL_XLA_VISION_BACKEND).
  • Molmo 2: the vision path runs on the resident IREE module in an xla-iree build; with MLXCEL_XLA_VISION_BACKEND=host, or without xla-iree, the session has no image path.
  • Qwen3-VL: the vision tower, the main merger, and every DeepStack merger run on the resident IREE module (qwen3_vl_runtime.rs).
  • Youtu-VL: the vision tower and its built-in merger run on the resident IREE module (youtu_vl_runtime.rs).
  • Gemma 3 VLM: the embeddings-prefill entry takes the transformers image mask (causal plus bidirectional attention inside each image block) as an additive prefill mask (prepared_gemma3.rs).
  • Gemma 3n audio: the audio encoder and the language merge run as two resident IREE modules; only mel extraction stays on the host (gemma3n_audio_runtime.rs).
  • Qwen2-VL and Qwen2.5-VL: the vision tower runs only on the resident IREE module, so a build without xla-iree fails startup for these checkpoints, and MLXCEL_XLA_VISION_BACKEND=host is rejected. The Qwen2.5-VL graph implements the windowed vision attention (window permutation, window-isolated masks, and the checkpoint's full-attention layers) and rejects temporal grids it does not support.

Muse Glimmer VLM is refused for OpenXLA image execution with a message that points to the MLX backend; other VLM families have no OpenXLA image path.

Retained Design

The implementation uses a session-level compute-backend seam: token input, backend-owned KV state, decode, and sampling belong to one backend session rather than wrapping individual tensor operations. It emits StableHLO/MLIR and has IREE runtime code for CPU, CUDA, Metal, and HIP (ROCm) targets. The session-contract design is recorded in ADR 0004 (docs/adr/0004-compute-backend-session-seam-and-stablehlo-family.md) in the repository.

Environment Variables

MLXCEL_BACKEND, MLXCEL_XLA_DEVICE, MLXCEL_XLA_PRECISION, MLXCEL_XLA_QUANT, MLXCEL_XLA_ROCM_TARGET, MLXCEL_XLA_VISION_BACKEND, MLXCEL_XLA_CONTEXT_CAPACITY, and MLXCEL_XLA_IREE_COMPILE are consumed only by code behind the feature gates, so they are inert in a stock binary and live in one built with xla-backend or xla-iree. See Environment variables.

For Backend Developers

Treat src/lib/mlxcel-xla/README.md and that ADR as implementation notes. Before advertising OpenXLA as a product feature, build both the seam-only and IREE configurations and re-run real-checkpoint token-parity tests on each claimed target. Operator documentation and deployment automation select MLX only.