feat(perf): GPU perf batch 3 — W4A16 AWQ matmul, FSDP2 hooks, RMSNorm+SwiGLU fused kernel

W4A16 AWQ on-the-fly dequantize+GEMM (rtx-compress)
- `w4a16_matmul.rs`: `AWQQuantizedWeightExt` trait + `matmul_cpu()` — group-aligned
  inner loop, f64 accumulation, low-nibble-first INT4 unpacking matching mx_kernels.cu
- `cuda_kernels/w4a16_gemm.cu`: `w4a16_dequant_gemm` kernel, one thread per (batch, out_col),
  8-INT4-per-iteration inner loop with `__ldg()` cache hints, BF16 scale decode, f32 accumulate
- `quantization/mod.rs`: exports `w4a16_matmul_cpu`, `AWQQuantizedWeightExt`
- Fixed 2 pre-existing pruning compile errors
- 12 tests: nibble unpack, identity weights, shape, vs-dequant (tol=1e-3), batch=1, zeros

FSDP2 forward/backward hooks (rtx-distributed)
- `fsdp2.rs`: `update_local_shard()` on `Fsdp2ShardedParam`; sync `all_gather()` +
  `reduce_scatter_gradient()` using `ProcessGroup::{all_gather,reduce_scatter}`
- `pre_forward_hook()` — all-gathers every param (or copies shard in single-process)
- `post_backward_hook()` — reduce-scatters gradients, zero_grad, re-shards cache
- `step(optimizer_fn)` — applies optimizer closure to each local shard
- `make_fsdp2_module()` top-level factory; `Fsdp2MemoryStats` gains 5 new fields
  incl. `memory_reduction_ratio ≈ world_size`
- 6 new tests (end-to-end training step included); total 444 pass

RMSNorm+SwiGLU fused CUDA kernel (rtx-fusion)
- `cuda/rms_norm_swiglu_fused.cu`: `rms_norm_kernel` + `rms_norm_swiglu_fused`;
  shared-mem warp reduction (block_x floats), launch: grid=(batch,1,1), block=(min(hidden,1024),1,1)
- `cuda_kernels/rms_norm_fused.rs`: CPU reference `rms_norm_cpu`/`swiglu_cpu`/
  `rms_norm_swiglu_cpu`; `#[cfg(feature="cuda")] RmsNormFusedKernel` NVRTC wrapper
- `codegen/cubecl.rs`: replaced RmsNorm comment stub with cfg-gated NVRTC dispatch
- `backend.rs` + `tensor.rs`: added 15 missing `Backend` trait impls (sin/cos/relu/conv2d/…)
  that blocked test compilation
- `Cargo.toml`: added rtx-fusion to workspace members
- 8 new tests (PyTorch-formula verified: x=[1,2,3,4] → [0.365, 0.730, 1.095, 1.461]);
  total 103 pass

Test results: 12 + 444 + 103 = 559 tests, 0 failures

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

crates/training/rtx-compress/src/quantization/w4a16_matmul.rs

@@ -0,0 +1,707 @@
+//! 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]);
+    }
+}