[shoumikhin]

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.
• Nothing changes in the exported program, so new programs still load on older runtimes, and older programs get the fix too.

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 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 [inductor] Support host-computed float scalar args for Triton templates 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. 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 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 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, 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. 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.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24326

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit bdea161 with merge base aed328d ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 24m 41s —— View job

---

This PR makes CUDA graph replay reject changed input shapes and reruns a failed capture call without a graph; I found no blockers, only minor issues: an undocumented replay contract, Error logs on calls that now succeed, two run_once structure nits, a pre-existing stale-pointer case, and two small test gaps.

PR Review: #24326

Summary

Replay now compares input shapes instead of byte counts. When a capture fails before its graph launches, the same call runs again without a graph. Three independent reviewers found no blockers. Everything below is minor, and the first item is a design question.

API Design

• A shape change after capture is still a hard error (backends/cuda/runtime/cuda_backend.cpp:773-778).
  • A failed capture now falls back to running without a graph, but a replay with a new input shape still returns InvalidArgument. For a method exported with a Dim, every later call whose shape differs from the capture call's fails. One example is test_output_shape_changes_at_capture call 6, which uses batch 1 under Dim(min=1).
  • This is not a regression, because a change in byte count already failed on main.
  • The contract isn't written down anywhere. The Replay bullet of the phase comment (cuda_delegate_handle.h:118-121) doesn't mention it, and no Markdown doc covers enable_cuda_graph_for_method. Please either document it there or fall back instead.
  • A fallback needs two more changes:
    • The check has to run over all inputs before the first copy is queued. Today it runs inside the copy loop.
    • Replay has to resize outputs back to their captured shapes. It relies on them staying fixed (:1098-1099), but an eager call resizes them (utils.h:114-116).

Code Quality

• The replay refusal no longer says what differs (cuda_backend.cpp:777).
  • The old message printed the expected and actual byte counts; the new one prints only the input index.
  • Report both dim counts when they differ, and otherwise the first differing dim with both sizes. Fix this →
• The main mask case still logs two Error lines on a call that now succeeds (shims/index.cu:337-340, cuda_backend.cpp:1064-1068).
  • On the 4th call of every mask-index method, the shim refuses through ET_CHECK_OR_RETURN_ERROR, and then the run() failure check logs again. Both lines are Error and come before the Info fallback line.
  • The PR moved the two layout messages (:1090, :1105) to Info for this reason. Log these two at Info during capture too; the eager retry still reports real failures at Error.
  • In the layout and resize cases, the fallback is now announced twice at Info: once at :1090 or :1105, then again at :820. Fix this →
• run_once splits the public overrides (cuda_backend.cpp:831, :1227). It opens a new private:/public: pair between execute and destroy, but the class already has private helper sections at :169 and :1287. Fix this →
• graph_ran is a side-channel out-parameter (cuda_backend.cpp:835-841).
  • Only the capture branch sets it. The retry at :825 passes it again with is_capture_step = false, where nothing writes it.
  • Splitting the capture attempt from the eager run would let the capture function report directly whether its graph launched.
• Pre-existing: stale graph_allocations after a failed first launch (cuda_backend.cpp:1124, next to the new graph_ran = true).
  • note_graph_allocations() runs before the first cudaGraphLaunch. If that launch fails, CaptureGuard destroys the graph but keeps graph_allocations, and release() later calls cudaFree on addresses of a graph that never ran.
  • Suspected effects: a stray error left for the next cudaGetLastError(), or freeing another graph's allocation if the address was reused.
  • Recording the allocations after the launch succeeds fixes this. It would also let CaptureGuard call release() instead of hand-copying it; this PR had to extend that copy and the original separately (:878, cuda_delegate_handle.h:233). Fix this →

Testing

• The recapture test doesn't cover static_input_sizes (backends/cuda/runtime/test/test_cuda_kv_cache.cpp:933-952).
  • It seeds and checks only static_input_ptrs and static_input_nbytes.
  • If release() stopped clearing the new field, a recapture after KV growth would append a second set of shapes, and replay would check inputs against the first capture's shapes. Fix this →
• assertIn("error 0x12", str(out)) also matches a set_inputs() rejection (test_output_stride_rearrange.py:402, :432).
  • The pybindings use the same "failed with error 0x…" wording for method->set_inputs() (extension/pybindings/pybindings.cpp:3771, :3940) as for execute() (:3995).
  • Match "execute() failed with error 0x12" so that only the replay check passes. Fix this →

Recommendation

Approve. Nothing blocks merging. The replay contract needs a decision, plus a comment documenting it if the hard error stays. The rest are small cleanups.

---

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 27m 4s —— View job

---

This PR makes CUDA graph replay reject changed input shapes and runs a failed capture call again without a graph; nothing blocks merging, but the hard error on a new input shape after capture should be documented or replaced with a fallback, and there are smaller logging, structure and test issues.

PR Review: #24326

Summary

Replay now checks input shapes instead of byte counts. When a capture fails before its graph launches, the same call runs again without a graph. Three independent reviewers covered all 8 changed files and found no blockers. The first item below is a design question and the rest are minor. This was a static review: there's no GPU on this runner, so nothing was built or run.

API Design

• After capture, a new input shape is a hard error, and nothing documents that (backends/cuda/runtime/cuda_backend.cpp:773-778).
  • A failed capture now falls back to eager. A replay whose input shape differs from the capture call's still returns InvalidArgument, even inside the exported Dim bounds, and it does so on every later call with that shape.
  • test_output_shape_changes_at_capture now uses that refusal on call 6 (batch 1 under Dim(min=1)) as its proof that graphs are on, so it locks the behavior in.
  • This isn't a regression, because a change in byte count already failed on main.
  • The Replay bullet of the phase comment (cuda_delegate_handle.h:118-121) doesn't mention the rule, and no .md file covers enable_cuda_graph_for_method. Please either document it there or fall back to eager.
  • A fallback also needs replay to resize outputs back to their captured shapes. Replay relies on them keeping that shape (cuda_backend.cpp:1098-1099), but an eager call resizes them (utils.h:114-116).

Code Quality

• The replay refusal no longer says what differs (cuda_backend.cpp:777).
  • The old message printed the expected and actual byte counts, and the new one prints only the input index.
  • Print both shapes with tensor_shape_to_c_string, as cuda_kv_pool.cpp:48-51 does. Fix this →
• A mask-index method still logs two Error lines on a call that now succeeds (shims/index.cu:337-340, cuda_backend.cpp:1064-1068).
  • On the capture call, both the shim's refusal ("a mask index cannot run inside a CUDA graph capture") and the AOTInductorModelContainerRun failed check log at Error. Then the retry returns Ok.
  • The PR moved the layout and resize messages (:1090, :1105) to Info for this reason. These two should be Info during capture too; the eager retry still reports real failures at Error.
  • In the layout and resize cases, the fallback is now announced twice at Info: at :1090 or :1105, then again at :820. Fix this →
• The retry depends on an out-parameter and on a side effect of CaptureGuard (cuda_backend.cpp:814-826, :835-841, :888, :1150).
  • execute() decides whether to retry from bool& graph_ran. One line of a 390-line function sets it, and it is set when the launch is queued, not when the graph has run. The retry at :825 passes it to a call that never writes it.
  • The retry only runs eagerly because CaptureGuard's destructor has already set the phase to Disabled (:888). That selects handle->run (:1040-1044) and skips the warmup decrement (:1173). If the guard stopped doing that, the retry would run as a warmup step and the next call would capture again.
  • A separate capture function that returns whether its graph launched would make both explicit.
• run_once splits the public overrides (cuda_backend.cpp:831, :1227).
  • It opens a new private:/public: pair between execute and destroy, while the class already has private sections at :169 and :1287. Fix this →
• Low confidence: when the .so lacks run_single_threaded, the retry probably fails too (cuda_backend.cpp:817-825, cuda_delegate_handle.h:290-295).
  • Without that export, the capture step calls run(), which records AOTI's completion event inside the capture.
  • If the capture then fails after run() returned (output layout, resize, end capture, instantiate, launch), the retry's run() queries that event. The header comment says that query fails.
  • The mask case should be unaffected, because the shim fails before run() records the event. This is no worse than main, but the comment at :818-819 claims more than holds. Either skip the retry when run_single_threaded is null, or note the limit there.
• Pre-existing: stale graph_allocations after a failed first launch (cuda_backend.cpp:1124, next to the new graph_ran = true at :1150).
  • note_graph_allocations() runs before the first cudaGraphLaunch. If that launch fails, CaptureGuard destroys the graph but keeps the list, and release() later calls cudaFree on allocations of a graph that never ran.
  • Suspected effects: a stray error left for a later cudaGetLastError(), or freeing another graph's allocation if the address was reused.
  • Recording the allocations once the launch succeeds fixes this. It would also let CaptureGuard call release() instead of a hand-copied version of it. This PR had to extend both copies separately (:878, cuda_delegate_handle.h:233). Fix this →

Testing

• No test checks the retry's safety claim with state (backends/cuda/tests/test_output_stride_rearrange.py:404-414).
  • _MaskedRows and _MaskedAdd are stateless, so a retry that applied an update twice would still pass.
  • Add a module that increments a registered buffer before x[mask] and returns something that depends on it. The AOTI partitioner keeps mutated buffers in the delegate (aoti_partitioner.py:104), and the existing per-call comparison with eager would catch a double update. Fix this →
• The recapture test doesn't cover static_input_sizes (backends/cuda/runtime/test/test_cuda_kv_cache.cpp:933-934, :951-952).
  • It seeds and checks only static_input_ptrs and static_input_nbytes.
  • If release() stopped clearing the new field, a recapture would append a second set of shapes, and replay would compare inputs against the first set. Fix this →
• assertIn("error 0x12", str(out)) also matches a set_inputs() rejection (test_output_stride_rearrange.py:402, :432).
  • The pybindings use the same "failed with error 0x…" wording for set_inputs() (extension/pybindings/pybindings.cpp:3771, :3940) as for execute() (:3995).
  • Match "execute() failed with error 0x12" so that only the replay check passes. Fix this →
• test_two_dim_mask draws an unseeded random mask (backends/cuda/tests/test_index_shims.py:88), so a failure that depends on the mask can't be reproduced. Seed it, as test_output_stride_rearrange.py does. Fix this →

Recommendation

Approve. Nothing blocks merging. The replay shape rule needs a decision, plus a comment documenting it if the hard error stays. The rest are small cleanups.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 58m 0s —— View job

---

Claude finished @shoumikhin's task

• Gather context (diff against the merge base with main, commits, PR discussion)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

Approve: nothing blocks merging, though the five minor findings below are worth follow-ups (a race in the new --as-needed link patch, the masked-copy rewrite skipping subgraphs, thin coverage of the new scatter shims, and two documentation gaps), and 61 of the 75 commits landed after the existing approval.

Detailed review of #24326

PR Review: #24326

Summary

The PR does two things:

• It makes the CUDA backend's CUDA graphs safe for programs whose behaviour depends on the data: a failed capture reruns without a graph, a replay with a changed input shape is refused, and a method with a constant in host memory runs without a graph.
• It fixes ops and dtypes that failed to export or load: scatter, scatter_reduce, masked copies (now written as where), _assert_async, _reshape_copy, varlen_attn and float64. It also fixes --as-needed linking and the Triton SDPA's last key tile.

There are no blocking issues, and the five findings below are all minor.

Thread Safety

• _link_libraries_as_needed loses its flag when two threads overlap (backends/cuda/cuda_backend.py:191-214).

  • Every entry saves whatever cpp_builder.get_cpp_torch_device_options is at that moment (line 191) and puts it back on exit (line 214).
  • Suppose thread B enters while thread A is inside, so B saves A's wrapper. If A exits first, it reinstalls the original function, and the rest of B's compile links without --as-needed.
  • B's .so then records libgomp. It fails to load in a process without torch on a machine with no libgomp, which is exactly what this helper exists to prevent.
  • When B exits, it leaves A's wrapper installed for good. That is harmless because the wrapper checks the flag, but the module stays patched.
  • The docstring says "Only threads inside this context get the flag", which is true only when contexts do not overlap. _allow_triton_templates_on_small_gpus (lines 137-164, already on main) is built the same way.
  • This only happens when two threads lower at the same time.
  • A count of active contexts behind a lock fixes both helpers: install the wrapper when the count goes from 0 to 1, restore the original when it returns to 0, and let the thread-local flag decide for each thread. Installing the wrapper once and only toggling the flag also works.

  Fix this →

Testing

• The new scatter shims are tested only along a tensor's last dim (backends/cuda/tests/test_index_shims.py:171-227).

  • test_scatter_with_tensor_source runs only bf16, int64 (with an int32 index) and bool.
  • Its docstring says that deterministic mode, which the test turns on, also sends every dtype to the shim. Even so, float32 and float16 never run.
  • Every scatter and scatter_reduce case scatters a 2-D tensor along its last dim, plus one 1-D mean.
  • So nothing covers the general stride walk in scatter_offsets_kernel (backends/cuda/runtime/shims/index.cu:890-923). Two cases would: a scatter along the middle dim of a 3-D tensor, and one with x transposed inside forward.

  Fix this →

Performance

• The masked-copy rewrite skips subgraphs (backends/cuda/passes/replace_masked_copy_with_where.py:90).

  • call() walks only graph_module.graph.nodes.
  • So an out[mask] = src[mask] inside a torch.cond branch, or any other higher-order-op body, still goes through the two index shims. Each of them waits on the host for the mask count, which is the wait the pass exists to remove.
  • RemoveAssertAsyncPass, which this PR also adds, walks every graph in graph_module.modules() (backends/cuda/passes/remove_assert_async.py:33).

  Fix this →

Code Quality

• _written_buffer_options does more than its docstring says (backends/aoti/aoti_backend.py:56-68).

  • The docstring says it keeps "the buffers a program writes out of constant folding".
  • In fact, joint_graph_constant_folding=False turns off Inductor's uniform-value folding for the whole graph. That happens in every CUDA and Metal program that writes any buffer, including programs with an in-graph KV cache.
  • The message of commit e26e5c8 says this ("there this only stops the folding of the other constants"). Whoever acts on the TODO will read the docstring, though, not the commit message.
  • Nothing in the PR exercises the Metal side.

  Fix this →

• A known silent truncation of .item() values is recorded only in a test comment (backends/cuda/tests/test_index_shims.py:256-258).

  • The comment says AOTInductor's Triton launcher takes a float read with .item() as int64_t, so 1.5 becomes 1. The test uses 2.0 to avoid this.
  • The float64 and float16 item shims this PR adds (backends/cuda/runtime/shims/memory.cpp:478-479) put two more dtypes on that path.
  • A value that is not a whole number, read with .item() and passed to a Triton kernel, then comes out wrong with no error. There is no upstream link and no note for users.
  • This comes from the test's own comment; I did not reproduce it.

  Fix this →

Recommendation

Approve. Nothing blocks merging, and the items above can be follow-ups.

The existing approval (SS-JIA, 2026-10-30 14:37 UTC) predates 61 of the 75 commits. Those include:

• the scatter-with-source shim
• the masked-copy and assert-removal passes
• the _reshape_copy and varlen_attn decompositions
• the buffer-folding switch

They deserve a second look from the approver before merge.

[Copilot]

🟡 Changes recommended

The CUDA mean shim produces incorrect results for heavily repeated half-precision indices, and failed initialization on Windows leaves a loaded temporary DLL behind.

2 open issues

• Half-precision mean count overflows for repeated indices · New
• Failed initialization leaks the loaded DLL on Windows · New

_{Review effort: Balanced}

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

Thanks for both reviews. I pushed fixes:

1. CUDA graphs and .item(). The item shims now refuse to read a device value while the stream is capturing, as the mask index does. So the capture call runs again without a graph instead of aborting. A float read with .item() also reaches Triton kernels as a float now. AOTInductor declares that argument int64_t in the kernel launcher, which already broke float32 before this pull request. A compile-time patch takes the type from the kernel signature until PyTorch does the same ([inductor] Support host-computed float scalar args for Triton templates pytorch#201882). The float64 and float16 test now scales by 1.5.

2. Passes, threads and tests. The masked-copy pass walks torch.cond branches. The small-GPU and --as-needed patches count the threads inside them, so one thread leaving no longer removes the patch from another. The varlen_attn and _reshape_copy programs now run compiled in the tests. The scatter tests cover the middle dim of a 3-D tensor, strided and transposed operands, float32 and float16, edge cases and a CUDA graph replay. I left out float64, int16 and uint8, since they use the same kernel template as the tested types.

3. Folding switch. On the tiny Gemma 4 decoder (KV cache, torch 2.14, A100), 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. So I kept the switch and fixed the docstring. A version check would add code for no measured gain. Its TODO marks it for removal once the pinned PyTorch is 2.15. The numbers are in the description.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 44m 48s —— View job

---

Approve: there are no blockers, only four minor issues: the new scatter_reduce kernel is serial in its worst case, the temporary library cleanup doesn't work on Windows, the phase diagram leaves out recapture(), and no test covers a dynamic-size mask in the masked-copy pass.

Detailed review

PR Review: #24326

Summary

This PR fixes how CUDA graphs handle mask indexing, .item() and changed input shapes. It also rewrites masked copies as where, drops tensor checks before compiling, and adds float64, _reshape_copy, scatter/scatter_reduce and varlen_attn support. Nothing blocks the merge. All four issues below are minor. backends/cuda/runtime/aoti_cuda_shims.lib is binary and was not reviewed.

Code Quality

• On Windows, a failed init still leaves the temporary library behind (backends/cuda/runtime/cuda_backend.cpp:496-510).

  • Every failure after load_library (line 535) returns with the library still loaded:
    • the function-pointer lookup (550)
    • creating the container (555)
    • the off-graph KV checks (579-600)
    • loading the weights (605-614), which is the case test_missing_weights_blob.py covers
  • Windows can't delete a loaded DLL, so the guard's filesystem::remove fails, the file stays in %TEMP%, and every failed load logs an Error. Linux can unlink a mapped file, so it isn't affected.
  • The new test can't see this. It patches only TMPDIR (test_missing_weights_blob.py:173-174), but temp_directory_path() reads TMP/TEMP on Windows, so os.listdir(tmp) is always empty there. No Windows job in cuda.yml runs the test today.
  • Fix: until create_with_device succeeds, nothing from the library is alive yet, so calling close_library(lib_handle) on those paths should be safe. After that, destroy()'s note about keeping the library loaded applies. For those later failures, say in the guard comment that Windows keeps the file. In the test, also set TMP/TEMP, or skip the empty-directory check on Windows.

  Fix this →

• The phase diagram this PR extends still leaves out recapture() (backends/cuda/runtime/cuda_delegate_handle.h:105-120).

  • CudaGraphState::recapture() (243-251, called from cuda_kv_pool.cpp:421-422) moves Replay back to Warmup with a single warmup step.
  • The diagram has no Replay → Warmup edge, and the Warmup text says the first kCudaGraphWarmupSteps calls run eagerly.
  • The diagram now lists the failure edges, so add this edge and the one-step re-warmup too.

  Fix this →

Testing

• No test exports a mask whose size is dynamic (backends/cuda/passes/replace_masked_copy_with_where.py:177-179).

  • When mask.numel() is symbolic, the count's range comes from bound_sympy(...).upper, which a plain Dim leaves unbounded.
  • Whether every check export adds on the count holds over that range depends on the bounds torch gives the unbacked count.
  • If one check doesn't, _only_range_checked returns None, and nothing reports it. The where rewrite still happens, but the index stays, along with the count readback and host sync this pass exists to remove.
  • The only dynamic-shape test, test_check_on_another_size_keeps_the_index, makes y dynamic, not the mask.
  • Add masked-copy cases with an unbounded and a bounded Dim on the mask's dim, and assert that index.Tensor and its _assert_scalar checks are gone.

  Fix this →

Performance

• scatter_reduce takes time proportional to its longest run (backends/cuda/runtime/shims/index.cu:973-984).

  • After the sort, one thread walks the whole run of points that land on one output element.
  • Per index.h:111-113, Inductor sends every reduction except "sum" to this shim, not only in deterministic mode. A reduction onto a few elements therefore runs on a few threads. Examples:
    • global max pooling with one graph in the batch: an amax of N×F values into 1×F is F runs of N points
    • a high-degree node in a GNN
    • a histogram with few bins
  • This isn't a regression, since these reductions failed to lower before this PR, but the slowdown on such inputs is steep.
  • amax/amin don't depend on order, so they could reduce the sorted runs in parallel, for example with run offsets from cub::DeviceRunLengthEncode and cub::DeviceSegmentedReduce, or with a block per long run.
  • At minimum, state the cost in the comment at 945-952 and in index.h.

  Fix this →

Recommendation

Approve: nothing blocks, and all four items can be follow-ups. SS-JIA approved before most of the 87 commits landed, so it would be worth having a person look at the runtime changes made since then.