The VRAM Accounting That Almost Broke Phase 9: A Deep Dive into GPU Memory Lifetimes in the cuzk Proving Engine

Introduction

In the high-stakes world of Filecoin proof generation, every millisecond of GPU time counts, and every byte of VRAM is precious. The cuzk SNARK proving engine, a custom-built Groth16 prover for Filecoin's Proof-of-Replication (PoRep) protocol, operates on a 16 GiB GPU with a pipeline that routinely allocates 12 GiB or more for a single partition's polynomial data. When the assistant set out to implement Phase 9 — a PCIe transfer optimization designed to eliminate GPU idle gaps by pre-staging polynomial uploads — it encountered a cascade of out-of-memory (OOM) failures that exposed a fundamental tension between performance optimization and memory safety in concurrent GPU programming.

The subject message, message 2431 in the conversation, captures a pivotal moment of realization. The assistant has just moved the pre-staging allocation inside the GPU mutex (in message 2430) to serialize it across workers, but it now recognizes that this fix alone is insufficient. A deeper, subtler problem lurks in the lifetime of a single 8 GiB buffer called d_bc — a buffer that, if freed too late, silently starves the next worker of VRAM and triggers cascading OOM panics. This message is a masterclass in GPU memory accounting under concurrency, and it reveals how the smallest detail in buffer lifetime management can determine whether a performance optimization succeeds or fails.

The Context: Phase 9's Two-Pronged Optimization

To understand why this message matters, we must first understand what Phase 9 was trying to accomplish. The Phase 8 baseline (established in segment 24) had achieved a 13-17% throughput improvement through a dual-GPU-worker interlock, but TIMELINE analysis (segment 25) revealed two root causes of GPU idle gaps: non-pinned host memory causing slow PCIe transfers, and Pippenger MSM sync stalls. Phase 9 targeted both.

Change 1 (Tier 1) was the PCIe transfer optimization proper. The original code uploaded three large polynomials — a, b, and c — from host memory to GPU device memory using synchronous cudaMemcpy calls inside the GPU mutex. This meant the worker holding the mutex spent tens of milliseconds just waiting for data to traverse the PCIe bus, while the other worker waited idle for the mutex. Phase 9 changed this by: (1) pinning the host memory pages with cudaHostRegister to enable full PCIe bandwidth, (2) allocating device buffers outside the mutex, and (3) issuing asynchronous cudaMemcpyAsync transfers on a dedicated CUDA stream with event-based synchronization. The idea was that the uploads could overlap with compute — the NTT (Number Theoretic Transform) kernels could start as soon as each polynomial's transfer completed, signaled by CUDA events.

Change 2 (Tier 3) targeted the Pippenger MSM (Multi-Scalar Multiplication) sync stalls. The original code issued a hard sync() call after each batch of MSM operations, forcing the GPU to drain its pipeline before the host could read results. By introducing double-buffered host result buffers (res_buf[2], ones_buf[2]) and deferring the sync() call to the next iteration, the assistant allowed GPU compute to overlap with device-to-host (DtoH) transfers, eliminating the stall.

The initial implementation, however, crashed immediately with OOM errors when run with the intended dual-worker configuration (gpu_workers_per_device=2). The assistant's debugging in messages 2418-2429 revealed the first-order problem: both workers tried to pre-stage their 12 GiB allocations simultaneously, exceeding the 16 GiB VRAM. The fix in message 2430 was to move the pre-staging allocation inside the GPU mutex, serializing it so only one worker could allocate at a time.

The Subject Message: A Deeper VRAM Crisis Revealed

Message 2431 begins with the assistant realizing that the mutex-serialization fix, while necessary, is not sufficient. The message opens with a stark admission:

But wait — there's still a VRAM issue. d_a (4 GiB) + d_bc (8 GiB) = 12 GiB for pre-staging alone. Then the NTT phase completes, d_bc is no longer needed, but batch_add allocates its own dev_ptr_t for working memory. With 12 GiB already allocated (d_a still alive during H MSM), plus MSM working memory, we might be tight.

This sentence encapsulates the entire dilemma. The assistant has already fixed the allocation serialization problem, but it now realizes that the deallocation timing is equally critical. The d_bc buffer — 8 GiB of device memory holding the b and c polynomials — is allocated during pre-staging and, in the current code, freed only after the GPU mutex is released. This means that when worker 1 acquires the mutex and begins its own pre-staging, worker 0's d_bc is still occupying 8 GiB of VRAM. Worker 1 then tries to allocate its own 12 GiB on top of that, and the GPU — with only 16 GiB total — runs out of memory.

The Step-by-Step VRAM Accounting

What follows in the message is a meticulous, line-by-line trace of the GPU memory timeline. The assistant enumerates six steps of the code flow, accounting for every allocation and deallocation:

Step 1: Pre-stage allocation. d_a (4 GiB) and d_bc (8 GiB) are allocated. Total: 12 GiB.

Step 2: Enter the per-GPU thread. The d_a pointer is wrapped in a gpu_ptr_t<fr_t> RAII wrapper. The function execute_ntt_msm_h_prestaged is called, which uses both buffers for the NTT phase. During this phase, all 12 GiB are actively in use. After the NTT and coefficient-wise multiplication complete, the data in d_bc is consumed — but critically, the buffer is not freed. It remains allocated because the cleanup code is outside the per-GPU thread scope, after the mutex release.

The assistant then notes the H MSM phase: the msm_t constructor allocates MSM bucket memory (~0.5 GiB). At this point, d_a (4 GiB) is still needed as scalars, and d_bc (8 GiB) is still allocated. Total: 12.5 GiB. It fits, but barely.

Step 3: d_a freed. After execute_ntt_msm_h_prestaged returns, the gpu_ptr_t destructor frees d_a (4 GiB). But d_bc (8 GiB) remains allocated.

Step 4: Barrier wait. The CPU preprocessing completes. No GPU memory change.

Step 5: batch_add. Four calls to batch_add allocate dev_ptr_t objects of ~0.6 GiB each. With d_bc still occupying 8 GiB: 8 + 0.6 = 8.6 GiB. Should fit.

Step 6: Tail MSMs. MSM objects are allocated. 8 GiB + MSM working memory. "Might be tight," the assistant notes.

The key insight emerges: "the cleanup of d_bc happens BEFORE unlock now (with my latest fix). Let me re-examine... Actually wait, d_bc needs to be freed MUCH earlier — as soon as the NTT phase completes (inside the per_gpu thread)."

This is the critical realization. The assistant had assumed that moving the d_bc cleanup before the mutex unlock (as done in message 2430) would be sufficient. But the timeline analysis shows that even that is not enough. The d_bc buffer must be freed inside the per-GPU thread, immediately after the NTT phase completes, because the next worker cannot begin its pre-staging until the mutex is released, and the mutex cannot be released until the current worker's GPU work is done. If d_bc lives until mutex release, it overlaps with the next worker's allocation.

The Root Cause: A Lifetime Mismatch

The fundamental issue is a mismatch between the logical lifetime of d_bc and its actual lifetime in the code. Logically, d_bc is only needed during the NTT phase — once the coefficient-wise multiplication kernel (sub_mult_with_constant) finishes, the b and c polynomials have been transformed and their original data is no longer needed. But the code was structured to free d_bc at the end of the entire GPU operation, after the mutex release, because the cleanup was grouped with other post-mutex operations like host page unregistration.

In the original (non-pre-staged) code path, this wasn't a problem because d_bc was allocated as a local dev_ptr_t<fr_t> inside execute_ntt_msm_h, and C++ RAII freed it automatically when that function returned. But in the pre-staged path, d_bc is allocated before the per-GPU thread starts, and its lifetime is managed manually. The assistant had inadvertently extended its lifetime far beyond its logical necessity.

The fix the assistant identifies is to free d_bc "right after the NTT phase inside execute_ntt_msm_h_prestaged, or immediately after returning from it." This is the correct approach: the buffer should be freed as soon as its data is consumed, regardless of where in the pipeline that occurs. In the original code, this happened naturally through scope-based lifetime management. In the pre-staged code, it requires explicit intervention.

Broader Implications for GPU Memory Management

This episode illustrates several important principles for high-performance GPU programming, particularly in concurrent pipelines:

1. Memory accounting must be dynamic, not static. It's not enough to know that a single worker's peak memory fits within VRAM. You must account for the overlapping lifetimes of buffers across workers, including buffers that are "in flight" — allocated by one worker but not yet freed when the next worker begins.

2. RAII is not a panacea under manual lifetime extension. The original code's use of dev_ptr_t and gpu_ptr_t RAII wrappers provided correct lifetime management through scope. But when the assistant introduced pre-staging, it broke the scope-based lifetime by allocating buffers in an outer scope and passing them into the inner scope. The RAII destructor for d_a still fired correctly (when the gpu_ptr_t wrapper went out of scope inside the per-GPU thread), but d_bc had no such wrapper — it was a raw cudaMalloc pointer whose lifetime was managed by explicit cudaFree calls placed after the mutex release.

3. The mutex is a memory-ordering construct, not just a concurrency construct. The assistant initially thought of the GPU mutex purely as a mechanism to serialize kernel execution. But the mutex also implicitly governs memory lifetimes: any buffer that lives until mutex release will overlap with the next worker's allocations. The fix required recognizing that d_bc's lifetime should be bounded by the NTT phase, not by the mutex scope.

4. Performance optimization often reveals hidden lifetime dependencies. The pre-staging optimization was designed to overlap PCIe transfers with compute. But in doing so, it changed the allocation timing and lifetime of the buffers, which in turn affected the memory pressure seen by other workers. Every optimization that changes when memory is allocated or freed has ripple effects on the entire concurrent pipeline.

The Thinking Process: A Model of Systematic Debugging

What makes this message particularly valuable is the transparency of the assistant's reasoning. It doesn't just state the conclusion — it walks through the entire code flow step by step, accounting for every GiB at every phase. This is the kind of systematic thinking that separates effective optimization from guesswork.

The assistant starts with a hypothesis ("there's still a VRAM issue"), then tests it by enumerating the allocation timeline. It identifies the specific point where d_bc should be freed but isn't. It cross-references with the original code path to see how the original handled the same lifetime. And it arrives at a concrete, actionable fix: restructure the pre-staged path to free d_bc inside the per-GPU thread, immediately after the NTT phase.

This thinking also demonstrates the importance of understanding the semantics of each buffer. d_bc is not just "some GPU memory" — it specifically holds the b and c polynomials after NTT transformation. Knowing what the buffer contains, and when that data is no longer needed, is essential to determining the correct lifetime. A developer who treats all buffers as interchangeable blobs of VRAM would never arrive at this insight.

Conclusion

Message 2431 is a turning point in the Phase 9 implementation. The assistant has identified that the mutex-serialization fix from message 2430, while necessary, is incomplete. The real solution requires freeing d_bc inside the per-GPU thread, immediately after the NTT phase, rather than waiting until after the mutex release. This insight bridges the gap between a working single-worker implementation and a correct dual-worker implementation.

The broader lesson is that GPU memory management under concurrency is a holistic problem. You cannot optimize allocation without also optimizing deallocation. You cannot change the timing of one buffer without considering the ripple effects on all other buffers. And you cannot rely on RAII alone when you manually extend buffer lifetimes across scope boundaries.

In the subsequent messages (chunk 1 of segment 26), the assistant would implement this fix and successfully run the full production benchmark, achieving a system throughput of 41.0 seconds per proof. But the foundation for that success was laid in this single message of careful VRAM accounting — a reminder that in GPU programming, the smallest detail can make the difference between a crash and a breakthrough.