Skip to content

fix(xla): fail fast on MLXCEL_XLA_QUANT=packed for Metal targets#617

Merged
inureyes merged 1 commit into
mainfrom
fix/613-packed-metal-guard
Jul 2, 2026
Merged

fix(xla): fail fast on MLXCEL_XLA_QUANT=packed for Metal targets#617
inureyes merged 1 commit into
mainfrom
fix/613-packed-metal-guard

Conversation

@inureyes

@inureyes inureyes commented Jul 2, 2026

Copy link
Copy Markdown
Member

Summary

Closes #613 (follow-up to #575). The packed int8 path (MLXCEL_XLA_QUANT=packed, #568: resident ui32 weight + f16 scales, dequantized in the StableHLO graph) compiles for metal-spirv but its prefill invoke faults at runtime.

Root cause

With the improved XLA_CHECK diagnostics (below), the fault is:

metal_device.m:650: INTERNAL; Metal command buffer failed (status 5);
while invoking native function hal.fence.await;
[ 0] bytecode prefill@1:84136

The graph lowers (no iree-compile error) and runs token-exact on the GB10 CUDA runtime, so this is IREE metal HAL / metal-spirv runtime behavior on the large ui32-unpack + dequant + matmul command buffer, not an mlxcel emitter defect. A fused quantized-matmul is upstream IREE work (#574), and the packed/int8 bandwidth win is un-demonstrable on the compute-bound Metal decode regardless (per ADR 0004 and the #575 f16 result). So the correct outcome here is robustness on Metal, not a throughput fix.

Changes

  • check_packed_supported(device, packed_active) rejects an active packed build on a metal device at the emit/compile boundary (before iree-compile and the weight upload), with an actionable message pointing at the CUDA / CPU targets. Wired into both compile paths (compile_vmfbs single-seq and the batch XlaRaggedEngine load). Backed by a pure packed_runs_on(device, packed_active) predicate with a unit test.
  • XLA_CHECK now prints the full annotated IREE status (iree_status_fprint), so any runtime fault names its cause (op / dispatch / driver message) instead of an opaque status code. This is what surfaced the metal_device.m command-buffer error.
  • Docs: ADR 0004 and the mlxcel-xla README updated with the root cause and the guard.

Validation (M1 Ultra)

  • MLXCEL_XLA_QUANT=packed + MLXCEL_XLA_DEVICE=metal: Error: MLXCEL_XLA_QUANT=packed is not supported on the Metal target: ... (clean, no command-buffer fault).
  • f16 on Metal (no packed): still generates (metal default, token-exact per feat: extend the fused low-precision path to the Metal target #575).
  • MLXCEL_XLA_QUANT=packed + MLXCEL_XLA_DEVICE=local-task (CPU): still runs (guard is metal-only, packed unaffected off-Metal).
  • Unit tests packed_runs_on_rejects_only_active_packed_on_metal and precision_lowers_on_rejects_only_bf16_on_metal pass; cargo fmt --all -- --check clean.

The packed int8 path (in-graph ui32-unpack + dequant, #568) compiles for metal-spirv but its prefill invoke faults at runtime: the metal HAL driver reports Metal command buffer failed (status 5) at hal.fence.await (metal_device.m), surfaced as IREE INTERNAL. The same graph runs token-exact on the GB10 CUDA runtime, so this is IREE metal HAL / metal-spirv runtime behavior on the large dequant-then-matmul command buffer, not an emitter defect; a fused quantized-matmul is upstream IREE work (#574) and the bandwidth win is un-demonstrable on the compute-bound Metal decode anyway.

Add check_packed_supported, which rejects an active packed build on a metal device at the emit/compile boundary (before iree-compile and the weight upload) with an actionable message pointing at the CUDA / CPU targets, instead of the opaque command-buffer fault. Wired into both single-seq (compile_vmfbs) and batch load paths; CUDA and CPU (local-task) packed runs are unaffected, and a non-packed Metal build never trips it. The guard logic is a pure packed_runs_on(device, packed_active) predicate with a unit test.

Also print the full annotated IREE status in XLA_CHECK (iree_status_fprint) so a runtime fault names its cause instead of an opaque status code; this is what surfaced the metal_device.m command-buffer error above.

Closes #613.
@inureyes inureyes added area:inference Generation, sampling, decoding (incl. speculative, DRY) type:bug Bug fixes, error corrections, or issue resolutions priority:low Low priority labels Jul 2, 2026
@inureyes
inureyes merged commit c800266 into main Jul 2, 2026
5 checks passed
@inureyes
inureyes deleted the fix/613-packed-metal-guard branch July 2, 2026 05:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) priority:low Low priority type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug(xla): packed int8 (MLXCEL_XLA_QUANT=packed) prefill invoke faults on metal-spirv

1 participant