Designing Phase 10: The Two-Lock Architecture for GPU Proving Pipeline Optimization

In the high-stakes world of Filecoin storage proving, where proofs must be generated continuously against tight deadlines, every millisecond of GPU time counts. The cuzk project—a persistent GPU-resident SNARK proving engine—had already undergone eight optimization phases, each peeling back another layer of the performance onion. Phase 9 had just revealed a surprising new bottleneck: after slashing GPU kernel times by 51%, the bottleneck had shifted from PCIe transfers to CPU memory bandwidth contention. Now, in a single dense message (message 2583), the assistant sat down to design Phase 10: a two-lock architecture intended to better overlap CPU and GPU work and break through the new performance wall.

This message is a masterclass in systems design reasoning under real hardware constraints. It shows the assistant working through lock ordering, VRAM capacity analysis, deadlock scenarios, and CUDA memory management semantics to arrive at a clean, simplified design. Let us dive deep into what this message contains, why it was written, and what it reveals about the art of GPU pipeline optimization.

The Optimization Journey That Led Here

To understand message 2583, we must first understand the context that produced it. The cuzk project is a pipelined SNARK proving engine for Filecoin's Proof-of-Replication (PoRep) protocol. It generates Groth16 proofs across multiple GPU partitions, each requiring substantial computation and memory—approximately 200 GiB peak memory for the full pipeline, with individual GPU partitions consuming up to 13.8 GiB of VRAM on 16 GiB devices.

Phase 8 had introduced a dual-worker GPU interlock: each GPU got its own std::mutex, and two workers per device could interleave their execution—one running GPU kernels while the other did CPU preprocessing. This achieved a 2.4x improvement over baseline, pushing throughput to 37.4 seconds per proof.

Phase 9 added PCIe transfer optimization: by pre-staging NTT input data to VRAM before acquiring the GPU mutex, the assistant reduced GPU kernel time from 3.75 seconds to 1.82 seconds per partition—a 51% improvement in GPU compute. But the steady-state throughput only improved 14%, from 37.4 to 32.1 seconds per proof in isolation, and degraded further to ~41 seconds at high concurrency. Something else was limiting performance.

The culprit was CPU memory bandwidth contention. The prep_msm function (which splits vectors for multi-scalar multiplication) took 1.91 seconds per partition, and b_g2_msm (G2 MSM) took 0.48 seconds. Together, these CPU operations consumed 2.39 seconds per partition—longer than the GPU kernels. At high concurrency (c=15+), ten synthesis workers competed with these CPU MSM operations for 8-channel DDR5 memory bandwidth, inflating CPU times by 12–27% and capping throughput at ~41 seconds per proof.

The bottleneck had moved. The GPU was now waiting for the CPU.

The User's Proposal

The user recognized this shift and proposed a novel solution in the preceding message: add more GPU workers (3–4 per device) with two locks instead of one. The existing compute_mtx would serialize GPU kernel execution, while a new mem_mtx would handle VRAM allocation, deallocation, and pre-staging. The idea was that with more workers, the CPU preprocessing of one worker could overlap with the GPU kernel execution of another, hiding the 2.39-second CPU critical path.

But this proposal faced a fundamental challenge: VRAM capacity. On a 16 GiB GPU, peak usage during H-MSM (the heaviest kernel) was approximately 13.8 GiB. There was only ~2 GiB headroom. Worker B could not pre-stage 12 GiB of buffers while Worker A's buffers were still live. The alloc and free operations had to be serialized, which threatened to negate the benefit of additional workers.

The assistant's task in message 2583 was to work out the detailed design, identify potential deadlocks, and produce a viable implementation plan.

The Assistant's Design Process

Message 2583 begins with the assistant having gathered all necessary information. The preceding messages show the assistant reading the CUDA source code (groth16_cuda.cu), tracing VRAM allocations through a subagent task, examining the engine's worker spawning logic in engine.rs, and analyzing the FFI boundary in lib.rs. Now, with the full picture assembled, the assistant produces a detailed design plan.

The message opens with a clear statement of intent: "Good. Now I have the full picture. Here's the detailed plan."

First Iteration: The Cooperative Free Slot

The first design iteration proposes a gpu_locks struct containing two mutexes plus a "cooperative free slot" mechanism:

struct gpu_locks {
    std::mutex compute_mtx;  // Held during GPU kernel execution
    std::mutex mem_mtx;      // Held during VRAM alloc/free
    // Cooperative free slot: previous worker's buffers to free
    fr_t* pending_free_dbc = nullptr;
    fr_t* pending_free_da = nullptr;
    cudaEvent_t pending_free_ev_a = nullptr;
    cudaEvent_t pending_free_ev_b = nullptr;
    cudaEvent_t pending_free_ev_c = nullptr;
    cudaStream_t pending_free_stream = nullptr;
};

The idea is that each worker, after finishing its GPU kernels, would register its freed buffers in a shared slot. The next worker, when acquiring mem_mtx, would free the predecessor's buffers before allocating its own. This avoids the resource deadlock scenario where Worker A holds compute_mtx waiting for mem_mtx to free VRAM, while Worker B holds mem_mtx waiting for VRAM that Worker A has not yet released.

The assistant then walks through the full worker flow with this design:

lock(mem_mtx):
  1. Free predecessor's pending buffers (if any)
  2. cudaDeviceSynchronize + pool trim
  3. cudaMemGetInfo
  4. cudaMalloc d_a + d_bc (12 GiB)
  5. async upload + events
unlock(mem_mtx)

lock(compute_mtx):
  6. stream.wait(events), NTT, batch_add, tail_msm
  7. cudaFree(d_bc) inside compute_mtx
  8. d_a freed when gpu_ptr_t goes out of scope
unlock(compute_mtx)

lock(mem_mtx):
  9. Register events/stream in pending_free slot
unlock(mem_mtx)

The Critical Insight

Then comes the pivotal moment in the message. The assistant pauses and re-examines the VRAM lifecycle:

"Wait — this simplifies things significantly. If d_a and d_bc are both freed inside compute_mtx, then we don't need a cooperative free slot at all! The next worker's mem_mtx alloc just needs to ensure the previous worker's kernels are done (via cudaDeviceSynchronize inside mem_mtx)."

This is the key insight. The assistant realizes that d_a (4 GiB) and d_bc (8 GiB)—the two largest VRAM allocations—are freed synchronously via cudaFree inside the compute_mtx region. cudaFree is a synchronous host-side call: when it returns, the memory is genuinely free and available for reallocation. There is no need for a cooperative handoff mechanism.

The assistant then traces through the entire VRAM lifecycle more carefully:

The Simplified Design

The final design is remarkably clean:

[no lock]     cudaHostRegister + prep_msm_thread starts
lock(mem_mtx)
  pool_trim + cudaMemGetInfo
  cudaMalloc 12 GiB + async upload + events
unlock(mem_mtx)
lock(compute_mtx)
  stream.wait(events), NTT kernels
  cudaFree(d_bc)
  H-MSM (uses d_a)
  gpu_ptr_t(d_a) destructor → cudaFree
  batch_add + tail_msm
unlock(compute_mtx)
[no lock]     destroy events + stream
              cudaHostUnregister
              prep_msm_thread.join()
              Epilogue

No nested locks, no cooperative free slot, no deadlock risk. The two locks are acquired and released in strict sequence: mem_mtx first, then compute_mtx. Because d_a and d_bc are freed synchronously inside compute_mtx, the VRAM is genuinely available when the next worker acquires mem_mtx.

But the assistant identifies one remaining question: is cudaMemPoolTrimTo inside mem_mtx sufficient without cudaDeviceSynchronize? The pool-cached memory from cudaFreeAsync calls inside compute_mtx might not be reclaimable until the device has actually processed the free operations. The assistant's reasoning is that by the time compute_mtx is released, all stream operations have been synchronized (the per-GPU thread calls sync() implicitly through various MSM operations). But this is an assumption that needs verification.

Assumptions and Potential Pitfalls

The design in message 2583 makes several assumptions that deserve scrutiny:

1. Synchronous cudaFree guarantees immediate availability. This is correct for cudaFree—it is a synchronous host-side operation that blocks until the memory is freed. However, CUDA driver behavior can vary across driver versions and GPU architectures. Some implementations may defer actual VRAM reclamation until the next synchronization point.

2. Pool trim without device synchronization is sufficient. This is the most questionable assumption. cudaMemPoolTrimTo trims the memory pool but may not reclaim memory that is still referenced by in-flight operations. If any GPU operation from the previous compute_mtx region has not fully completed, the pool trim might not free the expected memory. This is precisely the issue that will later cause problems in the implementation phase (as seen in Chunk 1 of the segment analysis, where device-global synchronization conflicts emerge).

3. The two-lock design actually achieves overlap. The critical question is whether the CPU work (prep_msm, b_g2_msm) can genuinely overlap with another worker's GPU kernels. The CPU work runs without any lock, so in theory it can overlap. But both CPU and GPU operations contend for memory bandwidth (DDR5 for CPU, VRAM bandwidth for GPU), and they share the same PCIe bus. The actual overlap might be less than expected due to these shared resources.

4. Three workers per device is feasible. With only ~2 GiB headroom on a 16 GiB GPU, adding a third worker increases memory pressure. If any worker's allocations are slightly larger than expected, or if fragmentation occurs, the system could run out of memory. The assistant's design relies on the fact that only one worker holds mem_mtx at a time, but the cumulative VRAM footprint across workers in different phases could still exceed capacity.

These assumptions are not flaws in the design—they are the normal uncertainties of systems optimization. The assistant's approach is to design the simplest viable solution, implement it, and then measure. The question at the end of the message shows the assistant is aware of the key risk and is seeking confirmation before proceeding.

The Thinking Process Visible in the Reasoning

What makes message 2583 particularly valuable is the visible thinking process. The assistant does not just present a final design—it walks through the reasoning step by step, showing the evolution from a complex solution to a simple one.

The message shows the assistant:

  1. Starting with a complex design (cooperative free slot) that handles the worst-case deadlock scenario
  2. Identifying a problem (the resource deadlock where Worker A holds compute_mtx waiting for mem_mtx, while Worker B holds mem_mtx waiting for VRAM that Worker A has not freed)
  3. Having an insight (d_a and d_bc are freed synchronously via cudaFree inside compute_mtx, so the VRAM is genuinely available when compute_mtx is released)
  4. Simplifying the design dramatically (eliminating the cooperative free slot entirely)
  5. Identifying remaining uncertainties (the pool trim behavior without cudaDeviceSynchronize) This is the hallmark of good systems design: iterate toward simplicity. The initial design with cooperative free slots was correct but complex. The insight about synchronous frees allowed a much simpler solution. The assistant demonstrates the engineering discipline of not over-engineering—if the simpler solution works, use it. The message also shows the assistant tracing through the concrete details of the codebase. Rather than reasoning abstractly about locks, the assistant maps each lock region to specific lines of code and specific allocations. This grounded reasoning is what makes the design credible.

Input Knowledge Required

To fully understand message 2583, the reader needs:

Output Knowledge Created

Message 2583 creates several pieces of new knowledge:

Conclusion

Message 2583 represents a pivotal moment in the cuzk optimization journey. It is the transition from Phase 9 (identifying the CPU memory bandwidth bottleneck) to Phase 10 (designing the solution). The assistant's careful reasoning through lock ordering, VRAM constraints, and CUDA memory semantics produces a clean design that avoids the deadlock risks of more complex approaches.

The message is also a testament to the value of deep systems knowledge. The insight that synchronous cudaFree eliminates the need for a cooperative free slot comes from understanding the semantics of CUDA memory management at a fundamental level. This is not knowledge that can be easily looked up—it comes from experience with GPU programming and an understanding of how CUDA drivers actually behave.

In the next phase, the assistant will implement this design and discover whether the assumptions hold. As the segment analysis reveals, the implementation will encounter unexpected issues: cudaDeviceSynchronize and cudaMemPoolTrimTo inside mem_mtx are device-global operations that conflict with another worker holding compute_mtx, effectively serializing the two locks and destroying the intended overlap. The design will need further refinement. But that is the nature of systems optimization—each iteration reveals new constraints, and the design evolves to meet them.

Message 2583 captures the moment of design insight, before the implementation reality sets in. It is a snapshot of a developer thinking through a hard problem, working from complexity toward simplicity, and preparing to build something that will push the boundaries of what their hardware can do.