[memory] BUFFERED loader: default io_depth staging footprint; cached pinned-buffer deleter lifetime

[malaiwah]

Found during a source-level memory review of a deployment image that pins InstantTensor at commit 85e7c5f5 ("add thread pool"). All code references and permalinks below are against that commit; both observations were re-verified line-by-line before filing. These are proposals from a static review — happy to be corrected if we misread the lifecycle.

Context for the numbers: we reviewed usage at tensor-parallel world_size=4 with the default 8 MiB chunks on 96 GB GPUs, loading a large sharded checkpoint via the AIO/URING (BUFFERED) path.

1. Default io_depth = 512 // world_size makes the load-phase staging footprint 4 GiB VRAM + 1 GiB pinned host per rank

Where

instanttensor/_impl.py:539-546 (permalink):

            else:
                # AIO/AIO_BUFFERED/URING/URING_BUFFERED
                if chunk_size is None:
                    chunk_size = 8*1024*1024
                if concurrency is None:
                    concurrency = 1 # max(1 // self.world_size, 1)
                if io_depth is None:
                    io_depth = max(512 // self.world_size, 3) # aio read + cudaMemcpyAsync + ncclAllGather

csrc/loader_common.cpp:75-95 (permalink):

    this->rank_chunk_size = this->thread_chunk_size * this->num_threads;
    this->world_chunk_size = this->rank_chunk_size * this->world_size;

    size_t inflight_device_buffer_size = this->io_depth * this->world_chunk_size;
    if (this->buffer_size < inflight_device_buffer_size) this->buffer_size = inflight_device_buffer_size;
    ...
    CUDA_CHECK(cudaMalloc(&this->device_buffer, this->buffer_size));
    ...
    if (this->need_host_buffer) {
        size_t inflight_host_buffer_size = this->io_depth * this->rank_chunk_size;

What / impact (recomputed)

At world_size=4 with all-default parameters on the AIO/URING path:

• io_depth = max(512//4, 3) = 128, chunk_size = 8 MiB, concurrency (num_threads) = 1
• rank_chunk_size = 8 MiB, world_chunk_size = 32 MiB
• device staging floor = io_depth × world_chunk_size = 128 × 32 MiB = 4.00 GiB of VRAM per rank (plus ~60 KiB alignment padding; _determine_buffer_size uses the same product as buffer_size_for_io, so the Python-side buffer_size never lowers it)
• pinned host buffer = io_depth × rank_chunk_size = 128 × 8 MiB = 1.00 GiB per rank (aligned_alloc + cudaHostRegister, BUFFERED backends)

The max_free_mem_usage guard (_impl.py:565-586) cannot moderate this in the common case: with the default ratio 0.5 the shrink only fires when free VRAM is below 2 × 4 GiB = 8 GiB, so on healthy GPUs every rank silently commits 5 GiB of staging for the whole load. And when free memory is tight, the fallback shrinks to io_depth=3 or hard-fails — issue #13 shows exactly this sequence ending in cudaMalloc ... out of memory at loader_common.cpp:79 (its second-model-load OOM has a separate proximate cause, but the size being requested comes from this formula).

The buffers are transient — Loader::close → destroy_buffer frees both (default env) — so this defines the load-phase peak rather than a leak. Still, the peak is what a co-resident allocator (e.g. an inference engine loading two models, or reserving KV cache right after load) has to survive, and 128-deep × 32 MiB in-flight is far past the point of diminishing returns for a single-threaded AIO/URING pipeline; the bandwidth-delay product of NVMe→host→device suggests depth 16–32 achieves the same throughput at 1/4–1/8 the footprint.

Suggested fix

Cap the automatic default (e.g. io_depth = clamp(512 // world_size, 3, 32)), or scale the default against avail_bytes up front instead of only shrinking on the 50%-of-free cliff. Users who want the deeper queue can still set INSTANTTENSOR_IO_DEPTH.

Verification
Run any multi-rank load with INSTANTTENSOR_DEBUG=1 and no overrides; the Config: line printed from Loader::open shows io_depth=128, host_buffer_size=1073741824, and device_buffer_size ≥ 4294967296 (the 4 GiB in-flight floor plus alignment padding; larger if the tensor-derived size dominates) at world_size=4. Compare wall-clock load time against INSTANTTENSOR_IO_DEPTH=32.

2. INSTANTTENSOR_CACHE_BUFFER=1: cached pinned-buffer deleter captures this of a Loader that is already destroyed — use-after-free at cache eviction / process exit

Where

The deleter is created in Loader::init_buffer, csrc/loader_common.cpp:108-114 (permalink):

            this->host_buffer_entry.deleter = [=](void *ptr) {
                if (this->backend == Backend::URING || this->backend == Backend::URING_BUFFERED) {
                    this->deregister_host_buffer_uring();
                }
                CUDA_CHECK(cudaHostUnregister(ptr));
                free(ptr);
            };

[=] captures the raw this pointer. With INSTANTTENSOR_CACHE_BUFFER=1, destroy_buffer (loader_common.cpp:131-138) moves the entry — deleter included — into the process-global cache instead of invoking it:

    if (this->need_host_buffer) {
        if (_env_cache_buffer()) {
            host_buffer_cache->put(std::move(this->host_buffer_entry));
        }

host_buffer_cache is a global (csrc/instant_tensor/io_context.hpp:89), and its destructor runs every stored deleter (io_context.hpp:45-50):

    ~HostBufferCache() {
        std::lock_guard<std::mutex> lock(mutex);
        for(auto& entry : cached_host_buffers) {
            entry.deleter(entry.ptr);
        }
    }

Lifecycle trace (why this is dead by then)

1. LoaderManager::open (csrc/main.cpp:54-55) spawns the loader on a detached thread:

   std::thread loader_thread(run_loader, std::unique_ptr<SPSCQueue<RPCRequest>>(input_queue), std::unique_ptr<SPSCQueue<RPCResponse>>(output_queue));
   loader_thread.detach();

2. run_loader (loader_common.cpp:601-610) holds the Loader as a stack local: Loader loader(std::move(input_queue), std::move(output_queue)); loader.run();. Nothing else owns or references the Loader object — LoaderManager only keeps the raw queue pointers, and even those are erased in LoaderManager::close.

3. The CLOSE RPC runs Loader::close → sets stop = true and calls destroy_buffer, which put()s the entry (with the this-capturing deleter) into the global cache.

4. Loader::run's loop observes stop, returns; run_loader returns; the stack-local Loader is destroyed and the detached thread exits. From this point the captured this dangles.

5. The stale deleter is eventually invoked with the dangling this:

   • at ~HostBufferCache — triggered either by the cleanup() atexit hook (main.cpp:135-139, registered in _impl.py:28) or by static destruction; or
   • via reuse: a subsequent Loader that hits the cache in init_buffer (loader_common.cpp:96-97) inherits the entry with the first loader's deleter (the deleter is only assigned on a cache miss), so whichever path finally runs it still dereferences the first, long-dead Loader.

   The deleter reads this->backend, and on the URING path calls this->deregister_host_buffer_uring(), which touches this->uring_register_buffer and io_uring_unregister_buffers(&this->uring_ring) on freed thread-stack memory. Even in the "benign" misread case it's undefined behavior; in the URING case it operates on a destroyed ring.

Secondary effect: with the flag on, the pinned buffer (1 GiB per rank at the defaults above) intentionally stays registered for process lifetime; for a load-once workflow that is a permanent pinned-memory hold, worth documenting alongside the fix.

This is latent/opt-in — nothing sets INSTANTTENSOR_CACHE_BUFFER by default — but anyone enabling it to speed up repeated opens gets UB at exit or on entry reuse.

Suggested fix

Make the deleter self-contained: capture by value only what it needs, e.g. backend and, for URING, do the ring deregistration in destroy_buffer (before put()) rather than in the deleter — deregistering the buffer from the loader's ring must happen while the loader is alive anyway, so the cached entry's deleter can reduce to [](void *ptr){ CUDA_CHECK(cudaHostUnregister(ptr)); free(ptr); } with no captures.

Verification

Build with ASan (or run under valgrind), set INSTANTTENSOR_CACHE_BUFFER=1, do one safe_open/close cycle with an AIO_BUFFERED/URING_BUFFERED backend, and exit the process: the atexit cleanup() → ~HostBufferCache → entry.deleter(ptr) chain reports a read of freed memory at the this->backend load in the lambda (loader_common.cpp:109).

[malaiwah]

Independent second-pass re-review of both items by a reviewer who did not author them, against the same tree (85e7c5f5); every quoted line was re-read and the arithmetic re-derived from scratch. One overclaim in item 1 is retracted below.

Claim	Verdict	Note
1. Default io_depth = 512//world_size ⇒ 4 GiB device + 1 GiB pinned host per rank at ws=4	HOLDS-WITH-CORRECTION	Byte math reproduces exactly (4,295,028,784 B device / 1,073,741,824 B host). Two corrections: the pinned host buffer is not BUFFERED-specific, and the #13 linkage overclaims (below).
1a. max_free_mem_usage guard cannot moderate it on a healthy GPU	HOLDS	Shrink fires only when free < 8 GiB at ws=4; re-derived independently.
1b. Buffers are a load-phase peak, not a leak (env unset)	HOLDS	Loader::close → destroy_buffer frees both; verified the frontend always reaches it via __exit__.
1c. "the size being requested [in #13] comes from this formula"	RETRACT	#13's failing cudaMalloc was post-shrink (~288 MiB), not the 4 GiB default. Same line, different magnitude — no causal link.
2. Cached pinned-buffer deleter captures a destroyed Loader's this ⇒ UAF	HOLDS	Ownership trace completed exhaustively (below); the lambda really does dereference this. Sharpened: it is a UAF or a data race with ~Loader, and it is not URING-specific.
2a. Suggested fix (capture-free deleter, deregister in destroy_buffer)	HOLDS — and it is order-safe, for a reason worth stating explicitly: Loader::close runs destroy_buffer() before close_file().

Reachability, for the record: weight_utils.py::instanttensor_weights_iterator calls safe_open(files, framework="pt", device=..., process_group=...) with no chunk_size/concurrency/io_depth/buffer_size/max_free_mem_usage, so every default here is live at ws=4. Both items are load-phase or opt-in; nothing is inside a captured CUDA graph, so there is no per-step amortization argument either way.

Claim 1 — math re-derived, and one correction

Independent recomputation, chunk=8 MiB, concurrency=1, ws=4, io_depth=max(512//4,3)=128:

• rank_chunk_size = 8 MiB, world_chunk_size = 32 MiB
• device: 128 × 32 MiB = 4,294,967,296 B, plus padded_size = 3*(16 + 4096 + 4096*1*4) = 61,488 B → 4,295,028,784 B
• pinned host: 128 × 8 MiB = 1,073,741,824 B
• guard: 8 MiB × 1 × 128 × 4 = 4 GiB > free×0.5 ⇔ free < 8 GiB. Confirmed.

Also confirmed that io_depth is genuinely the binding in-flight cap, so the sizing is self-consistent rather than over-provisioned (this was the first thing I tried to disprove):

chunk_id_t wait_chunk_id = max(chunk_id - (chunk_id_t)MAX_PREFETCH_CHUNKS, chunk_id - (chunk_id_t)this->io_depth);

MAX_PREFETCH_CHUNKS == 1024 (common.hpp:36), so the effective depth is min(1024, io_depth) = io_depth. Lowering io_depth therefore lowers real in-flight work, not just an unused reservation — which is exactly why it needs an A/B rather than being a free win.

Correction (scope): the title and the "(BUFFERED backends)" parenthetical are too narrow. need_host_buffer = true is set unconditionally in both open_file_aio (loader_io_aio.cpp:22) and open_file_uring (loader_io_uring.cpp:134), i.e. for AIO/URING (O_DIRECT) as well as the *_BUFFERED variants. The 1 GiB pinned staging exists on the upstream default path (default_backend = [URING, AIO]), not only when INSTANTTENSOR_BACKEND=BUFFERED is set.

Retraction (#13 linkage): #13's log shows both shrinks having already fired (io_depth 128→3, concurrency 16→3) before the cudaMalloc at loader_common.cpp:79 failed. Working the shrink formulas backwards from those two warnings gives avail_bytes ∈ [288, 384) MiB and a request of 8 MiB × 3 × 3 × 4 = 288 MiB — two orders of magnitude below the 4 GiB this claim is about. So #13 is not evidence for the default-depth footprint; it is evidence about the guard's floor and its stale snapshot: at ws=4 the floor is 8 MiB × 1 × 3 × 4 = 96 MiB device + 24 MiB pinned, which no INSTANTTENSOR_MAX_FREE_MEM_USAGE value can go below (hence 0.3 not helping), and mem_get_info() is sampled in _determine_io_params while the cudaMalloc happens later in _open(). Sibling note: #13's concurrency=16 is not reachable from this commit at ws=4 (the AIO/URING branch defaults to 1), so that report was produced by a different revision or an explicit INSTANTTENSOR_CONCURRENCY. The honest statement is "same allocation site, different magnitude, and it points at the floor + snapshot staleness" — not "the size comes from this formula".

The verification recipe in the issue is accurate, incidentally — loader_common.cpp:407 does print device_buffer_size= and host_buffer_size= under INSTANTTENSOR_DEBUG=1.

Claim 2 — ownership trace, completed

I tried to find any owner of the Loader that outlives the cache entry, and there is none:

• Whole-tree search for the type: the only object of type Loader anywhere is the stack local at loader_common.cpp:603. No shared_ptr<Loader>, no Loader* member, no static/global registry, no manager map entry.
• LoaderManager::loader_handle_to_queue stores RPCHandle{int, SPSCQueue<RPCRequest>*, SPSCQueue<RPCResponse>*} (main.cpp:6-10) — raw queue pointers, whose ownership lives in the Loader's unique_ptr members, and the map entry is erased in close() anyway.
• All internal threads that capture this (worker_threads, cuda_thread, last_page_reader_thread, wait_thread, io_depth_sample_thread) are joined in destroy_threads(), which Loader::close runs before destroy_buffer() — so at put() time no live thread holds this either. The captured this in the cached deleter is the sole surviving reference.
• The lambda does dereference it: this->backend is a load from the object, and deregister_host_buffer_uring() reads this->uring_register_buffer and passes &this->uring_ring. HostBufferCacheEntry::deleter is std::function<void(void*)> (types.hpp:122), so the move into the cache preserves the captured pointer verbatim. Built with -std=c++17 (setup.py:85), so the implicit-this capture in [=] is not even deprecated — no compiler diagnostic will catch this; ASan will.

Three refinements to the filed lifecycle:

1. It is a race, not only a stale read. _env_cache_buffer() is a function-local static (common.hpp:109-112), so a process either caches for its whole life or never does — the deleter is therefore invoked exactly once, in ~HostBufferCache, always carrying the first cache-missing loader's this. But ~HostBufferCache is reached from cleanup() on the main thread (main.cpp:135-139, atexit-registered at _impl.py:28), while the loader thread returns from run() only after pushing the CLOSE response. Nothing orders ~Loader before cleanup(). So the deleter can run on an object that is dead, or being destroyed concurrently. Both are UB; the second is also unreproducible.
2. Not URING-only. The deleter is created in shared code, so the this->backend read happens on the AIO and MMAP/in-memory paths too (loader_io_inmem.cpp:23 also sets need_host_buffer). The URING branch is the one with teeth: io_uring_unregister_buffers(&this->uring_ring) issues io_uring_register() with a ring_fd read out of freed thread-stack memory — glibc's thread-stack cache makes it likely that the read succeeds with reused garbage rather than faulting, so the worst case is a register op aimed at whatever fd number that garbage happens to be, not a clean crash.
3. Cosmetically, "the deleter is only assigned on a cache miss" is right and is the reason the entry cannot heal itself on reuse: the second loader re-registers the cached buffer into its own ring at loader_common.cpp:119-121 but never replaces the deleter.

Fix-safety notes

Claim 2 fix — safe, and I'd add one reason to prefer it. Loader::close orders destroy_threads() → destroy_buffer() → close_file() (loader_common.cpp:413-421), and destroy_uring_context() (the io_uring_queue_exit) lives in close_file(). So doing the ring deregistration inside destroy_buffer() before put() runs while the ring is still alive and all worker threads are already joined — correct on both the cached and uncached paths. Independently of the UAF, this is the only semantically right place: the registration is a property of this loader's ring, and each reusing loader re-registers the buffer in its own ring, so per-loader deregistration must not be deferred into a process-global cache at all. Resulting [](void *ptr){ CUDA_CHECK(cudaHostUnregister(ptr)); free(ptr); } is capture-free, matches aligned_alloc/cudaHostRegister on the base pointer (which is what the entry stores, and cache hits may hand back a larger entry — so unregistering by base pointer, not by requested size, remains correct), no numerics, no ordering hazard, no CUDA-graph interaction (load-phase only). One adjacent hazard the fix does not address and probably should not try to: if cleanup() never runs (os._exit, hard abort), ~HostBufferCache reaches cudaHostUnregister during static destruction after the CUDA context may be gone; that failure mode predates this issue.

Claim 1 fix — defensible, but adaptive is the right shape, and one invariant to check.

• The ≥3 floor must survive any clamp (aio read + cudaMemcpyAsync + ncclAllGather is a real 3-stage pipeline); clamp(512//ws, 3, 32) does. Executor queues are statically MAX_PREFETCH_CHUNKS=1024, so only raising io_depth could ever collide there.
• Device bytes shrink only while the io term dominates. buffer_size = max(recommended_buffer_size_for_tensors(...), chunk*conc*depth*ws) (_impl.py:599-601); the tensor-derived term is a floor the library sets deliberately for load/consume overlap and already returns max(buffer_size, max_tensor_size) (_impl.py:301), which is what keeps compute_layout's tensor_size > buffer_size throw unreachable. So a clamp cannot break that invariant, but it also cannot save VRAM below that floor — the pinned host GiB, by contrast, shrinks strictly linearly with io_depth. Worth stating in the fix so nobody expects a guaranteed 4 GiB saving.
• On the throughput question: for URING_BUFFERED every SQE is forced async (sqe->flags |= IOSQE_ASYNC, loader_io_uring.cpp:286, with the comment explaining why), and the SQ is sized io_depth * num_threads in io_uring_queue_init. Since the kernel derives the bounded io-wq worker cap from the SQ size and CPU count ([INFERENCE], io_init_wq_offload: min(sq_entries, 4*num_online_cpus())), depth 128 authorizes ~128 io-wq workers per rank × 4 ranks all doing 8 MiB page-cache copies. Deep queues here plausibly cost throughput rather than buy it — but that cuts the other way for URING/AIO with O_DIRECT against a wide NVMe array, which is precisely the case a fixed constant would regress. Hence: prefer sizing the default from mem_get_info() up front (e.g. target ≤1 GiB of device staging, floor 3) over a hardcoded 32, and keep INSTANTTENSOR_IO_DEPTH as the escape hatch.
• What to measure (both cheap, no new instrumentation needed): (a) INSTANTTENSOR_DEBUG=1 already prints Average io_depth from the sampler thread (loader_common.cpp:430) — if the achieved average is far below 128, the queue depth is nominal and the footprint is pure waste; (b) load wall clock at INSTANTTENSOR_IO_DEPTH ∈ {8, 16, 32, 128} on both a page-cache-warm and an O_DIRECT cold run, since those two regimes are the ones that disagree.

Priority triage

• Trivial + safe, land as-is: claim 2's capture-free deleter plus moving deregister_host_buffer_uring() into destroy_buffer(). No behaviour change with the flag off, removes UB with it on.
• Needs measurement before changing the default: claim 1. The footprint numbers are solid; the choice of cap is not, and the O_DIRECT/wide-array case is the one that can regress. The Average io_depth line makes this a one-run experiment.
• Drop: the sentence linking Speculative decoding draft model fails during load due to memory allocation constraints #13's OOM size to this formula, and the "BUFFERED backends" scoping of the pinned buffer. Both are wrong as written; the underlying observations survive without them.