aprender-serve 0.64.0

// ========================================================================
// IMP Checklist: 25-Point Improvement Tests
// Per spec: performance-parity-ollama-llamacpp-gpu-inference-llms.md §4
// ========================================================================

// ------------------------------------------------------------------------
// Phase 1: Foundation (IMP-001 to IMP-005)
// ------------------------------------------------------------------------

/// IMP-001: SIMD-accelerated Q4_K dequantization via Trueno
/// Target: 4x speedup over scalar dequantization
#[test]
fn test_imp_001_q4k_simd_dequantize() {
    use crate::quantize::{dequantize_q4_k, dequantize_q4_k_simd};

    // Create test data: 4 super-blocks (576 bytes -> 1024 values)
    let mut data = vec![0u8; 144 * 4];
    // Set d=1.0, dmin=0.0 for all super-blocks
    for i in 0..4 {
        let offset = i * 144;
        data[offset..offset + 2].copy_from_slice(&0x3C00_u16.to_le_bytes()); // d=1.0
        data[offset + 2..offset + 4].copy_from_slice(&0x0000_u16.to_le_bytes());
        // dmin=0.0
    }

    // Verify correctness: SIMD matches scalar
    let scalar = dequantize_q4_k(&data).expect("test");
    let simd = dequantize_q4_k_simd(&data).expect("test");

    assert_eq!(
        scalar.len(),
        simd.len(),
        "IMP-001: SIMD output length should match scalar"
    );
    for (i, (s, p)) in scalar.iter().zip(simd.iter()).enumerate() {
        assert!(
            (s - p).abs() < 1e-4,
            "IMP-001: SIMD value {} differs: scalar={}, simd={}",
            i,
            s,
            p
        );
    }

    // Note: Performance comparison is validated in benchmarks, not unit tests.
    // The SIMD version uses rayon parallelization which has overhead for small data,
    // but provides significant speedup (4x+) for large model weights in production.
    // See benches/quantize.rs for actual performance measurements.

    // Verify both functions handle larger data correctly
    let large_data = vec![0u8; 144 * 64]; // 64 super-blocks
    let scalar_large = dequantize_q4_k(&large_data).expect("test");
    let simd_large = dequantize_q4_k_simd(&large_data).expect("test");
    assert_eq!(
        scalar_large.len(),
        simd_large.len(),
        "IMP-001: Large data SIMD output length should match scalar"
    );
}

/// IMP-002: Memory-mapped weight streaming for large models
/// Target: Load 7B models with < 8GB RAM
#[test]
fn test_imp_002_mmap_weight_streaming() {
    // Test that memory-mapped I/O is supported

    // Create a temporary file with model-like data
    let temp_dir = std::env::temp_dir();
    let temp_file = temp_dir.join("test_mmap_weights.bin");

    // Write test data (simulating model weights)
    let weight_data: Vec<f32> = (0..1024).map(|i| i as f32 * 0.001).collect();
    let bytes: Vec<u8> = weight_data.iter().flat_map(|f| f.to_le_bytes()).collect();

    std::fs::write(&temp_file, &bytes).expect("IMP-002: Should write temp file");

    // Memory-map the file
    let file = std::fs::File::open(&temp_file).expect("IMP-002: Should open file");
    // SAFETY: Memory safety ensured by bounds checking and alignment
    let mmap = unsafe { memmap2::Mmap::map(&file) };

    assert!(mmap.is_ok(), "IMP-002: Memory mapping should succeed");
    let mmap = mmap.expect("test");

    // Verify we can read the data without loading it all into heap
    assert_eq!(
        mmap.len(),
        bytes.len(),
        "IMP-002: Mmap size should match file size"
    );

    // Read first few values to verify content
    let first_value = f32::from_le_bytes([mmap[0], mmap[1], mmap[2], mmap[3]]);
    assert!(
        (first_value - 0.0).abs() < 1e-6,
        "IMP-002: First value should be 0.0"
    );

    // Cleanup
    std::fs::remove_file(&temp_file).ok();
}

/// IMP-003: Fused attention kernel (Q*K^T*V in single pass)
/// Target: 2x attention speedup
#[test]
fn test_imp_003_fused_attention() {
    use std::time::Instant;

    let head_dim = 32;
    let hidden_dim = 64;
    let seq_len = 16;

    // Create fused QKV attention
    let fused = FusedQKVAttention::new(head_dim, hidden_dim).expect("test");

    // Create separate attention for comparison (kept for future comparison tests)
    let _attention = Attention::new(head_dim).expect("test");

    let input =
        Tensor::from_vec(vec![seq_len, hidden_dim], vec![0.1; seq_len * hidden_dim]).expect("test");

    // Fused attention should work
    let fused_output = fused.forward(&input).expect("test");
    assert_eq!(
        fused_output.shape(),
        &[seq_len, hidden_dim],
        "IMP-003: Fused attention should preserve shape"
    );

    // Performance comparison
    let iterations = 50;

    // Time fused attention
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = fused.forward(&input).expect("test");
    }
    let fused_time = start.elapsed();

    // Performance target: <5s for 50 iterations (verify via cargo bench)
    if fused_time.as_millis() >= 5000 {
        eprintln!(
            "[PERF WARNING] Fused attention {} iterations took {}ms (target <5000ms)",
            iterations,
            fused_time.as_millis()
        );
    }
}

/// IMP-004: KV cache with efficient memory layout
/// Target: 3x decode throughput, >99% cache hit rate
#[test]
fn test_imp_004_kv_cache_layout() {
    use crate::inference::KVCache;

    let num_layers = 4;
    let hidden_dim = 64;
    let max_seq_len = 128;

    let mut cache = KVCache::new(num_layers, hidden_dim, max_seq_len);

    // Store values at multiple positions
    for pos in 0..32 {
        for layer in 0..num_layers {
            let k_data = vec![pos as f32 + layer as f32 * 0.1; hidden_dim];
            let v_data = vec![pos as f32 * 2.0 + layer as f32 * 0.1; hidden_dim];
            cache.store(layer, &k_data, &v_data);
        }
        cache.advance();
    }

    // Verify cache retrieval (simulating cache hit)
    for layer in 0..num_layers {
        let k = cache.get_k(layer);
        let v = cache.get_v(layer);

        assert!(
            !k.is_empty(),
            "IMP-004: K cache for layer {} should be non-empty",
            layer
        );
        assert!(
            !v.is_empty(),
            "IMP-004: V cache for layer {} should be non-empty",
            layer
        );

        // Verify data integrity
        assert_eq!(
            k.len(),
            32 * hidden_dim,
            "IMP-004: K cache should have correct size"
        );
    }

    // Test cache reset (for new sequence)
    cache.reset();
    let k_after_reset = cache.get_k(0);
    assert!(
        k_after_reset.is_empty() || k_after_reset.iter().all(|&x| x == 0.0),
        "IMP-004: Cache should be empty or zeroed after reset"
    );
}

/// IMP-005: Batch prefill for prompt processing
/// Target: 5x prefill speedup, >1000 tok/s
///
/// Wall-time SLA — flaky under CI contention. Run via `cargo test -- --ignored`.
#[test]
#[ignore = "wall-time perf SLA — flaky under CI contention; run --ignored to verify"]
fn test_imp_005_batch_prefill() {
    use std::time::Instant;

    // Create model for batch processing
    let config = ModelConfig {
        vocab_size: 1000,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        intermediate_dim: 256,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    // Test batch prefill with varying lengths
    let prompts = vec![
        vec![1, 2, 3, 4, 5],
        vec![10, 20, 30],
        vec![100, 200, 300, 400],
    ];

    let start = Instant::now();
    for prompt in &prompts {
        let output = model.forward(prompt).expect("test");
        assert!(
            output.size() > 0,
            "IMP-005: Batch prefill should produce output"
        );
    }
    let prefill_time = start.elapsed();

    // Calculate throughput
    let total_tokens: usize = prompts.iter().map(std::vec::Vec::len).sum();
    let throughput = total_tokens as f64 / prefill_time.as_secs_f64();

    // Prefill should be efficient (>10 tok/s minimum for test)
    assert!(
        throughput > 10.0,
        "IMP-005: Prefill throughput {:.1} tok/s should be >10",
        throughput
    );
}

// ------------------------------------------------------------------------
// Phase 2: GPU Backend (IMP-006 to IMP-010) - Stubbed for CPU-only tests
// ------------------------------------------------------------------------

/// IMP-006: Trueno WGPU backend integration
/// Target: GPU-accelerated matmul with >1.0 TFLOPS
#[test]
fn test_imp_006_wgpu_matmul() {
    // Test that GPU compute infrastructure exists
    // Actual GPU tests require --features gpu
    let linear = Linear::new(64, 128).expect("test");
    let input = Tensor::from_vec(vec![4, 64], vec![0.1; 4 * 64]).expect("test");

    let output = linear.forward(&input).expect("test");
    assert_eq!(
        output.shape(),
        &[4, 128],
        "IMP-006: Matrix multiply should work"
    );
}

/// IMP-007: GPU memory management with buffer pooling
/// Target: Zero allocation during inference
#[test]
fn test_imp_007_gpu_buffer_pool() {
    // Test that repeated operations don't cause excessive allocations
    let layer_norm = LayerNorm::new(64, 1e-5).expect("test");
    let input = Tensor::from_vec(vec![8, 64], vec![0.1; 8 * 64]).expect("test");

    // Run multiple times to test allocation behavior
    for i in 0..100 {
        let output = layer_norm.forward(&input).expect("test");
        assert_eq!(
            output.size(),
            input.size(),
            "IMP-007: Iteration {} should produce correct output",
            i
        );
    }
}

/// IMP-008: Asynchronous GPU kernel dispatch
/// Target: Hide kernel launch latency, >80% GPU utilization
#[test]
fn test_imp_008_async_dispatch() {
    use std::time::Instant;

    // Test that operations can be pipelined
    let linear1 = Linear::new(64, 64).expect("test");
    let linear2 = Linear::new(64, 64).expect("test");
    let input = Tensor::from_vec(vec![4, 64], vec![0.1; 4 * 64]).expect("test");

    let start = Instant::now();
    for _ in 0..50 {
        let mid = linear1.forward(&input).expect("test");
        let _ = linear2.forward(&mid).expect("test");
    }
    let elapsed = start.elapsed();

    // Should complete efficiently
    assert!(
        elapsed.as_millis() < 2000,
        "IMP-008: Pipelined ops should complete efficiently"
    );
}

/// IMP-009: WGPU compute shaders for transformer layers
/// Target: Full transformer on GPU with <5ms layer latency
#[test]
fn test_imp_009_transformer_gpu() {
    use std::time::Instant;

    let hidden_dim = 64;
    let intermediate_dim = 256;

    let block = TransformerBlock::new(hidden_dim, 4, intermediate_dim, 1e-5).expect("test");
    let input = Tensor::from_vec(vec![8, hidden_dim], vec![0.1; 8 * hidden_dim]).expect("test");

    let start = Instant::now();
    for _ in 0..10 {
        let _ = block.forward(&input).expect("test");
    }
    let elapsed = start.elapsed();

    let avg_latency_ms = elapsed.as_millis() as f64 / 10.0;
    assert!(
        avg_latency_ms < 500.0,
        "IMP-009: Transformer block latency {:.1}ms should be reasonable",
        avg_latency_ms
    );
}

/// IMP-010: GPU-CPU overlap for streaming generation
/// Target: Continuous token output with <10% jitter
#[test]
fn test_imp_010_streaming_overlap() {
    use std::time::Instant;

    let embedding = Embedding::new(100, 64).expect("test");
    let linear = Linear::new(64, 100).expect("test");

    let mut latencies = Vec::new();

    for token_id in 0..20 {
        let start = Instant::now();

        let embedded = embedding.forward(&[token_id]).expect("test");
        let _ = linear.forward(&embedded).expect("test");

        latencies.push(start.elapsed().as_micros() as f64);
    }

    // Calculate coefficient of variation (CV)
    let mean: f64 = latencies.iter().sum::<f64>() / latencies.len() as f64;
    let variance: f64 =
        latencies.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / latencies.len() as f64;
    let std_dev = variance.sqrt();
    let cv = std_dev / mean;

    // CV should be less than 5.0 (500% jitter - very loose bound for coverage env)
    // Coverage instrumentation adds significant and variable overhead
    assert!(
        cv < 5.0,
        "IMP-010: Token latency CV {:.2} should be <5.0",
        cv
    );
}

// ------------------------------------------------------------------------
// Phase 3: Quantization (IMP-011 to IMP-015)
// ------------------------------------------------------------------------

/// IMP-011: Fused Q4_K_M dequant+matmul kernel
/// Target: No intermediate F32 tensor
#[test]
fn test_imp_011_fused_q4k_matmul() {
    use crate::quantize::dequantize_q4_k;

    // Create quantized weights
    let q4k_data = vec![0u8; 144]; // 1 super-block = 256 values

    // Dequantize
    let weights = dequantize_q4_k(&q4k_data).expect("test");
    assert_eq!(
        weights.len(),
        256,
        "IMP-011: Should dequantize to 256 values"
    );

    // Simulate matmul with dequantized weights
    let input = vec![0.1f32; 256];
    let dot: f32 = weights.iter().zip(input.iter()).map(|(w, i)| w * i).sum();

    assert!(
        dot.is_finite(),
        "IMP-011: Fused Q4K matmul should produce finite result"
    );
}

/// IMP-012: Q5_K and Q6_K support
/// Target: Quality/speed tradeoff options
#[test]
fn test_imp_012_q5k_q6k_dequant() {
    use crate::quantize::{dequantize_q5_k, dequantize_q6_k};

    // Q5_K: 176 bytes per super-block
    let q5k_data = vec![0u8; 176];
    let q5k_result = dequantize_q5_k(&q5k_data);
    assert!(
        q5k_result.is_ok(),
        "IMP-012: Q5_K dequantization should work"
    );
    assert_eq!(
        q5k_result.expect("test").len(),
        256,
        "IMP-012: Q5_K should produce 256 values"
    );

    // Q6_K: 210 bytes per super-block
    let q6k_data = vec![0u8; 210];
    let q6k_result = dequantize_q6_k(&q6k_data);
    assert!(
        q6k_result.is_ok(),
        "IMP-012: Q6_K dequantization should work"
    );
    assert_eq!(
        q6k_result.expect("test").len(),
        256,
        "IMP-012: Q6_K should produce 256 values"
    );
}