Add TensorRT weight streaming support to the ExecuTorch delegate

Set the weight streaming budget in the delegate init, after the engine is deserialized and before the execution context is created. The budget is carried from export as a weight_streaming_budget compile spec with values auto, disabled, or a byte count. There is no change to the pte format and no change to ExecuTorch. When no budget value is present the runtime makes no budget call, so existing programs behave exactly as before.

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,70 @@ 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 decimal int64
+  // (negative disables streaming, non-negative is an explicit byte budget). 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 not built with WEIGHT_STREAMING");
+  }
+
   Error err = initialize_engine_io(*handle);
   if (err != Error::Ok) {
     return err;

cpp/src/torch_tensorrt/executorch/WeightStreamingBudget.cpp

@@ -0,0 +1,53 @@
+#include "torch_tensorrt/executorch/WeightStreamingBudget.h"
+
+#include <cerrno>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <limits>
+
+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
+
+  // The value is a decimal int64 and is not NUL terminated, so copy it into a
+  // bounded stack buffer and terminate it ourselves. The widest int64 is 20
+  // characters plus an optional sign, so anything that does not fit is broken
+  // rather than truncated. An empty value (nbytes == 0) is rejected: producers
+  // never emit an empty budget.
+  constexpr std::size_t kBufSize = 32;
+  if (value == nullptr || nbytes == 0 || nbytes >= kBufSize) {
+    return result;
+  }
+  char buf[kBufSize];
+  std::memcpy(buf, value, nbytes);
+  buf[nbytes] = '\0';
+
+  errno = 0;
+  char* end = nullptr;
+  const long long parsed = std::strtoll(buf, &end, 10);
+  // Require the entire value to be one integer with no leftover characters.
+  // Because the buffer is exactly nbytes long, this single check also rejects an
+  // embedded NUL (e.g. "12\0x") and any trailing garbage (e.g. "12x" or "12 34").
+  if (end != buf + nbytes) {
+    return result;
+  }
+  if (errno == ERANGE) {
+    return result;
+  }
+  // The budget is a non-negative byte count. Reject negatives, and on platforms
+  // where long long is wider than 64 bits reject anything above the int64_t range.
+  if (parsed < 0 || parsed > std::numeric_limits<int64_t>::max()) {
+    return result;
+  }
+
+  result.valid = true;
+  result.bytes = static_cast<int64_t>(parsed);
+  return result;
+}
+
+} // namespace executorch_backend
+} // namespace torch_tensorrt