[pytorchbot]

Stacked on #24298. This branch starts with the #24222 cherry-pick, which is #24298. Please merge #24298 first, then update this branch. Until then GitHub also shows the #24222 files and their behavior changes here; they are described and reviewed in #24298. The change this PR adds is the #24299 cherry-pick and the #24326 follow-ups below, in the CUDA backend and the runtime.

What is wrong today

Robot policies and follow-up checks on #24222 found these bugs in the CUDA backend and the runtime:

1. CUDA mask indexing: x[mask] and x[mask] = values have no C shims, so lowering fails with "missing fallback kernels". This blocks SmolVLA.
2. CUDA graphs: a channels-last output turns CUDA graphs off, although a raw copy is correct.
3. CUDA graph shapes: when the captured call has a new batch size, it and every replay return the previous output shape.
4. Output buffers: Method::set_output_data_ptr checks a buffer against the output's current size. After a dynamic output shrinks, a too-small buffer is accepted and a later, larger call writes past it.

What this change does

One commit per fix:

1. New shims for index and index_put follow ATen's indexing rules. The Windows import library gains the new names. Inside a CUDA graph capture, a mask index fails with a clear error.
2. The check compares output strides with the strides the ETensor's dim order gives.
3. The capture step resizes each output to the captured shape.
4. The buffer is checked against the output's upper bound. Code that binds a smaller buffer after a call now gets InvalidArgument.

Tests

Each new test fails without its fix.

• CUDA (A100): mask reads and writes; the raw-copy check; a capture call with a new batch size.
• Runtime: a too-small output buffer is refused after a dynamic output shrinks; the runtime tests pass.

On LeRobot policies, SmolVLA now lowers to CUDA and runs in 43 ms per call on an A100, against 504 ms in eager PyTorch (cosine similarity 0.99996).

Follow-up fixes from #24326

This pick also carries fixes from #24326. A CUDA graph replay now refuses an input whose shape changed but kept its byte count. Before, such a call returned the captured output shape with no error, and the channels-last change above lets more outputs reach that path. The wheel's shared library check now lists the new index shims. 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.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit b1f5f7e with merge base 4db778e ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Test CoreML Backend / test-coreml / test-backend-macos (coreml, models) / macos-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[shoumikhin]

@claude /pr-review

[claude]

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

---

PR Review: #24320

Needs Discussion: with enable_cuda_graph_for_method, any method that uses a mask index (the SmolVLA case this PR targets) fails its 4th execute(), which is the CUDA graph capture call, and no test covers that path.

Detailed review

Summary

This cherry-pick of #24299 and #24326 makes four changes:

• adds CUDA C shims for mask indexing and index_put;
• lets CUDA graphs raw-copy an output whose strides match its dim order;
• resizes outputs at capture and refuses a replay after an input changes shape;
• checks Method::set_output_data_ptr buffers against the output's upper bound.

After fact-checking, one major issue (a mask index fails the CUDA graph capture call) and one minor test gap remain.

Diff base: the shallow checkout has no merge base with origin/release/1.6, so this review diffs against ba053ba, the parent of the release tip, which 553ffd1 merges in. That diff matches GitHub's 97 files.

Scope: 78 of those files are touched only by the #24222 cherry-pick this PR is stacked on (6207bdf, exir and several backends), and they got a lighter pass.

Skimmed or not reviewed:

• backends/qualcomm/tests/test_passes.py
• backends/vulkan/test/test_vulkan_delegate.py
• backends/xnnpack/test/**
• backends/transforms/test/test_create_delete_constant_placeholder.py
• most exir/**/test* files
• backends/cuda/runtime/aoti_cuda_shims.lib (export names only)

Code Quality

• With CUDA graphs, a mask index fails the capture call. backends/cuda/runtime/shims/index.cu:333-340, backends/cuda/runtime/cuda_backend.cpp:1040-1044. (Major; reported by all three reviewers.)
  • Trigger: a method that uses x[mask] or x[mask] = v (the SmolVLA case), loaded with enable_cuda_graph_for_method.
  • Calls 1-3 are warmup and run eagerly, so the mask index works and nothing flags it.
  • Call 4 is the capture. mask_positions returns NotSupported, run() fails, and execute returns Error::Internal, the same code a real failure gives. CaptureGuard then disables graphs, so calls 5 and later run eagerly.
  • Result: every such model fails exactly one execute() by design, and the caller cannot tell it apart from a real failure.
  • Existing pattern: when an output cannot be raw-copied, warmup already switches the method to eager without failing a call (cuda_backend.cpp:1180-1188). A mask index needs the same: detect it during warmup, or rerun the failed capture call eagerly once the guard has disabled graphs.
  • Missing test: the description says "mask reads and writes, also inside a capture" are tested. But test_mask_index_and_index_put (backends/cuda/tests/test_cuda_export.py:673) never sets enable_cuda_graph_for_method. So neither this path nor the unwinding of a capture whose run() itself fails is tested.
  • This code is already on main, so fix it there first. Fix this →

Testing

• No test covers floating-point accumulate in aoti_torch_index_put_out. backends/cuda/runtime/shims/index.cu:1022-1031, index.cu:237-269. (Minor; one reviewer, confirmed by reading the test.)
  • Untested paths: Float, Half and BFloat16 summed in float, and put_sorted_kernel when indices repeat.
  • Current coverage: the only accumulate case in test_mask_index_and_index_put is the int64 counts.
  • Unchecked claims: the comments promise ATen's accumulation type, a single rounding, and index-order association, and no test checks any of them.
  • Low confidence: I could not confirm that Inductor ever sends a floating-point accumulate to this fallback. If it never does, these cases are dead code and should be removed instead. Fix this →

Recommendation

Needs Discussion

The capture failure is not a regression, since a mask-index model could not be lowered to CUDA at all before this PR. It is still a deterministic error in the scenario this PR enables (SmolVLA) once CUDA graphs are on.

• Decision needed: ship 1.6 with the failure, or fix it on main and add the fix to this cherry-pick.
• Note: once the Fix crashes, wrong outputs and lowering failures in several backends and in exir #24222 cherry-pick this PR is stacked on lands, this diff shrinks to its own 19 files.

Review steps

• Gather context: diff against origin/release/1.6 (via ba053ba) and commit history
• Run three independent reviews, each starting from different files
• Consolidate findings and review the stacked-layer files no reviewer reached (Core ML, Arm, Cortex-M, MLX, NXP, OpenVINO, Qualcomm, stacked CUDA/AOTI parts, cummin)
• Fact-check findings (dropped nine as invalid or not worth raising on a cherry-pick)
• Post final review