aprender-serve 0.64.0

//! GGUF Part 16: PARITY-035 (Chunked Prefill for Long Contexts)
//!
//! Extracted from gguf_monolith.rs (PMAT-802)
//!
//! ## Test Groups
//!
//! - PARITY-035: Chunked Prefill Tests (IMP-320) (5 tests)

#![allow(clippy::needless_range_loop)]

// PARITY-035: Chunked Prefill Tests (IMP-320)
// =========================================================================

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity035a_chunked_prefill_creation() {
    println!("=== PARITY-035a: Chunked Prefill Creation ===\n");

    use crate::gguf::{ChunkedPrefill, ChunkedPrefillConfig};

    let prompt: Vec<u32> = (0..2048).collect();
    let config = ChunkedPrefillConfig::with_chunk_size(512);
    let prefill = ChunkedPrefill::new(&prompt, config);

    assert_eq!(prefill.total_chunks(), 4);
    assert_eq!(prefill.total_tokens(), 2048);
    assert!(prefill.has_more_chunks());

    println!("  ChunkedPrefill created:");
    println!("    Prompt length: 2048 tokens");
    println!("    Chunk size: 512 tokens");
    println!("    Total chunks: {}", prefill.total_chunks());

    println!("\n  Status: VERIFIED");
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity035b_chunk_iteration() {
    println!("=== PARITY-035b: Chunk Iteration ===\n");

    use crate::gguf::{ChunkedPrefill, ChunkedPrefillConfig};

    let prompt: Vec<u32> = (0..1500).collect();
    let config = ChunkedPrefillConfig::with_chunk_size(512);
    let mut prefill = ChunkedPrefill::new(&prompt, config);

    let mut chunk_sizes = Vec::new();
    while let Some(chunk) = prefill.next_chunk() {
        chunk_sizes.push(chunk.len());
        prefill.complete_chunk(10.0); // Simulate 10ms per chunk
    }

    // 1500 tokens / 512 = 2 full chunks + 1 partial
    assert_eq!(chunk_sizes.len(), 3);
    assert_eq!(chunk_sizes[0], 512);
    assert_eq!(chunk_sizes[1], 512);
    assert_eq!(chunk_sizes[2], 476); // Remaining tokens

    println!("  Chunks processed: {:?}", chunk_sizes);
    println!("  Total chunks: {}", chunk_sizes.len());
    assert!(!prefill.has_more_chunks());

    println!("\n  Status: VERIFIED");
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity035c_progress_tracking() {
    println!("=== PARITY-035c: Progress Tracking ===\n");

    use crate::gguf::{ChunkedPrefill, ChunkedPrefillConfig};

    let prompt: Vec<u32> = (0..2048).collect();
    let config = ChunkedPrefillConfig::with_chunk_size(512);
    let mut prefill = ChunkedPrefill::new(&prompt, config);

    // Process first chunk
    let _ = prefill.next_chunk();
    prefill.complete_chunk(100.0);

    let progress = prefill.progress();
    assert_eq!(progress.chunk_idx, 0);
    assert_eq!(progress.total_chunks, 4);
    assert_eq!(progress.tokens_processed, 512);
    assert_eq!(progress.total_tokens, 2048);

    println!("  After first chunk:");
    println!(
        "    Progress: {}/{} chunks",
        progress.chunk_idx + 1,
        progress.total_chunks
    );
    println!(
        "    Tokens: {}/{}",
        progress.tokens_processed, progress.total_tokens
    );
    println!("    Cumulative time: {:.1}ms", progress.cumulative_time_ms);

    // Process remaining chunks
    while let Some(_chunk) = prefill.next_chunk() {
        prefill.complete_chunk(100.0);
    }

    let final_progress = prefill.progress();
    assert_eq!(final_progress.tokens_processed, 2048);
    assert_eq!(final_progress.cumulative_time_ms, 400.0);

    println!("\n  After all chunks:");
    println!("    Total time: {:.1}ms", final_progress.cumulative_time_ms);

    println!("\n  Status: VERIFIED");
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity035d_ttft_improvement() {
    println!("=== PARITY-035d: TTFT Improvement ===\n");

    use crate::gguf::{ChunkedPrefill, ChunkedPrefillConfig};

    // Simulate 8K context
    let context_length = 8192;
    let prompt: Vec<u32> = (0..context_length as u32).collect();

    // Without chunking: process all at once
    // Typical prefill speed: ~2000 tok/s
    let prefill_tps = 2000.0;
    let full_prefill_ms = context_length as f64 / prefill_tps * 1000.0;

    println!("  Without chunking:");
    println!("    Context: {} tokens", context_length);
    println!("    Prefill speed: {:.0} tok/s", prefill_tps);
    println!(
        "    TTFT: {:.1}ms (must wait for full prefill)",
        full_prefill_ms
    );

    // With chunking: first token after first chunk
    let chunk_size = 512;
    let config = ChunkedPrefillConfig::with_chunk_size(chunk_size);
    let mut prefill = ChunkedPrefill::new(&prompt, config);

    let first_chunk_ms = chunk_size as f64 / prefill_tps * 1000.0;

    // Simulate processing
    while let Some(_chunk) = prefill.next_chunk() {
        let chunk_time = chunk_size as f64 / prefill_tps * 1000.0;
        prefill.complete_chunk(chunk_time);
    }

    println!("\n  With chunking ({}tok chunks):", chunk_size);
    println!("    Total chunks: {}", prefill.total_chunks());
    println!("    TTFT: {:.1}ms (after first chunk)", first_chunk_ms);

    let ttft_speedup = full_prefill_ms / first_chunk_ms;
    println!("\n  TTFT improvement: {:.1}x faster", ttft_speedup);

    // 8K / 512 = 16 chunks, so TTFT should be 16x faster
    assert!(
        ttft_speedup >= 14.0,
        "Should be at least 14x TTFT improvement"
    );

    println!(
        "\n  Status: VERIFIED - {:.1}x TTFT improvement for 8K context",
        ttft_speedup
    );
}

#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_parity035e_stats_and_throughput() {
    println!("=== PARITY-035e: Stats and Throughput ===\n");

    use crate::gguf::{ChunkedPrefill, ChunkedPrefillConfig};

    let prompt: Vec<u32> = (0..4096).collect();
    let config = ChunkedPrefillConfig::with_chunk_size(512);
    let mut prefill = ChunkedPrefill::new(&prompt, config);

    // Simulate realistic timing (256ms per 512-token chunk at 2000 tok/s)
    while let Some(_chunk) = prefill.next_chunk() {
        prefill.complete_chunk(256.0);
    }

    let stats = prefill.stats();

    println!("  Chunked Prefill Statistics:");
    println!("    Total chunks: {}", stats.total_chunks);
    println!("    Chunk size: {}", stats.chunk_size);
    println!("    Total tokens: {}", stats.total_tokens);
    println!("    Total time: {:.1}ms", stats.total_time_ms);
    println!("    Avg chunk time: {:.1}ms", stats.avg_chunk_time_ms);
    println!("    TTFT: {:.1}ms", stats.ttft_ms);
    println!("    Throughput: {:.0} tok/s", stats.tokens_per_second);

    assert_eq!(stats.total_chunks, 8);
    assert_eq!(stats.total_tokens, 4096);
    assert_eq!(stats.ttft_ms, 256.0);

    // 4096 tokens / 2048ms = 2000 tok/s
    assert!(
        stats.tokens_per_second >= 1900.0,
        "Should maintain ~2000 tok/s"
    );

    // IMP-320 target: TTFT < 500ms for 8K context
    // For 4K context with 512-token chunks, TTFT = 256ms
    assert!(stats.ttft_ms < 500.0, "TTFT should be < 500ms");

    println!("\n  IMP-320 Target: TTFT < 500ms for 8K context");
    println!("  Achieved: {:.0}ms TTFT for 4K context", stats.ttft_ms);

    println!("\n  Status: VERIFIED - meets TTFT target");
}