Skip to content

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

Merged
chilo-ms merged 1 commit into
microsoft:mainfrom
umangb-09:CompileApi-v2
Jun 1, 2026
Merged

[NvTensorRTRTX EP] Skip GPU JIT during compile-only sessions#28503
chilo-ms merged 1 commit into
microsoft:mainfrom
umangb-09:CompileApi-v2

Conversation

@umangb-09

Copy link
Copy Markdown
Contributor

Description

Add an internal session config entry, "session.compile_only", set by CompileModel() before
session initialization. The NvTensorRTRTX EP reads it in
NvExecutionProviderInfo::FromProviderOptions() and, when set, skips deserializeCudaEngine() /
createExecutionContext() in CreateNodeComputeInfoFromGraph().

The EP context node is still saved — that path uses the serialized engine buffer directly and does
not depend on the deserialized engine. A stub compute function is registered to satisfy the
framework; it returns NOT_IMPLEMENTED if called, which cannot happen in practice because
compile-only sessions are destroyed without inference.

Motivation and Context

OrtCompileAPI::CompileModel() creates an InferenceSession solely to drive EP::Compile() and
write out the EPContext model, then destroys it without running inference. During that session, the
NvTensorRTRTX EP was performing a full deserializeCudaEngine() and createExecutionContext()
uploading engine weights to the GPU and JIT-ing the engine, only to free everything when the session
was destroyed.

When the user then loads the EPContext model in a real session, the same JIT and upload happen again.
Net effect on the typical "compile, then load and run" flow:

ONNX model
    → CompileModel()         [JIT + GPU upload #1 — discarded]
    → EP context model saved to disk
    → Session from EP context model
                             [JIT + GPU upload #2 — necessary]
    → Inference

JIT and GPU upload run twice.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR introduces an internal compile-only session flag (session.compile_only) that is set by OrtCompileAPI::CompileModel() to let EPs (specifically NvTensorRTRTX) skip expensive GPU engine deserialization and execution context creation during compile-only sessions that are never used for inference.

Changes:

  • Set internal session config session.compile_only=1 when CompileModel() creates the temporary InferenceSession.
  • Plumb the flag into NvTensorRTRTX EP option parsing and provider state (compile_only_mode_).
  • In NvTensorRTRTX EP, skip deserializeCudaEngine() / createExecutionContext() during graph compilation and register a stub compute function for compile-only sessions.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
onnxruntime/core/session/utils.cc Clones session options in CompileModel() and sets session.compile_only for the compilation-only session.
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h Adds the internal kOrtSessionOptionCompileOnly config key string.
onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h Adds compile_only_mode to EP info struct.
onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc Reads session.compile_only from session config into compile_only_mode.
onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h Adds compile_only_mode_ member to the EP instance.
onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc Skips GPU deserialization/context creation in compile-only mode and returns a stub NodeComputeInfo.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc
Comment thread onnxruntime/core/session/utils.cc Outdated
@chilo-ms
chilo-ms enabled auto-merge (squash) June 1, 2026 15:45
@chilo-ms
chilo-ms merged commit 6572dd0 into microsoft:main Jun 1, 2026
91 of 92 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants