aprender-serve 0.64.0

//! GGUF Part 13: PARITY-078 - PARITY-086 (Phase 4/5: FlashAttention-2 & Stream-K)
//!
//! Extracted from gguf_monolith.rs (PMAT-802)
//!
//! ## Test Groups
//!
//! - PARITY-078: Work Partitioning (6 tests)
//! - PARITY-079: Non-matmul FLOP Reduction (6 tests)
//! - PARITY-080: Tensor Core Integration (6 tests)
//! - PARITY-081: Phase 4 Integration Summary (6 tests)
//! - PARITY-082: Stream-K Work Decomposition (6 tests)
//! - PARITY-083: Irregular Matrix Handling (6 tests)
//! - PARITY-084: Production Serving Integration (6 tests)
//! - PARITY-085: Benchmark Validation (6 tests)
//! - PARITY-086: Phase 5 Final Summary (6 tests)
//!
//! Phase 4 Target: FlashAttention-2 implementation
//! Phase 5 Target: Stream-K work decomposition & production polish

// ==================== PARITY-078: Work Partitioning ====================
// Improved parallelism via sequence and batch parallelization

/// PARITY-078a: Sequence parallelism
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_078a_sequence_parallelism() {
    println!("PARITY-078a: Sequence Parallelism");
    println!("==================================");
    println!();

    // FlashAttention-2 key insight: parallelize over sequence dimension
    // Each thread block handles one Q tile (Br rows)
    // Multiple blocks process different parts of sequence

    let seq_len = 4096u32;
    let br = 128u32;
    let n_q_blocks = seq_len.div_ceil(br);

    println!("  Sequence Parallelization:");
    println!("  --------------------------");
    println!("    Sequence length: {}", seq_len);
    println!("    Block size (Br): {}", br);
    println!("    Thread blocks for Q: {}", n_q_blocks);
    println!();

    // RTX 4090 SM utilization
    let sms = 128u32; // RTX 4090 has 128 SMs
    let blocks_per_head = n_q_blocks;
    let n_heads = 32u32;
    let total_blocks = blocks_per_head * n_heads;

    println!("  GPU Utilization (RTX 4090, {} SMs):", sms);
    println!("    Blocks per head: {}", blocks_per_head);
    println!("    Total heads: {}", n_heads);
    println!("    Total blocks: {}", total_blocks);
    println!("    Blocks per SM: {:.1}", total_blocks as f32 / sms as f32);
    println!();

    // Wave efficiency
    let waves = total_blocks.div_ceil(sms);
    let last_wave_occupancy = (total_blocks % sms) as f32 / sms as f32 * 100.0;

    println!("  Wave Efficiency:");
    println!("    Full waves: {}", total_blocks / sms);
    println!("    Total waves: {}", waves);
    println!(
        "    Last wave occupancy: {:.1}%",
        if last_wave_occupancy == 0.0 {
            100.0
        } else {
            last_wave_occupancy
        }
    );

    assert!(
        total_blocks >= sms,
        "PARITY-078a: Should have enough blocks to fill all SMs"
    );
    assert!(true, "PARITY-078a: Sequence parallelism documented");
}

/// PARITY-078b: Batch parallelism
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_078b_batch_parallelism() {
    println!("PARITY-078b: Batch Parallelism");
    println!("===============================");
    println!();

    // Batch dimension adds more parallelism
    let batch_size = 8u32;
    let seq_len = 2048u32;
    let n_heads = 32u32;
    let br = 128u32;

    let blocks_per_head = seq_len.div_ceil(br);
    let blocks_per_request = blocks_per_head * n_heads;
    let total_blocks = blocks_per_request * batch_size;

    println!("  Batch Configuration:");
    println!("  ---------------------");
    println!("    Batch size: {}", batch_size);
    println!("    Sequence length: {}", seq_len);
    println!("    Heads: {}", n_heads);
    println!();

    println!("  Parallelism Breakdown:");
    println!("    Blocks per head: {}", blocks_per_head);
    println!("    Blocks per request: {}", blocks_per_request);
    println!("    Total blocks: {}", total_blocks);
    println!();

    // Grid dimensions for CUDA
    let grid_x = blocks_per_head;
    let grid_y = n_heads;
    let grid_z = batch_size;

    println!("  CUDA Grid Dimensions:");
    println!("    grid.x (seq blocks): {}", grid_x);
    println!("    grid.y (heads): {}", grid_y);
    println!("    grid.z (batch): {}", grid_z);
    println!(
        "    Total: {} × {} × {} = {}",
        grid_x, grid_y, grid_z, total_blocks
    );

    let sms = 128u32;
    let occupancy = (total_blocks as f32 / sms as f32).min(1.0) * 100.0;
    println!();
    println!("  SM Occupancy: {:.1}%", occupancy);

    assert!(
        total_blocks > sms * 2,
        "PARITY-078b: Batched workload should saturate SMs"
    );
    assert!(true, "PARITY-078b: Batch parallelism documented");
}

/// PARITY-078c: Head parallelism
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_078c_head_parallelism() {
    println!("PARITY-078c: Head Parallelism");
    println!("==============================");
    println!();

    // Each attention head is independent - perfect parallelism
    let n_heads = 32u32;
    let head_dim = 128u32;
    let total_hidden = n_heads * head_dim;

    println!("  Attention Head Configuration:");
    println!("  ------------------------------");
    println!("    Number of heads: {}", n_heads);
    println!("    Head dimension: {}", head_dim);
    println!("    Total hidden dim: {}", total_hidden);
    println!();

    // Memory per head
    let seq_len = 2048u32;
    let q_per_head = seq_len * head_dim * 4;
    let k_per_head = seq_len * head_dim * 4;
    let v_per_head = seq_len * head_dim * 4;
    let o_per_head = seq_len * head_dim * 4;

    println!("  Memory per Head (seq_len={}):", seq_len);
    println!("    Q: {} KB", q_per_head / 1024);
    println!("    K: {} KB", k_per_head / 1024);
    println!("    V: {} KB", v_per_head / 1024);
    println!("    O: {} KB", o_per_head / 1024);
    println!(
        "    Total per head: {} MB",
        (q_per_head + k_per_head + v_per_head + o_per_head) / 1024 / 1024
    );
    println!();

    // Parallelism options
    println!("  Parallelization Strategies:");
    println!("    1. One block per head: {} blocks", n_heads);
    println!("    2. Multiple blocks per head: {} × seq_blocks", n_heads);
    println!("    3. Sub-head parallelism (tensor cores): 16×16 tiles");

    assert!(n_heads >= 8, "PARITY-078c: Modern models have 8+ heads");
    assert!(true, "PARITY-078c: Head parallelism documented");
}

/// PARITY-078d: Work stealing for load balancing
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_078d_work_stealing() {
    println!("PARITY-078d: Work Stealing for Load Balancing");
    println!("==============================================");
    println!();

    // Problem: Causal attention has triangular workload
    // Early rows: few K/V tiles to process
    // Late rows: many K/V tiles to process

    let seq_len = 2048u32;
    let bc = 64u32;
    let n_kv_blocks = seq_len / bc;

    println!("  Causal Attention Workload Distribution:");
    println!("  ----------------------------------------");
    println!("    Sequence length: {}", seq_len);
    println!("    K/V block size: {}", bc);
    println!("    K/V blocks: {}", n_kv_blocks);
    println!();

    // Work distribution (causal)
    let first_row_work = 1u32;
    let last_row_work = n_kv_blocks;
    let total_work = n_kv_blocks * (n_kv_blocks + 1) / 2;
    let avg_work = total_work as f32 / n_kv_blocks as f32;

    println!("  Work per Q Block:");
    println!("    First Q block: {} K/V blocks", first_row_work);
    println!("    Last Q block: {} K/V blocks", last_row_work);
    println!("    Total work: {} tile-ops", total_work);
    println!("    Average: {:.1} tiles per Q block", avg_work);
    println!();

    // Load imbalance
    let imbalance = last_row_work as f32 / avg_work;
    println!("  Load Imbalance:");
    println!("    Worst case / average: {:.2}x", imbalance);
    println!();

    // Work stealing solution (per FA2 paper)
    println!("  Work Stealing Strategy:");
    println!("    1. Global work counter (atomic)");
    println!("    2. Each warp fetches next tile");
    println!("    3. Dynamic assignment balances load");
    println!("    4. ~1.3x speedup on causal attention");

    assert!(
        imbalance > 1.5,
        "PARITY-078d: Causal attention has significant imbalance"
    );
    assert!(true, "PARITY-078d: Work stealing documented");
}

/// PARITY-078e: Split-K decomposition
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_078e_split_k_decomposition() {
    println!("PARITY-078e: Split-K Decomposition");
    println!("====================================");
    println!();

    // Split-K: Distribute K/V dimension across multiple thread blocks
    // Each block computes partial output, then reduce

    let seq_len = 4096u32;
    let bc = 64u32;
    let n_kv_blocks = seq_len / bc;
    let split_k = 4u32;

    println!("  Split-K Configuration:");
    println!("  -----------------------");
    println!("    K/V blocks total: {}", n_kv_blocks);
    println!("    Split factor: {}", split_k);
    println!("    Blocks per split: {}", n_kv_blocks / split_k);
    println!();

    // Memory for partial outputs
    let head_dim = 128u32;
    let br = 128u32;
    let partial_o = br * head_dim * 4;
    let partial_m = br * 4; // max values
    let partial_l = br * 4; // sum values

    println!("  Partial Output Storage (per split):");
    println!("    O partial: {} KB", partial_o / 1024);
    println!("    m partial: {} B", partial_m);
    println!("    l partial: {} B", partial_l);
    println!(
        "    Total × {}: {} KB",
        split_k,
        (partial_o + partial_m + partial_l) * split_k / 1024
    );
    println!();

    // Reduction phase
    println!("  Reduction Formula:");
    println!("    m_new = max(m_1, m_2, ..., m_k)");
    println!("    l_new = Σ l_i × exp(m_i - m_new)");
    println!("    O_new = Σ O_i × exp(m_i - m_new) / l_new");
    println!();

    // When to use split-K
    println!("  When to Use Split-K:");
    println!("    ✓ Long sequences (>4K)");
    println!("    ✓ Small batch sizes");
    println!("    ✓ Few attention heads");
    println!("    ✗ Short sequences (overhead > benefit)");

    assert!(split_k >= 2, "PARITY-078e: Split-K factor should be >= 2");
    assert!(true, "PARITY-078e: Split-K decomposition documented");
}

/// PARITY-078f: Work partitioning summary
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_078f_work_partitioning_summary() {
    println!("PARITY-078f: Work Partitioning Summary");
    println!("=======================================");
    println!();

    println!("  ╔═══════════════════════════════════════════════════════════════╗");
    println!("  ║          PARITY-078: Work Partitioning Complete               ║");
    println!("  ╠═══════════════════════════════════════════════════════════════╣");
    println!("  ║                                                               ║");
    println!("  ║  Parallelism Dimensions:                                      ║");
    println!("  ║  ─────────────────────                                        ║");
    println!("  ║  1. Sequence: grid.x = seq_len / Br                           ║");
    println!("  ║  2. Heads: grid.y = n_heads                                   ║");
    println!("  ║  3. Batch: grid.z = batch_size                                ║");
    println!("  ║  4. Split-K: Additional blocks for long sequences            ║");
    println!("  ║                                                               ║");
    println!("  ║  Load Balancing:                                              ║");
    println!("  ║  ───────────────                                              ║");
    println!("  ║  • Work stealing for causal attention                         ║");
    println!("  ║  • Dynamic tile assignment                                    ║");
    println!("  ║  • 1.3x speedup on triangular workloads                       ║");
    println!("  ║                                                               ║");
    println!("  ╚═══════════════════════════════════════════════════════════════╝");
    println!();

    println!("  NEXT: PARITY-079 - Non-matmul FLOP reduction");

    assert!(true, "PARITY-078f: Summary complete");
}

// ==================== PARITY-079: Non-matmul FLOP Reduction ====================
// Reduce overhead from softmax, rescaling, and memory operations

/// PARITY-079a: Softmax FLOP analysis
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_079a_softmax_flop_analysis() {
    println!("PARITY-079a: Softmax FLOP Analysis");
    println!("====================================");
    println!();

    // Softmax: exp(x - max(x)) / sum(exp(x - max(x)))
    // For each row of attention scores

    let seq_len = 2048u32;
    let n_heads = 32u32;
    let batch = 1u32;

    // Per-row operations
    let max_ops_per_row = seq_len; // find max
    let sub_ops_per_row = seq_len; // x - max
    let exp_ops_per_row = seq_len; // exp(...)
    let sum_ops_per_row = seq_len; // sum
    let div_ops_per_row = seq_len; // normalize

    let softmax_ops_per_row =
        max_ops_per_row + sub_ops_per_row + exp_ops_per_row + sum_ops_per_row + div_ops_per_row;
    let total_rows = batch * n_heads * seq_len;
    let total_softmax_flops = total_rows as u64 * softmax_ops_per_row as u64;

    println!("  Softmax Operations per Row (seq_len={}):", seq_len);
    println!("  ----------------------------------------");
    println!("    Max reduction: {} ops", max_ops_per_row);
    println!("    Subtraction: {} ops", sub_ops_per_row);
    println!("    Exponential: {} ops", exp_ops_per_row);
    println!("    Sum reduction: {} ops", sum_ops_per_row);
    println!("    Division: {} ops", div_ops_per_row);
    println!("    Total: {} ops/row", softmax_ops_per_row);
    println!();

    // Compare to matmul
    let head_dim = 128u32;
    let qk_flops = 2u64 * seq_len as u64 * seq_len as u64 * head_dim as u64;
    let av_flops = 2u64 * seq_len as u64 * seq_len as u64 * head_dim as u64;
    let matmul_flops = (qk_flops + av_flops) * n_heads as u64 * batch as u64;

    println!("  FLOP Comparison (batch={}, heads={}):", batch, n_heads);
    println!("    Softmax: {:.2} GFLOP", total_softmax_flops as f64 / 1e9);
    println!("    MatMul: {:.2} GFLOP", matmul_flops as f64 / 1e9);
    println!(
        "    Softmax / MatMul: {:.1}%",
        total_softmax_flops as f64 / matmul_flops as f64 * 100.0
    );
    println!();

    // Softmax is typically 1-5% of total FLOPs but can dominate memory bandwidth

    assert!(
        total_softmax_flops < matmul_flops / 10,
        "PARITY-079a: Softmax should be <10% of matmul FLOPs"
    );
    assert!(true, "PARITY-079a: Softmax FLOP analysis complete");
}

/// PARITY-079b: Online softmax optimization
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_079b_online_softmax() {
    println!("PARITY-079b: Online Softmax Optimization");
    println!("=========================================");
    println!();

    // Traditional softmax: 3 passes over data
    // 1. Find max
    // 2. Compute exp(x - max) and sum
    // 3. Normalize

    // Online softmax: 1 pass (FlashAttention)
    // Track running max and rescale on-the-fly

    println!("  Traditional Softmax (3 passes):");
    println!("  --------------------------------");
    println!("    Pass 1: m = max(x)");
    println!("    Pass 2: s = sum(exp(x - m))");
    println!("    Pass 3: y = exp(x - m) / s");
    println!("    Memory traffic: 3N reads + N writes");
    println!();

    println!("  Online Softmax (1 pass):");
    println!("  -------------------------");
    println!("    for i in range(N):");
    println!("        m_new = max(m_old, x[i])");
    println!("        s = s * exp(m_old - m_new) + exp(x[i] - m_new)");
    println!("        o = o * exp(m_old - m_new) / s");
    println!("    Memory traffic: N reads + N writes");
    println!();

    // Memory bandwidth savings
    let seq_len = 2048u32;
    let elem_size = 4u32;
    let traditional_bytes = (3 * seq_len + seq_len) * elem_size;
    let online_bytes = (seq_len + seq_len) * elem_size;
    let savings = traditional_bytes as f32 / online_bytes as f32;

    println!("  Memory Traffic (seq_len={}):", seq_len);
    println!("    Traditional: {} KB", traditional_bytes / 1024);
    println!("    Online: {} KB", online_bytes / 1024);
    println!("    Savings: {:.1}x", savings);

    assert!(
        savings > 1.5,
        "PARITY-079b: Online softmax should save >1.5x memory"
    );
    assert!(true, "PARITY-079b: Online softmax documented");
}

include!("parity_079c.rs");
include!("parity_080e.rs");
include!("parity_082b.rs");
include!("parity_083d.rs");
include!("parity_084f.rs");
include!("parity_086c.rs");