Repository navigation
fix(dynamo): keep a cast to uint8 inside the TensorRT engine - #5022
Open
shoumikhin wants to merge 2 commits into
Open
shoumikhin wants to merge 2 commits into
shoumikhin wants to merge 2 commits into
Conversation
Picking one frame out of a uint8 batch, as in `frames[1].float() / 255`, split the graph: the partitioner refused aten.select.int on a uint8 tensor, so the select ran in PyTorch between two engines. TensorRT allows uint8 only at engine inputs and outputs, but select only picks existing values. Like slice and index_select, it now goes in UINT8_DATA_MOVEMENT_OPS: the interpreter widens its uint8 input to int32, and a uint8 engine output is cast back to uint8. The select validator that refused uint8 is removed, since the gather no longer sees uint8, and so is the select entry in the interpreter's value-picking list, which can no longer be reached.
Rounding video frames to uint8 levels, as in `frames.round().clamp(0, 255).to(torch.uint8)`, split the graph: the cast to uint8 ran in PyTorch between two engines, even when the uint8 values were only sliced, concatenated and cast back to float. TensorRT allows uint8 only at engine inputs and outputs, so the cast now makes the widened int32 value that the engine already uses for uint8 data. It matches PyTorch exactly: PyTorch converts through int64, truncating toward zero, then wraps modulo 256 (255.9 gives 255, 256.7 gives 0, -1.5 gives 255). The TensorRT cast to int64 truncates the same way, and the existing wrap_to_uint8 helper does the wrap. A bool is already 0 or 1, so it is cast straight to int32. A uint8 engine output is still cast back to uint8 at the output. UINT8_CAST_SOURCES lists the source dtypes this applies to: bool, int32, int64, float16, bfloat16 and float32. The partitioner check and the _to_copy validator read it, and only the _to_copy converters call the new impl.cast.to_uint8, so other callers of impl.cast.to_copy are unchanged. A cast from any other dtype still runs in PyTorch. For int8 that is required: TensorRT treats an int8 input as quantized, so the math after the cast fails the build. With truncate_double, the _to_copy and aten.copy converters take the float64 rule from the graph's source dtype, as their validators do, since a uint8 source is int32 inside the engine. Values outside the int64 range, infinity and NaN are unspecified, since PyTorch's own result is undefined for them and differs between CPU and CUDA.
github-actions
Bot
added
component: tests
shoumikhin
force-pushed
the
uint8-cast-in-engine
branch
from
November 14, 2026 17:22
93d0280 to
b9b70bb
Compare
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
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.
This sits on top of #5019, so the first commit here is #5019. Only the last commit is new.
What is wrong today
Video models often round frames to uint8:
The values are only sliced, joined and cast back to float, yet the cast runs in PyTorch and splits the graph.
TensorRT allows uint8 only at engine inputs and outputs, so the partitioner refused the cast, and the
_to_copyvalidator did not list uint8.What this change does
Inside an engine, uint8 data already lives as int32. The cast now produces that int32 value, so it stays in the engine. Other users of
impl.cast.to_copyare unchanged.The values match PyTorch exactly, except values outside the int64 range, infinity and NaN, which PyTorch leaves undefined. PyTorch converts float to uint8 through int64: it truncates toward zero, then wraps modulo 256. So 255.9 gives 255, 256.7 gives 0 and -1.5 gives 255. The converter does the same: a TensorRT cast to int64 (also truncating toward zero), the existing
wrap_to_uint8helper, then a cast to int32. A bool (0 or 1) is cast straight to int32.A uint8 engine output is cast back to uint8.
UINT8_CAST_SOURCESlists the source dtypes this covers (bool, int32, int64, float16, bfloat16, float32); others still run in PyTorch. For int8 that is required: TensorRT treats int8 as quantized, and the build fails.With
truncate_double=True, a float64 copy of the uint8 result still runs as float32, decided from the graph's dtype.Tested
test_casts.pycases expect one engine and an exact match: a cast from each listed dtype (float32 must wrap 256.7, -1.5 and 3e9), atorch.condbranch, and the example above with frames in NHWC order.+ 200) is refused.truncate_double=True: a uint8 cast then.double()or a float64aten.copybuilds one engine.