789 lines
23 KiB
Rust
789 lines
23 KiB
Rust
//! Weight Streaming for Large Models
|
|
//!
|
|
//! This module provides on-demand weight loading for models larger than GPU memory,
|
|
//! using an LRU cache to keep frequently used layers in GPU memory while streaming
|
|
//! others from host memory or disk.
|
|
|
|
use crate::device::MetalDevice;
|
|
use crate::error::{MetalError, Result};
|
|
use crate::memory::{MetalBuffer, MetalBufferUsage};
|
|
use std::collections::{HashMap, VecDeque};
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
use parking_lot::{Mutex, RwLock};
|
|
use tracing::{debug, trace, warn};
|
|
|
|
/// Unique identifier for a layer's weights
|
|
pub type LayerId = u64;
|
|
|
|
/// Source for weight data
|
|
#[derive(Debug, Clone)]
|
|
pub enum WeightSource {
|
|
/// Memory-mapped file
|
|
MmappedFile {
|
|
/// Path to the file
|
|
path: PathBuf,
|
|
/// Offset in file
|
|
offset: usize,
|
|
/// Size in bytes
|
|
size: usize,
|
|
},
|
|
/// Host memory
|
|
HostMemory {
|
|
/// Data in host memory
|
|
data: Arc<Vec<f16>>,
|
|
},
|
|
/// Already in GPU memory (no streaming needed)
|
|
GpuResident,
|
|
}
|
|
|
|
/// Half-precision float for compact storage
|
|
#[allow(non_camel_case_types)]
|
|
pub type f16 = u16; // Use u16 as f16 placeholder
|
|
|
|
/// Configuration for weight streaming
|
|
#[derive(Debug, Clone)]
|
|
pub struct StreamingConfig {
|
|
/// Maximum GPU memory budget for weights (bytes)
|
|
pub gpu_budget: usize,
|
|
/// Number of layers to prefetch ahead
|
|
pub prefetch_lookahead: usize,
|
|
/// Enable async prefetching
|
|
pub async_prefetch: bool,
|
|
/// Eviction policy
|
|
pub eviction_policy: EvictionPolicy,
|
|
/// Pin certain layers (never evict)
|
|
pub pinned_layers: Vec<LayerId>,
|
|
}
|
|
|
|
impl Default for StreamingConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
gpu_budget: 4 * 1024 * 1024 * 1024, // 4GB
|
|
prefetch_lookahead: 2,
|
|
async_prefetch: true,
|
|
eviction_policy: EvictionPolicy::Lru,
|
|
pinned_layers: Vec::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Policy for evicting weights from GPU cache
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
pub enum EvictionPolicy {
|
|
/// Least Recently Used
|
|
Lru,
|
|
/// Least Frequently Used
|
|
Lfu,
|
|
/// First In First Out
|
|
Fifo,
|
|
}
|
|
|
|
/// Statistics for weight streaming
|
|
#[derive(Debug, Default, Clone)]
|
|
pub struct StreamingStats {
|
|
/// Total weight loads from source
|
|
pub total_loads: u64,
|
|
/// Cache hits (already in GPU)
|
|
pub cache_hits: u64,
|
|
/// Cache misses (needed to load)
|
|
pub cache_misses: u64,
|
|
/// Total bytes loaded
|
|
pub bytes_loaded: u64,
|
|
/// Total bytes evicted
|
|
pub bytes_evicted: u64,
|
|
/// Average load time (microseconds)
|
|
pub avg_load_time_us: f64,
|
|
/// Current cache size (bytes)
|
|
pub current_cache_size: usize,
|
|
/// Number of cached layers
|
|
pub cached_layers: usize,
|
|
}
|
|
|
|
impl StreamingStats {
|
|
/// Get cache hit rate
|
|
pub fn hit_rate(&self) -> f64 {
|
|
let total = self.cache_hits + self.cache_misses;
|
|
if total == 0 {
|
|
0.0
|
|
} else {
|
|
self.cache_hits as f64 / total as f64
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Entry in the weight cache
|
|
struct CacheEntry {
|
|
/// GPU buffer containing weights
|
|
buffer: MetalBuffer<f16>,
|
|
/// Size in bytes
|
|
size: usize,
|
|
/// Last access time
|
|
last_access: Instant,
|
|
/// Access count (for LFU)
|
|
access_count: u64,
|
|
/// Whether this entry is pinned
|
|
pinned: bool,
|
|
/// Source for reloading
|
|
source: WeightSource,
|
|
}
|
|
|
|
/// Weight Streaming Manager
|
|
pub struct WeightStreamingManager {
|
|
/// Device
|
|
device: Arc<MetalDevice>,
|
|
/// Configuration
|
|
config: StreamingConfig,
|
|
/// LRU cache of weights in GPU memory
|
|
cache: RwLock<HashMap<LayerId, CacheEntry>>,
|
|
/// LRU order tracking
|
|
lru_order: Mutex<VecDeque<LayerId>>,
|
|
/// Current cache size in bytes
|
|
current_size: Mutex<usize>,
|
|
/// Statistics
|
|
stats: RwLock<StreamingStats>,
|
|
/// Prefetch queue
|
|
prefetch_queue: Mutex<VecDeque<LayerId>>,
|
|
/// Weight sources
|
|
sources: RwLock<HashMap<LayerId, WeightSource>>,
|
|
}
|
|
|
|
impl WeightStreamingManager {
|
|
/// Create a new weight streaming manager
|
|
pub fn new(device: Arc<MetalDevice>, config: StreamingConfig) -> Self {
|
|
Self {
|
|
device,
|
|
config,
|
|
cache: RwLock::new(HashMap::new()),
|
|
lru_order: Mutex::new(VecDeque::new()),
|
|
current_size: Mutex::new(0),
|
|
stats: RwLock::new(StreamingStats::default()),
|
|
prefetch_queue: Mutex::new(VecDeque::new()),
|
|
sources: RwLock::new(HashMap::new()),
|
|
}
|
|
}
|
|
|
|
/// Create with default configuration
|
|
pub fn default_config(device: Arc<MetalDevice>) -> Self {
|
|
Self::new(device, StreamingConfig::default())
|
|
}
|
|
|
|
/// Register a weight source
|
|
pub fn register_weights(&self, layer_id: LayerId, source: WeightSource) {
|
|
let mut sources = self.sources.write();
|
|
sources.insert(layer_id, source);
|
|
}
|
|
|
|
/// Register weights from host memory
|
|
pub fn register_from_host(&self, layer_id: LayerId, data: Vec<f16>) {
|
|
self.register_weights(
|
|
layer_id,
|
|
WeightSource::HostMemory {
|
|
data: Arc::new(data),
|
|
},
|
|
);
|
|
}
|
|
|
|
/// Register weights from file
|
|
pub fn register_from_file(&self, layer_id: LayerId, path: PathBuf, offset: usize, size: usize) {
|
|
self.register_weights(
|
|
layer_id,
|
|
WeightSource::MmappedFile { path, offset, size },
|
|
);
|
|
}
|
|
|
|
/// Get weights for a layer, loading from source if needed
|
|
pub fn get_weights(&self, layer_id: LayerId) -> Result<WeightRef> {
|
|
// Check if already cached
|
|
{
|
|
let cache = self.cache.read();
|
|
if cache.contains_key(&layer_id) {
|
|
// Update stats and access time
|
|
self.record_hit(layer_id);
|
|
return Ok(WeightRef {
|
|
manager: self,
|
|
layer_id,
|
|
});
|
|
}
|
|
}
|
|
|
|
// Cache miss - need to load
|
|
self.load_weights(layer_id)?;
|
|
Ok(WeightRef {
|
|
manager: self,
|
|
layer_id,
|
|
})
|
|
}
|
|
|
|
/// Check if weights are cached
|
|
pub fn is_cached(&self, layer_id: LayerId) -> bool {
|
|
self.cache.read().contains_key(&layer_id)
|
|
}
|
|
|
|
/// Load weights from source into GPU cache
|
|
fn load_weights(&self, layer_id: LayerId) -> Result<()> {
|
|
let source = {
|
|
let sources = self.sources.read();
|
|
sources.get(&layer_id).cloned()
|
|
};
|
|
|
|
let source = source.ok_or_else(|| {
|
|
MetalError::internal(format!("No source registered for layer {}", layer_id))
|
|
})?;
|
|
|
|
let start = Instant::now();
|
|
|
|
let (buffer, size) = match &source {
|
|
WeightSource::HostMemory { data } => {
|
|
let size = data.len() * std::mem::size_of::<f16>();
|
|
|
|
// Evict if needed
|
|
self.evict_for_space(size)?;
|
|
|
|
// Create GPU buffer
|
|
let buffer = MetalBuffer::<f16>::new(
|
|
&self.device,
|
|
data.len(),
|
|
MetalBufferUsage::Shared,
|
|
)?;
|
|
|
|
// Copy data
|
|
buffer.write(data)?;
|
|
|
|
(buffer, size)
|
|
}
|
|
WeightSource::MmappedFile { path, offset, size } => {
|
|
// Evict if needed
|
|
self.evict_for_space(*size)?;
|
|
|
|
// In a real implementation, we would mmap the file
|
|
// For now, we read it conventionally
|
|
let data = std::fs::read(path).map_err(|e| MetalError::internal(e.to_string()))?;
|
|
|
|
let num_elements = (*size) / std::mem::size_of::<f16>();
|
|
let buffer = MetalBuffer::<f16>::new(
|
|
&self.device,
|
|
num_elements,
|
|
MetalBufferUsage::Shared,
|
|
)?;
|
|
|
|
// Copy relevant portion
|
|
let start_elem = (*offset) / std::mem::size_of::<f16>();
|
|
let end_elem = start_elem + num_elements;
|
|
|
|
if end_elem * 2 <= data.len() {
|
|
let slice: Vec<f16> = data[*offset..*offset + *size]
|
|
.chunks(2)
|
|
.map(|c| u16::from_le_bytes([c[0], c.get(1).copied().unwrap_or(0)]))
|
|
.collect();
|
|
buffer.write(&slice)?;
|
|
}
|
|
|
|
(buffer, *size)
|
|
}
|
|
WeightSource::GpuResident => {
|
|
return Err(MetalError::internal("Cannot load GPU-resident weights"));
|
|
}
|
|
};
|
|
|
|
let duration = start.elapsed();
|
|
|
|
// Add to cache
|
|
{
|
|
let mut cache = self.cache.write();
|
|
cache.insert(
|
|
layer_id,
|
|
CacheEntry {
|
|
buffer,
|
|
size,
|
|
last_access: Instant::now(),
|
|
access_count: 1,
|
|
pinned: self.config.pinned_layers.contains(&layer_id),
|
|
source,
|
|
},
|
|
);
|
|
}
|
|
|
|
// Update LRU order
|
|
{
|
|
let mut lru = self.lru_order.lock();
|
|
lru.push_back(layer_id);
|
|
}
|
|
|
|
// Update size tracking
|
|
{
|
|
let mut current = self.current_size.lock();
|
|
*current += size;
|
|
}
|
|
|
|
// Update stats
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.total_loads += 1;
|
|
stats.cache_misses += 1;
|
|
stats.bytes_loaded += size as u64;
|
|
stats.cached_layers += 1;
|
|
stats.current_cache_size += size;
|
|
|
|
// Update average load time
|
|
let total = stats.total_loads;
|
|
stats.avg_load_time_us = (stats.avg_load_time_us * (total - 1) as f64
|
|
+ duration.as_micros() as f64)
|
|
/ total as f64;
|
|
}
|
|
|
|
debug!(
|
|
"Loaded weights for layer {} ({} bytes) in {:?}",
|
|
layer_id, size, duration
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Evict weights to make space
|
|
fn evict_for_space(&self, needed: usize) -> Result<()> {
|
|
let current = *self.current_size.lock();
|
|
|
|
if current + needed <= self.config.gpu_budget {
|
|
return Ok(());
|
|
}
|
|
|
|
let to_free = current + needed - self.config.gpu_budget;
|
|
let mut freed = 0;
|
|
|
|
while freed < to_free {
|
|
let evict_id = self.select_for_eviction();
|
|
|
|
match evict_id {
|
|
Some(id) => {
|
|
let evicted_size = self.evict(id)?;
|
|
freed += evicted_size;
|
|
}
|
|
None => {
|
|
warn!("Cannot free enough space for weights: need {}, freed {}", to_free, freed);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Select a layer for eviction based on policy
|
|
fn select_for_eviction(&self) -> Option<LayerId> {
|
|
match self.config.eviction_policy {
|
|
EvictionPolicy::Lru => {
|
|
let lru = self.lru_order.lock();
|
|
let cache = self.cache.read();
|
|
|
|
// Find first non-pinned entry
|
|
lru.iter()
|
|
.find(|&&id| {
|
|
cache.get(&id).map(|e| !e.pinned).unwrap_or(false)
|
|
})
|
|
.copied()
|
|
}
|
|
EvictionPolicy::Lfu => {
|
|
let cache = self.cache.read();
|
|
cache
|
|
.iter()
|
|
.filter(|(_, e)| !e.pinned)
|
|
.min_by_key(|(_, e)| e.access_count)
|
|
.map(|(&id, _)| id)
|
|
}
|
|
EvictionPolicy::Fifo => {
|
|
let lru = self.lru_order.lock();
|
|
let cache = self.cache.read();
|
|
|
|
lru.iter()
|
|
.find(|&&id| {
|
|
cache.get(&id).map(|e| !e.pinned).unwrap_or(false)
|
|
})
|
|
.copied()
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Evict a specific layer
|
|
fn evict(&self, layer_id: LayerId) -> Result<usize> {
|
|
let size = {
|
|
let mut cache = self.cache.write();
|
|
if let Some(entry) = cache.remove(&layer_id) {
|
|
entry.size
|
|
} else {
|
|
return Ok(0);
|
|
}
|
|
};
|
|
|
|
// Remove from LRU order
|
|
{
|
|
let mut lru = self.lru_order.lock();
|
|
lru.retain(|&id| id != layer_id);
|
|
}
|
|
|
|
// Update size tracking
|
|
{
|
|
let mut current = self.current_size.lock();
|
|
*current = current.saturating_sub(size);
|
|
}
|
|
|
|
// Update stats
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.bytes_evicted += size as u64;
|
|
stats.cached_layers = stats.cached_layers.saturating_sub(1);
|
|
stats.current_cache_size = stats.current_cache_size.saturating_sub(size);
|
|
}
|
|
|
|
trace!("Evicted layer {} ({} bytes)", layer_id, size);
|
|
|
|
Ok(size)
|
|
}
|
|
|
|
/// Record a cache hit
|
|
fn record_hit(&self, layer_id: LayerId) {
|
|
// Update access time and count
|
|
{
|
|
let mut cache = self.cache.write();
|
|
if let Some(entry) = cache.get_mut(&layer_id) {
|
|
entry.last_access = Instant::now();
|
|
entry.access_count += 1;
|
|
}
|
|
}
|
|
|
|
// Move to back of LRU queue
|
|
{
|
|
let mut lru = self.lru_order.lock();
|
|
lru.retain(|&id| id != layer_id);
|
|
lru.push_back(layer_id);
|
|
}
|
|
|
|
// Update stats
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.cache_hits += 1;
|
|
}
|
|
}
|
|
|
|
/// Request async prefetch for upcoming layers
|
|
pub fn prefetch(&self, layer_ids: &[LayerId]) {
|
|
let mut queue = self.prefetch_queue.lock();
|
|
for &id in layer_ids {
|
|
if !self.is_cached(id) && !queue.contains(&id) {
|
|
queue.push_back(id);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Prefetch layers ahead of current
|
|
pub fn prefetch_ahead(&self, current_layer: LayerId) {
|
|
let lookahead = self.config.prefetch_lookahead;
|
|
let layers: Vec<LayerId> = (current_layer + 1..=current_layer + lookahead as u64).collect();
|
|
self.prefetch(&layers);
|
|
}
|
|
|
|
/// Process prefetch queue
|
|
pub fn process_prefetch(&self) -> Result<usize> {
|
|
let mut processed = 0;
|
|
|
|
loop {
|
|
let layer_id = {
|
|
let mut queue = self.prefetch_queue.lock();
|
|
queue.pop_front()
|
|
};
|
|
|
|
match layer_id {
|
|
Some(id) => {
|
|
if !self.is_cached(id) {
|
|
if let Err(e) = self.load_weights(id) {
|
|
warn!("Prefetch failed for layer {}: {}", id, e);
|
|
} else {
|
|
processed += 1;
|
|
}
|
|
}
|
|
}
|
|
None => break,
|
|
}
|
|
}
|
|
|
|
Ok(processed)
|
|
}
|
|
|
|
/// Pin a layer (prevent eviction)
|
|
pub fn pin_layer(&self, layer_id: LayerId) {
|
|
let mut cache = self.cache.write();
|
|
if let Some(entry) = cache.get_mut(&layer_id) {
|
|
entry.pinned = true;
|
|
}
|
|
}
|
|
|
|
/// Unpin a layer
|
|
pub fn unpin_layer(&self, layer_id: LayerId) {
|
|
let mut cache = self.cache.write();
|
|
if let Some(entry) = cache.get_mut(&layer_id) {
|
|
entry.pinned = false;
|
|
}
|
|
}
|
|
|
|
/// Clear cache
|
|
pub fn clear(&self) {
|
|
let mut cache = self.cache.write();
|
|
cache.clear();
|
|
|
|
let mut lru = self.lru_order.lock();
|
|
lru.clear();
|
|
|
|
let mut current = self.current_size.lock();
|
|
*current = 0;
|
|
|
|
let mut stats = self.stats.write();
|
|
stats.cached_layers = 0;
|
|
stats.current_cache_size = 0;
|
|
}
|
|
|
|
/// Get statistics
|
|
pub fn stats(&self) -> StreamingStats {
|
|
self.stats.read().clone()
|
|
}
|
|
|
|
/// Get configuration
|
|
pub fn config(&self) -> &StreamingConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Get current cache size
|
|
pub fn cache_size(&self) -> usize {
|
|
*self.current_size.lock()
|
|
}
|
|
|
|
/// Get number of cached layers
|
|
pub fn cached_count(&self) -> usize {
|
|
self.cache.read().len()
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Debug for WeightStreamingManager {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
let stats = self.stats();
|
|
f.debug_struct("WeightStreamingManager")
|
|
.field("cached_layers", &stats.cached_layers)
|
|
.field("cache_size", &stats.current_cache_size)
|
|
.field("hit_rate", &format!("{:.1}%", stats.hit_rate() * 100.0))
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
/// Reference to cached weights
|
|
pub struct WeightRef<'a> {
|
|
manager: &'a WeightStreamingManager,
|
|
layer_id: LayerId,
|
|
}
|
|
|
|
impl<'a> WeightRef<'a> {
|
|
/// Get the buffer
|
|
pub fn buffer(&self) -> Option<impl std::ops::Deref<Target = MetalBuffer<f16>> + '_> {
|
|
let cache = self.manager.cache.read();
|
|
if cache.contains_key(&self.layer_id) {
|
|
Some(parking_lot::RwLockReadGuard::map(cache, |c| {
|
|
&c.get(&self.layer_id).unwrap().buffer
|
|
}))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Get the layer ID
|
|
pub fn layer_id(&self) -> LayerId {
|
|
self.layer_id
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_streaming_config_default() {
|
|
let config = StreamingConfig::default();
|
|
assert_eq!(config.gpu_budget, 4 * 1024 * 1024 * 1024);
|
|
assert_eq!(config.prefetch_lookahead, 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_streaming_stats_hit_rate() {
|
|
let mut stats = StreamingStats::default();
|
|
stats.cache_hits = 80;
|
|
stats.cache_misses = 20;
|
|
assert!((stats.hit_rate() - 0.8).abs() < 0.001);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_streaming_manager_creation() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = WeightStreamingManager::default_config(device);
|
|
|
|
assert_eq!(manager.cached_count(), 0);
|
|
assert_eq!(manager.cache_size(), 0);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_register_and_load() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = WeightStreamingManager::default_config(device);
|
|
|
|
// Register weights from host
|
|
let data: Vec<f16> = (0..1024).map(|i| i as f16).collect();
|
|
manager.register_from_host(1, data);
|
|
|
|
// Load weights
|
|
let _weight_ref = manager.get_weights(1).unwrap();
|
|
assert!(manager.is_cached(1));
|
|
assert_eq!(manager.cached_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_cache_hit() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = WeightStreamingManager::default_config(device);
|
|
|
|
let data: Vec<f16> = (0..1024).map(|i| i as f16).collect();
|
|
manager.register_from_host(1, data);
|
|
|
|
// First access - miss
|
|
let _ref1 = manager.get_weights(1).unwrap();
|
|
|
|
// Second access - hit
|
|
let _ref2 = manager.get_weights(1).unwrap();
|
|
|
|
let stats = manager.stats();
|
|
assert_eq!(stats.cache_misses, 1);
|
|
assert_eq!(stats.cache_hits, 1);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_eviction() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let config = StreamingConfig {
|
|
gpu_budget: 4096, // Very small budget
|
|
..Default::default()
|
|
};
|
|
let manager = WeightStreamingManager::new(device, config);
|
|
|
|
// Register multiple layers that exceed budget
|
|
for i in 0..5 {
|
|
let data: Vec<f16> = (0..512).map(|j| (i * 512 + j) as f16).collect();
|
|
manager.register_from_host(i as LayerId, data);
|
|
let _ = manager.get_weights(i as LayerId);
|
|
}
|
|
|
|
// Should have evicted some
|
|
let stats = manager.stats();
|
|
assert!(stats.bytes_evicted > 0);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_pinning() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let config = StreamingConfig {
|
|
gpu_budget: 2048,
|
|
..Default::default()
|
|
};
|
|
let manager = WeightStreamingManager::new(device, config);
|
|
|
|
// Load and pin first layer
|
|
let data: Vec<f16> = (0..256).map(|i| i as f16).collect();
|
|
manager.register_from_host(1, data);
|
|
let _ = manager.get_weights(1).unwrap();
|
|
manager.pin_layer(1);
|
|
|
|
// Load more layers to trigger eviction
|
|
for i in 2..5 {
|
|
let data: Vec<f16> = (0..256).map(|j| (i * 256 + j) as f16).collect();
|
|
manager.register_from_host(i as LayerId, data);
|
|
let _ = manager.get_weights(i as LayerId);
|
|
}
|
|
|
|
// Pinned layer should still be cached
|
|
assert!(manager.is_cached(1));
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_prefetch() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = WeightStreamingManager::default_config(device);
|
|
|
|
// Register layers
|
|
for i in 0..5 {
|
|
let data: Vec<f16> = (0..256).map(|j| (i * 256 + j) as f16).collect();
|
|
manager.register_from_host(i as LayerId, data);
|
|
}
|
|
|
|
// Request prefetch
|
|
manager.prefetch(&[1, 2, 3]);
|
|
|
|
// Process prefetch
|
|
let processed = manager.process_prefetch().unwrap();
|
|
assert_eq!(processed, 3);
|
|
|
|
// All should be cached
|
|
assert!(manager.is_cached(1));
|
|
assert!(manager.is_cached(2));
|
|
assert!(manager.is_cached(3));
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_clear() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = WeightStreamingManager::default_config(device);
|
|
|
|
// Load some weights
|
|
for i in 0..3 {
|
|
let data: Vec<f16> = (0..256).map(|j| (i * 256 + j) as f16).collect();
|
|
manager.register_from_host(i as LayerId, data);
|
|
let _ = manager.get_weights(i as LayerId);
|
|
}
|
|
|
|
assert!(manager.cached_count() > 0);
|
|
|
|
manager.clear();
|
|
assert_eq!(manager.cached_count(), 0);
|
|
assert_eq!(manager.cache_size(), 0);
|
|
}
|
|
}
|