Skip to content

Preview branch that addresses a number of reported bugs #4138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
narendasan wants to merge 11 commits into
base: main
Choose a base branch
from
Open

Preview branch that addresses a number of reported bugs #4138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
3873e31
squash the commit
cehongwang Jan 29, 2026
a88200f
support attn_bias for efficient sdpa
zewenli98 Mar 17, 2026
8d81c4e
feat(//py/torch_tensorrt/dynamo): Adding comprehensive support for co…
narendasan Feb 19, 2026
0b449a7
feat(//py/torch_tensorrt/dynamo): Allow the refit system to cache com…
narendasan Mar 5, 2026
ed435c1
docs: Add documentation on how complex numerics works
narendasan Mar 5, 2026
f1b04cd
Instead of keying on shapes we add metadata prior to subgraph replace…
narendasan Mar 6, 2026
3fcd86b
docs: update for new metadata approach
narendasan Mar 6, 2026
e16b612
feat: Complex operations which are not supported will now fallback to…
narendasan Mar 11, 2026
00d0389
chore: addressing PR AIs
narendasan Mar 17, 2026
fb4fc99
remove the old complex subgraph detection
narendasan Mar 18, 2026
4e9c0a6
feat: Updating the profiler to the new kineto based one in PyTorch
narendasan Mar 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/build-test-linux-x86_64.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ jobs:
set -euo pipefail
pushd .
cd tests/py/dynamo
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_converter_tests_results.xml --dist=loadscope --maxfail=20 conversion/
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_converter_tests_results.xml --maxfail=20 conversion/
popd

L0-dynamo-core-tests:
Expand Down Expand Up @@ -236,6 +236,7 @@ jobs:
cd tests/py/dynamo
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_tests_results.xml runtime/test_001_*
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_partitioning_tests_results.xml partitioning/test_001_*
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_hlo_tests_results.xml hlo/

popd

Expand Down
3 changes: 2 additions & 1 deletion .github/workflows/build-test-linux-x86_64_rtx.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ jobs:
set -euo pipefail
pushd .
cd tests/py/dynamo
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_converter_tests_results.xml --dist=loadscope --maxfail=20 conversion/
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_converter_tests_results.xml --maxfail=20 conversion/
popd

L0-dynamo-core-tests:
Expand Down Expand Up @@ -204,6 +204,7 @@ jobs:
pushd .
cd tests/py/dynamo
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_tests_results.xml runtime/test_001_*
python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_hlo_tests_results.xml hlo/
popd

L1-dynamo-compile-tests:
Expand Down
1 change: 1 addition & 0 deletions .github/workflows/build-test-windows.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,7 @@ jobs:
cd tests/py/dynamo
../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_tests_results.xml runtime/test_001_*
../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_partitioning_tests_results.xml partitioning/test_001_*
../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_hlo_tests_results.xml hlo/
popd

L1-dynamo-compile-tests:
Expand Down
1 change: 1 addition & 0 deletions .github/workflows/build-test-windows_rtx.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,7 @@ jobs:
pushd .
cd tests/py/dynamo
../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_tests_results.xml runtime/test_001_*
../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_hlo_tests_results.xml hlo/
popd

L1-dynamo-compile-tests:
Expand Down
58 changes: 12 additions & 46 deletions core/runtime/execute_engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#include "c10/cuda/CUDAGuard.h"
#include "c10/cuda/CUDAStream.h"

#include "ATen/record_function.h"
#include "torch/csrc/jit/runtime/custom_operator.h"
#include "torch/torch.h"

Expand Down Expand Up @@ -238,12 +239,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr

// Intialize inputs and outputs to be available throughout the succeeding scopes
{ // Input Setup
std::unique_ptr<torch::autograd::profiler::RecordProfile> input_profiler_guard;
if (compiled_engine->profile_execution) {
input_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->input_profile_path);
}

RECORD_USER_SCOPE("torch_tensorrt_execute_engine::InputSetup");
setup_input_tensors(inputs, compiled_engine, cudagraphs_enabled, need_cudagraphs_record, inputShapeTensorValues);
// Check if input shapes can be inferred.
int32_t const io_size{compiled_engine->cuda_engine->getNbIOTensors()};
Expand All @@ -257,11 +253,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
}

{ // Output Setup
std::unique_ptr<torch::autograd::profiler::RecordProfile> output_profiler_guard;
if (compiled_engine->profile_execution) {
output_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->output_profile_path);
}
RECORD_USER_SCOPE("torch_tensorrt_execute_engine::OutputSetup");
if (can_use_pre_allocated_outputs) {
outputs = compiled_engine->pre_allocated_outputs;
} else {
Expand Down Expand Up @@ -303,14 +295,9 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
}

{ // Engine Execution (execute on engine stream)
RECORD_USER_SCOPE("torch_tensorrt_execute_engine::Enqueue");
c10::cuda::CUDAStreamGuard stream_guard(compiled_engine->engine_stream);

std::unique_ptr<torch::autograd::profiler::RecordProfile> enqueue_profiler_guard;
if (compiled_engine->profile_execution) {
enqueue_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->enqueue_profile_path);
}

// Block engine stream until results are available on caller stream
at::cuda::CUDAEvent caller_exec_complete;
caller_exec_complete.record(compiled_engine->caller_stream);
Expand Down Expand Up @@ -371,12 +358,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
std::list<std::vector<int64_t>> inputShapeTensorValues;

{ // Input Setup
std::unique_ptr<torch::autograd::profiler::RecordProfile> input_profiler_guard;
if (compiled_engine->profile_execution) {
input_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->input_profile_path);
}

RECORD_USER_SCOPE("torch_tensorrt_execute_engine::InputSetup");
setup_input_tensors(inputs, compiled_engine, false, false, inputShapeTensorValues);
// Check if input shapes can be inferred.
int32_t const io_size{compiled_engine->cuda_engine->getNbIOTensors()};
Expand All @@ -390,11 +372,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
}

{ // OutputAllocator Setup
std::unique_ptr<torch::autograd::profiler::RecordProfile> output_allocator_profiler_guard;
if (compiled_engine->profile_execution) {
output_allocator_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->output_profile_path);
}
RECORD_USER_SCOPE("torch_tensorrt_execute_engine::OutputAllocatorSetup");
create_output_allocator(compiled_engine);
}

Expand All @@ -412,14 +390,9 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
}

{ // Engine Execution (execute on engine stream)
RECORD_USER_SCOPE("torch_tensorrt_execute_engine::Enqueue");
c10::cuda::CUDAStreamGuard stream_guard(compiled_engine->engine_stream);

std::unique_ptr<torch::autograd::profiler::RecordProfile> enqueue_profiler_guard;
if (compiled_engine->profile_execution) {
enqueue_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->enqueue_profile_path);
}

// Block engine stream until results are available on caller stream
at::cuda::CUDAEvent caller_exec_complete;
caller_exec_complete.record(compiled_engine->caller_stream);
Expand All @@ -435,13 +408,10 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
trt_exec_complete.record(compiled_engine->engine_stream);
trt_exec_complete.block(compiled_engine->caller_stream);

std::unique_ptr<torch::autograd::profiler::RecordProfile> output_profiler_guard;
if (compiled_engine->profile_execution) {
output_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->output_profile_path);
}
std::vector<at::Tensor> outputs;
for (size_t i = 0; i < compiled_engine->out_binding_names.size(); i++) {
{ // Output Collection
RECORD_USER_SCOPE("torch_tensorrt_execute_engine::OutputCollection");
for (size_t i = 0; i < compiled_engine->out_binding_names.size(); i++) {
auto name = compiled_engine->out_binding_names[i];
auto dims = compiled_engine->output_allocator->getShapes().at(name);
auto dtype =
Expand All @@ -460,6 +430,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
output = output.reshape(-1).view(dtype).slice(0, 0, prod).reshape(shape);
outputs.push_back(output);
}
} // End Output Collection

if (compiled_engine->profile_execution) {
LOG_INFO(std::endl << *compiled_engine->trt_engine_profiler);
Expand Down Expand Up @@ -493,12 +464,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
bool cudagraphs_enabled = (CUDAGRAPHS_MODE == SUBGRAPH_CUDAGRAPHS);

if (MULTI_DEVICE_SAFE_MODE) {
std::unique_ptr<torch::autograd::profiler::RecordProfile> device_profiler_guard;
if (compiled_engine->profile_execution) {
device_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->device_profile_path);
}

RECORD_USER_SCOPE("torch_tensorrt_execute_engine::DeviceSelection");
RTDevice curr_device = get_current_device();
LOG_DEBUG("Current Device: " << curr_device);

Expand Down
Loading
Loading