The Heart of the Machine: Integrating a Unified Memory Manager into cuzk's GPU Proving Engine

Introduction

In the architecture of any high-performance computing system, there comes a moment when a well-designed component must prove it can survive contact with the real codebase. The unified memory manager for cuzk—a CUDA-accelerated zero-knowledge proving daemon for Filecoin—had been carefully specified in a 1072-line design document, its core types (MemoryBudget, MemoryReservation, PceCache) implemented and tested, its configuration schema rationalized, and its supporting infrastructure (SrsManager, pipeline.rs) updated to be budget-aware. All of this work, spanning segments 14 and 15, had created the parts of a new memory management architecture. But the parts do not make a machine. The machine is made when those parts are wired together into the central orchestrator that drives the entire proving pipeline.

That central orchestrator is engine.rs—a ~2837-line file that is the heart of the cuzk proving engine. It initializes the GPU, manages the dispatcher loop that distributes proof requests across partitions, coordinates SRS and PCE loading, spawns GPU worker tasks, and implements the split-prove API that releases GPU memory in two phases. Integrating the unified memory manager into engine.rs was the most complex and consequential part of the refactoring. This article tells the story of that integration, as captured in a single opencode coding session: the systematic transformation of engine.rs from a static concurrency-limiting architecture to a dynamic, budget-based admission control system.

The Architecture Before Integration

To understand what the integration accomplished, it is necessary to understand what it replaced. The pre-refactoring engine.rs managed memory through a collection of mechanisms that had accreted over time, each solving a local problem without awareness of the global memory picture.

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.

The start() method of the engine was the epicenter of these problems. It contained two large preload blocks—one iterating over all known CircuitId values to load every SRS parameter, another loading every PCE circuit from disk—that together consumed ~70 GiB of memory before the engine was ready to accept a single proof request. 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 Budget

The new architecture, specified in the memory manager design document and implemented across segments 14 and 15, 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 key types were already in place. 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 had been 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 had been 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.

But these types and patterns existed as isolated components. The task of segment 16 was to wire them into engine.rs—to replace the old preload blocks, the semaphore, the static cache references, and the monolithic worker loop with the new budget-based equivalents. This was not a mechanical substitution; it required understanding the ownership semantics, async boundaries, error paths, and concurrency patterns of a complex production system.

The Integration: A Systematic Transformation

The assistant approached the engine.rs integration with a methodical, checklist-driven methodology. The edits were numbered and tracked against specific line numbers in the original file, each building on the previous one in a carefully ordered dependency chain.

Rewriting the Start Method

The first edit ([msg 2165]) targeted the start() method—the initialization path that runs when the engine boots. Previously, this method contained two large blocks: one that iterated over all known CircuitId values and eagerly loaded every SRS parameter, and another that loaded every PCE circuit from disk. These blocks consumed ~70 GiB of memory before the engine was ready to accept a single proof request.

The new code removed both preload blocks entirely. Instead, it registered an evictor callback with the MemoryBudget—a closure that the budget invokes when it needs to free memory under pressure. The evictor calls SrsManager::evict() and PceCache::evict() to release idle entries. 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 (trying to load SRS into a budget that was already fully committed) or silently ignore the new memory manager (if the preload bypassed the budget entirely).

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, capped at 32 to prevent unbounded channel growth.

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. This was a profound shift: instead of a fixed concurrency limit that operators had to manually tune, the system now had a dynamic limit that adjusted automatically based on actual memory pressure. If the system had 256 GiB of RAM, it could run more concurrent partitions than if it had 128 GiB—without any configuration change.

The Dispatcher Transformation

Edits 4 through 7 ([msg 2168]–[msg 2176]) transformed the dispatcher loop—the central routing logic that receives proof requests and distributes them to the appropriate dispatch path. Both process_batch() and dispatch_batch() received new signatures: they now accept &MemoryBudget and &PceCache instead of the old semaphore and worker-count parameters.

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. This systematic hunting for every reference is the hallmark of a thorough refactoring—a single missed call site would cause a compilation error, and a single missed semaphore reference would mean the old admission control path still existed alongside the new one, creating a confusing dual-system state.

Partition Dispatch: PoRep and SnapDeals

The PoRep and SnapDeals per-partition dispatch paths underwent the most intricate changes ([msg 2178]–[msg 2191]). These paths handle proofs that are split across multiple GPU partitions—the most memory-intensive workload in the system. A single 32 GiB PoRep C2 proof, for example, is divided into partitions that each require ~13.6 GiB of working memory.

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 Architectural Significance

Several themes emerge from this segment 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 segment 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.