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,123 @@
+/*
+ * 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/per-device policy and the enqueue-handoff rule.
+// Allocation and event creation arrive as callables rather than being made here.
+
+#include <cuda_runtime.h>
+
+#include <cstddef>
+#include <unordered_map>
+#include <utility>
+
+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;
+};
+
+// Claims a device's handoff for a caller about to enqueue against its shared
+// scratch, creating the marker's event on first use.
+//
+// `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(
+    std::unordered_map<int, SharedScratchMarker>& markers,
+    int device_id,
+    CreateEvent create_event) {
+  SharedScratchMarker& marker = markers[device_id];
+  if (marker.event == nullptr) {
+    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 {marker.event, marker.pending};
+}
+
+// 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(std::unordered_map<int, SharedScratchMarker>& markers, int device_id) {
+  SharedScratchMarker& marker = markers[device_id];
+  if (marker.event != nullptr) {
+    marker.pending = true;
+  }
+  return marker.event;
+}
+
+// Bookkeeping for a per-device pool of device-memory buffers that grows
+// monotonically to the largest requested size.
+//
+// `alloc` returns nullptr on failure; the slot is then left untouched.
+// Allocating before releasing is what makes that true, and it costs peak
+// residency: while a slot grows, the old and the new buffer are both resident.
+// `release` must leave no in-flight enqueue pointing at the buffer it frees --
+// the CUDA caller syncs the device first.
+template <typename Alloc, typename Release>
+void* shared_scratch_get_or_grow(
+    std::unordered_map<int, std::pair<void*, std::size_t>>& pool,
+    int device_id,
+    std::size_t need,
+    std::size_t& out_size,
+    Alloc alloc,
+    Release release) {
+  auto& slot = pool[device_id];
+  if (slot.first != nullptr && slot.second >= need) {
+    out_size = slot.second;
+    return slot.first;
+  }
+  void* p = alloc(need);
+  if (p == nullptr) {
+    return nullptr;
+  }
+  if (slot.first != nullptr) {
+    release(slot.first);
+  }
+  slot = {p, need};
+  out_size = need;
+  return p;
+}
+
+} // namespace executorch_backend
+} // namespace torch_tensorrt

cpp/include/torch_tensorrt/executorch/TensorRTBackend.h

@@ -75,6 +75,10 @@ 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;
   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
@@ -109,6 +113,13 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface {
       ::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,14 @@ 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 an enqueue against that buffer must not overlap another one. The
+  backend orders consecutive enqueues itself, whether they run on one stream or
+  on two. What it cannot order is two `execute()` calls submitted concurrently on
+  one device: the caller must submit them one at a time, whether or not they
+  share a stream. 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 +112,34 @@ 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 `(N-1)` times the
+per-engine scratch. 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 first: the pool carries an ordering obligation the
+backend cannot discharge for you. The buffer is never released, so the device
+keeps the largest scratch it was ever asked for until the process exits.
+
 ## Standalone Backend Archive
 
 Use this path only when you need `libexecutorch_trt_backend.a` without building