feat: extend the OpenXLA packed-quant path to bf16 scales/biases #603

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 (test: precision accuracy gate and per-device default precision for the OpenXLA backend (#449) #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 epic: realize low-precision / quantized performance on the OpenXLA (IREE) backend #570.
• Split out of track(upstream-iree): fused dequant->matmul for the packed int8 bandwidth win #574 (which remains open, tracking the upstream fused quantized-matmul).
• Related: feat: int8/int4 packed weight quantization on the OpenXLA path (#516) #568 (packed path), fix(xla): load bf16 affine quant scales in the OpenXLA weight loader #569 (dequant-at-load bf16 broadening), spike: quantized-matmul / dequant-to-matmul fusion investigation on IREE (CUDA first) #573 (fusion spike).