Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,839 @@
|
||||
//! `SageAttention`: Quantized Attention for Efficient Inference
|
||||
//!
|
||||
//! `SageAttention` achieves 2-5x speedup over `FlashAttention` by using INT8
|
||||
//! quantization for the Q*K^T matrix multiplication while maintaining accuracy
|
||||
//! through careful per-token/per-head scaling.
|
||||
//!
|
||||
//! # Key Features
|
||||
//!
|
||||
//! - INT8 quantization for Q and K matrices
|
||||
//! - Per-token dynamic quantization scales
|
||||
//! - FP16/FP32 precision for softmax and V multiplication
|
||||
//! - Smooth quantization for better accuracy
|
||||
//! - Compatible with GQA/MQA patterns
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! use rtx_transformers::layers::{SageAttention, SageAttentionConfig};
|
||||
//!
|
||||
//! let config = SageAttentionConfig {
|
||||
//! num_heads: 32,
|
||||
//! head_dim: 128,
|
||||
//! use_smooth_quant: true,
|
||||
//! ..Default::default()
|
||||
//! };
|
||||
//!
|
||||
//! let sage_attn = SageAttention::new(config, &device)?;
|
||||
//! let output = sage_attn.forward(&query, &key, &value, mask)?;
|
||||
//! ```
|
||||
//!
|
||||
//! # References
|
||||
//!
|
||||
//! - `SageAttention`: <https://arxiv.org/abs/2410.02367>
|
||||
//! - Smooth Quantization: <https://arxiv.org/abs/2211.10438>
|
||||
|
||||
use crate::{Result, TransformerError};
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
|
||||
/// `SageAttention` configuration
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SageAttentionConfig {
|
||||
/// Number of attention heads
|
||||
pub num_heads: usize,
|
||||
/// Dimension per head
|
||||
pub head_dim: usize,
|
||||
/// Number of key-value heads (for GQA/MQA)
|
||||
pub num_kv_heads: Option<usize>,
|
||||
/// Whether to use causal masking
|
||||
pub causal: bool,
|
||||
/// Dropout probability
|
||||
pub dropout: f32,
|
||||
/// Whether to use smooth quantization (better accuracy)
|
||||
pub use_smooth_quant: bool,
|
||||
/// Smooth quantization alpha (migration strength)
|
||||
pub smooth_alpha: f32,
|
||||
/// Quantization bits for Q and K (default: 8)
|
||||
pub quant_bits: u8,
|
||||
/// Whether to use per-token quantization (vs per-tensor)
|
||||
pub per_token_quant: bool,
|
||||
/// Whether to use symmetric quantization
|
||||
pub symmetric_quant: bool,
|
||||
/// Scale factor for attention scores (default: `1/sqrt(head_dim)`)
|
||||
pub scale: Option<f32>,
|
||||
/// Maximum sequence length for KV cache
|
||||
pub max_seq_len: usize,
|
||||
/// Block size for tiled computation
|
||||
pub block_size: usize,
|
||||
}
|
||||
|
||||
impl Default for SageAttentionConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
num_heads: 32,
|
||||
head_dim: 128,
|
||||
num_kv_heads: None,
|
||||
causal: true,
|
||||
dropout: 0.0,
|
||||
use_smooth_quant: true,
|
||||
smooth_alpha: 0.5,
|
||||
quant_bits: 8,
|
||||
per_token_quant: true,
|
||||
symmetric_quant: true,
|
||||
scale: None,
|
||||
max_seq_len: 8192,
|
||||
block_size: 64,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SageAttentionConfig {
|
||||
/// Create config for a specific model size
|
||||
#[must_use]
|
||||
pub fn for_model(num_heads: usize, head_dim: usize) -> Self {
|
||||
Self {
|
||||
num_heads,
|
||||
head_dim,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Enable GQA with specified number of KV heads
|
||||
#[must_use]
|
||||
pub fn with_gqa(mut self, num_kv_heads: usize) -> Self {
|
||||
self.num_kv_heads = Some(num_kv_heads);
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable MQA (single KV head)
|
||||
#[must_use]
|
||||
pub fn with_mqa(mut self) -> Self {
|
||||
self.num_kv_heads = Some(1);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set causal masking
|
||||
#[must_use]
|
||||
pub fn causal(mut self, causal: bool) -> Self {
|
||||
self.causal = causal;
|
||||
self
|
||||
}
|
||||
|
||||
/// Disable smooth quantization for faster but less accurate inference
|
||||
#[must_use]
|
||||
pub fn fast_mode(mut self) -> Self {
|
||||
self.use_smooth_quant = false;
|
||||
self.per_token_quant = false;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantization parameters for a tensor
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct QuantParams {
|
||||
/// Scale factor (per token or per tensor)
|
||||
pub scales: Vec<f32>,
|
||||
/// Zero points (for asymmetric quantization)
|
||||
pub zero_points: Vec<i32>,
|
||||
/// Whether this is per-token quantization
|
||||
pub per_token: bool,
|
||||
}
|
||||
|
||||
impl QuantParams {
|
||||
/// Create symmetric per-tensor quantization params
|
||||
#[must_use]
|
||||
pub fn symmetric_per_tensor(scale: f32) -> Self {
|
||||
Self {
|
||||
scales: vec![scale],
|
||||
zero_points: vec![0],
|
||||
per_token: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create symmetric per-token quantization params
|
||||
#[must_use]
|
||||
pub fn symmetric_per_token(scales: Vec<f32>) -> Self {
|
||||
let zero_points = vec![0; scales.len()];
|
||||
Self {
|
||||
scales,
|
||||
zero_points,
|
||||
per_token: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `SageAttention` module for efficient quantized attention
|
||||
pub struct SageAttention {
|
||||
config: SageAttentionConfig,
|
||||
device: Device,
|
||||
/// Smooth quantization scaling factors (learned or calibrated)
|
||||
smooth_scales: Option<Tensor>,
|
||||
/// Attention scale factor
|
||||
scale: f32,
|
||||
}
|
||||
|
||||
impl SageAttention {
|
||||
/// Create a new `SageAttention` module
|
||||
pub fn new(config: SageAttentionConfig, device: &Device) -> Result<Self> {
|
||||
let scale = config
|
||||
.scale
|
||||
.unwrap_or(1.0 / (config.head_dim as f32).sqrt());
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
device: device.clone(),
|
||||
smooth_scales: None,
|
||||
scale,
|
||||
})
|
||||
}
|
||||
|
||||
/// Calibrate smooth quantization scales from sample data
|
||||
pub fn calibrate_smooth_scales(&mut self, sample_activations: &[Tensor]) -> Result<()> {
|
||||
if sample_activations.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Compute per-channel max absolute values across samples
|
||||
let mut max_vals = vec![0.0f32; self.config.head_dim];
|
||||
|
||||
for activation in sample_activations {
|
||||
let data = activation
|
||||
.to_vec_f32()
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
for (i, chunk) in data.chunks(self.config.head_dim).enumerate() {
|
||||
for (j, &val) in chunk.iter().enumerate() {
|
||||
let abs_val = val.abs();
|
||||
if abs_val > max_vals[j % self.config.head_dim] {
|
||||
max_vals[j % self.config.head_dim] = abs_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Compute smooth scales: s = max_act^alpha / max_weight^(1-alpha)
|
||||
// For now, use simplified scaling based on activation magnitudes
|
||||
let alpha = self.config.smooth_alpha;
|
||||
let smooth_scales: Vec<f32> = max_vals
|
||||
.iter()
|
||||
.map(|&v| if v > 1e-8 { v.powf(alpha) } else { 1.0 })
|
||||
.collect();
|
||||
|
||||
self.smooth_scales = Some(
|
||||
Tensor::from_vec(smooth_scales, &[self.config.head_dim], &self.device)
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?,
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Quantize a tensor to INT8 with per-token scaling
|
||||
fn quantize_per_token(&self, tensor: &Tensor) -> Result<(Vec<i8>, QuantParams)> {
|
||||
let shape = tensor.shape();
|
||||
let dims = shape.dims();
|
||||
let data = tensor
|
||||
.to_vec_f32()
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
|
||||
let batch_size = dims[0];
|
||||
let seq_len = dims[1];
|
||||
let hidden_dim = dims[2..].iter().product::<usize>();
|
||||
let token_size = hidden_dim;
|
||||
|
||||
let mut quantized = Vec::with_capacity(data.len());
|
||||
let mut scales = Vec::with_capacity(batch_size * seq_len);
|
||||
|
||||
for b in 0..batch_size {
|
||||
for s in 0..seq_len {
|
||||
let offset = (b * seq_len + s) * token_size;
|
||||
let token_data = &data[offset..offset + token_size];
|
||||
|
||||
// Find max absolute value for this token
|
||||
let max_abs = token_data.iter().map(|&x| x.abs()).fold(0.0f32, f32::max);
|
||||
|
||||
let scale = if max_abs > 1e-8 { max_abs / 127.0 } else { 1.0 };
|
||||
|
||||
scales.push(scale);
|
||||
|
||||
// Quantize token values
|
||||
for &val in token_data {
|
||||
let q = (val / scale).round().clamp(-127.0, 127.0) as i8;
|
||||
quantized.push(q);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((quantized, QuantParams::symmetric_per_token(scales)))
|
||||
}
|
||||
|
||||
/// Quantize a tensor to INT8 with per-tensor scaling
|
||||
fn quantize_per_tensor(&self, tensor: &Tensor) -> Result<(Vec<i8>, QuantParams)> {
|
||||
let data = tensor
|
||||
.to_vec_f32()
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
|
||||
// Find global max absolute value
|
||||
let max_abs = data.iter().map(|&x| x.abs()).fold(0.0f32, f32::max);
|
||||
|
||||
let scale = if max_abs > 1e-8 { max_abs / 127.0 } else { 1.0 };
|
||||
|
||||
let quantized: Vec<i8> = data
|
||||
.iter()
|
||||
.map(|&val| (val / scale).round().clamp(-127.0, 127.0) as i8)
|
||||
.collect();
|
||||
|
||||
Ok((quantized, QuantParams::symmetric_per_tensor(scale)))
|
||||
}
|
||||
|
||||
/// Compute INT8 matrix multiplication with dequantization
|
||||
fn quantized_matmul(
|
||||
&self,
|
||||
q_data: &[i8],
|
||||
q_params: &QuantParams,
|
||||
k_data: &[i8],
|
||||
k_params: &QuantParams,
|
||||
m: usize,
|
||||
k: usize,
|
||||
n: usize,
|
||||
) -> Vec<f32> {
|
||||
let mut output = vec![0.0f32; m * n];
|
||||
|
||||
// Perform INT32 accumulation for INT8 * INT8
|
||||
for i in 0..m {
|
||||
for j in 0..n {
|
||||
let mut acc: i32 = 0;
|
||||
for l in 0..k {
|
||||
acc += i32::from(q_data[i * k + l]) * i32::from(k_data[j * k + l]);
|
||||
}
|
||||
|
||||
// Dequantize with appropriate scales
|
||||
let q_scale = if q_params.per_token {
|
||||
q_params.scales[i.min(q_params.scales.len() - 1)]
|
||||
} else {
|
||||
q_params.scales[0]
|
||||
};
|
||||
|
||||
let k_scale = if k_params.per_token {
|
||||
k_params.scales[j.min(k_params.scales.len() - 1)]
|
||||
} else {
|
||||
k_params.scales[0]
|
||||
};
|
||||
|
||||
output[i * n + j] = acc as f32 * q_scale * k_scale * self.scale;
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/// Apply causal mask to attention scores
|
||||
fn apply_causal_mask(&self, scores: &mut [f32], seq_len: usize) {
|
||||
for i in 0..seq_len {
|
||||
for j in (i + 1)..seq_len {
|
||||
scores[i * seq_len + j] = f32::NEG_INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute softmax along the last dimension
|
||||
fn softmax(&self, scores: &mut [f32], seq_len: usize) {
|
||||
for row in scores.chunks_mut(seq_len) {
|
||||
// Find max for numerical stability
|
||||
let max_val = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
|
||||
// Compute exp(x - max) and sum
|
||||
let mut sum = 0.0f32;
|
||||
for val in row.iter_mut() {
|
||||
*val = (*val - max_val).exp();
|
||||
sum += *val;
|
||||
}
|
||||
|
||||
// Normalize
|
||||
for val in row.iter_mut() {
|
||||
*val /= sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward pass with quantized attention
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `query` - Query tensor [batch, `seq_len`, `num_heads`, `head_dim`]
|
||||
/// * `key` - Key tensor [batch, `seq_len`, `num_kv_heads`, `head_dim`]
|
||||
/// * `value` - Value tensor [batch, `seq_len`, `num_kv_heads`, `head_dim`]
|
||||
/// * `attention_mask` - Optional attention mask
|
||||
///
|
||||
/// # Returns
|
||||
/// Output tensor [batch, `seq_len`, `num_heads`, `head_dim`] and attention weights
|
||||
pub fn forward(
|
||||
&self,
|
||||
query: &Tensor,
|
||||
key: &Tensor,
|
||||
value: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
) -> Result<SageAttentionOutput> {
|
||||
let q_shape = query.shape();
|
||||
let batch_size = q_shape[0];
|
||||
let seq_len = q_shape[1];
|
||||
let num_heads = q_shape[2];
|
||||
let head_dim = q_shape[3];
|
||||
|
||||
let k_shape = key.shape();
|
||||
let kv_seq_len = k_shape[1];
|
||||
let num_kv_heads = k_shape[2];
|
||||
|
||||
// Handle GQA/MQA by repeating KV heads
|
||||
let kv_repeat = num_heads / num_kv_heads;
|
||||
|
||||
// Apply smooth quantization if enabled
|
||||
let query_smooth = if self.config.use_smooth_quant {
|
||||
if let Some(ref scales) = self.smooth_scales {
|
||||
query
|
||||
.mul(scales)
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?
|
||||
} else {
|
||||
query.clone()
|
||||
}
|
||||
} else {
|
||||
query.clone()
|
||||
};
|
||||
|
||||
let key_smooth = if self.config.use_smooth_quant {
|
||||
if let Some(ref scales) = self.smooth_scales {
|
||||
key.mul(scales)
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?
|
||||
} else {
|
||||
key.clone()
|
||||
}
|
||||
} else {
|
||||
key.clone()
|
||||
};
|
||||
|
||||
// Quantize Q and K
|
||||
let (q_quant, q_params) = if self.config.per_token_quant {
|
||||
self.quantize_per_token(&query_smooth)?
|
||||
} else {
|
||||
self.quantize_per_tensor(&query_smooth)?
|
||||
};
|
||||
|
||||
let (k_quant, k_params) = if self.config.per_token_quant {
|
||||
self.quantize_per_token(&key_smooth)?
|
||||
} else {
|
||||
self.quantize_per_tensor(&key_smooth)?
|
||||
};
|
||||
|
||||
// Get value data (kept in FP32)
|
||||
let v_data = value
|
||||
.to_vec_f32()
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
|
||||
// Process each batch and head
|
||||
let mut output_data = vec![0.0f32; batch_size * seq_len * num_heads * head_dim];
|
||||
let mut attention_weights_data =
|
||||
vec![0.0f32; batch_size * num_heads * seq_len * kv_seq_len];
|
||||
|
||||
for b in 0..batch_size {
|
||||
for h in 0..num_heads {
|
||||
let kv_h = h / kv_repeat; // Map to KV head
|
||||
|
||||
// Extract Q for this batch/head
|
||||
let mut q_slice: Vec<i8> = Vec::with_capacity(seq_len * head_dim);
|
||||
for s in 0..seq_len {
|
||||
let start = b * seq_len * num_heads * head_dim
|
||||
+ s * num_heads * head_dim
|
||||
+ h * head_dim;
|
||||
for d in 0..head_dim {
|
||||
q_slice.push(q_quant[start + d]);
|
||||
}
|
||||
}
|
||||
|
||||
// Extract K for this batch/kv_head (transposed for matmul)
|
||||
let mut k_slice: Vec<i8> = Vec::with_capacity(kv_seq_len * head_dim);
|
||||
for s in 0..kv_seq_len {
|
||||
let start = b * kv_seq_len * num_kv_heads * head_dim
|
||||
+ s * num_kv_heads * head_dim
|
||||
+ kv_h * head_dim;
|
||||
for d in 0..head_dim {
|
||||
k_slice.push(k_quant[start + d]);
|
||||
}
|
||||
}
|
||||
|
||||
// Compute Q * K^T using INT8 matmul
|
||||
let mut scores = self.quantized_matmul(
|
||||
&q_slice,
|
||||
&QuantParams::symmetric_per_token(
|
||||
(0..seq_len)
|
||||
.map(|s| {
|
||||
if q_params.per_token {
|
||||
q_params.scales
|
||||
[(b * seq_len + s).min(q_params.scales.len() - 1)]
|
||||
} else {
|
||||
q_params.scales[0]
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
&k_slice,
|
||||
&QuantParams::symmetric_per_token(
|
||||
(0..kv_seq_len)
|
||||
.map(|s| {
|
||||
if k_params.per_token {
|
||||
k_params.scales
|
||||
[(b * kv_seq_len + s).min(k_params.scales.len() - 1)]
|
||||
} else {
|
||||
k_params.scales[0]
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
seq_len,
|
||||
head_dim,
|
||||
kv_seq_len,
|
||||
);
|
||||
|
||||
// Apply causal mask if needed
|
||||
if self.config.causal && seq_len == kv_seq_len {
|
||||
self.apply_causal_mask(&mut scores, seq_len);
|
||||
}
|
||||
|
||||
// Apply optional attention mask
|
||||
if let Some(mask) = attention_mask {
|
||||
let mask_data = mask
|
||||
.to_vec_f32()
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
for (i, score) in scores.iter_mut().enumerate() {
|
||||
let mask_idx = i % mask_data.len();
|
||||
if mask_data[mask_idx] < 0.5 {
|
||||
*score = f32::NEG_INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Softmax
|
||||
self.softmax(&mut scores, kv_seq_len);
|
||||
|
||||
// Store attention weights
|
||||
let attn_offset = (b * num_heads + h) * seq_len * kv_seq_len;
|
||||
attention_weights_data[attn_offset..attn_offset + seq_len * kv_seq_len]
|
||||
.copy_from_slice(&scores);
|
||||
|
||||
// Extract V for this batch/kv_head
|
||||
let mut v_slice = Vec::with_capacity(kv_seq_len * head_dim);
|
||||
for s in 0..kv_seq_len {
|
||||
let start = b * kv_seq_len * num_kv_heads * head_dim
|
||||
+ s * num_kv_heads * head_dim
|
||||
+ kv_h * head_dim;
|
||||
for d in 0..head_dim {
|
||||
v_slice.push(v_data[start + d]);
|
||||
}
|
||||
}
|
||||
|
||||
// Compute attention @ V (FP32 matmul)
|
||||
for s in 0..seq_len {
|
||||
for d in 0..head_dim {
|
||||
let mut sum = 0.0f32;
|
||||
for k in 0..kv_seq_len {
|
||||
sum += scores[s * kv_seq_len + k] * v_slice[k * head_dim + d];
|
||||
}
|
||||
let out_idx = b * seq_len * num_heads * head_dim
|
||||
+ s * num_heads * head_dim
|
||||
+ h * head_dim
|
||||
+ d;
|
||||
output_data[out_idx] = sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Create output tensors
|
||||
let output = Tensor::from_vec(
|
||||
output_data,
|
||||
&[batch_size, seq_len, num_heads, head_dim],
|
||||
&self.device,
|
||||
)
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
|
||||
let attention_weights = Tensor::from_vec(
|
||||
attention_weights_data,
|
||||
&[batch_size, num_heads, seq_len, kv_seq_len],
|
||||
&self.device,
|
||||
)
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
|
||||
Ok(SageAttentionOutput {
|
||||
output,
|
||||
attention_weights: Some(attention_weights),
|
||||
quantization_error: self.estimate_quantization_error(query, key)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// Estimate quantization error for monitoring
|
||||
fn estimate_quantization_error(&self, query: &Tensor, key: &Tensor) -> Result<f32> {
|
||||
// Sample a small portion to estimate error
|
||||
let q_data = query
|
||||
.to_vec_f32()
|
||||
.map_err(|e| TransformerError::TensorOp(e.to_string()))?;
|
||||
|
||||
if q_data.is_empty() {
|
||||
return Ok(0.0);
|
||||
}
|
||||
|
||||
// Quantize and dequantize a sample
|
||||
let sample_size = q_data.len().min(1000);
|
||||
let sample = &q_data[..sample_size];
|
||||
|
||||
let max_abs = sample.iter().map(|&x| x.abs()).fold(0.0f32, f32::max);
|
||||
let scale = if max_abs > 1e-8 { max_abs / 127.0 } else { 1.0 };
|
||||
|
||||
let mut error_sum = 0.0f32;
|
||||
for &val in sample {
|
||||
let quantized = (val / scale).round().clamp(-127.0, 127.0) as i8;
|
||||
let dequantized = f32::from(quantized) * scale;
|
||||
error_sum += (val - dequantized).abs();
|
||||
}
|
||||
|
||||
Ok(error_sum / sample_size as f32)
|
||||
}
|
||||
|
||||
/// Get configuration
|
||||
#[must_use]
|
||||
pub fn config(&self) -> &SageAttentionConfig {
|
||||
&self.config
|
||||
}
|
||||
}
|
||||
|
||||
/// Output from `SageAttention` forward pass
|
||||
#[derive(Debug)]
|
||||
pub struct SageAttentionOutput {
|
||||
/// Attention output tensor
|
||||
pub output: Tensor,
|
||||
/// Optional attention weights for visualization
|
||||
pub attention_weights: Option<Tensor>,
|
||||
/// Estimated quantization error
|
||||
pub quantization_error: f32,
|
||||
}
|
||||
|
||||
/// Builder for `SageAttention` with fluent API
|
||||
pub struct SageAttentionBuilder {
|
||||
config: SageAttentionConfig,
|
||||
}
|
||||
|
||||
impl SageAttentionBuilder {
|
||||
/// Create a new builder with default config
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
config: SageAttentionConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set number of heads
|
||||
#[must_use]
|
||||
pub fn num_heads(mut self, num_heads: usize) -> Self {
|
||||
self.config.num_heads = num_heads;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set head dimension
|
||||
#[must_use]
|
||||
pub fn head_dim(mut self, head_dim: usize) -> Self {
|
||||
self.config.head_dim = head_dim;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set number of KV heads for GQA
|
||||
#[must_use]
|
||||
pub fn num_kv_heads(mut self, num_kv_heads: usize) -> Self {
|
||||
self.config.num_kv_heads = Some(num_kv_heads);
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable causal masking
|
||||
#[must_use]
|
||||
pub fn causal(mut self, causal: bool) -> Self {
|
||||
self.config.causal = causal;
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable smooth quantization
|
||||
#[must_use]
|
||||
pub fn smooth_quant(mut self, alpha: f32) -> Self {
|
||||
self.config.use_smooth_quant = true;
|
||||
self.config.smooth_alpha = alpha;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set quantization bits
|
||||
#[must_use]
|
||||
pub fn quant_bits(mut self, bits: u8) -> Self {
|
||||
self.config.quant_bits = bits;
|
||||
self
|
||||
}
|
||||
|
||||
/// Use per-token quantization
|
||||
#[must_use]
|
||||
pub fn per_token(mut self, per_token: bool) -> Self {
|
||||
self.config.per_token_quant = per_token;
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the `SageAttention` module
|
||||
pub fn build(self, device: &Device) -> Result<SageAttention> {
|
||||
SageAttention::new(self.config, device)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SageAttentionBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn create_test_device() -> Device {
|
||||
Device::Cpu
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sage_attention_config_default() {
|
||||
let config = SageAttentionConfig::default();
|
||||
assert_eq!(config.num_heads, 32);
|
||||
assert_eq!(config.head_dim, 128);
|
||||
assert!(config.use_smooth_quant);
|
||||
assert!(config.causal);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sage_attention_config_gqa() {
|
||||
let config = SageAttentionConfig::for_model(32, 128).with_gqa(8);
|
||||
assert_eq!(config.num_kv_heads, Some(8));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sage_attention_config_mqa() {
|
||||
let config = SageAttentionConfig::for_model(32, 128).with_mqa();
|
||||
assert_eq!(config.num_kv_heads, Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sage_attention_creation() {
|
||||
let device = create_test_device();
|
||||
let config = SageAttentionConfig::for_model(4, 64);
|
||||
let sage = SageAttention::new(config, &device);
|
||||
assert!(sage.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sage_attention_forward() {
|
||||
let device = create_test_device();
|
||||
let config = SageAttentionConfig {
|
||||
num_heads: 2,
|
||||
head_dim: 4,
|
||||
causal: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let sage = SageAttention::new(config, &device).unwrap();
|
||||
|
||||
// Create test tensors [batch=1, seq=4, heads=2, head_dim=4]
|
||||
let query = Tensor::randn(&[1, 4, 2, 4], &device).unwrap();
|
||||
let key = Tensor::randn(&[1, 4, 2, 4], &device).unwrap();
|
||||
let value = Tensor::randn(&[1, 4, 2, 4], &device).unwrap();
|
||||
|
||||
let output = sage.forward(&query, &key, &value, None);
|
||||
assert!(output.is_ok());
|
||||
|
||||
let result = output.unwrap();
|
||||
assert_eq!(result.output.shape(), &[1, 4, 2, 4]);
|
||||
assert!(result.quantization_error >= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sage_attention_gqa() {
|
||||
let device = create_test_device();
|
||||
let config = SageAttentionConfig {
|
||||
num_heads: 4,
|
||||
head_dim: 4,
|
||||
num_kv_heads: Some(2), // GQA with 4 heads, 2 KV heads
|
||||
causal: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let sage = SageAttention::new(config, &device).unwrap();
|
||||
|
||||
let query = Tensor::randn(&[1, 4, 4, 4], &device).unwrap();
|
||||
let key = Tensor::randn(&[1, 4, 2, 4], &device).unwrap();
|
||||
let value = Tensor::randn(&[1, 4, 2, 4], &device).unwrap();
|
||||
|
||||
let output = sage.forward(&query, &key, &value, None);
|
||||
assert!(output.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sage_attention_builder() {
|
||||
let device = create_test_device();
|
||||
|
||||
let sage = SageAttentionBuilder::new()
|
||||
.num_heads(8)
|
||||
.head_dim(64)
|
||||
.causal(true)
|
||||
.smooth_quant(0.5)
|
||||
.per_token(true)
|
||||
.build(&device);
|
||||
|
||||
assert!(sage.is_ok());
|
||||
|
||||
let sage = sage.unwrap();
|
||||
assert_eq!(sage.config().num_heads, 8);
|
||||
assert_eq!(sage.config().head_dim, 64);
|
||||
assert!(sage.config().use_smooth_quant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quantization_params() {
|
||||
let params = QuantParams::symmetric_per_tensor(0.1);
|
||||
assert_eq!(params.scales.len(), 1);
|
||||
assert!(!params.per_token);
|
||||
|
||||
let params = QuantParams::symmetric_per_token(vec![0.1, 0.2, 0.3]);
|
||||
assert_eq!(params.scales.len(), 3);
|
||||
assert!(params.per_token);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quantize_per_token() {
|
||||
let device = create_test_device();
|
||||
let config = SageAttentionConfig::for_model(2, 4);
|
||||
let sage = SageAttention::new(config, &device).unwrap();
|
||||
|
||||
let tensor = Tensor::randn(&[1, 2, 4], &device).unwrap();
|
||||
let (quantized, params) = sage.quantize_per_token(&tensor).unwrap();
|
||||
|
||||
assert_eq!(quantized.len(), 8); // 1 * 2 * 4
|
||||
assert_eq!(params.scales.len(), 2); // 1 * 2 tokens
|
||||
assert!(params.per_token);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quantize_per_tensor() {
|
||||
let device = create_test_device();
|
||||
let config = SageAttentionConfig::for_model(2, 4);
|
||||
let sage = SageAttention::new(config, &device).unwrap();
|
||||
|
||||
let tensor = Tensor::randn(&[1, 2, 4], &device).unwrap();
|
||||
let (quantized, params) = sage.quantize_per_tensor(&tensor).unwrap();
|
||||
|
||||
assert_eq!(quantized.len(), 8);
|
||||
assert_eq!(params.scales.len(), 1);
|
||||
assert!(!params.per_token);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fast_mode_config() {
|
||||
let config = SageAttentionConfig::for_model(32, 128).fast_mode();
|
||||
assert!(!config.use_smooth_quant);
|
||||
assert!(!config.per_token_quant);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user