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 of those engines asks 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.

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 is byte-identical to a binary built from origin/main in all nine comparisons (three models, three repetitions each), including canonicalised stderr — the real regression risk, given the feature ships off.

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 is never freed, 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.

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 no overlapping hunks.

Where I would spend review attention

This is three 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; and a third answering the second — the device lock held across the enqueue, the retired-buffer handling, and a concurrency regression test.

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. TensorRT answers a failed updateDeviceMemorySizeForShapes() and an engine that genuinely needs no activation scratch with the same value, and setDeviceMemoryV2(nullptr, 0) is itself rejected and returns nothing to test — so a context handed a zero keeps whatever buffer it last held, which a pool growth may already have freed. Forcing that path against a freed buffer reproduces an illegal memory access. Each engine therefore records what ICudaEngine::getDeviceMemorySizeV2() reports at its own init, and a zero fails the execute() only when that recorded requirement is non-zero. An engine that needs no scratch is given no buffer at all, so it has nothing to claim and nothing for the next claimant to order against.

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.

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. 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. 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. Growth happens only on an engine's first run and only for an engine larger than every engine before it, so loading the largest engine first avoids it entirely.
• 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 skips silently without one. 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, and the two-thread concurrency regression. Without a device it skips all twelve of its tests and exits zero, and the workflow's --test_output=errors keeps the skip reason out of the log, so a green run on a device-less host says nothing about the pool. What reddens that job if the runner loses its device is the reference export later in the same step, which calls .cuda().
• One mutant of the zero-scratch guard survives every test. Deleting the guard while keeping the init capture passes all twelve backend tests, because provoking it needs a failed updateDeviceMemorySizeForShapes(), which cannot be induced from inside the process. It dies only to an out-of-tree fault-injecting probe. The two neighbouring mistakes are covered: a guard that fires on every zero is killed by AnEngineNeedingNoActivationScratchRunsWithThePoolEnabled, and a missing init capture by EachEngineRecordsItsOwnActivationScratchRequirement.
• 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