The Critical Thread: Wiring a GPU Completion Counter for PI-Controlled Dispatch

Subject Message (msg 3441): [assistant] Add the counter clone alongside the notify clone: [edit] /tmp/czk/extern/cuzk/cuzk-core/src/engine.rs Edit applied successfully.

At first glance, this message appears almost trivial — a single edit command, two lines of code, a clone operation alongside an existing one. Yet this small edit represents the final wiring of a feedback sensor in a sophisticated control system designed to solve one of the hardest problems in GPU-accelerated proving: keeping the GPU pipeline full without overwhelming CPU memory. To understand why this message was written, we must trace the thread back through an iterative journey of failed semaphores, bursty P-controllers, and the emergence of a PI-controlled dispatch pacer.

The Problem: GPU Starvation in a Deep Synthesis Pipeline

The CuZK proving engine operates in two stages: CPU-bound synthesis (20–60 seconds per partition) followed by fast GPU proving (~1 second). The fundamental challenge is that the GPU must never be idle waiting for synthesized work, but neither can the system allow unbounded memory growth from too many pre-synthesized partitions sitting in the queue. This is a classic producer-consumer problem with a severe asymmetry — the producer (synthesis) is 20–60× slower than the consumer (GPU proving), and the pipeline delay from dispatch to GPU completion is measured in minutes.

Previous attempts to solve this had failed. A semaphore-based approach limited total in-flight partitions but failed to maintain a stable pipeline because it didn't target the right metric: the number of synthesized partitions waiting for the GPU. A P-controller that dispatched in bursts on each GPU completion proved too aggressive, instantly filling all allocation slots. A dampened P-controller still suffered from instability because the raw waiting count was a noisy, delayed feedback signal.

The PI-Controlled Pacer: A Control-Theoretic Approach

The assistant's reasoning in [msg 3429] reveals a deep engagement with control theory. The key insight was to treat the dispatch system as a feedback loop where the dispatch rate is the control input, the waiting count is the output, and the synthesis time creates a pure delay of 20–60 seconds. This framing led to a PI (Proportional-Integral) controller design with extremely conservative gains — Kp of 0.1 and Ki of 0.01 — because the long delay meant any aggressive correction would cause oscillation.

But the assistant recognized a critical limitation: PI control alone, operating on the noisy waiting-count signal, would be too slow to react. The solution was a feed-forward architecture. Instead of relying solely on error correction, the pacer would estimate the GPU's consumption rate directly and use that as a baseline dispatch rate, with PI corrections only for fine-tuning. This is the control engineering equivalent of "measure the disturbance directly rather than waiting for its effects."

Why This Message: Wiring the Feedback Sensor

The feed-forward architecture requires a measurement of the GPU completion rate. This is what the gpu_completion_count AtomicU64 provides. Every time a GPU worker finishes a job, it increments this counter. The pacer then reads the counter periodically, computes an Exponential Moving Average (EMA) of the inter-completion interval, and uses that as the feed-forward dispatch rate.

The subject message adds the clone of this counter at the GPU worker spawn point — line 2773 of engine.rs, right alongside the existing gpu_done_for_worker clone. This is the moment the sensor gets wired into every GPU worker's scope. Without this clone, the workers would have no access to the shared counter, and the pacer would be flying blind.

The edit itself is deceptively simple. The assistant reads the file at line 2770–2773 ([msg 3440]) and sees:

let gpu_work_queue = gpu_work_queue.clone();
let mut shutdown_rx = self.shutdown_rx.clone();
let st = self.status_tracker.clone();
let gpu_done_for_worker = gpu_done_notify.clone();

The edit adds a fifth clone: let gpu_completion_for_worker = gpu_completion_count.clone();. This is the pattern established by all the other shared-state clones — each worker gets its own Arc handle to the atomic counter, allowing safe concurrent access across threads.

Input Knowledge Required

To understand this message, one must grasp several layers of context:

The CuZK pipeline architecture: CPU synthesis feeds a GPU queue; GPU workers pull from this queue, prove, and signal completion via Notify. The gpu_done_notify is the existing completion signal used by the old semaphore/dispatch logic.
The shared-state pattern: The engine spawns multiple GPU worker tasks, each receiving cloned Arc handles to shared state (work queue, shutdown channel, status tracker, notify). The gpu_completion_count follows this exact pattern.
The control system design: The PI-controlled pacer uses two feedback signals — the smoothed waiting count (for PI correction) and the GPU completion rate (for feed-forward). The completion counter is the raw sensor for the latter.
Atomic operations and concurrency: AtomicU64 with fetch_add(1, Relaxed) is used because the counter only needs to be monotonically increasing; ordering is handled by the existing Notify synchronization.

Output Knowledge Created

This message produces a single, critical change: every GPU worker now carries a reference to the shared completion counter. The counter is incremented in the happy-path finalizer (as seen in [msg 3444]) before the notify_one() call that wakes the dispatcher. This means:

The pacer can now compute the GPU completion rate by reading the counter and comparing it against wall-clock time.
The EMA of the inter-completion interval becomes a live estimate of GPU throughput.
The feed-forward dispatch rate can be set to match the GPU's actual consumption rate, rather than relying on a fixed target. Without this wiring, the pacer would have no way to know how fast the GPU is actually consuming work. It would be forced to fall back to pure PI control on the waiting count, which — as the assistant's reasoning shows — would be dangerously slow given the 20–60 second pipeline delay.

Assumptions and Decisions

The assistant made several implicit assumptions in this edit:

The counter is sufficient for rate estimation: The assumption is that a simple atomic counter, incremented on each GPU completion, provides enough information to compute a stable rate estimate via EMA. This assumes that GPU completion times are relatively stationary over the EMA window, which may not hold during startup transients or when switching between proof types (which trigger pinned buffer re-allocations).
Relaxed ordering is safe: Using Relaxed ordering for the atomic increment assumes that the Notify mechanism provides sufficient synchronization for the dispatcher to observe a consistent view. This is reasonable because the dispatcher is woken by notify_one() and reads the counter after waking, but it does mean the counter could lag slightly behind reality.
The clone pattern is correct: Following the existing pattern of cloning shared state at worker spawn assumes that the gpu_completion_count has the same lifetime and ownership semantics as the other shared handles. This is correct — it's an Arc<AtomicU64> created in the dispatcher's scope and cloned into each worker.
Only the happy path matters: The counter is only incremented on the happy path (successful GPU proving). Error paths and shutdown paths do not increment it. This means the measured GPU rate will drop during error bursts, which is actually desirable — it naturally slows dispatch when the GPU is failing.

The Broader Context: Iterative Refinement

This message is the culmination of a long iterative process documented across segments 24 and 25 of the conversation. The team moved from a simple reactive semaphore (segment 24) to a P-controller with burst dispatch, then to a dampened P-controller, and finally to the PI-controlled pacer with feed-forward. Each iteration revealed new failure modes:

The semaphore limited in-flight work but couldn't maintain a stable queue depth.
The P-controller dispatched too aggressively because it responded to the deficit all at once.
The dampened P-controller was better but still unstable due to the noisy waiting-count signal. The PI pacer with feed-forward represents a qualitative shift: instead of reacting to the queue depth error, it measures the GPU's actual throughput and matches it. The completion counter is the sensor that makes this possible. The subsequent addition of a synthesis throughput cap (in the same segment) further refines the system by clamping the dispatch rate to not exceed the measured synthesis completion rate, preventing CPU contention when synthesis is the bottleneck.

Conclusion

The subject message — "Add the counter clone alongside the notify clone" — is a masterclass in how the smallest edits can carry the heaviest engineering weight. It is not merely a code change; it is the act of wiring a feedback sensor into a control system, transforming the dispatcher from a reactive burst engine into a rate-matching pacer. The message itself is terse, but the reasoning behind it spans control theory, concurrent programming, pipeline architecture, and the hard-won lessons of iterative deployment. In the broader narrative of the CuZK proving engine, this edit marks the moment the system gained the ability to see its own GPU throughput — and, with that vision, to finally achieve stable, high-utilization scheduling.