콘텐츠로 이동

OpenXLA 백엔드 소스 상태

mlxcel은 Apple Silicon과 CUDA에서 MLX를 프로덕션 compute backend로, ROCm에서는 실험적 backend로 사용합니다. 그 옆에 src/lib/mlxcel-xla/ 아래의 StableHLO/IREE 구현이 함께 있고, cfg(feature = "xla-backend")와 cfg(feature = "xla-iree")로 보호된 backend seam을 통해 닿습니다. 루트 Cargo.toml은 이 feature들과 mlxcel-xla 의존성을 선언하며, 전부 기본값이 off입니다.

빌드 구성

  • 기본 cargo build, Metal build, CUDA build는 OpenXLA 경로를 아예 컴파일하지 않으며 MLX backend만 들어갑니다.
  • 기본 바이너리에서 MLXCEL_BACKEND=xla를 설정해도 아무 일도 일어나지 않습니다. 선택자 자체가 feature gate 뒤에 있기 때문입니다.
  • cargo build --features xla-backend는 seam과 StableHLO 그래프 emitter를 컴파일합니다. 네이티브 툴체인이 필요 없어 기본 바이너리를 빌드할 수 있는 곳이면 어디서나 빌드되지만, 실행 엔진은 없습니다.
  • cargo build --features xla-iree는 실행 런타임까지 더합니다. build script가 미리 빌드된 IREE 배포본에 대해 C shim을 컴파일하므로 빌드 시점에 IREE_DIST, IREE_CUDA_HOME(+IREE_CUDA_COMPILE), IREE_MACOS_HOME(+IREE_MACOS_COMPILE), IREE_ROCM_HOME(+IREE_ROCM_COMPILE) 중 하나가 필요합니다. 뒤의 세 소스 빌드 트리는 scripts/iree/setup-cuda.sh, scripts/iree/setup-macos.sh, scripts/iree/setup-rocm.sh(make iree-rocm도 가능)가 만듭니다.
  • xla-diagnostics, xla-micro-oracle, xla-diagnostics-cpu는 xla-iree 위에 계측을 얹습니다. xla-diagnostics는 CPU 진단을 CUDA 빌드와 묶고, xla-diagnostics-rocm은 AMD GPU에서 중간 tensor oracle 검사를 하도록 ROCm 빌드와 묶습니다(rocm과 xla-diagnostics-cpu).
  • ROCm에서는 MLXCEL_XLA_DEVICE=hip과 MLXCEL_XLA_ROCM_TARGET의 gfx target 하나로 IREE HIP target을 실행합니다. gfx1151에서 검증을 마쳤으며, 기본 f16 precision에서 Llama-3.2-1B-Instruct-4bit와 Qwen3-0.6B-4bit의 64토큰 greedy 디코드가 CPU f32 oracle과 토큰 단위로 일치합니다(docs/benchmark_results/xla_iree_rocm_qualification.md).

XLA 소스 디렉터리와 README는 개발 자료로 남습니다. 이 feature들을 켜는 빌드는 로컬 opt-in 빌드이지 CI나 릴리스 구성이 아닙니다.

모델 범위

OpenXLA 경로는 StableHLO emitter로 정해진 텍스트 아키텍처 집합을 실행하고, CLI와 서버가 공유하는 prepared-prefill 경로에서 일부 VLM family의 vision 경로도 실행합니다.

  • LLaVA: vision tower와 projector를 상주 IREE 모듈에서 실행하고, MLX host 경로를 fallback으로 둡니다(MLXCEL_XLA_VISION_BACKEND).
  • Molmo 2: xla-iree 빌드에서는 vision 경로를 상주 IREE 모듈에서 실행합니다. MLXCEL_XLA_VISION_BACKEND=host이거나 xla-iree 없이 빌드하면 세션에 이미지 경로가 없습니다.
  • Qwen3-VL: vision tower, 메인 merger, 모든 DeepStack merger를 상주 IREE 모듈에서 실행합니다(qwen3_vl_runtime.rs).
  • Youtu-VL: vision tower와 내장 merger를 상주 IREE 모듈에서 실행합니다(youtu_vl_runtime.rs).
  • Gemma 3 VLM: embeddings-prefill 진입점이 transformers 이미지 마스크(causal attention에 이미지 블록 내부 양방향 attention을 더한 마스크)를 additive prefill mask로 받습니다(prepared_gemma3.rs).
  • Gemma 3n 오디오: 오디오 인코더와 언어 병합을 두 개의 상주 IREE 모듈로 실행하며 mel 추출만 호스트에 남습니다(gemma3n_audio_runtime.rs).
  • Qwen2-VL과 Qwen2.5-VL: vision tower를 상주 IREE 모듈에서만 실행하므로 xla-iree 없이 빌드하면 이 checkpoint는 시작에 실패하고, MLXCEL_XLA_VISION_BACKEND=host는 거부됩니다. Qwen2.5-VL graph는 window 단위 vision attention(window permutation, window별로 분리한 mask, checkpoint가 지정한 full-attention layer)을 구현하며, 지원하지 않는 temporal grid는 거부합니다.

Muse Glimmer VLM은 OpenXLA 이미지 실행을 거부하고 MLX backend를 쓰라는 메시지를 냅니다. 그 밖의 VLM family에는 OpenXLA 이미지 경로가 없습니다.

남아 있는 설계

구현은 session-level compute-backend seam을 사용합니다. 개별 tensor operation을 감싸는 대신 token input, backend-owned KV state, decode, sampling을 하나의 backend session이 소유합니다. StableHLO/MLIR을 emit하고 CPU, CUDA, Metal, HIP(ROCm) target용 IREE runtime 코드를 보유합니다. session contract 설계는 저장소의 ADR 0004(docs/adr/0004-compute-backend-session-seam-and-stablehlo-family.md)에 기록되어 있습니다.

환경 변수

MLXCEL_BACKEND, MLXCEL_XLA_DEVICE, MLXCEL_XLA_PRECISION, MLXCEL_XLA_QUANT, MLXCEL_XLA_ROCM_TARGET, MLXCEL_XLA_VISION_BACKEND, MLXCEL_XLA_CONTEXT_CAPACITY, MLXCEL_XLA_IREE_COMPILE은 feature gate 뒤의 코드에서만 소비됩니다. 따라서 기본 바이너리에서는 동작하지 않고, xla-backend나 xla-iree로 빌드한 바이너리에서 살아납니다. 환경 변수를 참조하십시오.

백엔드 개발자용 안내

src/lib/mlxcel-xla/README.md와 그 ADR은 구현 메모로 취급하십시오. OpenXLA를 제품 기능으로 알리기 전에 seam-only 구성과 IREE 구성을 모두 빌드하고 각 지원 target에서 실제 체크포인트 token parity를 다시 검증하십시오. operator 문서와 배포 자동화는 MLX만 선택합니다.