## Coverage Report

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

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

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

    Line
    Count
    Source
    1
    //! Batched forward pass variants
    2
    //!
    3
    //! Contains forward_batch, forward_batch_gpu, forward_batch_with_cache,
    4
    //! and supporting batch matmul/attention helpers.
    5
    6
    use crate::error::{RealizarError, Result};
    7
    use crate::gguf::ops;
    8
    use crate::gguf::{
    9
        DispatchMetrics, OwnedQuantizedKVCache, OwnedQuantizedModel, OwnedQuantizedTensor,
    10
        OwnedQKVWeights, QuantizedGenerateConfig, TokenBuffer, GGUF_TYPE_Q4_K, GGUF_TYPE_Q5_K,
    11
        GGUF_TYPE_Q6_K,
    12
    };
    13
    14
    impl OwnedQuantizedModel {
    15
        /// Batched forward pass for prompt prefill (PARITY-002)
    16
        ///
    17
        /// Processes all prompt tokens at once, enabling GPU acceleration
    18
        /// for the attention computation when the batch is large enough.
    19
        ///
    20
        /// # Arguments
    21
        /// * `tokens` - All prompt tokens to process at once
    22
        /// * `cache` - KV cache for storing computed K/V tensors
    23
        /// * `metrics` - Dispatch metrics tracker for CPU/GPU decision recording
    24
        ///
    25
        /// # Returns
    26
        /// Logits for next token prediction (from the last token position)
    27
        ///
    28
        /// # Errors
    29
        /// Returns error if tensor operations fail
    30
        #[cfg(feature = "gpu")]
    31
    5
        pub fn forward_batch_with_cache(
    32
    5
            &self,
    33
    5
            tokens: &[u32],
    34
    5
            cache: &mut OwnedQuantizedKVCache,
    35
    5
            metrics: &std::sync::Arc<DispatchMetrics>,
    36
    5
        ) -> Result<Vec<f32>> {
    37
    5
            if tokens.is_empty() {
    38
    0
                return Err(RealizarError::InvalidShape {
    39
    0
                    reason: "Tokens cannot be empty".to_string(),
    40
    0
                });
    41
    5
            }
    42
    43
    5
            let seq_len = tokens.len();
    44
    5
            let hidden_dim = self.config.hidden_dim;
    45
    46
            // 1. Embed all tokens at once: [seq_len, hidden_dim]
    47
    5
            let mut hidden_states: Vec<Vec<f32>> = tokens
    48
    5
                .iter()
    49
    84
                .map5(|&token_id| self.embed(&[token_id]))
    50
    5
                .collect();
    51
    52
            // 2. Process through transformer layers
    53
    5
            for (layer_idx, layer) in self.layers.iter().enumerate() {
    54
                // Collect Q, K, V for all positions
    55
    5
                let mut all_q: Vec<Vec<f32>> = Vec::with_capacity(seq_len);
    56
    5
                let mut all_k: Vec<Vec<f32>> = Vec::with_capacity(seq_len);
    57
    5
                let mut all_v: Vec<Vec<f32>> = Vec::with_capacity(seq_len);
    58
    59
    84
                for (pos, hidden) in hidden_states.iter()5.enumerate5() {
    60
                    // 2a. Attention layer norm
    61
    84
                    let normed = ops::layer_norm(
    62
    84
                        hidden,
    63
    84
                        &layer.attn_norm_weight,
    64
    84
                        layer.attn_norm_bias.as_deref(),
    65
    84
                        self.config.eps,
    66
                    );
    67
    68
                    // 2b. QKV projection
    69
    84
                    let mut qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0;
    70
    84
                    if let Some(ref bias0) = layer.qkv_bias {
    71
    0
                        ops::add_bias(&mut qkv, bias);
    72
    84
                    }
    73
    74
                    // 2c. Extract Q, K, V and apply RoPE
    75
                    // Note: This uses hidden_dim for all (assumes non-GQA or fused QKV)
    76
    84
                    let mut q = qkv[0..hidden_dim].to_vec();
    77
    84
                    let mut k = qkv[hidden_dim..2 * hidden_dim].to_vec();
    78
    84
                    let v = qkv[2 * hidden_dim..3 * hidden_dim].to_vec();
    79
    80
    84
                    self.apply_rope(&mut q, pos, self.config.num_heads);
    81
    84
                    self.apply_rope(&mut k, pos, self.config.num_heads); // Same as Q for non-GQA
    82
    83
    84
                    all_q.push(q);
    84
    84
                    all_k.push(k);
    85
    84
                    all_v.push(v);
    86
                }
    87
    88
                // 2d. Compute batched attention
    89
                // For PARITY-002: This is where GPU can accelerate!
    90
                // Attention scores: Q @ K^T is [seq_len, seq_len]
    91
    5
                let attn_outputs = self
    92
    5
                    .batched_attention_with_cache(&all_q, &all_k, &all_v, cache, layer_idx, metrics)?0;
    93
    94
                // 2e. Store all K/V in cache
    95
    84
                for (k, v) in all_k.iter()5.zip5(all_v.iter()5) {
    96
    84
                    cache.append(layer_idx, k, v);
    97
    84
                }
    98
    99
                // 2f. Attention output projection + residual
    100
    84
                for (pos, attn_out) in attn_outputs.iter()5.enumerate5() {
    101
    84
                    let mut attn_output = self.fused_matmul(attn_out, &layer.attn_output_weight)?0;
    102
    84
                    if let Some(ref bias0) = layer.attn_output_bias {
    103
    0
                        ops::add_bias(&mut attn_output, bias);
    104
    84
                    }
    105
    106
                    // Residual connection
    107
    17.6k
                    for i in 0..hidden_dim84 {
    108
    17.6k
                        hidden_states[pos][i] += attn_output[i];
    109
    17.6k
                    }
    110
                }
    111
    112
                // 2g. FFN for all positions
    113
    89
                for hidden84 in &mut hidden_states {
    114
    84
                    let mut ffn_hidden = self.fused_matmul(hidden, &layer.ffn_up_weight)?0;
    115
    84
                    if let Some(ref bias0) = layer.ffn_up_bias {
    116
    0
                        ops::add_bias(&mut ffn_hidden, bias);
    117
    84
                    }
    118
    84
                    ops::gelu(&mut ffn_hidden);
    119
    120
    84
                    let mut ffn_output = self.fused_matmul(&ffn_hidden, &layer.ffn_down_weight)?0;
    121
    84
                    if let Some(ref bias0) = layer.ffn_down_bias {
    122
    0
                        ops::add_bias(&mut ffn_output, bias);
    123
    84
                    }
    124
    125
                    // Residual
    126
    17.6k
                    for i in 0..hidden_dim84 {
    127
    17.6k
                        hidden[i] += ffn_output[i];
    128
    17.6k
                    }
    129
                }
    130
            }
    131
    132
            // Advance cache position for all processed tokens
    133
    84
            for _ in 0..seq_len5 {
    134
    84
                cache.advance();
    135
    84
            }
    136
    137
            // 3. Final layer norm and LM head for LAST token only
    138
    5
            let last_hidden = &hidden_states[seq_len - 1];
    139
    5
            let normed = ops::layer_norm(
    140
    5
                last_hidden,
    141
    5
                &self.output_norm_weight,
    142
    5
                self.output_norm_bias.as_deref(),
    143
    5
                self.config.eps,
    144
            );
    145
    146
            // 4. LM head projection
    147
    5
            let mut logits = self.fused_matmul(&normed, &self.lm_head_weight)?0;
    148
    5
            if let Some(ref bias0) = self.lm_head_bias {
    149
    0
                ops::add_bias(&mut logits, bias);
    150
    5
            }
    151
    152
    5
            Ok(logits)
    153
    5
        }
    154
    155
        /// Batched attention computation with GPU acceleration (PARITY-002)
    156
        ///
    157
        /// Computes attention for all positions at once, enabling GPU dispatch
    158
        /// when the workload (seq_len * hidden_dim * seq_len) exceeds the threshold.
    159
        ///
    160
        /// KEY OPTIMIZATION: Uses GPU matmul for Q @ K^T when workload is large enough.
    161
        /// This is the critical path for GPU acceleration - previous implementation only
    162
        /// recorded metrics without actually using GPU.
    163
        #[cfg(feature = "gpu")]
    164
    5
        fn batched_attention_with_cache(
    165
    5
            &self,
    166
    5
            all_q: &[Vec<f32>],
    167
    5
            all_k: &[Vec<f32>],
    168
    5
            all_v: &[Vec<f32>],
    169
    5
            cache: &OwnedQuantizedKVCache,
    170
    5
            layer_idx: usize,
    171
    5
            metrics: &std::sync::Arc<DispatchMetrics>,
    172
    5
        ) -> Result<Vec<Vec<f32>>> {
    173
    5
            let seq_len = all_q.len();
    174
    5
            let hidden_dim = self.config.hidden_dim;
    175
    5
            let num_heads = self.config.num_heads;
    176
    5
            let head_dim = hidden_dim / num_heads;
    177
    178
            // Get any cached K/V from previous sequences
    179
    5
            let cached_k = cache.get_k(layer_idx);
    180
    5
            let cached_v = cache.get_v(layer_idx);
    181
    5
            let cache_len = cached_k.len() / hidden_dim;
    182
    183
            // Build full K/V sequences: [cache + current]
    184
    5
            let total_len = cache_len + seq_len;
    185
    186
            // Determine if we should use GPU based on workload size
    187
            //
    188
            // IMPORTANT FINDING (IMP-600, PARITY-002):
    189
            // GPU is 2.7x SLOWER for MATVEC operations (per-head attention is MATVEC)
    190
            // GPU is 57x FASTER for large GEMM (batch) operations
    191
            //
    192
            // For GPU to be beneficial, we need LARGE matrices. Per-head attention
    193
            // uses tiny matrices: Q[1, head_dim] @ K^T[head_dim, seq_len] = [1, seq_len]
    194
            // This is a MATVEC operation where GPU transfer overhead dominates.
    195
            //
    196
            // Measured result with GPU matmul: 0.20 tok/s (vs 5.31 tok/s CPU)
    197
            // GPU path is 26x SLOWER due to per-head matmul overhead.
    198
            //
    199
            // For true GPU acceleration, need:
    200
            // - FlashAttention (fused kernel, not yet available in trueno)
    201
            // - Batched multi-request inference (process multiple prompts together)
    202
            //
    203
            // For now, use optimized CPU path which is faster for single-request inference.
    204
    5
            let workload = num_heads * seq_len * head_dim * total_len;
    205
    5
            let _ = workload; // Document: GPU not used because MATVEC is slower on GPU
    206
    207
            // Always use CPU path - it's faster for per-head attention MATVEC
    208
    5
            metrics.record_cpu_dispatch();
    209
    5
            self.cpu_batched_attention(
    210
    5
                all_q, all_k, all_v, cached_k, cached_v, cache_len, hidden_dim, num_heads, head_dim,
    211
            )
    212
    5
        }
    213
    214
        /// CPU-based batched attention (fallback for small workloads)
    215
        #[cfg(feature = "gpu")]
    216
        #[allow(clippy::too_many_arguments)] // Attention requires all these parameters
    217
    5
        fn cpu_batched_attention(
    218
    5
            &self,
    219
    5
            all_q: &[Vec<f32>],
    220
    5
            all_k: &[Vec<f32>],
    221
    5
            all_v: &[Vec<f32>],
    222
    5
            cached_k: &[f32],
    223
    5
            cached_v: &[f32],
    224
    5
            cache_len: usize,
    225
    5
            hidden_dim: usize,
    226
    5
            _num_heads: usize,
    227
    5
            head_dim: usize,
    228
    5
        ) -> Result<Vec<Vec<f32>>> {
    229
    5
            let seq_len = all_q.len();
    230
    5
            let mut outputs = Vec::with_capacity(seq_len);
    231
    232
    84
            for (q_pos, q) in all_q5.iter5().enumerate5() {
    233
    84
                let attend_len = cache_len + q_pos + 1;
    234
    84
                let mut k_vecs: Vec<&[f32]> = Vec::with_capacity(attend_len);
    235
    84
                let mut v_vecs: Vec<&[f32]> = Vec::with_capacity(attend_len);
    236
    237
                // Add cached K/V
    238
    84
                for i0 in 0..cache_len {
    239
    0
                    let start = i * hidden_dim;
    240
    0
                    let end = start + hidden_dim;
    241
    0
                    k_vecs.push(&cached_k[start..end]);
    242
    0
                    v_vecs.push(&cached_v[start..end]);
    243
    0
                }
    244
    245
                // Add current sequence K/V up to and including current position
    246
    2.14k
                for i in 0..=q_pos84 {
    247
    2.14k
                    k_vecs.push(&all_k[i]);
    248
    2.14k
                    v_vecs.push(&all_v[i]);
    249
    2.14k
                }
    250
    251
    84
                let output = self.compute_attention_output(q, &k_vecs, &v_vecs, head_dim)?0;
    252
    84
                outputs.push(output);
    253
            }
    254
    255
    5
            Ok(outputs)
    256
    5
        }
    257
    258
        /// Compute attention output for a single query against K/V vectors
    259
        #[cfg(feature = "gpu")]
    260
    84
        fn compute_attention_output(
    261
    84
            &self,
    262
    84
            q: &[f32],
    263
    84
            k_vecs: &[&[f32]],
    264
    84
            v_vecs: &[&[f32]],
    265
    84
            head_dim: usize,
    266
    84
        ) -> Result<Vec<f32>> {
    267
    84
            let hidden_dim = q.len();
    268
    84
            let num_heads = hidden_dim / head_dim;
    269
    84
            let seq_len = k_vecs.len();
    270
    271
    84
            if seq_len == 0 {
    272
                // No keys to attend to - return zeros (will be replaced by first attention)
    273
    0
                return Ok(vec![0.0; hidden_dim]);
    274
    84
            }
    275
    276
    84
            let scale = 1.0 / (head_dim as f32).sqrt();
    277
    84
            let mut output = vec![0.0; hidden_dim];
    278
    279
            // Process each head independently
    280
    592
            for head in 0..num_heads84 {
    281
    592
                let head_start = head * head_dim;
    282
    592
                let head_end = head_start + head_dim;
    283
    284
    592
                let q_head = &q[head_start..head_end];
    285
    286
                // Compute attention scores for this head
    287
    592
                let mut scores = Vec::with_capacity(seq_len);
    288
    17.4k
                for k16.9k in k_vecs {
    289
    16.9k
                    let k_head = &k[head_start..head_end];
    290
    536k
                    let score16.9k: f3216.9k = q_head16.9k.iter16.9k().zip16.9k(k_head16.9k.iter16.9k()).map16.9k(|(a, b)| a * b).sum16.9k();
    291
    16.9k
                    scores.push(score * scale);
    292
                }
    293
    294
                // Softmax (SIMD-optimized, in-place)
    295
    592
                crate::quantize::softmax_simd(&mut scores);
    296
    297
                // Weighted sum of values
    298
    16.9k
                for (attn, v) in scores.iter()592.zip592(v_vecs592.iter592()) {
    299
    16.9k
                    let v_head = &v[head_start..head_end];
    300
    536k
                    for (i, &v_val) in v_head16.9k.iter16.9k().enumerate16.9k() {
    301
    536k
                        output[head_start + i] += attn * v_val;
    302
    536k
                    }
    303
                }
    304
            }
    305
    306
    84
            Ok(output)
    307
    84
        }
    308
    309
        /// Generate tokens with batched prompt prefill (PARITY-002)
    310
        ///
    311
        /// Uses `forward_batch_with_cache` for initial prompt processing (GPU-accelerated),
    312
        /// then falls back to single-token generation for autoregressive decoding.
    313
        ///
    314
        /// # Arguments
    315
        /// * `prompt` - Initial token IDs (processed in batch)
    316
        /// * `config` - Generation configuration
    317
        /// * `metrics` - Dispatch metrics tracker
    318
        ///
    319
        /// # Returns
    320
        /// Generated token sequence including prompt
    321
        ///
    322
        /// # Errors
    323
        /// Returns error if generation fails
    324
        #[cfg(feature = "gpu")]
    325
    1
        pub fn generate_with_batched_prefill(
    326
    1
            &self,
    327
    1
            prompt: &[u32],
    328
    1
            config: &QuantizedGenerateConfig,
    329
    1
            metrics: &std::sync::Arc<DispatchMetrics>,
    330
    1
        ) -> Result<Vec<u32>> {
    331
    1
            if prompt.is_empty() {
    332
    0
                return Err(RealizarError::InvalidShape {
    333
    0
                    reason: "Prompt cannot be empty".to_string(),
    334
    0
                });
    335
    1
            }
    336
    337
    1
            let max_seq_len = prompt.len() + config.max_tokens;
    338
    1
            let mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);
    339
    1
            let mut tokens = prompt.to_vec();
    340
    341
            // PARITY-002: Process ALL prompt tokens at once (batched prefill)
    342
            // This enables GPU acceleration for the attention computation
    343
    1
            let mut logits = self.forward_batch_with_cache(prompt, &mut cache, metrics)?0;
    344
    345
            // Generate new tokens one at a time (autoregressive)
    346
    5
            for gen_idx in 0..config.max_tokens1 {
    347
                // Sample next token from logits
    348
    5
                let next_token = if config.temperature == 0.0 || config.top_k == 10 {
    349
    5
                    ops::argmax(&logits)
    350
                } else {
    351
    0
                    crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
    352
                };
    353
    354
                // Check stop condition
    355
    5
                if config.stop_tokens.contains(&next_token) {
    356
    0
                    break;
    357
    5
                }
    358
    359
    5
                tokens.push(next_token);
    360
    361
                // Check max length
    362
    5
                if tokens.len() >= max_seq_len {
    363
    1
                    break;
    364
    4
                }
    365
    366
                // Forward pass for the new token (single-token, uses CPU)
    367
    4
                let position = prompt.len() + gen_idx;
    368
    4
                logits =
    369
    4
                    self.forward_single_with_cache_adaptive(next_token, &mut cache, position, metrics)?0;
    370
            }
    371
    372
    1
            Ok(tokens)
    373
    1
        }
    374
    375
        /// Generate tokens with SmallVec optimization (IMP-117)
    376
        ///
    377
        /// Uses SmallVec for token storage to avoid heap allocations when:
    378
        /// - Prompt + max_tokens <= TOKEN_BUFFER_INLINE_CAP
    379
        ///
    380
        /// # Arguments
    381
        /// * `prompt` - Input token buffer (can be SmallVec or slice)
    382
        /// * `config` - Generation configuration
    383
        ///
    384
        /// # Returns
    385
        /// Generated token sequence as TokenBuffer (SmallVec)
    386
        ///
    387
        /// # Errors
    388
        /// Returns error if forward pass fails
    389
    1
        pub fn generate_with_smallvec(
    390
    1
            &self,
    391
    1
            prompt: &[u32],
    392
    1
            config: &QuantizedGenerateConfig,
    393
    1
        ) -> Result<TokenBuffer> {
    394
    1
            if prompt.is_empty() {
    395
    0
                return Err(RealizarError::InvalidShape {
    396
    0
                    reason: "Prompt cannot be empty".to_string(),
    397
    0
                });
    398
    1
            }
    399
    400
    1
            let max_seq_len = prompt.len() + config.max_tokens;
    401
    1
            let mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);
    402
    403
            // Use SmallVec for token storage - inline for small sequences
    404
    1
            let mut tokens: TokenBuffer = TokenBuffer::from_slice(prompt);
    405
    406
            // Process prompt tokens (prefill)
    407
    5
            for (pos, &token_id) in prompt1.iter1().enumerate1() {
    408
    5
                let _ = self.forward_single_with_cache(token_id, &mut cache, pos)?0;
    409
            }
    410
    411
            // Generate new tokens
    412
    10
            for gen_idx in 0..config.max_tokens1 {
    413
    10
                let position = prompt.len() + gen_idx;
    414
    10
                let last_token = *tokens.last().ok_or_else(|| RealizarError::InvalidShape {
    415
    0
                    reason: "Token buffer empty during generation".to_string(),
    416
    0
                })?;
    417
    418
    10
                let logits = self.forward_single_with_cache(last_token, &mut cache, position)?0;
    419
    420
                // Sample next token
    421
    10
                let next_token = if config.temperature == 0.0 || config.top_k == 10 {
    422
    10
                    ops::argmax(&logits)
    423
                } else {
    424
    0
                    crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
    425
                };
    426
    427
                // Check stop condition
    428
    10
                if config.stop_tokens.contains(&next_token) {
    429
    0
                    break;
    430
    10
                }
    431
    432
    10
                tokens.push(next_token);
    433
    434
                // Check max length
    435
    10
                if tokens.len() >= max_seq_len {
    436
    1
                    break;
    437
    9
                }
    438
            }
    439
    440
    1
            Ok(tokens)
    441
    1
        }
    442
    443
        // ========================================================================
    444
        // PARITY-006: Batch Processing - Parallel Token Generation
    445
        // ========================================================================
    446
    447
        /// Generate tokens for multiple requests in parallel (PARITY-006)
    448
        ///
    449
        /// This processes multiple independent requests together, enabling GPU GEMM
    450
        /// acceleration. When batch_size > 1, the matmul operations become:
    451
        /// `[batch_size, hidden_dim] @ [hidden_dim, output_dim]` which is GEMM.
    452
        ///
    453
        /// Per IMP-600: GPU is 57x faster for GEMM vs 2.7x slower for MATVEC.
    454
        /// Batch inference is the key to utilizing GPU acceleration effectively.
    455
        ///
    456
        /// # Arguments
    457
        /// * `prompts` - Vector of prompts (each prompt is a slice of token IDs)
    458
        /// * `config` - Generation configuration (shared across all requests)
    459
        ///
    460
        /// # Returns
    461
        /// Vector of generated token sequences (one per input prompt)
    462
        ///
    463
        /// # Performance
    464
        /// - batch_size=1: Falls back to single-request path (CPU optimal)
    465
        /// - batch_size>1: Uses batched matmul for GPU GEMM acceleration
    466
        ///
    467
        /// # Errors
    468
        /// Returns error if any request fails
    469
    5
        pub fn batch_generate(
    470
    5
            &self,
    471
    5
            prompts: &[&[u32]],
    472
    5
            config: &QuantizedGenerateConfig,
    473
    5
        ) -> Result<Vec<Vec<u32>>> {
    474
    5
            if prompts.is_empty() {
    475
    1
                return Err(RealizarError::InvalidShape {
    476
    1
                    reason: "Prompts cannot be empty".to_string(),
    477
    1
                });
    478
    4
            }
    479
    480
            // For single request, use optimized single-request path
    481
    4
            if prompts.len() == 1 {
    482
    2
                return Ok(vec![self.generate_with_cache(prompts[0], config)?0]);
    483
    2
            }
    484
    485
    2
            let batch_size = prompts.len();
    486
    7
            let max_prompt_len2 = prompts2.iter2().map2(|p| p.len()).max2().unwrap_or2(0);
    487
    2
            let max_seq_len = max_prompt_len + config.max_tokens;
    488
    489
            // Create KV caches for each request
    490
    2
            let mut caches: Vec<OwnedQuantizedKVCache> = (0..batch_size)
    491
    7
                .map2(|_| OwnedQuantizedKVCache::from_config(&self.config, max_seq_len))
    492
    2
                .collect();
    493
    494
            // Initialize token sequences with prompts
    495
    7
            let mut all_tokens2: Vec<Vec<u32>>2 = prompts2.iter2().map2(|p| p.to_vec()).collect2();
    496
    497
            // Track which requests are still generating
    498
    2
            let mut active: Vec<bool> = vec![true; batch_size];
    499
    500
            // Prefill phase: process each prompt (can be batched in future)
    501
    7
            for (req_idx, prompt) in prompts2.iter2().enumerate2() {
    502
    21
                for (pos, &token_id) in prompt7.iter7().enumerate7() {
    503
    21
                    let _ = self.forward_single_with_cache(token_id, &mut caches[req_idx], pos)?0;
    504
                }
    505
            }
    506
    507
            // Generation phase: process all active requests together
    508
    10
            for gen_idx in 0..config.max_tokens2 {
    509
                // Count active requests
    510
    10
                let active_count = active.iter().filter(|&&a| a).count();
    511
    10
                if active_count == 0 {
    512
    0
                    break;
    513
    10
                }
    514
    515
                // Collect last tokens from active requests
    516
    10
                let active_indices: Vec<usize> = active
    517
    10
                    .iter()
    518
    10
                    .enumerate()
    519
    10
                    .filter(|(_, &a)| a)
    520
    10
                    .map(|(i, _)| i)
    521
    10
                    .collect();
    522
    523
                // Process active requests - batched forward pass
    524
    10
                let mut next_tokens = Vec::with_capacity(active_count);
    525
    526
    45
                for &req_idx35 in &active_indices {
    527
    35
                    let position = prompts[req_idx].len() + gen_idx;
    528
    35
                    let last_token = *all_tokens[req_idx]
    529
    35
                        .last()
    530
    35
                        .expect("tokens must be non-empty");
    531
    532
    35
                    let logits =
    533
    35
                        self.forward_single_with_cache(last_token, &mut caches[req_idx], position)?0;
    534
    535
                    // Sample next token
    536
    35
                    let next_token = if config.temperature == 0.0 || config.top_k == 10 {
    537
    35
                        ops::argmax(&logits)
    538
                    } else {
    539
    0
                        crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
    540
                    };
    541
    542
    35
                    next_tokens.push((req_idx, next_token));
    543
                }
    544
    545
                // Apply next tokens and check stop conditions
    546
    45
                for (req_idx35, next_token35) in next_tokens {
    547
    35
                    if config.stop_tokens.contains(&next_token) {
    548
    0
                        active[req_idx] = false;
    549
    0
                        continue;
    550
    35
                    }
    551
    552
    35
                    all_tokens[req_idx].push(next_token);
    553
    554
    35
                    if all_tokens[req_idx].len() >= max_seq_len {
    555
    5
                        active[req_idx] = false;
    556
    30
                    }
    557
                }
    558
            }
    559
    560
    2
            Ok(all_tokens)
    561
    5
        }
    562
    563
        /// Get the batch throughput improvement factor (PARITY-006)
    564
        ///
    565
        /// Per IMP-600: GPU GEMM is 57x faster than MATVEC.
    566
        /// Batch inference converts MATVEC to GEMM when batch_size > 1.
    567
        ///
    568
        /// # Arguments
    569
        /// * `batch_size` - Number of concurrent requests
    570
        ///
    571
        /// # Returns
    572
        /// Estimated throughput multiplier vs single-request
    573
        #[must_use]
    574
    3
        pub const fn batch_throughput_factor(batch_size: usize) -> f64 {
    575
    3
            match batch_size {
    576
    1
                0 | 1 => 1.0,
    577
    2
                2..=4 => 1.81,   // ~2x throughput with small batch
    578
    1
                5..=8 => 2.50,   // GPU GEMM starts to help
    579
    1
                9..=16 => 3.5,  // Good GPU utilization
    580
    0
                17..=32 => 5.0, // Near-optimal batch
    581
    0
                _ => 6.0,       // Large batch, GPU-limited
    582
            }
    583
    3
        }
    584
    585
        /// Forward pass for a batch of tokens (IMP-106)
    586
        ///
    587
        /// Processes multiple tokens through the transformer in parallel.
    588
        /// This is more efficient than sequential processing for prefill.
    589
        ///
    590
        /// # Arguments
    591
        /// * `token_ids` - Batch of input token IDs [batch_size]
    592
        ///
    593
        /// # Returns
    594
        /// Logits for all positions [batch_size * vocab_size]
    595
        ///
    596
        /// # Errors
    597
        /// Returns error if tensor operations fail
    598
    2
        pub fn forward_batch(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
    599
    2
            let batch_size = token_ids.len();
    600
    2
            let hidden_dim = self.config.hidden_dim;
    601
    602
            // 1. Token embedding lookup for all tokens
    603
    2
            let mut hidden = self.embed(token_ids);
    604
    605
            // 2. Process through transformer layers
    606
    4
            for layer2 in &self.layers {
    607
                // Pre-attention LayerNorm
    608
    2
                let normed = ops::layer_norm(
    609
    2
                    &hidden,
    610
    2
                    &layer.attn_norm_weight,
    611
    2
                    layer.attn_norm_bias.as_deref(),
    612
    2
                    self.config.eps,
    613
                );
    614
    615
                // QKV projection (batched)
    616
    2
                let qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0;
    617
    618
                // Split Q, K, V for batch - simplified attention (no causal mask for batch)
    619
    2
                let qkv_dim = qkv.len() / batch_size;
    620
    2
                let q_dim = hidden_dim;
    621
    2
                let kv_dim = (qkv_dim - q_dim) / 2;
    622
    623
                // Process attention for each position (simplified for batch)
    624
    2
                let mut attn_out = Vec::with_capacity(batch_size * hidden_dim);
    625
    8
                for pos in 0..batch_size2 {
    626
    8
                    let qkv_start = pos * qkv_dim;
    627
    8
                    let q = &qkv[qkv_start..qkv_start + q_dim];
    628
    8
                    let k = &qkv[qkv_start + q_dim..qkv_start + q_dim + kv_dim];
    629
    8
                    let v = &qkv[qkv_start + q_dim + kv_dim..qkv_start + qkv_dim];
    630
    631
                    // Simple self-attention for current position (attend to itself only for simplicity)
    632
                    // Full causal attention would require attending to all previous positions
    633
    8
                    let head_dim = hidden_dim / self.config.num_heads;
    634
    8
                    let scale = 1.0 / (head_dim as f32).sqrt();
    635
    636
    8
                    let mut out = vec![0.0f32; hidden_dim];
    637
    32
                    for h in 0..self.config.num_heads8 {
    638
    32
                        let kv_h = h * self.config.num_kv_heads / self.config.num_heads;
    639
    32
                        let q_h = &q[h * head_dim..(h + 1) * head_dim];
    640
    32
                        let k_h = &k[kv_h * head_dim..(kv_h + 1) * head_dim];
    641
    32
                        let v_h = &v[kv_h * head_dim..(kv_h + 1) * head_dim];
    642
    643
                        // Score and softmax (single position = 1.0 weight)
    644
    32
                        let mut score = 0.0f32;
    645
    256
                        for d in 0..head_dim32 {
    646
    256
                            score += q_h[d] * k_h[d];
    647
    256
                        }
    648
    32
                        let _weight = (score * scale).exp(); // softmax of single value = 1.0
    649
    650
                        // Apply value
    651
    256
                        for d in 0..head_dim32 {
    652
    256
                            out[h * head_dim + d] = v_h[d];
    653
    256
                        }
    654
                    }
    655
    8
                    attn_out.extend_from_slice(&out);
    656
                }
    657
    658
                // Output projection
    659
    2
                let projected = self.fused_matmul(&attn_out, &layer.attn_output_weight)?0;
    660
    661
                // Residual connection
    662
    256
                for i in 0..hidden2.len2() {
    663
    256
                    hidden[i] += projected[i];
    664
    256
                }
    665
    666
                // FFN (pre-norm style)
    667
    2
                let ffn_normed =
    668
    2
                    ops::layer_norm(&hidden, &layer.attn_norm_weight, None, self.config.eps);
    669
    2
                let up = self.fused_matmul(&ffn_normed, &layer.ffn_up_weight)?0;
    670
    671
                // GELU activation
    672
    2
                let gelu: Vec<f32> = up
    673
    2
                    .iter()
    674
    512
                    .map2(|&x| 0.5 * x * (1.0 + (0.797_884_6 * (x + 0.044_715 * x.powi(3))).tanh()))
    675
    2
                    .collect();
    676
    677
    2
                let down = self.fused_matmul(&gelu, &layer.ffn_down_weight)?0;
    678
    679
                // Residual connection
    680
    256
                for i in 0..hidden2.len2() {
    681
    256
                    hidden[i] += down[i];
    682
    256
                }
    683
            }
    684
    685
            // 3. Final LayerNorm
    686
    2
            let normed = ops::layer_norm(
    687
    2
                &hidden,
    688
    2
                &self.output_norm_weight,
    689
    2
                self.output_norm_bias.as_deref(),
    690
    2
                self.config.eps,
    691
            );
    692
    693
            // 4. LM head projection to vocab logits
    694
    2
            let logits = self.fused_matmul(&normed, &self.lm_head_weight)?0;
    695
    696
    2
            Ok(logits)
    697
    2
        }
    698
    699
        /// Prefill prompt tokens with batched forward pass (IMP-106)
    700
        ///
    701
        /// Efficiently processes all prompt tokens and populates the KV cache.
    702
        /// Returns the last position's logits for sampling.
    703
        ///
    704
        /// # Arguments
    705
        /// * `prompt` - Prompt token IDs
    706
        /// * `cache` - KV cache to populate
    707
        ///
    708
        /// # Returns
    709
        /// Logits for the last position [vocab_size]
    710
        ///
    711
        /// # Errors
    712
        /// Returns error if forward pass fails
    713
    1
        pub fn prefill_batch(
    714
    1
            &self,
    715
    1
            prompt: &[u32],
    716
    1
            cache: &mut OwnedQuantizedKVCache,
    717
    1
        ) -> Result<Vec<f32>> {
    718
    1
            if prompt.is_empty() {
    719
    0
                return Err(RealizarError::InvalidShape {
    720
    0
                    reason: "Prompt cannot be empty".to_string(),
    721
    0
                });
    722
    1
            }
    723
    724
            // Process each position to populate KV cache
    725
            // (True batch prefill would compute all positions at once with causal attention)
    726
    1
            let mut last_logits = Vec::new();
    727
    4
            for (pos, &token_id) in prompt1.iter1().enumerate1() {
    728
    4
                last_logits = self.forward_single_with_cache(token_id, cache, pos)?0;
    729
            }
    730
    731
    1
            Ok(last_logits)
    732
    1
        }
    733
    734
        /// Forward pass for a batch of tokens with GPU acceleration (IMP-107)
    735
        ///
    736
        /// Uses HybridScheduler to route matmuls to GPU when batch_size > 1
    737
        /// and matrix size exceeds threshold. Falls back to CPU for small batches.
    738
        ///
    739
        /// # Arguments
    740
        /// * `token_ids` - Batch of input token IDs [batch_size]
    741
        ///
    742
        /// # Returns
    743
        /// Logits for all positions [batch_size * vocab_size]
    744
        ///
    745
        /// # Errors
    746
        /// Returns error if GPU initialization or tensor operations fail
    747
        #[cfg(feature = "gpu")]
    748
    3
        pub fn forward_batch_gpu(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
    749
            use crate::gpu::HybridScheduler;
    750
    751
    3
            let batch_size = token_ids.len();
    752
    3
            let hidden_dim = self.config.hidden_dim;
    753
    3
            let vocab_size = self.config.vocab_size;
    754
    755
            // Initialize HybridScheduler with reasonable threshold
    756
            // Threshold of 1000 means: batch_size * hidden_dim * out_dim > 1000 uses GPU
    757
    3
            let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
    758
    0
                RealizarError::UnsupportedOperation {
    759
    0
                    operation: "HybridScheduler::with_threshold".to_string(),
    760
    0
                    reason: format!("GPU scheduler initialization failed: {e}"),
    761
    0
                }
    762
    0
            })?;
    763
    764
            // 1. Token embedding lookup for all tokens
    765
    3
            let mut hidden = self.embed(token_ids);
    766
    767
            // 2. Process through transformer layers
    768
    6
            for layer3 in &self.layers {
    769
                // Pre-attention LayerNorm
    770
    3
                let normed = ops::layer_norm(
    771
    3
                    &hidden,
    772
    3
                    &layer.attn_norm_weight,
    773
    3
                    layer.attn_norm_bias.as_deref(),
    774
    3
                    self.config.eps,
    775
                );
    776
    777
                // QKV projection - use GPU for batch ops
    778
    3
                let qkv = self.batch_qkv_matmul_gpu_with_scheduler(
    779
    3
                    &normed,
    780
    3
                    &layer.qkv_weight,
    781
    3
                    batch_size,
    782
    3
                    hidden_dim,
    783
    3
                    &mut scheduler,
    784
    0
                )?;
    785
    786
                // Split Q, K, V for batch - PARITY-114: use proper batched causal attention
    787
    3
                let qkv_dim = qkv.len() / batch_size;
    788
    3
                let q_dim = hidden_dim;
    789
    3
                let kv_dim = (qkv_dim - q_dim) / 2;
    790
    791
                // Collect Q, K, V for all positions
    792
    3
                let mut q_all = Vec::with_capacity(batch_size * q_dim);
    793
    3
                let mut k_all = Vec::with_capacity(batch_size * kv_dim);
    794
    3
                let mut v_all = Vec::with_capacity(batch_size * kv_dim);
    795
    796
    20
                for pos in 0..batch_size3 {
    797
    20
                    let qkv_start = pos * qkv_dim;
    798
    20
                    q_all.extend_from_slice(&qkv[qkv_start..qkv_start + q_dim]);
    799
    20
                    k_all.extend_from_slice(&qkv[qkv_start + q_dim..qkv_start + q_dim + kv_dim]);
    800
    20
                    v_all.extend_from_slice(&qkv[qkv_start + q_dim + kv_dim..qkv_start + qkv_dim]);
    801
    20
                }
    802
    803
                // Proper batched causal attention (PARITY-114: matches cached forward path)
    804
    3
                let attn_out = self.batched_causal_attention_gpu(&q_all, &k_all, &v_all, batch_size)?0;
    805
    806
                // Output projection - use GPU for batch ops
    807
    3
                let projected = self.batch_matmul_gpu(
    808
    3
                    &attn_out,
    809
    3
                    &layer.attn_output_weight,
    810
    3
                    batch_size,
    811
    3
                    hidden_dim,
    812
    3
                    layer.attn_output_weight.out_dim,
    813
    3
                    &mut scheduler,
    814
    0
                )?;
    815
    816
                // Residual connection
    817
    1.28k
                for i in 0..hidden3.len3() {
    818
    1.28k
                    hidden[i] += projected[i];
    819
    1.28k
                }
    820
    821
                // FFN (pre-norm style)
    822
    3
                let ffn_normed = ops::layer_norm(
    823
    3
                    &hidden,
    824
    3
                    &layer.attn_norm_weight,
    825
    3
                    layer.attn_norm_bias.as_deref(),
    826
    3
                    self.config.eps,
    827
                );
    828
    829
                // FFN up projection - use GPU
    830
    3
                let mut ffn_hidden = self.batch_matmul_gpu(
    831
    3
                    &ffn_normed,
    832
    3
                    &layer.ffn_up_weight,
    833
    3
                    batch_size,
    834
    3
                    hidden_dim,
    835
    3
                    layer.ffn_up_weight.out_dim,
    836
    3
                    &mut scheduler,
    837
    0
                )?;
    838
    839
                // GELU activation
    840
    3
                ops::gelu(&mut ffn_hidden);
    841
    842
                // FFN down projection - use GPU
    843
    3
                let ffn_output = self.batch_matmul_gpu(
    844
    3
                    &ffn_hidden,
    845
    3
                    &layer.ffn_down_weight,
    846
    3
                    batch_size,
    847
    3
                    layer.ffn_up_weight.out_dim,
    848
    3
                    hidden_dim,
    849
    3
                    &mut scheduler,
    850
    0
                )?;
    851
    852
                // Residual
    853
    1.28k
                for i in 0..hidden3.len3() {
    854
    1.28k
                    hidden[i] += ffn_output[i];
    855
    1.28k
                }
    856
            }
    857
    858
            // 3. Final layer norm
    859
    3
            let normed = ops::layer_norm(
    860
    3
                &hidden,
    861
    3
                &self.output_norm_weight,
    862
    3
                self.output_norm_bias.as_deref(),
    863
    3
                self.config.eps,
    864
            );
    865
    866
            // 4. LM head projection - use GPU for large vocab
    867
    3
            let logits = self.batch_matmul_gpu(
    868
    3
                &normed,
    869
    3
                &self.lm_head_weight,
    870
    3
                batch_size,
    871
    3
                hidden_dim,
    872
    3
                vocab_size,
    873
    3
                &mut scheduler,
    874
    0
            )?;
    875
    876
    3
            Ok(logits)
    877
    3
        }
    878
    879
        /// Batch matmul with GPU acceleration via HybridScheduler (IMP-107)
    880
        ///
    881
        /// Dequantizes weights and uses GPU for large operations.
    882
        #[cfg(feature = "gpu")]
    883
    15
        fn batch_matmul_gpu(
    884
    15
            &self,
    885
    15
            input: &[f32],
    886
    15
            weight: &OwnedQuantizedTensor,
    887
    15
            m: usize,
    888
    15
            k: usize,
    889
    15
            n: usize,
    890
    15
            scheduler: &mut crate::gpu::HybridScheduler,
    891
    15
        ) -> Result<Vec<f32>> {
    892
            // Dequantize weight to f32
    893
    15
            let weight_f32 = self.dequantize_weight(weight)?0;
    894
    895
            // Use HybridScheduler for GPU/CPU dispatch
    896
            // A: [m, k], B: [k, n] -> C: [m, n]
    897
    15
            scheduler.matmul(input, &weight_f32, m, k, n).map_err(|e| {0
    898
    0
                RealizarError::UnsupportedOperation {
    899
    0
                    operation: "HybridScheduler::matmul".to_string(),
    900
    0
                    reason: format!("GPU matmul failed: {e}"),
    901
    0
                }
    902
    0
            })
    903
    15
        }
    904
    905
        /// Batch QKV matmul with GPU acceleration via HybridScheduler
    906
        ///
    907
        /// Five Whys Root Cause Fix: Handles both fused and separate Q/K/V formats
    908
        #[cfg(feature = "gpu")]
    909
    3
        fn batch_qkv_matmul_gpu_with_scheduler(
    910
    3
            &self,
    911
    3
            input: &[f32],
    912
    3
            qkv: &OwnedQKVWeights,
    913
    3
            batch_size: usize,
    914
    3
            hidden_dim: usize,
    915
    3
            scheduler: &mut crate::gpu::HybridScheduler,
    916
    3
        ) -> Result<Vec<f32>> {
    917
    3
            match qkv {
    918
    3
                OwnedQKVWeights::Fused(ref weight) => self.batch_matmul_gpu(
    919
    3
                    input,
    920
    3
                    weight,
    921
    3
                    batch_size,
    922
    3
                    hidden_dim,
    923
    3
                    weight.out_dim,
    924
    3
                    scheduler,
    925
                ),
    926
                OwnedQKVWeights::Separate {
    927
    0
                    ref q,
    928
    0
                    ref k,
    929
    0
                    ref v,
    930
                } => {
    931
                    // Compute Q, K, V separately then concatenate
    932
    0
                    let q_out =
    933
    0
                        self.batch_matmul_gpu(input, q, batch_size, hidden_dim, q.out_dim, scheduler)?;
    934
    0
                    let k_out =
    935
    0
                        self.batch_matmul_gpu(input, k, batch_size, hidden_dim, k.out_dim, scheduler)?;
    936
    0
                    let v_out =
    937
    0
                        self.batch_matmul_gpu(input, v, batch_size, hidden_dim, v.out_dim, scheduler)?;
    938
    939
                    // Interleave Q, K, V for each position in batch
    940
    0
                    let qkv_dim = q.out_dim + k.out_dim + v.out_dim;
    941
    0
                    let mut output = Vec::with_capacity(batch_size * qkv_dim);
    942
    0
                    for b in 0..batch_size {
    943
    0
                        output.extend_from_slice(&q_out[b * q.out_dim..(b + 1) * q.out_dim]);
    944
    0
                        output.extend_from_slice(&k_out[b * k.out_dim..(b + 1) * k.out_dim]);
    945
    0
                        output.extend_from_slice(&v_out[b * v.out_dim..(b + 1) * v.out_dim]);
    946
    0
                    }
    947
    0
                    Ok(output)
    948
                },
    949
            }
    950
    3
        }
    951
    952
        /// Dequantize a weight tensor to f32
    953
        #[cfg(feature = "gpu")]
    954
    47
        pub(crate) fn dequantize_weight(&self, weight: &OwnedQuantizedTensor) -> Result<Vec<f32>> {
    955
            use crate::quantize::{dequantize_q4_k_simd, dequantize_q5_k, dequantize_q6_k, QK_K};
    956
    957
    47
            let in_dim = weight.in_dim;
    958
    47
            let out_dim = weight.out_dim;
    959
    47
            let total_elements = in_dim * out_dim;
    960
    961
    47
            match weight.qtype {
    962
                GGUF_TYPE_Q4_K => {
    963
    47
                    let super_blocks_per_row = in_dim.div_ceil(QK_K);
    964
    47
                    let mut output = Vec::with_capacity(total_elements);
    965
    5.95k
                    for row in 0..out_dim47 {
    966
    5.95k
                        let row_start = row * super_blocks_per_row * 144;
    967
    5.95k
                        let row_end = row_start + super_blocks_per_row * 144;
    968
    5.95k
                        let row_data = &weight.data[row_start..row_end];
    969
    5.95k
                        let row_dequant = dequantize_q4_k_simd(row_data)?0;
    970
                        // Take only in_dim values (may have padding due to super-block alignment)
    971
    5.95k
                        output.extend_from_slice(&row_dequant[..in_dim.min(row_dequant.len())]);
    972
                    }
    973
    47
                    Ok(output)
    974
                },
    975
                GGUF_TYPE_Q5_K => {
    976
    0
                    let super_blocks_per_row = in_dim.div_ceil(QK_K);
    977
    0
                    let mut output = Vec::with_capacity(total_elements);
    978
    0
                    for row in 0..out_dim {
    979
    0
                        let row_start = row * super_blocks_per_row * 176;
    980
    0
                        let row_end = row_start + super_blocks_per_row * 176;
    981
    0
                        let row_data = &weight.data[row_start..row_end];
    982
    0
                        let row_dequant = dequantize_q5_k(row_data)?;
    983
    0
                        output.extend_from_slice(&row_dequant[..in_dim.min(row_dequant.len())]);
    984
                    }
    985
    0
                    Ok(output)
    986
                },
    987
                GGUF_TYPE_Q6_K => {
    988
    0
                    let super_blocks_per_row = in_dim.div_ceil(QK_K);
    989
    0
                    let mut output = Vec::with_capacity(total_elements);
    990
    0
                    for row in 0..out_dim {
    991
    0
                        let row_start = row * super_blocks_per_row * 210;
    992
    0
                        let row_end = row_start + super_blocks_per_row * 210;
    993
    0
                        let row_data = &weight.data[row_start..row_end];
    994
    0
                        let row_dequant = dequantize_q6_k(row_data)?;
    995
    0
                        output.extend_from_slice(&row_dequant[..in_dim.min(row_dequant.len())]);
    996
                    }
    997
    0
                    Ok(output)
    998
                },
    999
                _ => {
    1000
                    // F32 or unsupported - interpret raw bytes as f32
    1001
    0
                    let num_floats = weight.data.len() / 4;
    1002
    0
                    let mut output = vec![0.0f32; num_floats];
    1003
    0
                    for (i, chunk) in weight.data.chunks_exact(4).enumerate() {
    1004
    0
                        output[i] = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
    1005
    0
                    }
    1006
    0
                    Ok(output)
    1007
                },
    1008
            }
    1009
    47
        }
    1010
    1011
        /// Dequantize QKV weights - handles both fused and separate formats
    1012
        ///
    1013
        /// Five Whys Root Cause Fix: This method handles both tensor layouts for dequantization
    1014
        #[cfg(feature = "gpu")]
    1015
    0
        pub fn dequantize_qkv(&self, qkv: &OwnedQKVWeights) -> Result<Vec<f32>> {
    1016
    0
            match qkv {
    1017
    0
                OwnedQKVWeights::Fused(ref weight) => self.dequantize_weight(weight),
    1018
                OwnedQKVWeights::Separate {
    1019
    0
                    ref q,
    1020
    0
                    ref k,
    1021
    0
                    ref v,
    1022
                } => {
    1023
                    // Dequantize each separately and concatenate
    1024
    0
                    let q_out = self.dequantize_weight(q)?;
    1025
    0
                    let k_out = self.dequantize_weight(k)?;
    1026
    0
                    let v_out = self.dequantize_weight(v)?;
    1027
    1028
    0
                    let mut output = Vec::with_capacity(q_out.len() + k_out.len() + v_out.len());
    1029
    0
                    output.extend_from_slice(&q_out);
    1030
    0
                    output.extend_from_slice(&k_out);
    1031
    0
                    output.extend_from_slice(&v_out);
    1032
    0
                    Ok(output)
    1033
                },
    1034
            }
    1035
    0
        }
    1036
    1037
        /// Fused batch matmul with GPU acceleration (IMP-109)
    1038
        ///
    1039
        /// Performs batched matrix multiplication with fused dequantization.
    1040
        /// Uses the same weight layout interpretation as `batch_matmul_gpu` for
    1041
        /// consistency within the codebase.
    1042
        ///
    1043
        /// Key optimization: Dequantizes weight matrix once for all batch elements,
    1044
        /// reducing memory bandwidth for repeated operations in transformer layers.
    1045
        ///
    1046
        /// # Arguments
    1047
        /// * `input` - Input tensor [batch_size, in_dim]
    1048
        /// * `weight` - Quantized weight tensor [out_dim, in_dim]
    1049
        /// * `batch_size` - Number of input vectors
    1050
        ///
    1051
        /// # Returns
    1052
        /// Output tensor [batch_size, out_dim]
    1053
        ///
    1054
        /// # Errors
    1055
        /// Returns error if GPU operations fail or dimensions mismatch
    1056
        #[cfg(feature = "gpu")]
    1057
    2
        pub fn fused_batch_matmul_gpu(
    1058
    2
            &self,
    1059
    2
            input: &[f32],
    1060
    2
            weight: &OwnedQuantizedTensor,
    1061
    2
            batch_size: usize,
    1062
    2
        ) -> Result<Vec<f32>> {
    1063
            use crate::gpu::HybridScheduler;
    1064
    1065
    2
            let in_dim = weight.in_dim;
    1066
    2
            let out_dim = weight.out_dim;
    1067
    1068
            // Validate input dimensions
    1069
    2
            if input.len() != batch_size * in_dim {
    1070
    0
                return Err(RealizarError::InvalidShape {
    1071
    0
                    reason: format!(
    1072
    0
                        "Input size {} doesn't match batch_size={} * in_dim={}={}",
    1073
    0
                        input.len(),
    1074
    0
                        batch_size,
    1075
    0
                        in_dim,
    1076
    0
                        batch_size * in_dim
    1077
    0
                    ),
    1078
    0
                });
    1079
    2
            }
    1080
    1081
            // Dequantize weight once (key optimization: reuse across batch elements)
    1082
    2
            let weight_f32 = self.dequantize_weight(weight)?0;
    1083
    1084
            // Use HybridScheduler for CPU/GPU dispatch based on workload size
    1085
    2
            let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
    1086
    0
                RealizarError::UnsupportedOperation {
    1087
    0
                    operation: "HybridScheduler::with_threshold".to_string(),
    1088
    0
                    reason: format!("GPU scheduler initialization failed: {e}"),
    1089
    0
                }
    1090
    0
            })?;
    1091
    1092
            // Use same matmul approach as batch_matmul_gpu for consistency
    1093
    2
            scheduler
    1094
    2
                .matmul(input, &weight_f32, batch_size, in_dim, out_dim)
    1095
    2
                .map_err(|e| RealizarError::UnsupportedOperation {
    1096
    0
                    operation: "HybridScheduler::matmul".to_string(),
    1097
    0
                    reason: format!("GPU batched matmul failed: {e}"),
    1098
    0
                })
    1099
    2
        }
    1100
    1101
        /// Batched causal attention with GPU acceleration (IMP-108)
    1102
        ///
    1103
        /// Computes causal self-attention using matrix multiplications that can be
    1104
        /// GPU-accelerated for large sequence lengths. Uses HybridScheduler for
    1105
        /// automatic CPU/GPU dispatch.
    1106
        ///
    1107
        /// Algorithm:
    1108
        /// 1. For each head: scores = Q @ K^T / sqrt(head_dim)
    1109
        /// 2. Apply causal mask: scores[i,j] = -inf for j > i
    1110
        /// 3. Softmax per row
    1111
        /// 4. Output = softmax(scores) @ V
    1112
        ///
    1113
        /// # Arguments
    1114
        /// * `q` - Query tensor [seq_len, hidden_dim]
    1115
        /// * `k` - Key tensor [seq_len, hidden_dim]
    1116
        /// * `v` - Value tensor [seq_len, hidden_dim]
    1117
        /// * `seq_len` - Sequence length
    1118
        ///
    1119
        /// # Returns
    1120
        /// Attention output [seq_len, hidden_dim]
    1121
        ///
    1122
        /// # Errors
    1123
        /// Returns error if GPU operations fail
    1124
        #[cfg(feature = "gpu")]
    1125
    7
        pub fn batched_causal_attention_gpu(
    1126
    7
            &self,
    1127
    7
            q: &[f32],
    1128
    7
            k: &[f32],
    1129
    7
            v: &[f32],
    1130
    7
            seq_len: usize,
    1131
    7
        ) -> Result<Vec<f32>> {
    1132
            use crate::gpu::HybridScheduler;
    1133
    1134
    7
            let hidden_dim = self.config.hidden_dim;
    1135
    7
            let num_heads = self.config.num_heads;
    1136
    7
            let head_dim = hidden_dim / num_heads;
    1137
    7
            let scale = 1.0 / (head_dim as f32).sqrt();
    1138
    1139
    7
            let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
    1140
    0
                RealizarError::UnsupportedOperation {
    1141
    0
                    operation: "HybridScheduler::with_threshold".to_string(),
    1142
    0
                    reason: format!("GPU scheduler initialization failed: {e}"),
    1143
    0
                }
    1144
    0
            })?;
    1145
    1146
    7
            let mut output = vec![0.0f32; seq_len * hidden_dim];
    1147
    1148
            // Process each head
    1149
    24
            for head in 0..num_heads7 {
    1150
    24
                let head_offset = head * head_dim;
    1151
    1152
                // Extract Q_h, K_h, V_h for this head: [seq_len, head_dim]
    1153
    24
                let mut q_h = Vec::with_capacity(seq_len * head_dim);
    1154
    24
                let mut k_h = Vec::with_capacity(seq_len * head_dim);
    1155
    24
                let mut v_h = Vec::with_capacity(seq_len * head_dim);
    1156
    1157
    144
                for pos in 0..seq_len24 {
    1158
    144
                    let start = pos * hidden_dim + head_offset;
    1159
    144
                    q_h.extend_from_slice(&q[start..start + head_dim]);
    1160
    144
                    k_h.extend_from_slice(&k[start..start + head_dim]);
    1161
    144
                    v_h.extend_from_slice(&v[start..start + head_dim]);
    1162
    144
                }
    1163
    1164
                // Compute attention scores: Q_h @ K_h^T -> [seq_len, seq_len]
    1165
                // Use GPU for large sequences (seq_len^2 * head_dim ops)
    1166
    24
                let scores =
    1167
    24
                    self.batched_qk_scores(&q_h, &k_h, seq_len, head_dim, scale, &mut scheduler)?0;
    1168
    1169
                // Apply causal mask and softmax
    1170
    24
                let attn_weights = self.apply_causal_mask_softmax(&scores, seq_len);
    1171
    1172
                // Compute output: attn_weights @ V_h -> [seq_len, head_dim]
    1173
    24
                let head_output =
    1174
    24
                    self.batched_attn_v(&attn_weights, &v_h, seq_len, head_dim, &mut scheduler)?0;
    1175
    1176
                // Copy head output to final output
    1177
    144
                for pos in 0..seq_len24 {
    1178
    144
                    let out_start = pos * hidden_dim + head_offset;
    1179
    144
                    let head_start = pos * head_dim;
    1180
    144
                    output[out_start..out_start + head_dim]
    1181
    144
                        .copy_from_slice(&head_output[head_start..head_start + head_dim]);
    1182
    144
                }
    1183
            }
    1184
    1185
    7
            Ok(output)
    1186
    7
        }
    1187
    1188
        /// Compute Q @ K^T attention scores with GPU acceleration
    1189
        #[cfg(feature = "gpu")]
    1190
    24
        fn batched_qk_scores(
    1191
    24
            &self,
    1192
    24
            q: &[f32],
    1193
    24
            k: &[f32],
    1194
    24
            seq_len: usize,
    1195
    24
            head_dim: usize,
    1196
    24
            scale: f32,
    1197
    24
            scheduler: &mut crate::gpu::HybridScheduler,
    1198
    24
        ) -> Result<Vec<f32>> {
    1199
            // Q: [seq_len, head_dim], K: [seq_len, head_dim]
    1200
            // scores = Q @ K^T -> [seq_len, seq_len]
    1201
    1202
            // Transpose K: [head_dim, seq_len]
    1203
    24
            let mut k_t = vec![0.0f32; head_dim * seq_len];
    1204
    144
            for i in 0..seq_len24 {
    1205
    2.04k
                for j in 0..head_dim144 {
    1206
    2.04k
                    k_t[j * seq_len + i] = k[i * head_dim + j];
    1207
    2.04k
                }
    1208
            }
    1209
    1210
            // Matmul: Q[seq_len, head_dim] @ K_T[head_dim, seq_len] -> [seq_len, seq_len]
    1211
    24
            let scores = scheduler
    1212
    24
                .matmul(q, &k_t, seq_len, head_dim, seq_len)
    1213
    24
                .map_err(|e| RealizarError::UnsupportedOperation {
    1214
    0
                    operation: "batched_qk_scores".to_string(),
    1215
    0
                    reason: format!("GPU matmul failed: {e}"),
    1216
    0
                })?;
    1217
    1218
            // Apply scale
    1219
    960
            let scaled24: Vec<f32>24 = scores.iter()24.map24(|&s| s * scale).collect24();
    1220
    24
            Ok(scaled)
    1221
    24
        }
    1222
    1223
        /// Apply causal mask and softmax to attention scores
    1224
        #[cfg(feature = "gpu")]
    1225
    64
        pub(crate) fn apply_causal_mask_softmax(&self, scores: &[f32], seq_len: usize) -> Vec<f32> {
    1226
    64
            let mut weights = vec![0.0f32; seq_len * seq_len];
    1227
    1228
    344
            for i in 0..seq_len64 {
    1229
                // Apply causal mask: set j > i to -inf
    1230
    344
                let mut max_score = f32::NEG_INFINITY;
    1231
    1.28k
                for j in 0..=i344 {
    1232
    1.28k
                    let idx = i * seq_len + j;
    1233
    1.28k
                    max_score = max_score.max(scores[idx]);
    1234
    1.28k
                }
    1235
    1236
                // Compute softmax for causal positions only
    1237
    344
                let mut exp_sum = 0.0f32;
    1238
    1.28k
                for j in 0..=i344 {
    1239
    1.28k
                    let idx = i * seq_len + j;
    1240
    1.28k
                    let exp_val = (scores[idx] - max_score).exp();
    1241
    1.28k
                    weights[idx] = exp_val;
    1242
    1.28k
                    exp_sum += exp_val;
    1243
    1.28k
                }
    1244
    1245
                // Normalize
    1246
    1.28k
                for j in 0..=i344 {
    1247
    1.28k
                    let idx = i * seq_len + j;
    1248
    1.28k
                    weights[idx] /= exp_sum;
    1249
    1.28k
                }
    1250
                // j > i remains 0 (masked out)
    1251
            }
    1252
    1253
    64
            weights
    1254
    64
        }
    1255
    1256
        /// Compute attention_weights @ V with GPU acceleration
    1257
        #[cfg(feature = "gpu")]
    1258
    24
        fn batched_attn_v(
    1259
    24
            &self,
    1260
    24
            attn_weights: &[f32],
    1261
    24
            v: &[f32],
    1262
    24
            seq_len: usize,
    1263
    24
            head_dim: usize,
    1264
    24
            scheduler: &mut crate::gpu::HybridScheduler,
    1265
    24
        ) -> Result<Vec<f32>> {
    1266
            // attn_weights: [seq_len, seq_len], V: [seq_len, head_dim]
    1267
            // output = attn_weights @ V -> [seq_len, head_dim]
    1268
    24
            scheduler
    1269
    24
                .matmul(attn_weights, v, seq_len, seq_len, head_dim)
    1270
    24
                .map_err(|e| RealizarError::UnsupportedOperation {
    1271
    0
                    operation: "batched_attn_v".to_string(),
    1272
    0
                    reason: format!("GPU matmul failed: {e}"),
    1273
    0
                })
    1274
    24
        }
    1275
    1276
        // =========================================================================
    1277
        // IMP-110: Multi-Head Parallel Attention
    1278
        // =========================================================================
    1279
    1280
        /// Reshape tensor from [seq_len, hidden_dim] to [num_heads, seq_len, head_dim]
    1281
        ///
    1282
        /// IMP-110b: Prepares Q/K/V tensors for parallel multi-head processing.
    1283
        /// Original layout stores all head features contiguously per position.
    1284
        /// New layout groups by head for batched matmul operations.
    1285
        ///
    1286
        /// # Arguments
    1287
        /// * `input` - Input tensor [seq_len, hidden_dim]
    1288
        /// * `seq_len` - Sequence length
    1289
        /// * `num_heads` - Number of attention heads
    1290
        /// * `head_dim` - Dimension per head (hidden_dim / num_heads)
    1291
        ///
    1292
        /// # Returns
    1293
        /// Reshaped tensor [num_heads, seq_len, head_dim]
    1294
        #[cfg(feature = "gpu")]
    1295
    42
        pub fn reshape_for_parallel_heads(
    1296
    42
            &self,
    1297
    42
            input: &[f32],
    1298
    42
            seq_len: usize,
    1299
    42
            num_heads: usize,
    1300
    42
            head_dim: usize,
    1301
    42
        ) -> Result<Vec<f32>> {
    1302
    42
            let hidden_dim = num_heads * head_dim;
    1303
    42
            let expected_len = seq_len * hidden_dim;
    1304
    1305
    42
            if input.len() != expected_len {
    1306
    0
                return Err(RealizarError::InvalidShape {
    1307
    0
                    reason: format!(
    1308
    0
                        "Input size {} doesn't match seq_len={} * hidden_dim={}={}",
    1309
    0
                        input.len(),
    1310
    0
                        seq_len,
    1311
    0
                        hidden_dim,
    1312
    0
                        expected_len
    1313
    0
                    ),
    1314
    0
                });
    1315
    42
            }
    1316
    1317
    42
            let mut reshaped = vec![0.0f32; num_heads * seq_len * head_dim];
    1318
    1319
            // Transform: input[pos * hidden_dim + h * head_dim + d]
    1320
            //         -> reshaped[h * seq_len * head_dim + pos * head_dim + d]
    1321
    204
            for h in 0..num_heads42 {
    1322
    5.90k
                for pos in 0..seq_len204 {
    1323
    91.0k
                    for d in 0..head_dim5.90k {
    1324
    91.0k
                        let orig_idx = pos * hidden_dim + h * head_dim + d;
    1325
    91.0k
                        let new_idx = h * seq_len * head_dim + pos * head_dim + d;
    1326
    91.0k
                        reshaped[new_idx] = input[orig_idx];
    1327
    91.0k
                    }
    1328
                }
    1329
            }
    1330
    1331
    42
            Ok(reshaped)
    1332
    42
        }
    1333
    1334
        /// Compute batched Q@K^T scores for all heads in parallel
    1335
        ///
    1336
        /// IMP-110c: Computes attention scores for all heads in a single batch.
    1337
        /// Takes Q, K in original [seq_len, hidden_dim] layout and computes
    1338
        /// Q@K^T for each head.
    1339
        ///
    1340
        /// # Arguments
    1341
        /// * `q` - Query tensor [seq_len, hidden_dim]
    1342
        /// * `k` - Key tensor [seq_len, hidden_dim]
    1343
        /// * `seq_len` - Sequence length
    1344
        /// * `num_heads` - Number of attention heads
    1345
        /// * `head_dim` - Dimension per head
    1346
        /// * `scale` - Attention scale (1/sqrt(head_dim))
    1347
        ///
    1348
        /// # Returns
    1349
        /// Batched scores [num_heads, seq_len, seq_len]
    1350
        #[cfg(feature = "gpu")]
    1351
    3
        pub fn parallel_batched_qk_scores(
    1352
    3
            &self,
    1353
    3
            q: &[f32],
    1354
    3
            k: &[f32],
    1355
    3
            seq_len: usize,
    1356
    3
            num_heads: usize,
    1357
    3
            head_dim: usize,
    1358
    3
            scale: f32,
    1359
    3
        ) -> Result<Vec<f32>> {
    1360
            use crate::gpu::HybridScheduler;
    1361
    1362
            // Reshape Q and K to [num_heads, seq_len, head_dim]
    1363
    3
            let q_reshaped = self.reshape_for_parallel_heads(q, seq_len, num_heads, head_dim)?0;
    1364
    3
            let k_reshaped = self.reshape_for_parallel_heads(k, seq_len, num_heads, head_dim)?0;
    1365
    1366
    3
            let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
    1367
    0
                RealizarError::UnsupportedOperation {
    1368
    0
                    operation: "HybridScheduler::with_threshold".to_string(),
    1369
    0
                    reason: format!("GPU scheduler initialization failed: {e}"),
    1370
    0
                }
    1371
    0
            })?;
    1372
    1373
            // For each head: Q_h @ K_h^T -> [seq_len, seq_len]
    1374
            // Total output: [num_heads, seq_len, seq_len]
    1375
    3
            let mut all_scores = Vec::with_capacity(num_heads * seq_len * seq_len);
    1376
    1377
    12
            for h in 0..num_heads3 {
    1378
    12
                let head_start = h * seq_len * head_dim;
    1379
    12
                let q_h = &q_reshaped[head_start..head_start + seq_len * head_dim];
    1380
    12
                let k_h = &k_reshaped[head_start..head_start + seq_len * head_dim];
    1381
    1382
                // Transpose K_h: [seq_len, head_dim] -> [head_dim, seq_len]
    1383
    12
                let mut k_t = vec![0.0f32; head_dim * seq_len];
    1384
    80
                for i in 0..seq_len12 {
    1385
    1.15k
                    for j in 0..head_dim80 {
    1386
    1.15k
                        k_t[j * seq_len + i] = k_h[i * head_dim + j];
    1387
    1.15k
                    }
    1388
                }
    1389
    1390
                // Q_h @ K_h^T: [seq_len, head_dim] @ [head_dim, seq_len] -> [seq_len, seq_len]
    1391
    12
                let scores = scheduler
    1392
    12
                    .matmul(q_h, &k_t, seq_len, head_dim, seq_len)
    1393
    12
                    .map_err(|e| RealizarError::UnsupportedOperation {
    1394
    0
                        operation: "parallel_batched_qk_scores".to_string(),
    1395
    0
                        reason: format!("GPU matmul failed: {e}"),
    1396
    0
                    })?;
    1397
    1398
                // Apply scale and accumulate
    1399
    588
                for s576 in &scores {
    1400
    576
                    all_scores.push(s * scale);
    1401
    576
                }
    1402
            }
    1403
    1404
    3
            Ok(all_scores)
    1405
    3
        }
    1406
    1407
        /// Multi-head attention with parallel head processing
    1408
        ///
    1409
        /// IMP-110a: Processes all attention heads in parallel batches instead
    1410
        /// of iterating head-by-head. This enables better GPU utilization.
    1411
        ///
    1412
        /// # Arguments
    1413
        /// * `q` - Query tensor [seq_len, hidden_dim]
    1414
        /// * `k` - Key tensor [seq_len, hidden_dim]
    1415
        /// * `v` - Value tensor [seq_len, hidden_dim]
    1416
        /// * `seq_len` - Sequence length
    1417
        ///
    1418
        /// # Returns
    1419
        /// Attention output [seq_len, hidden_dim]
    1420
        #[cfg(feature = "gpu")]
    1421
    2
        pub fn parallel_multihead_attention_gpu(
    1422
    2
            &self,
    1423
    2
            q: &[f32],
    1424
    2
            k: &[f32],
    1425
    2
            v: &[f32],
    1426
    2
            seq_len: usize,
    1427
    2
        ) -> Result<Vec<f32>> {
    1428
            use crate::gpu::HybridScheduler;
    1429
    1430
    2
            let hidden_dim = self.config.hidden_dim;
    1431
    2
            let num_heads = self.config.num_heads;
    1432
    2
            let head_dim = hidden_dim / num_heads;
    1433
    2
            let scale = 1.0 / (head_dim as f32).sqrt();
    1434
    1435
            // Get batched scores for all heads: [num_heads, seq_len, seq_len]
    1436
    2
            let batched_scores =
    1437
    2
                self.parallel_batched_qk_scores(q, k, seq_len, num_heads, head_dim, scale)?0;
    1438
    1439
            // Apply causal mask and softmax per head
    1440
    2
            let mut batched_weights = vec![0.0f32; num_heads * seq_len * seq_len];
    1441
    8
            for h in 0..num_heads2 {
    1442
    8
                let head_offset = h * seq_len * seq_len;
    1443
    8
                let head_scores = &batched_scores[head_offset..head_offset + seq_len * seq_len];
    1444
    8
                let head_weights = self.apply_causal_mask_softmax(head_scores, seq_len);
    1445
    8
                batched_weights[head_offset..head_offset + seq_len * seq_len]
    1446
    8
                    .copy_from_slice(&head_weights);
    1447
    8
            }
    1448
    1449
            // Reshape V to [num_heads, seq_len, head_dim]
    1450
    2
            let v_reshaped = self.reshape_for_parallel_heads(v, seq_len, num_heads, head_dim)?0;
    1451
    1452
            // Compute attention output for all heads
    1453
    2
            let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
    1454
    0
                RealizarError::UnsupportedOperation {
    1455
    0
                    operation: "HybridScheduler::with_threshold".to_string(),
    1456
    0
                    reason: format!("GPU scheduler initialization failed: {e}"),
    1457
    0
                }
    1458
    0
            })?;
    1459
    1460
            // Output: [seq_len, hidden_dim]
    1461
    2
            let mut output = vec![0.0f32; seq_len * hidden_dim];
    1462
    1463
    8
            for h in 0..num_heads2 {
    1464
    8
                let weights_offset = h * seq_len * seq_len;
    1465
    8
                let v_offset = h * seq_len * head_dim;
    1466
    1467
    8
                let head_weights = &batched_weights[weights_offset..weights_offset + seq_len * seq_len];
    1468
    8
                let v_h = &v_reshaped[v_offset..v_offset + seq_len * head_dim];
    1469
    1470
                // weights @ V_h: [seq_len, seq_len] @ [seq_len, head_dim] -> [seq_len, head_dim]
    1471
    8
                let head_output = scheduler
    1472
    8
                    .matmul(head_weights, v_h, seq_len, seq_len, head_dim)
    1473
    8
                    .map_err(|e| RealizarError::UnsupportedOperation {
    1474
    0
                        operation: "parallel_multihead_attention_gpu".to_string(),
    1475
    0
                        reason: format!("GPU matmul failed: {e}"),
    1476
    0
                    })?;
    1477
    1478
                // Copy to output in original layout
    1479
    64
                for pos in 0..seq_len8 {
    1480
    64
                    let out_start = pos * hidden_dim + h * head_dim;
    1481
    64
                    let head_start = pos * head_dim;
    1482
    64
                    output[out_start..out_start + head_dim]
    1483
    64
                        .copy_from_slice(&head_output[head_start..head_start + head_dim]);
    1484
    64
                }
    1485
            }
    1486
    1487
    2
            Ok(output)
    1488
    2
        }
    1489
    1490
        // =========================================================================
    1491
        // IMP-111: Flash Attention-style Tiled Computation
    1492
        // =========================================================================
    1493
    1494
        /// Standard softmax (reference implementation)
    1495
        ///
    1496
        /// IMP-111a: Reference implementation for testing online softmax.
    1497
        /// Computes softmax in the standard way: exp(x - max) / sum(exp(x - max))
    1498
    9
        pub fn standard_softmax(&self, scores: &[f32]) -> Vec<f32> {
    1499
    9
            if scores.is_empty() {
    1500
    0
                return Vec::new();
    1501
    9
            }
    1502
    1503
            // Find max for numerical stability
    1504
    9
            let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1505
    1506
            // Compute exp(x - max) and sum
    1507
    80
            let exp_scores9: Vec<f32>9 = scores9.iter9().map9(|&s| (s - max_score).exp()).collect9();
    1508
    9
            let sum: f32 = exp_scores.iter().sum();
    1509
    1510
            // Normalize
    1511
    80
            exp_scores.iter()9.map9(|&e| e / sum).collect9()
    1512
    9
        }
    1513
    1514
        /// Online softmax with tiled processing (O(1) memory per tile)
    1515
        ///
    1516
        /// IMP-111a: Implements the "online softmax" algorithm that processes
    1517
        /// data in tiles without materializing the full softmax denominator.
    1518
        ///
    1519
        /// Algorithm:
    1520
        /// 1. Process tiles, tracking running max (m) and denominator (d)
    1521
        /// 2. When new tile has larger max, rescale previous denominator
    1522
        /// 3. Final pass normalizes all values
    1523
        ///
    1524
        /// # Arguments
    1525
        /// * `scores` - Input scores to apply softmax
    1526
        /// * `tile_size` - Size of each tile for processing
    1527
        ///
    1528
        /// # Returns
    1529
        /// Softmax probabilities
    1530
    1
        pub fn online_softmax(&self, scores: &[f32], tile_size: usize) -> Result<Vec<f32>> {
    1531
    1
            if scores.is_empty() {
    1532
    0
                return Ok(Vec::new());
    1533
    1
            }
    1534
    1535
    1
            let n = scores.len();
    1536
    1
            let tile_size = tile_size.max(1);
    1537
    1538
            // Running statistics
    1539
    1
            let mut global_max = f32::NEG_INFINITY;
    1540
    1
            let mut global_sum = 0.0f32;
    1541
    1542
            // First pass: compute global max and sum using online algorithm
    1543
    4
            for tile_start in (0..n)1.step_by1(tile_size1) {
    1544
    4
                let tile_end = (tile_start + tile_size).min(n);
    1545
    1546
                // Find local max in this tile
    1547
    4
                let local_max = scores[tile_start..tile_end]
    1548
    4
                    .iter()
    1549
    4
                    .cloned()
    1550
    4
                    .fold(f32::NEG_INFINITY, f32::max);
    1551
    1552
    4
                if local_max > global_max {
    1553
    2
                    // Rescale previous sum when we find a new max
    1554
    2
                    let rescale = (global_max - local_max).exp();
    1555
    2
                    global_sum *= rescale;
    1556
    2
                    global_max = local_max;
    1557
    2
                }
    1558
    1559
                // Add this tile's contribution to sum
    1560
    16
                for &s in &scores4[tile_start..tile_end]4 {
    1561
    16
                    global_sum += (s - global_max).exp();
    1562
    16
                }
    1563
            }
    1564
    1565
            // Second pass: compute final softmax values
    1566
    1
            let mut result = Vec::with_capacity(n);
    1567
    17
            for &s16 in scores {
    1568
    16
                result.push((s - global_max).exp() / global_sum);
    1569
    16
            }
    1570
    1571
    1
            Ok(result)
    1572
    1
        }
    1573
    1574
        /// Standard single-head attention (reference implementation)
    1575
        ///
    1576
        /// IMP-111b: Reference implementation that materializes full attention matrix.
    1577
        /// Used to verify tiled attention correctness.
    1578
        ///
    1579
        /// # Arguments
    1580
        /// * `q` - Query tensor [seq_len, head_dim]
    1581
        /// * `k` - Key tensor [seq_len, head_dim]
    1582
        /// * `v` - Value tensor [seq_len, head_dim]
    1583
        /// * `seq_len` - Sequence length
    1584
        /// * `head_dim` - Dimension per head
    1585
        /// * `scale` - Attention scale (1/sqrt(head_dim))
    1586
    1
        pub fn standard_single_head_attention(
    1587
    1
            &self,
    1588
    1
            q: &[f32],
    1589
    1
            k: &[f32],
    1590
    1
            v: &[f32],
    1591
    1
            seq_len: usize,
    1592
    1
            head_dim: usize,
    1593
    1
            scale: f32,
    1594
    1
        ) -> Result<Vec<f32>> {
    1595
            // Compute attention scores: Q @ K^T -> [seq_len, seq_len]
    1596
    1
            let mut scores = vec![0.0f32; seq_len * seq_len];
    1597
    8
            for i in 0..seq_len1 {
    1598
    64
                for j in 0..seq_len8 {
    1599
    64
                    let mut dot = 0.0f32;
    1600
    512
                    for d in 0..head_dim64 {
    1601
    512
                        dot += q[i * head_dim + d] * k[j * head_dim + d];
    1602
    512
                    }
    1603
    64
                    scores[i * seq_len + j] = dot * scale;
    1604
                }
    1605
            }
    1606
    1607
            // Apply softmax per row
    1608
    1
            let mut weights = vec![0.0f32; seq_len * seq_len];
    1609
    8
            for i in 0..seq_len1 {
    1610
    8
                let row_start = i * seq_len;
    1611
    8
                let row = &scores[row_start..row_start + seq_len];
    1612
    8
                let softmax = self.standard_softmax(row);
    1613
    8
                weights[row_start..row_start + seq_len].copy_from_slice(&softmax);
    1614
    8
            }
    1615
    1616
            // Compute output: weights @ V -> [seq_len, head_dim]
    1617
    1
            let mut output = vec![0.0f32; seq_len * head_dim];
    1618
    8
            for i in 0..seq_len1 {
    1619
    64
                for d in 0..head_dim8 {
    1620
    64
                    let mut acc = 0.0f32;
    1621
    512
                    for j in 0..seq_len64 {
    1622
    512
                        acc += weights[i * seq_len + j] * v[j * head_dim + d];
    1623
    512
                    }
    1624
    64
                    output[i * head_dim + d] = acc;
    1625
                }
    1626
            }
    1627
    1628
    1
            Ok(output)
    1629
    1
        }
    1630
    1631
        /// Tiled single-head attention (non-causal)
    1632
        ///
    1633
        /// IMP-111b: Flash Attention-style tiled computation.
    1634
        /// Processes K/V in tiles, maintaining running softmax statistics.
    1635
        #[allow(clippy::too_many_arguments)]
    1636
    1
        pub fn tiled_single_head_attention(
    1637
    1
            &self,
    1638
    1
            q: &[f32],
    1639
    1
            k: &[f32],
    1640
    1
            v: &[f32],
    1641
    1
            seq_len: usize,
    1642
    1
            head_dim: usize,
    1643
    1
            scale: f32,
    1644
    1
            tile_size: usize,
    1645
    1
        ) -> Result<Vec<f32>> {
    1646
    1
            let tile_size = tile_size.max(1);
    1647
    1
            let mut output = vec![0.0f32; seq_len * head_dim];
    1648
    1649
            // Process each query position
    1650
    8
            for i in 0..seq_len1 {
    1651
    8
                let q_i = &q[i * head_dim..(i + 1) * head_dim];
    1652
    1653
                // Running statistics for online softmax
    1654
    8
                let mut running_max = f32::NEG_INFINITY;
    1655
    8
                let mut running_sum = 0.0f32;
    1656
    8
                let mut running_output = vec![0.0f32; head_dim];
    1657
    1658
                // Process K/V in tiles
    1659
    16
                for tile_start in (0..seq_len)8.step_by8(tile_size8) {
    1660
    16
                    let tile_end = (tile_start + tile_size).min(seq_len);
    1661
    1662
                    // Compute scores for this tile: q_i @ K_tile^T
    1663
    16
                    let mut tile_scores = Vec::with_capacity(tile_end - tile_start);
    1664
    64
                    for j in tile_start16..tile_end16 {
    1665
    64
                        let mut dot = 0.0f32;
    1666
    512
                        for d in 0..head_dim64 {
    1667
    512
                            dot += q_i[d] * k[j * head_dim + d];
    1668
    512
                        }
    1669
    64
                        tile_scores.push(dot * scale);
    1670
                    }
    1671
    1672
                    // Find tile max
    1673
    16
                    let tile_max = tile_scores
    1674
    16
                        .iter()
    1675
    16
                        .cloned()
    1676
    16
                        .fold(f32::NEG_INFINITY, f32::max);
    1677
    1678
                    // Update running statistics
    1679
    16
                    let new_max = running_max.max(tile_max);
    1680
    1681
                    // Rescale previous output and sum
    1682
    16
                    if new_max > running_max && running_sum > 0.011 {
    1683
    3
                        let rescale = (running_max - new_max).exp();
    1684
    3
                        running_sum *= rescale;
    1685
    27
                        for out_val24 in &mut running_output {
    1686
    24
                            *out_val *= rescale;
    1687
    24
                        }
    1688
    13
                    }
    1689
    16
                    running_max = new_max;
    1690
    1691
                    // Accumulate this tile's contribution
    1692
    64
                    for (idx, &score) in tile_scores.iter()16.enumerate16() {
    1693
    64
                        let j = tile_start + idx;
    1694
    64
                        let weight = (score - running_max).exp();
    1695
    64
                        running_sum += weight;
    1696
    512
                        for d in 0..head_dim64 {
    1697
    512
                            running_output[d] += weight * v[j * head_dim + d];
    1698
    512
                        }
    1699
                    }
    1700
                }
    1701
    1702
                // Normalize output
    1703
    64
                for d in 0..head_dim8 {
    1704
    64
                    output[i * head_dim + d] = running_output[d] / running_sum;
    1705
    64
                }
    1706
            }
    1707
    1708
    1
            Ok(output)
    1709
    1
        }
    1710
    1711
        /// Tiled causal attention
    1712
        ///
    1713
        /// IMP-111c: Flash Attention with causal masking.
    1714
        /// For position i, only attends to positions 0..=i.
    1715
        #[allow(clippy::too_many_arguments)]
    1716
    29
        pub fn tiled_causal_attention(
    1717
    29
            &self,
    1718
    29
            q: &[f32],
    1719
    29
            k: &[f32],
    1720
    29
            v: &[f32],
    1721
    29
            seq_len: usize,
    1722
    29
            head_dim: usize,
    1723
    29
            scale: f32,
    1724
    29
            tile_size: usize,
    1725
    29
        ) -> Result<Vec<f32>> {
    1726
    29
            let tile_size = tile_size.max(1);
    1727
    29
            let mut output = vec![0.0f32; seq_len * head_dim];
    1728
    1729
            // Process each query position
    1730
    540
            for i in 0..seq_len29 {
    1731
    540
                let q_i = &q[i * head_dim..(i + 1) * head_dim];
    1732
    1733
                // Running statistics for online softmax
    1734
    540
                let mut running_max = f32::NEG_INFINITY;
    1735
    540
                let mut running_sum = 0.0f32;
    1736
    540
                let mut running_output = vec![0.0f32; head_dim];
    1737
    1738
                // Only process K/V up to position i (causal)
    1739
    540
                let causal_len = i + 1;
    1740
    1741
                // Process K/V in tiles
    1742
    2.00k
                for tile_start in (0..causal_len)540.step_by540(tile_size540) {
    1743
    2.00k
                    let tile_end = (tile_start + tile_size).min(causal_len);
    1744
    1745
                    // Compute scores for this tile: q_i @ K_tile^T
    1746
    2.00k
                    let mut tile_scores = Vec::with_capacity(tile_end - tile_start);
    1747
    6.83k
                    for j in tile_start2.00k..tile_end2.00k {
    1748
    6.83k
                        let mut dot = 0.0f32;
    1749
    103k
                        for d in 0..head_dim6.83k {
    1750
    103k
                            dot += q_i[d] * k[j * head_dim + d];
    1751
    103k
                        }
    1752
    6.83k
                        tile_scores.push(dot * scale);
    1753
                    }
    1754
    1755
                    // Find tile max
    1756
    2.00k
                    let tile_max = tile_scores
    1757
    2.00k
                        .iter()
    1758
    2.00k
                        .cloned()
    1759
    2.00k
                        .fold(f32::NEG_INFINITY, f32::max);
    1760
    1761
                    // Update running statistics
    1762
    2.00k
                    let new_max = running_max.max(tile_max);
    1763
    1764
                    // Rescale previous output and sum
    1765
    2.00k
                    if new_max > running_max && running_sum > 0.0920 {
    1766
    380
                        let rescale = (running_max - new_max).exp();
    1767
    380
                        running_sum *= rescale;
    1768
    6.11k
                        for out_val5.73k in &mut running_output {
    1769
    5.73k
                            *out_val *= rescale;
    1770
    5.73k
                        }
    1771
    1.62k
                    }
    1772
    2.00k
                    running_max = new_max;
    1773
    1774
                    // Accumulate this tile's contribution
    1775
    6.83k
                    for (idx, &score) in tile_scores.iter()2.00k.enumerate2.00k() {
    1776
    6.83k
                        let j = tile_start + idx;
    1777
    6.83k
                        let weight = (score - running_max).exp();
    1778
    6.83k
                        running_sum += weight;
    1779
    103k
                        for d in 0..head_dim6.83k {
    1780
    103k
                            running_output[d] += weight * v[j * head_dim + d];
    1781
    103k
                        }
    1782
                    }
    1783
                }
    1784
    1785
                // Normalize output
    1786
    540
                if running_sum > 0.0 {
    1787
    7.84k
                    for d in 0..head_dim540 {
    1788
    7.84k
                        output[i * head_dim + d] = running_output[d] / running_sum;
    1789
    7.84k
                    }
    1790
    0
                }
    1791
            }
    1792
    1793
    29
            Ok(output)
    1794
    29
        }
    1795
    }