## 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/gguf/ops.rs

    Line
    Count
    Source
    1
    //! Pure mathematical operations for GGUF inference
    2
    //!
    3
    //! This module contains standalone math functions used by both CPU and GPU
    4
    //! inference paths. By extracting these to a shared module, we enable:
    5
    //!
    6
    //! - Code reuse between `OwnedQuantizedModel` (CPU) and `OwnedQuantizedModelCuda` (GPU)
    7
    //! - Easier testing of mathematical correctness
    8
    //! - Clear separation of concerns
    9
    //!
    10
    //! ## Functions
    11
    //!
    12
    //! - `rms_norm`: RMSNorm normalization (LLaMA, Qwen, Mistral)
    13
    //! - `gelu`: GELU activation function
    14
    //! - `silu`: SiLU/Swish activation function
    15
    //! - `add_bias`: Add bias vector to output
    16
    //! - `argmax`: Find index of maximum value
    17
    //! - `softmax`: Numerically stable softmax
    18
    19
    use trueno::Vector as TruenoVector;
    20
    21
    // =============================================================================
    22
    // Normalization Operations
    23
    // =============================================================================
    24
    25
    /// RMSNorm (Root Mean Square Layer Normalization)
    26
    ///
    27
    /// Used by LLaMA, TinyLlama, Qwen, Mistral instead of LayerNorm.
    28
    /// Formula: output = x / sqrt(mean(x^2) + eps) * weight
    29
    ///
    30
    /// # Arguments
    31
    /// * `input` - Input tensor [seq_len * hidden_dim]
    32
    /// * `weight` - Normalization weights [hidden_dim]
    33
    /// * `eps` - Small constant for numerical stability (typically 1e-5 or 1e-6)
    34
    ///
    35
    /// # Returns
    36
    /// Normalized output [seq_len * hidden_dim]
    37
    1
    pub fn rms_norm(input: &[f32], weight: &[f32], eps: f32) -> Vec<f32> {
    38
    1
        let hidden_dim = weight.len();
    39
    1
        let seq_len = input.len() / hidden_dim;
    40
    1
        let mut output = Vec::with_capacity(input.len());
    41
    42
    1
        let weight_vec = TruenoVector::from_slice(weight);
    43
    44
    1
        for i in 0..seq_len {
    45
    1
            let start = i * hidden_dim;
    46
    1
            let end = start + hidden_dim;
    47
    1
            let x = &input[start..end];
    48
    49
    1
            let x_vec = TruenoVector::from_slice(x);
    50
    51
            // SIMD: sum of squares
    52
    1
            let sum_sq = x_vec
    53
    1
                .sum_of_squares()
    54
    1
                .unwrap_or_else(|_| x0.iter0().map0(|v| v0 * v0).sum0::<f32>());
    55
    56
    1
            let mean_sq = sum_sq / hidden_dim as f32;
    57
    1
            let inv_rms = 1.0 / (mean_sq + eps).sqrt();
    58
    59
            // SIMD: scale by inv_rms, then multiply by weight
    60
    1
            match x_vec
    61
    1
                .scale(inv_rms)
    62
    1
                .and_then(|scaled| scaled.mul(&weight_vec))
    63
            {
    64
    1
                Ok(result) => {
    65
    1
                    output.extend_from_slice(result.as_slice());
    66
    1
                }
    67
                Err(_) => {
    68
                    // Fallback to scalar
    69
    0
                    for j in 0..hidden_dim {
    70
    0
                        output.push(x[j] * inv_rms * weight[j]);
    71
    0
                    }
    72
                }
    73
            }
    74
        }
    75
    76
    1
        output
    77
    1
    }
    78
    79
    /// RMSNorm to pre-allocated buffer (zero-allocation path)
    80
    ///
    81
    /// # Arguments
    82
    /// * `input` - Input tensor [hidden_dim] (single position)
    83
    /// * `weight` - Normalization weights [hidden_dim]
    84
    /// * `eps` - Small constant for numerical stability
    85
    /// * `output` - Pre-allocated output buffer [hidden_dim]
    86
    0
    pub fn rms_norm_into(input: &[f32], weight: &[f32], eps: f32, output: &mut [f32]) {
    87
    0
        let hidden_dim = weight.len();
    88
    0
        let x = &input[..hidden_dim];
    89
    90
    0
        let x_vec = TruenoVector::from_slice(x);
    91
    0
        let weight_vec = TruenoVector::from_slice(weight);
    92
    93
    0
        let sum_sq = x_vec
    94
    0
            .sum_of_squares()
    95
    0
            .unwrap_or_else(|_| x.iter().map(|v| v * v).sum::<f32>());
    96
    97
    0
        let mean_sq = sum_sq / hidden_dim as f32;
    98
    0
        let inv_rms = 1.0 / (mean_sq + eps).sqrt();
    99
    100
    0
        match x_vec
    101
    0
            .scale(inv_rms)
    102
    0
            .and_then(|scaled| scaled.mul(&weight_vec))
    103
        {
    104
    0
            Ok(result) => {
    105
    0
                output[..hidden_dim].copy_from_slice(result.as_slice());
    106
    0
            }
    107
            Err(_) => {
    108
    0
                for j in 0..hidden_dim {
    109
    0
                    output[j] = x[j] * inv_rms * weight[j];
    110
    0
                }
    111
            }
    112
        }
    113
    0
    }
    114
    115
    /// Layer normalization with optional bias
    116
    ///
    117
    /// PMAT-094: This is actually RMSNorm for LLaMA-style models.
    118
    /// Kept for API compatibility with models that expect layer_norm signature.
    119
    ///
    120
    /// # Arguments
    121
    /// * `input` - Input tensor [seq_len * hidden_dim]
    122
    /// * `weight` - Normalization weights [hidden_dim]
    123
    /// * `bias` - Optional bias [hidden_dim]
    124
    /// * `eps` - Small constant for numerical stability
    125
    1.36k
    pub fn layer_norm(input: &[f32], weight: &[f32], bias: Option<&[f32]>, eps: f32) -> Vec<f32> {
    126
    1.36k
        let hidden_dim = weight.len();
    127
    1.36k
        let seq_len = input.len() / hidden_dim;
    128
    1.36k
        let mut output = Vec::with_capacity(input.len());
    129
    130
    1.47k
        for i in 0..seq_len1.36k {
    131
    1.47k
            let start = i * hidden_dim;
    132
    1.47k
            let end = start + hidden_dim;
    133
    1.47k
            let x = &input[start..end];
    134
    135
            // RMSNorm: compute root mean square (no mean subtraction!)
    136
    95.5k
            let sum_sq1.47k: f321.47k = x1.47k.iter1.47k().map1.47k(|v| v * v).sum1.47k();
    137
    1.47k
            let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();
    138
    139
    95.5k
            for j in 0..hidden_dim1.47k {
    140
    95.5k
                let normalized = x[j] / rms;
    141
    95.5k
                let mut val = normalized * weight[j];
    142
    95.5k
                if let Some(b2) = bias {
    143
    2
                    val += b[j];
    144
    95.5k
                }
    145
    95.5k
                output.push(val);
    146
            }
    147
        }
    148
    149
    1.36k
        output
    150
    1.36k
    }
    151
    152
    /// Layer normalization to pre-allocated buffer
    153
    0
    pub fn layer_norm_into(
    154
    0
        input: &[f32],
    155
    0
        weight: &[f32],
    156
    0
        bias: Option<&[f32]>,
    157
    0
        eps: f32,
    158
    0
        output: &mut [f32],
    159
    0
    ) {
    160
    0
        let hidden_dim = weight.len();
    161
    0
        let x = &input[..hidden_dim];
    162
    163
    0
        let sum_sq: f32 = x.iter().map(|v| v * v).sum();
    164
    0
        let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();
    165
    166
    0
        for j in 0..hidden_dim {
    167
    0
            let normalized = x[j] / rms;
    168
    0
            output[j] = normalized * weight[j];
    169
    0
            if let Some(b) = bias {
    170
    0
                output[j] += b[j];
    171
    0
            }
    172
        }
    173
    0
    }
    174
    175
    // =============================================================================
    176
    // Activation Functions
    177
    // =============================================================================
    178
    179
    /// GELU (Gaussian Error Linear Unit) activation
    180
    ///
    181
    /// Approximation: GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
    182
    ///
    183
    /// # Arguments
    184
    /// * `input` - Input tensor (modified in-place)
    185
    #[inline]
    186
    717
    pub fn gelu(input: &mut [f32]) {
    187
        const SQRT_2_OVER_PI: f32 = 0.797_884_6;
    188
        const C: f32 = 0.044_715;
    189
    190
    107k
        for x in input717.iter_mut717() {
    191
    107k
            let inner = SQRT_2_OVER_PI * (*x + C * *x * *x * *x);
    192
    107k
            *x = 0.5 * *x * (1.0 + inner.tanh());
    193
    107k
        }
    194
    717
    }
    195
    196
    /// SiLU (Sigmoid Linear Unit) / Swish activation
    197
    ///
    198
    /// SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))
    199
    /// Used in SwiGLU FFN (LLaMA, Mistral, etc.)
    200
    ///
    201
    /// # Arguments
    202
    /// * `input` - Input tensor (modified in-place)
    203
    #[inline]
    204
    2
    pub fn silu(input: &mut [f32]) {
    205
    2
        for x in input.iter_mut() {
    206
    2
            *x = *x * (1.0 / (1.0 + (-*x).exp()));
    207
    2
        }
    208
    2
    }
    209
    210
    // =============================================================================
    211
    // Utility Operations
    212
    // =============================================================================
    213
    214
    /// Add bias vector to output tensor
    215
    ///
    216
    /// # Arguments
    217
    /// * `output` - Output tensor [seq_len * out_dim] (modified in-place)
    218
    /// * `bias` - Bias vector [out_dim]
    219
    #[inline]
    220
    1
    pub fn add_bias(output: &mut [f32], bias: &[f32]) {
    221
    1
        let out_dim = bias.len();
    222
    1
        let seq_len = output.len() / out_dim;
    223
    2
        for s in 0..seq_len1 {
    224
    4
            for o in 0..out_dim2 {
    225
    4
                output[s * out_dim + o] += bias[o];
    226
    4
            }
    227
        }
    228
    1
    }
    229
    230
    /// Find index of maximum value (greedy decoding)
    231
    ///
    232
    /// # Arguments
    233
    /// * `logits` - Logit values [vocab_size]
    234
    ///
    235
    /// # Returns
    236
    /// Index of the maximum value
    237
    #[inline]
    238
    398
    pub fn argmax(logits: &[f32]) -> u32 {
    239
    398
        let mut max_idx = 0u32;
    240
    398
        let mut max_val = f32::NEG_INFINITY;
    241
    39.6k
        for (i, &val) in logits398.iter398().enumerate398() {
    242
    39.6k
            if val > max_val {
    243
    401
                max_val = val;
    244
    401
                max_idx = i as u32;
    245
    39.2k
            }
    246
        }
    247
    398
        max_idx
    248
    398
    }
    249
    250
    /// Numerically stable softmax
    251
    ///
    252
    /// Computes softmax(x) = exp(x - max(x)) / sum(exp(x - max(x)))
    253
    ///
    254
    /// # Arguments
    255
    /// * `logits` - Input logits (modified in-place to probabilities)
    256
    2
    pub fn softmax(logits: &mut [f32]) {
    257
        // Find max for numerical stability
    258
    2
        let max_val = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    259
    260
        // Compute exp(x - max) and sum
    261
    2
        let mut sum = 0.0f32;
    262
    6
        for x in logits2.iter_mut2() {
    263
    6
            *x = (*x - max_val).exp();
    264
    6
            sum += *x;
    265
    6
        }
    266
    267
        // Normalize
    268
    2
        let inv_sum = 1.0 / sum;
    269
    6
        for x in logits2.iter_mut2() {
    270
    6
            *x *= inv_sum;
    271
    6
        }
    272
    2
    }
    273
    274
    // =============================================================================
    275
    // Tests
    276
    // =============================================================================
    277
    278
    #[cfg(test)]
    279
    mod tests {
    280
        use super::*;
    281
    282
        #[test]
    283
    1
        fn test_gelu_zero() {
    284
    1
            let mut input = vec![0.0];
    285
    1
            gelu(&mut input);
    286
    1
            assert!((input[0] - 0.0).abs() < 1e-6);
    287
    1
        }
    288
    289
        #[test]
    290
    1
        fn test_gelu_positive() {
    291
    1
            let mut input = vec![1.0];
    292
    1
            gelu(&mut input);
    293
            // GELU(1) ≈ 0.8413
    294
    1
            assert!((input[0] - 0.8413).abs() < 0.01);
    295
    1
        }
    296
    297
        #[test]
    298
    1
        fn test_silu_zero() {
    299
    1
            let mut input = vec![0.0];
    300
    1
            silu(&mut input);
    301
    1
            assert!((input[0] - 0.0).abs() < 1e-6);
    302
    1
        }
    303
    304
        #[test]
    305
    1
        fn test_silu_positive() {
    306
    1
            let mut input = vec![1.0];
    307
    1
            silu(&mut input);
    308
            // SiLU(1) = 1 * sigmoid(1) ≈ 0.7311
    309
    1
            assert!((input[0] - 0.7311).abs() < 0.01);
    310
    1
        }
    311
    312
        #[test]
    313
    1
        fn test_argmax() {
    314
    1
            let logits = vec![0.1, 0.5, 0.3, 0.9, 0.2];
    315
    1
            assert_eq!(argmax(&logits), 3);
    316
    1
        }
    317
    318
        #[test]
    319
    1
        fn test_argmax_negative() {
    320
    1
            let logits = vec![-1.0, -0.5, -2.0];
    321
    1
            assert_eq!(argmax(&logits), 1);
    322
    1
        }
    323
    324
        #[test]
    325
    1
        fn test_softmax_uniform() {
    326
    1
            let mut logits = vec![0.0, 0.0, 0.0];
    327
    1
            softmax(&mut logits);
    328
    4
            for &p3 in &logits {
    329
    3
                assert!((p - 1.0 / 3.0).abs() < 1e-6);
    330
            }
    331
    1
        }
    332
    333
        #[test]
    334
    1
        fn test_softmax_sums_to_one() {
    335
    1
            let mut logits = vec![1.0, 2.0, 3.0];
    336
    1
            softmax(&mut logits);
    337
    1
            let sum: f32 = logits.iter().sum();
    338
    1
            assert!((sum - 1.0).abs() < 1e-6);
    339
    1
        }
    340
    341
        #[test]
    342
    1
        fn test_add_bias() {
    343
    1
            let mut output = vec![1.0, 2.0, 3.0, 4.0];
    344
    1
            let bias = vec![0.1, 0.2];
    345
    1
            add_bias(&mut output, &bias);
    346
    1
            assert!((output[0] - 1.1).abs() < 1e-6);
    347
    1
            assert!((output[1] - 2.2).abs() < 1e-6);
    348
    1
            assert!((output[2] - 3.1).abs() < 1e-6);
    349
    1
            assert!((output[3] - 4.2).abs() < 1e-6);
    350
    1
        }
    351
    352
        #[test]
    353
    1
        fn test_rms_norm_unit_weight() {
    354
    1
            let input = vec![1.0, 2.0, 3.0, 4.0];
    355
    1
            let weight = vec![1.0, 1.0, 1.0, 1.0];
    356
    1
            let output = rms_norm(&input, &weight, 1e-5);
    357
    358
            // Check output is normalized
    359
    4
            let sum_sq1: f321 = output.iter()1.map1(|x| x * x).sum1();
    360
    1
            let rms = (sum_sq / 4.0).sqrt();
    361
    1
            assert!((rms - 1.0).abs() < 0.1); // Approximately unit RMS
    362
    1
        }
    363
    364
        #[test]
    365
    1
        fn test_layer_norm_no_bias() {
    366
    1
            let input = vec![1.0, 2.0, 3.0, 4.0];
    367
    1
            let weight = vec![1.0, 1.0, 1.0, 1.0];
    368
    1
            let output = layer_norm(&input, &weight, None, 1e-5);
    369
    1
            assert_eq!(output.len(), 4);
    370
    1
        }
    371
    372
        #[test]
    373
    1
        fn test_layer_norm_with_bias() {
    374
    1
            let input = vec![1.0, 2.0];
    375
    1
            let weight = vec![1.0, 1.0];
    376
    1
            let bias = vec![0.5, 0.5];
    377
    1
            let output = layer_norm(&input, &weight, Some(&bias), 1e-5);
    378
    1
            assert_eq!(output.len(), 2);
    379
            // Output should have bias added
    380
    1
            assert!(output[0] > 0.0);
    381
    1
            assert!(output[1] > 0.0);
    382
    1
        }
    383
    }