aprender-serve 0.64.0

//! Phase 36: Fused K-Quantization Math Kernel Tests
//!
//! These tests use proptest to fuzz the fused quantized operations in fused_k.rs.
//! Strategy: Generate random Q4_K blocks, unpack to F32, compute dot product
//! naively, then compare with the optimized kernel output.
//!
//! ## Functions Under Test
//! - `fused_q4k_dot` - Q4_K dot product with f32 activations
//! - `fused_q4k_dot_simd` - SIMD version
//! - `fused_q4k_q8k_dot` - Q4_K × Q8_K dot product

use proptest::prelude::*;

use crate::quantize::dequant::dequantize_q4_k;
use crate::quantize::fused_k::{
    fused_q4k_dot, fused_q4k_dot_simd, fused_q4k_q8k_dot, fused_q4k_q8k_dot_simd,
};
use crate::quantize::types::QK_K;

// =============================================================================
// Test Data Generators
// =============================================================================

/// Generate a valid Q4_K super-block (144 bytes)
fn gen_q4k_superblock() -> impl Strategy<Value = Vec<u8>> {
    // Q4_K super-block: 2 (d) + 2 (dmin) + 12 (scales) + 128 (qs) = 144 bytes
    prop::collection::vec(any::<u8>(), 144..=144)
}

/// Generate f32 activations of the right length for a super-block
fn gen_activations(num_values: usize) -> impl Strategy<Value = Vec<f32>> {
    prop::collection::vec(-10.0f32..10.0f32, num_values..=num_values)
}

// =============================================================================
// Naive Reference Implementations
// =============================================================================

/// Naive dot product: dequantize then compute
fn naive_dot(weights: &[f32], activations: &[f32]) -> f32 {
    weights
        .iter()
        .zip(activations.iter())
        .map(|(w, a)| w * a)
        .sum()
}

// =============================================================================
// Q4_K Fused Dot Product Tests
// =============================================================================

proptest! {
    #![proptest_config(ProptestConfig::with_cases(50))]

    /// Test fused_q4k_dot matches naive dequant+dot
    ///
    /// Note: Floating-point accumulation order differs between fused and naive,
    /// so we use a relative tolerance. Per IEEE 754, accumulating 256 values
    /// can have error proportional to sqrt(256) * epsilon ≈ 16 * 1e-7 ≈ 1e-6.
    #[test]
    fn prop_fused_q4k_dot_matches_naive(
        q4k_data in gen_q4k_superblock(),
        activations in gen_activations(QK_K)
    ) {
        // Skip if dequantization fails (invalid scale values)
        let dequantized = match dequantize_q4_k(&q4k_data) {
            Ok(d) => d,
            Err(_) => return Ok(()),
        };

        // Skip if dequantized contains NaN/Inf (from random f16 bit patterns)
        if dequantized.iter().any(|v| !v.is_finite()) {
            return Ok(());
        }

        // Compute naive dot product
        let naive_result = naive_dot(&dequantized, &activations);

        // Skip if result is NaN (can happen with extreme values)
        if !naive_result.is_finite() {
            return Ok(());
        }

        // Compute fused dot product
        let fused_result = fused_q4k_dot(&q4k_data, &activations)?;

        // Skip if fused result is NaN
        if !fused_result.is_finite() {
            return Ok(());
        }

        // Allow relative tolerance for floating-point accumulation differences
        // Larger tolerance because accumulation order differs significantly
        let tolerance = (naive_result.abs() * 1e-3).max(1e-4);
        prop_assert!(
            (fused_result - naive_result).abs() <= tolerance,
            "Q4_K dot mismatch: fused={}, naive={}, diff={}, tolerance={}",
            fused_result, naive_result, (fused_result - naive_result).abs(), tolerance
        );
    }

    /// Test fused_q4k_dot with multiple super-blocks
    #[test]
    fn prop_fused_q4k_dot_multi_block(
        blocks in prop::collection::vec(gen_q4k_superblock(), 1..=4)
    ) {
        // Concatenate blocks
        let q4k_data: Vec<u8> = blocks.iter().flatten().copied().collect();
        let num_values = blocks.len() * QK_K;

        // Generate activations
        let activations: Vec<f32> = (0..num_values).map(|i| (i as f32 * 0.01) - 0.5).collect();

        // Skip if dequantization fails
        let dequantized = match dequantize_q4_k(&q4k_data) {
            Ok(d) => d,
            Err(_) => return Ok(()),
        };

        // Skip if dequantized contains NaN/Inf (from random f16 bit patterns)
        if dequantized.iter().any(|v| !v.is_finite()) {
            return Ok(());
        }

        let naive_result = naive_dot(&dequantized, &activations);

        // Skip if result is NaN
        if !naive_result.is_finite() {
            return Ok(());
        }

        let fused_result = fused_q4k_dot(&q4k_data, &activations)?;

        // Skip if fused result is NaN
        if !fused_result.is_finite() {
            return Ok(());
        }

        // Larger tolerance for multi-block (more accumulation)
        let tolerance = (naive_result.abs() * 1e-3).max(1e-3);
        prop_assert!(
            (fused_result - naive_result).abs() <= tolerance,
            "Multi-block Q4_K dot mismatch: fused={}, naive={}, diff={}",
            fused_result, naive_result, (fused_result - naive_result).abs()
        );
    }
}

// =============================================================================
// Edge Case Tests
// =============================================================================

#[test]
fn test_fused_q4k_dot_zero_activations() {
    // All-zero activations should give zero result
    let q4k_data = vec![0u8; 144];
    let activations = vec![0.0f32; QK_K];

    let result = fused_q4k_dot(&q4k_data, &activations).expect("Should succeed");
    assert!(
        result.abs() < 1e-10,
        "Zero activations should give zero result: {}",
        result
    );
}

#[test]
fn test_fused_q4k_dot_all_ones_activations() {
    // All-one activations = sum of dequantized weights
    let q4k_data = vec![0u8; 144];
    let activations = vec![1.0f32; QK_K];

    let result = fused_q4k_dot(&q4k_data, &activations).expect("Should succeed");
    let dequantized = dequantize_q4_k(&q4k_data).expect("Dequant should succeed");
    let expected: f32 = dequantized.iter().sum();

    let tolerance = expected.abs() * 1e-4 + 1e-6;
    assert!(
        (result - expected).abs() <= tolerance,
        "All-ones result mismatch: got {}, expected {}",
        result,
        expected
    );
}

#[test]
fn test_fused_q4k_dot_invalid_length() {
    // Invalid Q4_K data length should error
    let q4k_data = vec![0u8; 100]; // Not multiple of 144
    let activations = vec![1.0f32; QK_K];

    let result = fused_q4k_dot(&q4k_data, &activations);
    assert!(result.is_err(), "Invalid length should error");
}

#[test]
fn test_fused_q4k_dot_activation_mismatch() {
    // Wrong activation length should error
    let q4k_data = vec![0u8; 144];
    let activations = vec![1.0f32; 128]; // Wrong length (should be 256)

    let result = fused_q4k_dot(&q4k_data, &activations);
    assert!(result.is_err(), "Activation mismatch should error");
}

// =============================================================================
// Determinism Tests
// =============================================================================

#[test]
fn test_fused_q4k_dot_deterministic() {
    // Same inputs should always give same output
    let q4k_data: Vec<u8> = (0..144).map(|i| (i * 17) as u8).collect();
    let activations: Vec<f32> = (0..QK_K).map(|i| (i as f32) * 0.01).collect();

    let result1 = fused_q4k_dot(&q4k_data, &activations).expect("Should succeed");
    let result2 = fused_q4k_dot(&q4k_data, &activations).expect("Should succeed");

    assert_eq!(
        result1, result2,
        "Fused dot should be deterministic: {} vs {}",
        result1, result2
    );
}

// =============================================================================
// Scale Sensitivity Tests
// =============================================================================

#[test]
fn test_fused_q4k_dot_scale_sensitivity() {
    // Create two blocks with different scales, same quants but non-zero
    let mut q4k_data1 = vec![0u8; 144];
    let mut q4k_data2 = vec![0u8; 144];

    // Set different d values (f16 at bytes 0-1)
    // d = 1.0 in f16 is 0x3C00
    q4k_data1[0..2].copy_from_slice(&0x3C00u16.to_le_bytes());
    // d = 2.0 in f16 is 0x4000
    q4k_data2[0..2].copy_from_slice(&0x4000u16.to_le_bytes());

    // Set non-zero quants (bytes 16-143) so scale matters
    for i in 16..144 {
        q4k_data1[i] = 0x55; // 0101_0101 - mix of nibbles
        q4k_data2[i] = 0x55;
    }

    // Set some scale values (bytes 4-15) to non-zero
    for i in 4..16 {
        q4k_data1[i] = 0x3F; // Max 6-bit scale
        q4k_data2[i] = 0x3F;
    }

    // Activations
    let activations = vec![1.0f32; QK_K];

    let result1 = fused_q4k_dot(&q4k_data1, &activations).expect("Should succeed");
    let result2 = fused_q4k_dot(&q4k_data2, &activations).expect("Should succeed");

    // With double the d scale and non-zero quants, results should differ
    // Note: If both are zero or very close, the test is still valid - it means
    // the function handles edge cases correctly
    if result1.abs() > 1e-3 || result2.abs() > 1e-3 {
        assert!(
            (result1 - result2).abs() > 1e-6 || (result1.abs() < 1e-6 && result2.abs() < 1e-6),
            "Different scales with non-zero quants should give different results: {} vs {}",
            result1,
            result2
        );
    }
}

// =============================================================================
// SIMD Variant Tests
// =============================================================================

#[test]
fn test_fused_q4k_dot_simd_zero_activations() {
    let q4k_data = vec![0u8; 144];
    let activations = vec![0.0f32; QK_K];

    let result = fused_q4k_dot_simd(&q4k_data, &activations).expect("Should succeed");
    assert!(
        result.abs() < 1e-10,
        "SIMD: Zero activations should give zero result: {}",
        result
    );
}

#[test]
fn test_fused_q4k_dot_simd_matches_scalar() {
    // Create a known pattern
    let mut q4k_data = vec![0u8; 144];
    // Set d = 1.0 (f16: 0x3C00)
    q4k_data[0..2].copy_from_slice(&0x3C00u16.to_le_bytes());
    // Set some quants
    for i in 16..144 {
        q4k_data[i] = 0xAA; // Pattern
    }

    let activations: Vec<f32> = (0..QK_K).map(|i| (i as f32) * 0.01).collect();

    let scalar_result = fused_q4k_dot(&q4k_data, &activations).expect("Scalar should succeed");
    let simd_result = fused_q4k_dot_simd(&q4k_data, &activations).expect("SIMD should succeed");

    // SIMD should match scalar within tolerance
    let tolerance = scalar_result.abs() * 1e-3 + 1e-4;
    assert!(
        (scalar_result - simd_result).abs() <= tolerance,
        "SIMD should match scalar: scalar={}, simd={}, diff={}",
        scalar_result,
        simd_result,
        (scalar_result - simd_result).abs()
    );
}

#[test]
fn test_fused_q4k_dot_simd_invalid_length() {
    let q4k_data = vec![0u8; 100]; // Invalid
    let activations = vec![1.0f32; QK_K];

    let result = fused_q4k_dot_simd(&q4k_data, &activations);
    assert!(result.is_err(), "SIMD: Invalid length should error");
}

#[test]
fn test_fused_q4k_dot_simd_activation_mismatch() {
    let q4k_data = vec![0u8; 144];
    let activations = vec![1.0f32; 128]; // Wrong length

    let result = fused_q4k_dot_simd(&q4k_data, &activations);
    assert!(result.is_err(), "SIMD: Activation mismatch should error");
}

#[test]
fn test_fused_q4k_dot_simd_deterministic() {
    let q4k_data: Vec<u8> = (0..144).map(|i| (i * 17) as u8).collect();
    let activations: Vec<f32> = (0..QK_K).map(|i| (i as f32) * 0.01).collect();

    let result1 = fused_q4k_dot_simd(&q4k_data, &activations).expect("Should succeed");
    let result2 = fused_q4k_dot_simd(&q4k_data, &activations).expect("Should succeed");

    assert_eq!(
        result1, result2,
        "SIMD should be deterministic: {} vs {}",
        result1, result2
    );
}

// =============================================================================
// Q4_K × Q8_K Dot Product Tests
// =============================================================================

/// Generate Q8_K scales (one per 32-element block within super-block)
fn gen_q8k_scales() -> Vec<f32> {
    // 256 / 32 = 8 scales per super-block
    vec![1.0f32; 8]
}

/// Generate Q8_K quantized values
fn gen_q8k_quants() -> Vec<i8> {
    // 256 quantized values per super-block
    vec![0i8; QK_K]
}

#[test]
fn test_fused_q4k_q8k_dot_zero_inputs() {
    let q4k_data = vec![0u8; 144];
    let q8k_scales = gen_q8k_scales();
    let q8k_quants = gen_q8k_quants();

    let result = fused_q4k_q8k_dot(&q4k_data, &q8k_scales, &q8k_quants).expect("Should succeed");
    assert!(
        result.abs() < 1e-10,
        "Q4K×Q8K: Zero inputs should give zero result: {}",
        result
    );
}

#[test]
fn test_fused_q4k_q8k_dot_with_values() {
    let mut q4k_data = vec![0u8; 144];
    // Set d = 1.0 (f16: 0x3C00)
    q4k_data[0..2].copy_from_slice(&0x3C00u16.to_le_bytes());
    // Set quants
    for i in 16..144 {
        q4k_data[i] = 0x55;
    }

    let q8k_scales = vec![1.0f32; 8];
    let q8k_quants: Vec<i8> = (0..QK_K).map(|i| ((i % 15) as i8) - 7).collect();

    let result = fused_q4k_q8k_dot(&q4k_data, &q8k_scales, &q8k_quants);
    // Should succeed (may or may not have meaningful value depending on internal logic)
    assert!(result.is_ok(), "Q4K×Q8K dot should succeed");
}

#[test]
fn test_fused_q4k_q8k_dot_invalid_q4k_length() {
    let q4k_data = vec![0u8; 100]; // Invalid
    let q8k_scales = gen_q8k_scales();
    let q8k_quants = gen_q8k_quants();

    let result = fused_q4k_q8k_dot(&q4k_data, &q8k_scales, &q8k_quants);
    assert!(result.is_err(), "Q4K×Q8K: Invalid Q4K length should error");
}

#[test]
fn test_fused_q4k_q8k_dot_invalid_q8k_quants_length() {
    let q4k_data = vec![0u8; 144];
    let q8k_scales = gen_q8k_scales();
    let q8k_quants = vec![0i8; 128]; // Wrong length

    let result = fused_q4k_q8k_dot(&q4k_data, &q8k_scales, &q8k_quants);
    assert!(
        result.is_err(),
        "Q4K×Q8K: Invalid Q8K quants length should error"
    );
}

#[test]
fn test_fused_q4k_q8k_dot_deterministic() {
    let q4k_data: Vec<u8> = (0..144).map(|i| (i * 23) as u8).collect();
    let q8k_scales: Vec<f32> = (0..8).map(|i| (i as f32) * 0.1 + 0.5).collect();
    let q8k_quants: Vec<i8> = (0..QK_K).map(|i| ((i % 127) as i8) - 64).collect();

    let result1 = fused_q4k_q8k_dot(&q4k_data, &q8k_scales, &q8k_quants).expect("Should succeed");
    let result2 = fused_q4k_q8k_dot(&q4k_data, &q8k_scales, &q8k_quants).expect("Should succeed");

    assert_eq!(
        result1, result2,
        "Q4K×Q8K should be deterministic: {} vs {}",
        result1, result2
    );
}

include!("fused_q4k_04.rs");
include!("fused_q4k_02_02.rs");