osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
700 lines
26 KiB
Rust
700 lines
26 KiB
Rust
//! 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]);
|
||
}
|
||
}
|