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

Introduction

In the high-stakes world of Filecoin proof generation, memory is the invisible constraint that governs everything. A single 32 GiB PoRep (Proof-of-Replication) proof consumes over 44 GiB of CUDA-pinned memory for its Structured Reference String (SRS), another 26 GiB on the heap for its Pre-Compiled Circuit (PCE), and roughly 13.6 GiB per partition during synthesis and GPU proving. When multiple proof types—PoRep, WindowPoSt, WinningPoSt, SnapDeals—are processed concurrently, the memory footprint can balloon past 200 GiB, pushing even a 256 GiB machine to its breaking point.

Message 2071 in this coding session captures a pivotal moment: the assistant, having completed a deep forensic audit of the entire memory lifecycle across the cuzk GPU proving engine, now synthesizes that analysis into a concrete architectural design. Guided by the user's explicit requirements—a single unified memory budget, LRU eviction with a 5-minute minimum retention window, on-demand SRS/PCE loading, and two-phase working memory release—the assistant works through the design in real time, reasoning about trade-offs, verifying assumptions against the codebase, and laying the groundwork for the comprehensive cuzk-memory-manager.md specification that will follow.

This message is not merely a plan. It is the crystallization of dozens of hours of forensic tracing, the resolution of conflicting design pressures, and the moment where a fragile, static concurrency limit is replaced with a robust, memory-aware admission control system. To understand why this message was written and what it accomplishes, we must first understand the problem it solves.

The Problem: A Memory-Unaware Throttle

Before this message, the cuzk engine controlled partition concurrency through a single mechanism: partition_workers, a static semaphore count set at startup. This semaphore limited how many partitions could be synthesized and proved simultaneously, but it had no awareness of how much memory those partitions actually consumed. Worse, the working_memory_budget configuration option—which could have provided memory awareness—was entirely dead code, never enforced anywhere in the pipeline.

This design was fragile in several ways. First, the baseline memory consumption grew silently as additional proof types were loaded. SRS files for PoRep (44 GiB), WindowPoSt (57 GiB), and SnapDeals (33 GiB) were all pinned in CUDA memory simultaneously, never evicted. PCE structures for each circuit type accumulated in static OnceLock variables, also never freed. An operator who configured partition_workers = 12 for PoRep-only operation would find that same setting caused OOM crashes when WindowPoSt proofs began arriving, because the baseline had grown by 57 GiB without any adjustment to the concurrency limit.

Second, the PCE cold-start extraction path was completely ungated. On a first-ever startup, the engine would spawn a background thread to extract the PCE from a full circuit synthesis—a process with a ~40 GiB transient memory spike—while simultaneously allowing partition synthesis to proceed at full concurrency. On a 256 GiB machine with partition_workers = 12, this could produce: 12 partitions × 13.6 GiB = 163 GiB working set, plus 44 GiB SRS, plus 40 GiB PCE extraction = 247 GiB total, triggering an OOM kill.

Third, the SRS on-demand loading path had no coordination with the memory system. If a new proof type arrived mid-operation and its SRS wasn't preloaded, the engine would call ensure_loaded inside the partition dispatch path, triggering a 30-60 second load with a transient peak of ~88 GiB (mmap + cudaHostAlloc overlap), potentially causing an OOM while partition synthesis was already consuming memory.

These were not hypothetical problems. The user had experienced crashes, and the assistant's forensic audit had confirmed the root cause: the engine had no unified, memory-aware admission control system.

The Design Mandate: What the User Asked For

The user's requirements, articulated in message 2069 and refined through clarifying questions, were remarkably clear and specific:

Single unified budget: SRS, PCE, and working memory should all draw from the same pool, derived from total system RAM. No separate pinned vs. heap budgets.
Auto-detection from system RAM: The total budget should be determined automatically from /proc/meminfo MemTotal, with a configurable override for edge cases. No manual working_memory_budget configuration.
LRU eviction with 5-minute minimum age, triggered by budget pressure: SRS and PCE entries should be evicted only when (a) they haven't been used for at least 5 minutes, AND (b) the budget is under pressure (i.e., synthesis tasks are waiting for memory). This is a nuanced policy: entries younger than 5 minutes are protected from eviction even under memory pressure, preventing thrashing. But entries that have been idle for longer are fair game if the budget is tight.
Two-phase working memory release: The budget should be released in two stages. When the a/b/c vectors are dropped synchronously inside prove_start (approximately 12.5 GiB), that portion should be immediately returned to the budget. The remaining ~1.1 GiB (shell + aux) should be released after the proof completes and the async deallocation thread finishes.
Remove configurable preload: All SRS should be loaded on-demand, not preloaded at startup. The preload configuration option should be removed. These requirements formed the design mandate. The assistant's task in message 2071 was to translate this mandate into a concrete, implementable architecture.

The Reasoning Process: From Mandate to Architecture

The assistant's reasoning in message 2071 is a masterclass in architectural thinking. It moves through several distinct phases, each building on the previous one.

Phase 1: Confirming the Eviction Policy

The first insight is about the eviction policy. The user's answer to the clarifying question—"Evict under pressure, noting that allocating more pipeline workers is preferred"—reveals a subtle but important design point. The eviction logic should not be a simple timer that evicts after 5 minutes unconditionally. Instead, it should be a pressure-triggered eviction: the budget manager checks whether evictable entries exist only when a synthesis task is waiting for memory and the budget is exhausted. If an entry hasn't been used for 5 minutes and memory pressure exists, it gets evicted. If no pressure exists, the entry stays in memory indefinitely, avoiding unnecessary reload latency.

This is essentially an LRU cache with a minimum retention window. Entries younger than 5 minutes are protected. Entries older than 5 minutes are eviction candidates, but eviction only happens when the budget is under pressure. This design balances two competing goals: maximizing throughput (by keeping frequently-used data in memory) and maximizing concurrency (by freeing memory for additional pipeline workers when demand is high).

The assistant recognizes this immediately: "if SRS/PCE hasn't been used for 5 min AND the budget is under pressure (i.e., synthesis tasks are waiting for memory), then evict to free up budget for working set."

Phase 2: Designing the Core Components

With the eviction policy confirmed, the assistant begins sketching the core components:

MemoryBudget: A central struct that tracks total allocation and current usage. It provides acquire() and release() methods and coordinates eviction when the budget is exhausted.
MemoryReservation: An RAII (Resource Acquisition Is Initialization) type that holds a claim on the budget. When dropped, it releases the claimed memory. This ensures that budget releases happen automatically, even in error paths.
Eviction integration: The budget must be able to trigger eviction from both the SRS manager and the PCE cache when it cannot satisfy an acquisition request. The assistant considers several approaches for the eviction integration. Initially, it thinks about a trait-based approach where SRS and PCE implement a common Evictable interface. But then it realizes this is over-engineering: "the trait-based eviction approach adds unnecessary complexity. Since we only have two eviction sources—the SRS manager and PCE cache—both living within cuzk-core, I can simplify by having the budget hold direct references to these instead." This is a classic engineering trade-off: abstraction vs. simplicity. The assistant correctly chooses simplicity for a system with only two eviction sources, avoiding the complexity of a generic trait that would need to be implemented by future, unknown eviction sources.

Phase 3: The PCE Cache Transformation

A critical design decision involves the PCE cache. Currently, PCE structures are stored in static OnceLock<PreCompiledCircuit<Fr>> variables—write-once, never-cleared globals. To support eviction, this must change to something like Mutex<Option<Arc<PreCompiledCircuit<Fr>>>>.

The assistant recognizes that this transformation is safe because in-flight syntheses hold their own Arc clones. When the cache evicts an entry, it simply removes its reference from the cache. Any worker that holds an Arc to the PCE can continue using it until its proof completes. The data is only freed when the last Arc reference is dropped.

This is a textbook use of Arc for concurrent cache management, and the assistant correctly identifies it as the right approach.

Phase 4: The Two-Phase Release Model

The two-phase working memory release requires careful tracing of the actual code flow. The assistant maps out the partition lifecycle:

[acquire 13.6 GiB] → synthesis → prove_start → [release 12.5 GiB]
                                                    ↓
                                    pending proof (shell + aux ~1.1 GiB)
                                                    ↓
                                             prove_finish → [release 1.1 GiB]

The key question is: where exactly do the a/b/c vectors get freed? The assistant traces this to supraseal.rs:231-234, where prover.a = Vec::new() replaces the old vector, causing its destructor to run synchronously. This means the budget release for a/b/c can happen immediately after prove_start returns, in the same blocking thread.

For the remaining ~1.1 GiB, the assistant considers whether to wait for the async deallocation thread to complete. It makes a pragmatic decision: "The dealloc thread runs asynchronously in the background, and since the budget represents a logical reservation rather than physical memory, releasing it after the proof finishes makes sense even if the OS hasn't reclaimed the pages yet." This is a reasonable engineering judgment—the ~1.1 GiB over-commit window is small and brief, and the complexity of wiring a signal from the dealloc thread back into the budget system is not justified.

Phase 5: The Bellperson Integration Challenge

A subtle design challenge involves bellperson, the underlying constraint synthesis library. The a/b/c intermediate data is owned by bellperson's prover structures. When prove_start clears these vectors, the drop happens inside bellperson code. But the budget release must happen in cuzk's pipeline layer.

The assistant considers several approaches:

Move the drop logic into the pipeline layer
Add a callback that bellperson invokes after freeing the vectors
Have bellperson return the vectors so cuzk can manage the release The assistant concludes that the cleanest approach is to call the budget release in pipeline.rs immediately after prove_start returns, since that's when the memory is actually freed. This avoids modifying bellperson at all—the pipeline layer simply observes that the vectors have been cleared and releases the corresponding budget reservation.

Phase 6: Verifying Assumptions with Grep

Before writing the plan, the assistant executes two grep commands to verify its understanding of the existing codebase. The first grep searches for partition_semaphore and related terms to confirm the semaphore-based concurrency model. The second grep searches for the PCE OnceLock variables (POREP_32G_PCE, WINNING_POST_PCE, etc.) to confirm the current caching pattern.

These greps serve two purposes. First, they validate the assistant's mental model against the actual code. Second, they provide concrete line numbers and file paths that will be referenced in the design specification. This is a critical step: the assistant is not designing in the abstract but against a specific, known codebase.

Assumptions Made by the Assistant

The assistant's reasoning relies on several assumptions, some explicit and some implicit:

Glibc behavior: The assistant assumes that dropping large Vec allocations causes glibc to eventually return the memory to the OS, but acknowledges that this may not happen immediately without malloc_trim(0). This is a well-known characteristic of glibc's memory allocator, which maintains per-thread arenas and does not eagerly release memory to the kernel.
Acceptable reload latency: The assistant assumes that a 30-60 second SRS reload latency is acceptable for the first proof after eviction. This is a reasonable assumption given that proof generation itself takes minutes, but it should be validated against operational requirements.
Budget as logical reservation: The assistant assumes that the memory budget can track logical ownership rather than physical RSS. This means the budget may briefly over-commit by ~1.1 GiB while the dealloc thread runs. The assistant judges this acceptable because the over-commit is small relative to the total budget and short-lived.
Single unified budget is sufficient: The assistant assumes that a single budget derived from total system RAM is sufficient, without separate tracking of pinned vs. heap memory. This is a simplification that works for most systems but could be problematic on machines with strict CUDA pinned memory limits.
Arc-based safety for eviction: The assistant assumes that using Arc for PCE and SRS references provides safe concurrent access during eviction. This is correct: the cache removes its reference, but in-flight workers retain their clones until completion.

Potential Mistakes or Incorrect Assumptions

While the assistant's reasoning is generally sound, there are a few areas where the assumptions could be questioned:

The ~1.1 GiB over-commit window: The assistant assumes that releasing the budget before the dealloc thread completes is safe because the over-commit is small. However, if multiple proofs finish simultaneously, these over-commits could accumulate. If 12 partitions all finish at roughly the same time, the total over-commit could be ~13 GiB, which is significant. The assistant does not address this edge case.
Auto-detection from system RAM: The assistant assumes that /proc/meminfo MemTotal is the right source for auto-detection. However, on containerized or virtualized systems, this may report the host's total RAM rather than the container's limit. The assistant acknowledges this by suggesting a configurable override, but does not discuss cgroup-aware detection.
Eviction under pressure vs. proactive eviction: The assistant's design evicts only under pressure. This means that a long-idle SRS entry will stay in memory indefinitely if no new synthesis tasks arrive. This is efficient for memory but could be surprising to operators who expect the 5-minute timeout to be unconditional.
The trait-based vs. direct reference decision: The assistant correctly chooses simplicity, but this decision couples MemoryBudget directly to SrsManager and PceCache. If a third eviction source is added later, the budget will need to be modified. This is a reasonable trade-off for now.

Input Knowledge Required

To fully understand this message, the reader needs:

Understanding of CUDA memory management: Specifically, the distinction between regular heap memory and CUDA pinned memory (cudaHostAlloc), and the implications for system RAM consumption.
Knowledge of Filecoin proof types: PoRep, WindowPoSt, WinningPoSt, and SnapDeals, and their relative memory footprints. The reader must understand that different proof types have different SRS sizes (44 GiB for PoRep, 57 GiB for WindowPoSt, etc.).
Familiarity with the cuzk codebase architecture: The engine's partition dispatch pipeline, the SRS manager, the PCE caching layer, and the GPU worker loop. The assistant references specific line numbers and file paths.
Understanding of Rust concurrency primitives: Arc, OnceLock, Mutex, Semaphore, and how they interact with memory management.
Knowledge of glibc memory allocator behavior: Specifically, the fact that free() does not immediately return memory to the OS, and the role of malloc_trim().
Familiarity with LRU cache design: The concept of eviction policies, minimum retention windows, and pressure-triggered eviction.

Output Knowledge Created

This message produces several forms of output knowledge:

A validated design architecture: The assistant confirms that the user's requirements are feasible and internally consistent, and translates them into a concrete component architecture.
Specific implementation guidance: The assistant identifies the exact files and code locations that need to change, including line numbers for the semaphore creation, PCE cache access, and a/b/c vector clearing.
Edge case analysis: The assistant identifies potential issues like the bellperson integration challenge and the dealloc thread timing, and proposes specific solutions.
A decision record: The message documents the reasoning behind key design decisions—why a single unified budget was chosen over separate pools, why eviction is pressure-triggered rather than unconditional, why the budget uses logical ownership rather than physical RSS tracking.
Verification of existing code patterns: The grep commands confirm the assistant's understanding of the current codebase, providing a validated foundation for the design specification that will follow.

The Thinking Process: A Window into Architectural Design

The most valuable aspect of message 2071 is the visible thinking process. The assistant does not simply present a finished design. It works through the design incrementally, considering alternatives, rejecting over-engineered approaches, and verifying assumptions against the codebase.

We see the assistant initially consider a trait-based eviction interface, then reject it as unnecessary complexity. We see it trace the a/b/c vector clearing through the codebase to confirm the exact point where memory is freed. We see it consider the bellperson integration challenge from multiple angles before settling on the simplest approach. We see it weigh the trade-off between logical budget tracking and physical RSS tracking, and make a pragmatic judgment about the ~1.1 GiB over-commit window.

This is not a linear process. The assistant loops back to reconsider earlier decisions as new information emerges. The trait-based approach is considered, then rejected, then the implications of that rejection are traced through the design. The two-phase release model is refined as the assistant traces the actual code flow more carefully.

The message also reveals the assistant's awareness of the broader system context. It does not design in isolation but considers how the memory budget will interact with SRS loading, PCE extraction, partition dispatch, GPU proving, and async deallocation. It considers failure modes (OOM during SRS load, transient spikes during PCE extraction) and ensures the design addresses them.

Conclusion

Message 2071 is the pivotal design message in this coding session. It takes the user's high-level requirements—unified budget, LRU eviction, on-demand loading, two-phase release—and translates them into a concrete, implementable architecture. The assistant's reasoning demonstrates deep understanding of the codebase, careful consideration of trade-offs, and pragmatic engineering judgment.

The design that emerges from this message will be codified in the cuzk-memory-manager.md specification, which will define the MemoryBudget and MemoryReservation core types, the evictable PceCache, the budget-gated SRS loading, the streamlined configuration, and the exact integration points in the engine's partition dispatch and GPU worker loops.

But more than the specification itself, this message captures the process of architectural design—the iterative refinement, the trade-off analysis, the verification against real code. It is a case study in how to translate user requirements into a robust, memory-aware admission control system for a high-performance GPU proving engine.

The fragile static concurrency limit is dead. In its place rises a system that understands memory, adapts to changing workloads, and gracefully handles the competing demands of multiple proof types. And message 2071 is where that system was born.