TensorRT: LeRobot policies that do not run yet, and what blocks them

This is a meta report, not one edit: read #4953 first, then reproduce item 3's block-causal mask snippet and item 5's uint8 permute snippet. For item 1 start in complex_decomposition_adapter.py around the second constant_fold; for item 4 trace view_decomposition in _decompositions.py and the rank limit in _global_partitioner.py. Done means each listed blocker is either fixed and re-measured or explicitly ruled out.

We exported the LeRobot policies to ExecuTorch with the Torch-TensorRT backend, and ran each bundle with lerobot-rollout on a fake robot. This issue lists what still keeps policies from working on the TensorRT path, so it can be picked up. It is the TensorRT companion of #24325.

Results

Milliseconds per action chunk, batch size 1, median, measured inside lerobot-rollout.

• Stack: ExecuTorch release/1.6 with several unmerged fixes, including #24326, Torch-TensorRT release/2.15 with #4950 and #4953 (up to its "views as aten.reshape" commit), torch 2.15.0, TensorRT 11.4.
• GPUs: H100, A100 (shared with other jobs), Jetson AGX Thor, Jetson Orin Nano 8 GB.
• In every run with a number, PyTorch was never loaded, all program inputs and outputs stayed on the GPU, and there were zero CPU operators. Nsight Systems counted 2 to 4 uploads (the observation) and at most one download (the actions) per chunk. †: copies not traced.
• "Programs": what one bundle holds. The CUDA backend program draws the starting noise on the GPU.
• T<n>: Torch-TensorRT item n below. L: needs a LeRobot export change. P: PyTorch. R: refused by the exporter on purpose. p: queued, not measured yet. n/m: not measured in this run.

Policy	TensorRT H100	TensorRT A100	TensorRT Jetson Thor	TensorRT Jetson Orin Nano	Programs
act	3.4	4.3	8.5	26.0	TRT 1
diffusion	533	740	1259	3103 †	TRT 3 + CUDA 1
dm05	L	L	n/m	n/m
eo1	L	L	L	n/m
evo1	R	R	R	n/m
fastwam	T1	T1	n/m	n/m
fineart_vla	n/m	n/m	n/m	n/m
flux3	n/m	n/m	n/m	n/m
gaussian_actor	0.6	1.1	0.6	1.3 †	TRT 1 + CUDA 1
groot	18.6 *	30.0	58.5	197	TRT 10 + CUDA 1
lawam	P	P	n/m	n/m
lingbot_va	R	R	n/m	n/m
lingbot_vla_v2	L, then T4	L, then T4	n/m	n/m
molmoact2	L	L	n/m	n/m
multi_task_dit	79.7	98.7	193	n/m	TRT 3 + CUDA 1
pi0	121	271	n/m	p	TRT 1 + CUDA 1
pi05	31.7 *	59.0	91.2	446	TRT 9 + CUDA 1
pi0_fast	R	R	R	n/m
smolvla	18.0	29.1	37.3	176	TRT 4 + CUDA 1
tdmpc	R	R	R	R
vla_jepa	L	L	n/m	n/m
vqbet	2.3	3.3	3.5	7.5 †	TRT 1 + CUDA 1
wall_oss_05	L	L	n/m	n/m
wall_x	P	P	n/m	n/m
xvla	66.4	133	398 †	p	TRT 1 + CUDA 1

* H100 runs pi05 and groot with their LIBERO fine-tunes. The other GPUs use SO-101 fine-tunes.

Accuracy over 16 frames against eager PyTorch, as a fraction of the action range: the float32 policies, where measured, are within 6e-6 (pi0: 3e-5, with eager's own spread at 1e-2). smolvla, pi05 and groot (bf16) are at 0.0095 to 0.043, which is 0.7x to 1.8x eager's own spread. The LIBERO fine-tunes on the H100 are further off: pi05 0.085 (3x), groot 0.06 (eager repeats exactly).

On the Orin Nano with the plain release stack, pi05 exported with 478 MiB free at its worst moment, and groot with 1514 MiB. The pi0 base model (3B, float32) ran out of memory during export on the plain release stack. Its run here is pending.

#24325's TensorRT columns show more policies passing. They came from an earlier sweep (torch 2.16 nightly, ExecuTorch and Torch-TensorRT main) with extra LeRobot export branches, and they timed only the compiled model call, without saving and loading the program. We think most of the gap comes from the LeRobot side and the measuring method, not the library versions, but we did not measure that directly.

Already fixed or in review

• pytorch/TensorRT#4936 (merged, in release/2.15): index_put for the Qwen3-VL DeepStack mask writes, which groot needs.
• pytorch/TensorRT#4953 (in review): saving a program with no inputs (assert fake_mode is not None in lift). This moves the noise draw to the GPU (pi05: 3 uploads per chunk instead of 4). Its runtime build number commit was not checked by a policy run.
• pytorch/TensorRT#4950 (in review): pins ExecuTorch 1.6.0 on release/2.15.
• #24326 (in review), for the CUDA program in the bundles: the CUDA random shims are seeded from std::random_device. Before, 4 processes started together drew the same noise. Now all 4 differ. For CUDA-only bundles (TensorRT was already fine): boolean-mask copies use where (SmolVLA: 7 device-to-host copies per chunk down to 1), and tensor checks are dropped before compile (the Gaussian actor now passes). float64 weights and aten._reshape_copy support were pushed after our build and are not measured yet.

What blocks TensorRT

1. Complex lowering mixes a CPU constant with CUDA scalars (fastwam)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

It fails in the second constant_fold of complex_decomposition_adapter, on where(<98x98 bool constant on CPU>, scalar_tensor(0.0, cuda), scalar_tensor(-inf, cuda)). It needs offload_module_to_cpu=True, which LeRobot sets to leave about 600 MiB more free on an 8 GB Jetson. Small graphs with complex RoPE and a mask of the same size compile, so there is no minimal repro yet. Status: discovered, not attempted.

2. Engine build fails on a uint8 cast (fastwam)

With offload_module_to_cpu=False the folds pass, then:

MyelinCheckException: formats.cpp:3055: CHECK(def_tactic_desc) failed
Error Code 10: Internal Error (Could not find any implementation for node
  {ForeignNode[Cast ITensor arg1_1 from DataType.UINT8 to DataType.FLOAT ...]})

arg1_1 is the uint8 camera frame, so this may share a cause with item 5. Status: discovered, not minimal, not re-run with #4953.

3. Silent wrong results: a block-causal mask becomes plain causal when the graph has complex numbers

n = torch.arange(98, device="cuda")
mask = (n[:, None] // 49) >= (n[None] // 49)   # each frame sees itself and earlier frames
q, k = complex_rope(q), complex_rope(k)         # view_as_complex * freqs -> view_as_real
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

TensorRT differs from PyTorch by 1.01, and matches PyTorch with a plain causal mask to 1.3e-6. Without complex RoPE it is right. FastWAM's video mask is wrong the same way (1.17). As #4953 explains it, the complex lowering retraces SDPA so the mask becomes attn_bias, and force_causal_efficient_attention, on by default, replaces every square bias with is_causal=True. Status: being fixed in #4953 (only a proven causal bias is replaced). Not in our build yet, so we do not trust TensorRT results for policies with complex RoPE.

4. Views above 8 dimensions reach the CUDA backend as aten._reshape_copy (lingbot_vla_v2)

Method forward missing fallback kernels: aten::_reshape_copy

view_decomposition turns every aten.view into aten._reshape_copy before partitioning. A view above the rank limit, such as the 10-D Qwen3-VL patchify x.view(2,1,2,3,8,2,16,8,2,16), stays outside the engine and goes to the CUDA backend, which has no kernel for it. Status: the first fix in #4953 (views go back to reshape) was in our build. A follow-up in #4953 and a _reshape_copy decomposition in #24326 were not. Not re-measured: LingBot now stops earlier, in LeRobot.

5. A shape op on a uint8 input before a cast fails

img.permute(2, 0, 1).float()   # img is a uint8 graph input
# RuntimeError: Could not get tensor shape.

Casting first works. Transformers' torch image processors permute or stack uint8 frames before converting, so on plain release/2.15 no Qwen-VL pixel path reaches TensorRT. Status: fixed in #4953 (uint8 shape ops stay in PyTorch). That commit was in our build, but the repro was not re-run with it.

Not bugs

On the Orin Nano, the CUDA backend's NoValidChoicesError in Inductor autotune was memory pressure (142 MiB free). On an idle board, ACT and VQ-BeT export and pass. No is_big_gpu change was needed.

Not a TensorRT problem

LeRobot export changes (these rows pass in #24325, which used extra branches):

• eo1, wall_oss_05: need the fixed Qwen2.5-VL image grid and per-policy changes. They stop at .numpy() is not supported for tensor subclasses in the slow image processor.
• vla_jepa, lingbot_vla_v2: need the fixed Qwen3-VL grid. They stop at data-dependent guards (vla_jepa, lingbot).
• molmoact2: span mask from tensor ops, crops as one sequence. dm05: the Gemma3 processor out of the program, state in a fixed-length prompt (here).
• Refused on purpose: evo1 (uniform noise draws), tdmpc (history and planner state), pi0_fast (FAST decoding in Python), lingbot_va (streaming KV cache). For evo1, tdmpc and pi0_fast, #24325 exported only the model call.
• multi_task_dit on the Orin Nano and fineart_vla: our test checkpoints had no processor files. flux3: NATTEN was not installed here, and #24325 found that its custom ops have no ExecuTorch kernels.

PyTorch (seen only because the program is saved and loaded):

• torch.export.load loses val on autocast regions: lawam (here), and next evo1, eo1, molmoact2.
• torch.export.save cannot serialize flat_apply over ModIndex (transformers' or_mask_function masks): wall_x (here), and dm05 next.

ExecuTorch, CUDA-only bundles: complex numbers do not reach a .pte (fastwam), and scatter_reduce with amax has no CUDA shim (lingbot_vla_v2; no fix yet).

cc @Gasoonjia @digantdesai