diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 9d61165927d8c..9942ce3d0bf8b 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/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] diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index a07fa5b5ab395..26226c1d5ceb7 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/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(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(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(trt_engine->createRuntimeConfig()); - if (trt_runtime_config && cuda_graph_enable_) { - trt_runtime_config->setDynamicShapesKernelSpecializationStrategy(nvinfer1::DynamicShapesKernelSpecializationStrategy::kEAGER); + trt_runtime_config = std::unique_ptr(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(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(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(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(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& input_map, std::unordered_map& output_map, std::vector& 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 trt_engine; tensorrt_ptr::unique_pointer_exec_ctx trt_context; std::unordered_map input_indexes; // TRT engine input name -> ORT kernel context input index diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index 7382ec6e19f19..9ec287ba131d5 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/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_; diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc index 74e16079a7cad..afb9d8107392c 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ b/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) { diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h index 65847745b01ce..779d3f6b67c1a 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h +++ b/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{""}; diff --git a/onnxruntime/core/session/model_compilation_options.cc b/onnxruntime/core/session/model_compilation_options.cc index efaf28fbeefc0..a393bb42fe2cb 100644 --- a/onnxruntime/core/session/model_compilation_options.cc +++ b/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 }