314 lines
9.4 KiB
Rust
314 lines
9.4 KiB
Rust
#![cfg(feature = "disabled_tests")]
|
|
|
|
use anyhow::Result;
|
|
use rtx_compress::kv_cache::{CompressedKVCache, CompressionMethod, KVCacheConfig};
|
|
use rtx_tensor::{Device, Shape, Tensor};
|
|
|
|
#[test]
|
|
fn test_kv_cache_compression_ratio() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create sample KV cache: [batch, num_heads, seq_len, head_dim]
|
|
let keys = Tensor::randn(&[2, 12, 1024, 64], &device)?;
|
|
let values = Tensor::randn(&[2, 12, 1024, 64], &device)?;
|
|
|
|
let config = KVCacheConfig {
|
|
compression_method: CompressionMethod::ProductQuantization {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
use_opq: true,
|
|
},
|
|
compression_ratio_target: 4.0,
|
|
quality_threshold: 0.95, // 95% quality retention
|
|
max_cache_size_mb: 1024,
|
|
enable_sliding_window: true,
|
|
window_size: 2048,
|
|
enable_attention_scoring: true,
|
|
prefetch_batch_size: 32,
|
|
enable_auto_tuning: false,
|
|
};
|
|
|
|
let mut cache = CompressedKVCache::new(config)?;
|
|
|
|
// Insert KV pairs
|
|
cache.insert(0, &keys, &values)?;
|
|
|
|
// Verify compression metrics
|
|
let stats = cache.compression_stats();
|
|
assert!(
|
|
stats.compression_ratio >= 3.5,
|
|
"Compression ratio should be >= 3.5x, got {}",
|
|
stats.compression_ratio
|
|
);
|
|
assert!(
|
|
stats.quality_score >= 0.90,
|
|
"Quality should be >= 90%, got {}",
|
|
stats.quality_score
|
|
);
|
|
|
|
// Test retrieval
|
|
let (retrieved_keys, retrieved_values) = cache.get(0, 0, 1024)?;
|
|
|
|
// Verify shapes match
|
|
assert_eq!(retrieved_keys.shape(), keys.shape());
|
|
assert_eq!(retrieved_values.shape(), values.shape());
|
|
|
|
// Verify quality
|
|
let key_mse = (&keys - &retrieved_keys)?
|
|
.pow(2.0)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
let value_mse = (&values - &retrieved_values)?
|
|
.pow(2.0)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
assert!(key_mse < 0.01, "Key MSE should be < 0.01, got {}", key_mse);
|
|
assert!(
|
|
value_mse < 0.01,
|
|
"Value MSE should be < 0.01, got {}",
|
|
value_mse
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_kv_cache_incremental_updates() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let config = KVCacheConfig {
|
|
compression_method: CompressionMethod::VectorQuantization {
|
|
codebook_size: 1024,
|
|
update_frequency: 100,
|
|
},
|
|
compression_ratio_target: 3.0,
|
|
quality_threshold: 0.92,
|
|
max_cache_size_mb: 1024,
|
|
enable_sliding_window: false,
|
|
window_size: 1024,
|
|
enable_attention_scoring: false,
|
|
prefetch_batch_size: 16,
|
|
enable_auto_tuning: false,
|
|
};
|
|
|
|
let mut cache = CompressedKVCache::new(config)?;
|
|
|
|
// Start with small sequence
|
|
let keys1 = Tensor::randn(&[1, 8, 64, 32], &device)?;
|
|
let values1 = Tensor::randn(&[1, 8, 64, 32], &device)?;
|
|
|
|
cache.insert(0, &keys1, &values1)?;
|
|
|
|
// Add more tokens incrementally
|
|
for i in 1..10 {
|
|
let new_keys = Tensor::randn(&[1, 8, 1, 32], &device)?;
|
|
let new_values = Tensor::randn(&[1, 8, 1, 32], &device)?;
|
|
|
|
cache.append(0, &new_keys, &new_values)?;
|
|
|
|
// Verify total sequence length
|
|
let (retrieved_k, retrieved_v) = cache.get(0, 0, 64 + i)?;
|
|
assert_eq!(retrieved_k.shape()[2], 64 + i);
|
|
assert_eq!(retrieved_v.shape()[2], 64 + i);
|
|
}
|
|
|
|
// Verify final compression stats
|
|
let stats = cache.compression_stats();
|
|
assert!(stats.compression_ratio >= 2.5);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_kv_cache_memory_management() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let config = KVCacheConfig {
|
|
compression_method: CompressionMethod::MixedPrecision {
|
|
fp16_layers: vec![0, 1, 2],
|
|
int8_layers: vec![3, 4, 5],
|
|
int4_layers: vec![6, 7],
|
|
},
|
|
compression_ratio_target: 6.0,
|
|
quality_threshold: 0.88,
|
|
max_cache_size_mb: 1024,
|
|
enable_sliding_window: false,
|
|
window_size: 1024,
|
|
enable_attention_scoring: false,
|
|
prefetch_batch_size: 16,
|
|
enable_auto_tuning: false,
|
|
};
|
|
|
|
let mut cache = CompressedKVCache::new(config)?;
|
|
cache.set_memory_limit(1_000_000); // 1MB limit
|
|
|
|
// Fill cache beyond memory limit
|
|
for seq_id in 0..100 {
|
|
let keys = Tensor::randn(&[1, 4, 128, 16], &device)?;
|
|
let values = Tensor::randn(&[1, 4, 128, 16], &device)?;
|
|
|
|
cache.insert(seq_id, &keys, &values)?;
|
|
}
|
|
|
|
// Verify memory usage is within limits
|
|
let memory_usage = cache.memory_usage();
|
|
assert!(
|
|
memory_usage <= 1_100_000, // Allow 10% overhead
|
|
"Memory usage {} exceeds limit",
|
|
memory_usage
|
|
);
|
|
|
|
// Verify LRU eviction works
|
|
assert!(cache.contains_sequence(99)); // Most recent should be present
|
|
assert!(!cache.contains_sequence(0)); // Oldest should be evicted
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_kv_cache_batch_operations() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let config = KVCacheConfig {
|
|
compression_method: CompressionMethod::ProductQuantization {
|
|
num_subquantizers: 4,
|
|
codebook_size: 128,
|
|
use_opq: false,
|
|
},
|
|
compression_ratio_target: 4.0,
|
|
quality_threshold: 0.93,
|
|
max_cache_size_mb: 512,
|
|
enable_sliding_window: false,
|
|
window_size: 512,
|
|
enable_attention_scoring: false,
|
|
prefetch_batch_size: 8,
|
|
enable_auto_tuning: false,
|
|
};
|
|
|
|
let mut cache = CompressedKVCache::new(config)?;
|
|
|
|
// Insert batch of sequences
|
|
let batch_keys = Tensor::randn(&[8, 6, 256, 24], &device)?;
|
|
let batch_values = Tensor::randn(&[8, 6, 256, 24], &device)?;
|
|
|
|
cache.insert_batch(&[0, 1, 2, 3, 4, 5, 6, 7], &batch_keys, &batch_values)?;
|
|
|
|
// Retrieve batch
|
|
let (retrieved_keys, retrieved_values) = cache.get_batch(&[0, 2, 4, 6], 0, 256)?;
|
|
|
|
assert_eq!(retrieved_keys.shape().dims(), &[4, 6, 256, 24]);
|
|
assert_eq!(retrieved_values.shape().dims(), &[4, 6, 256, 24]);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_kv_cache_attention_pattern_optimization() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let config = KVCacheConfig {
|
|
compression_method: CompressionMethod::AdaptiveQuantization {
|
|
base_bits: 8,
|
|
attention_threshold: 0.1,
|
|
importance_decay: 0.95,
|
|
},
|
|
compression_ratio_target: 5.0,
|
|
quality_threshold: 0.90,
|
|
max_cache_size_mb: 1024,
|
|
enable_sliding_window: false,
|
|
window_size: 1024,
|
|
enable_attention_scoring: true,
|
|
prefetch_batch_size: 16,
|
|
enable_auto_tuning: false,
|
|
};
|
|
|
|
let mut cache = CompressedKVCache::new(config)?;
|
|
|
|
let keys = Tensor::randn(&[1, 4, 512, 32], &device)?;
|
|
let values = Tensor::randn(&[1, 4, 512, 32], &device)?;
|
|
|
|
// Simulate attention patterns (higher attention = more important)
|
|
let attention_weights = Tensor::randn(&[1, 4, 512], &device)?.softmax(-1)?;
|
|
|
|
cache.insert_with_attention(0, &keys, &values, &attention_weights)?;
|
|
|
|
// Important tokens should have better quality
|
|
let (retrieved_k, retrieved_v) = cache.get(0, 0, 512)?;
|
|
|
|
// Compute per-token reconstruction error
|
|
let key_errors = (keys - retrieved_k)?.pow(2.0)?.mean(&[-1])?;
|
|
let attention_flat = attention_weights.mean(&[1])?; // Average across heads
|
|
|
|
// Verify inverse correlation between attention and error
|
|
let correlation = compute_correlation(&attention_flat, &key_errors)?;
|
|
assert!(
|
|
correlation < -0.1,
|
|
"Higher attention should correlate with lower error"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// Helper function to compute correlation
|
|
fn compute_correlation(x: &Tensor, y: &Tensor) -> Result<f32> {
|
|
let x_mean = x.mean(&[], false)?;
|
|
let y_mean = y.mean(&[], false)?;
|
|
|
|
let x_centered = (x - &x_mean)?;
|
|
let y_centered = (y - &y_mean)?;
|
|
|
|
let numerator = (&x_centered * &y_centered)?.sum(Some(0))?;
|
|
let x_var = x_centered.pow(2.0)?.sum(Some(0))?;
|
|
let y_var = y_centered.pow(2.0)?.sum(Some(0))?;
|
|
|
|
let correlation = &numerator / &(&x_var * &y_var)?.pow(0.5)?;
|
|
correlation.to_scalar::<f32>()
|
|
}
|
|
|
|
#[test]
|
|
fn test_kv_cache_persistence() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let config = KVCacheConfig {
|
|
compression_method: CompressionMethod::ProductQuantization {
|
|
num_subquantizers: 2,
|
|
codebook_size: 64,
|
|
use_opq: false,
|
|
},
|
|
compression_ratio_target: 3.0,
|
|
quality_threshold: 0.90,
|
|
max_cache_size_mb: 256,
|
|
enable_sliding_window: false,
|
|
window_size: 512,
|
|
enable_attention_scoring: false,
|
|
prefetch_batch_size: 4,
|
|
enable_auto_tuning: false,
|
|
};
|
|
|
|
let mut cache = CompressedKVCache::new(config.clone())?;
|
|
|
|
let keys = Tensor::randn(&[1, 2, 128, 16], &device)?;
|
|
let values = Tensor::randn(&[1, 2, 128, 16], &device)?;
|
|
|
|
cache.insert(42, &keys, &values)?;
|
|
|
|
// Save cache
|
|
let saved_data = cache.save_to_bytes()?;
|
|
|
|
// Load into new cache
|
|
let mut cache2 = CompressedKVCache::new(config)?;
|
|
cache2.load_from_bytes(&saved_data)?;
|
|
|
|
// Verify data is preserved
|
|
let (retrieved_k, retrieved_v) = cache2.get(42, 0, 128)?;
|
|
|
|
let key_diff = (&keys - &retrieved_k)?.abs()?.max()?.to_scalar::<f32>()?;
|
|
let value_diff = (&values - &retrieved_v)?.abs()?.max()?.to_scalar::<f32>()?;
|
|
|
|
assert!(key_diff < 0.1, "Persisted keys should match original");
|
|
assert!(value_diff < 0.1, "Persisted values should match original");
|
|
|
|
Ok(())
|
|
}
|