feat(core): Memory-Mapped Model Loading

[m96-chan]

Overview

Implement memory-mapped file loading for large SafeTensors models.

Motivation

Large models (70B+) exceed VRAM capacity. Memory mapping enables:

• Streaming weights from disk to GPU
• Layer-by-layer loading without full RAM usage
• Fast model switching without full reload

Features

• mmap-based SafeTensors reader
• Lazy tensor loading (load on first access)
• Pinned memory staging for faster H2D
• Layer eviction for multi-model scenarios
• Windows/Linux support

Design

from pygpukit.llm import LazyModel

# Model weights stay on disk until accessed
model = LazyModel.from_safetensors("path/to/model", lazy=True)

# Only loads embedding layer
embeddings = model.embed_tokens(input_ids)

# Loads layer 0 on demand
hidden = model.blocks[0](embeddings)

Related

• SafeTensors mmap support
• HuggingFace Accelerate disk offload

[m96-chan]

Implementation Plan

Approach: Rust-based LazyTensor

Extend existing rust/pygpukit-core/src/llm/tensor_loader.rs with lazy GPU loading.

Current Infrastructure (Already Implemented)

• memmap2::Mmap - Memory-mapped file access
• SafeTensorsFile - Header parsing, tensor metadata
• TensorData - Zero-copy view into mmap
• memory/ - LRU eviction, size-class allocator

New Components

// rust/pygpukit-core/src/llm/lazy_tensor.rs

pub struct LazyTensor {
    file: Arc<SafeTensorsFile>,      // Shared mmap reference
    name: String,
    info: TensorInfo,
    gpu_ptr: Option<CUdeviceptr>,    // Cached GPU allocation
    pinned_ptr: Option<*mut u8>,     // Pinned staging buffer
    last_access: Instant,            // For LRU eviction
}

impl LazyTensor {
    /// Load to GPU on first access
    pub fn to_gpu(&mut self, stream: CUstream) -> Result<CUdeviceptr> {
        if let Some(ptr) = self.gpu_ptr {
            self.last_access = Instant::now();
            return Ok(ptr);  // Cache hit
        }
        
        // 1. Get mmap slice (OS pages in on access)
        let data = self.file.tensor(&self.name)?;
        
        // 2. Allocate pinned staging buffer
        // 3. Copy mmap -> pinned
        // 4. Async copy pinned -> GPU (cuMemcpyHtoDAsync)
        // 5. Cache gpu_ptr
    }
    
    /// Evict from GPU (keep mmap reference)
    pub fn evict(&mut self) -> usize {
        if let Some(ptr) = self.gpu_ptr.take() {
            // Return freed bytes for memory accounting
        }
    }
}

pub struct LazyModelLoader {
    files: Vec<Arc<SafeTensorsFile>>,  // Multiple .safetensors
    tensors: HashMap<String, LazyTensor>,
    memory_budget: usize,              // Max VRAM usage
    lru_cache: LruCache<String>,       // Eviction order
}

Data Flow

Python: model.blocks[0].attn.q_proj.weight
         ↓ PyO3
Rust:   LazyTensor.to_gpu(stream)
         ├── mmap slice (OS pages in ~4KB chunks)
         ├── cuMemAllocHost (pinned staging)
         ├── memcpy mmap → pinned
         └── cuMemcpyHtoDAsync pinned → GPU
         ↓
C++:    GEMM kernel execution

Eviction Strategy

impl LazyModelLoader {
    fn ensure_memory(&mut self, required: usize) -> Result<()> {
        while self.used_vram + required > self.memory_budget {
            // Evict least recently used tensor
            let victim = self.lru_cache.pop_lru()?;
            self.used_vram -= self.tensors[&victim].evict();
        }
        Ok(())
    }
}

Tasks

• LazyTensor struct with GPU caching
• Pinned memory staging pool
• LRU eviction integration
• LazyModelLoader for multi-file models
• PyO3 bindings for LazyModel
• Python LazyModel.from_safetensors() API
• Windows/Linux testing

Why Rust?

Aspect	Rust	C++
mmap	✅ memmap2 existing	Need to implement
Memory pool	✅ LRU existing	Need to implement
Thread safety	✅ Ownership model	Manual
SafeTensors	✅ Crate existing	Need to implement

Starting implementation.

[m96-chan]

Closed via PR #175