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

    Line
    Count
    Source
    1
    //! Batch Generation and Single-Token Forward (PMAT-802)
    2
    //!
    3
    //! Extracted from model.rs: incremental generation, single-token forward, and helpers.
    4
    5
    use crate::error::{RealizarError, Result};
    6
    use super::super::{exceeds_gpu_buffer_limit, cpu_matmul_transposed_simd, cpu_matmul};
    7
    use super::model::{GpuModel, GpuModelConfig};
    8
    9
    /// Generate tokens using GPU-accelerated forward pass with incremental decoding
    10
    ///
    11
    /// # Arguments
    12
    ///
    13
    /// * `model` - GPU model reference
    14
    /// * `prompt` - Initial token IDs
    15
    /// * `max_tokens` - Maximum tokens to generate
    16
    ///
    17
    /// # Returns
    18
    ///
    19
    /// Generated tokens (including prompt)
    20
    ///
    21
    /// # Errors
    22
    ///
    23
    /// Returns error if generation fails
    24
    0
    pub fn generate_gpu(model: &mut GpuModel, prompt: &[usize], max_tokens: usize) -> Result<Vec<usize>> {
    25
    0
        let mut tokens = prompt.to_vec();
    26
    0
        let vocab_size = model.config.vocab_size;
    27
    28
        // Process prompt first (full forward)
    29
    0
        let logits = model.forward_gpu(&tokens)?;
    30
    31
        // Get first prediction
    32
    0
        let last_pos_start = (tokens.len() - 1) * vocab_size;
    33
    0
        let last_logits = &logits[last_pos_start..last_pos_start + vocab_size];
    34
    35
    0
        let next_token = argmax(last_logits);
    36
    0
        tokens.push(next_token);
    37
    38
        // Generate remaining tokens one at a time (incremental)
    39
        // Use optimized greedy path for large vocabularies
    40
    0
        if vocab_size > 8192 {
    41
            // Large vocab: use fused LM head + argmax
    42
    0
            for _ in 1..max_tokens {
    43
    0
                let next_token = forward_single_token_greedy(model, &tokens)?;
    44
    0
                tokens.push(next_token);
    45
            }
    46
        } else {
    47
            // Small vocab: standard path
    48
    0
            for _ in 1..max_tokens {
    49
    0
                let logits = forward_single_token(model, &tokens)?;
    50
    0
                let next_token = argmax(&logits);
    51
    0
                tokens.push(next_token);
    52
            }
    53
        }
    54
    55
    0
        Ok(tokens)
    56
    0
    }
    57
    58
    /// Fast single-token forward pass for incremental generation
    59
    ///
    60
    /// Only processes the last token position, avoiding O(n²) recomputation.
    61
    0
    pub fn forward_single_token(model: &mut GpuModel, tokens: &[usize]) -> Result<Vec<f32>> {
    62
    0
        let hidden_dim = model.config.hidden_dim;
    63
    0
        let vocab_size = model.config.vocab_size;
    64
    65
        // Embed only the last token
    66
    0
        let last_token = *tokens.last().ok_or_else(|| RealizarError::InvalidShape {
    67
    0
            reason: "Token list empty".to_string(),
    68
    0
        })?;
    69
    70
    0
        if last_token >= vocab_size {
    71
    0
            return Err(RealizarError::InvalidShape {
    72
    0
                reason: format!("Token {} out of bounds", last_token),
    73
    0
            });
    74
    0
        }
    75
    76
    0
        let offset = last_token * hidden_dim;
    77
    0
        let mut hidden: Vec<f32> = model.embedding_weights[offset..offset + hidden_dim].to_vec();
    78
    79
        // Process through blocks (simplified for single token)
    80
    0
        for block_idx in 0..model.block_weights.len() {
    81
    0
            hidden = forward_block_single(model, &hidden, block_idx)?;
    82
        }
    83
    84
        // Final layer norm
    85
    0
        hidden = GpuModel::layer_norm_static(
    86
    0
            &hidden,
    87
    0
            &model.final_norm_weight,
    88
    0
            &model.final_norm_bias,
    89
    0
            hidden_dim,
    90
    0
            model.config.eps,
    91
        );
    92
    93
        // IMP-090, IMP-096: Use CPU fallback with SIMD for large vocab
    94
    0
        let lm_head_elements = hidden_dim * vocab_size;
    95
    0
        let output = if exceeds_gpu_buffer_limit(lm_head_elements) {
    96
            // IMP-096: CPU path with transposed weights + SIMD + fused bias
    97
            // Uses parallel dot products with perfect cache behavior
    98
    0
            cpu_matmul_transposed_simd(
    99
    0
                &hidden,
    100
    0
                &model.lm_head_weight_t,
    101
    0
                &model.lm_head_bias,
    102
    0
                hidden_dim,
    103
    0
                vocab_size,
    104
            )
    105
        } else {
    106
            // GPU path for smaller vocab
    107
    0
            let logits =
    108
    0
                model.scheduler
    109
    0
                    .matmul(&hidden, &model.lm_head_weight, 1, hidden_dim, vocab_size)?;
    110
            // Add bias
    111
    0
            logits
    112
    0
                .iter()
    113
    0
                .zip(model.lm_head_bias.iter())
    114
    0
                .map(|(&x, &b)| x + b)
    115
    0
                .collect()
    116
        };
    117
    118
    0
        Ok(output)
    119
    0
    }
    120
    121
    /// Single-token forward pass optimized for greedy sampling
    122
    ///
    123
    /// Returns the argmax token directly.
    124
    0
    pub fn forward_single_token_greedy(model: &mut GpuModel, tokens: &[usize]) -> Result<usize> {
    125
    0
        let hidden_dim = model.config.hidden_dim;
    126
    0
        let vocab_size = model.config.vocab_size;
    127
    128
        // Embed only the last token
    129
    0
        let last_token = *tokens.last().ok_or_else(|| RealizarError::InvalidShape {
    130
    0
            reason: "Token list empty".to_string(),
    131
    0
        })?;
    132
    133
    0
        if last_token >= vocab_size {
    134
    0
            return Err(RealizarError::InvalidShape {
    135
    0
                reason: format!("Token {} out of bounds", last_token),
    136
    0
            });
    137
    0
        }
    138
    139
    0
        let offset = last_token * hidden_dim;
    140
    0
        let mut hidden: Vec<f32> = model.embedding_weights[offset..offset + hidden_dim].to_vec();
    141
    142
        // Process through blocks (simplified for single token)
    143
    0
        for block_idx in 0..model.block_weights.len() {
    144
    0
            hidden = forward_block_single(model, &hidden, block_idx)?;
    145
        }
    146
    147
        // Final layer norm
    148
    0
        hidden = GpuModel::layer_norm_static(
    149
    0
            &hidden,
    150
    0
            &model.final_norm_weight,
    151
    0
            &model.final_norm_bias,
    152
    0
            hidden_dim,
    153
    0
            model.config.eps,
    154
        );
    155
    156
        // Use optimized CPU path with transposed weights for large vocab
    157
        // This uses row-major access pattern which is ~3-5x faster than column access
    158
        // IMP-090: Also use CPU path if vocab would exceed GPU buffer limits
    159
    0
        let lm_head_elements = hidden_dim * vocab_size;
    160
    0
        if vocab_size > 8192 || exceeds_gpu_buffer_limit(lm_head_elements) {
    161
            // CPU path with transposed weights: perfect cache behavior
    162
    0
            Ok(optimized_lm_head_argmax_transposed(
    163
    0
                &hidden,
    164
    0
                &model.lm_head_weight_t,
    165
    0
                &model.lm_head_bias,
    166
    0
                hidden_dim,
    167
    0
                vocab_size,
    168
    0
            ))
    169
        } else {
    170
            // GPU/small vocab path
    171
    0
            let logits =
    172
    0
                model.scheduler
    173
    0
                    .matmul(&hidden, &model.lm_head_weight, 1, hidden_dim, vocab_size)?;
    174
    0
            let output: Vec<f32> = logits
    175
    0
                .iter()
    176
    0
                .zip(model.lm_head_bias.iter())
    177
    0
                .map(|(&x, &b)| x + b)
    178
    0
                .collect();
    179
    0
            Ok(argmax(&output))
    180
        }
    181
    0
    }
    182
    183
    /// Single token forward through a transformer block (CPU-optimized for m=1)
    184
    ///
    185
    /// For single-token generation, CPU operations are faster than GPU due to transfer overhead.
    186
    #[allow(clippy::unnecessary_wraps)]
    187
    0
    pub fn forward_block_single(model: &mut GpuModel, input: &[f32], block_idx: usize) -> Result<Vec<f32>> {
    188
    0
        let hidden_dim = model.config.hidden_dim;
    189
    0
        let intermediate_dim = model.config.intermediate_dim;
    190
    0
        let kv_dim = model.config.kv_dim();
    191
    0
        let qkv_dim = model.config.qkv_dim();
    192
    193
        // Get block weights
    194
    0
        let block = &model.block_weights[block_idx];
    195
    196
        // Pre-norm
    197
    0
        let normed = GpuModel::layer_norm_static(
    198
    0
            input,
    199
    0
            &block.attn_norm_weight,
    200
    0
            &block.attn_norm_bias,
    201
    0
            hidden_dim,
    202
    0
            model.config.eps,
    203
        );
    204
    205
        // QKV projection for single token (GQA: qkv_dim = hidden_dim + 2*kv_dim)
    206
        // Use CPU matmul directly - GPU overhead not worth it for m=1
    207
    0
        let qkv_weight = &model.block_weights[block_idx].qkv_weight;
    208
    0
        let qkv = cpu_matmul(&normed, qkv_weight, 1, hidden_dim, qkv_dim);
    209
    210
        // Split QKV and apply simplified self-attention (single token)
    211
        // q and k unused for single-token (no cross-attention needed)
    212
        // GQA: V has kv_dim size, but we need hidden_dim output
    213
    0
        let v = &qkv[hidden_dim + kv_dim..];
    214
    215
        // For single token: attention output = v (self-attention with one token)
    216
        // GQA: V has kv_dim, need to repeat heads to get hidden_dim
    217
    0
        let num_kv_heads = model.config.num_kv_heads;
    218
    0
        let heads_per_kv = model.config.num_heads / num_kv_heads;
    219
    0
        let head_dim = model.config.head_dim();
    220
    221
    0
        let attn_out: Vec<f32> = if heads_per_kv == 1 {
    222
            // Standard MHA: no repetition needed
    223
    0
            v.to_vec()
    224
        } else {
    225
            // GQA: repeat each KV head to serve multiple Q heads
    226
    0
            let mut expanded = Vec::with_capacity(hidden_dim);
    227
    0
            for kv_h in 0..num_kv_heads {
    228
    0
                let v_head = &v[kv_h * head_dim..(kv_h + 1) * head_dim];
    229
    0
                for _ in 0..heads_per_kv {
    230
    0
                    expanded.extend_from_slice(v_head);
    231
    0
                }
    232
            }
    233
    0
            expanded
    234
        };
    235
    236
        // Output projection (CPU - m=1)
    237
    0
        let out_weight = &model.block_weights[block_idx].out_weight;
    238
    0
        let out_bias = &model.block_weights[block_idx].out_bias;
    239
    0
        let projected = cpu_matmul(&attn_out, out_weight, 1, hidden_dim, hidden_dim);
    240
    241
        // Residual 1
    242
    0
        let residual1: Vec<f32> = input
    243
    0
            .iter()
    244
    0
            .zip(projected.iter())
    245
    0
            .enumerate()
    246
    0
            .map(|(i, (&inp, &proj))| inp + proj + out_bias[i])
    247
    0
            .collect();
    248
    249
        // FFN pre-norm
    250
    0
        let ffn_norm_weight = &model.block_weights[block_idx].ffn_norm_weight;
    251
    0
        let ffn_norm_bias = &model.block_weights[block_idx].ffn_norm_bias;
    252
    0
        let ffn_normed = GpuModel::layer_norm_static(
    253
    0
            &residual1,
    254
    0
            ffn_norm_weight,
    255
    0
            ffn_norm_bias,
    256
    0
            hidden_dim,
    257
    0
            model.config.eps,
    258
        );
    259
    260
        // FFN fc1 (CPU - m=1)
    261
    0
        let ffn_fc1_weight = &model.block_weights[block_idx].ffn_fc1_weight;
    262
    0
        let ffn_fc1_bias = &model.block_weights[block_idx].ffn_fc1_bias;
    263
    264
        // FFN: SwiGLU when gate weight exists, otherwise GELU
    265
    0
        let activated: Vec<f32> = if let Some(ref gate_weight) = model.block_weights[block_idx].ffn_gate_weight {
    266
            // SwiGLU: silu(gate(x)) * up(x)
    267
    0
            let up_out = cpu_matmul(&ffn_normed, ffn_fc1_weight, 1, hidden_dim, intermediate_dim);
    268
    0
            let gate_out = cpu_matmul(&ffn_normed, gate_weight, 1, hidden_dim, intermediate_dim);
    269
    270
            // SwiGLU: silu(gate) * up
    271
    0
            up_out
    272
    0
                .iter()
    273
    0
                .zip(gate_out.iter())
    274
    0
                .map(|(&u, &g)| {
    275
    0
                    let silu_g = g / (1.0 + (-g).exp());
    276
    0
                    silu_g * u
    277
    0
                })
    278
    0
                .collect()
    279
        } else {
    280
            // Standard GELU FFN
    281
    0
            let fc1_out = cpu_matmul(&ffn_normed, ffn_fc1_weight, 1, hidden_dim, intermediate_dim);
    282
    283
    0
            fc1_out
    284
    0
                .iter()
    285
    0
                .enumerate()
    286
    0
                .map(|(i, &x)| {
    287
    0
                    let x = x + ffn_fc1_bias[i];
    288
    0
                    0.5 * x
    289
    0
                        * (1.0
    290
    0
                            + ((2.0f32 / std::f32::consts::PI).sqrt() * (x + 0.044_715 * x.powi(3)))
    291
    0
                                .tanh())
    292
    0
                })
    293
    0
                .collect()
    294
        };
    295
    296
        // FFN fc2 (CPU - m=1)
    297
    0
        let ffn_fc2_weight = &model.block_weights[block_idx].ffn_fc2_weight;
    298
    0
        let ffn_fc2_bias = &model.block_weights[block_idx].ffn_fc2_bias;
    299
    0
        let fc2_out = cpu_matmul(&activated, ffn_fc2_weight, 1, intermediate_dim, hidden_dim);
    300
    301
        // Residual 2
    302
    0
        let output: Vec<f32> = residual1
    303
    0
            .iter()
    304
    0
            .zip(fc2_out.iter())
    305
    0
            .enumerate()
    306
    0
            .map(|(i, (&r, &fc))| r + fc + ffn_fc2_bias[i])
    307
    0
            .collect();
    308
    309
    0
        Ok(output)
    310
    0
    }
    311
    312
    /// Argmax helper for sampling - vectorized for large vocabularies
    313
    #[allow(clippy::items_after_statements)]
    314
    523
    pub fn argmax(logits: &[f32]) -> usize {
    315
        // For small vocab, use simple iterator
    316
    523
        if logits.len() <= 1024 {
    317
    513
            return logits
    318
    513
                .iter()
    319
    513
                .enumerate()
    320
    128k
                .max_by513(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    321
    513
                .map_or(0, |(i, _)| i);
    322
    10
        }
    323
    324
        // For large vocab (32K+), use chunked parallel argmax
    325
        const CHUNK_SIZE: usize = 4096;
    326
    327
        // Find max in each chunk
    328
    10
        let chunk_maxes: Vec<(usize, f32)> = logits
    329
    10
            .chunks(CHUNK_SIZE)
    330
    10
            .enumerate()
    331
    41
            .map10(|(chunk_idx, chunk)| {
    332
    41
                let (local_idx, &max_val) = chunk
    333
    41
                    .iter()
    334
    41
                    .enumerate()
    335
    154k
                    .max_by41(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    336
    41
                    .expect("chunk is non-empty by construction");
    337
    41
                (chunk_idx * CHUNK_SIZE + local_idx, max_val)
    338
    41
            })
    339
    10
            .collect();
    340
    341
        // Find global max
    342
    10
        chunk_maxes
    343
    10
            .into_iter()
    344
    31
            .max_by10(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    345
    10
            .map_or(0, |(idx, _)| idx)
    346
    523
    }
    347
    348
    /// Optimized LM head + argmax using transposed weights with vectorized dot products
    349
    ///
    350
    /// Uses transposed weights [vocab_size, hidden_dim] for row-major access pattern.
    351
    /// Inner loop is vectorized by the compiler via slice operations.
    352
    #[allow(clippy::many_single_char_names, clippy::items_after_statements)]
    353
    6
    pub fn optimized_lm_head_argmax_transposed(
    354
    6
        hidden: &[f32],
    355
    6
        weight_t: &[f32], // Transposed: [vocab_size, hidden_dim]
    356
    6
        bias: &[f32],
    357
    6
        hidden_dim: usize,
    358
    6
        vocab_size: usize,
    359
    6
    ) -> usize {
    360
        use rayon::prelude::*;
    361
    362
        // Process in larger chunks for better parallelism
    363
        const CHUNK_SIZE: usize = 4096;
    364
    365
        // Find argmax in parallel
    366
    6
        (0..vocab_size)
    367
    6
            .into_par_iter()
    368
    6
            .step_by(CHUNK_SIZE)
    369
    7
            .map6(|chunk_start| {
    370
    7
                let chunk_end = (chunk_start + CHUNK_SIZE).min(vocab_size);
    371
    7
                let mut best_local_idx = chunk_start;
    372
    7
                let mut best_local_val = f32::NEG_INFINITY;
    373
    374
    5.01k
                for j in chunk_start7..chunk_end7 {
    375
                    // Row-major access: weight_t[j, :] is contiguous in memory
    376
    5.01k
                    let row = &weight_t[j * hidden_dim..(j + 1) * hidden_dim];
    377
    378
                    // Vectorized dot product - compiler can auto-vectorize this
    379
    320k
                    let dot5.01k: f325.01k = row5.01k.iter5.01k().zip5.01k(hidden5.01k.iter5.01k()).map5.01k(|(&w, &h)| w * h).sum5.01k();
    380
    381
    5.01k
                    let logit = dot + bias[j];
    382
    383
    5.01k
                    if logit > best_local_val {
    384
    11
                        best_local_val = logit;
    385
    11
                        best_local_idx = j;
    386
    5.00k
                    }
    387
                }
    388
    7
                (best_local_idx, best_local_val)
    389
    7
            })
    390
    6
            .reduce(
    391
                || (0, f32::NEG_INFINITY),
    392
    8
                |a, b| if a.1 > b.1 { a1 } else { b7 },
    393
            )
    394
            .0
    395
    6
    }
    396
    397
    /// Optimized GQA attention using GPU for matmul operations (IMP-089)
    398
    10
    pub fn optimized_gqa_attention(model: &mut GpuModel, qkv: &[f32], seq_len: usize) -> Result<Vec<f32>> {
    399
    10
        let hidden_dim = model.config.hidden_dim;
    400
    10
        let num_heads = model.config.num_heads;
    401
    10
        let num_kv_heads = model.config.num_kv_heads;
    402
    10
        let head_dim = model.config.head_dim();
    403
    10
        let kv_dim = model.config.kv_dim();
    404
    10
        let heads_per_kv = num_heads / num_kv_heads;
    405
    406
        // Split QKV (GQA: K/V have kv_dim per position)
    407
    10
        let q = &qkv[..seq_len * hidden_dim];
    408
    10
        let k = &qkv[seq_len * hidden_dim..seq_len * hidden_dim + seq_len * kv_dim];
    409
    10
        let v = &qkv[seq_len * hidden_dim + seq_len * kv_dim..];
    410
    411
    10
        let scale = 1.0 / (head_dim as f32).sqrt();
    412
    10
        let mut output = vec![0.0f32; seq_len * hidden_dim];
    413
    414
        // Process each head
    415
    40
        for head in 0..num_heads10 {
    416
    40
            let kv_head = head / heads_per_kv;
    417
    418
            // Extract Q for this head
    419
    40
            let mut q_head = Vec::with_capacity(seq_len * head_dim);
    420
    128
            for i in 0..seq_len40 {
    421
    128
                let start = i * hidden_dim + head * head_dim;
    422
    128
                q_head.extend_from_slice(&q[start..start + head_dim]);
    423
    128
            }
    424
    425
            // Extract K, V for the corresponding KV head (shared by multiple Q heads)
    426
    40
            let mut k_head = Vec::with_capacity(seq_len * head_dim);
    427
    40
            let mut v_head = Vec::with_capacity(seq_len * head_dim);
    428
    128
            for i in 0..seq_len40 {
    429
    128
                let start = i * kv_dim + kv_head * head_dim;
    430
    128
                k_head.extend_from_slice(&k[start..start + head_dim]);
    431
    128
                v_head.extend_from_slice(&v[start..start + head_dim]);
    432
    128
            }
    433
    434
            // Compute attention scores: Q @ K^T using GPU matmul
    435
    40
            let mut attn_scores = vec![f32::NEG_INFINITY; seq_len * seq_len];
    436
    40
            let scores = model
    437
    40
                .scheduler
    438
    40
                .matmul_transpose_b(&q_head, &k_head, seq_len, head_dim, seq_len)?0;
    439
    440
            // Apply causal mask and scale
    441
    128
            for i in 0..seq_len40 {
    442
    304
                for j in 0..=i128 {
    443
    304
                    attn_scores[i * seq_len + j] = scores[i * seq_len + j] * scale;
    444
    304
                }
    445
            }
    446
    447
            // Softmax per row
    448
    128
            for i in 0..seq_len40 {
    449
    128
                let row_start = i * seq_len;
    450
    128
                let row = &mut attn_scores[row_start..row_start + seq_len];
    451
    452
    128
                let max_val = row[..=i].iter().copied().fold(f32::NEG_INFINITY, f32::max);
    453
    454
    128
                let mut sum = 0.0f32;
    455
    304
                for item in row128.iter_mut128().take128(i + 1128) {
    456
    304
                    *item = (*item - max_val).exp();
    457
    304
                    sum += *item;
    458
    304
                }
    459
    460
    304
                for item in row128.iter_mut128().take128(i + 1128) {
    461
    304
                    *item /= sum;
    462
    304
                }
    463
    176
                for item in row128.iter_mut128().skip128(i + 1128) {
    464
    176
                    *item = 0.0;
    465
    176
                }
    466
            }
    467
    468
            // Compute output: attn @ V using GPU matmul
    469
    40
            let head_output =
    470
    40
                model.scheduler
    471
    40
                    .matmul(&attn_scores, &v_head, seq_len, seq_len, head_dim)?0;
    472
    473
            // Copy to output
    474
    128
            for i in 0..seq_len40 {
    475
    128
                let out_start = i * hidden_dim + head * head_dim;
    476
    128
                let head_start = i * head_dim;
    477
    128
                output[out_start..out_start + head_dim]
    478
    128
                    .copy_from_slice(&head_output[head_start..head_start + head_dim]);
    479
    128
            }
    480
        }
    481
    482
    10
        Ok(output)
    483
    10
    }
    484
    485
    /// Simplified attention (fallback, for M3 benchmarking)
    486
    #[allow(dead_code, clippy::unnecessary_wraps)]
    487
    3
    pub fn simplified_attention(config: &GpuModelConfig, qkv: &[f32], seq_len: usize) -> Result<Vec<f32>> {
    488
    3
        let hidden_dim = config.hidden_dim;
    489
    3
        let head_dim = hidden_dim / config.num_heads;
    490
    491
        // Split QKV
    492
    3
        let q = &qkv[..seq_len * hidden_dim];
    493
    3
        let k = &qkv[seq_len * hidden_dim..seq_len * 2 * hidden_dim];
    494
    3
        let v = &qkv[seq_len * 2 * hidden_dim..];
    495
    496
        // Simplified scaled dot-product attention per head
    497
    3
        let scale = 1.0 / (head_dim as f32).sqrt();
    498
    3
        let mut output = vec![0.0f32; seq_len * hidden_dim];
    499
    500
    8
        for head in 0..config.num_heads3 {
    501
    14
            for i in 0..seq_len8 {
    502
                // Compute attention weights for position i
    503
    14
                let mut weights = Vec::with_capacity(seq_len);
    504
    14
                let mut max_score = f32::NEG_INFINITY;
    505
    506
    22
                for j in 0..=i14 {
    507
                    // Causal: only attend to previous positions
    508
    22
                    let mut score = 0.0f32;
    509
    124
                    for d in 0..head_dim22 {
    510
    124
                        let q_idx = i * hidden_dim + head * head_dim + d;
    511
    124
                        let k_idx = j * hidden_dim + head * head_dim + d;
    512
    124
                        score += q[q_idx] * k[k_idx];
    513
    124
                    }
    514
    22
                    score *= scale;
    515
    22
                    max_score = max_score.max(score);
    516
    22
                    weights.push(score);
    517
                }
    518
    519
                // Softmax
    520
    14
                let mut sum = 0.0f32;
    521
    36
                for w22 in &mut weights {
    522
    22
                    *w = (*w - max_score).exp();
    523
    22
                    sum += *w;
    524
    22
                }
    525
    36
                for w22 in &mut weights {
    526
    22
                    *w /= sum;
    527
    22
                }
    528
    529
                // Weighted sum of values
    530
    96
                for d in 0..head_dim14 {
    531
    96
                    let out_idx = i * hidden_dim + head * head_dim + d;
    532
    124
                    for (j, &w) in weights.iter()96.enumerate96() {
    533
    124
                        let v_idx = j * hidden_dim + head * head_dim + d;
    534
    124
                        output[out_idx] += w * v[v_idx];
    535
    124
                    }
    536
                }
    537
            }
    538
        }
    539
    540
    3
        Ok(output)
    541
    3
    }