Repository navigation
fix(dynamo): several correctness and speed fixes found compiling large models - #4936
Open
shoumikhin wants to merge 13 commits into
Open
shoumikhin wants to merge 13 commits into
shoumikhin wants to merge 13 commits into
Conversation
github-actions
Bot
added
component: tests
shoumikhin
force-pushed
the
dynamo-converter-lowering-fixes
branch
from
October 27, 2026 22:52
6e1c878 to
fb99559
Compare
shoumikhin
changed the title
github-actions
Bot
added
the
component: lowering
shoumikhin
changed the title
shoumikhin
force-pushed
the
dynamo-converter-lowering-fixes
branch
from
October 28, 2026 02:42
417861f to
88ccbc4
Compare
shoumikhin
force-pushed
the
dynamo-converter-lowering-fixes
branch
2 times, most recently
from
October 28, 2026 17:50
9a91c2f to
06ed284
Compare
lanluo-nvidia
self-requested a review
October 28, 2026 18:00
lanluo-nvidia
left a comment
lanluo-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
AI says the broadcast expansion introduce by this pr might have the following issue:
Bound the broadcast size in integer scatter_add — _decompositions.py:368–370
(
TensorRT/py/torch_tensorrt/dynamo/lowering/_decompositions.py
Lines 368 to 370 in 46cad42
destination position with every index, without limiting the resulting tensor’s size. Fake-tensor tracing of a valid 1,048,576-element destination with 4,096 updates produces
4,294,967,296-element intermediates—32 GiB each for int64. The previous implementation’s largest intermediate had 1,048,576 elements. This exceeds TensorRT’s documented volume
limit on Turing (https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/operators/index.html) and risks excessive memory use elsewhere. Gate this path by the broadcast
volume, or chunk the comparison and reduction.
shoumikhin
force-pushed
the
dynamo-converter-lowering-fixes
branch
from
October 28, 2026 19:23
06ed284 to
dad9ec8
Compare
Contributor
Author
|
Thanks, good catch! The new path now only kicks in when the comparison stays reasonably small. Bigger scatters, or ones whose size we can't bound up front, fall back to the old loop, so they're never worse than before. Your example takes the loop now. It's slower there but safe, and making that case fast too (for example by chunking) feels like a good follow-up rather than something to add to this PR. The models this was meant for still stay on the fast path, so the big compile and speed wins hold. I also added a test for both sides of the limit. |
shoumikhin
force-pushed
the
dynamo-converter-lowering-fixes
branch
2 times, most recently
from
October 29, 2026 13:36
c970681 to
58bd172
Compare
In Python and PyTorch, a slice bound past either end of a dimension is clamped to that end. On a dimension of size 10, `x[-15:]` is the whole dimension and `x[:-15]` is empty. When the sliced dimension has a static size, the slice converter normalized its bounds with `get_positive_dim`, which wraps a negative value with a modulo. That is right for a dimension index but not for a slice bound: `x[-15:]` became `x[5:]` and `x[:-15]` became `x[:5]`. The engine built and ran, and silently returned the wrong rows. A slice whose start lands after its stop (an empty slice in PyTorch) could fail the engine build with a negative output length. This shows up in real models. In transformers, the sliding-window layers of the dynamic key/value cache (used by Gemma2 and Gemma3, for example) store `keys[:, :, -window + 1:]`. With a static prompt shorter than the window, the stored cache lost the first prompt tokens for most prompt lengths, in every sliding-window layer. Clamp static bounds the same way the converter already clamps a start computed at runtime, and floor the output length at zero. Literal bounds on a dynamic dimension take a different branch and are not changed here.
…oner The fast (adjacency) partitioner can take hours on large exported graphs. On a vision-language-action model whose lowered graph has 345k nodes, it did not finish in 75 minutes. TRTPartitioner.starter_nodes puts every callable node in the frontier. The inherited _SplitterBase.put_nodes_into_subgraphs then looks through that frontier for a node whose dependencies are all placed, once per placed node, so the work can grow with the square of the node count. Override put_nodes_into_subgraphs with a dependency counter: each node keeps the number of its dependencies not yet placed, and becomes ready when the count reaches zero. The time is now proportional to the nodes plus their dependency entries. A subgraph still takes every node of its kind that becomes ready, so the subgraphs hold the same nodes as before. Only the order inside a subgraph can change, and tagging and splitting do not depend on it. On that model, building the subgraphs now takes about 1 second; on just its first 80k nodes the old loop took about 5 minutes. On a synthetic 62k-node graph it goes from 134 s to 0.1 s, with the same subgraphs and the same split module. The new test checks that the subgraphs match the base class result and that the dependency sets are read a bounded number of times.
…omposed
With decompose_attention=True, scaled_dot_product_attention is replaced by
matmul and softmax, which return a contiguous tensor. The fused kernels lay
their output out differently: flash and cuDNN attention follow the query's
dimension order, and memory-efficient attention returns a transpose of an
(N, L, H, E) buffer.
An exported graph can depend on that layout. In the common pattern
out = F.scaled_dot_product_attention(q, k, v)
out = out.transpose(1, 2).contiguous().view(n, l, -1)
the transpose of the fused output is often already contiguous, so export
does not record the .contiguous(). After the decomposition the transpose
is not contiguous any more, and torch_tensorrt.dynamo.compile fails in
run_decompositions with "Cannot view a tensor with shape ... and strides
... as ...". The Florence-2 vision tower hits this.
The decomposition cannot tell which backend produced the layout, but the
attention node records it. A new pre-lowering pass, which runs only when
attention is decomposed, copies each attention output whose recorded
dimension order is not the default into that order (permute, contiguous
clone, permute back, then a clone that keeps this layout). It also visits
subgraphs: torch.cond branches and torch.no_grad() or autocast regions.
The original attention node becomes that last clone, so graph outputs,
in-place ops on the output and in-place updates through its views keep
working on the same tensor.
torch.compile with decompose_attention=True still fails the same way: its
backend does not run the pre-lowering passes, so it needs a separate fix.
post_lowering lowers each torch.cond branch by calling itself on the branch module, but it ignored the module that call returned. Most passes edit the graph in place, so this went unnoticed. The complex decomposition does not: it retraces the branch and returns a new module. That module was dropped and the parent kept the original branch, complex ops included, which TensorRT cannot convert. A conditional with complex math inside a branch therefore stayed in PyTorch, and require_full_compilation=True failed. post_lowering now installs the returned module in place of the branch, under every name the parent uses for it, and lowers a branch shared by several conditionals once. The cond node was traced against the old branch, so the new module is installed only if its inputs and outputs keep their dtypes and memory layout: strides, storage offset and storage size. A branch that returns a complex tensor changes its outputs to the real layout, so it is left as it was. So is a branch whose retraced output comes back in a different memory layout, for example contiguous where the original was transposed, because code after the conditional can read the output through its strides. Symbolic values that cannot be proven equal count as different, so data-dependent shapes keep the original branch instead of failing. A branch that takes a complex operand, such as a complex buffer, still runs in PyTorch as before.
With offload_module_to_cpu=True, constant folding keeps the folded constants on the CPU to save GPU memory, while their metadata still records the GPU. complex_decomposition_adapter then retraces the graph with those CPU tensors next to GPU activations, and the retrace fails: FakeTensorDeviceMismatchError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! The adapter falls back to the legacy complex pass. That pass fails the same way on a small rotary embedding model, and on a large model it did not finish within 45 minutes. So models with complex rotary embeddings compiled with offload_module_to_cpu either fail or take the slow path. The adapter now copies each offloaded constant to the device its metadata records while the retrace runs, then puts the CPU tensors back, both in the module it was given and in the module it returns, and frees the copies. Where constants live after lowering does not change. A successful retrace returns a new module, which no longer holds the parameters that only constant folding read. compile() and cross_compile_for_windows() now offload the module lowering started from as well as the lowered one, so those parameters still leave the GPU.
With truncate_double=True, TensorRT runs float64 work in float32, and the float64 ops around a cast already convert. The cast to float64 itself did not: the _to_copy validator only accepts dtypes TensorRT has, so every x.to(torch.float64) stayed in PyTorch and cut the graph. A rotary embedding that rotates in float64 in every attention block, as Wan 2.2's rope_apply does, came out split into many engines, with a PyTorch cast between each pair. The validator now accepts a float64 cast of a 16, 32 or 64-bit float when truncate_double is set, and impl.cast.to_copy casts it to float32 instead, the same way the cumsum and cat converters already treat a float64 request. Other float64 requests that reach impl.cast.to_copy, such as a copy into a float64 buffer, get the same treatment when the source is a float, bool or 8-bit integer, which float32 holds exactly, while a wider integer source still fails as before, because float32 cannot hold large integers exactly. _to_copy casts to float64 from any other dtype still stay in PyTorch, as before; TensorRT cannot cast some of them, such as float8, at all. So do casts whose PyTorch result is not contiguous (channels_last, or preserving the layout of a non-contiguous input), because the TensorRT copy is always contiguous. With dynamic shapes, a cast that names a memory_format stays in PyTorch too, even in the middle of a graph: the validator cannot know where partitioning will put it, and if it ends up alone on an engine input, the placeholder copy converter refuses an explicit memory_format and the build would fail. This relies on pytorch#4898, which gives every engine output position its own tensor. Without it, an engine that returns both a tensor and its float64 cast gets one binding for the two outputs, and the caller silently receives wrong values in the wrong shape.
When attention is decomposed, the decomposition marks the aranges behind the attention mask so that constant folding leaves them alone. The mark is node metadata. When a graph has complex operations, the complex decomposition adapter retraces the whole graph and then folds constants again. The retrace builds new nodes without that metadata, and the graph-level rule cannot find the pattern again because the attention op is already gone. So the aranges are folded after all, and a decomposed 1024x1024 causal mask becomes an engine constant (4 MB in float32) instead of being computed. A second run_decompositions, for example on a program the user already decomposed, loses the marks the same way. Tag the aranges in node.meta["custom"] instead, which retraces copy to the new nodes, and let the attention_mask_arange rule select tagged nodes. Disabling the rule still makes them foldable.
The Python runtime copies each shape-tensor input into a NumPy array and hands TensorRT the array's address. The array was a loop variable, so it was freed as soon as the next input was handled, before the engine ran. A later allocation could reuse that memory: with three shape inputs, the third array often lands where the first was, and TensorRT then reads the third value for the first input. The engine fails with a shape error or runs with the wrong shapes. torch.compile with dynamic shapes passes symbolic sizes as shape inputs, so it is the usual way to hit this. Keep the arrays in a list on the engine until the next call sets up its inputs.
…scattered extent The scatter_add decomposition loops over every slice of src along dim and emits a select, unsqueeze, zeros constant, scatter and add for each one, so a scatter_add over n slices becomes about 8n graph nodes and n constants the size of the input. A mixture-of-experts router that counts tokens per expert with an int64 scatter_add grows by about 290,000 nodes, and the TensorRT build then takes hours. TensorRT scatter overwrites colliding indices, which is why the loop scatters one slice at a time. For int32 and int64, compare every index with every output position along dim and sum the matching src values instead. Integer addition gives the same result in any order, so the output is exact. As with TensorRT gather and scatter, indices must be in range: an out-of-range index is ignored rather than raising. The graph no longer grows with n. Only the part of src covered by index is used, as in PyTorch. The sum multiplies by the 0/1 match mask because the where converter narrows int64 to int32. The compare adds one axis, so an index of rank 8, the TensorRT limit, keeps the loop. The comparison holds index.numel() * input.shape[dim] elements, so above 2^26 of them, well under TensorRT's tensor volume limit, the loop is kept too, as it is when the bound cannot be proven, for example for a dynamic dim with no upper bound; those cases keep the large graphs described above. The loop now runs over the extent of index along dim, which is all PyTorch reads, so a longer src no longer raises an IndexError, and filter_decomposition_table checks that extent. Every other dtype keeps the loop decomposition. Floating point sums would round differently in another order, bool scatter_add is a logical or rather than a count, and the converters cannot compute the narrow integer types. When index has a dynamic size along dim, the op stays in PyTorch or the loop specializes it. With a constant index, constant folding would turn the whole comparison into one engine constant. A new constant-fold exclusion rule, scatter_add_positions, keeps the positions arange out of folding, like attention_mask_arange does for attention masks. The decomposition tags the arange in node.meta["custom"], which FX retraces keep, so the rule still finds it after the complex decomposition retrace or a second run_decompositions. The rule applies only when src varies at runtime (it reads an input or a random op): with constant src and index, the update sum still folds to an output-sized constant. For a 128x64 constant index into a 4096x64 int64 output, the engine is about 90 KB, compared with about 3.2 MB for the loop. A few updates into a very wide output can run slower: 16 rows into a 262144x8 output took 0.47 ms instead of 0.36 ms.
A boolean mask whose True count is only known at runtime makes the number of updated rows a runtime value N in the index_put converter. Values and indices with a static row count then failed or gave wrong results: - h[mask, :] += v on a 3-D h (mask over two dims), as transformers' Qwen3-VL adds its visual embeddings in _deepstack_process, failed the build with "Cannot broadcast (128, 2048) to (<ITensor>, 2048)". This breaks compiling GR00T, LaWAM and VLA-JEPA. - h[mask] = v on a 2-D h compiled, but every row came from v's first row. - h[mask, idx] = v with a static idx failed the same way as the first case, and once values were fixed, wrote every row at idx[0]. The cause is that impl.slice.expand broadcasts a static dim whose target is a runtime value, which is only right for a dim of size 1. PyTorch requires values and indices to have N rows there, so the converter now keeps such a static first dim instead of expanding it to N. Dims of size 1 still broadcast, other dims and static shapes are unchanged. A mask whose True count differs from the number of value rows now fails at runtime, as PyTorch does, instead of writing part of the rows.
Torch-TensorRT marks an engine as needing an output allocator as soon as one of its operations has a data-dependent shape, such as aten.index.Tensor with a boolean mask. Its runtimes use that mark to keep CUDA graphs off, which is right: TensorRT cannot capture such an engine even when every output shape is known. ExecuTorch export read the same mark as "an output has a data-dependent shape" and refused the engine with "ExecuTorch export does not support TensorRT engines that require an output allocator". That is too broad. Qwen3-VL adds its visual embeddings with h[mask, :] += v: the selected rows have a data-dependent count, but the engine returns a tensor shaped like h. GR00T and LaWAM compiled to TensorRT but could not be saved to ExecuTorch. The ExecuTorch backend allocates static outputs and falls back to plain enqueue when a CUDA graph capture fails, so it can run such an engine. For a marked engine, export now asks the live engine for its output shapes at the node's example input shapes (infer_outputs, the call the Python runtime's fake kernel makes) and keeps refusing it only if an output shape is still unknown. Otherwise the engine's ExecuTorch record carries no mark, so the backend accepts it. Nothing is serialized for the check; it creates the engine's execution context if no run has created it yet. An engine whose input shapes have no example value, or that cannot report its outputs, is still refused, and so is a marked engine handed to the backend without this export step.
TensorRT allows uint8 tensors only at engine inputs and at outputs made by a cast or identity layer. The select converter emits a gather, so selecting from a uint8 input, such as frames[1] on a uint8 camera history before casting it to float, made a uint8 intermediate and failed the engine build: Tensor [GATHER]-[aten_ops.select.int]-[/select_gather]_output has datatype UInt8 but is not network-level IO A select whose uint8 result is itself an engine output failed the same way. VQ-BeT hits this with its uint8 image history. aten.gather already keeps a uint8 input in PyTorch for the same reason, and select now does too. Slice, permute, view and index on a uint8 input still fail; they need the float cast moved ahead of them, which is a separate change.
select_validator read only node.meta["val"]. A graph typed by torch.fx ShapeProp instead of torch.export, for example one built with symbolic_trace, carries tensor_meta but no val, so the check missed a uint8 input there. It now falls back to tensor_meta, as gather_validator already does. This builds on the previous commit, which added the check.
shoumikhin
force-pushed
the
dynamo-converter-lowering-fixes
branch
from
October 29, 2026 18:26
1bd5cd1 to
67540ed
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.
Twelve fixes found compiling LeRobot policies. Each fix cherry-picks alone onto
release/2.15(fix 12 is two commits).1. Out-of-range slice bounds
On a size-10 dimension,
x[-15:]becamex[5:]; static bounds now clamp like Python.2. Quadratic fast partitioner
Building subgraphs for 345,000 nodes took over 75 minutes; now about one second.
3. Decomposed attention layout
Export can drop
.contiguous()after fused attention, breaking a laterviewon decomposed attention; a pass restores it.4.
torch.condbranches after loweringpost_loweringdropped lowered branches; they are kept when dtypes and memory layouts match.5. Offloaded constants in the complex decomposition
With
offload_module_to_cpu=True, the retrace failed on folded CPU constants, now moved to their device for it.6. float64 casts with
truncate_double=Truex.to(torch.float64)split graphs; on float inputs it now becomes a float32 cast.7. Attention-mask fold exclusions
Retraces dropped marks keeping mask aranges unfolded, making a causal mask a 4 MB constant; marks now live in
node.meta["custom"].8. Freed shape-tensor inputs
The Python runtime freed shape inputs too early, so one could read another's value.
9. Integer
scatter_addwithout a loopThe loop made a graph of 345,000 nodes. Eligible int32 and int64 cases now compare and sum, up to 2^26 comparisons;
scatter_add_positionskeeps that comparison unfolded whensrcvaries at runtime.10.
index_putwith a data-dependent boolean maskh[mask, :] += vfailed to build;h[mask] = vwrote every row fromv[0]. Static multirow operands no longer broadcast to the runtime count.11. ExecuTorch export with data-dependent shapes inside
Export refused engines with any data-dependent operation; now it accepts them when every output shape is known. They run without CUDA graph capture, so they are slower: GR00T on H100: 20.2 ms against 13.9 ms with fixed-shape model code.
12.
selecton a uint8 tensorTensorRT allows uint8 only at engine inputs and outputs, so
frames[1]on uint8 failed the build; it stays in PyTorch, likegather.Testing
release/2.15.scatter_addcalls: compile 10 s vs 160 s, run 0.24 ms vs 20 ms.