The Multi-GPU Pivot: Rewriting the cuzk Engine for Parallel Proof Generation

Introduction

In the development of any high-performance computing system, there comes a moment when the architecture must be fundamentally rethought to support parallelism. For the cuzk proving daemon — a system designed to accelerate Filecoin's Groth16 proof generation across heterogeneous GPU hardware — that moment arrived in message <msg id=329>. This single message, a brief but consequential code-writing action, represents the architectural centerpiece of Phase 1 implementation: the transformation of a single-worker engine into a multi-GPU worker pool.

The message itself is deceptively compact. The assistant writes:

Now let me rewrite the engine to support multi-GPU workers. The key change is: 1. running becomes a map from worker_id → (JobId, ProofKind) 2. We detect GPUs from config or auto-detect, and spawn N workers 3. Each worker gets its own CUDA_VISIBLE_DEVICES env var 4. The scheduler now has per-worker channels [write] /home/theuser/curio/extern/cuzk/cuzk-core/src/engine.rs Wrote file successfully.

Then, almost as an afterthought, the LSP diagnostics output reveals pre-existing errors in an unrelated Go file — a reminder that this system lives at the intersection of multiple languages and build toolchains. But the real story is in those four bullet points, which encode weeks of architectural reasoning, deep knowledge of GPU programming patterns, and a careful balance between immediate functionality and future flexibility.

The Context: Why Multi-GPU Support Was Necessary

To understand the significance of this message, one must trace the path that led to it. The cuzk project began as a response to a specific bottleneck: the Filecoin C2 proof generation pipeline consumed approximately 200 GiB of peak memory, largely due to the SRS (Structured Reference String) parameters being loaded per-proof. The initial Phase 0 implementation ([msg 291] onward) established a working single-GPU daemon that proved the concept of SRS residency — keeping parameters loaded in GPU memory across proof requests to avoid reload overhead, yielding a 20.5% speedup.

But Phase 0 was a proof of concept, not a production system. Phase 1 demanded support for all four Filecoin proof types — WinningPoSt, WindowPoSt, SnapDeals, and the existing PoRep — and crucially, the ability to utilize multiple GPUs simultaneously. A Filecoin storage provider might have a machine with 4, 8, or even 16 GPUs; a single-worker daemon would leave the vast majority of that compute capacity idle. The multi-GPU rewrite was therefore not an optimization — it was a prerequisite for any real-world deployment.

The assistant had been systematically working through a todo list across the preceding messages ([msg 291] through [msg 328]). The protobuf definition had been extended with repeated bytes vanilla_proofs to support the multi-proof nature of PoSt requests. The types.rs had been updated with new fields. The prover.rs had been rewritten with real implementations for generate_winning_post_with_vanilla, generate_single_window_post_with_vanilla, and generate_empty_sector_update_proof_with_vanilla, complete with manual enum mappings from the gRPC uint64 values to the FFI's #[repr(i32)] discriminants. The service layer and bench tool had been updated. All that remained was the engine itself — the central coordinator that owns the scheduler, workers, and SRS manager.

The Four-Point Plan: Deconstructing the Architecture

The assistant's four bullet points are not arbitrary; each addresses a specific structural requirement of multi-GPU operation.

Point 1: running Becomes a Map

The original engine tracked a single running job. This was adequate for a single-worker design: at any moment, at most one proof was being generated. But with multiple workers, the engine needs to track which job is running on which worker. The map from worker_id → (JobId, ProofKind) serves multiple purposes: it enables status queries that report per-worker activity, it supports graceful shutdown by knowing which workers are busy, and it provides the foundation for future GPU affinity-based scheduling (where a job might prefer a specific GPU because its SRS parameters are already resident there).

The choice of ProofKind alongside JobId in the map value is telling. The assistant is not just tracking that a worker is busy, but what kind of proof it's generating. This foreshadows the Phase 2 priority scheduling where different proof types might have different urgency levels — WinningPoSt proofs, for instance, are time-critical for Filecoin consensus, while PoRep proofs can tolerate lower priority.

Point 2: GPU Detection and Worker Spawning

The second point — detecting GPUs from config or auto-detection, and spawning N workers — addresses the fundamental challenge of heterogeneous environments. A cuzk deployment might run on bare metal with direct GPU access, in a container with a subset of GPUs exposed, or in a cloud instance with GPU passthrough. The engine must handle all these cases.

The auto-detection mechanism implicitly assumes the ability to enumerate CUDA-capable devices, likely via nvidia-smi or the CUDA runtime API. The config override provides an escape hatch for environments where auto-detection is unreliable or where the operator wants to reserve some GPUs for other workloads. This dual approach reflects a pragmatic engineering philosophy: provide sensible defaults, but always allow explicit configuration.

The decision to spawn one worker per GPU (rather than, say, one worker per job with dynamic GPU assignment) is significant. It means each worker owns its GPU for the duration of the daemon's lifetime, avoiding the overhead of CUDA context initialization per proof. This aligns with the SRS residency strategy — keeping parameters loaded in GPU memory across proof requests requires that the GPU be dedicated to the daemon, not shared ephemerally.

Point 3: CUDA_VISIBLE_DEVICES Isolation

The third point — each worker gets its own CUDA_VISIBLE_DEVICES environment variable — is the most technically nuanced. CUDA uses this environment variable to restrict which physical GPUs a process can see. By setting it per-worker, the assistant ensures that each worker process (or thread) can only access its assigned GPU, preventing accidental cross-GPU memory access or scheduling conflicts.

This approach has a critical implication: the workers must be separate OS processes, or at minimum, the environment variable must be set before any CUDA calls are made. In Rust, this typically means spawning worker threads before initializing any CUDA context, or using std::process::Command to spawn child processes. The assistant's earlier research into GPU inference engine architectures (vLLM, Triton, TensorRT-LLM) in Segment 3 likely informed this design choice — those systems use similar process-level isolation for GPU workers.

The use of CUDA_VISIBLE_DEVICES rather than CUDA streams or contexts within a single process is a deliberate trade-off. Process-level isolation provides stronger guarantees (a crash in one worker doesn't affect others) and simpler resource accounting, at the cost of higher memory overhead (each process loads its own CUDA runtime) and more complex inter-worker communication. For a proving daemon where each proof generation is independent and memory-bound, this trade-off favors isolation.

Point 4: Per-Worker Scheduler Channels

The fourth point — the scheduler now has per-worker channels — completes the architectural picture. The original scheduler used a single shared queue with a BinaryHeap for priority ordering. With per-worker channels, the scheduler can dispatch jobs to specific workers based on GPU affinity, load balancing, or priority.

The channel-based design is a classic actor pattern: each worker runs an event loop that receives jobs from its channel, processes them, and sends results back. The scheduler, running in a separate async task, selects which worker gets which job. This decouples scheduling policy from execution, allowing the scheduling algorithm to be refined in Phase 2 without changing the worker implementation.

Notably, the assistant chose to retain the shared priority queue with BinaryHeap as the job dispatch mechanism, with per-worker channels as the delivery mechanism. This is evident from the later analysis in the chunk summary, which notes: "A key architectural decision was made to defer GPU affinity-based scheduling to Phase 2, as the current process-global GROTH_PARAM_MEMORY_CACHE makes per-GPU SRS tracking unnecessary for this phase. The shared priority queue with BinaryHeap was retained as the correct design for Phase 1."

Assumptions and Their Implications

Every architectural decision rests on assumptions, and this message is no exception. Several implicit assumptions deserve scrutiny.

Assumption 1: GPU count is static. The engine detects GPUs at startup and spawns one worker per GPU. This assumes that the GPU topology doesn't change during the daemon's lifetime — no hot-plugging, no GPU failures, no dynamic allocation. In cloud environments with GPU passthrough, this is generally safe. On bare metal with NVIDIA's MIG (Multi-Instance GPU) technology, the situation is more complex, as MIG partitions can be reconfigured dynamically.

Assumption 2: Workers are interchangeable. The per-worker channel design implies that any worker can handle any proof type. This is true for the current implementation, where all workers share the same process-global GROTH_PARAM_MEMORY_CACHE. However, Phase 2's GPU affinity scheduling will challenge this assumption — a worker that has SRS parameters for 32 GiB sectors might be preferred for PoRep jobs, while another worker with different parameters might handle PoSt. The worker state tracking added in this message (the worker_id → (JobId, ProofKind) map) lays the groundwork for this differentiation.

Assumption 3: CUDA_VISIBLE_DEVICES provides sufficient isolation. This assumes that the CUDA runtime respects this environment variable and that no other mechanism (e.g., CUDA context sharing, peer-to-peer access) bypasses it. In practice, CUDA_VISIBLE_DEVICES is well-tested and reliable, but it only controls which devices are visible, not how much memory each device uses. A memory-hungry worker could still exhaust GPU memory and cause CUDA out-of-memory errors that might affect other workers indirectly (e.g., through driver-level resource contention).

Assumption 4: The LSP errors in filecoin-ffi/proofs.go are pre-existing and unrelated. The assistant dismisses the Go LSP errors with "Those Go LSP errors are pre-existing (CGO-related, not related to our changes)" in the subsequent message ([msg 330]). This is correct — the errors relate to CGO type conversion issues in the Go FFI layer, not in the Rust code being modified. However, the presence of these errors in the diagnostic output serves as a reminder of the system's complexity: the cuzk daemon sits atop a multi-language stack (Go, Rust, C, CUDA), and changes in one layer can have subtle effects on others.

The Thinking Process: What the Message Reveals

The assistant's reasoning, visible in the bullet-point plan, reveals a methodical approach to architectural transformation. Rather than diving into code immediately, the assistant first enumerates the key structural changes, then executes them in a single [write] operation that rewrites the entire engine.rs file.

This "plan then execute" pattern is characteristic of the assistant's approach throughout the cuzk project. Earlier messages show extensive research into the codebase — reading every source file, examining FFI signatures, studying Go enum mappings — before writing any code. The engine rewrite is the culmination of that research, where the accumulated understanding of the system's architecture is crystallized into a concrete implementation.

The decision to rewrite the entire file rather than applying incremental edits is significant. It indicates that the changes are pervasive enough to warrant a fresh implementation rather than surgical modifications. The [write] command replaces the entire file content, which the assistant had previously read and understood. This is a high-risk, high-reward approach: it ensures consistency across the new design, but any oversight in the rewrite could introduce subtle bugs.

The assistant's confidence in this approach is justified by the extensive preparation. The preceding messages show a deep investigation of the scheduler internals, the worker loop, the SRS manager, and the status reporting infrastructure. The new engine design incorporates all these components while adding the multi-GPU dimension.

Input Knowledge Required

To fully understand this message, one must be familiar with several domains:

Output Knowledge Created

This message produces a single concrete artifact: the rewritten engine.rs file at /home/theuser/curio/extern/cuzk/cuzk-core/src/engine.rs. However, the implications extend far beyond that file:

  1. A multi-GPU architecture that can utilize all available GPUs simultaneously, each running an independent worker loop.
  2. A foundation for GPU affinity scheduling, with per-worker state tracking that can be extended in Phase 2.
  3. Process-level GPU isolation via CUDA_VISIBLE_DEVICES, preventing cross-worker interference.
  4. A channel-based dispatch mechanism that decouples scheduling policy from execution.
  5. A configurable worker count, supporting both auto-detection and explicit configuration.

Conclusion

Message <msg id=329> represents the architectural pivot point of the cuzk Phase 1 implementation. In four bullet points and a file write, the assistant transforms a single-worker proof-of-concept into a multi-GPU production-ready system. The brevity of the message belies the depth of reasoning behind it — weeks of codebase exploration, FFI analysis, and architectural planning distilled into a single, decisive action.

The LSP errors that follow the successful write serve as a fitting coda: they remind us that this system lives at the intersection of multiple languages and toolchains, where changes in one layer must be carefully coordinated with others. The assistant's ability to distinguish between pre-existing errors and new issues, to plan before executing, and to balance immediate functionality with future flexibility, is what makes this message a masterclass in systems engineering.

The multi-GPU engine, once committed, becomes the foundation for everything that follows: the gen-vanilla command for end-to-end testing, the Phase 2 GPU affinity scheduler, and ultimately, a production proving daemon capable of serving real Filecoin storage providers with heterogeneous GPU hardware.