stream-ai-runtime 0.1.0__py3-none-any.whl

stream_ai_runtime-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,8 @@
+stream_runtime/__init__.py,sha256=X0UZqQqGaUhxkpCvrNDKkRjeckrARZdBp6Ac5RDdZok,120
+stream_runtime/buffer_manager.py,sha256=f_zPpsT57-62o_LqhNDAEYIdG-C8OW1nefOTmoqsHx0,1385
+stream_runtime/config.py,sha256=i8Nhp-RUN-tQA2AYYBb5e9wrf1cDwi7XHaqa7RsElmM,127
+stream_runtime/engine.py,sha256=r2UW3X1-dtlXeTV-OHH_D4siJ4qxUfJ9vFdyiX57lOY,1040
+stream_ai_runtime-0.1.0.dist-info/METADATA,sha256=QEVWm5rB-GL_LF_ckG6zQgfCC5os65aSMq-rFr84MJw,128
+stream_ai_runtime-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+stream_ai_runtime-0.1.0.dist-info/top_level.txt,sha256=a8-xOiPZcRLIs3shDUsk0O6-OrhFewv5TOpKFLSHONQ,15
+stream_ai_runtime-0.1.0.dist-info/RECORD,,

stream_ai_runtime-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

stream_ai_runtime-0.1.0.dist-info/top_level.txt

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

stream_runtime/__init__.py

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

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_runtime/config.py

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

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