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 every context on a device from
one buffer, sized to the largest engine's requirement. 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().

That default path is byte-identical to a binary built from origin/main over 18 runs,
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:

• four execution contexts of one engine holding two fp32 8-head attention blocks over
  [1,2048,512]: 1188MB → 372MB, 3 × 272MB reclaimed
• one Method holding six one-block engines interleaved with six CUDA delegates:
  1656MB → 316MB, 5 × 268MB reclaimed

Outputs are identical between the two modes in both cases. The commit message carries the
full rationale, the ordering argument, and the two consequences a caller feels once the
option is on.

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

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.

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

• set_option is not unit-tested. No target in tests/cpp/executorch/ links the
  backend. Three behaviours live only there — skipping a foreign key, storing a valid
  boolean, rejecting a non-boolean. The store is exercised by the memory A/B; the other two
  are covered nowhere. CudaBackend's equivalents are equally uncovered. Closing this needs
  a new backend-linked cc_test, which I have not added.
• Concurrent execute() on one device is not covered, and the code says so. The
  requirement is that all delegate enqueues on a device are ordered on one stream; a
  default-on version would need the scratch keyed per stream rather than one buffer per
  device.
• 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.

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.

[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.

[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.

[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.

[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.

[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.

[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.

[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.