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

Source
//! Batched forward pass variants
//!
//! Contains forward_batch, forward_batch_gpu, forward_batch_with_cache,
//! and supporting batch matmul/attention helpers.

use crate::error::{RealizarError, Result};
use crate::gguf::ops;
use crate::gguf::{
    DispatchMetrics, OwnedQuantizedKVCache, OwnedQuantizedModel, OwnedQuantizedTensor,
    OwnedQKVWeights, QuantizedGenerateConfig, TokenBuffer, GGUF_TYPE_Q4_K, GGUF_TYPE_Q5_K,
    GGUF_TYPE_Q6_K,
};

impl OwnedQuantizedModel {
    /// Batched forward pass for prompt prefill (PARITY-002)
    ///
    /// Processes all prompt tokens at once, enabling GPU acceleration
    /// for the attention computation when the batch is large enough.
    ///
    /// # Arguments
    /// * `tokens` - All prompt tokens to process at once
    /// * `cache` - KV cache for storing computed K/V tensors
    /// * `metrics` - Dispatch metrics tracker for CPU/GPU decision recording
    ///
    /// # Returns
    /// Logits for next token prediction (from the last token position)
    ///
    /// # Errors
    /// Returns error if tensor operations fail
    #[cfg(feature = "gpu")]
    pub fn forward_batch_with_cache(
        &self,
        tokens: &[u32],
        cache: &mut OwnedQuantizedKVCache,
        metrics: &std::sync::Arc<DispatchMetrics>,
    ) -> Result<Vec<f32>> {
        if tokens.is_empty() {
            return Err(RealizarError::InvalidShape {
                reason: "Tokens cannot be empty".to_string(),
            });
        }

        let seq_len = tokens.len();
        let hidden_dim = self.config.hidden_dim;

        // 1. Embed all tokens at once: [seq_len, hidden_dim]
        let mut hidden_states: Vec<Vec<f32>> = tokens
            .iter()
            .map5(|&token_id| self.embed(&[token_id]))
            .collect();

        // 2. Process through transformer layers
        for (layer_idx, layer) in self.layers.iter().enumerate() {
            // Collect Q, K, V for all positions
            let mut all_q: Vec<Vec<f32>> = Vec::with_capacity(seq_len);
            let mut all_k: Vec<Vec<f32>> = Vec::with_capacity(seq_len);
            let mut all_v: Vec<Vec<f32>> = Vec::with_capacity(seq_len);

            for (pos, hidden) in hidden_states.iter()5.enumerate5() {
                // 2a. Attention layer norm
                let normed = ops::layer_norm(
                    hidden,
                    &layer.attn_norm_weight,
                    layer.attn_norm_bias.as_deref(),
                    self.config.eps,
                );

                // 2b. QKV projection
                let mut qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0;
                if let Some(ref bias0) = layer.qkv_bias {
                    ops::add_bias(&mut qkv, bias);
                }

                // 2c. Extract Q, K, V and apply RoPE
                // Note: This uses hidden_dim for all (assumes non-GQA or fused QKV)
                let mut q = qkv[0..hidden_dim].to_vec();
                let mut k = qkv[hidden_dim..2 * hidden_dim].to_vec();
                let v = qkv[2 * hidden_dim..3 * hidden_dim].to_vec();

                self.apply_rope(&mut q, pos, self.config.num_heads);
                self.apply_rope(&mut k, pos, self.config.num_heads); // Same as Q for non-GQA

                all_q.push(q);
                all_k.push(k);
                all_v.push(v);
            }

            // 2d. Compute batched attention
            // For PARITY-002: This is where GPU can accelerate!
            // Attention scores: Q @ K^T is [seq_len, seq_len]
            let attn_outputs = self
                .batched_attention_with_cache(&all_q, &all_k, &all_v, cache, layer_idx, metrics)?0;

            // 2e. Store all K/V in cache
            for (k, v) in all_k.iter()5.zip5(all_v.iter()5) {
                cache.append(layer_idx, k, v);
            }

            // 2f. Attention output projection + residual
            for (pos, attn_out) in attn_outputs.iter()5.enumerate5() {
                let mut attn_output = self.fused_matmul(attn_out, &layer.attn_output_weight)?0;
                if let Some(ref bias0) = layer.attn_output_bias {
                    ops::add_bias(&mut attn_output, bias);
                }

                // Residual connection
                for i in 0..hidden_dim84 {
                    hidden_states[pos][i] += attn_output[i];
                }
            }

            // 2g. FFN for all positions
            for hidden84 in &mut hidden_states {
                let mut ffn_hidden = self.fused_matmul(hidden, &layer.ffn_up_weight)?0;
                if let Some(ref bias0) = layer.ffn_up_bias {
                    ops::add_bias(&mut ffn_hidden, bias);
                }
                ops::gelu(&mut ffn_hidden);

                let mut ffn_output = self.fused_matmul(&ffn_hidden, &layer.ffn_down_weight)?0;
                if let Some(ref bias0) = layer.ffn_down_bias {
                    ops::add_bias(&mut ffn_output, bias);
                }

                // Residual
                for i in 0..hidden_dim84 {
                    hidden[i] += ffn_output[i];
                }
            }
        }

        // Advance cache position for all processed tokens
        for _ in 0..seq_len5 {
            cache.advance();
        }

        // 3. Final layer norm and LM head for LAST token only
        let last_hidden = &hidden_states[seq_len - 1];
        let normed = ops::layer_norm(
            last_hidden,
            &self.output_norm_weight,
            self.output_norm_bias.as_deref(),
            self.config.eps,
        );

        // 4. LM head projection
        let mut logits = self.fused_matmul(&normed, &self.lm_head_weight)?0;
        if let Some(ref bias0) = self.lm_head_bias {
            ops::add_bias(&mut logits, bias);
        }

        Ok(logits)
    }

    /// Batched attention computation with GPU acceleration (PARITY-002)
    ///
    /// Computes attention for all positions at once, enabling GPU dispatch
    /// when the workload (seq_len * hidden_dim * seq_len) exceeds the threshold.
    ///
    /// KEY OPTIMIZATION: Uses GPU matmul for Q @ K^T when workload is large enough.
    /// This is the critical path for GPU acceleration - previous implementation only
    /// recorded metrics without actually using GPU.
    #[cfg(feature = "gpu")]
    fn batched_attention_with_cache(
        &self,
        all_q: &[Vec<f32>],
        all_k: &[Vec<f32>],
        all_v: &[Vec<f32>],
        cache: &OwnedQuantizedKVCache,
        layer_idx: usize,
        metrics: &std::sync::Arc<DispatchMetrics>,
    ) -> Result<Vec<Vec<f32>>> {
        let seq_len = all_q.len();
        let hidden_dim = self.config.hidden_dim;
        let num_heads = self.config.num_heads;
        let head_dim = hidden_dim / num_heads;

        // Get any cached K/V from previous sequences
        let cached_k = cache.get_k(layer_idx);
        let cached_v = cache.get_v(layer_idx);
        let cache_len = cached_k.len() / hidden_dim;

        // Build full K/V sequences: [cache + current]
        let total_len = cache_len + seq_len;

        // Determine if we should use GPU based on workload size
        //
        // IMPORTANT FINDING (IMP-600, PARITY-002):
        // GPU is 2.7x SLOWER for MATVEC operations (per-head attention is MATVEC)
        // GPU is 57x FASTER for large GEMM (batch) operations
        //
        // For GPU to be beneficial, we need LARGE matrices. Per-head attention
        // uses tiny matrices: Q[1, head_dim] @ K^T[head_dim, seq_len] = [1, seq_len]
        // This is a MATVEC operation where GPU transfer overhead dominates.
        //
        // Measured result with GPU matmul: 0.20 tok/s (vs 5.31 tok/s CPU)
        // GPU path is 26x SLOWER due to per-head matmul overhead.
        //
        // For true GPU acceleration, need:
        // - FlashAttention (fused kernel, not yet available in trueno)
        // - Batched multi-request inference (process multiple prompts together)
        //
        // For now, use optimized CPU path which is faster for single-request inference.
        let workload = num_heads * seq_len * head_dim * total_len;
        let _ = workload; // Document: GPU not used because MATVEC is slower on GPU

        // Always use CPU path - it's faster for per-head attention MATVEC
        metrics.record_cpu_dispatch();
        self.cpu_batched_attention(
            all_q, all_k, all_v, cached_k, cached_v, cache_len, hidden_dim, num_heads, head_dim,
        )
    }

    /// CPU-based batched attention (fallback for small workloads)
    #[cfg(feature = "gpu")]
    #[allow(clippy::too_many_arguments)] // Attention requires all these parameters
    fn cpu_batched_attention(
        &self,
        all_q: &[Vec<f32>],
        all_k: &[Vec<f32>],
        all_v: &[Vec<f32>],
        cached_k: &[f32],
        cached_v: &[f32],
        cache_len: usize,
        hidden_dim: usize,
        _num_heads: usize,
        head_dim: usize,
    ) -> Result<Vec<Vec<f32>>> {
        let seq_len = all_q.len();
        let mut outputs = Vec::with_capacity(seq_len);

        for (q_pos, q) in all_q5.iter5().enumerate5() {
            let attend_len = cache_len + q_pos + 1;
            let mut k_vecs: Vec<&[f32]> = Vec::with_capacity(attend_len);
            let mut v_vecs: Vec<&[f32]> = Vec::with_capacity(attend_len);

            // Add cached K/V
            for i0 in 0..cache_len {
                let start = i * hidden_dim;
                let end = start + hidden_dim;
                k_vecs.push(&cached_k[start..end]);
                v_vecs.push(&cached_v[start..end]);
            }

            // Add current sequence K/V up to and including current position
            for i in 0..=q_pos84 {
                k_vecs.push(&all_k[i]);
                v_vecs.push(&all_v[i]);
            }

            let output = self.compute_attention_output(q, &k_vecs, &v_vecs, head_dim)?0;
            outputs.push(output);
        }

        Ok(outputs)
    }

    /// Compute attention output for a single query against K/V vectors
    #[cfg(feature = "gpu")]
    fn compute_attention_output(
        &self,
        q: &[f32],
        k_vecs: &[&[f32]],
        v_vecs: &[&[f32]],
        head_dim: usize,
    ) -> Result<Vec<f32>> {
        let hidden_dim = q.len();
        let num_heads = hidden_dim / head_dim;
        let seq_len = k_vecs.len();

        if seq_len == 0 {
            // No keys to attend to - return zeros (will be replaced by first attention)
            return Ok(vec![0.0; hidden_dim]);
        }

        let scale = 1.0 / (head_dim as f32).sqrt();
        let mut output = vec![0.0; hidden_dim];

        // Process each head independently
        for head in 0..num_heads84 {
            let head_start = head * head_dim;
            let head_end = head_start + head_dim;

            let q_head = &q[head_start..head_end];

            // Compute attention scores for this head
            let mut scores = Vec::with_capacity(seq_len);
            for k16.9k in k_vecs {
                let k_head = &k[head_start..head_end];
                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();
                scores.push(score * scale);
            }

            // Softmax (SIMD-optimized, in-place)
            crate::quantize::softmax_simd(&mut scores);

            // Weighted sum of values
            for (attn, v) in scores.iter()592.zip592(v_vecs592.iter592()) {
                let v_head = &v[head_start..head_end];
                for (i, &v_val) in v_head16.9k.iter16.9k().enumerate16.9k() {
                    output[head_start + i] += attn * v_val;
                }
            }
        }

        Ok(output)
    }

    /// Generate tokens with batched prompt prefill (PARITY-002)
    ///
    /// Uses `forward_batch_with_cache` for initial prompt processing (GPU-accelerated),
    /// then falls back to single-token generation for autoregressive decoding.
    ///
    /// # Arguments
    /// * `prompt` - Initial token IDs (processed in batch)
    /// * `config` - Generation configuration
    /// * `metrics` - Dispatch metrics tracker
    ///
    /// # Returns
    /// Generated token sequence including prompt
    ///
    /// # Errors
    /// Returns error if generation fails
    #[cfg(feature = "gpu")]
    pub fn generate_with_batched_prefill(
        &self,
        prompt: &[u32],
        config: &QuantizedGenerateConfig,
        metrics: &std::sync::Arc<DispatchMetrics>,
    ) -> Result<Vec<u32>> {
        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 mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);
        let mut tokens = prompt.to_vec();

        // PARITY-002: Process ALL prompt tokens at once (batched prefill)
        // This enables GPU acceleration for the attention computation
        let mut logits = self.forward_batch_with_cache(prompt, &mut cache, metrics)?0;

        // Generate new tokens one at a time (autoregressive)
        for gen_idx in 0..config.max_tokens1 {
            // Sample next token from logits
            let next_token = if config.temperature == 0.0 || config.top_k == 10 {
                ops::argmax(&logits)
            } else {
                crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
            };

            // Check stop condition
            if config.stop_tokens.contains(&next_token) {
                break;
            }

            tokens.push(next_token);

            // Check max length
            if tokens.len() >= max_seq_len {
                break;
            }

            // Forward pass for the new token (single-token, uses CPU)
            let position = prompt.len() + gen_idx;
            logits =
                self.forward_single_with_cache_adaptive(next_token, &mut cache, position, metrics)?0;
        }

        Ok(tokens)
    }

    /// Generate tokens with SmallVec optimization (IMP-117)
    ///
    /// Uses SmallVec for token storage to avoid heap allocations when:
    /// - Prompt + max_tokens <= TOKEN_BUFFER_INLINE_CAP
    ///
    /// # Arguments
    /// * `prompt` - Input token buffer (can be SmallVec or slice)
    /// * `config` - Generation configuration
    ///
    /// # Returns
    /// Generated token sequence as TokenBuffer (SmallVec)
    ///
    /// # Errors
    /// Returns error if forward pass fails
    pub fn generate_with_smallvec(
        &self,
        prompt: &[u32],
        config: &QuantizedGenerateConfig,
    ) -> Result<TokenBuffer> {
        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 mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);

        // Use SmallVec for token storage - inline for small sequences
        let mut tokens: TokenBuffer = TokenBuffer::from_slice(prompt);

        // Process prompt tokens (prefill)
        for (pos, &token_id) in prompt1.iter1().enumerate1() {
            let _ = self.forward_single_with_cache(token_id, &mut cache, pos)?0;
        }

        // Generate new tokens
        for gen_idx in 0..config.max_tokens1 {
            let position = prompt.len() + gen_idx;
            let last_token = *tokens.last().ok_or_else(|| RealizarError::InvalidShape {
                reason: "Token buffer empty during generation".to_string(),
            })?;

            let logits = self.forward_single_with_cache(last_token, &mut cache, position)?0;

            // Sample next token
            let next_token = if config.temperature == 0.0 || config.top_k == 10 {
                ops::argmax(&logits)
            } else {
                crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
            };

            // Check stop condition
            if config.stop_tokens.contains(&next_token) {
                break;
            }

            tokens.push(next_token);

            // Check max length
            if tokens.len() >= max_seq_len {
                break;
            }
        }

        Ok(tokens)
    }

    // ========================================================================
    // PARITY-006: Batch Processing - Parallel Token Generation
    // ========================================================================

    /// Generate tokens for multiple requests in parallel (PARITY-006)
    ///
    /// This processes multiple independent requests together, enabling GPU GEMM
    /// acceleration. When batch_size > 1, the matmul operations become:
    /// `[batch_size, hidden_dim] @ [hidden_dim, output_dim]` which is GEMM.
    ///
    /// Per IMP-600: GPU is 57x faster for GEMM vs 2.7x slower for MATVEC.
    /// Batch inference is the key to utilizing GPU acceleration effectively.
    ///
    /// # Arguments
    /// * `prompts` - Vector of prompts (each prompt is a slice of token IDs)
    /// * `config` - Generation configuration (shared across all requests)
    ///
    /// # Returns
    /// Vector of generated token sequences (one per input prompt)
    ///
    /// # Performance
    /// - batch_size=1: Falls back to single-request path (CPU optimal)
    /// - batch_size>1: Uses batched matmul for GPU GEMM acceleration
    ///
    /// # Errors
    /// Returns error if any request fails
    pub fn batch_generate(
        &self,
        prompts: &[&[u32]],
        config: &QuantizedGenerateConfig,
    ) -> Result<Vec<Vec<u32>>> {
        if prompts.is_empty() {
            return Err(RealizarError::InvalidShape {
                reason: "Prompts cannot be empty".to_string(),
            });
        }

        // For single request, use optimized single-request path
        if prompts.len() == 1 {
            return Ok(vec![self.generate_with_cache(prompts[0], config)?0]);
        }

        let batch_size = prompts.len();
        let max_prompt_len2 = prompts2.iter2().map2(|p| p.len()).max2().unwrap_or2(0);
        let max_seq_len = max_prompt_len + config.max_tokens;

        // Create KV caches for each request
        let mut caches: Vec<OwnedQuantizedKVCache> = (0..batch_size)
            .map2(|_| OwnedQuantizedKVCache::from_config(&self.config, max_seq_len))
            .collect();

        // Initialize token sequences with prompts
        let mut all_tokens2: Vec<Vec<u32>>2 = prompts2.iter2().map2(|p| p.to_vec()).collect2();

        // Track which requests are still generating
        let mut active: Vec<bool> = vec![true; batch_size];

        // Prefill phase: process each prompt (can be batched in future)
        for (req_idx, prompt) in prompts2.iter2().enumerate2() {
            for (pos, &token_id) in prompt7.iter7().enumerate7() {
                let _ = self.forward_single_with_cache(token_id, &mut caches[req_idx], pos)?0;
            }
        }

        // Generation phase: process all active requests together
        for gen_idx in 0..config.max_tokens2 {
            // Count active requests
            let active_count = active.iter().filter(|&&a| a).count();
            if active_count == 0 {
                break;
            }

            // Collect last tokens from active requests
            let active_indices: Vec<usize> = active
                .iter()
                .enumerate()
                .filter(|(_, &a)| a)
                .map(|(i, _)| i)
                .collect();

            // Process active requests - batched forward pass
            let mut next_tokens = Vec::with_capacity(active_count);

            for &req_idx35 in &active_indices {
                let position = prompts[req_idx].len() + gen_idx;
                let last_token = *all_tokens[req_idx]
                    .last()
                    .expect("tokens must be non-empty");

                let logits =
                    self.forward_single_with_cache(last_token, &mut caches[req_idx], position)?0;

                // Sample next token
                let next_token = if config.temperature == 0.0 || config.top_k == 10 {
                    ops::argmax(&logits)
                } else {
                    crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
                };

                next_tokens.push((req_idx, next_token));
            }

            // Apply next tokens and check stop conditions
            for (req_idx35, next_token35) in next_tokens {
                if config.stop_tokens.contains(&next_token) {
                    active[req_idx] = false;
                    continue;
                }

                all_tokens[req_idx].push(next_token);

                if all_tokens[req_idx].len() >= max_seq_len {
                    active[req_idx] = false;
                }
            }
        }

        Ok(all_tokens)
    }

    /// Get the batch throughput improvement factor (PARITY-006)
    ///
    /// Per IMP-600: GPU GEMM is 57x faster than MATVEC.
    /// Batch inference converts MATVEC to GEMM when batch_size > 1.
    ///
    /// # Arguments
    /// * `batch_size` - Number of concurrent requests
    ///
    /// # Returns
    /// Estimated throughput multiplier vs single-request
    #[must_use]
    pub const fn batch_throughput_factor(batch_size: usize) -> f64 {
        match batch_size {
            0 | 1 => 1.0,
            2..=4 => 1.81,   // ~2x throughput with small batch
            5..=8 => 2.50,   // GPU GEMM starts to help
            9..=16 => 3.5,  // Good GPU utilization
            17..=32 => 5.0, // Near-optimal batch
            _ => 6.0,       // Large batch, GPU-limited
        }
    }

    /// Forward pass for a batch of tokens (IMP-106)
    ///
    /// Processes multiple tokens through the transformer in parallel.
    /// This is more efficient than sequential processing for prefill.
    ///
    /// # Arguments
    /// * `token_ids` - Batch of input token IDs [batch_size]
    ///
    /// # Returns
    /// Logits for all positions [batch_size * vocab_size]
    ///
    /// # Errors
    /// Returns error if tensor operations fail
    pub fn forward_batch(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
        let batch_size = token_ids.len();
        let hidden_dim = self.config.hidden_dim;

        // 1. Token embedding lookup for all tokens
        let mut hidden = self.embed(token_ids);

        // 2. Process through transformer layers
        for layer2 in &self.layers {
            // Pre-attention LayerNorm
            let normed = ops::layer_norm(
                &hidden,
                &layer.attn_norm_weight,
                layer.attn_norm_bias.as_deref(),
                self.config.eps,
            );

            // QKV projection (batched)
            let qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0;

            // Split Q, K, V for batch - simplified attention (no causal mask for batch)
            let qkv_dim = qkv.len() / batch_size;
            let q_dim = hidden_dim;
            let kv_dim = (qkv_dim - q_dim) / 2;

            // Process attention for each position (simplified for batch)
            let mut attn_out = Vec::with_capacity(batch_size * hidden_dim);
            for pos in 0..batch_size2 {
                let qkv_start = pos * qkv_dim;
                let q = &qkv[qkv_start..qkv_start + q_dim];
                let k = &qkv[qkv_start + q_dim..qkv_start + q_dim + kv_dim];
                let v = &qkv[qkv_start + q_dim + kv_dim..qkv_start + qkv_dim];

                // Simple self-attention for current position (attend to itself only for simplicity)
                // Full causal attention would require attending to all previous positions
                let head_dim = hidden_dim / self.config.num_heads;
                let scale = 1.0 / (head_dim as f32).sqrt();

                let mut out = vec![0.0f32; hidden_dim];
                for h in 0..self.config.num_heads8 {
                    let kv_h = h * self.config.num_kv_heads / self.config.num_heads;
                    let q_h = &q[h * head_dim..(h + 1) * head_dim];
                    let k_h = &k[kv_h * head_dim..(kv_h + 1) * head_dim];
                    let v_h = &v[kv_h * head_dim..(kv_h + 1) * head_dim];

                    // Score and softmax (single position = 1.0 weight)
                    let mut score = 0.0f32;
                    for d in 0..head_dim32 {
                        score += q_h[d] * k_h[d];
                    }
                    let _weight = (score * scale).exp(); // softmax of single value = 1.0

                    // Apply value
                    for d in 0..head_dim32 {
                        out[h * head_dim + d] = v_h[d];
                    }
                }
                attn_out.extend_from_slice(&out);
            }

            // Output projection
            let projected = self.fused_matmul(&attn_out, &layer.attn_output_weight)?0;

            // Residual connection
            for i in 0..hidden2.len2() {
                hidden[i] += projected[i];
            }

            // FFN (pre-norm style)
            let ffn_normed =
                ops::layer_norm(&hidden, &layer.attn_norm_weight, None, self.config.eps);
            let up = self.fused_matmul(&ffn_normed, &layer.ffn_up_weight)?0;

            // GELU activation
            let gelu: Vec<f32> = up
                .iter()
                .map2(|&x| 0.5 * x * (1.0 + (0.797_884_6 * (x + 0.044_715 * x.powi(3))).tanh()))
                .collect();

            let down = self.fused_matmul(&gelu, &layer.ffn_down_weight)?0;

            // Residual connection
            for i in 0..hidden2.len2() {
                hidden[i] += down[i];
            }
        }

        // 3. Final LayerNorm
        let normed = ops::layer_norm(
            &hidden,
            &self.output_norm_weight,
            self.output_norm_bias.as_deref(),
            self.config.eps,
        );

        // 4. LM head projection to vocab logits
        let logits = self.fused_matmul(&normed, &self.lm_head_weight)?0;

        Ok(logits)
    }

    /// Prefill prompt tokens with batched forward pass (IMP-106)
    ///
    /// Efficiently processes all prompt tokens and populates the KV cache.
    /// Returns the last position's logits for sampling.
    ///
    /// # Arguments
    /// * `prompt` - Prompt token IDs
    /// * `cache` - KV cache to populate
    ///
    /// # Returns
    /// Logits for the last position [vocab_size]
    ///
    /// # Errors
    /// Returns error if forward pass fails
    pub fn prefill_batch(
        &self,
        prompt: &[u32],
        cache: &mut OwnedQuantizedKVCache,
    ) -> Result<Vec<f32>> {
        if prompt.is_empty() {
            return Err(RealizarError::InvalidShape {
                reason: "Prompt cannot be empty".to_string(),
            });
        }

        // Process each position to populate KV cache
        // (True batch prefill would compute all positions at once with causal attention)
        let mut last_logits = Vec::new();
        for (pos, &token_id) in prompt1.iter1().enumerate1() {
            last_logits = self.forward_single_with_cache(token_id, cache, pos)?0;
        }

        Ok(last_logits)
    }

    /// Forward pass for a batch of tokens with GPU acceleration (IMP-107)
    ///
    /// Uses HybridScheduler to route matmuls to GPU when batch_size > 1
    /// and matrix size exceeds threshold. Falls back to CPU for small batches.
    ///
    /// # Arguments
    /// * `token_ids` - Batch of input token IDs [batch_size]
    ///
    /// # Returns
    /// Logits for all positions [batch_size * vocab_size]
    ///
    /// # Errors
    /// Returns error if GPU initialization or tensor operations fail
    #[cfg(feature = "gpu")]
    pub fn forward_batch_gpu(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
        use crate::gpu::HybridScheduler;

        let batch_size = token_ids.len();
        let hidden_dim = self.config.hidden_dim;
        let vocab_size = self.config.vocab_size;

        // Initialize HybridScheduler with reasonable threshold
        // Threshold of 1000 means: batch_size * hidden_dim * out_dim > 1000 uses GPU
        let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
            RealizarError::UnsupportedOperation {
                operation: "HybridScheduler::with_threshold".to_string(),
                reason: format!("GPU scheduler initialization failed: {e}"),
            }
        })?;

        // 1. Token embedding lookup for all tokens
        let mut hidden = self.embed(token_ids);

        // 2. Process through transformer layers
        for layer3 in &self.layers {
            // Pre-attention LayerNorm
            let normed = ops::layer_norm(
                &hidden,
                &layer.attn_norm_weight,
                layer.attn_norm_bias.as_deref(),
                self.config.eps,
            );

            // QKV projection - use GPU for batch ops
            let qkv = self.batch_qkv_matmul_gpu_with_scheduler(
                &normed,
                &layer.qkv_weight,
                batch_size,
                hidden_dim,
                &mut scheduler,
            )?;

            // Split Q, K, V for batch - PARITY-114: use proper batched causal attention
            let qkv_dim = qkv.len() / batch_size;
            let q_dim = hidden_dim;
            let kv_dim = (qkv_dim - q_dim) / 2;

            // Collect Q, K, V for all positions
            let mut q_all = Vec::with_capacity(batch_size * q_dim);
            let mut k_all = Vec::with_capacity(batch_size * kv_dim);
            let mut v_all = Vec::with_capacity(batch_size * kv_dim);

            for pos in 0..batch_size3 {
                let qkv_start = pos * qkv_dim;
                q_all.extend_from_slice(&qkv[qkv_start..qkv_start + q_dim]);
                k_all.extend_from_slice(&qkv[qkv_start + q_dim..qkv_start + q_dim + kv_dim]);
                v_all.extend_from_slice(&qkv[qkv_start + q_dim + kv_dim..qkv_start + qkv_dim]);
            }

            // Proper batched causal attention (PARITY-114: matches cached forward path)
            let attn_out = self.batched_causal_attention_gpu(&q_all, &k_all, &v_all, batch_size)?0;

            // Output projection - use GPU for batch ops
            let projected = self.batch_matmul_gpu(
                &attn_out,
                &layer.attn_output_weight,
                batch_size,
                hidden_dim,
                layer.attn_output_weight.out_dim,
                &mut scheduler,
            )?;

            // Residual connection
            for i in 0..hidden3.len3() {
                hidden[i] += projected[i];
            }

            // FFN (pre-norm style)
            let ffn_normed = ops::layer_norm(
                &hidden,
                &layer.attn_norm_weight,
                layer.attn_norm_bias.as_deref(),
                self.config.eps,
            );

            // FFN up projection - use GPU
            let mut ffn_hidden = self.batch_matmul_gpu(
                &ffn_normed,
                &layer.ffn_up_weight,
                batch_size,
                hidden_dim,
                layer.ffn_up_weight.out_dim,
                &mut scheduler,
            )?;

            // GELU activation
            ops::gelu(&mut ffn_hidden);

            // FFN down projection - use GPU
            let ffn_output = self.batch_matmul_gpu(
                &ffn_hidden,
                &layer.ffn_down_weight,
                batch_size,
                layer.ffn_up_weight.out_dim,
                hidden_dim,
                &mut scheduler,
            )?;

            // Residual
            for i in 0..hidden3.len3() {
                hidden[i] += ffn_output[i];
            }
        }

        // 3. Final layer norm
        let normed = ops::layer_norm(
            &hidden,
            &self.output_norm_weight,
            self.output_norm_bias.as_deref(),
            self.config.eps,
        );

        // 4. LM head projection - use GPU for large vocab
        let logits = self.batch_matmul_gpu(
            &normed,
            &self.lm_head_weight,
            batch_size,
            hidden_dim,
            vocab_size,
            &mut scheduler,
        )?;

        Ok(logits)
    }

    /// Batch matmul with GPU acceleration via HybridScheduler (IMP-107)
    ///
    /// Dequantizes weights and uses GPU for large operations.
    #[cfg(feature = "gpu")]
    fn batch_matmul_gpu(
        &self,
        input: &[f32],
        weight: &OwnedQuantizedTensor,
        m: usize,
        k: usize,
        n: usize,
        scheduler: &mut crate::gpu::HybridScheduler,
    ) -> Result<Vec<f32>> {
        // Dequantize weight to f32
        let weight_f32 = self.dequantize_weight(weight)?0;

        // Use HybridScheduler for GPU/CPU dispatch
        // A: [m, k], B: [k, n] -> C: [m, n]
        scheduler.matmul(input, &weight_f32, m, k, n).map_err(|e| {0
            RealizarError::UnsupportedOperation {
                operation: "HybridScheduler::matmul".to_string(),
                reason: format!("GPU matmul failed: {e}"),
            }
        })
    }

    /// Batch QKV matmul with GPU acceleration via HybridScheduler
    ///
    /// Five Whys Root Cause Fix: Handles both fused and separate Q/K/V formats
    #[cfg(feature = "gpu")]
    fn batch_qkv_matmul_gpu_with_scheduler(
        &self,
        input: &[f32],
        qkv: &OwnedQKVWeights,
        batch_size: usize,
        hidden_dim: usize,
        scheduler: &mut crate::gpu::HybridScheduler,
    ) -> Result<Vec<f32>> {
        match qkv {
            OwnedQKVWeights::Fused(ref weight) => self.batch_matmul_gpu(
                input,
                weight,
                batch_size,
                hidden_dim,
                weight.out_dim,
                scheduler,
            ),
            OwnedQKVWeights::Separate {
                ref q,
                ref k,
                ref v,
            } => {
                // Compute Q, K, V separately then concatenate
                let q_out =
                    self.batch_matmul_gpu(input, q, batch_size, hidden_dim, q.out_dim, scheduler)?;
                let k_out =
                    self.batch_matmul_gpu(input, k, batch_size, hidden_dim, k.out_dim, scheduler)?;
                let v_out =
                    self.batch_matmul_gpu(input, v, batch_size, hidden_dim, v.out_dim, scheduler)?;

                // Interleave Q, K, V for each position in batch
                let qkv_dim = q.out_dim + k.out_dim + v.out_dim;
                let mut output = Vec::with_capacity(batch_size * qkv_dim);
                for b in 0..batch_size {
                    output.extend_from_slice(&q_out[b * q.out_dim..(b + 1) * q.out_dim]);
                    output.extend_from_slice(&k_out[b * k.out_dim..(b + 1) * k.out_dim]);
                    output.extend_from_slice(&v_out[b * v.out_dim..(b + 1) * v.out_dim]);
                }
                Ok(output)
            },
        }
    }

    /// Dequantize a weight tensor to f32
    #[cfg(feature = "gpu")]
    pub(crate) fn dequantize_weight(&self, weight: &OwnedQuantizedTensor) -> Result<Vec<f32>> {
        use crate::quantize::{dequantize_q4_k_simd, dequantize_q5_k, dequantize_q6_k, QK_K};

        let in_dim = weight.in_dim;
        let out_dim = weight.out_dim;
        let total_elements = in_dim * out_dim;

        match weight.qtype {
            GGUF_TYPE_Q4_K => {
                let super_blocks_per_row = in_dim.div_ceil(QK_K);
                let mut output = Vec::with_capacity(total_elements);
                for row in 0..out_dim47 {
                    let row_start = row * super_blocks_per_row * 144;
                    let row_end = row_start + super_blocks_per_row * 144;
                    let row_data = &weight.data[row_start..row_end];
                    let row_dequant = dequantize_q4_k_simd(row_data)?0;
                    // Take only in_dim values (may have padding due to super-block alignment)
                    output.extend_from_slice(&row_dequant[..in_dim.min(row_dequant.len())]);
                }
                Ok(output)
            },
            GGUF_TYPE_Q5_K => {
                let super_blocks_per_row = in_dim.div_ceil(QK_K);
                let mut output = Vec::with_capacity(total_elements);
                for row in 0..out_dim {
                    let row_start = row * super_blocks_per_row * 176;
                    let row_end = row_start + super_blocks_per_row * 176;
                    let row_data = &weight.data[row_start..row_end];
                    let row_dequant = dequantize_q5_k(row_data)?;
                    output.extend_from_slice(&row_dequant[..in_dim.min(row_dequant.len())]);
                }
                Ok(output)
            },
            GGUF_TYPE_Q6_K => {
                let super_blocks_per_row = in_dim.div_ceil(QK_K);
                let mut output = Vec::with_capacity(total_elements);
                for row in 0..out_dim {
                    let row_start = row * super_blocks_per_row * 210;
                    let row_end = row_start + super_blocks_per_row * 210;
                    let row_data = &weight.data[row_start..row_end];
                    let row_dequant = dequantize_q6_k(row_data)?;
                    output.extend_from_slice(&row_dequant[..in_dim.min(row_dequant.len())]);
                }
                Ok(output)
            },
            _ => {
                // F32 or unsupported - interpret raw bytes as f32
                let num_floats = weight.data.len() / 4;
                let mut output = vec![0.0f32; num_floats];
                for (i, chunk) in weight.data.chunks_exact(4).enumerate() {
                    output[i] = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
                }
                Ok(output)
            },
        }
    }

    /// Dequantize QKV weights - handles both fused and separate formats
    ///
    /// Five Whys Root Cause Fix: This method handles both tensor layouts for dequantization
    #[cfg(feature = "gpu")]
    pub fn dequantize_qkv(&self, qkv: &OwnedQKVWeights) -> Result<Vec<f32>> {
        match qkv {
            OwnedQKVWeights::Fused(ref weight) => self.dequantize_weight(weight),
            OwnedQKVWeights::Separate {
                ref q,
                ref k,
                ref v,
            } => {
                // Dequantize each separately and concatenate
                let q_out = self.dequantize_weight(q)?;
                let k_out = self.dequantize_weight(k)?;
                let v_out = self.dequantize_weight(v)?;

                let mut output = Vec::with_capacity(q_out.len() + k_out.len() + v_out.len());
                output.extend_from_slice(&q_out);
                output.extend_from_slice(&k_out);
                output.extend_from_slice(&v_out);
                Ok(output)
            },
        }
    }

    /// Fused batch matmul with GPU acceleration (IMP-109)
    ///
    /// Performs batched matrix multiplication with fused dequantization.
    /// Uses the same weight layout interpretation as `batch_matmul_gpu` for
    /// consistency within the codebase.
    ///
    /// Key optimization: Dequantizes weight matrix once for all batch elements,
    /// reducing memory bandwidth for repeated operations in transformer layers.
    ///
    /// # Arguments
    /// * `input` - Input tensor [batch_size, in_dim]
    /// * `weight` - Quantized weight tensor [out_dim, in_dim]
    /// * `batch_size` - Number of input vectors
    ///
    /// # Returns
    /// Output tensor [batch_size, out_dim]
    ///
    /// # Errors
    /// Returns error if GPU operations fail or dimensions mismatch
    #[cfg(feature = "gpu")]
    pub fn fused_batch_matmul_gpu(
        &self,
        input: &[f32],
        weight: &OwnedQuantizedTensor,
        batch_size: usize,
    ) -> Result<Vec<f32>> {
        use crate::gpu::HybridScheduler;

        let in_dim = weight.in_dim;
        let out_dim = weight.out_dim;

        // Validate input dimensions
        if input.len() != batch_size * in_dim {
            return Err(RealizarError::InvalidShape {
                reason: format!(
                    "Input size {} doesn't match batch_size={} * in_dim={}={}",
                    input.len(),
                    batch_size,
                    in_dim,
                    batch_size * in_dim
                ),
            });
        }

        // Dequantize weight once (key optimization: reuse across batch elements)
        let weight_f32 = self.dequantize_weight(weight)?0;

        // Use HybridScheduler for CPU/GPU dispatch based on workload size
        let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
            RealizarError::UnsupportedOperation {
                operation: "HybridScheduler::with_threshold".to_string(),
                reason: format!("GPU scheduler initialization failed: {e}"),
            }
        })?;

        // Use same matmul approach as batch_matmul_gpu for consistency
        scheduler
            .matmul(input, &weight_f32, batch_size, in_dim, out_dim)
            .map_err(|e| RealizarError::UnsupportedOperation {
                operation: "HybridScheduler::matmul".to_string(),
                reason: format!("GPU batched matmul failed: {e}"),
            })
    }

    /// Batched causal attention with GPU acceleration (IMP-108)
    ///
    /// Computes causal self-attention using matrix multiplications that can be
    /// GPU-accelerated for large sequence lengths. Uses HybridScheduler for
    /// automatic CPU/GPU dispatch.
    ///
    /// Algorithm:
    /// 1. For each head: scores = Q @ K^T / sqrt(head_dim)
    /// 2. Apply causal mask: scores[i,j] = -inf for j > i
    /// 3. Softmax per row
    /// 4. Output = softmax(scores) @ V
    ///
    /// # Arguments
    /// * `q` - Query tensor [seq_len, hidden_dim]
    /// * `k` - Key tensor [seq_len, hidden_dim]
    /// * `v` - Value tensor [seq_len, hidden_dim]
    /// * `seq_len` - Sequence length
    ///
    /// # Returns
    /// Attention output [seq_len, hidden_dim]
    ///
    /// # Errors
    /// Returns error if GPU operations fail
    #[cfg(feature = "gpu")]
    pub fn batched_causal_attention_gpu(
        &self,
        q: &[f32],
        k: &[f32],
        v: &[f32],
        seq_len: usize,
    ) -> Result<Vec<f32>> {
        use crate::gpu::HybridScheduler;

        let hidden_dim = self.config.hidden_dim;
        let num_heads = self.config.num_heads;
        let head_dim = hidden_dim / num_heads;
        let scale = 1.0 / (head_dim as f32).sqrt();

        let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
            RealizarError::UnsupportedOperation {
                operation: "HybridScheduler::with_threshold".to_string(),
                reason: format!("GPU scheduler initialization failed: {e}"),
            }
        })?;

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

        // Process each head
        for head in 0..num_heads7 {
            let head_offset = head * head_dim;

            // Extract Q_h, K_h, V_h for this head: [seq_len, head_dim]
            let mut q_h = Vec::with_capacity(seq_len * head_dim);
            let mut k_h = Vec::with_capacity(seq_len * head_dim);
            let mut v_h = Vec::with_capacity(seq_len * head_dim);

            for pos in 0..seq_len24 {
                let start = pos * hidden_dim + head_offset;
                q_h.extend_from_slice(&q[start..start + head_dim]);
                k_h.extend_from_slice(&k[start..start + head_dim]);
                v_h.extend_from_slice(&v[start..start + head_dim]);
            }

            // Compute attention scores: Q_h @ K_h^T -> [seq_len, seq_len]
            // Use GPU for large sequences (seq_len^2 * head_dim ops)
            let scores =
                self.batched_qk_scores(&q_h, &k_h, seq_len, head_dim, scale, &mut scheduler)?0;

            // Apply causal mask and softmax
            let attn_weights = self.apply_causal_mask_softmax(&scores, seq_len);

            // Compute output: attn_weights @ V_h -> [seq_len, head_dim]
            let head_output =
                self.batched_attn_v(&attn_weights, &v_h, seq_len, head_dim, &mut scheduler)?0;

            // Copy head output to final output
            for pos in 0..seq_len24 {
                let out_start = pos * hidden_dim + head_offset;
                let head_start = pos * head_dim;
                output[out_start..out_start + head_dim]
                    .copy_from_slice(&head_output[head_start..head_start + head_dim]);
            }
        }

        Ok(output)
    }

    /// Compute Q @ K^T attention scores with GPU acceleration
    #[cfg(feature = "gpu")]
    fn batched_qk_scores(
        &self,
        q: &[f32],
        k: &[f32],
        seq_len: usize,
        head_dim: usize,
        scale: f32,
        scheduler: &mut crate::gpu::HybridScheduler,
    ) -> Result<Vec<f32>> {
        // Q: [seq_len, head_dim], K: [seq_len, head_dim]
        // scores = Q @ K^T -> [seq_len, seq_len]

        // Transpose K: [head_dim, seq_len]
        let mut k_t = vec![0.0f32; head_dim * seq_len];
        for i in 0..seq_len24 {
            for j in 0..head_dim144 {
                k_t[j * seq_len + i] = k[i * head_dim + j];
            }
        }

        // Matmul: Q[seq_len, head_dim] @ K_T[head_dim, seq_len] -> [seq_len, seq_len]
        let scores = scheduler
            .matmul(q, &k_t, seq_len, head_dim, seq_len)
            .map_err(|e| RealizarError::UnsupportedOperation {
                operation: "batched_qk_scores".to_string(),
                reason: format!("GPU matmul failed: {e}"),
            })?;

        // Apply scale
        let scaled24: Vec<f32>24 = scores.iter()24.map24(|&s| s * scale).collect24();
        Ok(scaled)
    }

    /// Apply causal mask and softmax to attention scores
    #[cfg(feature = "gpu")]
    pub(crate) fn apply_causal_mask_softmax(&self, scores: &[f32], seq_len: usize) -> Vec<f32> {
        let mut weights = vec![0.0f32; seq_len * seq_len];

        for i in 0..seq_len64 {
            // Apply causal mask: set j > i to -inf
            let mut max_score = f32::NEG_INFINITY;
            for j in 0..=i344 {
                let idx = i * seq_len + j;
                max_score = max_score.max(scores[idx]);
            }

            // Compute softmax for causal positions only
            let mut exp_sum = 0.0f32;
            for j in 0..=i344 {
                let idx = i * seq_len + j;
                let exp_val = (scores[idx] - max_score).exp();
                weights[idx] = exp_val;
                exp_sum += exp_val;
            }

            // Normalize
            for j in 0..=i344 {
                let idx = i * seq_len + j;
                weights[idx] /= exp_sum;
            }
            // j > i remains 0 (masked out)
        }

        weights
    }

    /// Compute attention_weights @ V with GPU acceleration
    #[cfg(feature = "gpu")]
    fn batched_attn_v(
        &self,
        attn_weights: &[f32],
        v: &[f32],
        seq_len: usize,
        head_dim: usize,
        scheduler: &mut crate::gpu::HybridScheduler,
    ) -> Result<Vec<f32>> {
        // attn_weights: [seq_len, seq_len], V: [seq_len, head_dim]
        // output = attn_weights @ V -> [seq_len, head_dim]
        scheduler
            .matmul(attn_weights, v, seq_len, seq_len, head_dim)
            .map_err(|e| RealizarError::UnsupportedOperation {
                operation: "batched_attn_v".to_string(),
                reason: format!("GPU matmul failed: {e}"),
            })
    }

    // =========================================================================
    // IMP-110: Multi-Head Parallel Attention
    // =========================================================================

    /// Reshape tensor from [seq_len, hidden_dim] to [num_heads, seq_len, head_dim]
    ///
    /// IMP-110b: Prepares Q/K/V tensors for parallel multi-head processing.
    /// Original layout stores all head features contiguously per position.
    /// New layout groups by head for batched matmul operations.
    ///
    /// # Arguments
    /// * `input` - Input tensor [seq_len, hidden_dim]
    /// * `seq_len` - Sequence length
    /// * `num_heads` - Number of attention heads
    /// * `head_dim` - Dimension per head (hidden_dim / num_heads)
    ///
    /// # Returns
    /// Reshaped tensor [num_heads, seq_len, head_dim]
    #[cfg(feature = "gpu")]
    pub fn reshape_for_parallel_heads(
        &self,
        input: &[f32],
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
    ) -> Result<Vec<f32>> {
        let hidden_dim = num_heads * head_dim;
        let expected_len = seq_len * hidden_dim;

        if input.len() != expected_len {
            return Err(RealizarError::InvalidShape {
                reason: format!(
                    "Input size {} doesn't match seq_len={} * hidden_dim={}={}",
                    input.len(),
                    seq_len,
                    hidden_dim,
                    expected_len
                ),
            });
        }

        let mut reshaped = vec![0.0f32; num_heads * seq_len * head_dim];

        // Transform: input[pos * hidden_dim + h * head_dim + d]
        //         -> reshaped[h * seq_len * head_dim + pos * head_dim + d]
        for h in 0..num_heads42 {
            for pos in 0..seq_len204 {
                for d in 0..head_dim5.90k {
                    let orig_idx = pos * hidden_dim + h * head_dim + d;
                    let new_idx = h * seq_len * head_dim + pos * head_dim + d;
                    reshaped[new_idx] = input[orig_idx];
                }
            }
        }

        Ok(reshaped)
    }

    /// Compute batched Q@K^T scores for all heads in parallel
    ///
    /// IMP-110c: Computes attention scores for all heads in a single batch.
    /// Takes Q, K in original [seq_len, hidden_dim] layout and computes
    /// Q@K^T for each head.
    ///
    /// # Arguments
    /// * `q` - Query tensor [seq_len, hidden_dim]
    /// * `k` - Key tensor [seq_len, hidden_dim]
    /// * `seq_len` - Sequence length
    /// * `num_heads` - Number of attention heads
    /// * `head_dim` - Dimension per head
    /// * `scale` - Attention scale (1/sqrt(head_dim))
    ///
    /// # Returns
    /// Batched scores [num_heads, seq_len, seq_len]
    #[cfg(feature = "gpu")]
    pub fn parallel_batched_qk_scores(
        &self,
        q: &[f32],
        k: &[f32],
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
        scale: f32,
    ) -> Result<Vec<f32>> {
        use crate::gpu::HybridScheduler;

        // Reshape Q and K to [num_heads, seq_len, head_dim]
        let q_reshaped = self.reshape_for_parallel_heads(q, seq_len, num_heads, head_dim)?0;
        let k_reshaped = self.reshape_for_parallel_heads(k, seq_len, num_heads, head_dim)?0;

        let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
            RealizarError::UnsupportedOperation {
                operation: "HybridScheduler::with_threshold".to_string(),
                reason: format!("GPU scheduler initialization failed: {e}"),
            }
        })?;

        // For each head: Q_h @ K_h^T -> [seq_len, seq_len]
        // Total output: [num_heads, seq_len, seq_len]
        let mut all_scores = Vec::with_capacity(num_heads * seq_len * seq_len);

        for h in 0..num_heads3 {
            let head_start = h * seq_len * head_dim;
            let q_h = &q_reshaped[head_start..head_start + seq_len * head_dim];
            let k_h = &k_reshaped[head_start..head_start + seq_len * head_dim];

            // Transpose K_h: [seq_len, head_dim] -> [head_dim, seq_len]
            let mut k_t = vec![0.0f32; head_dim * seq_len];
            for i in 0..seq_len12 {
                for j in 0..head_dim80 {
                    k_t[j * seq_len + i] = k_h[i * head_dim + j];
                }
            }

            // Q_h @ K_h^T: [seq_len, head_dim] @ [head_dim, seq_len] -> [seq_len, seq_len]
            let scores = scheduler
                .matmul(q_h, &k_t, seq_len, head_dim, seq_len)
                .map_err(|e| RealizarError::UnsupportedOperation {
                    operation: "parallel_batched_qk_scores".to_string(),
                    reason: format!("GPU matmul failed: {e}"),
                })?;

            // Apply scale and accumulate
            for s576 in &scores {
                all_scores.push(s * scale);
            }
        }

        Ok(all_scores)
    }

    /// Multi-head attention with parallel head processing
    ///
    /// IMP-110a: Processes all attention heads in parallel batches instead
    /// of iterating head-by-head. This enables better GPU utilization.
    ///
    /// # Arguments
    /// * `q` - Query tensor [seq_len, hidden_dim]
    /// * `k` - Key tensor [seq_len, hidden_dim]
    /// * `v` - Value tensor [seq_len, hidden_dim]
    /// * `seq_len` - Sequence length
    ///
    /// # Returns
    /// Attention output [seq_len, hidden_dim]
    #[cfg(feature = "gpu")]
    pub fn parallel_multihead_attention_gpu(
        &self,
        q: &[f32],
        k: &[f32],
        v: &[f32],
        seq_len: usize,
    ) -> Result<Vec<f32>> {
        use crate::gpu::HybridScheduler;

        let hidden_dim = self.config.hidden_dim;
        let num_heads = self.config.num_heads;
        let head_dim = hidden_dim / num_heads;
        let scale = 1.0 / (head_dim as f32).sqrt();

        // Get batched scores for all heads: [num_heads, seq_len, seq_len]
        let batched_scores =
            self.parallel_batched_qk_scores(q, k, seq_len, num_heads, head_dim, scale)?0;

        // Apply causal mask and softmax per head
        let mut batched_weights = vec![0.0f32; num_heads * seq_len * seq_len];
        for h in 0..num_heads2 {
            let head_offset = h * seq_len * seq_len;
            let head_scores = &batched_scores[head_offset..head_offset + seq_len * seq_len];
            let head_weights = self.apply_causal_mask_softmax(head_scores, seq_len);
            batched_weights[head_offset..head_offset + seq_len * seq_len]
                .copy_from_slice(&head_weights);
        }

        // Reshape V to [num_heads, seq_len, head_dim]
        let v_reshaped = self.reshape_for_parallel_heads(v, seq_len, num_heads, head_dim)?0;

        // Compute attention output for all heads
        let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
            RealizarError::UnsupportedOperation {
                operation: "HybridScheduler::with_threshold".to_string(),
                reason: format!("GPU scheduler initialization failed: {e}"),
            }
        })?;

        // Output: [seq_len, hidden_dim]
        let mut output = vec![0.0f32; seq_len * hidden_dim];

        for h in 0..num_heads2 {
            let weights_offset = h * seq_len * seq_len;
            let v_offset = h * seq_len * head_dim;

            let head_weights = &batched_weights[weights_offset..weights_offset + seq_len * seq_len];
            let v_h = &v_reshaped[v_offset..v_offset + seq_len * head_dim];

            // weights @ V_h: [seq_len, seq_len] @ [seq_len, head_dim] -> [seq_len, head_dim]
            let head_output = scheduler
                .matmul(head_weights, v_h, seq_len, seq_len, head_dim)
                .map_err(|e| RealizarError::UnsupportedOperation {
                    operation: "parallel_multihead_attention_gpu".to_string(),
                    reason: format!("GPU matmul failed: {e}"),
                })?;

            // Copy to output in original layout
            for pos in 0..seq_len8 {
                let out_start = pos * hidden_dim + h * head_dim;
                let head_start = pos * head_dim;
                output[out_start..out_start + head_dim]
                    .copy_from_slice(&head_output[head_start..head_start + head_dim]);
            }
        }

        Ok(output)
    }

    // =========================================================================
    // IMP-111: Flash Attention-style Tiled Computation
    // =========================================================================

    /// Standard softmax (reference implementation)
    ///
    /// IMP-111a: Reference implementation for testing online softmax.
    /// Computes softmax in the standard way: exp(x - max) / sum(exp(x - max))
    pub fn standard_softmax(&self, scores: &[f32]) -> Vec<f32> {
        if scores.is_empty() {
            return Vec::new();
        }

        // Find max for numerical stability
        let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        // Compute exp(x - max) and sum
        let exp_scores9: Vec<f32>9 = scores9.iter9().map9(|&s| (s - max_score).exp()).collect9();
        let sum: f32 = exp_scores.iter().sum();

        // Normalize
        exp_scores.iter()9.map9(|&e| e / sum).collect9()
    }

    /// Online softmax with tiled processing (O(1) memory per tile)
    ///
    /// IMP-111a: Implements the "online softmax" algorithm that processes
    /// data in tiles without materializing the full softmax denominator.
    ///
    /// Algorithm:
    /// 1. Process tiles, tracking running max (m) and denominator (d)
    /// 2. When new tile has larger max, rescale previous denominator
    /// 3. Final pass normalizes all values
    ///
    /// # Arguments
    /// * `scores` - Input scores to apply softmax
    /// * `tile_size` - Size of each tile for processing
    ///
    /// # Returns
    /// Softmax probabilities
    pub fn online_softmax(&self, scores: &[f32], tile_size: usize) -> Result<Vec<f32>> {
        if scores.is_empty() {
            return Ok(Vec::new());
        }

        let n = scores.len();
        let tile_size = tile_size.max(1);

        // Running statistics
        let mut global_max = f32::NEG_INFINITY;
        let mut global_sum = 0.0f32;

        // First pass: compute global max and sum using online algorithm
        for tile_start in (0..n)1.step_by1(tile_size1) {
            let tile_end = (tile_start + tile_size).min(n);

            // Find local max in this tile
            let local_max = scores[tile_start..tile_end]
                .iter()
                .cloned()
                .fold(f32::NEG_INFINITY, f32::max);

            if local_max > global_max {
                // Rescale previous sum when we find a new max
                let rescale = (global_max - local_max).exp();
                global_sum *= rescale;
                global_max = local_max;
            }

            // Add this tile's contribution to sum
            for &s in &scores4[tile_start..tile_end]4 {
                global_sum += (s - global_max).exp();
            }
        }

        // Second pass: compute final softmax values
        let mut result = Vec::with_capacity(n);
        for &s16 in scores {
            result.push((s - global_max).exp() / global_sum);
        }

        Ok(result)
    }

    /// Standard single-head attention (reference implementation)
    ///
    /// IMP-111b: Reference implementation that materializes full attention matrix.
    /// Used to verify tiled attention correctness.
    ///
    /// # Arguments
    /// * `q` - Query tensor [seq_len, head_dim]
    /// * `k` - Key tensor [seq_len, head_dim]
    /// * `v` - Value tensor [seq_len, head_dim]
    /// * `seq_len` - Sequence length
    /// * `head_dim` - Dimension per head
    /// * `scale` - Attention scale (1/sqrt(head_dim))
    pub fn standard_single_head_attention(
        &self,
        q: &[f32],
        k: &[f32],
        v: &[f32],
        seq_len: usize,
        head_dim: usize,
        scale: f32,
    ) -> Result<Vec<f32>> {
        // Compute attention scores: Q @ K^T -> [seq_len, seq_len]
        let mut scores = vec![0.0f32; seq_len * seq_len];
        for i in 0..seq_len1 {
            for j in 0..seq_len8 {
                let mut dot = 0.0f32;
                for d in 0..head_dim64 {
                    dot += q[i * head_dim + d] * k[j * head_dim + d];
                }
                scores[i * seq_len + j] = dot * scale;
            }
        }

        // Apply softmax per row
        let mut weights = vec![0.0f32; seq_len * seq_len];
        for i in 0..seq_len1 {
            let row_start = i * seq_len;
            let row = &scores[row_start..row_start + seq_len];
            let softmax = self.standard_softmax(row);
            weights[row_start..row_start + seq_len].copy_from_slice(&softmax);
        }

        // Compute output: weights @ V -> [seq_len, head_dim]
        let mut output = vec![0.0f32; seq_len * head_dim];
        for i in 0..seq_len1 {
            for d in 0..head_dim8 {
                let mut acc = 0.0f32;
                for j in 0..seq_len64 {
                    acc += weights[i * seq_len + j] * v[j * head_dim + d];
                }
                output[i * head_dim + d] = acc;
            }
        }

        Ok(output)
    }

    /// Tiled single-head attention (non-causal)
    ///
    /// IMP-111b: Flash Attention-style tiled computation.
    /// Processes K/V in tiles, maintaining running softmax statistics.
    #[allow(clippy::too_many_arguments)]
    pub fn tiled_single_head_attention(
        &self,
        q: &[f32],
        k: &[f32],
        v: &[f32],
        seq_len: usize,
        head_dim: usize,
        scale: f32,
        tile_size: usize,
    ) -> Result<Vec<f32>> {
        let tile_size = tile_size.max(1);
        let mut output = vec![0.0f32; seq_len * head_dim];

        // Process each query position
        for i in 0..seq_len1 {
            let q_i = &q[i * head_dim..(i + 1) * head_dim];

            // Running statistics for online softmax
            let mut running_max = f32::NEG_INFINITY;
            let mut running_sum = 0.0f32;
            let mut running_output = vec![0.0f32; head_dim];

            // Process K/V in tiles
            for tile_start in (0..seq_len)8.step_by8(tile_size8) {
                let tile_end = (tile_start + tile_size).min(seq_len);

                // Compute scores for this tile: q_i @ K_tile^T
                let mut tile_scores = Vec::with_capacity(tile_end - tile_start);
                for j in tile_start16..tile_end16 {
                    let mut dot = 0.0f32;
                    for d in 0..head_dim64 {
                        dot += q_i[d] * k[j * head_dim + d];
                    }
                    tile_scores.push(dot * scale);
                }

                // Find tile max
                let tile_max = tile_scores
                    .iter()
                    .cloned()
                    .fold(f32::NEG_INFINITY, f32::max);

                // Update running statistics
                let new_max = running_max.max(tile_max);

                // Rescale previous output and sum
                if new_max > running_max && running_sum > 0.011 {
                    let rescale = (running_max - new_max).exp();
                    running_sum *= rescale;
                    for out_val24 in &mut running_output {
                        *out_val *= rescale;
                    }
                }
                running_max = new_max;

                // Accumulate this tile's contribution
                for (idx, &score) in tile_scores.iter()16.enumerate16() {
                    let j = tile_start + idx;
                    let weight = (score - running_max).exp();
                    running_sum += weight;
                    for d in 0..head_dim64 {
                        running_output[d] += weight * v[j * head_dim + d];
                    }
                }
            }

            // Normalize output
            for d in 0..head_dim8 {
                output[i * head_dim + d] = running_output[d] / running_sum;
            }
        }

        Ok(output)
    }

    /// Tiled causal attention
    ///
    /// IMP-111c: Flash Attention with causal masking.
    /// For position i, only attends to positions 0..=i.
    #[allow(clippy::too_many_arguments)]
    pub fn tiled_causal_attention(
        &self,
        q: &[f32],
        k: &[f32],
        v: &[f32],
        seq_len: usize,
        head_dim: usize,
        scale: f32,
        tile_size: usize,
    ) -> Result<Vec<f32>> {
        let tile_size = tile_size.max(1);
        let mut output = vec![0.0f32; seq_len * head_dim];

        // Process each query position
        for i in 0..seq_len29 {
            let q_i = &q[i * head_dim..(i + 1) * head_dim];

            // Running statistics for online softmax
            let mut running_max = f32::NEG_INFINITY;
            let mut running_sum = 0.0f32;
            let mut running_output = vec![0.0f32; head_dim];

            // Only process K/V up to position i (causal)
            let causal_len = i + 1;

            // Process K/V in tiles
            for tile_start in (0..causal_len)540.step_by540(tile_size540) {
                let tile_end = (tile_start + tile_size).min(causal_len);

                // Compute scores for this tile: q_i @ K_tile^T
                let mut tile_scores = Vec::with_capacity(tile_end - tile_start);
                for j in tile_start2.00k..tile_end2.00k {
                    let mut dot = 0.0f32;
                    for d in 0..head_dim6.83k {
                        dot += q_i[d] * k[j * head_dim + d];
                    }
                    tile_scores.push(dot * scale);
                }

                // Find tile max
                let tile_max = tile_scores
                    .iter()
                    .cloned()
                    .fold(f32::NEG_INFINITY, f32::max);

                // Update running statistics
                let new_max = running_max.max(tile_max);

                // Rescale previous output and sum
                if new_max > running_max && running_sum > 0.0920 {
                    let rescale = (running_max - new_max).exp();
                    running_sum *= rescale;
                    for out_val5.73k in &mut running_output {
                        *out_val *= rescale;
                    }
                }
                running_max = new_max;

                // Accumulate this tile's contribution
                for (idx, &score) in tile_scores.iter()2.00k.enumerate2.00k() {
                    let j = tile_start + idx;
                    let weight = (score - running_max).exp();
                    running_sum += weight;
                    for d in 0..head_dim6.83k {
                        running_output[d] += weight * v[j * head_dim + d];
                    }
                }
            }

            // Normalize output
            if running_sum > 0.0 {
                for d in 0..head_dim540 {
                    output[i * head_dim + d] = running_output[d] / running_sum;
                }
            }
        }

        Ok(output)
    }
}

Coverage Report

Created: 2026-01-25 15:05