From 2b538467a9aac2600c9f2e488d1e0bef527ee2c0 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 10 Sep 2026 11:14:21 +0900 Subject: [PATCH 1/3] feat(qwen3_5): load vendor fine-grained FP8 checkpoints The vLLM-style FP8 releases of the Qwen3.5 family store every converted projection as raw E4M3 bytes in `.weight` plus a bf16 inverse scale per 128x128 block in `.weight_scale_inv`. MLX's safetensors reader maps `F8_E4M3` to `uint8`, so those bytes already arrive intact and nothing errors: the checkpoint loads, generates, and is wrong by a per-block factor on every projection. That silence is the reason this is a correctness fix rather than a capability gap. `models::fp8_block` adds the detection and the load-time conversion. `qwen_fp8_block_quantization` reads `quantization_config` from the top level or `text_config` and accepts `quant_method: "fp8"` with `fmt` absent or `"e4m3"` and `weight_block_size: [128, 128]`; any other geometry, a non-E4M3 format, and a missing `weight_block_size` (per-tensor FP8, a different sidecar layout) are named errors rather than an assumed 128. `requantize_block_fp8_weights` decodes the bytes through a 256-entry device LUT, pads to whole blocks, applies the block scale, slices back, and requantizes to MLX-native mxfp8 (group size 32, 8 bits), which the existing quantized-matmul path executes with no per-token dequantization. Conversion is one tensor at a time and each packed pair is evaluated before the next tensor starts, so the f32 reconstruction of the largest projection is the only transient. Tensors with no sidecar (norms, embeddings, `conv1d`, `A_log`, `dt_bias`, and the whole vision tower on the released checkpoints) stay dense, and `UnifiedLinear` already picks per tensor from the presence of `.scales`. The pre-pass is wired into all three Qwen3.5 entry points: the text loader, the VLM loader (before the weight map is split, so sidecars are still paired), and the adapter-side special loader. Each also merges the effective `quantization` block into the text config, which the FP8 releases do not carry, so the model constructor builds quantized layers. Validated with `cargo test --profile test-fast --features cuda --lib models::fp8_block` (10/10) under `MLX_ENABLE_TF32=0`. The central test compares the device path bitwise, on both planes, against `quantize_weights_with_mode(decode(bytes) * expanded_scales, 32, 8, "mxfp8")` built independently on the host, over a 130x160 tensor that is block-aligned on neither axis. Refs #1361 --- src/loading/special.rs | 17 + src/loading/vlm_qwen.rs | 19 ++ src/models/fp8_block.rs | 433 +++++++++++++++++++++++++ src/models/fp8_block_tests.rs | 580 ++++++++++++++++++++++++++++++++++ src/models/mod.rs | 9 + src/models/qwen3_5.rs | 21 ++ src/models/sanitize.rs | 2 +- 7 files changed, 1080 insertions(+), 1 deletion(-) create mode 100644 src/models/fp8_block.rs create mode 100644 src/models/fp8_block_tests.rs diff --git a/src/loading/special.rs b/src/loading/special.rs index c5cc51895..9ddb92a4c 100644 --- a/src/loading/special.rs +++ b/src/loading/special.rs @@ -55,6 +55,16 @@ pub(crate) fn qwen35_text_config(config: &serde_json::Value) -> Result models::requantize_block_fp8_weights(owned, detected.block_rows) + .map_err(|error| anyhow::anyhow!("{error}"))?, + None => owned, + }; let owned = models::qwen3_5::sanitize_moe_weights(owned, &args); let model = models::Qwen35Model::from_weights(&owned, &args) .map_err(|err| anyhow::anyhow!("{}", err))?; diff --git a/src/loading/vlm_qwen.rs b/src/loading/vlm_qwen.rs index 284d53648..60bbd02ae 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -415,7 +415,21 @@ fn load_qwen3_5_vlm_with_variant( parse_required_vlm_subconfig(&full_config, "vision_config", "Qwen3.5 vision config")?; inherit_qwen_vision_quantization(&mut vision_config, &full_config); + // Vendor fine-grained FP8 releases (`quant_method: "fp8"`, one bf16 + // inverse scale per 128x128 block) declare no `quantization` block and + // ship raw E4M3 bytes. Reconstruct and requantize them to MLX-native + // mxfp8 before the weight map is split, so the sidecars are still paired + // with their tensors and the vision tower (which the released + // checkpoints leave unconverted) simply has none to consume. + let fp8_block = models::qwen_fp8_block_quantization(&full_config) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let raw_weights = load_vlm_weights_common(model_path, None)?; + let raw_weights = match fp8_block { + Some(detected) => models::requantize_block_fp8_weights(raw_weights, detected.block_rows) + .map_err(|error| anyhow::anyhow!("{error}"))?, + None => raw_weights, + }; let mut text_weights = mlxcel_core::weights::WeightMap::new(); let mut vision_weights = mlxcel_core::weights::WeightMap::new(); @@ -450,6 +464,11 @@ fn load_qwen3_5_vlm_with_variant( ); } + if let Some(detected) = fp8_block { + models::merge_fp8_block_quantization(&mut text_config_val, detected) + .map_err(|error| anyhow::anyhow!("{error}"))?; + } + let text_config: models::qwen3_5::Qwen35Config = serde_json::from_value(text_config_val) .map_err(|e| anyhow::anyhow!("Failed to parse Qwen3.5 text config: {}", e))?; text_config.validate_supported()?; diff --git a/src/models/fp8_block.rs b/src/models/fp8_block.rs new file mode 100644 index 000000000..66d34a984 --- /dev/null +++ b/src/models/fp8_block.rs @@ -0,0 +1,433 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Vendor fine-grained FP8 checkpoints: block-scale detection and load-time +//! requantization to MLX's native `mxfp8`. +//! +//! The vLLM-style FP8 releases of the Qwen3.5 family (`Qwen/Qwen3.8-27B-FP8` +//! and siblings) store every converted projection as two tensors: +//! +//! * `.weight` — raw `E4M3` bytes, one per weight element, +//! * `.weight_scale_inv` — a `bfloat16` multiplier per 128x128 block. +//! +//! The true weight is `e4m3_decode(byte) * scale_inv[block]`. MLX has no +//! float8 dtype and its safetensors reader maps `F8_E4M3` to `uint8`, so the +//! bytes arrive intact but the scale is never applied: without this pass the +//! checkpoint loads and generates, and every projection is wrong by its own +//! per-block factor. +//! +//! [`requantize_block_fp8_weights`] reconstructs each pair on device and +//! requantizes it to `mxfp8` (`group_size` 32, `bits` 8), which is the closest +//! native layout MLX executes: 8 bits per weight plus one shared E8M0 exponent +//! per 32 values, with no per-token dequantization in the forward pass. The +//! reconstruction is transient and per tensor, so peak memory stays close to +//! the checkpoint's own footprint rather than its dequantized size. +//! +//! This module deliberately lives outside `sanitize.rs`: that file is already +//! 3k lines, and the FP8 path is a self-contained pre-pass with its own tests. +//! +//! Used by: the Qwen3.5 text loader (`Qwen35Model::load`), the Qwen3.5 VLM +//! loader (`load_qwen3_5_vlm_with_variant`), and the adapter-side special +//! loader (`try_load_special_model_from_weights`). + +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr, dtype}; +use serde_json::{Value, json}; + +use super::sanitize::f8_e4m3_to_f32; + +/// Suffix identifying a vendor FP8 block-scale sidecar. +pub const FP8_BLOCK_SCALE_SUFFIX: &str = ".weight_scale_inv"; + +/// MLX fixes `mxfp8` at group size 32 and 8 bits (`fp_quantize` in +/// `mlx/ops.cpp` throws for anything else), so these are constants rather than +/// tunables. +pub const MXFP8_GROUP_SIZE: i32 = 32; +/// See [`MXFP8_GROUP_SIZE`]. +pub const MXFP8_BITS: i32 = 8; +/// See [`MXFP8_GROUP_SIZE`]. +pub const MXFP8_MODE: &str = "mxfp8"; + +/// The only block geometry this path accepts. +/// +/// Every public Qwen3.5-family FP8 release declares `[128, 128]`. A different +/// geometry is refused rather than assumed, because the block extent is what +/// pairs a scale with its weight elements and a wrong pairing silently +/// mis-scales the whole tensor. +pub const SUPPORTED_FP8_BLOCK: usize = 128; + +/// A detected fine-grained FP8 layout and the MLX quantization it maps onto. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Fp8BlockQuantization { + /// Block extent along the output (row) axis. + pub block_rows: usize, + /// Block extent along the input (column) axis. + pub block_cols: usize, + /// Group size of the MLX quantization the weights are converted to. + pub group_size: i32, + /// Bit width of the MLX quantization the weights are converted to. + pub bits: i32, + /// MLX quantization mode name. + pub mode: &'static str, +} + +impl Fp8BlockQuantization { + /// The `quantization` block to merge into a model config so the model + /// constructor builds quantized layers for the converted tensors. + /// + /// `mode` is included for readers that thread an explicit mode + /// (`QuantizationArgs`); the Qwen3.5 config carries only `group_size` and + /// `bits` and lets `infer_quantization_mode` derive `mxfp8` from + /// `bits == 8` plus the absent `.biases` plane. + #[must_use] + pub fn effective_config(&self) -> Value { + json!({ + "group_size": self.group_size, + "bits": self.bits, + "mode": self.mode, + }) + } +} + +/// Read the `quantization_config` block and report whether this checkpoint is +/// a fine-grained (block-scaled) FP8 release this loader can reconstruct. +/// +/// Accepts the block at the top level or inside `text_config`, since the VLM +/// wrappers place it in either. Returns: +/// +/// * `Ok(None)` when the checkpoint declares no FP8 quantization at all, which +/// covers every dense and affine-quantized checkpoint, +/// * `Ok(Some(..))` for `quant_method: "fp8"` with `fmt` absent or `"e4m3"` +/// and `weight_block_size: [128, 128]`, +/// * `Err` for an FP8 declaration whose geometry or float format is something +/// other than that. Failing the load is the point: the alternative is +/// applying a 128x128 pairing to scales that do not have it, which produces +/// fluent but wrong output rather than an error. +pub fn qwen_fp8_block_quantization( + full_config: &Value, +) -> Result, String> { + let Some(quant_config) = full_config.get("quantization_config").or_else(|| { + full_config + .get("text_config") + .and_then(|text| text.get("quantization_config")) + }) else { + return Ok(None); + }; + + let Some(quant_config) = quant_config.as_object() else { + return Ok(None); + }; + + let method = quant_config + .get("quant_method") + .and_then(Value::as_str) + .unwrap_or_default(); + if !method.eq_ignore_ascii_case("fp8") { + return Ok(None); + } + + match quant_config.get("fmt").and_then(Value::as_str) { + None => {} + Some(fmt) if fmt.eq_ignore_ascii_case("e4m3") => {} + Some(fmt) => { + return Err(format!( + "quantization_config declares quant_method \"fp8\" with fmt \"{fmt}\": mlxcel \ + decodes only the E4M3 float8 format, and decoding E4M3 bytes as another format \ + (or the reverse) produces a wrong exponent for every weight without failing" + )); + } + } + + let Some(block) = quant_config.get("weight_block_size") else { + return Err( + "quantization_config declares quant_method \"fp8\" but carries no weight_block_size: \ + mlxcel reconstructs only the fine-grained block-scaled layout (a bf16 scale per \ + 128x128 block in a *.weight_scale_inv sidecar). Per-tensor and per-channel FP8 \ + scales use a different sidecar layout and are not supported" + .to_string(), + ); + }; + + let dims = block + .as_array() + .ok_or_else(|| { + format!("quantization_config.weight_block_size must be an array, got {block}") + })? + .iter() + .map(|dim| { + dim.as_u64().ok_or_else(|| { + format!( + "quantization_config.weight_block_size entry {dim} is not a positive integer" + ) + }) + }) + .collect::, _>>()?; + + let block_extent = u64::try_from(SUPPORTED_FP8_BLOCK).unwrap_or(u64::MAX); + if dims.as_slice() != [block_extent, block_extent] { + return Err(format!( + "quantization_config.weight_block_size is {block}, but mlxcel reconstructs only \ + [{SUPPORTED_FP8_BLOCK}, {SUPPORTED_FP8_BLOCK}] fine-grained FP8 blocks. The block \ + extent is what pairs each scale with its weight elements, so applying the wrong one \ + mis-scales every tensor without failing" + )); + } + + Ok(Some(Fp8BlockQuantization { + block_rows: SUPPORTED_FP8_BLOCK, + block_cols: SUPPORTED_FP8_BLOCK, + group_size: MXFP8_GROUP_SIZE, + bits: MXFP8_BITS, + mode: MXFP8_MODE, + })) +} + +/// Merge the effective `quantization` block into a model config when the +/// checkpoint is a fine-grained FP8 release. +/// +/// Mirrors how `qwen35_text_config` merges a top-level `quantization` into +/// `text_config`: an explicit `quantization` already present always wins, so a +/// pre-converted checkpoint is untouched. Returns whether a block was +/// inserted. +pub fn merge_fp8_block_quantization( + config: &mut Value, + detected: Fp8BlockQuantization, +) -> Result { + if config.get("quantization").is_some() { + return Ok(false); + } + let object = config.as_object_mut().ok_or_else(|| { + "Cannot merge the FP8 quantization block into a non-object config".to_string() + })?; + object.insert("quantization".to_string(), detected.effective_config()); + Ok(true) +} + +/// Whether `weights` carries at least one FP8 block-scale sidecar. +/// +/// Cheap enough to gate the pre-pass on: a detected `quantization_config` says +/// what the publisher intended, this says what the shards actually hold. +#[must_use] +pub fn has_block_fp8_weights(weights: &WeightMap) -> bool { + weights + .keys() + .any(|key| key.ends_with(FP8_BLOCK_SCALE_SUFFIX)) +} + +/// A 256-entry f32 lookup table mapping an E4M3 byte to its value. +/// +/// Built once per call and consumed by `take` so the decode runs on device: +/// the alternative, a host loop over every byte, would walk 28 GB through the +/// CPU and materialize an f32 copy four times that size. +fn e4m3_lookup_table() -> UniquePtr { + let table: Vec = (0..=u8::MAX).map(f8_e4m3_to_f32).collect(); + mlxcel_core::from_slice_f32(&table, &[256]) +} + +/// Reconstruct one block-scaled FP8 tensor into a plain f32 array. +/// +/// `raw` holds the E4M3 bytes as `uint8`, `scale_inv` the per-block bf16 +/// multipliers. The result is `[rows, cols]` f32, matching the original weight +/// shape. Exposed so tests can compare this route against an independently +/// expressed reference (see `fp8_block_tests.rs`). +pub fn reconstruct_block_fp8_tensor( + name: &str, + raw: &MlxArray, + scale_inv: &MlxArray, + block: usize, + lut: &MlxArray, +) -> Result, String> { + let weight_shape = mlxcel_core::array_shape(raw); + if weight_shape.len() != 2 { + return Err(format!( + "FP8 block-scaled tensor {name} has rank {} ({weight_shape:?}); only rank-2 \ + projections carry a weight_scale_inv sidecar", + weight_shape.len() + )); + } + let weight_dtype = mlxcel_core::array_dtype(raw); + if weight_dtype != dtype::UINT8 { + return Err(format!( + "FP8 block-scaled tensor {name} loaded as MLX dtype {weight_dtype} rather than uint8 \ + ({}): the raw E4M3 bytes are needed to pair each element with its block scale, and a \ + pre-decoded float tensor has already lost that pairing", + dtype::UINT8 + )); + } + + let (rows, cols) = (weight_shape[0], weight_shape[1]); + if rows <= 0 || cols <= 0 { + return Err(format!( + "FP8 block-scaled tensor {name} has a degenerate shape {weight_shape:?}" + )); + } + if cols % MXFP8_GROUP_SIZE != 0 { + return Err(format!( + "FP8 block-scaled tensor {name} has {cols} input features, which is not a multiple of \ + the mxfp8 group size {MXFP8_GROUP_SIZE}. MLX packs one E8M0 scale per {MXFP8_GROUP_SIZE} \ + values along the last axis and cannot represent a partial group" + )); + } + + let block_i32 = i32::try_from(block) + .map_err(|_| format!("FP8 block extent {block} exceeds i32 for tensor {name}"))?; + // `i32::div_ceil` is still unstable on the pinned toolchain; both operands + // are positive here, so the classic form is exact. + let row_blocks = (rows + block_i32 - 1) / block_i32; + let col_blocks = (cols + block_i32 - 1) / block_i32; + + let scale_shape = mlxcel_core::array_shape(scale_inv); + if scale_shape != vec![row_blocks, col_blocks] { + return Err(format!( + "FP8 block-scale sidecar for {name} has shape {scale_shape:?}, expected \ + [{row_blocks}, {col_blocks}] for a {rows}x{cols} weight at block extent {block}" + )); + } + + // Decode the E4M3 bytes through the LUT on device. + let indices = mlxcel_core::astype(raw, dtype::UINT32); + let decoded = mlxcel_core::take(lut, &indices, 0); + + // Pad to whole blocks so the reshape is exact. The padding is zero and is + // sliced back off after scaling, so it can only affect the trailing block + // of a tensor whose extent is not a multiple of `block`. + let padded_rows = row_blocks * block_i32; + let padded_cols = col_blocks * block_i32; + let padded = if padded_rows == rows && padded_cols == cols { + decoded + } else { + mlxcel_core::pad( + &decoded, + &[0, padded_rows - rows, 0, padded_cols - cols], + 0.0, + ) + }; + + let blocked = mlxcel_core::reshape(&padded, &[row_blocks, block_i32, col_blocks, block_i32]); + let scales = mlxcel_core::astype(scale_inv, dtype::FLOAT32); + let scales = mlxcel_core::reshape(&scales, &[row_blocks, 1, col_blocks, 1]); + let scaled = mlxcel_core::multiply(&blocked, &scales); + let flat = mlxcel_core::reshape(&scaled, &[padded_rows, padded_cols]); + + Ok(if padded_rows == rows && padded_cols == cols { + flat + } else { + mlxcel_core::slice(&flat, &[0, 0], &[rows, cols]) + }) +} + +/// Reconstruct every block-scaled FP8 tensor in `weights` and replace it with +/// an MLX-native `mxfp8` pair. +/// +/// For each `.weight` / `.weight_scale_inv` pair the sidecar is +/// removed, `.weight` is replaced by the packed `uint32` plane, and +/// `.scales` is added. Tensors without a sidecar (norms, embeddings, +/// `conv1d`, `A_log`, `dt_bias`, and the whole vision tower on the released +/// checkpoints) are left untouched, and `UnifiedLinear` picks the quantized or +/// the dense path per tensor from the presence of `.scales`. +/// +/// Conversion is one tensor at a time and each quantized pair is evaluated +/// before the next tensor starts, so the f32 reconstruction of the largest +/// projection is the only transient and is released before the next one is +/// built. Doing the whole map lazily would instead hold every reconstruction +/// alive until the first `eval`, which for a 27B checkpoint is four times the +/// checkpoint's own size. +pub fn requantize_block_fp8_weights( + mut weights: WeightMap, + block: usize, +) -> Result { + if block == 0 { + return Err("FP8 block extent must be positive".to_string()); + } + + let mut scale_keys: Vec = weights + .keys() + .filter(|key| key.ends_with(FP8_BLOCK_SCALE_SUFFIX)) + .cloned() + .collect(); + if scale_keys.is_empty() { + return Ok(weights); + } + // Deterministic order so a failure reports the same tensor on every run. + scale_keys.sort(); + + let started = std::time::Instant::now(); + let total = scale_keys.len(); + tracing::info!( + tensors = total, + block, + "reconstructing fine-grained FP8 weights and requantizing to mxfp8" + ); + + let lut = e4m3_lookup_table(); + + for scale_key in scale_keys { + let prefix = scale_key + .strip_suffix(FP8_BLOCK_SCALE_SUFFIX) + .ok_or_else(|| format!("Internal error: {scale_key} lost its FP8 sidecar suffix"))? + .to_string(); + let weight_key = format!("{prefix}.weight"); + + let scale_inv = weights + .remove(&scale_key) + .ok_or_else(|| format!("Internal error: FP8 sidecar {scale_key} disappeared"))?; + let raw = weights.remove(&weight_key).ok_or_else(|| { + format!( + "FP8 block-scale sidecar {scale_key} has no matching {weight_key} in the \ + checkpoint; the scale cannot be applied to anything" + ) + })?; + + let restored = reconstruct_block_fp8_tensor(&weight_key, &raw, &scale_inv, block, &lut)?; + // Drop the E4M3 bytes and the sidecar before quantizing so the peak + // holds the reconstruction and the packed result, not all four. + drop(raw); + drop(scale_inv); + + let quantized = mlxcel_core::quantize_weights_with_mode( + &restored, + MXFP8_GROUP_SIZE, + MXFP8_BITS, + MXFP8_MODE, + ); + if mlxcel_core::quantized_weights_has_biases(&quantized) { + return Err(format!( + "MLX returned affine biases for a {MXFP8_MODE} quantization of {weight_key}; \ + the block-float path must not carry a zero-point plane" + )); + } + let packed = mlxcel_core::quantized_weights_w(&quantized); + let scales = mlxcel_core::quantized_weights_scales(&quantized); + + // Force realization here rather than at first forward: this is what + // keeps the reconstruction transient and the peak bounded. + mlxcel_core::eval(&packed); + mlxcel_core::eval(&scales); + + weights.insert(weight_key, packed); + weights.insert(format!("{prefix}.scales"), scales); + } + + tracing::info!( + tensors = total, + elapsed_ms = started.elapsed().as_millis() as u64, + "fine-grained FP8 requantization complete" + ); + + Ok(weights) +} + +#[cfg(test)] +#[path = "fp8_block_tests.rs"] +mod fp8_block_tests; diff --git a/src/models/fp8_block_tests.rs b/src/models/fp8_block_tests.rs new file mode 100644 index 000000000..2fb34c625 --- /dev/null +++ b/src/models/fp8_block_tests.rs @@ -0,0 +1,580 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{ + Fp8BlockQuantization, MXFP8_BITS, MXFP8_GROUP_SIZE, MXFP8_MODE, SUPPORTED_FP8_BLOCK, + has_block_fp8_weights, merge_fp8_block_quantization, qwen_fp8_block_quantization, + requantize_block_fp8_weights, +}; +use crate::models::sanitize::{f8_e4m3_to_f32, f32_to_f8_e4m3}; +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr, dtype}; +use serde_json::json; + +/// Deterministic xorshift so a failing case is reproducible from the test name +/// alone. `rand` is not a dependency of this crate's model layer. +struct Rng(u64); + +impl Rng { + fn new(seed: u64) -> Self { + Self(seed | 1) + } + + fn next_u32(&mut self) -> u32 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + (x >> 32) as u32 + } + + fn below(&mut self, bound: u32) -> u32 { + self.next_u32() % bound + } +} + +/// A bf16 value and its exact f32 widening. +/// +/// bf16 is the top 16 bits of an f32, so the widening is exact and the host +/// reference below multiplies by the very same number the device path reads. +fn bf16_value(rng: &mut Rng) -> (u16, f32) { + // f32 exponent field 120..=130 spans 2^-7 .. 2^3, which brackets the + // magnitudes real FP8 block scales take (block_amax / 448). + let exponent = 120 + rng.below(11) as u16; + let mantissa = rng.below(128) as u16; + let bits = (exponent << 7) | mantissa; + (bits, f32::from_bits(u32::from(bits) << 16)) +} + +/// An E4M3 byte that is not one of the two NaN encodings. +fn e4m3_byte(rng: &mut Rng) -> u8 { + loop { + let byte = rng.below(256) as u8; + if byte != 0x7F && byte != 0xFF { + return byte; + } + } +} + +fn read_f32(array: &MlxArray) -> Vec { + let contiguous = mlxcel_core::astype(array, dtype::FLOAT32); + mlxcel_core::eval(&contiguous); + mlxcel_core::array_to_raw_bytes(&contiguous) + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +fn raw_bytes(array: &UniquePtr) -> Vec { + let inner = array.as_ref().expect("array handle"); + mlxcel_core::eval(inner); + mlxcel_core::array_to_raw_bytes(inner) +} + +/// `WeightMap` holds `UniquePtr`, which is not `Debug`, so +/// `Result::expect_err` cannot be used on a requantization result. +fn expect_requantize_error(weights: WeightMap, reason: &str) -> String { + match requantize_block_fp8_weights(weights, SUPPORTED_FP8_BLOCK) { + Ok(_) => panic!("expected a failure: {reason}"), + Err(error) => error, + } +} + +struct BlockFp8Fixture { + weights: WeightMap, + /// `decode(bytes) * expanded_scales`, computed on the host. + expanded: Vec, + rows: usize, + cols: usize, +} + +/// Build a synthetic `.weight` / `.weight_scale_inv` pair plus +/// the host-side reconstruction of what it means. +/// +/// The host reconstruction expands each block scale over its whole 128x128 +/// tile and multiplies element by element. That is deliberately a different +/// expression of the same arithmetic from the device path under test, which +/// pads to whole blocks, reshapes to `[rb, 128, cb, 128]`, and broadcasts. The +/// two share only the E4M3 decode table. +fn block_fp8_fixture(prefix: &str, rows: usize, cols: usize, seed: u64) -> BlockFp8Fixture { + let block = SUPPORTED_FP8_BLOCK; + let row_blocks = rows.div_ceil(block); + let col_blocks = cols.div_ceil(block); + + let mut rng = Rng::new(seed); + let bytes: Vec = (0..rows * cols).map(|_| e4m3_byte(&mut rng)).collect(); + + let mut scale_bits = Vec::with_capacity(row_blocks * col_blocks); + let mut scale_values = Vec::with_capacity(row_blocks * col_blocks); + for _ in 0..row_blocks * col_blocks { + let (bits, value) = bf16_value(&mut rng); + scale_bits.extend_from_slice(&bits.to_le_bytes()); + scale_values.push(value); + } + + let mut expanded = Vec::with_capacity(rows * cols); + for r in 0..rows { + for c in 0..cols { + let scale = scale_values[(r / block) * col_blocks + (c / block)]; + expanded.push(f8_e4m3_to_f32(bytes[r * cols + c]) * scale); + } + } + + let mut weights = WeightMap::new(); + weights.insert( + format!("{prefix}.weight"), + mlxcel_core::from_bytes( + &bytes, + &[rows as i32, cols as i32], + mlxcel_core::dtype::UINT8, + ), + ); + weights.insert( + format!("{prefix}.weight_scale_inv"), + mlxcel_core::from_bytes_f16(&scale_bits, &[row_blocks as i32, col_blocks as i32], true), + ); + + BlockFp8Fixture { + weights, + expanded, + rows, + cols, + } +} + +/// The reconstruct-then-requantize path must land on exactly the tensors MLX +/// produces from the naive expanded-scale product, on both planes. +/// +/// This is the whole correctness claim for the reconstruction formula: block +/// pairing, padding, the broadcast axis order, and the trailing slice are all +/// wrong-answer-shaped rather than crash-shaped, and only a bitwise comparison +/// against an independently written reference catches them. +#[test] +fn fp8_block_requantize_matches_direct_path() { + // 130x160 is deliberately not block-aligned on either axis: rows need + // padding to 256 and columns to 256, so the trailing partial blocks and + // the slice back to the original extent are both exercised. 160 is still a + // multiple of the mxfp8 group size 32, as every real projection is. + let fixture = block_fp8_fixture("model.layers.0.mlp.down_proj", 130, 160, 0x51D3_9E11); + let rows = fixture.rows as i32; + let cols = fixture.cols as i32; + + assert!(has_block_fp8_weights(&fixture.weights)); + + let converted = requantize_block_fp8_weights(fixture.weights, SUPPORTED_FP8_BLOCK) + .expect("block FP8 requantization"); + + assert!( + !converted.contains_key("model.layers.0.mlp.down_proj.weight_scale_inv"), + "the block-scale sidecar must be consumed, not left for the model constructor" + ); + let packed = converted + .get("model.layers.0.mlp.down_proj.weight") + .expect("packed weight"); + let scales = converted + .get("model.layers.0.mlp.down_proj.scales") + .expect("mxfp8 scales"); + + // MLX packs mxfp8 as uint32 words of four 8-bit values, with one uint8 + // E8M0 exponent per group of 32. + assert_eq!( + mlxcel_core::array_shape(packed.as_ref().expect("packed")), + vec![rows, cols / 4] + ); + assert_eq!( + mlxcel_core::array_dtype(packed.as_ref().expect("packed")), + dtype::UINT32 + ); + assert_eq!( + mlxcel_core::array_shape(scales.as_ref().expect("scales")), + vec![rows, cols / MXFP8_GROUP_SIZE] + ); + assert_eq!( + mlxcel_core::array_dtype(scales.as_ref().expect("scales")), + dtype::UINT8 + ); + assert!( + !converted.contains_key("model.layers.0.mlp.down_proj.biases"), + "mxfp8 is a block-float mode and must not gain an affine zero-point plane" + ); + + // Reference: quantize the host-expanded product directly. + let reference_dense = mlxcel_core::from_slice_f32(&fixture.expanded, &[rows, cols]); + let reference = mlxcel_core::quantize_weights_with_mode( + &reference_dense, + MXFP8_GROUP_SIZE, + MXFP8_BITS, + MXFP8_MODE, + ); + let reference_w = mlxcel_core::quantized_weights_w(&reference); + let reference_scales = mlxcel_core::quantized_weights_scales(&reference); + + assert_eq!( + raw_bytes(packed), + raw_bytes(&reference_w), + "packed mxfp8 plane diverged from quantize(decode(bytes) * expanded_scales)" + ); + assert_eq!( + raw_bytes(scales), + raw_bytes(&reference_scales), + "mxfp8 scale plane diverged from quantize(decode(bytes) * expanded_scales)" + ); + + // Dequantizing must reproduce the reconstruction within the mxfp8 step. + // E4M3 carries four significant bits, so a value sharing its group's + // binade with the group maximum rounds by at most 2^-4 of that maximum. + let dequantized = unsafe { + mlxcel_core::dequantize( + packed.as_ref().expect("packed"), + scales.as_ref().expect("scales"), + std::ptr::null(), + MXFP8_GROUP_SIZE, + MXFP8_BITS, + MXFP8_MODE, + ) + }; + let recovered = read_f32(&dequantized); + assert_eq!(recovered.len(), fixture.expanded.len()); + + let group = MXFP8_GROUP_SIZE as usize; + for (index, (expected, actual)) in fixture.expanded.iter().zip(&recovered).enumerate() { + let group_start = (index / group) * group; + let group_max = fixture.expanded[group_start..group_start + group] + .iter() + .fold(0f32, |m, v| m.max(v.abs())); + let bound = group_max / 16.0 + f32::EPSILON; + let error = (expected - actual).abs(); + assert!( + error <= bound, + "element {index}: mxfp8 error {error} exceeded {bound} (group max {group_max})" + ); + } +} + +/// A block-aligned tensor takes the no-padding branch, and it must agree with +/// the reference there too. Without this the padded case could pass while the +/// common case silently took a different route. +#[test] +fn fp8_block_requantize_matches_direct_path_when_block_aligned() { + let fixture = block_fp8_fixture("model.layers.1.self_attn.q_proj", 256, 128, 0x0BAD_F00D); + let rows = fixture.rows as i32; + let cols = fixture.cols as i32; + + let converted = requantize_block_fp8_weights(fixture.weights, SUPPORTED_FP8_BLOCK) + .expect("block FP8 requantization"); + let packed = converted + .get("model.layers.1.self_attn.q_proj.weight") + .expect("packed weight"); + let scales = converted + .get("model.layers.1.self_attn.q_proj.scales") + .expect("mxfp8 scales"); + + let reference_dense = mlxcel_core::from_slice_f32(&fixture.expanded, &[rows, cols]); + let reference = mlxcel_core::quantize_weights_with_mode( + &reference_dense, + MXFP8_GROUP_SIZE, + MXFP8_BITS, + MXFP8_MODE, + ); + + assert_eq!( + raw_bytes(packed), + raw_bytes(&mlxcel_core::quantized_weights_w(&reference)) + ); + assert_eq!( + raw_bytes(scales), + raw_bytes(&mlxcel_core::quantized_weights_scales(&reference)) + ); +} + +/// The E4M3 decode must run through the byte table rather than any float +/// reinterpretation: every one of the 254 finite encodings has to survive the +/// device LUT lookup exactly. +#[test] +fn fp8_block_requantize_decodes_every_finite_e4m3_byte() { + // One row of 256 columns holds every byte; the block scale is 1.0 so the + // reconstruction is the decode alone. + let bytes: Vec = (0..=u8::MAX).collect(); + let mut weights = WeightMap::new(); + weights.insert( + "probe.weight".to_string(), + mlxcel_core::from_bytes(&bytes, &[1, 256], mlxcel_core::dtype::UINT8), + ); + // bf16 1.0 is 0x3F80; the row pads to one block and the 256 columns span + // two, so the sidecar is [1, 2] and both scales are 1.0. + let unit_scales: Vec = [0x3F80u16, 0x3F80] + .iter() + .flat_map(|bits| bits.to_le_bytes()) + .collect(); + weights.insert( + "probe.weight_scale_inv".to_string(), + mlxcel_core::from_bytes_f16(&unit_scales, &[1, 2], true), + ); + + let lut = super::e4m3_lookup_table(); + let restored = super::reconstruct_block_fp8_tensor( + "probe.weight", + weights["probe.weight"].as_ref().expect("weight"), + weights["probe.weight_scale_inv"] + .as_ref() + .expect("scale_inv"), + SUPPORTED_FP8_BLOCK, + lut.as_ref().expect("lut"), + ) + .expect("reconstruction"); + + let decoded = read_f32(&restored); + for byte in 0..=u8::MAX { + let expected = f8_e4m3_to_f32(byte); + let actual = decoded[byte as usize]; + if expected.is_nan() { + assert!(actual.is_nan(), "byte {byte:#04x} should decode to NaN"); + } else { + assert_eq!( + actual.to_bits(), + expected.to_bits(), + "byte {byte:#04x} decoded to {actual} instead of {expected}" + ); + } + } + + // The encoder round-trips every finite decode, so the table is the exact + // inverse the checkpoint was written with. + for byte in 0..=u8::MAX { + let value = f8_e4m3_to_f32(byte); + if value.is_nan() { + continue; + } + assert_eq!( + f32_to_f8_e4m3(value), + byte, + "byte {byte:#04x} failed to round-trip" + ); + } +} + +/// mxfp8 packs one E8M0 exponent per 32 values along the last axis, so a width +/// that is not a multiple of 32 has no representation. Refuse it by name +/// instead of letting MLX abort inside the first quantized matmul. +#[test] +fn fp8_block_requantize_rejects_odd_width() { + let fixture = block_fp8_fixture("model.layers.0.mlp.up_proj", 128, 100, 0x1234_5678); + let error = expect_requantize_error( + fixture.weights, + "a width of 100 is not a multiple of the mxfp8 group size", + ); + assert!( + error.contains("100") && error.contains("group size"), + "unhelpful error: {error}" + ); +} + +/// A sidecar with no weight, or with a shape that does not tile the weight, is +/// a corrupt pairing. Both mis-scale silently if accepted. +#[test] +fn fp8_block_requantize_rejects_broken_pairings() { + let mut orphan = WeightMap::new(); + orphan.insert( + "model.layers.0.mlp.gate_proj.weight_scale_inv".to_string(), + mlxcel_core::from_bytes_f16(&0x3F80u16.to_le_bytes(), &[1, 1], true), + ); + let error = expect_requantize_error(orphan, "a sidecar without its weight must fail"); + assert!(error.contains("no matching"), "unhelpful error: {error}"); + + let fixture = block_fp8_fixture("model.layers.0.mlp.gate_proj", 256, 256, 0x9999); + let mut wrong_scale = fixture.weights; + wrong_scale.insert( + "model.layers.0.mlp.gate_proj.weight_scale_inv".to_string(), + mlxcel_core::from_bytes_f16(&[0x80, 0x3F, 0x80, 0x3F], &[1, 2], true), + ); + let error = + expect_requantize_error(wrong_scale, "a [1, 2] sidecar cannot tile a 256x256 weight"); + assert!(error.contains("[2, 2]"), "unhelpful error: {error}"); +} + +/// A weight already decoded to a float dtype has lost the byte-to-block +/// pairing, so accepting it would apply the scales to the wrong numbers. +#[test] +fn fp8_block_requantize_rejects_predecoded_weights() { + let mut weights = WeightMap::new(); + weights.insert( + "model.layers.0.self_attn.o_proj.weight".to_string(), + mlxcel_core::zeros(&[128, 128], dtype::FLOAT16), + ); + weights.insert( + "model.layers.0.self_attn.o_proj.weight_scale_inv".to_string(), + mlxcel_core::from_bytes_f16(&0x3F80u16.to_le_bytes(), &[1, 1], true), + ); + let error = expect_requantize_error(weights, "a pre-decoded float weight must fail"); + assert!(error.contains("uint8"), "unhelpful error: {error}"); +} + +/// Tensors without a sidecar stay exactly as they were: norms, embeddings, the +/// gated-delta conv and the whole vision tower arrive dense in these +/// checkpoints, and `UnifiedLinear` decides per tensor from `.scales`. +#[test] +fn fp8_block_requantize_leaves_unpaired_tensors_dense() { + let fixture = block_fp8_fixture("model.layers.0.mlp.down_proj", 128, 128, 0xFEED); + let mut weights = fixture.weights; + weights.insert( + "model.norm.weight".to_string(), + mlxcel_core::ones(&[128], dtype::BFLOAT16), + ); + weights.insert( + "model.layers.0.linear_attn.conv1d.weight".to_string(), + mlxcel_core::ones(&[256, 4, 1], dtype::BFLOAT16), + ); + + let converted = + requantize_block_fp8_weights(weights, SUPPORTED_FP8_BLOCK).expect("requantization"); + + assert_eq!( + mlxcel_core::array_dtype(converted["model.norm.weight"].as_ref().expect("norm")), + dtype::BFLOAT16 + ); + assert_eq!( + mlxcel_core::array_shape( + converted["model.layers.0.linear_attn.conv1d.weight"] + .as_ref() + .expect("conv1d") + ), + vec![256, 4, 1] + ); + assert!(!converted.contains_key("model.norm.scales")); + assert!(!has_block_fp8_weights(&converted)); +} + +/// A checkpoint with no sidecars must come back untouched, so the pre-pass is +/// safe to run unconditionally on the Qwen3.5 paths. +#[test] +fn fp8_block_requantize_is_identity_without_sidecars() { + let mut weights = WeightMap::new(); + weights.insert( + "model.embed_tokens.weight".to_string(), + mlxcel_core::ones(&[8, 8], dtype::BFLOAT16), + ); + let converted = + requantize_block_fp8_weights(weights, SUPPORTED_FP8_BLOCK).expect("requantization"); + assert_eq!(converted.len(), 1); + assert!(converted.contains_key("model.embed_tokens.weight")); +} + +#[test] +fn qwen_fp8_config_detection() { + let supported = json!({ + "model_type": "qwen3_5", + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + }); + let detected = qwen_fp8_block_quantization(&supported) + .expect("detection") + .expect("a [128, 128] fp8 block layout is supported"); + assert_eq!( + detected, + Fp8BlockQuantization { + block_rows: 128, + block_cols: 128, + group_size: MXFP8_GROUP_SIZE, + bits: MXFP8_BITS, + mode: MXFP8_MODE, + } + ); + assert_eq!( + detected.effective_config(), + json!({"group_size": 32, "bits": 8, "mode": "mxfp8"}) + ); + + // The block extent is what pairs a scale with its elements, so an + // unvalidated geometry is an error rather than a silent 128 assumption. + let wrong_block = json!({ + "quantization_config": { + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [64, 64], + } + }); + let error = qwen_fp8_block_quantization(&wrong_block).expect_err("[64, 64] is unsupported"); + assert!(error.contains("weight_block_size"), "unhelpful: {error}"); + + // Another quantization method is simply not this path. + let int4 = json!({"quantization_config": {"quant_method": "int4", "bits": 4}}); + assert_eq!(qwen_fp8_block_quantization(&int4).expect("detection"), None); + + // No quantization metadata at all. + assert_eq!( + qwen_fp8_block_quantization(&json!({"model_type": "qwen3_5"})).expect("detection"), + None + ); + + // Per-tensor FP8 has a different sidecar layout and is refused by name. + let per_tensor = json!({"quantization_config": {"quant_method": "fp8", "fmt": "e4m3"}}); + let error = qwen_fp8_block_quantization(&per_tensor).expect_err("per-tensor fp8"); + assert!(error.contains("weight_block_size"), "unhelpful: {error}"); + + // A non-E4M3 float8 format would decode with the wrong exponent width. + let e5m2 = json!({ + "quantization_config": { + "quant_method": "fp8", + "fmt": "e5m2", + "weight_block_size": [128, 128], + } + }); + let error = qwen_fp8_block_quantization(&e5m2).expect_err("e5m2"); + assert!(error.contains("E4M3"), "unhelpful: {error}"); + + // The VLM wrappers place the block under text_config on some releases. + let nested = json!({ + "model_type": "qwen3_5", + "text_config": { + "model_type": "qwen3_5_text", + "quantization_config": { + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [128, 128], + } + } + }); + assert!( + qwen_fp8_block_quantization(&nested) + .expect("detection") + .is_some() + ); +} + +#[test] +fn merge_fp8_block_quantization_respects_an_existing_block() { + let detected = Fp8BlockQuantization { + block_rows: 128, + block_cols: 128, + group_size: MXFP8_GROUP_SIZE, + bits: MXFP8_BITS, + mode: MXFP8_MODE, + }; + + let mut fresh = json!({"model_type": "qwen3_5_text"}); + assert!(merge_fp8_block_quantization(&mut fresh, detected).expect("merge")); + assert_eq!(fresh["quantization"]["group_size"], json!(32)); + assert_eq!(fresh["quantization"]["bits"], json!(8)); + + // A pre-converted checkpoint already declares its own layout; the FP8 + // pre-pass must not overwrite it. + let mut converted = json!({"quantization": {"group_size": 64, "bits": 4}}); + assert!(!merge_fp8_block_quantization(&mut converted, detected).expect("merge")); + assert_eq!(converted["quantization"]["bits"], json!(4)); +} diff --git a/src/models/mod.rs b/src/models/mod.rs index f2d9cacef..d20021f1b 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -46,6 +46,11 @@ pub(crate) mod conv_decode; // every dynamic and every absent block, and `internlm2` dropped its // `rope_scaling` block at deserialization so the base never moved. pub mod dynamic_ntk_rope; +// `fp8_block` reads the vendor fine-grained FP8 checkpoints (E4M3 bytes plus a +// bf16 inverse scale per 128x128 block) and requantizes them to MLX-native +// mxfp8 at load. It is a pre-pass rather than part of `sanitize`, which is +// already 3k lines and carries no MLX-op-heavy conversion of this size. +pub mod fp8_block; pub mod gated_delta; // `rope_overrides` carries llama-server b10621's `--rope-scaling`, // `--rope-scale`, `--rope-freq-scale` and `--rope-freq-base` from the server @@ -354,6 +359,10 @@ pub(crate) use sanitize::{ // The only consumer outside `sanitize` is the diagnostics-gated Molmo2 vision // reference loader, so an unconditional re-export is dead in a default build // and `-D warnings` rejects it. +pub use fp8_block::{ + Fp8BlockQuantization, has_block_fp8_weights, merge_fp8_block_quantization, + qwen_fp8_block_quantization, requantize_block_fp8_weights, +}; #[cfg(any(test, feature = "xla-diagnostics", feature = "xla-diagnostics-cpu"))] pub(crate) use sanitize::load_weights_from_dir_with_filter; pub use sanitize::{ diff --git a/src/models/qwen3_5.rs b/src/models/qwen3_5.rs index 82cac1172..03c7bc1fa 100644 --- a/src/models/qwen3_5.rs +++ b/src/models/qwen3_5.rs @@ -1847,6 +1847,16 @@ impl Qwen35Model { .insert("quantization".to_string(), v["quantization"].clone()); } + // Vendor fine-grained FP8 releases carry no `quantization` block at + // all: they declare `quantization_config: {quant_method: "fp8", ...}` + // and ship raw E4M3 bytes with a per-128x128-block inverse scale. The + // weights are converted to MLX-native mxfp8 below, so the config has + // to declare that layout before the model constructor reads it. + let fp8_block = crate::models::qwen_fp8_block_quantization(&v)?; + if let Some(detected) = fp8_block { + crate::models::merge_fp8_block_quantization(&mut text_config_val, detected)?; + } + let config: Qwen35Config = serde_json::from_value(text_config_val) .map_err(|e| format!("Failed to parse config: {}", e))?; config.validate_supported().map_err(|e| e.to_string())?; @@ -1865,6 +1875,17 @@ impl Qwen35Model { println!("[Qwen3.5] Loading weights..."); let weights = crate::models::load_text_weights(model_dir, None)?; + // Reconstruct block-scaled FP8 tensors before any other sanitization: + // the sidecars are keyed off the original tensor names, and the pass + // replaces `.weight` with a packed plane plus `.scales`. + let weights = match fp8_block { + Some(detected) => { + println!("[Qwen3.5] Requantizing fine-grained FP8 weights to mxfp8..."); + crate::models::requantize_block_fp8_weights(weights, detected.block_rows)? + } + None => weights, + }; + // Strip language_model. prefix and sanitize let weights = sanitize_moe_weights(weights, &config); diff --git a/src/models/sanitize.rs b/src/models/sanitize.rs index 50edea6c2..73432934e 100644 --- a/src/models/sanitize.rs +++ b/src/models/sanitize.rs @@ -86,7 +86,7 @@ fn is_gemma4_unified_weight(name: &str) -> bool { /// F8_E4M3FN format: 1 sign bit, 4 exponent bits (bias=7), 3 mantissa bits. /// No infinity representation; the all-ones exponent with non-zero mantissa encodes NaN. /// Range: ±448.0. -fn f8_e4m3_to_f32(bits: u8) -> f32 { +pub(crate) fn f8_e4m3_to_f32(bits: u8) -> f32 { let sign = (bits >> 7) & 1; let exp = (bits >> 3) & 0xF; // 4-bit exponent let mant = bits & 0x7; // 3-bit mantissa From 28f719493c5d8625ac07bdb8403ccd6d3175f69b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 10 Sep 2026 11:17:41 +0900 Subject: [PATCH 2/3] docs(models): record the supported Qwen3.5 FP8 layout The Qwen3.5 entry said the FP8 and ModelOpt NVFP4 builds were both out of scope because mlxcel could not read vLLM-style `weight_scale_inv` block scales. The FP8 half is no longer true. Replace it with the layout that is now read, what the served model costs (8 bits per weight plus one E8M0 exponent per 32 values, close to the checkpoint's own footprint), which tensors stay dense, and which declarations are refused by name. ModelOpt NVFP4 stays out of scope on its own sidecar layout. Refs #1361 --- docs/supported-models.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/supported-models.md b/docs/supported-models.md index a63ff1fde..0d931a92b 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -101,7 +101,7 @@ Implemented VLM variants include: - Cohere Compass / North-Micro-Vision (`cohere_compass`): CohereLabs' small generative VLM. The vision half is a Qwen3-VL tower down to the config keys (Conv3d patch embed at patch 16 / temporal 2, a 48x48 learned position grid bilinearly resampled per image, 2D vision RoPE, `cu_seqlens`-packed attention, a 2x2 spatial merger, and three DeepStack mergers at blocks 8 / 16 / 24 injected after the first three decoder layers), so the encoder, image processor, prompt expansion and weight-prefix remap are the Qwen3-VL ones reused unchanged. The text half is a Command-style parallel decoder: one input LayerNorm feeds both the attention and the SwiGLU MLP and both land on the residual, there is no `post_attention_layernorm`, embeddings are tied, and the head output is multiplied by `logit_scale` (0.25). Its positional encoding is per layer type, keyed off `rope_parameters`: the `sliding_attention` layers rotate with the interleaved 3-axis MRoPE (section `[24, 20, 20]`, base 50000, split rotate-half) and see at most `sliding_window` (4096) keys, while the `full_attention` layers (indices 3, 7, 11, ... of 28) have a JSON `null` entry, meaning no positional encoding at all rather than a fallback to the default RoPE, and attend the whole prefix. `norm_type: "rms_norm"` is also implemented; a `transformer_block_type: "sequential"` or a non-interleaved `mrope_section` is refused at load rather than silently run as something else. The resize bounds come from the nested `image_processor` block of `processor_config.json` (`min_pixels` 16384, `max_pixels` 3868706), which is what HF's `AutoProcessor` builds the image processor from; the standalone `preprocessor_config.json` on the published checkpoint disagrees (65536 / 16777216) and is only the fallback. The distinction is not cosmetic: the stale bound upscales anything under 256x256, so a 224x224 input would carry 64 image tokens instead of 49 and desynchronize the prompt from the reference. Validated against `mlx-community/North-Micro-Vision-Instruct-4bit` (2.5B, text 4-bit affine group 64, vision tower bf16). Video input (`<|VIDEO_PAD|>`) and the `CohereCompassTextForSequenceClassification` reranker head are out of scope. - Pixtral and Mistral 3 VLM wrappers (Mistral 3 VLM supports both the standard Llama/Mistral text backbone and the Mistral4 MLA+MoE backbone; text and image-plus-text are validated on both, including the Mistral Small 4 119B checkpoint). Images keep their aspect ratio: each is downscaled so its longest side fits `size.longest_edge` and rounded up to a whole number of merged patches (`patch_size * spatial_merge_size`, so Pixtral rounds to 16 and Mistral 3 to 28), never upscaled. The `[IMG]` placeholder then expands to one patch token per merged cell laid out row by row, with `[IMG_BREAK]` between rows and `[IMG_END]` after the last, so non-square inputs keep the spatial structure the models were trained on. The encoder builds its 2D-RoPE grid per image from the actual patch shape, and multi-image batches with differing sizes run the tower once per image. - Qwen2-VL, Qwen2.5-VL, Qwen3-VL, and Qwen3-VL MoE: text, image, and video inputs are wired on both the CLI (`--video`) and OpenAI-compatible server (`video_url` content blocks). Video frames are decoded through the shared `ffmpeg` pipeline, sampled under Qwen processor frame bounds (`min_frames` and `max_frames` from `processor_config.json` when present, upstream defaults otherwise), padded only to `temporal_patch_size`, and rejected instead of silently clamped when a request exceeds `max_frames`. Server requests that omit a per-video FPS use the processor sidecar's `fps`; the CLI keeps its explicit `--fps` setting, whose default is 2.0. Mixed image/video prompts preserve the rendered placeholder order; a malformed prompt/media cardinality mismatch fails before generation. -- Qwen3.5-VL, Qwen3.6-VL MoE, and Qwen3.8-VL (`qwen3_5`, `qwen3_5_moe`): the Qwen 3.5-family paths serve all three public generations. Qwen3.8 declares `model_type: "qwen3_5"` / `Qwen3_5ForConditionalGeneration` and is architecturally identical to Qwen3.5 at the same size, down to a byte-identical weight-map key set, so no separate family exists in the tree. The stack is a hybrid text backbone (three gated-delta linear-attention layers for every full-attention layer, `full_attention_interval: 4`) under the Qwen3-VL vision tower, with interleaved MRoPE over an `[11, 11, 10]` section. Validated checkpoint: `mlx-community/Qwen3.8-27B-4bit` (16 GB, 4-bit affine, 64 layers, `head_dim` 256, 24 Q / 4 KV heads, vocab 248320 over a 248077-entry tokenizer). CLI text, image, and video, server text, image, and video, and multi-turn prompt-cache reuse all work at the code path level; decode measured 18.21 tok/s median over three text/image runs (spread 0.04) on Apple Silicon under background load, which is an order-of-magnitude baseline rather than a benchmark record. Video support shares the Qwen-VL path above: decoded clips produce `video_token_id` spans, Qwen visual grids keep a real temporal axis, and the MRoPE span scanner consumes both image and video token runs. Treat actual ffmpeg-backed Qwen3.8 video Q&A on target hardware as separate runtime qualification unless that exact checkpoint/run is cited. The three config keys the generation added (`output_gate_type`, `rope_parameters.mrope_interleaved`, and the top-level `language_model_only`) are read at load time and rejected with a named error when they ask for behavior mlxcel does not implement (`output_gate_type: "sigmoid"`, `mrope_interleaved: false`, `language_model_only: true`), rather than being dropped silently. `vision_start_token_id` is required from the config for this family and never defaulted: the whole family uses 248053, and a wrong start id mis-segments MRoPE vision spans without failing. Not supported on this family: MTP speculative decoding (the mlx-community conversions drop the `mtp.*` tensors and publish the drafter separately as a `qwen3_5_mtp` checkpoint, #1165) and audio. Despite the family naming Qwen3.8-27B is not an omni model: `config.json` carries no `audio_config`, `audio_tower`, `thinker_config`, or `talker_config`. The upstream `Qwen/Qwen3.8-27B` repo's own weight map holds `model.language_model.*`, `model.visual.*`, `mtp.*`, and `lm_head.weight`; the validated `mlx-community/Qwen3.8-27B-4bit` conversion is a different shape, `language_model.*` (with `lm_head.*` folded inside that prefix) and `vision_tower.*`, and drops `mtp.*` entirely. The FP8 and ModelOpt NVFP4 builds of the same model are out of scope: they use vLLM-style `weight_scale_inv` block scales that mlxcel cannot read. One footprint number is worth planning around because `head_dim: 256` is unusually large: measured KV state on the 27B is about 87 MiB of length-independent recurrent and conv state plus 64.1 KiB per token, against a 64.0 KiB (65,536 B/token = 16 full-attention layers x 4 KV heads x 256 head_dim x 2 x 2 bytes) architectural minimum for those layers; the two agree to within about 0.16%, so nothing is meaningfully over-allocated and the 48 gated-delta layers cost only the fixed term. At the native 262,144-token context that is 16 GiB of KV for a single sequence, on top of the weights. +- Qwen3.5-VL, Qwen3.6-VL MoE, and Qwen3.8-VL (`qwen3_5`, `qwen3_5_moe`): the Qwen 3.5-family paths serve all three public generations. Qwen3.8 declares `model_type: "qwen3_5"` / `Qwen3_5ForConditionalGeneration` and is architecturally identical to Qwen3.5 at the same size, down to a byte-identical weight-map key set, so no separate family exists in the tree. The stack is a hybrid text backbone (three gated-delta linear-attention layers for every full-attention layer, `full_attention_interval: 4`) under the Qwen3-VL vision tower, with interleaved MRoPE over an `[11, 11, 10]` section. Validated checkpoint: `mlx-community/Qwen3.8-27B-4bit` (16 GB, 4-bit affine, 64 layers, `head_dim` 256, 24 Q / 4 KV heads, vocab 248320 over a 248077-entry tokenizer). CLI text, image, and video, server text, image, and video, and multi-turn prompt-cache reuse all work at the code path level; decode measured 18.21 tok/s median over three text/image runs (spread 0.04) on Apple Silicon under background load, which is an order-of-magnitude baseline rather than a benchmark record. Video support shares the Qwen-VL path above: decoded clips produce `video_token_id` spans, Qwen visual grids keep a real temporal axis, and the MRoPE span scanner consumes both image and video token runs. Treat actual ffmpeg-backed Qwen3.8 video Q&A on target hardware as separate runtime qualification unless that exact checkpoint/run is cited. The three config keys the generation added (`output_gate_type`, `rope_parameters.mrope_interleaved`, and the top-level `language_model_only`) are read at load time and rejected with a named error when they ask for behavior mlxcel does not implement (`output_gate_type: "sigmoid"`, `mrope_interleaved: false`, `language_model_only: true`), rather than being dropped silently. `vision_start_token_id` is required from the config for this family and never defaulted: the whole family uses 248053, and a wrong start id mis-segments MRoPE vision spans without failing. Not supported on this family: MTP speculative decoding (the mlx-community conversions drop the `mtp.*` tensors and publish the drafter separately as a `qwen3_5_mtp` checkpoint, #1165) and audio. Despite the family naming Qwen3.8-27B is not an omni model: `config.json` carries no `audio_config`, `audio_tower`, `thinker_config`, or `talker_config`. The upstream `Qwen/Qwen3.8-27B` repo's own weight map holds `model.language_model.*`, `model.visual.*`, `mtp.*`, and `lm_head.weight`; the validated `mlx-community/Qwen3.8-27B-4bit` conversion is a different shape, `language_model.*` (with `lm_head.*` folded inside that prefix) and `vision_tower.*`, and drops `mtp.*` entirely. The vendor fine-grained FP8 builds (`Qwen/Qwen3.8-27B-FP8` and siblings) are supported: they declare `quantization_config: {"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [128, 128]}` and store each converted projection as raw E4M3 bytes in `*.weight` plus a bf16 inverse scale per 128x128 block in `*.weight_scale_inv`. The loader reconstructs each pair (`decode(byte) * block_scale`) and requantizes it to MLX-native mxfp8 at load, one tensor at a time, so the reconstruction stays transient; the served model is 8 bits per weight plus one E8M0 exponent per 32 values, roughly the checkpoint's own footprint rather than a dequantized copy. Tensors with no sidecar (norms, embeddings, `conv1d`, `A_log`, `dt_bias`, and the whole vision tower, which these releases list under `modules_to_not_convert`) stay dense. Any other `weight_block_size`, a non-E4M3 `fmt`, and per-tensor FP8 with no `weight_block_size` at all are named load errors rather than a silently mis-scaled model. The ModelOpt NVFP4 builds remain out of scope: they use a different sidecar layout (`weight_scale` / `weight_scale_2`). MTP is unaffected by this, since the family does not support MTP speculative decoding here in any case. One footprint number is worth planning around because `head_dim: 256` is unusually large: measured KV state on the 27B is about 87 MiB of length-independent recurrent and conv state plus 64.1 KiB per token, against a 64.0 KiB (65,536 B/token = 16 full-attention layers x 4 KV heads x 256 head_dim x 2 x 2 bytes) architectural minimum for those layers; the two agree to within about 0.16%, so nothing is meaningfully over-allocated and the 48 gated-delta layers cost only the fixed term. At the native 262,144-token context that is 16 GiB of KV for a single sequence, on top of the weights. - PaddleOCR-VL (`paddleocr_vl`): document-OCR VLM pairing a NaViT dynamic-resolution SigLIP-style vision encoder (Conv2d patch embedding, cached bilinearly interpolated learned position embeddings for repeated page grids, 2D vision RoPE, `cu_seqlens`-packed attention with uniform and length-bucketed batched fast paths, and a spatial-merge projector) with a lightweight ERNIE-4.5 text decoder that uses MRoPE. Best for plain OCR, tables, formulas, and chart understanding. - Step-3.7 (`step3p7`): StepFun's multimodal model. A 47-block `perception_encoder` ViT tower (Conv2d patchify, learned absolute position embeddings bilinearly resized for the patch grid, 2D vision RoPE, LayerScale, quick-GELU MLP) feeds two stride-2 conv downsamplers and a linear projector into the text hidden size. Each image runs a base pass (728 px, 52x52 grid collapsed to 169 tokens) and, when larger than 728 px, a tiled-patch pass (504 px windows, each 36x36 grid collapsed to 81 tokens); features are ordered patches-first then base per image and scattered into `` placeholders. The text decoder is the Step-3.5 MoE stack (`text_config.model_type: step3p5`), reused as the backbone. Text-only prompts work through the same checkpoint. - dots.ocr (`dots_ocr`): document-OCR VLM (rednote-hilab) pairing a 42-block `dots_vit` dynamic-resolution ViT with a Qwen2 text decoder. The tower shares the Qwen2-VL vision machinery (`cu_seqlens`-packed block-diagonal attention, 2D vision RoPE, merge-block patch ordering) but uses RMSNorm blocks, a SwiGLU vision MLP, a patch embed with bias and a trailing RMSNorm, and a `post_trunk_norm`; its merger projects each 2x2 patch block straight to the text width (no separate connector). Plain 1D RoPE text decode, no MRoPE. Best for layout analysis (bbox + category JSON), plain OCR, tables, formulas, and Markdown conversion. From a31c72f1335014471fa53fe3f46fb1433a75fc68 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 10 Sep 2026 11:25:02 +0900 Subject: [PATCH 3/3] docs(reports): add the technical report for PR #1742 Records what the released FP8 checkpoint actually holds (407 sidecars, all bf16, all exactly ceil(R/128) x ceil(C/128), all block-aligned, so the real checkpoint never exercises the padding branch), the numpy check that `weight_scale_inv` multiplies rather than divides despite its name, the measured 1.04x load-time peak against the 1.5x bound, and the two acceptance criteria that stay open: bf16 logit parity, because that 55.6 GB reference was deliberately not fetched, and the metal+accelerate command, which is not runnable on a Linux/CUDA host. Also replaces two em dashes in the module doc comment with colons. Refs #1361 --- ...qwen35-fp8-block-requantize-20260910.en.md | 177 ++++++++++++++++++ ...qwen35-fp8-block-requantize-20260910.ko.md | 177 ++++++++++++++++++ src/models/fp8_block.rs | 4 +- 3 files changed, 356 insertions(+), 2 deletions(-) create mode 100644 TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.en.md create mode 100644 TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.ko.md diff --git a/TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.en.md b/TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.en.md new file mode 100644 index 000000000..373f08a40 --- /dev/null +++ b/TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.en.md @@ -0,0 +1,177 @@ +# Technical Report: PR #1742 - feat(qwen3_5): load vendor fine-grained FP8 checkpoints + +**Date**: 2026-09-10 +**Author**: mlxcel maintainers +**Reviewer**: implementation review cycle +**Status**: Completed with one acceptance criterion deliberately unverified (the bf16 reference checkpoint was not fetched; the `metal,accelerate` acceptance command is not runnable on this host and was not run) +**Languages**: Rust +**Risk Level**: Medium (new load-time conversion on the Qwen3.5 paths; no behavior change for checkpoints that carry no `weight_scale_inv` sidecar) + +--- + +## Executive Summary + +The vendor FP8 releases of the Qwen3.5 family (`Qwen/Qwen3.8-27B-FP8` and siblings) store every converted projection as raw `E4M3` bytes in `.weight` plus a `bfloat16` inverse scale per 128x128 block in `.weight_scale_inv`. The issue that motivated this work assumed mlxcel would fail on the unexpected sidecar. It does not. MLX's safetensors reader maps `F8_E4M3` to `uint8` (`dtype_from_safetensor_str` in `mlx/io/safetensors.cpp`), so the bytes arrive intact, `UnifiedLinear` sees no `.scales` key and takes its dense path, and the model loads in a third of a second and generates. It is simply wrong by a per-block factor on every projection. + +That is the whole reason this is a correctness fix rather than a capability gap, and it is what the fix has to be evidenced against: a load that succeeds proves nothing here, because a load already succeeded. + +`models::fp8_block` reconstructs each pair on device and requantizes it to MLX-native `mxfp8` (group size 32, 8 bits). On the released 27B checkpoint the conversion adds about 30 seconds to load, holds peak allocator memory at 1.04x the final resident size, and produces coherent output. + +--- + +## 1. What the checkpoint actually holds + +Read directly from `models/mlx/qwen3.8-27b-fp8` rather than from the issue text: + +``` +config.json quantization_config: + {"activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", + "weight_block_size": [128, 128], "modules_to_not_convert": [... 882 entries ...]} +``` + +| Property | Measured | +|---|---| +| Tensors in the index | 1606 | +| `F8_E4M3` tensors | 407 | +| `*.weight_scale_inv` sidecars | 407 | +| `F8_E4M3` tensors with no sidecar | 0 | +| Sidecar in the same shard as its weight | 407 / 407 | +| Sidecar dtype | `BF16`, all of them | +| Sidecars whose shape is not exactly `ceil(R/128) x ceil(C/128)` | 0 | +| Weights whose width is not a multiple of 32 | 0 | +| Weights not block-aligned on both axes | 0 | + +The last row matters for testing. Every real tensor in this checkpoint is an exact multiple of 128 on both axes, so the real checkpoint never exercises the pad-and-slice branch. The unit test therefore uses a 130x160 tensor, which is aligned on neither. + +Converted tensors are the ten projections per decoder layer (`self_attn.{q,k,v,o}_proj`, `mlp.{gate,up,down}_proj`, `linear_attn.{in_proj_qkv,in_proj_z,out_proj}`) plus the seven in the bundled `mtp.*` head. Everything else stays `BF16`: `embed_tokens`, `lm_head`, all norms, `conv1d`, `A_log`, `dt_bias`, the low-rank `linear_attn.in_proj_a` / `in_proj_b`, and the entire 27-block vision tower. The 882-entry `modules_to_not_convert` list is informational: only a tensor that actually has a sidecar is converted, so the list never has to be parsed. + +### 1.1 The scale direction is the trap + +`weight_scale_inv` names an inverse, which invites a divide. It is a multiply. Checked in numpy on three real tensors, independently of the Rust path: + +| Tensor | `decode` alone | `decode * scale_inv` | `decode / scale_inv` | +|---|---|---|---| +| `layers.3.self_attn.q_proj` `[12288, 5120]` | max 448, std 92.7 | max 0.320, std 0.0172 | max 4.08e6 | +| `layers.1.mlp.down_proj` `[5120, 17408]` | max 448, std 76.8 | max 0.441, std 0.0108 | max 5.31e6 | +| `layers.1.linear_attn.in_proj_qkv` `[10240, 5120]` | max 448, std 90.9 | max 0.333, std 0.0154 | max 3.82e6 | + +The multiply lands on a normal transformer projection distribution; the divide is off by seven orders of magnitude. Note also that the raw decode maxes out at exactly 448 in every tensor, which is the E4M3 maximum: the publisher chose each block scale to fill the E4M3 range, which is what makes the block scale load-bearing rather than cosmetic. + +--- + +## 2. The conversion + +`requantize_block_fp8_weights(weights, block)` walks the sidecar keys in sorted order and, per pair: + +``` +indices = astype(raw_uint8, uint32) +decoded = take(lut_256_f32, indices, axis=0) # E4M3 decode, on device +padded = pad(decoded, rows -> rb*128, cols -> cb*128, 0.0) +blocked = reshape(padded, [rb, 128, cb, 128]) +scaled = blocked * reshape(astype(scale_inv, f32), [rb, 1, cb, 1]) +restored = slice(reshape(scaled, [rb*128, cb*128]), [0,0], [R, C]) +(q, s) = quantize(restored, group_size=32, bits=8, mode="mxfp8") +``` + +`.weight` becomes the packed `uint32` plane (`[R, C/4]`), `.scales` is added (`uint8`, `[R, C/32]`), and `.weight_scale_inv` is dropped. No `.biases` plane is produced, and the code refuses to continue if MLX ever returns one, because `mxfp8` is a block-float mode and an affine zero-point there would mean the mode was misrouted. + +Three points are worth stating explicitly. + +**The decode runs on device.** `sanitize.rs` already had a host-side `f8_e4m3_to_f32`, and the existing `F8_E4M3` branch in `tensor_view_to_array` uses it in a per-byte loop. At 26 GB of E4M3 bytes that loop would walk the whole checkpoint through the CPU and build an f32 `Vec` four times its size. Lifting the same function into a 256-entry `f32` table and gathering through it with `take` keeps the decode on the GPU and keeps the host copy at 1 KB. + +**f32, not bf16.** The issue's pseudocode decodes to bf16. f32 is used instead: E4M3 and bf16 scales are both exact in f32, so the product is the exactly-rounded true weight, and `mxfp8`'s own group-maximum search then sees the true maximum rather than a bf16-rounded one. The cost is a transient that is 4x the E4M3 bytes instead of 2x, which the measurement below shows is affordable. + +**Peak memory is bounded by construction, and measured.** Each pair is `eval`ed before the loop advances, and the E4M3 bytes and the sidecar are dropped before the quantize call. Without the `eval` MLX would keep every reconstruction lazy until the first forward pass, at which point the transient is not one tensor's f32 copy but the whole checkpoint's. + +### 2.1 What is refused, and why by name + +Each of these mis-scales a tensor without failing if it is accepted, which is precisely the failure mode this PR exists to remove: + +| Condition | Result | +|---|---| +| `weight_block_size` other than `[128, 128]` | named error | +| `fmt` present and not `e4m3` | named error | +| `quant_method: "fp8"` with no `weight_block_size` (per-tensor FP8) | named error | +| Weight rank other than 2 | named error | +| Sidecar shape that does not tile the weight | named error | +| Sidecar with no matching `.weight` | named error | +| Width not a multiple of 32 | named error | +| Weight already decoded to a float dtype | named error | +| No sidecar at all | tensor passes through unchanged | + +The last row is what lets the pre-pass run unconditionally on the Qwen3.5 paths. + +--- + +## 3. Where it is wired + +All three Qwen3.5 entry points, because all three reach `Qwen35Model::from_weights` and all three would otherwise build dense layers over unscaled bytes: + +- `Qwen35Model::load` (text): detect from the full config, merge the effective `quantization` into `text_config` before it is deserialized, convert after `load_text_weights` and before `sanitize_moe_weights`. +- `load_qwen3_5_vlm_with_variant` (VLM, the path the released checkpoint takes): convert immediately after `load_vlm_weights_common` and **before** the weight map is split into text and vision, so every sidecar is still adjacent to its weight. The vision tower simply has none. +- `try_load_special_model_from_weights` / `qwen35_text_config` (adapter): same two steps on the owned copy. + +The merged block is `{"group_size": 32, "bits": 8, "mode": "mxfp8"}`. `Qwen35Config` reads only `group_size` and `bits`; `UnifiedLinear::from_weights` then calls `infer_quantization_mode(has_biases: false, 32, 8)`, which returns `"mxfp8"`. Nothing had to be threaded through the model constructor, and the dense tensors keep taking the dense branch because that branch keys on the presence of `.scales` per tensor, not on the config. + +--- + +## 4. Validation + +Every command ran on Linux / aarch64 / GB10 (sm_121) with `--features cuda` and `MLX_ENABLE_TF32=0` set explicitly. MLX defaults `MLX_ENABLE_TF32=1`, which is a 10-bit mantissa and is not valid f32 evidence. + +### 4.1 The unit-level proof + +`cargo test --profile test-fast --features cuda --lib models::fp8_block` - 10 passed. + +The central test, `fp8_block_requantize_matches_direct_path`, compares the device path **bitwise on both planes** against a reference built on the host: expand each block scale over its whole tile, multiply element by element, and quantize the result with the same `quantize_weights_with_mode(..., 32, 8, "mxfp8")`. The two routes share only the E4M3 decode table. Block pairing, the padding, the broadcast axis order, and the trailing slice are all wrong-answer-shaped rather than crash-shaped, and a bitwise comparison is the only thing that catches them. The tensor is 130x160, aligned on neither axis, so the padding and the slice are both live. + +The same test also dequantizes the result and checks it against the reconstruction per group: E4M3 carries four significant bits, so an element sharing its group's binade with the group maximum rounds by at most `group_max / 16`. That bound holds for every element. + +Supporting cases: a block-aligned 256x128 tensor (the no-padding branch, also bitwise); all 256 E4M3 encodings through the device LUT compared to the host table bit for bit, plus a round-trip through `f32_to_f8_e4m3`; the seven refusals in the table above; dense tensors left untouched; and an identity check on a map with no sidecars. + +### 4.2 Real checkpoint + +`Qwen/Qwen3.8-27B-FP8`, 29 GB: + +``` +mlxcel generate -m models/mlx/qwen3.8-27b-fp8 -p "Write a limerick about a Rust compiler." -n 220 --show-reasoning +``` + +``` +Model loaded in 51.195s (resident: 23.36 GB, peak: 24.28 GB). +... +There once was a coder in Rust, +Whose compiler was grumpy and just, +It said, "You must borrow, +Or your code will not go, +'Cause lifetimes will make you be dust." +[Generated 207 tokens in 34.85s = 5.94 tok/s] +``` + +The model produced an on-topic reasoning trace that checked its own rhyme scheme, then a well-formed limerick, and stopped on its own at 207 of 220 tokens. A separate 64-token run loaded in 32.096 s at 23.36 GB resident and 24.37 GB peak. + +**Peak memory**: 24.37 / 23.36 = **1.04x**, against the 1.5x the issue asks for. These are MLX allocator figures (`mlxcel_core::memory::snapshot`, active and peak), which is the right instrument here: process RSS peaked at 6.9 GB under `/usr/bin/time -v` because MLX's CUDA allocations on this host do not appear in RSS. The number is measured, not asserted. + +### 4.3 Neighbouring suites + +`--lib models::qwen3_5` (28 passed, 11 ignored), `--lib models::sanitize` (71 passed), `--lib loading::special` (6 passed). `cargo fmt --all -- --check` and `cargo clippy --workspace --all-targets --features cuda -- -D warnings` pass. + +--- + +## 5. What is not verified + +**Logit parity against the bf16 original.** The issue asks that top-5 logits agree with `Qwen/Qwen3.8-27B` within the mxfp8 tolerance. That checkpoint is 55.6 GB and was deliberately not fetched on the validation host. The criterion is left unticked, and no substitute is offered for it: comparing against a different quantization of the same architecture, or against this PR's own dequantization, would answer a different question. What is proven instead is the reconstruction formula itself, bitwise, at the unit level, plus the scale direction on real tensors in numpy, plus coherent generation at model scale. + +**`cargo test --workspace --profile test-fast --features metal,accelerate`.** Not runnable on this host: Linux, aarch64, CUDA, no Metal and no Accelerate. Unrun, not passed. The CUDA equivalents above were run for the changed scope. + +**`qwen3_5_moe` FP8.** No such release was available. Detection covers it because it keys on the config, and a fused 3-D `experts.gate_up_proj` carrying block scales would be refused by name (rank 2 is required) rather than mis-converted. Per-expert 2-D experts would convert and reach the existing quantized switch-layer loaders, but that is reasoning, not a measurement. + +**MTP.** The drafter keeps its own weight loader inside `mlxcel-core`, which cannot reach `crate::models::fp8_block`, so the bundled `mtp.*` FP8 tensors are not converted for it. MTP speculative decoding is already unsupported for this family, and the VLM loader drops `mtp.*` entirely, so nothing regresses. + +--- + +## 6. Follow-ups worth filing + +- Move the FP8 block reconstruction into `mlxcel-core` so the MTP drafter and any other core-side loader can use it, if MTP support for this family is ever revisited. +- `requantize_block_fp8_weights` is family-agnostic already. DeepSeek V3 (UE8M0) and Mistral 4 FP8 could reuse it once their block geometry is confirmed; `mistral4.rs` currently drops `weight_scale_inv` keys outright, which is the same silent mis-scaling this PR removed from Qwen3.5. +- The conversion is about 30 s on a 27B checkpoint every load. A cached converted copy on disk would remove it, at the cost of a second full-size artifact. diff --git a/TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.ko.md b/TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.ko.md new file mode 100644 index 000000000..ba1e6fea1 --- /dev/null +++ b/TECHNICAL_REPORTS/1742-qwen35-fp8-block-requantize-20260910.ko.md @@ -0,0 +1,177 @@ +# 기술 보고서: PR #1742 - feat(qwen3_5): 벤더 fine-grained FP8 체크포인트 로딩 + +**날짜**: 2026-09-10 +**작성**: mlxcel maintainers +**검토**: 구현 리뷰 사이클 +**상태**: 완료. 인수 조건 하나는 의도적으로 미검증으로 남김 (bf16 원본 체크포인트를 받지 않음. 인수 조건에 적힌 `metal,accelerate` 명령은 이 호스트에서 실행 불가라 실행하지 않음) +**언어**: Rust +**위험도**: 중간 (Qwen3.5 경로에 로드 시점 변환이 새로 들어감. `weight_scale_inv` 사이드카가 없는 체크포인트의 동작은 그대로) + +--- + +## 요약 + +Qwen3.5 계열의 벤더 FP8 릴리스(`Qwen/Qwen3.8-27B-FP8` 및 형제 모델)는 변환 대상 프로젝션을 `.weight`의 raw `E4M3` 바이트와 `.weight_scale_inv`의 128x128 블록당 `bfloat16` 역스케일, 두 텐서로 나눠 저장한다. 이 작업의 출발점이 된 이슈는 mlxcel이 예상 못 한 사이드카에서 실패할 것이라고 봤다. 실패하지 않는다. MLX의 safetensors 리더는 `F8_E4M3`을 `uint8`로 매핑하고(`mlx/io/safetensors.cpp`의 `dtype_from_safetensor_str`), 그래서 바이트는 온전히 들어오며, `UnifiedLinear`는 `.scales` 키가 없으니 dense 경로를 타고, 모델은 0.3초 만에 로드돼 생성까지 한다. 다만 모든 프로젝션이 자기 블록 스케일만큼 틀려 있다. + +이것이 이 PR을 기능 추가가 아니라 정확성 수정으로 만드는 이유이고, 동시에 증거를 어디에 대야 하는지를 결정한다. 여기서 "로드에 성공했다"는 아무것도 증명하지 않는다. 이미 성공하고 있었기 때문이다. + +`models::fp8_block`은 각 쌍을 디바이스에서 복원한 뒤 MLX 네이티브 `mxfp8`(group size 32, 8 bits)로 재양자화한다. 공개된 27B 체크포인트에서 이 변환은 로드에 약 30초를 더하고, 할당기 피크를 최종 상주 크기의 1.04배로 유지하며, 정상적인 출력을 낸다. + +--- + +## 1. 체크포인트에 실제로 들어 있는 것 + +이슈 본문이 아니라 `models/mlx/qwen3.8-27b-fp8`을 직접 읽어 확인한 값이다. + +``` +config.json quantization_config: + {"activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", + "weight_block_size": [128, 128], "modules_to_not_convert": [... 882개 ...]} +``` + +| 항목 | 실측 | +|---|---| +| 인덱스의 텐서 수 | 1606 | +| `F8_E4M3` 텐서 | 407 | +| `*.weight_scale_inv` 사이드카 | 407 | +| 사이드카가 없는 `F8_E4M3` 텐서 | 0 | +| 자기 weight와 같은 샤드에 있는 사이드카 | 407 / 407 | +| 사이드카 dtype | 전부 `BF16` | +| 형상이 `ceil(R/128) x ceil(C/128)`이 아닌 사이드카 | 0 | +| 너비가 32의 배수가 아닌 weight | 0 | +| 두 축 모두 블록 정렬이 아닌 weight | 0 | + +마지막 줄이 테스트 설계를 결정한다. 이 체크포인트의 실제 텐서는 두 축 모두 정확히 128의 배수라서, 실물로는 pad-and-slice 분기가 한 번도 실행되지 않는다. 그래서 단위 테스트는 어느 축도 정렬돼 있지 않은 130x160 텐서를 쓴다. + +변환 대상은 디코더 레이어당 열 개의 프로젝션(`self_attn.{q,k,v,o}_proj`, `mlp.{gate,up,down}_proj`, `linear_attn.{in_proj_qkv,in_proj_z,out_proj}`)과 번들된 `mtp.*` 헤드의 일곱 개다. 나머지는 전부 `BF16`으로 남는다. `embed_tokens`, `lm_head`, 모든 norm, `conv1d`, `A_log`, `dt_bias`, 저랭크 `linear_attn.in_proj_a` / `in_proj_b`, 그리고 27블록짜리 비전 타워 전체가 여기 해당한다. 882개짜리 `modules_to_not_convert` 목록은 참고용이다. 실제로 사이드카를 가진 텐서만 변환하므로 이 목록은 파싱할 일이 없다. + +### 1.1 스케일 방향이 함정이다 + +`weight_scale_inv`라는 이름은 역수를 뜻하고, 그래서 나눗셈을 유도한다. 실제로는 곱셈이다. Rust 경로와 무관하게 numpy로 실제 텐서 세 개를 확인했다. + +| 텐서 | `decode`만 | `decode * scale_inv` | `decode / scale_inv` | +|---|---|---|---| +| `layers.3.self_attn.q_proj` `[12288, 5120]` | max 448, std 92.7 | max 0.320, std 0.0172 | max 4.08e6 | +| `layers.1.mlp.down_proj` `[5120, 17408]` | max 448, std 76.8 | max 0.441, std 0.0108 | max 5.31e6 | +| `layers.1.linear_attn.in_proj_qkv` `[10240, 5120]` | max 448, std 90.9 | max 0.333, std 0.0154 | max 3.82e6 | + +곱셈은 트랜스포머 프로젝션의 정상 분포에 떨어지고, 나눗셈은 일곱 자릿수만큼 벗어난다. 아울러 raw decode의 최댓값이 모든 텐서에서 정확히 448, 즉 E4M3의 최댓값이다. 배포자가 블록마다 E4M3 범위를 꽉 채우도록 스케일을 잡았다는 뜻이고, 그래서 블록 스케일은 장식이 아니라 값의 크기를 실제로 결정한다. + +--- + +## 2. 변환 + +`requantize_block_fp8_weights(weights, block)`은 사이드카 키를 정렬 순서로 훑으면서 쌍마다 다음을 수행한다. + +``` +indices = astype(raw_uint8, uint32) +decoded = take(lut_256_f32, indices, axis=0) # E4M3 디코드, 디바이스에서 +padded = pad(decoded, rows -> rb*128, cols -> cb*128, 0.0) +blocked = reshape(padded, [rb, 128, cb, 128]) +scaled = blocked * reshape(astype(scale_inv, f32), [rb, 1, cb, 1]) +restored = slice(reshape(scaled, [rb*128, cb*128]), [0,0], [R, C]) +(q, s) = quantize(restored, group_size=32, bits=8, mode="mxfp8") +``` + +`.weight`는 패킹된 `uint32` 평면(`[R, C/4]`)으로 바뀌고, `.scales`(`uint8`, `[R, C/32]`)가 추가되며, `.weight_scale_inv`는 제거된다. `.biases` 평면은 만들어지지 않고, MLX가 혹시라도 그것을 돌려주면 코드가 진행을 거부한다. `mxfp8`은 block-float 모드라 affine 영점이 나왔다는 것은 모드가 잘못 라우팅됐다는 뜻이기 때문이다. + +명시할 만한 지점이 셋 있다. + +**디코드는 디바이스에서 돈다.** `sanitize.rs`에는 이미 호스트 측 `f8_e4m3_to_f32`가 있고, `tensor_view_to_array`의 기존 `F8_E4M3` 분기가 바이트 단위 루프로 그것을 쓴다. E4M3 바이트가 26 GB인 상황에서 그 루프는 체크포인트 전체를 CPU로 통과시키고 네 배 크기의 f32 `Vec`을 만든다. 같은 함수를 256칸 `f32` 테이블로 올려 `take`로 게더하면 디코드가 GPU에 남고 호스트 사본은 1 KB로 끝난다. + +**bf16이 아니라 f32.** 이슈의 의사코드는 bf16으로 디코드한다. 대신 f32를 쓴다. E4M3 값도 bf16 스케일도 f32에서 정확하므로 곱은 정확히 반올림된 참값이 되고, 뒤이은 `mxfp8`의 그룹 최댓값 탐색이 bf16으로 뭉개진 값이 아니라 참 최댓값을 본다. 대가는 임시 버퍼가 E4M3 바이트의 2배가 아니라 4배가 되는 것인데, 아래 실측이 그 정도는 감당된다는 것을 보여준다. + +**피크 메모리는 구조적으로 묶여 있고, 실측했다.** 각 쌍은 루프가 다음으로 넘어가기 전에 `eval`되고, E4M3 바이트와 사이드카는 quantize 호출 전에 버려진다. `eval`이 없으면 MLX는 모든 복원을 첫 forward까지 lazy로 들고 있고, 그때의 임시 버퍼는 텐서 하나의 f32 사본이 아니라 체크포인트 전체분이 된다. + +### 2.1 무엇을 거부하고, 왜 이름 붙은 오류인가 + +아래는 전부 받아들이면 실패 없이 텐서 크기만 어긋나게 만드는 조건이다. 이 PR이 없애려는 실패 양상이 정확히 그것이다. + +| 조건 | 결과 | +|---|---| +| `[128, 128]`이 아닌 `weight_block_size` | 이름 붙은 오류 | +| `fmt`가 있고 `e4m3`이 아님 | 이름 붙은 오류 | +| `quant_method: "fp8"`인데 `weight_block_size` 없음 (per-tensor FP8) | 이름 붙은 오류 | +| weight 랭크가 2가 아님 | 이름 붙은 오류 | +| weight를 타일링하지 못하는 사이드카 형상 | 이름 붙은 오류 | +| 대응하는 `.weight`가 없는 사이드카 | 이름 붙은 오류 | +| 너비가 32의 배수가 아님 | 이름 붙은 오류 | +| 이미 float dtype으로 디코드된 weight | 이름 붙은 오류 | +| 사이드카가 아예 없음 | 텐서를 그대로 통과 | + +마지막 줄 덕분에 Qwen3.5 경로에서 이 사전 패스를 조건 없이 돌릴 수 있다. + +--- + +## 3. 어디에 연결했나 + +Qwen3.5의 세 진입점 전부다. 셋 다 `Qwen35Model::from_weights`에 도달하고, 그대로 두면 셋 다 스케일 안 걸린 바이트 위에 dense 레이어를 세운다. + +- `Qwen35Model::load` (텍스트): 전체 config에서 감지하고, `text_config`가 역직렬화되기 전에 유효 `quantization`을 병합하며, `load_text_weights` 뒤 `sanitize_moe_weights` 앞에서 변환한다. +- `load_qwen3_5_vlm_with_variant` (VLM, 공개 체크포인트가 타는 경로): `load_vlm_weights_common` 직후, 가중치 맵이 텍스트와 비전으로 **쪼개지기 전에** 변환한다. 그래야 모든 사이드카가 자기 weight 옆에 남아 있다. 비전 타워는 애초에 사이드카가 없다. +- `try_load_special_model_from_weights` / `qwen35_text_config` (어댑터): 소유 사본에 같은 두 단계를 적용한다. + +병합되는 블록은 `{"group_size": 32, "bits": 8, "mode": "mxfp8"}`이다. `Qwen35Config`는 `group_size`와 `bits`만 읽고, `UnifiedLinear::from_weights`가 `infer_quantization_mode(has_biases: false, 32, 8)`을 호출해 `"mxfp8"`을 얻는다. 모델 생성자에 새로 무엇을 꿰어 넣을 필요가 없었고, dense 텐서는 계속 dense 분기를 탄다. 그 분기가 config가 아니라 텐서별 `.scales` 존재 여부로 갈리기 때문이다. + +--- + +## 4. 검증 + +모든 명령은 Linux / aarch64 / GB10(sm_121)에서 `--features cuda`와 `MLX_ENABLE_TF32=0`을 명시해 실행했다. MLX의 기본값은 `MLX_ENABLE_TF32=1`이고 이는 가수 10비트라서 f32 증거가 되지 못한다. + +### 4.1 단위 수준의 증명 + +`cargo test --profile test-fast --features cuda --lib models::fp8_block` - 10개 통과. + +핵심 테스트 `fp8_block_requantize_matches_direct_path`는 디바이스 경로를 **두 평면 모두 비트 단위로** 호스트에서 만든 레퍼런스와 비교한다. 레퍼런스는 각 블록 스케일을 타일 전체로 확장해 원소별로 곱한 뒤 같은 `quantize_weights_with_mode(..., 32, 8, "mxfp8")`을 통과시켜 만든다. 두 경로가 공유하는 것은 E4M3 디코드 테이블뿐이다. 블록 짝짓기, 패딩, 브로드캐스트 축 순서, 뒤쪽 슬라이스는 전부 크래시가 아니라 틀린 답으로 나타나는 종류의 결함이고, 비트 비교만이 그것을 잡는다. 텐서는 130x160으로 어느 축도 정렬돼 있지 않아 패딩과 슬라이스가 둘 다 실제로 실행된다. + +같은 테스트가 결과를 dequantize해 그룹 단위로 복원값과 대조하기도 한다. E4M3은 유효 4비트를 실으므로, 그룹 최댓값과 같은 binade에 있는 원소는 최대 `group_max / 16`만큼 반올림된다. 모든 원소가 이 한계 안에 들었다. + +보조 케이스: 블록 정렬된 256x128 텐서(패딩 없는 분기, 역시 비트 단위 비교), 256가지 E4M3 인코딩 전부를 디바이스 LUT로 통과시켜 호스트 테이블과 비트 대조하고 `f32_to_f8_e4m3`로 왕복까지 확인, 위 표의 일곱 가지 거부 조건, dense 텐서 무변경, 사이드카가 없는 맵에 대한 항등성. + +### 4.2 실제 체크포인트 + +`Qwen/Qwen3.8-27B-FP8`, 29 GB: + +``` +mlxcel generate -m models/mlx/qwen3.8-27b-fp8 -p "Write a limerick about a Rust compiler." -n 220 --show-reasoning +``` + +``` +Model loaded in 51.195s (resident: 23.36 GB, peak: 24.28 GB). +... +There once was a coder in Rust, +Whose compiler was grumpy and just, +It said, "You must borrow, +Or your code will not go, +'Cause lifetimes will make you be dust." +[Generated 207 tokens in 34.85s = 5.94 tok/s] +``` + +모델은 자기 각운 구조를 스스로 점검하는 주제에 맞는 추론 흔적을 낸 뒤 형식이 맞는 limerick을 쓰고, 220토큰 중 207에서 스스로 멈췄다. 별도의 64토큰 실행은 32.096초에 로드됐고 상주 23.36 GB, 피크 24.37 GB였다. + +**피크 메모리**: 24.37 / 23.36 = **1.04배**로, 이슈가 요구한 1.5배 안이다. 이 수치는 MLX 할당기 값(`mlxcel_core::memory::snapshot`의 active와 peak)이고, 여기서는 그쪽이 맞는 계측기다. `/usr/bin/time -v`로 본 프로세스 RSS 피크는 6.9 GB인데, 이 호스트에서 MLX의 CUDA 할당이 RSS에 잡히지 않기 때문이다. 주장한 값이 아니라 측정한 값이다. + +### 4.3 인접 스위트 + +`--lib models::qwen3_5` (28 통과, 11 ignored), `--lib models::sanitize` (71 통과), `--lib loading::special` (6 통과). `cargo fmt --all -- --check`와 `cargo clippy --workspace --all-targets --features cuda -- -D warnings` 통과. + +--- + +## 5. 검증하지 못한 것 + +**bf16 원본 대비 로짓 일치.** 이슈는 첫 생성 위치의 top-5 로짓이 `Qwen/Qwen3.8-27B`와 mxfp8 허용오차 안에서 일치할 것을 요구한다. 그 체크포인트는 55.6 GB이고 검증 호스트에 의도적으로 받지 않았다. 이 항목은 체크하지 않은 채로 두며, 대체 증거를 제시하지도 않는다. 같은 아키텍처의 다른 양자화본과 비교하거나 이 PR 자신의 역양자화와 비교하는 것은 다른 질문에 답하는 일이기 때문이다. 대신 증명한 것은 복원 공식 자체를 단위 수준에서 비트 단위로, 스케일 방향을 실제 텐서에 numpy로, 그리고 모델 규모에서 정상 생성으로 확인한 것이다. + +**`cargo test --workspace --profile test-fast --features metal,accelerate`.** 이 호스트에서 실행 불가다. Linux, aarch64, CUDA이고 Metal도 Accelerate도 없다. 통과가 아니라 미실행이다. 변경 범위에 대해서는 위의 CUDA 대응 명령을 실행했다. + +**`qwen3_5_moe` FP8.** 해당 릴리스를 구할 수 없었다. 감지는 config를 보므로 범위에 들어가고, 블록 스케일을 가진 융합 3차원 `experts.gate_up_proj`는 잘못 변환되는 대신 이름 붙은 오류로 거부된다(랭크 2 필요). 전문가별 2차원 텐서라면 변환되어 기존 양자화 switch-layer 로더로 들어가겠지만, 그것은 추론이지 측정이 아니다. + +**MTP.** 드래프터는 `mlxcel-core` 안에 자기 가중치 로더를 갖고 있어 `crate::models::fp8_block`에 닿지 못하므로, 번들된 `mtp.*` FP8 텐서는 드래프터용으로 변환되지 않는다. 이 계열의 MTP 추측 디코딩은 이미 미지원이고 VLM 로더는 `mtp.*`를 통째로 버리므로 후퇴하는 것은 없다. + +--- + +## 6. 후속으로 낼 만한 것 + +- FP8 블록 복원을 `mlxcel-core`로 옮겨 MTP 드래프터를 비롯한 코어 측 로더가 쓸 수 있게 한다. 이 계열의 MTP 지원을 다시 볼 때 필요하다. +- `requantize_block_fp8_weights`는 이미 계열 비의존적이다. DeepSeek V3(UE8M0)와 Mistral 4 FP8도 블록 형상만 확인되면 재사용할 수 있다. `mistral4.rs`는 지금 `weight_scale_inv` 키를 그냥 버리는데, 이는 이 PR이 Qwen3.5에서 제거한 것과 같은 조용한 오스케일링이다. +- 변환은 27B 체크포인트 기준 로드마다 약 30초다. 변환본을 디스크에 캐시하면 없앨 수 있지만, 원본 크기의 산출물이 하나 더 생긴다. diff --git a/src/models/fp8_block.rs b/src/models/fp8_block.rs index 66d34a984..84519c1ac 100644 --- a/src/models/fp8_block.rs +++ b/src/models/fp8_block.rs @@ -18,8 +18,8 @@ //! The vLLM-style FP8 releases of the Qwen3.5 family (`Qwen/Qwen3.8-27B-FP8` //! and siblings) store every converted projection as two tensors: //! -//! * `.weight` — raw `E4M3` bytes, one per weight element, -//! * `.weight_scale_inv` — a `bfloat16` multiplier per 128x128 block. +//! * `.weight`: raw `E4M3` bytes, one per weight element, +//! * `.weight_scale_inv`: a `bfloat16` multiplier per 128x128 block. //! //! The true weight is `e4m3_decode(byte) * scale_inv[block]`. MLX has no //! float8 dtype and its safetensors reader maps `F8_E4M3` to `uint8`, so the