Freeing 12 GiB Per Partition: A Targeted Memory Optimization in the Phase 12 Split GPU Proving API

Introduction

In high-performance GPU proving pipelines for Filecoin's Proof-of-Replication (PoRep), memory is often the most constrained resource. The system under development—a Groth16 proof generation engine called cuzk—runs on a machine with 755 GiB of RAM, and every gigabyte counts. When the Phase 12 split GPU proving API was benchmarked with 12 partition workers (pw=12) and 15 concurrent jobs (j=15), the process was killed by the Out-Of-Memory (OOM) killer after RSS peaked at 668 GiB. This left only 87 GiB of headroom for the kernel, page tables, and other processes—a dangerously thin margin.

The user, who had been closely tracking these benchmarks, asked a pointed question in [msg 3055]: "the pending proof handle has nothing that can be freed early?" This question set off a precise forensic investigation by the assistant in [msg 3062], the subject of this article. What emerged was a textbook example of how deep understanding of data flow across a C++/Rust/CUDA boundary can yield significant memory savings without any change to the computational logic.

The Question and Its Context

The Phase 12 split API, committed at 99c31c2c, refactored the monolithic generate_groth16_proofs_c C++ function into two phases: prove_start (which returns a PendingProofHandle) and finalize (which completes the proof). This allowed the GPU worker thread to offload the ~1.7s b_g2_msm computation to a background finalizer task and immediately loop back to pick up the next synthesized partition. The throughput improved by ~2.4% (from 38.0s to 37.1s per proof), but the memory picture worsened: the PendingProofHandle now held synthesis data alive for the duration of the background b_g2_msm computation, adding ~1.7s of extra retention per partition.

The user's question cut to the heart of the problem: was all that data really needed for the full duration? The assistant recognized this as an opportunity to trace exactly which fields the background thread accessed, and which it did not.

The Investigation: Tracing Data Dependencies

The assistant's response in [msg 3062] is a model of systematic memory analysis. It begins by examining the ProvingAssignment struct definition in the Rust source:

pub struct ProvingAssignment<Scalar: PrimeField> {
    pub a_aux_density: DensityTracker,
    pub b_input_density: DensityTracker,
    pub b_aux_density: DensityTracker,
    pub a: Vec<Scalar>,
    pub b: Vec<Scalar>,
    pub c: Vec<Scalar>,
    pub input_assignment: Vec<Scalar>,
    pub aux_assignment: Vec<Scalar>,
    // ...
}

The assistant then cross-references this with the C++ prep_msm_thread code—the background thread that runs b_g2_msm while the GPU worker moves on. Using awk and grep commands against the CUDA source file groth16_cuda.cu, it searches for every reference to prover.a, prover.b, and prover.c within the thread's scope. The search reveals a critical finding: the a, b, and c fields—each a Vec&lt;Scalar&gt; of approximately 130 million elements at 32 bytes per element, totaling roughly 12 GiB per partition—are only accessed in the GPU kernel region (lines 847–1195 of the CUDA file). This region includes cudaHostRegister, cudaMemcpyAsync, and cudaHostUnregister calls. All of these complete before prove_start returns to the caller.

In contrast, the prep_msm_thread needs only:

The Insight: ~12 GiB of Immediately Freeable Memory

The assistant's conclusion is stated with clarity and precision:

We can free the a, b, c vectors immediately after prove_start returns! That's ~12 GiB per partition freed immediately instead of held for ~1.7s.

This is the core insight. For each partition in flight, 12 GiB of NTT evaluation data was being retained unnecessarily. With pw=12, up to 12 partitions could be synthesized concurrently, meaning up to 144 GiB of unnecessary memory retention. This alone could explain the jump from 367 GiB peak at pw=10 to 668 GiB at pw=12—the additional two workers, combined with the extra ~1.7s retention window from the split API, created a multiplicative effect.

The assistant also confirms that input_assignment and aux_assignment are already moved out (empty Vecs after synthesis), so they don't contribute to the retention problem. The density trackers are tiny (~48 MB total) and genuinely needed by the background thread, so they must stay.

The Thinking Process: What Makes This Analysis Effective

Several aspects of the assistant's reasoning are worth highlighting:

First, the assistant immediately translates the user's high-level question into a concrete data-flow investigation. Rather than speculating or guessing, it reads the actual source files and traces which fields are accessed by which code paths. The use of awk to scan the CUDA source for specific field references is a pragmatic, targeted approach—it doesn't require reading the entire 1,200-line file, just the relevant sections.

Second, the assistant correctly identifies the boundary between GPU-completed work and CPU-background work. The key insight is temporal: the cudaHostRegister/cudaMemcpyAsync/cudaHostUnregister sequence for a, b, c completes before the function returns. The GPU has already consumed these vectors by the time the C++ function hands control back to Rust. The background thread runs after this point, so it cannot need data that was already consumed by the GPU.

Third, the assistant validates negative cases. It explicitly checks that the prep_msm_thread does not reference prover.a, prover.b, or prover.c—the grep returns empty for the thread's scope. This negative evidence is as important as positive evidence.

Fourth, the assistant quantifies everything. Each field's size is estimated: 130M elements × 32 bytes = ~4 GiB per vector, three vectors = ~12 GiB total. This turns an abstract observation ("some memory can be freed") into a concrete, actionable number ("12 GiB per partition").

Assumptions and Potential Pitfalls

The analysis makes several implicit assumptions that are worth examining:

  1. The cudaHostUnregister calls truly indicate GPU completion. The assistant assumes that after cudaHostUnregister returns, the GPU is done reading the host memory. This is correct for properly synchronized CUDA streams—cudaHostUnregister blocks until all pending transfers involving the registered memory complete. However, if there are asynchronous operations that reference the data through other paths, freeing it could cause use-after-free. The assistant's grep confirms no other references exist within the prep_msm_thread, but it does not check for indirect references through other C++ functions or callbacks.
  2. The size estimate of 130M elements per vector is accurate. This depends on the specific circuit parameters (number of constraints in the PoRep circuit). The assistant uses this as a rough estimate, but the actual size could vary. Even if off by 20%, the savings would still be substantial (~10 GiB per partition).
  3. Freeing Rust Vecs is safe across the FFI boundary. The a, b, c vectors are Rust-owned Vec&lt;Scalar&gt; that were passed to C++ via pointers. After prove_start returns, the Rust side can drop these Vecs (setting them to empty or replacing them with Vec::new()). The C++ provers_owned struct holds copies of the Assignment struct, but those copies contain pointers into the original Rust Vecs' buffers. Wait—this is a critical point. The assistant's plan is to free the Rust Vecs, but the C++ provers_owned struct might still hold pointers to those buffers. Let me re-read the message... Actually, looking at the message more carefully, the assistant says "We can free the a, b, c vectors immediately after prove_start returns!" But the provers_owned in the C++ heap struct was created by copying the Assignment structs. The Assignment struct (from the C++ side) contains pointers like prover.a which point into the Rust Vec&lt;Scalar&gt; buffers. If those Vecs are dropped, the C++ pointers become dangling. However, the assistant's grep shows that the prep_msm_thread does NOT access prover.a, prover.b, or prover.c. So even though the pointers exist in the copied struct, they are never dereferenced. The question is whether the assistant considered this subtlety. The message doesn't explicitly address it, but the grep evidence is clear: the background thread never reads those fields.

Input and Output Knowledge

Input knowledge required to understand this message:

Broader Significance

This message represents a critical turning point in the Phase 12 optimization cycle. The team had hit a memory wall: pw=12 was impossible, and the optimal configuration (pw=10, gw=2, gt=32) was running at 367 GiB peak—nearly 50% of available RAM. The user's question and the assistant's analysis revealed that a significant fraction of that memory was being held unnecessarily.

The insight is also a testament to the value of the split API design. Before Phase 12, the monolithic generate_groth16_proofs_c function held all data until the proof was complete—there was no opportunity to free anything early. The split API created a natural boundary where ownership could be partitioned: data needed by the GPU kernel could be released as soon as the kernel finished, while data needed by the CPU background thread could be retained separately. The assistant's analysis simply identified which data fell into which category.

In the broader narrative of the optimization project, this message sets the stage for the next chunk of work (Chunk 1 of Segment 30), where the assistant implements the early deallocation and builds a global buffer tracker to diagnose remaining memory pressure. The 12 GiB per partition saving, combined with the subsequent semaphore and channel capacity optimizations, ultimately enables pw=12 to run without OOM, dropping peak RSS from 668 GiB to 294.7 GiB.

Conclusion

Message [msg 3062] is a masterclass in targeted memory analysis. In response to a single, well-framed question from the user, the assistant systematically traced data dependencies across a complex C++/Rust/CUDA codebase, identified a 12 GiB per partition optimization opportunity, and laid the groundwork for the next phase of memory pressure reduction. The analysis is notable for its precision—using grep and awk to prove negative access patterns—and its practical impact: freeing nearly 150 GiB of peak memory at pw=12. It demonstrates that in high-performance systems engineering, the most valuable optimizations often come not from changing what the code computes, but from understanding exactly when each piece of data is no longer needed.