## Coverage Report

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

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

    /home/noah/src/realizar/src/gguf/inference/forward/core.rs

    Line
    Count
    Source
    1
    //! Forward pass implementations for OwnedQuantizedModel
    2
    //!
    3
    //! Contains forward, forward_cached methods.
    4
    //! These are the core inference entry points.
    5
    6
    use crate::error::Result;
    7
    use crate::gguf::OwnedQKVWeights;
    8
    use crate::gguf::ops;
    9
    use crate::gguf::OwnedQuantizedModel;
    10
    11
    impl OwnedQuantizedModel {
    12
        /// Forward pass with fused Q4_K operations (IMP-100)
    13
        ///
    14
        /// This is 1.37x faster than dequantized f32 due to reduced memory bandwidth.
    15
        ///
    16
        /// # Arguments
    17
        ///
    18
        /// * `token_ids` - Input token IDs
    19
        ///
    20
        /// # Returns
    21
        ///
    22
        /// Logits for next token prediction [vocab_size]
    23
        ///
    24
        /// # Errors
    25
        ///
    26
        /// Returns error if tensor operations fail
    27
    4
        pub fn forward(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
    28
    4
            let hidden_dim = self.config.hidden_dim;
    29
            // Note: intermediate_dim is encoded in layer weight tensors (in_dim/out_dim)
    30
    4
            let _ = self.config.intermediate_dim;
    31
    32
            // 1. Token embedding lookup (f32, fast)
    33
    4
            let mut hidden = self.embed(token_ids);
    34
    35
            // Detect if model uses RMSNorm (LLaMA-style) or LayerNorm (phi-2 style)
    36
            // LLaMA models have ffn_gate_weight (SwiGLU) and no bias in norms
    37
    4
            let use_rmsnorm = self
    38
    4
                .layers
    39
    4
                .first()
    40
    4
                .is_some_and(|l| l.ffn_gate_weight.is_some() && l.attn_norm_bias0.is_none0());
    41
    42
            // 2. Process through transformer layers with FUSED Q4_K ops
    43
    4
            let cpu_debug_layers = std::env::var("CPU_DEBUG_LAYERS").is_ok();
    44
    4
            for (layer_idx, layer) in self.layers.iter().enumerate() {
    45
                // 2a. Attention layer norm (RMSNorm for LLaMA, LayerNorm for others)
    46
    4
                let normed = if use_rmsnorm {
    47
    0
                    ops::rms_norm(&hidden, &layer.attn_norm_weight, self.config.eps)
    48
                } else {
    49
    4
                    ops::layer_norm(
    50
    4
                        &hidden,
    51
    4
                        &layer.attn_norm_weight,
    52
    4
                        layer.attn_norm_bias.as_deref(),
    53
    4
                        self.config.eps,
    54
                    )
    55
                };
    56
    57
                // CORRECTNESS-011: CPU intermediate debug at L0
    58
    4
                if cpu_debug_layers && layer_idx < 20 {
    59
    0
                    eprintln!(
    60
    0
                        "[CPU-L{}] RMSNorm: first 3 = [{:.4}, {:.4}, {:.4}]",
    61
    0
                        layer_idx, normed[0], normed[1], normed[2]
    62
    0
                    );
    63
    4
                }
    64
    65
                // 2b. QKV projection with FUSED dequant+dot (1.37x faster)
    66
                // Note: qkv_dim may differ from 3*hidden_dim for GQA models
    67
    4
                let qkv_dim = layer.qkv_weight.out_dim();
    68
    4
                let q_dim = layer.qkv_weight.q_dim();
    69
                // For GQA, k_dim and v_dim may be smaller than q_dim
    70
    4
                let k_dim = match &layer.qkv_weight {
    71
    4
                    OwnedQKVWeights::Fused(_) => q_dim,
    72
    0
                    OwnedQKVWeights::Separate { k, .. } => k.out_dim,
    73
                };
    74
    4
                let v_dim = match &layer.qkv_weight {
    75
    4
                    OwnedQKVWeights::Fused(_) => q_dim,
    76
    0
                    OwnedQKVWeights::Separate { v, .. } => v.out_dim,
    77
                };
    78
    4
                let mut qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0;
    79
    4
                if let Some(ref bias0) = layer.qkv_bias {
    80
    0
                    ops::add_bias(&mut qkv, bias);
    81
    4
                }
    82
    83
                // CORRECTNESS-011: Q, K, V before RoPE (after bias)
    84
    4
                if cpu_debug_layers && (layer_idx < 20 || layer_idx == 40 || layer_idx == 50) {
    85
    0
                    eprintln!(
    86
    0
                        "[CPU-L{}] Q (before RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]",
    87
    0
                        layer_idx, qkv[0], qkv[1], qkv[2], qkv[3], qkv[4]
    88
    0
                    );
    89
    0
                    // K starts at q_dim offset
    90
    0
                    eprintln!(
    91
    0
                        "[CPU-L{}] K (before RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]",
    92
    0
                        layer_idx,
    93
    0
                        qkv[q_dim],
    94
    0
                        qkv[q_dim + 1],
    95
    0
                        qkv[q_dim + 2],
    96
    0
                        qkv[q_dim + 3],
    97
    0
                        qkv[q_dim + 4]
    98
    0
                    );
    99
    0
                    // V starts at q_dim + k_dim offset
    100
    0
                    let v_offset = q_dim + k_dim;
    101
    0
                    eprintln!(
    102
    0
                        "[CPU-L{}] V (before RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]",
    103
    0
                        layer_idx,
    104
    0
                        qkv[v_offset],
    105
    0
                        qkv[v_offset + 1],
    106
    0
                        qkv[v_offset + 2],
    107
    0
                        qkv[v_offset + 3],
    108
    0
                        qkv[v_offset + 4]
    109
    0
                    );
    110
    4
                }
    111
    112
                // 2c. Proper attention with RoPE and causal mask (IMP-101)
    113
    4
                let seq_len = token_ids.len();
    114
    115
                // Extract Q, K, V and apply RoPE to Q and K
    116
    4
                let mut q_all = Vec::with_capacity(seq_len * q_dim);
    117
    4
                let mut k_all = Vec::with_capacity(seq_len * k_dim);
    118
    4
                let mut v_all = Vec::with_capacity(seq_len * v_dim);
    119
    120
    8
                for s in 0..seq_len4 {
    121
    8
                    let qkv_start = s * qkv_dim;
    122
    123
                    // Extract Q, K, V for this position (QKV layout: [Q..., K..., V...])
    124
    8
                    let mut q = qkv[qkv_start..qkv_start + q_dim].to_vec();
    125
    8
                    let mut k = qkv[qkv_start + q_dim..qkv_start + q_dim + k_dim].to_vec();
    126
    8
                    let v = &qkv[qkv_start + q_dim + k_dim..qkv_start + q_dim + k_dim + v_dim];
    127
    128
                    // Apply RoPE to Q and K (position-dependent rotation)
    129
                    // GQA: Q has num_heads, K has num_kv_heads
    130
    8
                    self.apply_rope(&mut q, s, self.config.num_heads);
    131
    8
                    self.apply_rope(&mut k, s, self.config.num_kv_heads);
    132
    133
                    // CORRECTNESS-011: Q after RoPE at position 0
    134
    8
                    if cpu_debug_layers && layer_idx < 20 && s == 00 {
    135
    0
                        eprintln!(
    136
    0
                            "[CPU-L{}] Q (after RoPE): first 3 = [{:.4}, {:.4}, {:.4}]",
    137
    0
                            layer_idx, q[0], q[1], q[2]
    138
    0
                        );
    139
    0
                        eprintln!(
    140
    0
                            "[CPU-L{}] K (after RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]",
    141
    0
                            layer_idx, k[0], k[1], k[2], k[3], k[4]
    142
    0
                        );
    143
    0
                        eprintln!(
    144
    0
                            "[CPU-L{}] V: first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]",
    145
    0
                            layer_idx, v[0], v[1], v[2], v[3], v[4]
    146
    0
                        );
    147
    8
                    }
    148
    149
    8
                    q_all.extend_from_slice(&q);
    150
    8
                    k_all.extend_from_slice(&k);
    151
    8
                    v_all.extend_from_slice(v);
    152
                }
    153
    154
                // Compute scaled dot-product attention with causal mask
    155
    4
                let attn_out = self.causal_attention(&q_all, &k_all, &v_all, seq_len);
    156
    157
                // CORRECTNESS-011: Attention output
    158
    4
                if cpu_debug_layers && layer_idx < 20 {
    159
    0
                    eprintln!(
    160
    0
                        "[CPU-L{}] Attn output: first 3 = [{:.4}, {:.4}, {:.4}]",
    161
    0
                        layer_idx, attn_out[0], attn_out[1], attn_out[2]
    162
    0
                    );
    163
    4
                }
    164
    165
                // 2d. Attention output projection with FUSED ops
    166
                // Input is q_dim (attention output), projects back to hidden_dim
    167
    4
                let mut attn_output = self.fused_matmul(&attn_out, &layer.attn_output_weight)?0;
    168
    4
                if let Some(ref bias0) = layer.attn_output_bias {
    169
    0
                    ops::add_bias(&mut attn_output, bias);
    170
    4
                }
    171
    172
                // 2e. Residual connection
    173
    480
                for i in 0..hidden4.len4() {
    174
    480
                    hidden[i] += attn_output[i];
    175
    480
                }
    176
    177
                // 2f. Pre-FFN layer norm (LLaMA uses separate ffn_norm with RMSNorm)
    178
    4
                let ffn_input = if let Some(ref ffn_norm0) = layer.ffn_norm_weight {
    179
                    // LLaMA-style: separate FFN layer norm (use RMSNorm for LLaMA)
    180
    0
                    if use_rmsnorm {
    181
    0
                        ops::rms_norm(&hidden, ffn_norm, self.config.eps)
    182
                    } else {
    183
    0
                        ops::layer_norm(
    184
    0
                            &hidden,
    185
    0
                            ffn_norm,
    186
    0
                            layer.ffn_norm_bias.as_deref(),
    187
    0
                            self.config.eps,
    188
                        )
    189
                    }
    190
                } else {
    191
                    // phi-2 style: no separate FFN norm, use hidden directly
    192
                    // (some models apply attn_norm again, but we've already done residual)
    193
    4
                    hidden.clone()
    194
                };
    195
    196
                // 2g. FFN with SwiGLU or GELU activation
    197
    4
                let ffn_activated = if let Some(ref gate_weight0) = layer.ffn_gate_weight {
    198
                    // SwiGLU path (LLaMA, TinyLlama, Mistral, etc.)
    199
                    // output = down(gate(x) * silu(up(x)))
    200
    0
                    let mut ffn_up = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?;
    201
    0
                    if let Some(ref bias) = layer.ffn_up_bias {
    202
    0
                        ops::add_bias(&mut ffn_up, bias);
    203
    0
                    }
    204
    205
    0
                    let mut ffn_gate = self.fused_matmul(&ffn_input, gate_weight)?;
    206
    0
                    if let Some(ref bias) = layer.ffn_gate_bias {
    207
    0
                        ops::add_bias(&mut ffn_gate, bias);
    208
    0
                    }
    209
    210
                    // SwiGLU: down(silu(gate(x)) * up(x))
    211
                    // Apply SiLU to GATE projection, not up
    212
    0
                    ops::silu(&mut ffn_gate);
    213
    214
                    // Element-wise multiply: silu(gate) * up
    215
    0
                    for i in 0..ffn_gate.len() {
    216
    0
                        ffn_gate[i] *= ffn_up[i];
    217
    0
                    }
    218
    219
    0
                    ffn_gate
    220
                } else {
    221
                    // GELU path (phi-2, GPT-2, etc.)
    222
    4
                    let mut ffn_hidden = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?0;
    223
    4
                    if let Some(ref bias0) = layer.ffn_up_bias {
    224
    0
                        ops::add_bias(&mut ffn_hidden, bias);
    225
    4
                    }
    226
    4
                    ops::gelu(&mut ffn_hidden);
    227
    4
                    ffn_hidden
    228
                };
    229
    230
                // 2g. FFN down projection with FUSED ops
    231
    4
                let mut ffn_output = self.fused_matmul(&ffn_activated, &layer.ffn_down_weight)?0;
    232
    4
                if let Some(ref bias0) = layer.ffn_down_bias {
    233
    0
                    ops::add_bias(&mut ffn_output, bias);
    234
    4
                }
    235
    236
                // Residual connection
    237
    480
                for i in 0..hidden4.len4() {
    238
    480
                    hidden[i] += ffn_output[i];
    239
    480
                }
    240
    241
                // CORRECTNESS-011: Per-layer CPU debug output
    242
    4
                if cpu_debug_layers {
    243
    0
                    let seq_len = token_ids.len();
    244
    0
                    let last_hidden_start = (seq_len - 1) * hidden_dim;
    245
    0
                    let last_h = &hidden[last_hidden_start..last_hidden_start + hidden_dim];
    246
    0
                    let sum: f32 = last_h.iter().sum();
    247
    0
                    let sq_sum: f32 = last_h.iter().map(|x| x * x).sum();
    248
    0
                    let rms = (sq_sum / last_h.len() as f32).sqrt();
    249
    0
                    eprintln!(
    250
    0
                        "[CPU-L{}] After layer: first 3 = [{:.4}, {:.4}, {:.4}], sum = {:.4}, rms = {:.4}",
    251
    0
                        layer_idx, last_h[0], last_h[1], last_h[2], sum, rms
    252
                    );
    253
    4
                }
    254
            }
    255
    256
            // CORRECTNESS-011: CPU hidden state debug output (compare with GPU)
    257
    4
            if std::env::var("CPU_DEBUG").is_ok() {
    258
    0
                let seq_len = token_ids.len();
    259
    0
                let last_hidden_start = (seq_len - 1) * hidden_dim;
    260
    0
                let last_hidden_raw = &hidden[last_hidden_start..last_hidden_start + hidden_dim];
    261
    262
    0
                let sum: f32 = last_hidden_raw.iter().sum();
    263
    0
                let sq_sum: f32 = last_hidden_raw.iter().map(|x| x * x).sum();
    264
    0
                let rms = (sq_sum / last_hidden_raw.len() as f32).sqrt();
    265
    266
    0
                eprintln!("[CORRECTNESS-011] CPU Hidden before output_norm:");
    267
    0
                eprintln!(
    268
    0
                    "  first 5 = {:?}",
    269
    0
                    &last_hidden_raw[..5.min(last_hidden_raw.len())]
    270
                );
    271
    0
                eprintln!("  sum = {:.4}, rms = {:.4}", sum, rms);
    272
    0
                eprintln!("  (GPU shows: sum=466.2486, rms=39.4793)");
    273
    4
            }
    274
    275
            // 3. Final layer norm (RMSNorm for LLaMA, LayerNorm for others)
    276
    4
            let normed = if use_rmsnorm {
    277
    0
                ops::rms_norm(&hidden, &self.output_norm_weight, self.config.eps)
    278
            } else {
    279
    4
                ops::layer_norm(
    280
    4
                    &hidden,
    281
    4
                    &self.output_norm_weight,
    282
    4
                    self.output_norm_bias.as_deref(),
    283
    4
                    self.config.eps,
    284
                )
    285
            };
    286
    287
            // CORRECTNESS-011: CPU normed hidden state debug output
    288
    4
            if std::env::var("CPU_DEBUG").is_ok() {
    289
    0
                let seq_len = token_ids.len();
    290
    0
                let last_normed_start = (seq_len - 1) * hidden_dim;
    291
    0
                let last_normed = &normed[last_normed_start..last_normed_start + hidden_dim];
    292
    293
    0
                let sum: f32 = last_normed.iter().sum();
    294
    0
                let sq_sum: f32 = last_normed.iter().map(|x| x * x).sum();
    295
    0
                let rms = (sq_sum / last_normed.len() as f32).sqrt();
    296
    297
    0
                eprintln!("[CORRECTNESS-011] CPU Normed hidden:");
    298
    0
                eprintln!("  first 5 = {:?}", &last_normed[..5.min(last_normed.len())]);
    299
    0
                eprintln!("  sum = {:.4}, rms = {:.4}", sum, rms);
    300
    0
                eprintln!("  (GPU shows: sum=107.5945, rms=4.6616)");
    301
    4
            }
    302
    303
            // 4. LM head projection with FUSED ops (only last token)
    304
    4
            let seq_len = token_ids.len();
    305
    4
            let last_hidden_start = (seq_len - 1) * hidden_dim;
    306
    4
            let last_hidden = &normed[last_hidden_start..last_hidden_start + hidden_dim];
    307
    308
            // Compute logits using fused op
    309
    4
            let mut logits = self.fused_matmul(last_hidden, &self.lm_head_weight)?0;
    310
    311
    4
            if let Some(ref bias0) = self.lm_head_bias {
    312
    0
                ops::add_bias(&mut logits, bias);
    313
    4
            }
    314
    315
    4
            Ok(logits)
    316
    4
        }
    317
    318
        /// Forward pass with KV cache for efficient autoregressive decoding
    319
        ///
    320
        /// This method properly handles both architectures:
    321
        /// - LLaMA-style: RMSNorm, SwiGLU FFN, GQA attention
    322
        /// - phi-2 style: LayerNorm, GELU FFN, MHA attention
    323
        ///
    324
        /// Uses O(n) per-token cost instead of O(n²) by caching K/V.
    325
        ///
    326
        /// # Arguments
    327
        /// * `token_id` - Token to process
    328
        /// * `cache` - KV cache for all layers
    329
        /// * `position` - Position in sequence for RoPE
    330
        ///
    331
        /// # Returns
    332
        /// Logits for next token prediction [vocab_size]
    333
    16
        pub fn forward_cached(
    334
    16
            &self,
    335
    16
            token_id: u32,
    336
    16
            cache: &mut crate::gguf::OwnedQuantizedKVCache,
    337
    16
            position: usize,
    338
    16
        ) -> Result<Vec<f32>> {
    339
    16
            let hidden_dim = self.config.hidden_dim;
    340
    341
            // Detect architecture: LLaMA uses RMSNorm (no bias) and SwiGLU (has gate weight)
    342
    16
            let use_rmsnorm = self
    343
    16
                .layers
    344
    16
                .first()
    345
    16
                .is_some_and(|l| l.ffn_gate_weight.is_some() && l.attn_norm_bias0.is_none0());
    346
    347
            // 1. Token embedding lookup
    348
    16
            let mut hidden = self.embed(&[token_id]);
    349
    350
            // PAR-052: Debug output for OwnedQuantizedModel forward path
    351
    16
            let debug_forward = std::env::var("REALIZAR_DEBUG_FORWARD").is_ok();
    352
    16
            if debug_forward && position == 00 {
    353
    0
                eprintln!("[PAR-052] OwnedQuantizedModel::forward_cached");
    354
    0
                eprintln!("[PAR-052] Token ID: {}, Position: {}", token_id, position);
    355
    0
                eprintln!("[PAR-052] use_rmsnorm: {}", use_rmsnorm);
    356
    0
                eprintln!(
    357
    0
                    "[PAR-052] Embedding[0..8]: {:?}",
    358
    0
                    &hidden[..8.min(hidden.len())]
    359
    0
                );
    360
    0
                eprintln!("[PAR-052] Embedding sum: {:.6}", hidden.iter().sum::<f32>());
    361
    16
            }
    362
    363
            // 2. Process through transformer layers
    364
    16
            for (layer_idx, layer) in self.layers.iter().enumerate() {
    365
                // 2a. Attention layer norm (RMSNorm for LLaMA, LayerNorm for phi-2)
    366
    16
                let normed = if use_rmsnorm {
    367
    0
                    ops::rms_norm(&hidden, &layer.attn_norm_weight, self.config.eps)
    368
                } else {
    369
    16
                    ops::layer_norm(
    370
    16
                        &hidden,
    371
    16
                        &layer.attn_norm_weight,
    372
    16
                        layer.attn_norm_bias.as_deref(),
    373
    16
                        self.config.eps,
    374
                    )
    375
                };
    376
    377
                // PAR-052: Debug layer 0 normed and QKV values
    378
    16
                if debug_forward && layer_idx == 00 && position == 00 {
    379
    0
                    eprintln!(
    380
    0
                        "[PAR-052-L0] attn_norm[0..8]: {:?}",
    381
    0
                        &layer.attn_norm_weight[..8.min(layer.attn_norm_weight.len())]
    382
    0
                    );
    383
    0
                    eprintln!(
    384
    0
                        "[PAR-052-L0] normed[0..8]: {:?}",
    385
    0
                        &normed[..8.min(normed.len())]
    386
    0
                    );
    387
    0
                    eprintln!("[PAR-052-L0] normed sum: {:.6}", normed.iter().sum::<f32>());
    388
    16
                }
    389
    390
                // 2b. QKV projection
    391
    16
                let _qkv_dim = layer.qkv_weight.out_dim();
    392
    16
                let q_dim = layer.qkv_weight.q_dim();
    393
    16
                let k_dim = match &layer.qkv_weight {
    394
    16
                    OwnedQKVWeights::Fused(_) => q_dim,
    395
    0
                    OwnedQKVWeights::Separate { k, .. } => k.out_dim,
    396
                };
    397
    16
                let v_dim = match &layer.qkv_weight {
    398
    16
                    OwnedQKVWeights::Fused(_) => q_dim,
    399
    0
                    OwnedQKVWeights::Separate { v, .. } => v.out_dim,
    400
                };
    401
    402
    16
                let mut qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0;
    403
    16
                if let Some(ref bias0) = layer.qkv_bias {
    404
    0
                    ops::add_bias(&mut qkv, bias);
    405
    16
                }
    406
    407
                // PAR-052: Debug QKV after projection
    408
    16
                if debug_forward && layer_idx == 00 && position == 00 {
    409
    0
                    eprintln!(
    410
    0
                        "[PAR-052-L0] QKV dims: q={}, k={}, v={}, total={}",
    411
    0
                        q_dim,
    412
    0
                        k_dim,
    413
    0
                        v_dim,
    414
    0
                        qkv.len()
    415
    0
                    );
    416
    0
                    eprintln!("[PAR-052-L0] QKV sum: {:.6}", qkv.iter().sum::<f32>());
    417
    0
                    eprintln!("[PAR-052-L0] Q[0..8]: {:?}", &qkv[..8.min(q_dim)]);
    418
    16
                }
    419
    420
                // 2c. Extract Q, K, V and apply RoPE
    421
    16
                let mut q = qkv[0..q_dim].to_vec();
    422
    16
                let mut k = qkv[q_dim..q_dim + k_dim].to_vec();
    423
    16
                let v = qkv[q_dim + k_dim..q_dim + k_dim + v_dim].to_vec();
    424
    425
                // Apply RoPE with correct head counts for GQA
    426
    16
                self.apply_rope(&mut q, position, self.config.num_heads);
    427
    16
                self.apply_rope(&mut k, position, self.config.num_kv_heads);
    428
    429
                // PAR-052: Debug Q after RoPE
    430
    16
                if debug_forward && layer_idx == 00 && position == 00 {
    431
    0
                    eprintln!(
    432
    0
                        "[PAR-052-L0] Q after RoPE[0..8]: {:?}",
    433
    0
                        &q[..8.min(q.len())]
    434
    0
                    );
    435
    0
                    eprintln!(
    436
    0
                        "[PAR-052-L0] K after RoPE[0..4]: {:?}",
    437
    0
                        &k[..4.min(k.len())]
    438
    0
                    );
    439
    16
                }
    440
    441
                // 2d. Compute attention using cached K/V
    442
    16
                let k_cache = cache.get_k(layer_idx);
    443
    16
                let v_cache = cache.get_v(layer_idx);
    444
    445
    16
                let attn_out = if k_cache.is_empty() {
    446
                    // First token - just use V directly (self-attention with single token)
    447
                    // Expand V if GQA (num_kv_heads < num_heads)
    448
    3
                    if self.config.num_kv_heads < self.config.num_heads {
    449
    0
                        let head_dim = hidden_dim / self.config.num_heads;
    450
    0
                        let group_size = self.config.num_heads / self.config.num_kv_heads;
    451
    0
                        (0..self.config.num_heads)
    452
    0
                            .flat_map(|h| {
    453
    0
                                let kv_head = h / group_size;
    454
    0
                                let start = kv_head * head_dim;
    455
    0
                                v[start..start + head_dim].iter().copied()
    456
    0
                            })
    457
    0
                            .collect()
    458
                    } else {
    459
    3
                        v.clone()
    460
                    }
    461
                } else {
    462
                    // Use cached K/V for attention with GQA support
    463
    13
                    self.attention_with_cache_gqa(&q, k_cache, v_cache, &k, &v)
    464
                };
    465
    466
                // 2e. Store K and V in cache (store original size, not expanded)
    467
    16
                cache.append(layer_idx, &k, &v);
    468
    469
                // 2f. Attention output projection
    470
    16
                let mut attn_output = self.fused_matmul(&attn_out, &layer.attn_output_weight)?0;
    471
    16
                if let Some(ref bias0) = layer.attn_output_bias {
    472
    0
                    ops::add_bias(&mut attn_output, bias);
    473
    16
                }
    474
    475
                // 2g. Residual connection
    476
    1.02k
                for i in 0..hidden_dim16 {
    477
    1.02k
                    hidden[i] += attn_output[i];
    478
    1.02k
                }
    479
    480
                // 2h. Pre-FFN layer norm (LLaMA has separate ffn_norm)
    481
    16
                let ffn_input = if let Some(ref ffn_norm0) = layer.ffn_norm_weight {
    482
    0
                    if use_rmsnorm {
    483
    0
                        ops::rms_norm(&hidden, ffn_norm, self.config.eps)
    484
                    } else {
    485
    0
                        ops::layer_norm(
    486
    0
                            &hidden,
    487
    0
                            ffn_norm,
    488
    0
                            layer.ffn_norm_bias.as_deref(),
    489
    0
                            self.config.eps,
    490
                        )
    491
                    }
    492
                } else {
    493
    16
                    hidden.clone()
    494
                };
    495
    496
                // 2i. FFN with SwiGLU or GELU
    497
    16
                let ffn_output = if let Some(ref gate_weight0) = layer.ffn_gate_weight {
    498
                    // SwiGLU path (LLaMA)
    499
    0
                    let mut ffn_up = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?;
    500
    0
                    if let Some(ref bias) = layer.ffn_up_bias {
    501
    0
                        ops::add_bias(&mut ffn_up, bias);
    502
    0
                    }
    503
    504
    0
                    let mut ffn_gate = self.fused_matmul(&ffn_input, gate_weight)?;
    505
    0
                    if let Some(ref bias) = layer.ffn_gate_bias {
    506
    0
                        ops::add_bias(&mut ffn_gate, bias);
    507
    0
                    }
    508
    509
                    // SiLU on gate, then multiply with up
    510
    0
                    ops::silu(&mut ffn_gate);
    511
    0
                    for i in 0..ffn_gate.len() {
    512
    0
                        ffn_gate[i] *= ffn_up[i];
    513
    0
                    }
    514
    515
    0
                    let mut output = self.fused_matmul(&ffn_gate, &layer.ffn_down_weight)?;
    516
    0
                    if let Some(ref bias) = layer.ffn_down_bias {
    517
    0
                        ops::add_bias(&mut output, bias);
    518
    0
                    }
    519
    0
                    output
    520
                } else {
    521
                    // GELU path (phi-2)
    522
    16
                    let mut ffn_hidden = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?0;
    523
    16
                    if let Some(ref bias0) = layer.ffn_up_bias {
    524
    0
                        ops::add_bias(&mut ffn_hidden, bias);
    525
    16
                    }
    526
    16
                    ops::gelu(&mut ffn_hidden);
    527
    528
    16
                    let mut output = self.fused_matmul(&ffn_hidden, &layer.ffn_down_weight)?0;
    529
    16
                    if let Some(ref bias0) = layer.ffn_down_bias {
    530
    0
                        ops::add_bias(&mut output, bias);
    531
    16
                    }
    532
    16
                    output
    533
                };
    534
    535
                // 2j. Residual connection
    536
    1.02k
                for i in 0..hidden_dim16 {
    537
    1.02k
                    hidden[i] += ffn_output[i];
    538
    1.02k
                }
    539
            }
    540
    541
            // Advance cache position
    542
    16
            cache.advance();
    543
    544
            // 3. Final layer norm
    545
    16
            let normed = if use_rmsnorm {
    546
    0
                ops::rms_norm(&hidden, &self.output_norm_weight, self.config.eps)
    547
            } else {
    548
    16
                ops::layer_norm(
    549
    16
                    &hidden,
    550
    16
                    &self.output_norm_weight,
    551
    16
                    self.output_norm_bias.as_deref(),
    552
    16
                    self.config.eps,
    553
                )
    554
            };
    555
    556
            // PAR-052: Debug final hidden state
    557
    16
            if debug_forward && position == 00 {
    558
    0
                eprintln!(
    559
    0
                    "[PAR-052] Final hidden sum: {:.6}",
    560
    0
                    hidden.iter().sum::<f32>()
    561
    0
                );
    562
    0
                eprintln!(
    563
    0
                    "[PAR-052] Final hidden[0..8]: {:?}",
    564
    0
                    &hidden[..8.min(hidden.len())]
    565
    0
                );
    566
    0
                eprintln!(
    567
    0
                    "[PAR-052] After output_norm sum: {:.6}",
    568
    0
                    normed.iter().sum::<f32>()
    569
    0
                );
    570
    0
                eprintln!(
    571
    0
                    "[PAR-052] output_norm_weight[0..4]: {:?}",
    572
    0
                    &self.output_norm_weight[..4.min(self.output_norm_weight.len())]
    573
    0
                );
    574
    0
                eprintln!(
    575
    0
                    "[PAR-052] LM head weight dims: in={}, out={}",
    576
    0
                    self.lm_head_weight.in_dim, self.lm_head_weight.out_dim
    577
    0
                );
    578
    16
            }
    579
    580
            // 4. LM head projection
    581
    16
            let mut logits = self.fused_matmul(&normed, &self.lm_head_weight)?0;
    582
    16
            if let Some(ref bias0) = self.lm_head_bias {
    583
    0
                ops::add_bias(&mut logits, bias);
    584
    16
            }
    585
    586
            // PAR-052: Debug final logits
    587
    16
            if debug_forward && position == 00 {
    588
                // Find top-5 logits
    589
    0
                let mut indexed: Vec<(usize, f32)> = logits.iter().copied().enumerate().collect();
    590
    0
                indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    591
    0
                eprintln!("[PAR-052] Top-5 logits:");
    592
    0
                for (idx, val) in indexed.iter().take(5) {
    593
    0
                    eprintln!("  Token {}: {:.6}", idx, val);
    594
    0
                }
    595
    0
                eprintln!("[PAR-052] Logits sum: {:.6}", logits.iter().sum::<f32>());
    596
    16
            }
    597
    598
    16
            Ok(logits)
    599
    16
        }
    600
    }