feat: Allow for users / kv cache to add aliased I/O for inplace operations

.gitignore

@@ -45,6 +45,13 @@ py/torch_tensorrt/bin
 py/torch_tensorrt/BUILD
 py/torch_tensorrt/LICENSE
 py/torch_tensorrt/WORKSPACE
+# Build copies these into the package dir for wheel packaging; sources of
+# record live at the repo root (core/, csrc/).
+py/torch_tensorrt/CMakeLists.txt
+py/torch_tensorrt/README.md
+py/torch_tensorrt/cmake/
+py/torch_tensorrt/core/
+py/torch_tensorrt/examples/
 py/wheelhouse
 py/.eggs
 notebooks/.ipynb_checkpoints/

core/runtime/TRTEngine.cpp

@@ -97,7 +97,8 @@ TRTEngine::TRTEngine(
     bool requires_output_allocator,
     const std::string& serialized_metadata,
     const ResourceAllocationStrategy resource_allocation_strategy,
-    RuntimeSettings runtime_settings)
+    RuntimeSettings runtime_settings,
+    const std::unordered_map<std::string, AliasedIOSpec>& aliased_io)
     : TRTEngine(
           "deserialized_trt",
           serialized_engine,
@@ -109,7 +110,8 @@ TRTEngine::TRTEngine(
           requires_output_allocator,
           serialized_metadata,
           resource_allocation_strategy,
-          std::move(runtime_settings)) {}
+          std::move(runtime_settings),
+          aliased_io) {}
 
 TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
     : TRTEngine(
@@ -125,7 +127,8 @@ TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
           (static_cast<bool>(std::stoi(serialized_info[RESOURCE_ALLOCATION_STRATEGY_IDX]))
                ? ResourceAllocationStrategy::kDynamic
                : ResourceAllocationStrategy::kStatic),
-          RuntimeSettings{}) {
+          RuntimeSettings{},
+          deserialize_aliased_io(serialized_info[ALIASED_IO_IDX])) {
   // Single visible marker that this engine was instantiated through the C++ runtime
   // entry point (i.e. torch.classes.tensorrt.Engine), distinguishing it from the Python
   // TRTEngine path. Tests look for this string in captured stderr to verify the
@@ -148,7 +151,8 @@ TRTEngine::TRTEngine(
     bool requires_output_allocator,
     const std::string& serialized_metadata,
     const ResourceAllocationStrategy resource_allocation_strategy,
-    RuntimeSettings runtime_settings) {
+    RuntimeSettings runtime_settings,
+    const std::unordered_map<std::string, AliasedIOSpec>& aliased_io) {
   this->runtime_cfg = TRTRuntimeConfig(std::move(runtime_settings));
   TORCHTRT_CHECK(
       target_platform._platform != Platform::PlatformEnum::kUNKNOWN,
@@ -280,9 +284,62 @@ TRTEngine::TRTEngine(
   }
 
   has_dynamic_inputs = engine_has_dynamic_inputs(cuda_engine.get(), in_binding_names);
-  // Cache input binding metadata used by the runtime hot path, then reconstruct
-  // optimization-profile info from the TRT API so multi-profile selection works
-  // for any loaded engine.
+
+  // Store the build-time aliased_io map as the starting point. Then reconcile
+  // against ICudaEngine::getAliasedInputTensor — that API is the source of
+  // truth for KV-cache-style aliasing (TRT-enforced via IKVCacheUpdateLayer)
+  // and may report aliases the build-time path didn't record (e.g. for
+  // engines built outside Torch-TensorRT). User-declared aliases (kind=kUser)
+  // are preserved as-is since TRT doesn't know about them.
+  this->aliased_io = aliased_io;
+  for (const auto& out_name : this->out_binding_names) {
+    // TRT returns nullptr / empty string for non-aliased outputs; any thrown
+    // exception is a real error in the engine state and propagates.
+    const char* aliased_in = cuda_engine->getAliasedInputTensor(out_name.c_str());
+    if (aliased_in == nullptr || aliased_in[0] == '\0') {
+      continue;
+    }
+    auto it = this->aliased_io.find(out_name);
+    if (it == this->aliased_io.end()) {
+      this->aliased_io[out_name] = AliasedIOSpec{std::string(aliased_in), AliasKind::kKVCacheUpdate};
+      LOG_DEBUG("aliased_io reconciliation: discovered " << out_name << " -> " << aliased_in << " (kv_cache_update)");
+    } else if (it->second.kind == AliasKind::kUser) {
+      // A kUser alias is declared by Torch-TensorRT, not TRT. If the engine
+      // also reports an alias for this output the two views should agree on the
+      // source input; keep the kUser kind either way so the provenance recorded
+      // at build time is not silently rewritten to kv_cache_update.
+      if (it->second.input_binding_name != std::string(aliased_in)) {
+        LOG_WARNING(
+            "aliased_io: user-declared alias for output " << out_name << " (input: " << it->second.input_binding_name
+                                                          << ") disagrees with engine-reported input " << aliased_in
+                                                          << "; keeping the user-declared value.");
+      }
+    } else if (it->second.input_binding_name != std::string(aliased_in)) {
+      LOG_WARNING(
+          "aliased_io: build-time map disagrees with engine for output "
+          << out_name << " (build: " << it->second.input_binding_name << ", engine: " << aliased_in
+          << "); using engine value.");
+      it->second = AliasedIOSpec{std::string(aliased_in), AliasKind::kKVCacheUpdate};
+    }
+    // Validation: aliased outputs must not also require an output allocator,
+    // since aliasing requires the output shape to match the input's static
+    // shape, which is incompatible with the dynamic-allocation path.
+    TORCHTRT_CHECK(
+        !this->requires_output_allocator,
+        "Aliased output " << out_name
+                          << " is incompatible with dynamic output allocator. Aliasing requires fixed output shape.");
+  }
+
+  // Precompute the set of input binding names that are the alias source of some
+  // output so the per-call input-setup loop can test membership in O(1).
+  this->aliased_input_binding_names.clear();
+  for (const auto& kv : this->aliased_io) {
+    this->aliased_input_binding_names.insert(kv.second.input_binding_name);
+  }
+
+  // Cache input binding metadata used by the runtime hot path (depends on
+  // aliased_input_binding_names above), then reconstruct optimization-profile
+  // info from the TRT API so multi-profile selection works for any loaded engine.
   this->setup_input_binding_infos();
   this->setup_optimization_profiles();
 
@@ -527,7 +584,8 @@ FlattenedState TRTEngine::__obj_flatten__() {
       std::tuple("requires_output_allocator", serialized_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]),
       std::tuple("target_platform", serialized_info[TARGET_PLATFORM_IDX]),
       std::tuple("resource_allocation_strategy", serialized_info[RESOURCE_ALLOCATION_STRATEGY_IDX]),
-      std::tuple("requires_native_multidevice", serialized_info[REQUIRES_NATIVE_MULTIDEVICE_IDX]));
+      std::tuple("requires_native_multidevice", serialized_info[REQUIRES_NATIVE_MULTIDEVICE_IDX]),
+      std::tuple("aliased_io", serialized_info[ALIASED_IO_IDX]));
 }
 
 std::vector<std::string> TRTEngine::serialize() {
@@ -553,6 +611,7 @@ std::vector<std::string> TRTEngine::serialize() {
   serialized_info[RESOURCE_ALLOCATION_STRATEGY_IDX] =
       this->resource_allocation_strategy == ResourceAllocationStrategy::kDynamic ? "1" : "0";
   serialized_info[REQUIRES_NATIVE_MULTIDEVICE_IDX] = this->requires_native_multidevice ? "1" : "0";
+  serialized_info[ALIASED_IO_IDX] = serialize_aliased_io(this->aliased_io);
   // rank/world_size are runtime facts (may differ at load time); not serialized.
   // RuntimeSettings are intentionally NOT serialized: they're per-engine, in-memory
   // initialization values, not part of the engine's identity.
@@ -571,7 +630,8 @@ void TRTEngine::setup_input_binding_infos() {
     input_binding_infos.push_back(
         {name,
          util::TRTDataTypeToScalarType(cuda_engine->getTensorDataType(name.c_str())),
-         cuda_engine->isShapeInferenceIO(name.c_str())});
+         cuda_engine->isShapeInferenceIO(name.c_str()),
+         aliased_input_binding_names.find(name) != aliased_input_binding_names.end()});
   }
 }
 

core/runtime/TRTEngine.h

@@ -5,6 +5,8 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <unordered_map>
+#include <unordered_set>
 #include <utility>
 
 #include "ATen/core/function_schema.h"
@@ -37,6 +39,40 @@ namespace torch_tensorrt {
 namespace core {
 namespace runtime {
 
+// Origin of an aliased input/output binding pair. KV_CACHE_UPDATE is enforced
+// by TensorRT itself (via IKVCacheUpdateLayer; reported through
+// ICudaEngine::getAliasedInputTensor); USER is declared by the Torch-TensorRT
+// compile flow (TRT doesn't know about it; runtime validates and binds).
+enum class AliasKind : int8_t {
+  kKVCacheUpdate = 0,
+  kUser = 1,
+};
+
+struct AliasedIOSpec {
+  std::string input_binding_name;
+  AliasKind kind;
+};
+
+inline std::string alias_kind_to_string(AliasKind k) {
+  switch (k) {
+    case AliasKind::kKVCacheUpdate:
+      return "kv_cache_update";
+    case AliasKind::kUser:
+      return "user";
+  }
+  return "unknown";
+}
+
+inline AliasKind alias_kind_from_string(const std::string& s) {
+  if (s == "kv_cache_update")
+    return AliasKind::kKVCacheUpdate;
+  if (s == "user")
+    return AliasKind::kUser;
+  // Unknown kinds are conservatively treated as KV-cache-update — TRT enforces
+  // those without extra runtime work, so worst case the runtime silently no-ops.
+  return AliasKind::kKVCacheUpdate;
+}
+
 using FlattenedState = std::tuple<
     std::tuple<std::string, std::string>, // ABI_VERSION
     std::tuple<std::string, std::string>, // name
@@ -49,7 +85,8 @@ using FlattenedState = std::tuple<
     std::tuple<std::string, std::string>, // serialized metadata
     std::tuple<std::string, std::string>, // Platform
     std::tuple<std::string, std::string>, // Resource Allocation Strategy
-    std::tuple<std::string, std::string> // requires_native_multidevice
+    std::tuple<std::string, std::string>, // requires_native_multidevice
+    std::tuple<std::string, std::string> // aliased_io
     >;
 
 struct TorchTRTRuntimeStates {
@@ -74,7 +111,9 @@ struct TorchTRTRuntimeStates {
     if (new_cudagraphs && (!old_cudagraphs || shape_changed || context_changed)) {
       need_cudagraphs_record = true;
     }
-    // Pre-allocated output can be used when previous and current state are true without shape change
+    // Pre-allocated output can be used when previous and current state are true without shape change.
+    // Note: engines with aliased I/O disable pre-allocation entirely; that gating lives in
+    // execute_engine (which has access to the engine's aliased_io map) rather than here.
     if (old_pre_allocated_outputs && new_pre_allocated_output && !shape_changed) {
       can_use_pre_allocated_outputs = true;
     }
@@ -137,6 +176,20 @@ struct TRTEngine : torch::CustomClassHolder {
   std::vector<std::string> in_binding_names = {}; // ITO: PYT IDX
   std::vector<std::string> out_binding_names = {}; // ITO: PYT IDX
 
+  // For each output binding name that aliases an input binding, the alias spec.
+  // Populated either by build-time conversion records (forwarded from
+  // TRTInterpreterResult) or by reconciliation against the engine's own
+  // ICudaEngine::getAliasedInputTensor at construction time. The runtime
+  // consults this map in the output-binding loop to skip allocation and bind
+  // the same device pointer as the source input.
+  std::unordered_map<std::string, AliasedIOSpec> aliased_io = {};
+
+  // The set of input binding names that are the alias source of some output.
+  // Derived once from aliased_io at construction so the per-call input-setup
+  // loop can test membership in O(1) instead of scanning aliased_io on every
+  // input, every execution.
+  std::unordered_set<std::string> aliased_input_binding_names = {};
+
   bool hardware_compatible = false; // Whether the engine was compiled in hardware compatible mode
   std::string serialized_metadata; // This is a base64 encoded pkl object used to store metadata such as settings used
                                    // in compilation
@@ -154,7 +207,8 @@ struct TRTEngine : torch::CustomClassHolder {
       const std::string& serialized_metadata = "",
       const TRTEngine::ResourceAllocationStrategy resource_allocation_strategy =
           TRTEngine::ResourceAllocationStrategy::kStatic,
-      RuntimeSettings runtime_settings = RuntimeSettings{});
+      RuntimeSettings runtime_settings = RuntimeSettings{},
+      const std::unordered_map<std::string, AliasedIOSpec>& aliased_io = {});
 
   TRTEngine(std::vector<std::string> serialized_info);
 
@@ -170,7 +224,8 @@ struct TRTEngine : torch::CustomClassHolder {
       const std::string& serialized_metadata = "",
       const TRTEngine::ResourceAllocationStrategy resource_allocation_strategy =
           TRTEngine::ResourceAllocationStrategy::kStatic,
-      RuntimeSettings runtime_settings = RuntimeSettings{});
+      RuntimeSettings runtime_settings = RuntimeSettings{},
+      const std::unordered_map<std::string, AliasedIOSpec>& aliased_io = {});
 
   std::string to_str() const;
   static void verify_serialization_fmt(const std::vector<std::string>& serialized_info);
@@ -219,6 +274,10 @@ struct TRTEngine : torch::CustomClassHolder {
     std::string name;
     at::ScalarType expected_type;
     bool is_shape_tensor;
+    // True when this input is the alias source of some output binding. Precomputed
+    // here so the per-call input-setup loop avoids an aliased_input_binding_names
+    // lookup on every input, every execution.
+    bool is_aliased_input;
   };
   std::vector<InputBindingInfo> input_binding_infos = {};
   std::vector<at::Tensor> active_input_tensors = {};