fuel-transformers 0.19.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! GGUF-quantized Qwen3 ported to the lazy-graph API.
//!
//! Qwen3 with Q4_0 block-quantized Linear weights. The forward path is
//! identical to [`crate::models::lazy_qwen3::Qwen3Model`]; only the weight
//! storage of `attn_{q,k,v,o}` / `ffn_{gate,up,down}` / `output` changes
//! from F32/BF16 to `WeightStorage::Q4_0`. RmsNorm gains, the token
//! embedding table, per-head QK-norm gains, and Q/K/V biases (when
//! present) stay in F32.
//!
//! Qwen3 evolves Qwen2 with **per-head QK-norm**: a RmsNorm is applied
//! to Q and K AFTER the per-head reshape, along the `head_dim` axis.
//! Norm gains are `[head_dim]` (not `[hidden_size]` like OLMo2). When
//! `cfg.attention_bias = true`, Q/K/V projections also carry biases —
//! these stay as F32 just like the norm gains.
//!
//! Construction paths:
//! - [`QuantizedQwen3Model::load_from_mmapped`] — convenience: load f32
//!   Qwen3 weights from HF safetensors and quantize each Linear weight
//!   to Q4_0.
//! - [`QuantizedQwen3Model::from_f32_bake`] — take f32 source weights
//!   (same `[in, out]` layout as `Qwen3Weights`) and quantize on the
//!   fly. Used by tests and by callers that already have unquantized
//!   weights in memory.
//! - [`QuantizedQwen3Model::from_gguf`] — load directly from a Qwen3
//!   GGUF file, keeping Q4_0 tensors quantized and dequantizing other
//!   GGML dtypes to F32. Mirrors the SmolLM3 loader policy.

use crate::models::lazy_qwen3::{Qwen3Config, Qwen3LayerExtras, Qwen3Model, Qwen3Weights};
#[cfg(test)]
use fuel_ir::Shape;
use fuel_ir::error::Result;
#[cfg(test)]
use fuel_tensor::Device;
use fuel_tensor::{lazy::LayerWeights, lazy::Tensor, lazy::WeightStorage};
use std::sync::Arc;

/// GGUF-quantized Qwen3 causal language model. Wraps a plain
/// [`Qwen3Model`] whose Linear weights are `WeightStorage::Q4_0`.
///
/// All non-Linear tensors (token embedding, RmsNorm gains including
/// per-head QK-norm gains, optional Q/K/V biases) remain F32 — matching
/// the GGUF convention used by llama.cpp for Qwen3 releases.
#[derive(Debug, Clone)]
pub struct QuantizedQwen3Model {
    inner: Qwen3Model,
}

impl QuantizedQwen3Model {
    /// Forward over a token-ID sequence. Returns logits with shape
    /// `(1, seq, vocab_size)`.
    pub fn forward(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        self.inner.forward(tokens, start_pos)
    }

    /// Forward over pre-computed embeddings, skipping the token-embedding
    /// lookup. Embeds must have shape `(1, seq, hidden_size)`.
    pub fn forward_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        self.inner.forward_embeds(embeds, start_pos)
    }

    /// Per-token hidden states up to the final RmsNorm, shape
    /// `(1, seq, hidden_size)`.
    pub fn forward_hidden(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        self.inner.forward_hidden(tokens, start_pos)
    }

    /// Pre-embedded variant of [`Self::forward_hidden`].
    pub fn forward_hidden_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        self.inner.forward_hidden_embeds(embeds, start_pos)
    }

    /// Build per-token embeddings without running the decoder.
    pub fn embed_tokens_anchored(&self, anchor: &Tensor, tokens: &[u32]) -> Result<Tensor> {
        self.inner.embed_tokens_anchored(anchor, tokens)
    }

    /// Decode/prefill through a pre-allocated KV cache (GAP-029 increment 3).
    ///
    /// **Delegates ONE level, and that is the only level there is** —
    /// `Qwen3Model` is a leaf (`config` + `weights`). Recorded because the sibling
    /// `QuantizedLlama3Model` has precedent for reaching `self.inner.inner` and
    /// silently losing LLaMA-3.1 RoPE scaling; that trap needs two levels AND a
    /// scaled-frequency variant, and this family has neither.
    pub fn forward_with_kv_context(
        &self,
        tokens: &[u32],
        cache: &mut fuel_tensor::inference_context::KvCache,
        ctx: &mut fuel_tensor::inference_context::InferenceContext,
    ) -> Result<Vec<f32>> {
        self.inner.forward_with_kv_context(tokens, cache, ctx)
    }

    /// Plan-once persistent decode.
    pub fn forward_with_kv_context_persistent(
        &self,
        tokens: &[u32],
        cache: &mut fuel_tensor::inference_context::KvCache,
        ctx: &mut fuel_tensor::inference_context::InferenceContext,
        session: &mut Option<fuel_tensor::inference_context::DecodeSession>,
    ) -> Result<Vec<f32>> {
        self.inner
            .forward_with_kv_context_persistent(tokens, cache, ctx, session)
    }

    /// Persistent decode with the session owned by the `InferenceContext`.
    pub fn forward_decode_step(
        &self,
        tokens: &[u32],
        cache: &mut fuel_tensor::inference_context::KvCache,
        ctx: &mut fuel_tensor::inference_context::InferenceContext,
    ) -> Result<Vec<f32>> {
        self.inner.forward_decode_step(tokens, cache, ctx)
    }

    /// Model configuration.
    pub fn config(&self) -> &Qwen3Config {
        &self.inner.config
    }

    /// Underlying [`Qwen3Model`] for direct access to the lazy graph
    /// API. The wrapper exists solely to label the quantization origin.
    pub fn inner(&self) -> &Qwen3Model {
        &self.inner
    }

    /// Convenience: load f32 Qwen3 weights from HF safetensors and
    /// quantize each Linear weight to Q4_0. Equivalent to
    /// `Self::from_f32_bake(cfg, Qwen3Weights::load_from_mmapped(st, &cfg)?)`.
    pub fn load_from_mmapped(
        st: &fuel_loaders::safetensors::MmapedSafetensors,
        cfg: Qwen3Config,
    ) -> Result<Self> {
        let src = Qwen3Weights::load_from_mmapped(st, &cfg)?;
        Self::from_f32_bake(cfg, src)
    }

    /// Construct from in-memory f32 source weights, quantizing each
    /// Linear weight to Q4_0. Norm gains, biases, per-head QK-norm
    /// gains, and the token-embedding table stay in F32 — matching
    /// the GGUF convention used by llama.cpp for Qwen3 releases.
    ///
    /// Q4_0 blocks run along each Linear's *in_features* axis, so
    /// every dimension that serves as an in_features must be a
    /// multiple of the Q4_0 block size (32): `cfg.hidden_size`
    /// (attn_q/k/v/o, ffn_gate/up, output) and
    /// `cfg.intermediate_size` (ffn_down). `num_key_value_heads *
    /// head_dim` only ever appears as an out_features (attn_k /
    /// attn_v) and carries no divisibility requirement. Source
    /// weights follow the same `[in_features, out_features]`
    /// row-major layout as [`Qwen3Weights`].
    pub fn from_f32_bake(cfg: Qwen3Config, src: Qwen3Weights) -> Result<Self> {
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let kv = cfg.num_key_value_heads * cfg.head_dim;
        // GAP-279: Qwen3 architecturally decouples head_dim from hidden_size
        // (confirmed decoupler — e.g. Qwen3-0.6B is 16 heads x 128 head_dim =
        // 2048 against hidden_size 1024), so attn_q's out_features and
        // attn_o's in_features are q_dim, NOT hidden_size. Mirrors the
        // already-correct `lazy_quantized_qwen3_moe.rs::from_f32_bake`.
        let q_dim = cfg.num_attention_heads * cfg.head_dim;
        // Gate exactly the dims that serve as a Linear in_features —
        // Q4_0 blocks run along K only. kv (= num_key_value_heads *
        // head_dim) is only ever an out_features (attn_k / attn_v)
        // and needs no divisibility (mirrors the gemma3 gate).
        check_q4_0_divisible(
            "hidden_size (attn_q/k/v, ffn_gate/up, output in-features)",
            h,
        )?;
        check_q4_0_divisible("intermediate_size (ffn_down in-features)", i)?;
        check_q4_0_divisible("num_attention_heads * head_dim (attn_o in-features)", q_dim)?;

        let quantize_linear =
            |w: &WeightStorage, in_features: usize, out_features: usize| -> Result<WeightStorage> {
                let f32_in_out = match w {
                WeightStorage::F32(a) => a.to_vec(),
                _ => return Err(fuel_ir::error::Error::Msg(
                    "QuantizedQwen3Model::from_f32_bake: source weights must be WeightStorage::F32"
                        .into(),
                )
                .bt()),
            };
                if f32_in_out.len() != in_features * out_features {
                    return Err(fuel_ir::error::Error::Msg(format!(
                        "QuantizedQwen3Model::from_f32_bake: weight has {} elems, expected {}×{}",
                        f32_in_out.len(),
                        in_features,
                        out_features,
                    ))
                    .bt());
                }
                quantize_in_out_to_q4_0(&f32_in_out, in_features, out_features)
            };

        if src.layers.len() != src.layer_extras.len() {
            return Err(fuel_ir::error::Error::Msg(format!(
                "QuantizedQwen3Model::from_f32_bake: layers ({}) must have matching layer_extras ({})",
                src.layers.len(), src.layer_extras.len(),
            )).bt());
        }

        let mut layers: Vec<LayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
        for (idx, layer) in src.layers.into_iter().enumerate() {
            let attn_q = quantize_linear(&layer.attn_q, h, q_dim)
                .map_err(|e| layer_err(idx, "attn_q", e))?;
            let attn_k =
                quantize_linear(&layer.attn_k, h, kv).map_err(|e| layer_err(idx, "attn_k", e))?;
            let attn_v =
                quantize_linear(&layer.attn_v, h, kv).map_err(|e| layer_err(idx, "attn_v", e))?;
            let attn_o = quantize_linear(&layer.attn_o, q_dim, h)
                .map_err(|e| layer_err(idx, "attn_o", e))?;
            let ffn_gate = quantize_linear(&layer.ffn_gate, h, i)
                .map_err(|e| layer_err(idx, "ffn_gate", e))?;
            let ffn_up =
                quantize_linear(&layer.ffn_up, h, i).map_err(|e| layer_err(idx, "ffn_up", e))?;
            let ffn_down = quantize_linear(&layer.ffn_down, i, h)
                .map_err(|e| layer_err(idx, "ffn_down", e))?;
            layers.push(LayerWeights {
                attn_q,
                attn_q_bias: layer.attn_q_bias,
                attn_k,
                attn_k_bias: layer.attn_k_bias,
                attn_v,
                attn_v_bias: layer.attn_v_bias,
                attn_o,
                ffn_gate,
                ffn_up,
                ffn_down,
                attn_norm_gain: layer.attn_norm_gain,
                ffn_norm_gain: layer.ffn_norm_gain,
            });
        }

        let output = quantize_linear(&src.output, h, cfg.vocab_size)
            .map_err(|e| fuel_ir::error::Error::Msg(format!("output: {e}")).bt())?;

        let inner = Qwen3Model {
            config: cfg,
            weights: Qwen3Weights {
                instance: fuel_tensor::decode_shape::ModelInstanceId::next(),
                token_embedding: src.token_embedding,
                layers,
                layer_extras: src.layer_extras,
                final_norm_gain: src.final_norm_gain,
                output,
            },
        };
        Ok(Self { inner })
    }

    /// Load a GGUF-quantized Qwen3 checkpoint. Q4_0 weight tensors stay
    /// quantized; norms / biases / embeddings / per-head QK-norm gains
    /// dequantize to F32. Other GGML dtypes for Linear weights are
    /// dequantized to F32 — matching the SmolLM3 loader policy.
    /// Embedding and `lm_head` share storage when `output.weight` is
    /// absent (tied embeddings — common in smaller Qwen3 releases).
    ///
    /// llama.cpp GGUF tensor naming for Qwen3:
    /// - `token_embd.weight`, `output_norm.weight`, `output.weight`
    /// - `blk.{i}.attn_q.weight`     / `blk.{i}.attn_q.bias`
    /// - `blk.{i}.attn_k.weight`     / `blk.{i}.attn_k.bias`
    /// - `blk.{i}.attn_v.weight`     / `blk.{i}.attn_v.bias`
    /// - `blk.{i}.attn_output.weight`
    /// - `blk.{i}.attn_q_norm.weight` / `blk.{i}.attn_k_norm.weight`
    ///   (Qwen3-specific per-head QK-norm gains, shape `[head_dim]`)
    /// - `blk.{i}.ffn_gate.weight`   / `blk.{i}.ffn_up.weight`
    ///   `blk.{i}.ffn_down.weight`
    /// - `blk.{i}.attn_norm.weight`  / `blk.{i}.ffn_norm.weight`
    pub fn from_gguf<P: AsRef<std::path::Path>>(path: P, cfg: &Qwen3Config) -> Result<Self> {
        use fuel_loaders::quantized::gguf_mmap::MmapedContent;
        let mc = MmapedContent::from_path(path)?;
        let mmap_arc = mc.mmap();
        let mmap_bytes: &[u8] = &mmap_arc[..];
        Self::from_gguf_content_and_bytes(mc.content(), mmap_bytes, cfg)
    }

    /// Load a GGUF-quantized Qwen3 checkpoint from an in-memory byte
    /// buffer. Used by callers that already have the full GGUF file
    /// resident (e.g. WASM, where the host passes a `Vec<u8>` rather
    /// than a path). Semantics match [`Self::from_gguf`] — the byte
    /// buffer must contain the complete GGUF file (header + tensor
    /// data).
    pub fn from_gguf_bytes(bytes: &[u8], cfg: &Qwen3Config) -> Result<Self> {
        use fuel_loaders::quantized::gguf_file::Content;
        use std::io::Cursor;
        let mut cursor = Cursor::new(bytes);
        let content = Content::read(&mut cursor)?;
        Self::from_gguf_content_and_bytes(&content, bytes, cfg)
    }

    /// Shared body for the path- and bytes-based GGUF loaders. `bytes`
    /// must cover the entire GGUF file starting at byte 0 — tensor
    /// offsets are interpreted relative to `content.tensor_data_offset`.
    fn from_gguf_content_and_bytes(
        content: &fuel_loaders::quantized::gguf_file::Content,
        bytes: &[u8],
        cfg: &Qwen3Config,
    ) -> Result<Self> {
        let data_off = content.tensor_data_offset as usize;

        let get_tensor_bytes =
            |name: &str| -> Result<(&[u8], fuel_loaders::quantized::GgmlDType, Vec<usize>)> {
                let info = content.tensor_infos.get(name).ok_or_else(|| {
                    fuel_ir::error::Error::Msg(format!("gguf: missing tensor {name:?}"))
                })?;
                let elems = info.shape.elem_count();
                let block_size = info.ggml_dtype.block_size();
                let bytes_len = elems / block_size * info.ggml_dtype.type_size();
                let start = data_off + info.offset as usize;
                Ok((
                    &bytes[start..start + bytes_len],
                    info.ggml_dtype,
                    info.shape.dims().to_vec(),
                ))
            };

        let load_f32 = |name: &str| -> Result<Vec<f32>> {
            let (bytes, dt, _) = get_tensor_bytes(name)?;
            dequant_bytes_to_f32(bytes, dt, name)
        };

        let load_weight = |name: &str,

                           out_features: usize,
                           in_features: usize|

         -> Result<WeightStorage> {
            let (bytes, dt, dims) = get_tensor_bytes(name)?;
            let expected = out_features * in_features;
            let actual: usize = dims.iter().product();
            if actual != expected {
                return Err(fuel_ir::error::Error::Msg(format!(
                    "gguf: tensor {name:?} has {actual} elements, expected {expected} for [{out_features}, {in_features}]",
                )).bt());
            }
            match dt {
                fuel_loaders::quantized::GgmlDType::Q4_0 => Ok(WeightStorage::Q4_0 {
                    words: bytes_to_u32_arc(bytes),
                    bytes_len: bytes.len(),
                    in_features,
                    out_features,
                }),
                _ => {
                    // Dequantize to F32 and transpose [out, in] → [in, out].
                    let f32_out_in = dequant_bytes_to_f32(bytes, dt, name)?;
                    let mut f32_in_out = vec![0.0_f32; expected];
                    for o in 0..out_features {
                        for j in 0..in_features {
                            f32_in_out[j * out_features + o] = f32_out_in[o * in_features + j];
                        }
                    }
                    Ok(WeightStorage::F32(Arc::from(f32_in_out)))
                }
            }
        };

        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let kv = cfg.num_key_value_heads * cfg.head_dim;
        // GAP-279: see from_f32_bake's comment above — Qwen3 decouples
        // head_dim from hidden_size, so attn_q/attn_o are q_dim-wide, not
        // hidden_size-wide. Mirrors `lazy_quantized_qwen3_moe.rs`'s
        // already-correct GGUF loader.
        let q_dim = cfg.num_attention_heads * cfg.head_dim;

        let token_embedding = load_f32("token_embd.weight")?;
        if token_embedding.len() != cfg.vocab_size * h {
            return Err(fuel_ir::error::Error::Msg(format!(
                "gguf token_embd.weight: {} elems, expected {}×{}",
                token_embedding.len(),
                cfg.vocab_size,
                h,
            ))
            .bt());
        }

        let mut layers: Vec<LayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
        let mut layer_extras: Vec<Qwen3LayerExtras> = Vec::with_capacity(cfg.num_hidden_layers);
        for idx in 0..cfg.num_hidden_layers {
            let prefix = format!("blk.{idx}");
            let attn_q = load_weight(&format!("{prefix}.attn_q.weight"), q_dim, h)?;
            let attn_k = load_weight(&format!("{prefix}.attn_k.weight"), kv, h)?;
            let attn_v = load_weight(&format!("{prefix}.attn_v.weight"), kv, h)?;
            let attn_o = load_weight(&format!("{prefix}.attn_output.weight"), h, q_dim)?;
            let ffn_gate = load_weight(&format!("{prefix}.ffn_gate.weight"), i, h)?;
            let ffn_up = load_weight(&format!("{prefix}.ffn_up.weight"), i, h)?;
            let ffn_down = load_weight(&format!("{prefix}.ffn_down.weight"), h, i)?;
            let attn_norm_gain: Arc<[f32]> =
                Arc::from(load_f32(&format!("{prefix}.attn_norm.weight"))?);
            let ffn_norm_gain: Arc<[f32]> =
                Arc::from(load_f32(&format!("{prefix}.ffn_norm.weight"))?);
            // Optional Q/K/V biases — preserved verbatim from GGUF when
            // present. cfg.attention_bias tells us whether to expect
            // them, but be forgiving: if the GGUF lacks them, fall
            // through to None.
            let bias = |name: &str, len: usize| -> Option<Arc<[f32]>> {
                load_f32(name).ok().and_then(|v| {
                    if v.len() == len {
                        Some(Arc::from(v))
                    } else {
                        None
                    }
                })
            };
            let (attn_q_bias, attn_k_bias, attn_v_bias) = if cfg.attention_bias {
                (
                    bias(&format!("{prefix}.attn_q.bias"), q_dim),
                    bias(&format!("{prefix}.attn_k.bias"), kv),
                    bias(&format!("{prefix}.attn_v.bias"), kv),
                )
            } else {
                (None, None, None)
            };
            layers.push(LayerWeights {
                attn_q,
                attn_q_bias,
                attn_k,
                attn_k_bias,
                attn_v,
                attn_v_bias,
                attn_o,
                ffn_gate,
                ffn_up,
                ffn_down,
                attn_norm_gain,
                ffn_norm_gain,
            });
            // Qwen3-specific per-head QK-norm gains, shape `[head_dim]`.
            let q_norm_gain: Arc<[f32]> =
                Arc::from(load_f32(&format!("{prefix}.attn_q_norm.weight"))?);
            let k_norm_gain: Arc<[f32]> =
                Arc::from(load_f32(&format!("{prefix}.attn_k_norm.weight"))?);
            if q_norm_gain.len() != cfg.head_dim {
                return Err(fuel_ir::error::Error::Msg(format!(
                    "gguf {prefix}.attn_q_norm.weight: {} elems, expected head_dim={}",
                    q_norm_gain.len(),
                    cfg.head_dim,
                ))
                .bt());
            }
            if k_norm_gain.len() != cfg.head_dim {
                return Err(fuel_ir::error::Error::Msg(format!(
                    "gguf {prefix}.attn_k_norm.weight: {} elems, expected head_dim={}",
                    k_norm_gain.len(),
                    cfg.head_dim,
                ))
                .bt());
            }
            layer_extras.push(Qwen3LayerExtras {
                q_norm_gain,
                k_norm_gain,
            });
        }

        let final_norm_gain: Arc<[f32]> = Arc::from(load_f32("output_norm.weight")?);

        // Tied embeddings: smaller Qwen3 releases reuse the embedding
        // for `lm_head`. If `output.weight` is present we honor it;
        // otherwise we synthesize an F32 lm_head from the
        // (already-dequantized) token embedding by transposing
        // [vocab, h] → [h, vocab].
        let output = if content.tensor_infos.contains_key("output.weight") {
            load_weight("output.weight", cfg.vocab_size, h)?
        } else {
            let mut f32_in_out = vec![0.0_f32; h * cfg.vocab_size];
            for v in 0..cfg.vocab_size {
                for j in 0..h {
                    f32_in_out[j * cfg.vocab_size + v] = token_embedding[v * h + j];
                }
            }
            WeightStorage::F32(Arc::from(f32_in_out))
        };

        let inner = Qwen3Model {
            config: cfg.clone(),
            weights: Qwen3Weights {
                instance: fuel_tensor::decode_shape::ModelInstanceId::next(),
                token_embedding: Arc::from(token_embedding),
                layers,
                layer_extras,
                final_norm_gain,
                output,
            },
        };
        Ok(Self { inner })
    }
}

fn layer_err(idx: usize, name: &str, e: fuel_ir::error::Error) -> fuel_ir::error::Error {
    fuel_ir::error::Error::Msg(format!("layer {idx} {name}: {e}")).bt()
}

/// Build a [`Qwen3Config`] from parsed GGUF metadata + tensor infos.
/// Mirrors the `qwen3_cfg_from_gguf` helper in the
/// `fuel-examples/quantized-qwen3` binary so callers that already have
/// a parsed `Content` (e.g. WASM hosts that parse from a byte slice)
/// can derive the config without reimplementing key lookups.
///
/// `vocab_size` falls back to the first dim of `token_embd.weight`
/// when `qwen3.vocab_size` is absent (GGUF reverses dims on read so
/// `dims[0]` is the vocabulary axis). Tied-embedding status is
/// inferred from the presence of `output.weight`.
pub fn qwen3_config_from_gguf_content(
    content: &fuel_loaders::quantized::gguf_file::Content,
) -> Result<Qwen3Config> {
    let md = &content.metadata;
    let get = |k: &str| -> Result<&fuel_loaders::quantized::gguf_file::Value> {
        md.get(k).ok_or_else(|| {
            fuel_ir::error::Error::Msg(format!("gguf metadata: missing key {k:?}")).bt()
        })
    };
    let num_attention_heads = get("qwen3.attention.head_count")?.to_u32()? as usize;
    let num_key_value_heads = get("qwen3.attention.head_count_kv")?.to_u32()? as usize;
    let head_dim = get("qwen3.attention.key_length")?.to_u32()? as usize;
    let num_hidden_layers = get("qwen3.block_count")?.to_u32()? as usize;
    let hidden_size = get("qwen3.embedding_length")?.to_u32()? as usize;
    let intermediate_size = get("qwen3.feed_forward_length")?.to_u32()? as usize;
    let max_position_embeddings = get("qwen3.context_length")?.to_u32()? as usize;
    let rms_norm_eps = get("qwen3.attention.layer_norm_rms_epsilon")?.to_f32()? as f64;
    let rope_theta = get("qwen3.rope.freq_base")?.to_f32()? as f64;
    let vocab_size = match md.get("qwen3.vocab_size") {
        Some(v) => v.to_u32()? as usize,
        None => {
            let info = content
                .tensor_infos
                .get("token_embd.weight")
                .ok_or_else(|| {
                    fuel_ir::error::Error::Msg("gguf: missing token_embd.weight".into()).bt()
                })?;
            let dims = info.shape.dims();
            if dims.is_empty() {
                return Err(fuel_ir::error::Error::Msg(
                    "gguf: token_embd.weight has empty shape".into(),
                )
                .bt());
            }
            dims[0]
        }
    };
    let tie_word_embeddings = !content.tensor_infos.contains_key("output.weight");
    Ok(Qwen3Config {
        vocab_size,
        hidden_size,
        intermediate_size,
        num_hidden_layers,
        num_attention_heads,
        num_key_value_heads,
        head_dim,
        max_position_embeddings,
        sliding_window: None,
        max_window_layers: 0,
        use_sliding_window: false,
        rope_theta,
        rms_norm_eps,
        attention_bias: false,
        tie_word_embeddings,
    })
}

fn check_q4_0_divisible(name: &str, n: usize) -> Result<()> {
    const QK4_0: usize = 32;
    if !n.is_multiple_of(QK4_0) {
        return Err(fuel_ir::error::Error::Msg(format!(
            "QuantizedQwen3Model::from_f32_bake: {name} ({n}) must be divisible by Q4_0 block size ({QK4_0})"
        )).bt());
    }
    Ok(())
}

/// Quantize a `[in_features, out_features]` row-major F32 weight matrix
/// into a `WeightStorage::Q4_0` keeping GGUF's native `[out, in]` block
/// layout. The implementation does the `[in, out] → [out, in]` transpose
/// first, then runs the per-row Q4_0 quantization.
fn quantize_in_out_to_q4_0(
    f32_in_out: &[f32],
    in_features: usize,
    out_features: usize,
) -> Result<WeightStorage> {
    use fuel_quantized::{BlockQ4_0, GgmlType};
    const QK4_0: usize = 32;
    if !in_features.is_multiple_of(QK4_0) {
        return Err(fuel_ir::error::Error::Msg(format!(
            "Q4_0 quantize: in_features ({in_features}) must be divisible by {QK4_0}"
        ))
        .bt());
    }

    // Transpose [in, out] → [out, in] so each row is contiguous in K.
    let mut f32_out_in = vec![0.0_f32; out_features * in_features];
    for o in 0..out_features {
        for j in 0..in_features {
            f32_out_in[o * in_features + j] = f32_in_out[j * out_features + o];
        }
    }

    let n_blocks = out_features * in_features / QK4_0;
    let mut blocks: Vec<BlockQ4_0> = vec![BlockQ4_0::zeros(); n_blocks];
    BlockQ4_0::from_float(&f32_out_in, &mut blocks);

    // BlockQ4_0 is repr(C) and exactly 18 bytes; reinterpret as bytes.
    let bytes_len = n_blocks * std::mem::size_of::<BlockQ4_0>();
    let byte_slice: &[u8] = fuel_quantized::blocks_as_bytes(&blocks);
    // Q4_0 storage holds u32 words; pad bytes to a multiple of 4 by
    // copying into a Vec<u8> first.
    let padded_len = bytes_len.div_ceil(4) * 4;
    let mut padded = vec![0_u8; padded_len];
    padded[..bytes_len].copy_from_slice(byte_slice);
    let words: Vec<u32> = padded
        .as_chunks::<4>()
        .0
        .iter()
        .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect();
    Ok(WeightStorage::Q4_0 {
        words: Arc::from(words),
        bytes_len,
        in_features,
        out_features,
    })
}

fn bytes_to_u32_arc(bytes: &[u8]) -> Arc<[u32]> {
    let padded_len = bytes.len().div_ceil(4) * 4;
    let mut padded = vec![0_u8; padded_len];
    padded[..bytes.len()].copy_from_slice(bytes);
    let words: Vec<u32> = padded
        .as_chunks::<4>()
        .0
        .iter()
        .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect();
    Arc::from(words)
}

/// Dequantize a raw GGUF byte slice of the given GGML dtype to F32.
/// Mirrors the small SmolLM3 helper but lives here so this module
/// stays independent of the SmolLM3 internals.
fn dequant_bytes_to_f32(
    bytes: &[u8],
    dt: fuel_loaders::quantized::GgmlDType,
    name: &str,
) -> Result<Vec<f32>> {
    fuel_quantized::dequant_ggml_bytes(bytes, dt, name)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_cfg() -> Qwen3Config {
        Qwen3Config {
            vocab_size: 32,
            hidden_size: 32,
            intermediate_size: 64,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: 4,
            head_dim: 8,
            max_position_embeddings: 64,
            sliding_window: None,
            max_window_layers: 0,
            use_sliding_window: false,
            rope_theta: 10_000.0,
            rms_norm_eps: 1e-5,
            attention_bias: false,
            tie_word_embeddings: false,
        }
    }

    fn tiny_weights(cfg: &Qwen3Config) -> Qwen3Weights {
        let mut s: u32 = 24680;
        let mut next = move || -> f32 {
            s = s.wrapping_mul(1103515245).wrapping_add(12345);
            ((s >> 16) as u16 as f32 / 65535.0 - 0.5) * 0.05
        };
        let mut vec_of =
            |n: usize| -> Arc<[f32]> { Arc::from((0..n).map(|_| next()).collect::<Vec<_>>()) };
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let kv = cfg.num_key_value_heads * cfg.head_dim;
        let token_embedding = vec_of(cfg.vocab_size * h);
        let mut layers: Vec<LayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
        let mut layer_extras: Vec<Qwen3LayerExtras> = Vec::with_capacity(cfg.num_hidden_layers);
        for _ in 0..cfg.num_hidden_layers {
            layers.push(LayerWeights {
                attn_q: WeightStorage::F32(vec_of(h * h)),
                attn_q_bias: if cfg.attention_bias {
                    Some(vec_of(h))
                } else {
                    None
                },
                attn_k: WeightStorage::F32(vec_of(h * kv)),
                attn_k_bias: if cfg.attention_bias {
                    Some(vec_of(kv))
                } else {
                    None
                },
                attn_v: WeightStorage::F32(vec_of(h * kv)),
                attn_v_bias: if cfg.attention_bias {
                    Some(vec_of(kv))
                } else {
                    None
                },
                attn_o: WeightStorage::F32(vec_of(h * h)),
                ffn_gate: WeightStorage::F32(vec_of(h * i)),
                ffn_up: WeightStorage::F32(vec_of(h * i)),
                ffn_down: WeightStorage::F32(vec_of(i * h)),
                attn_norm_gain: Arc::from(vec![1.0_f32; h]),
                ffn_norm_gain: Arc::from(vec![1.0_f32; h]),
            });
            layer_extras.push(Qwen3LayerExtras {
                q_norm_gain: Arc::from(vec![1.0_f32; cfg.head_dim]),
                k_norm_gain: Arc::from(vec![1.0_f32; cfg.head_dim]),
            });
        }
        let final_norm_gain = Arc::from(vec![1.0_f32; h]);
        let output = WeightStorage::F32(vec_of(h * cfg.vocab_size));
        Qwen3Weights {
            instance: fuel_tensor::decode_shape::ModelInstanceId::next(),
            token_embedding,
            layers,
            layer_extras,
            final_norm_gain,
            output,
        }
    }

    /// GAP-279: `test_cfg()`'s default (`4 heads x 8 head_dim = 32 == hidden_size
    /// 32`) is coupled, which is exactly why the loader bug this test exists
    /// for (`load_weight`/`quantize_linear` calling `attn_q`/`attn_o` at
    /// `h, h` instead of `q_dim, h` / `h, q_dim`) was invisible to every
    /// existing test in this file. Builds weights sized to the real
    /// `q_dim = num_attention_heads * head_dim`, which `tiny_weights` above
    /// (hardcoded at `h*h`) cannot produce.
    fn tiny_weights_decoupled_q_dim(cfg: &Qwen3Config) -> Qwen3Weights {
        let mut s: u32 = 314159;
        let mut next = move || -> f32 {
            s = s.wrapping_mul(1103515245).wrapping_add(12345);
            ((s >> 16) as u16 as f32 / 65535.0 - 0.5) * 0.05
        };
        let mut vec_of =
            |n: usize| -> Arc<[f32]> { Arc::from((0..n).map(|_| next()).collect::<Vec<_>>()) };
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let kv = cfg.num_key_value_heads * cfg.head_dim;
        let q_dim = cfg.num_attention_heads * cfg.head_dim;
        let token_embedding = vec_of(cfg.vocab_size * h);
        let mut layers: Vec<LayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
        let mut layer_extras: Vec<Qwen3LayerExtras> = Vec::with_capacity(cfg.num_hidden_layers);
        for _ in 0..cfg.num_hidden_layers {
            layers.push(LayerWeights {
                attn_q: WeightStorage::F32(vec_of(h * q_dim)),
                attn_q_bias: None,
                attn_k: WeightStorage::F32(vec_of(h * kv)),
                attn_k_bias: None,
                attn_v: WeightStorage::F32(vec_of(h * kv)),
                attn_v_bias: None,
                attn_o: WeightStorage::F32(vec_of(q_dim * h)),
                ffn_gate: WeightStorage::F32(vec_of(h * i)),
                ffn_up: WeightStorage::F32(vec_of(h * i)),
                ffn_down: WeightStorage::F32(vec_of(i * h)),
                attn_norm_gain: Arc::from(vec![1.0_f32; h]),
                ffn_norm_gain: Arc::from(vec![1.0_f32; h]),
            });
            layer_extras.push(Qwen3LayerExtras {
                q_norm_gain: Arc::from(vec![1.0_f32; cfg.head_dim]),
                k_norm_gain: Arc::from(vec![1.0_f32; cfg.head_dim]),
            });
        }
        let final_norm_gain = Arc::from(vec![1.0_f32; h]);
        let output = WeightStorage::F32(vec_of(h * cfg.vocab_size));
        Qwen3Weights {
            instance: fuel_tensor::decode_shape::ModelInstanceId::next(),
            token_embedding,
            layers,
            layer_extras,
            final_norm_gain,
            output,
        }
    }

    /// Before this fix, `from_f32_bake` called `quantize_linear(&layer.attn_q,
    /// h, h)` / `quantize_linear(&layer.attn_o, h, h)`, which on a decoupled
    /// config rejects the weight outright: `tiny_weights_decoupled_q_dim`
    /// sizes `attn_q` at `h * q_dim` elements, so the old call's
    /// `in_features * out_features = h * h` check fails to match whenever
    /// `q_dim != h` — exactly the real-checkpoint failure this test
    /// reproduces red-first (see `fuel-transformers/tests/qwen3_cpu_yardstick.rs`
    /// for the real Qwen3-0.6B GGUF hitting the analogous `load_weight` check).
    ///
    /// `hidden_size`/`intermediate_size`/`q_dim` are all chosen as
    /// multiples of 32 (the Q4_0 block size `check_q4_0_divisible` enforces
    /// on every in_features dimension) — the smallest values that are both
    /// Q4_0-legal AND decoupled (`q_dim=64 != hidden_size=32`).
    #[test]
    fn from_f32_bake_handles_decoupled_head_dim() {
        let cfg = Qwen3Config {
            vocab_size: 24,
            hidden_size: 32,
            intermediate_size: 64,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            head_dim: 16,
            max_position_embeddings: 32,
            sliding_window: None,
            max_window_layers: 0,
            use_sliding_window: false,
            rope_theta: 10_000.0,
            rms_norm_eps: 1e-5,
            attention_bias: false,
            tie_word_embeddings: false,
        };
        assert_ne!(
            cfg.num_attention_heads * cfg.head_dim,
            cfg.hidden_size,
            "this test's whole point is a decoupled head_dim (64 != 32)"
        );
        let model =
            QuantizedQwen3Model::from_f32_bake(cfg.clone(), tiny_weights_decoupled_q_dim(&cfg))
                .expect("GAP-279: the quantized loader must accept a decoupled head_dim config");
        let logits = model.forward(&[1, 2, 3], 0).expect("forward").realize_f32();
        assert_eq!(logits.len(), 3 * cfg.vocab_size);
        for &v in &logits {
            assert!(v.is_finite());
        }
    }

    /// Added during the M1 investigation (2026-10-07): the apparent
    /// real-model decode divergence turned out to be a context-free CPU
    /// reference fixture (`run_greedy` in
    /// `fuel-transformers/tests/qwen3_cpu_yardstick.rs`), not a decode-path
    /// bug — this test (and its real-weight reproduction attempts) never
    /// found anything wrong. Kept anyway as GQA coverage this module
    /// otherwise lacks: the ONLY other test covering this API
    /// (`quantized_qwen3_windowed_decode_matches_quantized_forward`, right
    /// below) uses `test_cfg()`'s `num_attention_heads == num_key_value_heads`
    /// (no GQA, `n_rep=1`), so a real GQA regression here would go
    /// undetected without this one (`num_attention_heads: 8,
    /// num_key_value_heads: 4`, `n_rep=2`).
    #[test]
    fn quantized_qwen3_gqa_decode_matches_quantized_forward() {
        use fuel_ir::DType;
        use fuel_tensor::Device;
        use fuel_tensor::{
            inference_context::DecodeSession, inference_context::InferenceContext,
            inference_context::KvCache,
        };

        let cfg = Qwen3Config {
            num_attention_heads: 8,
            num_key_value_heads: 4,
            head_dim: 8,
            ..test_cfg()
        };
        assert_ne!(
            cfg.num_attention_heads, cfg.num_key_value_heads,
            "this test's whole point is GQA (n_rep=2); a config change here would \
             make it stop testing the M1 divergence"
        );
        let model =
            QuantizedQwen3Model::from_f32_bake(cfg.clone(), tiny_weights_decoupled_q_dim(&cfg))
                .unwrap();
        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5, 6];
        let prefill = 3;
        // Cache sized with slack beyond what any single step fills (the
        // pre-existing non-GQA precedent test uses exactly `tokens.len()`,
        // zero slack) — closer to how a real caller sizes a KvCache with
        // room for future decode steps.
        let max_seq_len = tokens.len() + 16;

        let dev = Device::cpu();
        let mut cache = KvCache::with_capacity(
            cfg.num_hidden_layers,
            cfg.num_key_value_heads,
            cfg.head_dim,
            max_seq_len,
            DType::F32,
            &dev,
        )
        .expect("with_capacity");
        let mut ctx = InferenceContext::new(dev);
        let mut session: Option<DecodeSession> = None;

        model
            .forward_with_kv_context_persistent(
                &tokens[..prefill],
                &mut cache,
                &mut ctx,
                &mut session,
            )
            .expect("prefill");

        for pos in prefill..tokens.len() {
            let got = model
                .forward_with_kv_context_persistent(
                    &tokens[pos..=pos],
                    &mut cache,
                    &mut ctx,
                    &mut session,
                )
                .expect("decode");
            let full = model.forward(&tokens[..=pos], 0).unwrap().realize_f32();
            let expected = &full[pos * cfg.vocab_size..(pos + 1) * cfg.vocab_size];
            let worst = got
                .iter()
                .zip(expected.iter())
                .map(|(a, b)| (a - b).abs())
                .fold(0.0_f32, f32::max);
            assert!(
                worst < 1e-5,
                "GQA (n_rep=2) decode at position {pos} diverged by {worst} from plain \
                 forward() on the SAME tiny weights — this is the M1 investigation's \
                 red-first reproduction; see this test's doc comment",
            );
        }
    }

    /// **GAP-029 increment 3 — the quantized wrapper's decode delegation.**
    ///
    /// The module's `test_cfg()` is deliberately dense (`sliding_window: None`),
    /// so this overrides it to a MIXED window: a quantized decode test on the
    /// dense config would exercise the delegation without ever touching the
    /// N=2 stacked-mask path. Three decode steps, so the assertions land on the
    /// per-token REBIND path rather than the held-graph build.
    #[test]
    fn quantized_qwen3_windowed_decode_matches_quantized_forward() {
        use fuel_ir::DType;
        use fuel_tensor::Device;
        use fuel_tensor::{
            inference_context::DecodeSession, inference_context::InferenceContext,
            inference_context::KvCache,
        };

        let cfg = Qwen3Config {
            sliding_window: Some(4),
            max_window_layers: 1,
            use_sliding_window: true,
            ..test_cfg()
        };
        let model = QuantizedQwen3Model::from_f32_bake(cfg.clone(), tiny_weights(&cfg)).unwrap();
        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5, 6];
        assert!(
            tokens.len() > 4,
            "non-vacuity: the window must actually bite"
        );
        let prefill = 3;

        let dev = Device::cpu();
        let mut cache = KvCache::with_capacity(
            cfg.num_hidden_layers,
            cfg.num_key_value_heads,
            cfg.head_dim,
            tokens.len(),
            DType::F32,
            &dev,
        )
        .expect("with_capacity");
        let mut ctx = InferenceContext::new(dev);
        let mut session: Option<DecodeSession> = None;

        model
            .forward_with_kv_context_persistent(
                &tokens[..prefill],
                &mut cache,
                &mut ctx,
                &mut session,
            )
            .expect("prefill");

        for pos in prefill..tokens.len() {
            let got = model
                .forward_with_kv_context_persistent(
                    &tokens[pos..=pos],
                    &mut cache,
                    &mut ctx,
                    &mut session,
                )
                .expect("decode");
            let full = model.forward(&tokens[..=pos], 0).unwrap().realize_f32();
            let expected = &full[pos * cfg.vocab_size..(pos + 1) * cfg.vocab_size];
            let worst = got
                .iter()
                .zip(expected.iter())
                .map(|(a, b)| (a - b).abs())
                .fold(0.0_f32, f32::max);
            assert!(
                worst < 1e-5,
                "quantized windowed decode at position {pos} diverged by {worst}; a \
                 single-mask port shows ~7e-3 here",
            );
        }
        assert_eq!(cache.cached_len, tokens.len());
    }

    #[test]
    fn forward_shape_finite_with_q4_0_weights() {
        let cfg = test_cfg();
        let src = tiny_weights(&cfg);
        let model = QuantizedQwen3Model::from_f32_bake(cfg.clone(), src).unwrap();
        // Sanity: all Linear projections in layer 0 are now Q4_0.
        let l0 = &model.inner().weights.layers[0];
        assert!(matches!(l0.attn_q, WeightStorage::Q4_0 { .. }));
        assert!(matches!(l0.attn_k, WeightStorage::Q4_0 { .. }));
        assert!(matches!(l0.attn_v, WeightStorage::Q4_0 { .. }));
        assert!(matches!(l0.attn_o, WeightStorage::Q4_0 { .. }));
        assert!(matches!(l0.ffn_gate, WeightStorage::Q4_0 { .. }));
        assert!(matches!(l0.ffn_up, WeightStorage::Q4_0 { .. }));
        assert!(matches!(l0.ffn_down, WeightStorage::Q4_0 { .. }));
        assert!(matches!(
            model.inner().weights.output,
            WeightStorage::Q4_0 { .. }
        ));
        // Per-head QK-norm gains stay F32 (they're not in WeightStorage,
        // but verify the count matches num_hidden_layers).
        assert_eq!(
            model.inner().weights.layer_extras.len(),
            cfg.num_hidden_layers
        );

        let logits = model.forward(&[1, 2, 3], 0).unwrap();
        assert_eq!(logits.shape().dims(), &[1, 3, cfg.vocab_size]);
        for &v in &logits.realize_f32() {
            assert!(v.is_finite(), "non-finite logit: {v}");
        }
    }

    #[test]
    fn forward_with_attention_bias() {
        // With attention_bias=true, Q/K/V biases are preserved verbatim
        // through quantization (they're F32, not WeightStorage).
        let mut cfg = test_cfg();
        cfg.attention_bias = true;
        let src = tiny_weights(&cfg);
        let model = QuantizedQwen3Model::from_f32_bake(cfg.clone(), src).unwrap();
        let l0 = &model.inner().weights.layers[0];
        assert!(l0.attn_q_bias.is_some());
        assert!(l0.attn_k_bias.is_some());
        assert!(l0.attn_v_bias.is_some());
        let logits = model.forward(&[1, 2, 3], 0).unwrap();
        assert_eq!(logits.shape().dims(), &[1, 3, cfg.vocab_size]);
        for &v in &logits.realize_f32() {
            assert!(v.is_finite(), "non-finite logit: {v}");
        }
    }

    #[test]
    fn forward_embeds_matches_forward_after_token_lookup() {
        let cfg = test_cfg();
        let src = tiny_weights(&cfg);
        let model = QuantizedQwen3Model::from_f32_bake(cfg.clone(), src).unwrap();
        let tokens: Vec<u32> = vec![1, 2, 3];
        let logits_ref = model.forward(&tokens, 0).unwrap().realize_f32();
        let anchor =
            Tensor::from_f32(vec![0.0_f32], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
        let embeds = model.embed_tokens_anchored(&anchor, &tokens).unwrap();
        let logits_via_embeds = model.forward_embeds(&embeds, 0).unwrap().realize_f32();
        let max_diff = logits_ref
            .iter()
            .zip(logits_via_embeds.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(
            max_diff < 1e-4,
            "Quantized Qwen3 forward vs forward_embeds must agree (max diff {max_diff})"
        );
    }
}