## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/realizar/src/inference/mod.rs

    Line
    Count
    Source
    1
    //! SIMD-accelerated inference engine
    2
    //!
    3
    //! Provides high-performance transformer inference using trueno's SIMD primitives.
    4
    //! Designed to compete with llama.cpp on CPU performance.
    5
    //!
    6
    //! ## Architecture
    7
    //!
    8
    //! ```text
    9
    //! GGUF Model → GGUFTransformer → TruenoInferenceEngine → Tokens
    10
    //! ```
    11
    //!
    12
    //! ## Modules
    13
    //!
    14
    //! - [`thread`] - Thread configuration for dynamic thread allocation
    15
    //! - [`simd`] - SIMD-accelerated operations (matmul, dot, activations)
    16
    //! - [`kv_cache`] - Key-value cache for autoregressive generation
    17
    //! - [`norm`] - Layer and RMS normalization
    18
    //! - [`rope`] - Rotary position embeddings
    19
    //! - [`engine`] - Main inference engine implementations
    20
    //! - [`quantized`] - Quantized weight storage and inference
    21
    //!
    22
    //! ## Performance Targets
    23
    //!
    24
    //! - Use trueno's SIMD Vector::dot for all dot products
    25
    //! - Use trueno's Matrix::matmul for weight projections
    26
    //! - Target >100 tokens/sec on CPU for 1B models
    27
    28
    // Submodules
    29
    mod kv_cache;
    30
    mod norm;
    31
    mod simd;
    32
    mod thread;
    33
    34
    // Re-exports for public API
    35
    pub use kv_cache::{attention_with_cache, attention_with_transposed_v, KVCache, OptimizedKVCache};
    36
    pub use norm::{apply_rope, simd_layer_norm, simd_rms_norm};
    37
    pub use simd::{
    38
        simd_add, simd_bf16_dot, simd_bf16_matmul, simd_bf16_to_f32, simd_dot, simd_f16_to_f32,
    39
        simd_gelu, simd_matmul, simd_mul, simd_silu, simd_softmax,
    40
    };
    41
    pub use thread::{
    42
        configure_optimal_thread_pool, configure_thread_pool, InferenceMode, ThreadConfig,
    43
    };
    44
    45
    use crate::error::{RealizarError, Result};
    46
    use crate::quantize::{fused_q4k_tiled_matvec, QK_K};
    47
    48
    // ============================================================================
    49
    // QUANTIZED WEIGHT STORAGE (Phase 3: Memory Bandwidth Optimization)
    50
    // ============================================================================
    51
    //
    52
    // Keeps weights in quantized format for 8x memory bandwidth reduction.
    53
    // Uses fused dequant+dot operations during inference.
    54
    // ============================================================================
    55
    56
    /// Quantized weight matrix stored in Q4_K format
    57
    ///
    58
    /// Uses fused dequantize+dot operations for 8x memory bandwidth reduction.
    59
    /// Each row is stored as raw Q4_K bytes, dequantized on-the-fly during matmul.
    60
    #[derive(Clone)]
    61
    pub struct Q4KWeight {
    62
        /// Raw Q4_K quantized data
    63
        pub data: Vec<u8>,
    64
        /// Input dimension (number of columns when dequantized)
    65
        pub in_dim: usize,
    66
        /// Output dimension (number of rows)
    67
        pub out_dim: usize,
    68
    }
    69
    70
    impl Q4KWeight {
    71
        /// Create a new quantized weight from raw Q4_K data
    72
        ///
    73
        /// # Arguments
    74
        ///
    75
        /// * `data` - Raw Q4_K quantized bytes
    76
        /// * `in_dim` - Number of input features (must be multiple of 256)
    77
        /// * `out_dim` - Number of output features
    78
        ///
    79
        /// # Errors
    80
        ///
    81
        /// Returns error if dimensions don't match the data size
    82
    2
        pub fn new(data: Vec<u8>, in_dim: usize, out_dim: usize) -> Result<Self> {
    83
            // Q4_K uses 256-element super-blocks, each taking 144 bytes
    84
    2
            let blocks_per_row = in_dim.div_ceil(QK_K);
    85
    2
            let bytes_per_row = blocks_per_row * 144; // Q4_K block size
    86
    2
            let expected_bytes = out_dim * bytes_per_row;
    87
    88
    2
            if data.len() != expected_bytes {
    89
    1
                return Err(RealizarError::InvalidShape {
    90
    1
                    reason: format!(
    91
    1
                        "Q4KWeight data size {} doesn't match expected {} for {}x{} matrix",
    92
    1
                        data.len(),
    93
    1
                        expected_bytes,
    94
    1
                        out_dim,
    95
    1
                        in_dim
    96
    1
                    ),
    97
    1
                });
    98
    1
            }
    99
    100
    1
            Ok(Self {
    101
    1
                data,
    102
    1
                in_dim,
    103
    1
                out_dim,
    104
    1
            })
    105
    2
        }
    106
    107
        /// Perform matrix-vector multiplication using fused Q4_K operations
    108
        ///
    109
        /// Uses tiled operations for cache efficiency. Each output element is
    110
        /// computed by fused dequantize+dot against the input vector.
    111
        ///
    112
        /// # Arguments
    113
        ///
    114
        /// * `input` - Input vector of length `in_dim`
    115
        ///
    116
        /// # Returns
    117
        ///
    118
        /// Output vector of length `out_dim`
    119
        ///
    120
        /// # Errors
    121
        ///
    122
        /// Returns error if input dimension doesn't match
    123
    0
        pub fn matvec(&self, input: &[f32]) -> Result<Vec<f32>> {
    124
    0
            if input.len() != self.in_dim {
    125
    0
                return Err(RealizarError::InvalidShape {
    126
    0
                    reason: format!(
    127
    0
                        "Input length {} doesn't match weight in_dim {}",
    128
    0
                        input.len(),
    129
    0
                        self.in_dim
    130
    0
                    ),
    131
    0
                });
    132
    0
            }
    133
    134
            // Use tiled Q4_K matmul for cache efficiency
    135
    0
            fused_q4k_tiled_matvec(&self.data, input, self.in_dim, self.out_dim, None)
    136
    0
        }
    137
    138
        /// Get memory usage in bytes
    139
        #[must_use]
    140
    2
        pub fn memory_bytes(&self) -> usize {
    141
    2
            self.data.len()
    142
    2
        }
    143
    144
        /// Get equivalent f32 memory usage for comparison
    145
        #[must_use]
    146
    2
        pub fn f32_equivalent_bytes(&self) -> usize {
    147
    2
            self.in_dim * self.out_dim * 4
    148
    2
        }
    149
    150
        /// Get compression ratio vs f32
    151
        #[must_use]
    152
    1
        pub fn compression_ratio(&self) -> f32 {
    153
    1
            self.f32_equivalent_bytes() as f32 / self.memory_bytes() as f32
    154
    1
        }
    155
    }
    156
    #[cfg(test)]
    157
    mod tests {
    158
        use super::*;
    159
    160
        // ------------------------------------------------------------------------
    161
        // SIMD Operation Tests
    162
        // ------------------------------------------------------------------------
    163
    164
        #[test]
    165
    1
        fn test_simd_dot_basic() {
    166
    1
            let a = vec![1.0, 2.0, 3.0, 4.0];
    167
    1
            let b = vec![1.0, 1.0, 1.0, 1.0];
    168
    1
            let result = simd_dot(&a, &b);
    169
    1
            assert!((result - 10.0).abs() < 1e-6);
    170
    1
        }
    171
    172
        #[test]
    173
    1
        fn test_simd_dot_zeros() {
    174
    1
            let a = vec![0.0; 100];
    175
    1
            let b = vec![1.0; 100];
    176
    1
            assert!((simd_dot(&a, &b)).abs() < 1e-6);
    177
    1
        }
    178
    179
        #[test]
    180
    1
        fn test_simd_add() {
    181
    1
            let mut a = vec![1.0, 2.0, 3.0];
    182
    1
            let b = vec![4.0, 5.0, 6.0];
    183
    1
            simd_add(&mut a, &b);
    184
    1
            assert_eq!(a, vec![5.0, 7.0, 9.0]);
    185
    1
        }
    186
    187
        #[test]
    188
    1
        fn test_simd_mul() {
    189
    1
            let mut a = vec![1.0, 2.0, 3.0];
    190
    1
            let b = vec![2.0, 3.0, 4.0];
    191
    1
            simd_mul(&mut a, &b);
    192
    1
            assert_eq!(a, vec![2.0, 6.0, 12.0]);
    193
    1
        }
    194
    195
        #[test]
    196
    1
        fn test_simd_softmax_basic() {
    197
    1
            let mut data = vec![1.0, 2.0, 3.0];
    198
    1
            simd_softmax(&mut data);
    199
    200
            // Check sum to 1
    201
    1
            let sum: f32 = data.iter().sum();
    202
    1
            assert!((sum - 1.0).abs() < 1e-6);
    203
    204
            // Check ordering preserved
    205
    1
            assert!(data[2] > data[1]);
    206
    1
            assert!(data[1] > data[0]);
    207
    1
        }
    208
    209
        #[test]
    210
    1
        fn test_simd_softmax_empty() {
    211
    1
            let mut data: Vec<f32> = vec![];
    212
    1
            simd_softmax(&mut data);
    213
    1
            assert!(data.is_empty());
    214
    1
        }
    215
    216
        #[test]
    217
    1
        fn test_simd_softmax_single() {
    218
    1
            let mut data = vec![5.0];
    219
    1
            simd_softmax(&mut data);
    220
    1
            assert!((data[0] - 1.0).abs() < 1e-6);
    221
    1
        }
    222
    223
        #[test]
    224
    1
        fn test_simd_silu() {
    225
    1
            let mut data = vec![0.0, 1.0, -1.0];
    226
    1
            simd_silu(&mut data);
    227
    1
            assert!((data[0]).abs() < 1e-6); // silu(0) = 0
    228
    1
            assert!(data[1] > 0.5); // silu(1) > 0.5
    229
    1
            assert!(data[2] < 0.0); // silu(-1) < 0
    230
    1
        }
    231
    232
        #[test]
    233
    1
        fn test_simd_gelu() {
    234
    1
            let mut data = vec![0.0, 1.0, -1.0];
    235
    1
            simd_gelu(&mut data);
    236
    1
            assert!((data[0]).abs() < 1e-6); // gelu(0) ≈ 0
    237
    1
            assert!(data[1] > 0.8); // gelu(1) ≈ 0.84
    238
    1
            assert!(data[2] < 0.0); // gelu(-1) < 0
    239
    1
        }
    240
    241
        // ------------------------------------------------------------------------
    242
        // KVCache Tests
    243
        // ------------------------------------------------------------------------
    244
    245
        #[test]
    246
    1
        fn test_kv_cache_new() {
    247
    1
            let cache = KVCache::new(4, 128, 512);
    248
    1
            assert_eq!(cache.len(), 0);
    249
    1
            assert!(cache.is_empty());
    250
    1
        }
    251
    252
        #[test]
    253
    1
        fn test_kv_cache_store_and_retrieve() {
    254
    1
            let mut cache = KVCache::new(2, 4, 10);
    255
    1
            let k = vec![1.0, 2.0, 3.0, 4.0];
    256
    1
            let v = vec![5.0, 6.0, 7.0, 8.0];
    257
    258
    1
            cache.store(0, &k, &v);
    259
    1
            cache.advance();
    260
    261
    1
            assert_eq!(cache.len(), 1);
    262
    1
            assert_eq!(cache.get_k(0), &k[..]);
    263
    1
            assert_eq!(cache.get_v(0), &v[..]);
    264
    1
        }
    265
    266
        #[test]
    267
    1
        fn test_kv_cache_reset() {
    268
    1
            let mut cache = KVCache::new(2, 4, 10);
    269
    1
            let k = vec![1.0; 4];
    270
    1
            let v = vec![2.0; 4];
    271
    272
    1
            cache.store(0, &k, &v);
    273
    1
            cache.advance();
    274
    1
            assert_eq!(cache.len(), 1);
    275
    276
    1
            cache.reset();
    277
    1
            assert_eq!(cache.len(), 0);
    278
    1
            assert!(cache.is_empty());
    279
    1
        }
    280
    281
        // ------------------------------------------------------------------------
    282
        // Normalization Tests
    283
        // ------------------------------------------------------------------------
    284
    285
        #[test]
    286
    1
        fn test_simd_layer_norm() {
    287
    1
            let input = vec![1.0, 2.0, 3.0, 4.0];
    288
    1
            let weight = vec![1.0; 4];
    289
    1
            let result = simd_layer_norm(&input, &weight, None, 1e-5);
    290
    291
            // Mean should be 2.5, normalized values should sum to ~0
    292
    1
            let sum: f32 = result.iter().sum();
    293
    1
            assert!(sum.abs() < 1e-5);
    294
    1
        }
    295
    296
        #[test]
    297
    1
        fn test_simd_rms_norm() {
    298
    1
            let input = vec![3.0, 4.0]; // RMS = sqrt((9+16)/2) = sqrt(12.5)
    299
    1
            let weight = vec![1.0, 1.0];
    300
    1
            let result = simd_rms_norm(&input, &weight, 1e-5);
    301
    302
            // Result should be normalized by RMS
    303
    1
            let rms = (12.5f32).sqrt();
    304
    1
            assert!((result[0] - 3.0 / rms).abs() < 1e-5);
    305
    1
            assert!((result[1] - 4.0 / rms).abs() < 1e-5);
    306
    1
        }
    307
    308
        // ------------------------------------------------------------------------
    309
        // RoPE Tests
    310
        // ------------------------------------------------------------------------
    311
    312
        #[test]
    313
    1
        fn test_rope_position_zero() {
    314
    1
            let mut x = vec![1.0, 0.0, 1.0, 0.0];
    315
    1
            apply_rope(&mut x, 4, 1, 0, 10000.0);
    316
            // At position 0, cos(0)=1, sin(0)=0, so no change
    317
    1
            assert!((x[0] - 1.0).abs() < 1e-5);
    318
    1
            assert!((x[2] - 1.0).abs() < 1e-5);
    319
    1
        }
    320
    321
        #[test]
    322
    1
        fn test_rope_changes_values() {
    323
    1
            let original = vec![1.0, 2.0, 3.0, 4.0];
    324
    1
            let mut x = original.clone();
    325
    1
            apply_rope(&mut x, 4, 1, 10, 10000.0);
    326
            // Values should change at non-zero positions
    327
    1
            assert!(x != original);
    328
    1
        }
    329
    330
        // ------------------------------------------------------------------------
    331
        // Q4KWeight Tests
    332
        // ------------------------------------------------------------------------
    333
    334
        #[test]
    335
    1
        fn test_q4k_weight_memory_stats() {
    336
            // Create minimal valid Q4_K data
    337
    1
            let in_dim = 256; // Minimum for Q4_K (one super-block)
    338
    1
            let out_dim = 1;
    339
    1
            let bytes_per_row = 144; // Q4_K block size for 256 elements
    340
    1
            let data = vec![0u8; out_dim * bytes_per_row];
    341
    342
    1
            let weight = Q4KWeight::new(data, in_dim, out_dim).expect("operation failed");
    343
    1
            assert_eq!(weight.memory_bytes(), bytes_per_row);
    344
    1
            assert_eq!(weight.f32_equivalent_bytes(), in_dim * out_dim * 4);
    345
    1
            assert!(weight.compression_ratio() > 1.0);
    346
    1
        }
    347
    348
        #[test]
    349
    1
        fn test_q4k_weight_invalid_size() {
    350
    1
            let data = vec![0u8; 100]; // Wrong size
    351
    1
            let result = Q4KWeight::new(data, 256, 1);
    352
    1
            assert!(result.is_err());
    353
    1
        }
    354
    }