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

cpp/BUILD

@@ -191,6 +191,28 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "tensorrt_executorch_shared_scratch_pool",
+    hdrs = [
+        "include/torch_tensorrt/executorch/SharedScratchPool.h",
+    ],
+    strip_include_prefix = "include",
+    target_compatible_with = select({
+        ":linux_x86_64": [],
+        ":sbsa": [],
+        "//conditions:default": ["@platforms//:incompatible"],
+    }),
+    deps = select({
+        ":linux_x86_64": [
+            "@cuda//:cuda_headers",
+        ],
+        ":sbsa": [
+            "@cuda//:cuda_headers",
+        ],
+        "//conditions:default": [],
+    }),
+)
+
 cc_library(
     name = "tensorrt_executorch_backend",
     srcs = [
@@ -211,6 +233,7 @@ cc_library(
     deps = [
         ":tensorrt_executorch_binding_names",
         ":tensorrt_executorch_blob_header",
+        ":tensorrt_executorch_shared_scratch_pool",
         ":tensorrt_executorch_weight_streaming_budget",
     ] + select({
         ":linux_x86_64": [
@@ -254,6 +277,7 @@ filegroup(
 filegroup(
     name = "executorch_api_headers",
     srcs = [
+        "include/torch_tensorrt/executorch/SharedScratchPool.h",
         "include/torch_tensorrt/executorch/TensorRTBackend.h",
         "include/torch_tensorrt/executorch/TensorRTBindingNames.h",
         "include/torch_tensorrt/executorch/TensorRTBlobHeader.h",

cpp/include/torch_tensorrt/executorch/SharedScratchPool.h

@@ -0,0 +1,183 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+// Bookkeeping for the TensorRT backend's shared per-device activation-scratch
+// pool: the grow/reuse policy, the enqueue-handoff rule, and the lock that scopes
+// both to a single device.
+// Allocation and event creation arrive as callables rather than being made here.
+
+#include <cuda_runtime.h>
+
+#include <cstddef>
+#include <mutex>
+#include <unordered_map>
+
+namespace torch_tensorrt {
+namespace executorch_backend {
+
+// Runtime backend option that backs execution-context activation scratch with a
+// shared per-device pool instead of giving every context its own. Boolean,
+// default false. Delivered as
+//   executorch::runtime::set_option("TensorRTBackend", options.view())
+// A context's allocation strategy is fixed when the context is created, so a
+// later call governs only the engines loaded after it, and a pooled context and
+// a private-scratch one coexist in one process.
+inline constexpr char kSharedActivationScratchKey[] = "use_shared_activation_scratch";
+
+// Per-device handoff marker for the shared scratch buffer: the pool-owned CUDA
+// event that the last enqueue against the buffer was recorded on.
+struct SharedScratchMarker {
+  cudaEvent_t event = nullptr; // never destroyed; one event serves the slot for the process lifetime
+  bool pending = false; // an enqueue against the buffer has been recorded on `event`
+};
+
+// What a caller about to enqueue against a device's shared scratch has to do:
+// when `needs_wait`, make its stream wait on `event` first; once the enqueue is
+// submitted, record it on `event`. `event` is null only when the slot has no
+// event and one could not be created.
+struct SharedScratchHandoff {
+  cudaEvent_t event = nullptr;
+  bool needs_wait = false;
+};
+
+// One device's shared scratch buffer and the marker ordering its handoff, behind
+// the lock that covers both.
+//
+// A claimant holds `mu` from the wait on the previous enqueue through its own
+// enqueue and the record of that enqueue on the marker. Holding it that far is
+// what makes the marker a complete account of who is using the buffer. Anything
+// less leaves an enqueue live in a window the marker does not cover, and a
+// claimant entering that window is handed the same buffer with nothing ordering
+// the two. `mu` covers one device, so a growth holds no lock a claim on
+// another device has to acquire.
+struct SharedScratchDevice {
+  std::mutex mu;
+  void* buffer = nullptr;
+  std::size_t capacity = 0;
+  SharedScratchMarker marker;
+};
+
+// Holds one SharedScratchDevice per device id.
+//
+// `get` locks only long enough to find or create the entry, and the reference it
+// returns stays usable once that lock is dropped: std::unordered_map keeps
+// references to elements valid across rehashing, and entries are never erased.
+// This one lock is shared by every device, which is why nothing but the lookup
+// runs under it.
+class SharedScratchPool {
+ public:
+  SharedScratchDevice& get(int device_id) {
+    std::lock_guard<std::mutex> lk(mu_);
+    return devices_[device_id];
+  }
+
+ private:
+  std::mutex mu_;
+  std::unordered_map<int, SharedScratchDevice> devices_;
+};
+
+// Claims a device's handoff for a caller about to enqueue against its shared
+// scratch, creating the marker's event on first use. Call with `dev.mu` held.
+//
+// `create_event` returns a CUDA event, or nullptr if one could not be created,
+// in which case the slot stays empty and the next call retries.
+//
+// The ordering between one enqueue and the next is carried by an event rather
+// than by the stream the previous enqueue used, because a stream handle cannot
+// carry it: synchronizing on a handle whose stream the caller has since
+// destroyed is a crash rather than an error return, CUDA recycles handle values
+// so a genuinely different stream can compare equal to the recorded one, and the
+// NULL stream is both a legal stream a caller can select and the only available
+// "no previous user" sentinel. An event names the work instead of the queue --
+// it stays valid after the stream that recorded it is destroyed, and waiting on
+// it from the stream that recorded it is already satisfied, so the common
+// single-stream case costs a host call and no device stall.
+template <typename CreateEvent>
+SharedScratchHandoff shared_scratch_claim_event(SharedScratchDevice& dev, CreateEvent create_event) {
+  if (dev.marker.event == nullptr) {
+    dev.marker.event = create_event();
+  }
+  // A slot with no event is never marked, so a failed creation reports nothing to
+  // wait for rather than a wait the caller has no event to perform.
+  return {dev.marker.event, dev.marker.pending};
+}
+
+// Call with `dev.mu` held.
+//
+// The mark precedes the record, so a failed record leaves the slot claiming an
+// enqueue the event does not cover -- the caller must then synchronize the stream
+// itself before returning the error.
+inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) {
+  if (dev.marker.event != nullptr) {
+    dev.marker.pending = true;
+  }
+  return dev.marker.event;
+}
+
+// A buffer a growth replaced, handed back for the caller to dispose of.
+//
+// A non-null `wait_for` is the marker's event, on which an enqueue that may still
+// be reading and writing `buffer` has been recorded; the caller must wait for
+// that event on the host before it frees, and must do so while it still holds
+// `dev.mu`. Once the lock is dropped the next claimant records its own enqueue on
+// the same event, and a wait made then would block on work that never touched
+// this buffer. A null `wait_for` means nothing was ever recorded against it.
+//
+// One event covers every enqueue the buffer ever served, but only because each of
+// them claims the handoff before enqueueing -- which orders its stream after the
+// event -- and records on the event afterwards, so the latest recording completes
+// only once all the earlier ones have. An enqueue that reaches the buffer without
+// doing both is covered by no wait here.
+//
+// The free itself belongs outside `dev.mu`: on CUDA it is a device-wide
+// synchronization, so performing it under the lock makes an unrelated claim on
+// this device wait for every stream on it.
+struct RetiredScratch {
+  void* buffer = nullptr;
+  cudaEvent_t wait_for = nullptr;
+};
+
+// Bookkeeping for a device's scratch buffer, which grows monotonically to the
+// largest requested size. Call with `dev.mu` held.
+//
+// `alloc` returns nullptr on failure; the buffer is then left untouched.
+// Allocating before releasing is what makes that true, and it costs peak
+// residency: while the buffer grows, the old and the new one are both resident.
+//
+// A growth reports the buffer it displaced through `out_retired`; see
+// RetiredScratch for what the caller owes it. Nothing is freed here, so a caller
+// that ignores `out_retired` leaks rather than frees a buffer an enqueue may
+// still be using.
+template <typename Alloc>
+void* shared_scratch_get_or_grow(
+    SharedScratchDevice& dev,
+    std::size_t need,
+    std::size_t& out_size,
+    Alloc alloc,
+    RetiredScratch& out_retired) {
+  if (dev.buffer != nullptr && dev.capacity >= need) {
+    out_size = dev.capacity;
+    return dev.buffer;
+  }
+  void* p = alloc(need);
+  if (p == nullptr) {
+    return nullptr;
+  }
+  if (dev.buffer != nullptr) {
+    out_retired.buffer = dev.buffer;
+    out_retired.wait_for = dev.marker.pending ? dev.marker.event : nullptr;
+  }
+  dev.buffer = p;
+  dev.capacity = need;
+  out_size = need;
+  return p;
+}
+
+} // namespace executorch_backend
+} // namespace torch_tensorrt

cpp/include/torch_tensorrt/executorch/TensorRTBackend.h

@@ -75,6 +75,14 @@ struct EngineHandle {
   size_t num_aliased_outputs = 0;
   int device_id = 0;
   bool unified_memory = false;
+  // Whether exec_ctx was created kUSER_MANAGED and draws its activation scratch
+  // from the shared per-device pool (kSharedActivationScratchKey,
+  // SharedScratchPool.h).
+  bool shared_scratch = false;
+  // The activation scratch the engine itself reports needing, read at init when
+  // shared_scratch is set. execute() needs it to tell a failed per-call query,
+  // which TensorRT also reports as zero, from an engine that genuinely needs none.
+  size_t engine_scratch_bytes = 0;
   std::mutex mu;
   // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or
   // destroying an execution context while one of its enqueues is in flight, so when
@@ -102,13 +110,26 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface {
   // past return, order any other stream against this one, and synchronize the stream
   // before reading device-resident outputs. The selected stream must be on the engine's
   // device, and calls on one handle must not overlap each other or its destruction.
+  // With the shared activation scratch pool (kSharedActivationScratchKey) one
+  // buffer per device backs every context created while the option was on. Calls
+  // on two such handles on one device may overlap: a per-device lock held across
+  // the enqueue serializes them, so they do not run concurrently on the device. A
+  // handle whose context was created while the option was off keeps its own
+  // scratch and is not subject to this.
   // Note that other CUDA delegates sharing the same guard may instead synchronize before
   // returning, so do not assume results are ready on return from this one.
   ::executorch::runtime::Error execute(
       ::executorch::runtime::BackendExecutionContext& context,
       ::executorch::runtime::DelegateHandle* handle,
       ::executorch::runtime::Span<::executorch::runtime::EValue*> args) const override;
 
+  // Applies the runtime backend options a caller passes to
+  // executorch::runtime::set_option("TensorRTBackend", ...). The only key read is
+  // kSharedActivationScratchKey (SharedScratchPool.h), a boolean.
+  ::executorch::runtime::Error set_option(
+      ET_UNUSED ::executorch::runtime::BackendOptionContext& context,
+      const ::executorch::runtime::Span<::executorch::runtime::BackendOption>& backend_options) override;
+
   void destroy(::executorch::runtime::DelegateHandle* handle) const override;
 };
 

cpp/src/torch_tensorrt/executorch/README.md

@@ -85,6 +85,17 @@ the removed `CudaStreamGuard`:
   complete, order any cross-stream producers/consumers with their own events,
   and synchronize the stream before reading outputs on the host.
 - With no guard active, the backend falls back to `cudaStreamPerThread`.
+- With the `use_shared_activation_scratch` backend option enabled, one buffer
+  per device backs the activation scratch of every execution context created
+  while it was on, so no two enqueues against it may overlap. The backend
+  enforces this itself: it holds a per-device lock from the claim on the buffer
+  through the enqueue and the completion event recorded on it, so two
+  `execute()` calls on one device are serialized at submission and the second's
+  stream waits on the first's enqueue. They may run on one stream or on two, and
+  they may be submitted concurrently from two threads — but they will not run
+  concurrently on the device, so the pool costs the parallelism between them.
+  Contexts created while the option was off keep their own scratch and are
+  unaffected.
 - The reference-runner smoke test runs inference inside a caller-stream guard on
   the discrete-GPU CI configuration, where all inputs and outputs are host-backed
   and therefore take the synchronized staging path. CI separately asserts that the
@@ -104,6 +115,52 @@ the removed `CudaStreamGuard`:
   asynchronous return described above is still uncovered and the interaction
   between a green context and the internal completion event remains untested.
 
+## Shared activation scratch
+
+A TensorRT execution context allocates its own activation scratch and holds it
+for as long as the context lives, so a model lowered to N single-layer engines
+pays N copies and can run out of device memory on the layer count alone. The
+`use_shared_activation_scratch` backend option — a boolean, off by default —
+instead backs every context on a device from one buffer, grown to the largest
+engine's requirement:
+
+```cpp
+#include <executorch/runtime/backend/interface.h>
+
+executorch::runtime::BackendOptions<1> options;
+options.set_option("use_shared_activation_scratch", true);
+executorch::runtime::set_option("TensorRTBackend", options.view());
+```
+
+Check what `executorch::runtime::set_option` returns: `Error::NotFound` means no
+backend is registered under that name, which is what a binary that has not linked
+the backend archive gets.
+
+N per-engine copies collapse to one, so the reclaimed memory is the sum of the N
+requirements less the largest of them. Set the option before loading the methods
+whose engines should use the pool, and read the `use_shared_activation_scratch`
+bullet of the caller-stream contract above: engines sharing a buffer do not run
+concurrently on the device. The pool never returns memory to the
+device, so the largest scratch it was ever asked for stays allocated until the
+process exits.
+
+The buffer grows when an engine asks for more than every engine before it did,
+and a growth is not free. It frees the buffer it replaces, and `cudaFree` waits
+for everything queued on the device, not only for the enqueues that used that
+buffer, so it can stall for far longer than the event wait that precedes it. The
+backend keeps that stall out from under the per-device lock, so it does not hold
+up another engine on the device, but it does fall after the growing call's own
+enqueue, so that one `execute()` waits for its own engine work too. A 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 reduces the pool to a
+single allocation.
+
+How much any one engine asks for is fixed when it is built, not when it runs.
+The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature makes an engine
+report what the shapes just bound need; without it, whether an engine does that
+or reports its profile maximum depends on how TensorRT planned it. Either way the
+pool can settle well above the live data, and nothing the runtime does changes it.
+
 ## Standalone Backend Archive
 
 Use this path only when you need `libexecutorch_trt_backend.a` without building