Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,474 @@
+//! SRAM management for Flash Attention tiling
+
+use crate::{
+    config::FlashAttentionConfig,
+    error::{FlashError, FlashResult},
+};
+use rtx_runtime::CudaBackend;
+use std::sync::Arc;
+use parking_lot::RwLock;
+use tracing::{debug, info};
+
+/// SRAM allocation block
+#[derive(Debug, Clone)]
+pub struct SRAMBlock {
+    /// Block ID for tracking
+    pub id: u64,
+    /// Starting offset in shared memory
+    pub offset: usize,
+    /// Size in bytes
+    pub size: usize,
+    /// Whether the block is currently in use
+    pub in_use: bool,
+    /// Block type (Q, K, V, or scratch)
+    pub block_type: SRAMBlockType,
+}
+
+/// Type of SRAM block
+#[derive(Debug, Clone, PartialEq)]
+pub enum SRAMBlockType {
+    /// Query block
+    Query,
+    /// Key block
+    Key,
+    /// Value block
+    Value,
+    /// Scratch space for intermediate computations
+    Scratch,
+    /// Attention scores
+    Scores,
+}
+
+/// SRAM usage statistics
+#[derive(Debug, Clone)]
+pub struct SRAMStats {
+    /// Total SRAM capacity in bytes
+    pub total_capacity: usize,
+    /// Currently allocated SRAM in bytes
+    pub allocated: usize,
+    /// Peak SRAM usage in bytes
+    pub peak_usage: usize,
+    /// Number of allocations
+    pub allocation_count: u64,
+    /// SRAM utilization efficiency (0.0 - 1.0)
+    pub efficiency: f32,
+    /// Number of allocation failures
+    pub allocation_failures: u64,
+}
+
+/// SRAM memory manager for Flash Attention tiling
+pub struct SRAMManager {
+    /// Configuration
+    config: FlashAttentionConfig,
+    /// CUDA backend
+    cuda_backend: Arc<CudaBackend>,
+    /// Allocated blocks
+    blocks: Arc<RwLock<Vec<SRAMBlock>>>,
+    /// Statistics
+    stats: Arc<RwLock<SRAMStats>>,
+    /// Next block ID
+    next_block_id: Arc<RwLock<u64>>,
+    /// Total SRAM capacity per block
+    total_capacity: usize,
+}
+
+impl SRAMManager {
+    /// Create a new SRAM manager
+    pub fn new(config: &FlashAttentionConfig, cuda_backend: &Arc<CudaBackend>) -> FlashResult<Self> {
+        info!("Initializing SRAM manager");
+
+        // Get device shared memory info
+        let device_info = cuda_backend.get_device_info()
+        .map_err(|e| FlashError::memory(format!("Failed to get device info: {e}")))?;
+
+        let total_capacity = device_info.shared_memory_per_block as usize;
+
+        // Validate configuration fits in SRAM
+        let required_sram = Self::calculate_required_sram(config)?;
+        if required_sram > total_capacity {
+            return Err(FlashError::config(format!(
+            "Configuration requires {required_sram} bytes of SRAM but only {total_capacity} bytes available"
+            )));
+        }
+
+        let stats = SRAMStats {
+            total_capacity,
+            allocated: 0,
+            peak_usage: 0,
+            allocation_count: 0,
+            efficiency: 0.0,
+            allocation_failures: 0,
+        };
+
+        Ok(Self {
+            config: config.clone(),
+            cuda_backend: cuda_backend.clone(),
+            blocks: Arc::new(RwLock::new(Vec::new())),
+            stats: Arc::new(RwLock::new(stats)),
+            next_block_id: Arc::new(RwLock::new(0)),
+            total_capacity,
+        })
+    }
+
+    /// Calculate required SRAM for the configuration
+    pub fn calculate_required_sram(config: &FlashAttentionConfig) -> FlashResult<usize> {
+        let element_size = match config.precision {
+            crate::config::PrecisionMode::FP32 => 4,
+            crate::config::PrecisionMode::FP16 | crate::config::PrecisionMode::BF16 => 2,
+            crate::config::PrecisionMode::FP8E4M3 { .. } | crate::config::PrecisionMode::FP8E5M2 { .. } => 1,
+            crate::config::PrecisionMode::Mixed { storage_precision, .. } => {
+                match storage_precision {
+                    crate::config::Precision::FP32 => 4,
+                    crate::config::Precision::FP16 | crate::config::Precision::BF16 => 2,
+                    crate::config::Precision::INT8
+                    | crate::config::Precision::FP8E4M3
+                    | crate::config::Precision::FP8E5M2 => 1,
+                }
+            }
+        };
+
+        // Memory for Q block
+        let q_memory = config.block_size_q * config.head_dim * element_size;
+
+        // Memory for K block
+        let k_memory = config.block_size_kv * config.head_dim * element_size;
+
+        // Memory for V block
+        let v_memory = config.block_size_kv * config.head_dim * element_size;
+
+        // Memory for attention scores (FP32 for numerical stability)
+        let scores_memory = config.block_size_q * config.block_size_kv * 4;
+
+        // Scratch memory for intermediate computations
+        let scratch_memory = config.block_size_q * config.head_dim * 4; // FP32 accumulation
+
+        // Add alignment padding (16-byte alignment)
+        let alignment = 16;
+        let total = q_memory + k_memory + v_memory + scores_memory + scratch_memory;
+        let aligned_total = total.div_ceil(alignment) * alignment;
+
+        debug!("SRAM requirements: Q={}, K={}, V={}, scores={}, scratch={}, total={}",
+        q_memory, k_memory, v_memory, scores_memory, scratch_memory, aligned_total);
+
+        Ok(aligned_total)
+    }
+
+    /// Allocate SRAM block
+    pub fn allocate_block(
+    &self,
+    size: usize,
+    block_type: SRAMBlockType,
+    alignment: usize,
+    ) -> FlashResult<SRAMBlock> {
+        let mut blocks = self.blocks.write();
+        let mut stats = self.stats.write();
+
+        // Find suitable offset with alignment
+        let offset = self.find_suitable_offset(&blocks, size, alignment)?;
+
+        // Create new block
+        let block_id = {
+            let mut next_id = self.next_block_id.write();
+            let id = *next_id;
+            *next_id += 1;
+            id
+        };
+
+        let block = SRAMBlock {
+            id: block_id,
+            offset,
+            size,
+            in_use: true,
+            block_type,
+        };
+
+        // Update statistics
+        stats.allocated += size;
+        stats.peak_usage = stats.peak_usage.max(stats.allocated);
+        stats.allocation_count += 1;
+        stats.efficiency = (stats.allocated as f32) / (self.total_capacity as f32);
+
+        blocks.push(block.clone());
+
+        debug!("Allocated SRAM block: id={}, offset={}, size={}, type={:?}",
+        block.id, block.offset, block.size, block.block_type);
+
+        Ok(block)
+    }
+
+    /// Deallocate SRAM block
+    pub fn deallocate_block(&self, block_id: u64) -> FlashResult<()> {
+        let mut blocks = self.blocks.write();
+        let mut stats = self.stats.write();
+
+        if let Some(pos) = blocks.iter().position(|b| b.id == block_id) {
+            let block = blocks.remove(pos);
+            stats.allocated -= block.size;
+            stats.efficiency = (stats.allocated as f32) / (self.total_capacity as f32);
+
+            debug!("Deallocated SRAM block: id={}, size={}", block_id, block.size);
+            Ok(())
+            } else {
+                Err(FlashError::memory(format!("SRAM block {block_id} not found")))
+            }
+        }
+
+        /// Get standard block layout for Flash Attention
+        pub fn get_standard_layout(&self) -> FlashResult<SRAMLayout> {
+            let element_size = match self.config.precision {
+                crate::config::PrecisionMode::FP32 => 4,
+                crate::config::PrecisionMode::FP16 | crate::config::PrecisionMode::BF16 => 2,
+                crate::config::PrecisionMode::FP8E4M3 { .. } | crate::config::PrecisionMode::FP8E5M2 { .. } => 1,
+                crate::config::PrecisionMode::Mixed { storage_precision, .. } => {
+                    match storage_precision {
+                        crate::config::Precision::FP32 => 4,
+                        crate::config::Precision::FP16 | crate::config::Precision::BF16 => 2,
+                        crate::config::Precision::INT8
+                        | crate::config::Precision::FP8E4M3
+                        | crate::config::Precision::FP8E5M2 => 1,
+                    }
+                }
+            };
+
+            let alignment = 16; // 16-byte alignment for vectorization
+            let mut offset = 0;
+
+            // Q block
+            let q_size = self.config.block_size_q * self.config.head_dim * element_size;
+            let q_offset = offset;
+            offset = Self::align_offset(offset + q_size, alignment);
+
+            // K block
+            let k_size = self.config.block_size_kv * self.config.head_dim * element_size;
+            let k_offset = offset;
+            offset = Self::align_offset(offset + k_size, alignment);
+
+            // V block
+            let v_size = self.config.block_size_kv * self.config.head_dim * element_size;
+            let v_offset = offset;
+            offset = Self::align_offset(offset + v_size, alignment);
+
+            // Scores block (FP32)
+            let scores_size = self.config.block_size_q * self.config.block_size_kv * 4;
+            let scores_offset = offset;
+            offset = Self::align_offset(offset + scores_size, alignment);
+
+            // Scratch block (FP32)
+            let scratch_size = self.config.block_size_q * self.config.head_dim * 4;
+            let scratch_offset = offset;
+            offset = Self::align_offset(offset + scratch_size, alignment);
+
+            if offset > self.total_capacity {
+                return Err(FlashError::memory(format!(
+                "SRAM layout requires {} bytes but only {} available",
+                offset, self.total_capacity
+                )));
+            }
+
+            Ok(SRAMLayout {
+                q_block: SRAMBlockInfo { offset: q_offset, size: q_size },
+                k_block: SRAMBlockInfo { offset: k_offset, size: k_size },
+                v_block: SRAMBlockInfo { offset: v_offset, size: v_size },
+                scores_block: SRAMBlockInfo { offset: scores_offset, size: scores_size },
+                scratch_block: SRAMBlockInfo { offset: scratch_offset, size: scratch_size },
+                total_size: offset,
+            })
+        }
+
+        /// Estimate SRAM requirements for a sequence length
+        pub fn estimate_requirements(&self, seq_len: usize) -> FlashResult<usize> {
+            // Calculate how many blocks we need
+            let _q_blocks = seq_len.div_ceil(self.config.block_size_q);
+            let _kv_blocks = seq_len.div_ceil(self.config.block_size_kv);
+
+            // SRAM is reused across blocks, so we only need space for one set of blocks
+            let standard_layout = self.get_standard_layout()?;
+
+            // Add some overhead for runtime management
+            let overhead = standard_layout.total_size / 10; // 10% overhead
+
+            Ok(standard_layout.total_size + overhead)
+        }
+
+        /// Get current SRAM statistics
+        pub fn stats(&self) -> SRAMStats {
+            self.stats.read().clone()
+        }
+
+        /// Get SRAM efficiency (0.0 - 1.0)
+        pub fn get_efficiency(&self) -> f32 {
+            self.stats.read().efficiency
+        }
+
+        /// Optimize SRAM allocation
+        pub fn optimize_allocation(&self) -> FlashResult<()> {
+            debug!("Optimizing SRAM allocation");
+
+            let mut blocks = self.blocks.write();
+
+            // Sort blocks by offset to identify gaps
+            blocks.sort_by_key(|b| b.offset);
+
+            // Defragment by moving blocks to eliminate gaps
+            let mut new_offset = 0;
+            for block in blocks.iter_mut() {
+                if !block.in_use {
+                    continue;
+                }
+
+                if block.offset != new_offset {
+                    debug!("Moving SRAM block {} from offset {} to {}",
+                    block.id, block.offset, new_offset);
+                    block.offset = new_offset;
+                }
+
+                new_offset = Self::align_offset(new_offset + block.size, 16);
+            }
+
+            // Remove unused blocks
+            blocks.retain(|b| b.in_use);
+
+            // Update efficiency
+            let mut stats = self.stats.write();
+            stats.efficiency = (stats.allocated as f32) / (self.total_capacity as f32);
+
+            info!("SRAM optimization complete, efficiency: {:.2}", stats.efficiency);
+            Ok(())
+        }
+
+        /// Calculate shared memory per block for kernel launch
+        pub fn calculate_shared_memory_per_block(&self) -> FlashResult<usize> {
+            let layout = self.get_standard_layout()?;
+            Ok(layout.total_size)
+        }
+
+        /// Find suitable offset for allocation
+        fn find_suitable_offset(
+        &self,
+        blocks: &[SRAMBlock],
+        size: usize,
+        alignment: usize,
+        ) -> FlashResult<usize> {
+            // Sort blocks by offset
+            let mut sorted_blocks: Vec<_> = blocks.iter().filter(|b| b.in_use).collect();
+            sorted_blocks.sort_by_key(|b| b.offset);
+
+            let mut offset = 0;
+
+            // Check gaps between blocks
+            for block in &sorted_blocks {
+                let aligned_offset = Self::align_offset(offset, alignment);
+                if aligned_offset + size <= block.offset {
+                    return Ok(aligned_offset);
+                }
+                offset = block.offset + block.size;
+            }
+
+            // Check if we can fit at the end
+            let aligned_offset = Self::align_offset(offset, alignment);
+            if aligned_offset + size <= self.total_capacity {
+                Ok(aligned_offset)
+                } else {
+                    Err(FlashError::memory(format!(
+                    "Cannot allocate {} bytes of SRAM (available: {})",
+                    size, self.total_capacity - offset
+                    )))
+                }
+            }
+
+            /// Align offset to the specified alignment
+            fn align_offset(offset: usize, alignment: usize) -> usize {
+                offset.div_ceil(alignment) * alignment
+            }
+        }
+
+        /// SRAM layout information
+        #[derive(Debug, Clone)]
+        pub struct SRAMLayout {
+            /// Query block layout
+            pub q_block: SRAMBlockInfo,
+            /// Key block layout
+            pub k_block: SRAMBlockInfo,
+            /// Value block layout
+            pub v_block: SRAMBlockInfo,
+            /// Attention scores block layout
+            pub scores_block: SRAMBlockInfo,
+            /// Scratch space block layout
+            pub scratch_block: SRAMBlockInfo,
+            /// Total SRAM usage
+            pub total_size: usize,
+        }
+
+        /// SRAM block information
+        #[derive(Debug, Clone)]
+        pub struct SRAMBlockInfo {
+            /// Offset within shared memory
+            pub offset: usize,
+            /// Size in bytes
+            pub size: usize,
+        }
+
+        impl SRAMLayout {
+            /// Get the CUDA kernel parameters for this layout
+            pub fn to_kernel_params(&self) -> SRAMKernelParams {
+                SRAMKernelParams {
+                    q_offset: self.q_block.offset as u32,
+                    k_offset: self.k_block.offset as u32,
+                    v_offset: self.v_block.offset as u32,
+                    scores_offset: self.scores_block.offset as u32,
+                    scratch_offset: self.scratch_block.offset as u32,
+                    total_sram: self.total_size as u32,
+                }
+            }
+        }
+
+        /// SRAM kernel parameters for CUDA kernels
+        #[derive(Debug, Clone)]
+        pub struct SRAMKernelParams {
+            pub q_offset: u32,
+            pub k_offset: u32,
+            pub v_offset: u32,
+            pub scores_offset: u32,
+            pub scratch_offset: u32,
+            pub total_sram: u32,
+        }
+
+        #[cfg(test)]
+        mod tests {
+            use super::*;
+            use crate::config::FlashAttentionConfig;
+
+            #[test]
+            fn test_sram_requirements_calculation() {
+                let config = FlashAttentionConfig::new(32, 128);
+                let required = SRAMManager::calculate_required_sram(&config).unwrap();
+
+                // Should be reasonable amount (less than 64KB for most configurations)
+                assert!(required > 0);
+                assert!(required < 65_536);
+            }
+
+            #[test]
+            fn test_sram_layout_generation() {
+                let config = FlashAttentionConfig::new(8, 64);
+
+                if let Ok(cuda_backend) = CudaBackend::new(rtx_runtime::DeviceId(0)) {
+                    let sram_manager = SRAMManager::new(&config, &Arc::new(cuda_backend)).unwrap();
+                    let layout = sram_manager.get_standard_layout().unwrap();
+
+                    assert!(layout.total_size > 0);
+                    assert!(layout.q_block.size > 0);
+                    assert!(layout.k_block.size > 0);
+                    assert!(layout.v_block.size > 0);
+                    assert!(layout.scores_block.size > 0);
+                    assert!(layout.scratch_block.size > 0);
+                }
+            }
+
+            #[test]
+            fn test_alignment() {
+                assert_eq!(SRAMManager::align_offset(10, 16), 16);
+                assert_eq!(SRAMManager::align_offset(16, 16), 16);
+                assert_eq!(SRAMManager::align_offset(17, 16), 32);
+            }
+        }