aprender-serve 0.64.0

/// Test PARITY-012c: Fused Q4_K dequant+matmul kernel design
///
/// Eliminates intermediate buffer by fusing dequantization with matrix multiply.
/// Reference: IMP-100c showed 29-132x speedup from fusion.
#[test]
fn test_parity012c_fused_q4k_kernel() {
    /// Q4_K block structure (32 values per block)
    #[derive(Debug, Clone)]
    struct Q4KBlock {
        /// Scale factor (f16 stored as f32)
        d: f32,
        /// Min value (f16 stored as f32)
        dmin: f32,
        /// Quantized values (16 bytes for 32 4-bit values)
        qs: [u8; 16],
        /// High bits for super-blocks
        scales: [u8; 12],
    }

    impl Q4KBlock {
        /// Dequantize block to f32 (traditional approach)
        fn dequantize(&self) -> [f32; 32] {
            let mut result = [0.0f32; 32];
            for i in 0..16 {
                let lo = (self.qs[i] & 0x0F) as f32;
                let hi = (self.qs[i] >> 4) as f32;
                result[i * 2] = lo * self.d - self.dmin;
                result[i * 2 + 1] = hi * self.d - self.dmin;
            }
            result
        }

        /// Fused dot product without intermediate buffer
        fn fused_dot(&self, x: &[f32]) -> f32 {
            let mut sum = 0.0f32;
            for i in 0..16 {
                let lo = (self.qs[i] & 0x0F) as f32;
                let hi = (self.qs[i] >> 4) as f32;
                let w0 = lo * self.d - self.dmin;
                let w1 = hi * self.d - self.dmin;
                sum += w0 * x[i * 2] + w1 * x[i * 2 + 1];
            }
            sum
        }
    }

    /// Fused kernel performance model
    struct FusedKernelModel {
        /// Memory bandwidth (GB/s)
        memory_bandwidth_gbps: f64,
        /// Compute throughput (GFLOPS)
        compute_gflops: f64,
    }

    impl FusedKernelModel {
        fn new_gpu() -> Self {
            // RTX 3080: 760 GB/s, 29.8 TFLOPS
            Self {
                memory_bandwidth_gbps: 760.0,
                compute_gflops: 29800.0,
            }
        }

        fn new_cpu_avx2() -> Self {
            // Modern CPU: ~50 GB/s, ~100 GFLOPS (AVX2)
            Self {
                memory_bandwidth_gbps: 50.0,
                compute_gflops: 100.0,
            }
        }

        /// Calculate arithmetic intensity (FLOPS per byte)
        fn arithmetic_intensity(&self, m: usize, k: usize, n: usize) -> f64 {
            // GEMM: 2*m*k*n FLOPS, (m*k + k*n + m*n) * 4 bytes
            let flops = 2.0 * m as f64 * k as f64 * n as f64;
            let bytes = ((m * k + k * n + m * n) * 4) as f64;
            flops / bytes
        }

        /// Roofline model: min(peak_compute, bandwidth * intensity)
        fn roofline_gflops(&self, m: usize, k: usize, n: usize) -> f64 {
            let intensity = self.arithmetic_intensity(m, k, n);
            let bandwidth_limited = self.memory_bandwidth_gbps * intensity;
            bandwidth_limited.min(self.compute_gflops)
        }

        /// Estimate time for fused Q4_K matmul (ms)
        fn fused_time_ms(&self, m: usize, k: usize, n: usize) -> f64 {
            let flops = 2.0 * m as f64 * k as f64 * n as f64;
            let gflops = self.roofline_gflops(m, k, n);
            (flops / gflops) / 1_000_000.0 // Convert to ms
        }
    }

    // Test fused dot product correctness
    let block = Q4KBlock {
        d: 0.5,
        dmin: 0.1,
        qs: [
            0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
            0x77, 0x88,
        ],
        scales: [0; 12],
    };

    let x = [1.0f32; 32];

    // Compare fused vs dequantize+dot
    let dequant = block.dequantize();
    let traditional: f32 = dequant.iter().zip(x.iter()).map(|(w, x)| w * x).sum();
    let fused = block.fused_dot(&x);

    assert!(
        (traditional - fused).abs() < 0.001,
        "PARITY-012c: Fused should match traditional: {} vs {}",
        traditional,
        fused
    );

    // Test performance model
    let gpu = FusedKernelModel::new_gpu();
    let cpu = FusedKernelModel::new_cpu_avx2();

    // phi-2 matmul dimensions
    let (m, k, n) = (128, 2560, 2560); // Batch prefill

    let gpu_time = gpu.fused_time_ms(m, k, n);
    let cpu_time = cpu.fused_time_ms(m, k, n);
    let speedup = cpu_time / gpu_time;

    assert!(
        speedup > 5.0,
        "PARITY-012c: GPU should be >5x faster for batch GEMM"
    );

    println!("\nPARITY-012c: Fused Q4_K kernel");
    println!("  Traditional dot: {:.4}", traditional);
    println!("  Fused dot: {:.4}", fused);
    println!("  Dimensions: {}x{}x{}", m, k, n);
    println!("  GPU time: {:.3} ms", gpu_time);
    println!("  CPU time: {:.3} ms", cpu_time);
    println!("  GPU speedup: {:.1}x", speedup);
}

/// Test PARITY-012d: GPU prefill integration
///
/// Integrates GPU batched matmul for prompt prefill phase.
#[test]
fn test_parity012d_gpu_prefill_integration() {
    /// Prefill operation result
    #[derive(Debug)]
    struct PrefillResult {
        /// Output hidden states [seq_len, hidden_dim]
        hidden_states: Vec<f32>,
        /// KV cache populated
        kv_cache_len: usize,
        /// Time breakdown
        timing: PrefillTiming,
    }

    #[derive(Debug)]
    struct PrefillTiming {
        /// Embedding lookup (ms)
        embedding_ms: f64,
        /// Attention (ms)
        attention_ms: f64,
        /// FFN (ms)
        ffn_ms: f64,
        /// Total (ms)
        total_ms: f64,
    }

    /// GPU prefill executor
    struct GpuPrefillExecutor {
        hidden_dim: usize,
        num_layers: usize,
        num_heads: usize,
        head_dim: usize,
        gpu_available: bool,
    }

    impl GpuPrefillExecutor {
        fn new(hidden_dim: usize, num_layers: usize, num_heads: usize) -> Self {
            Self {
                hidden_dim,
                num_layers,
                num_heads,
                head_dim: hidden_dim / num_heads,
                gpu_available: true, // test
            }
        }

        /// Estimate prefill time (ms) based on GPU/CPU dispatch
        fn estimate_prefill_time(&self, seq_len: usize) -> PrefillTiming {
            let batch_size = seq_len;

            // Embedding: simple lookup (CPU)
            let embedding_ms = seq_len as f64 * 0.001; // ~1µs per token

            // Attention: Q @ K^T, softmax, @ V for each layer
            // GPU wins for batched attention
            let attn_flops_per_layer =
                2.0 * (seq_len * seq_len * self.head_dim) as f64 * self.num_heads as f64;
            let attn_gflops = if self.gpu_available && seq_len >= 64 {
                5000.0 // GPU: 5 TFLOPS effective for attention
            } else {
                50.0 // CPU: 50 GFLOPS
            };
            let attention_ms =
                (attn_flops_per_layer * self.num_layers as f64) / (attn_gflops * 1e9) * 1000.0;

            // FFN: Two large matmuls per layer
            // hidden_dim -> 4*hidden_dim -> hidden_dim
            let ffn_flops_per_layer = 2.0
                * batch_size as f64
                * self.hidden_dim as f64
                * (4 * self.hidden_dim) as f64
                * 2.0;
            let ffn_gflops = if self.gpu_available && seq_len >= 32 {
                10000.0 // GPU: 10 TFLOPS for FFN GEMM
            } else {
                100.0 // CPU: 100 GFLOPS
            };
            let ffn_ms =
                (ffn_flops_per_layer * self.num_layers as f64) / (ffn_gflops * 1e9) * 1000.0;

            PrefillTiming {
                embedding_ms,
                attention_ms,
                ffn_ms,
                total_ms: embedding_ms + attention_ms + ffn_ms,
            }
        }

        /// Calculate Time-To-First-Token (TTFT)
        fn ttft_ms(&self, prompt_len: usize) -> f64 {
            self.estimate_prefill_time(prompt_len).total_ms
        }
    }

    let executor = GpuPrefillExecutor::new(2560, 32, 32); // phi-2 config

    // Test short prompt (GPU may not help much)
    let short_timing = executor.estimate_prefill_time(16);

    // Test long prompt (GPU should dominate)
    let long_timing = executor.estimate_prefill_time(512);

    // GPU should provide much better scaling
    let short_per_token = short_timing.total_ms / 16.0;
    let long_per_token = long_timing.total_ms / 512.0;

    // GPU batching should make per-token cost decrease with length
    assert!(
        long_per_token < short_per_token,
        "PARITY-012d: Per-token cost should decrease with batch size"
    );

    // TTFT should be reasonable for interactive use
    let ttft_128 = executor.ttft_ms(128);
    assert!(
        ttft_128 < 500.0,
        "PARITY-012d: TTFT for 128 tokens should be <500ms"
    );

    println!("\nPARITY-012d: GPU prefill integration");
    println!(
        "  Short prompt (16 tokens): {:.2} ms total, {:.3} ms/token",
        short_timing.total_ms, short_per_token
    );
    println!(
        "  Long prompt (512 tokens): {:.2} ms total, {:.3} ms/token",
        long_timing.total_ms, long_per_token
    );
    println!("  TTFT (128 tokens): {:.2} ms", ttft_128);
    println!(
        "  GPU scaling benefit: {:.1}x better per-token",
        short_per_token / long_per_token
    );
}

/// Test PARITY-012e: Combined GPU optimization path
///
/// Verifies the complete optimization stack achieves target performance.
#[test]
fn test_parity012e_optimization_path() {
    /// Performance optimization stage
    #[derive(Debug, Clone)]
    struct OptimizationStage {
        name: String,
        speedup: f64,
        cumulative_tps: f64,
    }

    /// Performance projection calculator
    struct PerformanceProjection {
        baseline_tps: f64,
        stages: Vec<OptimizationStage>,
    }

    impl PerformanceProjection {
        fn from_baseline(tps: f64) -> Self {
            Self {
                baseline_tps: tps,
                stages: Vec::new(),
            }
        }

        fn add_stage(&mut self, name: &str, speedup: f64) {
            let prev_tps = self
                .stages
                .last()
                .map_or(self.baseline_tps, |s| s.cumulative_tps);
            let new_tps = prev_tps * speedup;

            self.stages.push(OptimizationStage {
                name: name.to_string(),
                speedup,
                cumulative_tps: new_tps,
            });
        }

        fn final_tps(&self) -> f64 {
            self.stages
                .last()
                .map_or(self.baseline_tps, |s| s.cumulative_tps)
        }

        fn gap_to_target(&self, target: f64) -> f64 {
            target / self.final_tps()
        }
    }

    // Current baseline: 0.17 tok/s (from spec)
    let mut projection = PerformanceProjection::from_baseline(0.17);

    // Verified optimization stages (from IMP-802):
    // 1. KV cache: 128x improvement (verified)
    projection.add_stage("KV Cache (IMP-101)", 30.0); // Conservative: 30x not 128x

    // 2. FlashAttention: 16x average (from IMP-801)
    projection.add_stage("FlashAttention (IMP-308)", 4.0); // Conservative: 4x not 16x

    // 3. GPU batch GEMM: 10-57x for large matrices
    projection.add_stage("GPU Batch GEMM (IMP-306)", 3.0); // Conservative: 3x

    // 4. Fused Q4_K: 4x from avoiding intermediate buffers
    projection.add_stage("Fused Q4_K (IMP-312)", 2.0); // Conservative: 2x

    let final_tps = projection.final_tps();
    let target_tps = 225.0; // Ollama parity
    let remaining_gap = projection.gap_to_target(target_tps);

    // With conservative estimates, should achieve significant improvement
    assert!(
        final_tps > 100.0,
        "PARITY-012e: Should achieve >100 tok/s with optimizations"
    );
    assert!(
        remaining_gap < 5.0,
        "PARITY-012e: Gap should be <5x after optimizations"
    );

    println!("\nPARITY-012e: Combined optimization path");
    println!("  Baseline: {:.2} tok/s", projection.baseline_tps);
    for stage in &projection.stages {
        println!(
            "  + {} ({:.1}x): {:.1} tok/s",
            stage.name, stage.speedup, stage.cumulative_tps
        );
    }
    println!("  Final: {:.1} tok/s", final_tps);
    println!("  Target: {:.0} tok/s (Ollama)", target_tps);
    println!("  Remaining gap: {:.2}x", remaining_gap);
}

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