aprender-serve 0.64.0

/// Test PARITY-014d: Memory-performance tradeoff
///
/// Analyzes the tradeoff between dequantizing weights and GPU GEMM speedup.
#[test]
fn test_parity014d_memory_performance_tradeoff() {
    /// Memory-performance tradeoff analysis
    struct MemoryTradeoff {
        model_name: &'static str,
        quantized_size_mb: f64,
        dequantized_size_mb: f64,
        gpu_speedup: f64,
        memory_overhead: f64,
    }

    impl MemoryTradeoff {
        fn phi2() -> Self {
            // phi-2: 2.7B params, Q4_K_M ≈ 1.7GB
            Self {
                model_name: "phi-2 (2.7B)",
                quantized_size_mb: 1700.0,
                dequantized_size_mb: 1700.0 * 4.0, // 4x for f32
                gpu_speedup: 10.0,
                memory_overhead: 4.0,
            }
        }

        fn llama7b() -> Self {
            // LLaMA 7B: Q4_K_M ≈ 4GB
            Self {
                model_name: "LLaMA-7B",
                quantized_size_mb: 4000.0,
                dequantized_size_mb: 4000.0 * 4.0,
                gpu_speedup: 10.0,
                memory_overhead: 4.0,
            }
        }

        fn is_memory_acceptable(&self, gpu_vram_mb: f64) -> bool {
            self.dequantized_size_mb <= gpu_vram_mb * 0.8 // 80% of VRAM
        }

        fn speedup_per_memory(&self) -> f64 {
            self.gpu_speedup / self.memory_overhead
        }
    }

    let tradeoffs = vec![MemoryTradeoff::phi2(), MemoryTradeoff::llama7b()];

    println!("\nPARITY-014d: Memory-Performance Tradeoff Analysis");
    for t in &tradeoffs {
        println!("  {}:", t.model_name);
        println!("    Quantized: {:.0} MB", t.quantized_size_mb);
        println!("    Dequantized: {:.0} MB", t.dequantized_size_mb);
        println!("    GPU speedup: {:.0}x", t.gpu_speedup);
        println!("    Memory overhead: {:.0}x", t.memory_overhead);
        println!("    Speedup per memory: {:.1}", t.speedup_per_memory());
        println!("    Fits 8GB GPU: {}", t.is_memory_acceptable(8000.0));
        println!("    Fits 24GB GPU: {}", t.is_memory_acceptable(24000.0));
    }

    // Verify tradeoff analysis
    let phi2 = &tradeoffs[0];
    assert!(
        phi2.is_memory_acceptable(24000.0),
        "PARITY-014d: phi-2 dequantized should fit 24GB GPU"
    );
    assert!(
        phi2.speedup_per_memory() > 2.0,
        "PARITY-014d: GPU speedup should exceed memory cost"
    );

    println!("  Status: VERIFIED - Memory tradeoff analyzed");
}

/// Test PARITY-014e: End-to-end batch inference benchmark design
///
/// Designs the benchmark for measuring actual batch inference performance.
#[test]
fn test_parity014e_batch_benchmark_design() {
    /// Benchmark configuration
    struct BatchBenchmarkConfig {
        batch_sizes: Vec<usize>,
        prompt_lengths: Vec<usize>,
        generation_length: usize,
        num_iterations: usize,
    }

    /// Expected benchmark results
    struct BenchmarkExpectation {
        batch_size: usize,
        expected_tps_min: f64,
        expected_tps_max: f64,
        gap_to_ollama: f64,
    }

    impl BatchBenchmarkConfig {
        fn standard() -> Self {
            Self {
                batch_sizes: vec![1, 4, 8, 16, 32, 64],
                prompt_lengths: vec![8, 32, 128],
                generation_length: 32,
                num_iterations: 5,
            }
        }
    }

    let config = BatchBenchmarkConfig::standard();
    let expectations = vec![
        BenchmarkExpectation {
            batch_size: 1,
            expected_tps_min: 4.0,
            expected_tps_max: 6.0,
            gap_to_ollama: 40.0,
        },
        BenchmarkExpectation {
            batch_size: 8,
            expected_tps_min: 5.0,
            expected_tps_max: 8.0,
            gap_to_ollama: 30.0,
        },
        BenchmarkExpectation {
            batch_size: 32,
            expected_tps_min: 8.0,
            expected_tps_max: 15.0,
            gap_to_ollama: 15.0,
        },
        BenchmarkExpectation {
            batch_size: 64,
            expected_tps_min: 10.0,
            expected_tps_max: 20.0,
            gap_to_ollama: 12.0,
        },
    ];

    println!("\nPARITY-014e: Batch Benchmark Design");
    println!("  Configuration:");
    println!("    Batch sizes: {:?}", config.batch_sizes);
    println!("    Prompt lengths: {:?}", config.prompt_lengths);
    println!("    Generation length: {}", config.generation_length);
    println!("    Iterations: {}", config.num_iterations);

    println!("\n  Expected Performance:");
    for exp in &expectations {
        println!(
            "    batch={}: {:.0}-{:.0} tok/s, gap={:.0}x",
            exp.batch_size, exp.expected_tps_min, exp.expected_tps_max, exp.gap_to_ollama
        );
    }

    // Verify expectations are reasonable
    for exp in &expectations {
        assert!(
            exp.expected_tps_max > exp.expected_tps_min,
            "PARITY-014e: Max TPS should exceed min"
        );
        assert!(
            exp.gap_to_ollama > 1.0,
            "PARITY-014e: Gap to Ollama should be >1x"
        );
    }

    println!("\n  Next steps for actual benchmark:");
    println!("  1. Run: cargo run --release --example batch_inference_benchmark");
    println!("  2. Compare against Ollama batch inference");
    println!("  3. Profile hotspots for further optimization");
    println!("  Status: VERIFIED - Benchmark design complete");
}

// ========================================================================

// PARITY-015: Actual GPU Batch Forward Implementation
// ========================================================================
//
// Spec ref: docs/specifications/performance-parity-ollama-llamacpp-gpu-inference-llms.md
// Focus: Implement actual GPU-accelerated batch forward pass
//
// Key implementation:
// 1. Batch hidden states: [batch_size, hidden_dim]
// 2. Use GPU matmul via HybridScheduler
// 3. For quantized weights: dequantize once, cache, use GPU GEMM
//
// Tests:
// - PARITY-015a: Verify GPU matmul works with batched input
// - PARITY-015b: Dequantized weight caching strategy
// - PARITY-015c: Batched layer norm implementation
// - PARITY-015d: End-to-end batch forward timing
// - PARITY-015e: Integration verification

/// Test PARITY-015a: GPU matmul with batched input
///
/// Verifies that HybridScheduler correctly handles batched matmul.
#[test]
fn test_parity015a_gpu_batch_matmul_actual() {
    use crate::gpu::HybridScheduler;

    // SCALED DOWN from phi-2 FFN dimensions (batch 32, hidden 2560, intermediate
    // 10240). This test asserts a SHAPE property of the batched matmul, and a shape
    // property does not need production-sized tensors. At phi-2 size the matmul is
    // 32*2560*10240 = 839M MACs; `HybridScheduler::new()` succeeds without a GPU
    // (`GpuCompute::auto()` falls back to CPU), so on a CPU-only runner this ran
    // through `cpu_matmul` and 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): 8*128*512 = 524_288, so the
    // GPU-vs-CPU dispatch decision this test exercises (and prints) is UNCHANGED
    // on a GPU-equipped host.
    //
    // The 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 batched input: [batch_size, hidden_dim]
    let input: Vec<f32> = (0..batch_size * hidden_dim)
        .map(|i| (i as f32 * 0.001).sin())
        .collect();

    // Create weight matrix: [hidden_dim, intermediate_dim]
    let weight: Vec<f32> = (0..hidden_dim * intermediate_dim)
        .map(|i| (i as f32 * 0.0001).cos() * 0.01)
        .collect();

    // Test with HybridScheduler
    if let Ok(mut scheduler) = HybridScheduler::new() {
        let should_gpu = scheduler.should_use_gpu(batch_size, hidden_dim, intermediate_dim);
        println!("\nPARITY-015a: GPU Batch Matmul Actual Test");
        println!("  Input: [{}x{}]", batch_size, hidden_dim);
        println!("  Weight: [{}x{}]", hidden_dim, intermediate_dim);
        println!("  Output: [{}x{}]", batch_size, intermediate_dim);
        println!("  Should use GPU: {}", should_gpu);
        println!("  GPU available: {}", scheduler.has_gpu());

        // Perform actual matmul
        let result = scheduler.matmul(&input, &weight, batch_size, hidden_dim, intermediate_dim);

        match result {
            Ok(output) => {
                assert_eq!(
                    output.len(),
                    batch_size * intermediate_dim,
                    "PARITY-015a: Output should be [batch_size, intermediate_dim]"
                );

                println!("  Status: VERIFIED - GPU batch matmul works");
            },
            Err(e) => {
                // This arm used to swallow the error and PASS, which made the
                // assertion above unreachable for any defect surfacing as `Err`.
                // The "expected if no GPU" it used to claim was also false: with no
                // GPU, `HybridScheduler::matmul` takes the `cpu_matmul` path, which
                // returns `Ok` unconditionally. The input here is small and valid,
                // so an `Err` is a real defect, not a capability gap.
                panic!("PARITY-015a: matmul failed on valid input: {}", e);
            },
        }
    } else {
        println!("\nPARITY-015a: HybridScheduler not available");
    }
}

/// Test PARITY-015b: Dequantized weight caching strategy
///
/// Verifies strategy for caching dequantized weights for GPU GEMM.
#[test]
fn test_parity015b_dequant_cache_strategy() {
    use crate::quantize::dequantize_q4_k;

    /// Weight cache entry
    struct DequantizedWeight {
        data: Vec<f32>,
        in_dim: usize,
        out_dim: usize,
        memory_bytes: usize,
    }

    impl DequantizedWeight {
        fn new(quantized: &[u8], in_dim: usize, out_dim: usize) -> Option<Self> {
            let data = dequantize_q4_k(quantized).ok()?;
            let expected_elements = in_dim * out_dim;
            if data.len() >= expected_elements {
                Some(Self {
                    data: data[..expected_elements].to_vec(),
                    in_dim,
                    out_dim,
                    memory_bytes: expected_elements * 4,
                })
            } else {
                None
            }
        }

        fn as_slice(&self) -> &[f32] {
            &self.data
        }
    }

    /// Layer weight cache
    struct LayerWeightCache {
        ffn_up: Option<DequantizedWeight>,
        ffn_down: Option<DequantizedWeight>,
        total_bytes: usize,
    }

    impl LayerWeightCache {
        fn new() -> Self {
            Self {
                ffn_up: None,
                ffn_down: None,
                total_bytes: 0,
            }
        }

        fn memory_usage_mb(&self) -> f64 {
            self.total_bytes as f64 / 1_000_000.0
        }
    }

    // Simulate phi-2 layer cache (FFN weights only)
    let hidden_dim = 2560;
    let intermediate_dim = 10240;
    let num_layers = 32;

    let per_layer_bytes = (hidden_dim * intermediate_dim + intermediate_dim * hidden_dim) * 4;
    let total_bytes = per_layer_bytes * num_layers;

    println!("\nPARITY-015b: Dequantized Weight Caching Strategy");
    println!("  Model: phi-2 (32 layers)");
    println!("  FFN up: [{}x{}]", hidden_dim, intermediate_dim);
    println!("  FFN down: [{}x{}]", intermediate_dim, hidden_dim);
    println!(
        "  Per layer: {:.1} MB",
        per_layer_bytes as f64 / 1_000_000.0
    );
    println!("  Total cache: {:.1} MB", total_bytes as f64 / 1_000_000.0);
    println!("  Strategy: Cache on first batch inference call");

    // Verify cache sizing (8GB limit - fits on 24GB GPU with model)
    assert!(
        total_bytes < 8_000_000_000_usize,
        "PARITY-015b: Cache should fit in reasonable memory (8GB limit)"
    );

    // Cache efficiency analysis
    let quantized_bytes = total_bytes / 4; // Q4 is ~4x smaller
    let overhead = total_bytes as f64 / quantized_bytes as f64;
    println!(
        "  Quantized size: {:.1} MB",
        quantized_bytes as f64 / 1_000_000.0
    );
    println!("  Memory overhead: {:.1}x", overhead);

    println!("  Status: VERIFIED - Caching strategy defined");
}

/// Test PARITY-015c: Batched layer norm implementation
///
/// Verifies batched layer norm for GPU-accelerated forward pass.
#[test]
fn test_parity015c_batched_layer_norm() {
    /// Batched layer normalization
    fn batch_layer_norm(
        input: &[f32],        // [batch_size, hidden_dim] flattened
        weight: &[f32],       // [hidden_dim]
        bias: Option<&[f32]>, // [hidden_dim]
        batch_size: usize,
        hidden_dim: usize,
        eps: f32,
    ) -> Vec<f32> {
        let mut output = vec![0.0f32; batch_size * hidden_dim];

        for b in 0..batch_size {
            let start = b * hidden_dim;
            let end = start + hidden_dim;
            let x = &input[start..end];

            // Compute mean
            let mean: f32 = x.iter().sum::<f32>() / hidden_dim as f32;

            // Compute variance
            let var: f32 = x.iter().map(|&v| (v - mean).powi(2)).sum::<f32>() / hidden_dim as f32;

            let std = (var + eps).sqrt();

            // Normalize and scale
            for i in 0..hidden_dim {
                let normalized = (x[i] - mean) / std;
                output[start + i] = normalized * weight[i] + bias.map_or(0.0, |b| b[i]);
            }
        }

        output
    }

    // Test batched layer norm
    let batch_size = 4;
    let hidden_dim = 8;
    let eps = 1e-5;

    let input: Vec<f32> = (0..batch_size * hidden_dim)
        .map(|i| (i as f32 * 0.1).sin())
        .collect();
    let weight: Vec<f32> = vec![1.0; hidden_dim];
    let bias: Vec<f32> = vec![0.0; hidden_dim];

    let output = batch_layer_norm(&input, &weight, Some(&bias), batch_size, hidden_dim, eps);

    println!("\nPARITY-015c: Batched Layer Norm");
    println!("  Batch size: {}", batch_size);
    println!("  Hidden dim: {}", hidden_dim);
    println!("  Input: {:?}...", &input[..8.min(input.len())]);
    println!("  Output: {:?}...", &output[..8.min(output.len())]);

    // Verify output is normalized (mean ≈ 0, variance ≈ 1 for each batch)
    for b in 0..batch_size {
        let start = b * hidden_dim;
        let end = start + hidden_dim;
        let batch_out = &output[start..end];

        let mean: f32 = batch_out.iter().sum::<f32>() / hidden_dim as f32;
        let var: f32 =
            batch_out.iter().map(|&v| (v - mean).powi(2)).sum::<f32>() / hidden_dim as f32;

        assert!(
            mean.abs() < 0.1,
            "PARITY-015c: Batch {} mean should be ~0, got {}",
            b,
            mean
        );
        assert!(
            (var - 1.0).abs() < 0.2,
            "PARITY-015c: Batch {} variance should be ~1, got {}",
            b,
            var
        );
    }

    println!("  Status: VERIFIED - Batched layer norm correct");
}