//! KV Cache Management for GpuModel (PMAT-802)

//!

//! Extracted from model.rs to reduce module size.

//! Contains KV cache forward pass and generation logic.

use crate::error::{RealizarError, Result};

use super::super::{StreamingKVCache, exceeds_gpu_buffer_limit, cpu_matmul_transposed_simd};

use super::model::{GpuModel, GpuGenerateConfig};

/// Apply Rotary Position Embedding (RoPE) to Q or K vectors (Phase 21)

///

/// RoPE encodes position information by rotating pairs of elements

/// with position-dependent angles. This is CRITICAL for transformer attention.

///

/// # Arguments

/// * `x` - Mutable slice of Q or K vectors [seq_len * num_heads * head_dim]

/// * `seq_len` - Number of positions to encode

/// * `num_heads` - Number of attention heads in this tensor

/// * `head_dim` - Dimension per head

/// * `rope_theta` - Base frequency (typically 10000.0)

/// * `start_pos` - Starting position for RoPE (0 for prefill, cache_len for incremental)

fn apply_rope(

    x: &mut [f32],

    seq_len: usize,

    num_heads: usize,

    head_dim: usize,

    rope_theta: f32,

    start_pos: usize,

) {

    let half_dim = head_dim / 2;

    let head_dim_f32 = head_dim as f32;

    let total_dim = num_heads * head_dim;

    for pos in 0..
seq_len456
 {

        let position = start_pos + pos;

        let pos_f32 = position as f32;

        let pos_offset = pos * total_dim;

        for h in 0..
num_heads1.45k
 {

            let head_start = pos_offset + h * head_dim;

            let idx2_start = head_start + half_dim;

            for i in 0..
half_dim9.84k
 {

                let freq = 1.0 / rope_theta.powf(2.0 * i as f32 / head_dim_f32);

                let angle = pos_f32 * freq;

                let (sin_val, cos_val) = angle.sin_cos();

                let x1 = x[head_start + i];

                let x2 = x[idx2_start + i];

                // Apply rotation: [cos -sin; sin cos] * [x1; x2]

                x[head_start + i] = x1 * cos_val - x2 * sin_val;

                x[idx2_start + i] = x1 * sin_val + x2 * cos_val;

            }

        }

    }

}

/// Forward pass with KV cache population (IMP-031)

pub fn forward_gpu_with_cache(

    model: &mut GpuModel,

    token_ids: &[usize],

    kv_cache: &mut StreamingKVCache,

) -> Result<Vec<f32>> {

    if token_ids.is_empty() {

        return Err(RealizarError::InvalidShape {

            reason: "Token IDs cannot be empty".to_string(),

        });

    }

    let seq_len = token_ids.len();

    let hidden_dim = model.config.hidden_dim;

    // Step 1: Embed tokens

    let mut hidden = Vec::with_capacity(seq_len * hidden_dim);

    for &
token_id202
 in token_ids {

        if token_id >= model.config.vocab_size {

            return Err(RealizarError::InvalidShape {

                reason: format!(

                    "Token ID {} out of bounds (vocab_size={})",

                    token_id, model.config.vocab_size

                ),

            });

        }

        let offset = token_id * hidden_dim;

        hidden.extend_from_slice(&model.embedding_weights[offset..offset + hidden_dim]);

    }

    // Step 2: Pass through transformer blocks with KV cache population

    for block_idx in 0..
model.block_weights36
.
len36
() {

        hidden = forward_block_with_cache(model, &hidden, seq_len, block_idx, kv_cache)
?0
;

    }

    // Step 3: Final layer norm

    hidden = layer_norm_kv(model, &hidden);

    // Step 4: LM head projection - only for final position

    let final_hidden = &hidden[(seq_len - 1) * hidden_dim..seq_len * hidden_dim];

    let lm_head_elements = hidden_dim * model.config.vocab_size;

    let output = if exceeds_gpu_buffer_limit(lm_head_elements) {

        cpu_matmul_transposed_simd(

            final_hidden,

            &model.lm_head_weight_t,

            &model.lm_head_bias,

            hidden_dim,

            model.config.vocab_size,

        )

    } else {

        let logits = model.scheduler.matmul(

            final_hidden,

            &model.lm_head_weight,

            1,

            hidden_dim,

            model.config.vocab_size,

        )?;

        let mut output = logits;

        for (out_val, bias_val) in 
output.iter_mut()36
.
zip36
(
model.lm_head_bias.iter()36
) {

            *out_val += *bias_val;

        }

        output

    };

    Ok(output)

}

/// Incremental forward pass using cached KV (IMP-032)

pub fn forward_gpu_incremental(

    model: &mut GpuModel,

    token_id: usize,

    kv_cache: &mut StreamingKVCache,

) -> Result<Vec<f32>> {

    if token_id >= model.config.vocab_size {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Token ID {} out of bounds (vocab_size={})",

                token_id, model.config.vocab_size

            ),

        });

    }

    let hidden_dim = model.config.hidden_dim;

    // Step 1: Embed token

    let offset = token_id * hidden_dim;

    let mut hidden = model.embedding_weights[offset..offset + hidden_dim].to_vec();

    // Step 2: Pass through transformer blocks using KV cache

    for block_idx in 0..
model.block_weights35
.
len35
() {

        hidden = forward_block_incremental(model, &hidden, block_idx, kv_cache)
?0
;

    }

    // Step 3: Final layer norm

    hidden = layer_norm_kv(model, &hidden);

    // Step 4: LM head projection

    let lm_head_elements = hidden_dim * model.config.vocab_size;

    let output = if exceeds_gpu_buffer_limit(lm_head_elements) {

        cpu_matmul_transposed_simd(

            &hidden,

            &model.lm_head_weight_t,

            &model.lm_head_bias,

            hidden_dim,

            model.config.vocab_size,

        )

    } else {

        let logits = model.scheduler.matmul(

            &hidden,

            &model.lm_head_weight,

            1,

            hidden_dim,

            model.config.vocab_size,

        )?;

        let mut output = logits;

        for (out_val, bias_val) in 
output.iter_mut()35
.
zip35
(
model.lm_head_bias.iter()35
) {

            *out_val += *bias_val;

        }

        output

    };

    Ok(output)

}

/// Forward pass through a single block with KV cache population

fn forward_block_with_cache(

    model: &mut GpuModel,

    input: &[f32],

    seq_len: usize,

    block_idx: usize,

    kv_cache: &mut StreamingKVCache,

) -> Result<Vec<f32>> {

    let hidden_dim = model.config.hidden_dim;

    let intermediate_dim = model.config.intermediate_dim;

    let num_heads = model.config.num_heads;

    let num_kv_heads = model.config.num_kv_heads;

    let head_dim = model.config.head_dim();

    let kv_dim = model.config.kv_dim();

    let qkv_dim = model.config.qkv_dim();

    let block = &model.block_weights[block_idx];

    // Pre-norm

    let normed = GpuModel::layer_norm_static(

        input,

        &block.attn_norm_weight,

        &block.attn_norm_bias,

        hidden_dim,

        model.config.eps,

    );

    // QKV projection

    let mut qkv = model.scheduler.matmul(

        &normed,

        &model.block_weights[block_idx].qkv_weight,

        seq_len,

        hidden_dim,

        qkv_dim,

    )?;

    // Split Q, K, V (mutable for RoPE application)

    let q_end = seq_len * hidden_dim;

    let k_end = q_end + seq_len * kv_dim;

    // Phase 21: Apply RoPE to Q and K BEFORE caching

    // This is CRITICAL - without RoPE, attention has no position information

    let rope_theta = model.config.rope_theta;

    // Apply RoPE to Q (all heads)

    apply_rope(&mut qkv[..q_end], seq_len, num_heads, head_dim, rope_theta, 0);

    // Apply RoPE to K (KV heads)

    apply_rope(&mut qkv[q_end..k_end], seq_len, num_kv_heads, head_dim, rope_theta, 0);

    // Now split (after RoPE applied)

    let q = &qkv[..q_end];

    let k = &qkv[q_end..k_end];

    let v = &qkv[k_end..];

    // Cache K (with RoPE) and V

    for pos in 0..
seq_len96
 {

        let k_slice = &k[pos * kv_dim..(pos + 1) * kv_dim];

        let v_slice = &v[pos * kv_dim..(pos + 1) * kv_dim];

        kv_cache.append(block_idx, k_slice, v_slice);

    }

    // GQA attention

    let attn_out = gqa_attention_with_kv(model, q, k, v, seq_len, num_heads, num_kv_heads, head_dim)
?0
;

    // Output projection

    let projected = model.scheduler.matmul(

        &attn_out,

        &model.block_weights[block_idx].out_weight,

        seq_len,

        hidden_dim,

        hidden_dim,

    )?;

    // Residual 1

    let mut residual1: Vec<f32> = input

        .iter()

        .zip(projected.iter())

        .enumerate()

        .
map96
(|(i, (&inp, &proj))| {

            inp + proj + model.block_weights[block_idx].out_bias[i % hidden_dim]

        })

        .collect();

    // FFN pre-norm

    let ffn_normed = GpuModel::layer_norm_static(

        &residual1,

        &model.block_weights[block_idx].ffn_norm_weight,

        &model.block_weights[block_idx].ffn_norm_bias,

        hidden_dim,

        model.config.eps,

    );

    // FFN: SwiGLU when gate weight exists, otherwise GELU

    let activated: Vec<f32> = if let Some(
ref gate_weight0
) = model.block_weights[block_idx].ffn_gate_weight {

        // SwiGLU: silu(gate(x)) * up(x)

        let up_out = model.scheduler.matmul(

            &ffn_normed,

            &model.block_weights[block_idx].ffn_fc1_weight,

            seq_len,

            hidden_dim,

            intermediate_dim,

        )?;

        let gate_out = model.scheduler.matmul(

            &ffn_normed,

            gate_weight,

            seq_len,

            hidden_dim,

            intermediate_dim,

        )?;

        // SwiGLU: silu(gate) * up

        up_out

            .iter()

            .zip(gate_out.iter())

            .map(|(&u, &g)| {

                let silu_g = g / (1.0 + (-g).exp());

                silu_g * u

            })

            .collect()

    } else {

        // Standard GELU FFN

        let fc1_out = model.scheduler.matmul(

            &ffn_normed,

            &model.block_weights[block_idx].ffn_fc1_weight,

            seq_len,

            hidden_dim,

            intermediate_dim,

        )?;

        fc1_out

            .iter()

            .enumerate()

            .
map96
(|(i, &x)| {

                let x = x + model.block_weights[block_idx].ffn_fc1_bias[i % intermediate_dim];

                0.5 * x * (1.0 + ((2.0f32 / std::f32::consts::PI).sqrt() * (x + 0.044_715 * x.powi(3))).tanh())

            })

            .collect()

    };

    // FFN: fc2

    let fc2_out = model.scheduler.matmul(

        &activated,

        &model.block_weights[block_idx].ffn_fc2_weight,

        seq_len,

        intermediate_dim,

        hidden_dim,

    )?;

    // Residual 2

    for (i, x) in 
residual1.iter_mut()96
.
enumerate96
() {

        *x += fc2_out[i] + model.block_weights[block_idx].ffn_fc2_bias[i % hidden_dim];

    }

    Ok(residual1)

}

/// Incremental forward pass through a single block using cached KV

fn forward_block_incremental(

    model: &mut GpuModel,

    input: &[f32],

    block_idx: usize,

    kv_cache: &mut StreamingKVCache,

) -> Result<Vec<f32>> {

    let hidden_dim = model.config.hidden_dim;

    let intermediate_dim = model.config.intermediate_dim;

    let num_heads = model.config.num_heads;

    let num_kv_heads = model.config.num_kv_heads;

    let head_dim = model.config.head_dim();

    let kv_dim = model.config.kv_dim();

    let qkv_dim = model.config.qkv_dim();

    let block = &model.block_weights[block_idx];

    // Pre-norm (single position)

    let normed = GpuModel::layer_norm_static(

        input,

        &block.attn_norm_weight,

        &block.attn_norm_bias,

        hidden_dim,

        model.config.eps,

    );

    // QKV projection (single position)

    let mut qkv = model.scheduler.matmul(

        &normed,

        &model.block_weights[block_idx].qkv_weight,

        1,

        hidden_dim,

        qkv_dim,

    )?;

    // Get current position BEFORE caching (this is where new token goes)

    let (existing_k, _) = kv_cache.get_valid(block_idx);

    let current_pos = existing_k.len() / kv_dim;

    // Phase 21: Apply RoPE to Q and K at current position BEFORE caching

    let rope_theta = model.config.rope_theta;

    // Apply RoPE to Q (single position, all heads)

    apply_rope(&mut qkv[..hidden_dim], 1, num_heads, head_dim, rope_theta, current_pos);

    // Apply RoPE to K (single position, KV heads)

    apply_rope(&mut qkv[hidden_dim..hidden_dim + kv_dim], 1, num_kv_heads, head_dim, rope_theta, current_pos);

    // Split Q, K, V (single position, after RoPE)

    let q = &qkv[..hidden_dim];

    let k = &qkv[hidden_dim..hidden_dim + kv_dim];

    let v = &qkv[hidden_dim + kv_dim..];

    // Cache new K (with RoPE) and V

    kv_cache.append(block_idx, k, v);

    // Get all cached K/V for attention (now includes new K/V)

    let (all_k, all_v) = kv_cache.get_valid(block_idx);

    let cache_len = all_k.len() / kv_dim;

    // GQA incremental attention

    let attn_out = gqa_incremental_attention(model, q, all_k, all_v, cache_len, num_heads, num_kv_heads, head_dim)
?0
;

    // Output projection

    let projected = model.scheduler.matmul(

        &attn_out,

        &model.block_weights[block_idx].out_weight,

        1,

        hidden_dim,

        hidden_dim,

    )?;

    // Residual 1

    let mut residual1: Vec<f32> = input

        .iter()

        .zip(projected.iter())

        .enumerate()

        .
map132
(|(i, (&inp, &proj))| {

            inp + proj + model.block_weights[block_idx].out_bias[i]

        })

        .collect();

    // FFN pre-norm

    let ffn_normed = GpuModel::layer_norm_static(

        &residual1,

        &model.block_weights[block_idx].ffn_norm_weight,

        &model.block_weights[block_idx].ffn_norm_bias,

        hidden_dim,

        model.config.eps,

    );

    // FFN: SwiGLU when gate weight exists, otherwise GELU

    let activated: Vec<f32> = if let Some(
ref gate_weight0
) = model.block_weights[block_idx].ffn_gate_weight {

        // SwiGLU: silu(gate(x)) * up(x)

        let up_out = model.scheduler.matmul(

            &ffn_normed,

            &model.block_weights[block_idx].ffn_fc1_weight,

            1,

            hidden_dim,

            intermediate_dim,

        )?;

        let gate_out = model.scheduler.matmul(

            &ffn_normed,

            gate_weight,

            1,

            hidden_dim,

            intermediate_dim,

        )?;

        // SwiGLU: silu(gate) * up

        up_out

            .iter()

            .zip(gate_out.iter())

            .map(|(&u, &g)| {

                let silu_g = g / (1.0 + (-g).exp());

                silu_g * u

            })

            .collect()

    } else {

        // Standard GELU FFN

        let fc1_out = model.scheduler.matmul(

            &ffn_normed,

            &model.block_weights[block_idx].ffn_fc1_weight,

            1,

            hidden_dim,

            intermediate_dim,

        )?;

        fc1_out

            .iter()

            .enumerate()

            .
map132
(|(i, &x)| {

                let x = x + model.block_weights[block_idx].ffn_fc1_bias[i];

                0.5 * x * (1.0 + ((2.0f32 / std::f32::consts::PI).sqrt() * (x + 0.044_715 * x.powi(3))).tanh())

            })

            .collect()

    };

    // FFN: fc2

    let fc2_out = model.scheduler.matmul(

        &activated,

        &model.block_weights[block_idx].ffn_fc2_weight,

        1,

        intermediate_dim,

        hidden_dim,

    )?;

    // Residual 2

    for (i, x) in 
residual1.iter_mut()132
.
enumerate132
() {

        *x += fc2_out[i] + model.block_weights[block_idx].ffn_fc2_bias[i];

    }

    Ok(residual1)

}

/// GQA attention with KV (full sequence)

#[allow(clippy::too_many_arguments)]

fn gqa_attention_with_kv(

    _model: &GpuModel,

    q: &[f32],

    k: &[f32],

    v: &[f32],

    seq_len: usize,

    num_heads: usize,

    num_kv_heads: usize,

    head_dim: usize,

) -> Result<Vec<f32>> {

    let hidden_dim = num_heads * head_dim;

    let kv_dim = num_kv_heads * head_dim;

    let heads_per_kv = num_heads / num_kv_heads;

    let mut output = vec![0.0f32; seq_len * hidden_dim];

    let scale = 1.0 / (head_dim as f32).sqrt();

    for pos in 0..
seq_len96
 {

        for head in 0..
num_heads596
 {

            let kv_head = head / heads_per_kv;

            // Query for this head at this position

            let q_start = pos * hidden_dim + head * head_dim;

            let q_slice = &q[q_start..q_start + head_dim];

            // Compute attention scores for all positions up to current

            let mut scores = Vec::with_capacity(pos + 1);

            for kpos in 0..=
pos3.90k
 {

                let k_start = kpos * kv_dim + kv_head * head_dim;

                let k_slice = &k[k_start..k_start + head_dim];

                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
();

                scores.push(score * scale);

            }

            // Softmax

            let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);

            let 
exp_scores3.90k
: 
Vec<f32>3.90k
 = 
scores.iter()3.90k
.
map3.90k
(|&s| (s - max_score).exp()).
collect3.90k
();

            let sum: f32 = exp_scores.iter().sum();

            let 
weights3.90k
: 
Vec<f32>3.90k
 = 
exp_scores.iter()3.90k
.
map3.90k
(|&e| e / sum).
collect3.90k
();

            // Weighted sum of values

            let out_start = pos * hidden_dim + head * head_dim;

            for (kpos, &weight) in 
weights.iter()3.90k
.
enumerate3.90k
() {

                let v_start = kpos * kv_dim + kv_head * head_dim;

                for d in 0..
head_dim16.1k
 {

                    output[out_start + d] += weight * v[v_start + d];

                }

            }

        }

    }

    Ok(output)

}

/// GQA incremental attention (single query position)

#[allow(clippy::too_many_arguments)]

fn gqa_incremental_attention(

    _model: &GpuModel,

    q: &[f32],

    all_k: &[f32],

    all_v: &[f32],

    cache_len: usize,

    num_heads: usize,

    num_kv_heads: usize,

    head_dim: usize,

) -> Result<Vec<f32>> {

    let hidden_dim = num_heads * head_dim;

    let kv_dim = num_kv_heads * head_dim;

    let heads_per_kv = num_heads / num_kv_heads;

    let mut output = vec![0.0f32; hidden_dim];

    let scale = 1.0 / (head_dim as f32).sqrt();

    for head in 0..
num_heads132
 {

        let kv_head = head / heads_per_kv;

        let q_start = head * head_dim;

        let q_slice = &q[q_start..q_start + head_dim];

        // Attention scores for all cached positions

        let mut scores = Vec::with_capacity(cache_len);

        for kpos in 0..
cache_len1.02k
 {

            let k_start = kpos * kv_dim + kv_head * head_dim;

            let k_slice = &all_k[k_start..k_start + head_dim];

            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
();

            scores.push(score * scale);

        }

        // Softmax

        let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);

        let 
exp_scores1.02k
: 
Vec<f32>1.02k
 = 
scores.iter()1.02k
.
map1.02k
(|&s| (s - max_score).exp()).
collect1.02k
();

        let sum: f32 = exp_scores.iter().sum();

        let 
weights1.02k
: 
Vec<f32>1.02k
 = 
exp_scores.iter()1.02k
.
map1.02k
(|&e| e / sum).
collect1.02k
();

        // Weighted sum

        let out_start = head * head_dim;

        for (kpos, &weight) in 
weights.iter()1.02k
.
enumerate1.02k
() {

            let v_start = kpos * kv_dim + kv_head * head_dim;

            for d in 0..
head_dim23.2k
 {

                output[out_start + d] += weight * all_v[v_start + d];

            }

        }

    }

    Ok(output)

}

/// Generate tokens using KV cache (IMP-033)

pub fn generate_with_cache(

    model: &mut GpuModel,

    prompt: &[usize],

    config: &GpuGenerateConfig,

) -> Result<Vec<usize>> {

    if prompt.is_empty() {

        return Err(RealizarError::InvalidShape {

            reason: "Prompt cannot be empty".to_string(),

        });

    }

    let max_seq_len = prompt.len() + config.max_tokens;

    let head_dim = model.config.hidden_dim / model.config.num_heads;

    let mut kv_cache = StreamingKVCache::new(

        model.config.num_layers,

        max_seq_len,

        model.config.num_kv_heads,

        head_dim,

    );

    let mut tokens = prompt.to_vec();

    let logits = forward_gpu_with_cache(model, prompt, &mut kv_cache)
?0
;

    let mut next_token = if config.temperature == 0.0 || 
config.top_k == 10
 {

        argmax(&logits)

    } else {

        sample_topk(&logits, config.temperature, config.top_k)

    };

    if config.stop_tokens.contains(&next_token) {

        return Ok(tokens);

    }

    tokens.push(next_token);

    for _ in 1..config.max_tokens {

        let logits = forward_gpu_incremental(model, next_token, &mut kv_cache)
?0
;

        next_token = if config.temperature == 0.0 || 
config.top_k == 10
 {

            argmax(&logits)

        } else {

            sample_topk(&logits, config.temperature, config.top_k)

        };

        if config.stop_tokens.contains(&next_token) {

            break;

        }

        tokens.push(next_token);

    }

    Ok(tokens)

}

/// Layer norm helper for KV methods

fn layer_norm_kv(model: &GpuModel, input: &[f32]) -> Vec<f32> {

    GpuModel::layer_norm_static(

        input,

        &model.final_norm_weight,

        &model.final_norm_bias,

        model.config.hidden_dim,

        model.config.eps,

    )

}

/// Argmax helper

fn argmax(logits: &[f32]) -> usize {

    logits

        .iter()

        .enumerate()

        .
max_by32
(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))

        .map_or(0, |(idx, _)| idx)

}

/// Top-k sampling helper

fn sample_topk(logits: &[f32], temperature: f32, top_k: usize) -> usize {

    let scaled: Vec<f32> = logits.iter().map(|&x| x / temperature).collect();

    let max_logit = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);

    let exp_logits: Vec<f32> = scaled.iter().map(|&x| (x - max_logit).exp()).collect();

    let sum: f32 = exp_logits.iter().sum();

    let probs: Vec<f32> = exp_logits.iter().map(|&x| x / sum).collect();

    let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();

    indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    indexed.truncate(top_k);

    indexed.first().map_or(0, |&(idx, _)| idx)

}