The Persistence of Memory: Debugging a CUDA OOM Failure in the Phase 9 PCIe Optimization

In the high-stakes world of GPU-accelerated zero-knowledge proof generation, every millisecond counts. The cuzk SNARK proving engine had already been through eight optimization phases, each peeling back another layer of bottleneck to reveal the next. Phase 9 targeted a particularly stubborn source of GPU idle time: the 6 GiB of host-to-device (HtoD) transfers for the a/b/c polynomial vectors that preceded every Groth16 proof partition. By pinning host memory with cudaHostRegister and issuing asynchronous cudaMemcpyAsync transfers on a dedicated stream with event-based synchronization, the optimization aimed to overlap data movement with computation and eliminate the synchronous transfer stall that had been visible in GPU timeline traces.

The single-worker benchmarks were spectacular. NTT+MSM time dropped from ~2430 ms to ~690 ms — a 71.6% reduction. Overall GPU time per partition fell by 50–61%. System throughput improved from 37.4 seconds per proof to 32.1 seconds — a 14.2% gain. But the real target was the production configuration: two GPU workers per device, running concurrently to saturate the GPU's compute and copy engines. And that configuration was failing catastrophically.

The Subject Message: A Debugging Crossroads

Message 2426 captures a pivotal moment in the debugging process. After implementing what appeared to be a straightforward fix for the dual-worker OOM failure, the assistant restarts the daemon, launches a five-proof benchmark with concurrency=3 and gpu_workers_per_device=2, and watches all five proofs fail. The message is deceptively brief — a single bash command and its output — but it carries enormous diagnostic weight:

[assistant] Still failing. Let me check the new logs:
[bash] grep -i "error\|FAIL\|panic\|prestage\|CUZK_TIMING" /tmp/cuzk-phase9-daemon.log | grep -v "skipping\|PCE file" | tail -30
CUZK_TIMING: prep_msm_ms=2782
CUZK_TIMING: prestage_setup=fallback err=2
CUZK_TIMING: b_g2_msm_ms=1790 num_circuits=1
CUZK_TIMING: prep_msm_ms=3472
thread 'tokio-runtime-worker' panicked at /home/theuser/curio/extern/supraseal-c2/src/lib.rs:153:9:
cudaMalloc(&d_ptr, n * sizeof(T))@sppark-0.1.14/sppark/util/gpu_t.cuh:331 failed: "out of memory"
thread 'tokio-runtime-worker' panicked at /home/theuser/curio/extern/supraseal-c2/src/lib.rs:153:9:
cudaMalloc(&d_ptr, n * sizeof(T))@sppark-0.1.14/sppark...

The two-word opening — "Still failing" — is a quiet acknowledgment that the previous fix (moving the pre-staging cudaMalloc calls inside the GPU mutex to serialize allocations across workers) had not resolved the problem. The assistant turns immediately to the logs, deploying a carefully crafted grep pipeline designed to extract the signal from the noise. The pattern excludes "skipping" and "PCE file" messages to focus on the critical error indicators.

Reading the Diagnostic Output

The log output tells a grim story. The CUZK_TIMING: prestage_setup=fallback err=2 line is the most revealing. The prestage_setup timing label was added specifically for Phase 9 to track whether the new pre-staging path was being used or falling back. The err=2 value corresponds to cudaErrorMemoryAllocation — CUDA's out-of-memory error. This means the pre-staging code path attempted to allocate device memory for the a/b/c polynomial buffers and failed, then fell back to the original synchronous transfer path.

But the fallback path also fails. The panic at lib.rs:153 shows cudaMalloc failing with "out of memory" inside the sppark GPU utility library. This is the original code path — the one that was supposed to work even without pre-staging — also failing because the GPU's 16 GiB of VRAM is exhausted. The CUZK_TIMING: prep_msm_ms=2782 and prep_msm_ms=3472 lines show that some MSM preprocessing work completed before the OOM crash, but the overall pipeline cannot sustain two concurrent workers.

The presence of two panic messages from tokio-runtime-worker confirms that both GPU workers are hitting the OOM condition, not just one. This is a systemic failure, not an isolated race condition.

Why the First Fix Was Insufficient

The assistant's initial diagnosis in message 2419 had correctly identified that with gpu_workers_per_device=2, both workers tried to pre-stage simultaneously, each allocating 6 GiB for a total of 12 GiB — well within the 16 GiB VRAM budget, but leaving no room for the working memory needed during kernel execution. The fix moved the pre-staging allocation inside the GPU mutex, so that only one worker at a time could allocate its pre-stage buffers.

But this fix overlooked a critical detail: the pre-stage buffers persist for the entire duration of the worker's GPU time slice. When Worker A acquires the mutex, it allocates 6 GiB for d_a and d_bc, issues the async upload, and then runs its NTT and MSM kernels. During this time, Worker B is waiting for the mutex — but no memory is being freed. When Worker A finishes and releases the mutex, Worker B acquires it and tries to allocate its own 6 GiB of pre-stage buffers. But Worker A's buffers are still allocated (they're being used by the kernels that haven't completed yet, or they haven't been explicitly freed). The result: Worker A's 6 GiB + Worker B's 6 GiB + the kernel working memory (~7.5 GiB) = approximately 19.5 GiB, far exceeding the 16 GiB limit.

The fundamental issue is that the pre-stage buffers are allocated for the duration of the GPU work, not just for the upload phase. Simply serializing the allocation is not enough — the buffers must also be freed before the next worker's allocation, or the memory must be managed more carefully.

Assumptions and Their Consequences

The debugging process reveals several assumptions that proved incorrect:

Assumption 1: Serializing allocation is sufficient. The assistant assumed that by moving the cudaMalloc calls inside the mutex, only one worker would allocate at a time, and the total memory would stay within bounds. This overlooked the cumulative effect of persistent allocations across workers.

Assumption 2: The fallback path would work if pre-staging failed. The prestage_setup=fallback mechanism was designed as a graceful degradation — if the pre-staging allocation fails, fall back to the original synchronous transfer. But the fallback path itself requires device memory for the same buffers, and if the GPU is already OOM, the fallback will also fail.

Assumption 3: CUDA memory management is straightforward. The assistant had previously noted that CUDA's cudaMallocAsync/cudaFreeAsync memory pools do not release freed memory back to the synchronous cudaMalloc pool. This means that even if some memory is freed via the async allocator, it may not be available for subsequent synchronous allocations — a subtle but critical detail in a mixed-allocation environment.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

Output Knowledge Created

This message produces several critical pieces of knowledge:

  1. Evidence that the first fix is insufficient: The prestage_setup=fallback err=2 output proves that the mutex-based serialization alone does not prevent OOM.
  2. Confirmation of dual-worker failure: Both workers panic with OOM, not just one.
  3. Diagnostic narrowing: The problem is not in the pre-staging code path specifically — the fallback path also fails, indicating a deeper memory exhaustion issue.
  4. Direction for the next fix: The assistant must now address the persistence of pre-stage buffers across workers. The solution will involve freeing d_bc immediately after the NTT phase (before the mutex is released) and adding a memory-aware allocator that queries available VRAM before allocating.

The Thinking Process

The assistant's reasoning in this message is visible in the choice of grep command. The pattern error\|FAIL\|panic\|prestage\|CUZK_TIMING is carefully selected to capture all relevant diagnostic signals. The exclusion of skipping\|PCE file filters out routine log messages that would obscure the signal. The tail -30 ensures only the most recent log entries are examined — the assistant knows the failures happened recently and wants the freshest data.

The two-word preamble "Still failing" reveals the assistant's mental model: the fix was expected to work, and its failure is surprising. The assistant does not immediately jump to a new theory but instead gathers data, letting the logs speak first. This is a disciplined debugging approach — resist the urge to speculate, and let empirical evidence guide the next step.

The log output confirms the worst: the pre-staging path is falling back (indicating the optimization is not being applied), and even the fallback is crashing. The assistant now has a clear signal that the memory management strategy needs fundamental rethinking, not just a mutex rearrangement.

Conclusion

Message 2426 captures a moment of diagnostic clarity in a complex optimization effort. The Phase 9 PCIe transfer optimization had delivered impressive single-worker gains, but the dual-worker configuration exposed a subtle memory management flaw that a simple mutex-based fix could not address. The prestage_setup=fallback err=2 log line became the Rosetta Stone for understanding the failure — it told the assistant that the pre-staging path was not even being used, and the GPU was so memory-constrained that even the fallback could not survive.

The debugging journey documented in this message would ultimately lead to a three-part fix: serializing the pre-stage allocation under the mutex, freeing d_bc immediately after the NTT phase to reclaim memory before the mutex is released, and adding a memory-aware allocator that queries cudaMemGetInfo with a safety margin before attempting allocation. These changes would finally stabilize the dual-worker configuration and enable the full production benchmark that followed in the next chunk.

But in this moment — message 2426 — none of that is known yet. The assistant is staring at a log file, watching two GPU workers crash with out-of-memory errors, and beginning the process of understanding why a seemingly correct fix failed. It is a small but essential step in the iterative process of optimization: measure, diagnose, fix, and when the fix doesn't work, measure again.