Threading GPU Awareness Through the Stack: A Proper Fix for Multi-GPU Proving

Introduction

In the course of debugging a SnapDeals Out-of-Memory (OOM) crash on a dual-GPU proving host, a pivotal realization emerged: the existing "fix" of a shared mutex was a lazy hack that serialized all work onto GPU 0, wasting the second GPU entirely. Message 455 of the opencode session captures the moment when the assistant pivots from band-aid to architecture, planning the proper solution of threading a gpu_index parameter through the entire call chain from Rust engine to C++ CUDA kernel. This message is not a code edit itself, but a reasoning and planning message—a design deliberation that sets the stage for a multi-file refactor spanning C++, Rust FFI, bellperson prover functions, pipeline abstractions, and the engine's GPU worker code.

Context: How We Got Here

The conversation leading up to message 455 had been a deep debugging session. The team was working on CuZK, a GPU-accelerated proving engine for Filecoin proofs (PoRep, WindowPoSt, WinningPoSt, SnapDeals). A SnapDeals workload of 16 identical partitions was OOMing on a host with an RTX 4000 Ada (20 GB VRAM) despite each individual partition fitting comfortably within the memory budget. The logs showed a smoking gun: two workers entered the GPU proving code at the same millisecond, both attempting to allocate VRAM on GPU 0 simultaneously. One succeeded (allocating 4096 MiB for the d_a_cache), and the second crashed with cudaMallocAsync: out of memory.

The initial diagnosis had identified that the C++ GPU proving code always routes single-circuit proofs to GPU 0, regardless of which Rust worker submits them. This is because the C++ function generate_groth16_proofs_start_c computes n_gpus = min(ngpus(), num_circuits), and for partitioned proofs where num_circuits = 1, this yields n_gpus = 1, which means only one GPU thread is spawned with tid = 0, which calls select_gpu(0). Every worker, regardless of which GPU the Rust engine assigned it to, ends up proving on GPU 0.

The initial "fix" was a shared mutex in engine.rs that serialized all partition proofs onto GPU 0, preventing the OOM but at the cost of wasting the second GPU entirely. The user rightly called this out in message 444: "Isn't the shared lock just a lazy hack?" The assistant agreed, and message 455 is the beginning of the proper fix.

What Message 455 Actually Says

The message is brief but dense with intent. It opens with the assistant's reasoning:

Now let me see the full picture. The key change needed in C++ is: 1. Add int gpu_index parameter (use -1 to mean "auto / use all GPUs like before") 2. When gpu_index >= 0, force n_gpus = 1 and use select_gpu(gpu_index) instead of select_gpu(tid) where tid=0

This is the core design decision. The assistant then performs a read of the C++ source file to check the function signature of generate_groth16_proofs_start_c and its sync wrapper generate_groth16_proof_c, confirming the existing parameter list before planning the changes.

The quoted code shows the existing signature:

extern "C"
RustError::by_value generate_groth16_proofs_start_c(
    const Assignment<fr_t> provers[],
    size_t num_circuits,
    const fr_t r_s[], con...

The message ends mid-read, having gathered the information needed to proceed with implementation.

The Reasoning and Design Decisions

The assistant's reasoning reveals several key insights:

1. The root cause is architectural, not a simple race condition. The C++ code was designed with the assumption that multi-GPU work would always involve multiple circuits (e.g., num_circuits >= ngpus()). For partitioned proofs where each partition is a single circuit, this assumption breaks down. The code needs to be told which GPU to use explicitly.

2. The fix must be backward-compatible. The -1 sentinel value for gpu_index preserves the existing behavior for callers that don't know about GPU assignment (e.g., non-engine paths like tests or standalone tools). When gpu_index = -1, the C++ code falls back to the original min(ngpus(), num_circuits) logic.

3. The fix requires changes across the entire call stack. The assistant's todo list (from message 452) shows the full scope:

• C++ groth16_cuda.cu: add gpu_index parameter
• Rust FFI in supraseal-c2/src/lib.rs: update wrapper functions
• Bellperson prover functions (prove_start, prove_from_assignments)
• Pipeline layer (gpu_prove, gpu_prove_start)
• Engine's GPU worker code in engine.rs: revert the shared mutex hack 4. The per-GPU mutex design must be preserved. The original code had per-GPU mutexes allocated via alloc_gpu_mutex(). The shared mutex hack replaced these with a single mutex, serializing everything to GPU 0. The proper fix restores per-GPU mutexes and ensures each worker locks only the mutex for its assigned GPU.

Assumptions Made

The assistant makes several assumptions in this message:

Assumption 1: The C++ code's select_gpu() function accepts a GPU index and works correctly. The assistant assumes that calling select_gpu(gpu_index) with an explicit index will properly set the CUDA device for the current thread. This is a reasonable assumption given the existing code already calls select_gpu(tid) in the multi-circuit path.

Assumption 2: The Rust-side GPU assignment logic (which determines which worker goes to which GPU) is correct. The engine's GPU worker code already assigns workers to specific GPUs via CUDA_VISIBLE_DEVICES and per-GPU mutexes. The fix assumes this assignment logic is sound and only needs the C++ code to honor it.

Assumption 3: The -1 sentinel is a safe default. Non-engine callers (tests, benchmarks, standalone tools) that don't pass a GPU index will get the original behavior. This assumes those callers don't need multi-GPU awareness.

Assumption 4: The d_a_cache global variable won't cause cross-GPU issues. The assistant noted earlier that g_d_a_cache is a single global cache that can cause problems when workers alternate between GPUs. The fix doesn't address this directly, assuming it's a secondary concern.

Potential Mistakes and Incorrect Assumptions

The -1 sentinel approach, while pragmatic, introduces a subtle dual-path problem. The C++ code now has two ways to determine GPU selection: the old min(ngpus(), num_circuits) logic and the new explicit gpu_index parameter. If a caller accidentally passes -1 when it should pass an explicit index, the code silently falls back to GPU 0, potentially reintroducing the original bug. This is a design smell—it would be cleaner to require explicit GPU assignment for all engine paths and only keep the auto path for legacy callers.

The fix doesn't address the global d_a_cache issue. The g_d_a_cache static variable is shared across all GPU threads. When worker 0 finishes on GPU 0 and worker 1 starts on GPU 1, the cache may still hold a GPU 0 pointer, causing a cudaFree on the wrong device or forcing a reallocation. The assistant acknowledged this earlier (message 449: "bad for multi-GPU") but chose not to address it in this fix. This could cause performance degradation or subtle crashes under heavy load.

The assumption that the Rust-side GPU assignment is correct may be optimistic. The engine's worker code sets CUDA_VISIBLE_DEVICES before spawning workers, which masks GPU indices from the C++ code. If the C++ code uses select_gpu(0) after CUDA_VISIBLE_DEVICES has been set to "1", it will actually use the physical GPU 1 (which appears as device 0 in the masked view). The fix must ensure that gpu_index refers to the physical GPU index, not the masked one, or the masking must be removed. The assistant's later implementation (in subsequent messages) addresses this by passing the physical GPU ordinal.

Input Knowledge Required

To understand message 455, the reader needs:

1. Understanding of the CuZK proving architecture: The engine uses GPU workers that each handle a subset of partition proofs. Workers are assigned to specific GPUs via CUDA_VISIBLE_DEVICES and per-GPU mutexes.
2. Knowledge of the partitioned proof pipeline: PoRep and SnapDeals proofs are split into partitions, each proved independently. The C++ code's num_circuits parameter controls how many proofs are processed in a single call.
3. Familiarity with the C++/Rust FFI boundary: The extern "C" functions in groth16_cuda.cu are called from Rust via FFI wrappers in supraseal-c2/src/lib.rs. Adding a parameter requires changes on both sides.
4. Understanding of the shared mutex hack: The previous fix (message 433-443) replaced per-GPU mutexes with a single shared mutex, serializing all GPU work onto GPU 0.
5. Knowledge of CUDA device management: The select_gpu() function calls cudaSetDevice() to select the active GPU for the current thread. The gpu_index parameter directly controls which device is selected.

Output Knowledge Created

Message 455 produces several valuable outputs:

1. A clear design specification: The two-point plan (add gpu_index parameter, force single-GPU mode when specified) serves as the blueprint for the implementation that follows.
2. Confirmation of the function signature: By reading the C++ code, the assistant confirms the exact parameter list of generate_groth16_proofs_start_c, enabling precise edits.
3. A decision to revert the shared mutex hack: The assistant commits to restoring per-GPU mutexes, which is essential for multi-GPU load balancing.
4. Scope definition: The todo list from message 452 is implicitly confirmed, showing the full set of files that need modification.

The Thinking Process

The assistant's thinking in message 455 is a model of systematic debugging. It starts with the high-level design ("the key change needed in C++ is..."), then immediately grounds it in concrete code by reading the actual function signature. This pattern—alternating between abstract reasoning and concrete verification—characterizes the entire session.

The assistant resists the temptation to jump into coding. Instead, it pauses to "see the full picture," ensuring the design is sound before touching any files. This is particularly important given the scope of the change: threading a parameter through five layers of abstraction (C++ → Rust FFI → bellperson → pipeline → engine) requires careful planning to avoid inconsistencies.

The choice of -1 as a sentinel is telling. It reveals a design philosophy of backward compatibility: existing callers that don't know about GPU assignment continue to work unchanged. This is pragmatic but also reveals an assumption that the non-engine paths don't need multi-GPU support—an assumption that may need revisiting if those paths are ever used in multi-GPU environments.

Conclusion

Message 455 is a quiet but pivotal moment in the conversation. It's not a flashy code change or a dramatic breakthrough, but a moment of deliberate design—the assistant stepping back from the immediate pressure of a crashing proof pipeline to think about the right architecture. The shared mutex hack could have been left in place, and the SnapDeals workload would have worked (slowly, on one GPU). But the user's challenge—"Isn't the shared lock just a lazy hack?"—prompted a deeper analysis and a proper fix.

The message demonstrates a key skill in systems programming: knowing when to stop fixing symptoms and start fixing causes. The symptom was an OOM crash; the cause was GPU routing logic that didn't respect the Rust engine's assignment. The shared mutex addressed the symptom (only one worker in the GPU at a time) but not the cause (workers always targeting GPU 0). The gpu_index parameter addresses the cause, enabling the system to use both GPUs as intended.

In the broader arc of the conversation, message 455 is the turning point from reactive debugging to proactive architecture. The subsequent messages implement the changes planned here, deploy them to the test hosts, and verify that both GPUs are now load-balanced correctly. The SnapDeals workload that previously OOM'd now runs smoothly across both GPUs, and the shared mutex hack is consigned to the git history.