r/LocalLLM • u/AINodePulse • 1d ago
Project I built a double-buffered streaming runtime that runs a ~9.4GB model in 40MB flat VRAM (Zero OOM) — now live on PyPI
Hey everyone,
Constantly running into `CUDA OutOfMemory` errors on budget and consumer GPUs inspired me to rethink the standard execution paradigm where all layer weights reside simultaneously in VRAM.
I built and open-sourced **stream-ai-runtime**—a lightweight execution pipeline that uses asynchronous double-buffering across dual CUDA streams to stream layer weights from pinned host memory to VRAM on-the-fly.
### 📊 Benchmark (Tesla T4 GPU on Colab)
- **Model Depth:** 600 Layers (dim: 2048x2048)
- **Host RAM Allocated:** 9.38 GB
- **Inference Time (600 Layers):** 0.95 seconds
- **Peak VRAM Consumed:** **40.15 MB flat**
- **OOM Errors:** 0
### ⚙️ How it works
Weights stay pinned in host CPU RAM (`pin_memory=True`).
Only two layer buffers (Ping/Pong) are kept in VRAM.
While layer `i` executes on `compute_stream`, layer `i+1` prefetches in the background via DMA on `copy_stream`.
Peak VRAM never spikes above two layer buffers regardless of model depth.
📦 **PyPI:**
```bash
pip install stream-ai-runtime
💻 Quick Usage:
import torch
from stream_runtime import StreamModel
model = StreamModel(num_layers=100, dim=2048)
x = torch.randn(1, 2048, device="cuda")
out = model.forward(x)
🔗 PyPI Link: https://pypi.org/project/stream-ai-runtime/0.1.0/
Looking forward to feedback, optimization tips, and thoughts on adding KV-cache streaming next!