289 lines
11 KiB
Rust
289 lines
11 KiB
Rust
//! Sliding Window Attention Feature Demonstration
|
|
//!
|
|
//! This demonstrates all the key features requested:
|
|
//! 1. Configurable window sizes
|
|
//! 2. Efficient computation for long sequences
|
|
//! 3. Proper boundary handling at sequence start/end
|
|
//! 4. Option for overlapping windows (bidirectional)
|
|
//! 5. Integration with existing attention types (MQA/GQA)
|
|
|
|
use rtx_transformers::layers::{SlidingWindowAttention, SlidingWindowConfig, LayerNorm};
|
|
use rtx_transformers::architectures::TransformerConfig;
|
|
use rtx_transformers::{Result, TransformerError};
|
|
use rtx_tensor::{Tensor, Device, DType};
|
|
use rtx_autograd::{TensorAutograd, backward};
|
|
|
|
fn main() -> Result<()> {
|
|
println!("🚀 Sliding Window Attention Demonstration");
|
|
println!("==========================================");
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// Feature 1: Configurable window sizes
|
|
demonstrate_configurable_window_sizes(&device)?;
|
|
|
|
// Feature 2: Efficient computation for long sequences
|
|
demonstrate_long_sequence_efficiency(&device)?;
|
|
|
|
// Feature 3: Proper boundary handling
|
|
demonstrate_boundary_handling(&device)?;
|
|
|
|
// Feature 4: Overlapping windows (bidirectional vs causal)
|
|
demonstrate_overlapping_windows(&device)?;
|
|
|
|
// Feature 5: Integration with MQA/GQA
|
|
demonstrate_mqa_gqa_integration(&device)?;
|
|
|
|
// Bonus: Full transformer block simulation
|
|
demonstrate_transformer_block_integration(&device)?;
|
|
|
|
println!("\n✅ All sliding window attention features demonstrated successfully!");
|
|
println!("\n📊 Summary:");
|
|
println!(" • Configurable window sizes (16, 64, 256, etc.)");
|
|
println!(" • Efficient O(n*w) complexity instead of O(n²)");
|
|
println!(" • Proper boundary handling for short sequences");
|
|
println!(" • Causal and bidirectional attention modes");
|
|
println!(" • Full compatibility with MQA and GQA");
|
|
println!(" • Integration with existing transformer infrastructure");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn demonstrate_configurable_window_sizes(device: &Device) -> Result<()> {
|
|
println!("\n1. 🎛️ Configurable Window Sizes");
|
|
println!(" Testing different window sizes for efficiency trade-offs");
|
|
|
|
let window_sizes = vec![16, 64, 256];
|
|
let seq_len = 512;
|
|
|
|
for window_size in window_sizes {
|
|
let config = SlidingWindowConfig::new(8, 512, window_size)?;
|
|
let mut attention = SlidingWindowAttention::new(config, device)?;
|
|
attention.initialize_parameters()?;
|
|
|
|
let input = Tensor::ones_typed(&[1, seq_len, 512], DType::F32, device)?;
|
|
let output = attention.forward(&input, None, None)?;
|
|
|
|
let memory_savings = attention.memory_savings_ratio(seq_len);
|
|
let effective_span = attention.effective_attention_span(seq_len);
|
|
|
|
println!(" Window size {:3}: Memory savings {:5.1}%, Effective span {:3}",
|
|
window_size, memory_savings * 100.0, effective_span);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn demonstrate_long_sequence_efficiency(device: &Device) -> Result<()> {
|
|
println!("\n2. ⚡ Long Sequence Efficiency");
|
|
println!(" Comparing efficiency for increasingly long sequences");
|
|
|
|
let config = SlidingWindowConfig::new(12, 768, 128)?; // 128 token window
|
|
let mut attention = SlidingWindowAttention::new(config, device)?;
|
|
attention.initialize_parameters()?;
|
|
|
|
let sequence_lengths = vec![256, 512, 1024, 2048, 4096];
|
|
|
|
for seq_len in sequence_lengths {
|
|
let input = Tensor::ones_typed(&[1, seq_len, 768], DType::F32, device)?;
|
|
|
|
let start = std::time::Instant::now();
|
|
let output = attention.forward(&input, None, None)?;
|
|
let duration = start.elapsed();
|
|
|
|
let memory_savings = attention.memory_savings_ratio(seq_len);
|
|
|
|
println!(" Seq len {:4}: {:6.2}ms, {:5.1}% memory saved, O({}*128) vs O({}²)",
|
|
seq_len,
|
|
duration.as_secs_f64() * 1000.0,
|
|
memory_savings * 100.0,
|
|
seq_len,
|
|
seq_len);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn demonstrate_boundary_handling(device: &Device) -> Result<()> {
|
|
println!("\n3. 🎯 Boundary Handling");
|
|
println!(" Testing behavior at sequence boundaries");
|
|
|
|
let window_size = 64;
|
|
let config = SlidingWindowConfig::new(8, 512, window_size)?;
|
|
let mut attention = SlidingWindowAttention::new(config, device)?;
|
|
attention.initialize_parameters()?;
|
|
|
|
// Test cases: sequence lengths relative to window size
|
|
let test_cases = vec![
|
|
(32, "Short (< window)"),
|
|
(64, "Exact window size"),
|
|
(128, "Long (> window)"),
|
|
];
|
|
|
|
for (seq_len, description) in test_cases {
|
|
let input = Tensor::ones_typed(&[1, seq_len, 512], DType::F32, device)?;
|
|
let output = attention.forward(&input, None, None)?;
|
|
|
|
let effective_span = attention.effective_attention_span(seq_len);
|
|
|
|
println!(" {:18}: seq_len={:3}, effective_span={:3}, window_size={:2}",
|
|
description, seq_len, effective_span, window_size);
|
|
|
|
// Verify output shape is preserved
|
|
assert_eq!(output.shape().dims(), &[1, seq_len, 512]);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn demonstrate_overlapping_windows(device: &Device) -> Result<()> {
|
|
println!("\n4. 🔄 Overlapping Windows (Causal vs Bidirectional)");
|
|
println!(" Comparing causal and bidirectional attention modes");
|
|
|
|
let window_size = 32;
|
|
let seq_len = 128;
|
|
|
|
// Causal (traditional autoregressive)
|
|
let causal_config = SlidingWindowConfig::new(8, 512, window_size)?
|
|
.with_causal(true);
|
|
let mut causal_attention = SlidingWindowAttention::new(causal_config, device)?;
|
|
causal_attention.initialize_parameters()?;
|
|
|
|
// Bidirectional (encoder-style)
|
|
let bidirectional_config = SlidingWindowConfig::new(8, 512, window_size)?
|
|
.with_causal(false);
|
|
let mut bidirectional_attention = SlidingWindowAttention::new(bidirectional_config, device)?;
|
|
bidirectional_attention.initialize_parameters()?;
|
|
|
|
let input = Tensor::ones_typed(&[1, seq_len, 512], DType::F32, device)?;
|
|
|
|
let causal_output = causal_attention.forward(&input, None, None)?;
|
|
let bidirectional_output = bidirectional_attention.forward(&input, None, None)?;
|
|
|
|
let causal_span = causal_attention.effective_attention_span(seq_len);
|
|
let bidirectional_span = bidirectional_attention.effective_attention_span(seq_len);
|
|
|
|
println!(" Causal mode: effective_span={:2} (window_size + 1)", causal_span);
|
|
println!(" Bidirectional mode: effective_span={:2} (2 * window_size + 1)", bidirectional_span);
|
|
|
|
// Verify outputs are different
|
|
let causal_sum = causal_output.sum_all()?.to_scalar::<f32>()?;
|
|
let bidirectional_sum = bidirectional_output.sum_all()?.to_scalar::<f32>()?;
|
|
|
|
println!(" Output sums differ: causal={:.2}, bidirectional={:.2}",
|
|
causal_sum, bidirectional_sum);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn demonstrate_mqa_gqa_integration(device: &Device) -> Result<()> {
|
|
println!("\n5. 🤝 MQA/GQA Integration");
|
|
println!(" Testing integration with Multi-Query and Grouped-Query Attention");
|
|
|
|
let window_size = 64;
|
|
|
|
// Standard Multi-Head Attention (baseline)
|
|
let mha_config = SlidingWindowConfig::new(12, 768, window_size)?;
|
|
let mha_attention = SlidingWindowAttention::new(mha_config, device)?;
|
|
|
|
// Multi-Query Attention (1 KV head)
|
|
let mqa_config = SlidingWindowConfig::new(12, 768, window_size)?
|
|
.with_kv_heads(1)?;
|
|
let mqa_attention = SlidingWindowAttention::new(mqa_config, device)?;
|
|
|
|
// Grouped-Query Attention (4 KV heads)
|
|
let gqa_config = SlidingWindowConfig::new(12, 768, window_size)?
|
|
.with_kv_heads(4)?;
|
|
let gqa_attention = SlidingWindowAttention::new(gqa_config, device)?;
|
|
|
|
println!(" MHA: {} query heads, {} KV heads",
|
|
mha_attention.config().num_heads,
|
|
mha_attention.config().get_num_kv_heads());
|
|
|
|
println!(" MQA: {} query heads, {} KV heads ({}% KV memory vs MHA)",
|
|
mqa_attention.config().num_heads,
|
|
mqa_attention.config().get_num_kv_heads(),
|
|
(100 / 12)); // 1/12 of the KV memory
|
|
|
|
println!(" GQA: {} query heads, {} KV heads ({}% KV memory vs MHA)",
|
|
gqa_attention.config().num_heads,
|
|
gqa_attention.config().get_num_kv_heads(),
|
|
(4 * 100 / 12)); // 4/12 of the KV memory
|
|
|
|
// Test forward passes
|
|
let input = Tensor::ones_typed(&[1, 128, 768], DType::F32, device)?;
|
|
|
|
let mut mha_mut = SlidingWindowAttention::new(mha_config, device)?;
|
|
let mut mqa_mut = SlidingWindowAttention::new(mqa_config, device)?;
|
|
let mut gqa_mut = SlidingWindowAttention::new(gqa_config, device)?;
|
|
|
|
mha_mut.initialize_parameters()?;
|
|
mqa_mut.initialize_parameters()?;
|
|
gqa_mut.initialize_parameters()?;
|
|
|
|
let _mha_output = mha_mut.forward(&input, None, None)?;
|
|
let _mqa_output = mqa_mut.forward(&input, None, None)?;
|
|
let _gqa_output = gqa_mut.forward(&input, None, None)?;
|
|
|
|
println!(" ✓ All attention variants process the same input successfully");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn demonstrate_transformer_block_integration(device: &Device) -> Result<()> {
|
|
println!("\n6. 🏗️ Transformer Block Integration");
|
|
println!(" Simulating full transformer block with sliding window attention");
|
|
|
|
let config = SlidingWindowConfig::new(8, 512, 64)?;
|
|
let mut attention = SlidingWindowAttention::new(config, device)?;
|
|
attention.initialize_parameters()?;
|
|
|
|
// Layer norms for transformer block
|
|
let ln1 = LayerNorm::new(512, 1e-5, true, device)?;
|
|
let ln2 = LayerNorm::new(512, 1e-5, true, device)?;
|
|
|
|
let batch_size = 2;
|
|
let seq_len = 128;
|
|
let d_model = 512;
|
|
|
|
let input = Tensor::ones_typed(&[batch_size, seq_len, d_model], DType::F32, device)?;
|
|
let input_grad = input.set_requires_grad(true);
|
|
|
|
// Transformer block: x = x + attention(ln(x)); x = x + mlp(ln(x))
|
|
|
|
// Attention sublayer
|
|
let norm1 = ln1.forward(&input_grad)?;
|
|
let attn_out = attention.forward(&norm1, None, None)?;
|
|
let residual1 = input_grad.add(&attn_out)?;
|
|
|
|
// MLP sublayer (simplified)
|
|
let norm2 = ln2.forward(&residual1)?;
|
|
let mlp_out = norm2.mul_scalar(1.5)?.add(&norm2.mul_scalar(-0.5)?)?; // Simulate MLP
|
|
let final_output = residual1.add(&mlp_out)?;
|
|
|
|
println!(" Input shape: {:?}", input_grad.shape().dims());
|
|
println!(" Output shape: {:?}", final_output.shape().dims());
|
|
|
|
// Test gradient flow
|
|
let loss = final_output.mean_all()?;
|
|
let gradients = backward(&loss, &[&input_grad])?;
|
|
|
|
let grad_norm = gradients[0].sum_all()?.to_scalar::<f32>()?.abs();
|
|
println!(" Gradient norm: {:.6} (indicating proper gradient flow)", grad_norm);
|
|
|
|
println!(" ✓ Full transformer block with sliding window attention works!");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod demo_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_demo_runs_successfully() {
|
|
// Ensure the demo runs without panics
|
|
let result = main();
|
|
assert!(result.is_ok(), "Demo should run successfully: {:?}", result);
|
|
}
|
|
} |