Repository navigation
CUDA backend: per-program weights file names and antialiased resize - #24314
Merged
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24314
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 Cancelled Jobs, 1 Unrelated FailureAs of commit c46538e with merge base e62fd44 ( CANCELLED JOBS - The following jobs were cancelled. Please retry:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
shoumikhin
force-pushed
the
cuda-names-resize
branch
from
October 30, 2026 00:40
d6801e5 to
c46538e
Compare
shoumikhin
force-pushed
the
cuda-names-resize
branch
from
October 30, 2026 02:31
c46538e to
295688c
Compare
| ``` | ||
| This generates `.pte` and `.ptd` files that can be executed on CUDA devices. | ||
|
|
||
| The weights go to a data file named `aoti_cuda_blob.ptd`. To export several programs into one folder, give each program its own |
Contributor
There was a problem hiding this comment.
Its tech debt that we use .ptd at all we should just use the named data map inside the .pte.
JacobSzwejbka
approved these changes
Oct 30, 2026
shoumikhin
force-pushed
the
cuda-names-resize
branch
from
October 30, 2026 02:54
295688c to
e9bd579
Compare
The CUDA backend writes every program's weights to a file named
aoti_cuda_blob.ptd. Exporting two programs into one folder therefore
makes the second overwrite the first one's weights. Since the weights
are keyed by their content, the first program then fails to load
instead of running with the wrong weights. Renaming each file right
after its export works, because the program does not store the file
name, but it is an extra step that is easy to miss.
Add a data name option. A program lowered with
CudaPartitioner([..., CudaBackend.generate_data_name_compile_spec("encoder")])
writes its weights to encoder.ptd. Without the option the file is still
aoti_cuda_blob.ptd, so existing exports, scripts and runners are
unchanged, and existing files load as before.
The name becomes a file name, so it may use only up to 200 ASCII
letters, digits, '_', '-' and '.'. It may not start with '.' or '-'
(a leading '-' reads as an option, for example in merge_ptes
--input-ptd), may not be a Windows device name such as NUL, and may not
end with '.ptd'. One of the two data file writers adds '.ptd' only when
it is missing, so "encoder" and "encoder.ptd" would write the same
file. The check also runs where the name is used, so a hand-written
compile spec cannot point outside the output folder. All methods of one
program share one weights file, so they must use the same name. A
backend that keeps its data in the .pte, such as Metal, refuses the
option instead of ignoring it. merge_ptes ignores the name when it
compares variants, like cuda_include_ptx, so variants exported with
different names still merge.
The CUDA backend docs describe the option, its rules, and how to load
a program with its data file. The TensorRT page no longer says the
name is fixed or that a clash gives wrong answers, and its tip on
sharing the CUDA data file with the engines covers a custom name.
A new test in test_missing_weights_blob.py exports two programs with
their own names, and one with no constants, into one folder, checks
the data file names, then loads and runs the first two against eager
PyTorch. test_cuda_export.py checks the name rules, including
non-ASCII letters, path separators, a leading '-' and a '.ptd' ending,
checks that good names and the default are accepted, that a backend
which keeps its data in the .pte refuses a name, and the one name per
program rule. The existing merge test now merges variants with
different names.
shoumikhin
force-pushed
the
cuda-names-resize
branch
from
October 30, 2026 03:15
e9bd579 to
73976de
Compare
Contributor
Author
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 32m 19s —— View job Needs discussion: I found no problems in the per-program weights file names or the antialiased-resize rewrite, but the new Detailed reviewPR Review: #24314SummaryThe PR adds a The data-name and resize changes raised no findings. The metadata bit stops older runtimes from loading these programs, and it misses some methods it is meant to cover. Three independent reviewers read all 16 files, and I re-checked each finding below against the code. Nothing was run, because this job has no GPU and could not run Python. Backward Compatibility
Code Quality
RecommendationNeeds Discussion Before merging, decide where the What was done: three independent reviewers read all 16 changed files, and I fact-checked each finding against the code. What was not done: nothing was built or run (this job has no GPU and could not run Python), and no code was changed or pushed. |
F.interpolate(x, mode="bilinear" or "bicubic", antialias=True) becomes aten._upsample_bilinear2d_aa or aten._upsample_bicubic2d_aa. Inductor has no lowering for either, so it calls ATen's kernel, and AOTInductor has no C shim for them. Lowering to the CUDA backend failed with "missing fallback kernels". Image preprocessing often does this. The backend already rewrites ops like this in get_decomposition_table (conv1d as conv2d, for example), so this adds the bilinear and bicubic antialiased resizes there, for float32, float16 and bfloat16 inputs. The backend has no float64, and ATen's CUDA kernel has no integer version. An antialiased resize is separable, so the rewrite resizes the width, then the height. Along each, it computes every output's window and weights with ATen's formulas: the filter widened by the downscale factor, cut to its support, and normalized. Then it sums the window's inputs with those weights. Taps outside the window have zero weight and are masked out before the sum, so a NaN or inf reaches only the outputs whose window holds it, and the work grows with the window size, not the input size. Where a window edge falls exactly on an input, float rounding can put that input in or out of the window differently from ATen; its weight there is zero or close to it. Symbolic shape math keeps dynamic input and output sizes working. The number of taps per output is a shape, so it is computed with integers: in float32 the generated kernel can round a ratio such as 21 / 7 just above 3 and count more taps than the host allocated. For align_corners, the ratio for an output size of 1 comes from min and max, not a conditional, because torch.compile cannot handle a conditional size expression. The result keeps the input's memory format. The .vec overloads reuse torch's decomposition. The rewrite computes positions and weights in float32 with its own rounding, so its float32 results are not bit-identical to ATen's kernel: for inputs in [0, 1) they stayed within about 2e-5 of it on the sizes tried. Half and bfloat16 inputs are computed in float and rounded once at the end, so they can differ slightly from ATen, which keeps its weights in the input type. The tests are in a new file, test_antialiased_resize.py, because they run programs through the Python runtime. They lower bilinear and bicubic resizes whose input and output heights are dynamic, run them for heights that shrink, grow, go to 1, and have whole ratios, with a NaN in the input, and compare with eager CUDA. Other cases cover an uneven scale_factor, float16 and bfloat16 with a channels_last input, the rewrite's memory format, and torch.compile of an align_corners resize with dynamic sizes.
shoumikhin
force-pushed
the
cuda-names-resize
branch
from
October 30, 2026 03:35
3aec81d to
135614a
Compare
Contributor
Author
|
@pytorchbot cherry-pick --onto release/1.6 -c critical |
Collaborator
Cherry picking #24314The cherry pick PR is at #24324 and it is recommended to link a critical cherry pick PR with an issue. Details for Dev Infra teamRaised by workflow job |
shoumikhin
added a commit
that referenced
this pull request
Oct 30, 2026
…24314) This change fixes several problems in the CUDA backend. Each fix is its own commit. ## Name the weights file of each program The backend saves every program's weights to `aoti_cuda_blob.ptd`. Exporting several programs into one folder overwrites the weights, and the earlier programs then fail to load. Renaming each file after export works, but is easy to miss. A new option names the file: ```python partitioner = CudaPartitioner([ CudaBackend.generate_method_name_compile_spec("forward"), CudaBackend.generate_data_name_compile_spec("encoder"), ]) # writes encoder.ptd ``` Without the option nothing changes. The name is checked: up to 200 ASCII letters, digits, `_`, `-` and `.`, not starting with `.` or `-`, not ending with `.ptd`, and not a Windows device name such as `NUL`. All methods of a program must share it. A backend that keeps its data inside the `.pte`, such as Metal, refuses it. ## Lower antialiased resizes `F.interpolate(x, mode="bilinear" or "bicubic", antialias=True)` failed to lower with "missing fallback kernels". The backend now rewrites it into plain tensor math, with PyTorch's window and weight formulas, for dynamic sizes and float32, float16 and bfloat16. The results are not bit-identical: in float32 they stay within about 2e-5 of PyTorch's kernel for inputs between 0 and 1. ## Tests Ran on an A100 GPU (Linux x86_64): `test_cuda_export.py`, `test_antialiased_resize.py`, `test_missing_weights_blob.py`, `test_merge_ptes.py` and `test_cuda_weight_metadata.py`. Each test added for a fix fails when that fix is removed.
shoumikhin
added a commit
that referenced
this pull request
Oct 30, 2026
shoumikhin
added a commit
that referenced
this pull request
Oct 30, 2026
shoumikhin
added a commit
that referenced
this pull request
Oct 30, 2026
shoumikhin
added a commit
that referenced
this pull request
Oct 30, 2026
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 2, 2026
…ckend An operator Torch-TensorRT cannot convert used to fall back to a portable CPU kernel, with copies to and from the GPU around it. The executorch_tensorrt backend now hands those operators to ExecuTorch's CUDA backend, so a program runs entirely on the GPU. The CUDA backend writes its weights into a data file beside the program. Each part names that file after itself, through the backend's data-name compile spec, so parts saved into one bundle folder keep their own weights, and the file is recorded as the program's data file. The data-name compile spec comes from pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Tested on an NVIDIA GPU: a model with an operator TensorRT rejects compiles to TensorRT and CUDA delegates with no operator left on the CPU and its input and output in GPU memory, and matches PyTorch when run without PyTorch.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 2, 2026
executorch_cuda lowers each part with ExecuTorch's own CUDA backend, so a bundle runs on an NVIDIA GPU without TensorRT. ExecuTorch compiles the kernels ahead of time for the GPU that exports. It is one config class on ExecuTorchConfig, and its registration. ExecuTorch loads its programs, so the runtime, the engine and lerobot-rollout are unchanged. The CUDA backend keeps a program's kernels and weights in a separate data file, named after the part, as with executorch_tensorrt, so two parts never overwrite each other's weights and the bundle names one data file per program. Naming that file needs pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Inputs and outputs stay in GPU memory too, so a program reads an earlier one's outputs in place. A CUDA bundle records the GPU that built it, and is refused on another. The kernels are compiled in full float32, as lerobot-export compiles for every backend. PyTorch lets convolutions use TF32 by default, and the convolution kernels the CUDA backend generates follow that flag. On ACT on an H100 they then drifted about seven times further from PyTorch on the CPU than PyTorch's own GPU kernels do. In full float32 the actions match to about 1e-6, and a chunk still runs twice as fast as eager PyTorch on the GPU. It needs the CUDA build of ExecuTorch, and the CUDA compiler (nvcc) at export. Like the other ExecuTorch backends, it reads programs traced by PyTorch 2.15 or newer. Like the backends that take only part of a program, it folds the values a program computes from constants alone before lowering. SmolVLA's vision encoder picks its position ids with a boolean mask made from constants. Left in the program, that indexing makes the GPU read a size back to the host on every run; folded, it is gone. Folding has one cost today: SmolVLA's time embedding, computed in float64 from constants, becomes a float64 weight, and ExecuTorch's CUDA backend cannot load float64 weights yet, so SmolVLA waits for that fix on this backend. Tested on an H100 and on a Jetson AGX Thor: exported ACT, checked the bundle, and ran it with lerobot-rollout on a fake robot. On the H100, pi0 exported too, and its actions matched eager PyTorch to about 1e-5 on ten frames; in float32 it runs a little slower than eager. The check that a bundle runs without PyTorch now also runs a CUDA bundle.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 2, 2026
…ckend An operator Torch-TensorRT cannot convert used to fall back to a portable CPU kernel, with copies to and from the GPU around it. The executorch_tensorrt backend now hands those operators to ExecuTorch's CUDA backend, so a program runs entirely on the GPU. The CUDA backend writes its weights into a data file beside the program. Each part names that file after itself, through the backend's data-name compile spec, so parts saved into one bundle folder keep their own weights, and the file is recorded as the program's data file. The data-name compile spec comes from pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Tested on an NVIDIA GPU: a model with an operator TensorRT rejects compiles to TensorRT and CUDA delegates with no operator left on the CPU and its input and output in GPU memory, and matches PyTorch when run without PyTorch.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 2, 2026
executorch_cuda lowers each part with ExecuTorch's own CUDA backend, so a bundle runs on an NVIDIA GPU without TensorRT. ExecuTorch compiles the kernels ahead of time for the GPU that exports. It is one config class on ExecuTorchConfig, and its registration. ExecuTorch loads its programs, so the runtime, the engine and lerobot-rollout are unchanged. The CUDA backend keeps a program's kernels and weights in a separate data file, named after the part, as with executorch_tensorrt, so two parts never overwrite each other's weights and the bundle names one data file per program. Naming that file needs pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Inputs and outputs stay in GPU memory too, so a program reads an earlier one's outputs in place. A CUDA bundle records the GPU that built it, and is refused on another. The kernels are compiled in full float32, as lerobot-export compiles for every backend. PyTorch lets convolutions use TF32 by default, and the convolution kernels the CUDA backend generates follow that flag. On ACT on an H100 they then drifted about seven times further from PyTorch on the CPU than PyTorch's own GPU kernels do. In full float32 the actions match to about 1e-6, and a chunk still runs twice as fast as eager PyTorch on the GPU. It needs the CUDA build of ExecuTorch, and the CUDA compiler (nvcc) at export. Like the other ExecuTorch backends, it reads programs traced by PyTorch 2.15 or newer. Like the backends that take only part of a program, it folds the values a program computes from constants alone before lowering. SmolVLA's vision encoder picks its position ids with a boolean mask made from constants. Left in the program, that indexing makes the GPU read a size back to the host on every run; folded, it is gone. Folding has one cost today: SmolVLA's time embedding, computed in float64 from constants, becomes a float64 weight, and ExecuTorch's CUDA backend cannot load float64 weights yet, so SmolVLA waits for that fix on this backend. Tested on an H100 and on a Jetson AGX Thor: exported ACT, checked the bundle, and ran it with lerobot-rollout on a fake robot. On the H100, pi0 exported too, and its actions matched eager PyTorch to about 1e-5 on ten frames; in float32 it runs a little slower than eager. The check that a bundle runs without PyTorch now also runs a CUDA bundle.
shoumikhin
added a commit
that referenced
this pull request
Nov 2, 2026
… input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups (#24326) ## What is wrong today With `enable_cuda_graph_for_method` set, the CUDA backend runs a method a few times normally, captures a CUDA graph on the fourth call, and replays that graph after it. Some calls go wrong: - **A call that cannot be captured fails.** A mask index, such as `x[mask]` or `x[mask] = v`, reads a count back to the host, and a capture cannot record that. The capture call returned an error. An output that is first strided at the capture call failed the same way. A `.item()` on a device tensor, such as a `torch.cond` predicate computed on the GPU, ended the process. - **A replay can return a wrong result with no error.** Replay checked only each input's byte count. An input that changes from `(2, 3, 4, 4)` to `(2, 4, 3, 4)` keeps its byte count, so it replayed the captured shape and returned the wrong output. ## What this change does - Replay compares each input's shape with the captured one and refuses a different shape with `InvalidArgument`. The error names the dim and both sizes. A captured graph only works for the input shapes of the capture call, and the code comment on the graph phases now says so. - When a capture fails, the backend runs the same call again without a graph and returns its result. Graphs stay off for that method, as before. The retry is safe because a capture records kernels without running them, and it happens only if the new graph never started running. The `.item()` shims now refuse to read inside a capture, as the mask index does, so a capture that reads a device value is retried too instead of ending the process. - A method whose program keeps a constant in host memory, such as a `torch.cond` flag updated on every call, now runs without a graph. Host code runs only during the capture, so a graph would keep the branch of the capture call, and a retry after a failed capture would take the other one. The log names the constant found in host memory, or says the constants could not be listed. - An expected capture failure, such as the mask index refusing to run inside a capture, is logged at Info instead of Error, since the call then succeeds. A real failure of the call without a graph is still logged at Error. AOTInductor's generated code still prints its own line for the refused call. - The allocations a captured graph makes are recorded only after its first launch succeeds. Before, a failed first launch left addresses that were never allocated, and a later cleanup passed them to `cudaFree`. A failed capture now cleans up with `CudaGraphState::release()` instead of a copy of it. - These CUDA graph changes do not change the exported program, so older programs get them too. A program that uses the new shims described below (scatter, scatter_reduce, float64 or float16) needs a runtime that has them. Smaller commits: - The CUDA random shims seed from `std::random_device` instead of the clock, so two processes started in the same second no longer draw the same numbers. If `std::random_device` throws, or the copy of the seed to the GPU fails, the shim logs an error and frees what it allocated, and the next call tries again. - The wheel's shared library check now lists the mask index and `index_put` shims. - With PyTorch 2.14, a program that writes a one-element buffer, such as a counter, kept returning its first update, because Inductor folded the update into a constant. The AOTI backends (CUDA and Metal) now turn off Inductor's joint-graph constant folding, for the whole program, when a program writes a buffer. On a small Gemma 4 decoder that writes its KV cache, folding changes two nodes per graph; with it off, export and decode times stayed within run-to-run noise and the generated tokens were identical. - Unit tests for `aoti_torch_get_size`. - Comments on the index shims and on `set_output_data_ptr` that described accumulation, the sort count limit and the output index wrongly are corrected. - The mask index tests move to a new `test_index_shims.py`, since they need the Python runtime. They add a 2-D mask, and `index_put` with repeated indices in deterministic mode (half and bfloat16 accumulation, last write wins, an int32 index). - `test_output_shape_changes_at_capture` ends with a call that only a live graph rejects, so it fails if graphs are off. The KV cache recapture test checks that dropping the graph also clears the captured input shapes. ## Follow-ups to #24314 - Delete `test_align_corners_compiles_with_dynamic_sizes`. It could not fail: it passed with no resize rewrite at all and with the min and max in the align_corners ratio replaced by a plain `if`, because `torch.compile` compiles again for an output size of 1. With torch 2.14 and `CI` set it also failed for an unrelated reason, and it is the one failing test of the CUDA unit test job on main. `test_bicubic_antialias_resize` already runs an exported program at an output height of 1, where a plain `if` breaks. - The comment on that min and max now gives the real reason: an exported program keeps the branch it was traced with, so a plain `if` divides by zero at an output size of 1. - A backend that keeps its data in the `.pte`, such as Metal, now refuses a `data_name` compile spec built by hand too, before compiling anything. Before, only `generate_data_name_compile_spec` refused it, and a hand-built spec was ignored with no name check. - The CUDA docs warn that Windows and macOS usually treat data names that differ only in case as the same file. ## Copies through a boolean mask no longer read the mask count back to the host `out[mask] = src[mask]` ran as a mask index and a mask `index_put`. Each of them read the mask's count back to the host and waited for the GPU, on every call. Transformers' SmolVLM vision embeddings do this for their position ids. When the values are the same mask's picks from a tensor shaped like the destination, every element goes back to where it came from. So a new pass in the CUDA backend writes the copy as `torch.where(mask, src, out)`, which gives the same result, bit for bit. It does so inside `torch.cond` branches and other subgraphs too. It rewrites only when `where` would lay out its result like the destination, since `where` takes its layout from its inputs, the mask first, while `index_put` keeps the destination's layout. It also drops the index and the checks on its size when nothing else reads them and each check tests only the index's count and holds for any mask. A user check on the count, such as `torch._check(x[mask].shape[0] >= 1)`, keeps the index. The pass compares the two masks with a work list, so masks with long histories do not hit Python's recursion limit. On an A100, five calls of the SmolVLM position id code went from 10 count reads and 10 stream waits to none. ## Tensor checks no longer stop the CUDA export The argument checks in `torch.distributions` export as `aten._assert_async.msg`. The CUDA backend compiled each one into a device assert inside a Triton kernel. While compiling, AOTInductor runs every kernel once to pick its launch settings, on made-up buffers: random floats, and zeros for integers and bools. So a check that reads a bool from an earlier kernel saw false and failed, even though it holds for every real input. The export then stopped with "device-side assert triggered". A Gaussian policy head, `MultivariateNormal(mean, scale_tril=...)` with 14 actions, hits this. A new pass in the CUDA backend now removes these checks before compiling, so the CUDA delegate no longer runs them. The rest of ExecuTorch already removes the same checks before a program is saved. A failed device assert could not come back as an error anyway: it breaks the CUDA context for the rest of the process. Checks on sizes (`_assert_scalar`) stay, since they run on the host. ## float64 weights load in the CUDA backend A program with a float64 weight lowered and saved, then failed to load with `Malformed CUDA FQN weight metadata`. This happened with a float64 buffer, and with float64 math that constant propagation folds into one constant, such as a sinusoidal time embedding. The export wrote the weight correctly, but the runtime had no float64 dtype: the weight cache refused it, and the runtime's tensor type and dtype shims had none either. - The runtime's tensor type now has `Double`, with its 8 byte size. The same type wraps a method's inputs and outputs, so a float64 input or output works too. - The shims export `aoti_torch_dtype_float64`, which compiled code calls to name the dtype of a float64 constant. - The weight cache accepts `Double`, and uses 8 bytes when it checks that a weight fits its storage. - The shims that still refused float64 accept it: a float64 `.item()` (`aoti_torch_item_float64`, which did not exist, so the first execute ended the process), the single-element read those entry points share, `index_put` with accumulate, and the sort shim. A float16 `.item()` had the same problem with `aoti_torch_item_float16` and works now too. - A float read with `.item()` reached Triton kernels truncated to an integer, in float32 too, so a scale of 1.5 computed as 1 with no error. AOTInductor declares every symbolic argument of a kernel's launcher `int64_t`. A compile-time patch in the CUDA backend now declares it with the float type the kernel takes, the change that the open pytorch/pytorch#201882 makes in PyTorch. ## `aten._reshape_copy` lowers to the CUDA backend Torch-TensorRT turns `aten.view` into `aten._reshape_copy`. A view that TensorRT leaves to PyTorch reaches the CUDA backend in that form, and the export failed with `missing fallback kernels: aten::_reshape_copy`, since Inductor has no lowering for it and torch registers no decomposition. The CUDA backend now decomposes it the way ATen computes it: a contiguous copy, viewed to the new shape. ## `varlen_attn` lowers to the CUDA backend `torch.nn.attention.varlen.varlen_attn` runs attention over several sequences packed into one tensor. Lowering it failed with `missing fallback kernels: torch_attn::_varlen_attn`, since it calls Flash or cuDNN attention, which this backend does not ship. The CUDA backend now decomposes it into ops Inductor lowers: attention over all the packed tokens, with a mask that keeps each query in its own sequence. The output is contiguous, like the op's output for a query in the usual (tokens, heads, dim) order, so reshaping it before the output projection works. The decomposition matches the op for causal masks, sliding windows, grouped-query attention, and query and key sequences of different lengths. For a query sequence longer than `max_k + 1`, Flash attention 2 (PyTorch's default) and Flash attention 3 give different results for some windows, because version 2 changes the window based on `max_k`. So a window other than `(left, 0)` lowers only when `max_q <= max_k + 1` is known at export time, as in self-attention, and raises a clear error otherwise. `seqused_k` and `block_table` (paged KV) raise a clear error too. It computes the weights with a softmax, which subtracts each row's largest score first, so the weights still add up to one when scores are large. Tokens before the first sequence or past the last one, such as the padding of a fixed-size buffer, do not change any sequence's output, even when they hold NaN. Unlike in Flash attention, a NaN or inf in one sequence's values does reach the other sequences. It holds a float32 score per head for every pair of packed tokens and runs much slower than Flash attention, so it suits packs of a few thousand tokens, such as a few camera images. ## `scatter_reduce` with amax, amin, prod and mean lowers to the CUDA backend Lowering `x.scatter_reduce(dim, index, src, "amax")` failed with `missing fallback kernels: aoti_torch_cuda_scatter_reduce_two_out`, and so did "amin", "prod" and "mean". Inductor writes a scatter itself only when it can add atomically, and leaves the other reductions to an ATen kernel that this backend did not ship. A sum into int64 or bool, a sum into bfloat16 on GPUs older than sm_90, and deterministic mode use the same kernel. - A new shim next to the index shims implements it the way ATen does: `include_self`, elements the index does not name, the count and integer rounding of the mean, NaN in amax and amin, and bool for sum only. Half and bfloat16 sums add in float, as ATen does in deterministic mode. A product rounds to the tensor's dtype at every step and a mean rounds its sum before dividing, as ATen does, so half overflows and integers wrap where ATen's do. A mean also counts its terms in the tensor's dtype, as ATen does, so 128 int8 ones average to 1. - It sorts the writes by the element they land on and reduces each element's values in index order, so the result is deterministic. It also works inside a CUDA graph. The sort is now one helper that `index_put` uses too. - It is slower than ATen's atomic kernel: 1.5 to 3.2 times on an A100 at 512 output rows, 4096 source rows and 1024 columns, and up to 7 times when many values land on few elements. A decomposition into ops Inductor lowers natively was measured first. It has to compare every output element with every source element, and it was 8 to 89 times slower than ATen's kernel. - `aoti_cuda_shims.lib` and the wheel's shared library check list the new shim. The import library now also lists the shims the DLL exported but it left out, such as `aoti_torch_dtype_float16` and the `aoti_torch_item_` getters, so a float16 model lowered for Windows links. ## CUDA programs load without torch on machines with no libgomp On Linux, Inductor adds OpenMP (`-fopenmp`, `-lgomp`) to every link. Unless the compiler passes `--as-needed` to the linker, as Ubuntu's GCC does by default, the compiled library lists `libgomp.so.1` as a dependency even when it calls no OpenMP function. PyTorch ships a libgomp, so this went unnoticed while torch was imported. A process without torch failed to load the program with `libgomp.so.1: cannot open shared object file` on any machine with no system libgomp, such as a minimal container. The CUDA backend now links with `-Wl,--as-needed`, so the linker keeps only the libraries the program uses. A program that does call OpenMP keeps libgomp. PyTorch already does the same for MPS builds on macOS, and pytorch/pytorch#202902 adds it to Inductor for Linux. The torch-free tests in `test_pybindings_cuda_dlpack.py` now run without a system libgomp, and the ones that check which libraries got loaded also check for libgomp. On an A100, a small MLP and CNN exported with this change load without torch and match eager where the loader cannot find libgomp. This patch, and the one that lets Triton templates run on GPUs with few SMs, now stay in place while any thread lowers. Before, when two threads lowered at once, the first one to finish removed the patch from the other, and the rest of that compile ran without it. ## A failed CUDA load no longer leaves its library in the temporary directory To load a method, the CUDA backend writes the program's compiled library to a temporary file, such as `/tmp/executorch_cuda_<pid>_0.so`. Only `destroy()` removed that file, and `destroy()` never runs when the load fails, so every load that failed after writing the file left it behind. A scope guard now removes the file unless the load succeeds. `test_missing_weights_blob.py` now checks that a load that fails for a missing weights blob leaves nothing in the temporary directory. Without the fix, that check fails. The XNNPACK high-rank permute fix moved to #24380. ## The Triton SDPA leaves keys past the end out of the softmax With `triton_kernel_mode` on (the default), the CUDA backend runs `scaled_dot_product_attention` as its Triton kernel, `triton::sdpa`. For power-of-two head sizes, its main kernel reads keys in tiles of 32, 64 or 128. When the number of keys is not a multiple of the tile, the last tile reaches past the last key. The kernel loaded those positions as zeros but never removed them from the softmax, so without an attention mask or a causal mask each one got a score of 0 and took part of the attention weight. Calls that pass `kv_len` with neither a mask nor `is_causal`, such as the DFlash draft model's global attention layers, had the same problem. On an A100, with the query at 0 and the values at 1, where the right output is 1, one key gave 0.031, 33 keys gave 0.516 and 80 keys gave 0.625. With random inputs and one key, the output was off by 3.4 when the largest value was 3.67. All of these now match eager PyTorch. The kernel now gives every position past the last key, or past `kv_len`, a score of minus infinity, as the other attention kernels in `sdpa.py` already do. The first version of the kernel did this, and a later rewrite dropped it. Checking every tile cost up to 8% where no tile reaches past the end, so only calls that need the check pay for it. A mask or a causal bound already drops those keys, so calls with one skip it. The launcher also passes the largest tile size that divides the number of keys, and tiles that divide it skip the check at compile time. So calls with a mask or a causal bound, and tiles that divide a key count known at export time, compile to the same code as before. Without a mask or a causal bound, a `kv_len` or a symbolic key count checks every tile. `test_triton_sdpa.py` drops its note that called this a kernel limitation, runs its basic test with a single key, and runs every tile without a mask over 80 and 96 keys and over a `kv_len` of 80. These fail without the fix. ## `scatter` with a tensor source lowers to the CUDA backend Lowering `x.scatter(dim, index, src)` failed with a `missing fallback kernels` error that names `aoti_torch_cuda_scatter_src_out` when `src` is bfloat16, int64 or bool. Inductor writes a scatter from a tensor itself only in a dtype it can add atomically, so it leaves int64, bool and, on GPUs older than sm_90, bfloat16 to an ATen kernel, and every dtype in deterministic mode. A mixture-of-experts layer that writes its expert outputs back with `torch.zeros_like(h).scatter(1, index, src)` in bfloat16 hit this on an A100. A new shim next to the `scatter_reduce` one implements it with the same checks, the same sorting of the writes by the element they land on, and the same kernel. The kernel gets a new mode in which only the last write to each element stores its value. So when values land on the same element, the last one in index order wins, as ATen does in deterministic mode; ATen's default CUDA kernel keeps any one of them. As in that default kernel, an index out of range, negative ones included, triggers a device-side assert. In deterministic mode, ATen's scatter counts a negative index from the end instead. The shim takes float, double, half, bfloat16, the integer dtypes and bool. Like ATen, both scatter shims now take an int32 index as well as an int64 one. Mostly because of the sort, the shim is about 10 to 15 times slower than ATen's default kernel for a mixture-of-experts style bfloat16 scatter of 4 to 34 million values on an A100. The Windows import library `aoti_cuda_shims.lib` and the wheel's shared library check now list the new shim. On an A100 with PyTorch 2.14 and 2.15, scatters in float32, float16, bfloat16, int64 and bool lower, run and match eager PyTorch exactly. So does a small mixture-of-experts layer built the same way. `test_scatter_with_tensor_source` fails without the fix, and its int32 index case fails when the shims take only int64. ## Review order Start with the CUDA graph changes: `cuda_backend.cpp` and `cuda_delegate_handle.h` in `backends/cuda/runtime`, the mask and `.item()` shims in `runtime/shims/index.cu` and `memory.cpp`, and their tests in `test_output_stride_rearrange.py`. Then the lowering side: the passes in `backends/cuda/passes`, the decompositions and compile-time patches in `backends/cuda/cuda_backend.py`, and the folding switch in `backends/aoti/aoti_backend.py`. Then the scatter and scatter_reduce shims in `index.cu` with `test_index_shims.py`. The remaining commits are small follow-ups that stand on their own. ## Testing On an A100, with the CUDA backend built from this branch: - `test_index_shims.py`, `test_output_stride_rearrange.py`, `test_cuda_export.py`, `test_antialiased_resize.py`, `test_merge_ptes.py` and `test_cuda_weight_metadata.py` pass. - With the backend built from main, the new CUDA graph tests fail. - The replay error names the dim, for example "input 0 has size 4 in dim 1, captured with 3". - `test_cuda_kv_cache` passes, and fails when `release()` keeps the captured input shapes. - With the min and max replaced by a plain `if`, `test_bicubic_antialias_resize` fails. - With the new data name refusal removed, or moved back behind the compile, the data name test fails. - A `torch.cond` on a host-memory flag matches eager on every call; with a graph it returned the captured branch. - A random program run from two processes started together drew the same numbers in 2 of 3 tries before the seed change, and different ones in 3 of 3 after. - Mask index programs exported from this branch match eager PyTorch on every call, both on a runtime built from main with CUDA graphs off and on this branch's runtime with CUDA graphs on. - The shared shim unit tests pass in a standalone googletest build. - `test_replace_masked_copy_with_where.py` passes, and fails when the pass does nothing, is not registered, or skips its accumulate, same-mask, boolean mask, shared picks or count check. It writes into a destination filled with 7, so it fails if `where` loses the values the mask does not pick. A contiguous destination with a channels_last source, or with a channels_last mask, keeps its `index_put`; before, the rewrite failed to lower when a later `view(-1)` needed the destination's layout. It also fails when the pass walks only the top-level graph, since a masked copy in a `torch.cond` branch then keeps its index, and when it takes two draws of a random mask for the same mask. Two masks computed with over a thousand ops each raised RecursionError before. A check that reads the index's count but tests another size keeps the index; before, the pass removed it. - A masked copy lowered through the CUDA partitioner matches eager for masks that pick some, none and all positions. The test records the graph the pass returns, and fails when the pass is not registered or leaves the index in place. - A module that adds one to a device buffer on every call and reads a mask's count returns the call number on every call, so the capture call that runs again without a graph updates the buffer once. With the retry made to run the call twice, the capture call returns one too many. With PyTorch 2.14 it returned one on every call even without a graph, because Inductor folded the written buffer into a fill; it passes now with 2.14 and with 2.15 and 2.16 nightlies. - On the tiny Gemma 4 decoder from the Gemma 4 31B example's tests, which writes a KV cache, with PyTorch 2.14 on an A100: with the folding off, export took 193 s against 190.5 s with it on, a decode call took 945 µs against 1019 µs (medians of three exports and three runs on a busy machine), and the 41 generated tokens were the same. - The mask tests check that the runtime log reports one failed capture, so they fail when CUDA graphs are off or a failed capture is tried again. `test_host_read_runs_without_a_graph` also runs a module that scales by a `.item()` read; without the capture check in the item shims, the process aborts. - With the debug check on the phase before the retry turned into a release check, the mask tests pass. - `test_distribution_argument_check` passes with PyTorch 2.14 and fails with "device-side assert triggered" without the new pass. It also checks that the backend lists the pass, since without the pass the export fails only when autotuning happens to trip the assert. It writes the `MultivariateNormal` check as `torch._assert_async`, since PyTorch 2.14 cannot export the distribution itself. It also checks that the pass drops the check and the math that only fed it. A Gaussian policy head with hidden sizes 16, 64 and 256 exports, runs, and matches eager PyTorch. - A float64 buffer and a folded float64 constant load and match eager PyTorch. Before, both failed to load. `test_float64_weight_loads` fails without the fix. - `test_reshape_copy` in `test_cuda_export.py` fails without the fix. The one in `test_index_shims.py` runs `_reshape_copy` on a transposed, a sliced and a contiguous input and matches eager PyTorch. - `test_varlen_attn` in `test_cuda_export.py` fails without the fix. It lowers a program that flattens the output, compares the decomposition with the op on CUDA only (on ROCm, varlen_attn runs AOTriton, which lines an explicit window up with the start of each sequence, while the decomposition follows Flash attention on CUDA, which lines it up with the end), including a window in self-attention and NaN values outside the key sequences, and checks that a window the Flash attention versions disagree on raises. It also fails when the output is left non-contiguous, when the window check is removed or refuses every such window, and when the values outside the key sequences are not zeroed. Both `test_varlen_attn` tests include queries and keys filled with 1024, where every key gets the same weight; computing the weights as `exp(scores - logsumexp)` fails them, since at that size float32 rounds the log-sum-exp to a whole number and the weights no longer add up to one. The `test_varlen_attn` in `test_index_shims.py` lowers and runs programs that call `varlen_attn` with no window, a causal mask and the windows (8, 0) and (4, 2), flatten each output, and take `cu_seq_q` and `cu_seq_k` only at run time. They match eager PyTorch. - `test_missing_weights_blob.py`, `test_cuda_weight_reuse.py` and `test_device_placement.py` pass. - `test_scatter_reduce` now runs the sum too, which for int64 goes to the new shim. The test that a missing scatter kernel fails at export now tries to lower a scatter of a scalar in deterministic mode. The backend still lacks that kernel, so the test still checks the scatter fallback now that the scatter_reduce and scatter shims exist. It fails when the check stops watching that fallback. `test_scatter_mean_counts_in_its_dtype` fails with the old count and when a count that wraps to zero is not turned into one. `test_float64_and_float16_values` (a sort and `.item()` in each dtype) and the float64 masked accumulate fail with the old shims, and all pass now. The `.item()` scale is 1.5, so that test also fails without the patch that passes a float to Triton kernels as a float. - `test_scatter_reduce` and `test_sort_shim.py` pass. `test_scatter_reduce` fails without the fix, and also when the shim ignores `include_self`, rounds the integer mean toward zero, or lets a NaN lose in amax and amin. `test_scatter_strides` scatters along the middle dim of a 3-D tensor, given as a negative dim, from a strided source and from transposed operands, and fails when the shims walk only the last two dims. `test_scatter_reduce_edge_cases` covers a bool sum, int32 and half products that wrap and overflow, and an empty index. `test_scatter_reduce_sum_adds_in_float` sums 1000 half and bfloat16 values in deterministic mode, and fails when the sum rounds to the dtype at every step. `test_scatter_reduce_replays_in_a_graph` captures a `scatter_reduce` program in a CUDA graph and matches eager on every replay. - The weight cache, runtime tensor and shim unit tests pass in a standalone googletest build. Without the fix, the weight cache test of a float64 entry fails. - `test_small_gpu_override_survives_a_thread_that_leaves_first` overlaps two threads inside the small-GPU patch, and fails with the old save and restore. - With the seed copy made to fail once, the random shim logs it, fails that call, and seeds on the next one; before, it went on with an unseeded state. Not tested: a failed first graph launch (it cannot be triggered from a test), a `std::random_device` that throws, Windows, macOS case-insensitive file systems, the buffer folding change on Metal, and the ROCm build. No test covers the seed change, since a test cannot make two processes start in the same second. ### Checked on real robot policies On the release/1.6 stack with this change, on A100 and H100: - SmolVLA on the CUDA backend: one device-to-host copy per action chunk (only the action download), matching eager within the usual tolerance. - A Gaussian policy whose `torch.distributions` check used to stop the CUDA export: exports and runs, 0.76 ms per chunk. - A float64 weight that used to fail at load with "Malformed CUDA FQN weight metadata": loads and runs, max difference 0. - Four processes started together: four different random draws, against two out of four before.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
…ckend An operator Torch-TensorRT cannot convert used to fall back to a portable CPU kernel, with copies to and from the GPU around it. The executorch_tensorrt backend now hands those operators to ExecuTorch's CUDA backend, so a program runs entirely on the GPU. The CUDA backend writes its weights into a data file beside the program. Each part names that file after itself, through the backend's data-name compile spec, so parts saved into one bundle folder keep their own weights, and the file is recorded as the program's data file. The data-name compile spec comes from pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Tested on an NVIDIA GPU: a model with an operator TensorRT rejects compiles to TensorRT and CUDA delegates with no operator left on the CPU and its input and output in GPU memory, and matches PyTorch when run without PyTorch.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
executorch_cuda lowers each part with ExecuTorch's own CUDA backend, so a bundle runs on an NVIDIA GPU without TensorRT. ExecuTorch compiles the kernels ahead of time for the GPU that exports. It is one config class on ExecuTorchConfig, and its registration. ExecuTorch loads its programs, so the runtime, the engine and lerobot-rollout are unchanged. The CUDA backend keeps a program's kernels and weights in a separate data file, named after the part, as with executorch_tensorrt, so two parts never overwrite each other's weights and the bundle names one data file per program. Naming that file needs pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Inputs and outputs stay in GPU memory too, so a program reads an earlier one's outputs in place. A CUDA bundle records the GPU that built it, and is refused on another. The kernels are compiled in full float32, as lerobot-export compiles for every backend. PyTorch lets convolutions use TF32 by default, and the convolution kernels the CUDA backend generates follow that flag. On ACT on an H100 they then drifted about seven times further from PyTorch on the CPU than PyTorch's own GPU kernels do. In full float32 the actions match to about 1e-6, and a chunk still runs twice as fast as eager PyTorch on the GPU. It needs the CUDA build of ExecuTorch, and the CUDA compiler (nvcc) at export. Like the other ExecuTorch backends, it reads programs traced by PyTorch 2.15 or newer. Like the backends that take only part of a program, it folds the values a program computes from constants alone before lowering. SmolVLA's vision encoder picks its position ids with a boolean mask made from constants. Left in the program, that indexing makes the GPU read a size back to the host on every run; folded, it is gone. Folding has one cost today: SmolVLA's time embedding, computed in float64 from constants, becomes a float64 weight, and ExecuTorch's CUDA backend cannot load float64 weights yet, so SmolVLA waits for that fix on this backend. Tested on an H100 and on a Jetson AGX Thor: exported ACT, checked the bundle, and ran it with lerobot-rollout on a fake robot. On the H100, pi0 exported too, and its actions matched eager PyTorch to about 1e-5 on ten frames; in float32 it runs a little slower than eager. The check that a bundle runs without PyTorch now also runs a CUDA bundle.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
… part's weights file The executorch_cuda and executorch_tensorrt backends name the weights file of each program part after the part, so the parts of a bundle do not share one file name. That takes the CUDA backend's `generate_data_name_compile_spec` (pytorch/executorch#24314). The ExecuTorch 1.6.0 release candidate does not have it yet, and an export there stopped deep in the compile step with `AttributeError: type object 'CudaBackend' has no attribute 'generate_data_name_compile_spec'`. The export now refuses such an ExecuTorch with an error that names the installed version and the missing change. Like the refusal of a Torch-TensorRT without `fallback_pytorch_plugin_ops`, it is a ValueError, so the export does not add a "pip install" hint for a package that is already installed. The TensorRT backend asks for the CUDA partitioner before TensorRT builds the part's engines, so it refuses at once instead of after the build. `cuda_backend_names_its_weights()` answers the same question for tests that need the CUDA backend. Tested: a unit test refuses the CUDA partitioner of such an ExecuTorch with ExecuTorch 1.6.0 (release candidate) installed, and another shows the TensorRT backend refuses before calling Torch-TensorRT's compile; without the change it reaches the compile.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
…d the installed versions Three test files each built their own skip marks for the ExecuTorch backends, and they checked different things. None of them checked that the CUDA backend can name each part's weights file, so with the ExecuTorch 1.6.0 release candidate the `executorch_cuda` tests failed instead of skipping. A skipped test also said only "needs ExecuTorch and a CUDA GPU", which does not tell a reader which package is too old. `tests/export/backends.py` now holds one mark per backend: - `needs_executorch`: ExecuTorch 1.6 and PyTorch 2.15. - `needs_executorch_cuda`: also a CUDA GPU and an ExecuTorch whose CUDA backend names each part's weights file (pytorch/executorch#24314). - `needs_executorch_tensorrt`: also a Torch-TensorRT with `fallback_pytorch_plugin_ops` (pytorch/TensorRT#4953) and torch-tensorrt-executorch-runtime. Each reason names these fixes and the installed versions of PyTorch, ExecuTorch, Torch-TensorRT and its runtime. The reason always names ExecuTorch, and a CUDA GPU for a GPU backend, so `tests/conftest.py` marks the test `executorch` and `gpu` whether it skips or not. `test_export.py`, `test_export_history.py` and `test_executorch_tensorrt.py` use these marks. Tested: with the ExecuTorch 1.6.0 and Torch-TensorRT 2.15.0 release candidates on a GPU, every GPU export test skips and names pytorch/executorch#24314 or pytorch/TensorRT#4953; with ExecuTorch and Torch-TensorRT builds that have both fixes, the same tests run and pass.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
…he GPU backends after approval No CI job installed ExecuTorch, so every test that exports with a real backend skipped, and nothing checked that `lerobot-export` and `lerobot-rollout --bundle.path` still work with XNNPACK, ExecuTorch's CUDA backend or TensorRT. Two jobs now install the compilers over the locked environment and run the tests marked `executorch` (`tests/conftest.py` adds the mark to every test that skips without ExecuTorch): - Fast Export Tests, in Fast Tests, on every pull request, forks included: `ubuntu-latest`, the extras of the exported policies (`pi`, `smolvla`, `groot`, `diffusion`), then ExecuTorch's CPU build, then `pytest tests -m executorch`. The GPU tests among them skip there. - GPU Export Tests, in Full Tests, next to GPU Tests and on the same image and runner: ExecuTorch's and Torch-TensorRT's CUDA 13.2 builds, which match PyPI's torch 2.15.0, and `nvidia-cuda-nvcc` 13.2, the CUDA compiler the CUDA backend runs at export (`CUDA_HOME` points at it). Then `pytest tests -m "executorch and gpu"`. It runs when GPU Tests do: after approval, on main and by hand, never for a fork. ExecuTorch and Torch-TensorRT stay out of `uv.lock`, for the reasons the `export` note in `pyproject.toml` gives, so each job pins the exact wheels in one environment variable. The ExecuTorch 1.6 and Torch-TensorRT 2.15 releases are not on PyPI yet, so the variables point at their release candidates on PyTorch's test index; they become `executorch==1.6.0` and `torch-tensorrt==2.15.0` once those are published. The candidates lack two changes the GPU backends need (pytorch/executorch#24314 and pytorch/TensorRT#4953), and no torch-tensorrt-executorch-runtime takes ExecuTorch 1.6.0 yet, so today the GPU tests skip and say which change each one needs. They run as soon as released wheels have them. CONTRIBUTING.md and AGENTS.md say where these tests run and which mark a new export test takes. Tested: ran each job's commands in fresh environments installed only from the package indexes, with the locked export extras in place of the CI image. - Fast Export Tests (4 CPUs, a busy machine): 20 passed in 46 minutes. GR00T with XNNPACK failed as the previous commit describes; it skips now. - GPU Export Tests on an A100: all 34 GPU tests skip with their reasons, 57 seconds. - Not from the indexes: the same GPU tests with ExecuTorch and Torch-TensorRT built from source with the two changes, on an H100 with cold caches: 29 passed, 5 skipped (tiny pi0 and pi0.5 with TensorRT, and one test that downloads CLIP), 42 minutes. Not tested: the L4 runner and the CI image themselves.
JacobSzwejbka
added a commit
that referenced
this pull request
Nov 3, 2026
Supersedes #24324 with a clean release pick from #24314. Adds validated per-program CUDA weights filenames and lowerings for antialiased bilinear/bicubic resize. Retains the data-name validation, resize-test removal and documentation follow-ups already carried by #24324. Preserves the release behavior that mismatched CUDA weights fail to load. The resize rewrite follows the source PR tolerances rather than promising bit-identical floating-point results. ## Scope and landing order Land after #24488 in the verified landing order. This branch starts directly at release/1.6 commit `dfb78b9ecd933a5fee33ebc94f69ed4b71cbb1e7`. Its GitHub diff contains 10 files for this pick; it does not carry duplicate #24298 history or the excluded unmerged QNN work. Dependencies are documented here rather than embedded as other cherry-pick commits. ## Verification All changed Python files compile and pass Ruff F821 checks. All seven replacements merge cleanly in the documented sequence, including simulated squash merges. No QNN files are changed. Native CUDA, OpenVINO and portable C++ suites were not run locally on this macOS host; PR CI remains required before landing. Prepared with OpenAI Codex.
JacobSzwejbka
pushed a commit
that referenced
this pull request
Nov 3, 2026
… input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups (#24512) ## What is wrong today With `enable_cuda_graph_for_method` set, the CUDA backend runs a method a few times normally, captures a CUDA graph on the fourth call, and replays that graph after it. Some calls go wrong: - **A call that cannot be captured fails.** A mask index, such as `x[mask]` or `x[mask] = v`, reads a count back to the host, and a capture cannot record that. The capture call returned an error. An output that is first strided at the capture call failed the same way. A `.item()` on a device tensor, such as a `torch.cond` predicate computed on the GPU, ended the process. - **A replay can return a wrong result with no error.** Replay checked only each input's byte count. An input that changes from `(2, 3, 4, 4)` to `(2, 4, 3, 4)` keeps its byte count, so it replayed the captured shape and returned the wrong output. ## What this change does - Replay compares each input's shape with the captured one and refuses a different shape with `InvalidArgument`. The error names the dim and both sizes. A captured graph only works for the input shapes of the capture call, and the code comment on the graph phases now says so. - When a capture fails, the backend runs the same call again without a graph and returns its result. Graphs stay off for that method, as before. The retry is safe because a capture records kernels without running them, and it happens only if the new graph never started running. The `.item()` shims now refuse to read inside a capture, as the mask index does, so a capture that reads a device value is retried too instead of ending the process. - A method whose program keeps a constant in host memory, such as a `torch.cond` flag updated on every call, now runs without a graph. Host code runs only during the capture, so a graph would keep the branch of the capture call, and a retry after a failed capture would take the other one. The log names the constant found in host memory, or says the constants could not be listed. - An expected capture failure, such as the mask index refusing to run inside a capture, is logged at Info instead of Error, since the call then succeeds. A real failure of the call without a graph is still logged at Error. AOTInductor's generated code still prints its own line for the refused call. - The allocations a captured graph makes are recorded only after its first launch succeeds. Before, a failed first launch left addresses that were never allocated, and a later cleanup passed them to `cudaFree`. A failed capture now cleans up with `CudaGraphState::release()` instead of a copy of it. - These CUDA graph changes do not change the exported program, so older programs get them too. A program that uses the new shims described below (scatter, scatter_reduce, float64 or float16) needs a runtime that has them. Smaller commits: - The CUDA random shims seed from `std::random_device` instead of the clock, so two processes started in the same second no longer draw the same numbers. If `std::random_device` throws, or the copy of the seed to the GPU fails, the shim logs an error and frees what it allocated, and the next call tries again. - The wheel's shared library check now lists the mask index and `index_put` shims. - With PyTorch 2.14, a program that writes a one-element buffer, such as a counter, kept returning its first update, because Inductor folded the update into a constant. The AOTI backends (CUDA and Metal) now turn off Inductor's joint-graph constant folding, for the whole program, when a program writes a buffer. On a small Gemma 4 decoder that writes its KV cache, folding changes two nodes per graph; with it off, export and decode times stayed within run-to-run noise and the generated tokens were identical. - Unit tests for `aoti_torch_get_size`. - Comments on the index shims and on `set_output_data_ptr` that described accumulation, the sort count limit and the output index wrongly are corrected. - The mask index tests move to a new `test_index_shims.py`, since they need the Python runtime. They add a 2-D mask, and `index_put` with repeated indices in deterministic mode (half and bfloat16 accumulation, last write wins, an int32 index). - `test_output_shape_changes_at_capture` ends with a call that only a live graph rejects, so it fails if graphs are off. The KV cache recapture test checks that dropping the graph also clears the captured input shapes. ## Follow-ups to #24314 - Delete `test_align_corners_compiles_with_dynamic_sizes`. It could not fail: it passed with no resize rewrite at all and with the min and max in the align_corners ratio replaced by a plain `if`, because `torch.compile` compiles again for an output size of 1. With torch 2.14 and `CI` set it also failed for an unrelated reason, and it is the one failing test of the CUDA unit test job on main. `test_bicubic_antialias_resize` already runs an exported program at an output height of 1, where a plain `if` breaks. - The comment on that min and max now gives the real reason: an exported program keeps the branch it was traced with, so a plain `if` divides by zero at an output size of 1. - A backend that keeps its data in the `.pte`, such as Metal, now refuses a `data_name` compile spec built by hand too, before compiling anything. Before, only `generate_data_name_compile_spec` refused it, and a hand-built spec was ignored with no name check. - The CUDA docs warn that Windows and macOS usually treat data names that differ only in case as the same file. ## Copies through a boolean mask no longer read the mask count back to the host `out[mask] = src[mask]` ran as a mask index and a mask `index_put`. Each of them read the mask's count back to the host and waited for the GPU, on every call. Transformers' SmolVLM vision embeddings do this for their position ids. When the values are the same mask's picks from a tensor shaped like the destination, every element goes back to where it came from. So a new pass in the CUDA backend writes the copy as `torch.where(mask, src, out)`, which gives the same result, bit for bit. It does so inside `torch.cond` branches and other subgraphs too. It rewrites only when `where` would lay out its result like the destination, since `where` takes its layout from its inputs, the mask first, while `index_put` keeps the destination's layout. It also drops the index and the checks on its size when nothing else reads them and each check tests only the index's count and holds for any mask. A user check on the count, such as `torch._check(x[mask].shape[0] >= 1)`, keeps the index. The pass compares the two masks with a work list, so masks with long histories do not hit Python's recursion limit. On an A100, five calls of the SmolVLM position id code went from 10 count reads and 10 stream waits to none. ## Tensor checks no longer stop the CUDA export The argument checks in `torch.distributions` export as `aten._assert_async.msg`. The CUDA backend compiled each one into a device assert inside a Triton kernel. While compiling, AOTInductor runs every kernel once to pick its launch settings, on made-up buffers: random floats, and zeros for integers and bools. So a check that reads a bool from an earlier kernel saw false and failed, even though it holds for every real input. The export then stopped with "device-side assert triggered". A Gaussian policy head, `MultivariateNormal(mean, scale_tril=...)` with 14 actions, hits this. A new pass in the CUDA backend now removes these checks before compiling, so the CUDA delegate no longer runs them. The rest of ExecuTorch already removes the same checks before a program is saved. A failed device assert could not come back as an error anyway: it breaks the CUDA context for the rest of the process. Checks on sizes (`_assert_scalar`) stay, since they run on the host. ## float64 weights load in the CUDA backend A program with a float64 weight lowered and saved, then failed to load with `Malformed CUDA FQN weight metadata`. This happened with a float64 buffer, and with float64 math that constant propagation folds into one constant, such as a sinusoidal time embedding. The export wrote the weight correctly, but the runtime had no float64 dtype: the weight cache refused it, and the runtime's tensor type and dtype shims had none either. - The runtime's tensor type now has `Double`, with its 8 byte size. The same type wraps a method's inputs and outputs, so a float64 input or output works too. - The shims export `aoti_torch_dtype_float64`, which compiled code calls to name the dtype of a float64 constant. - The weight cache accepts `Double`, and uses 8 bytes when it checks that a weight fits its storage. - The shims that still refused float64 accept it: a float64 `.item()` (`aoti_torch_item_float64`, which did not exist, so the first execute ended the process), the single-element read those entry points share, `index_put` with accumulate, and the sort shim. A float16 `.item()` had the same problem with `aoti_torch_item_float16` and works now too. - A float read with `.item()` reached Triton kernels truncated to an integer, in float32 too, so a scale of 1.5 computed as 1 with no error. AOTInductor declares every symbolic argument of a kernel's launcher `int64_t`. A compile-time patch in the CUDA backend now declares it with the float type the kernel takes, the change that the open pytorch/pytorch#201882 makes in PyTorch. ## `aten._reshape_copy` lowers to the CUDA backend Torch-TensorRT turns `aten.view` into `aten._reshape_copy`. A view that TensorRT leaves to PyTorch reaches the CUDA backend in that form, and the export failed with `missing fallback kernels: aten::_reshape_copy`, since Inductor has no lowering for it and torch registers no decomposition. The CUDA backend now decomposes it the way ATen computes it: a contiguous copy, viewed to the new shape. ## `varlen_attn` lowers to the CUDA backend `torch.nn.attention.varlen.varlen_attn` runs attention over several sequences packed into one tensor. Lowering it failed with `missing fallback kernels: torch_attn::_varlen_attn`, since it calls Flash or cuDNN attention, which this backend does not ship. The CUDA backend now decomposes it into ops Inductor lowers: attention over all the packed tokens, with a mask that keeps each query in its own sequence. The output is contiguous, like the op's output for a query in the usual (tokens, heads, dim) order, so reshaping it before the output projection works. The decomposition matches the op for causal masks, sliding windows, grouped-query attention, and query and key sequences of different lengths. For a query sequence longer than `max_k + 1`, Flash attention 2 (PyTorch's default) and Flash attention 3 give different results for some windows, because version 2 changes the window based on `max_k`. So a window other than `(left, 0)` lowers only when `max_q <= max_k + 1` is known at export time, as in self-attention, and raises a clear error otherwise. `seqused_k` and `block_table` (paged KV) raise a clear error too. It computes the weights with a softmax, which subtracts each row's largest score first, so the weights still add up to one when scores are large. Tokens before the first sequence or past the last one, such as the padding of a fixed-size buffer, do not change any sequence's output, even when they hold NaN. Unlike in Flash attention, a NaN or inf in one sequence's values does reach the other sequences. It holds a float32 score per head for every pair of packed tokens and runs much slower than Flash attention, so it suits packs of a few thousand tokens, such as a few camera images. ## `scatter_reduce` with amax, amin, prod and mean lowers to the CUDA backend Lowering `x.scatter_reduce(dim, index, src, "amax")` failed with `missing fallback kernels: aoti_torch_cuda_scatter_reduce_two_out`, and so did "amin", "prod" and "mean". Inductor writes a scatter itself only when it can add atomically, and leaves the other reductions to an ATen kernel that this backend did not ship. A sum into int64 or bool, a sum into bfloat16 on GPUs older than sm_90, and deterministic mode use the same kernel. - A new shim next to the index shims implements it the way ATen does: `include_self`, elements the index does not name, the count and integer rounding of the mean, NaN in amax and amin, and bool for sum only. Half and bfloat16 sums add in float, as ATen does in deterministic mode. A product rounds to the tensor's dtype at every step and a mean rounds its sum before dividing, as ATen does, so half overflows and integers wrap where ATen's do. A mean also counts its terms in the tensor's dtype, as ATen does, so 128 int8 ones average to 1. - It sorts the writes by the element they land on and reduces each element's values in index order, so the result is deterministic. It also works inside a CUDA graph. The sort is now one helper that `index_put` uses too. - It is slower than ATen's atomic kernel: 1.5 to 3.2 times on an A100 at 512 output rows, 4096 source rows and 1024 columns, and up to 7 times when many values land on few elements. A decomposition into ops Inductor lowers natively was measured first. It has to compare every output element with every source element, and it was 8 to 89 times slower than ATen's kernel. - `aoti_cuda_shims.lib` and the wheel's shared library check list the new shim. The import library now also lists the shims the DLL exported but it left out, such as `aoti_torch_dtype_float16` and the `aoti_torch_item_` getters, so a float16 model lowered for Windows links. ## CUDA programs load without torch on machines with no libgomp On Linux, Inductor adds OpenMP (`-fopenmp`, `-lgomp`) to every link. Unless the compiler passes `--as-needed` to the linker, as Ubuntu's GCC does by default, the compiled library lists `libgomp.so.1` as a dependency even when it calls no OpenMP function. PyTorch ships a libgomp, so this went unnoticed while torch was imported. A process without torch failed to load the program with `libgomp.so.1: cannot open shared object file` on any machine with no system libgomp, such as a minimal container. The CUDA backend now links with `-Wl,--as-needed`, so the linker keeps only the libraries the program uses. A program that does call OpenMP keeps libgomp. PyTorch already does the same for MPS builds on macOS, and pytorch/pytorch#202902 adds it to Inductor for Linux. The torch-free tests in `test_pybindings_cuda_dlpack.py` now run without a system libgomp, and the ones that check which libraries got loaded also check for libgomp. On an A100, a small MLP and CNN exported with this change load without torch and match eager where the loader cannot find libgomp. This patch, and the one that lets Triton templates run on GPUs with few SMs, now stay in place while any thread lowers. Before, when two threads lowered at once, the first one to finish removed the patch from the other, and the rest of that compile ran without it. ## A failed CUDA load no longer leaves its library in the temporary directory To load a method, the CUDA backend writes the program's compiled library to a temporary file, such as `/tmp/executorch_cuda_<pid>_0.so`. Only `destroy()` removed that file, and `destroy()` never runs when the load fails, so every load that failed after writing the file left it behind. A scope guard now removes the file unless the load succeeds. `test_missing_weights_blob.py` now checks that a load that fails for a missing weights blob leaves nothing in the temporary directory. Without the fix, that check fails. The XNNPACK high-rank permute fix moved to #24380. ## The Triton SDPA leaves keys past the end out of the softmax With `triton_kernel_mode` on (the default), the CUDA backend runs `scaled_dot_product_attention` as its Triton kernel, `triton::sdpa`. For power-of-two head sizes, its main kernel reads keys in tiles of 32, 64 or 128. When the number of keys is not a multiple of the tile, the last tile reaches past the last key. The kernel loaded those positions as zeros but never removed them from the softmax, so without an attention mask or a causal mask each one got a score of 0 and took part of the attention weight. Calls that pass `kv_len` with neither a mask nor `is_causal`, such as the DFlash draft model's global attention layers, had the same problem. On an A100, with the query at 0 and the values at 1, where the right output is 1, one key gave 0.031, 33 keys gave 0.516 and 80 keys gave 0.625. With random inputs and one key, the output was off by 3.4 when the largest value was 3.67. All of these now match eager PyTorch. The kernel now gives every position past the last key, or past `kv_len`, a score of minus infinity, as the other attention kernels in `sdpa.py` already do. The first version of the kernel did this, and a later rewrite dropped it. Checking every tile cost up to 8% where no tile reaches past the end, so only calls that need the check pay for it. A mask or a causal bound already drops those keys, so calls with one skip it. The launcher also passes the largest tile size that divides the number of keys, and tiles that divide it skip the check at compile time. So calls with a mask or a causal bound, and tiles that divide a key count known at export time, compile to the same code as before. Without a mask or a causal bound, a `kv_len` or a symbolic key count checks every tile. `test_triton_sdpa.py` drops its note that called this a kernel limitation, runs its basic test with a single key, and runs every tile without a mask over 80 and 96 keys and over a `kv_len` of 80. These fail without the fix. ## `scatter` with a tensor source lowers to the CUDA backend Lowering `x.scatter(dim, index, src)` failed with a `missing fallback kernels` error that names `aoti_torch_cuda_scatter_src_out` when `src` is bfloat16, int64 or bool. Inductor writes a scatter from a tensor itself only in a dtype it can add atomically, so it leaves int64, bool and, on GPUs older than sm_90, bfloat16 to an ATen kernel, and every dtype in deterministic mode. A mixture-of-experts layer that writes its expert outputs back with `torch.zeros_like(h).scatter(1, index, src)` in bfloat16 hit this on an A100. A new shim next to the `scatter_reduce` one implements it with the same checks, the same sorting of the writes by the element they land on, and the same kernel. The kernel gets a new mode in which only the last write to each element stores its value. So when values land on the same element, the last one in index order wins, as ATen does in deterministic mode; ATen's default CUDA kernel keeps any one of them. As in that default kernel, an index out of range, negative ones included, triggers a device-side assert. In deterministic mode, ATen's scatter counts a negative index from the end instead. The shim takes float, double, half, bfloat16, the integer dtypes and bool. Like ATen, both scatter shims now take an int32 index as well as an int64 one. Mostly because of the sort, the shim is about 10 to 15 times slower than ATen's default kernel for a mixture-of-experts style bfloat16 scatter of 4 to 34 million values on an A100. The Windows import library `aoti_cuda_shims.lib` and the wheel's shared library check now list the new shim. On an A100 with PyTorch 2.14 and 2.15, scatters in float32, float16, bfloat16, int64 and bool lower, run and match eager PyTorch exactly. So does a small mixture-of-experts layer built the same way. `test_scatter_with_tensor_source` fails without the fix, and its int32 index case fails when the shims take only int64. ## Review order Start with the CUDA graph changes: `cuda_backend.cpp` and `cuda_delegate_handle.h` in `backends/cuda/runtime`, the mask and `.item()` shims in `runtime/shims/index.cu` and `memory.cpp`, and their tests in `test_output_stride_rearrange.py`. Then the lowering side: the passes in `backends/cuda/passes`, the decompositions and compile-time patches in `backends/cuda/cuda_backend.py`, and the folding switch in `backends/aoti/aoti_backend.py`. Then the scatter and scatter_reduce shims in `index.cu` with `test_index_shims.py`. The remaining commits are small follow-ups that stand on their own. ## Testing On an A100, with the CUDA backend built from this branch: - `test_index_shims.py`, `test_output_stride_rearrange.py`, `test_cuda_export.py`, `test_antialiased_resize.py`, `test_merge_ptes.py` and `test_cuda_weight_metadata.py` pass. - With the backend built from main, the new CUDA graph tests fail. - The replay error names the dim, for example "input 0 has size 4 in dim 1, captured with 3". - `test_cuda_kv_cache` passes, and fails when `release()` keeps the captured input shapes. - With the min and max replaced by a plain `if`, `test_bicubic_antialias_resize` fails. - With the new data name refusal removed, or moved back behind the compile, the data name test fails. - A `torch.cond` on a host-memory flag matches eager on every call; with a graph it returned the captured branch. - A random program run from two processes started together drew the same numbers in 2 of 3 tries before the seed change, and different ones in 3 of 3 after. - Mask index programs exported from this branch match eager PyTorch on every call, both on a runtime built from main with CUDA graphs off and on this branch's runtime with CUDA graphs on. - The shared shim unit tests pass in a standalone googletest build. - `test_replace_masked_copy_with_where.py` passes, and fails when the pass does nothing, is not registered, or skips its accumulate, same-mask, boolean mask, shared picks or count check. It writes into a destination filled with 7, so it fails if `where` loses the values the mask does not pick. A contiguous destination with a channels_last source, or with a channels_last mask, keeps its `index_put`; before, the rewrite failed to lower when a later `view(-1)` needed the destination's layout. It also fails when the pass walks only the top-level graph, since a masked copy in a `torch.cond` branch then keeps its index, and when it takes two draws of a random mask for the same mask. Two masks computed with over a thousand ops each raised RecursionError before. A check that reads the index's count but tests another size keeps the index; before, the pass removed it. - A masked copy lowered through the CUDA partitioner matches eager for masks that pick some, none and all positions. The test records the graph the pass returns, and fails when the pass is not registered or leaves the index in place. - A module that adds one to a device buffer on every call and reads a mask's count returns the call number on every call, so the capture call that runs again without a graph updates the buffer once. With the retry made to run the call twice, the capture call returns one too many. With PyTorch 2.14 it returned one on every call even without a graph, because Inductor folded the written buffer into a fill; it passes now with 2.14 and with 2.15 and 2.16 nightlies. - On the tiny Gemma 4 decoder from the Gemma 4 31B example's tests, which writes a KV cache, with PyTorch 2.14 on an A100: with the folding off, export took 193 s against 190.5 s with it on, a decode call took 945 µs against 1019 µs (medians of three exports and three runs on a busy machine), and the 41 generated tokens were the same. - The mask tests check that the runtime log reports one failed capture, so they fail when CUDA graphs are off or a failed capture is tried again. `test_host_read_runs_without_a_graph` also runs a module that scales by a `.item()` read; without the capture check in the item shims, the process aborts. - With the debug check on the phase before the retry turned into a release check, the mask tests pass. - `test_distribution_argument_check` passes with PyTorch 2.14 and fails with "device-side assert triggered" without the new pass. It also checks that the backend lists the pass, since without the pass the export fails only when autotuning happens to trip the assert. It writes the `MultivariateNormal` check as `torch._assert_async`, since PyTorch 2.14 cannot export the distribution itself. It also checks that the pass drops the check and the math that only fed it. A Gaussian policy head with hidden sizes 16, 64 and 256 exports, runs, and matches eager PyTorch. - A float64 buffer and a folded float64 constant load and match eager PyTorch. Before, both failed to load. `test_float64_weight_loads` fails without the fix. - `test_reshape_copy` in `test_cuda_export.py` fails without the fix. The one in `test_index_shims.py` runs `_reshape_copy` on a transposed, a sliced and a contiguous input and matches eager PyTorch. - `test_varlen_attn` in `test_cuda_export.py` fails without the fix. It lowers a program that flattens the output, compares the decomposition with the op on CUDA only (on ROCm, varlen_attn runs AOTriton, which lines an explicit window up with the start of each sequence, while the decomposition follows Flash attention on CUDA, which lines it up with the end), including a window in self-attention and NaN values outside the key sequences, and checks that a window the Flash attention versions disagree on raises. It also fails when the output is left non-contiguous, when the window check is removed or refuses every such window, and when the values outside the key sequences are not zeroed. Both `test_varlen_attn` tests include queries and keys filled with 1024, where every key gets the same weight; computing the weights as `exp(scores - logsumexp)` fails them, since at that size float32 rounds the log-sum-exp to a whole number and the weights no longer add up to one. The `test_varlen_attn` in `test_index_shims.py` lowers and runs programs that call `varlen_attn` with no window, a causal mask and the windows (8, 0) and (4, 2), flatten each output, and take `cu_seq_q` and `cu_seq_k` only at run time. They match eager PyTorch. - `test_missing_weights_blob.py`, `test_cuda_weight_reuse.py` and `test_device_placement.py` pass. - `test_scatter_reduce` now runs the sum too, which for int64 goes to the new shim. The test that a missing scatter kernel fails at export now tries to lower a scatter of a scalar in deterministic mode. The backend still lacks that kernel, so the test still checks the scatter fallback now that the scatter_reduce and scatter shims exist. It fails when the check stops watching that fallback. `test_scatter_mean_counts_in_its_dtype` fails with the old count and when a count that wraps to zero is not turned into one. `test_float64_and_float16_values` (a sort and `.item()` in each dtype) and the float64 masked accumulate fail with the old shims, and all pass now. The `.item()` scale is 1.5, so that test also fails without the patch that passes a float to Triton kernels as a float. - `test_scatter_reduce` and `test_sort_shim.py` pass. `test_scatter_reduce` fails without the fix, and also when the shim ignores `include_self`, rounds the integer mean toward zero, or lets a NaN lose in amax and amin. `test_scatter_strides` scatters along the middle dim of a 3-D tensor, given as a negative dim, from a strided source and from transposed operands, and fails when the shims walk only the last two dims. `test_scatter_reduce_edge_cases` covers a bool sum, int32 and half products that wrap and overflow, and an empty index. `test_scatter_reduce_sum_adds_in_float` sums 1000 half and bfloat16 values in deterministic mode, and fails when the sum rounds to the dtype at every step. `test_scatter_reduce_replays_in_a_graph` captures a `scatter_reduce` program in a CUDA graph and matches eager on every replay. - The weight cache, runtime tensor and shim unit tests pass in a standalone googletest build. Without the fix, the weight cache test of a float64 entry fails. - `test_small_gpu_override_survives_a_thread_that_leaves_first` overlaps two threads inside the small-GPU patch, and fails with the old save and restore. - With the seed copy made to fail once, the random shim logs it, fails that call, and seeds on the next one; before, it went on with an unseeded state. Not tested: a failed first graph launch (it cannot be triggered from a test), a `std::random_device` that throws, Windows, macOS case-insensitive file systems, the buffer folding change on Metal, and the ROCm build. No test covers the seed change, since a test cannot make two processes start in the same second. ### Checked on real robot policies On the release/1.6 stack with this change, on A100 and H100: - SmolVLA on the CUDA backend: one device-to-host copy per action chunk (only the action download), matching eager within the usual tolerance. - A Gaussian policy whose `torch.distributions` check used to stop the CUDA export: exports and runs, 0.76 ms per chunk. - A float64 weight that used to fail at load with "Malformed CUDA FQN weight metadata": loads and runs, max difference 0. - Four processes started together: four different random draws, against two out of four before. --------- Co-authored-by: Anthony Shoumikhin <anthony@shoumikh.in> Co-authored-by: Anthony Shoumikhin <shoumikhin@meta.com>
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
…ckend An operator Torch-TensorRT cannot convert used to fall back to a portable CPU kernel, with copies to and from the GPU around it. The executorch_tensorrt backend now hands those operators to ExecuTorch's CUDA backend, so a program runs entirely on the GPU. The CUDA backend writes its weights into a data file beside the program. Each part names that file after itself, through the backend's data-name compile spec, so parts saved into one bundle folder keep their own weights, and the file is recorded as the program's data file. The data-name compile spec comes from pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Tested on an NVIDIA GPU: a model with an operator TensorRT rejects compiles to TensorRT and CUDA delegates with no operator left on the CPU and its input and output in GPU memory, and matches PyTorch when run without PyTorch.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
executorch_cuda lowers each part with ExecuTorch's own CUDA backend, so a bundle runs on an NVIDIA GPU without TensorRT. ExecuTorch compiles the kernels ahead of time for the GPU that exports. It is one config class on ExecuTorchConfig, and its registration. ExecuTorch loads its programs, so the runtime, the engine and lerobot-rollout are unchanged. The CUDA backend keeps a program's kernels and weights in a separate data file, named after the part, as with executorch_tensorrt, so two parts never overwrite each other's weights and the bundle names one data file per program. Naming that file needs pytorch/executorch#24314, which ExecuTorch 1.6 gets with pytorch/executorch#24324 (in review). Inputs and outputs stay in GPU memory too, so a program reads an earlier one's outputs in place. A CUDA bundle records the GPU that built it, and is refused on another. The kernels are compiled in full float32, as lerobot-export compiles for every backend. PyTorch lets convolutions use TF32 by default, and the convolution kernels the CUDA backend generates follow that flag. On ACT on an H100 they then drifted about seven times further from PyTorch on the CPU than PyTorch's own GPU kernels do. In full float32 the actions match to about 1e-6, and a chunk still runs twice as fast as eager PyTorch on the GPU. It needs the CUDA build of ExecuTorch, and the CUDA compiler (nvcc) at export. Like the other ExecuTorch backends, it reads programs traced by PyTorch 2.15 or newer. Like the backends that take only part of a program, it folds the values a program computes from constants alone before lowering. SmolVLA's vision encoder picks its position ids with a boolean mask made from constants. Left in the program, that indexing makes the GPU read a size back to the host on every run; folded, it is gone. Folding has one cost today: SmolVLA's time embedding, computed in float64 from constants, becomes a float64 weight, and ExecuTorch's CUDA backend cannot load float64 weights yet, so SmolVLA waits for that fix on this backend. Tested on an H100 and on a Jetson AGX Thor: exported ACT, checked the bundle, and ran it with lerobot-rollout on a fake robot. On the H100, pi0 exported too, and its actions matched eager PyTorch to about 1e-5 on ten frames; in float32 it runs a little slower than eager. The check that a bundle runs without PyTorch now also runs a CUDA bundle.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
… part's weights file The executorch_cuda and executorch_tensorrt backends name the weights file of each program part after the part, so the parts of a bundle do not share one file name. That takes the CUDA backend's `generate_data_name_compile_spec` (pytorch/executorch#24314). The ExecuTorch 1.6.0 release candidate does not have it yet, and an export there stopped deep in the compile step with `AttributeError: type object 'CudaBackend' has no attribute 'generate_data_name_compile_spec'`. The export now refuses such an ExecuTorch with an error that names the installed version and the missing change. Like the refusal of a Torch-TensorRT without `fallback_pytorch_plugin_ops`, it is a ValueError, so the export does not add a "pip install" hint for a package that is already installed. The TensorRT backend asks for the CUDA partitioner before TensorRT builds the part's engines, so it refuses at once instead of after the build. `cuda_backend_names_its_weights()` answers the same question for tests that need the CUDA backend. Tested: a unit test refuses the CUDA partitioner of such an ExecuTorch with ExecuTorch 1.6.0 (release candidate) installed, and another shows the TensorRT backend refuses before calling Torch-TensorRT's compile; without the change it reaches the compile.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
…d the installed versions Three test files each built their own skip marks for the ExecuTorch backends, and they checked different things. None of them checked that the CUDA backend can name each part's weights file, so with the ExecuTorch 1.6.0 release candidate the `executorch_cuda` tests failed instead of skipping. A skipped test also said only "needs ExecuTorch and a CUDA GPU", which does not tell a reader which package is too old. `tests/export/backends.py` now holds one mark per backend: - `needs_executorch`: ExecuTorch 1.6 and PyTorch 2.15. - `needs_executorch_cuda`: also a CUDA GPU and an ExecuTorch whose CUDA backend names each part's weights file (pytorch/executorch#24314). - `needs_executorch_tensorrt`: also a Torch-TensorRT with `fallback_pytorch_plugin_ops` (pytorch/TensorRT#4953) and torch-tensorrt-executorch-runtime. Each reason names these fixes and the installed versions of PyTorch, ExecuTorch, Torch-TensorRT and its runtime. The reason always names ExecuTorch, and a CUDA GPU for a GPU backend, so `tests/conftest.py` marks the test `executorch` and `gpu` whether it skips or not. `test_export.py`, `test_export_history.py` and `test_executorch_tensorrt.py` use these marks. Tested: with the ExecuTorch 1.6.0 and Torch-TensorRT 2.15.0 release candidates on a GPU, every GPU export test skips and names pytorch/executorch#24314 or pytorch/TensorRT#4953; with ExecuTorch and Torch-TensorRT builds that have both fixes, the same tests run and pass.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 3, 2026
…he GPU backends after approval No CI job installed ExecuTorch, so every test that exports with a real backend skipped, and nothing checked that `lerobot-export` and `lerobot-rollout --bundle.path` still work with XNNPACK, ExecuTorch's CUDA backend or TensorRT. Two jobs now install the compilers over the locked environment and run the tests marked `executorch` (`tests/conftest.py` adds the mark to every test that skips without ExecuTorch): - Fast Export Tests, in Fast Tests, on every pull request, forks included: `ubuntu-latest`, the extras of the exported policies (`pi`, `smolvla`, `groot`, `diffusion`), then ExecuTorch's CPU build, then `pytest tests -m executorch`. The GPU tests among them skip there. - GPU Export Tests, in Full Tests, next to GPU Tests and on the same image and runner: ExecuTorch's and Torch-TensorRT's CUDA 13.2 builds, which match PyPI's torch 2.15.0, and `nvidia-cuda-nvcc` 13.2, the CUDA compiler the CUDA backend runs at export (`CUDA_HOME` points at it). Then `pytest tests -m "executorch and gpu"`. It runs when GPU Tests do: after approval, on main and by hand, never for a fork. ExecuTorch and Torch-TensorRT stay out of `uv.lock`, for the reasons the `export` note in `pyproject.toml` gives, so each job pins the exact wheels in one environment variable. The ExecuTorch 1.6 and Torch-TensorRT 2.15 releases are not on PyPI yet, so the variables point at their release candidates on PyTorch's test index; they become `executorch==1.6.0` and `torch-tensorrt==2.15.0` once those are published. The candidates lack two changes the GPU backends need (pytorch/executorch#24314 and pytorch/TensorRT#4953), and no torch-tensorrt-executorch-runtime takes ExecuTorch 1.6.0 yet, so today the GPU tests skip and say which change each one needs. They run as soon as released wheels have them. CONTRIBUTING.md and AGENTS.md say where these tests run and which mark a new export test takes. Tested: ran each job's commands in fresh environments installed only from the package indexes, with the locked export extras in place of the CI image. - Fast Export Tests (4 CPUs, a busy machine): 20 passed in 46 minutes. GR00T with XNNPACK failed as the previous commit describes; it skips now. - GPU Export Tests on an A100: all 34 GPU tests skip with their reasons, 57 seconds. - Not from the indexes: the same GPU tests with ExecuTorch and Torch-TensorRT built from source with the two changes, on an H100 with cold caches: 29 passed, 5 skipped (tiny pi0 and pi0.5 with TensorRT, and one test that downloads CLIP), 42 minutes. Not tested: the L4 runner and the CI image themselves.
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 change fixes several problems in the CUDA backend. Each fix is its own commit.
Name the weights file of each program
The backend saves every program's weights to
aoti_cuda_blob.ptd. Exporting several programs into one folder overwrites the weights, and the earlier programs then fail to load. Renaming each file after export works, but is easy to miss. A new option names the file:Without the option nothing changes. The name is checked: up to 200 ASCII letters, digits,
_,-and., not starting with.or-, not ending with.ptd, and not a Windows device name such asNUL. All methods of a program must share it. A backend that keeps its data inside the.pte, such as Metal, refuses it.Lower antialiased resizes
F.interpolate(x, mode="bilinear" or "bicubic", antialias=True)failed to lower with "missing fallback kernels". The backend now rewrites it into plain tensor math, with PyTorch's window and weight formulas, for dynamic sizes and float32, float16 and bfloat16. The results are not bit-identical: in float32 they stay within about 2e-5 of PyTorch's kernel for inputs between 0 and 1.Tests
Ran on an A100 GPU (Linux x86_64):
test_cuda_export.py,test_antialiased_resize.py,test_missing_weights_blob.py,test_merge_ptes.pyandtest_cuda_weight_metadata.py. Each test added for a fix fails when that fix is removed.