From Fragmented Caches to Unified Control: Implementing the cuzk Memory Manager

Introduction

Memory management in GPU-accelerated proving engines is a problem of competing constraints. The hardware demands large, contiguous allocations of pinned (page-locked) memory for direct GPU access. The proof protocols require loading multi-gigabyte Structured Reference Strings (SRS) and Pre-Compiled Constraint Evaluators (PCEs) for each proof type. The operator wants to maximize throughput by running as many concurrent proofs as the hardware can support. And the operating system will terminate the process if total memory exceeds physical RAM plus swap. Balancing these forces is the central challenge of production-grade proving infrastructure.

The cuzk (CUDA ZK proving daemon) engine had reached a point where its memory management was no longer sustainable. Configuration parameters existed but were never enforced. Caches grew without bound. Static concurrency limits bore no relationship to actual memory pressure. The system worked—most of the time—but it was fragile, unpredictable, and incapable of adapting to different hardware configurations or workload mixes.

This article examines Segment 15 of the opencode coding session, in which the assistant implemented the core of a new unified memory manager for cuzk, following a detailed specification document (cuzk-memory-manager.md) designed in the preceding segment [1]. The implementation spanned seven major source files, involved dozens of function signature changes, and fundamentally shifted the engine's memory model from a collection of independent, uncoordinated budgets to a single byte-level pool with on-demand loading and LRU eviction.

The Pre-Existing State: A System of Half-Measures

Before this segment, cuzk's memory management was characterized by what the assistant described as "fragmented and advisory" controls. The working_memory_budget configuration parameter existed in the config file but was never read by any runtime code—it was dead code, giving operators the illusion of control while having zero effect. The pinned_budget parameter was purely advisory: it logged a warning when exceeded but still proceeded to load SRS data regardless. The partition_workers setting was a static count that limited concurrent partition proofs, but it was not memory-aware—setting it too high caused out-of-memory (OOM) crashes, while setting it too low starved the GPU of work.

The SRS manager stored data in a HashMap and had no internal eviction mechanism. The only way to release SRS memory was through an external gRPC RPC call that was rarely used in production. PCE data was stored in four static OnceLock<PreCompiledCircuit<Fr>> globals in pipeline.rs, written once at startup and never cleared. If the system needed to support PoRep (~44 GiB SRS), SnapDeals (~33 GiB), and WindowPoSt (~57 GiB) simultaneously, the baseline memory footprint could reach approximately 203 GiB, consuming resources that could otherwise be used for productive proving work.

The assistant's earlier analysis, documented in the preceding segments, had traced every allocation and deallocation point through the 32 GiB PoRep pipeline, measured the transient spike during SRS loading (where an mmap'd file and a CUDA pinned allocation coexist briefly, creating an ~88 GiB peak), and identified the ungated PCE cold-start extraction that could spawn a ~40 GiB transient allocation without warning. This analysis formed the empirical basis for the specification that this segment would implement.

The Architecture: A Unified Byte-Level Budget

The new memory manager, as specified in cuzk-memory-manager.md and implemented in this segment, replaces the fragmented old system with a single unified budget. All major memory consumers—SRS pinned memory, PCE heap memory, and synthesis working memory—share one byte-level pool. The budget is auto-detected from system RAM by reading /proc/meminfo and subtracting a configurable safety margin (default 5 GiB). Operators can override the auto-detection with an explicit total_budget value if needed.

The core types are MemoryBudget and MemoryReservation, defined in the new memory.rs module. MemoryBudget tracks total and available bytes using atomic operations and provides async acquire() and synchronous try_acquire() methods. MemoryReservation is an RAII-style guard that releases its budget when dropped, ensuring that temporary allocations (like working memory for partition synthesis) are automatically accounted for. The module also defines estimation constants for each proof type, encoding the memory requirements for SRS, PCE, and working set per partition.

The design distinguishes between two kinds of allocations: permanent (SRS and PCE entries that persist until evicted) and temporary (working memory that is released after synthesis or proving). Permanent allocations bypass the MemoryReservation mechanism and directly manipulate the budget's internal counters through methods like into_permanent(), which converts a temporary reservation into a permanent allocation that does not get released on drop. This distinction is critical because SRS and PCE entries have lifetimes measured in minutes or hours, while working memory is held for seconds.

The Implementation Sequence: Building in Dependency Order

The assistant executed the implementation in a carefully ordered sequence of steps, each building on the previous one. The plan proceeded in strict dependency order: foundational types first, then configuration, then the components that consume them, and finally the engine integration that wires everything together.

Step 1: Creating memory.rs — The Foundation

The first step was to create the memory.rs module, containing MemoryBudget, MemoryReservation, detect_system_memory(), and all estimation constants. This module has no dependencies on other cuzk modules, making it the natural starting point. The assistant wrote the entire file in a single edit, defining the atomic counters, the async acquire loop with tokio Notify for waking blocked tasks, the RAII reservation guard, and the system memory detection logic that reads /proc/meminfo.

The estimation constants encode the memory footprint of each proof type. For example, PoRep SRS requires approximately 44 GiB of pinned memory, while WindowPoSt SRS requires approximately 57 GiB. These constants are used by the budget system to determine whether sufficient memory is available before loading begins. The detect_system_memory() function reads /proc/meminfo and parses the MemTotal line, providing an automatic baseline that adapts to different machine sizes.

Step 2: Registering the Module

The second step was a single-line edit to lib.rs: adding pub mod memory;. This small change is architecturally essential—without it, the memory module would be invisible to the rest of the crate. It is the gate that opens the dependency chain, allowing config.rs, srs_manager.rs, pipeline.rs, and engine.rs to reference MemoryBudget and MemoryReservation.

Step 3: Updating config.rs — Unified Budget Configuration

The third step was a multi-edit rewrite of config.rs. The assistant replaced the old dead-code fields (pinned_budget, working_memory_budget, partition_workers, preload) with the new unified budget configuration: total_budget, safety_margin, and eviction_min_idle. A parse_duration helper was added to support human-readable duration strings like "5min" for the eviction idle threshold.

Crucially, the assistant implemented deprecation warnings for the old fields. When the config parser encounters pinned_budget, working_memory_budget, partition_workers, or preload in the raw TOML string, it logs a warning directing operators to the new fields. This backward compatibility mechanism ensures that existing configurations continue to work while gently guiding users toward the new system. The test suite was also updated to reflect the new configuration structure.

The parse_duration helper deserves special mention. It parses strings like "5min", "30s", "2h" into Duration values, providing a human-friendly interface for the eviction idle threshold. The implementation handles common units (seconds, minutes, hours) and falls back to a default on parse failure, ensuring robustness.

Step 4: Rewriting srs_manager.rs — Budget-Aware SRS Loading

The fourth step was arguably the most strategically important single edit in the entire segment. The assistant replaced the entire #[cfg(feature = "cuda-supraseal")] implementation of SrsManager with a budget-aware version. The new SrsManager adds:

Step 5: Replacing Static PCE Caches with PceCache

The fifth step was the replacement of the four static OnceLock<PreCompiledCircuit<Fr>> globals in pipeline.rs with a new PceCache struct. This was the linchpin of the entire memory manager: without evictable PCE storage, the unified budget would have no way to reclaim the ~26 GiB per circuit type that PCEs occupy.

The PceCache struct provides eviction support, budget integration, and reference-counted access via Arc<PreCompiledCircuit<Fr>>. Where the old OnceLock statics offered write-once immutability with no lifecycle management, PceCache allows entries to be loaded on demand, evicted under pressure, and accessed through shared ownership that lets callers hold temporary references even after the cache has evicted the entry.

The assistant then updated all four extract_and_cache_pce_from_* functions to accept &PceCache instead of writing to statics, and updated synthesize_auto to take an optional &PceCache parameter. All nine call sites of synthesize_auto were updated in a systematic refactoring that used a combination of individual edits and a bulk sed command to pass None as the PCE cache placeholder. The old preload_pce_from_disk function was removed entirely, as loading is now on-demand rather than eager.

Step 6: Wiring the Engine

The sixth step began the integration of the memory manager into engine.rs, the central orchestrator of the proving pipeline. The assistant laid out a five-point plan:

  1. Add reservation field to SynthesizedJob
  2. Update Engine::new() to create SrsManager with the budget
  3. Update Engine::start() to create the budget, PceCache, wire the evictor, remove preload, and update channel sizing
  4. Update partition dispatch to replace partition_semaphore with budget.acquire()
  5. Update the GPU worker loop to implement two-phase reservation release The first two sub-steps were completed in this segment: the reservation field was added to SynthesizedJob, and Engine::new() was updated to wire the budget into SrsManager. The remaining changes—startup, partition dispatch, GPU worker loop, and evictor wiring—were left pending but the foundation was fully laid.

The Async/Blocking Boundary Challenge

One of the most subtle design challenges the assistant confronted was the transfer of budget reservations across the async/blocking boundary. The SRS loading flow works as follows: an async caller (the GPU worker loop) determines that a particular SRS parameter set is needed, checks budget availability, and spawns a blocking task to actually load the SRS from disk via C++ FFI. The challenge is that the async caller cannot hold a MemoryReservation across the spawn_blocking boundary because the reservation is an RAII guard that would be dropped when the async context yields.

The assistant's reasoning shows an iterative refinement through multiple approaches. The final solution uses a try_acquire method inside the blocking context: ensure_loaded checks the cache first, then attempts a non-blocking try_acquire to grab budget only if actually loading. If try_acquire fails, it returns an error, the async caller releases the lock, does an async acquire to wait for budget, and retries. This pattern correctly handles the race condition where another task loads the SRS between the check and the spawn, and avoids the need to call async methods from blocking contexts.

This design is a good example of the kind of systems-level thinking that the assistant brought to the implementation. The async/blocking boundary is a well-known source of bugs in Rust async code, and the assistant's solution is both correct and pragmatic.

The Eviction Policy: Conservative by Design

The eviction policy implemented in this segment is deliberately conservative. Entries are only evicted under memory pressure—never proactively. The eviction_min_idle threshold (default 5 minutes) ensures that recently used entries are not evicted prematurely. Additionally, the eviction logic checks the Arc::strong_count of each entry before evicting, ensuring that no in-flight proof is currently using the data being evicted.

This conservatism reflects the high cost of reloading SRS and PCE data. SRS loading involves mmap'ing a ~44 GiB file, allocating CUDA pinned memory, and deserializing points—a process that creates a transient ~88 GiB peak. PCE extraction can require ~40 GiB of transient memory for building a RecordingCS. The system should avoid reloading unless memory is genuinely constrained, and the 5-minute idle threshold provides a buffer against thrashing.

The eviction logic in SrsManager::evictable_entries() iterates over all cached entries, filtering those that have been idle longer than the threshold. The entries are sorted by idle time (oldest first), and the method returns the minimal set of entries whose combined budget would satisfy the requested amount. This ensures that eviction is both minimal (evicting only what is needed) and fair (evicting the least recently used entries first).

What Was Accomplished and What Remains

By the end of this segment, the assistant had:

  1. Created the memory.rs module with all core types and estimation constants
  2. Updated config.rs with the unified budget configuration and deprecation warnings
  3. Rewritten SrsManager to be budget-aware with eviction support
  4. Replaced static OnceLock PCE caches with a PceCache struct
  5. Updated all four extract_and_cache_pce_from_* functions to use PceCache
  6. Updated synthesize_auto with an optional PceCache parameter at all nine call sites
  7. Added the reservation field to SynthesizedJob in engine.rs
  8. Wired the budget into SrsManager in Engine::new() The remaining work in engine.rs—startup logic, partition dispatch, GPU worker loop, and evictor wiring—was explicitly noted as pending. The segment summary states: "The remaining engine.rs changes (startup, partition dispatch, GPU worker loop, evictor wiring) are still pending but the foundation is fully laid." This is a deliberate and sensible approach. The engine integration is the most complex part of the refactoring, involving interactions with the GPU worker loop, the partition dispatch semaphore, and the startup sequence. By completing the foundational changes first and deferring the engine integration, the assistant ensured that each component could be verified independently before being wired into the critical proving path.

The Broader Context

This segment is part of a larger arc of work spanning multiple segments. Earlier segments had fixed a WindowPoSt PCE crash (Segment 0), resolved GPU race conditions in multi-GPU setups (Segments 2–3), built a Docker container for mainnet proving (Segments 4–5), deployed a vast.ai management system (Segments 6–7), resolved OOM failures in low-RAM instances (Segment 8), hardened the platform with benchmark error reporting (Segments 9–10), investigated and fixed critical PSProve PoRep failures (Segments 11–13), and designed the memory manager specification (Segment 14).

The memory manager implementation in this segment is the culmination of all that preceding work. The OOM failures in low-RAM instances (Segment 8) demonstrated the need for dynamic, hardware-aware configuration. The PSProve investigation (Segments 11–13) revealed the fragility of the static cache system. The specification design (Segment 14) synthesized these lessons into a coherent architecture. Now, in Segment 15, that architecture is being brought to life.

Conclusion

Segment 15 represents a systematic, methodical implementation of a complex architectural transformation. The assistant built the new memory manager in dependency order, starting with foundational types and ending with engine integration. Each step was verified before proceeding to the next, and the design decisions—from the async/blocking budget transfer pattern to the conservative eviction policy—were carefully reasoned and documented.

The result is a memory manager that transforms cuzk from a system with fragmented, advisory, and non-evictable memory management into one with a unified byte-level budget, on-demand loading, and LRU eviction under pressure. The foundation is fully laid for the remaining engine integration work, and the codebase now has the infrastructure to prevent OOM crashes, maximize memory utilization, and adapt dynamically to different workloads and machine sizes.

The approach taken in this segment—building the foundation first, deferring the most complex integration until the supporting infrastructure is verified—is a textbook example of how to execute a large-scale refactoring in a production system. It minimizes risk, enables incremental testing, and ensures that each component can be validated independently before being wired into the critical path.

References

[1] "The Unified Memory Manager: A Systematic Refactoring of the cuzk GPU Proving Engine" — Chunk article for Segment 15, Chunk 0, covering the detailed implementation of the memory manager core.