The Byte-Based Checkin: A Pivot from Ad-Hoc Capping to Principled Memory Management

Subject Message: [assistant] Update checkin to use byte-based check: [edit] /tmp/czk/extern/cuzk/cuzk-core/src/pinned_pool.rs Edit applied successfully.

At first glance, message [msg 4160] appears trivial—a one-line note about updating a function to use "byte-based checks." But this message is the fulcrum of a significant architectural pivot. It represents the moment the assistant abandoned an ad-hoc, performance-hostile fix and committed to a principled, budget-aware memory management strategy. To understand why this single checkin update matters, one must trace the chain of reasoning that led to it.

The Precipitating Crisis

The story begins with a crash. The RTX 5090 vast.ai instance (C.32897009) had been OOM-killed during Phase 2 of a benchmark run ([msg 4154]). The assistant's initial diagnosis identified the pinned memory pool as the culprit: the pool could grow unboundedly, consuming all available memory and triggering the cgroup OOM killer. The first fix was a hard cap of 30 buffers, computed from (max_gpu_queue_depth + gpu_workers_per_device * num_gpus) * 3 ([msg 4121]). On a machine with 755 GiB of memory, 30 buffers at roughly 3.9 GiB each would cap the pool at ~116 GiB—a seemingly safe margin.

But the user rejected this approach decisively ([msg 4155]). Their objection was not about the number 30 but about the principle: "The pinned pool must allow for all synthesis to run in parallel, on larger systems this should be allowed to be even 20 synths in parallel which all can have 3 (a/b/c?) buffers." On a 755 GiB machine, 20 parallel syntheses could need 60+ buffers. A hard cap of 30 would force half of those syntheses to fall back to unpinned (heap) memory, catastrophically degrading H2D transfer performance. The cap was unprincipled—it would harm large systems while only incidentally protecting small ones.

The Byte-Based Pivot

The assistant's response in [msg 4156] shows the conceptual shift: "Instead of a hard buffer count cap, I should derive the cap from the memory budget — so on large machines the pool can grow to accommodate all parallel syntheses, while on small machines it's naturally constrained." This is the key insight: the cap should be proportional to available memory, not to an arbitrary buffer count.

The assistant then executed a series of edits across three files:

pinned_pool.rs: Changed the PinnedPool struct to store max_bytes instead of max_buffers. Updated the constructor, the checkout function (which allocates or reuses buffers), and—crucially—the checkin function (message [msg 4160]).
status.rs: Updated the status snapshot to report max_bytes instead of max_buffers ([msg 4162], [msg 4163]).
engine.rs: Changed the pool construction to compute max_bytes from the memory budget rather than from GPU pipeline parameters ([msg 4164]). The build initially failed due to a stale reference to max_buffers in a log message ([msg 4165]), which was quickly fixed ([msg 4166], [msg 4167]). The final build succeeded ([msg 4168]).

Why `checkin` Matters

The checkin function is where the pool decides whether to retain or free a buffer when it is returned after GPU proving. Under the old buffer-count regime, checkin would free buffers if the number of free buffers exceeded max_buffers. Under the new byte-based regime, checkin must check whether the total allocated bytes exceed max_bytes. This is a fundamentally different decision: it accounts for buffer size, not just buffer count. A pool full of large buffers might need to free some even if the count is low, while a pool full of tiny buffers might safely keep many more.

The byte-based checkin is the enforcement mechanism for the budget-derived cap. Without it, the pool could still grow unboundedly in terms of bytes even if the buffer count is modest (e.g., 10 buffers of 100 GiB each = 1 TiB). The checkin is the gatekeeper that ensures the pool self-limits.

Input Knowledge Required

To understand this message, one needs:

The pool lifecycle: Buffers are allocated on first checkout, reused on subsequent checkouts, and potentially freed on checkin if the pool is over capacity. The live_count tracks total allocated buffers (not yet freed), while free tracks reusable buffers sitting idle.
The memory budget system: The MemoryBudget manages a total budget with permanent and working components. Per-partition reservations are RAII guards that reserve memory during synthesis and release it when the partition completes. The pinned pool's cudaHostAlloc buffers are invisible to the budget—this is the root cause of the OOM.
The GPU pipeline: Synthesis uses pinned a/b/c buffers for fast H2D transfer. If pinned buffers are unavailable, synthesis falls back to heap memory, which requires a bounce-buffer copy at 1-4 GB/s instead of 50 GB/s.
The crash history: The RTX 5090 instance was OOM-killed because the pool grew unboundedly while the budget believed memory was available (since pool allocations were invisible to it).

Output Knowledge Created

This message produces:

A byte-constrained pool: The pool now self-limits based on total allocated bytes, not buffer count. This is strictly more correct because it accounts for varying buffer sizes.
A foundation for budget integration: The byte-based cap is a stepping stone toward full budget awareness. The assistant later designs a two-phase reservation model where the pool acquires budget when allocating and partitions reduce their reservations when checking out pinned buffers ([chunk 30.1]).
A scalable constraint: On large machines with generous budgets, max_bytes is high and the pool grows freely. On small machines, the budget-derived cap naturally limits pool growth without hard-coded buffer counts.

Assumptions and Potential Mistakes

The assistant assumes that a single max_bytes cap derived from the total budget is sufficient. This may not be true: the pool's memory is pinned (non-pageable), while the budget accounts for regular heap memory. Pinned memory cannot be swapped, so reserving too much of the budget for pinned buffers could starve heap allocations even if the total is within budget. The assistant later recognizes this and designs a more sophisticated two-phase model.

Another assumption is that the checkin byte check alone is sufficient enforcement. In practice, the pool also needs to prevent new allocations when over budget—the checkout function was updated in [msg 4159] to check bytes before allocating. The combination of pre-allocation check (checkout) and post-use cleanup (checkin) provides a complete enforcement loop.

The Thinking Process

The assistant's reasoning in [msg 4156] reveals the conceptual leap: "The key insight: the pinned pool's total size should never exceed some fraction of the total memory budget. Since pinned a/b/c replaces heap a/b/c that's already counted in the budget, the pool size should be bounded by the budget's total, not by arbitrary buffer counts." This is correct reasoning—pinned memory is a substitute for heap memory in the synthesis path, so it should be bounded by the same budget that bounds heap.

The choice to update checkin specifically (rather than only checkout) shows an understanding of the pool's lifecycle: checkout is the gate for new allocations, but checkin is where the pool releases memory when under pressure. Both must be byte-aware for the system to work correctly.

Conclusion

Message [msg 4160] is deceptively small. It represents the rejection of a quick fix (the buffer count cap) and the adoption of a principled, budget-derived constraint. The byte-based checkin is the enforcement mechanism that makes the pool self-limiting in proportion to available memory. While the assistant would later design an even more sophisticated two-phase reservation model, this message marks the critical pivot from ad-hoc capping to systematic memory management. It is a lesson in why memory constraints must be derived from the system's actual budget, not from arbitrary pipeline parameters.