fix(dynamo): save programs with no inputs, keep attention masks that are not plain causal, fix uint8 shape ops and views left to PyTorch, show the runtime TensorRT build number, keep plugins that run PyTorch out of ExecuTorch programs

core/runtime/execute_engine.cpp

@@ -457,6 +457,12 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
       int32_t const io_size{compiled_engine->cuda_engine->getNbIOTensors()};
       std::vector<char const*> names(io_size);
       int32_t const nbNames = ctx->inferShapes(names.size(), names.data());
+      // TensorRT returns -1 and writes no names when it rejects the shapes, such as on a failed assertion.
+      TORCHTRT_CHECK(
+          nbNames >= 0,
+          "Engine " << compiled_engine->name
+                    << " rejected the input shapes, see the TensorRT error above (for example a failed assertion)");
+      names.resize(static_cast<size_t>(nbNames));
       TORCHTRT_CHECK(
           nbNames == 0,
           "The shapes of the inputs: " << names << " of engine " << compiled_engine->name
@@ -669,6 +675,12 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
       int32_t const io_size{compiled_engine->cuda_engine->getNbIOTensors()};
       std::vector<char const*> names(io_size);
       int32_t const nbNames = ctx->inferShapes(names.size(), names.data());
+      // TensorRT returns -1 and writes no names when it rejects the shapes, such as on a failed assertion.
+      TORCHTRT_CHECK(
+          nbNames >= 0,
+          "Engine " << compiled_engine->name
+                    << " rejected the input shapes, see the TensorRT error above (for example a failed assertion)");
+      names.resize(static_cast<size_t>(nbNames));
       TORCHTRT_CHECK(
           nbNames == 0,
           "The shapes of the inputs: " << names << " of engine " << compiled_engine->name

cpp/BUILD

@@ -416,12 +416,14 @@ cc_library(
             "@executorch//:executorch_headers",
             "@executorch//:extension_cuda",
             "@tensorrt//:nvinfer",
+            "@tensorrt//:nvinferplugin",
         ],
         ":sbsa": [
             "@cuda//:cudart",
             "@executorch//:executorch_headers",
             "@executorch//:extension_cuda",
             "@tensorrt_sbsa//:nvinfer",
+            "@tensorrt_sbsa//:nvinferplugin",
         ],
         "//conditions:default": [],
     }),

cpp/src/torch_tensorrt/executorch/CMakeLists.txt

@@ -207,6 +207,7 @@ endif()
 set(_executorch_backend_tensorrt_link_libraries
     CUDA::cudart
     TensorRT::nvinfer
+    TensorRT::nvinfer_plugin
     Threads::Threads
     extension_cuda
 )

cpp/src/torch_tensorrt/executorch/EngineHandle.h

@@ -94,6 +94,8 @@ struct EngineHandle {
   bool pin_ignored_reported = false;
   // Same reasoning: an unaligned buffer is usually the same view on every call.
   bool unaligned_io_reported = false;
+  // And a program gives an input the same dim order on every call.
+  bool reordered_input_reported = false;
 
   ~EngineHandle();
 };

cpp/src/torch_tensorrt/executorch/README.md

@@ -123,8 +123,8 @@ cmake --build build-torchtrt-executorch --target executorch_trt_backend -j
 - Link the archive into your program with `-Wl,--whole-archive`, together with
   `libextension_cuda.so`, ExecuTorch's portable kernels (`libportable_ops_lib.a`, also with
   `-Wl,--whole-archive`, and `libportable_kernels.a`), the ExecuTorch core (`libexecutorch_core.a`),
-  TensorRT's `nvinfer`, the CUDA runtime and pthreads. The delegate and the kernels register
-  themselves from static initializers, and a plain static link drops them.
+  TensorRT's `nvinfer` and `nvinfer_plugin`, the CUDA runtime and pthreads. The delegate and the
+  kernels register themselves from static initializers, and a plain static link drops them.
 - The portable kernels include the two copy operators that a default export puts at the edges of
   the program, `et_copy::_h2d_copy.out` and `et_copy::_d2h_copy.out`. Without them, loading fails
   with `Error::OperatorMissing`, and a build with ExecuTorch logging on also logs
@@ -160,6 +160,16 @@ in place. `torch_tensorrt.save()` plans memory on 256-byte boundaries. ExecuTorc
 16 bytes, so a program finalized with `to_executorch()` directly should pass
 `MemoryPlanningPass(alignment=256)`, or its planned buffers are copied too.
 
+TensorRT also reads every input in row-major order. ExecuTorch can give an input another dim order,
+for example a channels-last program input, or the output of efficient attention in the CUDA backend.
+The delegate copies such an input into row-major order before every call, so the answer is right,
+and logs this once per engine. The cost depends on the layout. Attention's output and channels-last
+need only a few copy calls: about 0.25 ms for a 640x480 float32 image on an A100. A transposed
+matrix needs one call per row: about 6 ms for 1024x1024. From host memory the copy is much slower,
+tens of milliseconds or more for an image, so export such an input in row-major order where you can.
+An input that the engine updates in place must be row-major, or the call fails with
+`Error::InvalidArgument`.
+
 ### CUDA streams
 
 - With no stream chosen, the delegate runs on `cudaStreamPerThread`, the default stream of the

cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp

@@ -36,6 +36,7 @@
 #include <vector>
 
 #include <NvInfer.h>
+#include <NvInferPlugin.h>
 #include <cuda_runtime.h>
 
 #include <executorch/extension/cuda/caller_stream.h>
@@ -133,26 +134,29 @@ Error report_load_error(const char* operation) {
 // TensorRT reads an engine only with the version that built it, unless it was built version
 // compatible, and its own error then names internal serialization tags rather than versions.
 void log_tensorrt_versions(const std::string& exported_version) {
-  // TensorRT 10 and later encode the version as major * 10000 + minor * 100 + patch.
-  const int32_t version = getInferLibVersion();
-  const int major = version / 10000;
-  const int minor = version / 100 % 100;
-  const int patch = version % 100;
+  // The exporter records tensorrt.__version__, which includes the build number, so print the
+  // runtime's build number too. Otherwise equal versions read as different.
+  const int major = getInferLibMajorVersion();
+  const int minor = getInferLibMinorVersion();
+  const int patch = getInferLibPatchVersion();
+  const int build = getInferLibBuildVersion();
   if (exported_version.empty()) {
     ET_LOG(
         Error,
-        "TensorRTBackend::init: this runtime is TensorRT %d.%d.%d, and the program does not record which TensorRT exported the engine",
+        "TensorRTBackend::init: this runtime is TensorRT %d.%d.%d.%d, and the program does not record which TensorRT exported the engine",
         major,
         minor,
-        patch);
+        patch,
+        build);
   } else {
     ET_LOG(
         Error,
-        "TensorRTBackend::init: the engine was exported with TensorRT %s, and this runtime is TensorRT %d.%d.%d",
+        "TensorRTBackend::init: the engine was exported with TensorRT %s, and this runtime is TensorRT %d.%d.%d.%d",
         exported_version.c_str(),
         major,
         minor,
-        patch);
+        patch,
+        build);
   }
 }
 
@@ -293,6 +297,9 @@ nvinfer1::IRuntime* TensorRTBackend::shared_runtime() {
     }
     if (logger == nullptr) {
       logger = new TRTLogger();
+      // Engines can use TensorRT's standard plugins, such as ScatterElements for an
+      // accumulating index_put, and their creators are only registered on request.
+      initLibNvInferPlugins(logger, "");
     }
     runtime = nvinfer1::createInferRuntime(*logger);
     if (runtime != nullptr) {
@@ -740,6 +747,67 @@ bool is_tensor_address_aligned(const void* ptr) {
   return reinterpret_cast<uintptr_t>(ptr) % kTensorAddressAlignment == 0;
 }
 
+// Size-1 dims are skipped, since they move no element, and a tensor built without strides is row-major.
+int innermost_non_row_major_dim(const exec_aten::Tensor& t) {
+  const auto strides = t.strides();
+  if (strides.data() == nullptr) {
+    return -1;
+  }
+  int64_t row_major_stride = 1;
+  for (int d = static_cast<int>(t.dim()) - 1; d >= 0; --d) {
+    if (t.size(d) != 1 && strides[d] != row_major_stride) {
+      return d;
+    }
+    row_major_stride *= t.size(d);
+  }
+  return -1;
+}
+
+// Each row of the 2-D copies is a run of the row-major inner dims, so no CUDA kernel is needed. An outer dim whose
+// stride spans all the rows so far joins them, so channels-last costs one copy per channel, not one per image row.
+cudaError_t copy_to_row_major(void* dst, const exec_aten::Tensor& src, cudaStream_t stream) {
+  const int dim = innermost_non_row_major_dim(src);
+  const size_t element_size = static_cast<size_t>(src.element_size());
+  size_t row_bytes = element_size;
+  for (int d = dim + 1; d < src.dim(); ++d) {
+    row_bytes *= static_cast<size_t>(src.size(d));
+  }
+  const size_t pitch = static_cast<size_t>(src.strides()[dim]);
+  size_t rows = static_cast<size_t>(src.size(dim));
+  int top = dim;
+  while (top > 0 && static_cast<size_t>(src.strides()[top - 1]) == rows * pitch) {
+    rows *= static_cast<size_t>(src.size(--top));
+  }
+  size_t blocks = 1;
+  for (int d = 0; d < top; ++d) {
+    blocks *= static_cast<size_t>(src.size(d));
+  }
+  const auto* src_bytes = static_cast<const uint8_t*>(src.const_data_ptr());
+  auto* dst_bytes = static_cast<uint8_t*>(dst);
+  for (size_t block = 0; block < blocks; ++block) {
+    size_t src_offset = 0;
+    size_t index = block;
+    for (int d = top - 1; d >= 0; --d) {
+      const size_t size = static_cast<size_t>(src.size(d));
+      src_offset += index % size * static_cast<size_t>(src.strides()[d]);
+      index /= size;
+    }
+    const cudaError_t err = cudaMemcpy2DAsync(
+        dst_bytes + block * rows * row_bytes,
+        row_bytes,
+        src_bytes + src_offset * element_size,
+        pitch * element_size,
+        row_bytes,
+        rows,
+        cudaMemcpyDefault,
+        stream);
+    if (err != cudaSuccess) {
+      return err;
+    }
+  }
+  return cudaSuccess;
+}
+
 bool is_cuda_accessible_ptr(const void* ptr) {
   if (ptr == nullptr) {
     return false;
@@ -1898,6 +1966,15 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
           name.c_str());
       return Error::InvalidArgument;
     }
+    // TensorRT reads row-major, and ExecuTorch can plan another dim order, such as efficient attention's output.
+    const bool row_major = innermost_non_row_major_dim(et_in) < 0;
+    if (engine->input_is_alias_target[i] && !row_major) {
+      ET_LOG(
+          Error,
+          "TensorRTBackend::execute: aliased input '%s' must be row-major, because the engine updates it in place",
+          name.c_str());
+      return Error::InvalidArgument;
+    }
 
     void* bind_ptr = nullptr;
     if (et_in.nbytes() == 0) {
@@ -1909,11 +1986,18 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
         engine->cached_input_sizes[i] = 1;
       }
       bind_ptr = engine->cached_input_ptrs[i];
-    } else if (device_reachable && is_tensor_address_aligned(et_in.const_data_ptr())) {
+    } else if (device_reachable && row_major && is_tensor_address_aligned(et_in.const_data_ptr())) {
       bind_ptr = et_in.mutable_data_ptr();
     } else {
-      if (device_reachable) {
+      if (device_reachable && row_major) {
         report_unaligned(name, et_in.const_data_ptr());
+      } else if (!row_major && !engine->reordered_input_reported) {
+        engine->reordered_input_reported = true;
+        ET_LOG(
+            Info,
+            "TensorRTBackend::execute: input '%s' is not in row-major dim order, and TensorRT reads row-major, so "
+            "it is copied into row-major order on every call (reported once per engine)",
+            name.c_str());
       }
       const size_t needed = et_in.nbytes();
       if (needed > engine->cached_input_sizes[i]) {
@@ -1929,7 +2013,8 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
         engine->cached_input_sizes[i] = needed;
       }
       bind_ptr = engine->cached_input_ptrs[i];
-      cuda_err = cudaMemcpyAsync(bind_ptr, et_in.const_data_ptr(), needed, cudaMemcpyDefault, stream);
+      cuda_err = row_major ? cudaMemcpyAsync(bind_ptr, et_in.const_data_ptr(), needed, cudaMemcpyDefault, stream)
+                           : copy_to_row_major(bind_ptr, et_in, stream);
       drain_on_early_return.arm();
       if (cuda_err != cudaSuccess) {
         ET_LOG(

docsrc/user_guide/compilation/compilation_settings.rst

@@ -111,9 +111,18 @@ Optimization Tuning
        Disable only when you need strict IEEE FP32 semantics.
    * - ``attn_bias_is_causal``
      - ``True``
-     - Whether the attn_bias in efficient SDPA is causal. Default is True. This can 
-       accelerate models from HF because attn_bias is always a causal mask in HF. 
-       If you want to use non-causal attn_bias, you can set this to False.
+     - Which attn_bias of efficient SDPA to replace with TensorRT's built-in causal
+       mask, which is faster and more accurate. ``True``: a bias proven causal, one
+       that depends on no graph input and holds 0 on and below the diagonal and -inf
+       (or the dtype's lowest value) above, and also a bias that cannot be proven,
+       such as one built from an input, with a dynamic sequence length, or from a
+       module buffer under ``torch.compile``; you assert that it is causal. ``None``:
+       only a bias proven causal, for models whose input-built masks may not be
+       causal, such as masks with padding. A bias that is computed and found not
+       causal is always kept. A mask longer than 4096, or one whose check fails, is
+       not proven: ``True`` replaces it and ``None`` keeps it. ``False``: never
+       replace. A replaced bias is not read again, so refitting the weights it came
+       from does not change the engine.
 
 ----
 
@@ -450,13 +459,16 @@ rebuilt from scratch:
 
 ``max_aux_streams``, ``version_compatible``,
 ``optimization_level``, ``disable_tf32``, ``sparse_weights``,
-``engine_capability``, ``hardware_compatible``, ``refit_identical_engine_weights``,
+``engine_capability``, ``hardware_compatible``, ``target_compute_capabilities``,
+``refit_identical_engine_weights``,
 ``immutable_weights``, ``enable_weight_streaming``, ``tiling_optimization_level``,
 ``l2_limit_for_tiling``, ``enable_autocast``, ``autocast_low_precision_type``,
 ``autocast_excluded_nodes``, ``autocast_excluded_ops``,
 ``autocast_max_output_threshold``, ``autocast_max_depth_of_reduction``,
-``autocast_calibration_dataloader``, ``decompose_attention``, ``attn_bias_is_causal``,
-``fold_shared_kv_attention``.
+``autocast_calibration_dataloader``, ``decompose_attention``,
+``disabled_constant_fold_exclusions``, ``attn_bias_is_causal``,
+``fold_shared_kv_attention``, ``build_route``, ``fallback_pytorch_plugin_ops``.
+On TensorRT older than 10.14, ``strip_engine_weights`` also invalidates cached engines.
 
 Settings not in this list (e.g., ``debug``, ``dryrun``, ``pass_through_build_failures``)
 can be changed without invalidating the cache.

docsrc/user_guide/runtime_performance/executorch.rst

@@ -166,6 +166,26 @@ in `Run from Python`_ and `Run from C++`_. If the CUDA part has weights, loading
     the same directory overwrites the first one's weights. The first model can then return wrong
     answers with no error, or fail to load. Save each model into its own directory.
 
+Keep plugins that run PyTorch out of the engine
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Some converters build a TensorRT plugin that runs PyTorch when the engine runs: ``aten.index_put``
+with ``accumulate=True`` uses Torch-TensorRT's ScatterAdd plugin, and a plugin made with
+``torch_tensorrt.dynamo.conversion.plugins.custom_op`` that does not use an AOT kernel runs the
+PyTorch op from Python. The ExecuTorch runtime does not load PyTorch, so ``save()`` refuses such an
+engine and names the operators. Compile with ``fallback_pytorch_plugin_ops=True`` to keep those
+operators out of TensorRT. A ``torch.cond`` with one of them in a branch then runs in PyTorch too.
+With the setting, an accumulating ``index_put`` with one 1-D index stays in TensorRT instead,
+through TensorRT's own ScatterElements plugin, which needs no PyTorch.
+
+.. code-block:: python
+
+    trt_gm = torch_tensorrt.dynamo.compile(
+        exported,
+        arg_inputs=example_inputs,
+        fallback_pytorch_plugin_ops=True,
+    )
+
 Keep inputs and outputs on the GPU
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

examples/torchtrt_executorch_example/export_multi_profile.py

@@ -93,9 +93,7 @@ def stamp(phase: str) -> None:
 # narrower and shallower. ``sliding_window`` keeps the 1B model's 512, which is
 # wider than MAX_SEQ, so as in the real model the window never binds over the
 # exported range and every layer attends to the whole prefix. Narrowing it below
-# MAX_SEQ would make the sliding layers genuinely windowed, and the engine would
-# then need ``attn_bias_is_causal=False`` to keep the mask instead of assuming
-# plain causality.
+# MAX_SEQ would make the sliding layers genuinely windowed.
 MINI_CONFIG = dict(
     vocab_size=2048,
     hidden_size=320,

py/torch-tensorrt-executorch-runtime/native/BUILD.bazel

@@ -44,7 +44,13 @@ cmake(
     deps = [
         "@cuda//:cudart",
     ] + select({
-        ":aarch64_linux": ["@tensorrt_sbsa//:nvinfer"],
-        "//conditions:default": ["@tensorrt//:nvinfer"],
+        ":aarch64_linux": [
+            "@tensorrt_sbsa//:nvinfer",
+            "@tensorrt_sbsa//:nvinferplugin",
+        ],
+        "//conditions:default": [
+            "@tensorrt//:nvinfer",
+            "@tensorrt//:nvinferplugin",
+        ],
     }),
 )

py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt

@@ -74,6 +74,7 @@ target_link_libraries(executorch_backend_tensorrt PRIVATE
   executorch::extension_cuda
   CUDA::cudart
   TensorRT::nvinfer
+  TensorRT::nvinfer_plugin
   Threads::Threads)
 
 # Build to the standard ExecuTorch's own exported configuration asks for. The delegate is compiled

py/torch-tensorrt-executorch-runtime/native/check_imports_executorch_runtime.sh

@@ -99,7 +99,7 @@ fi
 for entry in $(needed_entries "${dyn}"); do
     case "${entry}" in
         libexecutorch.so|libexecutorch_extension_cuda.so|libexecutorch_threadpool.so) ;;
-        libnvinfer.so.*|libcudart.so.*) ;;
+        libnvinfer.so.*|libnvinfer_plugin.so.*|libcudart.so.*) ;;
         libstdc++.so.*|libc.so.*|libm.so.*|libgcc_s.so.*|libdl.so.*|librt.so.*|libpthread.so.*) ;;
         ld-linux-*.so.*|libatomic.so.*|libz.so.*) ;;
         *) fail "${target} links ${entry}, which is not one of the libraries this wheel may depend on" ;;