aprender-serve 0.64.0

//! GGUF Part 14: PARITY-018 - PARITY-025 (GPU Batch FFN & Request Infrastructure)
//!
//! Extracted from gguf_monolith.rs (PMAT-802)
//!
//! ## Test Groups
//!
//! - PARITY-018: Production GPU Batch FFN Integration (5 tests)
//! - PARITY-019: Production DequantizedWeightCache Integration (5 tests)
//! - PARITY-020: Batch Generation with GPU FFN (5 tests)
//! - PARITY-021: GPU Batch FFN Integration in Forward Pass (5 tests)
//! - PARITY-023: Request Batching Infrastructure (5 tests)
//! - PARITY-024: Batch Attention Tests (5 tests)
//! - PARITY-025: Batch Embedding and LM Head Tests (5 tests)

#![allow(clippy::needless_range_loop)]

#[cfg(feature = "gpu")]
use crate::gguf::{BatchGenerationStats, BatchingConfig};
use crate::gguf::{DequantizedFFNWeights, DequantizedWeightCache, QuantizedGenerateConfig};

// ========================================================================
// PARITY-018: Production GPU Batch FFN Integration
// ============================================================================
//
// Objective: Integrate GPU batch FFN into OwnedQuantizedModelCachedSync
//
// From PARITY-017:
// - gpu_batch_ffn() works: 10-13 GFLOPS
// - Integration points identified
// - Dequant cache: 200 MB/layer, 6.4 GB for phi-2
//
// Implementation:
// 1. Add DequantizedWeightCache to OwnedQuantizedModelCachedSync
// 2. Add batch_ffn_gpu() method
// 3. Add batch_generate_gpu() method
// ============================================================================

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity018a_dequantized_weight_cache_production() {
    use std::collections::HashMap;
    use std::sync::RwLock;

    // Production-ready DequantizedWeightCache
    // Uses RwLock for concurrent read access during batch inference

    struct DequantizedFFNWeightsLocal {
        up: Vec<f32>,   // [hidden, intermediate]
        down: Vec<f32>, // [intermediate, hidden]
    }

    struct DequantizedWeightCacheLocal {
        layers: RwLock<HashMap<usize, DequantizedFFNWeightsLocal>>,
        hidden_dim: usize,
        intermediate_dim: usize,
        num_layers: usize,
    }

    impl DequantizedWeightCacheLocal {
        fn new(hidden_dim: usize, intermediate_dim: usize, num_layers: usize) -> Self {
            Self {
                layers: RwLock::new(HashMap::new()),
                hidden_dim,
                intermediate_dim,
                num_layers,
            }
        }

        /// Dequantize all layers upfront (warmup phase)
        fn warmup<F>(&self, dequant_fn: F)
        where
            F: Fn(usize) -> (Vec<f32>, Vec<f32>),
        {
            let mut cache = self.layers.write().expect("test");
            for layer_idx in 0..self.num_layers {
                cache.entry(layer_idx).or_insert_with(|| {
                    let (up, down) = dequant_fn(layer_idx);
                    DequantizedFFNWeightsLocal { up, down }
                });
            }
        }

        /// Get dequantized weights (read-only, concurrent access)
        fn get(&self, layer_idx: usize) -> Option<(Vec<f32>, Vec<f32>)> {
            let cache = self.layers.read().expect("test");
            cache
                .get(&layer_idx)
                .map(|w| (w.up.clone(), w.down.clone()))
        }

        fn is_warmed_up(&self) -> bool {
            let cache = self.layers.read().expect("test");
            cache.len() == self.num_layers
        }

        fn memory_bytes(&self) -> usize {
            let cache = self.layers.read().expect("test");
            cache.len() * (self.hidden_dim * self.intermediate_dim * 2) * std::mem::size_of::<f32>()
        }
    }

    // Test with phi-2 dimensions
    let hidden_dim = 2560;
    let intermediate_dim = 10240;
    let num_layers = 32;

    let cache = DequantizedWeightCacheLocal::new(hidden_dim, intermediate_dim, num_layers);

    // Verify initial state
    assert!(
        !cache.is_warmed_up(),
        "PARITY-018a: Should not be warmed up initially"
    );
    assert_eq!(
        cache.memory_bytes(),
        0,
        "PARITY-018a: Initial memory should be 0"
    );

    // Warmup (simulate dequantization)
    cache.warmup(|_layer_idx| {
        let up = vec![0.01f32; hidden_dim * intermediate_dim];
        let down = vec![0.01f32; intermediate_dim * hidden_dim];
        (up, down)
    });

    // Verify warmed up
    assert!(
        cache.is_warmed_up(),
        "PARITY-018a: Should be warmed up after warmup()"
    );

    let expected_bytes =
        num_layers * (hidden_dim * intermediate_dim * 2) * std::mem::size_of::<f32>();
    assert_eq!(
        cache.memory_bytes(),
        expected_bytes,
        "PARITY-018a: Memory should match"
    );

    // Verify concurrent read access
    let weights = cache.get(0);
    assert!(
        weights.is_some(),
        "PARITY-018a: Should be able to get layer 0"
    );
    let (up, down) = weights.expect("test");
    assert_eq!(up.len(), hidden_dim * intermediate_dim);
    assert_eq!(down.len(), intermediate_dim * hidden_dim);

    println!("\nPARITY-018a: Production DequantizedWeightCache");
    println!("  Layers: {}", num_layers);
    println!(
        "  Memory: {:.1} GB",
        cache.memory_bytes() as f64 / (1024.0 * 1024.0 * 1024.0)
    );
    println!("  Warmed up: {}", cache.is_warmed_up());
    println!("  Status: VERIFIED - Production cache works");
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity018b_batch_ffn_gpu_method() {
    use crate::gpu::HybridScheduler;

    // Test batch_ffn_gpu as a standalone method
    // This will be integrated into OwnedQuantizedModelCachedSync

    fn batch_ffn_gpu(
        hidden_states: &[f32], // [batch, hidden]
        up_weight: &[f32],     // [hidden, intermediate]
        down_weight: &[f32],   // [intermediate, hidden]
        up_bias: Option<&[f32]>,
        down_bias: Option<&[f32]>,
        batch_size: usize,
        hidden_dim: usize,
        intermediate_dim: usize,
        scheduler: &mut HybridScheduler,
    ) -> Vec<f32> {
        // Up projection
        let mut intermediate = scheduler
            .matmul(
                hidden_states,
                up_weight,
                batch_size,
                hidden_dim,
                intermediate_dim,
            )
            .expect("Up projection failed");

        // Add up bias if present
        if let Some(bias) = up_bias {
            for b in 0..batch_size {
                for i in 0..intermediate_dim {
                    intermediate[b * intermediate_dim + i] += bias[i];
                }
            }
        }

        // GELU activation
        for x in &mut intermediate {
            let x64 = *x as f64;
            *x = (x64 * 0.5 * (1.0 + (x64 * 0.7978845608 * (1.0 + 0.044715 * x64 * x64)).tanh()))
                as f32;
        }

        // Down projection
        let mut output = scheduler
            .matmul(
                &intermediate,
                down_weight,
                batch_size,
                intermediate_dim,
                hidden_dim,
            )
            .expect("Down projection failed");

        // Add down bias if present
        if let Some(bias) = down_bias {
            for b in 0..batch_size {
                for i in 0..hidden_dim {
                    output[b * hidden_dim + i] += bias[i];
                }
            }
        }

        output
    }

    // SCALED DOWN from phi-2 dimensions (batch 32, hidden 2560, intermediate 10240).
    // This test asserts a SHAPE property of the up -> bias -> GELU -> down chain,
    // and a shape property does not need production-sized tensors. At phi-2 size
    // each projection is 32*2560*10240 = 839M MACs; `HybridScheduler::new()`
    // succeeds without a GPU (`GpuCompute::auto()` falls back to CPU), so on a
    // CPU-only runner both projections ran through `cpu_matmul` and this test
    // alone cost minutes of workspace-test wall clock — times three under
    // nextest's `retries = 2`. These dimensions do the same work 1600x smaller.
    //
    // The dimensions are chosen so m*k*n stays ABOVE `HybridScheduler`'s
    // `gpu_threshold` (64*64*64 = 262_144): up and down are both
    // 8*128*512 = 524_288, so the GPU-vs-CPU dispatch decision this test
    // exercises is UNCHANGED on a GPU-equipped host.
    //
    // The throughput/GFLOPS reporting was removed rather than rescaled: at this
    // size the number is meaningless, and a printed rate invites someone to cite
    // it. Benchmarks belong in `benches/`, not in `--lib` tests.
    let batch_size = 8;
    let hidden_dim = 128;
    let intermediate_dim = 512;

    // Create test data
    let hidden_states: Vec<f32> = (0..batch_size * hidden_dim)
        .map(|i| (i as f32 * 0.001).sin() * 0.1)
        .collect();
    let up_weight: Vec<f32> = (0..hidden_dim * intermediate_dim)
        .map(|i| (i as f32 * 0.0001).cos() * 0.01)
        .collect();
    let down_weight: Vec<f32> = (0..intermediate_dim * hidden_dim)
        .map(|i| (i as f32 * 0.0001).sin() * 0.01)
        .collect();

    println!("\nPARITY-018b: batch_ffn_gpu Method");

    if let Ok(mut scheduler) = HybridScheduler::new() {
        let output = batch_ffn_gpu(
            &hidden_states,
            &up_weight,
            &down_weight,
            None,
            None,
            batch_size,
            hidden_dim,
            intermediate_dim,
            &mut scheduler,
        );

        assert_eq!(
            output.len(),
            batch_size * hidden_dim,
            "PARITY-018b: Output should be [batch, hidden]"
        );

        println!("  Input: [{}x{}]", batch_size, hidden_dim);
        println!("  Output: [{}x{}]", batch_size, hidden_dim);
        println!("  Status: VERIFIED - batch_ffn_gpu works");
    } else {
        println!("  Status: SKIP - GPU not available");
    }
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity018c_batch_generate_gpu_flow() {
    // Test the batch_generate_gpu flow without actual model

    struct BatchRequest {
        tokens: Vec<u32>,
        position: usize,
        active: bool,
    }

    struct BatchGenerateGPU {
        gpu_threshold: usize,
        requests: Vec<BatchRequest>,
    }

    impl BatchGenerateGPU {
        fn new(prompts: &[&[u32]], gpu_threshold: usize) -> Self {
            let requests = prompts
                .iter()
                .map(|p| BatchRequest {
                    tokens: p.to_vec(),
                    position: p.len(),
                    active: true,
                })
                .collect();
            Self {
                gpu_threshold,
                requests,
            }
        }

        fn active_count(&self) -> usize {
            self.requests.iter().filter(|r| r.active).count()
        }

        fn should_use_gpu(&self) -> bool {
            self.active_count() >= self.gpu_threshold
        }

        fn step(&mut self) -> (usize, bool) {
            let active = self.active_count();
            let use_gpu = self.should_use_gpu();

            // Simulate generation step
            for req in &mut self.requests {
                if req.active {
                    req.tokens.push(0); // Dummy token
                    req.position += 1;
                    if req.position > 100 {
                        req.active = false;
                    }
                }
            }

            (active, use_gpu)
        }
    }

    // Test with 64 prompts (should use GPU)
    let prompts: Vec<Vec<u32>> = (0..64).map(|i| vec![1, 2, 3, i as u32]).collect();
    let prompt_refs: Vec<&[u32]> = prompts.iter().map(std::vec::Vec::as_slice).collect();

    let mut batch = BatchGenerateGPU::new(&prompt_refs, 32);

    println!("\nPARITY-018c: batch_generate_gpu Flow");
    println!("  Prompts: {}", prompts.len());
    println!("  GPU threshold: {}", batch.gpu_threshold);

    let mut gpu_steps = 0;
    let mut cpu_steps = 0;

    for _ in 0..10 {
        let (active, use_gpu) = batch.step();
        if use_gpu {
            gpu_steps += 1;
        } else {
            cpu_steps += 1;
        }
        println!("  Step: active={}, use_gpu={}", active, use_gpu);
    }

    assert!(
        gpu_steps > 0,
        "PARITY-018c: Should have GPU steps with 64 prompts"
    );
    println!("  GPU steps: {}, CPU steps: {}", gpu_steps, cpu_steps);
    println!("  Status: VERIFIED - Flow works correctly");
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity018d_integration_with_owned_quantized_model() {
    // Verify that OwnedQuantizedModelCachedSync has the necessary infrastructure
    // for GPU batch FFN integration

    use crate::gpu::HybridScheduler;

    println!("\nPARITY-018d: Integration with OwnedQuantizedModelCachedSync");

    // Check that HybridScheduler can be created
    if let Ok(scheduler) = HybridScheduler::new() {
        println!("  HybridScheduler: available");
        println!("  GPU available: {}", scheduler.has_gpu());
        println!("  GPU threshold: {}", scheduler.gpu_threshold());

        // The integration would add:
        // 1. dequant_cache: Option<DequantizedWeightCache> field
        // 2. batch_ffn_gpu() method
        // 3. batch_generate_gpu() method

        let integration_checklist = [
            ("OwnedQuantizedModelCachedSync struct", true),
            ("HybridScheduler caching", true),
            ("DequantizedWeightCache (to add)", false),
            ("batch_ffn_gpu method (to add)", false),
            ("batch_generate_gpu method (to add)", false),
        ];

        println!("\n  Integration Checklist:");
        for (item, done) in integration_checklist {
            let status = if done { "✓" } else { "○" };
            println!("    {} {}", status, item);
        }

        // Count completed items
        let completed = integration_checklist
            .iter()
            .filter(|(_, done)| *done)
            .count();
        let total = integration_checklist.len();

        println!(
            "\n  Progress: {}/{} ({}%)",
            completed,
            total,
            completed * 100 / total
        );
        println!("  Status: VERIFIED - Infrastructure exists, need to add GPU batch methods");
    } else {
        println!("  Status: SKIP - GPU not available");
    }
}

include!("parity018e_performance_target.rs");
include!("parity020a_batch.rs");
include!("parity021c_gpu.rs");
include!("parity024c_gpu.rs");