[NvTensorRTRTX EP] Skip GPU JIT during compile-only sessions

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -456,6 +456,15 @@ static const char* const kOrtSessionOptionStopShareEpContexts = "ep.stop_share_e
 static const char* const kOrtSessionOptionsEpContextModelExternalInitializersFileName =
     "ep.context_model_external_initializers_file_name";
 
+// Internal-only flag set by OrtCompileAPI::CompileModel() to signal EPs that this session
+// is being used for compilation only and will never be used for inference.
+// EPs can use this to skip GPU deserialization and execution context creation, which would
+// otherwise be wasteful since the session is destroyed immediately after compilation.
+// This is NOT a user-facing option and must not be set directly by application code.
+// "0": normal session (default)
+// "1": compile-only session (set internally by OrtCompileAPI::CompileModel)
+static const char* const kOrtSessionOptionCompileOnly = "session.compile_only";
+
 // Gemm fastmath mode provides fp32 gemm acceleration with bfloat16 based matmul.
 // Option values:
 // - "0": Gemm FastMath mode is not enabled. [DEFAULT]

onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc

@@ -1047,6 +1047,7 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info)
   dump_ep_context_model_ = info.dump_ep_context_model;
   ep_context_file_path_ = info.ep_context_file_path;
   ep_context_embed_mode_ = info.ep_context_embed_mode;
+  compile_only_mode_ = info.compile_only_mode;
   enable_engine_cache_for_ep_context_model();
   cache_prefix_ = info.engine_cache_prefix;
   // use a more global cache if given
@@ -2961,35 +2962,39 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr
                             << serialized_engine->size() << " bytes";
     }
 
-    trt_engine = std::unique_ptr<nvinfer1::ICudaEngine>(runtime_->deserializeCudaEngine(serialized_engine->data(), serialized_engine->size()));
-    if (trt_engine == nullptr) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
-                             "NvTensorRTRTX EP failed to deserialize engine for fused node: " + fused_node.Name());
-    }
+    // In compile-only mode the session will not be used for inference. Skip deserialization
+    // and GPU context creation — the serialized engine is already saved as an EP context node.
+    if (!compile_only_mode_) {
+      trt_engine = std::unique_ptr<nvinfer1::ICudaEngine>(runtime_->deserializeCudaEngine(serialized_engine->data(), serialized_engine->size()));
+      if (trt_engine == nullptr) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
+                               "NvTensorRTRTX EP failed to deserialize engine for fused node: " + fused_node.Name());
+      }
 
-    trt_runtime_config = std::unique_ptr<nvinfer1::IRuntimeConfig>(trt_engine->createRuntimeConfig());
-    if (trt_runtime_config && cuda_graph_enable_) {
-      trt_runtime_config->setDynamicShapesKernelSpecializationStrategy(nvinfer1::DynamicShapesKernelSpecializationStrategy::kEAGER);
+      trt_runtime_config = std::unique_ptr<nvinfer1::IRuntimeConfig>(trt_engine->createRuntimeConfig());
+      if (trt_runtime_config && cuda_graph_enable_) {
+        trt_runtime_config->setDynamicShapesKernelSpecializationStrategy(nvinfer1::DynamicShapesKernelSpecializationStrategy::kEAGER);
 #if TRT_MAJOR_RTX > 1 || (TRT_MAJOR_RTX == 1 && TRT_MINOR_RTX >= 3)
-      auto cuda_strategy_flag = trt_runtime_config->setCudaGraphStrategy(nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE);
-      LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] CUDA graph strategy with RTX Graph capture enabled : " << cuda_strategy_flag;
+        auto cuda_strategy_flag = trt_runtime_config->setCudaGraphStrategy(nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE);
+        LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] CUDA graph strategy with RTX Graph capture enabled : " << cuda_strategy_flag;
 #else
-      LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] CUDA graph is enabled but RTX Graph capture is not available. "
-                            << "The current TRT RTX version does not support RTX Graph. "
-                            << "Please upgrade to TRT RTX >= 1.3 to use RTX Graph capture feature for optimal CUDA graph performance.";
+        LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] CUDA graph is enabled but RTX Graph capture is not available. "
+                              << "The current TRT RTX version does not support RTX Graph. "
+                              << "Please upgrade to TRT RTX >= 1.3 to use RTX Graph capture feature for optimal CUDA graph performance.";
 #endif
-    }
-    trt_runtime_config->setExecutionContextAllocationStrategy(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED);
-    if (!runtime_cache_.empty()) {
-      runtime_cache_file = (runtime_cache_ / fused_node.Name()).string();
-      trt_runtime_cache = std::unique_ptr<nvinfer1::IRuntimeCache>(trt_runtime_config->createRuntimeCache());
-      auto cache_data = file_utils::ReadFile(runtime_cache_file);
-      if (!trt_runtime_cache->deserialize(cache_data.data(), cache_data.size())) {
-        trt_runtime_cache = std::unique_ptr<nvinfer1::IRuntimeCache>(trt_runtime_config->createRuntimeCache());
-        LOGS_DEFAULT(INFO) << "TensorRT RTX failed to deserialize the runtime cache, will overwrite with new one" << std::endl;
       }
-      if (!trt_runtime_config->setRuntimeCache(*trt_runtime_cache)) {
-        LOGS_DEFAULT(INFO) << "TensorRT RTX failed to set the runtime cache" << std::endl;
+      trt_runtime_config->setExecutionContextAllocationStrategy(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED);
+      if (!runtime_cache_.empty()) {
+        runtime_cache_file = (runtime_cache_ / fused_node.Name()).string();
+        trt_runtime_cache = std::unique_ptr<nvinfer1::IRuntimeCache>(trt_runtime_config->createRuntimeCache());
+        auto cache_data = file_utils::ReadFile(runtime_cache_file);
+        if (!trt_runtime_cache->deserialize(cache_data.data(), cache_data.size())) {
+          trt_runtime_cache = std::unique_ptr<nvinfer1::IRuntimeCache>(trt_runtime_config->createRuntimeCache());
+          LOGS_DEFAULT(INFO) << "TensorRT RTX failed to deserialize the runtime cache, will overwrite with new one" << std::endl;
+        }
+        if (!trt_runtime_config->setRuntimeCache(*trt_runtime_cache)) {
+          LOGS_DEFAULT(INFO) << "TensorRT RTX failed to set the runtime cache" << std::endl;
+        }
       }
     }
 
@@ -3033,6 +3038,24 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr
     }
   }
 
+  // In compile-only mode the engine was built and saved but GPU deserialization was skipped.
+  // Return a stub compute function — it will never be called since the session is
+  // destroyed immediately after compilation without running any inference.
+  if (compile_only_mode_) {
+    NodeComputeInfo stub_compute_info;
+    stub_compute_info.create_state_func = [](ComputeContext*, FunctionState* state) -> int {
+      *state = nullptr;
+      return 0;
+    };
+    stub_compute_info.release_state_func = [](FunctionState) {};
+    stub_compute_info.compute_func = [](FunctionState, const OrtApi*, OrtKernelContext*) -> Status {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
+                             "NvTensorRTRTX EP: inference is not available in a compile-only session");
+    };
+    node_compute_funcs.push_back(std::move(stub_compute_info));
+    return Status::OK();
+  }
+
   if (weight_stripped_engine_refit_) {
     LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refit engine from main ONNX file after engine build";
     auto status = RefitEngine(model_path_,
@@ -3320,6 +3343,24 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra
                                                                        std::unordered_map<std::string, size_t>& input_map,
                                                                        std::unordered_map<std::string, size_t>& output_map,
                                                                        std::vector<NodeComputeInfo>& node_compute_funcs) {
+  // Compile-only sessions never run inference. The input is already an EPContext model,
+  // so no engine save is needed either — skip deserialization and execution context creation
+  // and register a stub compute function that will not be invoked.
+  if (compile_only_mode_) {
+    NodeComputeInfo stub_compute_info;
+    stub_compute_info.create_state_func = [](ComputeContext*, FunctionState* state) -> int {
+      *state = nullptr;
+      return 0;
+    };
+    stub_compute_info.release_state_func = [](FunctionState) {};
+    stub_compute_info.compute_func = [](FunctionState, const OrtApi*, OrtKernelContext*) -> Status {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
+                             "NvTensorRTRTX EP: inference is not available in a compile-only session");
+    };
+    node_compute_funcs.push_back(std::move(stub_compute_info));
+    return Status::OK();
+  }
+
   std::unique_ptr<nvinfer1::ICudaEngine> trt_engine;
   tensorrt_ptr::unique_pointer_exec_ctx trt_context;
   std::unordered_map<std::string, size_t> input_indexes;   // TRT engine input name -> ORT kernel context input index

onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h

@@ -413,6 +413,9 @@ class NvExecutionProvider : public IExecutionProvider {
 
   // For create/dump EP context node model
   bool dump_ep_context_model_ = false;
+  // Set when the EP is instantiated by OrtCompileAPI::CompileModel(). Causes
+  // CreateNodeComputeInfoFromGraph to skip GPU deserialization and context creation.
+  bool compile_only_mode_ = false;
   std::string ep_context_file_path_;
   int ep_context_embed_mode_ = 0;
   std::string ctx_model_path_;

onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc

@@ -86,6 +86,8 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi
   }
   info.ep_context_file_path = session_options.GetConfigOrDefault(kOrtSessionOptionEpContextFilePath, "");
 
+  info.compile_only_mode = (session_options.GetConfigOrDefault(kOrtSessionOptionCompileOnly, "0") == "1");
+
   // If embed mode is not specified, default to 1 if dump_ep_context_model is true, otherwise 0
   auto embed_mode = std::stoi(session_options.GetConfigOrDefault(kOrtSessionOptionEpContextEmbedMode, "-1"));
   if (embed_mode == -1) {

onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h

@@ -51,6 +51,9 @@ struct NvExecutionProviderInfo {
   bool dump_ep_context_model{false};
   std::string ep_context_file_path{""};
   int ep_context_embed_mode{0};
+  // Set internally by OrtCompileAPI::CompileModel(). When true, the EP skips GPU
+  // deserialization and execution context creation since the session will not run inference.
+  bool compile_only_mode{false};
   std::string engine_cache_prefix{""};
   std::string op_types_to_exclude{""};
 

onnxruntime/core/session/model_compilation_options.cc

@@ -31,6 +31,10 @@ ModelCompilationOptions::ModelCompilationOptions(const onnxruntime::Environment&
   ORT_ENFORCE(session_options_.value.config_options.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1").IsOK());
   ORT_ENFORCE(session_options_.value.config_options.AddConfigEntry(kOrtSessionOptionsDisableModelCompile, "0").IsOK());
 
+  // Signal EPs that this session is used solely for compilation (no inference).
+  // EPs that opt in can use this to skip GPU deserialization / execution context creation.
+  ORT_ENFORCE(session_options_.value.config_options.AddConfigEntry(kOrtSessionOptionCompileOnly, "1").IsOK());
+
   session_options_.value.graph_optimization_level = TransformerLevel::Default;  // L0: required transformers only
 }