Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,273 @@
|
||||
//! Block memory manager for Flash Attention
|
||||
|
||||
use crate::{
|
||||
config::FlashAttentionConfig,
|
||||
error::{FlashError, FlashResult},
|
||||
};
|
||||
use rtx_runtime::CudaBackend;
|
||||
use std::sync::Arc;
|
||||
use parking_lot::RwLock;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// Memory block for GPU allocation
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MemoryBlock {
|
||||
/// Block ID
|
||||
pub id: u64,
|
||||
/// GPU memory pointer
|
||||
pub ptr: usize,
|
||||
/// Size in bytes
|
||||
pub size: usize,
|
||||
/// Whether the block is currently allocated
|
||||
pub allocated: bool,
|
||||
}
|
||||
|
||||
/// Block manager statistics
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BlockStats {
|
||||
/// Total number of blocks
|
||||
pub total_blocks: usize,
|
||||
/// Number of allocated blocks
|
||||
pub allocated_blocks: usize,
|
||||
/// Total managed memory in bytes
|
||||
pub total_memory: usize,
|
||||
/// Currently allocated memory in bytes
|
||||
pub allocated_memory: usize,
|
||||
/// Fragmentation ratio (0.0 - 1.0)
|
||||
pub fragmentation_ratio: f32,
|
||||
}
|
||||
|
||||
/// Block memory manager
|
||||
pub struct BlockManager {
|
||||
/// Configuration
|
||||
config: FlashAttentionConfig,
|
||||
/// CUDA backend
|
||||
cuda_backend: Arc<CudaBackend>,
|
||||
/// Managed blocks
|
||||
blocks: Arc<RwLock<Vec<MemoryBlock>>>,
|
||||
/// Statistics
|
||||
stats: Arc<RwLock<BlockStats>>,
|
||||
/// Next block ID
|
||||
next_block_id: Arc<RwLock<u64>>,
|
||||
}
|
||||
|
||||
impl BlockManager {
|
||||
/// Create a new block manager
|
||||
pub fn new(config: &FlashAttentionConfig, cuda_backend: &Arc<CudaBackend>) -> FlashResult<Self> {
|
||||
info!("Initializing block manager");
|
||||
|
||||
let stats = BlockStats {
|
||||
total_blocks: 0,
|
||||
allocated_blocks: 0,
|
||||
total_memory: 0,
|
||||
allocated_memory: 0,
|
||||
fragmentation_ratio: 0.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)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Allocate a memory block
|
||||
pub fn allocate_block(&self, size: usize) -> FlashResult<MemoryBlock> {
|
||||
debug!("Allocating memory block of size {}", size);
|
||||
|
||||
// Try to find an existing free block
|
||||
if let Some(block) = self.find_free_block(size)? {
|
||||
self.mark_block_allocated(block.id)?;
|
||||
return Ok(block);
|
||||
}
|
||||
|
||||
// Allocate new block
|
||||
let cuda_slice: cudarc::driver::CudaSlice<u8> = self.cuda_backend.allocate_memory(size)
|
||||
.map_err(|e| FlashError::memory(format!("Failed to allocate GPU memory: {e}")))?;
|
||||
let ptr = &cuda_slice as *const _ as usize;
|
||||
|
||||
let block_id = {
|
||||
let mut next_id = self.next_block_id.write();
|
||||
let id = *next_id;
|
||||
*next_id += 1;
|
||||
id
|
||||
};
|
||||
|
||||
let block = MemoryBlock {
|
||||
id: block_id,
|
||||
ptr,
|
||||
size,
|
||||
allocated: true,
|
||||
};
|
||||
|
||||
// Add to managed blocks
|
||||
{
|
||||
let mut blocks = self.blocks.write();
|
||||
blocks.push(block.clone());
|
||||
}
|
||||
|
||||
// Update statistics
|
||||
{
|
||||
let mut stats = self.stats.write();
|
||||
stats.total_blocks += 1;
|
||||
stats.allocated_blocks += 1;
|
||||
stats.total_memory += size;
|
||||
stats.allocated_memory += size;
|
||||
}
|
||||
|
||||
debug!("Allocated memory block: id={}, ptr=0x{:x}, size={}",
|
||||
block.id, block.ptr, block.size);
|
||||
|
||||
Ok(block)
|
||||
}
|
||||
|
||||
/// Deallocate a memory block
|
||||
pub fn deallocate_block(&self, block_id: u64) -> FlashResult<()> {
|
||||
debug!("Deallocating memory block {}", block_id);
|
||||
|
||||
let mut blocks = self.blocks.write();
|
||||
if let Some(block) = blocks.iter_mut().find(|b| b.id == block_id) {
|
||||
if !block.allocated {
|
||||
return Err(FlashError::memory(format!("Block {block_id} is already deallocated")));
|
||||
}
|
||||
|
||||
block.allocated = false;
|
||||
|
||||
// Update statistics
|
||||
let mut stats = self.stats.write();
|
||||
stats.allocated_blocks -= 1;
|
||||
stats.allocated_memory -= block.size;
|
||||
stats.fragmentation_ratio = self.calculate_fragmentation(&blocks);
|
||||
|
||||
debug!("Deallocated memory block {}", block_id);
|
||||
Ok(())
|
||||
} else {
|
||||
Err(FlashError::memory(format!("Block {block_id} not found")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Get block manager statistics
|
||||
pub fn stats(&self) -> BlockStats {
|
||||
self.stats.read().clone()
|
||||
}
|
||||
|
||||
/// Defragment memory by consolidating free blocks
|
||||
pub fn defragment(&self) -> FlashResult<()> {
|
||||
info!("Starting memory defragmentation");
|
||||
|
||||
let mut blocks = self.blocks.write();
|
||||
|
||||
// Remove deallocated blocks
|
||||
let mut removed_memory = 0;
|
||||
blocks.retain(|block| {
|
||||
if !block.allocated {
|
||||
removed_memory += block.size;
|
||||
// Note: In a real implementation, we would call cuda_backend.deallocate_memory(block.ptr)
|
||||
false
|
||||
} else {
|
||||
true
|
||||
}
|
||||
});
|
||||
|
||||
// Update statistics
|
||||
{
|
||||
let mut stats = self.stats.write();
|
||||
stats.total_blocks = blocks.len();
|
||||
stats.total_memory -= removed_memory;
|
||||
stats.fragmentation_ratio = self.calculate_fragmentation(&blocks);
|
||||
}
|
||||
|
||||
info!("Defragmentation complete: removed {} bytes", removed_memory);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Find a free block of at least the requested size
|
||||
fn find_free_block(&self, size: usize) -> FlashResult<Option<MemoryBlock>> {
|
||||
let blocks = self.blocks.read();
|
||||
for block in blocks.iter() {
|
||||
if !block.allocated && block.size >= size {
|
||||
return Ok(Some(block.clone()));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Mark a block as allocated
|
||||
fn mark_block_allocated(&self, block_id: u64) -> FlashResult<()> {
|
||||
let mut blocks = self.blocks.write();
|
||||
if let Some(block) = blocks.iter_mut().find(|b| b.id == block_id) {
|
||||
block.allocated = true;
|
||||
|
||||
// Update statistics
|
||||
let mut stats = self.stats.write();
|
||||
stats.allocated_blocks += 1;
|
||||
stats.allocated_memory += block.size;
|
||||
|
||||
Ok(())
|
||||
} else {
|
||||
Err(FlashError::memory(format!("Block {block_id} not found")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Calculate memory fragmentation ratio
|
||||
fn calculate_fragmentation(&self, blocks: &[MemoryBlock]) -> f32 {
|
||||
if blocks.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let free_blocks = blocks.iter().filter(|b| !b.allocated).count();
|
||||
let total_blocks = blocks.len();
|
||||
|
||||
if total_blocks == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(free_blocks as f32) / (total_blocks as f32)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::FlashAttentionConfig;
|
||||
|
||||
#[test]
|
||||
fn test_block_manager_creation() {
|
||||
let config = FlashAttentionConfig::new(8, 64);
|
||||
|
||||
if let Ok(cuda_backend) = CudaBackend::new(rtx_runtime::DeviceId(0)) {
|
||||
let block_manager = BlockManager::new(&config, &Arc::new(cuda_backend)).unwrap();
|
||||
let stats = block_manager.stats();
|
||||
|
||||
assert_eq!(stats.total_blocks, 0);
|
||||
assert_eq!(stats.allocated_blocks, 0);
|
||||
assert_eq!(stats.total_memory, 0);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fragmentation_calculation() {
|
||||
let config = FlashAttentionConfig::new(8, 64);
|
||||
|
||||
if let Ok(cuda_backend) = CudaBackend::new(rtx_runtime::DeviceId(0)) {
|
||||
let block_manager = BlockManager::new(&config, &Arc::new(cuda_backend)).unwrap();
|
||||
|
||||
// Test with empty blocks
|
||||
let empty_blocks = vec![];
|
||||
let fragmentation = block_manager.calculate_fragmentation(&empty_blocks);
|
||||
assert_eq!(fragmentation, 0.0);
|
||||
|
||||
// Test with mixed blocks
|
||||
let mixed_blocks = vec![
|
||||
MemoryBlock { id: 1, ptr: 0x1000, size: 1024, allocated: true },
|
||||
MemoryBlock { id: 2, ptr: 0x2000, size: 1024, allocated: false },
|
||||
MemoryBlock { id: 3, ptr: 0x3000, size: 1024, allocated: true },
|
||||
MemoryBlock { id: 4, ptr: 0x4000, size: 1024, allocated: false },
|
||||
];
|
||||
let fragmentation = block_manager.calculate_fragmentation(&mixed_blocks);
|
||||
assert_eq!(fragmentation, 0.5); // 2 free out of 4 total
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user