fix(xla): fail fast on MLXCEL_XLA_QUANT=packed for Metal targets#617
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes #613 (follow-up to #575). The packed int8 path (
MLXCEL_XLA_QUANT=packed, #568: residentui32weight + f16 scales, dequantized in the StableHLO graph) compiles formetal-spirvbut its prefill invoke faults at runtime.Root cause
With the improved
XLA_CHECKdiagnostics (below), the fault is:The graph lowers (no
iree-compileerror) and runs token-exact on the GB10 CUDA runtime, so this is IREE metal HAL /metal-spirvruntime 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 ametaldevice at the emit/compile boundary (beforeiree-compileand the weight upload), with an actionable message pointing at the CUDA / CPU targets. Wired into both compile paths (compile_vmfbssingle-seq and the batchXlaRaggedEngineload). Backed by a purepacked_runs_on(device, packed_active)predicate with a unit test.XLA_CHECKnow 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 themetal_device.mcommand-buffer error.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).f16on 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).packed_runs_on_rejects_only_active_packed_on_metalandprecision_lowers_on_rejects_only_bf16_on_metalpass;cargo fmt --all -- --checkclean.