//! W4A16 Weight-Only Quantized Matrix Multiplication
//!
//! This module implements fused dequantize+GEMM for AWQ INT4 weight-only quantization.
//! Weights remain in INT4 storage (4-bit precision, 2× memory savings over FP16) while
//! activations are BF16/f32. The kernel dequantizes weight blocks **on-the-fly** during
//! the inner product loop, eliminating the need to materialize a full-precision weight
//! matrix before every forward pass.
//!
//! # Memory Layout
//!
//! The existing [`AWQQuantizedWeight`] stores weights as `Vec<i8>` where each element
//! holds one INT4 value in the low four bits (signed, range −8…7).  Groups of
//! `group_size` consecutive weights (in row-major, flattened order) share a single
//! `f32` scale and `i8` zero-point.
//!
//! Dequantization of a single weight element:
//!
//! ```text
//! w_fp32 = (q_i4 - zero) * scale * act_scale[in_col]
//! ```
//!
//! # Nibble Packing Strategy
//!
//! When manually packing INT4 values into bytes (e.g., for GPU staging), this module
//! uses **low-nibble-first** order:
//!
//! ```text
//! byte = (val[2k] & 0x0F) | ((val[2k+1] & 0x0F) << 4)
//! ```
//!
//! Unpacking:
//! - low nibble:  `byte & 0x0F`         → element at even index
//! - high nibble: `(byte >> 4) & 0x0F`  → element at odd index
//!
//! This matches the convention used by the MX kernel in `cuda_kernels/mx_kernels.cu`.
//!
//! # Example
//!
//! ```rust,ignore
//! use rtx_compress::quantization::w4a16_matmul::w4a16_matmul_cpu;
//!
//! let result = weights.matmul_cpu(&activations, batch_size)?;
//! ```

use crate::{
    Result,
    error::{CompressionError, QuantizationError},
};

use super::advanced::AWQQuantizedWeight;

// ---------------------------------------------------------------------------
// BF16 helpers
// ---------------------------------------------------------------------------

/// Decode a little-endian BF16 byte pair to f32.
///
/// BF16 occupies the high 16 bits of an IEEE 754 f32.  We reconstruct the f32
/// by zero-extending the two bytes into the upper half of a 32-bit word.
#[inline(always)]
fn bf16_bytes_to_f32(lo: u8, hi: u8) -> f32 {
    // Interpret the two bytes as a u16 in little-endian order.
    let bits16 = u16::from_le_bytes([lo, hi]);
    // BF16 → f32: shift the 16-bit pattern into the high half of a u32.
    f32::from_bits((bits16 as u32) << 16)
}

// ---------------------------------------------------------------------------
// Nibble pack / unpack (public for test visibility)
// ---------------------------------------------------------------------------

/// Pack two 4-bit signed values (each in the range −8…7, stored as i8) into a
/// single byte using **low-nibble-first** order.
///
/// # Panics
///
/// Panics in debug builds if either value is outside `−8..=7`.
#[inline(always)]
pub fn pack_nibbles(low: i8, high: i8) -> u8 {
    debug_assert!((-8..=7).contains(&low), "low nibble out of INT4 range");
    debug_assert!((-8..=7).contains(&high), "high nibble out of INT4 range");
    ((low as u8) & 0x0F) | (((high as u8) & 0x0F) << 4)
}

/// Extract the **low** nibble (even-index element) from a packed byte and
/// sign-extend it to i8.
#[inline(always)]
pub fn unpack_low_nibble(byte: u8) -> i8 {
    let nibble = byte & 0x0F;
    // Sign-extend 4-bit → 8-bit: if bit 3 is set the value is negative.
    if nibble & 0x08 != 0 {
        (nibble | 0xF0) as i8
    } else {
        nibble as i8
    }
}

/// Extract the **high** nibble (odd-index element) from a packed byte and
/// sign-extend it to i8.
#[inline(always)]
pub fn unpack_high_nibble(byte: u8) -> i8 {
    let nibble = (byte >> 4) & 0x0F;
    if nibble & 0x08 != 0 {
        (nibble | 0xF0) as i8
    } else {
        nibble as i8
    }
}

// ---------------------------------------------------------------------------
// CPU reference implementation
// ---------------------------------------------------------------------------

/// W4A16 fused dequantize-then-GEMM on the CPU.
///
/// Computes `output[b, out_row] = Σ_i dequant(weights[out_row, i]) * activations[b, i]`
/// without materialising the full-precision weight matrix.
///
/// # Arguments
///
/// * `weights`     – AWQ-quantized weight matrix, shape `[out_features, in_features]`.
/// * `activations` – Flat row-major f32 buffer of shape `[batch_size, in_features]`.
/// * `batch_size`  – Number of input vectors (rows of the activation buffer).
///
/// # Returns
///
/// Flat row-major f32 buffer of shape `[batch_size, out_features]`.
///
/// # Errors
///
/// Returns [`CompressionError`] if the weight shape is invalid or the activation
/// buffer length is inconsistent with `batch_size`.
pub fn w4a16_matmul_cpu(
    weights: &AWQQuantizedWeight,
    activations: &[f32],
    batch_size: usize,
) -> Result<Vec<f32>> {
    // ------------------------------------------------------------------
    // Validate inputs
    // ------------------------------------------------------------------
    if weights.shape.len() != 2 {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "w4a16_matmul_cpu: weight shape must be 2-D, got {} dimensions",
                weights.shape.len()
            )),
        ));
    }

    let out_features = weights.shape[0];
    let in_features = weights.shape[1];

    if in_features == 0 || out_features == 0 {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(
                "w4a16_matmul_cpu: weight dimensions must be non-zero".to_string(),
            ),
        ));
    }

    let expected_act_len = batch_size * in_features;
    if activations.len() != expected_act_len {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "w4a16_matmul_cpu: activation buffer length {} \
                 does not match batch_size={} × in_features={}",
                activations.len(),
                batch_size,
                in_features
            )),
        ));
    }

    let group_size = weights.group_size;
    if group_size == 0 {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(
                "w4a16_matmul_cpu: group_size must be non-zero".to_string(),
            ),
        ));
    }

    // Precompute group count for bounds checking.
    let num_weight_elements = out_features * in_features;
    let num_groups = (num_weight_elements + group_size - 1) / group_size;

    if weights.scales.len() < num_groups {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "w4a16_matmul_cpu: scales vec too short: got {}, need {}",
                weights.scales.len(),
                num_groups
            )),
        ));
    }
    if weights.zeros.len() < num_groups {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "w4a16_matmul_cpu: zeros vec too short: got {}, need {}",
                weights.zeros.len(),
                num_groups
            )),
        ));
    }
    if weights.quantized_data.len() < num_weight_elements {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "w4a16_matmul_cpu: quantized_data too short: got {}, need {}",
                weights.quantized_data.len(),
                num_weight_elements
            )),
        ));
    }

    // ------------------------------------------------------------------
    // Output buffer [batch_size, out_features]
    // ------------------------------------------------------------------
    let mut output = vec![0.0f32; batch_size * out_features];

    // ------------------------------------------------------------------
    // Fused dequant + dot-product
    //
    // Outer loops: (batch_row, out_row).
    // Inner loop:  groups of `group_size` input features.
    //
    // Group accounting:  the existing AWQQuantizer assigns groups linearly
    // across the *entire* flattened weight array (not per-row).  Group index
    // for element at (out_row, in_col) is therefore:
    //
    //   flat_idx   = out_row * in_features + in_col
    //   group_idx  = flat_idx / group_size
    //
    // This matches the indexing in AWQQuantizedWeight::dequantize().
    // ------------------------------------------------------------------
    for b in 0..batch_size {
        let act_row_offset = b * in_features;

        for out_row in 0..out_features {
            let mut accumulator = 0.0f64;
            let weight_row_offset = out_row * in_features;

            // Walk in_features in group-aligned chunks for cache efficiency.
            let mut in_col = 0usize;
            while in_col < in_features {
                // Compute group metadata for the first element of this slice.
                let flat_idx = weight_row_offset + in_col;
                let group_idx = flat_idx / group_size;

                // How many elements remain in this group starting at in_col?
                // The group boundary within the row:
                //   next_group_flat = (group_idx + 1) * group_size
                //   next_group_in_col = next_group_flat - weight_row_offset
                //   (clamped to in_features)
                let next_group_flat = (group_idx + 1) * group_size;
                let next_in_col =
                    (next_group_flat.saturating_sub(weight_row_offset)).min(in_features);

                let scale = weights.scales[group_idx];
                let zero = weights.zeros[group_idx] as f32;

                // Inner loop: all elements sharing the same scale/zero.
                for col in in_col..next_in_col {
                    let q = weights.quantized_data[weight_row_offset + col] as f32;

                    // Per-column activation scale (inverse of the activation-derived
                    // scale applied during AWQ quantisation).
                    let act_scale = weights.activation_scales.get(col).copied().unwrap_or(1.0);

                    let w_fp32 = (q - zero) * scale * act_scale;
                    let a_fp32 = activations[act_row_offset + col];
                    accumulator += (w_fp32 * a_fp32) as f64;
                }

                in_col = next_in_col;
            }

            output[b * out_features + out_row] = accumulator as f32;
        }
    }

    Ok(output)
}

// ---------------------------------------------------------------------------
// Extension trait — attaches matmul_cpu to AWQQuantizedWeight
// ---------------------------------------------------------------------------

/// Extension trait that adds W4A16 fused matmul to [`AWQQuantizedWeight`].
pub trait AWQQuantizedWeightExt {
    /// Fused dequantize + GEMM on the CPU.
    ///
    /// Equivalent to `w4a16_matmul_cpu(self, activations, batch_size)`.
    fn matmul_cpu(&self, activations: &[f32], batch_size: usize) -> Result<Vec<f32>>;
}

impl AWQQuantizedWeightExt for AWQQuantizedWeight {
    fn matmul_cpu(&self, activations: &[f32], batch_size: usize) -> Result<Vec<f32>> {
        w4a16_matmul_cpu(self, activations, batch_size)
    }
}

// ---------------------------------------------------------------------------
// GPU stub (CUDA feature gate)
// ---------------------------------------------------------------------------

/// W4A16 GPU matmul stub — returns `NotImplemented` until the kernel is wired.
///
/// # Planned kernel: `w4a16_dequant_gemm`
///
/// Thread layout:
/// - One warpgroup (4 warps = 128 threads) per output tile of 128 outputs × 64 inputs.
/// - Each warp unrolls the inner loop in groups of 8:
///   - Load 4 bytes (8 nibbles) from `weights_packed` via `__ldg()` (read-only cache).
///   - Unpack 8 INT4 values: `lo[k] = (byte >> 0) & 0x0F`, `hi[k] = (byte >> 4) & 0x0F`,
///     sign-extend each to i8.
///   - Load the corresponding group scale and zero-point once per group boundary.
///   - Dequantize: `w_bf16[k] = __float2bfloat16((q[k] - zero) * scale * act_scale[col])`.
///   - Multiply: `acc_f32 += __bfloat162float(w_bf16[k]) * __bfloat162float(a_bf16[k])`.
///   - Warp-reduce the 128 partial sums with `__reduce_add_sync`.
/// - Write one BF16 output per warpgroup lane-0: `output[b, out] = __float2bfloat16(acc_f32)`.
///
/// Memory traffic analysis at FP16 throughput (RTX 5060 Ti, 448 GB/s):
/// - Weights: `out × in / 2` bytes (INT4 packed) — 2× vs FP16.
/// - Scales/zeros: `out × in / group_size × 2` bytes — negligible at group_size=128.
/// - Activations: `batch × in × 2` bytes (BF16).
/// - Output: `batch × out × 2` bytes (BF16).
///
/// Expected: ~2× bandwidth reduction vs FP16 matmul → ~2× throughput improvement
/// on memory-bound large-model inference.
#[cfg(feature = "cuda")]
pub fn w4a16_matmul_gpu(
    _weights: &AWQQuantizedWeight,
    _activations: &[f32],
    _batch_size: usize,
) -> Result<Vec<f32>> {
    Err(CompressionError::Quantization(
        QuantizationError::InvalidConfig(
            "w4a16 GPU kernel not yet linked — compile src/quantization/cuda/w4a16_gemm.cu \
             and wire the PTX through build.rs"
                .to_string(),
        ),
    ))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quantization::advanced::AWQQuantizedWeight;

    // -----------------------------------------------------------------------
    // Helper: build a minimal AWQQuantizedWeight with uniform scale/zero.
    //
    // Parameters
    // ----------
    // out_features, in_features — matrix dimensions
    // quant_values  — INT4 values (one per element, row-major)
    // scale         — single group scale (we set group_size = total elements so
    //                 there is exactly one group, making assertions simple)
    // zero          — zero-point applied to all elements
    // act_scales    — per-column activation-derived inverse scales
    // -----------------------------------------------------------------------
    fn make_weight(
        out_features: usize,
        in_features: usize,
        quant_values: Vec<i8>,
        scale: f32,
        zero: i8,
        act_scales: Vec<f32>,
    ) -> AWQQuantizedWeight {
        let total = out_features * in_features;
        // One group covering all elements keeps the arithmetic transparent.
        let group_size = total.max(1);
        AWQQuantizedWeight {
            quantized_data: quant_values,
            scales: vec![scale],
            zeros: vec![zero],
            activation_scales: act_scales,
            group_size,
            bit_width: 4,
            shape: vec![out_features, in_features],
        }
    }

    // -----------------------------------------------------------------------
    // Nibble pack / unpack
    // -----------------------------------------------------------------------

    #[test]
    fn test_nibble_unpack_low_nibble() {
        // 0xAB → low nibble = 0xB = 11 (unsigned), but sign-extended 4-bit:
        // 0xB = 0b1011 → bit3 set → negative → 0b1111_1011 = -5 as i8.
        // The value 11 in 4-bit signed is -5.
        let byte = 0xABu8;
        let low = unpack_low_nibble(byte);
        assert_eq!(low, -5i8, "0xB should sign-extend to -5");
    }

    #[test]
    fn test_nibble_unpack_high_nibble() {
        // 0xAB → high nibble = 0xA = 10 (unsigned); 4-bit signed: -6.
        let byte = 0xABu8;
        let high = unpack_high_nibble(byte);
        assert_eq!(high, -6i8, "0xA should sign-extend to -6");
    }

    #[test]
    fn test_nibble_pack_round_trip() {
        // Positive values
        let packed = pack_nibbles(3, 5);
        assert_eq!(unpack_low_nibble(packed), 3);
        assert_eq!(unpack_high_nibble(packed), 5);

        // Negative values
        let packed2 = pack_nibbles(-1, -8);
        assert_eq!(unpack_low_nibble(packed2), -1);
        assert_eq!(unpack_high_nibble(packed2), -8);

        // Zero values
        let packed3 = pack_nibbles(0, 0);
        assert_eq!(unpack_low_nibble(packed3), 0);
        assert_eq!(unpack_high_nibble(packed3), 0);
    }

    // -----------------------------------------------------------------------
    // Identity weights: all dequant to 1.0 → output = row-sum of activations
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_identity_weights() {
        // 1 output row, 4 input features.
        // We want dequant(q) = 1.0 for every weight.
        // Choose: zero = 0, scale = 1.0, q = 1 for all weights.
        // act_scale = 1.0 (identity).
        let out_f = 1usize;
        let in_f = 4usize;
        let w = make_weight(out_f, in_f, vec![1; 4], 1.0, 0, vec![1.0; 4]);

        let activations = vec![2.0f32, 3.0, 4.0, 5.0]; // sum = 14
        let result = w4a16_matmul_cpu(&w, &activations, 1).unwrap();

        assert_eq!(result.len(), 1);
        // output[0] = 1*2 + 1*3 + 1*4 + 1*5 = 14
        assert!((result[0] - 14.0).abs() < 1e-4, "got {}", result[0]);
    }

    // -----------------------------------------------------------------------
    // Output shape check
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_matmul_shape() {
        let out_f = 3usize;
        let in_f = 4usize;
        let total = out_f * in_f;
        let w = make_weight(out_f, in_f, vec![0i8; total], 1.0, 0, vec![1.0; in_f]);
        let activations = vec![1.0f32; 2 * in_f]; // batch_size = 2
        let result = w4a16_matmul_cpu(&w, &activations, 2).unwrap();

        // Expected shape: [batch_size=2, out_features=3] → 6 elements.
        assert_eq!(result.len(), 6, "output length must be batch*out_features");
    }

    // -----------------------------------------------------------------------
    // Zero activations → zero output
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_zero_activations() {
        let out_f = 2usize;
        let in_f = 4usize;
        let total = out_f * in_f;
        let w = make_weight(out_f, in_f, vec![3i8; total], 0.5, 0, vec![1.0; in_f]);
        let activations = vec![0.0f32; 1 * in_f]; // all zero
        let result = w4a16_matmul_cpu(&w, &activations, 1).unwrap();

        for (i, &v) in result.iter().enumerate() {
            assert!(v.abs() < 1e-6, "element {} should be zero, got {}", i, v);
        }
    }

    // -----------------------------------------------------------------------
    // Batch size 1 — single-sample inference
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_batch_size_1() {
        // 2 output rows, 2 input features.
        // W = [[1, 2], [3, 4]] (as INT4), scale=1, zero=0, act_scale=1.
        // act = [1, 1]
        // out[0] = 1*1 + 2*1 = 3
        // out[1] = 3*1 + 4*1 = 7
        let w = make_weight(2, 2, vec![1, 2, 3, 4], 1.0, 0, vec![1.0; 2]);
        let activations = vec![1.0f32, 1.0];
        let result = w4a16_matmul_cpu(&w, &activations, 1).unwrap();

        assert_eq!(result.len(), 2);
        assert!((result[0] - 3.0).abs() < 1e-4, "out[0]={}", result[0]);
        assert!((result[1] - 7.0).abs() < 1e-4, "out[1]={}", result[1]);
    }

    // -----------------------------------------------------------------------
    // Group size 128 (standard AWQ config)
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_group_size_128() {
        // 1 output row, 256 input features = 2 groups of 128.
        // Group 0: scale=2.0, zero=0; q values all = 1
        // Group 1: scale=3.0, zero=0; q values all = 1
        // activations all = 1.0
        // out = sum_{0..127}(1*2*1*1) + sum_{128..255}(1*3*1*1)
        //     = 128*2 + 128*3 = 256 + 384 = 640
        let out_f = 1usize;
        let in_f = 256usize;
        let group_size = 128usize;

        let quant_data = vec![1i8; in_f];
        let scales = vec![2.0f32, 3.0];
        let zeros = vec![0i8; 2];
        let act_scales = vec![1.0f32; in_f];
        let activations = vec![1.0f32; in_f];

        let w = AWQQuantizedWeight {
            quantized_data: quant_data,
            scales,
            zeros,
            activation_scales: act_scales,
            group_size,
            bit_width: 4,
            shape: vec![out_f, in_f],
        };

        let result = w4a16_matmul_cpu(&w, &activations, 1).unwrap();
        assert_eq!(result.len(), 1);
        assert!(
            (result[0] - 640.0).abs() < 1e-2,
            "expected 640, got {}",
            result[0]
        );
    }

    // -----------------------------------------------------------------------
    // w4a16_matmul_cpu vs dequantize().dot(act) — key correctness check
    //
    // We use a small, randomly-seeded matrix so the test is deterministic
    // and easy to reason about, then compare both paths within 1e-3 tolerance.
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_vs_dequant_then_matmul() {
        use rtx_tensor::{Device, Tensor};

        // 4 output rows, 8 input features, group_size = 4 (2 groups per row).
        let out_f = 4usize;
        let in_f = 8usize;
        let group_size = 4usize;
        let total = out_f * in_f;
        let num_groups = total / group_size; // = 8

        // Deterministic quant values in the INT4 range −8…7.
        // We cycle through a small sequence to get varied values.
        let quant_pattern: Vec<i8> = [1i8, -2, 3, -4, 5, -1, 2, 0]
            .iter()
            .cycle()
            .take(total)
            .copied()
            .collect();

        // Two distinct scales and zero-points to exercise non-trivial groups.
        let mut scales = Vec::with_capacity(num_groups);
        let mut zeros = Vec::with_capacity(num_groups);
        for g in 0..num_groups {
            scales.push(0.5 + 0.1 * (g as f32));
            zeros.push(if g % 2 == 0 { 1i8 } else { -1i8 });
        }

        // Uniform activation scales for simplicity (no per-column correction).
        let act_scales = vec![1.0f32; in_f];

        let w = AWQQuantizedWeight {
            quantized_data: quant_pattern,
            scales: scales.clone(),
            zeros: zeros.clone(),
            activation_scales: act_scales.clone(),
            group_size,
            bit_width: 4,
            shape: vec![out_f, in_f],
        };

        // Activations: [batch=2, in_f=8]
        let batch = 2usize;
        let activations: Vec<f32> = (0..batch * in_f).map(|i| (i as f32) * 0.25).collect();

        // --- Path A: fused w4a16_matmul_cpu ---
        let fused_result = w4a16_matmul_cpu(&w, &activations, batch).unwrap();

        // --- Path B: dequantize() then explicit matmul ---
        let dequant_tensor: Tensor = w.dequantize().unwrap();
        let dequant_data: Vec<f32> = dequant_tensor.to_vec().unwrap();
        // dequant_data is [out_f, in_f] row-major.
        let mut ref_result = vec![0.0f32; batch * out_f];
        for b in 0..batch {
            for o in 0..out_f {
                let mut acc = 0.0f32;
                for i in 0..in_f {
                    acc += dequant_data[o * in_f + i] * activations[b * in_f + i];
                }
                ref_result[b * out_f + o] = acc;
            }
        }

        // Compare: tolerance 1e-3 (both paths use f32; small rounding differences are fine).
        let tolerance = 1e-3f32;
        assert_eq!(
            fused_result.len(),
            ref_result.len(),
            "output length mismatch"
        );
        for (idx, (&fused, &reference)) in fused_result.iter().zip(ref_result.iter()).enumerate() {
            let diff = (fused - reference).abs();
            assert!(
                diff <= tolerance,
                "element {}: fused={:.6}, reference={:.6}, diff={:.6} > tolerance={:.6}",
                idx,
                fused,
                reference,
                diff,
                tolerance
            );
        }
    }

    // -----------------------------------------------------------------------
    // Batch size 4 — multi-sample
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_batch_size_4() {
        // Simple 2×2 weight matrix, batch=4.
        // W = [[2, 0], [0, 2]], scale=1, zero=0 → dequant = [[2,0],[0,2]] (identity*2)
        // activations[b] = [b, b]
        // out[b, 0] = 2*b + 0*b = 2b
        // out[b, 1] = 0*b + 2*b = 2b
        let w = make_weight(2, 2, vec![2, 0, 0, 2], 1.0, 0, vec![1.0; 2]);
        let activations: Vec<f32> = (0..4).flat_map(|b| vec![b as f32, b as f32]).collect();

        let result = w4a16_matmul_cpu(&w, &activations, 4).unwrap();
        assert_eq!(result.len(), 8);
        for b in 0..4usize {
            let expected = 2.0 * b as f32;
            assert!(
                (result[b * 2] - expected).abs() < 1e-4,
                "b={} out[0]: got {}, want {}",
                b,
                result[b * 2],
                expected
            );
            assert!(
                (result[b * 2 + 1] - expected).abs() < 1e-4,
                "b={} out[1]: got {}, want {}",
                b,
                result[b * 2 + 1],
                expected
            );
        }
    }

    // -----------------------------------------------------------------------
    // Error: activation buffer length mismatch
    // -----------------------------------------------------------------------

    #[test]
    fn test_w4a16_activation_length_error() {
        let w = make_weight(2, 4, vec![0i8; 8], 1.0, 0, vec![1.0; 4]);
        // Provide only 3 elements instead of batch=1 * in_f=4 = 4.
        let activations = vec![1.0f32; 3];
        let result = w4a16_matmul_cpu(&w, &activations, 1);
        assert!(result.is_err(), "should fail on activation length mismatch");
    }

    // -----------------------------------------------------------------------
    // AWQQuantizedWeightExt trait method
    // -----------------------------------------------------------------------

    #[test]
    fn test_awq_ext_matmul_cpu() {
        let w = make_weight(1, 2, vec![1, 1], 1.0, 0, vec![1.0; 2]);
        let acts = vec![3.0f32, 4.0];
        // Use extension trait method.
        let result = w.matmul_cpu(&acts, 1).unwrap();
        assert_eq!(result.len(), 1);
        // 1*3 + 1*4 = 7
        assert!((result[0] - 7.0).abs() < 1e-4, "got {}", result[0]);
    }
}