The Dispatcher Transformed: Budget-Based Admission Control in cuzk's GPU Proving Engine

In the middle of a sweeping refactoring of the cuzk GPU proving engine, the assistant issued a single, terse edit command:

### Edit 4: Update the dispatcher spawn block — remove partition_workers, add budget/pce_cache [edit] /tmp/czk/extern/cuzk/cuzk-core/src/engine.rs Edit applied successfully.

This message, [msg 2168], is the fourth in a sequence of six targeted edits to engine.rs—the central coordinator of the cuzk proving daemon. On its surface, it is a one-line status update confirming that an edit was applied. But to understand what this message means, one must zoom out to the architectural transformation underway: the replacement of a fragile static concurrency limit with a unified, memory-aware admission control system.

The Problem: Static Concurrency Meets Dynamic Memory Pressure

The cuzk proving engine is a high-throughput GPU-based system that generates zero-knowledge proofs for Filecoin storage proofs (PoRep, WindowPoSt, SnapDeals). Before this refactoring, the engine controlled GPU admission using a static partition_workers semaphore—a fixed integer configured at startup that capped how many proof partitions could be in flight simultaneously. This approach had a fundamental flaw: it was blind to actual memory pressure.

A 32 GiB PoRep proof consumes vastly different memory than a SnapDeals proof. SRS (Structured Reference String) parameters, PCE (Pre-Compiled Constraint Evaluator) circuits, and synthesis working sets all compete for the same system RAM. The old partition_workers semaphore could be set too high, causing OOM crashes when multiple large proofs ran concurrently, or too low, leaving GPU cycles idle while memory sat underutilized. The system had no mechanism to evict idle SRS or PCE entries under pressure, and the working_memory_budget config field was a dead letter—parsed but never enforced.

The Solution: A Unified Memory Budget

The memory manager specification (cuzk-memory-manager.md) designed a replacement: a single MemoryBudget auto-detected from system RAM, with byte-level tracking of all major consumers. SRS and PCE would be loaded on demand, consume quota, and be evicted under pressure via LRU eviction. The old partition_workers semaphore would be replaced by budget-based admission: a proof partition could only start if the budget had enough remaining capacity for its working set.

By the time the assistant reached Edit 4, the foundational pieces were already in place. The memory.rs module provided MemoryBudget, MemoryReservation, and system detection. The srs_manager.rs was rewritten to support budget-aware loading with ensure_loaded() accepting an Option<MemoryReservation>. The pipeline.rs module replaced the old static OnceLock-based PCE caches with a PceCache struct that integrated with the budget. Only engine.rs—the largest and most complex file—remained to be converted.

Edit 4 in Context: The Dispatcher Spawn Block

The assistant had already executed three edits before this one:

Edit 1 ([msg 2165]): Replaced the SRS preload block and PCE preload block in the start() method with evictor callback wiring, so the budget could request eviction of idle SRS/PCE entries when memory ran low.
Edit 2 ([msg 2166]): Updated channel capacity sizing to derive values from the budget instead of partition_workers, ensuring backpressure scales with available memory.
Edit 3 ([msg 2167]): Removed the partition_semaphore setup and updated log messages to reflect the new budget-based world. Edit 4 targets the dispatcher spawn block—the code that spawns the asynchronous dispatcher task responsible for routing proof requests from the scheduler to the GPU workers. This is the central nervous system of the engine. The dispatcher reads from a channel of incoming proof requests, determines whether each request should be handled monolithically or via the partitioned pipeline, and calls dispatch_batch() to send work to GPU workers. The old dispatcher spawn block carried partition_workers as a parameter, using it to size internal channels and control how many partitions could be dispatched concurrently. The edit replaced this with two new parameters: budget (an Arc<MemoryBudget>) and pce_cache (an Arc<PceCache>). This single change ripples through the entire dispatch logic. Instead of checking a static semaphore count, the dispatcher now calls budget.acquire() before spawning synthesis tasks, pre-acquires SRS budget in async context before spawn_blocking, and uses pce_cache.get() instead of the removed pipeline::get_pce() function.

The Thinking Process: Systematic, Spec-Driven Refactoring

The assistant's reasoning, visible across the preceding messages, reveals a methodical approach. Before making any edits, it read the entire engine.rs file (a file over 2,600 lines) and the full memory manager specification in parallel ([msg 2154] through [msg 2158]). It then verified the API signatures of the new components by reading memory.rs and srs_manager.rs ([msg 2159] through [msg 2161]). Crucially, it grepped for all remaining references to the old pipeline::get_pce() function ([msg 2163]), finding four call sites in engine.rs that would need updating—a proactive discovery that prevented a compilation error.

The assistant then announced its plan ([msg 2164]):

Now I have a thorough understanding of everything that needs to change. Let me start making the edits systematically. This is a large file with many interconnected changes, so I'll work section by section.

1. Remove SRS preload block and PCE preload block 2. Wire evictor callback 3. Update channel capacity sizing 4. Remove partition_semaphore

This plan was executed in order. Each edit built on the previous one, removing a dependency on the old system and wiring in the new one. The order was carefully chosen: the preload blocks had to go first because they referenced the old static caches; the channel sizing had to change before the semaphore was removed because the channels depended on the semaphore count; the dispatcher spawn block (Edit 4) had to change before the dispatch_batch signature (Edit 5) because the spawn block passes the parameters into the dispatcher closure.

Assumptions and Input Knowledge

This message assumes a substantial body of input knowledge. A reader must understand:

The architecture of engine.rs: the start() method that initializes the engine, the dispatcher loop that routes requests, the dispatch_batch() helper that sends work to GPU workers, and the GPU worker loop that performs synthesis and proving.
The old admission control system: the partition_workers config field, the partition_semaphore (a tokio::sync::Semaphore), and how they limited concurrent partition processing.
The new memory manager APIs: MemoryBudget::acquire() for working memory, SrsManager::ensure_loaded() with reservation support, PceCache::get() for on-demand PCE loading, and the evictor callback mechanism.
The sequence of edits: Edit 4 is not standalone—it depends on Edits 1-3 having already removed the preload blocks, updated channel sizing, and removed the semaphore. The assistant assumes these edits applied correctly and that the file is in a consistent intermediate state.
The Rust async runtime: the dispatcher spawns a tokio::task::spawn block, and the budget acquisition must happen in the correct async context.

Output Knowledge Created

This edit produces a critical intermediate state in the engine.rs transformation. After Edit 4, the dispatcher is no longer parameterized by partition_workers. Instead, it carries budget and pce_cache, making it ready for the subsequent edits that will update the dispatch_batch() signature (Edit 5, [msg 2169]) and the helper body (Edit 6, [msg 2170]). The four remaining pipeline::get_pce() call sites in the dispatcher loop are now one step closer to being replaced with pce_cache.get().

More broadly, this edit completes the decoupling of the engine's central routing logic from the old static semaphore. The dispatcher no longer knows about "partition workers" as a capacity metric—it only knows about memory budget. This is the conceptual heart of the refactoring: replacing a dimensionless integer (how many partitions) with a physical resource (how many bytes).

Conclusion

Message [msg 2168] is a deceptively simple status update that marks a pivotal moment in a complex architectural transformation. The assistant's systematic, spec-driven approach—reading the entire file, verifying API signatures, planning edits in dependency order, and executing them one by one—turned a risky multi-hundred-line refactoring into a controlled sequence of surgical changes. The dispatcher spawn block was the last major structural dependency on the old partition_workers system; after this edit, the remaining work was a matter of updating signatures and call sites. The unified memory budget was no longer a design on paper—it was wired into the engine's central nervous system.