feat(core): Memory-Mapped Model Loading #159

Activity

1. m96-chan

   added

   enhancementNew feature or requestNew feature or request

   on Dec 30, 2025
2. m96-chan commented on Dec 31, 2025

   

   m96-chan

   on Dec 31, 2025

   OwnerAuthor

   More actions

   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.
3. m96-chan

   added 2 commits that reference this issue on Dec 31, 2025

   feat(llm): add lazy model loading with streaming strategies (#159)

   ...

   7d7de55

   feat: add cuBLAS loader, kernel profiler, and lazy model loading (#134,

   ...

   26df666
4. m96-chan commented on Dec 31, 2025

   

   m96-chan

   on Dec 31, 2025

   OwnerAuthor

   More actions

   Closed via PR #175
5. m96-chan

   closed this as completedon Dec 31, 2025
6. m96-chan

   added 2 commits that reference this issue on Jan 1, 2026

   feat(diffusion): add FLUX.1 image generation module (#178)

   ...

   0b5c13e

   Release v0.2.19 (#188)

   ...

   7adbe5f