fix: accept host-resident shape-tensor inputs without a device round trip - #4718
Open
SrivastavaKshitij wants to merge 2 commits into
Open
SrivastavaKshitij wants to merge 2 commits into
SrivastavaKshitij wants to merge 2 commits into
Conversation
github-actions
Bot
added
component: core
Author
narendasan
reviewed
Sep 15, 2026
| // A device tensor is copied back with .cpu(), which synchronizes the stream; a host tensor | ||
| // is used as is, which does not, so passing shape inputs on the host avoids a per-call sync. | ||
| auto input_cpu = inputs[i].is_cuda() ? inputs[i].clone().contiguous().cpu().to(torch::kInt64) | ||
| : inputs[i].contiguous().to(torch::kInt64); |
Collaborator
There was a problem hiding this comment.
@apbose do you think we need the clone on the second arm?
Collaborator
There was a problem hiding this comment.
yeah I dont think the clone is required since we anyways deep copy the shape tensors to the engine owned vectors
apbose
reviewed
Sep 16, 2026
| if (!inp.defined() || inp.is_cuda()) { | ||
| continue; | ||
| } | ||
| if (i < compiled_engine->input_binding_infos.size() && compiled_engine->input_binding_infos[i].is_shape_tensor) { |
Collaborator
There was a problem hiding this comment.
in this, I think we should also check if the shape tensor is on cpu? basically adding && inp.is_cpu I don't think we have cases like that rn, but would be more fool-proof
| "calling the TensorRT engine (e.g. tensor.cuda() or tensor.to(device))."); | ||
| inp = inp.cuda(); | ||
| } | ||
|
|
Collaborator
There was a problem hiding this comment.
can we have a test for the host shape input path? the existing arange have shape inputs on cuda. We could use the ShapeInputModel from #4717 just to verify the correctness
Author
|
Thanks! All three addressed in
new tests pass; |
apbose
approved these changes
Sep 23, 2026
…trip A TensorRT shape-tensor input is read from host memory, but setup_input_tensors required every input on the device and copied a shape input back with .cpu(), and execute_engine moved any host input to the device first. A host shape value was therefore pushed device-side and pulled straight back, and the .cpu() copy synchronizes the stream. Accept a host-resident shape-tensor input as is; device inputs and all non-shape inputs are unchanged.
apbose
force-pushed
the
perf/host-shape-tensor-inputs
branch
from
September 23, 2026 03:48
2ec01df to
4c4613f
Compare
…input test - setup_input_tensors: remove clone() on the device arm; .cpu() already returns a fresh host tensor and values are deep-copied into active_shape_tensor_values. - execute_engine preamble: only leave a shape-tensor input in place when it is on CPU; a shape tensor on any other non-CUDA device falls through to the CUDA move. - add tests/py/dynamo/runtime/test_host_shape_inputs.py covering the host shape input path (arange converter tests only exercise a CUDA shape input).
apbose
force-pushed
the
perf/host-shape-tensor-inputs
branch
from
September 23, 2026 03:51
4c4613f to
b7b611f
Compare
Collaborator
|
approved. waiting on full CI to run after rebase |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
A TensorRT shape-tensor input is read from host memory, but
setup_input_tensorsrequired everyinput to be on the device and then copied a shape input back to the host with
.cpu(), andexecute_enginemoved any host input to the device before that. A shape value the caller alreadyholds on the host was therefore pushed device-side and pulled straight back, and the
.cpu()copysynchronizes the stream — the host blocks until the work already queued on the stream has finished.
This accepts a host-resident shape-tensor input and uses it as is. A device-resident shape input is
unchanged (still copied to the host); every non-shape input still must be on the device.
Before (a host shape value is moved to the device, then copied back — the slower path):
After (a host shape value is used directly):
The copy the host path avoids is a stream synchronization. With a backlog of GPU work queued on the
stream,
input.clone().contiguous().cpu().to(int64)blocks the host for the full backlog (~120 msin the benchmark) while
input.contiguous().to(int64)does not. On a GPU-bound model the wallclock is unchanged; the win is the removed host stall and the tighter host/device overlap.
Issues
Closes #4717
Type of change
Testing
The benchmark in the linked issue (
benchmark_host_shape_inputs.py, numbers above), one L4 innvcr.io/nvidia/pytorch:26.07-py3, built from this branch. Correctness: the engine's output matches eager for a shape input placed onthe host and on the device, across shape values 8/64/200. No regressions:
tests/py/dynamo/runtime/is 164 passed / 68 skipped, and
tests/py/dynamo/conversion/test_arange_aten.py(the shape-tensorinput converter) is 20 passed with 1 failure that also fails on
mainunchanged.Checklist: