## 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/apr_transformer/q4_simd.rs

    Line
    Count
    Source
    1
    //! SIMD-Accelerated Quantized APR Transformer (Q4_0)
    2
    //!
    3
    //! High-performance Q4_0 inference using SIMD-accelerated matmul primitives.
    4
    //! Extracted from apr_transformer/mod.rs (PMAT-802).
    5
    6
    #![allow(clippy::too_many_arguments)]
    7
    #![allow(clippy::similar_names)]
    8
    #![allow(dead_code)]
    9
    10
    use super::{AprTransformerConfig, AprKVCache};
    11
    use crate::error::{RealizarError, Result};
    12
    13
    // SIMD-Accelerated Quantized APR Transformer (Q4_0)
    14
    // ============================================================================
    15
    16
    /// Q4_0 quantized tensor for SIMD-accelerated inference
    17
    ///
    18
    /// Stores raw Q4_0 bytes (18 bytes per 32 values) with dimensions for matmul.
    19
    #[derive(Debug, Clone)]
    20
    pub struct QuantizedAprTensorQ4 {
    21
        /// Raw Q4_0 quantized data
    22
        pub data: Vec<u8>,
    23
        /// Input dimension (columns in weight matrix)
    24
        pub in_dim: usize,
    25
        /// Output dimension (rows in weight matrix)
    26
        pub out_dim: usize,
    27
    }
    28
    29
    impl QuantizedAprTensorQ4 {
    30
        /// Create a new Q4_0 tensor from raw data
    31
        #[must_use]
    32
    3
        pub fn new(data: Vec<u8>, in_dim: usize, out_dim: usize) -> Self {
    33
    3
            Self {
    34
    3
                data,
    35
    3
                in_dim,
    36
    3
                out_dim,
    37
    3
            }
    38
    3
        }
    39
    40
        /// Create empty tensor with proper Q4_0 allocation
    41
        #[must_use]
    42
    258
        pub fn zeros(in_dim: usize, out_dim: usize) -> Self {
    43
            const Q4_0_BLOCK_BYTES: usize = 18;
    44
            const Q4_0_BLOCK_SIZE: usize = 32;
    45
    258
            let num_elements = in_dim * out_dim;
    46
    258
            let num_blocks = num_elements.div_ceil(Q4_0_BLOCK_SIZE);
    47
    258
            let data = vec![0u8; num_blocks * Q4_0_BLOCK_BYTES];
    48
    258
            Self {
    49
    258
                data,
    50
    258
                in_dim,
    51
    258
                out_dim,
    52
    258
            }
    53
    258
        }
    54
    55
        /// Get expected data size in bytes
    56
        #[must_use]
    57
    10
        pub fn expected_bytes(num_elements: usize) -> usize {
    58
            const Q4_0_BLOCK_BYTES: usize = 18;
    59
            const Q4_0_BLOCK_SIZE: usize = 32;
    60
    10
            let num_blocks = num_elements.div_ceil(Q4_0_BLOCK_SIZE);
    61
    10
            num_blocks * Q4_0_BLOCK_BYTES
    62
    10
        }
    63
    }
    64
    65
    /// Q4_0 quantized layer for SIMD-accelerated inference
    66
    ///
    67
    /// Stores individual Q4_0 tensors for each weight matrix, enabling
    68
    /// direct use of `fused_q4_0_q8_0_parallel_matvec`.
    69
    #[derive(Debug, Clone)]
    70
    pub struct QuantizedAprLayerQ4 {
    71
        /// Attention norm weight (F32, small)
    72
        pub attn_norm_weight: Vec<f32>,
    73
        /// QKV projection weights (Q4_0)
    74
        pub qkv_weight: QuantizedAprTensorQ4,
    75
        /// Attention output projection (Q4_0)
    76
        pub attn_output_weight: QuantizedAprTensorQ4,
    77
        /// FFN up projection (Q4_0)
    78
        pub ffn_up_weight: QuantizedAprTensorQ4,
    79
        /// FFN down projection (Q4_0)
    80
        pub ffn_down_weight: QuantizedAprTensorQ4,
    81
        /// FFN gate projection for SwiGLU (Q4_0, optional)
    82
        pub ffn_gate_weight: Option<QuantizedAprTensorQ4>,
    83
        /// FFN norm weight (F32, optional)
    84
        pub ffn_norm_weight: Option<Vec<f32>>,
    85
    }
    86
    87
    /// SIMD-accelerated Quantized APR Transformer
    88
    ///
    89
    /// Stores weights in Q4_0 format and uses integer SIMD matmul
    90
    /// (`_mm256_maddubs_epi16`) for near-GGUF performance.
    91
    ///
    92
    /// # Performance
    93
    ///
    94
    /// Expected throughput: ~17 tok/s on TinyLlama-1.1B (1.36x vs GGUF)
    95
    /// With KV cache: ~25-34 tok/s expected (1.5-2x additional speedup)
    96
    #[derive(Debug, Clone)]
    97
    pub struct QuantizedAprTransformerQ4 {
    98
        /// Model configuration
    99
        pub config: AprTransformerConfig,
    100
        /// Token embedding (F32 for fast lookup)
    101
        pub token_embedding: Vec<f32>,
    102
        /// Quantized layers
    103
        pub layers: Vec<QuantizedAprLayerQ4>,
    104
        /// Output norm weight (F32)
    105
        pub output_norm_weight: Vec<f32>,
    106
        /// LM head weight (Q4_0)
    107
        pub lm_head_weight: QuantizedAprTensorQ4,
    108
    }
    109
    110
    /// Scratch buffer for zero-allocation inference
    111
    ///
    112
    /// Pre-allocates all intermediate buffers needed for a forward pass.
    113
    /// Reuse across multiple forward calls to eliminate per-token allocations.
    114
    #[derive(Debug)]
    115
    pub struct AprInferenceScratch {
    116
        /// Hidden state [hidden_dim]
    117
        pub hidden: Vec<f32>,
    118
        /// Normalized hidden state [hidden_dim]
    119
        pub normed: Vec<f32>,
    120
        /// QKV projection output [qkv_dim]
    121
        pub qkv_out: Vec<f32>,
    122
        /// Query vectors [q_dim]
    123
        pub q: Vec<f32>,
    124
        /// Key vectors [k_dim]
    125
        pub k: Vec<f32>,
    126
        /// Value vectors [v_dim]
    127
        pub v: Vec<f32>,
    128
        /// Attention output [hidden_dim]
    129
        pub attn_out: Vec<f32>,
    130
        /// FFN input [hidden_dim]
    131
        pub ffn_input: Vec<f32>,
    132
        /// FFN up projection [intermediate_dim]
    133
        pub ffn_up: Vec<f32>,
    134
        /// FFN gate projection [intermediate_dim]
    135
        pub ffn_gate: Vec<f32>,
    136
        /// FFN output [hidden_dim]
    137
        pub ffn_out: Vec<f32>,
    138
    }
    139
    140
    impl AprInferenceScratch {
    141
        /// Create scratch buffer sized for a model config
    142
        #[must_use]
    143
    9
        pub fn from_config(config: &AprTransformerConfig) -> Self {
    144
    9
            let hidden_dim = config.hidden_dim;
    145
    9
            let qkv_dim = hidden_dim * 3; // Conservative estimate
    146
    9
            let intermediate_dim = config.intermediate_dim;
    147
    148
    9
            Self {
    149
    9
                hidden: vec![0.0; hidden_dim],
    150
    9
                normed: vec![0.0; hidden_dim],
    151
    9
                qkv_out: vec![0.0; qkv_dim],
    152
    9
                q: vec![0.0; hidden_dim],
    153
    9
                k: vec![0.0; hidden_dim],
    154
    9
                v: vec![0.0; hidden_dim],
    155
    9
                attn_out: vec![0.0; hidden_dim],
    156
    9
                ffn_input: vec![0.0; hidden_dim],
    157
    9
                ffn_up: vec![0.0; intermediate_dim],
    158
    9
                ffn_gate: vec![0.0; intermediate_dim],
    159
    9
                ffn_out: vec![0.0; hidden_dim],
    160
    9
            }
    161
    9
        }
    162
    163
        /// Clear all buffers (set to zero)
    164
    11
        pub fn clear(&mut self) {
    165
    11
            self.hidden.fill(0.0);
    166
    11
            self.normed.fill(0.0);
    167
    11
            self.qkv_out.fill(0.0);
    168
    11
            self.q.fill(0.0);
    169
    11
            self.k.fill(0.0);
    170
    11
            self.v.fill(0.0);
    171
    11
            self.attn_out.fill(0.0);
    172
    11
            self.ffn_input.fill(0.0);
    173
    11
            self.ffn_up.fill(0.0);
    174
    11
            self.ffn_gate.fill(0.0);
    175
    11
            self.ffn_out.fill(0.0);
    176
    11
        }
    177
    }
    178
    179
    impl QuantizedAprTransformerQ4 {
    180
        /// Create from GGUF OwnedQuantizedModel (extracts Q4_0 bytes)
    181
        ///
    182
        /// # Arguments
    183
        ///
    184
        /// * `gguf` - Source GGUF model with Q4_0 weights
    185
        ///
    186
        /// # Returns
    187
        ///
    188
        /// Quantized APR transformer with same weights
    189
    0
        pub fn from_gguf(gguf: &crate::gguf::OwnedQuantizedModel) -> Self {
    190
            use crate::gguf::OwnedQKVWeights;
    191
    192
    0
            let config = AprTransformerConfig {
    193
    0
                architecture: gguf.config.architecture.clone(),
    194
    0
                hidden_dim: gguf.config.hidden_dim,
    195
    0
                num_layers: gguf.config.num_layers,
    196
    0
                num_heads: gguf.config.num_heads,
    197
    0
                num_kv_heads: gguf.config.num_kv_heads,
    198
    0
                vocab_size: gguf.config.vocab_size,
    199
    0
                intermediate_dim: gguf.config.intermediate_dim,
    200
    0
                context_length: gguf.config.context_length,
    201
    0
                rope_theta: gguf.config.rope_theta,
    202
    0
                eps: gguf.config.eps,
    203
    0
            };
    204
    205
    0
            let layers =
    206
    0
                gguf.layers
    207
    0
                    .iter()
    208
    0
                    .map(|l| {
    209
                        // Extract QKV weight data
    210
    0
                        let qkv_weight = match &l.qkv_weight {
    211
    0
                            OwnedQKVWeights::Fused(t) => {
    212
    0
                                QuantizedAprTensorQ4::new(t.data.clone(), t.in_dim, t.out_dim)
    213
                            },
    214
    0
                            OwnedQKVWeights::Separate { q, k, v } => {
    215
                                // Concatenate Q, K, V for fused format
    216
    0
                                let mut data =
    217
    0
                                    Vec::with_capacity(q.data.len() + k.data.len() + v.data.len());
    218
    0
                                data.extend_from_slice(&q.data);
    219
    0
                                data.extend_from_slice(&k.data);
    220
    0
                                data.extend_from_slice(&v.data);
    221
    0
                                QuantizedAprTensorQ4::new(
    222
    0
                                    data,
    223
    0
                                    q.in_dim,                          // hidden_dim
    224
    0
                                    q.out_dim + k.out_dim + v.out_dim, // qkv_dim
    225
                                )
    226
                            },
    227
                        };
    228
    229
                        QuantizedAprLayerQ4 {
    230
    0
                            attn_norm_weight: l.attn_norm_weight.clone(),
    231
    0
                            qkv_weight,
    232
    0
                            attn_output_weight: QuantizedAprTensorQ4::new(
    233
    0
                                l.attn_output_weight.data.clone(),
    234
    0
                                l.attn_output_weight.in_dim,
    235
    0
                                l.attn_output_weight.out_dim,
    236
                            ),
    237
    0
                            ffn_up_weight: QuantizedAprTensorQ4::new(
    238
    0
                                l.ffn_up_weight.data.clone(),
    239
    0
                                l.ffn_up_weight.in_dim,
    240
    0
                                l.ffn_up_weight.out_dim,
    241
                            ),
    242
    0
                            ffn_down_weight: QuantizedAprTensorQ4::new(
    243
    0
                                l.ffn_down_weight.data.clone(),
    244
    0
                                l.ffn_down_weight.in_dim,
    245
    0
                                l.ffn_down_weight.out_dim,
    246
                            ),
    247
    0
                            ffn_gate_weight: l.ffn_gate_weight.as_ref().map(|g| {
    248
    0
                                QuantizedAprTensorQ4::new(g.data.clone(), g.in_dim, g.out_dim)
    249
    0
                            }),
    250
    0
                            ffn_norm_weight: l.ffn_norm_weight.clone(),
    251
                        }
    252
    0
                    })
    253
    0
                    .collect();
    254
    255
    0
            let lm_head_weight = QuantizedAprTensorQ4::new(
    256
    0
                gguf.lm_head_weight.data.clone(),
    257
    0
                gguf.lm_head_weight.in_dim,
    258
    0
                gguf.lm_head_weight.out_dim,
    259
            );
    260
    261
    0
            Self {
    262
    0
                config,
    263
    0
                token_embedding: gguf.token_embedding.clone(),
    264
    0
                layers,
    265
    0
                output_norm_weight: gguf.output_norm_weight.clone(),
    266
    0
                lm_head_weight,
    267
    0
            }
    268
    0
        }
    269
    270
        /// Get model configuration
    271
        #[must_use]
    272
    1
        pub fn config(&self) -> &AprTransformerConfig {
    273
    1
            &self.config
    274
    1
        }
    275
    276
        /// Create a scratch buffer for zero-allocation inference
    277
        ///
    278
        /// # Example
    279
        ///
    280
        /// ```rust,ignore
    281
        /// let model = QuantizedAprTransformerQ4::from_gguf(&gguf);
    282
        /// let mut scratch = model.create_scratch();
    283
        ///
    284
        /// // Reuse scratch across multiple forward passes
    285
        /// for token_id in token_ids {
    286
        ///     let logits = model.forward_single_with_scratch(token_id, &mut scratch)?;
    287
        /// }
    288
        /// ```
    289
        #[must_use]
    290
    5
        pub fn create_scratch(&self) -> AprInferenceScratch {
    291
    5
            AprInferenceScratch::from_config(&self.config)
    292
    5
        }
    293
    294
        /// Forward pass using SIMD-accelerated Q4_0×Q8_0 matmul
    295
        ///
    296
        /// # Arguments
    297
        ///
    298
        /// * `token_ids` - Input token IDs
    299
        ///
    300
        /// # Returns
    301
        ///
    302
        /// Logits over vocabulary
    303
    12
        pub fn forward(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
    304
            use crate::quantize::fused_q4_0_q8_0_parallel_matvec;
    305
    306
    12
            if token_ids.is_empty() {
    307
    1
                return Err(RealizarError::InvalidShape {
    308
    1
                    reason: "Token sequence cannot be empty".to_string(),
    309
    1
                });
    310
    11
            }
    311
    312
    11
            let hidden_dim = self.config.hidden_dim;
    313
    11
            let num_heads = self.config.num_heads;
    314
    11
            let num_kv_heads = self.config.num_kv_heads;
    315
    11
            let head_dim = hidden_dim / num_heads;
    316
    11
            let eps = self.config.eps;
    317
    318
            // 1. Token embedding lookup (F32)
    319
    11
            let seq_len = token_ids.len();
    320
    11
            let mut hidden = Vec::with_capacity(seq_len * hidden_dim);
    321
    50
            for &token_id39 in token_ids {
    322
    39
                let offset = (token_id as usize) * hidden_dim;
    323
    39
                if offset + hidden_dim <= self.token_embedding.len() {
    324
    38
                    hidden.extend_from_slice(&self.token_embedding[offset..offset + hidden_dim]);
    325
    38
                } else {
    326
    1
                    hidden.extend(std::iter::repeat_n(0.0, hidden_dim));
    327
    1
                }
    328
            }
    329
    330
            // 2. Process through transformer layers
    331
    26
            for layer15 in &self.layers {
    332
                // Pre-attention RMS norm
    333
    15
                let mut normed = Vec::with_capacity(hidden.len());
    334
    45
                for s in 0..seq_len15 {
    335
    45
                    let start = s * hidden_dim;
    336
    45
                    let slice = &hidden[start..start + hidden_dim];
    337
    2.84k
                    let sq_sum45: f3245 = slice45.iter45().map45(|x| x * x).sum45();
    338
    45
                    let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    339
    2.84k
                    for (i, &x) in slice45.iter45().enumerate45() {
    340
    2.84k
                        normed.push(x / rms * layer.attn_norm_weight[i]);
    341
    2.84k
                    }
    342
                }
    343
    344
                // QKV projection using SIMD matmul
    345
    15
                let qkv_dim = layer.qkv_weight.out_dim;
    346
    15
                let mut qkv_out = Vec::with_capacity(seq_len * qkv_dim);
    347
    45
                for s in 0..seq_len15 {
    348
    45
                    let input = &normed[s * hidden_dim..(s + 1) * hidden_dim];
    349
    45
                    let qkv = fused_q4_0_q8_0_parallel_matvec(
    350
    45
                        &layer.qkv_weight.data,
    351
    45
                        input,
    352
    45
                        hidden_dim,
    353
    45
                        qkv_dim,
    354
    0
                    )?;
    355
    45
                    qkv_out.extend(qkv);
    356
                }
    357
    358
                // Proper attention with RoPE and causal mask
    359
    15
                let q_dim = num_heads * head_dim;
    360
    15
                let kv_dim = num_kv_heads * head_dim;
    361
    362
                // Extract Q, K, V and apply RoPE to Q and K
    363
    15
                let mut q_all = Vec::with_capacity(seq_len * q_dim);
    364
    15
                let mut k_all = Vec::with_capacity(seq_len * kv_dim);
    365
    15
                let mut v_all = Vec::with_capacity(seq_len * kv_dim);
    366
    367
    45
                for s in 0..seq_len15 {
    368
    45
                    let qkv_start = s * qkv_dim;
    369
    45
    370
    45
                    // Extract Q, K, V for this position (QKV layout: [Q..., K..., V...])
    371
    45
                    let mut q = qkv_out[qkv_start..qkv_start + q_dim].to_vec();
    372
    45
                    let mut k = qkv_out[qkv_start + q_dim..qkv_start + q_dim + kv_dim].to_vec();
    373
    45
                    let v = &qkv_out[qkv_start + q_dim + kv_dim..qkv_start + q_dim + 2 * kv_dim];
    374
    45
    375
    45
                    // Apply RoPE to Q and K (position-dependent rotation)
    376
    45
                    self.apply_rope(&mut q, s, num_heads);
    377
    45
                    self.apply_rope(&mut k, s, num_kv_heads);
    378
    45
    379
    45
                    q_all.extend_from_slice(&q);
    380
    45
                    k_all.extend_from_slice(&k);
    381
    45
                    v_all.extend_from_slice(v);
    382
    45
                }
    383
    384
                // Compute scaled dot-product attention with causal mask
    385
    15
                let attn_output = self.causal_attention(&q_all, &k_all, &v_all, seq_len);
    386
    387
                // Output projection using SIMD matmul
    388
    15
                let mut proj_out = Vec::with_capacity(seq_len * hidden_dim);
    389
    45
                for s in 0..seq_len15 {
    390
    45
                    let input = &attn_output[s * hidden_dim..(s + 1) * hidden_dim];
    391
    45
                    let proj = fused_q4_0_q8_0_parallel_matvec(
    392
    45
                        &layer.attn_output_weight.data,
    393
    45
                        input,
    394
    45
                        layer.attn_output_weight.in_dim,
    395
    45
                        layer.attn_output_weight.out_dim,
    396
    0
                    )?;
    397
    45
                    proj_out.extend(proj);
    398
                }
    399
    400
                // Residual connection
    401
    2.84k
                for i in 0..hidden15.len15() {
    402
    2.84k
                    hidden[i] += proj_out[i];
    403
    2.84k
                }
    404
    405
                // Pre-FFN norm (if present)
    406
    15
                let ffn_input = if let Some(ffn_norm) = &layer.ffn_norm_weight {
    407
    15
                    let mut normed_ffn = Vec::with_capacity(hidden.len());
    408
    45
                    for s in 0..seq_len15 {
    409
    45
                        let start = s * hidden_dim;
    410
    45
                        let slice = &hidden[start..start + hidden_dim];
    411
    2.84k
                        let sq_sum45: f3245 = slice45.iter45().map45(|x| x * x).sum45();
    412
    45
                        let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    413
    2.84k
                        for (i, &x) in slice45.iter45().enumerate45() {
    414
    2.84k
                            normed_ffn.push(x / rms * ffn_norm[i]);
    415
    2.84k
                        }
    416
                    }
    417
    15
                    normed_ffn
    418
                } else {
    419
    0
                    normed.clone()
    420
                };
    421
    422
                // FFN with SwiGLU (sequential to avoid nested parallelism overhead)
    423
    15
                let intermediate_dim = layer.ffn_up_weight.out_dim;
    424
    15
                let ffn_up = if let Some(gate2) = &layer.ffn_gate_weight {
    425
                    // SwiGLU: Sequential up + gate (both matmuls use internal parallelism)
    426
    2
                    let mut ffn_up_out = Vec::with_capacity(seq_len * intermediate_dim);
    427
    2
                    let mut ffn_gate_out = Vec::with_capacity(seq_len * intermediate_dim);
    428
    429
    6
                    for s in 0..seq_len2 {
    430
    6
                        let input = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    431
    432
                        // Up projection
    433
    6
                        let u = fused_q4_0_q8_0_parallel_matvec(
    434
    6
                            &layer.ffn_up_weight.data,
    435
    6
                            input,
    436
    6
                            hidden_dim,
    437
    6
                            intermediate_dim,
    438
    0
                        )?;
    439
    6
                        ffn_up_out.extend(u);
    440
    441
                        // Gate projection
    442
    6
                        let g = fused_q4_0_q8_0_parallel_matvec(
    443
    6
                            &gate.data,
    444
    6
                            input,
    445
    6
                            hidden_dim,
    446
    6
                            intermediate_dim,
    447
    0
                        )?;
    448
    6
                        ffn_gate_out.extend(g);
    449
                    }
    450
    451
                    // Apply SiLU to gate and multiply with up
    452
    768
                    for i in 0..ffn_up_out2.len2() {
    453
    768
                        let silu = ffn_gate_out[i] / (1.0 + (-ffn_gate_out[i]).exp());
    454
    768
                        ffn_up_out[i] *= silu;
    455
    768
                    }
    456
    2
                    ffn_up_out
    457
                } else {
    458
                    // Non-SwiGLU: Sequential up projection + GELU
    459
    13
                    let mut up = Vec::with_capacity(seq_len * intermediate_dim);
    460
    39
                    for s in 0..seq_len13 {
    461
    39
                        let input = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    462
    39
                        let u = fused_q4_0_q8_0_parallel_matvec(
    463
    39
                            &layer.ffn_up_weight.data,
    464
    39
                            input,
    465
    39
                            hidden_dim,
    466
    39
                            intermediate_dim,
    467
    0
                        )?;
    468
    39
                        up.extend(u);
    469
                    }
    470
                    // GELU activation (tanh approximation)
    471
                    const SQRT_2_OVER_PI: f32 = 0.797_884_6;
    472
                    const GELU_COEFF: f32 = 0.044_715;
    473
    4.94k
                    for x4.92k in &mut up {
    474
    4.92k
                        let t = (SQRT_2_OVER_PI * (*x + GELU_COEFF * *x * *x * *x)).tanh();
    475
    4.92k
                        *x = 0.5 * *x * (1.0 + t);
    476
    4.92k
                    }
    477
    13
                    up
    478
                };
    479
    480
                // FFN: down projection
    481
    15
                let mut ffn_down = Vec::with_capacity(seq_len * hidden_dim);
    482
    45
                for s in 0..seq_len15 {
    483
    45
                    let input = &ffn_up[s * intermediate_dim..(s + 1) * intermediate_dim];
    484
    45
                    let down = fused_q4_0_q8_0_parallel_matvec(
    485
    45
                        &layer.ffn_down_weight.data,
    486
    45
                        input,
    487
    45
                        intermediate_dim,
    488
    45
                        hidden_dim,
    489
    0
                    )?;
    490
    45
                    ffn_down.extend(down);
    491
                }
    492
    493
                // Residual connection
    494
    2.84k
                for i in 0..hidden15.len15() {
    495
    2.84k
                    hidden[i] += ffn_down[i];
    496
    2.84k
                }
    497
            }
    498
    499
            // 3. Final RMS norm
    500
    11
            let last_start = (seq_len - 1) * hidden_dim;
    501
    11
            let last_hidden = &hidden[last_start..last_start + hidden_dim];
    502
    672
            let sq_sum11: f3211 = last_hidden11.iter11().map11(|x| x * x).sum11();
    503
    11
            let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    504
    11
            let normed_final: Vec<f32> = last_hidden
    505
    11
                .iter()
    506
    11
                .enumerate()
    507
    672
                .map11(|(i, &x)| x / rms * self.output_norm_weight[i])
    508
    11
                .collect();
    509
    510
            // 4. LM head projection using SIMD matmul
    511
    11
            let vocab_size = self.config.vocab_size;
    512
    11
            let logits = fused_q4_0_q8_0_parallel_matvec(
    513
    11
                &self.lm_head_weight.data,
    514
    11
                &normed_final,
    515
    11
                hidden_dim,
    516
    11
                vocab_size,
    517
    0
            )?;
    518
    519
    11
            Ok(logits)
    520
    12
        }
    521
    522
        /// Create a KV cache for this model
    523
        #[must_use]
    524
    19
        pub fn create_kv_cache(&self) -> AprKVCache {
    525
    19
            AprKVCache::new(&self.config)
    526
    19
        }
    527
    528
        /// Forward pass for a single token using scratch buffer (zero allocation)
    529
        ///
    530
        /// This is the fastest path for autoregressive generation when combined
    531
        /// with `forward_with_cache_and_scratch`. It reuses pre-allocated buffers
    532
        /// to eliminate per-token allocations.
    533
        ///
    534
        /// # Arguments
    535
        ///
    536
        /// * `token_id` - Single token to process
    537
        /// * `scratch` - Pre-allocated scratch buffer (from `create_scratch()`)
    538
        ///
    539
        /// # Returns
    540
        ///
    541
        /// Logits over vocabulary
    542
    13
        pub fn forward_single_with_scratch(
    543
    13
            &self,
    544
    13
            token_id: u32,
    545
    13
            scratch: &mut AprInferenceScratch,
    546
    13
        ) -> Result<Vec<f32>> {
    547
            use crate::quantize::fused_q4_0_q8_0_parallel_matvec_into;
    548
    549
    13
            let hidden_dim = self.config.hidden_dim;
    550
    13
            let num_heads = self.config.num_heads;
    551
    13
            let num_kv_heads = self.config.num_kv_heads;
    552
    13
            let head_dim = hidden_dim / num_heads;
    553
    13
            let eps = self.config.eps;
    554
    555
            // 1. Token embedding lookup (write directly to scratch.hidden)
    556
    13
            let offset = (token_id as usize) * hidden_dim;
    557
    13
            if offset + hidden_dim <= self.token_embedding.len() {
    558
    12
                scratch.hidden[..hidden_dim]
    559
    12
                    .copy_from_slice(&self.token_embedding[offset..offset + hidden_dim]);
    560
    12
            } else {
    561
    1
                scratch.hidden[..hidden_dim].fill(0.0);
    562
    1
            }
    563
    564
            // 2. Process through transformer layers
    565
    26
            for layer13 in &self.layers {
    566
                // Pre-attention RMS norm (reuse scratch.normed)
    567
    832
                let sq_sum13: f3213 = scratch.hidden.iter()13.map13(|x| x * x).sum13();
    568
    13
                let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    569
    832
                for i in 0..hidden_dim13 {
    570
    832
                    scratch.normed[i] = scratch.hidden[i] / rms * layer.attn_norm_weight[i];
    571
    832
                }
    572
    573
                // QKV projection (zero-allocation - write directly to scratch.qkv_out)
    574
    13
                let qkv_dim = layer.qkv_weight.out_dim;
    575
    13
                fused_q4_0_q8_0_parallel_matvec_into(
    576
    13
                    &layer.qkv_weight.data,
    577
    13
                    &scratch.normed[..hidden_dim],
    578
    13
                    hidden_dim,
    579
    13
                    &mut scratch.qkv_out[..qkv_dim],
    580
    0
                )?;
    581
    582
                // Extract Q, K, V and apply RoPE (position=0 for single token)
    583
    13
                let q_dim = num_heads * head_dim;
    584
    13
                let kv_dim = num_kv_heads * head_dim;
    585
    586
    13
                scratch.q[..q_dim].copy_from_slice(&scratch.qkv_out[..q_dim]);
    587
    13
                scratch.k[..kv_dim].copy_from_slice(&scratch.qkv_out[q_dim..q_dim + kv_dim]);
    588
    13
                scratch.v[..kv_dim]
    589
    13
                    .copy_from_slice(&scratch.qkv_out[q_dim + kv_dim..q_dim + 2 * kv_dim]);
    590
    591
                // Apply RoPE at position 0
    592
    13
                self.apply_rope(&mut scratch.q[..q_dim], 0, num_heads);
    593
    13
                self.apply_rope(&mut scratch.k[..kv_dim], 0, num_kv_heads);
    594
    595
                // For single token, attention is trivial: output = V (softmax of 1 element = 1.0)
    596
    13
                let group_size = num_heads / num_kv_heads;
    597
    52
                for head in 0..num_heads13 {
    598
    52
                    let kv_head = head / group_size;
    599
    52
                    let v_offset = kv_head * head_dim;
    600
    52
                    let out_offset = head * head_dim;
    601
    52
                    scratch.attn_out[out_offset..out_offset + head_dim]
    602
    52
                        .copy_from_slice(&scratch.v[v_offset..v_offset + head_dim]);
    603
    52
                }
    604
    605
                // Output projection (write to scratch.ffn_out as temporary)
    606
    13
                fused_q4_0_q8_0_parallel_matvec_into(
    607
    13
                    &layer.attn_output_weight.data,
    608
    13
                    &scratch.attn_out[..hidden_dim],
    609
    13
                    layer.attn_output_weight.in_dim,
    610
    13
                    &mut scratch.ffn_out[..layer.attn_output_weight.out_dim],
    611
    0
                )?;
    612
    613
                // Residual connection (attn)
    614
    832
                for i in 0..hidden_dim13 {
    615
    832
                    scratch.hidden[i] += scratch.ffn_out[i];
    616
    832
                }
    617
    618
                // Pre-FFN norm
    619
    13
                if let Some(ffn_norm) = &layer.ffn_norm_weight {
    620
    832
                    let sq_sum13: f3213 = scratch.hidden.iter()13.map13(|x| x * x).sum13();
    621
    13
                    let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    622
    832
                    for i in 0..hidden_dim13 {
    623
    832
                        scratch.ffn_input[i] = scratch.hidden[i] / rms * ffn_norm[i];
    624
    832
                    }
    625
    0
                } else {
    626
    0
                    scratch.ffn_input[..hidden_dim].copy_from_slice(&scratch.normed[..hidden_dim]);
    627
    0
                }
    628
    629
                // FFN with SwiGLU
    630
    13
                let intermediate_dim = layer.ffn_up_weight.out_dim;
    631
    13
                if let Some(gate1) = &layer.ffn_gate_weight {
    632
                    // Up projection (zero-allocation)
    633
    1
                    fused_q4_0_q8_0_parallel_matvec_into(
    634
    1
                        &layer.ffn_up_weight.data,
    635
    1
                        &scratch.ffn_input[..hidden_dim],
    636
    1
                        hidden_dim,
    637
    1
                        &mut scratch.ffn_up[..intermediate_dim],
    638
    0
                    )?;
    639
    640
                    // Gate projection (zero-allocation)
    641
    1
                    fused_q4_0_q8_0_parallel_matvec_into(
    642
    1
                        &gate.data,
    643
    1
                        &scratch.ffn_input[..hidden_dim],
    644
    1
                        hidden_dim,
    645
    1
                        &mut scratch.ffn_gate[..intermediate_dim],
    646
    0
                    )?;
    647
    648
                    // SwiGLU: silu(gate) * up
    649
    128
                    for i in 0..intermediate_dim1 {
    650
    128
                        let silu = scratch.ffn_gate[i] / (1.0 + (-scratch.ffn_gate[i]).exp());
    651
    128
                        scratch.ffn_up[i] *= silu;
    652
    128
                    }
    653
                } else {
    654
                    // GELU path (zero-allocation)
    655
    12
                    fused_q4_0_q8_0_parallel_matvec_into(
    656
    12
                        &layer.ffn_up_weight.data,
    657
    12
                        &scratch.ffn_input[..hidden_dim],
    658
    12
                        hidden_dim,
    659
    12
                        &mut scratch.ffn_up[..intermediate_dim],
    660
    0
                    )?;
    661
    662
                    const SQRT_2_OVER_PI: f32 = 0.797_884_6;
    663
                    const GELU_COEFF: f32 = 0.044_715;
    664
    1.53k
                    for i in 0..intermediate_dim12 {
    665
    1.53k
                        let x = scratch.ffn_up[i];
    666
    1.53k
                        let t = (SQRT_2_OVER_PI * (x + GELU_COEFF * x * x * x)).tanh();
    667
    1.53k
                        scratch.ffn_up[i] = 0.5 * x * (1.0 + t);
    668
    1.53k
                    }
    669
                }
    670
    671
                // Down projection (write to scratch.ffn_out)
    672
    13
                fused_q4_0_q8_0_parallel_matvec_into(
    673
    13
                    &layer.ffn_down_weight.data,
    674
    13
                    &scratch.ffn_up[..intermediate_dim],
    675
    13
                    intermediate_dim,
    676
    13
                    &mut scratch.ffn_out[..hidden_dim],
    677
    0
                )?;
    678
    679
                // Residual connection (FFN)
    680
    832
                for i in 0..hidden_dim13 {
    681
    832
                    scratch.hidden[i] += scratch.ffn_out[i];
    682
    832
                }
    683
            }
    684
    685
            // 3. Final RMS norm
    686
    832
            let sq_sum13: f3213 = scratch.hidden.iter()13.map13(|x| x * x).sum13();
    687
    13
            let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    688
    832
            for i in 0..hidden_dim13 {
    689
    832
                scratch.normed[i] = scratch.hidden[i] / rms * self.output_norm_weight[i];
    690
    832
            }
    691
    692
            // 4. LM head projection (still allocates - logits must be returned)
    693
    13
            let vocab_size = self.config.vocab_size;
    694
    13
            let mut logits = vec![0.0f32; vocab_size];
    695
    13
            fused_q4_0_q8_0_parallel_matvec_into(
    696
    13
                &self.lm_head_weight.data,
    697
    13
                &scratch.normed[..hidden_dim],
    698
    13
                hidden_dim,
    699
    13
                &mut logits,
    700
    0
            )?;
    701
    702
    13
            Ok(logits)
    703
    13
        }
    704
    705
        /// Forward pass with KV cache for efficient autoregressive generation
    706
        ///
    707
        /// This method only computes attention for the new token(s), reusing
    708
        /// cached K/V from previous positions. Provides 1.5-2x speedup.
    709
        ///
    710
        /// # Arguments
    711
        ///
    712
        /// * `token_ids` - New token IDs to process (typically 1 for generation)
    713
        /// * `cache` - KV cache to use and update
    714
        ///
    715
        /// # Returns
    716
        ///
    717
        /// Logits over vocabulary for the last token
    718
    43
        pub fn forward_with_cache(
    719
    43
            &self,
    720
    43
            token_ids: &[u32],
    721
    43
            cache: &mut AprKVCache,
    722
    43
        ) -> Result<Vec<f32>> {
    723
            use crate::quantize::fused_q4_0_q8_0_parallel_matvec;
    724
    725
    43
            if token_ids.is_empty() {
    726
    1
                return Err(RealizarError::InvalidShape {
    727
    1
                    reason: "Token sequence cannot be empty".to_string(),
    728
    1
                });
    729
    42
            }
    730
    731
    42
            let hidden_dim = self.config.hidden_dim;
    732
    42
            let num_heads = self.config.num_heads;
    733
    42
            let num_kv_heads = self.config.num_kv_heads;
    734
    42
            let head_dim = hidden_dim / num_heads;
    735
    42
            let eps = self.config.eps;
    736
    737
            // Position in the sequence (including cached positions)
    738
    42
            let cache_len = cache.len();
    739
    42
            let new_seq_len = token_ids.len();
    740
    741
            // 1. Token embedding lookup (F32)
    742
    42
            let mut hidden = Vec::with_capacity(new_seq_len * hidden_dim);
    743
    142
            for &token_id100 in token_ids {
    744
    100
                let offset = (token_id as usize) * hidden_dim;
    745
    100
                if offset + hidden_dim <= self.token_embedding.len() {
    746
    100
                    hidden.extend_from_slice(&self.token_embedding[offset..offset + hidden_dim]);
    747
    100
                } else {
    748
    0
                    hidden.extend(std::iter::repeat_n(0.0, hidden_dim));
    749
    0
                }
    750
            }
    751
    752
            // 2. Process through transformer layers
    753
    47
            for (layer_idx, layer) in self.layers.iter()42.enumerate42() {
    754
                // Pre-attention RMS norm
    755
    47
                let mut normed = Vec::with_capacity(hidden.len());
    756
    113
                for s in 0..new_seq_len47 {
    757
    113
                    let start = s * hidden_dim;
    758
    113
                    let slice = &hidden[start..start + hidden_dim];
    759
    7.23k
                    let sq_sum113: f32113 = slice113.iter113().map113(|x| x * x).sum113();
    760
    113
                    let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    761
    7.23k
                    for (i, &x) in slice113.iter113().enumerate113() {
    762
    7.23k
                        normed.push(x / rms * layer.attn_norm_weight[i]);
    763
    7.23k
                    }
    764
                }
    765
    766
                // QKV projection using SIMD matmul (only for new tokens)
    767
    47
                let qkv_dim = layer.qkv_weight.out_dim;
    768
    47
                let mut qkv_out = Vec::with_capacity(new_seq_len * qkv_dim);
    769
    113
                for s in 0..new_seq_len47 {
    770
    113
                    let input = &normed[s * hidden_dim..(s + 1) * hidden_dim];
    771
    113
                    let qkv = fused_q4_0_q8_0_parallel_matvec(
    772
    113
                        &layer.qkv_weight.data,
    773
    113
                        input,
    774
    113
                        hidden_dim,
    775
    113
                        qkv_dim,
    776
    0
                    )?;
    777
    113
                    qkv_out.extend(qkv);
    778
                }
    779
    780
    47
                let q_dim = num_heads * head_dim;
    781
    47
                let kv_dim = num_kv_heads * head_dim;
    782
    783
                // Process new tokens: extract Q, K, V and apply RoPE
    784
    47
                let mut new_q = Vec::with_capacity(new_seq_len * q_dim);
    785
    113
                for s in 0..new_seq_len47 {
    786
    113
                    let qkv_start = s * qkv_dim;
    787
    113
                    let position = cache_len + s;
    788
    113
    789
    113
                    // Extract Q, K, V for this position
    790
    113
                    let mut q = qkv_out[qkv_start..qkv_start + q_dim].to_vec();
    791
    113
                    let mut k = qkv_out[qkv_start + q_dim..qkv_start + q_dim + kv_dim].to_vec();
    792
    113
                    let v =
    793
    113
                        qkv_out[qkv_start + q_dim + kv_dim..qkv_start + q_dim + 2 * kv_dim].to_vec();
    794
    113
    795
    113
                    // Apply RoPE with correct position
    796
    113
                    self.apply_rope(&mut q, position, num_heads);
    797
    113
                    self.apply_rope(&mut k, position, num_kv_heads);
    798
    113
    799
    113
                    new_q.extend_from_slice(&q);
    800
    113
    801
    113
                    // Append to cache (K and V with RoPE applied to K)
    802
    113
                    cache.append(layer_idx, &k, &v);
    803
    113
                }
    804
    805
                // Get full K and V from cache (includes new tokens)
    806
    47
                let (full_k, full_v) = cache.get(layer_idx);
    807
    47
                let total_seq_len = cache.len();
    808
    809
                // Compute attention: new Q attends to all cached K/V
    810
    47
                let attn_output = self.causal_attention_cached(
    811
    47
                    &new_q,
    812
    47
                    full_k,
    813
    47
                    full_v,
    814
    47
                    new_seq_len,
    815
    47
                    total_seq_len,
    816
    47
                    cache_len,
    817
                );
    818
    819
                // Output projection using SIMD matmul
    820
    47
                let mut proj_out = Vec::with_capacity(new_seq_len * hidden_dim);
    821
    113
                for s in 0..new_seq_len47 {
    822
    113
                    let input = &attn_output[s * hidden_dim..(s + 1) * hidden_dim];
    823
    113
                    let proj = fused_q4_0_q8_0_parallel_matvec(
    824
    113
                        &layer.attn_output_weight.data,
    825
    113
                        input,
    826
    113
                        layer.attn_output_weight.in_dim,
    827
    113
                        layer.attn_output_weight.out_dim,
    828
    0
                    )?;
    829
    113
                    proj_out.extend(proj);
    830
                }
    831
    832
                // Residual connection
    833
    7.23k
                for i in 0..hidden47.len47() {
    834
    7.23k
                    hidden[i] += proj_out[i];
    835
    7.23k
                }
    836
    837
                // Pre-FFN norm (if present)
    838
    47
                let ffn_input = if let Some(ffn_norm) = &layer.ffn_norm_weight {
    839
    47
                    let mut normed_ffn = Vec::with_capacity(hidden.len());
    840
    113
                    for s in 0..new_seq_len47 {
    841
    113
                        let start = s * hidden_dim;
    842
    113
                        let slice = &hidden[start..start + hidden_dim];
    843
    7.23k
                        let sq_sum113: f32113 = slice113.iter113().map113(|x| x * x).sum113();
    844
    113
                        let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    845
    7.23k
                        for (i, &x) in slice113.iter113().enumerate113() {
    846
    7.23k
                            normed_ffn.push(x / rms * ffn_norm[i]);
    847
    7.23k
                        }
    848
                    }
    849
    47
                    normed_ffn
    850
                } else {
    851
    0
                    normed.clone()
    852
                };
    853
    854
                // FFN with parallel up/gate for SwiGLU models
    855
    47
                let intermediate_dim = layer.ffn_up_weight.out_dim;
    856
    47
                let ffn_up = if let Some(gate4) = &layer.ffn_gate_weight {
    857
                    // SwiGLU: Parallel FFN up + gate
    858
    4
                    let (ffn_up_result, ffn_gate_result) = rayon::join(
    859
    4
                        || {
    860
    4
                            let mut up = Vec::with_capacity(new_seq_len * intermediate_dim);
    861
    8
                            for s in 0..new_seq_len4 {
    862
    8
                                let input = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    863
    8
                                if let Ok(u) = fused_q4_0_q8_0_parallel_matvec(
    864
    8
                                    &layer.ffn_up_weight.data,
    865
    8
                                    input,
    866
    8
                                    hidden_dim,
    867
    8
                                    intermediate_dim,
    868
    8
                                ) {
    869
    8
                                    up.extend(u);
    870
    8
                                }0
    871
                            }
    872
    4
                            up
    873
    4
                        },
    874
    4
                        || {
    875
    4
                            let mut g = Vec::with_capacity(new_seq_len * intermediate_dim);
    876
    8
                            for s in 0..new_seq_len4 {
    877
    8
                                let input = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    878
    8
                                if let Ok(gv) = fused_q4_0_q8_0_parallel_matvec(
    879
    8
                                    &gate.data,
    880
    8
                                    input,
    881
    8
                                    hidden_dim,
    882
    8
                                    intermediate_dim,
    883
    8
                                ) {
    884
    8
                                    g.extend(gv);
    885
    8
                                }0
    886
                            }
    887
    4
                            g
    888
    4
                        },
    889
                    );
    890
    891
    4
                    let mut up = ffn_up_result;
    892
    1.02k
                    for i in 0..up4.len4() {
    893
    1.02k
                        let silu = ffn_gate_result[i] / (1.0 + (-ffn_gate_result[i]).exp());
    894
    1.02k
                        up[i] *= silu;
    895
    1.02k
                    }
    896
    4
                    up
    897
                } else {
    898
                    // Non-SwiGLU: Sequential + GELU
    899
    43
                    let mut up = Vec::with_capacity(new_seq_len * intermediate_dim);
    900
    105
                    for s in 0..new_seq_len43 {
    901
    105
                        let input = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    902
    105
                        let u = fused_q4_0_q8_0_parallel_matvec(
    903
    105
                            &layer.ffn_up_weight.data,
    904
    105
                            input,
    905
    105
                            hidden_dim,
    906
    105
                            intermediate_dim,
    907
    0
                        )?;
    908
    105
                        up.extend(u);
    909
                    }
    910
                    const SQRT_2_OVER_PI: f32 = 0.797_884_6;
    911
                    const GELU_COEFF: f32 = 0.044_715;
    912
    13.4k
                    for x13.4k in &mut up {
    913
    13.4k
                        let t = (SQRT_2_OVER_PI * (*x + GELU_COEFF * *x * *x * *x)).tanh();
    914
    13.4k
                        *x = 0.5 * *x * (1.0 + t);
    915
    13.4k
                    }
    916
    43
                    up
    917
                };
    918
    919
                // FFN: down projection
    920
    47
                let mut ffn_down = Vec::with_capacity(new_seq_len * hidden_dim);
    921
    113
                for s in 0..new_seq_len47 {
    922
    113
                    let input = &ffn_up[s * intermediate_dim..(s + 1) * intermediate_dim];
    923
    113
                    let down = fused_q4_0_q8_0_parallel_matvec(
    924
    113
                        &layer.ffn_down_weight.data,
    925
    113
                        input,
    926
    113
                        intermediate_dim,
    927
    113
                        hidden_dim,
    928
    0
                    )?;
    929
    113
                    ffn_down.extend(down);
    930
                }
    931
    932
                // Residual connection
    933
    7.23k
                for i in 0..hidden47.len47() {
    934
    7.23k
                    hidden[i] += ffn_down[i];
    935
    7.23k
                }
    936
            }
    937
    938
            // 3. Final RMS norm (only for last token)
    939
    42
            let last_start = (new_seq_len - 1) * hidden_dim;
    940
    42
            let last_hidden = &hidden[last_start..last_start + hidden_dim];
    941
    2.68k
            let sq_sum42: f3242 = last_hidden42.iter42().map42(|x| x * x).sum42();
    942
    42
            let rms = (sq_sum / hidden_dim as f32 + eps).sqrt();
    943
    42
            let normed_final: Vec<f32> = last_hidden
    944
    42
                .iter()
    945
    42
                .enumerate()
    946
    2.68k
                .map42(|(i, &x)| x / rms * self.output_norm_weight[i])
    947
    42
                .collect();
    948
    949
            // 4. LM head projection using SIMD matmul
    950
    42
            let vocab_size = self.config.vocab_size;
    951
    42
            let logits = fused_q4_0_q8_0_parallel_matvec(
    952
    42
                &self.lm_head_weight.data,
    953
    42
                &normed_final,
    954
    42
                hidden_dim,
    955
    42
                vocab_size,
    956
    0
            )?;
    957
    958
    42
            Ok(logits)
    959
    43
        }
    960
    961
        /// Attention with KV cache - new Q attends to all cached K/V
    962
        ///
    963
        /// Parallelizes across attention heads for efficiency.
    964
    47
        fn causal_attention_cached(
    965
    47
            &self,
    966
    47
            new_q: &[f32],
    967
    47
            full_k: &[f32],
    968
    47
            full_v: &[f32],
    969
    47
            new_seq_len: usize,
    970
    47
            _total_seq_len: usize,
    971
    47
            cache_len: usize,
    972
    47
        ) -> Vec<f32> {
    973
            use rayon::prelude::*;
    974
    975
    47
            let num_heads = self.config.num_heads;
    976
    47
            let num_kv_heads = self.config.num_kv_heads;
    977
    47
            let head_dim = self.config.hidden_dim / num_heads;
    978
    47
            let scale = 1.0 / (head_dim as f32).sqrt();
    979
    47
            let group_size = num_heads / num_kv_heads;
    980
    981
    47
            let q_dim = num_heads * head_dim;
    982
    47
            let kv_dim = num_kv_heads * head_dim;
    983
    984
            const PARALLEL_HEAD_THRESHOLD: usize = 4;
    985
    986
    47
            if num_heads < PARALLEL_HEAD_THRESHOLD {
    987
                // Sequential path
    988
    0
                let mut output = vec![0.0f32; new_seq_len * q_dim];
    989
    0
                for head in 0..num_heads {
    990
    0
                    let kv_head = head / group_size;
    991
    0
                    let q_head_offset = head * head_dim;
    992
    0
                    let kv_head_offset = kv_head * head_dim;
    993
    994
    0
                    for i in 0..new_seq_len {
    995
    0
                        let pos = cache_len + i;
    996
    0
                        let mut scores = Vec::with_capacity(pos + 1);
    997
    0
                        let q_start = i * q_dim + q_head_offset;
    998
    999
                        // Attend to all positions up to current (causal)
    1000
    0
                        for j in 0..=pos {
    1001
    0
                            let k_start = j * kv_dim + kv_head_offset;
    1002
    0
                            let mut score = 0.0f32;
    1003
    0
                            for d in 0..head_dim {
    1004
    0
                                score += new_q[q_start + d] * full_k[k_start + d];
    1005
    0
                            }
    1006
    0
                            scores.push(score * scale);
    1007
                        }
    1008
    1009
                        // Softmax
    1010
    0
                        let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1011
    0
                        let mut exp_sum = 0.0f32;
    1012
    0
                        for s in &mut scores {
    1013
    0
                            *s = (*s - max_score).exp();
    1014
    0
                            exp_sum += *s;
    1015
    0
                        }
    1016
    0
                        for s in &mut scores {
    1017
    0
                            *s /= exp_sum;
    1018
    0
                        }
    1019
    1020
                        // Weighted sum
    1021
    0
                        let out_start = i * q_dim + q_head_offset;
    1022
    0
                        for (j, &weight) in scores.iter().enumerate() {
    1023
    0
                            let v_start = j * kv_dim + kv_head_offset;
    1024
    0
                            for d in 0..head_dim {
    1025
    0
                                output[out_start + d] += weight * full_v[v_start + d];
    1026
    0
                            }
    1027
                        }
    1028
                    }
    1029
                }
    1030
    0
                output
    1031
            } else {
    1032
                // Parallel path
    1033
    47
                let head_outputs: Vec<Vec<f32>> = (0..num_heads)
    1034
    47
                    .into_par_iter()
    1035
    188
                    .map47(|head| {
    1036
    188
                        let mut head_out = vec![0.0f32; new_seq_len * head_dim];
    1037
    188
                        let kv_head = head / group_size;
    1038
    188
                        let q_head_offset = head * head_dim;
    1039
    188
                        let kv_head_offset = kv_head * head_dim;
    1040
    1041
    452
                        for i in 0..new_seq_len188 {
    1042
    452
                            let pos = cache_len + i;
    1043
    452
                            let mut scores = Vec::with_capacity(pos + 1);
    1044
    452
                            let q_start = i * q_dim + q_head_offset;
    1045
    1046
    2.48k
                            for j in 0..=pos452 {
    1047
    2.48k
                                let k_start = j * kv_dim + kv_head_offset;
    1048
    2.48k
                                let mut score = 0.0f32;
    1049
    39.7k
                                for d in 0..head_dim2.48k {
    1050
    39.7k
                                    score += new_q[q_start + d] * full_k[k_start + d];
    1051
    39.7k
                                }
    1052
    2.48k
                                scores.push(score * scale);
    1053
                            }
    1054
    1055
    452
                            let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1056
    452
                            let mut exp_sum = 0.0f32;
    1057
    2.93k
                            for s2.48k in &mut scores {
    1058
    2.48k
                                *s = (*s - max_score).exp();
    1059
    2.48k
                                exp_sum += *s;
    1060
    2.48k
                            }
    1061
    2.93k
                            for s2.48k in &mut scores {
    1062
    2.48k
                                *s /= exp_sum;
    1063
    2.48k
                            }
    1064
    1065
    452
                            let out_start = i * head_dim;
    1066
    2.48k
                            for (j, &weight) in scores.iter()452.enumerate452() {
    1067
    2.48k
                                let v_start = j * kv_dim + kv_head_offset;
    1068
    39.7k
                                for d in 0..head_dim2.48k {
    1069
    39.7k
                                    head_out[out_start + d] += weight * full_v[v_start + d];
    1070
    39.7k
                                }
    1071
                            }
    1072
                        }
    1073
    188
                        head_out
    1074
    188
                    })
    1075
    47
                    .collect();
    1076
    1077
                // Merge
    1078
    47
                let mut output = vec![0.0f32; new_seq_len * q_dim];
    1079
    188
                for (head, head_out) in head_outputs47.into_iter47().enumerate47() {
    1080
    188
                    let head_offset = head * head_dim;
    1081
    452
                    for i in 0..new_seq_len188 {
    1082
    452
                        let src_start = i * head_dim;
    1083
    452
                        let dst_start = i * q_dim + head_offset;
    1084
    452
                        output[dst_start..dst_start + head_dim]
    1085
    452
                            .copy_from_slice(&head_out[src_start..src_start + head_dim]);
    1086
    452
                    }
    1087
                }
    1088
    47
                output
    1089
            }
    1090
    47
        }
    1091
    1092
        /// Get memory footprint in bytes
    1093
        #[must_use]
    1094
    5
        pub fn memory_size(&self) -> usize {
    1095
    5
            let embed_size = self.token_embedding.len() * 4;
    1096
    5
            let norm_size = self.output_norm_weight.len() * 4;
    1097
    5
            let lm_head_size = self.lm_head_weight.data.len();
    1098
    1099
    5
            let layer_size: usize = self
    1100
    5
                .layers
    1101
    5
                .iter()
    1102
    8
                .map5(|l| {
    1103
    8
                    l.attn_norm_weight.len() * 4
    1104
    8
                        + l.qkv_weight.data.len()
    1105
    8
                        + l.attn_output_weight.data.len()
    1106
    8
                        + l.ffn_up_weight.data.len()
    1107
    8
                        + l.ffn_down_weight.data.len()
    1108
    8
                        + l.ffn_gate_weight.as_ref().map_or(0, |g| g.data1.len1())
    1109
    8
                        + l.ffn_norm_weight.as_ref().map_or(0, |n| n.len() * 4)
    1110
    8
                })
    1111
    5
                .sum();
    1112
    1113
    5
            embed_size + norm_size + lm_head_size + layer_size
    1114
    5
        }
    1115
    1116
        /// Apply Rotary Position Embeddings (RoPE) to a tensor
    1117
        ///
    1118
        /// RoPE applies position-dependent rotation to pairs of dimensions,
    1119
        /// enabling the model to learn relative positional information.
    1120
    342
        fn apply_rope(&self, x: &mut [f32], position: usize, num_heads_in_x: usize) {
    1121
    342
            let head_dim = self.config.hidden_dim / self.config.num_heads;
    1122
    342
            let half_dim = head_dim / 2;
    1123
    342
            let theta = self.config.rope_theta;
    1124
    342
            let pos_f32 = position as f32;
    1125
    342
            let head_dim_f32 = head_dim as f32;
    1126
    1127
            // Apply rotation to each head with inline cos/sin computation
    1128
            // Avoids allocation by computing cos/sin on the fly
    1129
    1.37k
            for h in 0..num_heads_in_x342 {
    1130
    1.37k
                let head_start = h * head_dim;
    1131
    1.37k
                let idx2_start = head_start + half_dim;
    1132
    1133
    1.37k
                if idx2_start + half_dim > x.len() {
    1134
    0
                    continue;
    1135
    1.37k
                }
    1136
    1137
                // Process 4 elements at a time for better ILP
    1138
    1.37k
                let mut i = 0;
    1139
    4.09k
                while i + 4 <= half_dim {
    1140
    2.72k
                    // Compute 4 frequencies
    1141
    2.72k
                    let freq0 = 1.0 / theta.powf(2.0 * i as f32 / head_dim_f32);
    1142
    2.72k
                    let freq1 = 1.0 / theta.powf(2.0 * (i + 1) as f32 / head_dim_f32);
    1143
    2.72k
                    let freq2 = 1.0 / theta.powf(2.0 * (i + 2) as f32 / head_dim_f32);
    1144
    2.72k
                    let freq3 = 1.0 / theta.powf(2.0 * (i + 3) as f32 / head_dim_f32);
    1145
    2.72k
    1146
    2.72k
                    // Compute 4 angles
    1147
    2.72k
                    let angle0 = pos_f32 * freq0;
    1148
    2.72k
                    let angle1 = pos_f32 * freq1;
    1149
    2.72k
                    let angle2 = pos_f32 * freq2;
    1150
    2.72k
                    let angle3 = pos_f32 * freq3;
    1151
    2.72k
    1152
    2.72k
                    // Compute cos/sin (use sincos if available for better performance)
    1153
    2.72k
                    let (sin0, cos0) = angle0.sin_cos();
    1154
    2.72k
                    let (sin1, cos1) = angle1.sin_cos();
    1155
    2.72k
                    let (sin2, cos2) = angle2.sin_cos();
    1156
    2.72k
                    let (sin3, cos3) = angle3.sin_cos();
    1157
    2.72k
    1158
    2.72k
                    // Load x1 and x2 values
    1159
    2.72k
                    let x1_0 = x[head_start + i];
    1160
    2.72k
                    let x1_1 = x[head_start + i + 1];
    1161
    2.72k
                    let x1_2 = x[head_start + i + 2];
    1162
    2.72k
                    let x1_3 = x[head_start + i + 3];
    1163
    2.72k
    1164
    2.72k
                    let x2_0 = x[idx2_start + i];
    1165
    2.72k
                    let x2_1 = x[idx2_start + i + 1];
    1166
    2.72k
                    let x2_2 = x[idx2_start + i + 2];
    1167
    2.72k
                    let x2_3 = x[idx2_start + i + 3];
    1168
    2.72k
    1169
    2.72k
                    // Apply rotation: [cos -sin; sin cos] * [x1; x2]
    1170
    2.72k
                    x[head_start + i] = x1_0 * cos0 - x2_0 * sin0;
    1171
    2.72k
                    x[head_start + i + 1] = x1_1 * cos1 - x2_1 * sin1;
    1172
    2.72k
                    x[head_start + i + 2] = x1_2 * cos2 - x2_2 * sin2;
    1173
    2.72k
                    x[head_start + i + 3] = x1_3 * cos3 - x2_3 * sin3;
    1174
    2.72k
    1175
    2.72k
                    x[idx2_start + i] = x1_0 * sin0 + x2_0 * cos0;
    1176
    2.72k
                    x[idx2_start + i + 1] = x1_1 * sin1 + x2_1 * cos1;
    1177
    2.72k
                    x[idx2_start + i + 2] = x1_2 * sin2 + x2_2 * cos2;
    1178
    2.72k
                    x[idx2_start + i + 3] = x1_3 * sin3 + x2_3 * cos3;
    1179
    2.72k
    1180
    2.72k
                    i += 4;
    1181
    2.72k
                }
    1182
    1183
                // Handle remaining elements
    1184
    1.37k
                while i < half_dim {
    1185
    0
                    let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim_f32);
    1186
    0
                    let angle = pos_f32 * freq;
    1187
    0
                    let (sin_val, cos_val) = angle.sin_cos();
    1188
    0
    1189
    0
                    let x1 = x[head_start + i];
    1190
    0
                    let x2 = x[idx2_start + i];
    1191
    0
    1192
    0
                    x[head_start + i] = x1 * cos_val - x2 * sin_val;
    1193
    0
                    x[idx2_start + i] = x1 * sin_val + x2 * cos_val;
    1194
    0
    1195
    0
                    i += 1;
    1196
    0
                }
    1197
            }
    1198
    342
        }
    1199
    1200
        /// Compute scaled dot-product attention with causal mask and GQA support
    1201
        ///
    1202
        /// Implements multi-head attention with Grouped Query Attention (GQA),
    1203
        /// where multiple Q heads share the same K/V heads.
    1204
        ///
    1205
        /// Optimized for single-token inference (seq_len=1).
    1206
    15
        fn causal_attention(&self, q: &[f32], k: &[f32], v: &[f32], seq_len: usize) -> Vec<f32> {
    1207
    15
            let num_heads = self.config.num_heads;
    1208
    15
            let num_kv_heads = self.config.num_kv_heads;
    1209
    15
            let head_dim = self.config.hidden_dim / num_heads;
    1210
    15
            let scale = 1.0 / (head_dim as f32).sqrt();
    1211
    1212
            // GQA: multiple Q heads share each KV head
    1213
    15
            let group_size = num_heads / num_kv_heads;
    1214
    1215
            // Q has num_heads heads, K/V have num_kv_heads heads
    1216
    15
            let q_dim = num_heads * head_dim;
    1217
    15
            let kv_dim = num_kv_heads * head_dim;
    1218
    1219
            // Fast path for single token (common case in autoregressive generation)
    1220
            // With seq_len=1 and causal mask, each head just copies its V vector
    1221
            // (softmax of single element is 1.0)
    1222
    15
            if seq_len == 1 {
    1223
    7
                let mut output = vec![0.0f32; q_dim];
    1224
    28
                for head in 0..num_heads7 {
    1225
    28
                    let kv_head = head / group_size;
    1226
    28
                    let v_offset = kv_head * head_dim;
    1227
    28
                    let out_offset = head * head_dim;
    1228
    28
                    output[out_offset..out_offset + head_dim]
    1229
    28
                        .copy_from_slice(&v[v_offset..v_offset + head_dim]);
    1230
    28
                }
    1231
    7
                return output;
    1232
    8
            }
    1233
    1234
            // General case for seq_len > 1
    1235
            use rayon::prelude::*;
    1236
    1237
            // Parallel threshold - use parallel for 4+ heads
    1238
            const PARALLEL_HEAD_THRESHOLD: usize = 4;
    1239
    1240
    8
            if num_heads < PARALLEL_HEAD_THRESHOLD {
    1241
                // Sequential path for few heads
    1242
    1
                let mut output = vec![0.0f32; seq_len * q_dim];
    1243
    2
                for head in 0..num_heads1 {
    1244
    2
                    self.compute_head_attention(
    1245
    2
                        head,
    1246
    2
                        group_size,
    1247
    2
                        head_dim,
    1248
    2
                        scale,
    1249
    2
                        q,
    1250
    2
                        k,
    1251
    2
                        v,
    1252
    2
                        seq_len,
    1253
    2
                        q_dim,
    1254
    2
                        kv_dim,
    1255
    2
                        &mut output,
    1256
    2
                    );
    1257
    2
                }
    1258
    1
                output
    1259
            } else {
    1260
                // Parallel path - each head computes independently, then merge
    1261
    7
                let head_outputs: Vec<Vec<f32>> = (0..num_heads)
    1262
    7
                    .into_par_iter()
    1263
    32
                    .map7(|head| {
    1264
    32
                        let mut head_out = vec![0.0f32; seq_len * head_dim];
    1265
    32
                        let kv_head = head / group_size;
    1266
    32
                        let q_head_offset = head * head_dim;
    1267
    32
                        let kv_head_offset = kv_head * head_dim;
    1268
    1269
    152
                        for i in 0..seq_len32 {
    1270
    152
                            let mut scores = Vec::with_capacity(i + 1);
    1271
    152
                            let q_start = i * q_dim + q_head_offset;
    1272
    1273
    972
                            for j in 0..=i152 {
    1274
    972
                                let k_start = j * kv_dim + kv_head_offset;
    1275
    972
                                let mut score = 0.0f32;
    1276
    15.3k
                                for d in 0..head_dim972 {
    1277
    15.3k
                                    score += q[q_start + d] * k[k_start + d];
    1278
    15.3k
                                }
    1279
    972
                                scores.push(score * scale);
    1280
                            }
    1281
    1282
                            // Softmax
    1283
    152
                            let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1284
    152
                            let mut exp_sum = 0.0f32;
    1285
    1.12k
                            for s972 in &mut scores {
    1286
    972
                                *s = (*s - max_score).exp();
    1287
    972
                                exp_sum += *s;
    1288
    972
                            }
    1289
    1.12k
                            for s972 in &mut scores {
    1290
    972
                                *s /= exp_sum;
    1291
    972
                            }
    1292
    1293
                            // Weighted sum
    1294
    152
                            let out_start = i * head_dim;
    1295
    972
                            for (j, &weight) in scores.iter()152.enumerate152() {
    1296
    972
                                let v_start = j * kv_dim + kv_head_offset;
    1297
    15.3k
                                for d in 0..head_dim972 {
    1298
    15.3k
                                    head_out[out_start + d] += weight * v[v_start + d];
    1299
    15.3k
                                }
    1300
                            }
    1301
                        }
    1302
    32
                        head_out
    1303
    32
                    })
    1304
    7
                    .collect();
    1305
    1306
                // Merge head outputs into final output
    1307
    7
                let mut output = vec![0.0f32; seq_len * q_dim];
    1308
    32
                for (head, head_out) in head_outputs7.into_iter7().enumerate7() {
    1309
    32
                    let head_offset = head * head_dim;
    1310
    152
                    for i in 0..seq_len32 {
    1311
    152
                        let src_start = i * head_dim;
    1312
    152
                        let dst_start = i * q_dim + head_offset;
    1313
    152
                        output[dst_start..dst_start + head_dim]
    1314
    152
                            .copy_from_slice(&head_out[src_start..src_start + head_dim]);
    1315
    152
                    }
    1316
                }
    1317
    7
                output
    1318
            }
    1319
    15
        }
    1320
    1321
        /// Compute attention for a single head (helper for sequential path)
    1322
        #[allow(clippy::too_many_arguments)]
    1323
    2
        fn compute_head_attention(
    1324
    2
            &self,
    1325
    2
            head: usize,
    1326
    2
            group_size: usize,
    1327
    2
            head_dim: usize,
    1328
    2
            scale: f32,
    1329
    2
            q: &[f32],
    1330
    2
            k: &[f32],
    1331
    2
            v: &[f32],
    1332
    2
            seq_len: usize,
    1333
    2
            q_dim: usize,
    1334
    2
            kv_dim: usize,
    1335
    2
            output: &mut [f32],
    1336
    2
        ) {
    1337
    2
            let kv_head = head / group_size;
    1338
    2
            let q_head_offset = head * head_dim;
    1339
    2
            let kv_head_offset = kv_head * head_dim;
    1340
    1341
    4
            for i in 0..seq_len2 {
    1342
    4
                let mut scores = Vec::with_capacity(i + 1);
    1343
    4
                let q_start = i * q_dim + q_head_offset;
    1344
    1345
    6
                for j in 0..=i4 {
    1346
    6
                    let k_start = j * kv_dim + kv_head_offset;
    1347
    6
                    let mut score = 0.0f32;
    1348
    192
                    for d in 0..head_dim6 {
    1349
    192
                        score += q[q_start + d] * k[k_start + d];
    1350
    192
                    }
    1351
    6
                    scores.push(score * scale);
    1352
                }
    1353
    1354
    4
                let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1355
    4
                let mut exp_sum = 0.0f32;
    1356
    10
                for s6 in &mut scores {
    1357
    6
                    *s = (*s - max_score).exp();
    1358
    6
                    exp_sum += *s;
    1359
    6
                }
    1360
    10
                for s6 in &mut scores {
    1361
    6
                    *s /= exp_sum;
    1362
    6
                }
    1363
    1364
    4
                let out_start = i * q_dim + q_head_offset;
    1365
    6
                for (j, &weight) in scores.iter()4.enumerate4() {
    1366
    6
                    let v_start = j * kv_dim + kv_head_offset;
    1367
    192
                    for d in 0..head_dim6 {
    1368
    192
                        output[out_start + d] += weight * v[v_start + d];
    1369
    192
                    }
    1370
                }
    1371
            }
    1372
    2
        }
    1373
    }
    1374
    1375
    // =============================================================================