The Pivot: From Shared Mutex Hack to Proper Multi-GPU Architecture in CuZK

Introduction

In the course of debugging a production crash in a zero-knowledge proving system, a single message from an AI assistant marks a critical turning point — the moment a quick band-aid fix is recognized as fundamentally wrong, and the team commits to a proper architectural solution. This message, [msg 445], is the assistant's response to a sharp question from the user: "Isn't the shared lock just a lazy hack?" The answer is a candid "yes," followed by a re-examination of the root cause that leads to a multi-layer, cross-language refactor spanning C++ CUDA code, Rust FFI bindings, and the proving engine's worker dispatch logic.

The Crisis That Led Here

To understand why [msg 445] matters, we must first understand the crisis that preceded it. The team was building a high-performance zero-knowledge proving system called CuZK, designed to generate Groth16 proofs for Filecoin's proof types (PoRep, WindowPoSt, SnapDeals) using GPU acceleration. The system had recently implemented Pre-Compiled Constraint Evaluator (PCE) extraction — a technique that pre-computes portions of the constraint system to speed up proving — and had just deployed it to production.

The first sign of trouble came when a SnapDeals workload (16 identical partition proofs) was submitted to a remote host (p-dev-ngw-1) equipped with an RTX 4000 Ada GPU (20 GB VRAM). The result was an Out-of-Memory (OOM) crash during GPU proving. The logs told a damning story: two GPU workers entered the C++ proving code at the same millisecond, both attempted to allocate VRAM, and one of them failed with cudaMallocAsync failed: out of memory.

The assistant's initial diagnosis (in [msg 437] through [msg 443]) focused on whether the "shared mutex fix" — a recently introduced change that serialized GPU access through a single mutex — was deployed on the failing host. The shared mutex was a quick fix designed to prevent two workers from entering the GPU code simultaneously, thereby avoiding VRAM contention. The assistant's first recommendation was simply to deploy this shared mutex binary to the failing host.

The User's Challenge

But the user saw a deeper problem. In [msg 444], they asked a question that cut to the heart of the architecture:

"Why is GPU prove for the second GPU not running.. on the second GPU? That's the whole point. CuZK is meant to be a fairly sophisticated proving engine, so it must support multiple GPUs, gpu memory management, and GPU workers are supposed to be interlocking two phases of data transfer to gpu vs compute. Isn't the shared lock just a lazy hack?"

This question exposed a fundamental flaw in the shared mutex approach. The system had two GPUs available, and the Rust engine was dispatching workers to both of them. But the C++ proving code — specifically the generate_groth16_proofs_start_c function — always routed single-circuit proofs to GPU 0, regardless of which GPU the Rust worker was assigned to. The shared mutex didn't fix the routing problem; it just serialized all workers onto GPU 0, wasting the second GPU entirely and creating a VRAM bottleneck where none should exist.

The Pivot: Message 445

The full text of the subject message is worth quoting verbatim:

"You're right. The shared mutex is a hack that serializes everything to GPU 0 and wastes the second GPU entirely. The real fix should make the C++ code actually use the GPU that the Rust worker assigned.

>

Let me re-examine the root cause properly: [read] /tmp/czk/extern/supraseal-c2/cuda/groth16_cuda.cu <path>/tmp/czk/extern/supraseal-c2/cuda/groth16_cuda.cu</path> <type>file</type> <content>480: 481: semaphore_t barrier; 482: auto& caught_exception = pp->caught_exception; 483: size_t n_gpus = std::min(ngpus(), num_circuits); 484: 485: // Alias for the heap-owned provers copy. The prep_msm_thread captures 486: // this reference (which resolves to pp->provers_owned, heap-allocated) 487: // instead of the function parameter (which is on the stack). 488: auto& provers..."

The message is brief — just an admission of error followed by a targeted file read — but it represents a complete reversal of direction. The assistant's response begins with an unambiguous admission:

"You're right. The shared mutex is a hack that serializes everything to GPU 0 and wastes the second GPU entirely. The real fix should make the C++ code actually use the GPU that the Rust worker assigned."

This is not a defensive response. The assistant immediately accepts the critique and reframes the problem. The key insight is that the C++ code's GPU selection logic — n_gpus = min(ngpus(), num_circuits) at line 483 of groth16_cuda.cu — is fundamentally broken for partitioned proofs. When num_circuits=1 (as it is for partitioned proofs where each partition is proved individually), n_gpus evaluates to 1 regardless of how many physical GPUs are available. The subsequent loop at line 883 spawns a single thread with tid=0, which calls select_gpu(0), always routing the work to GPU 0.

The assistant then immediately acts on this insight by reading the C++ source code at the critical lines (480-488), focusing on the n_gpus calculation. This is the moment of diagnosis — the assistant has identified the root cause that the earlier analysis missed while it was focused on the shared mutex as a solution.

Why This Message Matters

[msg 445] is significant for several reasons. First, it represents a re-framing of the problem. The earlier analysis (messages [msg 437] through [msg 443]) had framed the issue as a missing deployment of the shared mutex fix. The assistant had even recommended deploying the shared mutex binary to the failing host. But the user's question revealed that the shared mutex was itself a symptom of a deeper architectural flaw — the C++ code's inability to route single-circuit proofs to the correct GPU.

Second, the message demonstrates intellectual honesty and adaptability. The assistant could have defended the shared mutex approach or argued that it was a reasonable interim solution. Instead, it immediately conceded the point and pivoted to a proper fix. This willingness to abandon a flawed approach is a hallmark of effective debugging.

Third, the message sets the stage for a major multi-layer refactor. The proper fix — threading a gpu_index parameter from the Rust engine through the FFI boundary into the C++ CUDA code — requires changes across at least five files: the C++ groth16_cuda.cu, the Rust FFI in supraseal-c2/src/lib.rs, the bellperson prover functions, the pipeline layer, and the engine's GPU worker code. This is not a small change; it's a fundamental architectural correction.

Assumptions and Thinking Process

The assistant's thinking in [msg 445] reveals several key assumptions. The primary assumption is that the C++ code's GPU selection logic (n_gpus = min(ngpus(), num_circuits)) is the root cause of the misrouting. This is correct — the code explicitly limits GPU usage to the minimum of available GPUs and circuit count, which for single-circuit proofs always yields 1. The assistant also assumes that the fix involves adding a gpu_index parameter to the C++ function signature and propagating it through the entire call chain.

One assumption that is implicit but important: the assistant assumes that the Rust engine already knows which GPU each worker should use. This is true — the engine assigns GPU ordinals to workers and even sets CUDA_VISIBLE_DEVICES environment variables. The problem is that this information is not passed to the C++ code, which independently decides which GPU to use based on its own flawed logic.

The thinking process visible in this message is one of re-examination. Rather than continuing down the path of deploying the shared mutex fix, the assistant steps back and asks: why is the second GPU not being used at all? This is the right question, and it leads directly to the root cause.

Input Knowledge and Output Knowledge

To understand [msg 445], the reader needs knowledge of: the CuZK proving system's architecture (Rust engine dispatching GPU workers, C++ CUDA code for proof generation), the concept of partitioned proofs (where each partition is proved as a single circuit), the shared mutex approach that was previously implemented, and the OOM crash that prompted the investigation.

The message creates new knowledge: the identification of line 483 (n_gpus = min(ngpus(), num_circuits)) as the root cause of GPU misrouting for single-circuit proofs. This is the key architectural insight that drives the subsequent refactor. The message also establishes the direction for the fix — threading gpu_index through the call chain — which the assistant proceeds to implement in the following messages.

Conclusion

[msg 445] is a turning point in the conversation — the moment when a quick hack is abandoned in favor of a proper architectural fix. It demonstrates the importance of questioning assumptions, the value of intellectual honesty in debugging, and the need to understand root causes rather than treating symptoms. The message is short — just a few lines of text and a file read — but it represents a fundamental shift in direction that leads to a correct, multi-GPU implementation of the CuZK proving engine. In the messages that follow, the assistant traces the full call chain, modifies five files across C++ and Rust, deploys the fix, and verifies that both GPUs are correctly load-balanced — a direct result of the pivot that began with this single message.