Designing a Memory-Aware Admission Control System for the cuzk GPU Proving Engine

Introduction

In any high-performance compute system, there comes a moment when the architecture must evolve from "it works" to "it works reliably under all conditions." For the cuzk GPU proving engine—a critical piece of Filecoin's storage proof infrastructure—that moment arrived in the form of a comprehensive memory management redesign that spanned an entire segment of an opencode coding session. What began as a user's open-ended request to "look at major allocation/deallocation points" culminated in a detailed architectural specification that replaces a fragile, static concurrency limit with a robust, memory-aware admission control system.

This article traces the arc of that transformation: the forensic audit that uncovered dead configuration options and quantified every buffer in the pipeline, the collaborative design process that refined the architecture through iterative Q&A, and the final specification that crystallized the design into an actionable blueprint. For anyone building memory-intensive GPU compute systems, this segment offers a masterclass in how to approach architectural redesign—starting with evidence, designing through trade-off analysis, and ending with a document that captures not just what to build, but why.

The Problem: A Proving Engine with a Memory Blind Spot

The cuzk engine is a GPU-accelerated Rust application that generates Groth16 zero-knowledge proofs for Filecoin's proof-of-replication (PoRep), WindowPoSt, and WinningPoSt protocols. It operates on a pipelined architecture: CPU-bound circuit synthesis produces intermediate proving state, which is then consumed by GPU-bound proof computation. This two-phase pipeline is designed to maximize throughput on expensive server hardware, but it creates a fundamental tension: how do you run many proof partitions concurrently without crossing the line into an out-of-memory (OOM) crash?

The existing system managed this tension with a set of static knobs. The primary throttle was a partition_workers semaphore—a simple counter that limited the number of concurrent partition synthesis tasks to a fixed value, typically 12. There was also a working_memory_budget configuration option that defaulted to "80GiB," a pinned_budget for SRS memory, and an SRS preload list that eagerly loaded every listed circuit at startup. In theory, these knobs allowed operators to tune memory usage. In practice, as the forensic audit would reveal, the working_memory_budget was entirely dead code—parsed from TOML configuration but never checked anywhere in the engine's dispatch logic. The only real throttle was the partition_workers semaphore, a memory-unaware counter that could not distinguish between a 13.6 GiB PoRep partition and a 0.4 GiB WinningPoSt partition.

This design was fragile on multiple levels. Operators had to manually calculate the correct partition_workers value for their hardware, a process that required deep knowledge of the memory footprint of each proof type. Getting it wrong meant either OOM crashes (if set too high) or GPU starvation (if set too low). There was no feedback loop between actual memory consumption and admission control, no mechanism to evict cached data under pressure, and no way to adapt to heterogeneous proof type mixes. The system was one configuration mistake away from disaster.

The Forensic Audit: Tracing Every Byte

The investigation began with the user's message at <msg id=2043>, which laid out a deliberate two-phase approach: "As step 1 look at what allocs/deallocs are done for 32G proofs over the life of a pipeline. Then we'll also think about an optimal strategy for memory management." The assistant responded by dispatching three parallel subagent tasks: one to search for documentation, one to trace memory allocations across the Rust codebase, and one to examine the C++ CUDA code for supraseal-c2. These ran concurrently, and their results fed into a comprehensive synthesis.

The assistant then read key source files directly—pipeline.rs, engine.rs, config.rs, srs_manager.rs, and the CUDA source—to build a complete picture. The resulting analysis, captured in <msg id=2054>, is a masterclass in systems-level forensic accounting. The assistant catalogued every allocation and deallocation point across the entire lifecycle of a 32 GiB PoRep proof, quantifying sizes with byte-level precision.

The Baseline Tax

The first major finding was the "baseline tax"—memory that lives for the entire process lifetime before any proof work begins. The Structured Reference String (SRS) consumes approximately 44 GiB of CUDA pinned host memory, allocated via cudaHostAlloc() and held for the process lifetime in SrsManager.loaded as Arc<SuprasealParameters>. This is page-locked RAM, registered with the CUDA driver for DMA transfers, meaning it cannot be swapped out. The Pre-Compiled Constraint Evaluator (PCE) adds another ~26 GiB on the heap, stored in static OnceLock<PreCompiledCircuit<Fr>> globals that are never freed. Together, SRS and PCE consume approximately 69 GiB of RAM before a single proof is even started.

The assistant also identified a critical transient spike during SRS loading: while the reader thread runs, both the mmap (~44 GiB virtual, demand-paged) and the pinned allocation (~44 GiB) exist simultaneously, creating a peak transient of approximately 88 GiB. This is the kind of detail that causes production incidents—invisible in steady-state monitoring, devastating when it occurs.

Per-Partition Working Memory

For each partition of a 32 GiB PoRep proof, the assistant traced the exact buffer sizes. The ProvingAssignment produces three evaluation vectors (a, b, c) each consuming approximately 4.17 GiB as Vec<Fr> (roughly 130 million field elements at 32 bytes each), plus an auxiliary assignment of ~0.74 GiB. Total per-partition: approximately 13.6 GiB.

The assistant then mapped the three-phase lifecycle. During Phase A (Synthesis), all ~13.6 GiB is allocated on the CPU heap. At Phase B (GPU Prove Start), raw pointers to a/b/c are passed to the C++ CUDA code, and crucially, the assistant identified the exact deallocation point: prover.a = Vec::new(), prover.b = Vec::new(), prover.c = Vec::new() at supraseal.rs:230-235. This frees approximately 12.5 GiB—92% of the per-partition working memory—the moment GPU proving begins. Only the auxiliary assignment (~0.74 GiB) and the C++ pending handle (~2 GiB estimated internal state) remain alive through Phase C (GPU Prove Finish), after which the async dealloc thread runs, serialized by DEALLOC_MTX.

The Throttle Audit

The assistant catalogued every mechanism that currently limits memory usage: the partition_semaphore (max concurrent partition synthesis tasks), the bounded channel (max synthesized proofs waiting for GPU), the synth_semaphore (max concurrent full-proof syntheses), the GPU mutex (one CUDA kernel per GPU at a time), the gpu_workers_per_device setting, the DEALLOC_MTX serialization, and the buffer flight counters (diagnostic only, NOT enforced).

The most damning entry in this audit was the working_memory_budget configuration option. It existed in the configuration, defaulted to "80GiB," but was not enforced anywhere. No code ever checked it. The only real throttle was partition_workers, a static count that the user had to manually tune based on their RAM size—a fragile and error-prone process.

The Peak Memory Calculation

The assistant calculated the worst-case peak memory for the common configuration of partition_workers=12 and gpu_workers_per_device=2:

Tracing SRS and PCE Lifecycles

The assistant dug deeper into the SRS and PCE loading paths. SRS loading begins at engine startup: Engine::start() calls SrsManager::ensure_loaded() in a spawn_blocking for each circuit ID in config.srs.preload. Inside the C++ layer, the constructor mmaps the 44 GiB .params file from disk, computes the total size of all point arrays, calls cudaHostAlloc() to allocate CUDA pinned host memory, then spawns a reader thread that deserializes points from the mmap into the pinned buffer. After the reader completes, the mmap is released. But during loading, both the mmap and the pinned allocation coexist, creating a transient peak of ~88 GiB.

PCE loading has two scenarios. In a warm start, the PCE file exists on disk and is deserialized at ~5 GB/s into a PreCompiledCircuit<Fr> struct, consuming ~26 GiB on the heap. The file content and the in-memory struct briefly coexist during deserialization. In a cold start, the first proof uses a slow synthesis path while a background thread extracts the PCE by building one partition circuit with RecordingCS, capturing the R1CS structure into CSR format. This extraction has its own ~13.6 GiB transient for the synthesis, plus the ~26 GiB PCE being built, totaling ~40 GiB of transient memory—all running on a background thread with no memory gate.

The assistant also confirmed a critical fact at <msg id=2068>: neither SRS nor PCE is ever evicted. SrsManager stores loaded entries in a HashMap and never removes them. PCE uses OnceLock which is write-once, never cleared. If the daemon processes both PoRep and SnapDeals proofs, the baseline grows: SRS PoRep 32G (44 GiB) + SRS SnapDeals 32G (33 GiB) + PCE PoRep 32G (26 GiB) + PCE SnapDeals 32G (~17 GiB) = ~120 GiB. On a 256 GiB machine, only ~126 GiB remains for working set—about 9 partitions instead of 12. If WindowPoSt also appears, baseline balloons to ~203 GiB, leaving room for only ~3 partitions.

Key Design Decisions

Through iterative Q&A across messages 2055-2073, the user and assistant collaboratively refined the architecture. Several key decisions emerged:

Unified memory budget. Rather than separate pools for pinned (SRS) and heap (PCE + working set) memory, the design uses a single unified budget auto-detected from system RAM. This simplifies the operator experience—fewer knobs to tune—and ensures that all memory consumers compete for the same finite resource.

LRU eviction with 5-minute minimum idle time, triggered under pressure. SRS and PCE entries track their last_used timestamp. When the budget is exhausted and a new allocation requires space, the evictor scans for entries that have been idle for at least 5 minutes and have an Arc::strong_count of 1 (meaning no in-flight proof is using them). The user explicitly confirmed this hybrid policy: "Evict under pressure, note that allocating more pipeline workers for new proofs is preferred over keeping unused SRS." This prevents thrashing while ensuring that unused cached data can be reclaimed when needed.

Two-phase working memory release. The a/b/c buffers (~12.5 GiB) are freed immediately after GPU prove start, and the budget is released synchronously at that point. The remaining ~1.1 GiB is released after proof finalization. This aligns budget accounting with actual memory lifecycle, allowing the next partition to begin synthesis sooner.

Removal of configurable preload and partition_workers. The SRS preload list is removed; all SRS loading is on-demand and gated by the budget. The partition_workers semaphore is replaced by the budget-based admission control, since the budget naturally limits partition concurrency by denying admission when memory is full.

Preservation of synthesis_concurrency. The user explicitly confirmed that synthesis_concurrency should remain as a separate CPU contention knob. Memory capacity (what the budget manages) and CPU thread contention (what synthesis_concurrency manages) are orthogonal concerns that require independent controls.

The Specification: From Design to Actionable Blueprint

The user's directive at <msg id=2074> was concise but consequential: "write down cuzk-memory-manager.md with all details needed for a new agents to implement correctly. Keep synth concurrency config, memory-bw/cpu contention is a separate constraint, partition-workers probably not needed." This message marked the transition from exploratory design to prescriptive documentation.

The assistant produced a comprehensive specification that defined the following components in detail:

Core Types in a New memory.rs Module

MemoryBudget is the central accounting structure. It holds a total_bytes value auto-detected from /proc/meminfo MemTotal minus a safety margin (default 5 GiB), an AtomicU64 tracking current usage, a tokio::sync::Notify for waking blocked acquirers, and an evictor callback that captures both the SRS manager and PCE cache. The acquire logic is a loop: try to increment usage atomically, and if the new total exceeds the budget, undo the increment and invoke the evictor. If the evictor frees enough memory, retry; otherwise, block on the Notify until some reservation is released.

MemoryReservation is an RAII handle that travels with a job through the pipeline. It holds a reference to the budget and an AtomicU64 tracking the remaining reservation amount. It provides a release(amount) method for partial release and implements Drop to release any remaining reservation as a safety net. This ensures that even if a code path forgets to explicitly release memory, the reservation is cleaned up when it goes out of scope.

Evictable PCE Cache

The four static OnceLock<PreCompiledCircuit<Fr>> globals are replaced by a PceCache struct containing a Mutex<HashMap<CircuitId, PceCacheEntry>> where each entry holds the Arc<PreCompiledCircuit<Fr>>, its size in bytes, and a last_used timestamp. The cache provides get() (returns an Arc clone, updates timestamp), insert() (acquires budget, stores entry), and evictable_entries() (returns entries idle ≥5 minutes with Arc::strong_count == 1). This design ensures safe concurrent access: in-flight proofs hold their own Arc clones, so eviction merely removes the cache's reference without affecting live proofs.

SRS Manager Changes

SrsManager gains last_used tracking updated on every get() or ensure_loaded() call. The ensure_loaded() method acquires budget before loading. A new evictable_entries() method returns entries meeting the idle and refcount criteria. The evict() method releases the budget after dropping the entry. The old advisory budget warning is removed—budget is now enforced.

Config Simplification

The configuration is dramatically simplified. Removed fields: srs.preload, memory.pinned_budget, memory.working_memory_budget, synthesis.partition_workers. Added fields: memory.total_budget (optional override string like "256GiB", default "auto"), memory.safety_margin (default "5GiB"), and eviction_min_idle (default "5min"). The system logs detected RAM, computed budget, and available working memory at startup.

Engine Integration Points

The specification identifies exact integration points. In partition dispatch, partition_sem.acquire_owned() is replaced with budget.acquire(PARTITION_EST_BYTES). The returned MemoryReservation is attached to the SynthesizedJob struct. In the GPU worker loop, after gpu_prove_start() returns, the worker calls reservation.release(ABC_EST_BYTES) to free the a/b/c portion. After gpu_prove_finish() returns, dropping the reservation releases the remaining ~1.1 GiB.

Memory Estimation Constants

The specification includes a table of hardcoded constants per proof type, rounded up to GiB:

| Proof type | Full synth output | After a/b/c freed | Shell + pending | |---|---|---|---| | PoRep 32G (1 partition) | 13.6 GiB | 12.5 GiB | 1.1 GiB | | PoRep 32G (batch 10) | 136 GiB | 125 GiB | 11 GiB | | WindowPoSt 32G | 13.2 GiB | 12.1 GiB | 1.0 GiB | | WinningPoSt 32G | 0.4 GiB | 0.35 GiB | 0.05 GiB | | SnapDeals 32G (1 partition) | 8.6 GiB | 7.9 GiB | 0.7 GiB |

Eviction Logic

The eviction callback collects evictable entries from both SRS manager and PCE cache, sorts them by last_used ascending (oldest first), and evicts entries until the needed amount is freed or no candidates remain. The 5-minute minimum idle time prevents thrashing on frequently accessed entries, while the Arc::strong_count == 1 check ensures that in-flight proofs are not disrupted.

Testing Strategy

The specification outlines a testing approach: unit tests for MemoryBudget acquire/release/partial-release with mocked evictors; integration tests verifying that the budget blocks synthesis when full and unblocks when memory is released; eviction tests confirming that idle entries are evicted under pressure while recently-used entries are protected; and a config migration test ensuring old config files produce sensible defaults.

The Broader Significance

This segment represents a complete design cycle for a critical subsystem: from problem identification through forensic analysis, collaborative design, and final specification. Several lessons emerge that generalize beyond this specific system.

First, measure before you design. The forensic audit was not optional—it was the foundation of everything that followed. The discovery that working_memory_budget was dead code, the quantification of the 69 GiB baseline tax, the identification of transient spikes during SRS loading and PCE extraction—these findings shaped every subsequent design decision. Without the audit, the design would have been based on assumptions rather than evidence.

Second, separate orthogonal concerns. The decision to keep synthesis_concurrency as a separate knob while removing partition_workers reflects a clear architectural principle: memory capacity and CPU contention are different constraints that require different controls. Mixing them in a single semaphore was the original design's fundamental flaw.

Third, specification matters. The final deliverable was not code changes but a specification document. This reflects a development philosophy that values design documentation as a separate artifact from implementation. The specification serves as knowledge preservation, implementation guidance, and a contract for parallel work. In an AI-assisted coding environment where different agent sessions may work on different parts of the system, this approach is particularly valuable.

Fourth, the user's role as architect. Throughout this segment, the user's questions and directives shaped the design at critical junctures. The question "How does SRS preload and PCE fit into this?" forced the expansion from a working-memory-only gate to a comprehensive memory management system. The question "Is PoRep SRS/PCE unloaded when e.g. Snap proof needs to be made?" revealed the unbounded baseline growth problem. The directive to write the specification transformed design exploration into actionable documentation. The user was not a passive consumer of the assistant's analysis but an active collaborator who understood the system deeply and knew when to probe, when to confirm, and when to demand closure.

Conclusion

The cuzk memory management redesign is a case study in how to approach architectural change in a complex, memory-intensive system. It began with a forensic audit that traced every allocation and deallocation point, producing a quantified model of the entire memory lifecycle. It proceeded through collaborative design that surfaced trade-offs, resolved ambiguities, and aligned on architectural principles. It culminated in a specification document that captures not just what to build, but why it should be built that way.

The result is a system that replaces a fragile static concurrency limit with a robust, memory-aware admission control system—one that can adapt to different proof types, evict unused cached data under pressure, and prevent OOM crashes without requiring operators to perform manual arithmetic. It is a design that respects the fundamental constraint of finite RAM while maximizing throughput within that constraint. And it is a testament to what becomes possible when forensic analysis meets collaborative design, captured in a specification that ensures the knowledge survives to guide implementation.

References