foxhunt/crates/ml/src/transformers/mod.rs

//! # State-of-the-Art Transformer Models for HFT
//!
//! This module implements cutting-edge transformer architectures optimized for
//! ultra-low latency financial market prediction targeting sub-100μs inference.
//!
//! ## Key Innovations for 2025 HFT Applications
//!
//! - **Minimal Architecture**: 1-2 layers, 1-2 heads, optimized for speed
//! - **FlashAttention 2.0**: Memory-efficient attention via Candle
//! - **Financial Features**: Market microstructure, order book, trade flow
//! - **GPU Acceleration**: Pre-allocated tensors, zero-copy operations
//! - **Quantization Ready**: Support for INT8/INT4 optimization
//! - **LoRA Fine-tuning**: Efficient market adaptation
//!
//! ## Architecture Philosophy
//!
//! Based on 2025 HFT requirements, these transformers prioritize:
//! 1. **Latency over Accuracy**: Sub-100μs inference is paramount
//! 2. **Hardware Optimization**: Custom kernels, CUDA graphs
//! 3. **Feature Engineering**: Alpha captured in features, not model complexity
//! 4. **Memory Efficiency**: Pre-allocated GPU memory pools
//!
//! ## Performance Targets
//!
//! - **Inference Latency**: <100μs end-to-end
//! - **Feature Processing**: <20μs for market data normalization
//! - **Model Forward Pass**: <50μs for transformer computation
//! - **Memory Usage**: <256MB GPU memory footprint

// Core modules that compile successfully
pub mod attention;

// Re-export core types that are implemented
pub use attention::AttentionMask;
// TODO: Re-export remaining types when implemented:
// pub use attention::{AttentionConfig, CrossModalAttention, MultiHeadAttention};

/// Transformer model types optimized for different `HFT` use cases
/// TransformerType component.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransformerType {
    /// Ultra-minimal transformer for <50μs inference
    Minimal,
    /// Temporal Fusion Transformer for multi-horizon forecasting
    TemporalFusion,
    /// Sparse transformer for efficiency with longer sequences
    Sparse,
    /// Cross-modal transformer for `price`/`volume`/news fusion
    CrossModal,
}

/// Model size presets optimized for different latency requirements
/// ModelSize component.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelSize {
    /// Ultra-fast: 1 layer, 1 head, 32 dims - target <25μs
    Nano,
    /// Fast: 1 layer, 2 heads, 64 dims - target <50μs
    Micro,
    /// Balanced: 2 layers, 2 heads, 128 dims - target <100μs
    Small,
    /// Custom size configuration
    Custom,
}

impl ModelSize {
    /// Get the configuration parameters for each model size
    pub const fn config(self) -> (usize, usize, usize) {
        match self {
            Self::Nano => (1, 1, 32),   // (layers, heads, dim)
            Self::Micro => (1, 2, 64),  // (layers, heads, dim)
            Self::Small => (2, 2, 128), // (layers, heads, dim)
            Self::Custom => (1, 1, 32), // Default to Nano
        }
    }

    /// Get expected inference latency in microseconds
    pub const fn expected_latency_us(self) -> u64 {
        match self {
            Self::Nano => 25,
            Self::Micro => 50,
            Self::Small => 100,
            Self::Custom => 50,
        }
    }
}

/// Device types for computation
/// DeviceType component.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DeviceType {
    /// `CPU` computation
    CPU,
    /// `CUDA` `GPU` computation
    Cuda,
    /// Metal `GPU` computation (Apple)
    Metal,
}

/// Configuration for `HFT`-optimized transformers
/// HFTTransformerConfig component.
#[derive(Debug, Clone, Copy)]
pub struct HFTTransformerConfig {
    /// Model type and architecture
    pub model_type: TransformerType,

    /// Model size preset
    pub model_size: ModelSize,

    /// Custom dimensions (if ModelSize::Custom)
    pub num_layers: usize,
    pub num_heads: usize,
    pub hidden_dim: usize,
    pub ff_dim: usize,

    /// Sequence length for market data
    pub seq_len: usize,

    /// Feature configuration
    pub feature_dim: usize,
    pub use_market_microstructure: bool,
    pub use_order_book_features: bool,
    pub use_trade_flow_features: bool,

    /// Optimization settings
    pub use_flash_attention: bool,
    pub use_sparse_attention: bool,
    pub attention_sparsity: f32,

    /// Memory optimization
    pub pre_allocate_tensors: bool,
    pub memory_pool_size: usize,

    /// Quantization
    pub use_quantization: bool,
    pub quantization_bits: u8, // 8, 4, or 2 bits

    /// Hardware settings
    pub device_type: DeviceType,
    pub use_cuda_graphs: bool,
    pub enable_profiling: bool,
}

impl Default for HFTTransformerConfig {
    fn default() -> Self {
        Self {
            model_type: TransformerType::Minimal,
            model_size: ModelSize::Micro,
            num_layers: 1,
            num_heads: 2,
            hidden_dim: 64,
            ff_dim: 256,
            seq_len: 64,
            feature_dim: 32,
            use_market_microstructure: true,
            use_order_book_features: true,
            use_trade_flow_features: true,
            use_flash_attention: true,
            use_sparse_attention: false,
            attention_sparsity: 0.1,
            pre_allocate_tensors: true,
            memory_pool_size: 1024 * 1024 * 64, // 64MB
            use_quantization: false,
            quantization_bits: 8,
            device_type: DeviceType::Cuda,
            use_cuda_graphs: false, // Enable after validation
            enable_profiling: false,
        }
    }
}

impl HFTTransformerConfig {
    /// Create configuration for ultra-low latency (Nano model)
    pub fn nano() -> Self {
        let (layers, heads, dim) = ModelSize::Nano.config();
        Self {
            model_size: ModelSize::Nano,
            num_layers: layers,
            num_heads: heads,
            hidden_dim: dim,
            ff_dim: dim * 2,
            seq_len: 32,
            feature_dim: 16,
            ..Default::default()
        }
    }

    /// Create configuration for balanced latency/accuracy (Micro model)
    pub fn micro() -> Self {
        let (layers, heads, dim) = ModelSize::Micro.config();
        Self {
            model_size: ModelSize::Micro,
            num_layers: layers,
            num_heads: heads,
            hidden_dim: dim,
            ff_dim: dim * 4,
            ..Default::default()
        }
    }

    /// Create configuration for maximum accuracy within 100μs (Small model)
    pub fn small() -> Self {
        let (layers, heads, dim) = ModelSize::Small.config();
        Self {
            model_size: ModelSize::Small,
            num_layers: layers,
            num_heads: heads,
            hidden_dim: dim,
            ff_dim: dim * 4,
            seq_len: 128,
            feature_dim: 64,
            ..Default::default()
        }
    }

    /// Enable all optimizations for production deployment
    pub fn production() -> Self {
        Self {
            use_flash_attention: true,
            pre_allocate_tensors: true,
            use_quantization: true,
            quantization_bits: 8,
            use_cuda_graphs: true,
            ..Self::micro()
        }
    }

    /// Configuration for benchmarking and validation
    pub fn benchmark() -> Self {
        Self {
            enable_profiling: true,
            ..Self::micro()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_model_size_config() {
        assert_eq!(ModelSize::Nano.config(), (1, 1, 32));
        assert_eq!(ModelSize::Micro.config(), (1, 2, 64));
        assert_eq!(ModelSize::Small.config(), (2, 2, 128));
    }

    #[test]
    fn test_latency_expectations() {
        assert_eq!(ModelSize::Nano.expected_latency_us(), 25);
        assert_eq!(ModelSize::Micro.expected_latency_us(), 50);
        assert_eq!(ModelSize::Small.expected_latency_us(), 100);
    }

    #[test]
    fn test_config_presets() {
        let nano = HFTTransformerConfig::nano();
        assert_eq!(nano.model_size, ModelSize::Nano);
        assert_eq!(nano.num_layers, 1);
        assert_eq!(nano.num_heads, 1);
        assert_eq!(nano.hidden_dim, 32);

        let production = HFTTransformerConfig::production();
        assert!(production.use_flash_attention);
        assert!(production.pre_allocate_tensors);
        assert!(production.use_quantization);
        assert!(production.use_cuda_graphs);
    }
}