feat(perf): GPU perf batch 4 — SmoothQuant INT8 forward, varlen FA, inference graph capture

SmoothQuant INT8 linear forward (rtx-compress)
- `advanced.rs`: `SmoothQuantizedLayer::forward_raw()` — smooth activations ÷ scales,
  INT8-quantize both sides (range −127…127 symmetric), INT8 GEMM w/ i32 accumulation,
  dequantize: acc * act_scale * weight_scale; `out_features()` / `in_features()` helpers
- `int8_matmul.rs`: `int8_matvec` + `int8_gemm` (i32 accumulation); 7 unit tests
- 4 forward_raw tests: shape, identity layer, scale effect, manual verification
- 11 new tests; total 112 pass

Variable-length packed flash attention (rtx-flash-attention) [commit 80d7c7f]
- `flash_attention_varlen.cu`: WGMMA-compatible CUDA kernel, BLOCK_Q/K=64, block=(128,1,1),
  grid=(ceil(max_seqlen/64), heads, 1); linear cu_seqlens scan for sequence-to-block
  mapping; 24KB shared memory (3 × 64 × 64 × 2 bytes); early-exit for past-end blocks
- `flash_varlen_forward.rs`: `varlen_attention_cpu` O(n²) reference + `#[cfg(cuda)]`
  `FlashVarlenKernel` NVRTC wrapper; `SdpaBackend::VarLen` added to backend_selector
- 8 CPU tests: single-seq matches regular attn, two seqs independent, causal mask,
  softmax sums to 1, empty sequence handled, output shape; total 50 pass

Inference CUDA graph capture (rtx-inference)
- `inference_graph.rs`: `InferenceGraphCapture` + `StepMode` {Warmup, Capture, Replay};
  state machine: N warmup steps → capture once → replay forever; `check_static_shape()`
  invalidates on batch/step change; `record_capture(graph_id)` stores graph
- `batch_processor.rs`: `graph_capture: Mutex<InferenceGraphCapture>` field (#[cfg(cuda)]);
  `advance()` wired at line 852 in `execute_batch_inference`; stream TODO matches
  training_loop.rs pattern; config gains `enable_decode_graphs`/`graph_warmup_steps`
- 8 pure-logic tests; total 93 pass; 4 integration test literals fixed

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

crates/production/rtx-inference/src/batch_processor.rs

@@ -7,6 +7,7 @@
 //! - Intelligent request coalescing and splitting
 //! - Real-time performance optimization
 
+use crate::inference_graph::InferenceGraphCapture;
 use crate::request::InferenceRequest;
 use crate::{FinishReason, InferenceError, InferenceResult, RequestMetrics, RequestResult};
 use rtx_tensor::{Device, Shape, Tensor};
@@ -37,6 +38,18 @@ pub struct BatchProcessorConfig {
     pub adaptive_batching: bool,
     /// Memory limit per batch
     pub max_batch_memory: usize,
+    /// Enable CUDA Graph capture for the decode step (requires `cuda` feature).
+    ///
+    /// When `true` the batch processor runs `graph_warmup_steps` warmup
+    /// iterations before attempting a single CUDA Graph capture. Subsequent
+    /// decode steps replay the captured graph instead of re-launching kernels.
+    /// Has no effect unless the `cuda` Cargo feature is enabled.
+    pub enable_decode_graphs: bool,
+    /// Number of warmup iterations before CUDA Graph capture is attempted.
+    ///
+    /// Only relevant when `enable_decode_graphs` is `true` and the `cuda`
+    /// feature is active.
+    pub graph_warmup_steps: usize,
 }
 
 impl Default for BatchProcessorConfig {
@@ -50,6 +63,8 @@ impl Default for BatchProcessorConfig {
             max_concurrent_batches: 4,
             adaptive_batching: true,
             max_batch_memory: 4 * 1024 * 1024 * 1024, // 4GB per batch
+            enable_decode_graphs: false,
+            graph_warmup_steps: 3,
         }
     }
 }
@@ -210,6 +225,18 @@ pub struct BatchProcessor {
 
     /// Statistics update task handle
     _stats_task: tokio::task::JoinHandle<()>,
+
+    /// CUDA Graph capture state machine for decode-step graph optimization.
+    ///
+    /// Present only when the `cuda` Cargo feature is active. Tracks warmup,
+    /// capture, and replay phases. Actual graph launch requires a stream handle
+    /// — see `TODO` in `execute_batch_inference`.
+    ///
+    /// Wrapped in a `parking_lot::Mutex` so it can be mutated through `&self`
+    /// (needed because `process_batch` / `execute_batch_inference` take `&self`
+    /// and may be called from cloned processors running concurrently).
+    #[cfg(feature = "cuda")]
+    graph_capture: parking_lot::Mutex<InferenceGraphCapture>,
 }
 
 impl BatchProcessor {
@@ -256,6 +283,13 @@ impl BatchProcessor {
             Self::start_batch_formation_task(lanes.clone(), config.clone(), device.clone());
         let stats_task = Self::start_stats_task(global_stats.clone(), lanes.clone());
 
+        // Build graph capture state machine from config (CUDA only).
+        #[cfg(feature = "cuda")]
+        let graph_capture = parking_lot::Mutex::new(
+            InferenceGraphCapture::new(config.graph_warmup_steps)
+                .enabled(config.enable_decode_graphs),
+        );
+
         let processor = Self {
             config,
             device,
@@ -267,6 +301,8 @@ impl BatchProcessor {
             cache_misses: Arc::new(AtomicU64::new(0)),
             _formation_task: formation_task,
             _stats_task: stats_task,
+            #[cfg(feature = "cuda")]
+            graph_capture,
         };
 
         info!("BatchProcessor initialized successfully");
@@ -762,6 +798,61 @@ impl BatchProcessor {
         let start_time = Instant::now();
         let mut results = Vec::with_capacity(batch.requests.len());
 
+        // ── CUDA Graph capture / replay ──────────────────────────────────────
+        // Check the graph capture state machine and route accordingly.
+        // Actual graph launch requires a stream handle; that wiring is deferred
+        // (same pattern as training_loop.rs).
+        // TODO: wire stream handle for actual graph launch
+        #[cfg(feature = "cuda")]
+        {
+            use crate::inference_graph::StepMode;
+            let mode = self.graph_capture.lock().step_mode();
+            let batch_size = batch.requests.len();
+            let seq_step = batch
+                .requests
+                .first()
+                .map(|r| r.request.input_tokens.len())
+                .unwrap_or(0);
+            match mode {
+                StepMode::Warmup => {
+                    trace!(
+                        "CUDA Graph: warmup step {} / {}",
+                        self.graph_capture.lock().step_count(),
+                        self.config.graph_warmup_steps,
+                    );
+                }
+                StepMode::Capture => {
+                    // Normal execution wraps the capture begin/end.
+                    // CudaGraphManager::begin_capture / end_capture require
+                    // a non-default stream which is not yet threaded through
+                    // BatchProcessor — deferred, same as TrainingLoop.
+                    debug!(
+                        "CUDA Graph: capture step (batch={}, seq_step={}) — \
+                         stream not yet wired, falling through to normal execution",
+                        batch_size, seq_step,
+                    );
+                }
+                StepMode::Replay => {
+                    if let Some(gid) = self.graph_capture.lock().captured_graph_id() {
+                        // Graph replay path — CudaGraphManager::launch(gid) would go here
+                        // once a stream handle is available.
+                        debug!(
+                            "CUDA Graph: replay graph_id={} \
+                             (stream not yet wired, falling through to normal execution)",
+                            gid,
+                        );
+                    } else {
+                        debug!(
+                            "CUDA Graph: replay mode but no graph captured yet, \
+                             falling through to normal execution"
+                        );
+                    }
+                }
+            }
+            self.graph_capture.lock().advance();
+        }
+        // ────────────────────────────────────────────────────────────────────
+
         // Prepare batch tensors
         let batch_tensors = self.prepare_batch_tensors(batch).await?;
 
@@ -1005,6 +1096,11 @@ impl Clone for BatchProcessor {
             cache_misses: Arc::clone(&self.cache_misses),
             _formation_task: tokio::spawn(async {}), // Placeholder handle
             _stats_task: tokio::spawn(async {}),     // Placeholder handle
+            #[cfg(feature = "cuda")]
+            graph_capture: parking_lot::Mutex::new(
+                InferenceGraphCapture::new(self.config.graph_warmup_steps)
+                    .enabled(self.config.enable_decode_graphs),
+            ),
         }
     }
 }

crates/production/rtx-inference/src/inference_graph.rs

@@ -0,0 +1,272 @@
+//! Inference CUDA Graph Capture State Machine
+//!
+//! Tracks warmup, capture, and replay phases for decode-step CUDA graph
+//! optimization. The state machine is pure logic — no GPU resources are held
+//! here. GPU resources (stream handles, graph IDs) are managed by
+//! [`CudaGraphManager`] and wired in by the caller.
+//!
+//! # Phase Transitions
+//!
+//! ```text
+//! Warmup (step_count < warmup_steps)
+//!   ──► Capture (step_count == warmup_steps && !capture_attempted)
+//!   ──► Replay (capture_attempted == true)
+//! ```
+//!
+//! When `enabled` is `false` the machine stays in `Warmup` indefinitely,
+//! which means normal (non-graph) execution is used for every step.
+
+/// The current execution mode as determined by the graph capture state machine.
+#[derive(Debug, Clone, PartialEq)]
+pub enum StepMode {
+    /// Normal execution — kernels are launched individually. Used during the
+    /// initial warmup phase to prime caches and stabilize shapes.
+    Warmup,
+    /// The next step should be captured into a CUDA Graph. Called exactly once
+    /// (at `step_count == warmup_steps`), then transitions to `Replay`.
+    Capture,
+    /// The captured graph is replayed instead of re-launching individual kernels.
+    Replay,
+}
+
+/// State machine for inference-time CUDA Graph capture and replay.
+///
+/// # Example
+///
+/// ```rust
+/// use rtx_inference::inference_graph::{InferenceGraphCapture, StepMode};
+///
+/// let mut capture = InferenceGraphCapture::new(3).enabled(true);
+/// assert_eq!(capture.step_mode(), StepMode::Warmup);
+///
+/// // Advance through warmup steps
+/// capture.advance(); // step 1
+/// capture.advance(); // step 2
+/// capture.advance(); // step 3 — still warmup (count is now 3, threshold is 3)
+/// assert_eq!(capture.step_mode(), StepMode::Capture);
+///
+/// // Record a successful capture
+/// capture.record_capture(42, 8, 1);
+/// assert!(capture.is_captured());
+/// assert_eq!(capture.step_mode(), StepMode::Replay);
+/// ```
+#[derive(Debug)]
+pub struct InferenceGraphCapture {
+    /// Captured CUDA Graph ID, or `None` if not yet captured.
+    graph_id: Option<u64>,
+    /// Whether a capture has already been attempted (prevents re-capture).
+    capture_attempted: bool,
+    /// Number of warmup steps before capture is attempted.
+    warmup_steps: usize,
+    /// Total number of steps advanced so far.
+    step_count: usize,
+    /// Whether CUDA Graph capture is enabled at all.
+    enabled: bool,
+    /// Batch size used during capture, for shape validation.
+    captured_batch_size: Option<usize>,
+    /// Sequence step length used during capture, for shape validation.
+    captured_seq_step: Option<usize>,
+}
+
+impl InferenceGraphCapture {
+    /// Create a new state machine with the given warmup step count.
+    ///
+    /// Graph capture is **disabled** by default. Call `.enabled(true)` to
+    /// activate it.
+    #[must_use]
+    pub fn new(warmup_steps: usize) -> Self {
+        Self {
+            graph_id: None,
+            capture_attempted: false,
+            warmup_steps,
+            step_count: 0,
+            enabled: false,
+            captured_batch_size: None,
+            captured_seq_step: None,
+        }
+    }
+
+    /// Builder method to set whether graph capture is enabled.
+    ///
+    /// When `enabled` is `false`, [`step_mode`](Self::step_mode) always
+    /// returns [`StepMode::Warmup`].
+    #[must_use]
+    pub fn enabled(mut self, enabled: bool) -> Self {
+        self.enabled = enabled;
+        self
+    }
+
+    /// Returns `true` if a graph has been successfully captured.
+    #[must_use]
+    pub fn is_captured(&self) -> bool {
+        self.graph_id.is_some()
+    }
+
+    /// Returns the total number of steps that have been advanced.
+    #[must_use]
+    pub fn step_count(&self) -> usize {
+        self.step_count
+    }
+
+    /// Determine the execution mode for the current step.
+    ///
+    /// - [`StepMode::Warmup`] — if capture is disabled, or `step_count < warmup_steps`.
+    /// - [`StepMode::Capture`] — if `step_count == warmup_steps` and no capture has
+    ///   been attempted yet.
+    /// - [`StepMode::Replay`] — if a capture has been attempted (successfully or not).
+    #[must_use]
+    pub fn step_mode(&self) -> StepMode {
+        if !self.enabled {
+            return StepMode::Warmup;
+        }
+        if self.step_count < self.warmup_steps {
+            StepMode::Warmup
+        } else if self.step_count == self.warmup_steps && !self.capture_attempted {
+            StepMode::Capture
+        } else {
+            StepMode::Replay
+        }
+    }
+
+    /// Advance the internal step counter by one and mark `capture_attempted`
+    /// once we are past the warmup threshold.
+    pub fn advance(&mut self) {
+        self.step_count += 1;
+        if self.step_count > self.warmup_steps {
+            self.capture_attempted = true;
+        }
+    }
+
+    /// Record a successful graph capture.
+    ///
+    /// # Arguments
+    ///
+    /// * `graph_id` — the ID returned by `CudaGraphManager::end_capture`.
+    /// * `batch_size` — batch dimension used during capture.
+    /// * `seq_step` — sequence-step dimension used during capture.
+    pub fn record_capture(&mut self, graph_id: u64, batch_size: usize, seq_step: usize) {
+        self.graph_id = Some(graph_id);
+        self.captured_batch_size = Some(batch_size);
+        self.captured_seq_step = Some(seq_step);
+        self.capture_attempted = true;
+    }
+
+    /// Check whether the provided shapes match the shapes used at capture time.
+    ///
+    /// Returns `true` if the shapes match (or if no capture has occurred yet),
+    /// `false` if they differ and the graph cannot be replayed safely.
+    #[must_use]
+    pub fn check_static_shape(&self, batch_size: usize, seq_step: usize) -> bool {
+        match (self.captured_batch_size, self.captured_seq_step) {
+            (Some(b), Some(s)) => b == batch_size && s == seq_step,
+            _ => true,
+        }
+    }
+
+    /// Return the captured graph ID, or `None` if no capture has occurred.
+    #[must_use]
+    pub fn captured_graph_id(&self) -> Option<u64> {
+        self.graph_id
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // -----------------------------------------------------------------------
+    // step_mode tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_step_mode_warmup_until_threshold() {
+        let mut cap = InferenceGraphCapture::new(3).enabled(true);
+        // Steps 0, 1, 2 are all Warmup
+        for _ in 0..3 {
+            assert_eq!(cap.step_mode(), StepMode::Warmup);
+            cap.advance();
+        }
+        // step_count is now 3 — at threshold; should be Capture, not Warmup
+        assert_ne!(cap.step_mode(), StepMode::Warmup);
+    }
+
+    #[test]
+    fn test_step_mode_capture_at_threshold() {
+        let mut cap = InferenceGraphCapture::new(3).enabled(true);
+        // Advance through exactly warmup_steps steps
+        for _ in 0..3 {
+            cap.advance();
+        }
+        assert_eq!(cap.step_mode(), StepMode::Capture);
+    }
+
+    #[test]
+    fn test_step_mode_replay_after_capture() {
+        let mut cap = InferenceGraphCapture::new(2).enabled(true);
+        cap.advance();
+        cap.advance();
+        assert_eq!(cap.step_mode(), StepMode::Capture);
+
+        cap.record_capture(1, 4, 1);
+        assert_eq!(cap.step_mode(), StepMode::Replay);
+    }
+
+    // -----------------------------------------------------------------------
+    // advance tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_advance_increments_step_count() {
+        let mut cap = InferenceGraphCapture::new(5).enabled(true);
+        assert_eq!(cap.step_count(), 0);
+        cap.advance();
+        assert_eq!(cap.step_count(), 1);
+        cap.advance();
+        assert_eq!(cap.step_count(), 2);
+    }
+
+    // -----------------------------------------------------------------------
+    // check_static_shape tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_static_shape_check_passes_same_shape() {
+        let mut cap = InferenceGraphCapture::new(1).enabled(true);
+        cap.advance();
+        cap.record_capture(7, 8, 1);
+        assert!(cap.check_static_shape(8, 1));
+    }
+
+    #[test]
+    fn test_static_shape_check_fails_different_batch() {
+        let mut cap = InferenceGraphCapture::new(1).enabled(true);
+        cap.advance();
+        cap.record_capture(7, 8, 1);
+        // Different batch size
+        assert!(!cap.check_static_shape(4, 1));
+    }
+
+    // -----------------------------------------------------------------------
+    // is_captured tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_is_captured_false_before_record() {
+        let cap = InferenceGraphCapture::new(3).enabled(true);
+        assert!(!cap.is_captured());
+    }
+
+    // -----------------------------------------------------------------------
+    // disabled tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_inference_graph_capture_disabled() {
+        let mut cap = InferenceGraphCapture::new(1).enabled(false);
+        // Even after advancing past the warmup threshold the mode stays Warmup
+        cap.advance();
+        cap.advance();
+        assert_eq!(cap.step_mode(), StepMode::Warmup);
+        assert!(!cap.is_captured());
+    }
+}

crates/production/rtx-inference/src/lib.rs

@@ -17,6 +17,8 @@
 
 pub mod batch_processor;
 pub mod cache;
+pub mod inference_graph;
+pub use inference_graph::{InferenceGraphCapture, StepMode};
 pub mod engine;
 pub mod error;
 pub mod model_loader;

crates/production/rtx-inference/tests/batch_engine_integration_test.rs

@@ -20,6 +20,8 @@ async fn test_batch_processor_initialization() {
         max_concurrent_batches: 4,
         adaptive_batching: true,
         max_batch_memory: 1024 * 1024 * 1024, // 1GB
+        enable_decode_graphs: false,
+        graph_warmup_steps: 3,
     };
 
     let device = Device::cuda(0).unwrap_or(Device::default());
@@ -63,6 +65,8 @@ async fn test_request_submission_and_processing() {
         max_concurrent_batches: 2,
         adaptive_batching: false,
         max_batch_memory: 1024 * 1024 * 512, // 512MB
+        enable_decode_graphs: false,
+        graph_warmup_steps: 3,
     };
 
     let device = Device::cuda(0).unwrap_or(Device::default());
@@ -106,6 +110,8 @@ async fn test_sla_lane_assignment() {
         max_concurrent_batches: 2,
         adaptive_batching: false,
         max_batch_memory: 1024 * 1024 * 256,
+        enable_decode_graphs: false,
+        graph_warmup_steps: 3,
     };
 
     let device = Device::cuda(0).unwrap_or(Device::default());
@@ -181,6 +187,8 @@ async fn test_batch_formation() {
         max_concurrent_batches: 2,
         adaptive_batching: true,
         max_batch_memory: 1024 * 1024 * 256,
+        enable_decode_graphs: false,
+        graph_warmup_steps: 3,
     };
 
     let device = Device::cuda(0).unwrap_or(Device::default());

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

@@ -1016,6 +1016,138 @@ impl SmoothQuantizedLayer {
     pub fn inverse_scales(&self) -> Vec<f32> {
         self.smoothing_scales.iter().map(|&s| 1.0 / s).collect()
     }
+
+    /// Number of output features (rows of the weight matrix).
+    #[inline]
+    pub fn out_features(&self) -> usize {
+        self.quantized_weights.shape[0]
+    }
+
+    /// Number of input features (columns of the weight matrix).
+    #[inline]
+    pub fn in_features(&self) -> usize {
+        self.quantized_weights.shape[1]
+    }
+
+    /// Run the quantized linear layer forward pass on a raw f32 buffer.
+    ///
+    /// Implements the SmoothQuant inference path:
+    /// 1. **Smooth activations** – divide each channel by its smoothing scale, migrating
+    ///    quantisation difficulty from activations to the pre-scaled weights.
+    /// 2. **Quantise activations to INT8** – compute a symmetric per-tensor scale
+    ///    `act_scale = max(|smoothed|) / 127` and round to the range `[−127, 127]`.
+    /// 3. **INT8 GEMM** – compute `acc[b, out] = Σ_k act_q[b,k] * (weight_q[out,k] − zp)`
+    ///    with i32 accumulation.
+    /// 4. **Dequantise** – `output[b, out] = acc[b, out] as f32 * act_scale * weight_scale`.
+    ///
+    /// # INT8 range
+    ///
+    /// Activations are clamped to **`−127…127`** (symmetric), matching the
+    /// `activation_bit_width = 8` convention used throughout this crate.
+    ///
+    /// # Dequantisation formula
+    ///
+    /// ```text
+    /// output[b, out] = (Σ_k act_q[b,k] * (weight_q[out,k] − zp)) * act_scale * weight_scale
+    /// ```
+    ///
+    /// # Arguments
+    ///
+    /// * `activations` – Row-major f32 slice of shape `[batch_size, in_features]`.
+    /// * `batch_size`  – Number of input rows.
+    ///
+    /// # Returns
+    ///
+    /// Row-major `Vec<f32>` of shape `[batch_size, out_features]`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the activation buffer length is inconsistent with `batch_size`.
+    pub fn forward_raw(
+        &self,
+        activations: &[f32],
+        batch_size: usize,
+    ) -> Result<Vec<f32>> {
+        let in_f = self.in_features();
+        let out_f = self.out_features();
+
+        let expected_len = batch_size * in_f;
+        if activations.len() != expected_len {
+            return Err(CompressionError::Quantization(
+                QuantizationError::InvalidConfig(format!(
+                    "SmoothQuantizedLayer::forward_raw: activation buffer length {} \
+                     does not match batch_size={} × in_features={}",
+                    activations.len(),
+                    batch_size,
+                    in_f
+                )),
+            ));
+        }
+
+        let zp = self.quantized_weights.zero_point as i32;
+        let weight_scale = self.quantized_weights.scale;
+
+        // ------------------------------------------------------------------
+        // Step 1 + 2: smooth then quantise activations to INT8
+        //
+        // We compute smoothed values first (in place into a temporary buffer),
+        // then find max_abs and derive act_scale before quantising.
+        // ------------------------------------------------------------------
+        let mut smoothed = vec![0.0f32; batch_size * in_f];
+        for b in 0..batch_size {
+            let row = b * in_f;
+            for c in 0..in_f {
+                // Divide by the per-channel smoothing scale (≥ ε for numerical safety).
+                let s = self.smoothing_scales[c].max(f32::EPSILON);
+                smoothed[row + c] = activations[row + c] / s;
+            }
+        }
+
+        // Per-tensor activation scale: max absolute value over the whole batch.
+        let max_abs = smoothed
+            .iter()
+            .map(|v| v.abs())
+            .fold(0.0f32, f32::max);
+
+        // Guard against zero-tensor inputs; any non-zero scale works here.
+        let act_scale = if max_abs < f32::EPSILON {
+            1.0f32
+        } else {
+            max_abs / 127.0f32
+        };
+
+        // Quantise each smoothed value to INT8 in range [−127, 127].
+        let act_q: Vec<i8> = smoothed
+            .iter()
+            .map(|&v| {
+                let q = (v / act_scale).round();
+                q.clamp(-127.0, 127.0) as i8
+            })
+            .collect();
+
+        // ------------------------------------------------------------------
+        // Step 3: INT8 GEMM with i32 accumulation
+        //
+        // output_i32[b, out_row] = Σ_k act_q[b,k] * (weight_q[out_row,k] − zp)
+        // ------------------------------------------------------------------
+        let mut output = vec![0.0f32; batch_size * out_f];
+        for b in 0..batch_size {
+            let act_row = b * in_f;
+            for out_row in 0..out_f {
+                let weight_row = out_row * in_f;
+                let mut acc = 0i32;
+                for k in 0..in_f {
+                    let wq = self.quantized_weights.data[weight_row + k] as i32 - zp;
+                    let aq = act_q[act_row + k] as i32;
+                    acc += aq * wq;
+                }
+                // Step 4: dequantise
+                output[b * out_f + out_row] = acc as f32 * act_scale * weight_scale;
+            }
+        }
+
+        Ok(output)
+    }
 }
 
 // =============================================================================
@@ -1123,4 +1255,169 @@ mod tests {
         assert!(ratio > 1.0); // Should be compressed
         assert!(ratio < 10.0); // But not impossibly so
     }
+
+    // -----------------------------------------------------------------------
+    // Helper: build a SmoothQuantizedLayer with fully controlled state
+    //
+    // Parameters:
+    //   out_features, in_features — weight matrix shape
+    //   weight_data   — INT8 quantized weights, row-major [out, in]
+    //   weight_scale  — single global scale for the weight tensor
+    //   weight_zp     — zero point for weight dequantization
+    //   smoothing     — per-channel smoothing scales, len = in_features
+    // -----------------------------------------------------------------------
+    fn make_smooth_layer(
+        out_features: usize,
+        in_features: usize,
+        weight_data: Vec<i8>,
+        weight_scale: f32,
+        weight_zp: i8,
+        smoothing: Vec<f32>,
+    ) -> SmoothQuantizedLayer {
+        SmoothQuantizedLayer {
+            quantized_weights: QuantizedTensorData {
+                data: weight_data,
+                scale: weight_scale,
+                zero_point: weight_zp,
+                shape: vec![out_features, in_features],
+            },
+            smoothing_scales: smoothing,
+            weight_bit_width: 8,
+            activation_bit_width: 8,
+        }
+    }
+
+    /// Output length must equal batch_size × out_features.
+    #[test]
+    fn test_smoothquant_forward_raw_shape() {
+        let out_f = 3usize;
+        let in_f = 4usize;
+        let layer = make_smooth_layer(
+            out_f,
+            in_f,
+            vec![0i8; out_f * in_f],
+            1.0,
+            0,
+            vec![1.0f32; in_f],
+        );
+        let activations = vec![0.0f32; 2 * in_f]; // batch=2
+        let out = layer.forward_raw(&activations, 2).unwrap();
+        assert_eq!(out.len(), 2 * out_f, "output length must be batch * out_features");
+    }
+
+    /// Identity layer (W=I, smoothing=1, zp=0, weight_scale=1) reproduces activations.
+    ///
+    /// With weight_scale=1 and act_scale derived from the input, roundtrip quantisation
+    /// introduces < 1% relative error for values in [1, 127].
+    #[test]
+    fn test_smoothquant_forward_raw_identity_layer() {
+        // 4×4 identity weight matrix (INT8), zero-point=0, weight_scale=1.
+        // smoothing_scales = [1.0; 4] → no smoothing effect.
+        // Activations: [1, 2, 3, 4] (batch=1).
+        // Expected output ≈ [1, 2, 3, 4] (within quantisation tolerance).
+        let in_f = 4usize;
+        let out_f = 4usize;
+        let identity: Vec<i8> = vec![
+            1, 0, 0, 0, //
+            0, 1, 0, 0, //
+            0, 0, 1, 0, //
+            0, 0, 0, 1, //
+        ];
+        let layer = make_smooth_layer(out_f, in_f, identity, 1.0, 0, vec![1.0f32; in_f]);
+        let activations = vec![1.0f32, 2.0, 3.0, 4.0];
+        let out = layer.forward_raw(&activations, 1).unwrap();
+
+        assert_eq!(out.len(), out_f);
+        // Tolerance: INT8 quantisation of activations loses at most 1/127 ≈ 0.8%.
+        for (i, (&computed, &expected)) in out.iter().zip(activations.iter()).enumerate() {
+            let rel_err = (computed - expected).abs() / (expected.abs().max(1e-6_f32));
+            assert!(
+                rel_err < 0.02,
+                "element {i}: got {computed:.4}, expected {expected:.4}, rel_err={rel_err:.4}"
+            );
+        }
+    }
+
+    /// smoothing_scales=[2.0]*n divides each activation channel by 2 before quantisation.
+    ///
+    /// With W=I, this halves the output compared to smoothing_scales=[1.0].
+    #[test]
+    fn test_smoothquant_forward_raw_scale_effect() {
+        let in_f = 4usize;
+        let out_f = 4usize;
+        let identity: Vec<i8> = vec![
+            1, 0, 0, 0, //
+            0, 1, 0, 0, //
+            0, 0, 1, 0, //
+            0, 0, 0, 1, //
+        ];
+
+        // Layer with smoothing_scales = [1.0; 4] (baseline)
+        let layer_no_smooth = make_smooth_layer(
+            out_f, in_f, identity.clone(), 1.0, 0, vec![1.0f32; in_f],
+        );
+        // Layer with smoothing_scales = [2.0; 4] (halves activations)
+        let layer_smooth = make_smooth_layer(
+            out_f, in_f, identity, 1.0, 0, vec![2.0f32; in_f],
+        );
+
+        let activations = vec![10.0f32, 20.0, 30.0, 40.0];
+
+        let out_baseline = layer_no_smooth.forward_raw(&activations, 1).unwrap();
+        let out_halved = layer_smooth.forward_raw(&activations, 1).unwrap();
+
+        for (i, (&baseline, &halved)) in out_baseline.iter().zip(out_halved.iter()).enumerate() {
+            let ratio = baseline / halved;
+            assert!(
+                (ratio - 2.0).abs() < 0.05,
+                "element {i}: baseline={baseline:.4} / halved={halved:.4} = ratio {ratio:.4}, expected ~2.0"
+            );
+        }
+    }
+
+    /// Manually computed reference: 2×2 layer, known weights/scales/input.
+    ///
+    /// Setup:
+    ///   W (INT8) = [[2, 0], [0, 2]], weight_scale=0.5, weight_zp=0
+    ///   smoothing_scales = [1.0, 1.0]
+    ///   activations = [10.0, 20.0] (batch=1)
+    ///
+    /// Step-by-step:
+    ///   1. Smooth: [10/1, 20/1] = [10.0, 20.0]
+    ///   2. act_scale = max(|10|, |20|) / 127 = 20/127 ≈ 0.15748
+    ///   3. Quantise: act_q = round([10/0.15748, 20/0.15748]).clamp(-127,127)
+    ///                      ≈ round([63.5, 127.0]) = [64, 127]  (or [63, 127] depending on rounding)
+    ///   4. INT8 GEMM (weight_zp=0):
+    ///      out[0] = 2*act_q[0] + 0*act_q[1]  = 2*64 = 128  (or 2*63=126)
+    ///      out[1] = 0*act_q[0] + 2*act_q[1]  = 2*127 = 254
+    ///   5. Dequant: output * act_scale * weight_scale
+    ///      out[0] ≈ 128 * 0.15748 * 0.5 ≈ 10.08  (expected ~10.0)
+    ///      out[1] ≈ 254 * 0.15748 * 0.5 ≈ 20.0
+    ///
+    /// Tolerance: 2% (quantisation rounding).
+    #[test]
+    fn test_smoothquant_forward_raw_matches_manual() {
+        let layer = make_smooth_layer(
+            2,
+            2,
+            vec![2i8, 0, 0, 2],
+            0.5,
+            0,
+            vec![1.0f32, 1.0],
+        );
+        let activations = vec![10.0f32, 20.0];
+        let out = layer.forward_raw(&activations, 1).unwrap();
+
+        assert_eq!(out.len(), 2);
+        assert!(
+            (out[0] - 10.0).abs() < 0.5,
+            "out[0] ≈ 10.0, got {:.4}",
+            out[0]
+        );
+        assert!(
+            (out[1] - 20.0).abs() < 0.5,
+            "out[1] ≈ 20.0, got {:.4}",
+            out[1]
+        );
+    }
 }

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

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

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

@@ -1,4 +1,5 @@
 pub mod advanced;
+pub mod int8_matmul;
 pub mod microscaling;
 pub mod mixed_precision;
 pub mod mx_gpu_kernels;
@@ -40,6 +41,9 @@ pub use w4a16_matmul::{
     w4a16_matmul_cpu,
 };
 
+// INT8 GEMM exports
+pub use int8_matmul::{int8_gemm, int8_matvec};
+
 // Advanced quantization exports (AWQ, GPTQ, SmoothQuant)
 pub use advanced::{
     // AWQ