mech.app
AI Agents

Thaw: Git-Branch Semantics for Running LLM Agents, How KV-Cache Snapshots Turn 340s Forks Into 0.88s

KV-cache snapshotting as infrastructure primitive: preserving inference state enables near-free agent branching for RL rollouts and parallel exploration.

Source: github.com
Thaw: Git-Branch Semantics for Running LLM Agents, How KV-Cache Snapshots Turn 340s Forks Into 0.88s

When an LLM agent explores multiple branches (RL rollouts, best-of-N sampling, parallel coding attempts), each branch typically re-runs prefill over the same shared context. You pay for the same prompt N times. Thaw snapshots a live inference session (weights, KV cache, scheduler state, prefix-hash table) and hydrates N children that diverge from the fork point without re-prefilling. Think of it as git branch for a running model.

The receipt from H100 80GB running Llama-3.1-8B: a pre-warmed pool boots once in 22.3s, then each fork round of 4 branches × 64 tokens runs in 0.88s median. Cold-boot equivalent would be ~340s per round. That is a ~400× amortized speedup. All rounds are bit-identical at the fork boundary, with full JSON receipts and reproducers in the repo.

NVIDIA shipped Dynamo Snapshot last week for fast pod cold-starts, but they explicitly free the KV cache before checkpoint. Thaw takes the opposite bet: preserve the KV cache so a fork is near-free. Different problem, opposite mechanic.

The Forking Problem in Agent Orchestration

Agent workflows that explore multiple paths hit a prefill bottleneck:

  • RL rollouts: Generate N trajectories from the same policy state.
  • Best-of-N sampling: Run N completions, pick the highest-scoring.
  • Parallel coding agents: Fork at a decision point, explore different implementations.

Traditional approach: boot N separate inference sessions, each re-prefilling the shared context. If your context is 2,048 tokens and you fork 8 branches, you prefill 16,384 tokens total. Only the divergent tokens (maybe 64 per branch) are unique.

Thaw’s approach: snapshot the inference session after prefill, restore N times, generate only the divergent tokens. You prefill 2,048 tokens once, then pay for 8 × 64 = 512 divergent tokens.

What Gets Snapshotted

Thaw captures four pieces of inference state:

  1. Model weights: The parameter tensors (already in GPU memory).
  2. KV cache: The key-value pairs from all previous tokens.
  3. Scheduler state: Which sequences are running, their positions, attention masks.
  4. Prefix-hash table: vLLM’s internal structure for prefix caching.

Traditional model checkpointing saves weights to disk for later training or serving. Thaw snapshots the entire runtime state of a live inference session, including ephemeral structures like the KV cache that are normally discarded between requests.

Memory and I/O Trade-Offs

Snapshot size: For Llama-3.1-8B with 2,048-token context:

  • Weights: ~16 GB (FP16)
  • KV cache: ~2 GB (depends on batch size, sequence length, layer count)
  • Scheduler + prefix-hash: ~100 MB

Total: ~18 GB per snapshot.

Hydration cost: Restoring a child session from snapshot is mostly a memory copy. On H100 with NVLink, copying 18 GB takes ~0.5s. Compare to cold-boot prefill of 2,048 tokens at ~85s (measured).

Memory overhead: If you fork 4 branches, you need 4 × KV cache space (~8 GB total for divergent tokens) plus the shared snapshot (~18 GB). Total: ~26 GB. Cold-booting 4 sessions would use 4 × 18 GB = 72 GB.

Table 1: Performance comparison: cold-boot vs. Thaw snapshot on H100 80GB, Llama-3.1-8B, 2,048-token shared context, 4 branches × 64 divergent tokens

ApproachPrefill Time (4 branches)Memory FootprintLatency per Fork Round
Cold-boot~340s72 GB~340s
Thaw snapshot22.3s (once)26 GB0.88s
Amortized speedup~400×2.8× smaller~400×

How vLLM and SGLang State Management Interacts

Both vLLM and SGLang manage inference state internally:

  • vLLM: Uses a SequenceGroup abstraction. Each sequence has a SequenceData object holding token IDs, a LogicalTokenBlock for KV cache allocation, and a SequenceStatus (running, finished, etc.). The scheduler tracks which blocks are in GPU memory.
  • SGLang: Similar architecture with a RequestState object and a RadixCache for prefix sharing.

Thaw hooks into these internals to serialize and restore state. The critical requirement: the fork boundary must be bit-identical. If the KV cache or scheduler state drifts between snapshot and restore, you get silent divergence (wrong tokens, broken attention masks).

What breaks if the fork boundary is not bit-identical:

  • Attention masks point to the wrong KV positions.
  • Prefix-hash lookups fail, forcing re-prefill.
  • Logits drift, causing different token sampling downstream.

Thaw validates bit-identity by comparing the first generated token after restore. If it does not match the expected value, the snapshot is invalid.

Implementation Shape

Thaw exposes two primitives:

from thaw_vllm import ThawEngine

# Boot once, prefill shared context
engine = ThawEngine(model="meta-llama/Llama-3.1-8B")
snapshot = engine.snapshot(prompt="Shared context here...")

# Fork 4 branches
children = snapshot.fork(n=4)
results = [child.generate(prompt=f"Branch {i}") for i, child in enumerate(children)]

Under the hood:

  1. snapshot() serializes weights, KV cache, scheduler state, and prefix-hash table to a memory buffer.
  2. fork(n=4) hydrates 4 child ThawEngine instances from the buffer.
  3. Each child runs generate() independently, starting from the fork point.

The snapshot is stored in pinned CPU memory (for multi-GPU setups) or GPU memory (for single-GPU). Hydration is a cudaMemcpy from the snapshot buffer to each child’s GPU memory.

Failure Modes

Snapshot invalidation: If the model weights change (e.g., LoRA adapter swap), the snapshot is invalid. Thaw does not detect this automatically. You must manually invalidate and re-snapshot.

OOM on fork: If you fork too many branches, you run out of GPU memory for divergent KV caches. Forking 16 branches on H100 80GB requires 16 × 2GB KV cache = 32GB, leaving 48GB for weights and scheduler. This is insufficient for Llama-3.1-70B (70GB weights alone). Thaw does not implement eviction or paging. You hit OOM and crash.

Scheduler state corruption: If vLLM’s internal scheduler state changes between snapshot and restore (e.g., a sequence finishes early), the restore fails. Thaw assumes the scheduler is in a stable state at the fork boundary.

Prefix-hash collisions: If two different prompts hash to the same prefix, vLLM’s prefix cache returns the wrong KV blocks. Thaw inherits this risk from vLLM.

Instrumentation Hooks

Thaw does not ship with built-in observability, but exposes hooks for custom monitoring:

  • snapshot.size(): Returns snapshot size in bytes (useful for tracking memory growth as context length increases).
  • fork.hydration_latency(): Returns per-child hydration time in seconds.
  • Memory pressure: Wrap fork() with torch.cuda.memory_allocated() to track how close you are to OOM.
  • Bit-identity validation: Hook into the first token comparison to log failures.

You can instrument these yourself by wrapping snapshot() and fork() with timing and memory profiling.

Deployment Shape

Thaw runs in-process with vLLM or SGLang. You do not deploy it separately. The snapshot lives in GPU or CPU memory, not on disk. If the process crashes, the snapshot is lost.

For production use, you would need to:

  • Persist snapshots to shared storage (S3, NFS) for crash recovery.
  • Implement snapshot versioning (track which model weights, which context length).
  • Add health checks for bit-identity validation.
  • Integrate with your orchestration layer (Kubernetes, Ray) to manage snapshot lifecycle.

Thaw does not provide any of this. It is a primitive, not a framework.

The Framework vs. Primitive Question

The author asks: should Thaw be a low-level primitive (current shape) or a framework node (LangGraph, TRL)?

Primitive advantages:

  • Works with any orchestration layer.
  • No opinions about agent architecture.
  • Easy to compose with existing tools.

Framework advantages:

  • Handles snapshot lifecycle automatically.
  • Integrates with observability (LangSmith, Weights & Biases).
  • Provides higher-level abstractions (e.g., “fork at decision point”).

The primitive shape is correct for now. Wrapping it in a framework would hide the plumbing and make it harder to debug when things break. Once the primitive stabilizes, framework integrations can be built on top.

Technical Verdict

Use Thaw when:

  • You run RL rollouts, best-of-N sampling, or parallel exploration where shared context exceeds 2,048 tokens and you fork 4+ branches per inference session.
  • Your context is large (>1K tokens) and your branches are short (<100 tokens).
  • You have GPU memory headroom for N × KV cache (e.g., H100 80GB can handle 8 branches of Llama-3.1-8B with 2K context).
  • You can tolerate in-process state (no crash recovery).

Avoid Thaw when:

  • Your divergent branches exceed 500 tokens each. The snapshot overhead becomes dominant and you lose the amortization benefit.
  • You need crash recovery or distributed snapshots across multiple nodes.
  • Your model weights change frequently (LoRA swaps, fine-tuning loops).
  • You are running on shared GPU infrastructure where memory is tight (e.g., 24GB consumer GPUs).

The 400× speedup is real, but it is amortized. If you only fork once, the 22.3s boot cost wipes out the savings. Thaw pays off when you fork repeatedly from the same snapshot.