Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,661 @@
|
||||
//! Distributed attention and KV cache management.
|
||||
//!
|
||||
//! This module implements distributed multi-head attention with sharded
|
||||
//! KV cache for memory-efficient inference of large language models.
|
||||
|
||||
use distllm_shared::{DataType, KVCacheConfig};
|
||||
|
||||
// ============================================================================
|
||||
// Distributed Attention
|
||||
// ============================================================================
|
||||
|
||||
/// Distributed multi-head attention with tensor parallelism support.
|
||||
#[derive(Debug)]
|
||||
pub struct DistributedAttention {
|
||||
/// Number of query heads per shard.
|
||||
pub num_heads: usize,
|
||||
/// Number of key-value heads per shard.
|
||||
pub num_kv_heads: usize,
|
||||
/// Dimension per attention head.
|
||||
pub head_dim: usize,
|
||||
/// Tensor parallel world size.
|
||||
pub tp_world_size: usize,
|
||||
/// Whether to use flash attention.
|
||||
pub use_flash_attention: bool,
|
||||
/// RoPE theta for positional encoding.
|
||||
pub rope_theta: f64,
|
||||
/// Maximum sequence length.
|
||||
pub max_seq_len: usize,
|
||||
/// Cached cos/sin for RoPE.
|
||||
rope_cache: Option<(Vec<f32>, Vec<f32>)>,
|
||||
}
|
||||
|
||||
impl DistributedAttention {
|
||||
/// Create a new distributed attention module.
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
num_heads: usize,
|
||||
num_kv_heads: usize,
|
||||
head_dim: usize,
|
||||
tp_world_size: usize,
|
||||
) -> Self {
|
||||
Self {
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
tp_world_size,
|
||||
use_flash_attention: true,
|
||||
rope_theta: 10000.0,
|
||||
max_seq_len: 4096,
|
||||
rope_cache: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize RoPE cache for given sequence length.
|
||||
pub fn init_rope_cache(&mut self, seq_len: usize) {
|
||||
let half_dim = self.head_dim / 2;
|
||||
let mut cos = Vec::with_capacity(seq_len * half_dim);
|
||||
let mut sin = Vec::with_capacity(seq_len * half_dim);
|
||||
|
||||
for pos in 0..seq_len {
|
||||
for i in 0..half_dim {
|
||||
let freq = 1.0 / self.rope_theta.powf(2.0 * i as f64 / self.head_dim as f64);
|
||||
let angle = pos as f64 * freq;
|
||||
cos.push(angle.cos() as f32);
|
||||
sin.push(angle.sin() as f32);
|
||||
}
|
||||
}
|
||||
|
||||
self.rope_cache = Some((cos, sin));
|
||||
}
|
||||
|
||||
/// Compute attention output (simulated).
|
||||
///
|
||||
/// Returns simulated output dimensions.
|
||||
#[must_use]
|
||||
pub fn forward(&self, seq_len: usize, _layer_id: usize) -> (usize, usize, usize) {
|
||||
// Output shape: (batch, seq_len, num_heads * head_dim)
|
||||
let output_dim = self.num_heads * self.head_dim;
|
||||
(1, seq_len, output_dim)
|
||||
}
|
||||
|
||||
/// Compute attention with flash attention optimization (simulated).
|
||||
#[must_use]
|
||||
pub fn flash_forward(
|
||||
&self,
|
||||
seq_len: usize,
|
||||
kv_len: usize,
|
||||
_layer_id: usize,
|
||||
) -> FlashAttentionOutput {
|
||||
// Flash attention reduces memory from O(n^2) to O(n)
|
||||
let memory_saved_ratio = if seq_len > 512 {
|
||||
1.0 - (512.0 / seq_len as f64)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
FlashAttentionOutput {
|
||||
output_shape: (1, seq_len, self.num_heads * self.head_dim),
|
||||
kv_len,
|
||||
memory_saved_ratio,
|
||||
num_blocks: seq_len.div_ceil(128), // Block size of 128
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply RoPE positional encoding (simulated).
|
||||
pub fn apply_rope(&self, _positions: &[usize]) -> bool {
|
||||
self.rope_cache.is_some()
|
||||
}
|
||||
|
||||
/// Get the number of heads after tensor parallel split.
|
||||
#[must_use]
|
||||
pub fn sharded_num_heads(&self) -> usize {
|
||||
self.num_heads
|
||||
}
|
||||
|
||||
/// Get the number of KV heads after tensor parallel split.
|
||||
#[must_use]
|
||||
pub fn sharded_num_kv_heads(&self) -> usize {
|
||||
self.num_kv_heads
|
||||
}
|
||||
|
||||
/// Calculate GQA (Grouped Query Attention) ratio.
|
||||
#[must_use]
|
||||
pub fn gqa_ratio(&self) -> usize {
|
||||
self.num_heads / self.num_kv_heads.max(1)
|
||||
}
|
||||
|
||||
/// Calculate attention memory in MB for given sequence length.
|
||||
#[must_use]
|
||||
pub fn attention_memory_mb(&self, seq_len: usize, dtype_bytes: usize) -> f64 {
|
||||
// Q, K, V, output matrices
|
||||
let qkv_size = 3 * seq_len * self.num_heads * self.head_dim * dtype_bytes;
|
||||
// Attention scores (with flash attention, this is tiled)
|
||||
let score_size = if self.use_flash_attention {
|
||||
128 * 128 * self.num_heads * dtype_bytes // Block size
|
||||
} else {
|
||||
seq_len * seq_len * self.num_heads * dtype_bytes
|
||||
};
|
||||
(qkv_size + score_size) as f64 / (1024.0 * 1024.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Output from flash attention computation.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FlashAttentionOutput {
|
||||
/// Output tensor shape.
|
||||
pub output_shape: (usize, usize, usize),
|
||||
/// KV cache length used.
|
||||
pub kv_len: usize,
|
||||
/// Memory saved compared to standard attention.
|
||||
pub memory_saved_ratio: f64,
|
||||
/// Number of blocks processed.
|
||||
pub num_blocks: usize,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Flash Attention
|
||||
// ============================================================================
|
||||
|
||||
/// Flash Attention implementation for memory-efficient attention.
|
||||
#[derive(Debug)]
|
||||
pub struct FlashAttention {
|
||||
/// Block size for tiling.
|
||||
pub block_size: usize,
|
||||
/// Whether to use causal masking.
|
||||
pub causal: bool,
|
||||
/// Softmax scaling factor.
|
||||
pub scale: f32,
|
||||
/// Dropout probability.
|
||||
pub dropout: f32,
|
||||
}
|
||||
|
||||
impl FlashAttention {
|
||||
/// Create a new flash attention module.
|
||||
#[must_use]
|
||||
pub fn new(head_dim: usize) -> Self {
|
||||
Self {
|
||||
block_size: 128,
|
||||
causal: true,
|
||||
scale: 1.0 / (head_dim as f32).sqrt(),
|
||||
dropout: 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Calculate number of blocks needed for sequence length.
|
||||
#[must_use]
|
||||
pub fn num_blocks(&self, seq_len: usize) -> usize {
|
||||
seq_len.div_ceil(self.block_size)
|
||||
}
|
||||
|
||||
/// Calculate memory usage in MB.
|
||||
#[must_use]
|
||||
pub fn memory_mb(&self, seq_len: usize, num_heads: usize, dtype_bytes: usize) -> f64 {
|
||||
// Flash attention only needs O(block_size^2) per block
|
||||
let block_mem = self.block_size * self.block_size * num_heads * dtype_bytes;
|
||||
let num_blocks = self.num_blocks(seq_len);
|
||||
// We process 2 blocks at a time (Q block and KV block)
|
||||
(2 * block_mem * num_blocks) as f64 / (1024.0 * 1024.0)
|
||||
}
|
||||
|
||||
/// Calculate FLOPs for attention computation.
|
||||
#[must_use]
|
||||
pub fn flops(
|
||||
&self,
|
||||
batch_size: usize,
|
||||
seq_len: usize,
|
||||
num_heads: usize,
|
||||
head_dim: usize,
|
||||
) -> u64 {
|
||||
// QK^T: 2 * batch * heads * seq * seq * dim
|
||||
// Softmax: 5 * batch * heads * seq * seq
|
||||
// AV: 2 * batch * heads * seq * seq * dim
|
||||
let qk_flops = 2 * batch_size * num_heads * seq_len * seq_len * head_dim;
|
||||
let softmax_flops = 5 * batch_size * num_heads * seq_len * seq_len;
|
||||
let av_flops = 2 * batch_size * num_heads * seq_len * seq_len * head_dim;
|
||||
(qk_flops + softmax_flops + av_flops) as u64
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FlashAttention {
|
||||
fn default() -> Self {
|
||||
Self::new(128)
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// KV Cache
|
||||
// ============================================================================
|
||||
|
||||
/// Key-Value cache for efficient autoregressive generation.
|
||||
#[derive(Debug)]
|
||||
pub struct KVCache {
|
||||
/// Configuration.
|
||||
pub config: KVCacheConfig,
|
||||
/// Current sequence length per layer.
|
||||
pub seq_lens: Vec<usize>,
|
||||
/// Total memory allocated in MB.
|
||||
pub allocated_memory_mb: f64,
|
||||
/// Whether paged attention is enabled.
|
||||
pub paged: bool,
|
||||
/// Block tables for paged attention.
|
||||
pub block_tables: Vec<Vec<usize>>,
|
||||
/// Free block list.
|
||||
pub free_blocks: Vec<usize>,
|
||||
/// Number of allocated blocks.
|
||||
pub num_allocated_blocks: usize,
|
||||
}
|
||||
|
||||
impl KVCache {
|
||||
/// Create a new KV cache.
|
||||
#[must_use]
|
||||
pub fn new(config: KVCacheConfig) -> Self {
|
||||
let seq_lens = vec![0; config.num_layers];
|
||||
let allocated_memory_mb = config.total_memory_gb() * 1024.0;
|
||||
let paged = config.paged_attention;
|
||||
|
||||
// Initialize block management for paged attention
|
||||
let total_blocks = if config.paged_attention {
|
||||
let tokens_per_block = config.block_size;
|
||||
let total_tokens = config.max_batch_size * config.max_seq_len;
|
||||
total_tokens.div_ceil(tokens_per_block)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let free_blocks = (0..total_blocks).collect();
|
||||
|
||||
Self {
|
||||
config,
|
||||
seq_lens,
|
||||
allocated_memory_mb,
|
||||
paged,
|
||||
block_tables: vec![],
|
||||
free_blocks,
|
||||
num_allocated_blocks: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Update cache for a layer with new tokens.
|
||||
pub fn update(&mut self, layer_id: usize, num_tokens: usize) {
|
||||
if layer_id < self.seq_lens.len() {
|
||||
self.seq_lens[layer_id] += num_tokens;
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current sequence length for a layer.
|
||||
#[must_use]
|
||||
pub fn get_seq_len(&self, layer_id: usize) -> usize {
|
||||
self.seq_lens.get(layer_id).copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Clear cache for all layers.
|
||||
pub fn clear(&mut self) {
|
||||
for len in &mut self.seq_lens {
|
||||
*len = 0;
|
||||
}
|
||||
if self.paged {
|
||||
self.block_tables.clear();
|
||||
self.free_blocks = (0..self.free_blocks.len() + self.num_allocated_blocks).collect();
|
||||
self.num_allocated_blocks = 0;
|
||||
}
|
||||
}
|
||||
|
||||
/// Allocate blocks for a new sequence.
|
||||
pub fn allocate_sequence(&mut self, seq_id: usize, initial_len: usize) -> bool {
|
||||
if !self.paged {
|
||||
return true;
|
||||
}
|
||||
|
||||
let blocks_needed = initial_len.div_ceil(self.config.block_size);
|
||||
if blocks_needed > self.free_blocks.len() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let allocated: Vec<usize> = self.free_blocks.drain(..blocks_needed).collect();
|
||||
self.num_allocated_blocks += allocated.len();
|
||||
|
||||
while self.block_tables.len() <= seq_id {
|
||||
self.block_tables.push(vec![]);
|
||||
}
|
||||
self.block_tables[seq_id] = allocated;
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Free blocks for a completed sequence.
|
||||
pub fn free_sequence(&mut self, seq_id: usize) {
|
||||
if !self.paged || seq_id >= self.block_tables.len() {
|
||||
return;
|
||||
}
|
||||
|
||||
let blocks = std::mem::take(&mut self.block_tables[seq_id]);
|
||||
self.num_allocated_blocks -= blocks.len();
|
||||
self.free_blocks.extend(blocks);
|
||||
}
|
||||
|
||||
/// Get memory usage ratio.
|
||||
#[must_use]
|
||||
pub fn memory_usage_ratio(&self) -> f64 {
|
||||
if self.paged {
|
||||
let total_blocks = self.free_blocks.len() + self.num_allocated_blocks;
|
||||
if total_blocks > 0 {
|
||||
self.num_allocated_blocks as f64 / total_blocks as f64
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
} else {
|
||||
let max_total: usize = self.seq_lens.len() * self.config.max_seq_len;
|
||||
let current_total: usize = self.seq_lens.iter().sum();
|
||||
if max_total > 0 {
|
||||
current_total as f64 / max_total as f64
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Calculate per-token memory in bytes.
|
||||
#[must_use]
|
||||
pub fn per_token_memory_bytes(&self) -> usize {
|
||||
let dtype_bytes = match self.config.dtype {
|
||||
DataType::Float32 => 4,
|
||||
DataType::Float16 | DataType::BFloat16 => 2,
|
||||
DataType::Int8 | DataType::FP8 => 1,
|
||||
DataType::Int4 => 1, // Rounded up
|
||||
};
|
||||
// K and V for all layers and heads
|
||||
2 * self.config.num_layers * self.config.num_kv_heads * self.config.head_dim * dtype_bytes
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Sharded KV Cache
|
||||
// ============================================================================
|
||||
|
||||
/// Sharded KV cache for distributed inference.
|
||||
#[derive(Debug)]
|
||||
pub struct ShardedKVCache {
|
||||
/// Local cache for this shard.
|
||||
pub local_cache: KVCache,
|
||||
/// Shard ID.
|
||||
pub shard_id: usize,
|
||||
/// Total number of shards.
|
||||
pub num_shards: usize,
|
||||
/// Layer range handled by this shard.
|
||||
pub layer_range: (usize, usize),
|
||||
}
|
||||
|
||||
impl ShardedKVCache {
|
||||
/// Create a new sharded KV cache.
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
config: KVCacheConfig,
|
||||
shard_id: usize,
|
||||
num_shards: usize,
|
||||
layer_start: usize,
|
||||
layer_end: usize,
|
||||
) -> Self {
|
||||
// Adjust config for this shard
|
||||
let shard_config = KVCacheConfig {
|
||||
num_layers: layer_end - layer_start,
|
||||
..config
|
||||
};
|
||||
|
||||
Self {
|
||||
local_cache: KVCache::new(shard_config),
|
||||
shard_id,
|
||||
num_shards,
|
||||
layer_range: (layer_start, layer_end),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this shard handles a given layer.
|
||||
#[must_use]
|
||||
pub fn handles_layer(&self, layer_id: usize) -> bool {
|
||||
layer_id >= self.layer_range.0 && layer_id < self.layer_range.1
|
||||
}
|
||||
|
||||
/// Get local layer index for a global layer ID.
|
||||
#[must_use]
|
||||
pub fn local_layer_index(&self, layer_id: usize) -> Option<usize> {
|
||||
if self.handles_layer(layer_id) {
|
||||
Some(layer_id - self.layer_range.0)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Update cache for a layer.
|
||||
pub fn update(&mut self, layer_id: usize, num_tokens: usize) {
|
||||
if let Some(local_idx) = self.local_layer_index(layer_id) {
|
||||
self.local_cache.update(local_idx, num_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get memory usage for this shard.
|
||||
#[must_use]
|
||||
pub fn memory_mb(&self) -> f64 {
|
||||
self.local_cache.allocated_memory_mb
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Attention Mask
|
||||
// ============================================================================
|
||||
|
||||
/// Attention mask types.
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Default)]
|
||||
pub enum AttentionMask {
|
||||
/// No mask (full attention).
|
||||
None,
|
||||
/// Causal mask (autoregressive).
|
||||
#[default]
|
||||
Causal,
|
||||
/// Sliding window attention.
|
||||
SlidingWindow { window_size: usize },
|
||||
/// Custom mask.
|
||||
Custom { mask: Vec<Vec<bool>> },
|
||||
}
|
||||
|
||||
impl AttentionMask {
|
||||
/// Create causal mask for given sequence length.
|
||||
#[must_use]
|
||||
pub fn causal(seq_len: usize) -> Vec<Vec<bool>> {
|
||||
(0..seq_len)
|
||||
.map(|i| (0..seq_len).map(|j| j <= i).collect())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Create sliding window mask.
|
||||
#[must_use]
|
||||
pub fn sliding_window(seq_len: usize, window_size: usize) -> Vec<Vec<bool>> {
|
||||
(0..seq_len)
|
||||
.map(|i| {
|
||||
(0..seq_len)
|
||||
.map(|j| j <= i && i.saturating_sub(window_size - 1) <= j)
|
||||
.collect()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get the effective attention span for a position.
|
||||
#[must_use]
|
||||
pub fn effective_span(&self, position: usize, seq_len: usize) -> usize {
|
||||
match self {
|
||||
Self::None => seq_len,
|
||||
Self::Causal => position + 1,
|
||||
Self::SlidingWindow { window_size } => (*window_size).min(position + 1),
|
||||
Self::Custom { mask } => {
|
||||
if position < mask.len() {
|
||||
mask[position].iter().filter(|&&m| m).count()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// Tests
|
||||
// ============================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_distributed_attention_creation() {
|
||||
let attn = DistributedAttention::new(64, 8, 128, 4);
|
||||
assert_eq!(attn.num_heads, 64);
|
||||
assert_eq!(attn.num_kv_heads, 8);
|
||||
assert_eq!(attn.gqa_ratio(), 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rope_cache_init() {
|
||||
let mut attn = DistributedAttention::new(32, 8, 128, 1);
|
||||
attn.init_rope_cache(1024);
|
||||
assert!(attn.rope_cache.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_attention_forward() {
|
||||
let attn = DistributedAttention::new(32, 8, 128, 1);
|
||||
let (batch, seq, dim) = attn.forward(512, 0);
|
||||
assert_eq!(batch, 1);
|
||||
assert_eq!(seq, 512);
|
||||
assert_eq!(dim, 32 * 128);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flash_attention_output() {
|
||||
let attn = DistributedAttention::new(32, 8, 128, 1);
|
||||
let output = attn.flash_forward(2048, 2048, 0);
|
||||
assert!(output.memory_saved_ratio > 0.0);
|
||||
assert!(output.num_blocks > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flash_attention_memory() {
|
||||
let flash = FlashAttention::new(128);
|
||||
let mem = flash.memory_mb(4096, 32, 2);
|
||||
// Should be much less than O(n^2)
|
||||
assert!(mem < 100.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flash_attention_flops() {
|
||||
let flash = FlashAttention::new(128);
|
||||
let flops = flash.flops(1, 1024, 32, 128);
|
||||
assert!(flops > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kv_cache_creation() {
|
||||
let config = KVCacheConfig::default();
|
||||
let cache = KVCache::new(config);
|
||||
assert_eq!(cache.seq_lens.len(), 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kv_cache_update() {
|
||||
let config = KVCacheConfig::default();
|
||||
let mut cache = KVCache::new(config);
|
||||
cache.update(0, 10);
|
||||
assert_eq!(cache.get_seq_len(0), 10);
|
||||
cache.update(0, 5);
|
||||
assert_eq!(cache.get_seq_len(0), 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kv_cache_clear() {
|
||||
let config = KVCacheConfig::default();
|
||||
let mut cache = KVCache::new(config);
|
||||
cache.update(0, 100);
|
||||
cache.update(1, 50);
|
||||
cache.clear();
|
||||
assert_eq!(cache.get_seq_len(0), 0);
|
||||
assert_eq!(cache.get_seq_len(1), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_paged_attention_allocation() {
|
||||
let config = KVCacheConfig {
|
||||
paged_attention: true,
|
||||
block_size: 16,
|
||||
max_seq_len: 256,
|
||||
max_batch_size: 4,
|
||||
..Default::default()
|
||||
};
|
||||
let mut cache = KVCache::new(config);
|
||||
assert!(cache.allocate_sequence(0, 32));
|
||||
assert_eq!(cache.num_allocated_blocks, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_paged_attention_free() {
|
||||
let config = KVCacheConfig {
|
||||
paged_attention: true,
|
||||
block_size: 16,
|
||||
max_seq_len: 256,
|
||||
max_batch_size: 4,
|
||||
..Default::default()
|
||||
};
|
||||
let mut cache = KVCache::new(config);
|
||||
cache.allocate_sequence(0, 32);
|
||||
let initial_free = cache.free_blocks.len();
|
||||
cache.free_sequence(0);
|
||||
assert_eq!(cache.free_blocks.len(), initial_free + 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sharded_kv_cache() {
|
||||
let config = KVCacheConfig {
|
||||
num_layers: 80,
|
||||
..Default::default()
|
||||
};
|
||||
let shard = ShardedKVCache::new(config, 0, 4, 0, 20);
|
||||
assert!(shard.handles_layer(10));
|
||||
assert!(!shard.handles_layer(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sharded_cache_local_index() {
|
||||
let config = KVCacheConfig::default();
|
||||
let shard = ShardedKVCache::new(config, 1, 4, 20, 40);
|
||||
assert_eq!(shard.local_layer_index(25), Some(5));
|
||||
assert_eq!(shard.local_layer_index(10), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_causal_mask() {
|
||||
let mask = AttentionMask::causal(4);
|
||||
assert!(mask[0][0]); // Can attend to position 0
|
||||
assert!(!mask[0][1]); // Cannot attend to future
|
||||
assert!(mask[3][0]); // Can attend to all previous
|
||||
assert!(mask[3][3]); // Can attend to self
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sliding_window_mask() {
|
||||
let mask = AttentionMask::sliding_window(10, 3);
|
||||
assert!(mask[5][5]); // Self
|
||||
assert!(mask[5][4]); // Previous
|
||||
assert!(mask[5][3]); // 2 positions back
|
||||
assert!(!mask[5][2]); // 3 positions back (outside window)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_effective_span() {
|
||||
let causal = AttentionMask::Causal;
|
||||
assert_eq!(causal.effective_span(5, 10), 6);
|
||||
|
||||
let window = AttentionMask::SlidingWindow { window_size: 3 };
|
||||
assert_eq!(window.effective_span(5, 10), 3);
|
||||
assert_eq!(window.effective_span(1, 10), 2);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user