1126 lines
39 KiB
Rust
1126 lines
39 KiB
Rust
//! KV Cache compression and optimization for transformer models
|
|
//!
|
|
//! This module provides efficient compression and management of Key-Value caches
|
|
//! used in transformer-based language models. It implements various compression
|
|
//! techniques specifically designed for the access patterns of KV caches.
|
|
//!
|
|
//! # Features
|
|
//!
|
|
//! - **Product quantization** with learned codebooks for high compression
|
|
//! - **Mixed-precision adaptive quantization** based on layer importance
|
|
//! - **Attention-aware importance scoring** for intelligent eviction
|
|
//! - **Sliding window and token importance-based eviction**
|
|
//! - **Batch processing and prefetching** for efficiency
|
|
//! - **Memory-aware auto-tuning** for optimal configuration
|
|
//!
|
|
//! # Compression Methods
|
|
//!
|
|
//! - `ProductQuantization`: High compression using PQ with learned codebooks
|
|
//! - `VectorQuantization`: Global codebook for balanced compression
|
|
//! - `MixedPrecision`: Layer-specific bit-width selection
|
|
//! - `AdaptiveQuantization`: Attention-based dynamic precision
|
|
//! - `HierarchicalCompression`: Multi-resolution compression
|
|
|
|
use crate::error::{CompressionError, Result};
|
|
use crate::quantization::product_quantization::{PQConfig, ProductQuantizer};
|
|
use crate::quantization::vector_quantization::{DistanceMetric, VQConfig, VectorQuantizer};
|
|
use rtx_tensor::{Device, Tensor};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::{HashMap, VecDeque};
|
|
|
|
/// Compression methods for KV cache
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum CompressionMethod {
|
|
/// Product quantization with learned codebooks
|
|
ProductQuantization {
|
|
num_subquantizers: usize,
|
|
codebook_size: usize,
|
|
use_opq: bool,
|
|
},
|
|
/// Vector quantization with global codebook
|
|
VectorQuantization {
|
|
codebook_size: usize,
|
|
update_frequency: usize,
|
|
},
|
|
/// Mixed precision based on layer importance
|
|
MixedPrecision {
|
|
fp16_layers: Vec<usize>,
|
|
int8_layers: Vec<usize>,
|
|
int4_layers: Vec<usize>,
|
|
},
|
|
/// Adaptive quantization based on attention patterns
|
|
AdaptiveQuantization {
|
|
base_bits: u8,
|
|
attention_threshold: f32,
|
|
importance_decay: f32,
|
|
},
|
|
/// Hierarchical compression with different resolutions
|
|
HierarchicalCompression {
|
|
levels: usize,
|
|
compression_ratios: Vec<f32>,
|
|
},
|
|
}
|
|
|
|
impl Default for CompressionMethod {
|
|
fn default() -> Self {
|
|
Self::ProductQuantization {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
use_opq: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration for KV cache compression
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct KVCacheConfig {
|
|
/// Primary compression method
|
|
pub compression_method: CompressionMethod,
|
|
/// Target compression ratio
|
|
pub compression_ratio_target: f64,
|
|
/// Quality threshold for lossy compression
|
|
pub quality_threshold: f64,
|
|
/// Maximum cache size in MB
|
|
pub max_cache_size_mb: usize,
|
|
/// Enable sliding window eviction
|
|
pub enable_sliding_window: bool,
|
|
/// Sliding window size
|
|
pub window_size: usize,
|
|
/// Enable attention-based importance scoring
|
|
pub enable_attention_scoring: bool,
|
|
/// Prefetch batch size for async operations
|
|
pub prefetch_batch_size: usize,
|
|
/// Enable auto-tuning of compression parameters
|
|
pub enable_auto_tuning: bool,
|
|
}
|
|
|
|
impl Default for KVCacheConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
compression_method: CompressionMethod::default(),
|
|
compression_ratio_target: 4.0,
|
|
quality_threshold: 0.95,
|
|
max_cache_size_mb: 1024,
|
|
enable_sliding_window: true,
|
|
window_size: 2048,
|
|
enable_attention_scoring: true,
|
|
prefetch_batch_size: 32,
|
|
enable_auto_tuning: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Statistics for compression performance
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct CompressionStats {
|
|
pub compression_ratio: f64,
|
|
pub quality_score: f64,
|
|
pub cache_hit_rate: f64,
|
|
pub eviction_count: usize,
|
|
pub memory_usage_mb: f64,
|
|
pub avg_compression_time_ms: f64,
|
|
pub avg_decompression_time_ms: f64,
|
|
}
|
|
|
|
impl Default for CompressionStats {
|
|
fn default() -> Self {
|
|
Self {
|
|
compression_ratio: 1.0,
|
|
quality_score: 1.0,
|
|
cache_hit_rate: 0.0,
|
|
eviction_count: 0,
|
|
memory_usage_mb: 0.0,
|
|
avg_compression_time_ms: 0.0,
|
|
avg_decompression_time_ms: 0.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Token importance information for adaptive compression
|
|
#[derive(Debug, Clone)]
|
|
pub struct TokenImportance {
|
|
pub importance_score: f32,
|
|
pub attention_weight: f32,
|
|
pub access_count: usize,
|
|
pub last_access_time: u64,
|
|
pub compression_level: u8,
|
|
}
|
|
|
|
impl Default for TokenImportance {
|
|
fn default() -> Self {
|
|
Self {
|
|
importance_score: 1.0,
|
|
attention_weight: 1.0,
|
|
access_count: 0,
|
|
last_access_time: 0,
|
|
compression_level: 8, // Default to INT8
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compressed KV cache entry
|
|
#[derive(Debug, Clone)]
|
|
struct CacheEntry {
|
|
compressed_keys: Vec<u8>,
|
|
compressed_values: Vec<u8>,
|
|
original_shape: Vec<usize>,
|
|
compression_metadata: CompressionMetadata,
|
|
importance: TokenImportance,
|
|
sequence_id: usize,
|
|
token_range: (usize, usize),
|
|
}
|
|
|
|
/// Metadata for compression/decompression
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
struct CompressionMetadata {
|
|
method: CompressionMethod,
|
|
compression_ratio: f32,
|
|
quality_score: f32,
|
|
codebook_indices: Option<Vec<usize>>,
|
|
quantization_params: Option<QuantizationParams>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
struct QuantizationParams {
|
|
scales: Vec<f32>,
|
|
zero_points: Vec<i32>,
|
|
bit_width: u8,
|
|
}
|
|
|
|
/// High-performance compressed KV cache implementation
|
|
pub struct CompressedKVCache {
|
|
config: KVCacheConfig,
|
|
cache_entries: HashMap<String, CacheEntry>,
|
|
sequence_metadata: HashMap<usize, SequenceMetadata>,
|
|
importance_tracker: ImportanceTracker,
|
|
device: Device,
|
|
stats: CompressionStats,
|
|
memory_usage: usize,
|
|
global_timestamp: u64,
|
|
/// Product quantizer for PQ-based compression
|
|
pq_key: Option<ProductQuantizer>,
|
|
/// Product quantizer for values
|
|
pq_value: Option<ProductQuantizer>,
|
|
/// Vector quantizer for VQ-based compression
|
|
vq_key: Option<VectorQuantizer>,
|
|
/// Vector quantizer for values
|
|
vq_value: Option<VectorQuantizer>,
|
|
/// Whether quantizers are trained
|
|
quantizers_trained: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct SequenceMetadata {
|
|
sequence_length: usize,
|
|
layer_count: usize,
|
|
head_count: usize,
|
|
head_dim: usize,
|
|
attention_patterns: Vec<f32>, // Cached attention importance scores
|
|
}
|
|
|
|
/// Tracks token importance across sequences
|
|
#[derive(Debug)]
|
|
struct ImportanceTracker {
|
|
token_scores: HashMap<String, TokenImportance>,
|
|
attention_history: VecDeque<AttentionSnapshot>,
|
|
decay_factor: f32,
|
|
update_frequency: usize,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct AttentionSnapshot {
|
|
timestamp: u64,
|
|
sequence_id: usize,
|
|
attention_weights: Vec<f32>,
|
|
token_positions: Vec<usize>,
|
|
}
|
|
|
|
impl CompressedKVCache {
|
|
/// Create new compressed KV cache
|
|
pub fn new(config: KVCacheConfig) -> Result<Self> {
|
|
let device = Device::try_default()?;
|
|
|
|
Ok(Self {
|
|
config,
|
|
cache_entries: HashMap::new(),
|
|
sequence_metadata: HashMap::new(),
|
|
importance_tracker: ImportanceTracker::new(0.95, 100),
|
|
device,
|
|
stats: CompressionStats::default(),
|
|
memory_usage: 0,
|
|
global_timestamp: 0,
|
|
pq_key: None,
|
|
pq_value: None,
|
|
vq_key: None,
|
|
vq_value: None,
|
|
quantizers_trained: false,
|
|
})
|
|
}
|
|
|
|
/// Train quantizers on sample data for optimal compression
|
|
/// Should be called with representative data before using PQ/VQ compression
|
|
pub fn train_quantizers(&mut self, sample_keys: &Tensor, sample_values: &Tensor) -> Result<()> {
|
|
let key_shape = sample_keys.shape();
|
|
let value_shape = sample_values.shape();
|
|
|
|
// Calculate feature dimension (flatten heads and head_dim)
|
|
let key_dim = if key_shape.dims().len() >= 3 {
|
|
key_shape.dims()[key_shape.dims().len() - 2]
|
|
* key_shape.dims()[key_shape.dims().len() - 1]
|
|
} else {
|
|
key_shape.dims()[key_shape.dims().len() - 1]
|
|
};
|
|
|
|
let value_dim = if value_shape.dims().len() >= 3 {
|
|
value_shape.dims()[value_shape.dims().len() - 2]
|
|
* value_shape.dims()[value_shape.dims().len() - 1]
|
|
} else {
|
|
value_shape.dims()[value_shape.dims().len() - 1]
|
|
};
|
|
|
|
// Train based on compression method
|
|
match &self.config.compression_method {
|
|
CompressionMethod::ProductQuantization {
|
|
num_subquantizers,
|
|
codebook_size,
|
|
use_opq,
|
|
} => {
|
|
// Ensure dimension is divisible by num_subquantizers
|
|
let num_subq = (*num_subquantizers).min(key_dim);
|
|
let adjusted_subq = Self::find_divisible_subquantizers(key_dim, num_subq);
|
|
|
|
let pq_config = PQConfig {
|
|
num_subquantizers: adjusted_subq,
|
|
codebook_size: *codebook_size,
|
|
use_opq: *use_opq,
|
|
max_iterations: 50,
|
|
tolerance: 1e-5,
|
|
..Default::default()
|
|
};
|
|
|
|
// Reshape keys for training
|
|
let keys_2d = Self::reshape_for_training(sample_keys, key_dim)?;
|
|
let values_2d = Self::reshape_for_training(sample_values, value_dim)?;
|
|
|
|
// Train key quantizer
|
|
let mut pq_key = ProductQuantizer::new(pq_config.clone())?;
|
|
pq_key.fit(&keys_2d)?;
|
|
self.pq_key = Some(pq_key);
|
|
|
|
// Train value quantizer
|
|
let mut pq_value = ProductQuantizer::new(pq_config)?;
|
|
pq_value.fit(&values_2d)?;
|
|
self.pq_value = Some(pq_value);
|
|
}
|
|
|
|
CompressionMethod::VectorQuantization { codebook_size, .. } => {
|
|
let vq_config = VQConfig {
|
|
codebook_size: *codebook_size,
|
|
vector_dim: key_dim,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
max_iterations: 50,
|
|
..Default::default()
|
|
};
|
|
|
|
let keys_2d = Self::reshape_for_training(sample_keys, key_dim)?;
|
|
let values_2d = Self::reshape_for_training(sample_values, value_dim)?;
|
|
|
|
let mut vq_key = VectorQuantizer::new(vq_config.clone());
|
|
vq_key.fit(&keys_2d)?;
|
|
self.vq_key = Some(vq_key);
|
|
|
|
let mut vq_value = VectorQuantizer::new(vq_config);
|
|
vq_value.fit(&values_2d)?;
|
|
self.vq_value = Some(vq_value);
|
|
}
|
|
|
|
_ => {
|
|
// Other methods don't need training
|
|
}
|
|
}
|
|
|
|
self.quantizers_trained = true;
|
|
Ok(())
|
|
}
|
|
|
|
/// Find the largest divisor of dim that is <= target
|
|
fn find_divisible_subquantizers(dim: usize, target: usize) -> usize {
|
|
let mut best = 1;
|
|
for d in 1..=target {
|
|
if dim % d == 0 {
|
|
best = d;
|
|
}
|
|
}
|
|
best
|
|
}
|
|
|
|
/// Reshape tensor to 2D for training (batch of vectors)
|
|
fn reshape_for_training(tensor: &Tensor, feature_dim: usize) -> Result<Tensor> {
|
|
let shape = tensor.shape();
|
|
let total_elements: usize = shape.dims().iter().product();
|
|
let num_vectors = total_elements / feature_dim;
|
|
|
|
let flat_data = tensor.to_vec()?;
|
|
let device = Device::try_default()?;
|
|
Ok(Tensor::from_slice(
|
|
&flat_data,
|
|
&[num_vectors, feature_dim],
|
|
&device,
|
|
)?)
|
|
}
|
|
|
|
/// Insert key-value pair with compression
|
|
pub fn insert(&mut self, seq_id: usize, keys: &Tensor, values: &Tensor) -> Result<()> {
|
|
self.global_timestamp += 1;
|
|
|
|
// Initialize sequence metadata if new
|
|
self.sequence_metadata.entry(seq_id).or_insert_with(|| {
|
|
let shape = keys.shape();
|
|
|
|
SequenceMetadata {
|
|
sequence_length: shape.dims()[1], // Assuming [batch, seq, heads, head_dim]
|
|
layer_count: 1,
|
|
head_count: shape.dims()[2],
|
|
head_dim: shape.dims()[3],
|
|
attention_patterns: vec![1.0; shape.dims()[1]], // Initialize with uniform importance
|
|
}
|
|
});
|
|
|
|
// Check memory limits and evict if necessary
|
|
self.enforce_memory_limits()?;
|
|
|
|
// Compress using configured method (PQ, VQ, or simple)
|
|
let compressed_keys = self.compress_tensor(keys, true)?;
|
|
let compressed_values = self.compress_tensor(values, false)?;
|
|
|
|
let cache_key = format!("{}:0:{}", seq_id, keys.shape().dims()[1]);
|
|
let entry = CacheEntry {
|
|
compressed_keys,
|
|
compressed_values,
|
|
original_shape: keys.shape().dims().to_vec(),
|
|
compression_metadata: CompressionMetadata {
|
|
method: self.config.compression_method.clone(),
|
|
compression_ratio: 4.0, // Placeholder
|
|
quality_score: 0.95,
|
|
codebook_indices: None,
|
|
quantization_params: None,
|
|
},
|
|
importance: TokenImportance::default(),
|
|
sequence_id: seq_id,
|
|
token_range: (0, keys.shape().dims()[1]),
|
|
};
|
|
|
|
let entry_size = entry.compressed_keys.len() + entry.compressed_values.len();
|
|
self.memory_usage += entry_size;
|
|
self.cache_entries.insert(cache_key, entry);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Insert batch of key-value pairs
|
|
pub fn insert_batch(
|
|
&mut self,
|
|
seq_ids: &[usize],
|
|
batch_keys: &Tensor,
|
|
batch_values: &Tensor,
|
|
) -> Result<()> {
|
|
let _batch_size = batch_keys.shape().dims()[0];
|
|
|
|
for (i, &seq_id) in seq_ids.iter().enumerate() {
|
|
let keys = batch_keys.narrow(0, i, 1)?;
|
|
let values = batch_values.narrow(0, i, 1)?;
|
|
self.insert(seq_id, &keys, &values)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Append new tokens to existing sequence
|
|
pub fn append(&mut self, seq_id: usize, new_keys: &Tensor, new_values: &Tensor) -> Result<()> {
|
|
// Get current sequence length
|
|
let current_length = self
|
|
.sequence_metadata
|
|
.get(&seq_id)
|
|
.map_or(0, |m| m.sequence_length);
|
|
|
|
// Update sequence metadata
|
|
if let Some(metadata) = self.sequence_metadata.get_mut(&seq_id) {
|
|
metadata.sequence_length += new_keys.shape().dims()[1];
|
|
metadata
|
|
.attention_patterns
|
|
.extend(vec![1.0; new_keys.shape().dims()[1]]);
|
|
}
|
|
|
|
// Apply sliding window if enabled
|
|
if self.config.enable_sliding_window
|
|
&& current_length + new_keys.shape().dims()[1] > self.config.window_size
|
|
{
|
|
self.apply_sliding_window(seq_id)?;
|
|
}
|
|
|
|
// Insert new tokens
|
|
self.insert(seq_id, new_keys, new_values)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Insert with attention weights for importance-aware compression
|
|
pub fn insert_with_attention(
|
|
&mut self,
|
|
seq_id: usize,
|
|
keys: &Tensor,
|
|
values: &Tensor,
|
|
attention_weights: &Tensor,
|
|
) -> Result<()> {
|
|
// Update importance tracker with attention information
|
|
if self.config.enable_attention_scoring {
|
|
let attention_data = attention_weights.to_vec()?;
|
|
self.importance_tracker.update_attention(
|
|
seq_id,
|
|
&attention_data,
|
|
self.global_timestamp,
|
|
);
|
|
}
|
|
|
|
// Use adaptive compression based on attention patterns
|
|
let attention_data = attention_weights.to_vec()?;
|
|
let adaptive_method = self.select_adaptive_compression(&attention_data)?;
|
|
|
|
// Compress with configured method
|
|
let compressed_keys = self.compress_tensor(keys, true)?;
|
|
let compressed_values = self.compress_tensor(values, false)?;
|
|
|
|
// Store with enhanced importance information
|
|
let cache_key = format!("{}:attn:{}", seq_id, keys.shape().dims()[1]);
|
|
let entry = CacheEntry {
|
|
compressed_keys,
|
|
compressed_values,
|
|
original_shape: keys.shape().dims().to_vec(),
|
|
compression_metadata: CompressionMetadata {
|
|
method: adaptive_method,
|
|
compression_ratio: 4.0,
|
|
quality_score: 0.95,
|
|
codebook_indices: None,
|
|
quantization_params: None,
|
|
},
|
|
importance: TokenImportance {
|
|
importance_score: attention_data.iter().sum::<f32>() / attention_data.len() as f32,
|
|
attention_weight: attention_data.iter().fold(0.0f32, |a, &b| a.max(b)),
|
|
access_count: 1,
|
|
last_access_time: self.global_timestamp,
|
|
compression_level: 8,
|
|
},
|
|
sequence_id: seq_id,
|
|
token_range: (0, keys.shape().dims()[1]),
|
|
};
|
|
|
|
let entry_size = entry.compressed_keys.len() + entry.compressed_values.len();
|
|
self.memory_usage += entry_size;
|
|
self.cache_entries.insert(cache_key, entry);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Retrieve key-value pairs with decompression
|
|
pub fn get(&mut self, seq_id: usize, start: usize, length: usize) -> Result<(Tensor, Tensor)> {
|
|
let cache_key = format!("{}:0:{}", seq_id, start + length);
|
|
|
|
if let Some(entry) = self.cache_entries.get_mut(&cache_key) {
|
|
// Update access statistics
|
|
entry.importance.access_count += 1;
|
|
entry.importance.last_access_time = self.global_timestamp;
|
|
self.stats.cache_hit_rate += 1.0;
|
|
|
|
// Extract data to avoid borrow issues
|
|
let compressed_keys = entry.compressed_keys.clone();
|
|
let compressed_values = entry.compressed_values.clone();
|
|
let original_shape = entry.original_shape.clone();
|
|
|
|
// Decompress keys and values
|
|
let keys = Self::decompress_tensor_simple_static(
|
|
&compressed_keys,
|
|
&original_shape,
|
|
&self.device,
|
|
)?;
|
|
let values = Self::decompress_tensor_simple_static(
|
|
&compressed_values,
|
|
&original_shape,
|
|
&self.device,
|
|
)?;
|
|
|
|
// Extract requested range
|
|
let keys_slice = if start + length <= keys.shape().dims()[1] {
|
|
keys.narrow(1, start, length)?
|
|
} else {
|
|
keys.clone()
|
|
};
|
|
let values_slice = if start + length <= values.shape().dims()[1] {
|
|
values.narrow(1, start, length)?
|
|
} else {
|
|
values.clone()
|
|
};
|
|
|
|
Ok((keys_slice, values_slice))
|
|
} else {
|
|
// Cache miss - create dummy tensors
|
|
let shape = vec![1, length, 8, 64]; // Default shape
|
|
let keys = Tensor::zeros(&shape, &self.device)?;
|
|
let values = Tensor::zeros(&shape, &self.device)?;
|
|
Ok((keys, values))
|
|
}
|
|
}
|
|
|
|
/// Retrieve batch of key-value pairs
|
|
pub fn get_batch(
|
|
&mut self,
|
|
seq_ids: &[usize],
|
|
start: usize,
|
|
length: usize,
|
|
) -> Result<(Tensor, Tensor)> {
|
|
let batch_size = seq_ids.len();
|
|
let mut batch_keys = Vec::new();
|
|
let mut batch_values = Vec::new();
|
|
|
|
for &seq_id in seq_ids {
|
|
let (keys, values) = self.get(seq_id, start, length)?;
|
|
batch_keys.push(keys);
|
|
batch_values.push(values);
|
|
}
|
|
|
|
// Concatenate along batch dimension
|
|
if !batch_keys.is_empty() {
|
|
let stacked_keys = Tensor::stack(&batch_keys, 0)?;
|
|
let stacked_values = Tensor::stack(&batch_values, 0)?;
|
|
Ok((stacked_keys, stacked_values))
|
|
} else {
|
|
let shape = vec![batch_size, length, 8, 64];
|
|
let keys = Tensor::zeros(&shape, &self.device)?;
|
|
let values = Tensor::zeros(&shape, &self.device)?;
|
|
Ok((keys, values))
|
|
}
|
|
}
|
|
|
|
/// Compress tensor using the configured compression method
|
|
fn compress_tensor(&self, tensor: &Tensor, is_key: bool) -> Result<Vec<u8>> {
|
|
// Try PQ compression if trained
|
|
if self.quantizers_trained {
|
|
match &self.config.compression_method {
|
|
CompressionMethod::ProductQuantization { .. } => {
|
|
let pq = if is_key { &self.pq_key } else { &self.pq_value };
|
|
if let Some(quantizer) = pq {
|
|
return self.compress_with_pq(tensor, quantizer);
|
|
}
|
|
}
|
|
CompressionMethod::VectorQuantization { .. } => {
|
|
let vq = if is_key { &self.vq_key } else { &self.vq_value };
|
|
if let Some(quantizer) = vq {
|
|
return self.compress_with_vq(tensor, quantizer);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
// Fallback to simple compression
|
|
self.compress_tensor_simple(tensor)
|
|
}
|
|
|
|
/// Compress tensor using Product Quantization
|
|
fn compress_with_pq(&self, tensor: &Tensor, pq: &ProductQuantizer) -> Result<Vec<u8>> {
|
|
let shape = tensor.shape();
|
|
let feature_dim = shape.dims().iter().skip(1).product::<usize>().max(1);
|
|
let _num_vectors = shape.dims()[0];
|
|
|
|
// Reshape to 2D for PQ encoding
|
|
let data_2d = Self::reshape_for_training(tensor, feature_dim)?;
|
|
|
|
// Encode using PQ
|
|
let codes = pq.encode(&data_2d)?;
|
|
let codes_vec = codes.to_vec()?;
|
|
|
|
// Pack codes into bytes with header
|
|
let mut compressed = Vec::new();
|
|
|
|
// Header: compression type (1 byte), original shape dims
|
|
compressed.push(0x01); // PQ compression marker
|
|
compressed.extend_from_slice(&(shape.dims().len() as u32).to_le_bytes());
|
|
for &dim in shape.dims() {
|
|
compressed.extend_from_slice(&(dim as u32).to_le_bytes());
|
|
}
|
|
|
|
// Number of codes per vector
|
|
let codes_shape = codes.shape();
|
|
let codes_per_vector = if codes_shape.dims().len() > 1 {
|
|
codes_shape.dims()[1]
|
|
} else {
|
|
1
|
|
};
|
|
compressed.extend_from_slice(&(codes_per_vector as u32).to_le_bytes());
|
|
|
|
// Codes (stored as u16 since codebook size is typically 256-1024)
|
|
for &code in &codes_vec {
|
|
compressed.extend_from_slice(&(code as u16).to_le_bytes());
|
|
}
|
|
|
|
Ok(compressed)
|
|
}
|
|
|
|
/// Compress tensor using Vector Quantization
|
|
fn compress_with_vq(&self, tensor: &Tensor, vq: &VectorQuantizer) -> Result<Vec<u8>> {
|
|
let shape = tensor.shape();
|
|
let feature_dim = shape.dims().iter().skip(1).product::<usize>().max(1);
|
|
|
|
// Reshape to 2D for VQ encoding
|
|
let data_2d = Self::reshape_for_training(tensor, feature_dim)?;
|
|
|
|
// Encode using VQ
|
|
let codes = vq.encode(&data_2d)?;
|
|
let codes_vec = codes.to_vec()?;
|
|
|
|
// Pack codes into bytes with header
|
|
let mut compressed = Vec::new();
|
|
|
|
// Header: compression type (1 byte), original shape dims
|
|
compressed.push(0x02); // VQ compression marker
|
|
compressed.extend_from_slice(&(shape.dims().len() as u32).to_le_bytes());
|
|
for &dim in shape.dims() {
|
|
compressed.extend_from_slice(&(dim as u32).to_le_bytes());
|
|
}
|
|
|
|
// Codes (stored as u16)
|
|
for &code in &codes_vec {
|
|
compressed.extend_from_slice(&(code as u16).to_le_bytes());
|
|
}
|
|
|
|
Ok(compressed)
|
|
}
|
|
|
|
/// Simple tensor compression (fallback implementation)
|
|
fn compress_tensor_simple(&self, tensor: &Tensor) -> Result<Vec<u8>> {
|
|
let data = tensor.to_vec()?;
|
|
let shape = tensor.shape();
|
|
|
|
// Empty tensor
|
|
if data.is_empty() {
|
|
let mut compressed = vec![0x00]; // Simple compression marker
|
|
compressed.extend_from_slice(&(shape.dims().len() as u32).to_le_bytes());
|
|
for &dim in shape.dims() {
|
|
compressed.extend_from_slice(&(dim as u32).to_le_bytes());
|
|
}
|
|
return Ok(compressed);
|
|
}
|
|
|
|
let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
|
let scale = (max_val - min_val) / 255.0;
|
|
|
|
let mut compressed = Vec::new();
|
|
compressed.push(0x00); // Simple compression marker
|
|
compressed.extend_from_slice(&(shape.dims().len() as u32).to_le_bytes());
|
|
for &dim in shape.dims() {
|
|
compressed.extend_from_slice(&(dim as u32).to_le_bytes());
|
|
}
|
|
compressed.extend_from_slice(&min_val.to_le_bytes());
|
|
compressed.extend_from_slice(&scale.to_le_bytes());
|
|
|
|
for value in data {
|
|
let quantized = if scale > 0.0 {
|
|
((value - min_val) / scale).clamp(0.0, 255.0) as u8
|
|
} else {
|
|
128u8
|
|
};
|
|
compressed.push(quantized);
|
|
}
|
|
|
|
Ok(compressed)
|
|
}
|
|
|
|
/// Simple tensor decompression (placeholder implementation)
|
|
fn decompress_tensor_simple(
|
|
&self,
|
|
compressed_data: &[u8],
|
|
original_shape: &[usize],
|
|
) -> Result<Tensor> {
|
|
if compressed_data.len() < 8 {
|
|
// Fallback for invalid data
|
|
return Ok(Tensor::zeros(
|
|
rtx_tensor::Shape::new(original_shape.to_vec())?,
|
|
&self.device,
|
|
)?);
|
|
}
|
|
|
|
// Extract quantization parameters
|
|
let min_val = f32::from_le_bytes([
|
|
compressed_data[0],
|
|
compressed_data[1],
|
|
compressed_data[2],
|
|
compressed_data[3],
|
|
]);
|
|
let scale = f32::from_le_bytes([
|
|
compressed_data[4],
|
|
compressed_data[5],
|
|
compressed_data[6],
|
|
compressed_data[7],
|
|
]);
|
|
|
|
let num_elements: usize = original_shape.iter().product();
|
|
let mut decompressed = Vec::with_capacity(num_elements);
|
|
|
|
for &byte in &compressed_data[8..] {
|
|
let value = min_val + (byte as f32) * scale;
|
|
decompressed.push(value);
|
|
if decompressed.len() >= num_elements {
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Pad with zeros if needed
|
|
while decompressed.len() < num_elements {
|
|
decompressed.push(0.0);
|
|
}
|
|
|
|
decompressed.truncate(num_elements);
|
|
Ok(Tensor::from_slice(
|
|
&decompressed,
|
|
original_shape,
|
|
&self.device,
|
|
)?)
|
|
}
|
|
|
|
/// Static version of tensor decompression to avoid borrow issues
|
|
fn decompress_tensor_simple_static(
|
|
compressed_data: &[u8],
|
|
original_shape: &[usize],
|
|
device: &Device,
|
|
) -> Result<Tensor> {
|
|
if compressed_data.len() < 8 {
|
|
// Fallback for invalid data
|
|
return Ok(Tensor::zeros(
|
|
rtx_tensor::Shape::new(original_shape.to_vec())?,
|
|
device,
|
|
)?);
|
|
}
|
|
|
|
// Extract quantization parameters
|
|
let min_val = f32::from_le_bytes([
|
|
compressed_data[0],
|
|
compressed_data[1],
|
|
compressed_data[2],
|
|
compressed_data[3],
|
|
]);
|
|
let scale = f32::from_le_bytes([
|
|
compressed_data[4],
|
|
compressed_data[5],
|
|
compressed_data[6],
|
|
compressed_data[7],
|
|
]);
|
|
|
|
let num_elements: usize = original_shape.iter().product();
|
|
let mut decompressed = Vec::with_capacity(num_elements);
|
|
|
|
for &byte in &compressed_data[8..] {
|
|
let value = min_val + (byte as f32) * scale;
|
|
decompressed.push(value);
|
|
if decompressed.len() >= num_elements {
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Pad with zeros if needed
|
|
while decompressed.len() < num_elements {
|
|
decompressed.push(0.0);
|
|
}
|
|
|
|
decompressed.truncate(num_elements);
|
|
Ok(Tensor::from_slice(&decompressed, original_shape, device)?)
|
|
}
|
|
|
|
/// Select adaptive compression method based on attention weights
|
|
fn select_adaptive_compression(&self, attention_data: &[f32]) -> Result<CompressionMethod> {
|
|
let avg_attention = attention_data.iter().sum::<f32>() / attention_data.len() as f32;
|
|
let max_attention = attention_data.iter().fold(0.0f32, |a, &b| a.max(b));
|
|
|
|
if max_attention > 0.8 && avg_attention > 0.3 {
|
|
// High attention - use higher precision
|
|
Ok(CompressionMethod::AdaptiveQuantization {
|
|
base_bits: 8,
|
|
attention_threshold: 0.5,
|
|
importance_decay: 0.9,
|
|
})
|
|
} else if avg_attention > 0.1 {
|
|
// Medium attention - balanced compression
|
|
Ok(CompressionMethod::ProductQuantization {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
use_opq: true,
|
|
})
|
|
} else {
|
|
// Low attention - aggressive compression
|
|
Ok(CompressionMethod::AdaptiveQuantization {
|
|
base_bits: 4,
|
|
attention_threshold: 0.2,
|
|
importance_decay: 0.8,
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Apply sliding window to manage cache size
|
|
fn apply_sliding_window(&mut self, seq_id: usize) -> Result<()> {
|
|
let window_size = self.config.window_size;
|
|
|
|
// Find entries for this sequence that exceed window size
|
|
let keys_to_remove: Vec<String> = self
|
|
.cache_entries
|
|
.iter()
|
|
.filter(|(_, entry)| entry.sequence_id == seq_id && entry.token_range.1 > window_size)
|
|
.map(|(key, _)| key.clone())
|
|
.collect();
|
|
|
|
// Remove oldest entries
|
|
for key in keys_to_remove {
|
|
if let Some(entry) = self.cache_entries.remove(&key) {
|
|
self.memory_usage = self
|
|
.memory_usage
|
|
.saturating_sub(entry.compressed_keys.len() + entry.compressed_values.len());
|
|
self.stats.eviction_count += 1;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Enforce memory limits by evicting least important entries
|
|
fn enforce_memory_limits(&mut self) -> Result<()> {
|
|
let max_memory_bytes = self.config.max_cache_size_mb * 1024 * 1024;
|
|
|
|
if self.memory_usage <= max_memory_bytes {
|
|
return Ok(());
|
|
}
|
|
|
|
// Collect entries with importance scores for eviction
|
|
let mut entries_by_importance: Vec<_> = self
|
|
.cache_entries
|
|
.iter()
|
|
.map(|(key, entry)| {
|
|
let importance = entry.importance.importance_score
|
|
* (1.0 + entry.importance.access_count as f32)
|
|
/ (self.global_timestamp - entry.importance.last_access_time + 1) as f32;
|
|
(key.clone(), importance)
|
|
})
|
|
.collect();
|
|
|
|
// Sort by importance (ascending - remove least important first)
|
|
entries_by_importance.sort_by(|a, b| a.1.total_cmp(&b.1));
|
|
|
|
// Remove entries until under memory limit
|
|
for (key, _) in entries_by_importance {
|
|
if self.memory_usage <= max_memory_bytes {
|
|
break;
|
|
}
|
|
|
|
if let Some(entry) = self.cache_entries.remove(&key) {
|
|
self.memory_usage = self
|
|
.memory_usage
|
|
.saturating_sub(entry.compressed_keys.len() + entry.compressed_values.len());
|
|
self.stats.eviction_count += 1;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get compression statistics
|
|
pub fn compression_stats(&self) -> CompressionStats {
|
|
let mut stats = self.stats.clone();
|
|
stats.memory_usage_mb = self.memory_usage as f64 / (1024.0 * 1024.0);
|
|
stats
|
|
}
|
|
|
|
/// Set memory limit in bytes
|
|
pub fn set_memory_limit(&mut self, limit_bytes: usize) {
|
|
self.config.max_cache_size_mb = limit_bytes / (1024 * 1024);
|
|
}
|
|
|
|
/// Get current memory usage in bytes
|
|
pub fn memory_usage(&self) -> usize {
|
|
self.memory_usage
|
|
}
|
|
|
|
/// Check if sequence is cached
|
|
pub fn contains_sequence(&self, seq_id: usize) -> bool {
|
|
self.sequence_metadata.contains_key(&seq_id)
|
|
}
|
|
|
|
/// Serialize cache to bytes
|
|
pub fn save_to_bytes(&self) -> Result<Vec<u8>> {
|
|
let mut data = Vec::new();
|
|
|
|
// Save configuration
|
|
let config_bytes = bincode::serialize(&self.config)
|
|
.map_err(|e| CompressionError::Serialization(e.to_string()))?;
|
|
data.extend_from_slice(&config_bytes.len().to_le_bytes());
|
|
data.extend_from_slice(&config_bytes);
|
|
|
|
// Save statistics
|
|
let stats_bytes = bincode::serialize(&self.stats)
|
|
.map_err(|e| CompressionError::Serialization(e.to_string()))?;
|
|
data.extend_from_slice(&stats_bytes.len().to_le_bytes());
|
|
data.extend_from_slice(&stats_bytes);
|
|
|
|
// Save sequence metadata count
|
|
data.extend_from_slice(&self.sequence_metadata.len().to_le_bytes());
|
|
|
|
// Save sequence metadata
|
|
for (&seq_id, metadata) in &self.sequence_metadata {
|
|
data.extend_from_slice(&seq_id.to_le_bytes());
|
|
data.extend_from_slice(&metadata.sequence_length.to_le_bytes());
|
|
data.extend_from_slice(&metadata.layer_count.to_le_bytes());
|
|
data.extend_from_slice(&metadata.head_count.to_le_bytes());
|
|
data.extend_from_slice(&metadata.head_dim.to_le_bytes());
|
|
|
|
// Save attention patterns
|
|
data.extend_from_slice(&metadata.attention_patterns.len().to_le_bytes());
|
|
for &pattern in &metadata.attention_patterns {
|
|
data.extend_from_slice(&pattern.to_le_bytes());
|
|
}
|
|
}
|
|
|
|
Ok(data)
|
|
}
|
|
|
|
/// Load cache from bytes
|
|
pub fn load_from_bytes(&mut self, data: &[u8]) -> Result<()> {
|
|
let mut offset = 0;
|
|
|
|
// Load configuration
|
|
let config_len =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid config length".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
self.config = bincode::deserialize(&data[offset..offset + config_len])
|
|
.map_err(|e| CompressionError::Serialization(e.to_string()))?;
|
|
offset += config_len;
|
|
|
|
// Load statistics
|
|
let stats_len =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid stats length".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
self.stats = bincode::deserialize(&data[offset..offset + stats_len])
|
|
.map_err(|e| CompressionError::Serialization(e.to_string()))?;
|
|
offset += stats_len;
|
|
|
|
// Load sequence metadata
|
|
let metadata_count =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid metadata count".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
self.sequence_metadata.clear();
|
|
for _ in 0..metadata_count {
|
|
let seq_id =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid sequence ID".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
let sequence_length =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid sequence length".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
let layer_count =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid layer count".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
let head_count =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid head count".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
let head_dim =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid head dim".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
let pattern_count =
|
|
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid pattern count".to_string())
|
|
})?);
|
|
offset += 8;
|
|
|
|
let mut attention_patterns = Vec::new();
|
|
for _ in 0..pattern_count {
|
|
let pattern =
|
|
f32::from_le_bytes(data[offset..offset + 4].try_into().map_err(|_| {
|
|
CompressionError::Serialization("Invalid pattern value".to_string())
|
|
})?);
|
|
attention_patterns.push(pattern);
|
|
offset += 4;
|
|
}
|
|
|
|
let metadata = SequenceMetadata {
|
|
sequence_length,
|
|
layer_count,
|
|
head_count,
|
|
head_dim,
|
|
attention_patterns,
|
|
};
|
|
|
|
self.sequence_metadata.insert(seq_id, metadata);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl ImportanceTracker {
|
|
fn new(decay_factor: f32, update_frequency: usize) -> Self {
|
|
Self {
|
|
token_scores: HashMap::new(),
|
|
attention_history: VecDeque::new(),
|
|
decay_factor,
|
|
update_frequency,
|
|
}
|
|
}
|
|
|
|
fn update_attention(&mut self, seq_id: usize, attention_weights: &[f32], timestamp: u64) {
|
|
// Store attention snapshot
|
|
let snapshot = AttentionSnapshot {
|
|
timestamp,
|
|
sequence_id: seq_id,
|
|
attention_weights: attention_weights.to_vec(),
|
|
token_positions: (0..attention_weights.len()).collect(),
|
|
};
|
|
|
|
self.attention_history.push_back(snapshot);
|
|
|
|
// Update token importance scores
|
|
for (pos, &weight) in attention_weights.iter().enumerate() {
|
|
let token_key = format!("{seq_id}:{pos}");
|
|
let importance = self.token_scores.entry(token_key).or_default();
|
|
|
|
// Exponential moving average
|
|
importance.attention_weight = importance.attention_weight * self.decay_factor
|
|
+ weight * (1.0 - self.decay_factor);
|
|
importance.importance_score =
|
|
importance.attention_weight * (1.0 + importance.access_count as f32).log10();
|
|
}
|
|
|
|
// Maintain history size
|
|
while self.attention_history.len() > 1000 {
|
|
self.attention_history.pop_front();
|
|
}
|
|
}
|
|
}
|