foxhunt/crates/ml/tests/tft_inference_latency_benchmark.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! TFT Inference Latency Benchmark for HFT Production
//!
//! **Objective**: Measure TFT inference latency and ensure P95 <5ms for production HFT.
//!
//! **Performance Targets**:
//! - P95 latency: <5ms (5000us)
//! - Mean latency: <2ms (2000us)
//! - Consistent performance: P99/P50 ratio <2.0
//!
//! **Comparison with Other Models**:
//! - DQN: P95 2.1ms
//! - PPO: P95 3.2ms
//! - MAMBA-2: P95 1.8ms
//! - TFT: P95 target <5ms
//!
//! **Optimization Strategies** (if >5ms):
//! 1. CUDA Kernel Fusion: Reduce kernel launch overhead
//! 2. Batch Size = 1: Single-sample inference (lowest latency)
//! 3. Mixed Precision: FP16 inference (2x faster)
//! 4. Model Quantization: INT8 inference (4x faster)
//! 5. Attention Optimization: Flash Attention (2-4x faster)

#![allow(unused_crate_dependencies)]

use ml::tft::{TFTConfig, TemporalFusionTransformer};
use ml::MLError;
use ml_supervised::gpu_tensor::GpuTensor;
use std::sync::Arc;
use std::time::Instant;
use tracing::info;
use tracing::warn;

/// Helper: Create realistic TFT input tensors for benchmarking.
/// Uses the TFT model's own stream for GPU allocation.
fn create_tft_inputs(
    config: &TFTConfig,
    stream: &Arc<cudarc::driver::CudaStream>,
) -> Result<(GpuTensor, GpuTensor, GpuTensor), MLError> {
    // Static features: [batch=1, num_static_features]
    let static_data = vec![0.5f32; config.num_static_features];
    let static_features =
        GpuTensor::from_vec(static_data, &[1, config.num_static_features], stream)?;

    // Historical features: [batch=1, seq_len, num_unknown_features]
    let hist_len = config.sequence_length;
    let hist_dim = config.num_unknown_features;
    let hist_data = vec![0.5f32; hist_len * hist_dim];
    let historical_features = GpuTensor::from_vec(hist_data, &[1, hist_len, hist_dim], stream)?;

    // Future features: [batch=1, prediction_horizon, num_known_features]
    let fut_len = config.prediction_horizon;
    let fut_dim = config.num_known_features;
    let fut_data = vec![0.5f32; fut_len * fut_dim];
    let future_features = GpuTensor::from_vec(fut_data, &[1, fut_len, fut_dim], stream)?;

    Ok((static_features, historical_features, future_features))
}

/// Primary benchmark: TFT inference latency with statistical analysis
#[test]
fn test_tft_inference_latency_p95_target() -> Result<(), MLError> {
    info!("=== TFT Inference Latency Benchmark ===");
    info!("Target: P95 <5ms (5000us) for production HFT");

    // Production-realistic TFT configuration
    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 128,
        num_heads: 8,
        num_layers: 3,
        prediction_horizon: 10,
        sequence_length: 50,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
        learning_rate: 1e-3,
        batch_size: 1,     // HFT: Single-sample inference for lowest latency
        dropout_rate: 0.0, // Inference mode: No dropout
        l2_regularization: 1e-4,
        use_flash_attention: true,
        memory_efficient: true,
        max_inference_latency_us: 5000,
        target_throughput_pps: 100_000,
    };

    let mut tft = TemporalFusionTransformer::new(config.clone())?;
    let stream = tft.stream().clone();

    // Prepare inputs
    let (static_features, historical_features, future_features) =
        create_tft_inputs(&config, &stream)?;

    // ==================== WARMUP PHASE ====================
    // Critical for CUDA: Ensure kernels are compiled and caches are warm
    info!("Warmup: 5 iterations (CUDA kernel compilation)");
    for _ in 0..5 {
        let _ = tft.forward(&static_features, &historical_features, &future_features)?;
    }

    // ==================== BENCHMARK PHASE ====================
    // 100 iterations for stable statistics
    info!("Benchmark: 100 iterations for stable statistics");
    let num_iterations = 100;
    let mut latencies_us = Vec::with_capacity(num_iterations);

    for _ in 0..num_iterations {
        let start = Instant::now();
        let _output = tft.forward(&static_features, &historical_features, &future_features)?;
        let elapsed_us = start.elapsed().as_micros() as u64;
        latencies_us.push(elapsed_us);
    }

    // ==================== STATISTICAL ANALYSIS ====================
    latencies_us.sort_unstable();

    let mean = latencies_us.iter().sum::<u64>() as f64 / latencies_us.len() as f64;
    let p50 = latencies_us[latencies_us.len() / 2];
    let p95 = latencies_us[latencies_us.len() * 95 / 100];
    let p99 = latencies_us[latencies_us.len() * 99 / 100];
    let min = latencies_us[0];
    let max = latencies_us[latencies_us.len() - 1];

    // Consistency metric: P99/P50 ratio (lower is better)
    let consistency_ratio = p99 as f64 / p50 as f64;

    // ==================== RESULTS ====================
    info!(
        mean_us = mean as u64,
        mean_ms = mean / 1000.0,
        p50_us = p50,
        p50_ms = p50 as f64 / 1000.0,
        p95_us = p95,
        p95_ms = p95 as f64 / 1000.0,
        p99_us = p99,
        p99_ms = p99 as f64 / 1000.0,
        min_us = min,
        max_us = max,
        consistency_ratio,
        "TFT Inference Latency Statistics"
    );

    // ==================== VALIDATION ====================
    // Primary target: P95 <5ms (5000us)
    if p95 < 5000 {
        info!(p95_us = p95, "PASS: P95 latency is <5ms target");
    } else {
        warn!(
            p95_us = p95,
            "WARNING: P95 latency exceeds 5ms target. Optimization strategies: \
             1. Enable Flash Attention (2-4x speedup) \
             2. Mixed Precision FP16 (2x speedup) \
             3. Model Quantization INT8 (4x speedup) \
             4. Reduce hidden_dim or num_layers \
             5. CUDA kernel fusion"
        );
    }

    // Secondary target: Mean <2ms
    if mean < 2000.0 {
        info!(mean_us = mean as u64, "PASS: Mean latency is <2ms");
    } else {
        info!(mean_us = mean as u64, "INFO: Mean latency exceeds 2ms (non-critical)");
    }

    // Consistency check: P99/P50 ratio <2.0
    if consistency_ratio < 2.0 {
        info!(consistency_ratio, "PASS: Consistency ratio is <2.0 (stable performance)");
    } else {
        warn!(consistency_ratio, "WARNING: Consistency ratio exceeds 2.0 (high variance)");
    }

    // Test assertion: P95 must be <5ms for production readiness
    assert!(
        p95 < 5000,
        "FAIL: TFT P95 latency {}us exceeds 5ms target ({}ms)",
        p95,
        p95 as f64 / 1000.0
    );

    Ok(())
}

/// Comparison benchmark: TFT vs other models (DQN, PPO, MAMBA-2)
#[test]
fn test_tft_latency_comparison_with_other_models() -> Result<(), MLError> {
    info!("=== Model Inference Latency Comparison ===");

    // TFT benchmark (from above)
    let tft_config = TFTConfig {
        hidden_dim: 128,
        num_heads: 8,
        num_layers: 3,
        prediction_horizon: 10,
        sequence_length: 50,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
        batch_size: 1,
        dropout_rate: 0.0,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(tft_config.clone())?;
    let stream = tft.stream().clone();
    let (static_features, historical_features, future_features) =
        create_tft_inputs(&tft_config, &stream)?;

    // Warmup
    for _ in 0..5 {
        let _ = tft.forward(&static_features, &historical_features, &future_features)?;
    }

    // Benchmark TFT
    let mut tft_latencies = Vec::new();
    for _ in 0..100 {
        let start = Instant::now();
        let _ = tft.forward(&static_features, &historical_features, &future_features)?;
        tft_latencies.push(start.elapsed().as_micros() as u64);
    }
    tft_latencies.sort_unstable();

    let tft_mean = tft_latencies.iter().sum::<u64>() as f64 / tft_latencies.len() as f64;
    let tft_p50 = tft_latencies[tft_latencies.len() / 2];
    let tft_p95 = tft_latencies[tft_latencies.len() * 95 / 100];
    let tft_p99 = tft_latencies[tft_latencies.len() * 99 / 100];

    // ==================== COMPARISON TABLE ====================
    // Reference: DQN P95=2.1ms, PPO P95=3.2ms, MAMBA-2 P95=1.8ms (all <5ms target)
    info!(
        tft_mean_us = tft_mean as u64,
        tft_p50_us = tft_p50,
        tft_p95_ms = tft_p95 as f64 / 1000.0,
        tft_p99_ms = tft_p99 as f64 / 1000.0,
        tft_meets_target = tft_p95 < 5000,
        "Model latency comparison: DQN=2.1ms, PPO=3.2ms, MAMBA2=1.8ms (all <5ms)"
    );

    // Verify TFT meets target
    assert!(
        tft_p95 < 5000,
        "TFT P95 latency {}us exceeds 5ms target",
        tft_p95
    );

    Ok(())
}

/// Test: TFT latency with different batch sizes (batch optimization)
#[test]
fn test_tft_batch_size_latency_tradeoff() -> Result<(), MLError> {
    info!("=== TFT Batch Size Latency Trade-off ===");

    let batch_sizes = vec![1, 2, 4, 8];

    for batch_size in batch_sizes {
        let config = TFTConfig {
            hidden_dim: 128,
            num_heads: 8,
            num_layers: 3,
            prediction_horizon: 10,
            sequence_length: 50,
            num_quantiles: 9,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
            batch_size,
            dropout_rate: 0.0,
            ..Default::default()
        };

        let mut tft = TemporalFusionTransformer::new(config.clone())?;
        let stream = tft.stream().clone();

        // Create batched inputs
        let static_data = vec![0.5f32; batch_size * config.num_static_features];
        let static_features = GpuTensor::from_vec(
            static_data,
            &[batch_size, config.num_static_features],
            &stream,
        )?;

        let hist_data =
            vec![0.5f32; batch_size * config.sequence_length * config.num_unknown_features];
        let historical_features = GpuTensor::from_vec(
            hist_data,
            &[
                batch_size,
                config.sequence_length,
                config.num_unknown_features,
            ],
            &stream,
        )?;

        let fut_data =
            vec![0.5f32; batch_size * config.prediction_horizon * config.num_known_features];
        let future_features = GpuTensor::from_vec(
            fut_data,
            &[
                batch_size,
                config.prediction_horizon,
                config.num_known_features,
            ],
            &stream,
        )?;

        // Warmup
        for _ in 0..5 {
            let _ = tft.forward(&static_features, &historical_features, &future_features)?;
        }

        // Benchmark
        let num_iterations = 50;
        let mut latencies = Vec::new();
        for _ in 0..num_iterations {
            let start = Instant::now();
            let _ = tft.forward(&static_features, &historical_features, &future_features)?;
            latencies.push(start.elapsed().as_micros() as u64);
        }

        let avg_latency_us = latencies.iter().sum::<u64>() / num_iterations as u64;
        let per_sample_us = avg_latency_us as f64 / batch_size as f64;
        let throughput = 1_000_000.0 / per_sample_us;

        info!(
            batch_size,
            avg_latency_us,
            per_sample_us = per_sample_us as u64,
            throughput = throughput as u64,
            "Batch size latency tradeoff"
        );
    }

    info!("Insight: Larger batches amortize overhead, increasing throughput. HFT recommendation: batch_size=1 for lowest latency (<5ms)");

    Ok(())
}

/// Test: TFT latency with Flash Attention enabled/disabled
#[test]
fn test_tft_flash_attention_speedup() -> Result<(), MLError> {
    info!("=== TFT Flash Attention Speedup ===");

    // Test both configurations
    let flash_configs = vec![("Standard Attention", false), ("Flash Attention", true)];

    let mut baseline_p95 = 0u64;

    for (name, use_flash) in flash_configs {
        let config = TFTConfig {
            hidden_dim: 128,
            num_heads: 8,
            num_layers: 3,
            prediction_horizon: 10,
            sequence_length: 50,
            num_quantiles: 9,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
            use_flash_attention: use_flash,
            batch_size: 1,
            dropout_rate: 0.0,
            ..Default::default()
        };

        let mut tft = TemporalFusionTransformer::new(config.clone())?;
        let stream = tft.stream().clone();
        let (static_features, historical_features, future_features) =
            create_tft_inputs(&config, &stream)?;

        // Warmup
        for _ in 0..5 {
            let _ = tft.forward(&static_features, &historical_features, &future_features)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..100 {
            let start = Instant::now();
            let _ = tft.forward(&static_features, &historical_features, &future_features)?;
            latencies.push(start.elapsed().as_micros() as u64);
        }
        latencies.sort_unstable();

        let p50 = latencies[latencies.len() / 2];
        let p95 = latencies[latencies.len() * 95 / 100];

        let speedup = if baseline_p95 > 0 {
            baseline_p95 as f64 / p95 as f64
        } else {
            baseline_p95 = p95;
            1.0
        };

        info!(
            config_name = name,
            p50_us = p50,
            p95_us = p95,
            speedup,
            "Flash attention comparison"
        );
    }

    info!("Flash Attention expected speedup: 2-4x for long sequences");

    Ok(())
}

/// Test: TFT latency with different model sizes (hidden_dim, num_layers)
#[test]
fn test_tft_model_size_latency_scaling() -> Result<(), MLError> {
    info!("=== TFT Model Size Latency Scaling ===");

    // Test configurations: (hidden_dim, num_layers, name)
    let model_configs = vec![
        (64, 2, "Small"),
        (128, 3, "Medium (Production)"),
        (256, 4, "Large"),
        (512, 6, "Extra Large"),
    ];

    for (hidden_dim, num_layers, name) in model_configs {
        let config = TFTConfig {
            hidden_dim,
            num_heads: 8,
            num_layers,
            prediction_horizon: 10,
            sequence_length: 50,
            num_quantiles: 9,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
            batch_size: 1,
            dropout_rate: 0.0,
            use_flash_attention: true,
            ..Default::default()
        };

        let mut tft = TemporalFusionTransformer::new(config.clone())?;
        let stream = tft.stream().clone();
        let (static_features, historical_features, future_features) =
            create_tft_inputs(&config, &stream)?;

        // Warmup
        for _ in 0..5 {
            let _ = tft.forward(&static_features, &historical_features, &future_features)?;
        }

        // Benchmark
        let mut latencies = Vec::new();
        for _ in 0..100 {
            let start = Instant::now();
            let _ = tft.forward(&static_features, &historical_features, &future_features)?;
            latencies.push(start.elapsed().as_micros() as u64);
        }
        latencies.sort_unstable();

        let p95 = latencies[latencies.len() * 95 / 100];
        let meets_target = p95 < 5000;

        info!(
            model_name = name,
            hidden_dim,
            num_layers,
            p95_us = p95,
            meets_target,
            "Model size latency scaling"
        );
    }

    info!("Latency scales with model size: Smaller models = lower latency. Production: Medium (128, 3 layers) balances accuracy and latency");

    Ok(())
}

/// Test: TFT memory usage during inference
#[test]
fn test_tft_inference_memory_usage() -> Result<(), MLError> {
    info!("=== TFT Inference Memory Usage ===");

    let config = TFTConfig {
        hidden_dim: 128,
        num_heads: 8,
        num_layers: 3,
        prediction_horizon: 10,
        sequence_length: 50,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
        batch_size: 1,
        dropout_rate: 0.0,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(config.clone())?;
    let stream = tft.stream().clone();
    let (static_features, historical_features, future_features) =
        create_tft_inputs(&config, &stream)?;

    // Single inference
    let _ = tft.forward(&static_features, &historical_features, &future_features)?;

    // Estimate memory usage
    let static_mem = config.num_static_features * 4; // f32
    let hist_mem = config.sequence_length * config.num_unknown_features * 4;
    let fut_mem = config.prediction_horizon * config.num_known_features * 4;
    let output_mem = config.prediction_horizon * config.num_quantiles * 4;

    let total_input_mem = static_mem + hist_mem + fut_mem;
    let total_mem = total_input_mem + output_mem;

    info!(
        static_mem_bytes = static_mem,
        static_mem_kb = static_mem as f64 / 1024.0,
        hist_mem_bytes = hist_mem,
        hist_mem_kb = hist_mem as f64 / 1024.0,
        fut_mem_bytes = fut_mem,
        fut_mem_kb = fut_mem as f64 / 1024.0,
        output_mem_bytes = output_mem,
        output_mem_kb = output_mem as f64 / 1024.0,
        total_mem_bytes = total_mem,
        total_mem_kb = total_mem as f64 / 1024.0,
        "Memory Usage Breakdown"
    );

    // Verify memory is reasonable
    assert!(
        total_mem < 10 * 1024 * 1024,
        "Memory usage exceeds 10MB target"
    );

    Ok(())
}