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

Introduction

In any high-performance computing system, memory management is the invisible hand that determines whether throughput soars or the process crashes with an OOM (Out of Memory) signal. The cuzk GPU proving engine—a critical component in the Filecoin network's proof pipeline—had reached a point where its memory management strategy had become a bottleneck. The system relied on a static partition_workers semaphore to limit concurrency, a mechanism that was fundamentally memory-unaware and fragile. Message 2055 in this coding session represents a pivotal moment: the assistant, after conducting a deep forensic audit of the entire memory lifecycle for 32 GiB PoRep proofs, presents a comprehensive design proposal for a new memory-aware admission control system. This message is not merely a suggestion—it is a carefully reasoned architectural specification that replaces a brittle count-based throttle with a robust, budget-based admission gate.

This article examines message 2055 in detail, exploring why it was written, the reasoning behind its design decisions, the assumptions it makes, the knowledge it builds upon, and the intellectual framework it establishes for future implementation.

The Context: Why This Message Was Written

Message 2055 did not emerge from a vacuum. It was the culmination of an intensive forensic audit spanning multiple code files across the cuzk codebase. In the preceding messages ([msg 2046] through [msg 2054]), the assistant had systematically traced every allocation and deallocation point in the 32 GiB PoRep proof pipeline. It had read pipeline.rs, engine.rs, supraseal.rs, srs_manager.rs, and even the CUDA C++ kernel code in groth16_cuda.cu. The result was a comprehensive memory lifecycle analysis that revealed several critical problems.

The most alarming finding was that the working_memory_budget configuration option—a setting that ostensibly allowed users to cap memory usage—was entirely dead code. It was parsed from the configuration file, stored in a struct, and then never read or enforced anywhere in the engine. The only real throttle was the partition_workers semaphore, a static integer that limited how many partition synthesis tasks could run concurrently. This approach had a fundamental flaw: it counted tasks, not bytes. A partition for a WinningPoSt proof might consume only 0.4 GiB, while a PoRep partition consumed 13.6 GiB, yet the semaphore treated them identically. Users had to manually calculate how many partitions would fit in their RAM, and getting this calculation wrong meant either OOM crashes (too many partitions) or GPU starvation (too few).

The assistant's analysis also uncovered a cascade of secondary issues: the bounded channel between synthesis and GPU proving had its capacity set to max(synthesis_lookahead, partition_workers), meaning a large partition_workers value created a large buffer where synthesized partitions could pile up waiting for the GPU. The deallocation path was serialized by a mutex, creating a bottleneck where freed memory remained mapped longer than necessary. And there was no mechanism to respond to actual system memory pressure—if another process consumed RAM, cuzk had no way to back off.

Message 2055 is the assistant's response to this diagnosis. It is the "what do we do about it" moment—a design proposal that addresses each identified weakness with a coherent architectural solution.

The Core Design: MemoryBudget Admission Gate

The centerpiece of the proposal is the MemoryBudget admission controller. The assistant describes it with a clear ASCII diagram:

                        ┌──────────────────────┐
  partition task ──────►│  MemoryBudget::acquire │──── permit ────► synthesis
                        │  (est_gib, timeout)    │
                        └────────┬───────────────┘
                                 │ blocks if
                                 │ current + est > budget
                                 │
                        ┌────────▼───────────────┐
                        │  MemoryBudget::release  │◄──── drop guard
                        │  (after dealloc done)   │
                        └────────────────────────┘

This design replaces the static partition_workers semaphore with a budget-based gate. Instead of asking "how many partitions can we run?", it asks "how much memory can we use?" Each synthesis request declares its estimated cost—13.6 GiB for a PoRep partition, 136 GiB for a batch of ten, 4 GiB for WinningPoSt—and the gate blocks until budget is available.

The reasoning behind this choice is clear from the message's emphasis on heterogeneous proof types. The existing system could not distinguish between a lightweight WinningPoSt proof and a heavyweight PoRep proof. A memory budget naturally handles this: a small proof can slip through even when many large proofs are in flight, because it actually fits within the remaining budget. This is a direct consequence of the assistant's earlier analysis, which had mapped the exact memory footprint of each proof type.

The Two-Phase Release Model

One of the most sophisticated aspects of the design is the two-phase working memory release. The assistant had discovered through its forensic audit that the synthesis output consists of three large vectors (a, b, c) totaling approximately 12.5 GiB, plus auxiliary data of about 1.1 GiB. Critically, the three large vectors are consumed and freed immediately after prove_start returns—the C++ GPU code extracts raw pointers and the Rust side drops the vectors. But the auxiliary data and the C++ pending proof handle survive until the deallocation thread completes.

The design exploits this lifecycle by splitting the memory reservation into two phases:

Integration Points: Five Surgical Changes

The assistant specifies exactly five integration points, demonstrating a deep understanding of the codebase's architecture:

  1. engine.rs partition dispatch (~line 1424): Replace partition_sem.acquire_owned() with memory_budget.acquire(PARTITION_EST_BYTES).
  2. engine.rs GPU worker (~line 2376): After gpu_prove_start returns, call memory_budget.release_partial(ABC_EST_BYTES).
  3. supraseal.rs dealloc thread (~line 284): After buf_dealloc_done(), call memory_budget.release(remaining).
  4. engine.rs monolithic synthesis (~line 2135): Gate on memory_budget.acquire(FULL_PROOF_EST_BYTES) for batch synthesis.
  5. config.rs: Add working_memory_budget_bytes() parser. Each integration point corresponds to a specific location identified during the forensic audit. The line numbers are approximate but grounded in actual code reading. This level of specificity transforms the proposal from a vague architectural vision into an actionable implementation plan. The assistant is not just saying "add a memory budget"—it is saying "here is exactly where to insert it, what to replace, and what to call."

The Size Estimates Table

The message includes a carefully constructed table of memory estimates for each proof type:

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

These numbers are not guesses. They are derived from the assistant's earlier analysis of the constraint counts, the CSR matrix sizes, and the ProvingAssignment struct layout. The PoRep 32 GiB estimate of 13.6 GiB per partition, for example, comes from the earlier breakdown: ~4.17 GiB each for a, b, c vectors (each ~130 million field elements at 32 bytes), plus ~0.74 GiB for aux assignment (~23 million elements), plus density bitvecs and other overhead. The assistant had traced these exact numbers from the code.

This table serves a dual purpose. First, it provides the hardcoded constants that the MemoryBudget controller will use for its estimates. Second, it demonstrates that the assistant has done the homework—these are not placeholder values but precise measurements grounded in the codebase's actual data structures.

Assumptions Embedded in the Design

Every design rests on assumptions, and message 2055 is no exception. Understanding these assumptions is critical for evaluating the proposal's strengths and limitations.

Assumption 1: Static estimates are sufficient. The design uses hardcoded constants for each proof type's memory footprint, rather than runtime measurement. This assumes that the memory consumption of synthesis is deterministic and predictable based on proof type alone. In practice, this is largely true for Groth16 proving, where constraint counts are fixed per circuit type. However, it does not account for variability due to different sector sizes (e.g., 2 KiB vs 32 GiB) or potential future circuit changes.

Assumption 2: The budget should be based on system RAM, not a configurable value. The assistant later clarifies (in the question at the end) that auto-detection from system RAM is an option. But the core design as presented assumes working_memory_budget from config is the source of truth. This assumes the user or operator knows their system's memory capacity and can set an appropriate budget.

Assumption 3: The two-phase release model is sufficient. The design splits the lifecycle into three phases (acquire, partial release, full release). This assumes that the a/b/c vectors are always freed immediately after prove_start, which the forensic audit confirmed. But it also assumes that no other memory allocations occur between these phases that could invalidate the budget tracking.

Assumption 4: The partition_workers semaphore can be entirely replaced. The design subsumes the existing semaphore, assuming that the memory budget gate alone provides sufficient concurrency control. This is a strong claim—it assumes that memory is the only resource constraint worth modeling. CPU contention, GPU queue depth, and I/O bandwidth are not addressed by this design.

Assumption 5: The budget tracks estimated usage, not actual RSS. The design uses AtomicU64 counters that track declared allocations, not actual RSS from /proc/self/statm. This is a deliberate choice (the assistant offers RSS-based tracking as an alternative at the end of the message), but it means the system cannot detect memory leaks or unexpected consumption from other components.

The Thinking Process Visible in the Reasoning

The assistant's reasoning, visible in the preceding messages, reveals a methodical approach to problem-solving. The forensic audit followed a clear pattern: read the code, trace the allocations, identify the gaps, and then design the solution.

The key insight—that working_memory_budget was dead code—came from reading engine.rs and discovering that the field was never referenced after configuration parsing. This is the kind of finding that only emerges from systematic code reading, not from high-level architectural analysis.

The two-phase release model emerged from tracing the exact sequence of events in supraseal.rs: the prove_start function returns, then the Rust code drops the a/b/c vectors, then later the dealloc thread runs. The assistant connected this sequence to the memory lifecycle and realized that the budget could be partially released mid-stream.

The decision to use hardcoded constants rather than runtime measurement reflects a pragmatic engineering judgment. Runtime measurement (e.g., reading RSS) would be more accurate but also more complex, with potential race conditions and measurement noise. Hardcoded constants, derived from the codebase's actual data structures, provide deterministic behavior that can be tested and validated.

The assistant's final question—"Do you want me to proceed with implementing this MemoryBudget approach? Or would you prefer a different direction?"—reveals an important aspect of the thinking process. The assistant is not prescribing a single solution. It is offering a well-reasoned proposal while acknowledging that alternatives exist (simpler wiring of the existing budget, or RSS-based tracking). This open-endedness invites collaboration and refinement.

Input Knowledge Required to Understand This Message

To fully grasp message 2055, a reader needs substantial context:

  1. The cuzk architecture: Understanding that cuzk is a GPU proving engine for Filecoin, that it uses a partitioned pipeline where each partition is synthesized independently and then proved on GPU, and that it supports multiple proof types (PoRep, WindowPoSt, WinningPoSt, SnapDeals).
  2. The memory lifecycle: Knowing that SRS (Structured Reference String) is ~44 GiB of pinned host memory, that PCE (Pre-Compiled Constraint Evaluator) is ~26 GiB of heap memory, and that each partition's synthesis produces ~13.6 GiB of working memory that is freed in stages.
  3. The existing throttling mechanisms: Understanding the partition_workers semaphore, the bounded channel, the GPU mutex, and the deallocation serialization.
  4. The config system: Knowing that working_memory_budget existed as a config option but was never enforced.
  5. The proof type differences: Understanding that WinningPoSt has ~3.5 million constraints while PoRep has ~130 million per partition, leading to vastly different memory footprints.
  6. Rust concurrency primitives: Understanding tokio::sync::Notify, AtomicU64, and async semaphores to evaluate the proposed MemoryBudget implementation.

Output Knowledge Created by This Message

Message 2055 creates several forms of knowledge:

  1. A concrete, implementable design: The MemoryBudget admission controller is specified well enough that a developer could implement it directly from the message. The Rust struct sketch, the method signatures, and the integration points provide a clear implementation roadmap.
  2. A memory lifecycle model: The two-phase release model and the size estimates table codify the assistant's forensic findings into a reusable model that can be referenced in future design discussions.
  3. A decision framework: By presenting the proposal alongside alternatives (simpler wiring, RSS-based tracking), the message creates a structured decision point. The user can evaluate trade-offs and choose a direction.
  4. Documentation of dead code: The message (building on the preceding analysis) documents that working_memory_budget is dead code, which is itself a valuable finding for the codebase's maintainers.
  5. A vocabulary for discussing memory management: Terms like "budget-based admission," "two-phase release," and "heterogeneous proof type awareness" become part of the project's shared language.

Mistakes and Potential Issues

While the design is well-reasoned, several potential issues deserve scrutiny:

The static estimate approach has a fragility problem. If a new proof type is added or an existing circuit is modified, the hardcoded constants must be updated. There is no mechanism to detect drift between the estimates and actual memory usage. A runtime assertion that checks actual allocation size against the estimate could catch discrepancies.

The design does not address VRAM pressure. The analysis identified that GPU VRAM is also a constrained resource, with the CUDA memory pool configured to never release memory (UINT64_MAX pool threshold). The MemoryBudget controller only manages host memory. A complete solution might need a GPU-side budget as well.

The release_partial API introduces complexity. The caller must know exactly how much to release at each phase. If the a/b/c vectors are not exactly 12.5 GiB (due to alignment, padding, or different sector sizes), the partial release could be inaccurate. The design assumes precise knowledge of allocation sizes.

The design does not handle memory fragmentation. The malloc_trim(0) calls in the existing code suggest that fragmentation is a concern. The budget tracks total allocated bytes but does not account for the fact that freed memory may not be immediately usable due to fragmentation.

The proposal removes partition_workers but keeps synthesis_concurrency. The assistant explicitly notes that synthesis_concurrency should remain as a separate CPU contention knob. This is a reasonable distinction, but it creates a system with two independent throttles (memory budget + CPU concurrency) that could interact in unexpected ways.

Conclusion

Message 2055 is a masterclass in system design grounded in forensic code analysis. It takes a complex, multi-layered problem—memory management in a GPU proving engine—and decomposes it into a clean architectural solution. The MemoryBudget admission controller is not a generic or abstract design; it is a specific response to the concrete problems uncovered during the audit: dead configuration, memory-unaware throttling, heterogeneous proof types, and a fragile static limit.

The message's power lies in its specificity. The size estimates table, the integration point line numbers, the Rust struct sketch, and the two-phase release model all demonstrate that the assistant has done the hard work of understanding the system before proposing changes. This is not a drive-by suggestion but a carefully researched design document.

The final question—inviting the user to choose between this approach and alternatives—transforms the message from a monologue into a dialogue. It acknowledges that design is about trade-offs, and that the best solution depends on priorities that the user must articulate. Whether the user chooses the MemoryBudget path, a simpler approach, or an RSS-based alternative, message 2055 has already done the essential work: it has framed the problem, mapped the territory, and drawn a clear path forward.