The Complete Integration: How a Unified Memory Manager Transformed cuzk's GPU Proving Engine

Introduction

In the high-stakes world of GPU-accelerated zero-knowledge proof generation for Filecoin, memory management is not a peripheral concern—it is the central constraint that shapes every architectural decision. The cuzk (CUDA ZK proving daemon) engine operates at the intersection of cryptographic protocol complexity, GPU compute intensity, and the unforgiving physics of system RAM. A single 32 GiB PoRep proof requires approximately 70 GiB of baseline resident memory (44 GiB for Structured Reference String parameters plus 26 GiB for Pre-Compiled Constraint Evaluator circuits), and each concurrent partition adds another ~13.6 GiB of working memory. The old system managed these resources through a fragile collection of static configuration knobs—partition_workers, pinned_budget, working_memory_budget, and preload—that were poorly understood, frequently misconfigured, and completely disconnected from actual memory pressure.

This article synthesizes the work captured in a single chunk of an opencode coding session: the complete integration of a unified memory manager into the cuzk GPU proving engine. Over the course of dozens of messages spanning hundreds of tool calls, the assistant systematically replaced a static concurrency-limiting architecture with a dynamic, budget-based admission control system. The transformation touched every major component of the proving pipeline—from the start() method that initializes the engine, through the dispatcher loop and partition dispatch paths, to the GPU worker loop where memory is actually consumed and released. This is the story of that transformation.

The Architecture of the Old System

Before the refactoring, the cuzk engine managed memory through a hodgepodge of mechanisms that had accreted over time. The primary admission control was a partition_semaphore—a Tokio semaphore whose capacity was set to the partition_workers configuration value. When a partition proof needed to run, it would acquire a permit from this semaphore; when the permit count was exhausted, no new partitions could start. This was a crude throttle: it prevented too many GPU proofs from running simultaneously, but it had no awareness of actual memory consumption. The system could have 70 GiB of SRS and PCE loaded, plus 13.6 GiB per partition, and the semaphore would happily dispatch partitions until it hit its worker limit, regardless of whether the system had sufficient RAM.

The SRS and PCE caches were equally problematic. SRS parameters were preloaded for all known circuit types at engine startup, consuming ~44 GiB of pinned GPU memory before any proof was even requested. PCE circuits were stored in four static OnceLock-based global variables, consuming ~26 GiB of heap memory. Neither cache had any eviction mechanism—once loaded, data stayed in memory forever. The working_memory_budget and pinned_budget configuration fields existed in the config schema but were never actually enforced in code; they were documentation artifacts that gave operators a false sense of control.

This architecture was fragile because it could not adapt. Operators had to manually tune partition_workers based on their hardware, and getting it wrong meant either underutilized GPUs (too few workers) or OOM-killed processes (too many). The system had no unified understanding of its own memory footprint.

The Design: A Unified Memory Budget

The specification for the new system, documented in a comprehensive 1072-line design document (see [msg 2151]), proposed a radical alternative: replace all the static knobs with a single byte-level budget auto-detected from system RAM. The core insight was that all memory consumers—SRS pinned allocations, PCE heap objects, and per-partition working vectors—draw from the same finite pool, and the system should admit new work only when sufficient budget remains.

The design introduced several new types. MemoryBudget is the central allocator: it tracks total available memory, provides acquire() and release() methods for working memory, and supports an evictor callback that the budget invokes when it needs to free space. MemoryReservation is a RAII guard that releases budget on drop or via explicit release(). The SrsManager was made budget-aware: ensure_loaded() now accepts an optional MemoryReservation and calls into_permanent() to transfer budget ownership from temporary reservations to the permanent cache. The four static OnceLock-based PCE caches were replaced by a PceCache struct supporting insertion, lookup, and LRU eviction.

The two-phase memory release pattern was the centerpiece of the GPU worker loop design. The Groth16 proving API exposes a split flow: gpu_prove_start performs the first phase of GPU computation and frees the a/b/c intermediate buffers (~12.5 GiB), while gpu_prove_finish completes the proof. By releasing the a/b/c portion immediately after gpu_prove_start returns, and dropping the remaining reservation only after gpu_prove_finish, the system reclaims 92% of the working memory earlier, making it available for the next partition sooner.

The Integration: Rewriting engine.rs

The integration of this design into engine.rs—the heart of the proving pipeline, spanning approximately 2837 lines—was the most complex part of the implementation. The assistant approached it systematically, working through a numbered checklist that tracked each change against specific line numbers in the original file.

The Start Method Rewrite

The first edit ([msg 2165]) targeted the start() method, removing the old SRS and PCE preload blocks and wiring the evictor callback into the MemoryBudget after GPU detection. Previously, the engine would iterate over all known CircuitId values and eagerly load every SRS parameter and PCE circuit into memory. The new code registers an evictor callback—a closure that the budget calls when it needs to free memory—and removes the preload loops entirely. SRS and PCE are now loaded on demand by the dispatch paths, each acquisition consuming a portion of the budget.

This edit was the foundation for everything that followed. Without it, the old preload blocks would conflict with the new budget-aware APIs, and the engine would either crash at startup or silently ignore the new memory manager.

Channel Capacity and Semaphore Removal

The second and third edits ([msg 2166], [msg 2167]) addressed the channel capacity sizing and the partition_semaphore. Previously, internal channel capacities were computed as partition_workers * some_factor—a heuristic that assumed more workers meant more throughput, ignoring the reality that memory was the binding constraint. The new code derives channel capacities from the budget divided by the estimated per-partition working set. The partition_semaphore—the old admission control mechanism—was removed entirely. Its function was now performed by budget.acquire(), which blocks only when insufficient memory is available.

The Dispatcher Transformation

Edits 4 through 7 ([msg 2168]–[msg 2176]) transformed the dispatcher loop. Both process_batch() and dispatch_batch() received new signatures: they now accept &MemoryBudget and &PceCache instead of the old semaphore and worker-count parameters. All five dispatch_batch call sites in the dispatcher loop were updated to pass these new arguments. This was not a mechanical change—it required understanding the lifetime and ownership semantics of the budget and cache across async boundaries, ensuring that references lived long enough for the spawned tasks to use them.

The assistant's methodical approach is visible in the grep that preceded these edits ([msg 2171]). Searching for dispatch_batch across the codebase revealed exactly five call sites, each of which needed updating. The assistant then worked through them one by one, reading each context to understand how the budget and cache should be passed.

Partition Dispatch: PoRep and SnapDeals

The PoRep and SnapDeals per-partition dispatch paths underwent the most intricate changes ([msg 2178]–[msg 2191]). Previously, these paths spawned tasks guarded by the semaphore and used pipeline::get_pce()—a global function backed by a static cache—for PCE extraction. The new code acquires working-memory budget via budget.acquire() before spawning tasks, pre-acquires SRS budget in async context before entering spawn_blocking, and uses the pce_cache instance for extraction instead of the removed get_pce() function.

The SRS loading calls were updated from mgr.ensure_loaded(&circuit_id) to mgr.ensure_loaded(&circuit_id, srs_reservation), passing an optional reservation that the SrsManager uses to inform eviction decisions. The PCE background extraction was similarly updated: instead of calling pipeline::get_pce() which accessed a global static cache, the code now calls pce_cache.get() on the shared PceCache instance.

A notable discovery during this phase was that the slotted pipeline path—a legacy fallback for single-sector PoRep C2 proofs—had become dead code ([msg 2193]). The new per-partition dispatch path catches all single-sector PoRep C2 cases, making the slotted path unreachable. The assistant chose to leave it in place "for safety" rather than removing it, updating only the ensure_loaded and PCE check calls within it. This pragmatic decision avoids the risk of inadvertently breaking something, at the cost of leaving dead code for future readers to wonder about.

The Monolithic Synthesis Path

The monolithic synthesis path ([msg 2195]–[msg 2199]) was the final and most intricate section to update. Unlike the partition dispatch paths, which handle proofs split across multiple partitions, the monolithic path handles proofs that run in a single synthesis step. The assistant added budget acquisition before the spawn_blocking call (since synthesis happens inside blocking context, budget must be acquired beforehand to avoid blocking a worker thread while waiting for memory). The resulting MemoryReservation was attached to the SynthesizedJob struct via a new reservation: Option<MemoryReservation> field.

The critical fix in this path was the panic handling ([msg 2198], [msg 2199]). If synthesis panics inside spawn_blocking, the reservation must still be released. Without an explicit drop(reservation) in the panic path, the reservation would be leaked—the bytes would never be returned to the budget, permanently reducing available memory. The fix was simple but essential: move the reservation variable into the closure and ensure it is dropped on both the success and panic paths.

The Two-Phase GPU Memory Release

The GPU worker loop edit ([msg 2200]–[msg 2211]) implemented the two-phase release pattern that was the centerpiece of the memory manager design. The assistant called this "the most critical part," and the subsequent messages show careful reasoning about ownership, error paths, and the synchronous fallback.

The pattern works as follows. When a synthesized job enters the GPU worker loop, the reservation and circuit_id are extracted from SynthesizedJob before the job is moved into spawn_blocking. After gpu_prove_start completes successfully, the a/b/c portion of the reservation is released immediately—these buffers are no longer needed once the GPU has copied them. The remaining reservation is moved into the finalizer task and dropped only after gpu_prove_finish completes.

The assistant handled several edge cases with care. The synchronous fallback path (CUZK_DISABLE_SPLIT_PROVE), which bypasses the split API and runs the entire proof in a single call, was updated to drop the reservation after the synchronous call completes ([msg 2205], [msg 2206]). The non-CUDA fallback path (#[cfg(not(feature = "cuda-supraseal"))]), which is a stub that returns an error, was updated to drop the reservation immediately ([msg 2210]). Error paths in the split API—where gpu_prove_start fails—were handled by dropping the reservation in the error arm.

A moment of confusion about Rust ownership semantics ([msg 2208]) shows the assistant reasoning through a subtle issue: in a match statement where reservation is an Option<MemoryReservation>, can each arm independently move the value? The assistant initially thought the error arms couldn't use reservation because the success arm moved it, then realized that match arms are mutually exclusive and each arm gets its own ownership of captured variables. This is a correct understanding of Rust's borrow checker, but the momentary confusion highlights the cognitive complexity of managing ownership in concurrent code.

Verification: Closing the Loop

After the engine edits were complete, the assistant ran a series of grep queries to verify that no stale references to old APIs remained ([msg 2214]–[msg 2220]). The grep searched for pipeline::get_pce, partition_workers, and partition_semaphore across engine.rs. The result was almost perfect: a single match on a comment reading "// RAM. No static partition_workers count needed." This was harmless documentation, not executable code. The verification confirmed that every call site, every variable reference, every import of the old APIs had been successfully replaced.

The assistant then updated its todo list ([msg 2221]), marking memory.rs, lib.rs, config.rs, srs_manager.rs, and pipeline.rs as completed. This administrative step served as a checkpoint, signaling the transition from the core engine integration to the remaining peripheral files.

The Cleanup Phase: Config and Bench

The refactoring was not complete until every consumer of the old APIs was updated. Two files remained: the example configuration file (cuzk.example.toml) and the standalone benchmarking utility (cuzk-bench/src/main.rs).

The Configuration File

The example config file ([msg 2222], [msg 2223]) was updated to reflect the new configuration surface. The old [srs] section with preload directives, the [memory] section with pinned_budget and working_memory_budget fields (which were never actually enforced), and the [synthesis] section with partition_workers were all removed. The new config introduces total_budget (the unified memory budget, typically auto-detected), safety_margin (memory reserved for the OS), and eviction_min_idle (the minimum idle duration before an SRS or PCE entry becomes eligible for eviction).

The assistant made a deliberate choice to remove the deprecated fields entirely rather than comment them out. This assumes that backward compatibility via serde's unknown-field-ignoring behavior is sufficient for existing deployments, and that the example file should represent the ideal, current state. Operators migrating old configs will either receive warnings from the warn_deprecated() methods in config.rs or have their old fields silently ignored—both safe paths.

The Benchmark Utility

The cuzk-bench binary required more extensive changes ([msg 2224]–[msg 2238]). Unlike the engine, which owns an SrsManager and a MemoryBudget, the bench is a thin command-line wrapper that calls pipeline functions directly. It does not create an Engine—it must create its own PceCache instance.

The assistant discovered two functions in the bench that needed updating: run_pce_bench and run_pce_pipeline. Both called extract_and_cache_pce_from_c1(), which now requires a &PceCache parameter, and both previously used pipeline::get_pce() for retrieving cached PCEs. The fix for each function was the same: create a local PceCache with a large budget (since the bench is not constrained by a real GPU memory budget), pass it to the extraction function, and use pce_cache.get() instead of get_pce().

The assistant's methodical approach is visible in the sequence of messages: it first updated run_pce_bench ([msg 2233], [msg 2234]), then read the file again to check for additional call sites ([msg 2235]), discovered run_pce_pipeline ([msg 2236], [msg 2237]), and applied the same fix ([msg 2238]). This systematic hunting for every last reference is the hallmark of a thorough refactoring.

The Silent Signal: Completion

The chunk ends with an empty user message ([msg 2239])—a continuation signal in the opencode interaction model. The user sends nothing but empty <conversation_data> tags, saying in effect: "I trust you to know what to do next. I trust you to do it correctly." The assistant's response ([msg 2240]) is a comprehensive status summary that inventories every file, every change, and every remaining task. This is the moment when the heavy lifting is done and the implementation transitions to the final verification and testing phase.

Themes and Patterns

Several themes emerge from this chunk that are worth highlighting as general lessons for complex software refactoring.

Systematic decomposition. The assistant broke the engine.rs integration into discrete, independently verifiable steps, each targeting a specific code region with a specific transformation. The edits followed a clear dependency order: remove preload blocks before wiring the evictor, remove the semaphore before converting dispatch, convert dispatch before implementing two-phase release. This ordering minimized the risk of cascading failures.

Verification as a discipline. After every major edit sequence, the assistant ran grep queries to verify that no stale references remained. This is not paranoia—it is engineering rigor. A single missed reference to pipeline::get_pce would cause a compilation error, and a single missed partition_semaphore would mean the old admission control path still existed alongside the new one, creating a confusing dual-system state.

Error path discipline. The two-phase release pattern required careful handling of every error path: panic in synthesis, failure in gpu_prove_start, the synchronous fallback, the non-CUDA fallback. The assistant ensured that reservations were dropped in every error arm, preventing memory leaks that would permanently reduce the available budget.

Documentation as integration. The config file update was not an afterthought—it was the moment when the memory manager refactoring became complete in a practical sense. Code changes alone are insufficient; documentation must be updated, configuration surfaces must be rationalized, and operators must be able to discover the new model.

Conclusion

The integration of the unified memory manager into the cuzk GPU proving engine was a profound architectural transformation. It replaced a fragile system of static concurrency limits, unaccounted caches, and dead configuration fields with a coherent, budget-based admission control system that understands its own memory footprint. The transformation touched every major component of the proving pipeline, from the start() method through the dispatcher loop, partition dispatch paths, monolithic synthesis path, and GPU worker loop. It updated configuration files, benchmarking utilities, and documentation.

The work in this chunk demonstrates a methodical, systematic approach to complex refactoring: understand the old code, design the new architecture, implement the supporting infrastructure, integrate into the central orchestrator, verify thoroughly, and update every downstream consumer. The result is a proving engine that can safely handle concurrent large proofs without crashing, adapt to varying memory conditions, and gracefully recover from errors—a critical improvement for a production system generating storage proofs on Filecoin.