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

    Line
    Count
    Source
    1
    //! Single-token forward pass with KV cache
    2
    //!
    3
    //! Contains forward_single_with_cache and forward_single_with_cache_adaptive.
    4
    //! These are the decode-phase entry points for autoregressive generation.
    5
    6
    use crate::error::Result;
    7
    use crate::gguf::ops;
    8
    use crate::gguf::{
    9
        DispatchMetrics, InferenceScratchBuffer, OwnedQuantizedKVCache, OwnedQuantizedModel,
    10
        GGUF_TYPE_Q4_K,
    11
    };
    12
    13
    impl OwnedQuantizedModel {
    14
        /// Forward pass for a single token using KV cache (IMP-101c)
    15
        ///
    16
        /// This is O(n) per token instead of O(n²) due to KV cache reuse.
    17
        ///
    18
        /// # Arguments
    19
        /// * `token_id` - Single input token ID
    20
        /// * `cache` - Mutable reference to KV cache
    21
        /// * `position` - Position in sequence for RoPE
    22
        ///
    23
        /// # Returns
    24
        /// Logits for next token prediction [vocab_size]
    25
        ///
    26
        /// # Errors
    27
        /// Returns error if tensor operations fail
    28
    249
        pub fn forward_single_with_cache(
    29
    249
            &self,
    30
    249
            token_id: u32,
    31
    249
            cache: &mut OwnedQuantizedKVCache,
    32
    249
            position: usize,
    33
    249
        ) -> Result<Vec<f32>> {
    34
    249
            let hidden_dim = self.config.hidden_dim;
    35
    36
            // 1. Token embedding lookup
    37
    249
            let mut hidden = self.embed(&[token_id]);
    38
    39
            // DEBUG: Print hidden state after embedding
    40
    249
            let debug_forward = std::env::var("REALIZAR_DEBUG_FORWARD").is_ok();
    41
    249
            if debug_forward {
    42
    0
                let hidden_sum: f32 = hidden.iter().sum();
    43
    0
                eprintln!("[DEBUG-FORWARD] Token={}, Position={}", token_id, position);
    44
    0
                eprintln!(
    45
    0
                    "[DEBUG-FORWARD] After embed: sum={:.6}, hidden[0..4]={:?}",
    46
    0
                    hidden_sum,
    47
    0
                    &hidden[..4.min(hidden.len())]
    48
    0
                );
    49
    249
            }
    50
    51
            // Detect if model uses RMSNorm (LLaMA-style) or LayerNorm (phi-2 style)
    52
            // LLaMA models have ffn_gate_weight (SwiGLU) and no bias in norms
    53
    249
            let use_rmsnorm = self
    54
    249
                .layers
    55
    249
                .first()
    56
    249
                .is_some_and(|l| l.ffn_gate_weight.is_some() && l.attn_norm_bias0.is_none0());
    57
    58
            // Pre-allocate attention output buffer - reused across all layers
    59
    249
            let mut attn_out_buffer = vec![0.0f32; hidden_dim];
    60
    61
            // 2. Process through transformer layers
    62
    249
            for (layer_idx, layer) in self.layers.iter().enumerate() {
    63
                // 2a+2b. Fused attention layer norm + QKV projection
    64
                // For RMSNorm models: fuse norm + matmul to eliminate intermediate allocation
    65
                // For LayerNorm models: use separate operations (has bias)
    66
    249
                let mut qkv = if use_rmsnorm {
    67
    0
                    self.fused_rmsnorm_qkv_matmul(
    68
    0
                        &hidden,
    69
    0
                        &layer.attn_norm_weight,
    70
    0
                        self.config.eps,
    71
    0
                        &layer.qkv_weight,
    72
    0
                    )?
    73
                } else {
    74
    249
                    let normed = ops::layer_norm(
    75
    249
                        &hidden,
    76
    249
                        &layer.attn_norm_weight,
    77
    249
                        layer.attn_norm_bias.as_deref(),
    78
    249
                        self.config.eps,
    79
                    );
    80
    249
                    self.qkv_matmul(&normed, &layer.qkv_weight)?0
    81
                };
    82
    249
                if let Some(ref bias0) = layer.qkv_bias {
    83
    0
                    ops::add_bias(&mut qkv, bias);
    84
    249
                }
    85
    86
                // 2c. Extract Q, K, V with GQA-aware sizes and apply RoPE
    87
                // Q: [hidden_dim] = [num_heads * head_dim]
    88
                // K: [kv_dim] = [num_kv_heads * head_dim]
    89
                // V: [kv_dim] = [num_kv_heads * head_dim]
    90
                // Optimization: apply RoPE in-place to avoid Q/K copies
    91
    249
                let num_kv_heads = self.config.num_kv_heads;
    92
    249
                let head_dim = hidden_dim / self.config.num_heads;
    93
    249
                let kv_dim = num_kv_heads * head_dim;
    94
    95
                // Apply RoPE in-place to Q and K within QKV buffer
    96
    249
                self.apply_rope(&mut qkv[0..hidden_dim], position, self.config.num_heads);
    97
    249
                self.apply_rope(
    98
    249
                    &mut qkv[hidden_dim..hidden_dim + kv_dim],
    99
    249
                    position,
    100
    249
                    num_kv_heads,
    101
                );
    102
    103
                // Use slices to avoid copies (only copy K for cache storage)
    104
    249
                let q = &qkv[0..hidden_dim];
    105
    249
                let k = &qkv[hidden_dim..hidden_dim + kv_dim];
    106
    249
                let v = &qkv[hidden_dim + kv_dim..hidden_dim + 2 * kv_dim];
    107
    108
                // 2d. Get cached K/V and compute attention with GQA support
    109
    249
                let k_cache = cache.get_k(layer_idx);
    110
    249
                let v_cache = cache.get_v(layer_idx);
    111
    112
                // Use pre-allocated attention output buffer (reused across layers)
    113
    249
                if k_cache.is_empty() {
    114
                    // First token - no cache yet, output is just weighted V
    115
                    // With single query and single K/V, need to expand V for all Q heads
    116
    26
                    let q_per_kv = self.config.num_heads / num_kv_heads;
    117
    104
                    for q_head in 0..self.config.num_heads26 {
    118
    104
                        let kv_head = q_head / q_per_kv;
    119
    104
                        let v_start = kv_head * head_dim;
    120
    104
                        let out_start = q_head * head_dim;
    121
    104
                        attn_out_buffer[out_start..out_start + head_dim]
    122
    104
                            .copy_from_slice(&v[v_start..v_start + head_dim]);
    123
    104
                    }
    124
                } else {
    125
                    // Use cached K/V for attention with GQA
    126
                    // Uses pre-allocated buffer to avoid 704 Vec allocations per token
    127
    223
                    self.attention_with_cache_gqa_into(q, k_cache, v_cache, k, v, &mut attn_out_buffer);
    128
    129
                    // CORRECTNESS-013: Debug CPU attention output for layer 0 at position 1+
    130
    223
                    if layer_idx == 0 && position >= 1 && std::env::var("CPU_DEBUG").is_ok() {
    131
    0
                        eprintln!(
    132
    0
                            "[CORRECTNESS-013-CPU] Layer 0 attention output at pos={}, first 10: {:?}",
    133
                            position,
    134
    0
                            &attn_out_buffer[..10.min(attn_out_buffer.len())]
    135
                        );
    136
    0
                        for h in 0..3 {
    137
    0
                            let start = h * head_dim;
    138
    0
                            eprintln!(
    139
    0
                                "[CORRECTNESS-013-CPU] Head {} first 5: {:?}",
    140
    0
                                h,
    141
    0
                                &attn_out_buffer[start..start + 5]
    142
    0
                            );
    143
    0
                        }
    144
    223
                    }
    145
                }
    146
    147
                // 2e. Store K and V in cache for future tokens
    148
    249
                cache.append(layer_idx, k, v);
    149
    150
                // 2f. Attention output projection
    151
    249
                let mut attn_output = self.fused_matmul(&attn_out_buffer, &layer.attn_output_weight)?0;
    152
    249
                if let Some(ref bias0) = layer.attn_output_bias {
    153
    0
                    ops::add_bias(&mut attn_output, bias);
    154
    249
                }
    155
    156
                // 2g. Residual connection
    157
    15.8k
                for i in 0..hidden_dim249 {
    158
    15.8k
                    hidden[i] += attn_output[i];
    159
    15.8k
                }
    160
    161
                // 2h+2i. FFN with optional layer norm and SwiGLU/GELU activation
    162
                // For RMSNorm + SwiGLU: fuse norm + up/gate matmuls to eliminate intermediate
    163
    249
                let ffn_activated = match (&layer.ffn_norm_weight, &layer.ffn_gate_weight) {
    164
                    // Fused path: RMSNorm + SwiGLU (LLaMA, TinyLlama, Mistral, etc.)
    165
    0
                    (Some(ref ffn_norm), Some(ref gate_weight)) if use_rmsnorm => {
    166
    0
                        let (mut ffn_up, mut ffn_gate) = self.fused_rmsnorm_ffn_up_gate(
    167
    0
                            &hidden,
    168
    0
                            ffn_norm,
    169
    0
                            self.config.eps,
    170
    0
                            &layer.ffn_up_weight,
    171
    0
                            gate_weight,
    172
    0
                        )?;
    173
    174
    0
                        if let Some(ref bias) = layer.ffn_up_bias {
    175
    0
                            ops::add_bias(&mut ffn_up, bias);
    176
    0
                        }
    177
    0
                        if let Some(ref bias) = layer.ffn_gate_bias {
    178
    0
                            ops::add_bias(&mut ffn_gate, bias);
    179
    0
                        }
    180
    181
                        // SwiGLU: silu(gate) * up
    182
    0
                        ops::silu(&mut ffn_gate);
    183
    0
                        for i in 0..ffn_gate.len() {
    184
    0
                            ffn_gate[i] *= ffn_up[i];
    185
    0
                        }
    186
    0
                        ffn_gate
    187
                    },
    188
    189
                    // Non-fused SwiGLU (LayerNorm models with gate)
    190
    0
                    (ffn_norm_opt, Some(ref gate_weight)) => {
    191
    0
                        let ffn_input = if let Some(ref ffn_norm) = ffn_norm_opt {
    192
    0
                            ops::layer_norm(
    193
    0
                                &hidden,
    194
    0
                                ffn_norm,
    195
    0
                                layer.ffn_norm_bias.as_deref(),
    196
    0
                                self.config.eps,
    197
                            )
    198
                        } else {
    199
    0
                            hidden.clone()
    200
                        };
    201
    202
    0
                        let mut ffn_up = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?;
    203
    0
                        if let Some(ref bias) = layer.ffn_up_bias {
    204
    0
                            ops::add_bias(&mut ffn_up, bias);
    205
    0
                        }
    206
    207
    0
                        let mut ffn_gate = self.fused_matmul(&ffn_input, gate_weight)?;
    208
    0
                        if let Some(ref bias) = layer.ffn_gate_bias {
    209
    0
                            ops::add_bias(&mut ffn_gate, bias);
    210
    0
                        }
    211
    212
                        // SwiGLU: silu(gate) * up
    213
    0
                        ops::silu(&mut ffn_gate);
    214
    0
                        for i in 0..ffn_gate.len() {
    215
    0
                            ffn_gate[i] *= ffn_up[i];
    216
    0
                        }
    217
    0
                        ffn_gate
    218
                    },
    219
    220
                    // GELU path (phi-2, GPT-2, etc.) - no gate weight
    221
    249
                    (ffn_norm_opt, None) => {
    222
    249
                        let ffn_input = if let Some(ref ffn_norm0) = ffn_norm_opt {
    223
    0
                            if use_rmsnorm {
    224
    0
                                ops::rms_norm(&hidden, ffn_norm, self.config.eps)
    225
                            } else {
    226
    0
                                ops::layer_norm(
    227
    0
                                    &hidden,
    228
    0
                                    ffn_norm,
    229
    0
                                    layer.ffn_norm_bias.as_deref(),
    230
    0
                                    self.config.eps,
    231
                                )
    232
                            }
    233
                        } else {
    234
    249
                            hidden.clone()
    235
                        };
    236
    237
    249
                        let mut ffn_hidden = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?0;
    238
    249
                        if let Some(ref bias0) = layer.ffn_up_bias {
    239
    0
                            ops::add_bias(&mut ffn_hidden, bias);
    240
    249
                        }
    241
    249
                        ops::gelu(&mut ffn_hidden);
    242
    249
                        ffn_hidden
    243
                    },
    244
                };
    245
    246
                // 2j. FFN down projection
    247
    249
                let mut ffn_output = self.fused_matmul(&ffn_activated, &layer.ffn_down_weight)?0;
    248
    249
                if let Some(ref bias0) = layer.ffn_down_bias {
    249
    0
                    ops::add_bias(&mut ffn_output, bias);
    250
    249
                }
    251
    252
                // Residual
    253
    15.8k
                for i in 0..hidden_dim249 {
    254
    15.8k
                    hidden[i] += ffn_output[i];
    255
    15.8k
                }
    256
    257
                // DEBUG: Print hidden state after first layer
    258
    249
                if debug_forward && layer_idx == 00 {
    259
    0
                    let hidden_sum: f32 = hidden.iter().sum();
    260
    0
                    eprintln!(
    261
    0
                        "[DEBUG-FORWARD] After layer 0: sum={:.6}, hidden[0..4]={:?}",
    262
    0
                        hidden_sum,
    263
    0
                        &hidden[..4.min(hidden.len())]
    264
    0
                    );
    265
    249
                }
    266
            }
    267
    268
            // Advance cache position after processing all layers
    269
    249
            cache.advance();
    270
    271
            // DEBUG: Print hidden state before LM head
    272
    249
            if debug_forward {
    273
    0
                let hidden_sum: f32 = hidden.iter().sum();
    274
    0
                let hidden_max = hidden.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    275
    0
                let hidden_min = hidden.iter().copied().fold(f32::INFINITY, f32::min);
    276
    0
                eprintln!(
    277
    0
                    "[DEBUG-FORWARD] Hidden after all layers: sum={:.4}, min={:.4}, max={:.4}",
    278
    0
                    hidden_sum, hidden_min, hidden_max
    279
    0
                );
    280
    0
                eprintln!(
    281
    0
                    "[DEBUG-FORWARD] Hidden[0..8]: {:?}",
    282
    0
                    &hidden[..8.min(hidden.len())]
    283
    0
                );
    284
    0
                eprintln!(
    285
    0
                    "[DEBUG-LM-HEAD] lm_head_weight: in_dim={}, out_dim={}, qtype={}, data_len={}",
    286
    0
                    self.lm_head_weight.in_dim,
    287
    0
                    self.lm_head_weight.out_dim,
    288
    0
                    self.lm_head_weight.qtype,
    289
    0
                    self.lm_head_weight.data.len()
    290
    0
                );
    291
    0
                eprintln!(
    292
    0
                    "[DEBUG-LM-HEAD] First 16 bytes of lm_head data: {:02x?}",
    293
    0
                    &self.lm_head_weight.data[..16.min(self.lm_head_weight.data.len())]
    294
    0
                );
    295
    0
                eprintln!(
    296
    0
                    "[DEBUG-LM-HEAD] output_norm_weight[0..4]: {:?}",
    297
    0
                    &self.output_norm_weight[..4.min(self.output_norm_weight.len())]
    298
    0
                );
    299
    249
            }
    300
    301
            // 3+4. Fused final layer norm + LM head projection
    302
            // For RMSNorm models: fuse norm + matmul to eliminate intermediate allocation
    303
    249
            let mut logits = if use_rmsnorm {
    304
    0
                self.fused_rmsnorm_lm_head(&hidden)?
    305
            } else {
    306
    249
                let normed = ops::layer_norm(
    307
    249
                    &hidden,
    308
    249
                    &self.output_norm_weight,
    309
    249
                    self.output_norm_bias.as_deref(),
    310
    249
                    self.config.eps,
    311
                );
    312
    249
                self.fused_matmul(&normed, &self.lm_head_weight)?0
    313
            };
    314
    315
            // DEBUG: Verify Q8_0 matmul by manual computation
    316
    249
            if debug_forward {
    317
                // Get the normalized hidden state
    318
    0
                let normed = ops::rms_norm(&hidden, &self.output_norm_weight, self.config.eps);
    319
    0
                eprintln!(
    320
    0
                    "[DEBUG-VERIFY] Normed hidden[0..8]: {:?}",
    321
    0
                    &normed[..8.min(normed.len())]
    322
                );
    323
    324
                // Manual dequantize row 0 of LM head weight
    325
                const Q8_0_BLOCK_BYTES: usize = 34;
    326
                const Q8_0_BLOCK_SIZE: usize = 32;
    327
    0
                let blocks_per_row = self.lm_head_weight.in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    328
    0
                let bytes_per_row = blocks_per_row * Q8_0_BLOCK_BYTES;
    329
    330
                // Dequantize row 0 (token 0's projection weights)
    331
    0
                let row0_data = &self.lm_head_weight.data[0..bytes_per_row];
    332
    0
                let mut row0_f32 = vec![0.0f32; self.lm_head_weight.in_dim];
    333
    0
                for block_idx in 0..blocks_per_row {
    334
    0
                    let block_start = block_idx * Q8_0_BLOCK_BYTES;
    335
    0
                    let block = &row0_data[block_start..block_start + Q8_0_BLOCK_BYTES];
    336
    0
                    let scale = half::f16::from_le_bytes([block[0], block[1]]).to_f32();
    337
    0
                    for j in 0..32 {
    338
    0
                        let idx = block_idx * 32 + j;
    339
    0
                        if idx >= self.lm_head_weight.in_dim {
    340
    0
                            break;
    341
    0
                        }
    342
    0
                        row0_f32[idx] = (block[2 + j] as i8 as f32) * scale;
    343
                    }
    344
                }
    345
    0
                eprintln!(
    346
    0
                    "[DEBUG-VERIFY] LM head row 0 (dequantized) first 8: {:?}",
    347
    0
                    &row0_f32[..8.min(row0_f32.len())]
    348
                );
    349
    350
                // Compute dot product manually
    351
    0
                let manual_logit0: f32 = normed.iter().zip(row0_f32.iter()).map(|(a, b)| a * b).sum();
    352
    0
                eprintln!("[DEBUG-VERIFY] Manual logits[0] = {:.6}", manual_logit0);
    353
    0
                eprintln!("[DEBUG-VERIFY] Computed logits[0] = {:.6}", logits[0]);
    354
    0
                eprintln!(
    355
    0
                    "[DEBUG-VERIFY] Difference = {:.6}",
    356
    0
                    (manual_logit0 - logits[0]).abs()
    357
                );
    358
    359
                // Check top tokens
    360
    0
                let mut indexed: Vec<(usize, f32)> =
    361
    0
                    logits.iter().enumerate().map(|(i, &v)| (i, v)).collect();
    362
    0
                indexed.sort_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
    363
    0
                eprintln!(
    364
    0
                    "[DEBUG-VERIFY] Top 5 tokens: {:?}",
    365
    0
                    &indexed[..5.min(indexed.len())]
    366
                );
    367
    249
            }
    368
    369
    249
            if let Some(ref bias0) = self.lm_head_bias {
    370
    0
                ops::add_bias(&mut logits, bias);
    371
    249
            }
    372
    373
    249
            Ok(logits)
    374
    249
        }
    375
    376
        /// Single-token forward pass with pre-allocated scratch buffers
    377
        ///
    378
        /// Uses OwnedInferenceScratchBuffer to eliminate per-token allocations.
    379
        /// For Qwen2.5-0.5B, this saves ~40KB of allocations per token.
    380
        ///
    381
        /// # Arguments
    382
        /// * `token_id` - Token to process
    383
        /// * `cache` - KV cache for incremental decoding
    384
        ///
    385
        /// Forward pass with adaptive CPU/GPU attention selection (IMP-124)
    386
        ///
    387
        /// This variant of `forward_single_with_cache` uses `adaptive_attention_with_cache`
    388
        /// to automatically select between CPU and GPU backends based on cache length.
    389
        /// It also records dispatch decisions to the provided metrics tracker.
    390
        ///
    391
        /// # Arguments
    392
        /// * `token_id` - Token to process
    393
        /// * `cache` - KV cache for incremental decoding
    394
        /// * `position` - Position in sequence
    395
        /// * `metrics` - Dispatch metrics tracker for CPU/GPU decision recording
    396
        ///
    397
        /// # Returns
    398
        /// Logits for next token prediction [vocab_size]
    399
        ///
    400
        /// # Errors
    401
        /// Returns error if tensor operations fail
    402
        #[cfg(feature = "gpu")]
    403
    345
        pub fn forward_single_with_cache_adaptive(
    404
    345
            &self,
    405
    345
            token_id: u32,
    406
    345
            cache: &mut OwnedQuantizedKVCache,
    407
    345
            position: usize,
    408
    345
            metrics: &std::sync::Arc<DispatchMetrics>,
    409
    345
        ) -> Result<Vec<f32>> {
    410
    345
            let hidden_dim = self.config.hidden_dim;
    411
    412
            // 1. Token embedding lookup
    413
    345
            let mut hidden = self.embed(&[token_id]);
    414
    415
            // Detect if model uses RMSNorm (LLaMA-style) or LayerNorm (phi-2 style)
    416
            // LLaMA models have ffn_gate_weight (SwiGLU) and no bias in norms
    417
    345
            let use_rmsnorm = self
    418
    345
                .layers
    419
    345
                .first()
    420
    345
                .is_some_and(|l| l.ffn_gate_weight.is_some() && l.attn_norm_bias0.is_none0());
    421
    422
            // GQA dimensions
    423
    345
            let num_kv_heads = self.config.num_kv_heads;
    424
    345
            let head_dim = hidden_dim / self.config.num_heads;
    425
    345
            let kv_dim = num_kv_heads * head_dim;
    426
    427
            // PARITY-113: Track CUDA kernel count for GPU dispatch metrics
    428
            #[cfg(feature = "cuda")]
    429
            let cuda_enabled = self.cuda_enabled();
    430
    431
            // 2. Process through transformer layers
    432
    359
            for (layer_idx, layer) in self.layers.iter()345.enumerate345() {
    433
                // 2a. Attention layer norm (RMSNorm for LLaMA, LayerNorm for others)
    434
    359
                let normed = if use_rmsnorm {
    435
    0
                    ops::rms_norm(&hidden, &layer.attn_norm_weight, self.config.eps)
    436
                } else {
    437
    359
                    ops::layer_norm(
    438
    359
                        &hidden,
    439
    359
                        &layer.attn_norm_weight,
    440
    359
                        layer.attn_norm_bias.as_deref(),
    441
    359
                        self.config.eps,
    442
                    )
    443
                };
    444
    445
                // 2b. QKV projection
    446
                // PARITY-113: Record GPU dispatch when CUDA path is used for matmul
    447
                #[cfg(feature = "cuda")]
    448
                if cuda_enabled {
    449
                    let start = std::time::Instant::now();
    450
                    let qkv_result = self.qkv_matmul(&normed, &layer.qkv_weight)?;
    451
                    metrics.record_gpu_dispatch();
    452
                    metrics.record_gpu_latency(start.elapsed());
    453
                    let mut qkv = qkv_result;
    454
                    if let Some(ref bias) = layer.qkv_bias {
    455
                        ops::add_bias(&mut qkv, bias);
    456
                    }
    457
    458
                    // 2c. Extract Q, K, V with GQA-aware sizes and apply RoPE
    459
                    let mut q = qkv[0..hidden_dim].to_vec();
    460
                    let mut k = qkv[hidden_dim..hidden_dim + kv_dim].to_vec();
    461
                    let v = qkv[hidden_dim + kv_dim..hidden_dim + 2 * kv_dim].to_vec();
    462
    463
                    self.apply_rope(&mut q, position, self.config.num_heads);
    464
                    self.apply_rope(&mut k, position, num_kv_heads);
    465
    466
                    // 2d. Get cached K/V and compute attention with GQA support
    467
                    let k_cache = cache.get_k(layer_idx);
    468
                    let v_cache = cache.get_v(layer_idx);
    469
    470
                    let attn_out = if k_cache.is_empty() {
    471
                        // First token - expand V for all Q heads (GQA)
    472
                        let mut expanded_v = vec![0.0f32; hidden_dim];
    473
                        let q_per_kv = self.config.num_heads / num_kv_heads;
    474
                        for q_head in 0..self.config.num_heads {
    475
                            let kv_head = q_head / q_per_kv;
    476
                            let v_start = kv_head * head_dim;
    477
                            let out_start = q_head * head_dim;
    478
                            expanded_v[out_start..out_start + head_dim]
    479
                                .copy_from_slice(&v[v_start..v_start + head_dim]);
    480
                        }
    481
                        expanded_v
    482
                    } else {
    483
                        let start = std::time::Instant::now();
    484
                        let result =
    485
                            self.adaptive_attention_with_cache(&q, k_cache, v_cache, &k, &v)?;
    486
                        metrics.record_gpu_dispatch();
    487
                        metrics.record_gpu_latency(start.elapsed());
    488
                        result
    489
                    };
    490
    491
                    // 2e. Store K and V in cache
    492
                    cache.append(layer_idx, &k, &v);
    493
    494
                    // 2f. Attention output projection
    495
                    let start = std::time::Instant::now();
    496
                    let mut attn_output = self.fused_matmul(&attn_out, &layer.attn_output_weight)?;
    497
                    metrics.record_gpu_dispatch();
    498
                    metrics.record_gpu_latency(start.elapsed());
    499
                    if let Some(ref bias) = layer.attn_output_bias {
    500
                        ops::add_bias(&mut attn_output, bias);
    501
                    }
    502
    503
                    // 2g. Residual connection
    504
                    for i in 0..hidden_dim {
    505
                        hidden[i] += attn_output[i];
    506
                    }
    507
    508
                    // 2h. Pre-FFN layer norm (LLaMA uses separate ffn_norm with RMSNorm)
    509
                    let ffn_input = if let Some(ref ffn_norm) = layer.ffn_norm_weight {
    510
                        if use_rmsnorm {
    511
                            ops::rms_norm(&hidden, ffn_norm, self.config.eps)
    512
                        } else {
    513
                            ops::layer_norm(
    514
                                &hidden,
    515
                                ffn_norm,
    516
                                layer.ffn_norm_bias.as_deref(),
    517
                                self.config.eps,
    518
                            )
    519
                        }
    520
                    } else {
    521
                        hidden.clone()
    522
                    };
    523
    524
                    // 2i. FFN with SwiGLU or GELU activation
    525
                    let ffn_activated = if let Some(ref gate_weight) = layer.ffn_gate_weight {
    526
                        // SwiGLU path
    527
                        let start = std::time::Instant::now();
    528
                        let mut ffn_up = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?;
    529
                        metrics.record_gpu_dispatch();
    530
                        metrics.record_gpu_latency(start.elapsed());
    531
                        if let Some(ref bias) = layer.ffn_up_bias {
    532
                            ops::add_bias(&mut ffn_up, bias);
    533
                        }
    534
    535
                        let start = std::time::Instant::now();
    536
                        let mut ffn_gate = self.fused_matmul(&ffn_input, gate_weight)?;
    537
                        metrics.record_gpu_dispatch();
    538
                        metrics.record_gpu_latency(start.elapsed());
    539
                        if let Some(ref bias) = layer.ffn_gate_bias {
    540
                            ops::add_bias(&mut ffn_gate, bias);
    541
                        }
    542
    543
                        ops::silu(&mut ffn_gate);
    544
                        for i in 0..ffn_gate.len() {
    545
                            ffn_gate[i] *= ffn_up[i];
    546
                        }
    547
                        ffn_gate
    548
                    } else {
    549
                        // GELU path
    550
                        let start = std::time::Instant::now();
    551
                        let mut ffn_hidden = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?;
    552
                        metrics.record_gpu_dispatch();
    553
                        metrics.record_gpu_latency(start.elapsed());
    554
                        if let Some(ref bias) = layer.ffn_up_bias {
    555
                            ops::add_bias(&mut ffn_hidden, bias);
    556
                        }
    557
                        ops::gelu(&mut ffn_hidden);
    558
                        ffn_hidden
    559
                    };
    560
    561
                    // 2j. FFN down projection
    562
                    let start = std::time::Instant::now();
    563
                    let mut ffn_output = self.fused_matmul(&ffn_activated, &layer.ffn_down_weight)?;
    564
                    metrics.record_gpu_dispatch();
    565
                    metrics.record_gpu_latency(start.elapsed());
    566
                    if let Some(ref bias) = layer.ffn_down_bias {
    567
                        ops::add_bias(&mut ffn_output, bias);
    568
                    }
    569
    570
                    // Residual
    571
                    for i in 0..hidden_dim {
    572
                        hidden[i] += ffn_output[i];
    573
                    }
    574
    575
                    continue;
    576
                }
    577
    578
                // CPU path (non-CUDA)
    579
    359
                let mut qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0;
    580
    359
                if let Some(ref bias0) = layer.qkv_bias {
    581
    0
                    ops::add_bias(&mut qkv, bias);
    582
    359
                }
    583
    584
                // 2c. Extract Q, K, V with GQA-aware sizes and apply RoPE
    585
    359
                let mut q = qkv[0..hidden_dim].to_vec();
    586
    359
                let mut k = qkv[hidden_dim..hidden_dim + kv_dim].to_vec();
    587
    359
                let v = qkv[hidden_dim + kv_dim..hidden_dim + 2 * kv_dim].to_vec();
    588
    589
    359
                self.apply_rope(&mut q, position, self.config.num_heads);
    590
    359
                self.apply_rope(&mut k, position, num_kv_heads);
    591
    592
                // 2d. Get cached K/V and compute attention with adaptive dispatch
    593
    359
                let k_cache = cache.get_k(layer_idx);
    594
    359
                let v_cache = cache.get_v(layer_idx);
    595
    596
    359
                let attn_out = if k_cache.is_empty() {
    597
                    // First token - expand V for all Q heads (GQA)
    598
    15
                    let mut expanded_v = vec![0.0f32; hidden_dim];
    599
    15
                    let q_per_kv = self.config.num_heads / num_kv_heads;
    600
    52
                    for q_head in 0..self.config.num_heads15 {
    601
    52
                        let kv_head = q_head / q_per_kv;
    602
    52
                        let v_start = kv_head * head_dim;
    603
    52
                        let out_start = q_head * head_dim;
    604
    52
                        expanded_v[out_start..out_start + head_dim]
    605
    52
                            .copy_from_slice(&v[v_start..v_start + head_dim]);
    606
    52
                    }
    607
    15
                    expanded_v
    608
                } else {
    609
                    // Use adaptive attention with metrics tracking
    610
    344
                    let cache_len = k_cache.len() / kv_dim;
    611
                    const GPU_CACHE_LEN_THRESHOLD: usize = 64;
    612
    613
    344
                    if cache_len >= GPU_CACHE_LEN_THRESHOLD {
    614
    62
                        let start = std::time::Instant::now();
    615
    62
                        let result =
    616
    62
                            self.adaptive_attention_with_cache(&q, k_cache, v_cache, &k, &v)?0;
    617
    62
                        metrics.record_gpu_dispatch();
    618
    62
                        metrics.record_gpu_latency(start.elapsed());
    619
    62
                        result
    620
                    } else {
    621
    282
                        let start = std::time::Instant::now();
    622
    282
                        let result = self.attention_with_cache_gqa(&q, k_cache, v_cache, &k, &v);
    623
    282
                        metrics.record_cpu_dispatch();
    624
    282
                        metrics.record_cpu_latency(start.elapsed());
    625
    282
                        result
    626
                    }
    627
                };
    628
    629
                // 2e. Store K and V in cache for future tokens
    630
    359
                cache.append(layer_idx, &k, &v);
    631
    632
                // 2f. Attention output projection
    633
    359
                let mut attn_output = self.fused_matmul(&attn_out, &layer.attn_output_weight)?0;
    634
    359
                if let Some(ref bias0) = layer.attn_output_bias {
    635
    0
                    ops::add_bias(&mut attn_output, bias);
    636
    359
                }
    637
    638
                // 2g. Residual connection
    639
    17.4k
                for i in 0..hidden_dim359 {
    640
    17.4k
                    hidden[i] += attn_output[i];
    641
    17.4k
                }
    642
    643
                // 2h. Pre-FFN layer norm (LLaMA uses separate ffn_norm with RMSNorm)
    644
    359
                let ffn_input = if let Some(ref ffn_norm0) = layer.ffn_norm_weight {
    645
    0
                    if use_rmsnorm {
    646
    0
                        ops::rms_norm(&hidden, ffn_norm, self.config.eps)
    647
                    } else {
    648
    0
                        ops::layer_norm(
    649
    0
                            &hidden,
    650
    0
                            ffn_norm,
    651
    0
                            layer.ffn_norm_bias.as_deref(),
    652
    0
                            self.config.eps,
    653
                        )
    654
                    }
    655
                } else {
    656
    359
                    hidden.clone()
    657
                };
    658
    659
                // 2i. FFN with SwiGLU or GELU activation
    660
    359
                let ffn_activated = if let Some(ref gate_weight0) = layer.ffn_gate_weight {
    661
                    // SwiGLU path
    662
    0
                    let mut ffn_up = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?;
    663
    0
                    if let Some(ref bias) = layer.ffn_up_bias {
    664
    0
                        ops::add_bias(&mut ffn_up, bias);
    665
    0
                    }
    666
    667
    0
                    let mut ffn_gate = self.fused_matmul(&ffn_input, gate_weight)?;
    668
    0
                    if let Some(ref bias) = layer.ffn_gate_bias {
    669
    0
                        ops::add_bias(&mut ffn_gate, bias);
    670
    0
                    }
    671
    672
    0
                    ops::silu(&mut ffn_gate);
    673
    0
                    for i in 0..ffn_gate.len() {
    674
    0
                        ffn_gate[i] *= ffn_up[i];
    675
    0
                    }
    676
    0
                    ffn_gate
    677
                } else {
    678
                    // GELU path
    679
    359
                    let mut ffn_hidden = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?0;
    680
    359
                    if let Some(ref bias0) = layer.ffn_up_bias {
    681
    0
                        ops::add_bias(&mut ffn_hidden, bias);
    682
    359
                    }
    683
    359
                    ops::gelu(&mut ffn_hidden);
    684
    359
                    ffn_hidden
    685
                };
    686
    687
                // 2j. FFN down projection
    688
    359
                let mut ffn_output = self.fused_matmul(&ffn_activated, &layer.ffn_down_weight)?0;
    689
    359
                if let Some(ref bias0) = layer.ffn_down_bias {
    690
    0
                    ops::add_bias(&mut ffn_output, bias);
    691
    359
                }
    692
    693
                // Residual
    694
    17.4k
                for i in 0..hidden_dim359 {
    695
    17.4k
                    hidden[i] += ffn_output[i];
    696
    17.4k
                }
    697
            }
    698
    699
            // Advance cache position after processing all layers
    700
    345
            cache.advance();
    701
    702
            // 3. Final layer norm (RMSNorm for LLaMA, LayerNorm for others)
    703
    345
            let normed = if use_rmsnorm {
    704
    0
                ops::rms_norm(&hidden, &self.output_norm_weight, self.config.eps)
    705
            } else {
    706
    345
                ops::layer_norm(
    707
    345
                    &hidden,
    708
    345
                    &self.output_norm_weight,
    709
    345
                    self.output_norm_bias.as_deref(),
    710
    345
                    self.config.eps,
    711
                )
    712
            };
    713
    714
            // 4. LM head projection
    715
            // PARITY-113: Record GPU dispatch for LM head when CUDA is enabled
    716
            #[cfg(feature = "cuda")]
    717
            if cuda_enabled {
    718
                let start = std::time::Instant::now();
    719
                let mut logits = self.fused_matmul(&normed, &self.lm_head_weight)?;
    720
                metrics.record_gpu_dispatch();
    721
                metrics.record_gpu_latency(start.elapsed());
    722
                if let Some(ref bias) = self.lm_head_bias {
    723
                    ops::add_bias(&mut logits, bias);
    724
                }
    725
                return Ok(logits);
    726
            }
    727
    728
    345
            let mut logits = self.fused_matmul(&normed, &self.lm_head_weight)?0;
    729
    345
            if let Some(ref bias0) = self.lm_head_bias {
    730
    0
                ops::add_bias(&mut logits, bias);
    731
    345
            }
    732
    733
    345
            Ok(logits)
    734
    345
        }
    735
    736
        /// Zero-allocation forward pass using scratch buffers (IMP-131)
    737
        ///
    738
        /// All intermediate results are written to pre-allocated scratch buffers.
    739
        /// Output logits are stored in `scratch.logits`.
    740
    0
        pub(crate) fn forward_single_with_scratch(
    741
    0
            &self,
    742
    0
            token_id: u32,
    743
    0
            cache: &mut OwnedQuantizedKVCache,
    744
    0
            position: usize,
    745
    0
            scratch: &mut InferenceScratchBuffer,
    746
    0
        ) -> Result<()> {
    747
    0
            let hidden_dim = self.config.hidden_dim;
    748
    0
            let intermediate_dim = self.config.intermediate_dim;
    749
    750
            // Detect architecture
    751
    0
            let use_rmsnorm = self
    752
    0
                .layers
    753
    0
                .first()
    754
    0
                .is_some_and(|l| l.ffn_gate_weight.is_some() && l.attn_norm_bias.is_none());
    755
    756
            // 1. Token embedding lookup → scratch.hidden
    757
    0
            self.embed_into(token_id, &mut scratch.hidden);
    758
    759
            // 2. Process through transformer layers
    760
    0
            for (layer_idx, layer) in self.layers.iter().enumerate() {
    761
                // 2a. Attention layer norm → scratch.normed
    762
    0
                if use_rmsnorm {
    763
    0
                    ops::rms_norm_into(
    764
    0
                        &scratch.hidden,
    765
    0
                        &layer.attn_norm_weight,
    766
    0
                        self.config.eps,
    767
    0
                        &mut scratch.normed,
    768
    0
                    );
    769
    0
                } else {
    770
    0
                    ops::layer_norm_into(
    771
    0
                        &scratch.hidden,
    772
    0
                        &layer.attn_norm_weight,
    773
    0
                        layer.attn_norm_bias.as_deref(),
    774
    0
                        self.config.eps,
    775
    0
                        &mut scratch.normed,
    776
    0
                    );
    777
    0
                }
    778
    779
                // 2b. QKV projection → scratch.qkv (zero-allocation via P1-REV)
    780
                // PAR-126: Fix GQA dimension issue - use config instead of q_dim() which
    781
                // incorrectly assumes Q=K=V for fused weights
    782
    0
                let num_kv_heads = self.config.num_kv_heads;
    783
    0
                let head_dim = hidden_dim / self.config.num_heads;
    784
    0
                let kv_dim = num_kv_heads * head_dim;
    785
                // Q uses all heads, K/V use only kv_heads (GQA)
    786
    0
                let q_dim = hidden_dim;
    787
    0
                let k_dim = kv_dim;
    788
    0
                let v_dim = kv_dim;
    789
    0
                let qkv_dim = q_dim + k_dim + v_dim;
    790
    791
                // PAR-126: Pre-quantize normalized hidden to Q8K for VNNI-accelerated matmul
    792
                // This allows reusing quantized activations for QKV projection
    793
                // NOTE: Q8K requires hidden_dim to be multiple of 256. For smaller models
    794
                // like 0.5B (hidden=896), fall back to f32 path.
    795
    0
                let use_q8k_path = hidden_dim.is_multiple_of(256);
    796
    797
    0
                if use_q8k_path {
    798
                    use crate::quantize::quantize_activations_q8k_into;
    799
    0
                    let hidden_sb = hidden_dim / 256;
    800
    0
                    quantize_activations_q8k_into(
    801
    0
                        &scratch.normed[..hidden_dim],
    802
    0
                        &mut scratch.q8k_hidden_scales[..hidden_sb],
    803
    0
                        &mut scratch.q8k_hidden_quants[..hidden_dim],
    804
    0
                    )?;
    805
    806
                    // Write directly to scratch.qkv, using Q8K-accelerated path
    807
    0
                    self.qkv_matmul_q8k_into(
    808
    0
                        &scratch.normed,
    809
    0
                        &layer.qkv_weight,
    810
    0
                        &mut scratch.qkv[..qkv_dim],
    811
    0
                        &scratch.q8k_hidden_scales[..hidden_sb],
    812
    0
                        &scratch.q8k_hidden_quants[..hidden_dim],
    813
    0
                    )?;
    814
                } else {
    815
                    // Fall back to f32 path for non-256-aligned hidden dims
    816
    0
                    self.qkv_matmul_into(
    817
    0
                        &scratch.normed,
    818
    0
                        &layer.qkv_weight,
    819
    0
                        &mut scratch.qkv[..qkv_dim],
    820
    0
                    )?;
    821
                }
    822
    823
                // Copy from scratch.qkv to individual Q, K, V buffers
    824
    0
                scratch.q[..q_dim].copy_from_slice(&scratch.qkv[..q_dim]);
    825
    0
                scratch.k[..k_dim].copy_from_slice(&scratch.qkv[q_dim..q_dim + k_dim]);
    826
    0
                scratch.v[..v_dim].copy_from_slice(&scratch.qkv[q_dim + k_dim..qkv_dim]);
    827
    828
                // Add bias if present
    829
    0
                if let Some(ref bias) = layer.qkv_bias {
    830
    0
                    for i in 0..q_dim {
    831
    0
                        scratch.q[i] += bias[i];
    832
    0
                    }
    833
    0
                    for i in 0..k_dim {
    834
    0
                        scratch.k[i] += bias[q_dim + i];
    835
    0
                    }
    836
    0
                    for i in 0..v_dim {
    837
    0
                        scratch.v[i] += bias[q_dim + k_dim + i];
    838
    0
                    }
    839
    0
                }
    840
    841
                // Apply RoPE
    842
    0
                self.apply_rope(&mut scratch.q[..q_dim], position, self.config.num_heads);
    843
    0
                self.apply_rope(&mut scratch.k[..k_dim], position, self.config.num_kv_heads);
    844
    845
                // 2c. Compute attention
    846
    0
                let k_cache = cache.get_k(layer_idx);
    847
    0
                let v_cache = cache.get_v(layer_idx);
    848
    849
    0
                if k_cache.is_empty() {
    850
                    // First token - expand V if GQA
    851
    0
                    if self.config.num_kv_heads < self.config.num_heads {
    852
    0
                        let head_dim = hidden_dim / self.config.num_heads;
    853
    0
                        let group_size = self.config.num_heads / self.config.num_kv_heads;
    854
    0
                        for h in 0..self.config.num_heads {
    855
    0
                            let kv_head = h / group_size;
    856
    0
                            let src_start = kv_head * head_dim;
    857
    0
                            let dst_start = h * head_dim;
    858
    0
                            scratch.attn_out[dst_start..dst_start + head_dim]
    859
    0
                                .copy_from_slice(&scratch.v[src_start..src_start + head_dim]);
    860
    0
                        }
    861
    0
                    } else {
    862
    0
                        scratch.attn_out[..hidden_dim].copy_from_slice(&scratch.v[..hidden_dim]);
    863
    0
                    }
    864
    0
                } else {
    865
    0
                    self.attention_with_cache_gqa_into(
    866
    0
                        &scratch.q[..q_dim],
    867
    0
                        k_cache,
    868
    0
                        v_cache,
    869
    0
                        &scratch.k[..k_dim],
    870
    0
                        &scratch.v[..v_dim],
    871
    0
                        &mut scratch.attn_out,
    872
    0
                    );
    873
    0
                }
    874
    875
                // Store K, V in cache
    876
    0
                cache.append(layer_idx, &scratch.k[..k_dim], &scratch.v[..v_dim]);
    877
    878
                // 2d. Attention output projection → scratch.attn_proj
    879
                // PAR-128: Use Q8K-accelerated path for attention output projection
    880
                // attn_out is hidden_dim sized, reuse hidden Q8K buffers
    881
    0
                let use_q8k_attn_out = use_q8k_path && layer.attn_output_weight.qtype == GGUF_TYPE_Q4_K;
    882
    883
    0
                if use_q8k_attn_out {
    884
                    use crate::quantize::{
    885
                        fused_q4k_q8k_parallel_matvec_into, quantize_activations_q8k_into,
    886
                    };
    887
    0
                    let hidden_sb = hidden_dim / 256;
    888
                    // Quantize attention output to Q8K (reuse hidden Q8K buffers)
    889
    0
                    quantize_activations_q8k_into(
    890
    0
                        &scratch.attn_out[..hidden_dim],
    891
    0
                        &mut scratch.q8k_hidden_scales[..hidden_sb],
    892
    0
                        &mut scratch.q8k_hidden_quants[..hidden_dim],
    893
    0
                    )?;
    894
    0
                    fused_q4k_q8k_parallel_matvec_into(
    895
    0
                        &layer.attn_output_weight.data,
    896
    0
                        &scratch.q8k_hidden_scales[..hidden_sb],
    897
    0
                        &scratch.q8k_hidden_quants[..hidden_dim],
    898
    0
                        layer.attn_output_weight.in_dim,
    899
    0
                        layer.attn_output_weight.out_dim,
    900
    0
                        &mut scratch.attn_proj,
    901
    0
                    )?;
    902
                } else {
    903
    0
                    self.fused_matmul_into(
    904
    0
                        &scratch.attn_out[..hidden_dim],
    905
    0
                        &layer.attn_output_weight,
    906
    0
                        &mut scratch.attn_proj,
    907
    0
                    )?;
    908
                }
    909
    0
                if let Some(ref bias) = layer.attn_output_bias {
    910
    0
                    for i in 0..hidden_dim {
    911
    0
                        scratch.attn_proj[i] += bias[i];
    912
    0
                    }
    913
    0
                }
    914
    915
                // 2e. Residual connection
    916
    0
                for i in 0..hidden_dim {
    917
    0
                    scratch.hidden[i] += scratch.attn_proj[i];
    918
    0
                }
    919
    920
                // 2f. Pre-FFN layer norm → scratch.normed
    921
    0
                if let Some(ref ffn_norm) = layer.ffn_norm_weight {
    922
    0
                    if use_rmsnorm {
    923
    0
                        ops::rms_norm_into(
    924
    0
                            &scratch.hidden,
    925
    0
                            ffn_norm,
    926
    0
                            self.config.eps,
    927
    0
                            &mut scratch.normed,
    928
    0
                        );
    929
    0
                    } else {
    930
    0
                        ops::layer_norm_into(
    931
    0
                            &scratch.hidden,
    932
    0
                            ffn_norm,
    933
    0
                            layer.ffn_norm_bias.as_deref(),
    934
    0
                            self.config.eps,
    935
    0
                            &mut scratch.normed,
    936
    0
                        );
    937
    0
                    }
    938
    0
                } else {
    939
    0
                    scratch.normed[..hidden_dim].copy_from_slice(&scratch.hidden[..hidden_dim]);
    940
    0
                }
    941
    942
                // 2g. FFN
    943
    0
                if let Some(ref gate_weight) = layer.ffn_gate_weight {
    944
                    // SwiGLU path (LLaMA)
    945
                    // PAR-126: Use Q8K-accelerated path only if hidden_dim is 256-aligned
    946
    0
                    if use_q8k_path {
    947
                        // Pre-quantize normed hidden to Q8K for VNNI-accelerated FFN matmul
    948
                        // Quantize once, reuse for both up and gate matmuls
    949
                        use crate::quantize::quantize_activations_q8k_into;
    950
    0
                        let hidden_sb = hidden_dim / 256;
    951
    0
                        quantize_activations_q8k_into(
    952
    0
                            &scratch.normed[..hidden_dim],
    953
    0
                            &mut scratch.q8k_hidden_scales[..hidden_sb],
    954
    0
                            &mut scratch.q8k_hidden_quants[..hidden_dim],
    955
    0
                        )?;
    956
    957
                        // Use fused FFN up+gate kernel to eliminate rayon::join overhead
    958
                        // This reduces parallel region spawns from 2 to 1 per layer
    959
    0
                        let up_weight = &layer.ffn_up_weight;
    960
    0
                        let q8k_scales = &scratch.q8k_hidden_scales[..hidden_sb];
    961
    0
                        let q8k_quants = &scratch.q8k_hidden_quants[..hidden_dim];
    962
    963
                        // Check if both weights are Q4K for fused path
    964
    0
                        if up_weight.qtype == GGUF_TYPE_Q4_K && gate_weight.qtype == GGUF_TYPE_Q4_K {
    965
                            use crate::quantize::fused_q4k_q8k_ffn_up_gate_into;
    966
    0
                            fused_q4k_q8k_ffn_up_gate_into(
    967
    0
                                &up_weight.data,
    968
    0
                                &gate_weight.data,
    969
    0
                                q8k_scales,
    970
    0
                                q8k_quants,
    971
    0
                                up_weight.in_dim,
    972
    0
                                up_weight.out_dim,
    973
    0
                                &mut scratch.ffn_up,
    974
    0
                                &mut scratch.ffn_gate,
    975
    0
                            )?;
    976
                        } else {
    977
                            // Fallback to separate matmuls if not both Q4K
    978
                            use crate::quantize::fused_q4k_q8k_parallel_matvec_into;
    979
    0
                            let (up_result, gate_result) = rayon::join(
    980
    0
                                || {
    981
    0
                                    if up_weight.qtype == GGUF_TYPE_Q4_K {
    982
    0
                                        fused_q4k_q8k_parallel_matvec_into(
    983
    0
                                            &up_weight.data,
    984
    0
                                            q8k_scales,
    985
    0
                                            q8k_quants,
    986
    0
                                            up_weight.in_dim,
    987
    0
                                            up_weight.out_dim,
    988
    0
                                            &mut scratch.ffn_up,
    989
                                        )
    990
                                    } else {
    991
    0
                                        self.fused_matmul_into(
    992
    0
                                            &scratch.normed[..hidden_dim],
    993
    0
                                            up_weight,
    994
    0
                                            &mut scratch.ffn_up,
    995
                                        )
    996
                                    }
    997
    0
                                },
    998
    0
                                || {
    999
    0
                                    if gate_weight.qtype == GGUF_TYPE_Q4_K {
    1000
    0
                                        fused_q4k_q8k_parallel_matvec_into(
    1001
    0
                                            &gate_weight.data,
    1002
    0
                                            q8k_scales,
    1003
    0
                                            q8k_quants,
    1004
    0
                                            gate_weight.in_dim,
    1005
    0
                                            gate_weight.out_dim,
    1006
    0
                                            &mut scratch.ffn_gate,
    1007
                                        )
    1008
                                    } else {
    1009
    0
                                        self.fused_matmul_into(
    1010
    0
                                            &scratch.normed[..hidden_dim],
    1011
    0
                                            gate_weight,
    1012
    0
                                            &mut scratch.ffn_gate,
    1013
                                        )
    1014
                                    }
    1015
    0
                                },
    1016
                            );
    1017
    0
                            up_result?;
    1018
    0
                            gate_result?;
    1019
                        }
    1020
                    } else {
    1021
                        // Fall back to f32 path for non-256-aligned hidden dims
    1022
    0
                        let up_weight = &layer.ffn_up_weight;
    1023
    0
                        let (up_result, gate_result) = rayon::join(
    1024
    0
                            || {
    1025
    0
                                self.fused_matmul_into(
    1026
    0
                                    &scratch.normed[..hidden_dim],
    1027
    0
                                    up_weight,
    1028
    0
                                    &mut scratch.ffn_up,
    1029
                                )
    1030
    0
                            },
    1031
    0
                            || {
    1032
    0
                                self.fused_matmul_into(
    1033
    0
                                    &scratch.normed[..hidden_dim],
    1034
    0
                                    gate_weight,
    1035
    0
                                    &mut scratch.ffn_gate,
    1036
                                )
    1037
    0
                            },
    1038
                        );
    1039
    0
                        up_result?;
    1040
    0
                        gate_result?;
    1041
                    }
    1042
    1043
    0
                    if let Some(ref bias) = layer.ffn_up_bias {
    1044
    0
                        for i in 0..intermediate_dim {
    1045
    0
                            scratch.ffn_up[i] += bias[i];
    1046
    0
                        }
    1047
    0
                    }
    1048
    0
                    if let Some(ref bias) = layer.ffn_gate_bias {
    1049
    0
                        for i in 0..intermediate_dim {
    1050
    0
                            scratch.ffn_gate[i] += bias[i];
    1051
    0
                        }
    1052
    0
                    }
    1053
    1054
                    // SiLU on gate, multiply with up
    1055
    0
                    ops::silu(&mut scratch.ffn_gate[..intermediate_dim]);
    1056
    0
                    for i in 0..intermediate_dim {
    1057
    0
                        scratch.ffn_gate[i] *= scratch.ffn_up[i];
    1058
    0
                    }
    1059
    1060
                    // PAR-127: Use Q8K-accelerated FFN down projection for Q4K weights
    1061
                    // Q6K uses f32 path since Q8K conversion overhead > bandwidth savings
    1062
    0
                    let use_q8k_down = intermediate_dim.is_multiple_of(256)
    1063
    0
                        && layer.ffn_down_weight.qtype == GGUF_TYPE_Q4_K;
    1064
    1065
    0
                    if use_q8k_down {
    1066
                        use crate::quantize::{
    1067
                            fused_q4k_q8k_parallel_matvec_into, quantize_activations_q8k_into,
    1068
                        };
    1069
    0
                        let inter_sb = intermediate_dim / 256;
    1070
    0
                        quantize_activations_q8k_into(
    1071
    0
                            &scratch.ffn_gate[..intermediate_dim],
    1072
    0
                            &mut scratch.q8k_inter_scales[..inter_sb],
    1073
    0
                            &mut scratch.q8k_inter_quants[..intermediate_dim],
    1074
    0
                        )?;
    1075
    0
                        fused_q4k_q8k_parallel_matvec_into(
    1076
    0
                            &layer.ffn_down_weight.data,
    1077
    0
                            &scratch.q8k_inter_scales[..inter_sb],
    1078
    0
                            &scratch.q8k_inter_quants[..intermediate_dim],
    1079
    0
                            layer.ffn_down_weight.in_dim,
    1080
    0
                            layer.ffn_down_weight.out_dim,
    1081
    0
                            &mut scratch.ffn_down,
    1082
    0
                        )?;
    1083
                    } else {
    1084
    0
                        self.fused_matmul_into(
    1085
    0
                            &scratch.ffn_gate[..intermediate_dim],
    1086
    0
                            &layer.ffn_down_weight,
    1087
    0
                            &mut scratch.ffn_down,
    1088
    0
                        )?;
    1089
                    }
    1090
    0
                    if let Some(ref bias) = layer.ffn_down_bias {
    1091
    0
                        for i in 0..hidden_dim {
    1092
    0
                            scratch.ffn_down[i] += bias[i];
    1093
    0
                        }
    1094
    0
                    }
    1095
                } else {
    1096
                    // GELU path (phi-2)
    1097
                    // PAR-129: Use Q8K-accelerated FFN for GELU models (Q4K only)
    1098
    0
                    let use_q8k_gelu_up = use_q8k_path && layer.ffn_up_weight.qtype == GGUF_TYPE_Q4_K;
    1099
    0
                    let use_q8k_gelu_down = intermediate_dim.is_multiple_of(256)
    1100
    0
                        && layer.ffn_down_weight.qtype == GGUF_TYPE_Q4_K;
    1101
    1102
    0
                    if use_q8k_gelu_up {
    1103
                        // Reuse already-quantized hidden from QKV (scratch.q8k_hidden_*)
    1104
                        use crate::quantize::fused_q4k_q8k_parallel_matvec_into;
    1105
    0
                        let hidden_sb = hidden_dim / 256;
    1106
    0
                        fused_q4k_q8k_parallel_matvec_into(
    1107
    0
                            &layer.ffn_up_weight.data,
    1108
    0
                            &scratch.q8k_hidden_scales[..hidden_sb],
    1109
    0
                            &scratch.q8k_hidden_quants[..hidden_dim],
    1110
    0
                            layer.ffn_up_weight.in_dim,
    1111
    0
                            layer.ffn_up_weight.out_dim,
    1112
    0
                            &mut scratch.ffn_up,
    1113
    0
                        )?;
    1114
                    } else {
    1115
    0
                        self.fused_matmul_into(
    1116
    0
                            &scratch.normed[..hidden_dim],
    1117
    0
                            &layer.ffn_up_weight,
    1118
    0
                            &mut scratch.ffn_up,
    1119
    0
                        )?;
    1120
                    }
    1121
    0
                    if let Some(ref bias) = layer.ffn_up_bias {
    1122
    0
                        for i in 0..intermediate_dim {
    1123
    0
                            scratch.ffn_up[i] += bias[i];
    1124
    0
                        }
    1125
    0
                    }
    1126
    0
                    ops::gelu(&mut scratch.ffn_up[..intermediate_dim]);
    1127
    1128
    0
                    if use_q8k_gelu_down {
    1129
                        use crate::quantize::{
    1130
                            fused_q4k_q8k_parallel_matvec_into, quantize_activations_q8k_into,
    1131
                        };
    1132
    0
                        let inter_sb = intermediate_dim / 256;
    1133
    0
                        quantize_activations_q8k_into(
    1134
    0
                            &scratch.ffn_up[..intermediate_dim],
    1135
    0
                            &mut scratch.q8k_inter_scales[..inter_sb],
    1136
    0
                            &mut scratch.q8k_inter_quants[..intermediate_dim],
    1137
    0
                        )?;
    1138
    0
                        fused_q4k_q8k_parallel_matvec_into(
    1139
    0
                            &layer.ffn_down_weight.data,
    1140
    0
                            &scratch.q8k_inter_scales[..inter_sb],
    1141
    0
                            &scratch.q8k_inter_quants[..intermediate_dim],
    1142
    0
                            layer.ffn_down_weight.in_dim,
    1143
    0
                            layer.ffn_down_weight.out_dim,
    1144
    0
                            &mut scratch.ffn_down,
    1145
    0
                        )?;
    1146
                    } else {
    1147
    0
                        self.fused_matmul_into(
    1148
    0
                            &scratch.ffn_up[..intermediate_dim],
    1149
    0
                            &layer.ffn_down_weight,
    1150
    0
                            &mut scratch.ffn_down,
    1151
    0
                        )?;
    1152
                    }
    1153
    0
                    if let Some(ref bias) = layer.ffn_down_bias {
    1154
    0
                        for i in 0..hidden_dim {
    1155
    0
                            scratch.ffn_down[i] += bias[i];
    1156
    0
                        }
    1157
    0
                    }
    1158
                }
    1159
    1160
                // 2h. FFN residual
    1161
    0
                for i in 0..hidden_dim {
    1162
    0
                    scratch.hidden[i] += scratch.ffn_down[i];
    1163
    0
                }
    1164
            }
    1165
    1166
            // 3. Final layer norm → scratch.normed
    1167
    0
            if use_rmsnorm {
    1168
    0
                ops::rms_norm_into(
    1169
    0
                    &scratch.hidden,
    1170
    0
                    &self.output_norm_weight,
    1171
    0
                    self.config.eps,
    1172
    0
                    &mut scratch.normed,
    1173
    0
                );
    1174
    0
            } else {
    1175
    0
                ops::layer_norm_into(
    1176
    0
                    &scratch.hidden,
    1177
    0
                    &self.output_norm_weight,
    1178
    0
                    self.output_norm_bias.as_deref(),
    1179
    0
                    self.config.eps,
    1180
    0
                    &mut scratch.normed,
    1181
    0
                );
    1182
    0
            }
    1183
    1184
            // 4. LM head → scratch.logits
    1185
    0
            self.fused_matmul_into(
    1186
    0
                &scratch.normed[..hidden_dim],
    1187
    0
                &self.lm_head_weight,
    1188
    0
                &mut scratch.logits,
    1189
    0
            )?;
    1190
    1191
    0
            Ok(())
    1192
    0
        }
    1193
    }