stream-ai-runtime 0.1.0__tar.gz

stream_ai_runtime-0.1.0/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: stream-ai-runtime
+Version: 0.1.0
+Summary: Zero-OOM streaming runtime for AI models
+Dynamic: summary

stream_ai_runtime-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

stream_ai_runtime-0.1.0/setup.py

@@ -0,0 +1,7 @@
+from setuptools import setup, find_packages
+setup(
+    name="stream-ai-runtime", 
+    version="0.1.0", 
+    packages=find_packages(),
+    description="Zero-OOM streaming runtime for AI models"
+)

stream_ai_runtime-0.1.0/stream_ai_runtime.egg-info/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: stream-ai-runtime
+Version: 0.1.0
+Summary: Zero-OOM streaming runtime for AI models
+Dynamic: summary

stream_ai_runtime-0.1.0/stream_ai_runtime.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+setup.py
+stream_ai_runtime.egg-info/PKG-INFO
+stream_ai_runtime.egg-info/SOURCES.txt
+stream_ai_runtime.egg-info/dependency_links.txt
+stream_ai_runtime.egg-info/top_level.txt
+stream_runtime/__init__.py
+stream_runtime/buffer_manager.py
+stream_runtime/config.py
+stream_runtime/engine.py

stream_ai_runtime-0.1.0/stream_ai_runtime.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

stream_ai_runtime-0.1.0/stream_ai_runtime.egg-info/top_level.txt

@@ -0,0 +1 @@
+stream_runtime

stream_ai_runtime-0.1.0/stream_runtime/__init__.py

@@ -0,0 +1,3 @@
+from .config import StreamRuntimeConfig
+from .buffer_manager import DoubleBufferManager
+from .engine import StreamModel

stream_ai_runtime-0.1.0/stream_runtime/buffer_manager.py

@@ -0,0 +1,30 @@
+import torch
+
+class DoubleBufferManager:
+    def __init__(self, shape, dtype=torch.float32, device="cuda"):
+        self.device = device
+        self.buffer_ping = torch.empty(shape, dtype=dtype, device=device)
+        self.buffer_pong = torch.empty(shape, dtype=dtype, device=device)
+        self.current_slot = 0
+        self.compute_stream = torch.cuda.Stream()
+        self.copy_stream = torch.cuda.Stream()
+        self.copy_done = [torch.cuda.Event(), torch.cuda.Event()]
+        self.compute_done = [torch.cuda.Event(), torch.cuda.Event()]
+
+    def prefetch_next_chunk(self, cpu_tensor):
+        next_slot = 1 - self.current_slot
+        with torch.cuda.stream(self.copy_stream):
+            self.compute_done[next_slot].wait(self.copy_stream)
+            target = self.buffer_pong if next_slot == 1 else self.buffer_ping
+            target.copy_(cpu_tensor, non_blocking=True)
+            self.copy_done[next_slot].record(self.copy_stream)
+
+    def get_ready_buffer(self):
+        with torch.cuda.stream(self.compute_stream):
+            self.copy_done[self.current_slot].wait(self.compute_stream)
+            ready_buf = self.buffer_ping if self.current_slot == 0 else self.buffer_pong
+            return ready_buf
+
+    def mark_compute_complete(self):
+        self.compute_done[self.current_slot].record(self.compute_stream)
+        self.current_slot = 1 - self.current_slot

stream_ai_runtime-0.1.0/stream_runtime/config.py

@@ -0,0 +1,6 @@
+from dataclasses import dataclass
+
+@dataclass
+class StreamRuntimeConfig:
+    chunk_size_mb: int = 128
+    device: str = "cuda"

stream_ai_runtime-0.1.0/stream_runtime/engine.py

@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+from .buffer_manager import DoubleBufferManager
+
+class StreamModel:
+    def __init__(self, num_layers=32, dim=2048):
+        self.num_layers = num_layers
+        self.dim = dim
+        self.cpu_weights = [torch.randn(dim, dim, pin_memory=True) for _ in range(num_layers)]
+        self.buffer_mgr = DoubleBufferManager((dim, dim))
+
+    def forward(self, x):
+        with torch.cuda.stream(self.buffer_mgr.copy_stream):
+            self.buffer_mgr.buffer_ping.copy_(self.cpu_weights[0], non_blocking=True)
+            self.buffer_mgr.copy_done[0].record(self.buffer_mgr.copy_stream)
+
+        for i in range(self.num_layers):
+            if i + 1 < self.num_layers:
+                self.buffer_mgr.prefetch_next_chunk(self.cpu_weights[i + 1])
+            buf = self.buffer_mgr.get_ready_buffer()
+            with torch.cuda.stream(self.buffer_mgr.compute_stream):
+                x = torch.matmul(x, buf)
+            self.buffer_mgr.mark_compute_complete()
+
+        torch.cuda.synchronize()
+        return x