perf(executorch): opt-in shared per-device activation-scratch pool

[Conarnar]

Description

A multi-layer model lowered to the ExecuTorch TensorRT delegate becomes N separate single-layer engines, and by default every execution context allocates its own activation scratch and holds it for as long as the context lives. Device memory therefore scales with the layer count, and multi-layer models OOM at runtime on the layer count alone.

This adds an opt-in shared per-device pool that backs the activation scratch of every context created while the option is on from one buffer, grown to the largest figure any call on that device has asked for. It ships disabled: with use_shared_activation_scratch unset, each context keeps its private kSTATIC scratch, no extra log output is emitted, and the only added work is one relaxed atomic load per engine init and a bool test or two per execute(). A context created while the option was off keeps its own scratch and is outside the pool entirely, and so is one whose engine reports needing no activation scratch under any shape.

How much any one engine asks for is decided when the engine is built, not when it runs. Whatever updateDeviceMemorySizeForShapes() answers is binding rather than advisory — setDeviceMemoryV2 refuses a smaller buffer, and an engine backed by less than it asked for writes past the end — but whether an engine reports what the shapes just bound need or reports its profile maximum depends on how TensorRT planned it. The builder's PreviewFeature::kRUNTIME_ACTIVATION_RESIZE_10_10 produces the former; without it either can happen. So the pool can settle well above the live data, and nothing a runtime does changes that.

That default path was compared against a binary built from the commit this PR branches from, over nine runs (three models, three repetitions each): canonicalised stderr identical every time, at 45 / 45 / 17 lines. That stderr is the whole of the evidence — this harness writes nothing to stdout — and the canonicaliser normalises timestamps, pointers, source line numbers and the harness's own device-wide memory reading, which moves by a tenth of a megabyte between two runs of the same binary. The control, the same binary with the option on, differs at 47 lines against 45, so the comparison can fail.

Measured on one 80GB A100 with TensorRT 11.2.1.2 and CUDA 13, reading cudaMemGetInfo after a cudaFree(0) baseline. N per-engine copies collapse to one, so what is reclaimed is the sum of the N per-engine requirements less the largest of them:

• four execution contexts of one engine holding two fp32 8-head attention blocks over [1,2048,512]: 1188MB → 372MB. Uniform engines, so that sum-less-largest is 3 × 272MB.
• one Method holding six one-block engines of the same shape interleaved with six CUDA delegates: 1656MB → 316MB. Also uniform: 5 × 268MB.
• the same shape of Method with its six engines at differing sizes: 792MB → 316MB. Sum 780,140,544 B less the largest 281,018,368 B is 476.0MB, which is what the A/B reads.

Outputs are identical between the two modes in every case.

Two consequences a caller feels once the option is on. The pool never shrinks — a growth does free the buffer it replaces, but the high-water mark never comes down — so a device keeps the largest scratch it was ever asked for until the process exits, where per-context kSTATIC scratch is released with its context. And a growth allocates the new buffer before releasing the old one, so both are resident for that moment — that ordering is what leaves the existing buffer usable when an allocation fails. A third, the device-wide wait that same growth's cudaFree makes, is the first of the known gaps below; it is the one a caller is most likely to feel, and it is documented rather than fixed.

No dependencies; this does not stack on anything. It is orthogonal to weight streaming (#4336), which targets engine weight memory rather than activation scratch. It composes with the zero-copy KV work if that lands too — measured together on the same model, with byte-identical generated ids, and the two branches still merge without a conflict.

Where I would spend review attention

This is eight commits: the pool itself; a second answering the first review — the zero-scratch guard, per-device locking in place of one process-wide mutex, and a backend-linked test target; a third answering the second — the device lock held across the enqueue, the retired-buffer handling, and a concurrency regression test; a fourth answering the third — the pool header out of the installed set, and the growth-frequency claims corrected; a fifth — cudaEventBlockingSync on the handoff event and three prose corrections; a sixth — the zero-answer sizing below, the stream-capture refusal, the cudaDeviceReset() note, and a fixture that empties the pool between tests, with a unit test for the hook it does that through; a seventh — the install of that buffer read back so one TensorRT refused never reaches enqueueV3, the capture guard moved ahead of everything a capture cannot take, the pool made a leaked singleton, and the reset hook's disposer taken out from under the locks; and an eighth — the capture advice corrected (it had said loading with the option off is enough to capture, which is not what CUDA does), TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1 actually passed by the CI invocation, the pool's two test-only entry points moved into a testonly translation unit so a release build exports neither, and a discriminating test for each of the three fixes the seventh commit had shipped without one.

The enqueue handoff, not the allocation. The pool itself is simple; the ordering is the part with teeth. Contexts share one buffer while their enqueues can still be in flight, so each device slot carries a pool-owned cudaEvent_t: wait on it before enqueueing, record after. An earlier revision tracked the last stream instead and was wrong three ways — synchronizing a destroyed stream segfaults rather than returning an error, CUDA recycles stream handle values so two distinct streams compare equal, and the NULL stream is a legal caller stream indistinguishable from "no previous user". All three are structural with an event, and ~EngineHandle in the same file had already made this choice and documented why.

A reported zero is ambiguous. updateDeviceMemorySizeForShapes() answers zero for three different things: a query that failed, a call whose bound shapes need no activation scratch — a valid empty input on a dynamic profile, which an earlier revision rejected outright — and an engine that needs none under any shape. Only the last of the three can be told apart before any call is made, and it is told apart once, at init: each engine records whether ICudaEngine::getDeviceMemorySizeV2() is non-zero, and one that reports zero is left out of the pool entirely, so it claims no per-device lock and does not serialize against the engines that do.

The other two arrive at execute() indistinguishable where the query is read, and neither can be handed nothing. setDeviceMemoryV2(nullptr, 0) is itself refused (condition: memory != nullptr || expectedSize == 0), so a context handed a zero keeps whatever buffer it last held, which a pool growth may already have freed — forcing that path reproduces an illegal memory access. And enqueueV3 refuses a kUSER_MANAGED context with no device memory installed as soon as the engine reports needing some under any shape, however little the shapes actually bound need: Parameter check failed, condition: noDeviceMemory. So neither is asked to size the pool: a zero is handed whatever the pool already holds, at its current size, which grows nothing. Only against a pool holding nothing does it allocate, and then one byte, leaving the first call with a real requirement to size the pool. The obvious alternative, standing the engine's own getDeviceMemorySizeV2() in for the zero, satisfies enqueueV3 too but covers the whole profile: on the dynamic engine the tests build that is 4 MiB against the 128 KiB the next call needs, and since the pool never shrinks a single empty batch would pin it at 32x for the process lifetime — on the path the pool exists to keep small.

Only one of those two is safe with a buffer that small, and the install is what separates them. Where the bound shapes genuinely need nothing the engine expects nothing, and one byte is accepted. Where the query failed the engine still expects what it always did, setDeviceMemoryV2 refuses the undersized buffer — and it returns void, so a caller that does not ask cannot tell the two apart. Measured on TensorRT 11.2.1.2 against this branch's own two-softmax net (expecting 8388608 bytes): install one byte and enqueueV3 returns true, cudaStreamSynchronize reports no error, the output matches byte for byte, and the engine writes the buffer the context was last given. Free that buffer, let the next cudaMalloc take the address, and 4090 of 4096 bytes of an unrelated allocation are overwritten with nothing reporting an error anywhere. So the install is read back, through a TensorRT IErrorRecorder attached for the duration of that one call: a correct install records nothing, an undersized one records kINVALID_ARGUMENT naming both sizes, and the genuinely-empty case — expected size zero, one byte installed — records nothing and keeps working. A refusal ends the call with Error::InvalidState and nothing reaches enqueueV3. The check runs on every pooled install, not only on those following a zero, so "no path enqueues on an install TensorRT rejected" is a property of the code rather than of an argument about when a zero can occur — and because the refusal is observable, the engine's profile-wide figure never has to be installed and the pool is still never pinned at the profile maximum.

Two hazards the pool does not absorb. Capturing a CUDA graph from the stream a pooled execute() runs on is refused with Error::NotSupported. The handoff between one enqueue and the next is a wait on an event recorded outside the capture, and cudaStreamWaitEvent on such an event fails with cudaErrorStreamCaptureIsolation and invalidates the capture under every capture mode, cudaStreamCaptureModeRelaxed included — the caller would learn of that only from cudaStreamEndCapture returning 901 and a null graph far from the cause. That wait, not the allocation, is what produced the 901 when the guard was removed and measured: the pool was already large enough, so no growth ran. A growth is mode-dependent by comparison — its cudaMalloc, its cudaEventSynchronize and its cudaFree each return cudaErrorStreamCaptureUnsupported and invalidate the capture under the Global and ThreadLocal modes, while Relaxed permits all three and then runs them uncaptured, leaving a replay pointed at a buffer the pool may since have freed. Measured directly against CUDA 13.0 on this device, one process per cell, rather than read off the documentation. The check is one cudaStreamIsCapturing on the selected stream, and it runs ahead of everything execute() does that a capture cannot take — not merely ahead of the pool's own calls, since the host wait on a previous enqueue and the cudaMalloc that grows a host-input staging buffer come before those and, outside Relaxed, would invalidate the capture first. Only the device query and the device switch run earlier, and a capture takes both. The pageable copy that follows the staging allocation is not one of the calls the guard has to get ahead of: on its own it is captured cleanly at 4 KiB and 64 MiB in all three modes, and replaying execute()'s own order it returns cudaErrorStreamCaptureInvalidated rather than the cudaErrorStreamCaptureUnsupported a prohibited call returns — that is the error for touching a capture the two calls before it had already killed, not for being one of them.

Loading the engine with the option off takes the pool's calls out of the way, and the refusal with them, but not the rest of what a capture cannot take. Under Global and ThreadLocal the wait on a previous enqueue and the staging cudaMalloc are still prohibited, and under every mode, Relaxed included, so is the cudaStreamSynchronize that ends any call staging through host memory, aliasing an output, or running with no caller stream — a capturing stream cannot be synchronized at all. So a call with the option off captures cleanly when it runs on the capturing stream under a CallerStreamGuard, binds only non-empty device-resident tensors, aliases no output, and follows no call that left an enqueue in flight. That is measured through the delegate rather than argued, and so is its consequence: such a call leaves its own enqueue in flight, so under Global and ThreadLocal the next one on the same handle fails the last condition. With the option off a handle is good for one captured call, not for capture. Two tests pin the two outcomes. Moving the guard out of the pooled branch, so that the refusal covered this path too, would refuse exactly the calls that do capture, for contexts that never touch the pool — mutating it to be unconditional fails both of those tests — so it stays where it is and the README carries the rule.

What the check cannot cover is a capture live on some other stream: under Global, and under ThreadLocal from the calling thread, the same calls invalidate that capture and the call is not refused. Asking cudaStreamIsCapturing(cudaStreamLegacy) instead does not close the gap — measured, it reports a capture on another stream only when that stream is a blocking one, and then reports it under Relaxed and under another thread's ThreadLocal too, where the pool's calls are permitted and refusing would be wrong. CUDA offers no "is a capture live in this process" query, so the rule is the caller's to keep, and the installed header and the README say that rather than promising more. cudaDeviceReset() is documented rather than guarded, in the installed header and the README: the pool holds its buffer and its handoff event for the process lifetime and a reset destroys both, and catching that would mean revalidating both on every call, which costs as much as the work it protects.

The lock scope, and how far it now reaches. The registry that finds a device's entry has one lock, held for a single find-or-insert with no CUDA call under it. The device's own lock is held from the claim through setDeviceMemoryV2, the enqueue, and the record of that enqueue on the handoff event — releasing it any earlier leaves an enqueue live in a window the event does not yet cover, and a second claimant entering that window is handed the same buffer with nothing ordering the two. That produced silently wrong output on every trial before the fix. The lock nests inside the per-handle EngineHandle::mu, which already spans the enqueue, and is never taken in the other order. Entries are never erased and std::unordered_map keeps references valid across rehashing, which is what lets the registry lock be dropped before the entry is used.

The per-handle capture. A context's allocation strategy is fixed at creation, so each EngineHandle records the setting in effect at its own init() and execute() consults that, never the global. That is what lets a later set_option govern only subsequent engines and lets pooled and private-scratch contexts coexist, with no freeze and no rejected calls.

SharedScratchPool.h is not installed. It was in executorch_api_headers for no reason: TensorRTBackend.h never included it and named the option key only in comments, so the installed set carried a mutex, a per-device registry and three functions whose only locking rule is a sentence, to deliver one string constant. The key now sits beside set_option, and the header moved to cpp/src/torch_tensorrt/executorch/. Dropping it from the package outright would have broken the standalone archive — libtorchtrt.tar.gz ships TensorRTBackend.cpp as buildable source and that source includes the header — so it moves rather than leaves. Every header remaining in cpp/include/ is installed, which is the invariant the original placement broke. The same question applies to symbols, not just headers: the pool's two test-only entry points — one of which frees the live pool with no wait for work in flight — were defined in TensorRTBackend.cpp and so exported from the archive the README tells a consumer to build. They now live in a SharedScratchPoolTestHooks translation unit that only a testonly Bazel target compiles, and nm/readelf on the release object shows neither symbol.

third_party/cuda/BUILD gains a target, the one file outside the delegate. The header needs the cudaEvent_t typedef — a compile-time dependency, not a runtime one — and the repo had no headers-only CUDA target. Depending on cudart instead put libcudart in the DT_NEEDED of a host-side test that makes no CUDA call, and broke it with exit 127.

Known gaps

• A growth stalls on everything queued on the device, and that is documented rather than fixed. The wait before freeing a replaced buffer is on the per-device handoff event, but cudaFree performs its own device-wide synchronization, so the free waits for every stream on the device and not only for the enqueues that used that buffer. It is the only call on the growth path that does: with a host function parked on a stream neither the pool nor the growing engine had ever used, the cudaMalloc for the new buffer and the wait on the retired one each returned at once and the cudaFree returned after 3000 ms, exactly when that host function was released. The free is issued outside the device lock, so it does not hold up another engine's claim, but it does fall after the growing call's own enqueue — that one execute() waits for its own engine work. The wait has no upper bound and a caller can make it permanent: if some stream on the device is waiting on work only this thread submits after execute() returns, the growing call never comes back, and the thread that would unblock it is the one inside cudaFree. The backend's own tests deadlocked on exactly that once, when a change of test order made a case that parks a host function on its own stream be the one that grew the pool. Neither fix on offer is local. cudaFreeAsync would bound the wait to the buffer's own users but pairs only with cudaMallocAsync, so it moves every pool allocation onto the stream-ordered allocator — the redesign an earlier round considered and rejected. Deferring the free to a point where the wait is already paid has no such point: the two calls before it on the same path each return immediately, so deferral moves the wait rather than bounding it. So the installed header and the README state it as a caller obligation instead. How often it happens depends on the engine. execute() re-queries the requirement on every call once the shapes are bound, so an engine whose answer does not vary with those shapes grows the pool at most once, on its first run — a program built only from those settles after its largest engine has run once. An engine whose answer does vary can grow it on any call needing more than every call before it, and with one of those in the program nothing bounds the number of allocations. Note this is run order, not load order: loading an engine allocates nothing.
• The option cannot be enabled from Python or from a .pte. It arrives only through the C++ set_option; there is no load-time runtime spec and no compile-spec fallback of the kind weight_streaming_budget has. Deferred rather than built here.
• The backend-linked test needs a real GPU, and the job that would give it one rarely runs. tests/cpp/executorch/test_shared_scratch_backend.cpp is the only target in that package that links the delegate, and it is what covers set_option, the per-engine capture, the pooled path, the install check, the capture refusal and the two-thread concurrency regression. Without a device it skips all nineteen of its cases and exits zero, which Bazel reports as a passing target: measured with CUDA_VISIBLE_DEVICES='', 19 of 19 skipped and exit 0. The suite prints how many cases it skipped and why, and TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1 turns a skip into a failure — same binary, same empty device list, exit 1 with 19 failures. The CI invocation now passes that variable, and the job it sits in already asks for a GPU runner and starts its container with --gpus all. What keeps these cases off most runs is not the runner but the lane gate on the whole ExecuTorch job, below.
• The zero-answer path is covered, but not by inducing a real query failure. AnEmptyInputRunsWithThePoolEnabled reaches it through a legitimate empty input on a dynamic profile and reads the pool's capacity at three points, pinning the first of them to the one-byte minimum exactly, which kills a zero handed no buffer at all, a zero sized from the engine's profile-wide figure, and a minimum raised to anything the next call's requirement would still bound. AnUndersizedScratchInstallIsSeenAsRefused then drives the same TensorRT call execute() makes over a live kUSER_MANAGED context: a correct install is accepted, a one-byte install for an engine expecting more is refused, and the empty-batch context still accepts one byte. A genuine failed updateDeviceMemorySizeForShapes() cannot be induced from inside the process — the size execute() installs is the pool's capacity, which is never below what the query just returned — so what is pinned end to end is the check that catches its consequence rather than the cause itself. That execute() still goes through the checked install is pinned separately by ExecuteInstallsPooledScratchThroughTheCheckedHelper, which leaves a counting IErrorRecorder on the delegate's own context and asserts it is replaced and restored exactly once per pooled run; replacing the checked call with a bare setDeviceMemoryV2 takes both counts to zero and no other case notices. What no test pins is the refusal reaching a caller as Error::InvalidState, because the refusal cannot be reached. The neighbouring mistakes are covered too: an engine kept in the pool when it needs no scratch under any shape is killed by AnEngineNeedingNoActivationScratchRunsWithThePoolEnabled, and a missing init capture by EachEngineRecordsWhetherItNeedsActivationScratch.
• A failed cudaEventSynchronize on a replaced buffer fails the call, and that path is not tested. It leaks the buffer rather than freeing it under a possibly-live enqueue, and returns Error::InvalidProgram instead of enqueueing against a context on a device already in a faulted state. Neither a failed event wait nor a failed cudaFree can be induced from inside the suite without corrupting the device for every case after it.
• cudaDeviceReset() is documented, not tested. Reproducing it would destroy the primary context for every test that ran after it in the same process, so the hazard is argued in the installed header and the README rather than exercised. The stream-capture refusal beside it is tested, twice: APooledEngineRefusesToRunWhileItsStreamIsCapturing under Relaxed, and APooledEngineRefusesACaptureBeforeAnythingCanInvalidateIt under the default Global, which is the one that can see where in execute() the refusal comes. Two more cover the path the refusal points a caller at: AnUnpooledEngineCapturesOnAFirstCallThatStagesNothing and AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight, both under Global. What none of the four covers is Relaxed with the option off, where the wait on a previous enqueue is permitted and runs outside the graph.
• The ExecuTorch runtime CI job has been skipped on every run at this head, gated on a channel that concurrency kept cancelling, and the gate ignores skipped channels — so every number here was local. tests/cpp/executorch carries 4 cc_test targets and 40 cases at the merge base with main; this PR takes it to 6 targets and 79 cases, the two new targets being test_shared_scratch_pool (20) and test_shared_scratch_backend (19). All 6 targets and all 79 cases have been compiled and run under bazel on a real A100 with --nocache_test_results and with TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1 set, so a skip would have been a failure; nothing skipped and all pass. The lane itself still needs one clean run.
• The stream-handle hazards are argued, not reproduced end to end. The destroyed-stream crash was reproduced through the delegate; handle recycling and the NULL-stream collision were measured at the CUDA level. Corruption from a missing wait was never reproduced through a real engine, on either design.
• Not measured: the ~210-engine target model (the saving is linear by construction and the growth policy is exercised), and weight streaming combined with the pool end to end.

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[shoumikhin]

took a proper look at this, most of it against 11.2.1.2 since that's the pin in MODULE.bazel. the design holds up and the event handoff is doing real work. one blocking thing on the zero path, rest is smaller stuff.

two things i chased that turned out to be fine, noting them so nobody else burns time on them: the pool does cover weight streaming scratch (updateDeviceMemorySizeForShapes tracks getDeviceMemorySizeV2 exactly, scratch included, checked with the budget moved around), and a caller stream from a green context records on the per-device event without complaint and actually orders the work. no concerns on either.

[shoumikhin]

this treats the error sentinel as a real answer, and the argument in the comment only covers the first call.

tested against 11.2.1.2:

• first call with nothing set, enqueueV3 does refuse. so the comment is right about that case.
• setDeviceMemoryV2(nullptr, 0) is itself rejected ("Cannot set memory to nullptr"), and it returns void, so the failure is invisible to us.
• on a later call the context silently keeps its previous pointer and enqueueV3 returns true.

that previous pointer can be freed memory. once another engine grows the pool you cudaFree the old buffer, so a spurious 0 here runs the engine against a dead allocation. i reproduced it: run once with a good buffer, free it the way the release lambda does, let an unrelated cudaMalloc take the address, then hit the zero path. all 4194304 bytes of the unrelated allocation got overwritten, and enqueueV3 still returned true with a correct output.

the zero branch also skips get_or_grow_shared_scratch entirely, so you lose the wait and the in-flight mark in the same step.

simplest fix is to return an error on 0.

[Conarnar]

Fixed.
It is actually possible for the engine to require 0 activation scratch so an extra check for that case is also needed.

[shoumikhin]

"the exact scratch requirement for this call" is only true with kRUNTIME_ACTIVATION_RESIZE_10_10 on, and nothing here enables it (we only set MULTIDEVICE_RUNTIME_10_16).

measured on 11.2.1.2: preview off, a query at batch 64 under a profile max of 256 returns exactly the profile-max size. preview on, same engine returns 8192 at batch 1 vs 33554432 at batch 4096.

not a safety issue since it oversizes, but the pool ends up sized to the profile max rather than to the call, and that's most of the savings story.

[Conarnar]

Fixed

[shoumikhin]

scratch_pool_mu is one mutex for all devices and it's held for the whole function, including cudaMalloc and the release lambda's cudaDeviceSynchronize + cudaFree.

so a growth on device 0 blocks a plain claim on device 1, which only needs its own map slot. the README says concurrent execute on different devices is fine, and that stops being true during a growth. the sync is unbounded as well, it waits on everything queued on the device, not just the scratch users.

per-device lock would fix both scopes, or move the cuda calls out from under the map lock.

[Conarnar]

Fixed.

[shoumikhin]

this only holds when every engine needs the same amount. the pool grows to max(s_i) and never shrinks, so the saving is sum(s_i) - max(s_i). for engines needing 1, 2 and 4 units that's 3, not 8.

the numbers in the description used uniform engines so it wouldn't show up there. the commit message has the same claim.

[Conarnar]

Fixed.

[shoumikhin]

both helpers mutate the caller's map with no locking, and the header never says the caller has to serialize. the backend gets away with it by holding scratch_pool_mu, but this header goes out in executorch_api_headers, so it's API and the next caller won't know.

either document the precondition, or wrap the maps in a type that owns the lock, or keep the header private.

[Conarnar]

Fixed, with your first two suggestions. The header could not be made private since TensorRTBackend.h needs it for set_option.

[shoumikhin]

name says stores nothing, but pool[device_id] default-inserts before alloc runs, so there is an entry, just one with a null pointer. the test only checks the return value and the retry, so it passes either way.

the header comment ("the slot is then left untouched") says the same thing. EventCreationFailureIsReportedAndRetried has the identical gap on the markers map. either narrow the names or assert pool.empty() and don't insert until the alloc succeeds.

[Conarnar]

Fixed.

[shoumikhin]

this only depends on tensorrt_executorch_shared_scratch_pool, never on the backend, so nothing here covers the wiring. i can revert the context to kSTATIC, or delete the updateDeviceMemorySizeForShapes/setDeviceMemoryV2 pair, or drop the wait/record calls, and all 11 tests stay green.

the description calls out the missing set_option and concurrency tests, but not that the plain single-threaded path has no automated coverage at all.

[Conarnar]

Fixed.

[shoumikhin]

while you're in this header, the execute() comment just above still ends with "calls on one handle must not overlap each other or its destruction". the pool adds a stronger rule (no two handles on the same device may overlap) and that only lives in the README right now. this is the installed header, so it should carry it.

[Conarnar]

Fixed

[shoumikhin]

The pool design reads well. Three things.

Two pooled engines running at once on one device give wrong numbers, silently. The device lock is released when get_or_grow_shared_scratch returns, so setDeviceMemoryV2 and enqueueV3 run without it. The new enqueue is only recorded on the event afterwards. A second thread claiming in that gap gets the same buffer. That thread does wait on the event, but the event carries an earlier enqueue, not the one now in flight. So the wait does not order the two against each other, and both engines write the same scratch.

I tried this with two real engines sharing one buffer. On every trial one engine's output was wrong, with no CUDA error and no TensorRT error. The written caveat points at growing the pool and freeing the buffer. This needs neither, so it is the normal state once the pool has settled.

Holding the device lock from the claim through the enqueue and the event record fixed it every time. It looks safe here: engine->mu is already held across enqueueV3 today, and core/runtime/execute_engine.cpp already puts a mutex around the enqueue and says the other context calls belong in the same scope.

The option cannot be turned on from Python or from a .pte. It only arrives through the C++ set_option. weight_streaming_budget handles this in the same file with a load-time runtime spec plus a compile-spec fallback.

The release lambda frees the old buffer even when the cudaEventSynchronize before it failed, which is the one case the wait exists to prevent.

One note: the ExecuTorch CI jobs are skipped on this commit because the matrix job was cancelled, so the two new test files have not run yet.

Happy to share the harness or the measurements.

[shoumikhin]

This lock ends when the function returns, so setDeviceMemoryV2 and enqueueV3 run without it, and the enqueue is only recorded on the event after that. Nothing else serializes two pooled engines on one device, since engine->mu is per handle.

So a second thread can claim inside that gap and get the same buffer back, because the reuse branch returns the existing buffer untouched whenever the capacity fits. It does call cudaStreamWaitEvent, but the event carries an earlier enqueue rather than the one now in flight, so the wait does not order the two against each other. Both engines then write the same scratch.

I tried this with two real engines sharing one buffer. On every trial one engine's output was wrong, with no CUDA error and no TensorRT error. Nothing grows and nothing is freed, so this is the ordinary state once the pool has settled on its largest size. The caveat in the comment above, and in the README, points at growth freeing a live buffer, which is the case that did not corrupt in my testing: cudaFree is implicitly synchronizing, and I measured it blocking for the whole remaining kernel. So the documented hazard is the survivable one and this one is not mentioned.

In a harness following the same call order, holding the claim through the enqueue and the event record fixed it every time. That means execute() owning the device lock across setDeviceMemoryV2, enqueueV3 and the mark, not a local change here. It looks safe: engine->mu is already held across enqueueV3 in execute() today, and core/runtime/execute_engine.cpp already puts a mutex around the enqueue and says the other context calls belong in the same scope. Worth checking the cost for a weight streaming engine, where the header says enqueueV3 becomes synchronous.

[Conarnar]

Fixed. Checked the weight-streaming cost you flagged: at budget 0, enqueueV3 took 0.439 ms against a 75.7 ms inference, and a second context's enqueue issued mid-stream returned in 0.267 ms. Didn't reproduce.

[shoumikhin]

This global is the only way in, and the only writer is the C++ set_option, so a program loaded from Python has no way to turn the pool on. A hand written compile spec for the key would not help either, since init() only looks for the weight streaming key.

weight_streaming_budget in this same file takes a load time runtime spec first and falls back to a compile spec baked in at export, and the comment there explains that the compile spec exists for loaders that cannot pass backend options yet. The new option has neither channel.

Not a blocker, since the C++ path works and the feature is off by default. But the users who hit the memory problem this solves are often the ones loading a .pte.

[Conarnar]

Deferred. There's no Python binding for set_option so for this to be done in a follow-up, the only route would be through compile spec during export.

[shoumikhin]

The result is discarded and cudaFree(old) runs on the next line either way. If the wait fails, the free is the exact thing the wait exists to prevent, so an enqueue may still be reading that buffer. Returning early instead would leak that one replaced buffer, which is a bounded cost.

Also worth knowing that no test reaches this path: every pooled engine in the backend test asks for the same size, so the reuse branch is always taken and the pool never grows. One extra engine with a larger shape would cover both the free and the wait.

[Conarnar]

Fixed. The larger engine covers the allocation and the free, but not the wait — cudaFree synchronizes device-wide anyway, so deleting the wait still leaves the test green. The comment says so rather than claiming it.

[shoumikhin]

Swapping the device sync for an event wait does not remove the device-wide wait, because cudaFree performs one itself. On an H100 with the handoff event already complete and unrelated work queued on another stream, the event wait returned in under a millisecond and this cudaFree took about 1.5 seconds, matching the queued work. It also runs with the device lock held, so a claim on that device waits behind unrelated work too.

Either move the free outside the lock and keep a retire list, or say in the commit message and the README that a growth still stalls on everything queued on the device.

[Conarnar]

Fixed

[shoumikhin]

Some notes inline.

[shoumikhin]

The line that says a growth happens only on an engine's first run does not hold for a dynamic shape engine, and the paragraph three lines below it says so itself. The per call query answers for the shapes just bound. I built an engine with a real dynamic profile and a user managed context, no preview feature, and the same engine asked for 33554432, 67108864, 100663296 and 134217728 bytes as the batch went from 1 to 4. That is three growths from one engine, so loading the largest engine first does not bound the allocations.

This matters beyond the text. The comment above the leak on a failed pre free wait says the cost is bounded because growth is rare, and the readme tells a caller the device wide free stall is avoidable by load order. Neither is true here. Could you either scope both claims to fixed shape engines, or say plainly that a dynamic shape engine can grow on any call?

[Conarnar]

Fixed, though I was not able to replicate your findings. Growth depends on whether an engine's reported requirement varies with the bound shapes, so I scoped the README by that rather than by fixed versus dynamic shape.

[shoumikhin]

The reason given for keeping this header public does not hold in the code as it stands. TensorRTBackend.h does not include it. It names the option key in three comments and nowhere else, and it compiles without it. The readme example passes the string literal, and the new backend test spells the key out on purpose.

So what ships to users, for one string constant, is a mutex, a per device registry, and three functions whose only lock rule is a sentence saying to call them with the device lock held. There is no assert and no annotation, and the lock and the state are separate public members so the type cannot enforce it. Fourteen of the sixteen cases in the pool unit test call those helpers with no lock at all, which is the first thing a new reader will copy.

Moving the key constant into TensorRTBackend.h next to set_option, and taking the pool header out of the installed set, closes this without adding anything.

[Conarnar]

Fixed, though SharedScratchPool.h was moved to cpp/src/torch_tensorrt/executorch/ as libtorchtrt.tar.gz ships TensorRTBackend.cpp as buildable source, which requires the header.

[shoumikhin]

The two new tests only run inside the ExecuTorch runtime job, and that job is skipped in every run at this head. It is gated on the standard channel, and that channel was cancelled by concurrency. The required gate only fails on a real failure, so skipped channels do not block anything.

That means every test number in the description is a local number, and the first time either file gets compiled in CI would be the push to main. Could you run the full lane on this branch once before merging?

Worth pairing with the device less skip you already documented. Right now a run that skips all twelve cases and a run that never happened both look green.

[Conarnar]

Ran //tests/cpp/executorch:executorch_backend_tests under bazel on a real device, at the pushed head and again with this stack applied. 6/6 both times, including the two new files.

[Conarnar]

Pushed two more commits.

The zero-scratch guard rejected a valid input. With the pool on, a dynamic engine given an
empty input (batch 0 inside a [0..128] profile, all dimensions specified) has
updateDeviceMemorySizeForShapes() legitimately answer 0 while the engine's own requirement is
non-zero. That is the shape the guard treated as a failed query, so execute() returned
Error::InvalidState and the engine never ran, on an input the backend otherwise supports (it
allocates one-byte placeholders for empty tensors). Reproduced on TensorRT 11.1, 10.13.3 and
11.2.1.2, with and without kRUNTIME_ACTIVATION_RESIZE_10_10.

A zero has three causes: a query that failed, a call whose bound shapes need no scratch, and an
engine that needs none under any shape. Only the last can be told apart, and it is told apart once
at init: such an engine is left out of the pool entirely, so it takes no per-device lock and does
not serialize against the engines that do. The other two arrive at execute() indistinguishable,
and both are handed whatever the pool already holds, at its current size, which grows nothing;
any live buffer is large enough for a call that needs none.

Neither can be handed nothing at all. enqueueV3 refuses a kUSER_MANAGED context with no device
memory installed as soon as the engine reports needing some under any shape, however little the
bound shapes need: Parameter check failed, condition: noDeviceMemory. So against a pool holding
nothing the call allocates a one-byte minimum and the first call with a real requirement sizes the
pool. Standing the engine's own getDeviceMemorySizeV2() in for the zero satisfies that check too,
but that figure covers the whole profile: on the dynamic engine the tests build it is 4 MiB against
the 128 KiB the next call needs, and since the pool never shrinks a single empty batch would pin it
at 32x for the process lifetime, on the path the pool exists to keep small.

Capturing a CUDA graph from a pooled execute()'s stream is now refused with
Error::NotSupported. The handoff between one enqueue and the next is a wait on an event recorded
outside the capture, and cudaStreamWaitEvent on such an event fails with
cudaErrorStreamCaptureIsolation and invalidates the capture under every capture mode,
cudaStreamCaptureModeRelaxed included. That wait, not the allocation, is what produced the 901
and the null graph from cudaStreamEndCapture when the guard was removed and measured; the pool
was already large enough that no growth ran. A growth is mode-dependent by comparison: its
cudaMalloc, cudaEventSynchronize and cudaFree each invalidate the capture under Global and
ThreadLocal, while Relaxed permits all three and runs them uncaptured, leaving a replay pointed at
a buffer the pool may since have freed. Measured directly on CUDA 13.0, one process per cell,
rather than read off the documentation. The check is one cudaStreamIsCapturing ahead of the
device lock; loading the engine with the option off still captures cleanly.

cudaDeviceReset() is documented rather than guarded, in the installed header and the README.
The pool holds its buffer and its handoff event for the process lifetime and a reset destroys the
primary context under both. Catching it would mean revalidating both on every call, on the pool's
hot path, and the check costs about what it protects. The header's execute() contract also now
states that a growth turns an otherwise asynchronous call into a device-wide cudaFree wait.

The handoff event passes cudaEventBlockingSync. Created with cudaEventDisableTiming alone,
cudaEventSynchronize busy-waits, measured at 99.5% of a core on an H100 versus 0.1% with the
flag, and that wait runs while the per-device lock is held, so one core spins and every other
engine on the device queues behind it. The engine-completion event in the same file already passed
the flag. The README example also handles set_option's Error, which is not ET_NODISCARD and
so compiled while dropping it, immediately above a paragraph telling the caller to check exactly
that return.

Two test defects, both of which weakened the suite rather than the backend. Nothing reset the
process-wide pool between tests, so the suite passed in declaration order and in no other: running
the empty-input case before ASecondPooledEnqueueWaitsForTheFirstOnAnotherStream turned the
latter's first run into a growth, whose cudaFree waits device-wide and so waited on the blocking
host function that test parks on its own stream, from the one thread that could release it. A
deadlock ending at the 60 s watchdog. --gtest_repeat=2 on the growth case failed on its own
second iteration, blaming the backend for a fixture artifact. A fixture now resets the pool in
setup and teardown through test-only hooks declared in SharedScratchPool.h, which ships under
src/ and is not installed, so no production semantics and no installed API changed. Separately,
device_bytes_in_use() returned zero on a failed cudaMemGetInfo; both memory comparisons
subtract two of its readings, and a zero for the first of a pair made the second look like the
whole cost of what was measured between them. It now reports the failure and the call sites assert
on it.

tests/cpp/executorch is 6 targets and 70 cases, run on a real device: 69 at the pushed head and
70 with the last commit, the extra being a unit test for the reset hook. Two tests that existed
only to pin a now-unused engine-wide byte count were replaced by one pinning the predicate the
pooled path branches on.