track(upstream-iree): fused dequant->matmul for the packed int8 bandwidth win
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #570. **Status: BLOCKED for fused-kernel implementation; a bounded latest-pin revalidation can start now.** This remains a tracking issue, not active fused-kernel implementation work. The in-tree representation is already shipped and correctness-verified; the throughput win it exists for still requires IREE target support on CUDA or `metal-spirv`. The 2026-07-22 upstream audit found no advertised data-tiling support for either target in the latest IREE candidate.
## What is already in-tree (the representation is ready)
- The packed int8 path (`MLXCEL_XLA_QUANT=packed`, #568): MLX 4/8-bit weights stay resident packed (`ui32` payload + f16 scales/biases) and are dequantized inside the StableHLO graph (bit-unpack, then `q*scale + bias`). Per-weight-dtype weight ABI landed with it.
- Dequant-at-load broadened to bf16 scales/biases (#569).
- The path is **token-exact** with the f32/f16 path on the CUDA target (the in-graph reconstruction is bit-identical to the host dequant), verified on GB10.
- Metal and precision guards so the unsupported knobs fail fast instead of faulting mid-run: `MLXCEL_XLA_QUANT=packed` on Metal (#613) and `MLXCEL_XLA_PRECISION=bf16` on Metal (#612).
## The blocker (measured)
The #573 spike ran real graphs through `iree-compile` + the runtime and concluded the fusion is upstream IREE work:
- **CUDA (GB10):** the packed decode is token-exact but about **4.3x slower** than f16 (~1.6 vs ~6.7 tok/s, Llama-3.2-1B 4-bit, greedy). IREE's CUDA codegen does **not** fuse the in-graph unpack+dequant into the matmul: the decode step is ~678 dispatches and the reconstructed f32 weight is materialized to DRAM every step, so the graph pays MORE bandwidth and compute, not less. The `aggressive-fusion` / `generalize-matmul` / `early-trunc-fusion` flags leave the dispatch count unchanged (678 -> 677).
- **Metal (M1 Ultra):** the packed prefill compiles for `metal-spirv` but the invoke faults at runtime in the metal HAL driver (`metal_device.m`: `Metal command buffer failed (status 5)` at `hal.fence.await`, surfaced as IREE `INTERNAL`). Found in #575, guarded off in #613.
So the memory-bandwidth payoff is not realized by authoring the dequant in the portable graph alone; it needs the TARGET to fuse dequant into the matmul (a quantized-matmul op, or an int8 `dot_general` lowering to the hardware int8 path). ADR 0004 already scopes per-backend fused kernel codegen as upstream IREE's responsibility.
## Upstream status refresh (2026-07-22)
The three previously linked IREE issues do not all represent active blockers:
- **`iree-org/iree#18513` — OPEN / PARTIAL / related foundation.** DispatchCreation data-tiling and grouped quantized-matmul fusion have advanced, including a regression showing dequant and contraction in one dispatch. However, IREE's documentation for `iree-3.12.0rc20260721` lists data-tiling support only for x86/AArch64/RISC-V CPU, ROCm GPU, and VMVX — not CUDA or `metal-spirv`. Keep this as background tracking, but its closure alone would not satisfy this issue.
- **`iree-org/iree#8330` — DONE / historical reference.** The lowering from `linalg.quantized_matmul` to `linalg.matmul` plus zero-point arithmetic landed in `iree-org/iree#8409` on 2022-02-27 and remains covered by IREE tests. It does not provide a fused packed-dequant GPU kernel, so it is not an active dependency.
- **`iree-org/iree#16937` — OPEN / context only.** `iree-org/iree#16927` added a shared-memory-aware cooperative-matrix schedule, but the reporter reproduced the SPIR-V failure afterward and a default Distribute-path follow-up was left unverified. More importantly, that Vulkan compile-time shared-memory failure is not the same failure class as mlxcel's Metal HAL runtime command-buffer fault. Keep it as diagnostic context, not as the Metal unblock tracker.
The direct upstream unblock remains target support that fuses dequant/unpack into the matmul (or lowers a grouped/quantized dot to a native int8 kernel) on CUDA and `metal-spirv`, so the packed graph does not materialize reconstructed f32 weights to DRAM each step.
The latest stable IREE release is still `v3.11.0`; the next validation target is the official candidate `iree-3.12.0rc20260721` at `dc9601f88654749456c7cee4ae87e13de2654e1e`. The candidate compiler/runtime must be version-matched on both CUDA and Metal.
We do not open an issue on the IREE repository (project policy: issues live on `lablup/mlxcel`). Re-check the target support and measured behavior on each MLX/IREE pin bump rather than using upstream issue closure alone.
## Next actionable step: latest-pin revalidation
This bounded work can start without treating the fused-kernel implementation as unblocked:
- [ ] **CUDA / GB10:** build the version-matched `iree-3.12.0rc20260721` compiler/runtime, rerun `scripts/xla/fusion_spike.sh`, and record packed/f16 dispatch counts, token exactness, tok/s, and whether generated code contains native int8 instructions (`dp4a` / integer MMA) instead of scalar widening.
- [ ] **Metal / M1 Ultra:** compile and invoke the minimal packed repro on the matching candidate, capture the faulting encoder/dispatch with Metal command-buffer diagnostics, and determine whether the failure is a bad kernel or a HAL command-buffer limit.
- [ ] Record both results here. If CUDA still has no fused/native-int8 path and Metal still faults, keep this issue blocked and update only the last-checked pin/date.
## Unblock criteria (what mlxcel does when upstream lands)
- Re-enable the packed path behind the existing `MLXCEL_XLA_QUANT=packed` gate; relax the Metal guard (#613) once `metal-spirv` executes the fused kernel.
- On GB10, int8 packed `tok/s > f16 tok/s` (reverses the #568 regression), token-exact or within-gate per the #515 harness, with the decode dispatch count materially reduced from ~678.
- Fold in the packed path accepting bf16 scales/biases (today it requires f16; #569 broadened only the dequant-at-load path), verified on a bf16-scale MLX checkpoint through the normal load path.
## Payoff caveat
Even once fused, the bandwidth win only matters on an int-native / NPU target. On the compute-bound Metal decode it is un-demonstrable regardless (the #575 f16 result shows the Metal decode tracks compute, not bandwidth; ADR 0004). So this unblocks the CUDA/NPU path first; Metal gets correctness, not throughput.
## References
#570 (epic), #573 (spike: in-tree vs upstream split), #568 / #569 (packed representation + bf16 dequant-at-load), #612 / #613 (Metal fail-fast guards), #575 / #611 (Metal transfer validation), ADR `docs/adr/0004-compute-backend-session-seam-and-stablehlo-family.md`.
---
**Refresh log**
- 2026-07-22 — Audited the linked IREE issues and latest official release state.
- Reclassified `iree-org/iree#18513` as partial/background because CUDA and `metal-spirv` are still outside the documented data-tiling target set.
- Reclassified completed `iree-org/iree#8330` as historical context, not an active dependency.
- Reclassified `iree-org/iree#16937` as diagnostic context, not a direct tracker for the Metal runtime fault.
- Added a concrete `iree-3.12.0rc20260721` CUDA/Metal revalidation checklist; the fused-kernel implementation remains blocked.
Contributor guide
Assessment
This issue has not been assessed yet.