lablup / lablup/mlxcel

test(xla): qualify the OpenXLA/IREE engine on ROCm/HIP

Open
#851 0 comments 0 reactions 0 assignees View on GitHub
area:benchmark area:inference priority:medium status:ready type:test
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Context

`mlxcel` has executable OpenXLA/StableHLO paths for IREE CPU (`local-task`), NVIDIA CUDA (`cuda`), and Apple Metal (`metal`), but it has no ROCm setup, static-link registration, compile-target selection, or recorded AMD GPU qualification result. The StableHLO emitter and the existing correctness/batch harnesses are device-independent, so this issue should add only the ROCm integration required to exercise the same engine and then produce an evidence-backed support decision.

IREE documents ROCm as a stable deployment configuration using the `rocm` compiler backend with the `hip` HAL device. The compiler requires an explicit GPU target such as `gfx942`, and the runtime must include the HIP driver. Use the repository's version-matched IREE pin for both compiler and runtime; at issue creation the intended validation target is `iree-3.12.0rc20260721` at `dc9601f88654749456c7cee4ae87e13de2654e1e`.

This is a qualification issue, not a cross-vendor performance contest. The primary question is whether the current OpenXLA/IREE engine is correct, stable, reproducible, and measurable on one real ROCm machine. Absolute CUDA/Metal comparisons are context only because the hardware differs.

Related work: #449 (OpenXLA backend), #496 (reusable correctness harness), #571 (IREE source-build workflow), and #574 (packed dequant-to-matmul tracking).

## Scope and expected integration points

- Add `scripts/iree/setup-rocm.sh`, mirroring the version-matched compiler/runtime behavior of `setup-cuda.sh`, with a versioned cache and `--info` / `--env` modes.
- Add an explicit `make iree-rocm` target. Linux auto-detection in `make iree` / `make iree-env` may be extended only if CUDA-versus-ROCm selection is deterministic; an explicit target must always remain available.
- Introduce ROCm-specific build variables such as `IREE_ROCM_HOME`, `IREE_ROCM_COMPILE`, and `MLXCEL_XLA_ROCM_TARGET`, while continuing to export the common `MLXCEL_XLA_IREE_COMPILE` path.
- Extend the root `build.rs`, `src/lib/mlxcel-xla/build.rs`, and `src/lib/mlxcel-xla/csrc/xla_iree.c` to link the source-built runtime, retain the HIP registration archive, register the HIP driver, and keep the existing CPU/CUDA/Metal paths unchanged.
- Extend `src/lib/mlxcel-xla/src/iree.rs` so `MLXCEL_XLA_DEVICE=hip` lowers with `--iree-hal-target-device=hip` and `--iree-rocm-target=`. Reject a missing or ambiguous target with an actionable error instead of choosing a generic architecture.
- Update `src/lib/mlxcel-xla/README.md`, `docs/installation.md`, and `docs/environment-variables.md` with the ROCm build, device, target, and multi-GPU selection workflow.
- Reuse `scripts/xla/validate_arch.sh`, `xla_oracle_check`, `xla_batch_bench`, and the normal server path. Generalize `scripts/xla/fusion_spike.sh` only as needed for the optional packed-quant investigation.

Do not add ROCm to the default build graph or require a ROCm host in ordinary CI. The `xla-iree` feature remains explicit and default-off.

## Phase 0: capture a reproducible machine baseline

- [ ] Record OS, kernel, container/runtime boundary if any, GPU SKU, VRAM, GPU count, PCI topology, ROCm version, HIP runtime/driver version, and the exact LLVM target reported by `amdgpu-arch`, `rocm_agent_enumerator`, or `rocminfo`.
- [ ] Verify access to `/dev/kfd` and `/dev/dri`, and record the effective groups/capabilities needed by the process. If running in a container, record the device mounts and security options.
- [ ] On a heterogeneous or multi-GPU host, select one GPU explicitly, record `HIP_VISIBLE_DEVICES`, and never infer a target from a different visible device.
- [ ] Record the repository commit, IREE tag and SHA, compiler version, runtime build flags, model identity/checksum, and all `MLXCEL_XLA_*` variables in a machine-readable environment log.

## Phase 1: build and prove the IREE ROCm toolchain

- [ ] Install the official compiler artifact for the pinned IREE release and validate its checksum. Source-build the runtime from the exact same IREE SHA with `IREE_BUILD_COMPILER=OFF`, `IREE_HAL_DRIVER_DEFAULTS=OFF`, `IREE_HAL_DRIVER_LOCAL_TASK=ON`, `IREE_HAL_DRIVER_LOCAL_SYNC=ON`, and `IREE_HAL_DRIVER_HIP=ON`.
- [ ] Make the setup idempotent and prevent an old CUDA/Metal/ROCm cache from being silently reused after a pin change.
- [ ] Verify that `iree-compile --iree-hal-list-target-backends` includes `rocm`, the linked runtime exposes the `hip` driver, and device enumeration reports the same `gfx*` target captured in Phase 0.
- [ ] Compile a minimal StableHLO add and matmul with `--iree-hal-target-device=hip --iree-rocm-target=`, run both through the HIP device, validate numeric output, and repeat enough times to catch immediate driver/reset instability.

Reference setup and target-selection documentation: https://iree.dev/guides/deployment-configurations/gpu-rocm/ and https://iree.dev/building-from-source/getting-started/.

## Phase 2: wire the normal mlxcel execution path

- [ ] `make iree-rocm` followed by `eval "$(scripts/iree/setup-rocm.sh --env)"` and `cargo build --release --features xla-iree` succeeds from a fresh checkout on the ROCm host.
- [ ] `MLXCEL_XLA_DEVICE=hip` creates the HIP device, compiles both prefill and decode VMFBs for the recorded target, uploads resident weights, and completes a normal `mlxcel generate` request.
- [ ] Wrong target, missing HIP runtime, inaccessible device, and unsupported precision failures surface a concise diagnostic containing the failing stage and remediation.
- [ ] Plain builds and the existing CPU/CUDA/Metal `xla-iree` paths are unchanged; run syntax/lint/unit checks that do not require those other GPUs.

Starter commands:

```bash
export MLXCEL_XLA_ROCM_TARGET=
make iree-rocm
eval "$(scripts/iree/setup-rocm.sh --env)"
cargo build --release --features xla-iree
MLXCEL_BACKEND=xla MLXCEL_XLA_DEVICE=hip MLXCEL_XLA_PRECISION=f32 ./target/release/mlxcel generate -m -p "The capital of France is" -n 64
```

## Phase 3: correctness and serving gates

Use the existing Llama-3.2-1B checkpoint and a fixed prompt/oracle set so the result is directly comparable with the CPU, CUDA, and Metal findings already in the repository.

- [ ] Run `cargo test -p mlxcel-xla --lib validation::tests -- --nocapture` as the structural pre-gate.
- [ ] Run `scripts/xla/validate_arch.sh --model --device hip --max-new 64 --batch 4 --requests 8 --maxcap 24` at `f32`; require token-exact single-sequence output against the HF fp32 oracle and reference-exact batched output.
- [ ] Repeat the trajectory and batch gates at `f16`. Compare against the recorded f32 trajectory for at least 64 greedy steps; if token-exactness is too strict, apply the existing #515 divergence/perplexity criterion and record the first divergent token plus the numeric margin.
- [ ] Exercise both bundled batch sizes (`B_max=4` and `B_max=8`) with mixed request lengths, slot reuse, and at least 100 total generated tokens per configuration.
- [ ] Start `mlxcel-server` with `MLXCEL_BACKEND=xla MLXCEL_XLA_DEVICE=hip`, send one non-streaming and one streaming `/v1/completions` request, and verify completion text, finish reason, usage accounting, and clean shutdown.
- [ ] Run a 30-minute repeated generate/batch smoke or an equivalent bounded loop; fail on hangs, GPU reset, monotonically growing VRAM, corrupted output, or unreleased device state.

## Phase 4: precision, packed quantization, and performance characterization

Base ROCm qualification requires f32 and f16. Treat bf16 and packed quantization as separate capability rows so their failure does not hide a usable base backend.

- [ ] Probe `bf16` compile and execution support for the recorded GPU target. If unsupported, add a fail-fast guard and document the target-specific limitation; if supported, run the same 64-token and batch correctness gates.
- [ ] Run `MLXCEL_XLA_QUANT=packed` for token correctness and adapt the #574 fusion probe to ROCm. Record packed-versus-f16 Flow dispatch counts, whether dequantized weights are materialized, and whether generated AMD code uses an appropriate native integer dot/matrix path.
- [ ] Include the ROCm-supported data-tiling configuration in the packed sweep and report whether it changes dispatch formation, correctness, VRAM traffic, or decode throughput. Update #574 with the result.
- [ ] Measure cold VMFB compile time, warm load time, prefill latency, single-sequence decode tok/s, batch throughput for B4/B8, and peak VRAM for f32, f16, and each supported optional mode.
- [ ] Use at least one warm-up followed by three measured runs per configuration. Report the median and range, raw commands, prompt/token counts, and whether VMFB cache was cold or warm. Compare modes only within the same ROCm host.

For compiler/runtime failures, follow IREE's GPU isolation workflow and preserve the smallest failing StableHLO/VMFB plus compiler diagnostics: https://iree.dev/developers/debugging/gpu/.

## Result artifact and support decision

Publish a checked-in findings document or a structured issue comment containing:

- The complete environment/pin table and exact reproduction commands.
- A matrix for minimal IREE smoke, mlxcel f32/f16, optional bf16/packed, single-sequence, B4/B8, and server results.
- Correctness evidence, timing/VRAM tables, and links to compact raw logs or reduced reproducers.
- Every workaround or required environment variable, including GPU visibility and `gfx*` target selection.
- One conclusion: **supported**, **conditionally supported**, or **blocked**, with the boundary stated precisely.

## Acceptance criteria

- [ ] A fresh ROCm machine can reproduce the pinned compiler/runtime setup with one documented command and no reuse of artifacts from another IREE version.
- [ ] The IREE `rocm` compiler target and `hip` runtime device pass the minimal StableHLO smoke on the selected GPU.
- [ ] The normal mlxcel load, prefill, decode, batch, and server paths execute through `MLXCEL_XLA_DEVICE=hip` without device-specific graph forks.
- [ ] f32 and f16 results meet the Phase 3 correctness gates, or the smallest failing stage is reduced and documented with a follow-up tracker.
- [ ] Performance and peak-VRAM measurements are reproducible and distinguish cold compilation from warm execution.
- [ ] bf16 and packed quantization each have an explicit supported/guarded/blocked result; the packed result is cross-posted to #574.
- [ ] Documentation includes setup, target selection, supported precision/quantization modes, known limitations, and the final support decision.

This qualification issue may close with a negative or conditional result only when the failure is reproducible, reduced to the smallest failing layer, and linked to a concrete follow-up. A raw failure log without a support boundary is not sufficient.

## References

- IREE ROCm deployment: https://iree.dev/guides/deployment-configurations/gpu-rocm/
- IREE deployment target/driver matrix: https://iree.dev/guides/deployment-configurations/
- IREE source-build options: https://iree.dev/building-from-source/getting-started/
- IREE GPU debugging playbook: https://iree.dev/developers/debugging/gpu/
- mlxcel OpenXLA design: `docs/adr/0004-compute-backend-session-seam-and-stablehlo-family.md`

Contributor guide

Open the contributing guide

Research direction

Start with scripts/iree/setup-cuda.sh, the root Makefile, root build.rs, src/lib/mlxcel-xla/build.rs, src/lib/mlxcel-xla/csrc/xla_iree.c, and src/lib/mlxcel-xla/src/iree.rs. Run the listed minimal StableHLO smoke and validation commands on a recorded ROCm machine before wiring the normal path. Done means reproducible HIP execution, correctness and serving evidence, measurements, and documented supported, conditional, or blocked results.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust, shell
Domain
backend, build-system, documentation, infrastructure, performance, testing
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.