Threading GPU Awareness Through the Stack: From Shared Mutex Hack to Proper Multi-GPU Architecture

Introduction

In the life of any complex software project, there comes a moment when a quick workaround must be confronted for what it truly is: a deferral of architectural debt. This segment of the opencode session captures exactly such a moment, and the methodical, multi-layered refactoring that followed. The CuZK proving engine — a high-performance GPU-accelerated zero-knowledge proof system — had been limping along with a shared mutex hack that serialized all partition proofs onto a single GPU, wasting half the hardware on a dual-GPU system. The fix worked in the narrow sense of preventing data races, but it was a lazy solution that the user rightly called out.

What followed was a systematic refactoring that threaded a gpu_index parameter through five layers of code spanning C++ CUDA kernels, Rust FFI bindings, Bellperson prover functions, pipeline orchestration, and engine worker dispatch. The fix was built, deployed to a remote dual-GPU test host, and rigorously verified through nvidia-smi, journalctl logs, and real-time nvtop monitoring. Both GPUs came alive with proper load balancing — workers 0 and 1 on GPU 0, workers 2 and 3 on GPU 1 — and the changes were committed with a detailed message across five files.

Immediately following the commit, the conversation pivoted to an entirely different domain: preparing a Docker container packaging Curio, CuZK, and CUDA dependencies for fetching 32G PoRep and Snap proving parameters. The assistant researched the existing build system, read Dockerfiles and Makefiles, spawned parallel subagent tasks, and solicited user configuration decisions (mainnet build, CUDA 13 base image, runtime parameter fetching). This pivot from deep systems debugging to infrastructure engineering demonstrates the breadth of work required to make a proving engine production-ready.

This article traces both arcs — the multi-GPU fix and the Docker preparation — and examines the engineering methodology that connects them.

The Shared Mutex Hack: A Quick Fix That Wasted Hardware

The story begins with a problem that had been diagnosed in earlier segments: the CuZK proving engine's C++ CUDA code always routed single-circuit proofs to GPU 0, regardless of which Rust worker submitted them. The function generate_groth16_proofs_start_c computed n_gpus = std::min(ngpus(), num_circuits). For partitioned proofs, num_circuits is always 1 (each partition is proven as a single circuit), so n_gpus was always 1, and the code always called select_gpu(0). This meant that every worker — whether assigned to GPU 0 or GPU 1 by the Rust engine — ended up executing its CUDA kernels on physical GPU 0. On a two-GPU system, workers 0 and 1 (assigned to GPU 0) and workers 2 and 3 (assigned to GPU 1) all targeted GPU 0 simultaneously, creating data races on GPU memory allocations and kernel launches.

The initial fix was pragmatic but shortsighted. Rather than fixing the GPU routing logic, the assistant added a shared mutex in engine.rs that serialized all partition proofs onto GPU 0. This prevented the data race — only one worker could enter the C++ GPU code at a time — but at the cost of wasting the second GPU entirely. The assistant deployed this fix, and it worked in the narrow sense: proofs no longer failed due to data corruption. But nvidia-smi showed GPU 0 at 100% utilization while GPU 1 sat at 0%.

The shared mutex hack was acknowledged as suboptimal but treated as a temporary measure. The team moved on to other work — implementing PCE extraction for all proof types, fixing a WindowPoSt crash, and investigating partitioned proof failures. But the architectural debt remained, silently accumulating interest.

The SnapDeals OOM: When the Hack Broke

The shared mutex hack's inadequacy was exposed by a new workload. The user reported a SnapDeals proof job crashing with an out-of-memory (OOM) error on a different host — p-dev-ngw-1, equipped with a single 20 GB RTX 4000 Ada GPU. The journalctl logs told a damning story: two GPU workers entered the C++ proving code simultaneously, both allocating memory on GPU 0, and the second allocation exhausted the 20 GB VRAM budget.

The critical log line was CUZK_TIMING: d_a_cache allocated 4096 MiB on gpu 0. Worker 1, assigned to GPU 1, was allocating a 4 GB cache on GPU 0 — the exact same bug that had caused the PoRep data race. The SnapDeals circuit was enormous: PCE extraction revealed 81 million constraints, 443 million non-zero entries, and a 15.8 GB memory footprint. With two workers competing for the same GPU's VRAM, the OOM was inevitable.

The user's question was pointed: "can this be related to recent changes or was it already there (low vs high mem gpu)?" This forced a reassessment. Was the SnapDeals crash a regression from the shared mutex fix, or a pre-existing limitation of the 20 GB GPU?

The Moment of Clarity: "Isn't the Shared Lock Just a Lazy Hack?"

The turning point came when the user directly challenged the shared mutex approach:

"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 was the moment the mask slipped. The assistant immediately agreed and pivoted from defending the shared mutex to planning the proper fix. The user added two crucial pieces of context: SnapDeals consists of 16 identical partitions, and those 16 proofs should obviously load-balance between the two GPUs.

The proper fix was clear: thread a gpu_index parameter through the entire call chain so that the C++ code uses select_gpu(gpu_index) instead of always defaulting to GPU 0. The design used -1 as a sentinel value meaning "auto-detect" for callers outside the engine path (tests, standalone tools), while any non-negative value forced selection of that specific GPU. This preserved backward compatibility while enabling the engine to pass each worker's assigned GPU ordinal.

The Implementation: A Methodical Bottom-Up Refactoring

The assistant executed the fix with systematic precision, working from the deepest layer upward across five files:

C++ CUDA code (groth16_cuda.cu): Added a gpu_index parameter to generate_groth16_proofs_start_c and generate_groth16_proofs_c. For single-circuit proofs, the code now calls select_gpu(gpu_index) instead of defaulting to GPU 0. Multi-circuit batched proofs, which already had their own GPU selection logic, were left unchanged. The d_a_cache was also converted from a global singleton to a per-GPU array, preventing cache thrashing when workers on different GPUs run concurrently.

Rust FFI layer (supraseal-c2/src/lib.rs): Updated the extern "C" declarations to include the new gpu_index: i32 parameter. Updated the public wrapper functions start_groth16_proof and generate_groth16_proof to accept and forward the parameter. The older generate_groth16_proofs function (used by non-engine callers) passes -1 for auto-selection.

Bellperson prover (supraseal.rs): Updated prove_start and prove_from_assignments to accept gpu_index: i32 and forward it through the FFI calls. This required threading the parameter through multiple function signatures.

Pipeline layer (pipeline.rs): Updated gpu_prove and gpu_prove_start to accept and forward gpu_index. Updated all internal call sites that previously passed std::ptr::null_mut() for the mutex to also pass -1 for auto GPU selection. A pre-existing bug was also discovered and fixed during the refactoring: a variable name mismatch (synth_duration vs synthesis_duration) that was exposed by the edits.

Engine layer (engine.rs): Reverted the shared mutex hack, restoring per-GPU mutexes. The GPU worker now passes gpu_ordinal as i32 as the gpu_index parameter, ensuring each worker targets its assigned GPU.

The assistant verified completeness using rg (ripgrep) to find all call sites of gpu_prove and ensure they were updated. The build succeeded, confirming that all type signatures, parameter changes, and FFI declarations were consistent across five layers of code spanning C++ and Rust.

Deployment: The Proof That Both GPUs Work

With the binary compiled, the assistant deployed it to the remote test host (cs-calib, a machine with two NVIDIA RTX A6000 GPUs). The deployment command at [msg 515] was a carefully orchestrated sequence: stop the daemon, rsync the new binary to a temporary location, copy it to the permanent path, restart the service, wait for initialization, and verify the status. The output showed the daemon active (running) with 7.1 GB of memory allocated and 328 tasks spawned.

The first single PoRep proof completed successfully in approximately 110 seconds — confirming basic functionality was intact. Then came the critical check: nvidia-smi at [msg 518] showed both GPUs with 12,986 MiB of memory allocated. Previously, GPU 1 had only 700 MB. This was the first concrete evidence that the gpu_index parameter threading was working: both GPUs were now being initialized with the d_a_cache and SRS parameters.

At [msg 519], the assistant designed a burst test — three concurrent proof submissions on a two-GPU system — to stress-test the load balancing. The reasoning was sound: if the fix was incorrect, concurrent proofs would cause contention, data races, or OOM errors. If the fix was correct, the engine's GPU workers would distribute the work across both devices.

The results were conclusive. All three proofs completed successfully. The journalctl logs at [msg 520] showed workers correctly load-balanced: workers 0 and 1 on GPU 0 (6–7 jobs each), workers 2 and 3 on GPU 1 (5–6 jobs each). The d_a_cache allocation logs at [msg 527] showed both GPUs receiving d_a_cache allocated 4096 MiB at startup — GPU 0 and GPU 1 both initialized their caches during the first proof. The user independently confirmed via nvtop at [msg 526] that both GPUs were seeing load, providing real-time visual validation that complemented the log-based analysis.

The Verification Process: Digging Deeper

The verification process was not without moments of doubt. At [msg 522], the assistant initially saw only gpu_tid=0 in the timing logs and worried that GPU 1 was not being used. However, further investigation revealed that gpu_tid is the loop index within a per-GPU thread — always 0 for single-circuit proofs — not the physical GPU device. The actual GPU is determined by the gpu_base + tid offset, where gpu_base is the worker's assigned GPU ordinal.

A narrow time window on the d_a_cache logs at [msg 523] initially showed only GPU 0, requiring a broader query (10 minutes instead of 3) to reveal both GPUs. At [msg 527], the assistant expanded the journalctl query and found the conclusive evidence: at 03:39:59, both d_a_cache allocated 4096 MiB on gpu 0 and d_a_cache allocated 4096 MiB on gpu 1 fired during the first proof. Both GPUs were being used properly.

These moments of doubt — captured in messages 522–525 — are a testament to the assistant's methodological rigor: rather than accepting the first answer, it dug deeper, cross-referenced multiple data sources, and only declared success when the evidence was unambiguous.

The Commit: Documentation as Engineering Artifact

With verification complete, the assistant committed the changes at [msg 533]. The commit message is a model of engineering documentation:

"cuzk: route partition proofs to correct GPU via gpu_index parameter

Thread a gpu_index parameter through the entire proving stack (C++ -> supraseal-c2 -> bellperson -> pipeline -> engine) so that single-circuit partition proofs run on the GPU assigned to the Rust worker instead of always landing on GPU 0.

Previously, the C++ code computed n_gpus = min(ngpus(), num_circuits), which for single-circuit proofs always resolved to GPU 0 via select_gpu(0). This made per-GPU mutexes ineffective on multi-GPU systems: workers assigned to different GPUs could run CUDA kernels simultaneously on GPU 0, causing proof corruption (the original shared-mutex workaround serialized everything to one GPU).

Now gpu_index >= 0 pins work to that specific GPU, while -1 preserves the original multi-GPU fan-out for batched proofs. Also converts the global d_a_cache singleton to a per-GPU array to avoid thrashing when workers on different GPUs run concurrently."

The message explicitly contrasts the old behavior with the new behavior, candidly acknowledges the failed first attempt ("the original shared-mutex workaround serialized everything to one GPU"), and documents the design rationale for the sentinel value convention. This intellectual honesty — documenting dead ends as well as solutions — is a hallmark of good engineering.

The diff touched five files with 75 insertions and 63 deletions. The net change was only 12 lines — a reminder that the best fixes are often small in diff but large in insight.

The Pivot: From GPU Debugging to Docker Infrastructure

Immediately following the commit, the conversation took a sharp turn. At [msg 536], the user issued a request belonging to an entirely different domain of software engineering:

"Prepare a docker container with curio, cuzk, all curio dependencies, and fetching 32G PoRep and Snap (only, ideally) proving params with e.g. curio fetch-params (fetch probably as autostart script). The container should start into bash, for now, and will be published to dockerhub."

This pivot reveals something important about the user's mental model: the multi-GPU fix was never an end in itself. It was a blocker that needed to be resolved before the real goal — packaging the entire proving stack into a portable, reproducible Docker image — could be pursued. The debugging phase was necessary infrastructure work; the Docker container was the deliverable.

The assistant did not immediately start writing a Dockerfile. Instead, it initiated a structured research phase spanning messages 537–543. This is a deliberate methodological choice: before constructing, understand.

The assistant created a todo list with research items, then spawned two parallel subagent tasks at [msg 538] to explore the curio build system and the cuzk build requirements. These subagents ran concurrently, each producing comprehensive reports on their respective domains. The curio analysis covered existing Dockerfiles, Makefiles, build scripts, and dependency chains. The cuzk analysis covered the Cargo workspace structure, Rust version requirements, CUDA dependencies, and the supraseal C++ build process.

But the assistant did not stop at subagent summaries. It then read the primary sources directly: the existing Dockerfile at [msg 540] to understand the multi-stage build patterns already in use; the 00-vars.mk and 10-deps.mk Makefiles to understand build variables and dependency targets; the 30-build.mk file at [msg 541] to understand the curio build command itself; and the cuzk.example.toml to understand the daemon configuration. At [msg 543], the assistant read build-blst.sh and the first 70 lines of supraseal/build.sh to understand the C++/CUDA compilation requirements, including the GCC version constraints imposed by the CUDA toolkit.

This systematic reading reveals a clear hierarchy of concerns: first understand the containerization pattern (Dockerfile), then understand the build variable conventions (vars.mk), then understand the dependency initialization (deps.mk), then understand the specific CUDA compilation requirements (build scripts). The assistant was building a mental model before building a container.

The Deliberate Pause: Asking Before Acting

At [msg 544], the assistant reached a critical inflection point. It had completed its research and had everything it needed to write a Dockerfile. But instead of proceeding, it paused. The reason is captured in the first sentence: "Now I have everything I need. Let me also quickly check what curio fetch-params actually downloads and whether we can filter for just PoRep and Snap params."

The assistant then presented a structured question to the user with three configuration decisions that would fundamentally shape the Dockerfile's architecture:

Network target (calibnet vs. mainnet): The curio build system uses Go build tags to target either the test network or production. This choice affects which parameters are fetched and which contracts are compiled in.
CUDA base image version (12.8 vs. 13): The assistant had been researching CUDA 12.8 as the "latest stable" version, but the user's answer — "cuda13 is a thing now" — revealed a more recent release. The base image choice ripples through the entire Dockerfile: it determines which apt packages are available, which nvcc version is used, and which runtime libraries must be included.
Parameter fetching strategy (runtime vs. baked-in): The proving parameters for 32G PoRep and Snap total approximately 100 GB. Baking them into the Docker image would make it impractically large. Fetching them at runtime via an entrypoint script keeps the image lean (~15 GB) but requires internet access on first start. The assistant labeled one option as "Recommended" (calibnet, CUDA 12.8, fetch at runtime) but presented alternatives with clear trade-offs. This framing shows mature engineering judgment: the assistant had already evaluated the options and formed a preference, but recognized that some decisions must be informed by human authority. The user's truncated response — "mainnet, cuda13 is a thing now..." — overrode the recommendation on both the network target and the CUDA version, while presumably accepting the runtime fetching strategy. This interaction is a model of effective human-AI collaboration: the assistant researches thoroughly, identifies decision points, presents structured options, and the user provides the domain-specific knowledge that the assistant cannot infer from code alone.

The Bridge Between Two Worlds

What makes this segment particularly compelling is not the individual accomplishments — the multi-GPU fix or the Docker research — but the seamless transition between them. The assistant demonstrated the ability to context-switch from deep systems debugging (thinking about CUDA device selection, mutex semantics, and FFI parameter threading) to infrastructure engineering (thinking about base images, multi-stage builds, and entrypoint scripts) without losing methodological rigor.

The common thread is the assistant's systematic approach: understand the system before modifying it, verify empirically rather than assuming correctness, and document decisions clearly. In the debugging phase, this meant reading code, tracing call chains, deploying to remote hosts, and cross-referencing log sources. In the infrastructure phase, this meant reading Dockerfiles, spawning research subagents, examining build scripts, and soliciting user input before writing code.

The segment also illustrates a deeper truth about software engineering: the boundary between "debugging" and "infrastructure" is often artificial. The multi-GPU fix was necessary because the proving engine would crash on multi-GPU systems. The Docker container was necessary because the proving engine needed to be distributed to users. Both tasks were in service of the same goal — making GPU-accelerated zero-knowledge proving work reliably in production — but they demanded different skills, different tools, and different modes of thinking.

Lessons Learned

Several themes emerge from this segment:

The danger of the quick fix. The shared mutex hack was not wrong — it prevented data races. But it was incomplete, and it masked the true architectural flaw. The SnapDeals OOM revealed that the hack was insufficient even for its intended purpose, and the user's challenge forced a proper solution.

The importance of threading parameters through the full call chain. The GPU assignment was correct at the Rust engine level but was silently discarded at every subsequent layer. The fix required adding the gpu_index parameter at every boundary, from C++ through FFI through Bellperson through pipeline through engine. Missing any single call site would have left the bug partially unfixed.

The value of systematic verification. The assistant used grep to verify every call site, built the project to confirm compilation, deployed to remote hosts for runtime validation, and cross-referenced multiple log sources before declaring success. Each verification step caught potential issues before they could cause failures in production.

The role of the user in driving architectural quality. The user's challenge — "Isn't the shared lock just a lazy hack?" — was the catalyst that transformed a workaround into a proper fix. In collaborative AI-assisted development, the human's ability to question assumptions and demand better solutions is irreplaceable.

The discipline of research before construction. When pivoting to the Docker task, the assistant did not immediately start writing a Dockerfile. It researched the existing build system, read primary sources, spawned parallel subagents, and solicited user input. This upfront investment in understanding prevented costly mistakes and ensured the Dockerfile would align with the project's existing patterns.

Conclusion

This segment of the opencode session captures a complete arc of systems engineering: from the recognition that a shared mutex hack was an architectural dead end, through the methodical refactoring that threaded GPU awareness through five layers of a multi-language proving stack, to the deployment and rigorous verification on a remote dual-GPU host, and finally the pivot toward Docker infrastructure for packaging the entire stack into a portable container.

The multi-GPU fix for CuZK's proving engine is a case study in the evolution from quick hack to proper architecture. What began as a shared mutex workaround for a GPU data race was transformed, through the pressure of a SnapDeals OOM crash and a user's pointed critique, into a systematic refactoring that actually uses both GPUs as intended. The fix was deployed, verified, and committed — restoring the proving engine's ability to load-balance across multiple GPUs.

The Docker pivot that followed demonstrates that the fix was never the end goal. It was a prerequisite for the broader mission of making CuZK production-ready across diverse hardware environments. The story is not yet complete — the Dockerfile has not been written, the container has not been built — but the foundation has been laid. The research is done, the user's configuration decisions have been captured, and the assistant is poised to construct the container.

What makes this segment satisfying is not the resolution of a single problem, but the demonstration of a repeatable methodology: understand, verify, document, pivot, research, ask, build. It is a template for how complex systems engineering should be done, whether by humans or by AI.