ferrox-quant 0.21.0

Quantized weight formats and fused dequant kernels for Ferrox
Documentation
//! The importance-matrix goldens: what the INSTALLED `llama-quantize`
//! b7650 wrote, with `--imatrix`, for two real rows of a real
//! checkpoint, in every K-quant format ferrox encodes.
//!
//! Sixteen super-blocks of `Qwen3-0.6B-BF16.gguf`: the whole of
//! `blk.0.attn_q.weight` row 0 and `blk.5.ffn_up.weight` row 7 (1024
//! columns each, four super-blocks per row), then eight single
//! super-blocks chosen as described below. The imatrix is
//! `llama-imatrix`'s own, computed on the CPU over 8 chunks of 512
//! tokens, stored here as the per-column quotient `in_sum2 / counts`
//! the quantizer consumes (f32 division, as `quantize.cpp:296` does
//! it). The goldens are the bytes of those super-blocks in
//! `llama-quantize --imatrix`'s Q4_K_M, Q5_K_M and Q6_K outputs, where
//! every tensor involved is Q4_K, Q5_K and Q6_K respectively.
//!
//! Real blocks and not the synthetic `k_quant_fixture`, for the reason
//! `testdata` gives: synthetic noise never lands on the near-ties that
//! make a fused multiply-add matter. These come from the same run that
//! was byte-identical across the whole checkpoint -- 311 of 311
//! tensors for all five targets -- so a golden going red here is a
//! real divergence from the tool, not a fixture artefact.
//!
//! # Which `make_qp_quants` sites the fixture pins, measured
//!
//! Two whole rows turn every STRUCTURAL mutation red (the plain weight
//! rule, the plain grid, `63/max` for stage 2, `sw` replaced by ones,
//! the `make_qp_quants` candidate loop cut, Q6_K's `qw` dropped) and
//! NONE of the twelve fused multiply-adds in `qp_quants.rs`. So each of
//! those sites was reverted in turn and the whole checkpoint scanned
//! for the first Q4_K super-block whose bytes moved; that block is
//! blocks 8..16 here, one per site, in the order of the `qp_quants.rs`
//! source: `x - scale*L` (1 block in the model moved), the candidate
//! `x - scale_is*l` (5), `sumlx` (1747), `suml2` (936), `slx` (1049),
//! `sl2` (556), the `new_l` `slx` (1068) and `sl2` (534). The
//! refinement loop itself moves 308,360 super-blocks, so any block pins
//! it.
//!
//! Four sites have NO differing super-block anywhere in the 2,098,688
//! Q4_K super-blocks of this checkpoint, and no fixture can pin them:
//! the two error accumulators `best_mse` and `mse` in `make_qp_quants`
//! (they only choose the candidate, and the choice never sat on the
//! rounding), and `sum_x2` and `sigma2 + x*x` in the imatrix weight
//! rule in `fit.rs`. They are spelled fused because the C is, and the
//! whole-model measurement is the only evidence either way. Nor does
//! the `uint8_t` wrap of a negative code in `make_qp_quants` or its
//! `nearest_int` tie rounding move a byte of this model.
//!
//! Every value is stored as bits (`u16` BF16, `u32` f32) so the fixture
//! cannot drift through a decimal round trip. The arrays live in
//! `data.rs`, generated by the script in the PR; this file is the doc,
//! the widening and the tests.

mod data;
use data::{
    BLOCKS_BF16, IMATRIX_F32, LLAMA_CPP_Q4_K_IMATRIX_GOLDEN, LLAMA_CPP_Q5_K_IMATRIX_GOLDEN,
    LLAMA_CPP_Q6_K_IMATRIX_GOLDEN,
};

/// The super-blocks widened to f32, the way `ferrox quantize` widens a
/// BF16 source (the top 16 bits of the f32, exactly).
pub(crate) fn blocks() -> Vec<f32> {
    BLOCKS_BF16
        .iter()
        .map(|&b| f32::from_bits(u32::from(b) << 16))
        .collect()
}

pub(crate) fn imatrix() -> Vec<f32> {
    IMATRIX_F32.iter().map(|&b| f32::from_bits(b)).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{encode_row_q4_k, encode_row_q5_k, encode_row_q6_k};

    fn check(name: &str, block_bytes: usize, got: &[u8], want: &[u8]) {
        assert_eq!(got.len(), want.len(), "{name}: length");
        for (b, (g, w)) in got
            .chunks(block_bytes)
            .zip(want.chunks(block_bytes))
            .enumerate()
        {
            assert_eq!(
                g, w,
                "{name}: super-block {b} disagrees with llama-quantize --imatrix"
            );
        }
    }

    /// `quantize_row_q4_K_impl`: the `qw * sqrt(sigma2 + x^2)` weights,
    /// the `(-0.9, 0.05, 36)` grid and the `make_qp_quants` stage 2,
    /// byte for byte. Super-blocks are independent, so a 1024-wide row
    /// encoded as four 256-wide calls is the same bytes; the fixture
    /// carries one imatrix slice per super-block.
    #[test]
    fn q4_k_with_an_imatrix_matches_llama_quantize_imatrix() {
        let x = blocks();
        let qw = imatrix();
        let mut got = Vec::new();
        for (row, w) in x.chunks(256).zip(qw.chunks(256)) {
            encode_row_q4_k(row, Some(w), &mut got).unwrap();
        }
        check("Q4_K", 144, &got, &LLAMA_CPP_Q4_K_IMATRIX_GOLDEN);
    }

    /// `quantize_row_q5_K_impl`, which is the Q4_K flow with `nmax = 31`
    /// and the clamp on the packed scale codes.
    #[test]
    fn q5_k_with_an_imatrix_matches_llama_quantize_imatrix() {
        let x = blocks();
        let qw = imatrix();
        let mut got = Vec::new();
        for (row, w) in x.chunks(256).zip(qw.chunks(256)) {
            encode_row_q5_k(row, Some(w), &mut got).unwrap();
        }
        check("Q5_K", 176, &got, &LLAMA_CPP_Q5_K_IMATRIX_GOLDEN);
    }

    /// `quantize_row_q6_K_impl`: the raw imatrix slice handed to
    /// `make_qx_quants` as `qw`, and nothing else changed.
    #[test]
    fn q6_k_with_an_imatrix_matches_llama_quantize_imatrix() {
        let x = blocks();
        let qw = imatrix();
        let mut got = Vec::new();
        for (row, w) in x.chunks(256).zip(qw.chunks(256)) {
            encode_row_q6_k(row, Some(w), &mut got).unwrap();
        }
        check("Q6_K", 210, &got, &LLAMA_CPP_Q6_K_IMATRIX_GOLDEN);
    }

    /// The imatrix changes the bytes. If it did not, the three goldens
    /// above would be pinning the plain encoder under another name.
    #[test]
    fn the_imatrix_actually_changes_every_format() {
        let x = blocks();
        let qw = imatrix();
        for (name, enc) in [
            (
                "Q4_K",
                encode_row_q4_k as fn(&[f32], Option<&[f32]>, &mut Vec<u8>) -> Option<()>,
            ),
            ("Q5_K", encode_row_q5_k),
            ("Q6_K", encode_row_q6_k),
        ] {
            let (mut plain, mut weighted) = (Vec::new(), Vec::new());
            enc(&x[..1024], None, &mut plain).unwrap();
            enc(&x[..1024], Some(&qw[..1024]), &mut weighted).unwrap();
            assert_ne!(
                plain, weighted,
                "{name}: the imatrix had no effect on this row"
            );
        }
    }
}