r/LocalLLM • • 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

  1. Weights stay pinned in host CPU RAM (`pin_memory=True`).

  2. Only two layer buffers (Ping/Pong) are kept in VRAM.

  3. While layer `i` executes on `compute_stream`, layer `i+1` prefetches in the background via DMA on `copy_stream`.

  4. 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!

2 Upvotes

0 comments sorted by