Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,259 @@
|
||||
//! Flash Attention configuration
|
||||
|
||||
/// Configuration for Flash Attention
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FlashAttentionConfig {
|
||||
/// Block size for Q dimension tiling (threads per threadgroup)
|
||||
pub block_q: usize,
|
||||
/// Block size for K/V dimension tiling
|
||||
pub block_kv: usize,
|
||||
/// Maximum head dimension supported
|
||||
pub max_head_dim: usize,
|
||||
/// Softmax scaling factor (typically 1/sqrt(head_dim))
|
||||
/// If None, will be computed automatically from head_dim
|
||||
pub softmax_scale: Option<f32>,
|
||||
/// Enable causal masking (for autoregressive models)
|
||||
pub causal: bool,
|
||||
/// Dropout probability (0.0 = no dropout)
|
||||
pub dropout_p: f32,
|
||||
/// Use float16 (half precision) for Q/K/V tensors
|
||||
/// Accumulation is always done in float32 for numerical stability
|
||||
pub use_f16: bool,
|
||||
}
|
||||
|
||||
impl Default for FlashAttentionConfig {
|
||||
fn default() -> Self {
|
||||
// Apple Silicon GPUs have 32KB threadgroup memory limit
|
||||
// Threadgroup memory includes:
|
||||
// - Q_shared: BLOCK_Q * HEAD_DIM * bytes_per_element
|
||||
// - K_shared: BLOCK_KV * HEAD_DIM * bytes_per_element
|
||||
// - V_shared: BLOCK_KV * HEAD_DIM * bytes_per_element
|
||||
// - Per-thread spill (o_acc array): BLOCK_Q * HEAD_DIM * 4 (always f32)
|
||||
// With block_q=16, block_kv=16, head_dim=64: ~16KB total (f32), ~12KB (f16)
|
||||
Self {
|
||||
block_q: 16,
|
||||
block_kv: 16,
|
||||
max_head_dim: 64,
|
||||
softmax_scale: None,
|
||||
causal: false,
|
||||
dropout_p: 0.0,
|
||||
use_f16: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FlashAttentionConfig {
|
||||
/// Create a new configuration with default settings
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Create a configuration for causal (autoregressive) attention
|
||||
pub fn causal() -> Self {
|
||||
Self {
|
||||
causal: true,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a configuration optimized for large head dimensions (128)
|
||||
/// Uses smaller block sizes to fit within 32KB threadgroup memory
|
||||
pub fn large_head_dim() -> Self {
|
||||
// For head_dim=128 with BLOCK_Q=8, BLOCK_KV=8:
|
||||
// Memory = (8 + 16) * 128 * 4 = 12,288 bytes (~12KB)
|
||||
Self {
|
||||
block_q: 8,
|
||||
block_kv: 8,
|
||||
max_head_dim: 128,
|
||||
softmax_scale: None,
|
||||
causal: false,
|
||||
dropout_p: 0.0,
|
||||
use_f16: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a configuration using float16 (half precision)
|
||||
///
|
||||
/// Uses f16 for Q/K/V storage but f32 accumulation for numerical stability.
|
||||
/// This reduces memory bandwidth and allows larger block sizes.
|
||||
pub fn f16() -> Self {
|
||||
// With f16, we can use larger blocks:
|
||||
// Memory = (16 + 32) * 64 * 2 = 6,144 bytes (~6KB for shared)
|
||||
// Plus f32 accumulators: 16 * 64 * 4 = 4KB
|
||||
Self {
|
||||
block_q: 16,
|
||||
block_kv: 16,
|
||||
max_head_dim: 64,
|
||||
softmax_scale: None,
|
||||
causal: false,
|
||||
dropout_p: 0.0,
|
||||
use_f16: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a causal configuration using float16
|
||||
pub fn f16_causal() -> Self {
|
||||
Self {
|
||||
causal: true,
|
||||
..Self::f16()
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a large head dimension config with float16
|
||||
pub fn f16_large_head_dim() -> Self {
|
||||
// For f16 with head_dim=128:
|
||||
// Memory = (8 + 16) * 128 * 2 = 6,144 bytes (~6KB for shared)
|
||||
Self {
|
||||
block_q: 8,
|
||||
block_kv: 8,
|
||||
max_head_dim: 128,
|
||||
softmax_scale: None,
|
||||
causal: false,
|
||||
dropout_p: 0.0,
|
||||
use_f16: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a causal configuration for large head dimensions (128)
|
||||
pub fn large_head_dim_causal() -> Self {
|
||||
Self {
|
||||
causal: true,
|
||||
..Self::large_head_dim()
|
||||
}
|
||||
}
|
||||
|
||||
/// Set maximum head dimension
|
||||
pub fn with_max_head_dim(mut self, max_head_dim: usize) -> Self {
|
||||
self.max_head_dim = max_head_dim;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set block sizes for Q and K/V tiling
|
||||
pub fn with_block_sizes(mut self, block_q: usize, block_kv: usize) -> Self {
|
||||
self.block_q = block_q;
|
||||
self.block_kv = block_kv;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set explicit softmax scale
|
||||
pub fn with_softmax_scale(mut self, scale: f32) -> Self {
|
||||
self.softmax_scale = Some(scale);
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable or disable causal masking
|
||||
pub fn with_causal(mut self, causal: bool) -> Self {
|
||||
self.causal = causal;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set dropout probability
|
||||
pub fn with_dropout(mut self, p: f32) -> Self {
|
||||
self.dropout_p = p.clamp(0.0, 1.0);
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable or disable float16 (half precision) mode
|
||||
pub fn with_f16(mut self, use_f16: bool) -> Self {
|
||||
self.use_f16 = use_f16;
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the softmax scale, computing from head_dim if not explicitly set
|
||||
pub fn get_softmax_scale(&self, head_dim: usize) -> f32 {
|
||||
self.softmax_scale
|
||||
.unwrap_or_else(|| 1.0 / (head_dim as f32).sqrt())
|
||||
}
|
||||
|
||||
/// Validate configuration against tensor dimensions
|
||||
pub fn validate(&self, head_dim: usize) -> Result<(), String> {
|
||||
if head_dim > self.max_head_dim {
|
||||
return Err(format!(
|
||||
"head_dim {} exceeds max_head_dim {}",
|
||||
head_dim, self.max_head_dim
|
||||
));
|
||||
}
|
||||
if self.block_q == 0 || self.block_kv == 0 {
|
||||
return Err("Block sizes must be non-zero".to_string());
|
||||
}
|
||||
if self.dropout_p < 0.0 || self.dropout_p > 1.0 {
|
||||
return Err("Dropout probability must be in [0, 1]".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_config() {
|
||||
let config = FlashAttentionConfig::default();
|
||||
assert_eq!(config.block_q, 16);
|
||||
assert_eq!(config.block_kv, 16);
|
||||
assert_eq!(config.max_head_dim, 64);
|
||||
assert!(!config.causal);
|
||||
assert_eq!(config.dropout_p, 0.0);
|
||||
assert!(!config.use_f16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_causal_config() {
|
||||
let config = FlashAttentionConfig::causal();
|
||||
assert!(config.causal);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_large_head_dim_config() {
|
||||
let config = FlashAttentionConfig::large_head_dim();
|
||||
assert_eq!(config.block_q, 8);
|
||||
assert_eq!(config.block_kv, 8);
|
||||
assert_eq!(config.max_head_dim, 128);
|
||||
assert!(!config.causal);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_large_head_dim_causal_config() {
|
||||
let config = FlashAttentionConfig::large_head_dim_causal();
|
||||
assert_eq!(config.max_head_dim, 128);
|
||||
assert!(config.causal);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_f16_config() {
|
||||
let config = FlashAttentionConfig::f16();
|
||||
assert!(config.use_f16);
|
||||
assert!(!config.causal);
|
||||
assert_eq!(config.max_head_dim, 64);
|
||||
|
||||
let config_causal = FlashAttentionConfig::f16_causal();
|
||||
assert!(config_causal.use_f16);
|
||||
assert!(config_causal.causal);
|
||||
|
||||
let config_large = FlashAttentionConfig::f16_large_head_dim();
|
||||
assert!(config_large.use_f16);
|
||||
assert_eq!(config_large.max_head_dim, 128);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_softmax_scale() {
|
||||
let config = FlashAttentionConfig::default();
|
||||
let scale = config.get_softmax_scale(64);
|
||||
assert!((scale - 0.125).abs() < 1e-6); // 1/sqrt(64) = 0.125
|
||||
|
||||
let scale_128 = config.get_softmax_scale(128);
|
||||
assert!((scale_128 - 0.08838835).abs() < 1e-6); // 1/sqrt(128)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validation() {
|
||||
let config = FlashAttentionConfig::default();
|
||||
assert!(config.validate(64).is_ok());
|
||||
assert!(config.validate(128).is_err()); // exceeds max_head_dim for default
|
||||
|
||||
let large_config = FlashAttentionConfig::large_head_dim();
|
||||
assert!(large_config.validate(128).is_ok());
|
||||
assert!(large_config.validate(256).is_err());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user