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]>
299 lines
11 KiB
Rust
299 lines
11 KiB
Rust
//! INT8 Matrix Multiplication with i32 Accumulation
|
||
//!
|
||
//! Provides a reusable CPU reference implementation of INT8 GEMM for use by
|
||
//! `SmoothQuantizedLayer::forward_raw` and as a future wiring point for GPU kernels.
|
||
//!
|
||
//! # INT8 Range
|
||
//!
|
||
//! This module uses the **symmetric** INT8 range **−127…127**, matching SmoothQuant's
|
||
//! convention. The full INT8 range −128…127 is intentionally avoided so that the
|
||
//! absolute maximum value is the same on both sides of zero (127), which simplifies
|
||
//! per-tensor scale computation as `scale = max_abs / 127.0`.
|
||
//!
|
||
//! # Accumulation
|
||
//!
|
||
//! All inner products accumulate into `i32` to avoid overflow. The maximum
|
||
//! theoretical dot-product magnitude for vectors of length N with all elements
|
||
//! at ±127 is `127 * 127 * N = 16129 * N`. An `i32` saturates at ~2.1 × 10⁹,
|
||
//! so it handles `N ≤ 131_072` safely — well beyond any realistic hidden dimension.
|
||
//!
|
||
//! # Dequantisation Formula
|
||
//!
|
||
//! Given:
|
||
//! - `acc_i32` : INT8 dot product accumulated into i32
|
||
//! - `act_scale` : f32 scale derived as `max(|smoothed activations|) / 127.0`
|
||
//! - `weight_scale` : f32 scale stored in `QuantizedTensorData::scale`
|
||
//! - `zp` : i8 zero-point stored in `QuantizedTensorData::zero_point`
|
||
//!
|
||
//! The output element is:
|
||
//!
|
||
//! ```text
|
||
//! output = (sum_k act_q[k] * (weight_q[k] - zp)) * act_scale * weight_scale
|
||
//! ```
|
||
//!
|
||
//! where the inner sum was accumulated in i32 before being cast to f32.
|
||
|
||
use crate::{
|
||
Result,
|
||
error::{CompressionError, QuantizationError},
|
||
};
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// INT8 matrix-vector product with i32 accumulation.
|
||
///
|
||
/// Computes `output[out_row] = Σ_k weights[out_row * in_features + k] * activations[k]`
|
||
/// for all output rows, accumulating into i32 to avoid overflow.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `weights` – Row-major INT8 weight matrix, shape `[out_features, in_features]`.
|
||
/// * `activations` – INT8 activation vector, length `in_features`.
|
||
/// * `out_features` – Number of output rows.
|
||
/// * `in_features` – Number of input columns.
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// `Vec<i32>` of length `out_features`.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns an error if the slice lengths are inconsistent with the declared dimensions.
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```rust
|
||
/// use rtx_compress::quantization::int8_matmul::int8_matvec;
|
||
///
|
||
/// // Identity-like: W = [[1,0],[0,1]], act = [3,5]
|
||
/// let weights = vec![1i8, 0, 0, 1];
|
||
/// let act = vec![3i8, 5];
|
||
/// let out = int8_matvec(&weights, &act, 2, 2).unwrap();
|
||
/// assert_eq!(out, vec![3i32, 5i32]);
|
||
/// ```
|
||
pub fn int8_matvec(
|
||
weights: &[i8],
|
||
activations: &[i8],
|
||
out_features: usize,
|
||
in_features: usize,
|
||
) -> Result<Vec<i32>> {
|
||
if weights.len() != out_features * in_features {
|
||
return Err(CompressionError::Quantization(
|
||
QuantizationError::InvalidConfig(format!(
|
||
"int8_matvec: weight slice length {} does not match \
|
||
out_features={} * in_features={}",
|
||
weights.len(),
|
||
out_features,
|
||
in_features
|
||
)),
|
||
));
|
||
}
|
||
if activations.len() != in_features {
|
||
return Err(CompressionError::Quantization(
|
||
QuantizationError::InvalidConfig(format!(
|
||
"int8_matvec: activation length {} does not match in_features={}",
|
||
activations.len(),
|
||
in_features
|
||
)),
|
||
));
|
||
}
|
||
|
||
let mut output = vec![0i32; out_features];
|
||
for out_row in 0..out_features {
|
||
let row_offset = out_row * in_features;
|
||
let mut acc = 0i32;
|
||
for k in 0..in_features {
|
||
acc += (weights[row_offset + k] as i32) * (activations[k] as i32);
|
||
}
|
||
output[out_row] = acc;
|
||
}
|
||
Ok(output)
|
||
}
|
||
|
||
/// Batched INT8 GEMM: `weights [out × in]` × `activations [batch × in]` → `[batch × out]`.
|
||
///
|
||
/// Output is row-major `[batch_size, out_features]` with i32 elements.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `weights` – Row-major INT8 weight matrix, shape `[out_features, in_features]`.
|
||
/// * `activations` – Row-major INT8 activation matrix, shape `[batch_size, in_features]`.
|
||
/// * `batch_size` – Number of input rows.
|
||
/// * `out_features` – Number of output rows.
|
||
/// * `in_features` – Number of input columns (must match both slices).
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Row-major `Vec<i32>` of shape `[batch_size, out_features]`.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns an error if slice lengths are inconsistent with declared dimensions.
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```rust
|
||
/// use rtx_compress::quantization::int8_matmul::int8_gemm;
|
||
///
|
||
/// // W = [[1,2],[3,4]], acts = [[1,0],[0,1]] (batch=2)
|
||
/// let weights = vec![1i8, 2, 3, 4];
|
||
/// let acts = vec![1i8, 0, 0, 1];
|
||
/// let out = int8_gemm(&weights, &acts, 2, 2, 2).unwrap();
|
||
/// // batch 0: [1*1+2*0, 3*1+4*0] = [1, 3]
|
||
/// // batch 1: [1*0+2*1, 3*0+4*1] = [2, 4]
|
||
/// assert_eq!(out, vec![1i32, 3, 2, 4]);
|
||
/// ```
|
||
pub fn int8_gemm(
|
||
weights: &[i8],
|
||
activations: &[i8],
|
||
batch_size: usize,
|
||
out_features: usize,
|
||
in_features: usize,
|
||
) -> Result<Vec<i32>> {
|
||
if weights.len() != out_features * in_features {
|
||
return Err(CompressionError::Quantization(
|
||
QuantizationError::InvalidConfig(format!(
|
||
"int8_gemm: weight slice length {} does not match \
|
||
out_features={} * in_features={}",
|
||
weights.len(),
|
||
out_features,
|
||
in_features
|
||
)),
|
||
));
|
||
}
|
||
if activations.len() != batch_size * in_features {
|
||
return Err(CompressionError::Quantization(
|
||
QuantizationError::InvalidConfig(format!(
|
||
"int8_gemm: activation slice length {} does not match \
|
||
batch_size={} * in_features={}",
|
||
activations.len(),
|
||
batch_size,
|
||
in_features
|
||
)),
|
||
));
|
||
}
|
||
|
||
let mut output = vec![0i32; batch_size * out_features];
|
||
for b in 0..batch_size {
|
||
let act_offset = b * in_features;
|
||
for out_row in 0..out_features {
|
||
let weight_offset = out_row * in_features;
|
||
let mut acc = 0i32;
|
||
for k in 0..in_features {
|
||
acc += (weights[weight_offset + k] as i32) * (activations[act_offset + k] as i32);
|
||
}
|
||
output[b * out_features + out_row] = acc;
|
||
}
|
||
}
|
||
Ok(output)
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Tests
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
// -----------------------------------------------------------------------
|
||
// int8_matvec tests
|
||
// -----------------------------------------------------------------------
|
||
|
||
/// Weight matrix = identity → output equals activation vector.
|
||
#[test]
|
||
fn test_int8_matvec_identity_weights() {
|
||
// W = [[1,0,0],[0,1,0],[0,0,1]] (3×3 identity)
|
||
let weights: Vec<i8> = vec![1, 0, 0, 0, 1, 0, 0, 0, 1];
|
||
let activations: Vec<i8> = vec![7, -3, 5];
|
||
let out = int8_matvec(&weights, &activations, 3, 3).unwrap();
|
||
assert_eq!(out, vec![7i32, -3, 5]);
|
||
}
|
||
|
||
/// All-zero activations produce a zero output regardless of weights.
|
||
#[test]
|
||
fn test_int8_matvec_zero_activation() {
|
||
let weights: Vec<i8> = vec![1, 2, 3, 4, 5, 6];
|
||
let activations: Vec<i8> = vec![0, 0, 0];
|
||
let out = int8_matvec(&weights, &activations, 2, 3).unwrap();
|
||
assert_eq!(out, vec![0i32, 0]);
|
||
}
|
||
|
||
/// Shape mismatch between declared in_features and activation slice returns Err.
|
||
#[test]
|
||
fn test_int8_matvec_shape_mismatch() {
|
||
// Weight matrix is 2×4 (8 elements), but activations has 3 elements, not 4.
|
||
let weights: Vec<i8> = vec![1, 2, 3, 4, 5, 6, 7, 8];
|
||
let activations: Vec<i8> = vec![1, 2, 3]; // wrong length
|
||
let result = int8_matvec(&weights, &activations, 2, 4);
|
||
assert!(
|
||
result.is_err(),
|
||
"expected Err on in_features mismatch, got Ok"
|
||
);
|
||
}
|
||
|
||
// -----------------------------------------------------------------------
|
||
// int8_gemm tests
|
||
// -----------------------------------------------------------------------
|
||
|
||
/// batch_size=2 with distinct activation rows produces independent outputs.
|
||
#[test]
|
||
fn test_int8_gemm_batch_size_2() {
|
||
// W = [[1,1],[1,1]] (all-ones 2×2), batch of 2: act0=[2,3], act1=[10,-1]
|
||
let weights: Vec<i8> = vec![1, 1, 1, 1];
|
||
// Row-major [batch=2, in=2]: [act0_col0, act0_col1, act1_col0, act1_col1]
|
||
let activations: Vec<i8> = vec![2, 3, 10, -1];
|
||
let out = int8_gemm(&weights, &activations, 2, 2, 2).unwrap();
|
||
// batch 0, out_row 0: 1*2 + 1*3 = 5
|
||
// batch 0, out_row 1: 1*2 + 1*3 = 5
|
||
// batch 1, out_row 0: 1*10 + 1*(-1) = 9
|
||
// batch 1, out_row 1: 1*10 + 1*(-1) = 9
|
||
assert_eq!(out, vec![5i32, 5, 9, 9]);
|
||
}
|
||
|
||
/// Known-value accumulation check: W=[[1,2],[3,4]], acts=[[1,0],[0,1]].
|
||
///
|
||
/// Output layout is [batch, out_features]:
|
||
/// batch 0: [1*1+2*0, 3*1+4*0] = [1, 3]
|
||
/// batch 1: [1*0+2*1, 3*0+4*1] = [2, 4]
|
||
#[test]
|
||
fn test_int8_gemm_accumulation_correct() {
|
||
let weights: Vec<i8> = vec![1, 2, 3, 4]; // [out=2, in=2]
|
||
let activations: Vec<i8> = vec![1, 0, 0, 1]; // [batch=2, in=2]
|
||
let out = int8_gemm(&weights, &activations, 2, 2, 2).unwrap();
|
||
assert_eq!(out, vec![1i32, 3, 2, 4]);
|
||
}
|
||
|
||
/// Max INT8 values (±127) accumulate into i32 without overflow.
|
||
///
|
||
/// For in_features=256, the maximum dot-product is 127 * 127 * 256 = 4,128,256,
|
||
/// well within i32::MAX (2,147,483,647).
|
||
#[test]
|
||
fn test_int8_overflow_clamped() {
|
||
let n = 256usize;
|
||
// Both weights and activations all at +127.
|
||
let weights: Vec<i8> = vec![127i8; n];
|
||
let activations: Vec<i8> = vec![127i8; n];
|
||
let out = int8_matvec(&weights, &activations, 1, n).unwrap();
|
||
let expected = 127i32 * 127 * n as i32; // = 4_128_256
|
||
assert_eq!(out[0], expected);
|
||
// Confirm no silent truncation to i8.
|
||
assert!(out[0] > i8::MAX as i32, "accumulation stayed in i32");
|
||
}
|
||
|
||
/// in_features mismatch between weight and activation slice lengths returns Err.
|
||
#[test]
|
||
fn test_int8_matvec_weight_shape_mismatch() {
|
||
// Declare out=2, in=4 but supply only 6 weight elements (≠ 8).
|
||
let weights: Vec<i8> = vec![1, 2, 3, 4, 5, 6];
|
||
let activations: Vec<i8> = vec![1, 2, 3, 4];
|
||
let result = int8_matvec(&weights, &activations, 2, 4);
|
||
assert!(
|
||
result.is_err(),
|
||
"expected Err on weight slice length mismatch"
|
||
);
|
||
}
|
||
}
|