thaw-vllm 0.1.0__tar.gz

thaw_vllm-0.1.0/PKG-INFO

@@ -0,0 +1,242 @@
+Metadata-Version: 2.4
+Name: thaw-vllm
+Version: 0.1.0
+Summary: Fast snapshot/restore for LLM inference. 17x faster cold starts, multi-GPU tensor parallel, KV cache snapshots.
+Author: Nils Matteson
+License: MIT
+Project-URL: Homepage, https://github.com/matteso1/thaw
+Project-URL: Repository, https://github.com/matteso1/thaw
+Keywords: llm,inference,vllm,gpu,cold-start,kv-cache
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Provides-Extra: vllm
+Requires-Dist: vllm>=0.18; extra == "vllm"
+Provides-Extra: serve
+Requires-Dist: vllm>=0.18; extra == "serve"
+Requires-Dist: fastapi; extra == "serve"
+Requires-Dist: uvicorn; extra == "serve"
+Provides-Extra: all
+Requires-Dist: vllm>=0.18; extra == "all"
+Requires-Dist: fastapi; extra == "all"
+Requires-Dist: uvicorn; extra == "all"
+
+# thaw
+
+**Fast snapshot/restore for LLM inference. 17x faster cold starts on 70B, multi-GPU tensor parallel, KV cache preservation.**
+
+vLLM cold-starts Llama-3-70B on 2x A100 in 546 seconds. thaw restores it in **31.8 seconds** — a **17.2x speedup**. Bit-identical outputs, verified by greedy decoding. Multi-GPU tensor parallel, Rust+CUDA pipelined DMA, and KV cache snapshots that no other tool offers.
+
+## Benchmarks
+
+**Llama-3-70B-Instruct (141 GB fp16) on 2x A100 SXM 80GB — tensor parallel:**
+
+| Method | Time | Speedup |
+|--------|------|---------|
+| Normal vLLM cold start | 546.5s | 1x |
+| **thaw restore (TP=2)** | **31.8s** | **17.2x** |
+| Weight restore only | 10.5s | 6.74 GB/s per rank |
+
+**Llama-3-8B-Instruct (16 GB fp16) — single GPU, H100 SXM:**
+
+| Method | Time | Throughput | Speedup |
+|--------|------|-----------|---------|
+| Normal vLLM cold start | 20.7s | — | 1x |
+| **thaw (NVMe)** | **3.7s** | 8.26 GB/s | **5.6x** |
+| **thaw (RAM hot path)** | **3.5s** | 10.69 GB/s | **5.9x** |
+
+**Agent fork — clone a running AI session (Llama-3-8B-Instruct, H100 SXM):**
+
+| Operation | Time | Notes |
+|-----------|------|-------|
+| Weight restore (Rust pipelined) | **1.1s** | **14.79 GB/s** — PCIe Gen5-saturating |
+| KV cache restore | **0.135s** | 65 blocks, 136 MB |
+| Total restore (incl. vLLM init) | **7.3s** | vs 16s normal cold start |
+| Fork 3 parallel completions | **1.6s avg** | All share 872-token cached prefix |
+
+All paths produce **bit-identical** inference output. KV cache restore preserves prefix cache across cold starts — new requests skip prefill entirely.
+
+<details>
+<summary>More GPUs and models</summary>
+
+| GPU | Model | Normal | thaw | Speedup |
+|-----|-------|--------|------|---------|
+| 2x A100 SXM 80GB | Llama-3-70B (TP=2) | 546.5s | 31.8s | **17.2x** |
+| H100 SXM 80GB | Llama-3-8B | 20.7s | 3.5s | **5.9x** |
+| RTX PRO 6000 (Blackwell) | Llama-3-8B | 28.6s | 3.2s | **8.9x** |
+| RTX A6000 | Llama-3-8B | 73.2s | 5.8s | **12.6x** |
+
+Larger models show bigger speedups because weight loading dominates more of the total cold start time.
+
+</details>
+
+## How it works
+
+```
+Normal vLLM cold start:
+  Download weights → deserialize safetensors → copy to GPU → init KV cache → ready
+  [==================================] 20.7s
+
+thaw restore:
+  Dummy init → DMA snapshot to GPU (pipelined, pinned memory, O_DIRECT)
+  [=====] 3.5s
+```
+
+**Freeze** captures all GPU state into binary snapshots — model weights (`.thaw`) and KV cache blocks (`.thawkv`).
+
+**Restore** initializes vLLM with dummy weights (fast — no disk I/O), then overwrites them from the snapshot using double-buffered pipelined DMA through pinned host memory. Two CUDA streams overlap PCIe transfers with disk reads. KV cache blocks are restored separately with their prefix cache hash mappings, so new requests immediately get cache hits.
+
+Two restore modes:
+- **Disk**: reads snapshot from NVMe with O_DIRECT, bypassing the kernel page cache. Throughput limited by NVMe bandwidth.
+- **RAM hot path**: snapshot pre-loaded in memory (tmpfs, shared memory, mmap). Pure PCIe DMA — 10.69 GB/s on H100. For production use where snapshots are pre-staged.
+
+**KV cache snapshots** capture the prefix-cached blocks that vLLM retains after generation. On restore, block data is DMA'd back to GPU and the prefix cache hash table is reconstructed. Requests with matching prefixes skip prefill — the most expensive part of inference.
+
+## Architecture
+
+```
+thaw/
+  crates/
+    thaw-core/       Rust. File format, region tables, I/O. No CUDA dep.
+    thaw-cuda-sys/   Rust. FFI bindings to CUDA runtime (cudaMallocHost,
+                     cudaMemcpyAsync, streams). Built via build.rs.
+    thaw-runtime/    Rust. Orchestration: freeze/restore pipelines, double-
+                     buffered DMA, O_DIRECT, MockCuda for Mac testing.
+    thaw-py/         Rust. PyO3 bindings exposing pipelined freeze/restore
+                     to Python. Builds a native .so via maturin.
+    thaw-cli/        Rust. GPU benchmark binary.
+  python/
+    thaw_vllm/       Python package (pip install thaw-vllm).
+      snapshot.py    Freeze/restore weights, Rust backend fallback.
+      kv_snapshot.py KV cache freeze/restore.
+      loader.py      vLLM ModelLoader: load_format="thaw".
+      server.py      OpenAI-compatible API server.
+      cli.py         CLI: thaw freeze, thaw serve, thaw info.
+    vllm_demo.py     End-to-end benchmark: normal vs thaw cold start.
+    kv_cache_demo.py KV cache snapshot/restore demo with correctness test.
+  demos/
+    agent_fork.py    Agent fork demo: clone session, fork parallel completions.
+```
+
+**Testing on Mac, shipping on GPU.** The `CudaBackend` trait abstracts all GPU operations. `MockCuda` (a HashMap-backed fake) lets 48 runtime tests run on any machine. The `cuda` feature flag activates real GPU paths only when needed.
+
+## Quick start
+
+```bash
+pip install thaw-vllm[serve]
+```
+
+```python
+import thaw_vllm
+from vllm import LLM, SamplingParams
+
+# Freeze: save model weights to a snapshot
+llm = LLM(model="meta-llama/Meta-Llama-3-8B", dtype="float16", enforce_eager=True)
+thaw_vllm.freeze_model_pipelined(model, "/path/to/weights.thaw")
+
+# Restore: two lines, 5.9x faster cold start
+llm = thaw_vllm.load("meta-llama/Meta-Llama-3-8B", "/path/to/weights.thaw")
+```
+
+Or use `load_format="thaw"` directly with vLLM:
+
+```python
+import thaw_vllm  # registers the loader
+llm = LLM(model="meta-llama/Meta-Llama-3-8B",
+          load_format="thaw",
+          model_loader_extra_config={"snapshot": "/path/to/weights.thaw"})
+```
+
+**Multi-GPU** — tensor parallel with per-rank snapshots:
+
+```python
+# Freeze: each GPU saves its shard
+llm = LLM(model="meta-llama/Meta-Llama-3-70B-Instruct", tensor_parallel_size=2, ...)
+thaw_vllm.freeze_model_tp(llm, "/path/to/weights.thaw")
+# Creates: weights.thaw (rank 0), weights.rank1.thaw (rank 1)
+
+# Restore: 17.2x faster than normal cold start
+llm = thaw_vllm.load("meta-llama/Meta-Llama-3-70B-Instruct", "/path/to/weights.thaw",
+                      tensor_parallel_size=2)
+```
+
+**Agent fork demo** — clone a running AI session, fork parallel completions:
+
+```bash
+python demos/agent_fork.py --snapshot weights.thaw
+python demos/agent_fork.py --snapshot weights.thaw --full-cycle  # destroy + restore
+```
+
+**CLI:**
+
+```bash
+thaw freeze --model meta-llama/Meta-Llama-3-8B --output weights.thaw
+thaw serve  --model meta-llama/Meta-Llama-3-8B --snapshot weights.thaw
+thaw info   weights.thaw
+```
+
+<details>
+<summary>Building with Rust+CUDA backend (optional, higher throughput)</summary>
+
+```bash
+git clone https://github.com/matteso1/thaw.git && cd thaw
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+source "$HOME/.cargo/env"
+pip install "maturin[patchelf]" vllm
+cd crates/thaw-py && maturin develop --release --features cuda && cd ../..
+pip install -e ".[serve]"
+```
+
+</details>
+
+## Competitive landscape
+
+The model loading space is active. Here's how thaw compares:
+
+| Project | Approach | Throughput | Limitations |
+|---------|----------|-----------|-------------|
+| **thaw** | Pipelined DMA, pinned memory, O_DIRECT + KV cache snapshot | 6.7-14.8 GB/s per GPU | — |
+| fastsafetensors (IBM) | GDS + 4x NVMe RAID0 | 26.4 GB/s | Requires GDS setup + RAID hardware |
+| NVIDIA Model Streamer | Multi-threaded concurrent streaming | ~2 GB/s (single SSD) | NVIDIA-maintained, less flexible |
+| CoreWeave Tensorizer | HTTP/S3 streaming + deserialization | ~4.6 GB/s local | Tied to CoreWeave ecosystem |
+| vLLM Sleep Mode | Offload to CPU RAM, reload | 0.26-3s | Not a cold start — requires prior warm load |
+| Modal GPU Snapshots | CUDA checkpoint/restore API | ~10x reduction | Alpha. Doesn't help with large model weight loading |
+| InferX | GPU runtime snapshotting | Claims 2s for 70B | No public code or benchmarks |
+
+**thaw's differentiation:**
+1. **KV cache snapshot/restore** — nobody else does this. Preserves prefix cache across cold starts, eliminates prefill. Enables agent forking, session migration, warm handoff.
+2. **Single NVMe performance** — most deployments don't have RAID0. thaw already matches or beats multi-threaded alternatives on one drive.
+3. **No special hardware** — no GDS, no RAID, no driver patches. Works on any CUDA 12+ GPU.
+
+See [docs/LANDSCAPE.md](./docs/LANDSCAPE.md) for detailed analysis.
+
+## Roadmap
+
+- [x] Weight snapshot/restore (pure Python path)
+- [x] Rust+CUDA pipelined freeze/restore (double-buffered DMA, O_DIRECT)
+- [x] RAM-backed restore path (PCIe-saturating, 10.69 GB/s)
+- [x] PyO3 bindings + vLLM integration shim
+- [x] H100 / A6000 / Blackwell benchmarks
+- [x] **KV cache snapshot/restore** — the moat (freeze/restore prefix-cached blocks, verified on Llama-3-8B)
+- [x] `pip install thaw-vllm` + CLI (`thaw freeze`, `thaw serve`, `thaw info`)
+- [x] `load_format="thaw"` — native vLLM ModelLoader integration
+- [x] OpenAI-compatible API server (`thaw serve`)
+- [x] Streaming support in API server (SSE, OpenAI-compatible)
+- [x] **Agent fork demo** — clone a running AI session, fork parallel completions from shared KV cache (full-cycle: 14.79 GB/s restore, 0.135s KV restore on H100 SXM)
+- [x] **Multi-GPU / tensor parallel** — 17.2x speedup on Llama-3-70B with 2x A100 (TP=2), bit-exact correctness verified
+- [ ] SGLang integration
+- [ ] Cloud snapshot storage (S3/GCS)
+- [ ] GPUDirect Storage support
+
+## Design
+
+Full technical architecture, file format spec, and rationale: [DESIGN.md](./DESIGN.md)
+
+## License
+
+MIT

thaw_vllm-0.1.0/README.md

@@ -0,0 +1,214 @@
+# thaw
+
+**Fast snapshot/restore for LLM inference. 17x faster cold starts on 70B, multi-GPU tensor parallel, KV cache preservation.**
+
+vLLM cold-starts Llama-3-70B on 2x A100 in 546 seconds. thaw restores it in **31.8 seconds** — a **17.2x speedup**. Bit-identical outputs, verified by greedy decoding. Multi-GPU tensor parallel, Rust+CUDA pipelined DMA, and KV cache snapshots that no other tool offers.
+
+## Benchmarks
+
+**Llama-3-70B-Instruct (141 GB fp16) on 2x A100 SXM 80GB — tensor parallel:**
+
+| Method | Time | Speedup |
+|--------|------|---------|
+| Normal vLLM cold start | 546.5s | 1x |
+| **thaw restore (TP=2)** | **31.8s** | **17.2x** |
+| Weight restore only | 10.5s | 6.74 GB/s per rank |
+
+**Llama-3-8B-Instruct (16 GB fp16) — single GPU, H100 SXM:**
+
+| Method | Time | Throughput | Speedup |
+|--------|------|-----------|---------|
+| Normal vLLM cold start | 20.7s | — | 1x |
+| **thaw (NVMe)** | **3.7s** | 8.26 GB/s | **5.6x** |
+| **thaw (RAM hot path)** | **3.5s** | 10.69 GB/s | **5.9x** |
+
+**Agent fork — clone a running AI session (Llama-3-8B-Instruct, H100 SXM):**
+
+| Operation | Time | Notes |
+|-----------|------|-------|
+| Weight restore (Rust pipelined) | **1.1s** | **14.79 GB/s** — PCIe Gen5-saturating |
+| KV cache restore | **0.135s** | 65 blocks, 136 MB |
+| Total restore (incl. vLLM init) | **7.3s** | vs 16s normal cold start |
+| Fork 3 parallel completions | **1.6s avg** | All share 872-token cached prefix |
+
+All paths produce **bit-identical** inference output. KV cache restore preserves prefix cache across cold starts — new requests skip prefill entirely.
+
+<details>
+<summary>More GPUs and models</summary>
+
+| GPU | Model | Normal | thaw | Speedup |
+|-----|-------|--------|------|---------|
+| 2x A100 SXM 80GB | Llama-3-70B (TP=2) | 546.5s | 31.8s | **17.2x** |
+| H100 SXM 80GB | Llama-3-8B | 20.7s | 3.5s | **5.9x** |
+| RTX PRO 6000 (Blackwell) | Llama-3-8B | 28.6s | 3.2s | **8.9x** |
+| RTX A6000 | Llama-3-8B | 73.2s | 5.8s | **12.6x** |
+
+Larger models show bigger speedups because weight loading dominates more of the total cold start time.
+
+</details>
+
+## How it works
+
+```
+Normal vLLM cold start:
+  Download weights → deserialize safetensors → copy to GPU → init KV cache → ready
+  [==================================] 20.7s
+
+thaw restore:
+  Dummy init → DMA snapshot to GPU (pipelined, pinned memory, O_DIRECT)
+  [=====] 3.5s
+```
+
+**Freeze** captures all GPU state into binary snapshots — model weights (`.thaw`) and KV cache blocks (`.thawkv`).
+
+**Restore** initializes vLLM with dummy weights (fast — no disk I/O), then overwrites them from the snapshot using double-buffered pipelined DMA through pinned host memory. Two CUDA streams overlap PCIe transfers with disk reads. KV cache blocks are restored separately with their prefix cache hash mappings, so new requests immediately get cache hits.
+
+Two restore modes:
+- **Disk**: reads snapshot from NVMe with O_DIRECT, bypassing the kernel page cache. Throughput limited by NVMe bandwidth.
+- **RAM hot path**: snapshot pre-loaded in memory (tmpfs, shared memory, mmap). Pure PCIe DMA — 10.69 GB/s on H100. For production use where snapshots are pre-staged.
+
+**KV cache snapshots** capture the prefix-cached blocks that vLLM retains after generation. On restore, block data is DMA'd back to GPU and the prefix cache hash table is reconstructed. Requests with matching prefixes skip prefill — the most expensive part of inference.
+
+## Architecture
+
+```
+thaw/
+  crates/
+    thaw-core/       Rust. File format, region tables, I/O. No CUDA dep.
+    thaw-cuda-sys/   Rust. FFI bindings to CUDA runtime (cudaMallocHost,
+                     cudaMemcpyAsync, streams). Built via build.rs.
+    thaw-runtime/    Rust. Orchestration: freeze/restore pipelines, double-
+                     buffered DMA, O_DIRECT, MockCuda for Mac testing.
+    thaw-py/         Rust. PyO3 bindings exposing pipelined freeze/restore
+                     to Python. Builds a native .so via maturin.
+    thaw-cli/        Rust. GPU benchmark binary.
+  python/
+    thaw_vllm/       Python package (pip install thaw-vllm).
+      snapshot.py    Freeze/restore weights, Rust backend fallback.
+      kv_snapshot.py KV cache freeze/restore.
+      loader.py      vLLM ModelLoader: load_format="thaw".
+      server.py      OpenAI-compatible API server.
+      cli.py         CLI: thaw freeze, thaw serve, thaw info.
+    vllm_demo.py     End-to-end benchmark: normal vs thaw cold start.
+    kv_cache_demo.py KV cache snapshot/restore demo with correctness test.
+  demos/
+    agent_fork.py    Agent fork demo: clone session, fork parallel completions.
+```
+
+**Testing on Mac, shipping on GPU.** The `CudaBackend` trait abstracts all GPU operations. `MockCuda` (a HashMap-backed fake) lets 48 runtime tests run on any machine. The `cuda` feature flag activates real GPU paths only when needed.
+
+## Quick start
+
+```bash
+pip install thaw-vllm[serve]
+```
+
+```python
+import thaw_vllm
+from vllm import LLM, SamplingParams
+
+# Freeze: save model weights to a snapshot
+llm = LLM(model="meta-llama/Meta-Llama-3-8B", dtype="float16", enforce_eager=True)
+thaw_vllm.freeze_model_pipelined(model, "/path/to/weights.thaw")
+
+# Restore: two lines, 5.9x faster cold start
+llm = thaw_vllm.load("meta-llama/Meta-Llama-3-8B", "/path/to/weights.thaw")
+```
+
+Or use `load_format="thaw"` directly with vLLM:
+
+```python
+import thaw_vllm  # registers the loader
+llm = LLM(model="meta-llama/Meta-Llama-3-8B",
+          load_format="thaw",
+          model_loader_extra_config={"snapshot": "/path/to/weights.thaw"})
+```
+
+**Multi-GPU** — tensor parallel with per-rank snapshots:
+
+```python
+# Freeze: each GPU saves its shard
+llm = LLM(model="meta-llama/Meta-Llama-3-70B-Instruct", tensor_parallel_size=2, ...)
+thaw_vllm.freeze_model_tp(llm, "/path/to/weights.thaw")
+# Creates: weights.thaw (rank 0), weights.rank1.thaw (rank 1)
+
+# Restore: 17.2x faster than normal cold start
+llm = thaw_vllm.load("meta-llama/Meta-Llama-3-70B-Instruct", "/path/to/weights.thaw",
+                      tensor_parallel_size=2)
+```
+
+**Agent fork demo** — clone a running AI session, fork parallel completions:
+
+```bash
+python demos/agent_fork.py --snapshot weights.thaw
+python demos/agent_fork.py --snapshot weights.thaw --full-cycle  # destroy + restore
+```
+
+**CLI:**
+
+```bash
+thaw freeze --model meta-llama/Meta-Llama-3-8B --output weights.thaw
+thaw serve  --model meta-llama/Meta-Llama-3-8B --snapshot weights.thaw
+thaw info   weights.thaw
+```
+
+<details>
+<summary>Building with Rust+CUDA backend (optional, higher throughput)</summary>
+
+```bash
+git clone https://github.com/matteso1/thaw.git && cd thaw
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+source "$HOME/.cargo/env"
+pip install "maturin[patchelf]" vllm
+cd crates/thaw-py && maturin develop --release --features cuda && cd ../..
+pip install -e ".[serve]"
+```
+
+</details>
+
+## Competitive landscape
+
+The model loading space is active. Here's how thaw compares:
+
+| Project | Approach | Throughput | Limitations |
+|---------|----------|-----------|-------------|
+| **thaw** | Pipelined DMA, pinned memory, O_DIRECT + KV cache snapshot | 6.7-14.8 GB/s per GPU | — |
+| fastsafetensors (IBM) | GDS + 4x NVMe RAID0 | 26.4 GB/s | Requires GDS setup + RAID hardware |
+| NVIDIA Model Streamer | Multi-threaded concurrent streaming | ~2 GB/s (single SSD) | NVIDIA-maintained, less flexible |
+| CoreWeave Tensorizer | HTTP/S3 streaming + deserialization | ~4.6 GB/s local | Tied to CoreWeave ecosystem |
+| vLLM Sleep Mode | Offload to CPU RAM, reload | 0.26-3s | Not a cold start — requires prior warm load |
+| Modal GPU Snapshots | CUDA checkpoint/restore API | ~10x reduction | Alpha. Doesn't help with large model weight loading |
+| InferX | GPU runtime snapshotting | Claims 2s for 70B | No public code or benchmarks |
+
+**thaw's differentiation:**
+1. **KV cache snapshot/restore** — nobody else does this. Preserves prefix cache across cold starts, eliminates prefill. Enables agent forking, session migration, warm handoff.
+2. **Single NVMe performance** — most deployments don't have RAID0. thaw already matches or beats multi-threaded alternatives on one drive.
+3. **No special hardware** — no GDS, no RAID, no driver patches. Works on any CUDA 12+ GPU.
+
+See [docs/LANDSCAPE.md](./docs/LANDSCAPE.md) for detailed analysis.
+
+## Roadmap
+
+- [x] Weight snapshot/restore (pure Python path)
+- [x] Rust+CUDA pipelined freeze/restore (double-buffered DMA, O_DIRECT)
+- [x] RAM-backed restore path (PCIe-saturating, 10.69 GB/s)
+- [x] PyO3 bindings + vLLM integration shim
+- [x] H100 / A6000 / Blackwell benchmarks
+- [x] **KV cache snapshot/restore** — the moat (freeze/restore prefix-cached blocks, verified on Llama-3-8B)
+- [x] `pip install thaw-vllm` + CLI (`thaw freeze`, `thaw serve`, `thaw info`)
+- [x] `load_format="thaw"` — native vLLM ModelLoader integration
+- [x] OpenAI-compatible API server (`thaw serve`)
+- [x] Streaming support in API server (SSE, OpenAI-compatible)
+- [x] **Agent fork demo** — clone a running AI session, fork parallel completions from shared KV cache (full-cycle: 14.79 GB/s restore, 0.135s KV restore on H100 SXM)
+- [x] **Multi-GPU / tensor parallel** — 17.2x speedup on Llama-3-70B with 2x A100 (TP=2), bit-exact correctness verified
+- [ ] SGLang integration
+- [ ] Cloud snapshot storage (S3/GCS)
+- [ ] GPUDirect Storage support
+
+## Design
+
+Full technical architecture, file format spec, and rationale: [DESIGN.md](./DESIGN.md)
+
+## License
+
+MIT

thaw_vllm-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "thaw-vllm"
+version = "0.1.0"
+description = "Fast snapshot/restore for LLM inference. 17x faster cold starts, multi-GPU tensor parallel, KV cache snapshots."
+readme = "README.md"
+license = {text = "MIT"}
+requires-python = ">=3.9"
+authors = [
+    {name = "Nils Matteson"},
+]
+keywords = ["llm", "inference", "vllm", "gpu", "cold-start", "kv-cache"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "torch>=2.0",
+]
+
+[project.optional-dependencies]
+vllm = ["vllm>=0.18"]
+serve = ["vllm>=0.18", "fastapi", "uvicorn"]
+all = ["vllm>=0.18", "fastapi", "uvicorn"]
+
+[project.scripts]
+thaw = "thaw_vllm.cli:main"
+
+[project.urls]
+Homepage = "https://github.com/matteso1/thaw"
+Repository = "https://github.com/matteso1/thaw"
+
+[tool.setuptools.packages.find]
+where = ["python"]
+
+[tool.setuptools.package-dir]
+"" = "python"

thaw_vllm-0.1.0/python/thaw_vllm/__init__.py

@@ -0,0 +1,87 @@
+# thaw_vllm — GPU snapshot/restore for vLLM model weights.
+#
+# This module provides two operations:
+#
+#   freeze_model(model, path)   — snapshot all model weights to a .thaw file
+#   restore_model(model, path)  — load weights from a .thaw file back onto GPU
+#
+# The file format is thaw's binary format (see crates/thaw-core). This
+# Python implementation reads/writes the same byte layout as the Rust
+# side, so files are interchangeable. The Python path uses PyTorch's
+# own CUDA memory operations under the hood — no custom FFI required.
+#
+# This is the MVP integration layer. In production, the Rust
+# implementation (via thaw-py/PyO3) replaces this for higher
+# throughput. The file format is identical either way.
+
+from thaw_vllm.snapshot import (
+    freeze_model,
+    freeze_model_pipelined,
+    freeze_model_tp,
+    restore_model,
+    restore_model_from_ram,
+    restore_model_pipelined,
+)
+from thaw_vllm.kv_snapshot import (
+    freeze_kv_cache,
+    freeze_kv_cache_tp,
+    restore_kv_cache,
+    restore_kv_cache_tp,
+)
+
+# Register load_format="thaw" with vLLM when available.
+try:
+    from thaw_vllm.loader import ThawModelLoader  # noqa: F401
+except ImportError:
+    # vLLM not installed — loader registration is optional.
+    pass
+
+
+def load(model: str, snapshot: str, kv_snapshot: str = None, **kwargs):
+    """One-line thaw-powered model loading.
+
+    Usage:
+        import thaw_vllm
+        llm = thaw_vllm.load("meta-llama/Meta-Llama-3-8B", "/path/to/weights.thaw")
+
+    Multi-GPU:
+        llm = thaw_vllm.load("meta-llama/Meta-Llama-3-70B", "/path/to/weights.thaw",
+                             tensor_parallel_size=4)
+    """
+    from vllm import LLM
+
+    kwargs.setdefault("enforce_eager", True)
+    kwargs.setdefault("dtype", "float16")
+
+    tp_size = kwargs.get("tensor_parallel_size", 1)
+
+    llm = LLM(
+        model=model,
+        load_format="thaw",
+        model_loader_extra_config={"snapshot": snapshot},
+        **kwargs,
+    )
+
+    if kv_snapshot:
+        if tp_size > 1:
+            restore_kv_cache_tp(llm, kv_snapshot)
+        else:
+            restore_kv_cache(llm, kv_snapshot)
+
+    return llm
+
+
+__all__ = [
+    "freeze_model",
+    "freeze_model_pipelined",
+    "freeze_model_tp",
+    "restore_model",
+    "restore_model_from_ram",
+    "restore_model_pipelined",
+    "freeze_kv_cache",
+    "freeze_kv_cache_tp",
+    "restore_kv_cache",
+    "restore_kv_cache_tp",
+    "load",
+    "ThawModelLoader",
+]