Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,306 @@
|
||||
//! Benchmarks demonstrating GQA performance characteristics
|
||||
//!
|
||||
//! This module provides benchmarks comparing GQA with MQA and theoretical MHA
|
||||
//! across different dimensions: memory usage, computation efficiency, and
|
||||
//! gradient computation overhead.
|
||||
|
||||
#[cfg(all(test, feature = "disabled_tests"))]
|
||||
mod benchmarks {
|
||||
use crate::layers::{GroupedQueryAttention, GQAConfig, MultiQueryAttention, MQAConfig};
|
||||
use crate::architectures::TransformerConfig;
|
||||
use rtx_tensor::{Tensor, Device, DType};
|
||||
use rtx_autograd::backward;
|
||||
use std::time::Instant;
|
||||
|
||||
#[test]
|
||||
fn benchmark_gqa_vs_mqa_memory() {
|
||||
// Compare memory usage across different attention mechanisms
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let seq_lengths = vec![512, 1024, 2048, 4096];
|
||||
let batch_size = 4;
|
||||
|
||||
println!("\n=== Memory Usage Comparison (bytes) ===");
|
||||
println!("{:>8} | {:>12} | {:>12} | {:>12} | {:>8}",
|
||||
"Seq Len", "MHA (est.)", "GQA (4grp)", "MQA", "GQA/MHA");
|
||||
println!("{:-<70}", "");
|
||||
|
||||
for seq_len in seq_lengths {
|
||||
// Simulate MHA memory (16 heads)
|
||||
let base_config = TransformerConfig::new(50257, 1024, 12, 16, 4096, 2048);
|
||||
let mha_memory = base_config.estimate_kv_cache_memory(seq_len) * batch_size;
|
||||
|
||||
// GQA with 4 groups
|
||||
let mut gqa_config = base_config.clone();
|
||||
gqa_config.set_mqa(false, 4).unwrap();
|
||||
let gqa_attention_config = GQAConfig::from_transformer_config(&gqa_config).unwrap();
|
||||
let gqa_layer = GroupedQueryAttention::new(gqa_attention_config, &device).unwrap();
|
||||
let gqa_memory = gqa_layer.compute_kv_cache_memory(seq_len, batch_size);
|
||||
|
||||
// MQA with 1 group
|
||||
let mut mqa_config = base_config.clone();
|
||||
mqa_config.set_mqa(true, 1).unwrap();
|
||||
let mqa_attention_config = MQAConfig::from_transformer_config(&mqa_config).unwrap();
|
||||
let mqa_layer = MultiQueryAttention::new(mqa_attention_config, &device).unwrap();
|
||||
let mqa_memory = mqa_layer.compute_kv_cache_memory(seq_len, batch_size);
|
||||
|
||||
let gqa_mha_ratio = gqa_memory as f64 / mha_memory as f64;
|
||||
|
||||
println!("{:>8} | {:>12} | {:>12} | {:>12} | {:>7.2}x",
|
||||
seq_len, mha_memory, gqa_memory, mqa_memory, 1.0/gqa_mha_ratio);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn benchmark_gqa_forward_pass_time() {
|
||||
// Benchmark forward pass performance across different configurations
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 2;
|
||||
let seq_len = 512;
|
||||
let d_model = 1024;
|
||||
|
||||
let configurations = vec![
|
||||
(8, "Conservative GQA"),
|
||||
(4, "Balanced GQA"),
|
||||
(2, "Aggressive GQA"),
|
||||
(1, "MQA"),
|
||||
];
|
||||
|
||||
println!("\n=== Forward Pass Performance ===");
|
||||
println!("{:>15} | {:>8} | {:>12} | {:>12}", "Configuration", "Groups", "Time (μs)", "Rel. Speed");
|
||||
println!("{:-<55}", "");
|
||||
|
||||
let mut baseline_time = None;
|
||||
|
||||
for (num_groups, config_name) in configurations {
|
||||
let input_data = vec![0.1f32; batch_size * seq_len * d_model];
|
||||
let hidden_states = Tensor::from_vec(
|
||||
input_data,
|
||||
&[batch_size, seq_len, d_model],
|
||||
&device
|
||||
).unwrap();
|
||||
|
||||
if num_groups == 1 {
|
||||
// MQA
|
||||
let mut transformer_config = TransformerConfig::new(50257, d_model, 12, 16, 4096, 2048);
|
||||
transformer_config.set_mqa(true, 1).unwrap();
|
||||
let mqa_config = MQAConfig::from_transformer_config(&transformer_config).unwrap();
|
||||
let mut mqa_layer = MultiQueryAttention::new(mqa_config, &device).unwrap();
|
||||
mqa_layer.initialize_parameters().unwrap();
|
||||
|
||||
let start = Instant::now();
|
||||
for _ in 0..100 {
|
||||
let _ = mqa_layer.forward(&hidden_states, None, None).unwrap();
|
||||
}
|
||||
let elapsed = start.elapsed();
|
||||
let time_per_call = elapsed.as_micros() / 100;
|
||||
|
||||
if baseline_time.is_none() {
|
||||
baseline_time = Some(time_per_call);
|
||||
}
|
||||
|
||||
let relative_speed = baseline_time.unwrap() as f64 / time_per_call as f64;
|
||||
println!("{:>15} | {:>8} | {:>12} | {:>11.2}x",
|
||||
config_name, num_groups, time_per_call, relative_speed);
|
||||
} else {
|
||||
// GQA
|
||||
let mut transformer_config = TransformerConfig::new(50257, d_model, 12, 16, 4096, 2048);
|
||||
transformer_config.set_mqa(false, num_groups).unwrap();
|
||||
let gqa_config = GQAConfig::from_transformer_config(&transformer_config).unwrap();
|
||||
let mut gqa_layer = GroupedQueryAttention::new(gqa_config, &device).unwrap();
|
||||
gqa_layer.initialize_parameters().unwrap();
|
||||
|
||||
let start = Instant::now();
|
||||
for _ in 0..100 {
|
||||
let _ = gqa_layer.forward(&hidden_states, None, None).unwrap();
|
||||
}
|
||||
let elapsed = start.elapsed();
|
||||
let time_per_call = elapsed.as_micros() / 100;
|
||||
|
||||
if baseline_time.is_none() {
|
||||
baseline_time = Some(time_per_call);
|
||||
}
|
||||
|
||||
let relative_speed = baseline_time.unwrap() as f64 / time_per_call as f64;
|
||||
println!("{:>15} | {:>8} | {:>12} | {:>11.2}x",
|
||||
config_name, num_groups, time_per_call, relative_speed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn benchmark_gqa_scaling_characteristics() {
|
||||
// Test how GQA performance scales with different parameters
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
|
||||
println!("\n=== GQA Scaling Characteristics ===");
|
||||
|
||||
// Test 1: Scaling with number of heads
|
||||
println!("\n1. Scaling with Number of Heads (fixed 4:1 ratio):");
|
||||
println!("{:>10} | {:>10} | {:>12} | {:>12} | {:>10}",
|
||||
"Heads", "Groups", "Memory (KB)", "FLOP Red.", "Efficiency");
|
||||
println!("{:-<65}", "");
|
||||
|
||||
let head_counts = vec![8, 12, 16, 24, 32];
|
||||
for num_heads in head_counts {
|
||||
let num_groups = num_heads / 4; // 4:1 ratio
|
||||
if num_groups == 0 { continue; }
|
||||
|
||||
let mut config = TransformerConfig::new(50257, 1024, 12, num_heads, 4096, 2048);
|
||||
config.set_mqa(false, num_groups).unwrap();
|
||||
|
||||
let gqa_config = GQAConfig::from_transformer_config(&config).unwrap();
|
||||
let gqa = GroupedQueryAttention::new(gqa_config, &device).unwrap();
|
||||
|
||||
let memory = gqa.compute_kv_cache_memory(1024, 1) / 1024; // KB
|
||||
let efficiency = gqa.efficiency_metrics();
|
||||
let flops = gqa.compute_flops_reduction(1024);
|
||||
|
||||
println!("{:>10} | {:>10} | {:>12} | {:>11.2}x | {:>9.2}x",
|
||||
num_heads, num_groups, memory, flops.flops_reduction,
|
||||
efficiency.memory_vs_mha);
|
||||
}
|
||||
|
||||
// Test 2: Scaling with sequence length
|
||||
println!("\n2. Scaling with Sequence Length (16 heads, 4 groups):");
|
||||
println!("{:>8} | {:>12} | {:>12} | {:>15}",
|
||||
"Seq Len", "Memory (MB)", "Compute Intensity", "Flash Memory");
|
||||
println!("{:-<55}", "");
|
||||
|
||||
let mut config = TransformerConfig::new(50257, 1024, 12, 16, 4096, 2048);
|
||||
config.set_mqa(false, 4).unwrap();
|
||||
let gqa_config = GQAConfig::from_transformer_config(&config).unwrap();
|
||||
let gqa = GroupedQueryAttention::new(gqa_config, &device).unwrap();
|
||||
|
||||
let seq_lengths = vec![512, 1024, 2048, 4096];
|
||||
for seq_len in seq_lengths {
|
||||
let memory = gqa.compute_kv_cache_memory(seq_len, 1) as f64 / (1024.0 * 1024.0); // MB
|
||||
let flops = gqa.compute_flops_reduction(seq_len);
|
||||
let flash_memory = gqa.estimate_flash_attention_memory(seq_len, 1) as f64 / (1024.0 * 1024.0); // MB
|
||||
|
||||
println!("{:>8} | {:>11.2} | {:>15.2} | {:>14.2}",
|
||||
seq_len, memory, flops.compute_intensity, flash_memory);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn benchmark_gqa_gradient_overhead() {
|
||||
// Measure gradient computation overhead for GQA vs MQA
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 1;
|
||||
let seq_len = 256;
|
||||
let d_model = 768;
|
||||
|
||||
println!("\n=== Gradient Computation Overhead ===");
|
||||
println!("{:>15} | {:>12} | {:>12} | {:>12}",
|
||||
"Configuration", "Forward (μs)", "Backward (μs)", "Total (μs)");
|
||||
println!("{:-<60}", "");
|
||||
|
||||
let configs = vec![
|
||||
(4, "GQA (3:1 ratio)"),
|
||||
(2, "GQA (6:1 ratio)"),
|
||||
(1, "MQA"),
|
||||
];
|
||||
|
||||
for (num_groups, config_name) in configs {
|
||||
let input_data = vec![0.5f32; batch_size * seq_len * d_model];
|
||||
let hidden_states = Tensor::from_vec(
|
||||
input_data,
|
||||
&[batch_size, seq_len, d_model],
|
||||
&device
|
||||
).unwrap();
|
||||
|
||||
if num_groups == 1 {
|
||||
// MQA
|
||||
let mut transformer_config = TransformerConfig::new(50257, d_model, 12, 12, 3072, 1024);
|
||||
transformer_config.set_mqa(true, 1).unwrap();
|
||||
let mqa_config = MQAConfig::from_transformer_config(&transformer_config).unwrap();
|
||||
let mut mqa_layer = MultiQueryAttention::new(mqa_config, &device).unwrap();
|
||||
mqa_layer.initialize_parameters().unwrap();
|
||||
|
||||
// Measure forward pass
|
||||
let hidden_states_grad = hidden_states.set_requires_grad(true);
|
||||
let start = Instant::now();
|
||||
let output = mqa_layer.forward(&hidden_states_grad, None, None).unwrap();
|
||||
let forward_time = start.elapsed().as_micros();
|
||||
|
||||
// Measure backward pass
|
||||
let loss = output.sum(None).unwrap();
|
||||
let start = Instant::now();
|
||||
let _ = backward(&loss, &[&hidden_states_grad]).unwrap();
|
||||
let backward_time = start.elapsed().as_micros();
|
||||
|
||||
println!("{:>15} | {:>12} | {:>12} | {:>12}",
|
||||
config_name, forward_time, backward_time, forward_time + backward_time);
|
||||
} else {
|
||||
// GQA
|
||||
let mut transformer_config = TransformerConfig::new(50257, d_model, 12, 12, 3072, 1024);
|
||||
transformer_config.set_mqa(false, num_groups).unwrap();
|
||||
let gqa_config = GQAConfig::from_transformer_config(&transformer_config).unwrap();
|
||||
let mut gqa_layer = GroupedQueryAttention::new(gqa_config, &device).unwrap();
|
||||
gqa_layer.initialize_parameters().unwrap();
|
||||
|
||||
// Measure forward pass
|
||||
let hidden_states_grad = hidden_states.set_requires_grad(true);
|
||||
let start = Instant::now();
|
||||
let output = gqa_layer.forward(&hidden_states_grad, None, None).unwrap();
|
||||
let forward_time = start.elapsed().as_micros();
|
||||
|
||||
// Measure backward pass
|
||||
let loss = output.sum(None).unwrap();
|
||||
let start = Instant::now();
|
||||
let _ = backward(&loss, &[&hidden_states_grad]).unwrap();
|
||||
let backward_time = start.elapsed().as_micros();
|
||||
|
||||
println!("{:>15} | {:>12} | {:>12} | {:>12}",
|
||||
config_name, forward_time, backward_time, forward_time + backward_time);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn benchmark_gqa_flash_attention_benefits() {
|
||||
// Demonstrate the benefits of GQA for Flash Attention scenarios
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
|
||||
println!("\n=== Flash Attention Benefits with GQA ===");
|
||||
|
||||
let model_sizes = vec![
|
||||
("GPT-2 Small", 768, 12),
|
||||
("GPT-2 Medium", 1024, 16),
|
||||
("GPT-2 Large", 1280, 20),
|
||||
];
|
||||
|
||||
for (model_name, d_model, num_heads) in model_sizes {
|
||||
println!("\n{}:", model_name);
|
||||
println!("{:>12} | {:>10} | {:>12} | {:>15} | {:>12}",
|
||||
"Strategy", "Groups", "Block Q/KV", "Memory (MB)", "Group Repl");
|
||||
println!("{:-<70}", "");
|
||||
|
||||
let grouping_strategies = vec![
|
||||
(num_heads / 2, "Conservative"),
|
||||
(num_heads / 4, "Balanced"),
|
||||
(num_heads / 8, "Aggressive"),
|
||||
];
|
||||
|
||||
for (num_groups, strategy) in grouping_strategies {
|
||||
if num_groups == 0 { continue; }
|
||||
|
||||
let mut config = TransformerConfig::new(50257, d_model, 12, num_heads, d_model * 4, 2048);
|
||||
config.set_mqa(false, num_groups).unwrap();
|
||||
|
||||
let gqa_config = GQAConfig::from_transformer_config(&config).unwrap();
|
||||
let gqa = GroupedQueryAttention::new(gqa_config, &device).unwrap();
|
||||
|
||||
let flash_config = gqa.get_flash_attention_config().unwrap();
|
||||
let (block_q, block_kv) = gqa.get_optimal_flash_block_sizes();
|
||||
let flash_memory = gqa.estimate_flash_attention_memory(2048, 4) as f64 / (1024.0 * 1024.0);
|
||||
|
||||
println!("{:>12} | {:>10} | {:>6}/{:<5} | {:>14.2} | {:>12}",
|
||||
strategy, num_groups, block_q, block_kv, flash_memory,
|
||||
flash_config.group_replication_factor);
|
||||
}
|
||||
}
|
||||
|
||||
println!("\nNote: Lower memory usage and optimal block sizes indicate better Flash Attention compatibility");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user