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

Introduction

In the world of high-performance GPU proving engines—particularly those used in blockchain and decentralized storage networks like Filecoin—memory management is not a peripheral concern; it is often the central bottleneck that determines whether a system crashes, stalls, or runs efficiently. Message 2073 in this coding session captures a pivotal moment: the assistant, having gathered requirements through iterative Q&A, presents a comprehensive architectural plan for replacing a fragile, static concurrency limit with a robust, memory-aware admission control system for the cuzk GPU proving engine.

This message is a design document in the truest sense. It is not a line-by-line implementation but rather a blueprint—a detailed specification that bridges the gap between high-level user requirements and concrete code changes. The assistant lays out new types, modified data structures, integration points across seven files, memory estimation constants, eviction policies, and configuration migration paths. To understand this message fully, one must appreciate the forensic audit that preceded it, the design tensions it resolves, and the architectural philosophy it embodies.

The Context: Why This Message Was Written

The Problem Landscape

The cuzk engine is a GPU-based proving system for Filecoin's proof-of-replication (PoRep) and proof-of-spacetime (PoSt) circuits. It processes large proofs—32 GiB sectors—by splitting the work into partitions, each requiring approximately 13.6 GiB of working memory during synthesis and proving. The engine also loads Supraseal Parameters (SRS), which are large (~44 GiB for PoRep) pinned-memory structures, and Pre-Compiled Constraint Evaluators (PCE), which are heap-allocated circuit representations (~26 GiB for PoRep).

The existing architecture had a critical flaw: it used a static partition_workers semaphore as the sole throttle on concurrent partition processing. This semaphore was set at startup based on a fixed count, with no awareness of how much memory was actually available. As the assistant discovered through code audit, the working_memory_budget configuration option—which should have been the mechanism for memory-aware gating—was entirely dead code. It was never enforced. The system would cheerfully try to load SRS and PCE for multiple circuit types (PoRep, WindowPoSt, SnapDeals) without ever evicting them, leading to a baseline memory footprint that could exceed 200 GiB on a 256 GiB machine, leaving only enough room for three partitions worth of working memory.

This was a ticking time bomb. An operator who configured partition_workers = 12 might find the system working fine with PoRep alone, only to see it OOM when WindowPoSt proofs also arrived and its 57 GiB SRS and 26 GiB PCE were loaded on top of the existing baseline.

The Requirements Gathering Process

The message we are analyzing did not emerge from a vacuum. It is the culmination of a multi-turn conversation that began with a simple question: "Is PoRep SRS/PCE unloaded when e.g. Snap proof needs to be made?" The assistant's response (message 2068) was a detailed forensic audit of the codebase, tracing the SRS lifecycle through SrsManager's HashMap (never evicted internally) and the PCE lifecycle through static OnceLock globals (write-once, never cleared). This audit quantified the problem in stark numbers: a 256 GiB machine running PoRep and SnapDeals would have a ~120 GiB baseline, leaving ~126 GiB for working memory—enough for about 9 partitions at 13.6 GiB each. If WindowPoSt also appeared, the baseline ballooned to ~203 GiB, leaving room for only ~3 partitions.

The user then articulated a set of design constraints (message 2069):

  1. Loaded PCE/SRS should consume memory quota from the budget
  2. Add a 5-minute eviction timeout for PCE/SRS slots
  3. Remove configurable preload—all SRS should be loaded on-demand
  4. Working memory should be released in two phases: a/b/c freed immediately after the GPU prove starts, and the remainder freed after the proof completes The assistant responded with a set of clarifying questions (message 2070), exploring design alternatives: single unified budget versus separate pinned/heap pools, auto-detection versus manual configuration, unconditional eviction versus pressure-triggered eviction, and the implications of SRS reload latency (30-60 seconds for a 44 GiB file). The user's answers (message 2071) were decisive: single unified budget, auto-detect from system RAM, evict under pressure with the nuance that freeing unused SRS to allow more pipeline workers is preferred, and two-phase release as described. Message 2073—the subject of this article—is the synthesis of everything that came before. It is the moment when the assistant transforms requirements into architecture.

The Architecture: How Decisions Were Made

The Unified Budget Model

The most fundamental design decision in this message is the choice of a single unified memory budget rather than separate pools for pinned memory (SRS) and heap memory (PCE + working set). This decision was presented as a question in message 2070 with two options, and the user chose option (a).

The reasoning behind this choice is subtle but important. Pinned memory (allocated via cudaHostAlloc) is indeed more constrained than regular heap memory—it is page-locked and non-swappable. However, from the perspective of the system operator, both types of memory consume physical RAM. A machine with 256 GiB of RAM does not care whether a particular allocation is pinned or not; it only cares about total utilization. By using a single budget, the system presents a simple mental model: "the machine has X GiB of RAM, the budget is X minus a safety margin, and all allocations draw from the same pool."

This decision also simplifies the configuration surface. Instead of asking operators to tune pinned_budget and working_memory_budget separately—a task that requires deep knowledge of the engine's internal memory partitioning—the system auto-detects total RAM and handles the rest. The assistant's design removes three configuration options (srs.preload, memory.pinned_budget, memory.working_memory_budget, synthesis.partition_workers) and adds only two (memory.total_budget as an optional override, memory.safety_margin as an optional tweak).

The Eviction Policy: LRU with a 5-Minute Minimum Age

The eviction policy is another carefully considered decision. The assistant proposes an LRU (Least Recently Used) cache with a 5-minute minimum retention window. This means:

  1. Entries younger than 5 minutes are protected from eviction, even under memory pressure.
  2. Entries older than 5 minutes are eviction candidates, but eviction only happens when the budget is under pressure (i.e., a synthesis task is waiting for memory).
  3. Among candidates, the oldest entry (by last_used timestamp) is evicted first. This policy balances several competing concerns. The 5-minute minimum window prevents thrashing—if a proof type is used every few minutes, its SRS and PCE stay cached, avoiding the 30-60 second reload penalty. The pressure-triggered eviction ensures that memory is only freed when it is actually needed, rather than being aggressively evicted only to be reloaded moments later. The LRU ordering ensures that the most recently used circuits—presumably the ones most likely to be needed again—are retained longest. The assistant also considers a critical safety detail: eviction only considers entries where Arc::strong_count == 1, meaning no in-flight proof holds a reference to the data. This prevents a race condition where the cache evicts an SRS or PCE that is actively being used by a concurrent proof. The Arc-based reference counting provides natural thread safety: the cache drops its reference, but the worker's cloned Arc keeps the data alive until the proof completes.

The Two-Phase Working Memory Release

The two-phase release model is a direct response to the user's requirement. The assistant traces the exact code path where a/b/c vectors are dropped (synchronously inside prove_start at supraseal.rs:231-234) and maps this to the first release point (~12.5 GiB freed). The remaining ~1.1 GiB (shell + aux + pending) is released after gpu_prove_finish returns.

This design is notable for its pragmatic handling of the asynchronous deallocation problem. The assistant notes that the dealloc thread runs asynchronously after gpu_prove_finish, and that the budget represents a logical reservation rather than physical RSS. By releasing the budget when the proof finishes (rather than waiting for the async dealloc thread to complete), the system accepts a brief over-commit of ~1.1 GiB. This is a deliberate trade-off: it avoids the complexity of having the dealloc thread signal back into the budget system, at the cost of a small, transient over-commit that is unlikely to cause problems in practice.

The Evictor Callback Pattern

One of the more elegant architectural decisions is the evictor callback. The assistant initially considered a trait-based approach for eviction sources but simplified it to a closure-based pattern. The MemoryBudget struct holds a Box<dyn Fn(u64) -> u64 + Send + Sync>—a callback that takes a needed_bytes parameter and returns the number of bytes actually freed.

This callback is set once during engine initialization and captures references to both the SrsManager and the PceCache. When the budget cannot satisfy an acquire request, it invokes the evictor, which iterates over both caches, collects evictable entries (older than 5 minutes with refcount == 1), sorts them by last_used ascending, and evicts the oldest first until enough memory is freed.

The beauty of this pattern is that it keeps the MemoryBudget generic and testable—it does not need to know about SRS or PCE specifics. The eviction logic is encapsulated in the callback, which is wired up by the engine that has access to all managers. This is a textbook example of dependency inversion.

Assumptions and Their Implications

Assumption: Auto-Detected RAM is Accurate

The assistant assumes that /proc/meminfo MemTotal provides a reliable basis for the memory budget, minus a safety margin. This is generally true for dedicated proving machines, but it has edge cases. If the machine runs other workloads (monitoring agents, log aggregators, SSH daemons, etc.), the actual available memory may be lower than MemTotal minus the safety margin. The configurable safety_margin parameter provides some flexibility, but the default of 5 GiB may be insufficient for machines with significant non-cuzk workloads.

Assumption: SRS Reload Latency is Acceptable

The assistant assumes that a 30-60 second SRS reload latency is acceptable when an evicted circuit type is needed again. This is a reasonable assumption for a proving engine where proofs arrive sporadically, but it could be problematic for latency-sensitive workloads. If a WindowPoSt proof arrives and the WindowPoSt SRS was evicted 30 seconds ago, the system will block for up to a minute to reload it. The 5-minute minimum retention window mitigates this, but it does not eliminate it.

Assumption: synthesis_concurrency Should Be Retained

The assistant explicitly asks whether synthesis_concurrency should also be removed, then recommends keeping it. The reasoning is that synthesis_concurrency controls CPU thread contention (rayon pool sharing), which is a separate concern from memory. This is a sound architectural distinction, but it does introduce a subtle interaction: if the budget limits partition concurrency to, say, 8, but synthesis_concurrency is set to 12, the synthesis_concurrency setting is effectively irrelevant for memory-constrained scenarios. The assistant acknowledges this but argues that the CPU contention knob is still useful for scenarios where memory is plentiful but CPU cores are limited.

Assumption: The Dealloc Thread is Fire-and-Forget

The assistant assumes that the async dealloc thread can be treated as fire-and-forget, accepting a brief over-commit of ~1.1 GiB. This is a pragmatic engineering trade-off, but it does mean that the budget system is not perfectly accurate at all times. A pathological scenario where many proofs finish simultaneously could lead to a cumulative over-commit of several GiB. The assistant implicitly assumes that this transient over-commit is bounded and unlikely to cause OOM, which is reasonable given the scale of the safety margin (5 GiB) relative to the per-proof over-commit (1.1 GiB).

Input Knowledge Required

To fully understand this message, the reader needs:

  1. Knowledge of the cuzk codebase structure: The assistant references specific files (engine.rs, pipeline.rs, srs_manager.rs, config.rs) and line numbers. Understanding the architecture requires knowing what these files contain and how they relate.
  2. Understanding of GPU proving memory lifecycles: The distinction between SRS (pinned memory, loaded once, shared across proofs), PCE (heap memory, extracted from circuit, also shared), and working memory (per-partition allocations for a/b/c vectors, shell, aux) is central to the design.
  3. Familiarity with Rust concurrency primitives: The design uses Arc, AtomicU64, tokio::sync::Notify, Mutex, OnceLock, and Box<dyn Fn> closures. Understanding how these interact is essential.
  4. Knowledge of the Filecoin proof types: The message references PoRep (Proof-of-Replication), WindowPoSt (Window Proof-of-Spacetime), WinningPoSt, and SnapDeals. Each has different memory requirements and SRS/PCE sizes.
  5. Understanding of the existing bug: The message builds on the discovery that working_memory_budget was dead code and that partition_workers was the only throttle. Without this context, the motivation for the redesign is unclear.

Output Knowledge Created

This message creates several forms of knowledge:

  1. A concrete, actionable implementation plan: The message specifies exactly what changes are needed in each file, what new types to create, what existing types to modify, and what configuration options to add or remove. This is not a vague architecture sketch; it is a detailed specification that could be handed to a developer for implementation.
  2. Memory estimation constants: The message provides specific numbers for each proof type's memory footprint (e.g., POREP_PARTITION_FULL_GIB: 13.6, WPOST_FULL_GIB: 13.2). These constants are essential for the budget system and represent a form of empirical knowledge about the system's behavior.
  3. A design rationale: The message documents why each decision was made—why a unified budget, why LRU with 5-minute minimum, why two-phase release, why the evictor callback pattern. This rationale is valuable for future maintainers who need to understand or modify the design.
  4. Configuration migration path: The message shows the old and new configuration formats, making it clear what operators need to change when upgrading.
  5. Edge case identification: The message identifies potential issues (SRS reload latency, the synthesis_concurrency interaction, the fire-and-forget dealloc thread) and explains the trade-offs made.

The Thinking Process

The assistant's reasoning is visible throughout the message, particularly in the way it structures the plan. The architecture overview diagram at the top establishes the high-level relationships. Then each component is described in detail, with specific attention to integration points.

One notable aspect of the thinking process is the assistant's willingness to question its own design. The final section asks: "should synthesis_concurrency and partition_workers both be removed, or just partition_workers?" This is not a rhetorical question—the assistant genuinely identifies a design tension and presents the reasoning for keeping synthesis_concurrency separate. This self-questioning is a hallmark of thorough engineering design.

Another revealing moment is the assistant's handling of the evictor callback. The initial design considered a trait-based approach but was simplified to a closure. The message does not show the intermediate steps, but the reasoning is implicit in the final design: the closure is simpler, more flexible, and avoids the overhead of defining and implementing a trait for each eviction source.

The assistant also demonstrates a strong understanding of the existing codebase. References to specific line numbers (supraseal.rs:231-234, engine.rs:866-926, engine.rs:1409-1568) show that the assistant has read and understood the code, not just the high-level architecture. This attention to detail is critical for a plan that aims to be directly actionable.

Potential Mistakes and Incorrect Assumptions

The synthesis_concurrency Retention

The assistant's recommendation to keep synthesis_concurrency as a separate knob is reasonable, but it may introduce confusion. Operators might not understand the distinction between a memory-based throttle (the budget) and a CPU-based throttle (synthesis concurrency). If both are set incorrectly, the system could be underutilized (budget allows 10 partitions, but synthesis concurrency limits to 2) or overcommitted (synthesis concurrency allows 10, but budget limits to 3). The assistant acknowledges this tension but does not fully resolve it.

The Eviction Callback's Scope

The evictor callback captures references to SrsManager and PceCache, but what about other potential memory consumers? If future development adds new large cached structures, the evictor callback would need to be updated. A more extensible design might use a registry pattern where memory consumers register themselves with the budget. The closure-based approach is simpler but less extensible.

The Safety Margin Default

The default safety margin of 5 GiB may be too small for production systems. Modern Linux systems can use several GiB for page cache, dentry cache, and other kernel structures, especially under memory pressure. A more conservative default (e.g., 10 GiB or 5% of total RAM) might be safer. The assistant does not provide a strong rationale for the 5 GiB default.

The Fire-and-Forget Dealloc Thread

As noted earlier, the assumption that the dealloc thread can be fire-and-forget is a pragmatic trade-off, but it does introduce a small window of over-commit. If the system is operating near its memory limit and multiple proofs finish simultaneously, the cumulative over-commit could theoretically exceed the safety margin. The assistant does not analyze this scenario quantitatively.

Conclusion

Message 2073 is a masterclass in architectural design under constraints. It transforms a set of user requirements—some precise ("add 5 min evict timeout"), some directional ("remove configurable preload")—into a coherent, detailed, and implementable system. The assistant demonstrates deep knowledge of the existing codebase, a clear understanding of memory management trade-offs, and the ability to make and justify design decisions.

The resulting design replaces a fragile static concurrency limit with a robust memory-aware admission control system. It simplifies the configuration surface while adding sophisticated features like LRU eviction with pressure-triggered activation. It handles edge cases like in-flight proof references through Arc reference counting. It provides a clear migration path for existing operators.

But perhaps most importantly, this message captures the reasoning behind the architecture. It is not just a plan; it is a documented set of design decisions, trade-offs, and rationales. For any developer who later needs to implement, modify, or debug this system, this message is invaluable. It answers not just "what should be built" but "why should it be built this way."

In the broader context of the coding session, this message represents a turning point. The earlier messages were diagnostic—auditing the existing code, identifying bugs, understanding memory lifecycles. This message is prescriptive—laying out a new architecture that solves the identified problems. The subsequent implementation work will be guided by this blueprint, and the final system will be judged against the design criteria established here.

For anyone studying how to design complex systems under real-world constraints, this message is a rich case study. It shows how to balance competing concerns, how to simplify without losing essential functionality, how to document design rationale, and how to bridge the gap between requirements and implementation. It is, in short, a model of thoughtful engineering communication.