## Coverage Report

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

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

    /home/noah/src/realizar/src/gpu/scheduler/kv.rs

    Line
    Count
    Source
    1
    //! KV Cache Management for GpuModel (PMAT-802)
    2
    //!
    3
    //! Extracted from model.rs to reduce module size.
    4
    //! Contains KV cache forward pass and generation logic.
    5
    6
    use crate::error::{RealizarError, Result};
    7
    use super::super::{StreamingKVCache, exceeds_gpu_buffer_limit, cpu_matmul_transposed_simd};
    8
    use super::model::{GpuModel, GpuGenerateConfig};
    9
    10
    /// Apply Rotary Position Embedding (RoPE) to Q or K vectors (Phase 21)
    11
    ///
    12
    /// RoPE encodes position information by rotating pairs of elements
    13
    /// with position-dependent angles. This is CRITICAL for transformer attention.
    14
    ///
    15
    /// # Arguments
    16
    /// * `x` - Mutable slice of Q or K vectors [seq_len * num_heads * head_dim]
    17
    /// * `seq_len` - Number of positions to encode
    18
    /// * `num_heads` - Number of attention heads in this tensor
    19
    /// * `head_dim` - Dimension per head
    20
    /// * `rope_theta` - Base frequency (typically 10000.0)
    21
    /// * `start_pos` - Starting position for RoPE (0 for prefill, cache_len for incremental)
    22
    456
    fn apply_rope(
    23
    456
        x: &mut [f32],
    24
    456
        seq_len: usize,
    25
    456
        num_heads: usize,
    26
    456
        head_dim: usize,
    27
    456
        rope_theta: f32,
    28
    456
        start_pos: usize,
    29
    456
    ) {
    30
    456
        let half_dim = head_dim / 2;
    31
    456
        let head_dim_f32 = head_dim as f32;
    32
    456
        let total_dim = num_heads * head_dim;
    33
    34
    1.45k
        for pos in 0..seq_len456 {
    35
    1.45k
            let position = start_pos + pos;
    36
    1.45k
            let pos_f32 = position as f32;
    37
    1.45k
            let pos_offset = pos * total_dim;
    38
    39
    9.84k
            for h in 0..num_heads1.45k {
    40
    9.84k
                let head_start = pos_offset + h * head_dim;
    41
    9.84k
                let idx2_start = head_start + half_dim;
    42
    43
    79.1k
                for i in 0..half_dim9.84k {
    44
    79.1k
                    let freq = 1.0 / rope_theta.powf(2.0 * i as f32 / head_dim_f32);
    45
    79.1k
                    let angle = pos_f32 * freq;
    46
    79.1k
                    let (sin_val, cos_val) = angle.sin_cos();
    47
    79.1k
    48
    79.1k
                    let x1 = x[head_start + i];
    49
    79.1k
                    let x2 = x[idx2_start + i];
    50
    79.1k
    51
    79.1k
                    // Apply rotation: [cos -sin; sin cos] * [x1; x2]
    52
    79.1k
                    x[head_start + i] = x1 * cos_val - x2 * sin_val;
    53
    79.1k
                    x[idx2_start + i] = x1 * sin_val + x2 * cos_val;
    54
    79.1k
                }
    55
            }
    56
        }
    57
    456
    }
    58
    59
    /// Forward pass with KV cache population (IMP-031)
    60
    36
    pub fn forward_gpu_with_cache(
    61
    36
        model: &mut GpuModel,
    62
    36
        token_ids: &[usize],
    63
    36
        kv_cache: &mut StreamingKVCache,
    64
    36
    ) -> Result<Vec<f32>> {
    65
    36
        if token_ids.is_empty() {
    66
    0
            return Err(RealizarError::InvalidShape {
    67
    0
                reason: "Token IDs cannot be empty".to_string(),
    68
    0
            });
    69
    36
        }
    70
    71
    36
        let seq_len = token_ids.len();
    72
    36
        let hidden_dim = model.config.hidden_dim;
    73
    74
        // Step 1: Embed tokens
    75
    36
        let mut hidden = Vec::with_capacity(seq_len * hidden_dim);
    76
    238
        for &token_id202 in token_ids {
    77
    202
            if token_id >= model.config.vocab_size {
    78
    0
                return Err(RealizarError::InvalidShape {
    79
    0
                    reason: format!(
    80
    0
                        "Token ID {} out of bounds (vocab_size={})",
    81
    0
                        token_id, model.config.vocab_size
    82
    0
                    ),
    83
    0
                });
    84
    202
            }
    85
    202
            let offset = token_id * hidden_dim;
    86
    202
            hidden.extend_from_slice(&model.embedding_weights[offset..offset + hidden_dim]);
    87
        }
    88
    89
        // Step 2: Pass through transformer blocks with KV cache population
    90
    96
        for block_idx in 0..model.block_weights36.len36() {
    91
    96
            hidden = forward_block_with_cache(model, &hidden, seq_len, block_idx, kv_cache)?0;
    92
        }
    93
    94
        // Step 3: Final layer norm
    95
    36
        hidden = layer_norm_kv(model, &hidden);
    96
    97
        // Step 4: LM head projection - only for final position
    98
    36
        let final_hidden = &hidden[(seq_len - 1) * hidden_dim..seq_len * hidden_dim];
    99
    36
        let lm_head_elements = hidden_dim * model.config.vocab_size;
    100
    36
        let output = if exceeds_gpu_buffer_limit(lm_head_elements) {
    101
    0
            cpu_matmul_transposed_simd(
    102
    0
                final_hidden,
    103
    0
                &model.lm_head_weight_t,
    104
    0
                &model.lm_head_bias,
    105
    0
                hidden_dim,
    106
    0
                model.config.vocab_size,
    107
            )
    108
        } else {
    109
    36
            let logits = model.scheduler.matmul(
    110
    36
                final_hidden,
    111
    36
                &model.lm_head_weight,
    112
                1,
    113
    36
                hidden_dim,
    114
    36
                model.config.vocab_size,
    115
    0
            )?;
    116
    36
            let mut output = logits;
    117
    9.21k
            for (out_val, bias_val) in output.iter_mut()36.zip36(model.lm_head_bias.iter()36) {
    118
    9.21k
                *out_val += *bias_val;
    119
    9.21k
            }
    120
    36
            output
    121
        };
    122
    123
    36
        Ok(output)
    124
    36
    }
    125
    126
    /// Incremental forward pass using cached KV (IMP-032)
    127
    35
    pub fn forward_gpu_incremental(
    128
    35
        model: &mut GpuModel,
    129
    35
        token_id: usize,
    130
    35
        kv_cache: &mut StreamingKVCache,
    131
    35
    ) -> Result<Vec<f32>> {
    132
    35
        if token_id >= model.config.vocab_size {
    133
    0
            return Err(RealizarError::InvalidShape {
    134
    0
                reason: format!(
    135
    0
                    "Token ID {} out of bounds (vocab_size={})",
    136
    0
                    token_id, model.config.vocab_size
    137
    0
                ),
    138
    0
            });
    139
    35
        }
    140
    141
    35
        let hidden_dim = model.config.hidden_dim;
    142
    143
        // Step 1: Embed token
    144
    35
        let offset = token_id * hidden_dim;
    145
    35
        let mut hidden = model.embedding_weights[offset..offset + hidden_dim].to_vec();
    146
    147
        // Step 2: Pass through transformer blocks using KV cache
    148
    132
        for block_idx in 0..model.block_weights35.len35() {
    149
    132
            hidden = forward_block_incremental(model, &hidden, block_idx, kv_cache)?0;
    150
        }
    151
    152
        // Step 3: Final layer norm
    153
    35
        hidden = layer_norm_kv(model, &hidden);
    154
    155
        // Step 4: LM head projection
    156
    35
        let lm_head_elements = hidden_dim * model.config.vocab_size;
    157
    35
        let output = if exceeds_gpu_buffer_limit(lm_head_elements) {
    158
    0
            cpu_matmul_transposed_simd(
    159
    0
                &hidden,
    160
    0
                &model.lm_head_weight_t,
    161
    0
                &model.lm_head_bias,
    162
    0
                hidden_dim,
    163
    0
                model.config.vocab_size,
    164
            )
    165
        } else {
    166
    35
            let logits = model.scheduler.matmul(
    167
    35
                &hidden,
    168
    35
                &model.lm_head_weight,
    169
                1,
    170
    35
                hidden_dim,
    171
    35
                model.config.vocab_size,
    172
    0
            )?;
    173
    35
            let mut output = logits;
    174
    8.96k
            for (out_val, bias_val) in output.iter_mut()35.zip35(model.lm_head_bias.iter()35) {
    175
    8.96k
                *out_val += *bias_val;
    176
    8.96k
            }
    177
    35
            output
    178
        };
    179
    180
    35
        Ok(output)
    181
    35
    }
    182
    183
    /// Forward pass through a single block with KV cache population
    184
    96
    fn forward_block_with_cache(
    185
    96
        model: &mut GpuModel,
    186
    96
        input: &[f32],
    187
    96
        seq_len: usize,
    188
    96
        block_idx: usize,
    189
    96
        kv_cache: &mut StreamingKVCache,
    190
    96
    ) -> Result<Vec<f32>> {
    191
    96
        let hidden_dim = model.config.hidden_dim;
    192
    96
        let intermediate_dim = model.config.intermediate_dim;
    193
    96
        let num_heads = model.config.num_heads;
    194
    96
        let num_kv_heads = model.config.num_kv_heads;
    195
    96
        let head_dim = model.config.head_dim();
    196
    96
        let kv_dim = model.config.kv_dim();
    197
    96
        let qkv_dim = model.config.qkv_dim();
    198
    199
    96
        let block = &model.block_weights[block_idx];
    200
    201
        // Pre-norm
    202
    96
        let normed = GpuModel::layer_norm_static(
    203
    96
            input,
    204
    96
            &block.attn_norm_weight,
    205
    96
            &block.attn_norm_bias,
    206
    96
            hidden_dim,
    207
    96
            model.config.eps,
    208
        );
    209
    210
        // QKV projection
    211
    96
        let mut qkv = model.scheduler.matmul(
    212
    96
            &normed,
    213
    96
            &model.block_weights[block_idx].qkv_weight,
    214
    96
            seq_len,
    215
    96
            hidden_dim,
    216
    96
            qkv_dim,
    217
    0
        )?;
    218
    219
        // Split Q, K, V (mutable for RoPE application)
    220
    96
        let q_end = seq_len * hidden_dim;
    221
    96
        let k_end = q_end + seq_len * kv_dim;
    222
    223
        // Phase 21: Apply RoPE to Q and K BEFORE caching
    224
        // This is CRITICAL - without RoPE, attention has no position information
    225
    96
        let rope_theta = model.config.rope_theta;
    226
    227
        // Apply RoPE to Q (all heads)
    228
    96
        apply_rope(&mut qkv[..q_end], seq_len, num_heads, head_dim, rope_theta, 0);
    229
    230
        // Apply RoPE to K (KV heads)
    231
    96
        apply_rope(&mut qkv[q_end..k_end], seq_len, num_kv_heads, head_dim, rope_theta, 0);
    232
    233
        // Now split (after RoPE applied)
    234
    96
        let q = &qkv[..q_end];
    235
    96
        let k = &qkv[q_end..k_end];
    236
    96
        let v = &qkv[k_end..];
    237
    238
        // Cache K (with RoPE) and V
    239
    596
        for pos in 0..seq_len96 {
    240
    596
            let k_slice = &k[pos * kv_dim..(pos + 1) * kv_dim];
    241
    596
            let v_slice = &v[pos * kv_dim..(pos + 1) * kv_dim];
    242
    596
            kv_cache.append(block_idx, k_slice, v_slice);
    243
    596
        }
    244
    245
        // GQA attention
    246
    96
        let attn_out = gqa_attention_with_kv(model, q, k, v, seq_len, num_heads, num_kv_heads, head_dim)?0;
    247
    248
        // Output projection
    249
    96
        let projected = model.scheduler.matmul(
    250
    96
            &attn_out,
    251
    96
            &model.block_weights[block_idx].out_weight,
    252
    96
            seq_len,
    253
    96
            hidden_dim,
    254
    96
            hidden_dim,
    255
    0
        )?;
    256
    257
        // Residual 1
    258
    96
        let mut residual1: Vec<f32> = input
    259
    96
            .iter()
    260
    96
            .zip(projected.iter())
    261
    96
            .enumerate()
    262
    62.7k
            .map96(|(i, (&inp, &proj))| {
    263
    62.7k
                inp + proj + model.block_weights[block_idx].out_bias[i % hidden_dim]
    264
    62.7k
            })
    265
    96
            .collect();
    266
    267
        // FFN pre-norm
    268
    96
        let ffn_normed = GpuModel::layer_norm_static(
    269
    96
            &residual1,
    270
    96
            &model.block_weights[block_idx].ffn_norm_weight,
    271
    96
            &model.block_weights[block_idx].ffn_norm_bias,
    272
    96
            hidden_dim,
    273
    96
            model.config.eps,
    274
        );
    275
    276
        // FFN: SwiGLU when gate weight exists, otherwise GELU
    277
    96
        let activated: Vec<f32> = if let Some(ref gate_weight0) = model.block_weights[block_idx].ffn_gate_weight {
    278
            // SwiGLU: silu(gate(x)) * up(x)
    279
    0
            let up_out = model.scheduler.matmul(
    280
    0
                &ffn_normed,
    281
    0
                &model.block_weights[block_idx].ffn_fc1_weight,
    282
    0
                seq_len,
    283
    0
                hidden_dim,
    284
    0
                intermediate_dim,
    285
    0
            )?;
    286
    0
            let gate_out = model.scheduler.matmul(
    287
    0
                &ffn_normed,
    288
    0
                gate_weight,
    289
    0
                seq_len,
    290
    0
                hidden_dim,
    291
    0
                intermediate_dim,
    292
    0
            )?;
    293
    294
            // SwiGLU: silu(gate) * up
    295
    0
            up_out
    296
    0
                .iter()
    297
    0
                .zip(gate_out.iter())
    298
    0
                .map(|(&u, &g)| {
    299
    0
                    let silu_g = g / (1.0 + (-g).exp());
    300
    0
                    silu_g * u
    301
    0
                })
    302
    0
                .collect()
    303
        } else {
    304
            // Standard GELU FFN
    305
    96
            let fc1_out = model.scheduler.matmul(
    306
    96
                &ffn_normed,
    307
    96
                &model.block_weights[block_idx].ffn_fc1_weight,
    308
    96
                seq_len,
    309
    96
                hidden_dim,
    310
    96
                intermediate_dim,
    311
    0
            )?;
    312
    313
    96
            fc1_out
    314
    96
                .iter()
    315
    96
                .enumerate()
    316
    125k
                .map96(|(i, &x)| {
    317
    125k
                    let x = x + model.block_weights[block_idx].ffn_fc1_bias[i % intermediate_dim];
    318
    125k
                    0.5 * x * (1.0 + ((2.0f32 / std::f32::consts::PI).sqrt() * (x + 0.044_715 * x.powi(3))).tanh())
    319
    125k
                })
    320
    96
                .collect()
    321
        };
    322
    323
        // FFN: fc2
    324
    96
        let fc2_out = model.scheduler.matmul(
    325
    96
            &activated,
    326
    96
            &model.block_weights[block_idx].ffn_fc2_weight,
    327
    96
            seq_len,
    328
    96
            intermediate_dim,
    329
    96
            hidden_dim,
    330
    0
        )?;
    331
    332
        // Residual 2
    333
    62.7k
        for (i, x) in residual1.iter_mut()96.enumerate96() {
    334
    62.7k
            *x += fc2_out[i] + model.block_weights[block_idx].ffn_fc2_bias[i % hidden_dim];
    335
    62.7k
        }
    336
    337
    96
        Ok(residual1)
    338
    96
    }
    339
    340
    /// Incremental forward pass through a single block using cached KV
    341
    132
    fn forward_block_incremental(
    342
    132
        model: &mut GpuModel,
    343
    132
        input: &[f32],
    344
    132
        block_idx: usize,
    345
    132
        kv_cache: &mut StreamingKVCache,
    346
    132
    ) -> Result<Vec<f32>> {
    347
    132
        let hidden_dim = model.config.hidden_dim;
    348
    132
        let intermediate_dim = model.config.intermediate_dim;
    349
    132
        let num_heads = model.config.num_heads;
    350
    132
        let num_kv_heads = model.config.num_kv_heads;
    351
    132
        let head_dim = model.config.head_dim();
    352
    132
        let kv_dim = model.config.kv_dim();
    353
    132
        let qkv_dim = model.config.qkv_dim();
    354
    355
    132
        let block = &model.block_weights[block_idx];
    356
    357
        // Pre-norm (single position)
    358
    132
        let normed = GpuModel::layer_norm_static(
    359
    132
            input,
    360
    132
            &block.attn_norm_weight,
    361
    132
            &block.attn_norm_bias,
    362
    132
            hidden_dim,
    363
    132
            model.config.eps,
    364
        );
    365
    366
        // QKV projection (single position)
    367
    132
        let mut qkv = model.scheduler.matmul(
    368
    132
            &normed,
    369
    132
            &model.block_weights[block_idx].qkv_weight,
    370
            1,
    371
    132
            hidden_dim,
    372
    132
            qkv_dim,
    373
    0
        )?;
    374
    375
        // Get current position BEFORE caching (this is where new token goes)
    376
    132
        let (existing_k, _) = kv_cache.get_valid(block_idx);
    377
    132
        let current_pos = existing_k.len() / kv_dim;
    378
    379
        // Phase 21: Apply RoPE to Q and K at current position BEFORE caching
    380
    132
        let rope_theta = model.config.rope_theta;
    381
    382
        // Apply RoPE to Q (single position, all heads)
    383
    132
        apply_rope(&mut qkv[..hidden_dim], 1, num_heads, head_dim, rope_theta, current_pos);
    384
    385
        // Apply RoPE to K (single position, KV heads)
    386
    132
        apply_rope(&mut qkv[hidden_dim..hidden_dim + kv_dim], 1, num_kv_heads, head_dim, rope_theta, current_pos);
    387
    388
        // Split Q, K, V (single position, after RoPE)
    389
    132
        let q = &qkv[..hidden_dim];
    390
    132
        let k = &qkv[hidden_dim..hidden_dim + kv_dim];
    391
    132
        let v = &qkv[hidden_dim + kv_dim..];
    392
    393
        // Cache new K (with RoPE) and V
    394
    132
        kv_cache.append(block_idx, k, v);
    395
    396
        // Get all cached K/V for attention (now includes new K/V)
    397
    132
        let (all_k, all_v) = kv_cache.get_valid(block_idx);
    398
    132
        let cache_len = all_k.len() / kv_dim;
    399
    400
        // GQA incremental attention
    401
    132
        let attn_out = gqa_incremental_attention(model, q, all_k, all_v, cache_len, num_heads, num_kv_heads, head_dim)?0;
    402
    403
        // Output projection
    404
    132
        let projected = model.scheduler.matmul(
    405
    132
            &attn_out,
    406
    132
            &model.block_weights[block_idx].out_weight,
    407
            1,
    408
    132
            hidden_dim,
    409
    132
            hidden_dim,
    410
    0
        )?;
    411
    412
        // Residual 1
    413
    132
        let mut residual1: Vec<f32> = input
    414
    132
            .iter()
    415
    132
            .zip(projected.iter())
    416
    132
            .enumerate()
    417
    16.3k
            .map132(|(i, (&inp, &proj))| {
    418
    16.3k
                inp + proj + model.block_weights[block_idx].out_bias[i]
    419
    16.3k
            })
    420
    132
            .collect();
    421
    422
        // FFN pre-norm
    423
    132
        let ffn_normed = GpuModel::layer_norm_static(
    424
    132
            &residual1,
    425
    132
            &model.block_weights[block_idx].ffn_norm_weight,
    426
    132
            &model.block_weights[block_idx].ffn_norm_bias,
    427
    132
            hidden_dim,
    428
    132
            model.config.eps,
    429
        );
    430
    431
        // FFN: SwiGLU when gate weight exists, otherwise GELU
    432
    132
        let activated: Vec<f32> = if let Some(ref gate_weight0) = model.block_weights[block_idx].ffn_gate_weight {
    433
            // SwiGLU: silu(gate(x)) * up(x)
    434
    0
            let up_out = model.scheduler.matmul(
    435
    0
                &ffn_normed,
    436
    0
                &model.block_weights[block_idx].ffn_fc1_weight,
    437
                1,
    438
    0
                hidden_dim,
    439
    0
                intermediate_dim,
    440
    0
            )?;
    441
    0
            let gate_out = model.scheduler.matmul(
    442
    0
                &ffn_normed,
    443
    0
                gate_weight,
    444
                1,
    445
    0
                hidden_dim,
    446
    0
                intermediate_dim,
    447
    0
            )?;
    448
    449
            // SwiGLU: silu(gate) * up
    450
    0
            up_out
    451
    0
                .iter()
    452
    0
                .zip(gate_out.iter())
    453
    0
                .map(|(&u, &g)| {
    454
    0
                    let silu_g = g / (1.0 + (-g).exp());
    455
    0
                    silu_g * u
    456
    0
                })
    457
    0
                .collect()
    458
        } else {
    459
            // Standard GELU FFN
    460
    132
            let fc1_out = model.scheduler.matmul(
    461
    132
                &ffn_normed,
    462
    132
                &model.block_weights[block_idx].ffn_fc1_weight,
    463
                1,
    464
    132
                hidden_dim,
    465
    132
                intermediate_dim,
    466
    0
            )?;
    467
    468
    132
            fc1_out
    469
    132
                .iter()
    470
    132
                .enumerate()
    471
    32.7k
                .map132(|(i, &x)| {
    472
    32.7k
                    let x = x + model.block_weights[block_idx].ffn_fc1_bias[i];
    473
    32.7k
                    0.5 * x * (1.0 + ((2.0f32 / std::f32::consts::PI).sqrt() * (x + 0.044_715 * x.powi(3))).tanh())
    474
    32.7k
                })
    475
    132
                .collect()
    476
        };
    477
    478
        // FFN: fc2
    479
    132
        let fc2_out = model.scheduler.matmul(
    480
    132
            &activated,
    481
    132
            &model.block_weights[block_idx].ffn_fc2_weight,
    482
            1,
    483
    132
            intermediate_dim,
    484
    132
            hidden_dim,
    485
    0
        )?;
    486
    487
        // Residual 2
    488
    16.3k
        for (i, x) in residual1.iter_mut()132.enumerate132() {
    489
    16.3k
            *x += fc2_out[i] + model.block_weights[block_idx].ffn_fc2_bias[i];
    490
    16.3k
        }
    491
    492
    132
        Ok(residual1)
    493
    132
    }
    494
    495
    /// GQA attention with KV (full sequence)
    496
    #[allow(clippy::too_many_arguments)]
    497
    96
    fn gqa_attention_with_kv(
    498
    96
        _model: &GpuModel,
    499
    96
        q: &[f32],
    500
    96
        k: &[f32],
    501
    96
        v: &[f32],
    502
    96
        seq_len: usize,
    503
    96
        num_heads: usize,
    504
    96
        num_kv_heads: usize,
    505
    96
        head_dim: usize,
    506
    96
    ) -> Result<Vec<f32>> {
    507
    96
        let hidden_dim = num_heads * head_dim;
    508
    96
        let kv_dim = num_kv_heads * head_dim;
    509
    96
        let heads_per_kv = num_heads / num_kv_heads;
    510
    511
    96
        let mut output = vec![0.0f32; seq_len * hidden_dim];
    512
    96
        let scale = 1.0 / (head_dim as f32).sqrt();
    513
    514
    596
        for pos in 0..seq_len96 {
    515
    3.90k
            for head in 0..num_heads596 {
    516
    3.90k
                let kv_head = head / heads_per_kv;
    517
    518
                // Query for this head at this position
    519
    3.90k
                let q_start = pos * hidden_dim + head * head_dim;
    520
    3.90k
                let q_slice = &q[q_start..q_start + head_dim];
    521
    522
                // Compute attention scores for all positions up to current
    523
    3.90k
                let mut scores = Vec::with_capacity(pos + 1);
    524
    16.1k
                for kpos in 0..=pos3.90k {
    525
    16.1k
                    let k_start = kpos * kv_dim + kv_head * head_dim;
    526
    16.1k
                    let k_slice = &k[k_start..k_start + head_dim];
    527
    528
    259k
                    let score16.1k: f3216.1k = q_slice16.1k.iter16.1k().zip16.1k(k_slice16.1k.iter16.1k()).map16.1k(|(&a, &b)| a * b).sum16.1k();
    529
    16.1k
                    scores.push(score * scale);
    530
                }
    531
    532
                // Softmax
    533
    3.90k
                let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    534
    16.1k
                let exp_scores3.90k: Vec<f32>3.90k = scores.iter()3.90k.map3.90k(|&s| (s - max_score).exp()).collect3.90k();
    535
    3.90k
                let sum: f32 = exp_scores.iter().sum();
    536
    16.1k
                let weights3.90k: Vec<f32>3.90k = exp_scores.iter()3.90k.map3.90k(|&e| e / sum).collect3.90k();
    537
    538
                // Weighted sum of values
    539
    3.90k
                let out_start = pos * hidden_dim + head * head_dim;
    540
    16.1k
                for (kpos, &weight) in weights.iter()3.90k.enumerate3.90k() {
    541
    16.1k
                    let v_start = kpos * kv_dim + kv_head * head_dim;
    542
    259k
                    for d in 0..head_dim16.1k {
    543
    259k
                        output[out_start + d] += weight * v[v_start + d];
    544
    259k
                    }
    545
                }
    546
            }
    547
        }
    548
    549
    96
        Ok(output)
    550
    96
    }
    551
    552
    /// GQA incremental attention (single query position)
    553
    #[allow(clippy::too_many_arguments)]
    554
    132
    fn gqa_incremental_attention(
    555
    132
        _model: &GpuModel,
    556
    132
        q: &[f32],
    557
    132
        all_k: &[f32],
    558
    132
        all_v: &[f32],
    559
    132
        cache_len: usize,
    560
    132
        num_heads: usize,
    561
    132
        num_kv_heads: usize,
    562
    132
        head_dim: usize,
    563
    132
    ) -> Result<Vec<f32>> {
    564
    132
        let hidden_dim = num_heads * head_dim;
    565
    132
        let kv_dim = num_kv_heads * head_dim;
    566
    132
        let heads_per_kv = num_heads / num_kv_heads;
    567
    568
    132
        let mut output = vec![0.0f32; hidden_dim];
    569
    132
        let scale = 1.0 / (head_dim as f32).sqrt();
    570
    571
    1.02k
        for head in 0..num_heads132 {
    572
    1.02k
            let kv_head = head / heads_per_kv;
    573
    574
    1.02k
            let q_start = head * head_dim;
    575
    1.02k
            let q_slice = &q[q_start..q_start + head_dim];
    576
    577
            // Attention scores for all cached positions
    578
    1.02k
            let mut scores = Vec::with_capacity(cache_len);
    579
    23.2k
            for kpos in 0..cache_len1.02k {
    580
    23.2k
                let k_start = kpos * kv_dim + kv_head * head_dim;
    581
    23.2k
                let k_slice = &all_k[k_start..k_start + head_dim];
    582
    583
    372k
                let score23.2k: f3223.2k = q_slice23.2k.iter23.2k().zip23.2k(k_slice23.2k.iter23.2k()).map23.2k(|(&a, &b)| a * b).sum23.2k();
    584
    23.2k
                scores.push(score * scale);
    585
            }
    586
    587
            // Softmax
    588
    1.02k
            let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    589
    23.2k
            let exp_scores1.02k: Vec<f32>1.02k = scores.iter()1.02k.map1.02k(|&s| (s - max_score).exp()).collect1.02k();
    590
    1.02k
            let sum: f32 = exp_scores.iter().sum();
    591
    23.2k
            let weights1.02k: Vec<f32>1.02k = exp_scores.iter()1.02k.map1.02k(|&e| e / sum).collect1.02k();
    592
    593
            // Weighted sum
    594
    1.02k
            let out_start = head * head_dim;
    595
    23.2k
            for (kpos, &weight) in weights.iter()1.02k.enumerate1.02k() {
    596
    23.2k
                let v_start = kpos * kv_dim + kv_head * head_dim;
    597
    372k
                for d in 0..head_dim23.2k {
    598
    372k
                    output[out_start + d] += weight * all_v[v_start + d];
    599
    372k
                }
    600
            }
    601
        }
    602
    603
    132
        Ok(output)
    604
    132
    }
    605
    606
    /// Generate tokens using KV cache (IMP-033)
    607
    1
    pub fn generate_with_cache(
    608
    1
        model: &mut GpuModel,
    609
    1
        prompt: &[usize],
    610
    1
        config: &GpuGenerateConfig,
    611
    1
    ) -> Result<Vec<usize>> {
    612
    1
        if prompt.is_empty() {
    613
    0
            return Err(RealizarError::InvalidShape {
    614
    0
                reason: "Prompt cannot be empty".to_string(),
    615
    0
            });
    616
    1
        }
    617
    618
    1
        let max_seq_len = prompt.len() + config.max_tokens;
    619
    1
        let head_dim = model.config.hidden_dim / model.config.num_heads;
    620
    1
        let mut kv_cache = StreamingKVCache::new(
    621
    1
            model.config.num_layers,
    622
    1
            max_seq_len,
    623
    1
            model.config.num_kv_heads,
    624
    1
            head_dim,
    625
        );
    626
    627
    1
        let mut tokens = prompt.to_vec();
    628
    1
        let logits = forward_gpu_with_cache(model, prompt, &mut kv_cache)?0;
    629
    630
    1
        let mut next_token = if config.temperature == 0.0 || config.top_k == 10 {
    631
    1
            argmax(&logits)
    632
        } else {
    633
    0
            sample_topk(&logits, config.temperature, config.top_k)
    634
        };
    635
    636
    1
        if config.stop_tokens.contains(&next_token) {
    637
    0
            return Ok(tokens);
    638
    1
        }
    639
    1
        tokens.push(next_token);
    640
    641
    1
        for _ in 1..config.max_tokens {
    642
    31
            let logits = forward_gpu_incremental(model, next_token, &mut kv_cache)?0;
    643
    644
    31
            next_token = if config.temperature == 0.0 || config.top_k == 10 {
    645
    31
                argmax(&logits)
    646
            } else {
    647
    0
                sample_topk(&logits, config.temperature, config.top_k)
    648
            };
    649
    650
    31
            if config.stop_tokens.contains(&next_token) {
    651
    0
                break;
    652
    31
            }
    653
    31
            tokens.push(next_token);
    654
        }
    655
    656
    1
        Ok(tokens)
    657
    1
    }
    658
    659
    /// Layer norm helper for KV methods
    660
    71
    fn layer_norm_kv(model: &GpuModel, input: &[f32]) -> Vec<f32> {
    661
    71
        GpuModel::layer_norm_static(
    662
    71
            input,
    663
    71
            &model.final_norm_weight,
    664
    71
            &model.final_norm_bias,
    665
    71
            model.config.hidden_dim,
    666
    71
            model.config.eps,
    667
        )
    668
    71
    }
    669
    670
    /// Argmax helper
    671
    32
    fn argmax(logits: &[f32]) -> usize {
    672
    32
        logits
    673
    32
            .iter()
    674
    32
            .enumerate()
    675
    8.16k
            .max_by32(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    676
    32
            .map_or(0, |(idx, _)| idx)
    677
    32
    }
    678
    679
    /// Top-k sampling helper
    680
    0
    fn sample_topk(logits: &[f32], temperature: f32, top_k: usize) -> usize {
    681
    0
        let scaled: Vec<f32> = logits.iter().map(|&x| x / temperature).collect();
    682
    0
        let max_logit = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    683
    0
        let exp_logits: Vec<f32> = scaled.iter().map(|&x| (x - max_logit).exp()).collect();
    684
    0
        let sum: f32 = exp_logits.iter().sum();
    685
    0
        let probs: Vec<f32> = exp_logits.iter().map(|&x| x / sum).collect();
    686
    687
    0
        let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
    688
    0
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    689
    0
        indexed.truncate(top_k);
    690
    0
        indexed.first().map_or(0, |&(idx, _)| idx)
    691
    0
    }