## 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/gpu/scheduler/model.rs

    Line
    Count
    Source
    1
    //! GPU Model and Configuration (PMAT-802)
    2
    //!
    3
    //! GpuModel, GpuModelConfig, GpuGenerateConfig, AttentionBuffers
    4
    5
    use crate::error::{RealizarError, Result};
    6
    use super::super::{
    7
        HybridScheduler, StreamingKVCache, exceeds_gpu_buffer_limit,
    8
        cpu_matmul_transposed_simd, cpu_matmul,
    9
    };
    10
    #[cfg(feature = "cuda")]
    11
    use super::core::CudaScheduler;
    12
    13
    /// GPU-accelerated model for M3 parity (128 tok/s target)
    14
    ///
    15
    /// Wraps standard Model and uses HybridScheduler for GPU-accelerated
    16
    /// matrix multiplications in the forward pass.
    17
    pub struct GpuModel {
    18
        /// Embedding weights (vocab_size x hidden_dim)
    19
        pub(crate) embedding_weights: Vec<f32>,
    20
        /// Linear layer weights for each block
    21
        /// Each block has: attn_q, attn_k, attn_v, attn_out, ffn_fc1, ffn_fc2
    22
        pub(crate) block_weights: Vec<BlockWeights>,
    23
        /// Final layer norm weights
    24
        pub(crate) final_norm_weight: Vec<f32>,
    25
        pub(crate) final_norm_bias: Vec<f32>,
    26
        /// LM head weights (hidden_dim x vocab_size)
    27
        pub(crate) lm_head_weight: Vec<f32>,
    28
        /// LM head weights transposed (vocab_size x hidden_dim) for fast CPU inference
    29
        pub(crate) lm_head_weight_t: Vec<f32>,
    30
        pub(crate) lm_head_bias: Vec<f32>,
    31
        /// GPU scheduler (HybridScheduler - may force CPU for m=1)
    32
        pub(crate) scheduler: HybridScheduler,
    33
        /// IMP-1003: Optional CUDA-only scheduler that ALWAYS uses GPU
    34
        /// When present, this scheduler is preferred over HybridScheduler for matmul
    35
        #[cfg(feature = "cuda")]
    36
        pub(crate) cuda_scheduler: Option<CudaScheduler>,
    37
        /// Model configuration
    38
        pub config: GpuModelConfig,
    39
        /// Pre-allocated attention buffers for optimized incremental decoding (M17)
    40
        pub(crate) attention_buffers: Option<AttentionBuffers>,
    41
    }
    42
    43
    /// Weights for a single transformer block
    44
    ///
    45
    /// Used by adapters (PMAT-106) to construct GpuModel from various formats.
    46
    pub struct BlockWeights {
    47
        /// Attention layer norm weight
    48
        pub attn_norm_weight: Vec<f32>,
    49
        /// Attention layer norm bias
    50
        pub attn_norm_bias: Vec<f32>,
    51
        /// Combined QKV projection weights (hidden_dim x 3*hidden_dim)
    52
        pub qkv_weight: Vec<f32>,
    53
        /// QKV projection bias (reserved for future use)
    54
        #[allow(dead_code)]
    55
        pub qkv_bias: Vec<f32>,
    56
        /// Output projection weight (hidden_dim x hidden_dim)
    57
        pub out_weight: Vec<f32>,
    58
        /// Output projection bias
    59
        pub out_bias: Vec<f32>,
    60
        /// FFN layer norm weight
    61
        pub ffn_norm_weight: Vec<f32>,
    62
        /// FFN layer norm bias
    63
        pub ffn_norm_bias: Vec<f32>,
    64
        /// FFN first layer weight (up projection)
    65
        pub ffn_fc1_weight: Vec<f32>,
    66
        /// FFN first layer bias
    67
        pub ffn_fc1_bias: Vec<f32>,
    68
        /// FFN second layer weight (down projection)
    69
        pub ffn_fc2_weight: Vec<f32>,
    70
        /// FFN second layer bias
    71
        pub ffn_fc2_bias: Vec<f32>,
    72
        /// FFN gate projection weight for SwiGLU (optional)
    73
        /// When present, FFN uses SwiGLU: down(SiLU(gate(x)) * up(x))
    74
        /// When None, FFN uses GELU: down(GELU(up(x)))
    75
        pub ffn_gate_weight: Option<Vec<f32>>,
    76
    }
    77
    78
    /// IMP-1007: Weight type for split-borrow matmul
    79
    ///
    80
    /// This enum specifies which weight matrix to use in matmul_split,
    81
    /// enabling zero-clone matmul operations by using Rust's split borrow pattern.
    82
    #[derive(Debug, Clone, Copy)]
    83
    pub enum WeightType {
    84
        /// QKV projection: [hidden_dim, qkv_dim]
    85
        Qkv,
    86
        /// Output projection: [hidden_dim, hidden_dim]
    87
        Output,
    88
        /// FFN FC1: [hidden_dim, intermediate_dim]
    89
        FfnFc1,
    90
        /// FFN FC2: [intermediate_dim, hidden_dim]
    91
        FfnFc2,
    92
        /// LM head: [hidden_dim, vocab_size]
    93
        LmHead,
    94
    }
    95
    96
    /// Configuration for GPU model
    97
    #[derive(Debug, Clone)]
    98
    pub struct GpuModelConfig {
    99
        /// Vocabulary size
    100
        pub vocab_size: usize,
    101
        /// Hidden dimension
    102
        pub hidden_dim: usize,
    103
        /// Number of attention heads (Q heads)
    104
        pub num_heads: usize,
    105
        /// Number of key-value heads for GQA (IMP-088)
    106
        /// For standard MHA: num_kv_heads == num_heads
    107
        /// For GQA (Qwen, Llama-3): num_kv_heads < num_heads
    108
        pub num_kv_heads: usize,
    109
        /// Number of transformer blocks
    110
        pub num_layers: usize,
    111
        /// FFN intermediate dimension
    112
        pub intermediate_dim: usize,
    113
        /// Layer normalization epsilon
    114
        pub eps: f32,
    115
        /// RoPE theta for rotary position embeddings (Phase 21)
    116
        /// Default: 10000.0 (standard LLaMA)
    117
        pub rope_theta: f32,
    118
    }
    119
    120
    impl GpuModelConfig {
    121
        /// Head dimension (hidden_dim / num_heads)
    122
        #[inline]
    123
    5.36k
        pub fn head_dim(&self) -> usize {
    124
    5.36k
            self.hidden_dim / self.num_heads
    125
    5.36k
        }
    126
    127
        /// K/V dimension for GQA (num_kv_heads * head_dim)
    128
        #[inline]
    129
    3.58k
        pub fn kv_dim(&self) -> usize {
    130
    3.58k
            self.num_kv_heads * self.head_dim()
    131
    3.58k
        }
    132
    133
        /// Total QKV projection output dimension
    134
        /// For MHA: 3 * hidden_dim
    135
        /// For GQA: hidden_dim + 2 * kv_dim
    136
        #[inline]
    137
    1.79k
        pub fn qkv_dim(&self) -> usize {
    138
    1.79k
            self.hidden_dim + 2 * self.kv_dim()
    139
    1.79k
        }
    140
    141
        /// Whether this is a GQA model (num_kv_heads < num_heads)
    142
        #[inline]
    143
    2
        pub fn is_gqa(&self) -> bool {
    144
    2
            self.num_kv_heads < self.num_heads
    145
    2
        }
    146
    }
    147
    148
    /// Configuration for GPU text generation (M14: E2E Inference)
    149
    #[derive(Debug, Clone)]
    150
    pub struct GpuGenerateConfig {
    151
        /// Maximum tokens to generate
    152
        pub max_tokens: usize,
    153
        /// Sampling temperature (0.0 = greedy)
    154
        pub temperature: f32,
    155
        /// Top-k sampling (1 = greedy)
    156
        pub top_k: usize,
    157
        /// Stop token IDs
    158
        pub stop_tokens: Vec<usize>,
    159
    }
    160
    161
    impl Default for GpuGenerateConfig {
    162
    2
        fn default() -> Self {
    163
    2
            Self {
    164
    2
                max_tokens: 64,
    165
    2
                temperature: 0.0,
    166
    2
                top_k: 1,
    167
    2
                stop_tokens: Vec::new(),
    168
    2
            }
    169
    2
        }
    170
    }
    171
    172
    impl GpuGenerateConfig {
    173
        /// Create config for deterministic (greedy) generation
    174
        #[must_use]
    175
    11
        pub fn deterministic(max_tokens: usize) -> Self {
    176
    11
            Self {
    177
    11
                max_tokens,
    178
    11
                temperature: 0.0,
    179
    11
                top_k: 1,
    180
    11
                stop_tokens: Vec::new(),
    181
    11
            }
    182
    11
        }
    183
    184
        /// Create config with temperature and top-k sampling
    185
        #[must_use]
    186
    3
        pub fn with_sampling(max_tokens: usize, temperature: f32, top_k: usize) -> Self {
    187
    3
            Self {
    188
    3
                max_tokens,
    189
    3
                temperature,
    190
    3
                top_k,
    191
    3
                stop_tokens: Vec::new(),
    192
    3
            }
    193
    3
        }
    194
    195
        /// Add stop tokens to config
    196
        #[must_use]
    197
    4
        pub fn with_stop_tokens(mut self, stop_tokens: Vec<usize>) -> Self {
    198
    4
            self.stop_tokens = stop_tokens;
    199
    4
            self
    200
    4
        }
    201
    }
    202
    203
    /// Pre-allocated attention buffers for optimized incremental decoding (M17)
    204
    ///
    205
    /// Eliminates per-token memory allocation during incremental generation by
    206
    /// reusing pre-allocated buffers for Q, attention scores, and output.
    207
    #[derive(Debug)]
    208
    pub struct AttentionBuffers {
    209
        /// Q buffer for single-token attention [hidden_dim]
    210
        pub q_buffer: Vec<f32>,
    211
        /// Attention scores buffer [num_heads * max_seq_len]
    212
        pub scores_buffer: Vec<f32>,
    213
        /// Attention output buffer [hidden_dim]
    214
        pub output_buffer: Vec<f32>,
    215
        /// K/V projection buffer [hidden_dim]
    216
        pub kv_proj_buffer: Vec<f32>,
    217
        /// Intermediate FFN buffer [intermediate_dim]
    218
        pub ffn_buffer: Vec<f32>,
    219
        /// Max sequence length these buffers support
    220
        pub max_seq_len: usize,
    221
    }
    222
    223
    impl AttentionBuffers {
    224
        /// Create pre-allocated attention buffers from model config
    225
        ///
    226
        /// # Arguments
    227
        ///
    228
        /// * `config` - Model configuration
    229
        /// * `max_seq_len` - Maximum sequence length to support
    230
        #[must_use]
    231
    8
        pub fn new(config: &GpuModelConfig, max_seq_len: usize) -> Self {
    232
    8
            Self {
    233
    8
                q_buffer: vec![0.0; config.hidden_dim],
    234
    8
                scores_buffer: vec![0.0; config.num_heads * max_seq_len],
    235
    8
                output_buffer: vec![0.0; config.hidden_dim],
    236
    8
                kv_proj_buffer: vec![0.0; config.hidden_dim],
    237
    8
                ffn_buffer: vec![0.0; config.intermediate_dim],
    238
    8
                max_seq_len,
    239
    8
            }
    240
    8
        }
    241
    242
        /// Reset all buffers to zero (for reuse)
    243
    1
        pub fn reset(&mut self) {
    244
    1
            self.q_buffer.fill(0.0);
    245
    1
            self.scores_buffer.fill(0.0);
    246
    1
            self.output_buffer.fill(0.0);
    247
    1
            self.kv_proj_buffer.fill(0.0);
    248
    1
            self.ffn_buffer.fill(0.0);
    249
    1
        }
    250
    }
    251
    252
    impl GpuModel {
    253
        /// Create a new GPU-accelerated model with random initialization
    254
        ///
    255
        /// # Errors
    256
        ///
    257
        /// Returns error if GPU initialization fails
    258
    19
        pub fn new(config: GpuModelConfig) -> Result<Self> {
    259
    19
            let scheduler = HybridScheduler::new()?0;
    260
    261
            // Initialize weights (small random values for testing)
    262
    19
            let embedding_weights = vec![0.01f32; config.vocab_size * config.hidden_dim];
    263
    264
    19
            let mut block_weights = Vec::with_capacity(config.num_layers);
    265
    46
            for _ in 0..config.num_layers19 {
    266
    46
                block_weights.push(BlockWeights {
    267
    46
                    attn_norm_weight: vec![1.0f32; config.hidden_dim],
    268
    46
                    attn_norm_bias: vec![0.0f32; config.hidden_dim],
    269
    46
                    qkv_weight: vec![0.01f32; config.hidden_dim * 3 * config.hidden_dim],
    270
    46
                    qkv_bias: vec![0.0f32; 3 * config.hidden_dim],
    271
    46
                    out_weight: vec![0.01f32; config.hidden_dim * config.hidden_dim],
    272
    46
                    out_bias: vec![0.0f32; config.hidden_dim],
    273
    46
                    ffn_norm_weight: vec![1.0f32; config.hidden_dim],
    274
    46
                    ffn_norm_bias: vec![0.0f32; config.hidden_dim],
    275
    46
                    ffn_fc1_weight: vec![0.01f32; config.hidden_dim * config.intermediate_dim],
    276
    46
                    ffn_fc1_bias: vec![0.0f32; config.intermediate_dim],
    277
    46
                    ffn_fc2_weight: vec![0.01f32; config.intermediate_dim * config.hidden_dim],
    278
    46
                    ffn_fc2_bias: vec![0.0f32; config.hidden_dim],
    279
    46
                    ffn_gate_weight: None, // No SwiGLU in test models
    280
    46
                });
    281
    46
            }
    282
    283
    19
            let final_norm_weight = vec![1.0f32; config.hidden_dim];
    284
    19
            let final_norm_bias = vec![0.0f32; config.hidden_dim];
    285
    19
            let lm_head_weight = vec![0.01f32; config.hidden_dim * config.vocab_size];
    286
    19
            let lm_head_bias = vec![0.0f32; config.vocab_size];
    287
    288
            // Pre-compute transposed LM head for fast CPU inference
    289
            // Original: [hidden_dim, vocab_size] -> Transposed: [vocab_size, hidden_dim]
    290
    19
            let lm_head_weight_t =
    291
    19
                Self::transpose_weights(&lm_head_weight, config.hidden_dim, config.vocab_size);
    292
    293
    19
            Ok(Self {
    294
    19
                embedding_weights,
    295
    19
                block_weights,
    296
    19
                final_norm_weight,
    297
    19
                final_norm_bias,
    298
    19
                lm_head_weight,
    299
    19
                lm_head_weight_t,
    300
    19
                lm_head_bias,
    301
    19
                scheduler,
    302
    19
                #[cfg(feature = "cuda")]
    303
    19
                cuda_scheduler: None,
    304
    19
                config,
    305
    19
                attention_buffers: None,
    306
    19
            })
    307
    19
        }
    308
    309
        /// IMP-1003: Create GPU model with CUDA-only scheduler
    310
        ///
    311
        /// Unlike `new()`, this constructor creates a model that ALWAYS uses CUDA
    312
        /// for matmul operations, even for m=1 (single-token generation).
    313
        ///
    314
        /// # Errors
    315
        ///
    316
        /// Returns error if GPU or CUDA initialization fails
    317
        #[cfg(feature = "cuda")]
    318
        pub fn new_with_cuda(config: GpuModelConfig) -> Result<Self> {
    319
            let scheduler = HybridScheduler::new()?;
    320
            let cuda_scheduler = Some(CudaScheduler::new()?);
    321
    322
            // Initialize weights (small random values for testing)
    323
            let embedding_weights = vec![0.01f32; config.vocab_size * config.hidden_dim];
    324
    325
            let mut block_weights = Vec::with_capacity(config.num_layers);
    326
            for _ in 0..config.num_layers {
    327
                block_weights.push(BlockWeights {
    328
                    attn_norm_weight: vec![1.0f32; config.hidden_dim],
    329
                    attn_norm_bias: vec![0.0f32; config.hidden_dim],
    330
                    qkv_weight: vec![0.01f32; config.hidden_dim * config.qkv_dim()],
    331
                    qkv_bias: vec![0.0f32; config.qkv_dim()],
    332
                    out_weight: vec![0.01f32; config.hidden_dim * config.hidden_dim],
    333
                    out_bias: vec![0.0f32; config.hidden_dim],
    334
                    ffn_norm_weight: vec![1.0f32; config.hidden_dim],
    335
                    ffn_norm_bias: vec![0.0f32; config.hidden_dim],
    336
                    ffn_fc1_weight: vec![0.01f32; config.hidden_dim * config.intermediate_dim],
    337
                    ffn_fc1_bias: vec![0.0f32; config.intermediate_dim],
    338
                    ffn_fc2_weight: vec![0.01f32; config.intermediate_dim * config.hidden_dim],
    339
                    ffn_fc2_bias: vec![0.0f32; config.hidden_dim],
    340
                    ffn_gate_weight: None, // No SwiGLU in test models
    341
                });
    342
            }
    343
    344
            let final_norm_weight = vec![1.0f32; config.hidden_dim];
    345
            let final_norm_bias = vec![0.0f32; config.hidden_dim];
    346
            let lm_head_weight = vec![0.01f32; config.hidden_dim * config.vocab_size];
    347
            let lm_head_bias = vec![0.0f32; config.vocab_size];
    348
    349
            let lm_head_weight_t =
    350
                Self::transpose_weights(&lm_head_weight, config.hidden_dim, config.vocab_size);
    351
    352
            Ok(Self {
    353
                embedding_weights,
    354
                block_weights,
    355
                final_norm_weight,
    356
                final_norm_bias,
    357
                lm_head_weight,
    358
                lm_head_weight_t,
    359
                lm_head_bias,
    360
                scheduler,
    361
                cuda_scheduler,
    362
                config,
    363
                attention_buffers: None,
    364
            })
    365
        }
    366
    367
        /// IMP-1003: Check if this model has CUDA scheduler enabled
    368
        #[cfg(feature = "cuda")]
    369
        #[must_use]
    370
        pub fn has_cuda_scheduler(&self) -> bool {
    371
            self.cuda_scheduler.is_some()
    372
        }
    373
    374
        /// IMP-1003: Perform matmul using CUDA scheduler (always GPU, even for m=1)
    375
        ///
    376
        /// # Errors
    377
        ///
    378
        /// Returns error if CUDA scheduler is not available or matmul fails
    379
        #[cfg(feature = "cuda")]
    380
        #[allow(clippy::many_single_char_names)]
    381
        pub fn cuda_matmul(
    382
            &mut self,
    383
            a: &[f32],
    384
            b: &[f32],
    385
            m: usize,
    386
            k: usize,
    387
            n: usize,
    388
        ) -> Result<Vec<f32>> {
    389
            if let Some(ref mut cuda_sched) = self.cuda_scheduler {
    390
                cuda_sched.matmul(a, b, m, k, n)
    391
            } else {
    392
                // Fallback to HybridScheduler
    393
                self.scheduler.matmul(a, b, m, k, n)
    394
            }
    395
        }
    396
    397
        /// IMP-1005: Unified matmul dispatch that prefers CudaScheduler when available
    398
        ///
    399
        /// This method is used throughout forward_gpu() and forward_block_idx() to
    400
        /// ensure CUDA is used for all matmul operations when cuda_scheduler is present.
    401
        ///
    402
        /// # Errors
    403
        ///
    404
        /// Returns error if matmul fails
    405
        #[allow(clippy::many_single_char_names)]
    406
    6.76k
        pub fn do_matmul(
    407
    6.76k
            &mut self,
    408
    6.76k
            a: &[f32],
    409
    6.76k
            b: &[f32],
    410
    6.76k
            m: usize,
    411
    6.76k
            k: usize,
    412
    6.76k
            n: usize,
    413
    6.76k
        ) -> Result<Vec<f32>> {
    414
            #[cfg(feature = "cuda")]
    415
            if let Some(ref mut cuda_sched) = self.cuda_scheduler {
    416
                return cuda_sched.matmul(a, b, m, k, n);
    417
            }
    418
            // Fallback to HybridScheduler (or always use it when cuda feature disabled)
    419
    6.76k
            self.scheduler.matmul(a, b, m, k, n)
    420
    6.76k
        }
    421
    422
        /// IMP-1007: Zero-clone matmul using split borrow pattern
    423
        ///
    424
        /// This method eliminates weight cloning by using Rust's split borrow pattern.
    425
        /// It directly borrows weights from block_weights while mutably borrowing schedulers.
    426
        ///
    427
        /// # Arguments
    428
        ///
    429
        /// * `input` - Input tensor
    430
        /// * `block_idx` - Block index for block weights (ignored for LmHead)
    431
        /// * `op` - Which matmul operation/weight to use
    432
        ///
    433
        /// # Errors
    434
        ///
    435
        /// Returns error if matmul fails
    436
    0
        pub fn matmul_split(
    437
    0
            &mut self,
    438
    0
            input: &[f32],
    439
    0
            block_idx: usize,
    440
    0
            op: WeightType,
    441
    0
        ) -> Result<Vec<f32>> {
    442
            // IMP-1007: Use split borrowing to avoid weight cloning
    443
            // Extract dimensions from config (Copy types, no borrow conflict)
    444
    0
            let hidden_dim = self.config.hidden_dim;
    445
    0
            let qkv_dim = self.config.qkv_dim();
    446
    0
            let intermediate_dim = self.config.intermediate_dim;
    447
    0
            let vocab_size = self.config.vocab_size;
    448
    449
            // Get weight reference and dimensions based on operation
    450
    0
            let (weight, m, k, n) = match op {
    451
    0
                WeightType::Qkv => (
    452
    0
                    &self.block_weights[block_idx].qkv_weight,
    453
    0
                    1,
    454
    0
                    hidden_dim,
    455
    0
                    qkv_dim,
    456
    0
                ),
    457
    0
                WeightType::Output => (
    458
    0
                    &self.block_weights[block_idx].out_weight,
    459
    0
                    1,
    460
    0
                    hidden_dim,
    461
    0
                    hidden_dim,
    462
    0
                ),
    463
    0
                WeightType::FfnFc1 => (
    464
    0
                    &self.block_weights[block_idx].ffn_fc1_weight,
    465
    0
                    1,
    466
    0
                    hidden_dim,
    467
    0
                    intermediate_dim,
    468
    0
                ),
    469
    0
                WeightType::FfnFc2 => (
    470
    0
                    &self.block_weights[block_idx].ffn_fc2_weight,
    471
    0
                    1,
    472
    0
                    intermediate_dim,
    473
    0
                    hidden_dim,
    474
    0
                ),
    475
    0
                WeightType::LmHead => (&self.lm_head_weight, 1, hidden_dim, vocab_size),
    476
            };
    477
    478
            // Clone weight to work around borrow checker - this is the safe fallback.
    479
            // For zero-clone operations, use matmul_zero_clone() instead (IMP-1007).
    480
    0
            let weight_clone = weight.clone();
    481
    482
            // Now call do_matmul with cloned weight
    483
    0
            self.do_matmul(input, &weight_clone, m, k, n)
    484
    0
        }
    485
    486
        /// IMP-1007: Zero-clone matmul helper using explicit scheduler extraction
    487
        ///
    488
        /// This is a more aggressive optimization that temporarily extracts the
    489
        /// cuda_scheduler to enable truly zero-clone matmul operations.
    490
        ///
    491
        /// # Safety
    492
        ///
    493
        /// This method uses `Option::take()` to temporarily move the scheduler,
    494
        /// which is safe but requires careful handling to restore it.
    495
        #[cfg(feature = "cuda")]
    496
        pub fn matmul_zero_clone(
    497
            &mut self,
    498
            input: &[f32],
    499
            block_idx: usize,
    500
            op: WeightType,
    501
        ) -> Result<Vec<f32>> {
    502
            // Extract dimensions
    503
            let hidden_dim = self.config.hidden_dim;
    504
            let qkv_dim = self.config.qkv_dim();
    505
            let intermediate_dim = self.config.intermediate_dim;
    506
            let vocab_size = self.config.vocab_size;
    507
    508
            // Temporarily take cuda_scheduler out of self
    509
            let mut cuda_sched: Option<CudaScheduler> = self.cuda_scheduler.take();
    510
    511
            // Now we can borrow block_weights freely
    512
            let (weight, m, k, n) = match op {
    513
                WeightType::Qkv => (
    514
                    &self.block_weights[block_idx].qkv_weight,
    515
                    1,
    516
                    hidden_dim,
    517
                    qkv_dim,
    518
                ),
    519
                WeightType::Output => (
    520
                    &self.block_weights[block_idx].out_weight,
    521
                    1,
    522
                    hidden_dim,
    523
                    hidden_dim,
    524
                ),
    525
                WeightType::FfnFc1 => (
    526
                    &self.block_weights[block_idx].ffn_fc1_weight,
    527
                    1,
    528
                    hidden_dim,
    529
                    intermediate_dim,
    530
                ),
    531
                WeightType::FfnFc2 => (
    532
                    &self.block_weights[block_idx].ffn_fc2_weight,
    533
                    1,
    534
                    intermediate_dim,
    535
                    hidden_dim,
    536
                ),
    537
                WeightType::LmHead => (&self.lm_head_weight, 1, hidden_dim, vocab_size),
    538
            };
    539
    540
            // Perform matmul with extracted scheduler
    541
            let result: Result<Vec<f32>> = if let Some(sched) = cuda_sched.as_mut() {
    542
                CudaScheduler::matmul(sched, input, weight, m, k, n)
    543
            } else {
    544
                self.scheduler.matmul(input, weight, m, k, n)
    545
            };
    546
    547
            // Restore cuda_scheduler
    548
            self.cuda_scheduler = cuda_sched;
    549
    550
            result
    551
        }
    552
    553
        // =========================================================================
    554
        // IMP-1008: RefCell-based zero-clone matmul (interior mutability pattern)
    555
        // =========================================================================
    556
    557
        /// IMP-1008: Zero-clone matmul using interior mutability
    558
        ///
    559
        /// This method takes `&self` instead of `&mut self` by wrapping scheduler
    560
        /// access in RefCell. This eliminates the need to clone weights.
    561
        ///
    562
        /// # Errors
    563
        ///
    564
        /// Returns error if matmul fails or RefCell is already borrowed.
    565
        #[cfg(feature = "cuda")]
    566
        #[allow(clippy::many_single_char_names)]
    567
        pub fn matmul_refcell(
    568
            &self,
    569
            a: &[f32],
    570
            b: &[f32],
    571
            m: usize,
    572
            k: usize,
    573
            n: usize,
    574
        ) -> Result<Vec<f32>> {
    575
            // IMP-1008: For RefCell pattern, we need to use a different approach
    576
            // Since cuda_scheduler is Option<CudaScheduler>, we use UnsafeCell
    577
            // pattern with explicit unsafe block to avoid changing struct layout.
    578
            //
    579
            // This is safe because:
    580
            // 1. We only access cuda_scheduler mutably here
    581
            // 2. No other code paths access it during matmul
    582
            // 3. This is single-threaded execution
    583
    584
            // Use raw pointer to bypass borrow checker (safe in single-threaded context)
    585
            // SAFETY: This is safe because:
    586
            // - We're in single-threaded context (LLM inference)
    587
            // - cuda_scheduler is only accessed through this method during matmul
    588
            // - The borrow is released before returning
    589
            let cuda_sched_ptr = std::ptr::addr_of!(self.cuda_scheduler).cast_mut();
    590
    591
            // SAFETY: Memory safety ensured by bounds checking and alignment
    592
            let result: Result<Vec<f32>> = unsafe {
    593
                if let Some(sched) = (*cuda_sched_ptr).as_mut() {
    594
                    CudaScheduler::matmul(sched, a, b, m, k, n)
    595
                } else {
    596
                    // Fallback to HybridScheduler (also needs mut access)
    597
                    let sched_ptr = std::ptr::addr_of!(self.scheduler).cast_mut();
    598
                    (*sched_ptr).matmul(a, b, m, k, n)
    599
                }
    600
            };
    601
            result
    602
        }
    603
    604
        /// IMP-1008: Forward single block without weight cloning
    605
        ///
    606
        /// Uses interior mutability pattern to avoid cloning weights on each matmul.
    607
        /// This method takes `&self` instead of `&mut self`.
    608
        ///
    609
        /// # Errors
    610
        ///
    611
        /// Returns error if forward pass fails.
    612
        #[cfg(feature = "cuda")]
    613
        pub fn forward_block_refcell(
    614
            &self,
    615
            input: &[f32],
    616
            block_idx: usize,
    617
            kv_cache: &mut StreamingKVCache,
    618
        ) -> Result<Vec<f32>> {
    619
            // Phase 21 Debug: trace first forward call only
    620
            static DEBUG_COUNTER: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
    621
            let call_count = DEBUG_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    622
            let debug_this_call = block_idx == 0 && call_count == 0;  // Only first call to block 0
    623
    624
            // Extract config values (Copy types, no borrow conflict)
    625
            let hidden_dim = self.config.hidden_dim;
    626
            let num_heads = self.config.num_heads;
    627
            let head_dim = self.config.head_dim();
    628
            let kv_dim = self.config.kv_dim();
    629
            let qkv_dim = self.config.qkv_dim();
    630
            let intermediate_dim = self.config.intermediate_dim;
    631
            let eps = self.config.eps;
    632
            let num_kv_heads = self.config.num_kv_heads;
    633
    634
            if debug_this_call {
    635
                eprintln!("[PHASE21] forward_block_refcell START block_idx={}", block_idx);
    636
                eprintln!("[PHASE21] input L2: {:.4}", input.iter().map(|x| x*x).sum::<f32>().sqrt());
    637
            }
    638
    639
            // IMP-1008: No cloning! Direct reference to weights
    640
            // Pre-attention layer norm (static function avoids &self borrow)
    641
            let normed = Self::layer_norm_static(
    642
                input,
    643
                &self.block_weights[block_idx].attn_norm_weight,
    644
                &self.block_weights[block_idx].attn_norm_bias,
    645
                hidden_dim,
    646
                eps,
    647
            );
    648
    649
            // QKV projection - NO CLONE!
    650
            let mut qkv = self.matmul_refcell(
    651
                &normed,
    652
                &self.block_weights[block_idx].qkv_weight,
    653
                1,
    654
                hidden_dim,
    655
                qkv_dim,
    656
            )?;
    657
    658
            if debug_this_call {
    659
                eprintln!("[PHASE21] QKV L2: {:.4}", qkv.iter().map(|x| x*x).sum::<f32>().sqrt());
    660
            }
    661
    662
            // Get current position BEFORE caching (Phase 21)
    663
            let (cached_k_ref, _) = kv_cache.get_valid(block_idx);
    664
            let current_pos = cached_k_ref.len() / kv_dim;
    665
    666
            // Phase 21: Apply RoPE to Q and K BEFORE caching
    667
            // Without RoPE, attention has no position information and produces garbage
    668
            let rope_theta = self.config.rope_theta;
    669
            Self::apply_rope_inline(&mut qkv[0..hidden_dim], num_heads, head_dim, rope_theta, current_pos);
    670
            Self::apply_rope_inline(&mut qkv[hidden_dim..hidden_dim + kv_dim], num_kv_heads, head_dim, rope_theta, current_pos);
    671
    672
            // Split QKV (GQA: K/V have kv_dim, not hidden_dim) - after RoPE
    673
            let q = qkv[0..hidden_dim].to_vec();
    674
            let k_new = qkv[hidden_dim..hidden_dim + kv_dim].to_vec();
    675
            let v_new = qkv[hidden_dim + kv_dim..].to_vec();
    676
    677
            // Get cached K/V and clone to avoid borrow issues with kv_cache
    678
            let (cached_k, cached_v) = kv_cache.get_valid(block_idx);
    679
            let keys_cached = cached_k.to_vec();
    680
            let vals_cached = cached_v.to_vec();
    681
    682
            // Append new K/V (with RoPE applied) to cache
    683
            kv_cache.append(block_idx, &k_new, &v_new);
    684
    685
            // Build full K/V (cached + new)
    686
            let kv_len = keys_cached.len() / kv_dim + 1;
    687
            let mut full_k = keys_cached;
    688
            full_k.extend_from_slice(&k_new);
    689
            let mut full_v = vals_cached;
    690
            full_v.extend_from_slice(&v_new);
    691
    692
            // GQA attention (IMP-089): static method to avoid borrow conflicts
    693
            let attn_output = Self::gqa_multihead_attention(
    694
                &q,
    695
                &full_k,
    696
                &full_v,
    697
                kv_len,
    698
                num_heads,
    699
                num_kv_heads,
    700
                head_dim,
    701
            );
    702
    703
            if debug_this_call {
    704
                eprintln!("[PHASE21] attn_output L2: {:.4}", attn_output.iter().map(|x| x*x).sum::<f32>().sqrt());
    705
            }
    706
    707
            // Output projection - NO CLONE!
    708
            let attn_proj = self.matmul_refcell(
    709
                &attn_output,
    710
                &self.block_weights[block_idx].out_weight,
    711
                1,
    712
                hidden_dim,
    713
                hidden_dim,
    714
            )?;
    715
    716
            // Add residual and bias
    717
            let out_bias = &self.block_weights[block_idx].out_bias;
    718
            let post_attn: Vec<f32> = input
    719
                .iter()
    720
                .zip(attn_proj.iter())
    721
                .zip(out_bias.iter())
    722
                .map(|((&i, &a), &b)| i + a + b)
    723
                .collect();
    724
    725
            // FFN with layer norm (static function)
    726
            let ffn_normed = Self::layer_norm_static(
    727
                &post_attn,
    728
                &self.block_weights[block_idx].ffn_norm_weight,
    729
                &self.block_weights[block_idx].ffn_norm_bias,
    730
                hidden_dim,
    731
                eps,
    732
            );
    733
    734
            // FFN: SwiGLU when gate weight exists, otherwise GELU
    735
            let fc1_activated: Vec<f32> = if let Some(ref gate_weight) = self.block_weights[block_idx].ffn_gate_weight {
    736
                // SwiGLU: silu(gate(x)) * up(x)
    737
                // Up projection
    738
                let up_out = self.matmul_refcell(
    739
                    &ffn_normed,
    740
                    &self.block_weights[block_idx].ffn_fc1_weight,
    741
                    1,
    742
                    hidden_dim,
    743
                    intermediate_dim,
    744
                )?;
    745
    746
                // Gate projection
    747
                let gate_out = self.matmul_refcell(
    748
                    &ffn_normed,
    749
                    gate_weight,
    750
                    1,
    751
                    hidden_dim,
    752
                    intermediate_dim,
    753
                )?;
    754
    755
                // SwiGLU: silu(gate) * up
    756
                // silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
    757
                up_out
    758
                    .iter()
    759
                    .zip(gate_out.iter())
    760
                    .map(|(&u, &g)| {
    761
                        let silu_g = g / (1.0 + (-g).exp());
    762
                        silu_g * u
    763
                    })
    764
                    .collect()
    765
            } else {
    766
                // Standard GELU FFN
    767
                let fc1_out = self.matmul_refcell(
    768
                    &ffn_normed,
    769
                    &self.block_weights[block_idx].ffn_fc1_weight,
    770
                    1,
    771
                    hidden_dim,
    772
                    intermediate_dim,
    773
                )?;
    774
    775
                // Add bias and GELU activation
    776
                let ffn_fc1_bias = &self.block_weights[block_idx].ffn_fc1_bias;
    777
                fc1_out
    778
                    .iter()
    779
                    .zip(ffn_fc1_bias.iter())
    780
                    .map(|(&x, &b)| {
    781
                        let x_b = x + b;
    782
                        x_b * 0.5 + x_b * 0.5 * (0.797_884_6 * (x_b + 0.044_715 * x_b.powi(3))).tanh()
    783
                    })
    784
                    .collect()
    785
            };
    786
    787
            // FFN FC2 (down projection) - NO CLONE!
    788
            let fc2_out = self.matmul_refcell(
    789
                &fc1_activated,
    790
                &self.block_weights[block_idx].ffn_fc2_weight,
    791
                1,
    792
                intermediate_dim,
    793
                hidden_dim,
    794
            )?;
    795
    796
            // Add bias and residual
    797
            let ffn_fc2_bias = &self.block_weights[block_idx].ffn_fc2_bias;
    798
            let output: Vec<f32> = post_attn
    799
                .iter()
    800
                .zip(fc2_out.iter())
    801
                .zip(ffn_fc2_bias.iter())
    802
                .map(|((&h, &f), &b)| h + f + b)
    803
                .collect();
    804
    805
            if debug_this_call {
    806
                eprintln!("[PHASE21] block output L2: {:.4}", output.iter().map(|x| x*x).sum::<f32>().sqrt());
    807
            }
    808
    809
            Ok(output)
    810
        }
    811
    812
        /// IMP-1008: Full incremental forward pass without weight cloning
    813
        ///
    814
        /// Uses interior mutability pattern throughout for zero-clone operation.
    815
        ///
    816
        /// # Errors
    817
        ///
    818
        /// Returns error if forward pass fails.
    819
        #[cfg(feature = "cuda")]
    820
        pub fn forward_refcell(
    821
            &self,
    822
            token_id: usize,
    823
            kv_cache: &mut StreamingKVCache,
    824
        ) -> Result<Vec<f32>> {
    825
            // Phase 21: Debug first call only
    826
            static FWD_COUNTER: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
    827
            let fwd_count = FWD_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    828
            let debug_this_fwd = fwd_count == 0;
    829
    830
            if token_id >= self.config.vocab_size {
    831
                return Err(RealizarError::InvalidShape {
    832
                    reason: format!(
    833
                        "Token ID {} out of bounds (vocab_size={})",
    834
                        token_id, self.config.vocab_size
    835
                    ),
    836
                });
    837
            }
    838
    839
            let hidden_dim = self.config.hidden_dim;
    840
    841
            // Embed single token
    842
            let offset = token_id * hidden_dim;
    843
            let mut hidden = self.embedding_weights[offset..offset + hidden_dim].to_vec();
    844
    845
            // Process through all blocks - NO CLONE!
    846
            for block_idx in 0..self.config.num_layers {
    847
                hidden = self.forward_block_refcell(&hidden, block_idx, kv_cache)?;
    848
            }
    849
    850
            // Final layer norm
    851
            hidden = self.layer_norm_refcell(&hidden, &self.final_norm_weight, &self.final_norm_bias);
    852
    853
            // LM head projection
    854
            let lm_head_elements = hidden_dim * self.config.vocab_size;
    855
            let output = if exceeds_gpu_buffer_limit(lm_head_elements) {
    856
                // CPU path with transposed weights + SIMD + fused bias
    857
                cpu_matmul_transposed_simd(
    858
                    &hidden,
    859
                    &self.lm_head_weight_t,
    860
                    &self.lm_head_bias,
    861
                    hidden_dim,
    862
                    self.config.vocab_size,
    863
                )
    864
            } else {
    865
                // GPU path - NO CLONE!
    866
                let vocab_size = self.config.vocab_size;
    867
                let logits =
    868
                    self.matmul_refcell(&hidden, &self.lm_head_weight, 1, hidden_dim, vocab_size)?;
    869
                // Add bias
    870
                logits
    871
                    .into_iter()
    872
                    .zip(self.lm_head_bias.iter())
    873
                    .map(|(l, &b)| l + b)
    874
                    .collect()
    875
            };
    876
    877
            if debug_this_fwd {
    878
                // Find argmax
    879
                let (argmax_idx, argmax_val) = output.iter()
    880
                    .enumerate()
    881
                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    882
                    .unwrap_or((0, &0.0));
    883
                eprintln!("[PHASE21] forward_refcell: final hidden L2: {:.4}", hidden.iter().map(|x| x*x).sum::<f32>().sqrt());
    884
                eprintln!("[PHASE21] forward_refcell: logits argmax: {} (val: {:.4})", argmax_idx, argmax_val);
    885
            }
    886
    887
            Ok(output)
    888
        }
    889
    890
        /// IMP-1008: Layer norm with RefCell pattern (takes &self)
    891
        #[cfg(feature = "cuda")]
    892
        fn layer_norm_refcell(&self, input: &[f32], weight: &[f32], bias: &[f32]) -> Vec<f32> {
    893
            Self::layer_norm_static(input, weight, bias, self.config.hidden_dim, self.config.eps)
    894
        }
    895
    896
        /// IMP-1008: Generate tokens without weight cloning
    897
        ///
    898
        /// Uses interior mutability pattern for zero-clone inference.
    899
        ///
    900
        /// # Errors
    901
        ///
    902
        /// Returns error if generation fails.
    903
        #[cfg(feature = "cuda")]
    904
        pub fn generate_refcell(
    905
            &self,
    906
            prompt: &[usize],
    907
            config: &GpuGenerateConfig,
    908
        ) -> Result<Vec<usize>> {
    909
            if prompt.is_empty() {
    910
                return Err(RealizarError::InvalidShape {
    911
                    reason: "Prompt cannot be empty".to_string(),
    912
                });
    913
            }
    914
    915
            let num_kv_heads = self.config.num_kv_heads;
    916
            let head_dim = self.config.head_dim();
    917
            let max_seq_len = prompt.len() + config.max_tokens;
    918
    919
            // Initialize KV cache
    920
            let mut kv_cache =
    921
                StreamingKVCache::new(self.config.num_layers, max_seq_len, num_kv_heads, head_dim);
    922
    923
            let mut tokens = prompt.to_vec();
    924
    925
            // Process prompt tokens to populate KV cache
    926
            for &token_id in prompt {
    927
                let _ = self.forward_refcell(token_id, &mut kv_cache)?;
    928
            }
    929
    930
            // Generate new tokens
    931
            for _ in 0..config.max_tokens {
    932
                let last_token = *tokens.last().unwrap_or(&0);
    933
                let logits = self.forward_refcell(last_token, &mut kv_cache)?;
    934
    935
                // Sample next token (greedy when temperature=0, otherwise top-k)
    936
                let next_token = if config.temperature == 0.0 || config.top_k == 1 {
    937
                    // Greedy decoding
    938
                    logits
    939
                        .iter()
    940
                        .enumerate()
    941
                        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    942
                        .map_or(0, |(idx, _)| idx)
    943
                } else {
    944
                    // Top-k sampling with temperature
    945
                    Self::sample_topk_generate(&logits, config.temperature, config.top_k)
    946
                };
    947
    948
                tokens.push(next_token);
    949
    950
                // Check for stop tokens
    951
                if config.stop_tokens.contains(&next_token) {
    952
                    break;
    953
                }
    954
            }
    955
    956
            Ok(tokens)
    957
        }
    958
    959
        /// Create GPU model from GGUF config (M13: Real Model Loading)
    960
        ///
    961
        /// This is a convenience constructor that creates a model with zero-initialized
    962
        /// weights from a config. Use `from_mapped_gguf()` to load actual weights.
    963
        ///
    964
        /// # Arguments
    965
        ///
    966
        /// * `config` - Model configuration
    967
        ///
    968
        /// # Errors
    969
        ///
    970
        /// Returns error if GPU initialization fails
    971
        ///
    972
        /// # Examples
    973
        ///
    974
        /// ```rust,ignore
    975
        /// let config = GpuModelConfig {
    976
        ///     vocab_size: 32000,
    977
        ///     hidden_dim: 4096,
    978
        ///     num_heads: 32,
    979
        ///     num_kv_heads: 32,
    980
        ///     num_layers: 32,
    981
        ///     intermediate_dim: 11008,
    982
        ///     eps: 1e-5,
    983
        ///     rope_theta: 10000.0,
    984
        /// };
    985
        /// let model = GpuModel::from_gguf_config(config)?;
    986
        /// ```
    987
    9
        pub fn from_gguf_config(config: GpuModelConfig) -> Result<Self> {
    988
            // Delegate to new() which handles initialization
    989
    9
            Self::new(config)
    990
    9
        }
    991
    992
        /// Load GPU model from memory-mapped GGUF file (M13: Real Model Loading)
    993
        ///
    994
        /// This is the primary method for loading real GGUF models to GPU.
    995
        /// It dequantizes weights on-the-fly and uploads them to GPU buffers.
    996
        ///
    997
        /// # Arguments
    998
        ///
    999
        /// * `mapped` - Memory-mapped GGUF model
    1000
        ///
    1001
        /// # Errors
    1002
        ///
    1003
        /// Returns error if:
    1004
        /// - Required tensors are missing
    1005
        /// - Tensor shapes don't match expected dimensions
    1006
        /// - GPU initialization or upload fails
    1007
        ///
    1008
        /// # Examples
    1009
        ///
    1010
        /// ```rust,ignore
    1011
        /// let mapped = MappedGGUFModel::from_path("model.gguf")?;
    1012
        /// let model = GpuModel::from_mapped_gguf(&mapped)?;
    1013
        /// let logits = model.forward_gpu_owned(&[1, 2, 3])?;
    1014
        /// ```
    1015
    0
        pub fn from_mapped_gguf(mapped: &crate::gguf::MappedGGUFModel) -> Result<Self> {
    1016
            use crate::gguf::GGUFConfig;
    1017
    1018
            // Extract config from GGUF metadata
    1019
    0
            let gguf_config = GGUFConfig::from_gguf(&mapped.model)?;
    1020
    1021
    0
            let config = GpuModelConfig {
    1022
    0
                vocab_size: gguf_config.vocab_size,
    1023
    0
                hidden_dim: gguf_config.hidden_dim,
    1024
    0
                num_heads: gguf_config.num_heads,
    1025
    0
                num_kv_heads: gguf_config.num_kv_heads, // IMP-088: GQA support
    1026
    0
                num_layers: gguf_config.num_layers,
    1027
    0
                intermediate_dim: gguf_config.intermediate_dim,
    1028
    0
                eps: gguf_config.eps,
    1029
    0
                rope_theta: gguf_config.rope_theta, // Phase 21: RoPE support
    1030
    0
            };
    1031
    1032
    0
            let scheduler = HybridScheduler::new()?;
    1033
    0
            let data = mapped.data();
    1034
    1035
            // Load token embeddings (always dequantized for fast lookup)
    1036
    0
            let embedding_weights = mapped.model.get_tensor_f32("token_embd.weight", data)?;
    1037
    1038
            // Load transformer blocks
    1039
    0
            let mut block_weights = Vec::with_capacity(config.num_layers);
    1040
    0
            for layer_idx in 0..config.num_layers {
    1041
    0
                let prefix = format!("blk.{}", layer_idx);
    1042
    1043
                // Attention norm (small, keep as f32)
    1044
    0
                let attn_norm_weight = mapped
    1045
    0
                    .model
    1046
    0
                    .get_tensor_f32(&format!("{}.attn_norm.weight", prefix), data)?;
    1047
    0
                let attn_norm_bias = mapped
    1048
    0
                    .model
    1049
    0
                    .get_tensor_f32(&format!("{}.attn_norm.bias", prefix), data)
    1050
    0
                    .unwrap_or_else(|_| vec![0.0f32; config.hidden_dim]);
    1051
    1052
                // QKV projection - try fused QKV first (LLaMA), then separate Q/K/V (Qwen)
    1053
    0
                let (qkv_weight, qkv_bias) = if let Ok(fused_qkv) = mapped
    1054
    0
                    .model
    1055
    0
                    .get_tensor_f32(&format!("{}.attn_qkv.weight", prefix), data)
    1056
                {
    1057
                    // Fused QKV (LLaMA-style)
    1058
    0
                    let bias = mapped
    1059
    0
                        .model
    1060
    0
                        .get_tensor_f32(&format!("{}.attn_qkv.bias", prefix), data)
    1061
    0
                        .unwrap_or_else(|_| vec![0.0f32; 3 * config.hidden_dim]);
    1062
    0
                    (fused_qkv, bias)
    1063
                } else {
    1064
                    // Separate Q/K/V (Qwen-style) - concatenate into fused format
    1065
                    // For GQA: Q has num_heads * head_dim, K/V have num_kv_heads * head_dim
    1066
    0
                    let head_dim = config.hidden_dim / config.num_heads;
    1067
    0
                    let kv_dim = config.num_kv_heads * head_dim; // K/V dimension for GQA
    1068
    1069
    0
                    let q_weight = mapped
    1070
    0
                        .model
    1071
    0
                        .get_tensor_f32(&format!("{}.attn_q.weight", prefix), data)?;
    1072
    0
                    let k_weight = mapped
    1073
    0
                        .model
    1074
    0
                        .get_tensor_f32(&format!("{}.attn_k.weight", prefix), data)?;
    1075
    0
                    let v_weight = mapped
    1076
    0
                        .model
    1077
    0
                        .get_tensor_f32(&format!("{}.attn_v.weight", prefix), data)?;
    1078
    1079
                    // Concatenate Q, K, V weights
    1080
    0
                    let mut qkv_weight =
    1081
    0
                        Vec::with_capacity(q_weight.len() + k_weight.len() + v_weight.len());
    1082
    0
                    qkv_weight.extend_from_slice(&q_weight);
    1083
    0
                    qkv_weight.extend_from_slice(&k_weight);
    1084
    0
                    qkv_weight.extend_from_slice(&v_weight);
    1085
    1086
                    // Load biases if available (use correct dimensions for GQA)
    1087
    0
                    let q_bias = mapped
    1088
    0
                        .model
    1089
    0
                        .get_tensor_f32(&format!("{}.attn_q.bias", prefix), data)
    1090
    0
                        .unwrap_or_else(|_| vec![0.0f32; config.hidden_dim]);
    1091
    0
                    let k_bias = mapped
    1092
    0
                        .model
    1093
    0
                        .get_tensor_f32(&format!("{}.attn_k.bias", prefix), data)
    1094
    0
                        .unwrap_or_else(|_| vec![0.0f32; kv_dim]); // GQA: K/V use num_kv_heads
    1095
    0
                    let v_bias = mapped
    1096
    0
                        .model
    1097
    0
                        .get_tensor_f32(&format!("{}.attn_v.bias", prefix), data)
    1098
    0
                        .unwrap_or_else(|_| vec![0.0f32; kv_dim]); // GQA: K/V use num_kv_heads
    1099
    1100
                    // Total bias size: Q (hidden_dim) + K (kv_dim) + V (kv_dim)
    1101
    0
                    let total_bias_dim = config.hidden_dim + 2 * kv_dim;
    1102
    0
                    let mut qkv_bias = Vec::with_capacity(total_bias_dim);
    1103
    0
                    qkv_bias.extend_from_slice(&q_bias);
    1104
    0
                    qkv_bias.extend_from_slice(&k_bias);
    1105
    0
                    qkv_bias.extend_from_slice(&v_bias);
    1106
    1107
    0
                    (qkv_weight, qkv_bias)
    1108
                };
    1109
    1110
                // Output projection
    1111
    0
                let out_weight = mapped
    1112
    0
                    .model
    1113
    0
                    .get_tensor_f32(&format!("{}.attn_output.weight", prefix), data)?;
    1114
    0
                let out_bias = mapped
    1115
    0
                    .model
    1116
    0
                    .get_tensor_f32(&format!("{}.attn_output.bias", prefix), data)
    1117
    0
                    .unwrap_or_else(|_| vec![0.0f32; config.hidden_dim]);
    1118
    1119
                // FFN norm
    1120
    0
                let ffn_norm_weight = mapped
    1121
    0
                    .model
    1122
    0
                    .get_tensor_f32(&format!("{}.ffn_norm.weight", prefix), data)
    1123
    0
                    .unwrap_or_else(|_| vec![1.0f32; config.hidden_dim]);
    1124
    0
                let ffn_norm_bias = mapped
    1125
    0
                    .model
    1126
    0
                    .get_tensor_f32(&format!("{}.ffn_norm.bias", prefix), data)
    1127
    0
                    .unwrap_or_else(|_| vec![0.0f32; config.hidden_dim]);
    1128
    1129
                // FFN projections
    1130
    0
                let ffn_fc1_weight = mapped
    1131
    0
                    .model
    1132
    0
                    .get_tensor_f32(&format!("{}.ffn_up.weight", prefix), data)?;
    1133
    0
                let ffn_fc1_bias = mapped
    1134
    0
                    .model
    1135
    0
                    .get_tensor_f32(&format!("{}.ffn_up.bias", prefix), data)
    1136
    0
                    .unwrap_or_else(|_| vec![0.0f32; config.intermediate_dim]);
    1137
    1138
    0
                let ffn_fc2_weight = mapped
    1139
    0
                    .model
    1140
    0
                    .get_tensor_f32(&format!("{}.ffn_down.weight", prefix), data)?;
    1141
    0
                let ffn_fc2_bias = mapped
    1142
    0
                    .model
    1143
    0
                    .get_tensor_f32(&format!("{}.ffn_down.bias", prefix), data)
    1144
    0
                    .unwrap_or_else(|_| vec![0.0f32; config.hidden_dim]);
    1145
    1146
                // Try to load gate weight for SwiGLU (optional)
    1147
    0
                let ffn_gate_weight = mapped
    1148
    0
                    .model
    1149
    0
                    .get_tensor_f32(&format!("{}.ffn_gate.weight", prefix), data)
    1150
    0
                    .ok();
    1151
    1152
    0
                block_weights.push(BlockWeights {
    1153
    0
                    attn_norm_weight,
    1154
    0
                    attn_norm_bias,
    1155
    0
                    qkv_weight,
    1156
    0
                    qkv_bias,
    1157
    0
                    out_weight,
    1158
    0
                    out_bias,
    1159
    0
                    ffn_norm_weight,
    1160
    0
                    ffn_norm_bias,
    1161
    0
                    ffn_fc1_weight,
    1162
    0
                    ffn_fc1_bias,
    1163
    0
                    ffn_fc2_weight,
    1164
    0
                    ffn_fc2_bias,
    1165
    0
                    ffn_gate_weight,
    1166
    0
                });
    1167
            }
    1168
    1169
            // Final layer norm
    1170
    0
            let final_norm_weight = mapped.model.get_tensor_f32("output_norm.weight", data)?;
    1171
    0
            let final_norm_bias = mapped
    1172
    0
                .model
    1173
    0
                .get_tensor_f32("output_norm.bias", data)
    1174
    0
                .unwrap_or_else(|_| vec![0.0f32; config.hidden_dim]);
    1175
    1176
            // LM head
    1177
    0
            let lm_head_weight = mapped.model.get_tensor_f32("output.weight", data)?;
    1178
    0
            let lm_head_bias = mapped
    1179
    0
                .model
    1180
    0
                .get_tensor_f32("output.bias", data)
    1181
    0
                .unwrap_or_else(|_| vec![0.0f32; config.vocab_size]);
    1182
    1183
            // Pre-compute transposed LM head for fast CPU inference
    1184
    0
            let lm_head_weight_t =
    1185
    0
                Self::transpose_weights(&lm_head_weight, config.hidden_dim, config.vocab_size);
    1186
    1187
    0
            Ok(Self {
    1188
    0
                embedding_weights,
    1189
    0
                block_weights,
    1190
    0
                final_norm_weight,
    1191
    0
                final_norm_bias,
    1192
    0
                lm_head_weight,
    1193
    0
                lm_head_weight_t,
    1194
    0
                lm_head_bias,
    1195
    0
                scheduler,
    1196
    0
                #[cfg(feature = "cuda")]
    1197
    0
                cuda_scheduler: None,
    1198
    0
                config,
    1199
    0
                attention_buffers: None,
    1200
    0
            })
    1201
    0
        }
    1202
    1203
        /// Create GpuModel from pre-extracted APR weights (PMAT-106)
    1204
        ///
    1205
        /// This constructor is used by `AprToGpuAdapter` to create a `GpuModel`
    1206
        /// from dequantized APR weights.
    1207
        ///
    1208
        /// # Arguments
    1209
        ///
    1210
        /// * `config` - GPU model configuration
    1211
        /// * `embedding_weights` - Token embedding weights
    1212
        /// * `block_weights` - Transformer block weights
    1213
        /// * `final_norm_weight` - Final layer norm weight
    1214
        /// * `final_norm_bias` - Final layer norm bias
    1215
        /// * `lm_head_weight` - LM head weight (row-major)
    1216
        /// * `lm_head_weight_t` - LM head weight transposed (for fast CPU inference)
    1217
        /// * `lm_head_bias` - LM head bias
    1218
        ///
    1219
        /// # Errors
    1220
        ///
    1221
        /// Returns error if GPU scheduler initialization fails
    1222
        #[allow(clippy::too_many_arguments)]
    1223
    0
        pub fn from_apr_weights(
    1224
    0
            config: GpuModelConfig,
    1225
    0
            embedding_weights: Vec<f32>,
    1226
    0
            block_weights: Vec<BlockWeights>,
    1227
    0
            final_norm_weight: Vec<f32>,
    1228
    0
            final_norm_bias: Vec<f32>,
    1229
    0
            lm_head_weight: Vec<f32>,
    1230
    0
            lm_head_weight_t: Vec<f32>,
    1231
    0
            lm_head_bias: Vec<f32>,
    1232
    0
        ) -> Result<Self> {
    1233
    0
            let scheduler = HybridScheduler::new()?;
    1234
    1235
            // Phase 21: Initialize CudaScheduler for GPU-accelerated matmul
    1236
            #[cfg(feature = "cuda")]
    1237
            let cuda_scheduler = match CudaScheduler::new() {
    1238
                Ok(cs) => {
    1239
                    eprintln!("[PHASE21] CudaScheduler initialized for APR model");
    1240
                    Some(cs)
    1241
                }
    1242
                Err(e) => {
    1243
                    eprintln!("[PHASE21] CudaScheduler init failed (using HybridScheduler fallback): {}", e);
    1244
                    None
    1245
                }
    1246
            };
    1247
    1248
    0
            Ok(Self {
    1249
    0
                embedding_weights,
    1250
    0
                block_weights,
    1251
    0
                final_norm_weight,
    1252
    0
                final_norm_bias,
    1253
    0
                lm_head_weight,
    1254
    0
                lm_head_weight_t,
    1255
    0
                lm_head_bias,
    1256
    0
                scheduler,
    1257
    0
                #[cfg(feature = "cuda")]
    1258
    0
                cuda_scheduler,
    1259
    0
                config,
    1260
    0
                attention_buffers: None,
    1261
    0
            })
    1262
    0
        }
    1263
    1264
        /// Get model configuration (M13: Real Model Loading)
    1265
        #[must_use]
    1266
    4
        pub fn config(&self) -> &GpuModelConfig {
    1267
    4
            &self.config
    1268
    4
        }
    1269
    1270
        // ============================================================================
    1271
        // Phase 8: Optimized Incremental Decoding (M17)
    1272
        // ============================================================================
    1273
    1274
        /// Create GPU model with pre-allocated attention buffers (M17)
    1275
        ///
    1276
        /// Allocates reusable buffers for incremental decoding, eliminating
    1277
        /// per-token memory allocation overhead.
    1278
        ///
    1279
        /// # Arguments
    1280
        ///
    1281
        /// * `config` - Model configuration
    1282
        /// * `max_seq_len` - Maximum sequence length to support
    1283
        ///
    1284
        /// # Errors
    1285
        ///
    1286
        /// Returns error if GPU initialization fails
    1287
    5
        pub fn with_attention_buffers(config: GpuModelConfig, max_seq_len: usize) -> Result<Self> {
    1288
    5
            let buffers = AttentionBuffers::new(&config, max_seq_len);
    1289
    5
            let mut model = Self::new(config)?0;
    1290
    5
            model.attention_buffers = Some(buffers);
    1291
    5
            Ok(model)
    1292
    5
        }
    1293
    1294
        /// Check if model has pre-allocated attention buffers (M17)
    1295
        #[must_use]
    1296
    1
        pub fn has_attention_buffers(&self) -> bool {
    1297
    1
            self.attention_buffers.is_some()
    1298
    1
        }
    1299
    1300
        /// Optimized text generation using pre-allocated buffers (M17)
    1301
        ///
    1302
        /// Uses the optimized incremental forward pass with pre-allocated buffers
    1303
        /// and batched multi-head attention for better performance.
    1304
        ///
    1305
        /// # Arguments
    1306
        ///
    1307
        /// * `prompt` - Initial token IDs
    1308
        /// * `config` - Generation configuration
    1309
        ///
    1310
        /// # Errors
    1311
        ///
    1312
        /// Returns error if generation fails
    1313
    32
        pub fn generate_optimized(
    1314
    32
            &mut self,
    1315
    32
            prompt: &[usize],
    1316
    32
            config: &GpuGenerateConfig,
    1317
    32
        ) -> Result<Vec<usize>> {
    1318
    32
            if prompt.is_empty() {
    1319
    1
                return Err(RealizarError::InvalidShape {
    1320
    1
                    reason: "Prompt cannot be empty".to_string(),
    1321
    1
                });
    1322
    31
            }
    1323
    1324
            // Initialize KV cache
    1325
            // IMP-093: For GQA, use num_kv_heads since K/V have fewer heads than Q
    1326
    31
            let head_dim = self.config.hidden_dim / self.config.num_heads;
    1327
    31
            let max_seq_len = self
    1328
    31
                .attention_buffers
    1329
    31
                .as_ref()
    1330
    31
                .map_or(512, |b| b.max_seq_len);
    1331
    31
            let mut kv_cache = StreamingKVCache::new(
    1332
    31
                self.config.num_layers,
    1333
    31
                max_seq_len,
    1334
    31
                self.config.num_kv_heads, // GQA: K/V have fewer heads
    1335
    31
                head_dim,
    1336
            );
    1337
    1338
    31
            let mut tokens = prompt.to_vec();
    1339
    1340
            // Process prompt with cache - returns logits for final position only [vocab_size]
    1341
    31
            let logits = self.forward_gpu_with_cache(prompt, &mut kv_cache)?0;
    1342
    1343
            // Sample first token (logits is already for last position only)
    1344
    31
            let mut next_token = if config.temperature == 0.0 || config.top_k == 11 {
    1345
    30
                Self::argmax(&logits)
    1346
            } else {
    1347
    1
                Self::sample_topk_generate(&logits, config.temperature, config.top_k)
    1348
            };
    1349
    1350
    31
            if config.stop_tokens.contains(&next_token) {
    1351
    2
                return Ok(tokens);
    1352
    29
            }
    1353
    1354
    29
            tokens.push(next_token);
    1355
    1356
            // Generate remaining tokens using optimized incremental forward
    1357
    29
            for _ in 1..config.max_tokens {
    1358
    472
                let logits = self.forward_gpu_incremental_optimized(next_token, &mut kv_cache)?0;
    1359
    1360
    472
                next_token = if config.temperature == 0.0 || config.top_k == 12 {
    1361
    470
                    Self::argmax(&logits)
    1362
                } else {
    1363
    2
                    Self::sample_topk_generate(&logits, config.temperature, config.top_k)
    1364
                };
    1365
    1366
    472
                if config.stop_tokens.contains(&next_token) {
    1367
    0
                    break;
    1368
    472
                }
    1369
    1370
    472
                tokens.push(next_token);
    1371
            }
    1372
    1373
    29
            Ok(tokens)
    1374
    32
        }
    1375
    1376
        /// Optimized incremental forward pass using pre-allocated buffers (M17)
    1377
        ///
    1378
        /// Single-token forward pass optimized by:
    1379
        /// - Reusing pre-allocated attention buffers
    1380
        /// - Direct KV cache access without copying
    1381
        /// - Batched multi-head attention computation
    1382
        ///
    1383
        /// # Arguments
    1384
        ///
    1385
        /// * `token_id` - Single token to process
    1386
        /// * `kv_cache` - Mutable reference to KV cache
    1387
        ///
    1388
        /// # Errors
    1389
        ///
    1390
        /// Returns error if forward pass fails
    1391
    524
        pub fn forward_gpu_incremental_optimized(
    1392
    524
            &mut self,
    1393
    524
            token_id: usize,
    1394
    524
            kv_cache: &mut StreamingKVCache,
    1395
    524
        ) -> Result<Vec<f32>> {
    1396
    524
            if token_id >= self.config.vocab_size {
    1397
    0
                return Err(RealizarError::InvalidShape {
    1398
    0
                    reason: format!(
    1399
    0
                        "Token ID {} out of bounds (vocab_size={})",
    1400
    0
                        token_id, self.config.vocab_size
    1401
    0
                    ),
    1402
    0
                });
    1403
    524
            }
    1404
    1405
    524
            let hidden_dim = self.config.hidden_dim;
    1406
    1407
            // Get embedding for single token
    1408
    524
            let offset = token_id * hidden_dim;
    1409
    524
            let mut hidden: Vec<f32> = self.embedding_weights[offset..offset + hidden_dim].to_vec();
    1410
    1411
            // Process through all blocks with optimized attention
    1412
    1.55k
            for block_idx in 0..self.block_weights524.len524() {
    1413
    1.55k
                hidden = self.forward_block_incremental_optimized(&hidden, block_idx, kv_cache)?0;
    1414
            }
    1415
    1416
            // Final layer norm
    1417
    524
            hidden = self.layer_norm(&hidden, &self.final_norm_weight, &self.final_norm_bias);
    1418
    1419
            // LM head projection (single token)
    1420
            // IMP-090, IMP-096: Use CPU fallback with SIMD for large vocab
    1421
    524
            let lm_head_elements = hidden_dim * self.config.vocab_size;
    1422
    524
            let output = if exceeds_gpu_buffer_limit(lm_head_elements) {
    1423
                // IMP-096: CPU path with transposed weights + SIMD + fused bias
    1424
    0
                cpu_matmul_transposed_simd(
    1425
    0
                    &hidden,
    1426
    0
                    &self.lm_head_weight_t,
    1427
    0
                    &self.lm_head_bias,
    1428
    0
                    hidden_dim,
    1429
    0
                    self.config.vocab_size,
    1430
                )
    1431
            } else {
    1432
                // IMP-1006: Use do_matmul to route to CudaScheduler when available
    1433
    524
                let lm_weight = self.lm_head_weight.clone();
    1434
    524
                let vocab_size = self.config.vocab_size;
    1435
    524
                let logits = self.do_matmul(&hidden, &lm_weight, 1, hidden_dim, vocab_size)?0;
    1436
                // Add bias
    1437
    524
                logits
    1438
    524
                    .into_iter()
    1439
    524
                    .zip(self.lm_head_bias.iter())
    1440
    134k
                    .map524(|(l, &b)| l + b)
    1441
    524
                    .collect()
    1442
            };
    1443
    1444
    524
            Ok(output)
    1445
    524
        }
    1446
    1447
        /// Optimized block forward with batched multi-head attention (M17, IMP-092)
    1448
        ///
    1449
        /// IMP-092: Eliminated weight cloning (~130MB per layer) by using explicit
    1450
        /// field borrowing. Previous version cloned 3.7GB per token across 28 layers.
    1451
    1.55k
        pub fn forward_block_incremental_optimized(
    1452
    1.55k
            &mut self,
    1453
    1.55k
            input: &[f32],
    1454
    1.55k
            block_idx: usize,
    1455
    1.55k
            kv_cache: &mut StreamingKVCache,
    1456
    1.55k
        ) -> Result<Vec<f32>> {
    1457
            // Extract config values (Copy types, no borrow conflict)
    1458
    1.55k
            let hidden_dim = self.config.hidden_dim;
    1459
    1.55k
            let num_heads = self.config.num_heads;
    1460
    1.55k
            let head_dim = self.config.head_dim();
    1461
    1.55k
            let kv_dim = self.config.kv_dim();
    1462
    1.55k
            let qkv_dim = self.config.qkv_dim();
    1463
    1.55k
            let intermediate_dim = self.config.intermediate_dim;
    1464
    1.55k
            let eps = self.config.eps;
    1465
    1.55k
            let num_kv_heads = self.config.num_kv_heads;
    1466
    1467
            // IMP-092: Use REFERENCES instead of cloning 130MB of weights per layer
    1468
            // Pre-attention layer norm (static function avoids &self borrow)
    1469
    1.55k
            let normed = Self::layer_norm_static(
    1470
    1.55k
                input,
    1471
    1.55k
                &self.block_weights[block_idx].attn_norm_weight,
    1472
    1.55k
                &self.block_weights[block_idx].attn_norm_bias,
    1473
    1.55k
                hidden_dim,
    1474
    1.55k
                eps,
    1475
            );
    1476
    1477
            // QKV projection for single token [1, hidden_dim] @ [hidden_dim, qkv_dim]
    1478
            // For GQA: qkv_dim = hidden_dim + 2*kv_dim (K/V have fewer heads)
    1479
            // IMP-1006: Use do_matmul to route to CudaScheduler when available
    1480
    1.55k
            let qkv_weight = self.block_weights[block_idx].qkv_weight.clone();
    1481
    1.55k
            let mut qkv = self.do_matmul(&normed, &qkv_weight, 1, hidden_dim, qkv_dim)?0;
    1482
    1483
            // Get current position BEFORE caching (Phase 21)
    1484
    1.55k
            let (cached_k_ref, _) = kv_cache.get_valid(block_idx);
    1485
    1.55k
            let current_pos = cached_k_ref.len() / kv_dim;
    1486
    1487
            // Phase 21: Apply RoPE to Q and K BEFORE caching
    1488
            // Without RoPE, attention has no position information and produces garbage
    1489
    1.55k
            let rope_theta = self.config.rope_theta;
    1490
    1.55k
            Self::apply_rope_inline(&mut qkv[0..hidden_dim], num_heads, head_dim, rope_theta, current_pos);
    1491
    1.55k
            Self::apply_rope_inline(&mut qkv[hidden_dim..hidden_dim + kv_dim], num_kv_heads, head_dim, rope_theta, current_pos);
    1492
    1493
            // Split QKV (GQA: K/V have kv_dim, not hidden_dim) - after RoPE
    1494
    1.55k
            let q = qkv[0..hidden_dim].to_vec();
    1495
    1.55k
            let k_new = qkv[hidden_dim..hidden_dim + kv_dim].to_vec();
    1496
    1.55k
            let v_new = qkv[hidden_dim + kv_dim..].to_vec();
    1497
    1498
            // Get cached K/V and clone to avoid borrow issues with kv_cache
    1499
    1.55k
            let (cached_k, cached_v) = kv_cache.get_valid(block_idx);
    1500
    1.55k
            let keys_cached = cached_k.to_vec();
    1501
    1.55k
            let vals_cached = cached_v.to_vec();
    1502
    1503
            // Append new K/V (with RoPE applied) to cache
    1504
    1.55k
            kv_cache.append(block_idx, &k_new, &v_new);
    1505
    1506
            // Build full K/V (cached + new)
    1507
            // GQA: K/V have kv_dim per position, not hidden_dim
    1508
    1.55k
            let kv_len = keys_cached.len() / kv_dim + 1;
    1509
    1.55k
            let mut full_k = keys_cached;
    1510
    1.55k
            full_k.extend_from_slice(&k_new);
    1511
    1.55k
            let mut full_v = vals_cached;
    1512
    1.55k
            full_v.extend_from_slice(&v_new);
    1513
    1514
            // GQA attention (IMP-089): static method to avoid borrow conflicts
    1515
    1.55k
            let attn_output = Self::gqa_multihead_attention(
    1516
    1.55k
                &q,
    1517
    1.55k
                &full_k,
    1518
    1.55k
                &full_v,
    1519
    1.55k
                kv_len,
    1520
    1.55k
                num_heads,
    1521
    1.55k
                num_kv_heads,
    1522
    1.55k
                head_dim,
    1523
            );
    1524
    1525
            // Output projection
    1526
            // IMP-1006: Use do_matmul to route to CudaScheduler when available
    1527
    1.55k
            let out_weight = self.block_weights[block_idx].out_weight.clone();
    1528
    1.55k
            let attn_proj = self.do_matmul(&attn_output, &out_weight, 1, hidden_dim, hidden_dim)?0;
    1529
    1530
            // Add residual and bias
    1531
    1.55k
            let out_bias = &self.block_weights[block_idx].out_bias;
    1532
    1.55k
            let mut post_attn: Vec<f32> = input
    1533
    1.55k
                .iter()
    1534
    1.55k
                .zip(attn_proj.iter())
    1535
    1.55k
                .zip(out_bias.iter())
    1536
    163k
                .map1.55k(|((&i, &a), &b)| i + a + b)
    1537
    1.55k
                .collect();
    1538
    1539
            // FFN with layer norm (static function)
    1540
    1.55k
            let ffn_normed = Self::layer_norm_static(
    1541
    1.55k
                &post_attn,
    1542
    1.55k
                &self.block_weights[block_idx].ffn_norm_weight,
    1543
    1.55k
                &self.block_weights[block_idx].ffn_norm_bias,
    1544
    1.55k
                hidden_dim,
    1545
    1.55k
                eps,
    1546
            );
    1547
    1548
            // FFN: SwiGLU when gate weight exists, otherwise GELU
    1549
            // IMP-1006: Use do_matmul to route to CudaScheduler when available
    1550
    1.55k
            let fc1_activated: Vec<f32> = if let Some(ref gate_weight0) = self.block_weights[block_idx].ffn_gate_weight {
    1551
                // SwiGLU: silu(gate(x)) * up(x)
    1552
    0
                let fc1_weight = self.block_weights[block_idx].ffn_fc1_weight.clone();
    1553
    0
                let gate_weight = gate_weight.clone();
    1554
    1555
    0
                let up_out = self.do_matmul(&ffn_normed, &fc1_weight, 1, hidden_dim, intermediate_dim)?;
    1556
    0
                let gate_out = self.do_matmul(&ffn_normed, &gate_weight, 1, hidden_dim, intermediate_dim)?;
    1557
    1558
                // SwiGLU: silu(gate) * up
    1559
    0
                up_out
    1560
    0
                    .iter()
    1561
    0
                    .zip(gate_out.iter())
    1562
    0
                    .map(|(&u, &g)| {
    1563
    0
                        let silu_g = g / (1.0 + (-g).exp());
    1564
    0
                        silu_g * u
    1565
    0
                    })
    1566
    0
                    .collect()
    1567
            } else {
    1568
                // Standard GELU FFN
    1569
    1.55k
                let fc1_weight = self.block_weights[block_idx].ffn_fc1_weight.clone();
    1570
    1.55k
                let fc1_out = self.do_matmul(&ffn_normed, &fc1_weight, 1, hidden_dim, intermediate_dim)?0;
    1571
    1572
    1.55k
                let ffn_fc1_bias = &self.block_weights[block_idx].ffn_fc1_bias;
    1573
    1.55k
                fc1_out
    1574
    1.55k
                    .iter()
    1575
    1.55k
                    .zip(ffn_fc1_bias.iter())
    1576
    326k
                    .map1.55k(|(&x, &b)| {
    1577
    326k
                        let x_b = x + b;
    1578
    326k
                        x_b * 0.5 + x_b * 0.5 * (0.797_884_6 * (x_b + 0.044_715 * x_b.powi(3))).tanh()
    1579
    326k
                    })
    1580
    1.55k
                    .collect()
    1581
            };
    1582
    1583
            // FFN FC2 (down projection)
    1584
            // IMP-1006: Use do_matmul to route to CudaScheduler when available
    1585
    1.55k
            let fc2_weight = self.block_weights[block_idx].ffn_fc2_weight.clone();
    1586
    1.55k
            let fc2_out =
    1587
    1.55k
                self.do_matmul(&fc1_activated, &fc2_weight, 1, intermediate_dim, hidden_dim)?0;
    1588
    1589
            // Add residual and bias
    1590
    1.55k
            let ffn_fc2_bias = &self.block_weights[block_idx].ffn_fc2_bias;
    1591
    163k
            for i in 0..hidden_dim1.55k {
    1592
    163k
                post_attn[i] += fc2_out[i] + ffn_fc2_bias[i];
    1593
    163k
            }
    1594
    1595
    1.55k
            Ok(post_attn)
    1596
    1.55k
        }
    1597
    1598
        /// Apply Rotary Position Embedding (RoPE) inline (Phase 21)
    1599
        ///
    1600
        /// RoPE encodes position information by rotating pairs of elements
    1601
        /// with position-dependent angles. This is CRITICAL for transformer attention.
    1602
        ///
    1603
        /// # Arguments
    1604
        /// * `x` - Mutable slice of Q or K vectors for a single position [num_heads * head_dim]
    1605
        /// * `num_heads` - Number of attention heads
    1606
        /// * `head_dim` - Dimension per head
    1607
        /// * `rope_theta` - Base frequency (typically 10000.0)
    1608
        /// * `position` - Token position for RoPE encoding
    1609
    3.10k
        fn apply_rope_inline(
    1610
    3.10k
            x: &mut [f32],
    1611
    3.10k
            num_heads: usize,
    1612
    3.10k
            head_dim: usize,
    1613
    3.10k
            rope_theta: f32,
    1614
    3.10k
            position: usize,
    1615
    3.10k
        ) {
    1616
    3.10k
            let half_dim = head_dim / 2;
    1617
    3.10k
            let head_dim_f32 = head_dim as f32;
    1618
    3.10k
            let pos_f32 = position as f32;
    1619
    1620
    20.4k
            for h in 0..num_heads3.10k {
    1621
    20.4k
                let head_start = h * head_dim;
    1622
    20.4k
                let idx2_start = head_start + half_dim;
    1623
    1624
    163k
                for i in 0..half_dim20.4k {
    1625
    163k
                    let freq = 1.0 / rope_theta.powf(2.0 * i as f32 / head_dim_f32);
    1626
    163k
                    let angle = pos_f32 * freq;
    1627
    163k
                    let (sin_val, cos_val) = angle.sin_cos();
    1628
    163k
    1629
    163k
                    let x1 = x[head_start + i];
    1630
    163k
                    let x2 = x[idx2_start + i];
    1631
    163k
    1632
    163k
                    // Apply rotation: [cos -sin; sin cos] * [x1; x2]
    1633
    163k
                    x[head_start + i] = x1 * cos_val - x2 * sin_val;
    1634
    163k
                    x[idx2_start + i] = x1 * sin_val + x2 * cos_val;
    1635
    163k
                }
    1636
            }
    1637
    3.10k
        }
    1638
    1639
        /// GQA multi-head attention (IMP-089, IMP-092, IMP-094)
    1640
        ///
    1641
        /// Grouped Query Attention where K/V have fewer heads than Q.
    1642
        /// Each KV head serves (num_heads / num_kv_heads) Q heads.
    1643
        ///
    1644
        /// IMP-094: Uses trueno SIMD-accelerated dot product and softmax
    1645
        /// for ~10x speedup over scalar implementation.
    1646
        ///
    1647
        /// Static method to avoid borrow conflicts with scheduler and weights.
    1648
    1.55k
        fn gqa_multihead_attention(
    1649
    1.55k
            q: &[f32], // Q: [num_heads * head_dim]
    1650
    1.55k
            k: &[f32], // K: [kv_len * num_kv_heads * head_dim]
    1651
    1.55k
            v: &[f32], // V: [kv_len * num_kv_heads * head_dim]
    1652
    1.55k
            kv_len: usize,
    1653
    1.55k
            num_heads: usize,    // Number of Q heads
    1654
    1.55k
            num_kv_heads: usize, // Number of K/V heads (for GQA, < num_heads)
    1655
    1.55k
            head_dim: usize,
    1656
    1.55k
        ) -> Vec<f32> {
    1657
            use trueno::Vector;
    1658
    1659
    1.55k
            let hidden_dim = num_heads * head_dim;
    1660
    1.55k
            let kv_dim = num_kv_heads * head_dim;
    1661
    1.55k
            let scale = 1.0 / (head_dim as f32).sqrt();
    1662
    1663
            // Number of Q heads per KV head
    1664
    1.55k
            let heads_per_kv = num_heads / num_kv_heads;
    1665
    1666
    1.55k
            let mut output = vec![0.0; hidden_dim];
    1667
    1668
            // Compute attention for all Q heads
    1669
    10.2k
            for h in 0..num_heads1.55k {
    1670
    10.2k
                let q_head = &q[h * head_dim..(h + 1) * head_dim];
    1671
                // IMP-094: Create trueno vector for SIMD dot product
    1672
    10.2k
                let q_vec = Vector::from_slice(q_head);
    1673
    1674
                // Map Q head to KV head (GQA: multiple Q heads share one KV head)
    1675
    10.2k
                let kv_head = h / heads_per_kv;
    1676
    1677
                // Compute attention scores for this head using SIMD dot product
    1678
    10.2k
                let mut scores = Vec::with_capacity(kv_len);
    1679
    225k
                for pos in 0..kv_len10.2k {
    1680
    225k
                    // K offset: pos * kv_dim + kv_head * head_dim
    1681
    225k
                    let k_offset = pos * kv_dim + kv_head * head_dim;
    1682
    225k
                    let cached_key = &k[k_offset..k_offset + head_dim];
    1683
    225k
    1684
    225k
                    // IMP-094: SIMD dot product via trueno
    1685
    225k
                    let k_vec = Vector::from_slice(cached_key);
    1686
    225k
                    let score = q_vec.dot(&k_vec).unwrap_or(0.0) * scale;
    1687
    225k
                    scores.push(score);
    1688
    225k
                }
    1689
    1690
                // IMP-094: SIMD softmax via trueno
    1691
    10.2k
                let scores_vec = Vector::from_slice(&scores);
    1692
    10.2k
                let attn_weights: Vec<f32> = scores_vec.softmax().map_or_else(
    1693
    0
                    |_| {
    1694
                        // Fallback to scalar softmax
    1695
    0
                        let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    1696
    0
                        let exp_scores: Vec<f32> =
    1697
    0
                            scores.iter().map(|&s| (s - max_score).exp()).collect();
    1698
    0
                        let sum_exp: f32 = exp_scores.iter().sum();
    1699
    0
                        exp_scores.iter().map(|&e| e / sum_exp).collect()
    1700
    0
                    },
    1701
    10.2k
                    |v| v.as_slice().to_vec(),
    1702
                );
    1703
    1704
                // Weighted sum of values (still scalar - SIMD benefit is marginal for small head_dim)
    1705
    225k
                for (pos, &weight) in attn_weights.iter()10.2k.enumerate10.2k() {
    1706
                    // V offset: pos * kv_dim + kv_head * head_dim
    1707
    225k
                    let v_offset = pos * kv_dim + kv_head * head_dim;
    1708
    225k
                    let v_head = &v[v_offset..v_offset + head_dim];
    1709
    1710
    3.61M
                    for d in 0..head_dim225k {
    1711
    3.61M
                        output[h * head_dim + d] += weight * v_head[d];
    1712
    3.61M
                    }
    1713
                }
    1714
            }
    1715
    1716
    1.55k
            output
    1717
    1.55k
        }
    1718
    1719
        // ============================================================================
    1720
        // Phase 9: Fused Kernels & Vectorization (M18)
    1721
        // ============================================================================
    1722
    1723
        /// Check if model has fused QKV projection (M18 - IMP-037)
    1724
        ///
    1725
        /// Fused QKV uses a single matmul instead of three separate projections.
    1726
        /// This is always true for GpuModel as QKV weights are stored combined.
    1727
        #[must_use]
    1728
    1
        pub fn has_fused_qkv(&self) -> bool {
    1729
            // QKV weights are stored as [hidden_dim, 3*hidden_dim] for fused projection
    1730
    1
            !self.block_weights.is_empty()
    1731
    1
                && self.block_weights[0].qkv_weight.len()
    1732
    1
                    == self.config.hidden_dim * 3 * self.config.hidden_dim
    1733
    1
        }
    1734
    1735
        /// Fused QKV projection (M18 - IMP-037)
    1736
        ///
    1737
        /// Performs Q, K, V projection in a single matmul operation.
    1738
        ///
    1739
        /// # Arguments
    1740
        ///
    1741
        /// * `input` - Input tensor [hidden_dim]
    1742
        ///
    1743
        /// # Returns
    1744
        ///
    1745
        /// Tuple of (Q, K, V) tensors, each [hidden_dim]
    1746
        ///
    1747
        /// # Errors
    1748
        ///
    1749
        /// Returns error if matmul fails
    1750
    1
        pub fn fused_qkv_projection(
    1751
    1
            &mut self,
    1752
    1
            input: &[f32],
    1753
    1
        ) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>)> {
    1754
    1
            let hidden_dim = self.config.hidden_dim;
    1755
    1
            let kv_dim = self.config.kv_dim();
    1756
    1
            let qkv_dim = self.config.qkv_dim();
    1757
    1758
            // Use first block's QKV weights for projection
    1759
    1
            let qkv_weight = &self.block_weights[0].qkv_weight;
    1760
    1761
            // Single matmul: [1, hidden_dim] @ [hidden_dim, qkv_dim] -> [1, qkv_dim]
    1762
            // For GQA: qkv_dim = hidden_dim + 2*kv_dim
    1763
    1
            let qkv = self
    1764
    1
                .scheduler
    1765
    1
                .matmul(input, qkv_weight, 1, hidden_dim, qkv_dim)?0;
    1766
    1767
            // Split into Q, K, V (GQA: K/V have kv_dim, not hidden_dim)
    1768
    1
            let q = qkv[0..hidden_dim].to_vec();
    1769
    1
            let k = qkv[hidden_dim..hidden_dim + kv_dim].to_vec();
    1770
    1
            let v = qkv[hidden_dim + kv_dim..].to_vec();
    1771
    1772
    1
            Ok((q, k, v))
    1773
    1
        }
    1774
    1775
        /// Generation with fused QKV projection (M18 - IMP-037)
    1776
        ///
    1777
        /// Uses fused QKV projection for improved performance.
    1778
        ///
    1779
        /// # Errors
    1780
        ///
    1781
        /// Returns error if generation fails due to invalid input or model state.
    1782
    1
        pub fn generate_with_fused_qkv(
    1783
    1
            &mut self,
    1784
    1
            prompt: &[usize],
    1785
    1
            config: &GpuGenerateConfig,
    1786
    1
        ) -> Result<Vec<usize>> {
    1787
            // Fused QKV is already used in generate_optimized via forward_block_incremental_optimized
    1788
            // This method provides explicit API for benchmarking
    1789
    1
            self.generate_optimized(prompt, config)
    1790
    1
        }
    1791
    1792
        /// Check if model has fused attention projection (M18 - IMP-039)
    1793
        #[must_use]
    1794
    1
        pub fn has_fused_attn_proj(&self) -> bool {
    1795
            // Attention output projection is stored in block_weights
    1796
    1
            !self.block_weights.is_empty()
    1797
    1
                && self.block_weights[0].out_weight.len()
    1798
    1
                    == self.config.hidden_dim * self.config.hidden_dim
    1799
    1
        }
    1800
    1801
        /// Forward pass with fused attention projection (M18 - IMP-039)
    1802
        ///
    1803
        /// Uses fused attention output projection for improved performance.
    1804
        ///
    1805
        /// # Errors
    1806
        ///
    1807
        /// Returns error if forward pass fails due to invalid token or cache state.
    1808
    10
        pub fn forward_with_fused_attn_proj(
    1809
    10
            &mut self,
    1810
    10
            token_id: usize,
    1811
    10
            kv_cache: &mut StreamingKVCache,
    1812
    10
        ) -> Result<Vec<f32>> {
    1813
            // Fused attention projection is already used in forward_gpu_incremental_optimized
    1814
            // This method provides explicit API for benchmarking
    1815
    10
            self.forward_gpu_incremental_optimized(token_id, kv_cache)
    1816
    10
        }
    1817
    1818
        /// Check if model has fused output residual capability (M19 - IMP-042)
    1819
        #[must_use]
    1820
    1
        pub fn has_fused_output_residual(&self) -> bool {
    1821
            // Fused output residual requires attention buffers and block weights
    1822
    1
            self.attention_buffers.is_some() && !self.block_weights.is_empty()
    1823
    1
        }
    1824
    1825
        /// Forward pass with fused output projection + residual (M19 - IMP-042)
    1826
        ///
    1827
        /// Combines the output projection matrix multiplication with residual
    1828
        /// connection in a single fused operation.
    1829
        ///
    1830
        /// # Errors
    1831
        ///
    1832
        /// Returns error if forward pass fails due to invalid token or cache state.
    1833
    11
        pub fn forward_with_fused_output_residual(
    1834
    11
            &mut self,
    1835
    11
            token_id: usize,
    1836
    11
            kv_cache: &mut StreamingKVCache,
    1837
    11
        ) -> Result<Vec<f32>> {
    1838
            // Currently uses the optimized forward path
    1839
            // The fused operation is implemented in forward_block_incremental_optimized
    1840
            // This method provides explicit API for benchmarking
    1841
    11
            self.forward_gpu_incremental_optimized(token_id, kv_cache)
    1842
    11
        }
    1843
    1844
        /// Forward pass taking ownership of token_ids (convenience wrapper)
    1845
        ///
    1846
        /// This is useful when you don't need to keep the token_ids after the call.
    1847
        ///
    1848
        /// # Arguments
    1849
        ///
    1850
        /// * `token_ids` - Input token IDs (as Vec for owned semantics in tests)
    1851
        ///
    1852
        /// # Errors
    1853
        ///
    1854
        /// Returns error if forward pass fails
    1855
    1
        pub fn forward_gpu_owned(&mut self, token_ids: &[usize]) -> Result<Vec<f32>> {
    1856
    1
            self.forward_gpu(token_ids)
    1857
    1
        }
    1858
    1859
        /// Generate text tokens using GPU-accelerated inference (M14: E2E Inference)
    1860
        ///
    1861
        /// Performs autoregressive token generation starting from a prompt.
    1862
        /// Uses GPU for forward passes and CPU for sampling.
    1863
        ///
    1864
        /// # Arguments
    1865
        ///
    1866
        /// * `prompt` - Initial token IDs to start generation from
    1867
        /// * `config` - Generation configuration (max tokens, temperature, etc.)
    1868
        ///
    1869
        /// # Returns
    1870
        ///
    1871
        /// Vector of generated token IDs (including the prompt)
    1872
        ///
    1873
        /// # Errors
    1874
        ///
    1875
        /// Returns error if:
    1876
        /// - Prompt is empty
    1877
        /// - Forward pass fails
    1878
        ///
    1879
        /// # Examples
    1880
        ///
    1881
        /// ```rust,ignore
    1882
        /// let config = GpuGenerateConfig::deterministic(32);
    1883
        /// let tokens = model.generate(&[1, 2, 3], &config)?;
    1884
        /// ```
    1885
    22
        pub fn generate(&mut self, prompt: &[usize], config: &GpuGenerateConfig) -> Result<Vec<usize>> {
    1886
            // IMP-1009: Use zero-clone RefCell path when CUDA is available
    1887
            // This provides ~7x speedup by eliminating weight cloning
    1888
            #[cfg(feature = "cuda")]
    1889
            if self.cuda_scheduler.is_some() {
    1890
                return self.generate_refcell(prompt, config);
    1891
            }
    1892
    1893
            // Fallback to clone-based path for non-CUDA or HybridScheduler
    1894
            // IMP-091: Uses KV cache for O(n) generation
    1895
    22
            self.generate_optimized(prompt, config)
    1896
    22
        }
    1897
    1898
        // =========================================================================
    1899
        // Phase 7: KV Cache Integration - Wrappers (extracted to kv.rs)
    1900
        // =========================================================================
    1901
    1902
        /// Forward pass with KV cache population (IMP-031) - delegates to kv module
    1903
    35
        pub fn forward_gpu_with_cache(
    1904
    35
            &mut self,
    1905
    35
            token_ids: &[usize],
    1906
    35
            kv_cache: &mut StreamingKVCache,
    1907
    35
        ) -> Result<Vec<f32>> {
    1908
    35
            super::kv::forward_gpu_with_cache(self, token_ids, kv_cache)
    1909
    35
        }
    1910
    1911
        /// Incremental forward pass using cached KV (IMP-032) - delegates to kv module
    1912
    4
        pub fn forward_gpu_incremental(
    1913
    4
            &mut self,
    1914
    4
            token_id: usize,
    1915
    4
            kv_cache: &mut StreamingKVCache,
    1916
    4
        ) -> Result<Vec<f32>> {
    1917
    4
            super::kv::forward_gpu_incremental(self, token_id, kv_cache)
    1918
    4
        }
    1919
    1920
        /// Generate with KV cache (IMP-033) - delegates to kv module
    1921
    1
        pub fn generate_with_cache(
    1922
    1
            &mut self,
    1923
    1
            prompt: &[usize],
    1924
    1
            config: &GpuGenerateConfig,
    1925
    1
        ) -> Result<Vec<usize>> {
    1926
    1
            super::kv::generate_with_cache(self, prompt, config)
    1927
    1
        }
    1928
    1929
        /// Top-k sampling with temperature (returns highest prob token in top-k for determinism)
    1930
    3
        fn sample_topk_generate(logits: &[f32], temperature: f32, top_k: usize) -> usize {
    1931
            // Apply temperature
    1932
    768
            let scaled3: Vec<f32>3 = logits3.iter3().map3(|&x| x / temperature).collect3();
    1933
    1934
            // Softmax with numerical stability
    1935
    3
            let max_logit = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    1936
    768
            let exp_logits3: Vec<f32>3 = scaled.iter()3.map3(|&x| (x - max_logit).exp()).collect3();
    1937
    3
            let sum: f32 = exp_logits.iter().sum();
    1938
    768
            let probs3: Vec<f32>3 = exp_logits.iter()3.map3(|&x| x / sum).collect3();
    1939
    1940
            // Get top-k indices by sorting
    1941
    3
            let mut indexed: Vec<(usize, f32)> =
    1942
    768
                probs.iter()3.enumerate3().map3(|(i, &p)| (i, p)).collect3();
    1943
    765
            indexed3.sort_by3(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    1944
    1945
            // Truncate to top_k and return highest probability token (deterministic)
    1946
    3
            indexed.truncate(top_k);
    1947
    3
            indexed.first().map_or(0, |&(idx, _)| idx)
    1948
    3
        }
    1949
    1950
        /// Transpose weight matrix from [rows, cols] to [cols, rows]
    1951
    19
        fn transpose_weights(weights: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    1952
    19
            let mut transposed = vec![0.0f32; rows * cols];
    1953
    1.72k
            for i in 0..rows19 {
    1954
    25.9M
                for j in 0..cols1.72k {
    1955
    25.9M
                    transposed[j * rows + i] = weights[i * cols + j];
    1956
    25.9M
                }
    1957
            }
    1958
    19
            transposed
    1959
    19
        }
    1960
    1961
        /// Check if GPU is being used
    1962
        #[must_use]
    1963
    1
        pub fn has_gpu(&self) -> bool {
    1964
    1
            self.scheduler.has_gpu()
    1965
    1
        }
    1966
    1967
        /// GPU-accelerated forward pass
    1968
        ///
    1969
        /// Uses HybridScheduler for matrix multiplications.
    1970
        ///
    1971
        /// # Arguments
    1972
        ///
    1973
        /// * `token_ids` - Input token IDs
    1974
        ///
    1975
        /// # Returns
    1976
        ///
    1977
        /// Logits tensor with shape `[seq_len, vocab_size]`
    1978
        ///
    1979
        /// # Errors
    1980
        ///
    1981
        /// Returns error if forward pass fails
    1982
    5
        pub fn forward_gpu(&mut self, token_ids: &[usize]) -> Result<Vec<f32>> {
    1983
    5
            if token_ids.is_empty() {
    1984
    0
                return Err(RealizarError::InvalidShape {
    1985
    0
                    reason: "Token IDs cannot be empty".to_string(),
    1986
    0
                });
    1987
    5
            }
    1988
    1989
    5
            let seq_len = token_ids.len();
    1990
    5
            let hidden_dim = self.config.hidden_dim;
    1991
    1992
            // Step 1: Embed tokens
    1993
    5
            let mut hidden = Vec::with_capacity(seq_len * hidden_dim);
    1994
    21
            for &token_id16 in token_ids {
    1995
    16
                if token_id >= self.config.vocab_size {
    1996
    0
                    return Err(RealizarError::InvalidShape {
    1997
    0
                        reason: format!(
    1998
    0
                            "Token ID {} out of bounds (vocab_size={})",
    1999
    0
                            token_id, self.config.vocab_size
    2000
    0
                        ),
    2001
    0
                    });
    2002
    16
                }
    2003
    16
                let offset = token_id * hidden_dim;
    2004
    16
                hidden.extend_from_slice(&self.embedding_weights[offset..offset + hidden_dim]);
    2005
            }
    2006
    2007
            // Step 2: Pass through transformer blocks
    2008
    10
            for block_idx in 0..self.block_weights5.len5() {
    2009
    10
                hidden = self.forward_block_idx(&hidden, seq_len, block_idx)?0;
    2010
            }
    2011
    2012
            // Step 3: Final layer norm
    2013
    5
            hidden = self.layer_norm(&hidden, &self.final_norm_weight, &self.final_norm_bias);
    2014
    2015
            // Step 4: LM head projection
    2016
            // [seq_len, hidden_dim] @ [hidden_dim, vocab_size] -> [seq_len, vocab_size]
    2017
            // Phase 22 FIX: Use lm_head_weight_t (transposed) which is [hidden_dim, vocab_size]
    2018
            // The original lm_head_weight is [vocab_size, hidden_dim] (APR convention)
    2019
            // IMP-090: Use CPU fallback for large vocab to avoid GPU buffer overflow
    2020
    5
            let lm_head_elements = hidden_dim * self.config.vocab_size;
    2021
    5
            let logits = if exceeds_gpu_buffer_limit(lm_head_elements) {
    2022
                // CPU fallback for large vocab (>256MB weight matrix)
    2023
    0
                cpu_matmul(
    2024
    0
                    &hidden,
    2025
    0
                    &self.lm_head_weight_t,
    2026
    0
                    seq_len,
    2027
    0
                    hidden_dim,
    2028
    0
                    self.config.vocab_size,
    2029
                )
    2030
            } else {
    2031
                // GPU path for smaller vocab (IMP-1005: use do_matmul for CUDA)
    2032
                // Clone weights to avoid borrow conflict with &mut self in do_matmul
    2033
    5
                let lm_weight = self.lm_head_weight_t.clone();
    2034
    5
                self.do_matmul(
    2035
    5
                    &hidden,
    2036
    5
                    &lm_weight,
    2037
    5
                    seq_len,
    2038
    5
                    hidden_dim,
    2039
    5
                    self.config.vocab_size,
    2040
    0
                )?
    2041
            };
    2042
    2043
            // Add bias
    2044
    5
            let mut output = logits;
    2045
    16
            for i in 0..seq_len5 {
    2046
    403k
                for j in 0..self.config.vocab_size16 {
    2047
    403k
                    output[i * self.config.vocab_size + j] += self.lm_head_bias[j];
    2048
    403k
                }
    2049
            }
    2050
    2051
    5
            Ok(output)
    2052
    5
        }
    2053
    2054
        /// Forward pass through a single transformer block by index
    2055
    10
        pub fn forward_block_idx(
    2056
    10
            &mut self,
    2057
    10
            input: &[f32],
    2058
    10
            seq_len: usize,
    2059
    10
            block_idx: usize,
    2060
    10
        ) -> Result<Vec<f32>> {
    2061
    10
            let hidden_dim = self.config.hidden_dim;
    2062
    10
            let intermediate_dim = self.config.intermediate_dim;
    2063
    10
            let qkv_dim = self.config.qkv_dim();
    2064
    2065
            // Get references to block weights (avoid cloning)
    2066
    10
            let block = &self.block_weights[block_idx];
    2067
    10
            let attn_norm_weight = &block.attn_norm_weight;
    2068
    10
            let attn_norm_bias = &block.attn_norm_bias;
    2069
    2070
            // Pre-norm (uses references, no clone)
    2071
    10
            let normed = Self::layer_norm_static(
    2072
    10
                input,
    2073
    10
                attn_norm_weight,
    2074
    10
                attn_norm_bias,
    2075
    10
                hidden_dim,
    2076
    10
                self.config.eps,
    2077
            );
    2078
    2079
            // IMP-1005: Clone weights to avoid borrow conflict with &mut self in do_matmul
    2080
    10
            let qkv_weight = self.block_weights[block_idx].qkv_weight.clone();
    2081
    2082
            // QKV projection (IMP-1005: use do_matmul for CUDA)
    2083
            // [seq_len, hidden_dim] @ [hidden_dim, qkv_dim] -> [seq_len, qkv_dim]
    2084
    10
            let qkv = self.do_matmul(&normed, &qkv_weight, seq_len, hidden_dim, qkv_dim)?0;
    2085
    2086
            // Optimized GQA attention with GPU matmul for scores
    2087
    10
            let attn_out = self.optimized_gqa_attention(&qkv, seq_len)?0;
    2088
    2089
            // IMP-1005: Clone weights to avoid borrow conflict
    2090
    10
            let out_weight = self.block_weights[block_idx].out_weight.clone();
    2091
    10
            let out_bias = self.block_weights[block_idx].out_bias.clone();
    2092
    2093
            // Output projection (IMP-1005: use do_matmul for CUDA)
    2094
    10
            let projected = self.do_matmul(&attn_out, &out_weight, seq_len, hidden_dim, hidden_dim)?0;
    2095
    2096
            // Residual 1 (vectorized)
    2097
    10
            let mut residual1: Vec<f32> = input
    2098
    10
                .iter()
    2099
    10
                .zip(projected.iter())
    2100
    10
                .enumerate()
    2101
    3.96k
                .map10(|(i, (&inp, &proj))| inp + proj + out_bias[i % hidden_dim])
    2102
    10
                .collect();
    2103
    2104
            // IMP-1005: Clone weights to avoid borrow conflict
    2105
    10
            let ffn_norm_weight = self.block_weights[block_idx].ffn_norm_weight.clone();
    2106
    10
            let ffn_norm_bias = self.block_weights[block_idx].ffn_norm_bias.clone();
    2107
    2108
            // FFN pre-norm
    2109
    10
            let ffn_normed = Self::layer_norm_static(
    2110
    10
                &residual1,
    2111
    10
                &ffn_norm_weight,
    2112
    10
                &ffn_norm_bias,
    2113
    10
                hidden_dim,
    2114
    10
                self.config.eps,
    2115
            );
    2116
    2117
            // IMP-1005: Clone weights to avoid borrow conflict
    2118
    10
            let ffn_fc1_weight = self.block_weights[block_idx].ffn_fc1_weight.clone();
    2119
    10
            let ffn_fc1_bias = self.block_weights[block_idx].ffn_fc1_bias.clone();
    2120
    10
            let ffn_gate_weight = self.block_weights[block_idx].ffn_gate_weight.clone();
    2121
    2122
            // FFN: SwiGLU when gate weight exists, otherwise GELU
    2123
    10
            let activated: Vec<f32> = if let Some(gate_weight0) = ffn_gate_weight {
    2124
                // SwiGLU: silu(gate(x)) * up(x)
    2125
    0
                let up_out = self.do_matmul(&ffn_normed, &ffn_fc1_weight, seq_len, hidden_dim, intermediate_dim)?;
    2126
    0
                let gate_out = self.do_matmul(&ffn_normed, &gate_weight, seq_len, hidden_dim, intermediate_dim)?;
    2127
    2128
                // SwiGLU: silu(gate) * up
    2129
    0
                up_out
    2130
    0
                    .iter()
    2131
    0
                    .zip(gate_out.iter())
    2132
    0
                    .map(|(&u, &g)| {
    2133
    0
                        let silu_g = g / (1.0 + (-g).exp());
    2134
    0
                        silu_g * u
    2135
    0
                    })
    2136
    0
                    .collect()
    2137
            } else {
    2138
                // Standard GELU FFN
    2139
    10
                let fc1_out = self.do_matmul(
    2140
    10
                    &ffn_normed,
    2141
    10
                    &ffn_fc1_weight,
    2142
    10
                    seq_len,
    2143
    10
                    hidden_dim,
    2144
    10
                    intermediate_dim,
    2145
    0
                )?;
    2146
    2147
                // GELU activation + bias (vectorized)
    2148
    10
                fc1_out
    2149
    10
                    .iter()
    2150
    10
                    .enumerate()
    2151
    7.93k
                    .map10(|(i, &x)| {
    2152
    7.93k
                        let x = x + ffn_fc1_bias[i % intermediate_dim];
    2153
                        // GELU approximation
    2154
    7.93k
                        0.5 * x
    2155
    7.93k
                            * (1.0
    2156
    7.93k
                                + ((2.0f32 / std::f32::consts::PI).sqrt() * (x + 0.044_715 * x.powi(3)))
    2157
    7.93k
                                    .tanh())
    2158
    7.93k
                    })
    2159
    10
                    .collect()
    2160
            };
    2161
    2162
            // IMP-1005: Clone weights to avoid borrow conflict
    2163
    10
            let ffn_fc2_weight = self.block_weights[block_idx].ffn_fc2_weight.clone();
    2164
    10
            let ffn_fc2_bias = self.block_weights[block_idx].ffn_fc2_bias.clone();
    2165
    2166
            // FFN: fc2 (IMP-1005: use do_matmul for CUDA)
    2167
    10
            let fc2_out = self.do_matmul(
    2168
    10
                &activated,
    2169
    10
                &ffn_fc2_weight,
    2170
    10
                seq_len,
    2171
    10
                intermediate_dim,
    2172
    10
                hidden_dim,
    2173
    0
            )?;
    2174
    2175
            // Residual 2 (vectorized, in-place)
    2176
    3.96k
            for (i, x) in residual1.iter_mut()10.enumerate10() {
    2177
    3.96k
                *x += fc2_out[i] + ffn_fc2_bias[i % hidden_dim];
    2178
    3.96k
            }
    2179
    2180
    10
            Ok(residual1)
    2181
    10
        }
    2182
    2183
        /// RMSNorm (Root Mean Square Layer Normalization)
    2184
        ///
    2185
        /// PMAT-094 FIX: Qwen2, LLaMA, Mistral use RMSNorm, NOT LayerNorm.
    2186
        /// Formula: output = x / sqrt(mean(x^2) + eps) * weight + bias
    2187
        #[allow(clippy::cast_precision_loss)]
    2188
    4.17k
        pub(crate) fn layer_norm_static(
    2189
    4.17k
            input: &[f32],
    2190
    4.17k
            weight: &[f32],
    2191
    4.17k
            bias: &[f32],
    2192
    4.17k
            hidden_dim: usize,
    2193
    4.17k
            eps: f32,
    2194
    4.17k
        ) -> Vec<f32> {
    2195
    4.17k
            let num_rows = input.len() / hidden_dim;
    2196
    4.17k
            let mut output = Vec::with_capacity(input.len());
    2197
    2198
    5.39k
            for row in 0..num_rows4.17k {
    2199
    5.39k
                let start = row * hidden_dim;
    2200
    5.39k
                let row_data = &input[start..start + hidden_dim];
    2201
    2202
                // RMSNorm: compute root mean square (no mean subtraction!)
    2203
    567k
                let sum_sq5.39k: f325.39k = row_data5.39k.iter5.39k().map5.39k(|&x| x * x).sum5.39k();
    2204
    5.39k
                let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();
    2205
    2206
                // Normalize and scale
    2207
    567k
                for (i, &x) in row_data5.39k.iter5.39k().enumerate5.39k() {
    2208
    567k
                    let normalized = x / rms;
    2209
    567k
                    output.push(normalized * weight[i] + bias[i]);
    2210
    567k
                }
    2211
            }
    2212
    2213
    4.17k
            output
    2214
    4.17k
        }
    2215
    2216
        /// Layer normalization (instance method)
    2217
    529
        fn layer_norm(&self, input: &[f32], weight: &[f32], bias: &[f32]) -> Vec<f32> {
    2218
    529
            Self::layer_norm_static(input, weight, bias, self.config.hidden_dim, self.config.eps)
    2219
    529
        }
    2220
    2221
        /// Generate tokens using GPU-accelerated forward pass with incremental decoding (wrapper)
    2222
    0
        pub fn generate_gpu(&mut self, prompt: &[usize], max_tokens: usize) -> Result<Vec<usize>> {
    2223
    0
            super::batch::generate_gpu(self, prompt, max_tokens)
    2224
    0
        }
    2225
    2226
        /// Argmax helper for sampling (wrapper)
    2227
    500
        fn argmax(logits: &[f32]) -> usize {
    2228
    500
            super::batch::argmax(logits)
    2229
    500
        }
    2230
    2231
        /// Optimized GQA attention using GPU for matmul operations (wrapper)
    2232
    10
        fn optimized_gqa_attention(&mut self, qkv: &[f32], seq_len: usize) -> Result<Vec<f32>> {
    2233
    10
            super::batch::optimized_gqa_attention(self, qkv, seq_len)
    2234
    10
        }
    2235
    }