Add TensorRT weight streaming support to the ExecuTorch delegate

Apply a TensorRT weight streaming budget in the delegate init(), after the engine
is deserialized and before the execution context is created (the budget cannot be
changed once a context exists). When the engine was built with weight streaming,
the delegate applies TensorRT's automatic budget by default, mirroring the PyTorch
runtimes. An explicit non-negative byte budget can be set via
torch_tensorrt.save(output_format="executorch", weight_streaming_budget=N) and is
carried as an ExecuTorch CompileSpec.

Non-streamable engines make no budget call, so existing programs are unchanged.
The one intended behavior change is that a streamable engine now applies the
automatic budget on load, which enables running models whose weights exceed GPU
memory.

Ref #4334

Author: shoumikhin

cpp/BUILD

@@ -96,6 +96,17 @@ cc_library(
     strip_include_prefix = "include",
 )
 
+cc_library(
+    name = "tensorrt_executorch_weight_streaming_budget",
+    srcs = [
+        "src/torch_tensorrt/executorch/WeightStreamingBudget.cpp",
+    ],
+    hdrs = [
+        "include/torch_tensorrt/executorch/WeightStreamingBudget.h",
+    ],
+    strip_include_prefix = "include",
+)
+
 cc_library(
     name = "tensorrt_executorch_binding_names",
     hdrs = [
@@ -127,6 +138,7 @@ cc_library(
     deps = [
         ":tensorrt_executorch_binding_names",
         ":tensorrt_executorch_blob_header",
+        ":tensorrt_executorch_weight_streaming_budget",
     ] + select({
         ":linux_x86_64": [
             "@executorch//:executorch_headers",
@@ -150,6 +162,7 @@ filegroup(
         "src/torch_tensorrt/executorch/README.md",
         "src/torch_tensorrt/executorch/TensorRTBackend.cpp",
         "src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp",
+        "src/torch_tensorrt/executorch/WeightStreamingBudget.cpp",
     ],
 )
 
@@ -169,5 +182,6 @@ filegroup(
         "include/torch_tensorrt/executorch/TensorRTBackend.h",
         "include/torch_tensorrt/executorch/TensorRTBindingNames.h",
         "include/torch_tensorrt/executorch/TensorRTBlobHeader.h",
+        "include/torch_tensorrt/executorch/WeightStreamingBudget.h",
     ],
 )

cpp/include/torch_tensorrt/executorch/WeightStreamingBudget.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace torch_tensorrt {
+namespace executorch_backend {
+
+// Compile spec key that carries the weight streaming budget from export into the
+// delegate. Must match WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY on the Python
+// side (py/torch_tensorrt/executorch/partitioner.py).
+inline constexpr char kWeightStreamingBudgetKey[] = "weight_streaming_budget";
+
+// Result of parsing a weight streaming budget compile spec value.
+//
+// The value is a non-negative decimal integer: an explicit GPU budget in bytes.
+// The automatic budget is intentionally not encoded here: when no budget spec is
+// present, the delegate applies TensorRT's automatic budget itself, mirroring the
+// PyTorch runtimes.
+//
+//   valid == false -> a value was present but could not be parsed (or was
+//                     negative); the caller should reject the program.
+//   valid == true  -> bytes holds the parsed non-negative byte budget.
+struct WsBudget {
+  bool valid = false;
+  int64_t bytes = 0;
+};
+
+// Parses a budget from a raw byte range. The value is NOT assumed to be NUL
+// terminated; only the first nbytes bytes are read.
+WsBudget parse_weight_streaming_budget(const void* value, std::size_t nbytes);
+
+} // namespace executorch_backend
+} // namespace torch_tensorrt

cpp/src/torch_tensorrt/executorch/CMakeLists.txt

@@ -23,6 +23,7 @@ find_package(Threads REQUIRED)
 set(_torchtrt_executorch_sources
     "${CMAKE_CURRENT_LIST_DIR}/TensorRTBackend.cpp"
     "${CMAKE_CURRENT_LIST_DIR}/TensorRTBlobHeader.cpp"
+    "${CMAKE_CURRENT_LIST_DIR}/WeightStreamingBudget.cpp"
 )
 
 add_library(executorch_trt_backend STATIC ${_torchtrt_executorch_sources})

cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp

@@ -8,6 +8,7 @@
 #include "torch_tensorrt/executorch/TensorRTBackend.h"
 #include "torch_tensorrt/executorch/TensorRTBindingNames.h"
 #include "torch_tensorrt/executorch/TensorRTBlobHeader.h"
+#include "torch_tensorrt/executorch/WeightStreamingBudget.h"
 
 #include <cstdint>
 #include <cstring>
@@ -225,8 +226,6 @@ Result<DelegateHandle*> TensorRTBackend::init(
     BackendInitContext& context,
     FreeableBuffer* processed,
     ArrayRef<CompileSpec> compile_specs) const {
-  (void)compile_specs;
-
   TORCHTRT_ET_CHECK_NOT_NULL(processed, Error::InvalidArgument, "TensorRTBackend::init: null processed buffer");
   TORCHTRT_ET_CHECK_NOT_NULL(processed->data(), Error::InvalidArgument, "TensorRTBackend::init: null processed buffer");
 
@@ -284,6 +283,71 @@ Result<DelegateHandle*> TensorRTBackend::init(
   TORCHTRT_ET_CHECK_NOT_NULL(
       handle->engine, Error::InvalidProgram, "TensorRTBackend::init: failed to deserialize TensorRT engine");
 
+  // Apply the weight streaming budget. This must run before the execution
+  // context is created below: TensorRT forbids changing the budget while a
+  // context is active. The budget is carried from export as an optional
+  // "weight_streaming_budget" compile spec whose value is a non-negative decimal
+  // byte budget; negative or malformed values are rejected as InvalidProgram.
+  // When no spec is present and the engine supports streaming, we apply
+  // TensorRT's automatic budget, which is what the PyTorch runtimes do on
+  // deserialize.
+  const CompileSpec* ws_spec = nullptr;
+  for (const auto& spec : compile_specs) {
+    if (spec.key != nullptr && std::strcmp(spec.key, kWeightStreamingBudgetKey) == 0) {
+      if (ws_spec != nullptr) {
+        // The budget must appear at most once; a second match means the spec
+        // list is inconsistent, so reject the program instead of guessing.
+        ET_LOG(Error, "TensorRTBackend::init: duplicate weight_streaming_budget compile spec");
+        return Error::InvalidProgram;
+      }
+      ws_spec = &spec;
+    }
+  }
+
+  WsBudget ws_request;
+  if (ws_spec != nullptr) {
+    ws_request = parse_weight_streaming_budget(ws_spec->value.buffer, ws_spec->value.nbytes);
+    if (!ws_request.valid) {
+      ET_LOG(Error, "TensorRTBackend::init: malformed weight_streaming_budget compile spec");
+      return Error::InvalidProgram;
+    }
+  }
+
+  const int64_t streamable = handle->engine->getStreamableWeightsSize();
+  if (streamable > 0) {
+    // getStreamableWeightsSize is > 0 only when the engine was built with
+    // BuilderFlag::kWEIGHT_STREAMING.
+    const bool is_explicit = ws_spec != nullptr;
+    int64_t budget;
+    if (is_explicit) {
+      // An explicit budget is a non-negative byte count, clamped to the
+      // streamable size (TensorRT also caps it, but clamp for a clear log).
+      budget = ws_request.bytes > streamable ? streamable : ws_request.bytes;
+    } else {
+      budget = handle->engine->getWeightStreamingAutomaticBudget();
+    }
+    if (!handle->engine->setWeightStreamingBudgetV2(budget)) {
+      if (!is_explicit && handle->engine->setWeightStreamingBudgetV2(0)) {
+        // The automatic budget could not be applied; fall back to maximum
+        // streaming (budget 0), which always fits.
+        ET_LOG(Info, "TensorRTBackend::init: automatic weight streaming budget failed; using maximum streaming");
+      } else {
+        ET_LOG(Error, "TensorRTBackend::init: setWeightStreamingBudgetV2(%lld) failed", (long long)budget);
+        return Error::InvalidProgram;
+      }
+    }
+    ET_LOG(
+        Info,
+        "TensorRTBackend::init: weight streaming budget=%lld streamable=%lld scratch=%lld",
+        (long long)handle->engine->getWeightStreamingBudgetV2(),
+        (long long)streamable,
+        (long long)handle->engine->getWeightStreamingScratchMemorySize());
+  } else if (ws_spec != nullptr) {
+    // A budget was requested but the engine cannot stream. The engine is still
+    // valid and runs fully resident, so log and continue rather than fail.
+    ET_LOG(Info, "TensorRTBackend::init: weight_streaming_budget ignored; engine was not built with enable_weight_streaming=True");
+  }
+
   Error err = initialize_engine_io(*handle);
   if (err != Error::Ok) {
     return err;

cpp/src/torch_tensorrt/executorch/WeightStreamingBudget.cpp

@@ -0,0 +1,35 @@
+#include "torch_tensorrt/executorch/WeightStreamingBudget.h"
+
+#include <charconv>
+#include <cstddef>
+#include <cstdint>
+#include <system_error>
+
+namespace torch_tensorrt {
+namespace executorch_backend {
+
+WsBudget parse_weight_streaming_budget(const void* value, std::size_t nbytes) {
+  WsBudget result; // valid == false until the value is fully parsed
+
+  if (value == nullptr || nbytes == 0) {
+    return result;
+  }
+  // The value is a non-negative decimal byte budget and is not NUL terminated.
+  // std::from_chars consumes only ASCII digits (no leading whitespace, sign, or
+  // base prefix), so a leftover byte (trailing garbage or an embedded NUL), an
+  // out-of-range value, or a negative leaves the result invalid.
+  const char* const first = static_cast<const char*>(value);
+  const char* const last = first + nbytes;
+  int64_t parsed = 0;
+  const std::from_chars_result fc = std::from_chars(first, last, parsed);
+  if (fc.ec != std::errc() || fc.ptr != last || parsed < 0) {
+    return result;
+  }
+
+  result.valid = true;
+  result.bytes = parsed;
+  return result;
+}
+
+} // namespace executorch_backend
+} // namespace torch_tensorrt