aprender-serve 0.64.0

//! GGUF Part 04: IMP-112 (HybridScheduler Caching) + IMP-111 (Flash Attention) +
//!               IMP-113 (Batched GPU Kernels) + IMP-114 (Flattened GEMM)
//!
//! Extracted from gguf_monolith.rs (PMAT-802)
//!
//! ## Test Coverage
//!
//! - IMP-112a-d: HybridScheduler caching, cached vs uncached, multiple ops, attention
//! - PARITY-114: CUDA GEMM correctness verification
//! - IMP-111a-d: Online softmax, tiled attention, causal mask, tile sizes
//! - IMP-113a-d: Batched GEMM single dispatch, attention correctness, dispatch count, softmax
//! - IMP-114a-d: Flattened batched GEMM, loop matching, attention, large batches

use crate::gguf::test_helpers::create_test_model_with_config;
use crate::gguf::{GGUFConfig, OwnedQuantizedModelCached};

// =========================================================================
// IMP-112: HybridScheduler Caching
// =========================================================================

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_imp_112a_cached_scheduler_initialization() {
    // IMP-112a: Verify cached scheduler initializes lazily and is reused
    // This tests that OwnedQuantizedModelCached provides scheduler caching
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        explicit_head_dim: None,
        query_pre_attn_scalar: None,
        bos_token_id: None,
        eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    // Create cached wrapper
    let cached_model = OwnedQuantizedModelCached::new(model);

    // First call should initialize scheduler
    let tokens = vec![1u32, 5, 10];
    let result1 = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("IMP-112a: First cached forward should succeed");

    // Verify output shape
    assert_eq!(
        result1.len(),
        tokens.len() * config.vocab_size,
        "IMP-112a: Should return correct output shape"
    );

    // Second call should reuse scheduler (much faster)
    let result2 = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("IMP-112a: Second cached forward should succeed");

    // Results should be identical (same scheduler, same computation)
    assert_eq!(result1.len(), result2.len());
    for i in 0..result1.len() {
        let diff = (result1[i] - result2[i]).abs();
        assert!(
            diff < 1e-6,
            "IMP-112a: Results should be identical on repeated calls, pos {}: diff={}",
            i,
            diff
        );
    }
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_imp_112b_cached_matches_uncached() {
    // IMP-112b: Verify cached scheduler produces identical results to uncached
    //
    // PARITY-114 RESOLVED: CUDA GEMM grid launch dimensions were swapped (M<->N).
    // The fix swaps Grid X and Grid Y in all GEMM launch configurations:
    // - Grid X = (n + 31) / 32 for columns (N dimension)
    // - Grid Y = (m + 31) / 32 for rows (M dimension)
    //
    // Both cached (CUDA) and uncached (wgpu) paths now produce matching results.
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        explicit_head_dim: None,
        query_pre_attn_scalar: None,
        bos_token_id: None,
        eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model.clone());

    let tokens = vec![1u32, 5, 10, 20];

    // Uncached forward (creates new scheduler each time)
    let uncached_result = model
        .forward_batch_gpu(&tokens)
        .expect("Uncached forward should succeed");

    // Cached forward (reuses scheduler)
    let cached_result = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("Cached forward should succeed");

    // Results should match
    assert_eq!(uncached_result.len(), cached_result.len());
    for i in 0..uncached_result.len() {
        let diff = (uncached_result[i] - cached_result[i]).abs();
        assert!(
            diff < 1e-4,
            "IMP-112b: Cached should match uncached, pos {}: uncached={}, cached={}, diff={}",
            i,
            uncached_result[i],
            cached_result[i],
            diff
        );
    }
}

/// PARITY-114 RESOLVED: CUDA path correctness verification
///
/// This test verifies the CUDA GEMM fix is correct by checking that
/// CUDA produces values in a reasonable range (not ~10x smaller as before).
///
/// Root cause was swapped grid dimensions in CUDA launch config:
/// - Grid X was (m+31)/32 but should be (n+31)/32 for columns
/// - Grid Y was (n+31)/32 but should be (m+31)/32 for rows
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_114_cuda_gemm_correctness() {
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        explicit_head_dim: None,
        query_pre_attn_scalar: None,
        bos_token_id: None,
        eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    let tokens = vec![1u32, 5, 10, 20];

    // Cached forward uses CUDA when available
    let result = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("PARITY-114: CUDA forward should succeed");

    // CUDA produces finite values
    assert_eq!(result.len(), tokens.len() * config.vocab_size);
    assert!(
        result.iter().all(|x| x.is_finite()),
        "PARITY-114: CUDA should produce finite values"
    );

    // PARITY-114 RESOLVED: Values should no longer be ~10x smaller
    // The fix swaps Grid X and Grid Y in CUDA launch configurations
    let max_val = result.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let min_val = result.iter().cloned().fold(f32::INFINITY, f32::min);
    eprintln!(
        "PARITY-114 RESOLVED: CUDA output range: [{:.4}, {:.4}]",
        min_val, max_val
    );
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_imp_112c_multiple_operations_same_scheduler() {
    // IMP-112c: Verify multiple different operations share the same scheduler
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        explicit_head_dim: None,
        query_pre_attn_scalar: None,
        bos_token_id: None,
        eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    // Multiple forward passes with different inputs
    let tokens1 = vec![1u32, 2, 3];
    let tokens2 = vec![10u32, 20, 30, 40];
    let tokens3 = vec![5u32];

    let result1 = cached_model
        .forward_batch_gpu_cached(&tokens1)
        .expect("IMP-112c: Forward 1 should succeed");
    let result2 = cached_model
        .forward_batch_gpu_cached(&tokens2)
        .expect("IMP-112c: Forward 2 should succeed");
    let result3 = cached_model
        .forward_batch_gpu_cached(&tokens3)
        .expect("IMP-112c: Forward 3 should succeed");

    // Verify shapes
    assert_eq!(result1.len(), 3 * config.vocab_size);
    assert_eq!(result2.len(), 4 * config.vocab_size);
    assert_eq!(result3.len(), config.vocab_size);

    // All results should be finite
    assert!(result1.iter().all(|x| x.is_finite()));
    assert!(result2.iter().all(|x| x.is_finite()));
    assert!(result3.iter().all(|x| x.is_finite()));
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_imp_112d_cached_attention_matches_uncached() {
    // IMP-112d: Verify cached parallel attention matches uncached
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        explicit_head_dim: None,
        query_pre_attn_scalar: None,
        bos_token_id: None,
        eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model.clone());

    let seq_len = 8;
    let hidden_dim = config.hidden_dim;

    // Create Q, K, V tensors
    let q: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    // Uncached attention
    let uncached_attn = model
        .parallel_multihead_attention_gpu(&q, &k, &v, seq_len)
        .expect("Uncached attention should succeed");

    // Cached attention
    let cached_attn = cached_model
        .parallel_multihead_attention_gpu_cached(&q, &k, &v, seq_len)
        .expect("Cached attention should succeed");

    // Results should match
    assert_eq!(uncached_attn.len(), cached_attn.len());
    for i in 0..uncached_attn.len() {
        let diff = (uncached_attn[i] - cached_attn[i]).abs();
        assert!(
            diff < 1e-4,
            "IMP-112d: Cached attention should match uncached, pos {}: diff={}",
            i,
            diff
        );
    }
}

// =========================================================================
// IMP-111: Flash Attention-style Tiled Computation
// =========================================================================

#[test]
fn test_imp_111a_online_softmax_correctness() {
    // IMP-111a: Verify online softmax matches standard softmax
    // Online softmax processes data in tiles, tracking running max and sum
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        explicit_head_dim: None,
        query_pre_attn_scalar: None,
        bos_token_id: None,
        eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    // Test data: attention scores for one row
    let scores: Vec<f32> = (0..16).map(|i| ((i % 7) as f32 - 3.0) * 0.5).collect();

    // Standard softmax (reference)
    let standard = model.standard_softmax(&scores);

    // Online softmax (tiled, O(1) memory per tile)
    let tile_size = 4;
    let online = model
        .online_softmax(&scores, tile_size)
        .expect("IMP-111a: Online softmax should succeed");

    // Results should match within numerical tolerance
    assert_eq!(standard.len(), online.len());
    for i in 0..standard.len() {
        let diff = (standard[i] - online[i]).abs();
        assert!(
            diff < 1e-5,
            "IMP-111a: Online softmax differs at {}: standard={}, online={}, diff={}",
            i,
            standard[i],
            online[i],
            diff
        );
    }

    // Verify both sum to 1
    let std_sum: f32 = standard.iter().sum();
    let online_sum: f32 = online.iter().sum();
    assert!(
        (std_sum - 1.0).abs() < 1e-5,
        "Standard softmax should sum to 1"
    );
    assert!(
        (online_sum - 1.0).abs() < 1e-5,
        "Online softmax should sum to 1"
    );
}

#[test]
fn test_imp_111b_tiled_attention_matches_standard() {
    // IMP-111b: Verify tiled attention produces same output as standard
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        explicit_head_dim: None,
        query_pre_attn_scalar: None,
        bos_token_id: None,
        eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let seq_len = 8;
    let head_dim = config.hidden_dim / config.num_heads; // 8

    // Create Q, K, V for single head
    let q: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    let scale = 1.0 / (head_dim as f32).sqrt();

    // Standard attention (materializes full attention matrix)
    let standard_output = model
        .standard_single_head_attention(&q, &k, &v, seq_len, head_dim, scale)
        .expect("Standard attention should succeed");

    // Tiled attention (O(1) memory for softmax per tile)
    let tile_size = 4;
    let tiled_output = model
        .tiled_single_head_attention(&q, &k, &v, seq_len, head_dim, scale, tile_size)
        .expect("IMP-111b: Tiled attention should succeed");

    // Results should match
    assert_eq!(standard_output.len(), tiled_output.len());
    for i in 0..standard_output.len() {
        let diff = (standard_output[i] - tiled_output[i]).abs();
        assert!(
            diff < 1e-4,
            "IMP-111b: Tiled attention differs at {}: standard={}, tiled={}, diff={}",
            i,
            standard_output[i],
            tiled_output[i],
            diff
        );
    }
}

include!("imp_111c.rs");
include!("imp_114a.rs");