## 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/layers/mod.rs

    Line
    Count
    Source
    1
    //! Neural network layers for transformer models
    2
    //!
    3
    //! Implements core building blocks for transformer architectures:
    4
    //! - Layer normalization
    5
    //! - Multi-head attention (MHA, MQA, GQA)
    6
    //! - Feed-forward networks
    7
    //! - Position embeddings: `RoPE` and `ALiBi`
    8
    //!
    9
    //! ## Example
    10
    //!
    11
    //! ```rust,ignore
    12
    //! use realizar::layers::LayerNorm;
    13
    //!
    14
    //! let layer_norm = LayerNorm::new(512, 1e-5)?;
    15
    //! let normalized = layer_norm.forward(&input)?;
    16
    //! ```
    17
    //!
    18
    //! ## Loading Models from Files
    19
    //!
    20
    //! Realizar supports loading models from GGUF and Safetensors formats.
    21
    //!
    22
    //! ### GGUF Example
    23
    //!
    24
    //! ```rust,ignore
    25
    //! use realizar::{gguf::GGUFModel, layers::{Model, ModelConfig}};
    26
    //!
    27
    //! // Load GGUF file
    28
    //! let file_data = std::fs::read("model.gguf")?;
    29
    //! let gguf = GGUFModel::from_bytes(&file_data)?;
    30
    //!
    31
    //! // Extract model config from metadata
    32
    //! let config = ModelConfig {
    33
    //!     vocab_size: 32000,
    34
    //!     hidden_dim: 4096,
    35
    //!     num_heads: 32,
    36
    //!     num_layers: 32,
    37
    //!     intermediate_dim: 11008,
    38
    //!     eps: 1e-5,
    39
    //! };
    40
    //!
    41
    //! // Create model and load weights
    42
    //! let mut model = Model::new(config)?;
    43
    //!
    44
    //! // Extract tensors by name (requires knowledge of naming convention)
    45
    //! let embedding_weights = gguf.get_tensor_f32("token_embd.weight", &file_data)?;
    46
    //! // ... load weights into model layers ...
    47
    //! ```
    48
    //!
    49
    //! ### Safetensors Example
    50
    //!
    51
    //! ```rust,ignore
    52
    //! use realizar::{safetensors::SafetensorsModel, layers::{Model, ModelConfig}};
    53
    //!
    54
    //! // Load Safetensors file
    55
    //! let file_data = std::fs::read("model.safetensors")?;
    56
    //! let safetensors = SafetensorsModel::from_bytes(&file_data)?;
    57
    //!
    58
    //! // Extract tensors
    59
    //! let embedding_weights = safetensors.get_tensor_f32("model.embed_tokens.weight")?;
    60
    //! // ... load weights into model layers ...
    61
    //! ```
    62
    //!
    63
    //! ### Tensor Naming Conventions
    64
    //!
    65
    //! Different model families use different tensor naming conventions:
    66
    //!
    67
    //! - **`LLaMA` models (GGUF):**
    68
    //!   - `token_embd.weight` - Token embeddings
    69
    //!   - `blk.{layer}.attn_q.weight` - Query projection for layer N
    70
    //!   - `blk.{layer}.attn_k.weight` - Key projection
    71
    //!   - `blk.{layer}.attn_v.weight` - Value projection
    72
    //!   - `blk.{layer}.ffn_up.weight` - FFN up projection
    73
    //!
    74
    //! - **`HuggingFace` models (Safetensors):**
    75
    //!   - `model.embed_tokens.weight` - Token embeddings
    76
    //!   - `model.layers.{layer}.self_attn.q_proj.weight` - Query projection
    77
    //!   - `model.layers.{layer}.self_attn.k_proj.weight` - Key projection
    78
    //!   - `model.layers.{layer}.mlp.up_proj.weight` - FFN up projection
    79
    //!
    80
    //! Consult model documentation for specific naming conventions.
    81
    82
    use crate::{
    83
        error::{RealizarError, Result},
    84
        tensor::Tensor,
    85
    };
    86
    87
    // PMAT-802: Extracted modules
    88
    mod position;
    89
    pub use position::{RoPE, RopeScalingType, ScaledRoPE, ALiBi};
    90
    mod model;
    91
    pub use model::{KVCache, TransformerBlock, Embedding, Model, ModelConfig};
    92
    mod attention;
    93
    pub use attention::{Attention, SlidingWindowAttention, FusedQKVAttention, MultiHeadAttention};
    94
    95
    /// Apply softmax activation function
    96
    ///
    97
    /// Softmax: `y[i] = exp(x[i]) / sum(exp(x[j]))` for all j
    98
    ///
    99
    /// Applies softmax normalization along the last dimension. Uses numerically stable
    100
    /// implementation with max subtraction to prevent overflow.
    101
    ///
    102
    /// Used in attention mechanisms for probability distributions.
    103
    ///
    104
    /// # Arguments
    105
    ///
    106
    /// * `input` - Input tensor
    107
    ///
    108
    /// # Returns
    109
    ///
    110
    /// Tensor with softmax applied along last dimension (values sum to 1.0)
    111
    ///
    112
    /// # Errors
    113
    ///
    114
    /// Returns error if input is empty
    115
    ///
    116
    /// # Examples
    117
    ///
    118
    /// ```rust,ignore
    119
    /// let input = Tensor::from_vec(vec![3], vec![1.0, 2.0, 3.0])?;
    120
    /// let output = softmax(&input)?;
    121
    /// // output sums to 1.0
    122
    /// ```
    123
    3.39k
    pub fn softmax(input: &Tensor<f32>) -> Result<Tensor<f32>> {
    124
    3.39k
        let data = input.data();
    125
    3.39k
        let shape = input.shape();
    126
    127
    3.39k
        if data.is_empty() {
    128
    0
            return Err(RealizarError::InvalidShape {
    129
    0
                reason: "Cannot apply softmax to empty tensor".to_string(),
    130
    0
            });
    131
    3.39k
        }
    132
    133
    3.39k
        if shape.is_empty() {
    134
    0
            return Err(RealizarError::InvalidShape {
    135
    0
                reason: "Cannot apply softmax to tensor with empty shape".to_string(),
    136
    0
            });
    137
    3.39k
        }
    138
    139
    3.39k
        let last_dim = shape[shape.len() - 1];
    140
    3.39k
        let num_groups = data.len() / last_dim;
    141
    3.39k
        let mut output = Vec::with_capacity(data.len());
    142
    143
        // Apply softmax to each group (row) independently
    144
    132k
        for group_idx in 0..num_groups3.39k {
    145
    132k
            let start = group_idx * last_dim;
    146
    132k
            let end = start + last_dim;
    147
    132k
            let group = &data[start..end];
    148
    149
            // Find max for numerical stability
    150
    132k
            let max_val = group.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    151
    152
            // Compute exp(x - max) for each element
    153
    18.2M
            let exp_vals132k: Vec<f32>132k = group132k.iter132k().map132k(|&x| (x - max_val).exp()).collect132k();
    154
    155
            // Sum of exponentials
    156
    132k
            let sum_exp: f32 = exp_vals.iter().sum();
    157
    158
            // Normalize to get probabilities
    159
    18.3M
            for &exp_val18.2M in &exp_vals {
    160
    18.2M
                output.push(exp_val / sum_exp);
    161
    18.2M
            }
    162
        }
    163
    164
    3.39k
        Tensor::from_vec(shape.to_vec(), output)
    165
    3.39k
    }
    166
    167
    /// Apply GELU activation function
    168
    ///
    169
    /// GELU (Gaussian Error Linear Unit): `y = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))`
    170
    ///
    171
    /// Used in transformer models like BERT, GPT-2, GPT-3.
    172
    ///
    173
    /// # Arguments
    174
    ///
    175
    /// * `input` - Input tensor
    176
    ///
    177
    /// # Returns
    178
    ///
    179
    /// Tensor with GELU applied element-wise
    180
    ///
    181
    /// # Errors
    182
    ///
    183
    /// Returns error if input is empty
    184
    ///
    185
    /// # Examples
    186
    ///
    187
    /// ```rust,ignore
    188
    /// let input = Tensor::from_vec(vec![3], vec![-1.0, 0.0, 1.0])?;
    189
    /// let output = gelu(&input)?;
    190
    /// ```
    191
    1.60k
    pub fn gelu(input: &Tensor<f32>) -> Result<Tensor<f32>> {
    192
    1.60k
        let data = input.data();
    193
    1.60k
        if data.is_empty() {
    194
    0
            return Err(RealizarError::InvalidShape {
    195
    0
                reason: "Cannot apply GELU to empty tensor".to_string(),
    196
    0
            });
    197
    1.60k
        }
    198
    199
        // Apply GELU activation using approximation
    200
    1.60k
        let output: Vec<f32> = data
    201
    1.60k
            .iter()
    202
    7.51M
            .map1.60k(|&x| {
    203
                // GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
    204
    7.51M
                let sqrt_2_over_pi = 0.797_884_6; // sqrt(2/π)
    205
    7.51M
                let c = 0.044_715;
    206
    7.51M
                let inner = sqrt_2_over_pi * (x + c * x * x * x);
    207
    7.51M
                0.5 * x * (1.0 + inner.tanh())
    208
    7.51M
            })
    209
    1.60k
            .collect();
    210
    211
    1.60k
        Tensor::from_vec(input.shape().to_vec(), output)
    212
    1.60k
    }
    213
    214
    /// Layer normalization
    215
    ///
    216
    /// Normalizes activations across the feature dimension using:
    217
    /// ```text
    218
    /// y = (x - mean(x)) / sqrt(variance(x) + eps) * gamma + beta
    219
    /// ```
    220
    ///
    221
    /// # References
    222
    ///
    223
    /// Layer Normalization: <https://arxiv.org/abs/1607.06450>
    224
    #[derive(Debug, Clone)]
    225
    pub struct LayerNorm {
    226
        /// Normalized shape (feature dimension)
    227
        normalized_shape: usize,
    228
        /// Epsilon for numerical stability
    229
        eps: f32,
    230
        /// Scale parameter (gamma)
    231
        weight: Vec<f32>,
    232
        /// Shift parameter (beta)
    233
        bias: Vec<f32>,
    234
    }
    235
    236
    impl LayerNorm {
    237
        /// Create a new layer normalization layer
    238
        ///
    239
        /// # Arguments
    240
        ///
    241
        /// * `normalized_shape` - Size of the feature dimension to normalize
    242
        /// * `eps` - Small constant for numerical stability (default: `1e-5`)
    243
        ///
    244
        /// # Errors
    245
        ///
    246
        /// Returns error if `normalized_shape` is zero
    247
        ///
    248
        /// # Examples
    249
        ///
    250
        /// ```rust,ignore
    251
        /// let layer_norm = LayerNorm::new(512, 1e-5)?;
    252
        /// ```
    253
    574
        pub fn new(normalized_shape: usize, eps: f32) -> Result<Self> {
    254
    574
            if normalized_shape == 0 {
    255
    4
                return Err(RealizarError::InvalidShape {
    256
    4
                    reason: "normalized_shape must be > 0".to_string(),
    257
    4
                });
    258
    570
            }
    259
    260
            // Initialize weight (gamma) to 1.0
    261
    570
            let weight = vec![1.0; normalized_shape];
    262
            // Initialize bias (beta) to 0.0
    263
    570
            let bias = vec![0.0; normalized_shape];
    264
    265
    570
            Ok(Self {
    266
    570
                normalized_shape,
    267
    570
                eps,
    268
    570
                weight,
    269
    570
                bias,
    270
    570
            })
    271
    574
        }
    272
    273
        /// Forward pass through layer normalization
    274
        ///
    275
        /// # Arguments
    276
        ///
    277
        /// * `input` - Input tensor with shape `[..., normalized_shape]`
    278
        ///
    279
        /// # Returns
    280
        ///
    281
        /// Normalized tensor with same shape as input
    282
        ///
    283
        /// # Errors
    284
        ///
    285
        /// Returns error if:
    286
        /// - Input is empty
    287
        /// - Last dimension doesn't match `normalized_shape`
    288
        ///
    289
        /// # Examples
    290
        ///
    291
        /// ```rust,ignore
    292
        /// let input = Tensor::from_vec(vec![2, 512], data)?;
    293
        /// let output = layer_norm.forward(&input)?;
    294
        /// assert_eq!(output.shape(), &[2, 512]);
    295
        /// ```
    296
    6.13k
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    297
    6.13k
            let shape = input.shape();
    298
    6.13k
            if shape.is_empty() {
    299
    0
                return Err(RealizarError::InvalidShape {
    300
    0
                    reason: "Input tensor cannot be empty".to_string(),
    301
    0
                });
    302
    6.13k
            }
    303
    304
    6.13k
            let last_dim = shape[shape.len() - 1];
    305
    6.13k
            if last_dim != self.normalized_shape {
    306
    2
                return Err(RealizarError::InvalidShape {
    307
    2
                    reason: format!(
    308
    2
                        "Last dimension {} doesn't match normalized_shape {}",
    309
    2
                        last_dim, self.normalized_shape
    310
    2
                    ),
    311
    2
                });
    312
    6.13k
            }
    313
    314
    6.13k
            let data = input.data();
    315
    6.13k
            let total_size = data.len();
    316
    6.13k
            let num_groups = total_size / self.normalized_shape;
    317
    318
    6.13k
            let mut output = Vec::with_capacity(total_size);
    319
    320
    356k
            for group_idx in 0..num_groups6.13k {
    321
    356k
                let start = group_idx * self.normalized_shape;
    322
    356k
                let end = start + self.normalized_shape;
    323
    356k
                let group = &data[start..end];
    324
    325
                // Compute mean
    326
                #[allow(clippy::cast_precision_loss)]
    327
    356k
                let mean: f32 = group.iter().sum::<f32>() / self.normalized_shape as f32;
    328
    329
                // Compute variance
    330
                #[allow(clippy::cast_precision_loss)]
    331
    356k
                let variance: f32 = group
    332
    356k
                    .iter()
    333
    13.8M
                    .map356k(|&x| {
    334
    13.8M
                        let diff = x - mean;
    335
    13.8M
                        diff * diff
    336
    13.8M
                    })
    337
    356k
                    .sum::<f32>()
    338
    356k
                    / self.normalized_shape as f32;
    339
    340
                // Normalize and apply affine transformation
    341
    13.8M
                for (i, &x) in group356k.iter356k().enumerate356k() {
    342
    13.8M
                    let normalized = (x - mean) / (variance + self.eps).sqrt();
    343
    13.8M
                    let transformed = normalized * self.weight[i] + self.bias[i];
    344
    13.8M
                    output.push(transformed);
    345
    13.8M
                }
    346
            }
    347
    348
            // Debug assertion for numerical stability
    349
    6.13k
            debug_assert!(
    350
    13.8M
                output.iter()6.13k.all6.13k(|&x| x.is_finite()),
    351
    0
                "LayerNorm produced NaN or Inf values - check input distribution"
    352
            );
    353
    354
    6.13k
            Tensor::from_vec(shape.to_vec(), output)
    355
    6.13k
        }
    356
    357
        /// Get the normalized shape
    358
        #[must_use]
    359
    4
        pub fn normalized_shape(&self) -> usize {
    360
    4
            self.normalized_shape
    361
    4
        }
    362
    363
        /// Get epsilon value
    364
        #[must_use]
    365
    3
        pub fn eps(&self) -> f32 {
    366
    3
            self.eps
    367
    3
        }
    368
    }
    369
    370
    /// Linear transformation layer
    371
    ///
    372
    /// Applies linear transformation: `y = x * W + b`
    373
    /// where W is weight matrix and b is bias vector.
    374
    ///
    375
    /// # References
    376
    ///
    377
    /// Standard fully-connected layer used in neural networks.
    378
    #[derive(Debug, Clone)]
    379
    pub struct Linear {
    380
        /// Input features
    381
        in_features: usize,
    382
        /// Output features
    383
        out_features: usize,
    384
        /// Weight matrix `[in_features, out_features]`
    385
        weight: Vec<f32>,
    386
        /// Bias vector `[out_features]`
    387
        bias: Vec<f32>,
    388
    }
    389
    390
    impl Linear {
    391
        /// Create a new linear layer
    392
        ///
    393
        /// # Arguments
    394
        ///
    395
        /// * `in_features` - Number of input features
    396
        /// * `out_features` - Number of output features
    397
        ///
    398
        /// # Errors
    399
        ///
    400
        /// Returns error if either dimension is zero
    401
        ///
    402
        /// # Examples
    403
        ///
    404
        /// ```rust,ignore
    405
        /// let linear = Linear::new(512, 2048)?;
    406
        /// ```
    407
    1.45k
        pub fn new(in_features: usize, out_features: usize) -> Result<Self> {
    408
    1.45k
            if in_features == 0 || out_features == 01.45k {
    409
    7
                return Err(RealizarError::InvalidShape {
    410
    7
                    reason: "in_features and out_features must be > 0".to_string(),
    411
    7
                });
    412
    1.44k
            }
    413
    414
            // Initialize weights to zero (will be loaded from model)
    415
    1.44k
            let weight = vec![0.0; in_features * out_features];
    416
            // Initialize bias to zero
    417
    1.44k
            let bias = vec![0.0; out_features];
    418
    419
    1.44k
            Ok(Self {
    420
    1.44k
                in_features,
    421
    1.44k
                out_features,
    422
    1.44k
                weight,
    423
    1.44k
                bias,
    424
    1.44k
            })
    425
    1.45k
        }
    426
    427
        /// Forward pass through linear layer
    428
        ///
    429
        /// # Arguments
    430
        ///
    431
        /// * `input` - Input tensor with shape `[batch, in_features]` or `[in_features]`
    432
        ///
    433
        /// # Returns
    434
        ///
    435
        /// Output tensor with shape `[batch, out_features]` or `[out_features]`
    436
        ///
    437
        /// # Errors
    438
        ///
    439
        /// Returns error if:
    440
        /// - Input last dimension doesn't match `in_features`
    441
        /// - Input is empty
    442
        ///
    443
        /// # Examples
    444
        ///
    445
        /// ```rust,ignore
    446
        /// let input = Tensor::from_vec(vec![2, 512], data)?;
    447
        /// let output = linear.forward(&input)?;
    448
        /// assert_eq!(output.shape(), &[2, 2048]);
    449
        /// ```
    450
    11.3k
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    451
    11.3k
            let shape = input.shape();
    452
    11.3k
            if shape.is_empty() {
    453
    0
                return Err(RealizarError::InvalidShape {
    454
    0
                    reason: "Input tensor cannot be empty".to_string(),
    455
    0
                });
    456
    11.3k
            }
    457
    458
    11.3k
            let last_dim = shape[shape.len() - 1];
    459
    11.3k
            if last_dim != self.in_features {
    460
    2
                return Err(RealizarError::InvalidShape {
    461
    2
                    reason: format!(
    462
    2
                        "Last dimension {} doesn't match in_features {}",
    463
    2
                        last_dim, self.in_features
    464
    2
                    ),
    465
    2
                });
    466
    11.3k
            }
    467
    468
    11.3k
            let data = input.data();
    469
    11.3k
            let total_size = data.len();
    470
    11.3k
            let num_rows = total_size / self.in_features;
    471
    472
    11.3k
            let mut output = Vec::with_capacity(num_rows * self.out_features);
    473
    474
            // For each input row, compute: output = input * weight + bias
    475
    807k
            for row_idx in 0..num_rows11.3k {
    476
    807k
                let input_start = row_idx * self.in_features;
    477
    807k
                let input_row = &data[input_start..input_start + self.in_features];
    478
    479
                // Matrix-vector multiplication: output[j] = sum(input[i] * weight[i][j]) + bias[j]
    480
    37.0M
                for j in 0..self.out_features807k {
    481
    37.0M
                    let mut sum = self.bias[j];
    482
    1.40G
                    for (i, &input_val) in input_row37.0M.iter37.0M().enumerate37.0M() {
    483
    1.40G
                        sum += input_val * self.weight[i * self.out_features + j];
    484
    1.40G
                    }
    485
    37.0M
                    output.push(sum);
    486
                }
    487
            }
    488
    489
            // Construct output shape
    490
    11.3k
            let mut output_shape = shape[..shape.len() - 1].to_vec();
    491
    11.3k
            output_shape.push(self.out_features);
    492
    493
            // Debug assertion for numerical stability - catch exploding activations early
    494
    11.3k
            debug_assert!(
    495
    37.0M
                output.iter()11.3k.all11.3k(|&x| x.is_finite()),
    496
    0
                "Linear layer produced NaN or Inf values - check for exploding gradients/activations"
    497
            );
    498
    499
    11.3k
            Tensor::from_vec(output_shape, output)
    500
    11.3k
        }
    501
    502
        /// Get input features
    503
        #[must_use]
    504
    9
        pub fn in_features(&self) -> usize {
    505
    9
            self.in_features
    506
    9
        }
    507
    508
        /// Get output features
    509
        #[must_use]
    510
    9
        pub fn out_features(&self) -> usize {
    511
    9
            self.out_features
    512
    9
        }
    513
    514
        /// Get weight matrix (for loading from model)
    515
        #[must_use]
    516
    37
        pub fn weight_mut(&mut self) -> &mut [f32] {
    517
    37
            &mut self.weight
    518
    37
        }
    519
    520
        /// Get bias vector (for loading from model)
    521
        #[must_use]
    522
    21
        pub fn bias_mut(&mut self) -> &mut [f32] {
    523
    21
            &mut self.bias
    524
    21
        }
    525
    }
    526
    527
    // =============================================================================
    528
    // BENCH-SPRINT-002: QuantizedLinear (Q4_K Inference)
    529
    // Per benchmark-model-runners-spec.md v2.0: Inline dequantization for 8x
    530
    // memory bandwidth reduction vs f32.
    531
    //
    532
    // References:
    533
    // - [21] Dettmers et al. (2022) - "LLM.int8(): 8-bit Matrix Multiplication"
    534
    // - [22] Frantar et al. (2022) - "GPTQ: Post-Training Quantization"
    535
    // =============================================================================
    536
    537
    /// Q4_K Quantized Linear Layer
    538
    ///
    539
    /// Memory-efficient linear layer using 4-bit K-quantization (Q4_K) format.
    540
    /// Achieves ~8x memory reduction vs f32 by storing weights as quantized bytes
    541
    /// and performing inline dequantization during matrix-vector multiplication.
    542
    ///
    543
    /// # Performance Characteristics
    544
    ///
    545
    /// - Memory: 4.5 bits/weight (vs 32 bits for f32) → ~7x reduction
    546
    /// - Compute: Fused dequant+dot avoids intermediate f32 tensor
    547
    /// - Bandwidth: Memory-bound, not compute-bound (per memory wall analysis)
    548
    ///
    549
    /// # Format
    550
    ///
    551
    /// Q4_K uses super-blocks of 256 values:
    552
    /// - 144 bytes per super-block
    553
    /// - Contains: d (scale), dmin (min), 12-byte scales, 128-byte quantized values
    554
    ///
    555
    /// # Example
    556
    ///
    557
    /// ```rust,ignore
    558
    /// // Create from raw Q4_K bytes loaded from GGUF model
    559
    /// let layer = QuantizedLinear::new(4096, 4096, q4k_bytes, bias)?;
    560
    /// let output = layer.forward(&activations)?;
    561
    /// ```
    562
    #[derive(Debug, Clone)]
    563
    pub struct QuantizedLinear {
    564
        /// Input features (must be multiple of 256 for Q4_K)
    565
        in_features: usize,
    566
        /// Output features
    567
        out_features: usize,
    568
        /// Q4_K quantized weight bytes [out_features * bytes_per_row]
    569
        weight_bytes: Vec<u8>,
    570
        /// Bias vector [out_features]
    571
        bias: Vec<f32>,
    572
        /// Bytes per output row (super_blocks_per_row * 144)
    573
        bytes_per_row: usize,
    574
    }
    575
    576
    impl QuantizedLinear {
    577
        /// Create a new Q4_K quantized linear layer
    578
        ///
    579
        /// # Arguments
    580
        ///
    581
        /// * `in_features` - Number of input features (should align to 256 for efficiency)
    582
        /// * `out_features` - Number of output features
    583
        /// * `weight_bytes` - Raw Q4_K quantized weight data
    584
        /// * `bias` - Bias vector
    585
        ///
    586
        /// # Errors
    587
        ///
    588
        /// Returns error if:
    589
        /// - Dimensions are zero
    590
        /// - Weight bytes don't match expected size
    591
        /// - Bias length doesn't match out_features
    592
    8
        pub fn new(
    593
    8
            in_features: usize,
    594
    8
            out_features: usize,
    595
    8
            weight_bytes: Vec<u8>,
    596
    8
            bias: Vec<f32>,
    597
    8
        ) -> Result<Self> {
    598
            // Q4_K: 144 bytes per super-block of 256 values
    599
            const SUPER_BLOCK_VALUES: usize = 256;
    600
            const SUPER_BLOCK_BYTES: usize = 144;
    601
    602
    8
            if in_features == 0 || out_features == 07 {
    603
    1
                return Err(RealizarError::InvalidShape {
    604
    1
                    reason: "in_features and out_features must be > 0".to_string(),
    605
    1
                });
    606
    7
            }
    607
    608
    7
            if bias.len() != out_features {
    609
    1
                return Err(RealizarError::InvalidShape {
    610
    1
                    reason: format!(
    611
    1
                        "Bias length {} doesn't match out_features {}",
    612
    1
                        bias.len(),
    613
    1
                        out_features
    614
    1
                    ),
    615
    1
                });
    616
    6
            }
    617
    618
    6
            let super_blocks_per_row = in_features.div_ceil(SUPER_BLOCK_VALUES);
    619
    6
            let bytes_per_row = super_blocks_per_row * SUPER_BLOCK_BYTES;
    620
    6
            let expected_bytes = out_features * bytes_per_row;
    621
    622
    6
            if weight_bytes.len() != expected_bytes {
    623
    0
                return Err(RealizarError::InvalidShape {
    624
    0
                    reason: format!(
    625
    0
                        "Weight bytes {} doesn't match expected {} ({}x{})",
    626
    0
                        weight_bytes.len(),
    627
    0
                        expected_bytes,
    628
    0
                        out_features,
    629
    0
                        bytes_per_row
    630
    0
                    ),
    631
    0
                });
    632
    6
            }
    633
    634
    6
            Ok(Self {
    635
    6
                in_features,
    636
    6
                out_features,
    637
    6
                weight_bytes,
    638
    6
                bias,
    639
    6
                bytes_per_row,
    640
    6
            })
    641
    8
        }
    642
    643
        /// Forward pass through quantized linear layer
    644
        ///
    645
        /// Uses fused dequantization+dot product for memory efficiency.
    646
        /// Per llama.cpp optimization: inline dequant avoids 8x memory traffic penalty.
    647
        ///
    648
        /// # Arguments
    649
        ///
    650
        /// * `input` - Input tensor with shape `[batch, in_features]` or `[in_features]`
    651
        ///
    652
        /// # Returns
    653
        ///
    654
        /// Output tensor with shape `[batch, out_features]` or `[out_features]`
    655
        ///
    656
        /// # Errors
    657
        ///
    658
        /// Returns error if:
    659
        /// - Input tensor is empty
    660
        /// - Input last dimension doesn't match `in_features`
    661
        /// - Quantization format error during fused dequant+dot
    662
    2
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    663
            use crate::quantize::fused_q4k_dot_simd;
    664
    665
    2
            let shape = input.shape();
    666
    2
            if shape.is_empty() {
    667
    0
                return Err(RealizarError::InvalidShape {
    668
    0
                    reason: "Input tensor cannot be empty".to_string(),
    669
    0
                });
    670
    2
            }
    671
    672
    2
            let last_dim = shape[shape.len() - 1];
    673
    2
            if last_dim != self.in_features {
    674
    0
                return Err(RealizarError::InvalidShape {
    675
    0
                    reason: format!(
    676
    0
                        "Last dimension {} doesn't match in_features {}",
    677
    0
                        last_dim, self.in_features
    678
    0
                    ),
    679
    0
                });
    680
    2
            }
    681
    682
    2
            let data = input.data();
    683
    2
            let total_size = data.len();
    684
    2
            let num_rows = total_size / self.in_features;
    685
    686
    2
            let mut output = Vec::with_capacity(num_rows * self.out_features);
    687
    688
            // For each input row, compute output using fused Q4_K dequant+dot
    689
    9
            for row_idx in 0..num_rows2 {
    690
    9
                let input_start = row_idx * self.in_features;
    691
    9
                let input_row = &data[input_start..input_start + self.in_features];
    692
    693
                // For each output column, use fused dequant+dot
    694
    36
                for j in 0..self.out_features9 {
    695
    36
                    let weight_start = j * self.bytes_per_row;
    696
    36
                    let weight_row =
    697
    36
                        &self.weight_bytes[weight_start..weight_start + self.bytes_per_row];
    698
    699
                    // Fused dequantization + dot product (SIMD-accelerated)
    700
    36
                    let dot = fused_q4k_dot_simd(weight_row, input_row)?0;
    701
    36
                    output.push(dot + self.bias[j]);
    702
                }
    703
            }
    704
    705
            // Construct output shape
    706
    2
            let mut output_shape = shape[..shape.len() - 1].to_vec();
    707
    2
            output_shape.push(self.out_features);
    708
    709
            // Handle degenerate case: 1D input produces 1D output
    710
    2
            if output_shape.is_empty() {
    711
    0
                output_shape.push(self.out_features);
    712
    2
            }
    713
    714
    2
            Tensor::from_vec(output_shape, output)
    715
    2
        }
    716
    717
        /// Get input features
    718
        #[must_use]
    719
    5
        pub fn in_features(&self) -> usize {
    720
    5
            self.in_features
    721
    5
        }
    722
    723
        /// Get output features
    724
        #[must_use]
    725
    5
        pub fn out_features(&self) -> usize {
    726
    5
            self.out_features
    727
    5
        }
    728
    729
        /// Get weight bytes (for model inspection)
    730
        #[must_use]
    731
    1
        pub fn weight_bytes(&self) -> &[u8] {
    732
    1
            &self.weight_bytes
    733
    1
        }
    734
    735
        /// Get bias vector
    736
        #[must_use]
    737
    1
        pub fn bias(&self) -> &[f32] {
    738
    1
            &self.bias
    739
    1
        }
    740
    741
        /// Memory usage in bytes (for diagnostics)
    742
        #[must_use]
    743
    1
        pub fn memory_bytes(&self) -> usize {
    744
    1
            self.weight_bytes.len() + self.bias.len() * std::mem::size_of::<f32>()
    745
    1
        }
    746
    }
    747
    748
    /// Fused LayerNorm + Linear layer
    749
    ///
    750
    /// Combines layer normalization and linear transformation in a single pass
    751
    /// to reduce memory bandwidth by avoiding intermediate tensor writes.
    752
    ///
    753
    /// # Algorithm
    754
    ///
    755
    /// Standard (two passes):
    756
    /// ```text
    757
    /// norm_out = LayerNorm(input)   // Full tensor written to memory
    758
    /// output = Linear(norm_out)      // Full tensor read from memory
    759
    /// ```
    760
    ///
    761
    /// Fused (single pass):
    762
    /// ```text
    763
    /// For each row:
    764
    ///   1. Compute mean and variance (in registers)
    765
    ///   2. For each output column:
    766
    ///      a. Normalize input in registers
    767
    ///      b. Apply linear transformation directly
    768
    /// ```
    769
    ///
    770
    /// This reduces memory traffic by ~50% for the LayerNorm→Linear pattern.
    771
    ///
    772
    /// # References
    773
    ///
    774
    /// - "Fused Operations for Deep Learning" - NVIDIA, 2019
    775
    #[derive(Debug, Clone)]
    776
    pub struct FusedLayerNormLinear {
    777
        /// Feature dimension (must match between LayerNorm and Linear input)
    778
        feature_dim: usize,
    779
        /// Output dimension
    780
        out_features: usize,
    781
        /// LayerNorm epsilon
    782
        eps: f32,
    783
        /// LayerNorm weight (gamma)
    784
        norm_weight: Vec<f32>,
    785
        /// LayerNorm bias (beta)
    786
        norm_bias: Vec<f32>,
    787
        /// Linear weight matrix [feature_dim, out_features]
    788
        linear_weight: Vec<f32>,
    789
        /// Linear bias vector [out_features]
    790
        linear_bias: Vec<f32>,
    791
    }
    792
    793
    impl FusedLayerNormLinear {
    794
        /// Create a new fused LayerNorm+Linear layer
    795
        ///
    796
        /// # Arguments
    797
        ///
    798
        /// * `feature_dim` - Input feature dimension (normalized dimension)
    799
        /// * `out_features` - Output dimension of linear layer
    800
        /// * `eps` - LayerNorm epsilon for numerical stability
    801
        ///
    802
        /// # Errors
    803
        ///
    804
        /// Returns error if feature_dim or out_features is zero
    805
    17
        pub fn new(feature_dim: usize, out_features: usize, eps: f32) -> Result<Self> {
    806
    17
            if feature_dim == 0 || out_features == 015 {
    807
    3
                return Err(RealizarError::InvalidShape {
    808
    3
                    reason: "feature_dim and out_features must be > 0".to_string(),
    809
    3
                });
    810
    14
            }
    811
    812
    14
            Ok(Self {
    813
    14
                feature_dim,
    814
    14
                out_features,
    815
    14
                eps,
    816
    14
                norm_weight: vec![1.0; feature_dim],
    817
    14
                norm_bias: vec![0.0; feature_dim],
    818
    14
                linear_weight: vec![0.0; feature_dim * out_features],
    819
    14
                linear_bias: vec![0.0; out_features],
    820
    14
            })
    821
    17
        }
    822
    823
        /// Forward pass with fused LayerNorm + Linear
    824
        ///
    825
        /// Computes `Linear(LayerNorm(input))` in a single pass without
    826
        /// materializing the intermediate normalized tensor.
    827
        ///
    828
        /// # Arguments
    829
        ///
    830
        /// * `input` - Input tensor `[batch, feature_dim]` or `[feature_dim]`
    831
        ///
    832
        /// # Returns
    833
        ///
    834
        /// Output tensor `[batch, out_features]` or `[out_features]`
    835
        ///
    836
        /// # Errors
    837
        ///
    838
        /// Returns error if input dimensions don't match feature_dim
    839
    106
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    840
    106
            let shape = input.shape();
    841
    106
            if shape.is_empty() {
    842
    0
                return Err(RealizarError::InvalidShape {
    843
    0
                    reason: "Input tensor cannot be empty".to_string(),
    844
    0
                });
    845
    106
            }
    846
    847
    106
            let last_dim = shape[shape.len() - 1];
    848
    106
            if last_dim != self.feature_dim {
    849
    2
                return Err(RealizarError::InvalidShape {
    850
    2
                    reason: format!(
    851
    2
                        "Last dimension {} doesn't match feature_dim {}",
    852
    2
                        last_dim, self.feature_dim
    853
    2
                    ),
    854
    2
                });
    855
    104
            }
    856
    857
    104
            let data = input.data();
    858
    104
            let num_rows = data.len() / self.feature_dim;
    859
    104
            let mut output = Vec::with_capacity(num_rows * self.out_features);
    860
    861
    3.26k
            for row_idx in 0..num_rows104 {
    862
    3.26k
                let row_start = row_idx * self.feature_dim;
    863
    3.26k
                let row = &data[row_start..row_start + self.feature_dim];
    864
    865
                // Step 1: Compute mean (in registers)
    866
                #[allow(clippy::cast_precision_loss)]
    867
    3.26k
                let mean: f32 = row.iter().sum::<f32>() / self.feature_dim as f32;
    868
    869
                // Step 2: Compute variance (in registers)
    870
                #[allow(clippy::cast_precision_loss)]
    871
    3.26k
                let variance: f32 = row
    872
    3.26k
                    .iter()
    873
    827k
                    .map3.26k(|&x| {
    874
    827k
                        let diff = x - mean;
    875
    827k
                        diff * diff
    876
    827k
                    })
    877
    3.26k
                    .sum::<f32>()
    878
    3.26k
                    / self.feature_dim as f32;
    879
    880
    3.26k
                let inv_std = 1.0 / (variance + self.eps).sqrt();
    881
    882
                // Step 3: Fused normalize + linear for each output column
    883
                // This avoids writing normalized values to memory
    884
    1.65M
                for j in 0..self.out_features3.26k {
    885
    1.65M
                    let mut sum = self.linear_bias[j];
    886
    423M
                    for (i, &x) in row1.65M.iter1.65M().enumerate1.65M() {
    887
    423M
                        // Normalize in registers
    888
    423M
                        let normalized = (x - mean) * inv_std;
    889
    423M
                        let transformed = normalized * self.norm_weight[i] + self.norm_bias[i];
    890
    423M
                        // Apply linear weight immediately
    891
    423M
                        sum += transformed * self.linear_weight[i * self.out_features + j];
    892
    423M
                    }
    893
    1.65M
                    output.push(sum);
    894
                }
    895
            }
    896
    897
    104
            let mut output_shape = shape[..shape.len() - 1].to_vec();
    898
    104
            output_shape.push(self.out_features);
    899
    900
    104
            Tensor::from_vec(output_shape, output)
    901
    106
        }
    902
    903
        /// Parallel forward pass using rayon
    904
        ///
    905
        /// Parallelizes over rows for multi-core utilization.
    906
        ///
    907
        /// # Errors
    908
        ///
    909
        /// Returns error if:
    910
        /// - Input tensor is empty
    911
        /// - Last dimension doesn't match feature_dim
    912
    104
        pub fn forward_parallel(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    913
            use rayon::prelude::*;
    914
    915
    104
            let shape = input.shape();
    916
    104
            if shape.is_empty() {
    917
    0
                return Err(RealizarError::InvalidShape {
    918
    0
                    reason: "Input tensor cannot be empty".to_string(),
    919
    0
                });
    920
    104
            }
    921
    922
    104
            let last_dim = shape[shape.len() - 1];
    923
    104
            if last_dim != self.feature_dim {
    924
    1
                return Err(RealizarError::InvalidShape {
    925
    1
                    reason: format!(
    926
    1
                        "Last dimension {} doesn't match feature_dim {}",
    927
    1
                        last_dim, self.feature_dim
    928
    1
                    ),
    929
    1
                });
    930
    103
            }
    931
    932
    103
            let data = input.data();
    933
    103
            let num_rows = data.len() / self.feature_dim;
    934
    935
    103
            let output: Vec<f32> = (0..num_rows)
    936
    103
                .into_par_iter()
    937
    3.26k
                .flat_map103(|row_idx| {
    938
    3.26k
                    let row_start = row_idx * self.feature_dim;
    939
    3.26k
                    let row = &data[row_start..row_start + self.feature_dim];
    940
    941
                    // Compute mean and variance
    942
                    #[allow(clippy::cast_precision_loss)]
    943
    3.26k
                    let mean: f32 = row.iter().sum::<f32>() / self.feature_dim as f32;
    944
    945
                    #[allow(clippy::cast_precision_loss)]
    946
    3.26k
                    let variance: f32 = row
    947
    3.26k
                        .iter()
    948
    827k
                        .map3.26k(|&x| {
    949
    827k
                            let diff = x - mean;
    950
    827k
                            diff * diff
    951
    827k
                        })
    952
    3.26k
                        .sum::<f32>()
    953
    3.26k
                        / self.feature_dim as f32;
    954
    955
    3.26k
                    let inv_std = 1.0 / (variance + self.eps).sqrt();
    956
    957
                    // Fused normalize + linear
    958
    3.26k
                    (0..self.out_features)
    959
    1.65M
                        .map3.26k(|j| {
    960
    1.65M
                            let mut sum = self.linear_bias[j];
    961
    423M
                            for (i, &x) in row1.65M.iter1.65M().enumerate1.65M() {
    962
    423M
                                let normalized = (x - mean) * inv_std;
    963
    423M
                                let transformed = normalized * self.norm_weight[i] + self.norm_bias[i];
    964
    423M
                                sum += transformed * self.linear_weight[i * self.out_features + j];
    965
    423M
                            }
    966
    1.65M
                            sum
    967
    1.65M
                        })
    968
    3.26k
                        .collect::<Vec<f32>>()
    969
    3.26k
                })
    970
    103
                .collect();
    971
    972
    103
            let mut output_shape = shape[..shape.len() - 1].to_vec();
    973
    103
            output_shape.push(self.out_features);
    974
    975
    103
            Tensor::from_vec(output_shape, output)
    976
    104
        }
    977
    978
        /// Get feature dimension
    979
        #[must_use]
    980
    4
        pub fn feature_dim(&self) -> usize {
    981
    4
            self.feature_dim
    982
    4
        }
    983
    984
        /// Get output features
    985
        #[must_use]
    986
    3
        pub fn out_features(&self) -> usize {
    987
    3
            self.out_features
    988
    3
        }
    989
    990
        /// Get mutable reference to LayerNorm weight (gamma)
    991
        #[must_use]
    992
    3
        pub fn norm_weight_mut(&mut self) -> &mut [f32] {
    993
    3
            &mut self.norm_weight
    994
    3
        }
    995
    996
        /// Get mutable reference to LayerNorm bias (beta)
    997
        #[must_use]
    998
    3
        pub fn norm_bias_mut(&mut self) -> &mut [f32] {
    999
    3
            &mut self.norm_bias
    1000
    3
        }
    1001
    1002
        /// Get mutable reference to Linear weight
    1003
        #[must_use]
    1004
    6
        pub fn linear_weight_mut(&mut self) -> &mut [f32] {
    1005
    6
            &mut self.linear_weight
    1006
    6
        }
    1007
    1008
        /// Get mutable reference to Linear bias
    1009
        #[must_use]
    1010
    5
        pub fn linear_bias_mut(&mut self) -> &mut [f32] {
    1011
    5
            &mut self.linear_bias
    1012
    5
        }
    1013
    }
    1014
    1015
    /// Feed-forward network (FFN)
    1016
    ///
    1017
    /// Two-layer feed-forward network with GELU activation:
    1018
    /// ```text
    1019
    /// FFN(x) = Linear2(GELU(Linear1(x)))
    1020
    /// ```
    1021
    ///
    1022
    /// Typically used in transformer blocks with:
    1023
    /// - `hidden_dim` = model dimension (e.g., 768, 512)
    1024
    /// - `intermediate_dim` = expansion (typically 4x `hidden_dim`)
    1025
    ///
    1026
    /// # References
    1027
    ///
    1028
    /// Standard transformer FFN from "Attention is All You Need"
    1029
    #[derive(Debug, Clone)]
    1030
    pub struct FeedForward {
    1031
        /// First linear layer (expansion)
    1032
        fc1: Linear,
    1033
        /// Second linear layer (projection)
    1034
        fc2: Linear,
    1035
        /// Hidden dimension
    1036
        hidden_dim: usize,
    1037
        /// Intermediate dimension
    1038
        intermediate_dim: usize,
    1039
    }
    1040
    1041
    impl FeedForward {
    1042
        /// Create a new feed-forward network
    1043
        ///
    1044
        /// # Arguments
    1045
        ///
    1046
        /// * `hidden_dim` - Input/output dimension
    1047
        /// * `intermediate_dim` - Intermediate dimension (typically 4x `hidden_dim`)
    1048
        ///
    1049
        /// # Errors
    1050
        ///
    1051
        /// Returns error if dimensions are zero
    1052
        ///
    1053
        /// # Examples
    1054
        ///
    1055
        /// ```rust,ignore
    1056
        /// let ffn = FeedForward::new(768, 3072)?; // GPT-2 style (4x expansion)
    1057
        /// ```
    1058
    206
        pub fn new(hidden_dim: usize, intermediate_dim: usize) -> Result<Self> {
    1059
    206
            let fc1203 = Linear::new(hidden_dim, intermediate_dim)?3;
    1060
    203
            let fc2 = Linear::new(intermediate_dim, hidden_dim)?0;
    1061
    1062
    203
            Ok(Self {
    1063
    203
                fc1,
    1064
    203
                fc2,
    1065
    203
                hidden_dim,
    1066
    203
                intermediate_dim,
    1067
    203
            })
    1068
    206
        }
    1069
    1070
        /// Forward pass through feed-forward network
    1071
        ///
    1072
        /// # Arguments
    1073
        ///
    1074
        /// * `input` - Input tensor with shape `[..., hidden_dim]`
    1075
        ///
    1076
        /// # Returns
    1077
        ///
    1078
        /// Output tensor with shape `[..., hidden_dim]`
    1079
        ///
    1080
        /// # Errors
    1081
        ///
    1082
        /// Returns error if input shape doesn't match `hidden_dim`
    1083
        ///
    1084
        /// # Examples
    1085
        ///
    1086
        /// ```rust,ignore
    1087
        /// let input = Tensor::from_vec(vec![2, 768], data)?;
    1088
        /// let output = ffn.forward(&input)?;
    1089
        /// assert_eq!(output.shape(), &[2, 768]);
    1090
        /// ```
    1091
    1.59k
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    1092
            // fc1: [hidden_dim] -> [intermediate_dim]
    1093
    1.59k
            let hidden = self.fc1.forward(input)?0;
    1094
    1095
            // GELU activation
    1096
    1.59k
            let activated = gelu(&hidden)?0;
    1097
    1098
            // fc2: [intermediate_dim] -> [hidden_dim]
    1099
    1.59k
            self.fc2.forward(&activated)
    1100
    1.59k
        }
    1101
    1102
        /// Get hidden dimension
    1103
        #[must_use]
    1104
    4
        pub fn hidden_dim(&self) -> usize {
    1105
    4
            self.hidden_dim
    1106
    4
        }
    1107
    1108
        /// Get intermediate dimension
    1109
        #[must_use]
    1110
    3
        pub fn intermediate_dim(&self) -> usize {
    1111
    3
            self.intermediate_dim
    1112
    3
        }
    1113
    1114
        /// Get mutable reference to first linear layer (for loading weights)
    1115
        #[must_use]
    1116
    16
        pub fn fc1_mut(&mut self) -> &mut Linear {
    1117
    16
            &mut self.fc1
    1118
    16
        }
    1119
    1120
        /// Get mutable reference to second linear layer (for loading weights)
    1121
        #[must_use]
    1122
    14
        pub fn fc2_mut(&mut self) -> &mut Linear {
    1123
    14
            &mut self.fc2
    1124
    14
        }
    1125
    }
    1126
    1127
    #[cfg(test)]
    1128
    1129
    #[cfg(test)]
    1130
    mod tests;