## Coverage Report

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

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

    /home/noah/src/realizar/src/gguf/inference/cached/sync.rs

    Line
    Count
    Source
    1
    //! Thread-safe cached model wrapper (Mutex-based)
    2
    //!
    3
    //! `OwnedQuantizedModelCachedSync` uses Mutex for interior mutability,
    4
    //! suitable for async HTTP servers and multi-threaded inference.
    5
    6
    use crate::error::{RealizarError, Result};
    7
    use crate::gguf::{
    8
        BatchGenerationStats, DispatchMetrics, OwnedQuantizedKVCache,
    9
        OwnedQuantizedModel, QuantizedGenerateConfig,
    10
    };
    11
    use super::weights::{DequantizedFFNWeights, DequantizedWeightCache};
    12
    13
    /// Thread-safe cached model wrapper with Mutex-based scheduler caching
    14
    ///
    15
    /// Uses `Mutex` for interior mutability to cache GPU schedulers. Safe for
    16
    /// multi-threaded HTTP serving with async handlers.
    17
    pub struct OwnedQuantizedModelCachedSync {
    18
        /// Inner model (not cached)
    19
        model: OwnedQuantizedModel,
    20
        /// Cached HybridScheduler for GPU operations (wgpu backend)
    21
        /// Uses Mutex for thread-safe interior mutability
    22
        scheduler: std::sync::Mutex<Option<crate::gpu::HybridScheduler>>,
    23
        /// PARITY-103: Cached CudaScheduler for direct CUDA operations
    24
        /// Bypasses wgpu 256MB buffer limit by using cuBLAS directly
    25
        #[cfg(feature = "cuda")]
    26
        cuda_scheduler: std::sync::Mutex<Option<crate::gpu::CudaScheduler>>,
    27
        /// Dequantized weight cache for GPU batch inference (PARITY-019)
    28
        /// Uses RwLock for concurrent read access during batch inference
    29
        dequant_cache: std::sync::RwLock<Option<DequantizedWeightCache>>,
    30
    }
    31
    32
    // Explicitly implement Send + Sync for HTTP server usage
    33
    #[cfg(feature = "gpu")]
    34
    unsafe impl Send for OwnedQuantizedModelCachedSync {}
    35
    #[cfg(feature = "gpu")]
    36
    unsafe impl Sync for OwnedQuantizedModelCachedSync {}
    37
    38
    #[cfg(feature = "gpu")]
    39
    impl OwnedQuantizedModelCachedSync {
    40
        /// Create a new thread-safe cached model wrapper
    41
        ///
    42
        /// The scheduler is lazily initialized on first GPU operation.
    43
        /// The dequantized weight cache is lazily initialized via `warmup_gpu_cache()`.
    44
        /// PARITY-103: Also initializes CudaScheduler when CUDA feature is enabled.
    45
        #[must_use]
    46
    39
        pub fn new(model: OwnedQuantizedModel) -> Self {
    47
    39
            Self {
    48
    39
                model,
    49
    39
                scheduler: std::sync::Mutex::new(None),
    50
    39
                #[cfg(feature = "cuda")]
    51
    39
                cuda_scheduler: std::sync::Mutex::new(None),
    52
    39
                dequant_cache: std::sync::RwLock::new(None),
    53
    39
            }
    54
    39
        }
    55
    56
        /// Get reference to inner model
    57
        #[must_use]
    58
    2.64k
        pub fn model(&self) -> &OwnedQuantizedModel {
    59
    2.64k
            &self.model
    60
    2.64k
        }
    61
    62
        /// Get or create the cached scheduler (thread-safe)
    63
        ///
    64
        /// # Errors
    65
        /// Returns error if scheduler creation fails or lock is poisoned
    66
    0
        fn get_scheduler(
    67
    0
            &self,
    68
    0
        ) -> Result<std::sync::MutexGuard<'_, Option<crate::gpu::HybridScheduler>>> {
    69
    0
            let mut scheduler_opt =
    70
    0
                self.scheduler
    71
    0
                    .lock()
    72
    0
                    .map_err(|_| RealizarError::UnsupportedOperation {
    73
    0
                        operation: "scheduler_lock".to_string(),
    74
    0
                        reason: "Scheduler mutex poisoned".to_string(),
    75
    0
                    })?;
    76
    77
            // Initialize if not already done
    78
    0
            if scheduler_opt.is_none() {
    79
                use crate::gpu::HybridScheduler;
    80
    0
                let new_scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {
    81
    0
                    RealizarError::UnsupportedOperation {
    82
    0
                        operation: "HybridScheduler::with_threshold".to_string(),
    83
    0
                        reason: format!("GPU scheduler initialization failed: {e}"),
    84
    0
                    }
    85
    0
                })?;
    86
    0
                *scheduler_opt = Some(new_scheduler);
    87
    0
            }
    88
    89
    0
            Ok(scheduler_opt)
    90
    0
        }
    91
    92
        /// PARITY-103: Get or create the cached CUDA scheduler (thread-safe)
    93
        ///
    94
        /// Bypasses wgpu 256MB buffer limit by using cuBLAS directly.
    95
        /// Returns None if CUDA is not available.
    96
        ///
    97
        /// # Errors
    98
        /// Returns error if lock is poisoned
    99
        #[cfg(feature = "cuda")]
    100
        fn get_cuda_scheduler(
    101
            &self,
    102
        ) -> Result<std::sync::MutexGuard<'_, Option<crate::gpu::CudaScheduler>>> {
    103
            use crate::gpu::CudaScheduler;
    104
    105
            let mut scheduler_opt =
    106
                self.cuda_scheduler
    107
                    .lock()
    108
                    .map_err(|_| RealizarError::UnsupportedOperation {
    109
                        operation: "cuda_scheduler_lock".to_string(),
    110
                        reason: "CUDA scheduler mutex poisoned".to_string(),
    111
                    })?;
    112
    113
            // Initialize if not already done
    114
            if scheduler_opt.is_none() {
    115
                match CudaScheduler::new() {
    116
                    Ok(new_scheduler) => {
    117
                        eprintln!("PARITY-103: CudaScheduler initialized successfully");
    118
                        *scheduler_opt = Some(new_scheduler);
    119
                    },
    120
                    Err(e) => {
    121
                        // CUDA not available, leave as None (will fallback to wgpu)
    122
                        eprintln!("PARITY-103: CudaScheduler::new() failed: {:?}", e);
    123
                    },
    124
                }
    125
            }
    126
    127
            Ok(scheduler_opt)
    128
        }
    129
    130
        /// PARITY-103: Batch matmul preferring CUDA over wgpu (thread-safe)
    131
        ///
    132
        /// Tries CudaScheduler first (no buffer limits), falls back to HybridScheduler (wgpu).
    133
        /// This bypasses the wgpu 256MB buffer limit that was blocking GPU batch inference.
    134
        #[cfg(feature = "cuda")]
    135
        fn batch_matmul_gpu_prefer_cuda(
    136
            &self,
    137
            input: &[f32],
    138
            weight_f32: &[f32],
    139
            batch_size: usize,
    140
            in_dim: usize,
    141
            out_dim: usize,
    142
        ) -> Result<Vec<f32>> {
    143
            // Validate input
    144
            if input.len() != batch_size * in_dim {
    145
                return Err(RealizarError::InvalidShape {
    146
                    reason: format!(
    147
                        "Input size {} doesn't match batch_size={} * in_dim={}",
    148
                        input.len(),
    149
                        batch_size,
    150
                        in_dim
    151
                    ),
    152
                });
    153
            }
    154
    155
            // Try CUDA first (no buffer size limits)
    156
            if let Ok(mut cuda_guard) = self.get_cuda_scheduler() {
    157
                if let Some(ref mut cuda_sched) = *cuda_guard {
    158
                    return cuda_sched
    159
                        .matmul(input, weight_f32, batch_size, in_dim, out_dim)
    160
                        .map_err(|e| RealizarError::UnsupportedOperation {
    161
                            operation: "batch_matmul_gpu_prefer_cuda".to_string(),
    162
                            reason: format!("CUDA matmul failed: {e}"),
    163
                        });
    164
                }
    165
            }
    166
    167
            // Fallback to wgpu (may hit 256MB limit for large batches)
    168
            let mut scheduler_guard = self.get_scheduler()?;
    169
            if let Some(ref mut scheduler) = *scheduler_guard {
    170
                return scheduler
    171
                    .matmul(input, weight_f32, batch_size, in_dim, out_dim)
    172
                    .map_err(|e| RealizarError::UnsupportedOperation {
    173
                        operation: "batch_matmul_gpu_prefer_cuda".to_string(),
    174
                        reason: format!("GPU matmul failed: {e}"),
    175
                    });
    176
            }
    177
    178
            Err(RealizarError::UnsupportedOperation {
    179
                operation: "batch_matmul_gpu_prefer_cuda".to_string(),
    180
                reason: "No GPU scheduler available".to_string(),
    181
            })
    182
        }
    183
    184
        /// PARITY-103: Batch matmul preferring CUDA (non-CUDA fallback)
    185
        #[cfg(not(feature = "cuda"))]
    186
    0
        fn batch_matmul_gpu_prefer_cuda(
    187
    0
            &self,
    188
    0
            input: &[f32],
    189
    0
            weight_f32: &[f32],
    190
    0
            batch_size: usize,
    191
    0
            in_dim: usize,
    192
    0
            out_dim: usize,
    193
    0
        ) -> Result<Vec<f32>> {
    194
            // Validate input
    195
    0
            if input.len() != batch_size * in_dim {
    196
    0
                return Err(RealizarError::InvalidShape {
    197
    0
                    reason: format!(
    198
    0
                        "Input size {} doesn't match batch_size={} * in_dim={}",
    199
    0
                        input.len(),
    200
    0
                        batch_size,
    201
    0
                        in_dim
    202
    0
                    ),
    203
    0
                });
    204
    0
            }
    205
    206
    0
            let mut scheduler_guard = self.get_scheduler()?;
    207
    0
            if let Some(ref mut scheduler) = *scheduler_guard {
    208
    0
                return scheduler
    209
    0
                    .matmul(input, weight_f32, batch_size, in_dim, out_dim)
    210
    0
                    .map_err(|e| RealizarError::UnsupportedOperation {
    211
    0
                        operation: "batch_matmul_gpu_prefer_cuda".to_string(),
    212
    0
                        reason: format!("GPU matmul failed: {e}"),
    213
    0
                    });
    214
    0
            }
    215
    216
    0
            Err(RealizarError::UnsupportedOperation {
    217
    0
                operation: "batch_matmul_gpu_prefer_cuda".to_string(),
    218
    0
                reason: "No GPU scheduler available".to_string(),
    219
    0
            })
    220
    0
        }
    221
    222
        /// Generate tokens with KV cache using thread-safe cached scheduler
    223
        ///
    224
        /// Delegates to the inner model's `generate_with_cache` method.
    225
        /// The scheduler caching benefits GPU batch operations; single-token
    226
        /// generation uses CPU path with KV cache for O(n) scaling.
    227
        ///
    228
        /// # Arguments
    229
        /// * `prompt` - Input token IDs
    230
        /// * `config` - Generation configuration
    231
        ///
    232
        /// # Returns
    233
        /// Generated token sequence including prompt
    234
        ///
    235
        /// # Errors
    236
        /// Returns error if generation fails
    237
    0
        pub fn generate_with_cache(
    238
    0
            &self,
    239
    0
            prompt: &[u32],
    240
    0
            config: &QuantizedGenerateConfig,
    241
    0
        ) -> Result<Vec<u32>> {
    242
            // Delegate to inner model - CPU path with KV cache is already efficient
    243
    0
            self.model.generate_with_cache(prompt, config)
    244
    0
        }
    245
    246
        /// Generate tokens with adaptive CPU/GPU attention (IMP-126)
    247
        ///
    248
        /// This variant of `generate_with_cache` uses adaptive CPU/GPU dispatch
    249
        /// based on cache length and records dispatch decisions to metrics.
    250
        ///
    251
        /// # Arguments
    252
        /// * `prompt` - Initial token IDs
    253
        /// * `config` - Generation configuration
    254
        /// * `metrics` - Dispatch metrics tracker for CPU/GPU decision recording
    255
        ///
    256
        /// # Returns
    257
        /// Generated token sequence including prompt
    258
        ///
    259
        /// # Errors
    260
        /// Returns error if generation fails
    261
        #[cfg(feature = "gpu")]
    262
    5
        pub fn generate_with_cache_adaptive(
    263
    5
            &self,
    264
    5
            prompt: &[u32],
    265
    5
            config: &QuantizedGenerateConfig,
    266
    5
            metrics: &std::sync::Arc<DispatchMetrics>,
    267
    5
        ) -> Result<Vec<u32>> {
    268
            // Delegate to inner model's adaptive generation
    269
    5
            self.model
    270
    5
                .generate_with_cache_adaptive(prompt, config, metrics)
    271
    5
        }
    272
    273
        /// Forward pass with cached scheduler (thread-safe)
    274
        ///
    275
        /// Uses the cached HybridScheduler for GPU operations.
    276
        ///
    277
        /// # Errors
    278
        /// Returns error if GPU operations fail
    279
        #[allow(clippy::let_underscore_untyped)] // Placeholder for future use
    280
    0
        pub fn forward_batch_gpu_cached(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
    281
    0
            let batch_size = token_ids.len();
    282
    0
            let vocab_size = self.model.config.vocab_size;
    283
    284
            // Get cached scheduler (for future GPU operations)
    285
    0
            let mut scheduler_guard = self.get_scheduler()?;
    286
    0
            let _ = scheduler_guard
    287
    0
                .as_mut()
    288
    0
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    289
    0
                    operation: "forward_batch_gpu_cached".to_string(),
    290
    0
                    reason: "Scheduler not initialized".to_string(),
    291
    0
                })?;
    292
    293
            // 1. Token embedding lookup
    294
    0
            let hidden = self.model.embed(token_ids);
    295
    296
            // 2. Process through layers
    297
    0
            for layer in &self.model.layers {
    298
    0
                // Simplified single-layer forward - reuse inner model logic
    299
    0
                // For full implementation, would need to port the complete forward pass
    300
    0
                let _ = layer;
    301
    0
            }
    302
    303
            // 3. Output normalization and LM head
    304
            // For now, return placeholder - full implementation requires porting forward logic
    305
    0
            let output = vec![0.0f32; batch_size * vocab_size];
    306
    0
            let _ = hidden;
    307
    308
    0
            Ok(output)
    309
    0
        }
    310
    311
        /// Adaptive fused attention for production serving (IMP-121)
    312
        ///
    313
        /// Thread-safe wrapper that automatically selects CPU or GPU based on
    314
        /// sequence length. Uses the cached scheduler for efficient GPU operations.
    315
        ///
    316
        /// # Arguments
    317
        /// * `q` - Query tensor [seq_len, head_dim]
    318
        /// * `k` - Key tensor [seq_len, head_dim]
    319
        /// * `v` - Value tensor [seq_len, head_dim]
    320
        /// * `seq_len` - Sequence length
    321
        /// * `head_dim` - Head dimension
    322
        /// * `scale` - Attention scale factor
    323
        ///
    324
        /// # Returns
    325
        /// Output tensor [seq_len, head_dim]
    326
    9
        pub fn adaptive_fused_attention(
    327
    9
            &self,
    328
    9
            q: &[f32],
    329
    9
            k: &[f32],
    330
    9
            v: &[f32],
    331
    9
            seq_len: usize,
    332
    9
            head_dim: usize,
    333
    9
            scale: f32,
    334
    9
        ) -> Result<Vec<f32>> {
    335
            // Threshold for GPU dispatch (from IMP-119 analysis)
    336
            const GPU_SEQ_LEN_THRESHOLD: usize = 64;
    337
    338
    9
            if seq_len >= GPU_SEQ_LEN_THRESHOLD {
    339
                // Long sequence: Use GPU path
    340
    4
                self.gpu_fused_causal_attention(q, k, v, seq_len, head_dim, scale)
    341
            } else {
    342
                // Short sequence: Use CPU path
    343
    5
                self.cpu_fused_causal_attention(q, k, v, seq_len, head_dim, scale)
    344
            }
    345
    9
        }
    346
    347
        /// CPU fused causal attention (thread-safe wrapper)
    348
    5
        fn cpu_fused_causal_attention(
    349
    5
            &self,
    350
    5
            q: &[f32],
    351
    5
            k: &[f32],
    352
    5
            v: &[f32],
    353
    5
            seq_len: usize,
    354
    5
            head_dim: usize,
    355
    5
            scale: f32,
    356
    5
        ) -> Result<Vec<f32>> {
    357
            // Use tiled implementation from inner model
    358
    5
            self.model
    359
    5
                .tiled_causal_attention(q, k, v, seq_len, head_dim, scale, 4)
    360
    5
        }
    361
    362
        /// GPU fused causal attention (thread-safe)
    363
    4
        fn gpu_fused_causal_attention(
    364
    4
            &self,
    365
    4
            q: &[f32],
    366
    4
            k: &[f32],
    367
    4
            v: &[f32],
    368
    4
            seq_len: usize,
    369
    4
            head_dim: usize,
    370
    4
            scale: f32,
    371
    4
        ) -> Result<Vec<f32>> {
    372
    4
            let mut scheduler_guard =
    373
    4
                self.scheduler
    374
    4
                    .lock()
    375
    4
                    .map_err(|_| RealizarError::UnsupportedOperation {
    376
    0
                        operation: "gpu_fused_causal_attention".to_string(),
    377
    0
                        reason: "Failed to acquire scheduler lock".to_string(),
    378
    0
                    })?;
    379
    380
            // Initialize scheduler if needed
    381
    4
            if scheduler_guard.is_none() {
    382
                use crate::gpu::HybridScheduler;
    383
    1
                let new_scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
    384
    0
                    RealizarError::UnsupportedOperation {
    385
    0
                        operation: "HybridScheduler::with_threshold".to_string(),
    386
    0
                        reason: format!("GPU scheduler initialization failed: {e}"),
    387
    0
                    }
    388
    0
                })?;
    389
    1
                *scheduler_guard = Some(new_scheduler);
    390
    3
            }
    391
    392
    4
            let scheduler =
    393
    4
                scheduler_guard
    394
    4
                    .as_mut()
    395
    4
                    .ok_or_else(|| RealizarError::UnsupportedOperation {
    396
    0
                        operation: "gpu_fused_causal_attention".to_string(),
    397
    0
                        reason: "Scheduler not initialized".to_string(),
    398
    0
                    })?;
    399
    400
            // Transpose K for matmul
    401
    4
            let mut k_transposed = vec![0.0f32; head_dim * seq_len];
    402
    256
            for pos in 0..seq_len4 {
    403
    4.09k
                for d in 0..head_dim256 {
    404
    4.09k
                    k_transposed[d * seq_len + pos] = k[pos * head_dim + d];
    405
    4.09k
                }
    406
            }
    407
    408
            // GPU Q @ K^T
    409
    4
            let scores = scheduler
    410
    4
                .matmul(q, &k_transposed, seq_len, head_dim, seq_len)
    411
    4
                .map_err(|e| RealizarError::UnsupportedOperation {
    412
    0
                    operation: "gpu_fused Q@K^T".to_string(),
    413
    0
                    reason: format!("GPU matmul failed: {}", e),
    414
    0
                })?;
    415
    416
            // CPU causal softmax
    417
    4
            let mut weights = vec![0.0f32; seq_len * seq_len];
    418
    256
            for i in 0..seq_len4 {
    419
    256
                let mut max_val = f32::NEG_INFINITY;
    420
    8.32k
                for j in 0..=i256 {
    421
    8.32k
                    let score = scores[i * seq_len + j] * scale;
    422
    8.32k
                    if score > max_val {
    423
    1.09k
                        max_val = score;
    424
    7.22k
                    }
    425
                }
    426
    256
                let mut sum = 0.0f32;
    427
    8.32k
                for j in 0..=i256 {
    428
    8.32k
                    let score = scores[i * seq_len + j] * scale;
    429
    8.32k
                    weights[i * seq_len + j] = (score - max_val).exp();
    430
    8.32k
                    sum += weights[i * seq_len + j];
    431
    8.32k
                }
    432
    256
                if sum > 0.0 {
    433
    8.32k
                    for j in 0..=i256 {
    434
    8.32k
                        weights[i * seq_len + j] /= sum;
    435
    8.32k
                    }
    436
    0
                }
    437
            }
    438
    439
            // GPU weights @ V
    440
    4
            scheduler
    441
    4
                .matmul(&weights, v, seq_len, seq_len, head_dim)
    442
    4
                .map_err(|e| RealizarError::UnsupportedOperation {
    443
    0
                    operation: "gpu_fused weights@V".to_string(),
    444
    0
                    reason: format!("GPU matmul failed: {}", e),
    445
    0
                })
    446
    4
        }
    447
    448
        /// Adaptive multihead attention for production serving (IMP-121)
    449
        ///
    450
        /// Thread-safe multi-head attention that automatically selects backend.
    451
        ///
    452
        /// # Arguments
    453
        /// * `q` - Query tensor [seq_len, hidden_dim]
    454
        /// * `k` - Key tensor [seq_len, hidden_dim]
    455
        /// * `v` - Value tensor [seq_len, hidden_dim]
    456
        /// * `seq_len` - Sequence length
    457
        ///
    458
        /// # Returns
    459
        /// Output tensor [seq_len, hidden_dim]
    460
    1
        pub fn adaptive_multihead_attention(
    461
    1
            &self,
    462
    1
            q: &[f32],
    463
    1
            k: &[f32],
    464
    1
            v: &[f32],
    465
    1
            seq_len: usize,
    466
    1
        ) -> Result<Vec<f32>> {
    467
    1
            let hidden_dim = self.model.config.hidden_dim;
    468
    1
            let num_heads = self.model.config.num_heads;
    469
    1
            let head_dim = hidden_dim / num_heads;
    470
    1
            let scale = 1.0 / (head_dim as f32).sqrt();
    471
    472
            // Reshape Q, K, V to [num_heads, seq_len, head_dim]
    473
    1
            let q_reshaped = self
    474
    1
                .model
    475
    1
                .reshape_for_parallel_heads(q, seq_len, num_heads, head_dim)?0;
    476
    1
            let k_reshaped = self
    477
    1
                .model
    478
    1
                .reshape_for_parallel_heads(k, seq_len, num_heads, head_dim)?0;
    479
    1
            let v_reshaped = self
    480
    1
                .model
    481
    1
                .reshape_for_parallel_heads(v, seq_len, num_heads, head_dim)?0;
    482
    483
    1
            let mut attn_output = vec![0.0f32; num_heads * seq_len * head_dim];
    484
    485
    4
            for h in 0..num_heads1 {
    486
    4
                let head_offset = h * seq_len * head_dim;
    487
    4
                let q_head = &q_reshaped[head_offset..head_offset + seq_len * head_dim];
    488
    4
                let k_head = &k_reshaped[head_offset..head_offset + seq_len * head_dim];
    489
    4
                let v_head = &v_reshaped[head_offset..head_offset + seq_len * head_dim];
    490
    491
    4
                let head_output =
    492
    4
                    self.adaptive_fused_attention(q_head, k_head, v_head, seq_len, head_dim, scale)?0;
    493
    494
    4
                attn_output[head_offset..head_offset + seq_len * head_dim]
    495
    4
                    .copy_from_slice(&head_output);
    496
            }
    497
    498
            // Reshape back to [seq_len, hidden_dim]
    499
    1
            let mut output = vec![0.0f32; seq_len * hidden_dim];
    500
    4
            for h in 0..num_heads1 {
    501
    4
                let head_start = h * seq_len * head_dim;
    502
    256
                for pos in 0..seq_len4 {
    503
    256
                    let src_start = head_start + pos * head_dim;
    504
    256
                    let dst_start = pos * hidden_dim + h * head_dim;
    505
    256
                    output[dst_start..dst_start + head_dim]
    506
    256
                        .copy_from_slice(&attn_output[src_start..src_start + head_dim]);
    507
    256
                }
    508
            }
    509
    510
    1
            Ok(output)
    511
    1
        }
    512
    513
        /// Warmup GPU weight cache for batch inference (PARITY-019)
    514
        ///
    515
        /// Pre-dequantizes all FFN weights to f32 for GPU GEMM operations.
    516
        /// Call this once at server startup to avoid dequantization during inference.
    517
        ///
    518
        /// # Memory Usage
    519
        /// - phi-2 (32 layers): ~6.4 GB
    520
        /// - Per layer: 2 × hidden_dim × intermediate_dim × 4 bytes
    521
        ///
    522
        /// # Returns
    523
        /// - Total memory allocated in bytes
    524
        /// - Number of layers cached
    525
        ///
    526
        /// # Errors
    527
        /// Returns error if dequantization fails
    528
    0
        pub fn warmup_gpu_cache(&self) -> Result<(usize, usize)> {
    529
    0
            let config = &self.model.config;
    530
    0
            let hidden_dim = config.hidden_dim;
    531
    0
            let intermediate_dim = config.intermediate_dim;
    532
    0
            let num_layers = self.model.layers.len();
    533
    534
            // Create cache with model dimensions
    535
    0
            let cache = DequantizedWeightCache::new(hidden_dim, intermediate_dim, num_layers);
    536
    537
            // Dequantize each layer's FFN weights
    538
            // Note: warmup closure can't return Result, so we use unwrap_or_default
    539
            // for robustness. In production, use warmup_gpu_cache_checked() for error handling.
    540
    0
            cache.warmup(|layer_idx| {
    541
    0
                let layer = &self.model.layers[layer_idx];
    542
    543
                // Dequantize using model's dequantize_weight method
    544
    0
                let up = self
    545
    0
                    .model
    546
    0
                    .dequantize_weight(&layer.ffn_up_weight)
    547
    0
                    .unwrap_or_default();
    548
    0
                let down = self
    549
    0
                    .model
    550
    0
                    .dequantize_weight(&layer.ffn_down_weight)
    551
    0
                    .unwrap_or_default();
    552
    553
    0
                (up, down)
    554
    0
            });
    555
    556
    0
            let memory_bytes = cache.memory_bytes();
    557
    0
            let cached_count = cache.cached_count();
    558
    559
            // Store in the cache field
    560
    0
            let mut cache_guard =
    561
    0
                self.dequant_cache
    562
    0
                    .write()
    563
    0
                    .map_err(|_| RealizarError::UnsupportedOperation {
    564
    0
                        operation: "warmup_gpu_cache".to_string(),
    565
    0
                        reason: "Cache lock poisoned".to_string(),
    566
    0
                    })?;
    567
    0
            *cache_guard = Some(cache);
    568
    569
    0
            Ok((memory_bytes, cached_count))
    570
    0
        }
    571
    572
        /// Check if GPU cache is warmed up
    573
    0
        pub fn is_gpu_cache_warm(&self) -> bool {
    574
    0
            self.dequant_cache
    575
    0
                .read()
    576
    0
                .map(|guard| guard.is_some())
    577
    0
                .unwrap_or(false)
    578
    0
        }
    579
    580
        /// Get GPU cache memory usage in bytes
    581
    0
        pub fn gpu_cache_memory(&self) -> usize {
    582
    0
            self.dequant_cache
    583
    0
                .read()
    584
    0
                .ok()
    585
    0
                .and_then(|guard| guard.as_ref().map(DequantizedWeightCache::memory_bytes))
    586
    0
                .unwrap_or(0)
    587
    0
        }
    588
    589
        /// Get dequantized weights for a layer (for GPU batch FFN)
    590
        ///
    591
        /// Returns None if cache not warmed up or layer not found.
    592
    0
        pub fn get_dequantized_ffn_weights(&self, layer_idx: usize) -> Option<DequantizedFFNWeights> {
    593
    0
            self.dequant_cache
    594
    0
                .read()
    595
    0
                .ok()
    596
    0
                .and_then(|guard| guard.as_ref().and_then(|c| c.get(layer_idx)))
    597
    0
        }
    598
    599
        /// Batch FFN forward pass using GPU (PARITY-019)
    600
        ///
    601
        /// Processes multiple tokens in parallel using GPU GEMM.
    602
        /// Requires cache to be warmed up via `warmup_gpu_cache()`.
    603
        ///
    604
        /// # Arguments
    605
        /// * `hidden_states` - Input tensor [batch_size × hidden_dim]
    606
        /// * `layer_idx` - Layer index for weight lookup
    607
        ///
    608
        /// # Returns
    609
        /// Output tensor [batch_size × hidden_dim]
    610
        ///
    611
        /// # Errors
    612
        /// Returns error if cache not warmed or GPU operations fail
    613
        /// PARITY-103: Batch FFN using CUDA when available
    614
        ///
    615
        /// Uses CudaScheduler first (no buffer limits), falls back to HybridScheduler (wgpu).
    616
        /// This bypasses the wgpu 256MB buffer limit that was blocking GPU batch inference.
    617
    0
        pub fn batch_ffn_gpu(&self, hidden_states: &[f32], layer_idx: usize) -> Result<Vec<f32>> {
    618
    0
            let config = &self.model.config;
    619
    0
            let hidden_dim = config.hidden_dim;
    620
    0
            let intermediate_dim = config.intermediate_dim;
    621
    0
            let batch_size = hidden_states.len() / hidden_dim;
    622
    623
    0
            if batch_size == 0 {
    624
    0
                return Err(RealizarError::UnsupportedOperation {
    625
    0
                    operation: "batch_ffn_gpu".to_string(),
    626
    0
                    reason: "Empty batch".to_string(),
    627
    0
                });
    628
    0
            }
    629
    630
            // Get cached weights
    631
    0
            let weights = self.get_dequantized_ffn_weights(layer_idx).ok_or_else(|| {
    632
    0
                RealizarError::UnsupportedOperation {
    633
    0
                    operation: "batch_ffn_gpu".to_string(),
    634
    0
                    reason: format!(
    635
    0
                        "Layer {} not cached. Call warmup_gpu_cache() first.",
    636
    0
                        layer_idx
    637
    0
                    ),
    638
    0
                }
    639
    0
            })?;
    640
    641
            // PARITY-103: Up projection preferring CUDA
    642
    0
            let mut intermediate = self.batch_matmul_gpu_prefer_cuda(
    643
    0
                hidden_states,
    644
    0
                &weights.up,
    645
    0
                batch_size,
    646
    0
                hidden_dim,
    647
    0
                intermediate_dim,
    648
    0
            )?;
    649
    650
            // Add up bias if present
    651
    0
            if let Some(ref bias) = weights.up_bias {
    652
    0
                for b in 0..batch_size {
    653
    0
                    for i in 0..intermediate_dim {
    654
    0
                        intermediate[b * intermediate_dim + i] += bias[i];
    655
    0
                    }
    656
                }
    657
    0
            }
    658
    659
            // GELU activation (CPU - fused in future)
    660
    0
            for x in &mut intermediate {
    661
    0
                let x64 = *x as f64;
    662
    0
                *x = (x64
    663
    0
                    * 0.5
    664
    0
                    * (1.0 + (x64 * 0.797_884_560_8 * (1.0 + 0.044_715 * x64 * x64)).tanh()))
    665
    0
                    as f32;
    666
    0
            }
    667
    668
            // PARITY-103: Down projection preferring CUDA
    669
    0
            let mut output = self.batch_matmul_gpu_prefer_cuda(
    670
    0
                &intermediate,
    671
    0
                &weights.down,
    672
    0
                batch_size,
    673
    0
                intermediate_dim,
    674
    0
                hidden_dim,
    675
    0
            )?;
    676
    677
            // Add down bias if present
    678
    0
            if let Some(ref bias) = weights.down_bias {
    679
    0
                for b in 0..batch_size {
    680
    0
                    for i in 0..hidden_dim {
    681
    0
                        output[b * hidden_dim + i] += bias[i];
    682
    0
                    }
    683
                }
    684
    0
            }
    685
    686
    0
            Ok(output)
    687
    0
        }
    688
    689
        /// PARITY-103: Batch QKV projection using CUDA when available
    690
        ///
    691
        /// Projects hidden states to Q, K, V for all requests in batch.
    692
        /// [batch, hidden] @ [hidden, 3*hidden] = [batch, 3*hidden]
    693
        ///
    694
        /// Uses CudaScheduler first (no buffer limits), falls back to HybridScheduler (wgpu).
    695
        ///
    696
        /// # Arguments
    697
        /// * `hidden_states` - Flattened hidden states [batch * hidden_dim]
    698
        /// * `layer_idx` - Layer index for weight lookup
    699
        ///
    700
        /// # Returns
    701
        /// Flattened QKV projections [batch * 3 * hidden_dim]
    702
        #[cfg(feature = "gpu")]
    703
    0
        pub fn batch_qkv_projection_gpu(
    704
    0
            &self,
    705
    0
            hidden_states: &[f32],
    706
    0
            layer_idx: usize,
    707
    0
        ) -> Result<Vec<f32>> {
    708
    0
            let hidden_dim = self.model.config.hidden_dim;
    709
    0
            let batch_size = hidden_states.len() / hidden_dim;
    710
    0
            let qkv_dim = 3 * hidden_dim;
    711
    712
    0
            if batch_size == 0 {
    713
    0
                return Ok(Vec::new());
    714
    0
            }
    715
    716
    0
            let layer = &self.model.layers[layer_idx];
    717
    718
            // Dequantize QKV weight for GPU GEMM
    719
    0
            let qkv_weight = self.model.dequantize_qkv(&layer.qkv_weight)?;
    720
    721
            // PARITY-103: QKV projection preferring CUDA
    722
    0
            let mut qkv = self.batch_matmul_gpu_prefer_cuda(
    723
    0
                hidden_states,
    724
    0
                &qkv_weight,
    725
    0
                batch_size,
    726
    0
                hidden_dim,
    727
    0
                qkv_dim,
    728
    0
            )?;
    729
    730
            // Add bias if present
    731
    0
            if let Some(ref bias) = layer.qkv_bias {
    732
    0
                for b in 0..batch_size {
    733
    0
                    for i in 0..qkv_dim {
    734
    0
                        qkv[b * qkv_dim + i] += bias[i];
    735
    0
                    }
    736
                }
    737
    0
            }
    738
    739
    0
            Ok(qkv)
    740
    0
        }
    741
    742
        /// Batch attention output projection using GPU GEMM (PARITY-024)
    743
        ///
    744
        /// Projects attention outputs for all requests in batch.
    745
        /// [batch, hidden] @ [hidden, hidden] = [batch, hidden]
    746
        ///
    747
        /// # Arguments
    748
        /// * `attention_outputs` - Flattened attention outputs [batch * hidden_dim]
    749
        /// * `layer_idx` - Layer index for weight lookup
    750
        ///
    751
        /// # Returns
    752
        /// Flattened projected outputs [batch * hidden_dim]
    753
        #[cfg(feature = "gpu")]
    754
    0
        pub fn batch_attention_output_gpu(
    755
    0
            &self,
    756
    0
            attention_outputs: &[f32],
    757
    0
            layer_idx: usize,
    758
    0
        ) -> Result<Vec<f32>> {
    759
    0
            let hidden_dim = self.model.config.hidden_dim;
    760
    0
            let batch_size = attention_outputs.len() / hidden_dim;
    761
    762
    0
            if batch_size == 0 {
    763
    0
                return Ok(Vec::new());
    764
    0
            }
    765
    766
    0
            let layer = &self.model.layers[layer_idx];
    767
    768
            // Dequantize output weight for GPU GEMM
    769
    0
            let output_weight = self.model.dequantize_weight(&layer.attn_output_weight)?;
    770
    771
            // PARITY-103: Output projection preferring CUDA (bypasses wgpu 256MB limit)
    772
            // [batch, hidden] @ [hidden, hidden] = [batch, hidden]
    773
    0
            let mut output = self.batch_matmul_gpu_prefer_cuda(
    774
    0
                attention_outputs,
    775
    0
                &output_weight,
    776
    0
                batch_size,
    777
    0
                hidden_dim,
    778
    0
                hidden_dim,
    779
    0
            )?;
    780
    781
            // Add bias if present
    782
    0
            if let Some(ref bias) = layer.attn_output_bias {
    783
    0
                for b in 0..batch_size {
    784
    0
                    for i in 0..hidden_dim {
    785
    0
                        output[b * hidden_dim + i] += bias[i];
    786
    0
                    }
    787
                }
    788
    0
            }
    789
    790
    0
            Ok(output)
    791
    0
        }
    792
    793
        /// Batch LM head projection using GPU GEMM (PARITY-025)
    794
        ///
    795
        /// Projects hidden states to vocabulary logits for all requests in batch.
    796
        /// [batch, hidden] @ [hidden, vocab] = [batch, vocab]
    797
        ///
    798
        /// # Arguments
    799
        /// * `hidden_states` - Flattened normalized hidden states [batch * hidden_dim]
    800
        ///
    801
        /// # Returns
    802
        /// Flattened logits [batch * vocab_size]
    803
        #[cfg(feature = "gpu")]
    804
    0
        pub fn batch_lm_head_gpu(&self, hidden_states: &[f32]) -> Result<Vec<f32>> {
    805
    0
            let hidden_dim = self.model.config.hidden_dim;
    806
    0
            let vocab_size = self.model.config.vocab_size;
    807
    0
            let batch_size = hidden_states.len() / hidden_dim;
    808
    809
    0
            if batch_size == 0 {
    810
    0
                return Ok(Vec::new());
    811
    0
            }
    812
    813
            // Dequantize LM head weight for GPU GEMM
    814
    0
            let lm_head_weight = self.model.dequantize_weight(&self.model.lm_head_weight)?;
    815
    816
            // PARITY-103: LM head projection preferring CUDA (bypasses wgpu 256MB limit)
    817
            // [batch, hidden] @ [hidden, vocab] = [batch, vocab]
    818
    0
            let mut logits = self.batch_matmul_gpu_prefer_cuda(
    819
    0
                hidden_states,
    820
    0
                &lm_head_weight,
    821
    0
                batch_size,
    822
    0
                hidden_dim,
    823
    0
                vocab_size,
    824
    0
            )?;
    825
    826
            // Add bias if present
    827
    0
            if let Some(ref bias) = self.model.lm_head_bias {
    828
    0
                for b in 0..batch_size {
    829
    0
                    for i in 0..vocab_size {
    830
    0
                        logits[b * vocab_size + i] += bias[i];
    831
    0
                    }
    832
                }
    833
    0
            }
    834
    835
    0
            Ok(logits)
    836
    0
        }
    837
    838
        /// Batch generation with GPU-accelerated FFN (PARITY-020)
    839
        ///
    840
        /// Processes multiple prompts in parallel using GPU batch operations.
    841
        /// The key optimization is converting MATVEC (single token) to GEMM (batch tokens).
    842
        ///
    843
        /// # Architecture
    844
        /// - Attention: CPU with KV cache (MATVEC is faster on CPU)
    845
        /// - FFN: GPU with batch GEMM (batch_size ≥ 32 uses GPU)
    846
        /// - Sampling: CPU (negligible compared to matmul)
    847
        ///
    848
        /// # Arguments
    849
        /// * `prompts` - Multiple prompts to process in parallel [num_prompts][seq_len]
    850
        /// * `config` - Generation configuration (shared across all prompts)
    851
        ///
    852
        /// # Returns
    853
        /// Generated sequences for each prompt [num_prompts][generated_len]
    854
        ///
    855
        /// # Errors
    856
        /// Returns error if GPU cache not warmed up or generation fails
    857
        ///
    858
        /// # Performance
    859
        /// - Single prompt: ~5 tok/s (CPU-bound, no batching benefit)
    860
        /// - 32 prompts: ~150 tok/s total (~4.7 tok/s per prompt)
    861
        /// - 64 prompts: ~280 tok/s total (~4.4 tok/s per prompt, memory-bound)
    862
    0
        pub fn batch_generate_gpu(
    863
    0
            &self,
    864
    0
            prompts: &[Vec<u32>],
    865
    0
            config: &QuantizedGenerateConfig,
    866
    0
        ) -> Result<Vec<Vec<u32>>> {
    867
    0
            if prompts.is_empty() {
    868
    0
                return Ok(Vec::new());
    869
    0
            }
    870
    871
            // Verify GPU cache is warmed up
    872
    0
            if !self.is_gpu_cache_warm() {
    873
    0
                return Err(RealizarError::UnsupportedOperation {
    874
    0
                    operation: "batch_generate_gpu".to_string(),
    875
    0
                    reason: "GPU cache not warmed up. Call warmup_gpu_cache() first.".to_string(),
    876
    0
                });
    877
    0
            }
    878
    879
    0
            let num_prompts = prompts.len();
    880
    0
            let max_seq_len = prompts.iter().map(Vec::len).max().unwrap_or(0) + config.max_tokens;
    881
    882
            // Initialize KV caches for each prompt
    883
    0
            let mut caches: Vec<OwnedQuantizedKVCache> = prompts
    884
    0
                .iter()
    885
    0
                .map(|_| OwnedQuantizedKVCache::from_config(&self.model.config, max_seq_len))
    886
    0
                .collect();
    887
    888
            // Initialize token sequences (copy prompts)
    889
    0
            let mut sequences: Vec<Vec<u32>> = prompts.to_vec();
    890
    891
            // Track generation progress per prompt
    892
    0
            let mut done: Vec<bool> = vec![false; num_prompts];
    893
    894
            // PARITY-097: Parallel prefill across prompts using rayon
    895
            // Each prompt's prefill is independent (different KV cache)
    896
            // Model is shared immutably (&self), caches are mutated independently
    897
            use rayon::prelude::*;
    898
    899
    0
            caches
    900
    0
                .par_iter_mut()
    901
    0
                .zip(prompts.par_iter())
    902
    0
                .try_for_each(|(cache, prompt)| {
    903
    0
                    for (pos, &token_id) in prompt.iter().enumerate() {
    904
    0
                        self.model.forward_single_with_cache(token_id, cache, pos)?;
    905
                    }
    906
    0
                    Ok::<_, RealizarError>(())
    907
    0
                })?;
    908
    909
            // Generation loop with batched FFN (PARITY-021: GPU optimization)
    910
    0
            for gen_idx in 0..config.max_tokens {
    911
                // Collect active prompts for this generation step
    912
    0
                let active_indices: Vec<usize> = (0..num_prompts).filter(|&i| !done[i]).collect();
    913
    914
    0
                if active_indices.is_empty() {
    915
    0
                    break;
    916
    0
                }
    917
    918
    0
                let active_count = active_indices.len();
    919
    920
                // Use batched forward when we have enough active prompts for GPU benefit
    921
                // GPU batch threshold is 32 (from IMP-600 analysis)
    922
                const GPU_BATCH_THRESHOLD: usize = 32;
    923
    924
    0
                if active_count >= GPU_BATCH_THRESHOLD {
    925
                    // PARITY-021: Batched forward with GPU FFN
    926
                    // Collect tokens, positions, and cache slices for active prompts
    927
    0
                    let batch_tokens: Vec<u32> = active_indices
    928
    0
                        .iter()
    929
    0
                        .map(|&idx| {
    930
    0
                            *sequences[idx]
    931
    0
                                .last()
    932
    0
                                .expect("sequence must have at least prompt tokens")
    933
    0
                        })
    934
    0
                        .collect();
    935
    936
    0
                    let batch_positions: Vec<usize> = active_indices
    937
    0
                        .iter()
    938
    0
                        .map(|&idx| prompts[idx].len() + gen_idx)
    939
    0
                        .collect();
    940
    941
                    // PARITY-096: Extract caches without cloning using std::mem::take
    942
                    // This avoids expensive cache cloning on every generation step
    943
    0
                    let mut batch_caches: Vec<OwnedQuantizedKVCache> = active_indices
    944
    0
                        .iter()
    945
    0
                        .map(|&idx| std::mem::take(&mut caches[idx]))
    946
    0
                        .collect();
    947
    948
                    // Forward batch with GPU FFN
    949
    0
                    let all_logits = self.forward_batch_with_gpu_ffn(
    950
    0
                        &batch_tokens,
    951
    0
                        &mut batch_caches,
    952
    0
                        &batch_positions,
    953
    0
                    )?;
    954
    955
                    // PARITY-096: Put caches back (move, not clone)
    956
    0
                    for (i, &idx) in active_indices.iter().enumerate() {
    957
    0
                        caches[idx] = std::mem::take(&mut batch_caches[i]);
    958
    0
                    }
    959
    960
                    // Sample and update sequences
    961
    0
                    for (i, &prompt_idx) in active_indices.iter().enumerate() {
    962
    0
                        let logits = &all_logits[i];
    963
    0
                        let next_token = if config.temperature == 0.0 || config.top_k == 1 {
    964
    0
                            OwnedQuantizedModel::argmax(logits)
    965
                        } else {
    966
    0
                            OwnedQuantizedModel::sample_topk(logits, config.temperature, config.top_k)
    967
                        };
    968
    969
    0
                        if config.stop_tokens.contains(&next_token) {
    970
    0
                            done[prompt_idx] = true;
    971
    0
                        } else {
    972
    0
                            sequences[prompt_idx].push(next_token);
    973
    0
                            if sequences[prompt_idx].len() >= max_seq_len {
    974
    0
                                done[prompt_idx] = true;
    975
    0
                            }
    976
                        }
    977
                    }
    978
                } else {
    979
                    // Sequential forward for small batches (CPU is faster)
    980
    0
                    for &prompt_idx in &active_indices {
    981
    0
                        let position = prompts[prompt_idx].len() + gen_idx;
    982
    0
                        let last_token = *sequences[prompt_idx]
    983
    0
                            .last()
    984
    0
                            .expect("sequence must have at least prompt tokens");
    985
    986
    0
                        let logits = self.model.forward_single_with_cache(
    987
    0
                            last_token,
    988
    0
                            &mut caches[prompt_idx],
    989
    0
                            position,
    990
    0
                        )?;
    991
    992
    0
                        let next_token = if config.temperature == 0.0 || config.top_k == 1 {
    993
    0
                            OwnedQuantizedModel::argmax(&logits)
    994
                        } else {
    995
    0
                            OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
    996
                        };
    997
    998
    0
                        if config.stop_tokens.contains(&next_token) {
    999
    0
                            done[prompt_idx] = true;
    1000
    0
                        } else {
    1001
    0
                            sequences[prompt_idx].push(next_token);
    1002
    0
                            if sequences[prompt_idx].len() >= max_seq_len {
    1003
    0
                                done[prompt_idx] = true;
    1004
    0
                            }
    1005
                        }
    1006
                    }
    1007
                }
    1008
            }
    1009
    1010
    0
            Ok(sequences)
    1011
    0
        }
    1012
    1013
        /// Batched forward pass with GPU FFN optimization (PARITY-021)
    1014
        ///
    1015
        /// Processes multiple tokens in parallel with GPU-accelerated FFN.
    1016
        /// Attention is still per-token with CPU KV cache, but FFN uses GPU GEMM.
    1017
        ///
    1018
        /// # Arguments
    1019
        /// * `token_ids` - Token IDs for each prompt [batch_size]
    1020
        /// * `caches` - Per-prompt KV caches
    1021
        /// * `positions` - Position for each prompt [batch_size]
    1022
        ///
    1023
        /// # Returns
    1024
        /// Logits for each prompt [batch_size][vocab_size]
    1025
        ///
    1026
        /// # GPU Dispatch
    1027
        /// - batch_size >= 32: GPU GEMM for FFN (10x speedup)
    1028
        /// - batch_size < 32: CPU fallback
    1029
    0
        pub fn forward_batch_with_gpu_ffn(
    1030
    0
            &self,
    1031
    0
            token_ids: &[u32],
    1032
    0
            caches: &mut [OwnedQuantizedKVCache],
    1033
    0
            positions: &[usize],
    1034
    0
        ) -> Result<Vec<Vec<f32>>> {
    1035
    0
            let batch_size = token_ids.len();
    1036
    0
            if batch_size == 0 {
    1037
    0
                return Ok(Vec::new());
    1038
    0
            }
    1039
    0
            if batch_size != caches.len() || batch_size != positions.len() {
    1040
    0
                return Err(RealizarError::InvalidShape {
    1041
    0
                    reason: format!(
    1042
    0
                        "Batch size mismatch: tokens={}, caches={}, positions={}",
    1043
    0
                        batch_size,
    1044
    0
                        caches.len(),
    1045
    0
                        positions.len()
    1046
    0
                    ),
    1047
    0
                });
    1048
    0
            }
    1049
    1050
    0
            let hidden_dim = self.model.config.hidden_dim;
    1051
    0
            let num_layers = self.model.layers.len();
    1052
    1053
            // Threshold for GPU dispatch (based on IMP-600 analysis)
    1054
            const GPU_BATCH_THRESHOLD: usize = 32;
    1055
    0
            let use_gpu = batch_size >= GPU_BATCH_THRESHOLD && self.is_gpu_cache_warm();
    1056
    1057
            // PARITY-098: Parallel embedding using rayon
    1058
            use rayon::prelude::*;
    1059
    0
            let mut hidden_states: Vec<Vec<f32>> = token_ids
    1060
    0
                .par_iter()
    1061
    0
                .map(|&tid| self.model.embed(&[tid]))
    1062
    0
                .collect();
    1063
    1064
            // 2. Process through transformer layers
    1065
    0
            for layer_idx in 0..num_layers {
    1066
    0
                let layer = &self.model.layers[layer_idx];
    1067
    1068
                // PARITY-024: GPU batch attention path vs CPU sequential path
    1069
    0
                if use_gpu {
    1070
                    // GPU path: batch QKV projection, per-prompt attention, batch output projection
    1071
    1072
                    // 2a. PARITY-098: Parallel batch layer norm
    1073
    0
                    let normed_batch: Vec<Vec<f32>> = hidden_states
    1074
    0
                        .par_iter()
    1075
    0
                        .map(|hidden| {
    1076
    0
                            self.model.layer_norm(
    1077
    0
                                hidden,
    1078
    0
                                &layer.attn_norm_weight,
    1079
    0
                                layer.attn_norm_bias.as_deref(),
    1080
    0
                                self.model.config.eps,
    1081
                            )
    1082
    0
                        })
    1083
    0
                        .collect();
    1084
    1085
                    // 2b. Batch QKV projection using GPU GEMM (PARITY-024)
    1086
    0
                    let batch_normed: Vec<f32> = normed_batch.iter().flatten().copied().collect();
    1087
    0
                    let batch_qkv = self.batch_qkv_projection_gpu(&batch_normed, layer_idx)?;
    1088
    1089
                    // 2c-2e. PARITY-099: Parallel attention computation per prompt
    1090
                    // Each prompt has its own KV cache, so we can parallelize
    1091
    0
                    let qkv_dim = 3 * hidden_dim;
    1092
    1093
    0
                    let attention_outputs: Vec<Vec<f32>> = caches
    1094
    0
                        .par_iter_mut()
    1095
    0
                        .enumerate()
    1096
    0
                        .map(|(prompt_idx, cache)| {
    1097
    0
                            let qkv_start = prompt_idx * qkv_dim;
    1098
    0
                            let qkv = &batch_qkv[qkv_start..qkv_start + qkv_dim];
    1099
    1100
                            // Extract Q, K, V
    1101
    0
                            let mut q = qkv[0..hidden_dim].to_vec();
    1102
    0
                            let mut k = qkv[hidden_dim..2 * hidden_dim].to_vec();
    1103
    0
                            let v = qkv[2 * hidden_dim..3 * hidden_dim].to_vec();
    1104
    1105
                            // Apply RoPE (position-dependent, must be per-prompt)
    1106
                            // Note: Uses num_heads for both (non-GQA code path)
    1107
    0
                            self.model.apply_rope(
    1108
    0
                                &mut q,
    1109
    0
                                positions[prompt_idx],
    1110
    0
                                self.model.config.num_heads,
    1111
                            );
    1112
    0
                            self.model.apply_rope(
    1113
    0
                                &mut k,
    1114
    0
                                positions[prompt_idx],
    1115
    0
                                self.model.config.num_heads,
    1116
                            );
    1117
    1118
                            // Attention with KV cache (must be per-prompt, different caches)
    1119
                            // PARITY-027: Use FlashAttention for long sequences (O(N) memory)
    1120
    0
                            let k_cache = cache.get_k(layer_idx);
    1121
    0
                            let v_cache = cache.get_v(layer_idx);
    1122
    1123
                            // FlashAttention threshold: use for sequences >= 512 tokens
    1124
                            const FLASH_ATTENTION_THRESHOLD: usize = 512;
    1125
    0
                            let cache_len = k_cache.len() / hidden_dim;
    1126
    0
                            let use_flash_attention = cache_len >= FLASH_ATTENTION_THRESHOLD;
    1127
    1128
    0
                            let attn_out = if k_cache.is_empty() {
    1129
    0
                                v.clone()
    1130
    0
                            } else if use_flash_attention {
    1131
                                // FlashAttention: O(N) memory, tiled computation
    1132
                                const FLASH_BLOCK_SIZE: usize = 64;
    1133
    0
                                self.model.flash_attention_tiled(
    1134
    0
                                    &q,
    1135
    0
                                    k_cache,
    1136
    0
                                    v_cache,
    1137
    0
                                    &k,
    1138
    0
                                    &v,
    1139
                                    FLASH_BLOCK_SIZE,
    1140
                                )
    1141
                            } else {
    1142
                                // Standard attention: O(N²) memory but faster for short sequences
    1143
    0
                                self.model
    1144
    0
                                    .attention_with_cache(&q, k_cache, v_cache, &k, &v)
    1145
                            };
    1146
    1147
                            // Store K and V in cache
    1148
    0
                            cache.append(layer_idx, &k, &v);
    1149
    0
                            attn_out
    1150
    0
                        })
    1151
    0
                        .collect();
    1152
    1153
                    // 2f. Batch attention output projection using GPU GEMM (PARITY-024)
    1154
    0
                    let batch_attn: Vec<f32> = attention_outputs.iter().flatten().copied().collect();
    1155
    0
                    let batch_output = self.batch_attention_output_gpu(&batch_attn, layer_idx)?;
    1156
    1157
                    // 2g. PARITY-100: Parallel residual connection
    1158
    0
                    hidden_states
    1159
    0
                        .par_iter_mut()
    1160
    0
                        .enumerate()
    1161
    0
                        .for_each(|(prompt_idx, hidden)| {
    1162
    0
                            let start = prompt_idx * hidden_dim;
    1163
    0
                            for i in 0..hidden_dim {
    1164
    0
                                hidden[i] += batch_output[start + i];
    1165
    0
                            }
    1166
    0
                        });
    1167
                } else {
    1168
                    // CPU sequential path (original implementation)
    1169
    0
                    for (prompt_idx, hidden) in hidden_states.iter_mut().enumerate() {
    1170
                        // Attention layer norm
    1171
    0
                        let normed = self.model.layer_norm(
    1172
    0
                            hidden,
    1173
    0
                            &layer.attn_norm_weight,
    1174
    0
                            layer.attn_norm_bias.as_deref(),
    1175
    0
                            self.model.config.eps,
    1176
                        );
    1177
    1178
                        // QKV projection
    1179
    0
                        let mut qkv = self.model.qkv_matmul(&normed, &layer.qkv_weight)?;
    1180
    0
                        if let Some(ref bias) = layer.qkv_bias {
    1181
    0
                            self.model.add_bias(&mut qkv, bias);
    1182
    0
                        }
    1183
    1184
                        // Extract Q, K, V and apply RoPE
    1185
                        // Note: Uses num_heads for both (non-GQA code path)
    1186
    0
                        let mut q = qkv[0..hidden_dim].to_vec();
    1187
    0
                        let mut k = qkv[hidden_dim..2 * hidden_dim].to_vec();
    1188
    0
                        let v = qkv[2 * hidden_dim..3 * hidden_dim].to_vec();
    1189
    1190
    0
                        self.model.apply_rope(
    1191
    0
                            &mut q,
    1192
    0
                            positions[prompt_idx],
    1193
    0
                            self.model.config.num_heads,
    1194
                        );
    1195
    0
                        self.model.apply_rope(
    1196
    0
                            &mut k,
    1197
    0
                            positions[prompt_idx],
    1198
    0
                            self.model.config.num_heads,
    1199
                        );
    1200
    1201
                        // Get cached K/V and compute attention
    1202
    0
                        let k_cache = caches[prompt_idx].get_k(layer_idx);
    1203
    0
                        let v_cache = caches[prompt_idx].get_v(layer_idx);
    1204
    1205
    0
                        let attn_out = if k_cache.is_empty() {
    1206
    0
                            v.clone()
    1207
                        } else {
    1208
    0
                            self.model
    1209
    0
                                .attention_with_cache(&q, k_cache, v_cache, &k, &v)
    1210
                        };
    1211
    1212
                        // Store K and V in cache
    1213
    0
                        caches[prompt_idx].append(layer_idx, &k, &v);
    1214
    1215
                        // Attention output projection
    1216
    0
                        let mut attn_output = self
    1217
    0
                            .model
    1218
    0
                            .fused_matmul(&attn_out, &layer.attn_output_weight)?;
    1219
    0
                        if let Some(ref bias) = layer.attn_output_bias {
    1220
    0
                            self.model.add_bias(&mut attn_output, bias);
    1221
    0
                        }
    1222
    1223
                        // Residual connection
    1224
    0
                        for i in 0..hidden_dim {
    1225
    0
                            hidden[i] += attn_output[i];
    1226
    0
                        }
    1227
                    }
    1228
                }
    1229
    1230
                // 2h. FFN - GPU batch or CPU sequential
    1231
    0
                if use_gpu {
    1232
                    // GPU batch FFN: collect hidden states, process together, scatter back
    1233
    0
                    let batch_hidden: Vec<f32> = hidden_states.iter().flatten().copied().collect();
    1234
    0
                    let ffn_output = self.batch_ffn_gpu(&batch_hidden, layer_idx)?;
    1235
    1236
                    // PARITY-100: Parallel scatter and residual
    1237
    0
                    hidden_states
    1238
    0
                        .par_iter_mut()
    1239
    0
                        .enumerate()
    1240
    0
                        .for_each(|(prompt_idx, hidden)| {
    1241
    0
                            let start = prompt_idx * hidden_dim;
    1242
    0
                            for i in 0..hidden_dim {
    1243
    0
                                hidden[i] += ffn_output[start + i];
    1244
    0
                            }
    1245
    0
                        });
    1246
                } else {
    1247
                    // CPU sequential FFN
    1248
    0
                    for hidden in &mut hidden_states {
    1249
    0
                        let mut ffn_hidden = self.model.fused_matmul(hidden, &layer.ffn_up_weight)?;
    1250
    0
                        if let Some(ref bias) = layer.ffn_up_bias {
    1251
    0
                            self.model.add_bias(&mut ffn_hidden, bias);
    1252
    0
                        }
    1253
    0
                        self.model.gelu(&mut ffn_hidden);
    1254
    1255
    0
                        let mut ffn_output = self
    1256
    0
                            .model
    1257
    0
                            .fused_matmul(&ffn_hidden, &layer.ffn_down_weight)?;
    1258
    0
                        if let Some(ref bias) = layer.ffn_down_bias {
    1259
    0
                            self.model.add_bias(&mut ffn_output, bias);
    1260
    0
                        }
    1261
    1262
                        // Residual
    1263
    0
                        for i in 0..hidden_dim {
    1264
    0
                            hidden[i] += ffn_output[i];
    1265
    0
                        }
    1266
                    }
    1267
                }
    1268
            }
    1269
    1270
            // PARITY-100: Parallel cache advance
    1271
    0
            caches.par_iter_mut().for_each(|cache| {
    1272
    0
                cache.advance();
    1273
    0
            });
    1274
    1275
            // 3. Final layer norm and LM head for each prompt
    1276
            // PARITY-025: Use GPU batch LM head when batch >= threshold
    1277
    0
            let vocab_size = self.model.config.vocab_size;
    1278
    1279
    0
            let all_logits: Vec<Vec<f32>> = if use_gpu {
    1280
                // GPU path: batch layer norm and LM head projection
    1281
    1282
                // 3a. PARITY-098: Parallel final layer norm
    1283
    0
                let normed_batch: Vec<Vec<f32>> = hidden_states
    1284
    0
                    .par_iter()
    1285
    0
                    .map(|hidden| {
    1286
    0
                        self.model.layer_norm(
    1287
    0
                            hidden,
    1288
    0
                            &self.model.output_norm_weight,
    1289
    0
                            self.model.output_norm_bias.as_deref(),
    1290
    0
                            self.model.config.eps,
    1291
                        )
    1292
    0
                    })
    1293
    0
                    .collect();
    1294
    1295
                // 3b. Batch LM head projection using GPU GEMM (PARITY-025)
    1296
    0
                let batch_normed: Vec<f32> = normed_batch.iter().flatten().copied().collect();
    1297
    0
                let batch_logits = self.batch_lm_head_gpu(&batch_normed)?;
    1298
    1299
                // 3c. PARITY-098: Parallel scatter logits back to per-prompt vectors
    1300
    0
                (0..batch_size)
    1301
    0
                    .into_par_iter()
    1302
    0
                    .map(|i| {
    1303
    0
                        let start = i * vocab_size;
    1304
    0
                        batch_logits[start..start + vocab_size].to_vec()
    1305
    0
                    })
    1306
    0
                    .collect()
    1307
            } else {
    1308
                // CPU path: sequential per-prompt processing
    1309
    0
                let mut result = Vec::with_capacity(batch_size);
    1310
    0
                for hidden in &hidden_states {
    1311
    0
                    let normed = self.model.layer_norm(
    1312
    0
                        hidden,
    1313
    0
                        &self.model.output_norm_weight,
    1314
    0
                        self.model.output_norm_bias.as_deref(),
    1315
    0
                        self.model.config.eps,
    1316
                    );
    1317
    1318
    0
                    let mut logits = self
    1319
    0
                        .model
    1320
    0
                        .fused_matmul(&normed, &self.model.lm_head_weight)?;
    1321
    0
                    if let Some(ref bias) = self.model.lm_head_bias {
    1322
    0
                        self.model.add_bias(&mut logits, bias);
    1323
    0
                    }
    1324
    0
                    result.push(logits);
    1325
                }
    1326
    0
                result
    1327
            };
    1328
    1329
    0
            Ok(all_logits)
    1330
    0
        }
    1331
    1332
        /// Get batch generation statistics
    1333
        ///
    1334
        /// Returns information about the batch processing capabilities.
    1335
    0
        pub fn batch_stats(&self) -> BatchGenerationStats {
    1336
    0
            let is_cached = self.is_gpu_cache_warm();
    1337
    0
            let memory_gb = self.gpu_cache_memory() as f64 / 1_000_000_000.0;
    1338
    0
            let num_layers = self.model.layers.len();
    1339
    0
            let hidden_dim = self.model.config.hidden_dim;
    1340
    0
            let intermediate_dim = self.model.config.intermediate_dim;
    1341
    1342
    0
            BatchGenerationStats {
    1343
    0
                gpu_cache_ready: is_cached,
    1344
    0
                cache_memory_gb: memory_gb,
    1345
    0
                num_layers,
    1346
    0
                hidden_dim,
    1347
    0
                intermediate_dim,
    1348
    0
                recommended_batch_size: 32, // GPU GEMM threshold
    1349
    0
                max_batch_size: 64,         // Memory-limited
    1350
    0
            }
    1351
    0
        }
    1352
    }