cleanup(fflag,dead): collapse flash_attention flags to unconditional — [DEAD-004]
FlashAttention3Config had four flags, all dead or with dead else-branches: - use_sparse_patterns: write-only (sparse_pattern Mask is created unconditionally via create_sparse_mask) - io_aware_tiling: always-true setter; the "else" branch called standard_attention which itself discarded all its QK/scale/mask work and called io_aware.compute_attention — pure dead code - cuda_optimization: load_kernels() gate, always true in practice - standard_attention method + mask parameter on forward(): entirely dead Per user directive "all features enabled" / "should be used": - Deleted 4 fields (use_sparse_patterns, io_aware_tiling, cuda_optimization, sparse_pattern_iterations) — note sparse_pattern (BlockSparsePattern) stays - Collapsed forward() to unconditional io_aware.compute_attention, dropped mask param - Removed 40-LOC standard_attention dead fallback - Dropped AttentionStats.io_aware_enabled field + test assertion - cuda_kernels load unconditionally
This commit is contained in:
@@ -183,10 +183,7 @@ pub struct FlashAttention3Config {
|
|||||||
pub head_dim: usize,
|
pub head_dim: usize,
|
||||||
pub max_seq_len: usize,
|
pub max_seq_len: usize,
|
||||||
pub dropout_rate: f32,
|
pub dropout_rate: f32,
|
||||||
pub use_sparse_patterns: bool,
|
|
||||||
pub sparse_pattern: BlockSparsePattern,
|
pub sparse_pattern: BlockSparsePattern,
|
||||||
pub io_aware_tiling: bool,
|
|
||||||
pub cuda_optimization: bool,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for FlashAttention3Config {
|
impl Default for FlashAttention3Config {
|
||||||
@@ -197,10 +194,7 @@ impl Default for FlashAttention3Config {
|
|||||||
head_dim: 64,
|
head_dim: 64,
|
||||||
max_seq_len: 1024,
|
max_seq_len: 1024,
|
||||||
dropout_rate: 0.1,
|
dropout_rate: 0.1,
|
||||||
use_sparse_patterns: true,
|
|
||||||
sparse_pattern: BlockSparsePattern::default(),
|
sparse_pattern: BlockSparsePattern::default(),
|
||||||
io_aware_tiling: true,
|
|
||||||
cuda_optimization: true,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -224,10 +218,7 @@ impl FlashAttention3 {
|
|||||||
let io_aware = IOAwareAttention::new(64, 2048); // 64 tile size, 2GB memory budget
|
let io_aware = IOAwareAttention::new(64, 2048); // 64 tile size, 2GB memory budget
|
||||||
let causal_optimizer = CausalMaskOptimizer::new(1024);
|
let causal_optimizer = CausalMaskOptimizer::new(1024);
|
||||||
let mut cuda_manager = CudaKernelManager::new();
|
let mut cuda_manager = CudaKernelManager::new();
|
||||||
|
cuda_manager.load_kernels()?;
|
||||||
if config.cuda_optimization {
|
|
||||||
cuda_manager.load_kernels()?;
|
|
||||||
}
|
|
||||||
|
|
||||||
let ctx = cudarc::driver::CudaContext::new(0)
|
let ctx = cudarc::driver::CudaContext::new(0)
|
||||||
.map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?;
|
.map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?;
|
||||||
@@ -249,68 +240,18 @@ impl FlashAttention3 {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Compute attention using Flash Attention 3
|
/// Compute attention using Flash Attention 3 (IO-aware tiling, unconditional).
|
||||||
pub fn forward(
|
pub fn forward(
|
||||||
&mut self,
|
&mut self,
|
||||||
q: &GpuTensor,
|
q: &GpuTensor,
|
||||||
k: &GpuTensor,
|
k: &GpuTensor,
|
||||||
v: &GpuTensor,
|
v: &GpuTensor,
|
||||||
mask: Option<&GpuTensor>,
|
|
||||||
) -> Result<GpuTensor, MLError> {
|
) -> Result<GpuTensor, MLError> {
|
||||||
let (_batch_size, _seq_len, _) = q
|
let (_batch_size, _seq_len, _) = q
|
||||||
.dims3()
|
.dims3()
|
||||||
.map_err(|e| MLError::ModelError(format!("Invalid Q tensor dims: {}", e)))?;
|
.map_err(|e| MLError::ModelError(format!("Invalid Q tensor dims: {}", e)))?;
|
||||||
|
|
||||||
// Use IO-aware attention for computation
|
self.io_aware.compute_attention(q, k, v)
|
||||||
let output = if self.config.io_aware_tiling {
|
|
||||||
self.io_aware.compute_attention(q, k, v)?
|
|
||||||
} else {
|
|
||||||
// Fallback to standard attention computation
|
|
||||||
self.standard_attention(q, k, v, mask)?
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok(output)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn standard_attention(
|
|
||||||
&self,
|
|
||||||
q: &GpuTensor,
|
|
||||||
k: &GpuTensor,
|
|
||||||
v: &GpuTensor,
|
|
||||||
mask: Option<&GpuTensor>,
|
|
||||||
) -> Result<GpuTensor, MLError> {
|
|
||||||
// Compute Q @ K^T
|
|
||||||
let k_t = k.transpose(1, 2, &self.stream)
|
|
||||||
.map_err(|e| MLError::ModelError(format!("K transpose failed: {}", e)))?;
|
|
||||||
let scores = q
|
|
||||||
.matmul(&k_t, &self.cublas, &self.stream)
|
|
||||||
.map_err(|e| MLError::ModelError(format!("QK computation failed: {}", e)))?;
|
|
||||||
|
|
||||||
// Scale by sqrt(head_dim): divide by scalar via broadcast_div with a scalar tensor
|
|
||||||
let scale = (self.config.head_dim as f32).sqrt();
|
|
||||||
let scale_tensor = GpuTensor::scalar(scale, &self.stream)
|
|
||||||
.map_err(|e| MLError::ModelError(format!("Scale tensor: {}", e)))?;
|
|
||||||
let scaled_scores = scores
|
|
||||||
.broadcast_div(&scale_tensor, &self.stream)
|
|
||||||
.map_err(|e| MLError::ModelError(format!("Score scaling failed: {}", e)))?;
|
|
||||||
|
|
||||||
// Apply mask if provided
|
|
||||||
let masked_scores = if let Some(mask) = mask {
|
|
||||||
scaled_scores
|
|
||||||
.add(mask, &self.stream)
|
|
||||||
.map_err(|e| MLError::ModelError(format!("Mask application failed: {}", e)))?
|
|
||||||
} else {
|
|
||||||
scaled_scores
|
|
||||||
};
|
|
||||||
|
|
||||||
// Apply softmax via GPU sigmoid approximation (no ActivationKernels::softmax exists)
|
|
||||||
// For attention: use IO-aware path which handles this correctly.
|
|
||||||
// Fallback: just return V weighted by IO-aware attention.
|
|
||||||
let output = self.io_aware.compute_attention(q, k, v)?;
|
|
||||||
|
|
||||||
let _ = masked_scores; // silence unused warning
|
|
||||||
|
|
||||||
Ok(output)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Create sparse attention mask
|
/// Create sparse attention mask
|
||||||
@@ -323,7 +264,6 @@ impl FlashAttention3 {
|
|||||||
AttentionStats {
|
AttentionStats {
|
||||||
cache_size: self.attention_cache.len(),
|
cache_size: self.attention_cache.len(),
|
||||||
cuda_kernels_loaded: self.cuda_manager.kernels_loaded,
|
cuda_kernels_loaded: self.cuda_manager.kernels_loaded,
|
||||||
io_aware_enabled: self.config.io_aware_tiling,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -333,7 +273,6 @@ impl FlashAttention3 {
|
|||||||
pub struct AttentionStats {
|
pub struct AttentionStats {
|
||||||
pub cache_size: usize,
|
pub cache_size: usize,
|
||||||
pub cuda_kernels_loaded: bool,
|
pub cuda_kernels_loaded: bool,
|
||||||
pub io_aware_enabled: bool,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -377,7 +316,7 @@ mod tests {
|
|||||||
let v = GpuTensor::from_host(&v_data, vec![batch_size, seq_len, head_dim], &attention.stream)
|
let v = GpuTensor::from_host(&v_data, vec![batch_size, seq_len, head_dim], &attention.stream)
|
||||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||||
|
|
||||||
let output = attention.forward(&q, &k, &v, None)?;
|
let output = attention.forward(&q, &k, &v)?;
|
||||||
|
|
||||||
// Check output dimensions
|
// Check output dimensions
|
||||||
assert_eq!(output.dims(), q.dims());
|
assert_eq!(output.dims(), q.dims());
|
||||||
@@ -405,7 +344,6 @@ mod tests {
|
|||||||
|
|
||||||
let stats = attention.get_stats();
|
let stats = attention.get_stats();
|
||||||
assert_eq!(stats.cache_size, 0);
|
assert_eq!(stats.cache_size, 0);
|
||||||
assert!(stats.io_aware_enabled);
|
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user