Threading the Needle: How a Single Edit in Bellperson Fixed Multi-GPU Proving

Message 476 in this opencode session is deceptively brief:

[edit] /tmp/czk/extern/bellperson/src/groth16/prover/supraseal.rs Edit applied successfully.

On its surface, this is nothing more than a tool-call confirmation — the assistant's edit tool reported that a file was modified without error. But this single line is the keystone of a multi-layered architectural fix that spanned five files across three languages (C++, Rust FFI, and Rust application code). Understanding why this particular edit was necessary, what it changed, and what assumptions it encoded reveals the full arc of a debugging journey that began with a mysterious out-of-memory crash on a 20 GB RTX 4000 Ada GPU and ended with a properly load-balanced multi-GPU proving pipeline.

The Context: From Lazy Hack to Proper Architecture

The story leading to message 476 begins with a partitioned proof workload for SnapDeals — 16 identical partitions that should have been distributed evenly across two GPUs. Instead, the system crashed with an OOM error. The initial diagnosis revealed a deeper problem: the C++ GPU proving code in groth16_cuda.cu always routed single-circuit proofs to GPU 0, regardless of which Rust worker submitted them. The line n_gpus = min(ngpus(), num_circuits) at line 483 of the CUDA file meant that whenever num_circuits=1 (as it is for partitioned proofs), only one GPU thread was spawned, and it always called select_gpu(0).

The first attempted fix was a shared mutex — a single std::mutex that serialized all partition proofs onto GPU 0. The user immediately recognized this as a "lazy hack" ([msg 444]), pointing out that CuZK was "meant to be a fairly sophisticated proving engine" that should support multiple GPUs with proper interlocking of data transfer and compute phases. The shared mutex wasted the second GPU entirely and, worse, still allowed OOMs because two workers could enter the GPU code simultaneously on the same device, exceeding the 20 GB VRAM budget for a single SnapDeals partition.

The proper solution was architectural: thread a gpu_index parameter through the entire call chain so that the C++ code uses the GPU assigned by the Rust engine instead of always defaulting to GPU 0. This required coordinated changes across five layers:

  1. C++ (groth16_cuda.cu): Add int gpu_index parameter; when gpu_index >= 0, force n_gpus = 1 and call select_gpu(gpu_index) instead of select_gpu(0).
  2. Rust FFI (supraseal-c2/src/lib.rs): Update the extern "C" declarations and the public wrapper functions (start_groth16_proof, generate_groth16_proof, generate_groth16_proofs) to accept and forward the new parameter.
  3. Bellperson (supraseal.rs): Update prove_start and prove_from_assignments to accept gpu_index and pass it through.
  4. Pipeline (pipeline.rs): Update gpu_prove and gpu_prove_start to thread the parameter.
  5. Engine (engine.rs): Revert the shared mutex hack, restore per-GPU mutexes, and pass each worker's assigned gpu_ordinal as the gpu_index. Message 476 represents step 3 — the edit to bellperson's prove_start function. It is the moment when the GPU index parameter crossed the boundary from the low-level FFI layer into the application-level proving abstraction.## Why Bellperson Was the Critical Layer The bellperson library sits at a crucial architectural junction. It is a fork of the Bellman zk-SNARK proving library, customized for Filecoin's proof requirements. The prove_start function in supraseal.rs is the entry point that the pipeline layer calls when it wants to initiate GPU proving. It receives the synthesized proof assignments, allocates the necessary GPU resources, calls the FFI to the C++ CUDA code, and returns a pending proof handle that can be finalized later. Before message 476, prove_start accepted a gpu_mtx: GpuMutexPtr parameter — a pointer to a C++ std::mutex — but had no concept of which GPU to use. The mutex was the only synchronization mechanism, and the GPU selection was entirely determined by the C++ code's internal logic (min(ngpus(), num_circuits)). This meant that even though the Rust engine carefully assigned each worker to a specific GPU (via CUDA_VISIBLE_DEVICES environment variable and ordinal tracking), that assignment was silently ignored by the proving code. The edit in message 476 added a gpu_index: i32 parameter to prove_start (and correspondingly to prove_from_assignments), with the convention that -1 means "auto" (use the old multi-GPU logic for multi-circuit proofs) and any non-negative value means "use this specific GPU." This simple integer parameter was the mechanism by which the Rust engine's worker assignment could finally reach the CUDA code that actually selects the GPU device.

The Design Decision: -1 as Auto

A subtle but important design choice is encoded in the gpu_index parameter's semantics. The value -1 signals "auto mode" — the C++ code falls back to its original behavior of spawning one GPU thread per circuit up to the number of available GPUs. This preserves backward compatibility for callers that don't have a specific GPU assignment, such as the older generate_groth16_proofs function used by non-engine paths.

For the engine's GPU workers, however, each worker passes its assigned GPU ordinal (0, 1, etc.). When gpu_index >= 0, the C++ code forces n_gpus = 1 and calls select_gpu(gpu_index) instead of select_gpu(tid). This means a worker assigned to GPU 1 will actually prove on GPU 1, using GPU 1's per-GPU mutex, GPU 1's memory pool, and GPU 1's d_a_cache.

This design assumes that the Rust engine's GPU assignment is authoritative and correct — an assumption validated by the earlier debugging that showed the engine was already setting CUDA_VISIBLE_DEVICES and tracking ordinals correctly, but the C++ code was overriding that choice. The fix essentially aligns the C++ behavior with the Rust engine's intent.

What Message 476 Doesn't Show

The edit tool confirmation in message 476 is opaque — it doesn't reveal what lines were changed. To understand the actual diff, one must look at the surrounding messages. In [msg 475], the assistant reads prove_start's signature and sees it lacks the gpu_index parameter. In [msg 477], the assistant confirms the edit was applied. The actual changes involved:

  1. Adding gpu_index: i32 to the prove_start function signature.
  2. Passing gpu_index through to the FFI call to start_groth16_proof.
  3. Similarly updating prove_from_assignments (the synchronous wrapper) in <msg id=480-481>. These changes are mechanical but critical — they are the conduit through which the GPU assignment flows from the Rust engine's worker scheduler, through the bellperson abstraction, across the FFI boundary, and finally into the CUDA kernel selection logic.## The Thinking Process: Bottom-Up Construction The assistant's approach to implementing this fix reveals a deliberate methodology. Rather than making all changes at once, the assistant worked bottom-up through the call stack: starting with the C++ CUDA code (<msg id=457-461>), then the Rust FFI layer (<msg id=466-473>), then bellperson ([msg 476]), then the pipeline (<msg id=485-486>), and finally the engine ([msg 489]). This ordering ensured that each layer's changes could be verified against the layer below before proceeding upward. The todo list maintained throughout this process ([msg 452], [msg 453], [msg 462], [msg 483], [msg 487]) shows the assistant tracking progress across five work items, marking each as "pending," "in_progress," or "completed" as the implementation advanced. This structured approach was essential for a change that touched five files across three languages and required consistent parameter naming and semantics at every level.

Assumptions and Potential Pitfalls

The implementation makes several assumptions worth examining:

  1. The d_a_cache global singleton is safe for concurrent GPU access. The assistant noted in [msg 456] that d_a_cache has a gpu_id field and handles switching between GPUs by freeing on the old GPU and reallocating on the new one. However, with concurrent workers on different GPUs, this global cache could thrash — two workers might repeatedly free and reallocate the same cache entry. The assistant acknowledged this ("I need to make it a per-GPU array") but deferred that fix, likely because the immediate OOM crash was the higher priority.
  2. The -1 sentinel value is unambiguous. Using -1 as "auto" relies on the convention that valid GPU ordinals are non-negative. This is safe for CUDA (device ordinals start at 0) but means the parameter type i32 (signed) is necessary rather than u32.
  3. The engine's GPU assignment is correct. The fix trusts that the Rust engine's worker scheduler assigns GPUs correctly and that the gpu_ordinal value passed through the chain matches the actual GPU device intended. Any mismatch between the ordinal and the CUDA_VISIBLE_DEVICES mapping would still cause problems, but that would be a bug in the engine layer, not in the proving path.

Conclusion

Message 476 is a single edit confirmation line that represents the culmination of a deep architectural debugging session. What began as an OOM crash on a SnapDeals workload was traced to a fundamental mismatch between the Rust engine's multi-GPU scheduling and the C++ proving code's single-GPU default. The shared mutex hack was a temporary bandage that the user correctly identified as insufficient. The proper fix — threading a gpu_index parameter through every layer of the proving stack — required coordinated changes across C++, Rust FFI, bellperson, pipeline, and engine code.

The bellperson edit in message 476 is the connective tissue that makes the entire fix work. Without it, the gpu_index parameter would reach the FFI boundary but never propagate to the application-level proving API that the pipeline and engine actually call. It is a small edit with outsized architectural significance — a reminder that in complex systems, the most critical changes are often the ones that bridge layers, not the ones that rewrite them.