Uh oh!
There was an error while loading. Please reload this page.
perf(runtime): skip TensorRT engine fakification during export by reporting tracing_mode "real" - #4489
Merged
lanluo-nvidia merged 1 commit intoAug 20, 2026
Conversation
torch.export fakifies every torchbind script object it traces through. For a TensorRT engine that is enormously expensive twice over: maybe_to_fake_obj calls __obj_flatten__, which runs TRTEngine::serialize() (a full cuda_engine->serialize() plus base64) and then base64-decodes it again in FakeTRTEngine.__obj_unflatten__, and FakeScriptObject.__init__ deep-copies the engine through its pickle, which serializes and then fully deserializes another copy onto the device. The fake stores the resulting blob and never reads it -- every field it uses is metadata. ExecuTorch re-runs decompositions during to_edge_transform_and_lower, so this happens roughly three times per engine, and the cost grows with engine count and size. torch offers an opt-out: tracing_with_real() returns True for objects whose tracing_mode() reports "real", and maybe_to_fake_obj then hands the real object to the meta kernel untouched. fake_tensorrt_execute_engine already handles that case -- it branches on real_obj and otherwise reads get_serialized_metadata(), taking output shapes from the stored symbolic expressions -- so nothing executes or mutates the engine during tracing. The Python TRTEngine has reported "real" for this reason since it was added; the TorchBind class simply never did. Measured on a 2-layer Gemma4-MoE hybrid TensorRT + CUDA ExecuTorch export (3 engines): 706s -> 180s wall, with the fakification phase going from 518s to 0.0s. The exported program is unchanged -- identical delegate count, operators, values, and byte-identical delegate payloads.
7 tasks
Conarnar added a commit
to Conarnar/TensorRT
that referenced
this pull request
Aug 18, 2026
Export resolves engine info for two purposes, and neither wants the base64 form serialize() produces. validate_engine_program reads only flags; replace_execute_engine wants the engine as a byte tensor. Both were paying an encode in C++ and, for the second, a matching decode here. Take each half from the accessor that provides it: metadata_only=True on get_engine_info_from_state and _resolve_engine_info, so a metadata reader never triggers engine serialization. _resolve_engine_tensor returns the engine as a uint8 tensor directly, which also drops the torch.frombuffer rebuild. Both fall back to the old path when the runtime lacks the accessors, so this does not require a matching runtime. They are missing only when the Torch-TensorRT C++ library is older than this Python package -- a source or editable build where only the Python half was rebuilt. The symptom would otherwise be silence, since the fallback is correct and merely slower, so it warns once per accessor naming the cause. _resolve_engine_object is factored out of _resolve_engine_info because the engine arg is a get_attr before ExecuTorch lifts constants and a placeholder after; handling only the first made the tensor accessor silently fall back to base64 on every graph that had been through staging, which is exactly the case that matters. Nothing is cached. With the accessors present a metadata read is a member read away from free, so there is nothing worth memoizing; a runtime without them re-serializes per read, which is the price of keeping this change small. backend.py and partitioner.py are deliberately untouched. Their engine reads go through the no_op_placeholder branch, which already carries the engine as a tensor argument and serializes nothing, so routing them through the accessors would add a parameter that changes no work. Measured on a base carrying pytorch#4473, with pytorch#4489 present in the loaded runtime, exporting a model with a 67MB engine; six runs per configuration, interleaved. The passes that read engine info go from 1.40s to 0.11s and the ExecuTorch lowering phase from 2.35s to 1.08s, a 54% reduction. Engine-state reads are nearly all of that: 1.07s to 0.07s, with the Python-side base64 decode going from 0.26s to nothing. Serializing the engine is only about 0.07s of a 0.54s __getstate__, so most of what disappears is the base64 encode and the string copies feeding it. A metadata-only read costs 10-20us.
Conarnar added a commit
to Conarnar/TensorRT
that referenced
this pull request
Aug 18, 2026
Export resolves engine info for two purposes, and neither wants the base64 form serialize() produces. validate_engine_program reads only flags; replace_execute_engine wants the engine as a byte tensor. Both were paying an encode in C++ and, for the second, a matching decode here. Take each half from the accessor that provides it: metadata_only=True on get_engine_info_from_state and _resolve_engine_info, so a metadata reader never triggers engine serialization. _resolve_engine_tensor returns the engine as a uint8 tensor directly, which also drops the torch.frombuffer rebuild. Both fall back to the old path when the runtime lacks the accessors, so this does not require a matching runtime. They are missing only when the Torch-TensorRT C++ library is older than this Python package -- a source or editable build where only the Python half was rebuilt. The symptom would otherwise be silence, since the fallback is correct and merely slower, so it warns once per accessor naming the cause. _resolve_engine_object is factored out of _resolve_engine_info because the engine arg is a get_attr before ExecuTorch lifts constants and a placeholder after; handling only the first made the tensor accessor silently fall back to base64 on every graph that had been through staging, which is exactly the case that matters. Nothing is cached. With the accessors present a metadata read is a member read away from free, so there is nothing worth memoizing; a runtime without them re-serializes per read, which is the price of keeping this change small. backend.py and partitioner.py are deliberately untouched. Their engine reads go through the no_op_placeholder branch, which already carries the engine as a tensor argument and serializes nothing, so routing them through the accessors would add a parameter that changes no work. Measured on a base carrying pytorch#4473, with pytorch#4489 present in the loaded runtime, exporting a model with a 67MB engine; six runs per configuration, interleaved. The passes that read engine info go from 1.40s to 0.11s and the ExecuTorch lowering phase from 2.35s to 1.08s, a 54% reduction. Engine-state reads are nearly all of that: 1.07s to 0.07s, with the Python-side base64 decode going from 0.26s to nothing. Serializing the engine is only about 0.07s of a 0.54s __getstate__, so most of what disappears is the base64 encode and the string copies feeding it. A metadata-only read costs 10-20us.
Conarnar added a commit
to Conarnar/TensorRT
that referenced
this pull request
Aug 19, 2026
Export resolves engine info for two purposes, and neither wants the base64 form serialize() produces. validate_engine_program reads only flags; replace_execute_engine wants the engine as a byte tensor. Both were paying an encode in C++ and, for the second, a matching decode here. Take each half from the accessor that provides it: metadata_only=True on get_engine_info_from_state and _resolve_engine_info, so a metadata reader never triggers engine serialization. _resolve_engine_tensor returns the engine as a uint8 tensor directly, which also drops the torch.frombuffer rebuild. Both fall back to the old path when the runtime lacks the accessors, so this does not require a matching runtime. They are missing only when the Torch-TensorRT C++ library is older than this Python package -- a source or editable build where only the Python half was rebuilt. The symptom would otherwise be silence, since the fallback is correct and merely slower, so it warns once per accessor naming the cause. _resolve_engine_object is factored out of _resolve_engine_info because the engine arg is a get_attr before ExecuTorch lifts constants and a placeholder after; handling only the first made the tensor accessor silently fall back to base64 on every graph that had been through staging, which is exactly the case that matters. Nothing is cached. With the accessors present a metadata read is a member read away from free, so there is nothing worth memoizing; a runtime without them re-serializes per read, which is the price of keeping this change small. backend.py and partitioner.py are deliberately untouched. Their engine reads go through the no_op_placeholder branch, which already carries the engine as a tensor argument and serializes nothing, so routing them through the accessors would add a parameter that changes no work. Measured on a base carrying pytorch#4473, with pytorch#4489 present in the loaded runtime, exporting a model with a 67MB engine; six runs per configuration, interleaved. The passes that read engine info go from 1.40s to 0.11s and the ExecuTorch lowering phase from 2.35s to 1.08s, a 54% reduction. Engine-state reads are nearly all of that: 1.07s to 0.07s, with the Python-side base64 decode going from 0.26s to nothing. Serializing the engine is only about 0.07s of a 0.54s __getstate__, so most of what disappears is the base64 encode and the string copies feeding it. A metadata-only read costs 10-20us.
Contributor
Reviewed this and it looks correct to me. Below is what I checked, then two things I would like to see before it lands. The mechanism is exactly as you describe. A typo would assert loudly rather than quietly degrade, which is a nice property. The Python engine has reported Every consumer that could notice the difference still works. The meta kernel branches on One spot to glance at: Two requests:
Also worth noting |
Uh oh!
There was an error while loading. Please reload this page.
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 freeto 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
torch.exportfakifies every TorchBind script object it traces through, and for a TensorRT engine that is expensive twice over:maybe_to_fake_objcalls__obj_flatten__, which runsTRTEngine::serialize()— a fullcuda_engine->serialize()plus base64 — andFakeTRTEngine.__obj_unflatten__then base64-decodes it again.FakeScriptObject.__init__deep-copies the real object, which for a TorchBind class goes throughdef_pickle: serialize, then a fulldeserializeCudaEngineof a second copy onto the device.The fake never reads any of it.
FakeTRTEngineassigns the blob toself.serialized_engineand every field it actually uses is metadata. ExecuTorch re-runs decompositions duringto_edge_transform_and_lower, so this happens roughly three times per engine, and the cost grows with engine count and size.torchalready provides the opt-out:tracing_with_real()returns True for an object whosetracing_mode()reports"real", andmaybe_to_fake_objthen hands the real object to the meta kernel untouched. This adds that method to the TorchBind engine.Why this is safe
fake_tensorrt_execute_enginealready handles being given a real engine: it branches onreal_objand otherwise callsget_serialized_metadata()directly, taking output shapes from the symbolic expressions stored in the engine's metadata. Nothing executes and nothing mutates the engine during tracing.The Python
TRTEnginehas reported"real"for exactly this reason since it was added — its docstring citestracing_with_real. The TorchBind class simply never did the same.Measurements
The mechanism is
torch.exportfakification, so any export carrying TensorRT engines benefits; the numbers below come from a two-layer Gemma4-MoE exported through the hybrid TensorRT + CUDA ExecuTorch path (3 engines), fresh Inductor cache, otherwise identical runs:The exported program is unchanged: same delegate count, operators and values, and byte-identical delegate payloads.
Testing
Verified by rebuilding and re-running the export above with no patches: fakification drops to zero and the artifact is byte-identical to the baseline. Existing
tests/py/dynamo/executorchandtests/py/dynamo/lowering/test_buffer_lifting.pypass (75).No automated test is included. Nothing in
tests/currently exercises fakification of an engine, and the meaningful property — that export stops serializing the engine per fakify call — is a timing assertion. The cheapest real check would be assertingtracing_mode() == "real"on a built engine, which requires a GPU. Happy to add that if you want it gated.