feat: extend the OpenXLA packed-quant path to bf16 scales/biases
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
The OpenXLA packed in-graph dequant path (`MLXCEL_XLA_QUANT=packed`, landed in #568) currently requires the MLX affine `scales`/`biases` to be **f16**. mlx-lm stores them as **bf16** for some checkpoints (Qwen3, Gemma3-27B, Qwen3-MoE). Issue #569 already broadened the dequant-at-LOAD path (dequantize to f32 at load) to accept bf16 scales, but the packed path (which uploads the raw scales and dequantizes in the StableHLO graph) still rejects bf16 and only accepts f16. This issue extends the packed path to accept bf16 scales/biases, so a bf16-scale MLX checkpoint can use the packed representation.
## Background
This was the residual in-tree deliverable of #574 and is deferred, low-priority breadth work.
The #573 spike (merged, PR #578) concluded that the dequant to matmul fusion on the CUDA target is upstream IREE work: there is no in-tree flag or representation that fuses it, and int8 has no tensor-core path in IREE 3.11.0rc. As a result the packed path is off by default and is currently dominated by f16-resident weights (#572) on CUDA.
This bf16-scale breadth is therefore low priority. It prepares the packed representation for more checkpoints so it is ready for when upstream IREE lands a fused CUDA quantized-matmul, and it was deferred out of #574.
## Scope
Concrete changes, all in `src/lib/mlxcel-xla`:
- **Emitter** (`emitter/model.rs`, `take_weight` packed branch): declare the `.scales` / `.biases` args as `bf16` (instead of the current hard-coded `f16`) when the checkpoint stores them as bf16. This requires threading the scale element type from the checkpoint to the emit path. The emitter currently emits purely from `config.json`, which does not carry the scale dtype, so thread it analogously to how precision is threaded via `resolve_precision`.
- **`Builder::dequant_affine`** (`emitter/builder.rs`): handle a bf16 scales/biases operand by converting bf16 to f32 in-graph before `q*scale + bias`, matching the existing f16 handling.
- **Loader** (`iree.rs`, `load_weights` QuantRaw path): accept `Dtype::BF16` scales/biases (today it errors unless `Dtype::F16`) and upload the raw bf16 bytes with a new `WDT_BF16` dtype code.
- **C shim** (`csrc/xla_iree.c`): add a `case` mapping the bf16 dtype code to `IREE_HAL_ELEMENT_TYPE_BFLOAT_16` (esize 2), alongside the existing f32/f16/u32 cases.
- Keep the emitter and loader in lock-step on the scale dtype (like the existing packed f16 path), so the uploaded buffer dtype always matches the emitted arg.
## Integration / Acceptance
- [ ] Wired into the normal load/serve path behind the same `MLXCEL_XLA_QUANT=packed` gate; a bf16-scale MLX checkpoint that satisfies `Config::supports_packed_quant` loads and runs.
- [ ] Token-exact or within-gate (#515 harness) versus the dequant-at-load path on a bf16-scale checkpoint, verified through the normal generation path on a GPU IREE target.
- [ ] The f16-scale packed path stays byte-identical (no regression).
## Dependencies / Links
- Part of #570.
- Split out of #574 (which remains open, tracking the upstream fused quantized-matmul).
- Related: #568 (packed path), #569 (dequant-at-load bf16 broadening), #573 (fusion spike).
Contributor guide
Research direction
Trace the packed path from resolve_precision and emitter/model.rs take_weight through emitter/builder.rs dequant_affine, then inspect iree.rs load_weights and csrc/xla_iree.c. Run the existing packed-path load or generation checks on an f16-scale checkpoint before changing behavior. Done means a bf16-scale checkpoint loads and runs through MLXCEL_XLA_QUANT=packed, matches the dequant-at-load path, and leaves the f16 path byte-identical.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- c, rust
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100