aprender-serve 0.66.0

//! Quantization and dequantization for model weights
//!
//! Implements quantization formats used by GGUF models:
//! - `F16`: 16-bit IEEE 754 half-precision
//! - `Q4_0`: 4-bit quantization (block size 32)
//! - `Q4_1`: 4-bit with scale and min (block size 32)
//! - `Q5_0`: 5-bit quantization (block size 32)
//! - `Q5_1`: 5-bit with scale and min (block size 32)
//! - `Q8_0`: 8-bit quantization (block size 32)
//! - `Q4_K`: 4-bit K-quantization (super-block size 256)
//! - `Q5_K`: 5-bit K-quantization (super-block size 256)
//! - `Q6_K`: 6-bit K-quantization (super-block size 256)
//!
//! ## `Q4_0` Format
//!
//! `Q4_0` stores weights in blocks of 32 values:
//! - 1 float32 scale factor per block
//! - 16 bytes of 4-bit quantized values (2 values per byte)
//! - Dequantization: `value = scale * quantized_value`
//!
//! ## `Q8_0` Format
//!
//! `Q8_0` stores weights in blocks of 32 values:
//! - 1 float32 scale factor per block
//! - 32 int8 quantized values
//! - Dequantization: `value = scale * quantized_value`
//!
//! ## `Q4_K` Format
//!
//! `Q4_K` uses super-blocks of 256 values divided into 8 blocks of 32 values:
//! - 1 half-precision super-block scale (`d`)
//! - 1 half-precision super-block min (`dmin`)
//! - 12 bytes of 6-bit block scales (packed)
//! - 128 bytes of 4-bit quantized values
//! - Dequantization: `value = d * scale * quantized - dmin * min`
//! - Achieves 4.5 bits per weight with better quality than `Q4_0`
//!
//! ## `Q5_K` Format
//!
//! `Q5_K` uses super-blocks of 256 values divided into 8 blocks of 32 values:
//! - 1 half-precision super-block scale (`d`)
//! - 1 half-precision super-block min (`dmin`)
//! - 12 bytes of 6-bit block scales (packed)
//! - 32 bytes of high bits (1 bit per value for 5-bit quantization)
//! - 128 bytes of low 4-bit quantized values
//! - Dequantization: `value = d * scale * quantized - dmin * min`
//! - Achieves 5.5 bits per weight (higher quality than `Q4_K`)
//!
//! ## `Q6_K` Format
//!
//! `Q6_K` uses super-blocks of 256 values divided into 16 blocks of 16 values:
//! - 1 half-precision super-block scale (`d`)
//! - 16 bytes of 8-bit block scales
//! - 64 bytes of high 2 bits (2 bits per value for 6-bit quantization)
//! - 128 bytes of low 4-bit quantized values
//! - Dequantization: `value = d * scale * quantized`
//! - Achieves 6.5625 bits per weight (highest quality K-quant format)

use crate::error::{RealizarError, Result};

// ============================================================================
// Shattered submodules (PMAT-802)
// ============================================================================

pub mod activation;
pub mod bsum_precompute;
pub mod contract_tests;
pub mod dequant;
pub mod encode;
pub mod format_trait;
pub mod fused_gate_up;
pub mod fused_k;
pub mod fused_q5k_q6k;
pub(crate) mod gemv_pool;
pub mod generic_dot;
pub mod generic_matvec;
pub mod parallel_dequant;
pub mod parallel_k;
pub mod simd;
pub mod types;

// Re-export types from submodules (PMAT-802)
pub use types::{
    detect_simd_backend, DequantStats, Q4_0Block, Q4_KBlock, Q5_KBlock, Q6_KBlock, Q8KSuperBlock,
    Q8_0Block, SimdBackend, BLOCK_SIZE, QK_K,
};

// Re-export dequantization functions (PMAT-802)
pub use dequant::{
    dequantize_f16, dequantize_q2_k, dequantize_q3_k, dequantize_q4_0, dequantize_q4_1,
    dequantize_q4_k, dequantize_q5_0, dequantize_q5_1, dequantize_q5_k, dequantize_q6_k,
    dequantize_q8_0, f16_to_f32,
};

// Re-export fused K-quant operations (PMAT-802)
pub mod direct_f32;
pub use direct_f32::fused_q4k_parallel_matvec_f32_into;
pub use fused_k::{fused_q4k_dot, fused_q4k_dot_simd, fused_q4k_q8k_dot, fused_q4k_q8k_dot_simd};
pub use fused_q5k_q6k::{
    fused_q4k_q8_dot, fused_q5k_dot, fused_q5k_dot_simd, fused_q6k_dot, fused_q6k_dot_simd,
};

// Re-export parallel K-quant operations (PMAT-802)
// LAYOUT-002: All kernels are ROW-MAJOR. No colmajor/auto aliases.
pub use parallel_k::{
    fused_q4k_parallel_matvec, fused_q4k_parallel_matvec_into, fused_q4k_q8k_ffn_up_gate_into,
    fused_q4k_q8k_parallel_matvec_into, fused_q4k_tiled_matvec, fused_q5k_parallel_matvec,
    fused_q5k_parallel_matvec_into, fused_q6k_parallel_matvec, fused_q6k_parallel_matvec_into,
};

// Re-export activation functions (PMAT-802)
pub use activation::{
    fused_rmsnorm_ffn_up_gate, fused_rmsnorm_q4_0_matmul, fused_swiglu_simd,
    quantize_activations_q8_0, quantize_rmsnorm_q8_0, quantize_rmsnorm_q8_0_into, softmax_simd,
};

// Re-export parallel dequant operations (PMAT-802)
pub use parallel_dequant::{
    apply_rope_rotation_simd, dequantize_q4_k_parallel, dequantize_q4_k_simd,
    dequantize_q8_0_parallel, dequantize_q8_0_simd,
};

// Re-export SIMD utilities (for tests and internal use)
pub use simd::{extract_scale_min, read_f16};

// Re-export format trait and generic kernels (Contract: quantized-dot-product-v1.yaml)
pub use format_trait::{Q4_0Fmt, Q8_0Fmt, QuantBlockFormat, QuantFamily, Q4K, Q5K, Q6K};
pub use generic_dot::{compute_bsums, generic_fused_dot_scalar};
pub use generic_matvec::{generic_parallel_matvec, generic_parallel_matvec_into};

// Re-export fused gate+up kernel (PMAT-FFN-FUSION)
pub use fused_gate_up::{
    fused_gate_up_q4k_into, fused_gate_up_q5k_into, fused_gate_up_q6k_into,
    generic_fused_gate_up_matvec_into,
};

// Re-export bsum precomputation (Contract: quantized-dot-product-v1.yaml, Step 3)
pub use bsum_precompute::{fused_q4k_q8k_parallel_matvec_with_bsums_into, precompute_q8k_bsums};

// Re-export encoding functions (Toyota Way: ONE source of truth)
// aprender imports these for format conversion - NEVER duplicates
pub use encode::{
    dequantize_q4_k_to_f32,
    dequantize_q5_k_to_f32,
    dequantize_q6_k_to_f32,
    // Q4_K
    quantize_q4_k,
    quantize_q4_k_matrix,
    // Q5_K
    quantize_q5_k,
    quantize_q5_k_matrix,
    // Q6_K
    quantize_q6_k,
    quantize_q6_k_matrix,
    // Transpose (LAYOUT-002)
    transpose_q4k_for_matmul,
    transpose_q5k_for_matmul,
    transpose_q6k_for_matmul,
    // Constants
    F16_MIN_NORMAL,
};

/// Pre-computed f16 to f32 lookup table (65536 entries = 256KB)
///
/// Eliminates per-block f16 conversion overhead in hot paths.
/// Per spec §4.1: f16 scale LUT should provide ~1.1x throughput improvement.
///
/// # Safety
/// The table is initialized once on first access and is immutable thereafter.
static F16_TO_F32_LUT: std::sync::LazyLock<Box<[f32; 65536]>> = std::sync::LazyLock::new(|| {
    let mut lut = Box::new([0.0f32; 65536]);
    for i in 0..65536u32 {
        lut[i as usize] = half::f16::from_bits(i as u16).to_f32();
    }
    lut
});

/// Fast f16 to f32 conversion using pre-computed LUT
///
/// Takes raw u16 bits (little-endian) and returns f32 value.
/// A direct table lookup, in place of the bit manipulation in
/// half::f16::from_bits().to_f32(). No receipt measures the difference on this
/// tree, so no speed factor is claimed for it.
#[inline]
pub(crate) fn f16_to_f32_lut(bits: u16) -> f32 {
    F16_TO_F32_LUT[bits as usize]
}

// BLOCK_SIZE, QK_K, Q4_0Block, Q8_0Block, Q8KSuperBlock moved to types.rs (PMAT-802)

/// Quantize f32 activations to Q8_K super-blocks (zero-allocation variant)
///
/// Pre-allocates output buffers for scales and quantized values.
/// Used for amortized quantization in hot inference path.
///
/// # Arguments
/// * `activations` - Input f32 values (must be multiple of 256)
/// * `scales` - Output scales buffer (len = activations.len() / 256)
/// * `quants` - Output int8 buffer (len = activations.len())
///
/// # Errors
/// Returns error if length is not a multiple of 256
pub fn quantize_activations_q8k_into(
    activations: &[f32],
    scales: &mut [f32],
    quants: &mut [i8],
) -> Result<()> {
    if !activations.len().is_multiple_of(256) {
        return Err(RealizarError::FormatError {
            reason: format!(
                "Q8_K quantization requires length multiple of 256, got {}",
                activations.len()
            ),
        });
    }

    let num_superblocks = activations.len() / 256;

    if scales.len() < num_superblocks {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Scales buffer too small: need {}, have {}",
                num_superblocks,
                scales.len()
            ),
        });
    }

    if quants.len() < activations.len() {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Quants buffer too small: need {}, have {}",
                activations.len(),
                quants.len()
            ),
        });
    }

    for (sb_idx, chunk) in activations.as_chunks::<256>().0.iter().enumerate() {
        Q8KSuperBlock::quantize_into(
            chunk,
            &mut scales[sb_idx],
            &mut quants[sb_idx * 256..(sb_idx + 1) * 256],
        );
    }

    Ok(())
}

/// Quantize a slice of f32 values to Q8_0 blocks
///
/// # Arguments
/// * `values` - F32 values (must be multiple of 32 in length)
///
/// # Returns
/// Vector of Q8_0Block, one per 32 values
///
/// # Errors
/// Returns error if length is not a multiple of 32
pub fn quantize_to_q8_blocks(values: &[f32]) -> Result<Vec<Q8_0Block>> {
    if !values.len().is_multiple_of(32) {
        return Err(RealizarError::FormatError {
            reason: format!(
                "Q8_0 quantization requires length multiple of 32, got {}",
                values.len()
            ),
        });
    }

    let blocks: Vec<Q8_0Block> = values
        .as_chunks::<32>()
        .0
        .iter()
        // `as_chunks::<32>()` yields `&[f32; 32]`, so the width is a type and
        // the fallible conversion is gone.
        .map(Q8_0Block::quantize)
        .collect();

    Ok(blocks)
}

/// Dequantize Q8_0 blocks back to f32 values
pub fn dequantize_q8_blocks(blocks: &[Q8_0Block]) -> Vec<f32> {
    let mut output = Vec::with_capacity(blocks.len() * 32);
    for block in blocks {
        output.extend_from_slice(&block.dequantize());
    }
    output
}

// Q4_KBlock, Q5_KBlock, Q6_KBlock moved to types.rs (PMAT-802)

/// PMAT-PERF-002: Pre-interleaved Q4_K weights for SIMD-friendly access
///
/// Weights reordered at load time to eliminate gather operations during inference.
/// This enables contiguous SIMD loads for Q4_K GEMV instead of scattered nibble
/// extraction. The size of that win is unmeasured here; the op counts under
/// `# Performance` below are structural, not benchmarked.
///
/// # Layout
///
/// Original Q4_K layout (training-friendly):
/// ```text
/// Super-block: [d, dmin, scales[12], qs[128]]
/// qs layout: byte[i] contains value[2i] in low nibble, value[2i+1] in high nibble
/// ```
///
/// Interleaved layout (inference-friendly):
/// ```text
/// Super-block: [d, dmin, scales[12], qs_interleaved[128]]
/// qs_interleaved: values reordered for 32-byte aligned SIMD loads
/// After AVX2 256-bit load + nibble extraction, values are in processing order
/// ```
///
/// # Performance
///
/// - Before: Nibble extraction requires shift/mask per byte (32 ops for 64 values)
/// - After: Single SIMD load gets 32 contiguous values (1 op for 32 values)
/// - Expected speedup: 2-4x for GEMV kernel
///
/// # References
///
/// - Intel AVX-512 Guide: contiguous loads vs VPGATHERDD (vendor guidance, not
///   a measurement of this code)
/// - llama.cpp: Pre-interleaved layout in ggml-quants.c
/// - CUTLASS: Tile-based weight layout for tensor cores
#[derive(Debug, Clone)]
pub struct InterleavedQ4K {
    /// Super-block scales (one per super-block, f32 from f16)
    pub d: Vec<f32>,
    /// Super-block mins (one per super-block, f32 from f16)
    pub dmin: Vec<f32>,
    /// Block scales (12 bytes per super-block, 6-bit packed)
    pub scales: Vec<u8>,
    /// Interleaved 4-bit quantized values
    /// Reordered so SIMD loads get contiguous values without gather
    pub qs: Vec<u8>,
    /// Number of super-blocks
    pub num_super_blocks: usize,
}

include!("product.rs");
include!("q4_0.rs");
include!("fused_q4_0_q8_0.rs");
include!("fused_q8_0_q8_0.rs");

// ─────────────────────────────────────────────────────────────────────────────
// L0-1b (#2971, PMAT-1070): crushed Q8_K blocks.
//
// Q8_K quantises activations with ONE scale per 256 elements (`max/127`). On a
// massive-activation token (Qwen2.5-1.5B, position 0: dim 408 = −146.7 while
// its 255 block-mates are ≤ 7.3) the scale is set by the one element and the
// others collapse onto a handful of int8 levels — the layer-26 gate/up outputs
// came out 13 % low on CPU while the GPU was within 0.5 % of the float64 truth
// (docs/audits/l0-1b-arms.md). The remedy is per matmul: when the activation
// vector carries a crushed block, the Q4_K drivers run the f32-activation dot
// (`fused_q4k_dot_simd`, the `DIRECT_FP32_GEMV` path) for that call.
// ─────────────────────────────────────────────────────────────────────────────

/// A 256-block is crushed when `max|x| / second-largest|x| >= CRUSHED_BLOCK_RATIO`.
///
/// basis: docs/audits/l0-1b-arms.md §Fallback criterion — on the 1.5B over 77
/// ordinary positions × 28 layers the normed residual-stream inputs never
/// exceed 6.0 (p99.9 4.7); the first token's crushed blocks are ≥ 20.
/// `max|x|/rms` cannot serve: it saturates at 16 (= √256) on ordinary blocks.
pub const CRUSHED_BLOCK_RATIO: f32 = 8.0;

/// True when any 256-element block of `activations` is crushed (see
/// [`CRUSHED_BLOCK_RATIO`]). A one-hot block is crushed; an all-zero block is not.
#[must_use]
pub fn has_crushed_block(activations: &[f32]) -> bool {
    activations.chunks(256).any(block_is_crushed)
}

fn block_is_crushed(block: &[f32]) -> bool {
    let (mut max, mut second) = (0.0f32, 0.0f32);
    for &v in block {
        let a = v.abs();
        if a > max {
            second = max;
            max = a;
        } else if a > second {
            second = a;
        }
    }
    if max == 0.0 {
        return false;
    }
    second == 0.0 || max / second >= CRUSHED_BLOCK_RATIO
}

/// Diagnostic: with `APR_CRUSHED_TRACE=1` every fallback prints one line
/// (`[crushed-block] in_dim=… out_dim=…`), so the fallback RATE on a prompt is
/// `grep -c` over stderr — the number the receipt's speed basis cites. Read once.
pub fn note_crushed_fallback(in_dim: usize, out_dim: usize) {
    static TRACE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    if *TRACE.get_or_init(|| std::env::var("APR_CRUSHED_TRACE").as_deref() == Ok("1")) {
        eprintln!("[crushed-block] in_dim={in_dim} out_dim={out_dim}");
    }
}

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

    fn block_with(max: f32, second: f32, fill: f32) -> Vec<f32> {
        let mut b = vec![fill; 256];
        b[408 % 256] = -max;
        b[17] = second;
        b
    }

    #[test]
    fn the_layer_26_block_is_crushed() {
        // pos 0, layer 26, ffn_norm block 1: max 146.68, second 7.28 → ratio 20.1
        assert!(has_crushed_block(&block_with(146.68, 7.28, 0.5)));
    }

    #[test]
    fn an_ordinary_block_is_not() {
        // the worst ordinary block measured: ratio 6.0
        assert!(!has_crushed_block(&block_with(6.0, 1.0, 0.3)));
    }

    #[test]
    fn a_one_hot_block_is_crushed_and_an_all_zero_block_is_not() {
        let mut one_hot = vec![0.0f32; 256];
        one_hot[5] = 3.0;
        assert!(has_crushed_block(&one_hot));
        assert!(!has_crushed_block(&vec![0.0f32; 256]));
    }

    #[test]
    fn only_the_crushed_block_counts_and_the_tail_is_a_block_too() {
        let mut x = block_with(6.0, 1.0, 0.3); // ordinary
        x.extend(block_with(9.0, 1.0, 0.3)); // ratio 9 → crushed
        assert!(has_crushed_block(&x));
        let mut tail = vec![0.3f32; 256];
        tail.extend([0.0f32; 10]); // a 10-element tail block, all zero
        assert!(!has_crushed_block(&tail));
    }
}